Upload folder using huggingface_hub (part 2)

flax_model/alphafold3/structure/cpp/mmcif_altlocs.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ALTLOCS_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ALTLOCS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "alphafold3/structure/cpp/mmcif_layout.h"
+
+namespace alphafold3 {
+
+// Returns the list of indices that should be kept after resolving alt-locs.
+// 1) Partial Residue. Each cycle of alt-locs are resolved separately with the
+//    highest occupancy alt-loc. Tie-breaks are resolved alphabetically. See
+//    tests for examples.
+// 2) Whole Residue. These are resolved in two passes.
+//    a) The residue with the highest occupancy is chosen.
+//    b) The locations for a given residue are resolved.
+//    All tie-breaks are resolved alphabetically. See tests for examples.
+//
+// Preconditions: layout and comp_ids, alt_ids, occupancies are all from same
+// mmCIF file and chain_indices are monotonically increasing and less than
+// layout.num_chains().
+//
+// comp_ids from '_atom_site.label_comp_id'.
+// alt_ids from '_atom_site.label_alt_id'.
+// occupancies from '_atom_site.occupancy'.
+std::vector<std::uint64_t> ResolveMmcifAltLocs(
+    const MmcifLayout& layout, absl::Span<const std::string> comp_ids,
+    absl::Span<const std::string> atom_ids,
+    absl::Span<const std::string> alt_ids,
+    absl::Span<const std::string> occupancies,
+    absl::Span<const std::size_t> chain_indices);
+
+}  // namespace alphafold3
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ALTLOCS_H_

flax_model/alphafold3/structure/cpp/mmcif_atom_site.pyi

@@ -0,0 +1,15 @@
+
+
+from collections.abc import Callable
+from flax_model.alphafold3.cpp import cif_dict
+
+
+def get_internal_to_author_chain_id_map(
+    mmcif: cif_dict.CifDict
+) -> dict[str,str]: ...
+
+
+def get_or_infer_type_symbol(
+    mmcif: cif_dict.CifDict,
+    atom_id_to_type_symbol: Callable[[str, str], str],
+) -> list[str]: ...

flax_model/alphafold3/structure/cpp/mmcif_atom_site_pybind.cc

@@ -0,0 +1,83 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <cstddef>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "pybind11/gil.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"
+#include "pybind11_abseil/absl_casters.h"
+
+namespace alphafold3 {
+namespace {
+namespace py = pybind11;
+
+// If present, returns the _atom_site.type_symbol. If not, infers it using
+// _atom_site.label_comp_id (residue name), _atom_site.label_atom_id (atom name)
+// and the CCD.
+py::list GetOrInferTypeSymbol(const CifDict& mmcif,
+                              const py::object& atom_id_to_type_symbol) {
+  const auto& type_symbol = mmcif["_atom_site.type_symbol"];
+  const int num_atom = mmcif["_atom_site.id"].size();
+  py::list patched_type_symbol(num_atom);
+  if (type_symbol.empty()) {
+    const auto& label_comp_id = mmcif["_atom_site.label_comp_id"];
+    const auto& label_atom_id = mmcif["_atom_site.label_atom_id"];
+    CHECK_EQ(label_comp_id.size(), num_atom);
+    CHECK_EQ(label_atom_id.size(), num_atom);
+    for (int i = 0; i < num_atom; i++) {
+      patched_type_symbol[i] =
+          atom_id_to_type_symbol(label_comp_id[i], label_atom_id[i]);
+    }
+  } else {
+    for (int i = 0; i < num_atom; i++) {
+      patched_type_symbol[i] = type_symbol[i];
+    }
+  }
+  return patched_type_symbol;
+}
+
+absl::flat_hash_map<absl::string_view, absl::string_view>
+GetInternalToAuthorChainIdMap(const CifDict& mmcif) {
+  const auto& label_asym_ids = mmcif["_atom_site.label_asym_id"];
+  const auto& auth_asym_ids = mmcif["_atom_site.auth_asym_id"];
+  CHECK_EQ(label_asym_ids.size(), auth_asym_ids.size());
+
+  absl::flat_hash_map<absl::string_view, absl::string_view> mapping;
+  for (size_t i = 0, num_rows = label_asym_ids.size(); i < num_rows; ++i) {
+    // Use only the first internal_chain_id occurrence to generate the mapping.
+    // It should not matter as there should not be a case where a single
+    // internal chain ID would map to more than one author chain IDs (i.e. the
+    // mapping should be injective). Since we need this method to be fast, we
+    // choose not to check it.
+    mapping.emplace(label_asym_ids[i], auth_asym_ids[i]);
+  }
+  return mapping;
+}
+
+}  // namespace
+
+namespace py = pybind11;
+
+void RegisterModuleMmcifAtomSite(pybind11::module m) {
+  m.def("get_or_infer_type_symbol", &GetOrInferTypeSymbol, py::arg("mmcif"),
+        py::arg("atom_id_to_type_symbol"));
+
+  m.def("get_internal_to_author_chain_id_map", &GetInternalToAuthorChainIdMap,
+        py::arg("mmcif"), py::call_guard<py::gil_scoped_release>());
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_atom_site_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ATOM_SITE_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ATOM_SITE_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleMmcifAtomSite(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_ATOM_SITE_PYBIND_H_

flax_model/alphafold3/structure/cpp/mmcif_layout.h

@@ -0,0 +1,146 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+
+namespace alphafold3 {
+
+// Holds the layout of a parsed mmCIF file.
+class MmcifLayout {
+ public:
+  MmcifLayout(std::vector<std::size_t> chain_ends,
+              std::vector<std::size_t> residues, std::size_t model_offset,
+              std::size_t num_models)
+      : chain_ends_(std::move(chain_ends)),
+        residue_ends_(std::move(residues)),
+        model_offset_(model_offset),
+        num_models_(num_models) {}
+
+  // Reads a layout from a valid parsed mmCIF. If a valid model_id is provided
+  // the offsets will select that model from the mmCIF.
+  // If no model_id is specified, we calculate the layout of the first model
+  // only. Therefore it is a requirement that each model has identical atom
+  // layouts. An error is returned if the atom counts do not between models.
+  static absl::StatusOr<MmcifLayout> Create(const CifDict& mmcif,
+                                            absl::string_view model_id = "");
+
+  std::string ToDebugString() const;
+
+  // Returns the start index and one past the last residue index of a given
+  // chain. A chain_index of n refers to the n-th chain in the mmCIF. The
+  // returned residue indices are 0-based enumerations of residues in the
+  // _atom_site records, and therefore do not include missing residues.
+  std::pair<std::size_t, std::size_t> residue_range(
+      std::size_t chain_index) const {
+    if (chain_index > 0) {
+      return {chain_ends_[chain_index - 1], chain_ends_[chain_index]};
+    } else {
+      return {0, chain_ends_[0]};
+    }
+  }
+
+  // Returns the start index and one past the last index of a given residue.
+  // A residue_index of n refers to the n-th residue in the mmCIF, not
+  // including residues that are unresolved (i.e. only using _atom_site).
+  std::pair<std::size_t, std::size_t> atom_range(
+      std::size_t residue_index) const {
+    if (residue_index > 0) {
+      return {residue_ends_[residue_index - 1], residue_ends_[residue_index]};
+    } else {
+      return {model_offset_, residue_ends_[residue_index]};
+    }
+  }
+
+  // If model_id was provided during construction then this is 1, otherwise
+  // it is the number of models present in the mmCIF.
+  std::size_t num_models() const { return num_models_; }
+  // The number of atoms in the chosen model.
+  std::size_t num_atoms() const {
+    return residue_ends_.empty() ? 0 : residue_ends_.back() - model_offset_;
+  }
+  // The number of chains in the chosen model.
+  std::size_t num_chains() const { return chain_ends_.size(); }
+  // The number of residues in the chosen model, not counting unresolved
+  // residues.
+  std::size_t num_residues() const { return residue_ends_.size(); }
+
+  // Returns the first atom index that is part of the specified chain.
+  // The chain is specified using chain_index, which is a 0-based
+  // enumeration of the chains in the _atom_site table.
+  std::size_t atom_site_from_chain_index(std::size_t chain_index) const {
+    if (chain_index == 0) {
+      return model_offset_;
+    }
+    return atom_site_from_residue_index(chain_ends_[chain_index - 1]);
+  }
+
+  // Returns the first atom index that is part of the specified residue.
+  // The residue is specified using residue_index, which is a 0-based
+  // enumeration of the residues in the _atom_site table.
+  std::size_t atom_site_from_residue_index(std::size_t residues_index) const {
+    if (residues_index == 0) {
+      return model_offset_;
+    }
+    return residue_ends_[residues_index - 1];
+  }
+
+  // One past last residue index of each chain. The residue index does not
+  // include unresolved residues and is a simple 0-based enumeration of the
+  // residues in _atom_site table.
+  const std::vector<std::size_t>& chains() const { return chain_ends_; }
+
+  // Indices of the first atom of each chain. Note that this returns atom
+  // indices (like residue_starts()), not residue indices (like chains()).
+  std::vector<std::size_t> chain_starts() const;
+
+  // One past last atom index of each residue.
+  const std::vector<std::size_t>& residues() const { return residue_ends_; }
+
+  // Indices of the first atom of each residue.
+  std::vector<std::size_t> residue_starts() const {
+    std::vector<std::size_t> residue_starts;
+    if (!residue_ends_.empty()) {
+      residue_starts.reserve(residue_ends_.size());
+      residue_starts.push_back(model_offset_);
+      residue_starts.insert(residue_starts.end(), residue_ends_.begin(),
+                            residue_ends_.end() - 1);
+    }
+    return residue_starts;
+  }
+
+  // The first atom index that is part of the specified model.
+  std::size_t model_offset() const { return model_offset_; }
+
+  void Filter(absl::Span<const std::uint64_t> keep_indices);
+
+ private:
+  std::vector<std::size_t> chain_ends_;
+  std::vector<std::size_t> residue_ends_;
+  std::size_t model_offset_;
+  std::size_t num_models_;
+};
+
+}  // namespace alphafold3
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_H_

flax_model/alphafold3/structure/cpp/mmcif_layout.pyi

@@ -0,0 +1,18 @@
+
+
+from flax_model.alphafold3.cpp import cif_dict
+
+class MmcifLayout:
+  def atom_range(self, residue_index: int) -> tuple[int, int]: ...
+  def chain_starts(self) -> list[int]: ...
+  def chains(self) -> list[int]: ...
+  def model_offset(self) -> int: ...
+  def num_atoms(self) -> int: ...
+  def num_chains(self) -> int: ...
+  def num_models(self) -> int: ...
+  def num_residues(self) -> int: ...
+  def residue_range(self, chain_index: int) -> tuple[int, int]: ...
+  def residue_starts(self) -> list[int]: ...
+  def residues(self) -> list[int]: ...
+
+def from_mmcif(mmcif: cif_dict.CifDict, model_id: str = ...) -> MmcifLayout: ...

flax_model/alphafold3/structure/cpp/mmcif_layout_lib.cc

@@ -0,0 +1,213 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "alphafold3/structure/cpp/mmcif_layout.h"
+
+namespace alphafold3 {
+
+std::string MmcifLayout::ToDebugString() const {
+  return absl::StrFormat(
+      "MmcifLayout(models=%d, chains=%d, num_residues=%d, atoms=%d)",
+      num_models(), num_chains(), num_residues(), num_atoms());
+}
+
+// Changes layout to match keep_indices removing empty chains/residues.
+void MmcifLayout::Filter(absl::Span<const std::uint64_t> keep_indices) {
+  if (num_chains() == 0) {
+    return;
+  }
+  // Update residue indices.
+  auto keep_it = absl::c_lower_bound(keep_indices, residue_ends_.front());
+  for (auto& residue : residue_ends_) {
+    while (keep_it != keep_indices.end() && *keep_it < residue) {
+      ++keep_it;
+    }
+    residue = std::distance(keep_indices.begin(), keep_it);
+  }
+  // Unique residue_ends_ with updating chains.
+  auto first = residue_ends_.begin();
+  auto tail = first;
+  std::size_t num_skipped = 0;
+  std::size_t current = 0;
+  for (std::size_t& chain_end : chain_ends_) {
+    for (auto e = residue_ends_.begin() + chain_end; first != e; ++first) {
+      std::size_t next = *first;
+      *tail = next;
+      if (current != next) {
+        current = next;
+        ++tail;
+      } else {
+        ++num_skipped;
+      }
+    }
+    chain_end -= num_skipped;
+  }
+  residue_ends_.erase(tail, residue_ends_.end());
+
+  current = 0;
+  chain_ends_.erase(std::remove_if(chain_ends_.begin(), chain_ends_.end(),
+                                   [&current](std::size_t next) {
+                                     bool result = current == next;
+                                     current = next;
+                                     return result;
+                                   }),
+                    chain_ends_.end());
+  model_offset_ = 0;
+}
+
+absl::StatusOr<MmcifLayout> MmcifLayout::Create(const CifDict& mmcif,
+                                                absl::string_view model_id) {
+  auto model_ids = mmcif["_atom_site.pdbx_PDB_model_num"];
+  auto chain_ids = mmcif["_atom_site.label_asym_id"];     // chain ID.
+  auto label_seq_ids = mmcif["_atom_site.label_seq_id"];  // residue ID.
+  auto auth_seq_ids = mmcif["_atom_site.auth_seq_id"];    // author residue ID.
+  auto insertion_codes = mmcif["_atom_site.pdbx_PDB_ins_code"];
+
+  if (model_ids.size() != chain_ids.size() ||
+      model_ids.size() != label_seq_ids.size() ||
+      (model_ids.size() != auth_seq_ids.size() && !auth_seq_ids.empty()) ||
+      (model_ids.size() != insertion_codes.size() &&
+       !insertion_codes.empty())) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid _atom_site table.",  //
+        " len(_atom_site.pdbx_PDB_model_num): ", model_ids.size(),
+        " len(_atom_site.label_asym_id): ", chain_ids.size(),
+        " len(_atom_site.label_seq_id): ", label_seq_ids.size(),
+        " len(_atom_site.auth_seq_id): ", auth_seq_ids.size(),
+        " len(_atom_site.pdbx_PDB_ins_code): ", insertion_codes.size()));
+  }
+  std::size_t num_atoms = model_ids.size();
+  if (num_atoms == 0) {
+    return MmcifLayout({}, {}, 0, 0);
+  }
+  std::size_t model_offset = 0;
+  std::size_t num_models;
+  std::size_t num_atoms_per_model;
+  if (model_id.empty()) {
+    absl::string_view first_model_id = model_ids.front();
+
+    // Binary search for where the first model ends.
+    num_atoms_per_model = std::distance(
+        model_ids.begin(),
+        absl::c_upper_bound(model_ids, first_model_id, std::not_equal_to<>{}));
+    if (num_atoms % num_atoms_per_model != 0) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Each model must have the same number of atoms: (", num_atoms, " % ",
+          num_atoms_per_model, " == ", num_atoms % num_atoms_per_model, ")."));
+    }
+    num_models = num_atoms / num_atoms_per_model;
+    // Test boundary conditions for each model hold.
+    for (std::size_t i = 1; i < num_models; ++i) {
+      if ((model_ids[i * num_atoms_per_model] !=
+           model_ids[(i + 1) * num_atoms_per_model - 1]) ||
+          (model_ids[i * num_atoms_per_model - 1] ==
+           model_ids[i * num_atoms_per_model])) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("Each model must have the same number of atoms: (",
+                         num_atoms, " % ", num_atoms_per_model,
+                         " == ", num_atoms % num_atoms_per_model, ")."));
+      }
+    }
+  } else {
+    num_models = 1;
+    model_offset =
+        std::distance(model_ids.begin(), absl::c_find(model_ids, model_id));
+    if (model_offset == model_ids.size()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Unknown model_id: ", model_id));
+    }
+    model_ids.remove_prefix(model_offset);
+    chain_ids.remove_prefix(model_offset);
+    label_seq_ids.remove_prefix(model_offset);
+    if (!auth_seq_ids.empty()) auth_seq_ids.remove_prefix(model_offset);
+    if (!insertion_codes.empty()) insertion_codes.remove_prefix(model_offset);
+
+    num_atoms_per_model = std::distance(
+        model_ids.begin(), std::upper_bound(model_ids.begin(), model_ids.end(),
+                                            model_id, std::not_equal_to<>{}));
+    num_atoms = num_atoms_per_model;
+  }
+  std::vector<std::size_t> residues;
+  std::vector<std::size_t> chains;
+  absl::string_view chain_id = chain_ids.front();
+  if (!auth_seq_ids.empty() && !insertion_codes.empty()) {
+    // If author residue IDs are present then these are preferred to
+    // label residue IDs because they work for multi-residue ligands (which
+    // are given constant "." label residue IDs).
+    // NB: Author residue IDs require both the auth_seq_id and the insertion
+    // code to be unique.
+    absl::string_view auth_seq_id = auth_seq_ids.front();
+    absl::string_view insertion_code = insertion_codes.front();
+    for (std::size_t i = 1; i < num_atoms_per_model; ++i) {
+      if (absl::string_view current_chain_id = chain_ids[i];
+          current_chain_id != chain_id) {
+        residues.push_back(i + model_offset);
+        chains.push_back(residues.size());
+        chain_id = current_chain_id;
+        auth_seq_id = auth_seq_ids[i];
+        insertion_code = insertion_codes[i];
+      } else if (absl::string_view current_seq_id = auth_seq_ids[i],
+                 current_insertion_code = insertion_codes[i];
+                 insertion_code != current_insertion_code ||
+                 auth_seq_id != current_seq_id) {
+        residues.push_back(i + model_offset);
+        auth_seq_id = current_seq_id;
+        insertion_code = current_insertion_code;
+      }
+    }
+  } else {
+    absl::string_view label_seq_id = label_seq_ids.front();
+    for (std::size_t i = 1; i < num_atoms_per_model; ++i) {
+      if (absl::string_view current_chain_id = chain_ids[i];
+          current_chain_id != chain_id) {
+        residues.push_back(i + model_offset);
+        chains.push_back(residues.size());
+        chain_id = current_chain_id;
+        label_seq_id = label_seq_ids[i];
+      } else if (absl::string_view current_seq_id = label_seq_ids[i];
+                 label_seq_id != current_seq_id) {
+        residues.push_back(i + model_offset);
+        label_seq_id = current_seq_id;
+      }
+    }
+  }
+  residues.push_back(num_atoms_per_model + model_offset);
+  chains.push_back(residues.size());
+  return MmcifLayout(std::move(chains), std::move(residues), model_offset,
+                     num_models);
+}
+
+std::vector<std::size_t> MmcifLayout::chain_starts() const {
+  std::vector<std::size_t> chain_starts;
+  chain_starts.reserve(chain_ends_.size());
+  for (std::size_t index = 0; index < chain_ends_.size(); ++index) {
+    chain_starts.push_back(atom_site_from_chain_index(index));
+  }
+  return chain_starts;
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_layout_pybind.cc

@@ -0,0 +1,49 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include "alphafold3/structure/cpp/mmcif_layout.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"
+
+namespace alphafold3 {
+
+namespace py = pybind11;
+
+void RegisterModuleMmcifLayout(pybind11::module m) {
+  py::class_<MmcifLayout>(m, "MmcifLayout")
+      .def("__str__", &MmcifLayout::ToDebugString)
+      .def("num_models", &MmcifLayout::num_models)
+      .def("num_chains", &MmcifLayout::num_chains)
+      .def("num_residues", &MmcifLayout::num_residues)
+      .def("num_atoms", &MmcifLayout::num_atoms)
+      .def("residue_range", &MmcifLayout::residue_range, py::arg("chain_index"))
+      .def("atom_range", &MmcifLayout::atom_range, py::arg("residue_index"))
+      .def("chains", &MmcifLayout::chains,
+           py::doc("Returns a list of indices one past the last residue of "
+                   "each chain."))
+      .def(
+          "chain_starts", &MmcifLayout::chain_starts,
+          py::doc("Returns a list of indices of the first atom of each chain."))
+      .def("residues", &MmcifLayout::residues,
+           py::doc("Returns a list of indices one past the last atom of each "
+                   "residue."))
+      .def("residue_starts", &MmcifLayout::residue_starts,
+           py::doc(
+               "Returns a list of indices of the first atom of each residue."))
+      .def("model_offset", &MmcifLayout::model_offset,
+           py::doc("Returns the first atom index that is part of the specified "
+                   "model."));
+
+  m.def("from_mmcif", &MmcifLayout::Create, py::arg("mmcif"),
+        py::arg("model_id") = "");
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_layout_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleMmcifLayout(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_LAYOUT_PYBIND_H_

flax_model/alphafold3/structure/cpp/mmcif_struct_conn.h

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_H_
+
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+
+namespace alphafold3 {
+
+// Returns a pair of atom indices for each row in the bonds table (aka
+// _struct_conn). The indices are simple 0-based indexes into the columns of
+// the _atom_site table in the input mmCIF, and do not necessarily correspond
+// to the values in _atom_site.id, or any other column.
+absl::StatusOr<std::pair<std::vector<std::size_t>, std::vector<std::size_t>>>
+GetBondAtomIndices(const CifDict& mmcif, absl::string_view model_id);
+
+}  // namespace alphafold3
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_H_

flax_model/alphafold3/structure/cpp/mmcif_struct_conn.pyi

@@ -0,0 +1,5 @@
+
+
+from flax_model.alphafold3.cpp import cif_dict
+
+def get_bond_atom_indices(mmcif_dict: cif_dict.CifDict, model_id: str) -> tuple[list[int],list[int]]: ...

flax_model/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc

@@ -0,0 +1,380 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <cstddef>
+#include <iterator>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "alphafold3/structure/cpp/mmcif_struct_conn.h"
+
+namespace alphafold3 {
+
+namespace {
+
+struct AtomId {
+  absl::string_view chain_id;
+  absl::string_view res_id_1;
+  absl::string_view res_id_2;
+  absl::string_view atom_name;
+  absl::string_view alt_id;
+
+  friend bool operator==(const AtomId&, const AtomId&) = default;
+  template <typename H>
+  friend H AbslHashValue(H h, const AtomId& m) {
+    return H::combine(std::move(h), m.chain_id, m.res_id_1, m.res_id_2,
+                      m.atom_name, m.alt_id);
+  }
+};
+
+using StringArrayRef = absl::Span<const std::string>;
+using BondIndexByAtom = absl::flat_hash_map<AtomId, std::vector<std::size_t>>;
+using BondAtomIndices = std::vector<std::size_t>;
+
+// Returns whether each container is the same size.
+template <typename C, typename... Cs>
+bool AreSameSize(const C& c, const Cs&... cs) {
+  return ((c.size() == cs.size()) && ...);
+}
+
+struct ColumnSpec {
+  absl::string_view chain_id_col;
+  absl::string_view res_id_1_col;
+  absl::string_view res_id_2_col;
+  absl::string_view atom_name_col;
+  std::optional<absl::string_view> alt_id_col;  // Not used by OpenMM.
+};
+
+class AtomColumns {
+ public:
+  static absl::StatusOr<AtomColumns> Create(const CifDict& mmcif,
+                                            const ColumnSpec& column_spec) {
+    StringArrayRef chain_id = mmcif[column_spec.chain_id_col];
+    StringArrayRef res_id_1 = mmcif[column_spec.res_id_1_col];
+    StringArrayRef res_id_2 = mmcif[column_spec.res_id_2_col];
+    StringArrayRef atom_name = mmcif[column_spec.atom_name_col];
+    if (!AreSameSize(chain_id, res_id_1, res_id_2, atom_name)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Atom columns are not the same size. ",                       //
+          "len(", column_spec.chain_id_col, ")=", chain_id.size(),      //
+          ", len(", column_spec.res_id_1_col, ")=", res_id_1.size(),    //
+          ", len(", column_spec.res_id_2_col, ")=", res_id_2.size(),    //
+          ", len(", column_spec.atom_name_col, ")=", atom_name.size(),  //
+          "."));
+    }
+    if (column_spec.alt_id_col.has_value()) {
+      StringArrayRef alt_id = mmcif[*column_spec.alt_id_col];
+      if (!AreSameSize(alt_id, chain_id)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Atom columns are not the same size. ",                   //
+            "len(", column_spec.chain_id_col, ")=", chain_id.size(),  //
+            ", len(", *column_spec.alt_id_col, ")=", alt_id.size(),   //
+            "."));
+      }
+      return AtomColumns(chain_id, res_id_1, res_id_2, atom_name, alt_id,
+                         column_spec);
+    } else {
+      return AtomColumns(chain_id, res_id_1, res_id_2, atom_name, std::nullopt,
+                         column_spec);
+    }
+  }
+
+  inline std::size_t size() const { return size_; }
+
+  absl::string_view GetNormalizedAltId(const std::size_t index) const {
+    constexpr absl::string_view kFullStop = ".";
+    if (alt_id_.has_value()) {
+      absl::string_view alt_id = (*alt_id_)[index];
+      return alt_id == "?" ? kFullStop : alt_id;
+    } else {
+      return kFullStop;
+    }
+  }
+
+  AtomId GetAtom(const std::size_t index) const {
+    return {.chain_id = chain_id_[index],
+            .res_id_1 = res_id_1_[index],
+            .res_id_2 = res_id_2_[index],
+            .atom_name = atom_name_[index],
+            .alt_id = GetNormalizedAltId(index)};
+  }
+
+  std::string GetAtomString(const std::size_t index) const {
+    std::string alt_id_col;
+    if (column_spec_.alt_id_col.has_value()) {
+      alt_id_col = *column_spec_.alt_id_col;
+    } else {
+      alt_id_col = "default label_alt_id";
+    }
+    return absl::StrCat(
+        column_spec_.chain_id_col, "=", chain_id_[index], ", ",    //
+        column_spec_.res_id_1_col, "=", res_id_1_[index], ", ",    //
+        column_spec_.res_id_2_col, "=", res_id_2_[index], ", ",    //
+        column_spec_.atom_name_col, "=", atom_name_[index], ", ",  //
+        alt_id_col, "=", GetNormalizedAltId(index));               //
+  }
+
+ private:
+  AtomColumns(StringArrayRef chain_id, StringArrayRef res_id_1,
+              StringArrayRef res_id_2, StringArrayRef atom_name,
+              std::optional<StringArrayRef> alt_id,
+              const ColumnSpec& column_spec)
+      : chain_id_(chain_id),
+        res_id_1_(res_id_1),
+        res_id_2_(res_id_2),
+        atom_name_(atom_name),
+        alt_id_(alt_id),
+        column_spec_(column_spec),
+        size_(chain_id.size()) {}
+  StringArrayRef chain_id_;
+  StringArrayRef res_id_1_;
+  StringArrayRef res_id_2_;
+  StringArrayRef atom_name_;
+  std::optional<StringArrayRef> alt_id_;
+  ColumnSpec column_spec_;
+  std::size_t size_;
+};
+
+// Adds the atom index to any rows in the bond table involving that atom.
+absl::Status FillInBondsForAtom(const BondIndexByAtom& bond_index_by_atom,
+                                const AtomId& atom,
+                                const std::size_t atom_index,
+                                BondAtomIndices& bond_atom_indices) {
+  if (auto bond_index_it = bond_index_by_atom.find(atom);
+      bond_index_it != bond_index_by_atom.end()) {
+    for (std::size_t bond_index : bond_index_it->second) {
+      if (bond_index < 0 || bond_index >= bond_atom_indices.size()) {
+        return absl::OutOfRangeError(
+            absl::StrCat("Bond index out of range: ", bond_index));
+      }
+      bond_atom_indices[bond_index] = atom_index;
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Checks that the CifDict has all of the columns in the column spec.
+bool HasAllColumns(const CifDict& mmcif, const ColumnSpec& columns) {
+  return mmcif.Contains(columns.chain_id_col) &&
+         mmcif.Contains(columns.res_id_1_col) &&
+         mmcif.Contains(columns.res_id_2_col) &&
+         mmcif.Contains(columns.atom_name_col) &&
+         (!columns.alt_id_col.has_value() ||
+          mmcif.Contains(*columns.alt_id_col));
+}
+
+// Fully specified ptnr1 atom.
+constexpr ColumnSpec kStructConnPtnr1ColumnsFull{
+    .chain_id_col = "_struct_conn.ptnr1_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr1_auth_seq_id",
+    .res_id_2_col = "_struct_conn.pdbx_ptnr1_PDB_ins_code",
+    .atom_name_col = "_struct_conn.ptnr1_label_atom_id",
+    .alt_id_col = "_struct_conn.pdbx_ptnr1_label_alt_id",
+};
+
+// Fully specified ptnr2 atom.
+constexpr ColumnSpec kStructConnPtnr2ColumnsFull{
+    .chain_id_col = "_struct_conn.ptnr2_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr2_auth_seq_id",
+    .res_id_2_col = "_struct_conn.pdbx_ptnr2_PDB_ins_code",
+    .atom_name_col = "_struct_conn.ptnr2_label_atom_id",
+    .alt_id_col = "_struct_conn.pdbx_ptnr2_label_alt_id",
+};
+
+// Columns used by OpenMM for ptnr1 atoms.
+constexpr ColumnSpec kStructConnPtnr1OpenMM{
+    .chain_id_col = "_struct_conn.ptnr1_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr1_label_seq_id",
+    .res_id_2_col = "_struct_conn.ptnr1_label_comp_id",
+    .atom_name_col = "_struct_conn.ptnr1_label_atom_id",
+    .alt_id_col = std::nullopt,
+};
+
+// Columns used by OpenMM for ptnr2 atoms.
+constexpr ColumnSpec kStructConnPtnr2OpenMM{
+    .chain_id_col = "_struct_conn.ptnr2_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr2_label_seq_id",
+    .res_id_2_col = "_struct_conn.ptnr2_label_comp_id",
+    .atom_name_col = "_struct_conn.ptnr2_label_atom_id",
+    .alt_id_col = std::nullopt,
+};
+
+// Fully specified atom sites.
+constexpr ColumnSpec kAtomSiteColumnsFull{
+    .chain_id_col = "_atom_site.label_asym_id",
+    .res_id_1_col = "_atom_site.auth_seq_id",
+    .res_id_2_col = "_atom_site.pdbx_PDB_ins_code",
+    .atom_name_col = "_atom_site.label_atom_id",
+    .alt_id_col = "_atom_site.label_alt_id",
+};
+
+// Atom site columns used to match OpenMM _struct_conn tables.
+constexpr ColumnSpec kAtomSiteColumnsOpenMM{
+    .chain_id_col = "_atom_site.label_asym_id",
+    .res_id_1_col = "_atom_site.label_seq_id",
+    .res_id_2_col = "_atom_site.label_comp_id",
+    .atom_name_col = "_atom_site.label_atom_id",
+    .alt_id_col = "_atom_site.label_alt_id",
+};
+
+}  // namespace
+
+absl::StatusOr<std::pair<BondAtomIndices, BondAtomIndices>> GetBondAtomIndices(
+    const CifDict& mmcif, absl::string_view model_id) {
+  ColumnSpec ptnr1_columns, ptnr2_columns, atom_site_columns;
+
+  if (HasAllColumns(mmcif, kStructConnPtnr1ColumnsFull) &&
+      HasAllColumns(mmcif, kStructConnPtnr2ColumnsFull)) {
+    ptnr1_columns = kStructConnPtnr1ColumnsFull;
+    ptnr2_columns = kStructConnPtnr2ColumnsFull;
+    atom_site_columns = kAtomSiteColumnsFull;
+  } else {
+    ptnr1_columns = kStructConnPtnr1OpenMM;
+    ptnr2_columns = kStructConnPtnr2OpenMM;
+    atom_site_columns = kAtomSiteColumnsOpenMM;
+  }
+
+  absl::StatusOr<AtomColumns> ptnr1_atoms =
+      AtomColumns::Create(mmcif, ptnr1_columns);
+  if (!ptnr1_atoms.ok()) {
+    return ptnr1_atoms.status();
+  }
+  absl::StatusOr<AtomColumns> ptnr2_atoms =
+      AtomColumns::Create(mmcif, ptnr2_columns);
+  if (!ptnr2_atoms.ok()) {
+    return ptnr2_atoms.status();
+  }
+  StringArrayRef struct_conn_id = mmcif["_struct_conn.id"];
+  if (!AreSameSize(struct_conn_id, *ptnr1_atoms, *ptnr2_atoms)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid '_struct_conn.' loop. ",                  //
+        "len(id) = ", struct_conn_id.size(), ", ",         //
+        "len(ptnr1_atoms) = ", ptnr1_atoms->size(), ", ",  //
+        "len(ptnr2_atoms) = ", ptnr2_atoms->size(), "."    //
+        ));
+  }
+
+  absl::StatusOr<AtomColumns> atoms =
+      AtomColumns::Create(mmcif, atom_site_columns);
+  if (!atoms.ok()) {
+    return atoms.status();
+  }
+  StringArrayRef atom_site_id = mmcif["_atom_site.id"];
+  StringArrayRef atom_site_model_id = mmcif["_atom_site.pdbx_PDB_model_num"];
+  if (!AreSameSize(atom_site_id, atom_site_model_id, *atoms)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid '_atom_site.' loop. ",                                //
+        "len(id)= ", atom_site_id.size(), ", ",                        //
+        "len(pdbx_PDB_model_num)= ", atom_site_model_id.size(), ", ",  //
+        "len(atoms)= ", atoms->size(), "."));                          //
+  }
+
+  // Build maps from atom ID tuples to the rows in _struct_conn where that
+  // atom appears (NB could be multiple).
+  const std::size_t struct_conn_size = struct_conn_id.size();
+  BondIndexByAtom ptnr1_rows_by_atom(struct_conn_size);
+  BondIndexByAtom ptnr2_rows_by_atom(struct_conn_size);
+  for (std::size_t i = 0; i < struct_conn_size; ++i) {
+    ptnr1_rows_by_atom[ptnr1_atoms->GetAtom(i)].push_back(i);
+    ptnr2_rows_by_atom[ptnr2_atoms->GetAtom(i)].push_back(i);
+  }
+
+  // Allocate two output arrays with one element per row in struct_conn, where
+  // each element will be the index of that atom in the atom_site table.
+  // Fill the arrays with atom_site_size, which is an invalid value, so that
+  // we can check at the end that each atom has been found.
+  const std::size_t atom_site_size = atom_site_id.size();
+  BondAtomIndices ptnr1_atom_indices(struct_conn_size, atom_site_size);
+  BondAtomIndices ptnr2_atom_indices(struct_conn_size, atom_site_size);
+
+  bool model_id_ecountered = false;
+  absl::flat_hash_set<absl::string_view> seen_alt_ids;
+  for (std::size_t atom_i = 0; atom_i < atom_site_size; ++atom_i) {
+    if (atom_site_model_id[atom_i] != model_id) {
+      if (!model_id_ecountered) {
+        continue;
+      } else {
+        // Models are contiguous so once we see a different model ID after
+        // encountering our model ID then we can exit early.
+        break;
+      }
+    } else {
+      model_id_ecountered = true;
+    }
+    AtomId atom = atoms->GetAtom(atom_i);
+    seen_alt_ids.insert(atom.alt_id);
+
+    if (auto fill_in_bonds_status1 = FillInBondsForAtom(
+            ptnr1_rows_by_atom, atom, atom_i, ptnr1_atom_indices);
+        !fill_in_bonds_status1.ok()) {
+      return fill_in_bonds_status1;
+    }
+    if (auto fill_in_bonds_status2 = FillInBondsForAtom(
+            ptnr2_rows_by_atom, atom, atom_i, ptnr2_atom_indices);
+        !fill_in_bonds_status2.ok()) {
+      return fill_in_bonds_status2;
+    }
+  }
+  // The seen_alt_ids check is a workaround for a known PDB issue: some mmCIFs
+  // (2evw, 2g0v, 2g0x, 2g0z, 2g10, 2g11, 2g12, 2g14, 2grz, 2ntw as of 2024)
+  // have multiple models and they set different whole-chain altloc in each
+  // model. The bond table however doesn't distinguish between models, so there
+  // are bonds that are valid only for some models. E.g. 2grz has model 1 with
+  // chain A with altloc A, and model 2 with chain A with altloc B. The bonds
+  // table lists a bond for each of these.
+
+  // Check that a ptnr1 atom was found for every bond.
+  if (auto row_it = absl::c_find(ptnr1_atom_indices, atom_site_size);
+      row_it != ptnr1_atom_indices.end()) {
+    if (seen_alt_ids.size() > 1 || seen_alt_ids.contains(".") ||
+        seen_alt_ids.contains("?")) {
+      std::size_t i = std::distance(ptnr1_atom_indices.begin(), row_it);
+      return absl::InvalidArgumentError(
+          absl::StrCat("Error parsing \"", mmcif.GetDataName(), "\". ",
+                       "Cannot find atom for bond ID ", struct_conn_id[i], ": ",
+                       ptnr1_atoms->GetAtomString(i)));
+    }
+  }
+
+  // Check that a ptnr2 atom was found for every bond.
+  if (auto row_it = absl::c_find(ptnr2_atom_indices, atom_site_size);
+      row_it != ptnr2_atom_indices.end()) {
+    if (seen_alt_ids.size() > 1 || seen_alt_ids.contains(".") ||
+        seen_alt_ids.contains("?")) {
+      std::size_t i = std::distance(ptnr2_atom_indices.begin(), row_it);
+      return absl::InvalidArgumentError(
+          absl::StrCat("Error parsing \"", mmcif.GetDataName(), "\". ",
+                       "Cannot find atom for bond ID ", struct_conn_id[i], ": ",
+                       ptnr2_atoms->GetAtomString(i)));
+    }
+  }
+
+  if (!model_id_ecountered) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Error parsing \"", mmcif.GetDataName(), "\". model_id \"", model_id,
+        "\" not found in _atom_site.pdbx_PDB_model_num."));
+  }
+
+  return std::make_pair(std::move(ptnr1_atom_indices),
+                        std::move(ptnr2_atom_indices));
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_struct_conn_pybind.cc

@@ -0,0 +1,68 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "alphafold3/structure/cpp/mmcif_struct_conn.h"
+#include "pybind11/gil.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"
+
+namespace alphafold3 {
+
+namespace py = pybind11;
+
+constexpr char kGetBondAtomIndices[] = R"(
+Extracts the indices of the atoms that participate in bonds.
+
+This function has a workaround for a known PDB issue: some mmCIFs have
+(2evw, 2g0v, 2g0x, 2g0z, 2g10, 2g11, 2g12, 2g14, 2grz, 2ntw as of 2024)
+multiple models and they set different whole-chain altloc in each model.
+The bond table however doesn't distinguish between models, so there are
+bonds that are valid only for some models. E.g. 2grz has model 1 with
+chain A with altloc A, and model 2 with chain A with altloc B. The bonds
+table lists a bond for each of these. This case is rather rare (10 cases
+in PDB as of 2024). For the offending bonds, the returned atom index is
+set to the size of the atom_site table, i.e. it is an invalid index.
+
+Args:
+  mmcif: The mmCIF object to process.
+  model_id: The ID of the model that the returned atoms will belong to. This
+    should be a value in the mmCIF's _atom_site.pdbx_PDB_model_num column.
+
+Returns:
+  Two lists of atom indices, `from_atoms` and `to_atoms`, each one having
+  length num_bonds (as defined by _struct_conn, the bonds table). The bond
+  i, defined by the i'th row in _struct_conn, is a bond from atom at index
+  from_atoms[i], to the atom at index to_atoms[i]. The indices are simple
+  0-based indexes into the columns of the _atom_site table in the input
+  mmCIF, and do not necessarily correspond to the values in _atom_site.id,
+  or any other column.
+)";
+
+void RegisterModuleMmcifStructConn(pybind11::module m) {
+  m.def(
+      "get_bond_atom_indices",
+      [](const CifDict& mmcif, absl::string_view model_id) {
+        auto result = GetBondAtomIndices(mmcif, model_id);
+        if (result.ok()) {
+          return *result;
+        }
+        throw py::value_error(std::string(result.status().message()));
+      },
+      py::arg("mmcif_dict"), py::arg("model_id"),
+      py::doc(kGetBondAtomIndices + 1),
+      py::call_guard<py::gil_scoped_release>());
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_struct_conn_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleMmcifStructConn(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_STRUCT_CONN_PYBIND_H_

flax_model/alphafold3/structure/cpp/mmcif_utils.pyi

@@ -0,0 +1,63 @@
+
+
+from collections.abc import Sequence
+
+import numpy as np
+
+from flax_model.alphafold3.cpp import cif_dict
+from flax_model.alphafold3.structure.python import mmcif_layout
+
+
+def filter(
+    mmcif: cif_dict.CifDict,
+    include_nucleotides: bool,
+    include_ligands: bool = ...,
+    include_water: bool = ...,
+    include_other: bool = ...,
+    model_id: str = ...,
+) -> tuple[np.ndarray[int], mmcif_layout.MmcifLayout]: ...
+
+
+def fix_residues(
+    layout: mmcif_layout.MmcifLayout,
+    comp_id: Sequence[str],
+    atom_id: Sequence[str],
+    atom_x: Sequence[float],
+    atom_y: Sequence[float],
+    atom_z: Sequence[float],
+    fix_arg: bool = ...,
+) -> None: ...
+
+
+def read_layout(
+    mmcif: cif_dict.CifDict, model_id: str = ...
+) -> mmcif_layout.MmcifLayout: ...
+
+
+def selected_ligand_residue_mask(
+    layout: mmcif_layout.MmcifLayout,
+    atom_site_label_asym_ids: list[str],
+    atom_site_label_seq_ids: list[str],
+    atom_site_auth_seq_ids: list[str],
+    atom_site_label_comp_ids: list[str],
+    atom_site_pdbx_pdb_ins_codes: list[str],
+    nonpoly_asym_ids: list[str],
+    nonpoly_auth_seq_ids: list[str],
+    nonpoly_pdb_ins_codes: list[str],
+    nonpoly_mon_ids: list[str],
+    branch_asym_ids: list[str],
+    branch_auth_seq_ids: list[str],
+    branch_pdb_ins_codes: list[str],
+    branch_mon_ids: list[str],
+) -> tuple[list[bool], list[bool]]: ...
+
+
+def selected_polymer_residue_mask(
+    layout: mmcif_layout.MmcifLayout,
+    atom_site_label_asym_ids: list[str],
+    atom_site_label_seq_ids: list[str],
+    atom_site_label_comp_ids: list[str],
+    poly_seq_asym_ids: list[str],
+    poly_seq_seq_ids: list[str],
+    poly_seq_mon_ids: list[str],
+) -> list[bool]: ...

flax_model/alphafold3/structure/cpp/mmcif_utils_pybind.cc

@@ -0,0 +1,787 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <Python.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <new>
+#include <optional>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "numpy/ndarrayobject.h"
+#include "numpy/ndarraytypes.h"
+#include "numpy/npy_common.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/memory/memory.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "alphafold3/structure/cpp/mmcif_altlocs.h"
+#include "alphafold3/structure/cpp/mmcif_layout.h"
+#include "pybind11/cast.h"
+#include "pybind11/gil.h"
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"
+#include "pybind11_abseil/absl_casters.h"
+
+namespace alphafold3 {
+namespace {
+namespace py = pybind11;
+
+struct PyObjectDeleter {
+  inline void operator()(PyObject* obj) const { Py_CLEAR(obj); }
+};
+
+using ScopedPyObject = std::unique_ptr<PyObject, PyObjectDeleter>;
+
+using StringArrayRef = absl::Span<const std::string>;
+using Indexer = absl::flat_hash_map<absl::string_view, std::size_t>;
+
+// Returns the reverse look-up map of name to index.
+Indexer MakeIndex(StringArrayRef col) {
+  Indexer index;
+  index.reserve(col.size());
+  for (std::size_t i = 0; i < col.size(); ++i) {
+    index[col[i]] = i;
+  }
+  return index;
+}
+
+// Returns whether each container is the same size.
+template <typename C, typename... Cs>
+bool AreSameSize(C c, const Cs&... cs) {
+  return ((c.size() == cs.size()) && ...);
+}
+
+// Stores references to columns in `_atom_site` ensuring they all exist and
+// are the same size.
+struct AtomSiteLoop {
+  explicit AtomSiteLoop(const CifDict& cif_dict)
+      : id(cif_dict["_atom_site.id"]),
+        model_id(cif_dict["_atom_site.pdbx_PDB_model_num"]),
+        chain_id(cif_dict["_atom_site.label_asym_id"]),
+        seq_id(cif_dict["_atom_site.label_seq_id"]),
+
+        comp_id(cif_dict["_atom_site.label_comp_id"]),
+        atom_id(cif_dict["_atom_site.label_atom_id"]),
+
+        alt_id(cif_dict["_atom_site.label_alt_id"]),
+        occupancy(cif_dict["_atom_site.occupancy"])
+
+  {
+    if (!AreSameSize(id, model_id, chain_id, seq_id, comp_id, atom_id, alt_id,
+                     occupancy)) {
+      throw py::value_error(
+          absl::StrCat("Invalid '_atom_site.' loop. ",                     //
+                       "len(id)=", id.size(), ", ",                        //
+                       "len(pdbx_PDB_model_num)=", model_id.size(), ", ",  //
+                       "len(label_asym_id)=", chain_id.size(), ", ",       //
+                       "len(label_seq_id)=", seq_id.size(), ", ",          //
+                       "len(label_comp_id)=", comp_id.size(), ", ",        //
+                       "len(atom_id)=", atom_id.size(), ", ",              //
+                       "len(label_alt_id)=", alt_id.size(), ", ",          //
+                       "len(occupancy)=", occupancy.size()));
+    }
+  }
+  StringArrayRef id;
+  StringArrayRef model_id;
+  StringArrayRef chain_id;
+  StringArrayRef seq_id;
+  StringArrayRef comp_id;
+  StringArrayRef atom_id;
+  StringArrayRef alt_id;
+  StringArrayRef occupancy;
+};
+
+// Stores references to columns in `_entity` ensuring they all exist and are the
+// same size.
+struct EntityLoop {
+  explicit EntityLoop(const CifDict& cif_dict)
+      : id(cif_dict["_entity.id"]), type(cif_dict["_entity.type"]) {
+    if (!AreSameSize(id, type)) {
+      throw py::value_error(absl::StrCat("Invalid '_entity.' loop. ",  //
+                                         "len(id)=", id.size(), ", ",  //
+                                         "len(type)=", type.size()));
+    }
+  }
+  StringArrayRef id;
+  StringArrayRef type;
+};
+
+// Stores references to columns in `_entity_poly` ensuring they all exist and
+// are the same size.
+struct EntityPolyLoop {
+  explicit EntityPolyLoop(const CifDict& cif_dict)
+      : entity_id(cif_dict["_entity_poly.entity_id"]),
+        type(cif_dict["_entity_poly.type"]) {
+    if (!AreSameSize(entity_id, type)) {
+      throw py::value_error(absl::StrCat("Invalid '_entity_poly.' loop. ",  //
+                                         "len(entity_id)=", entity_id.size(),
+                                         ", ",  //
+                                         "len(type)=", type.size()));
+    }
+  }
+  StringArrayRef entity_id;
+  StringArrayRef type;
+};
+
+// Returns a set of entity names removing ones not included by the flags
+// specified.
+absl::flat_hash_set<absl::string_view> SelectChains(const CifDict& mmcif,
+                                                    bool include_nucleotides,
+                                                    bool include_ligands,
+                                                    bool include_water,
+                                                    bool include_other) {
+  EntityLoop entity_loop(mmcif);
+  EntityPolyLoop entity_poly(mmcif);
+  absl::flat_hash_set<absl::string_view> permitted_polymers{"polypeptide(L)"};
+  absl::flat_hash_set<absl::string_view> forbidden_polymers;
+  for (absl::string_view type :
+       {"polydeoxyribonucleotide", "polyribonucleotide",
+        "polydeoxyribonucleotide/polyribonucleotide hybrid"}) {
+    if (include_nucleotides) {
+      permitted_polymers.emplace(type);
+    } else {
+      forbidden_polymers.emplace(type);
+    }
+  }
+
+  absl::flat_hash_set<absl::string_view> permitted_nonpoly_entity_types;
+  absl::flat_hash_set<absl::string_view> forbidden_nonpoly_entity_types;
+  for (absl::string_view type : {"non-polymer", "branched"}) {
+    if (include_ligands) {
+      permitted_nonpoly_entity_types.emplace(type);
+    } else {
+      forbidden_nonpoly_entity_types.emplace(type);
+    }
+  }
+  absl::string_view water_type = "water";
+  if (include_water) {
+    permitted_nonpoly_entity_types.emplace(water_type);
+  } else {
+    forbidden_nonpoly_entity_types.emplace(water_type);
+  }
+
+  StringArrayRef chain_ids = mmcif["_struct_asym.id"];
+  StringArrayRef entity_ids = mmcif["_struct_asym.entity_id"];
+  Indexer chain_index = MakeIndex(chain_ids);
+  Indexer entity_poly_index = MakeIndex(entity_poly.entity_id);
+  Indexer entity_id_to_index = MakeIndex(entity_loop.id);
+
+  absl::flat_hash_set<absl::string_view> keep_chain_id;
+  for (std::size_t i = 0; i < chain_ids.size(); ++i) {
+    absl::string_view chain_id = chain_ids[i];
+    absl::string_view entity_id = entity_ids[i];
+    if (entity_id_to_index.empty() ||
+        entity_loop.type[entity_id_to_index[entity_id]] == "polymer") {
+      if (auto it = entity_poly_index.find(entity_id);
+          it != entity_poly_index.end()) {
+        absl::string_view poly_type = entity_poly.type[it->second];
+        if (include_other) {
+          if (!forbidden_polymers.contains(poly_type)) {
+            keep_chain_id.insert(chain_id);
+          }
+        } else {
+          if (permitted_polymers.contains(poly_type)) {
+            keep_chain_id.insert(chain_id);
+          }
+        }
+      }
+    } else {
+      absl::string_view entity_type =
+          entity_loop.type[entity_id_to_index[entity_id]];
+      if (include_other) {
+        if (!forbidden_nonpoly_entity_types.contains(entity_type)) {
+          keep_chain_id.insert(chain_id);
+          continue;
+        }
+      } else {
+        if (permitted_nonpoly_entity_types.contains(entity_type)) {
+          keep_chain_id.insert(chain_id);
+          continue;
+        }
+      }
+    }
+  }
+  return keep_chain_id;
+}
+
+class ProcessResidue {
+ public:
+  explicit ProcessResidue(const char* residue)
+      : residue_(PyUnicode_InternFromString(residue)) {}
+  bool IsResidue(PyObject* residue) {
+    return ArePyObjectsEqual(residue_.get(), residue);
+  }
+
+  static bool ArePyObjectsEqual(PyObject* lhs, PyObject* rhs) {
+    switch (PyObject_RichCompareBool(lhs, rhs, Py_EQ)) {
+      case -1:
+        PyErr_Clear();
+        return false;
+      case 0:
+        return false;
+      default:
+        return true;
+    }
+  }
+
+ private:
+  ScopedPyObject residue_;
+};
+
+struct Position3 {
+  float x;
+  float y;
+  float z;
+};
+
+float DistanceSquared(Position3 v1, Position3 v2) {
+  float dx = v1.x - v2.x;
+  float dy = v1.y - v2.y;
+  float dz = v1.z - v2.z;
+  return dx * dx + dy * dy + dz * dz;
+}
+
+class FixArginine : public ProcessResidue {
+ public:
+  FixArginine()
+      : ProcessResidue("ARG"),
+        cd_(PyUnicode_InternFromString("CD")),
+        nh1_(PyUnicode_InternFromString("NH1")),
+        nh2_(PyUnicode_InternFromString("NH2")),
+        hh11_(PyUnicode_InternFromString("HH11")),
+        hh21_(PyUnicode_InternFromString("HH21")),
+        hh12_(PyUnicode_InternFromString("HH12")),
+        hh22_(PyUnicode_InternFromString("HH22")) {}
+  void Fix(absl::Span<PyObject*> atom_ids, absl::Span<const float> atom_x,
+           absl::Span<const float> atom_y, absl::Span<const float> atom_z) {
+    std::ptrdiff_t cd_index = -1;
+    std::ptrdiff_t nh1_index = -1;
+    std::ptrdiff_t nh2_index = -1;
+    std::ptrdiff_t hh11_index = -1;
+    std::ptrdiff_t hh21_index = -1;
+    std::ptrdiff_t hh12_index = -1;
+    std::ptrdiff_t hh22_index = -1;
+    for (std::ptrdiff_t index = 0; index < atom_ids.size(); ++index) {
+      PyObject* atom_id = atom_ids[index];
+      if (cd_index == -1 && ArePyObjectsEqual(atom_id, cd_.get())) {
+        cd_index = index;
+      } else if (nh1_index == -1 && ArePyObjectsEqual(atom_id, nh1_.get())) {
+        nh1_index = index;
+      } else if (nh2_index == -1 && ArePyObjectsEqual(atom_id, nh2_.get())) {
+        nh2_index = index;
+      } else if (hh11_index == -1 && ArePyObjectsEqual(atom_id, hh11_.get())) {
+        hh11_index = index;
+      } else if (hh21_index == -1 && ArePyObjectsEqual(atom_id, hh21_.get())) {
+        hh21_index = index;
+      } else if (hh12_index == -1 && ArePyObjectsEqual(atom_id, hh12_.get())) {
+        hh12_index = index;
+      } else if (hh22_index == -1 && ArePyObjectsEqual(atom_id, hh22_.get())) {
+        hh22_index = index;
+      }
+    }
+    if (cd_index < 0 || nh1_index < 0 || nh2_index < 0) {
+      return;
+    }
+    Position3 cd_pos(atom_x[cd_index], atom_y[cd_index], atom_z[cd_index]);
+    Position3 nh1_pos(atom_x[nh1_index], atom_y[nh1_index], atom_z[nh1_index]);
+    Position3 nh2_pos(atom_x[nh2_index], atom_y[nh2_index], atom_z[nh2_index]);
+    if (DistanceSquared(nh1_pos, cd_pos) <= DistanceSquared(nh2_pos, cd_pos)) {
+      return;
+    }
+    std::swap(atom_ids[nh1_index], atom_ids[nh2_index]);
+    if (hh11_index >= 0 && hh21_index >= 0) {
+      std::swap(atom_ids[hh11_index], atom_ids[hh21_index]);
+    } else if (hh11_index >= 0) {
+      Py_DECREF(atom_ids[hh11_index]);
+      Py_INCREF(hh21_.get());
+      atom_ids[hh11_index] = hh21_.get();
+    } else if (hh21_index >= 0) {
+      Py_DECREF(atom_ids[hh21_index]);
+      Py_INCREF(hh11_.get());
+      atom_ids[hh21_index] = hh11_.get();
+    }
+    if (hh12_index >= 0 && hh22_index >= 0) {
+      std::swap(atom_ids[hh12_index], atom_ids[hh22_index]);
+    } else if (hh12_index >= 0) {
+      Py_DECREF(atom_ids[hh12_index]);
+      Py_INCREF(hh22_.get());
+      atom_ids[hh12_index] = hh22_.get();
+    } else if (hh22_index >= 0) {
+      Py_DECREF(atom_ids[hh22_index]);
+      Py_INCREF(hh21_.get());
+      atom_ids[hh22_index] = hh21_.get();
+    }
+  }
+
+ private:
+  ScopedPyObject cd_;
+  ScopedPyObject nh1_;
+  ScopedPyObject nh2_;
+  ScopedPyObject hh11_;
+  ScopedPyObject hh21_;
+  ScopedPyObject hh12_;
+  ScopedPyObject hh22_;
+};
+
+// Returns the layout of the mmCIF `_atom_site` table.
+inline MmcifLayout ReadMmcifLayout(const CifDict& mmcif,
+                                   absl::string_view model_id = "") {
+  py::gil_scoped_release release;
+  auto mmcif_layout = MmcifLayout::Create(mmcif, model_id);
+  if (mmcif_layout.ok()) {
+    return *mmcif_layout;
+  }
+
+  throw py::value_error(std::string(mmcif_layout.status().message()));
+}
+
+std::pair<py::object, MmcifLayout> MmcifFilter(  //
+    const CifDict& mmcif,                        //
+    bool include_nucleotides,                    //
+    bool include_ligands,                        //
+    bool include_water,                          //
+    bool include_other,                          //
+    absl::string_view model_id) {
+  if (_import_array() < 0) {
+    throw py::import_error("Failed to import NumPy.");
+  }
+  auto layout = ReadMmcifLayout(mmcif, model_id);
+  std::unique_ptr<std::vector<std::uint64_t>> keep_indices;
+  size_t new_num_atoms;
+
+  {
+    py::gil_scoped_release release;
+
+    AtomSiteLoop atom_site(mmcif);
+
+    auto keep_chain_ids =
+        SelectChains(mmcif, include_nucleotides, include_ligands, include_water,
+                     include_other);
+
+    std::vector<std::size_t> chain_indices;
+    chain_indices.reserve(keep_chain_ids.size());
+    for (std::size_t i = 0; i < layout.num_chains(); ++i) {
+      if (keep_chain_ids.contains(
+              atom_site.chain_id[layout.atom_site_from_chain_index(i)])) {
+        chain_indices.push_back(i);
+      }
+    }
+
+    keep_indices =
+        absl::WrapUnique(new std::vector<std::uint64_t>(ResolveMmcifAltLocs(
+            layout, atom_site.comp_id, atom_site.atom_id, atom_site.alt_id,
+            atom_site.occupancy, chain_indices)));
+    new_num_atoms = keep_indices->size();
+
+    if (layout.num_models() > 1) {
+      keep_indices->reserve(layout.num_models() * new_num_atoms);
+      std::uint64_t* start = &(*keep_indices->begin());
+      std::size_t num_atom = keep_indices->size();
+      // Copy first model indices into all model indices offsetting each copy.
+      for (std::size_t i = 1; i < layout.num_models(); ++i) {
+        std::size_t offset = i * layout.num_atoms();
+        std::transform(start, start + num_atom,
+                       std::back_inserter(*keep_indices),
+                       [offset](std::size_t v) { return v + offset; });
+      }
+    }
+  }
+
+  layout.Filter(*keep_indices);
+
+  npy_intp shape[] = {static_cast<npy_intp>(layout.num_models()),
+                      static_cast<npy_intp>(new_num_atoms)};
+  PyObject* arr =
+      PyArray_SimpleNewFromData(2, shape, NPY_INT64, keep_indices->data());
+  // Create a capsule to hold the memory of the buffer so NumPy knows how to
+  // delete it when done with it.
+  PyObject* capsule = PyCapsule_New(
+      keep_indices.release(), nullptr, +[](PyObject* capsule_cleanup) {
+        void* memory = PyCapsule_GetPointer(capsule_cleanup, nullptr);
+        delete static_cast<std::vector<std::size_t>*>(memory);
+      });
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(arr), capsule);
+
+  return std::make_pair(py::reinterpret_steal<py::object>(arr),
+                        std::move(layout));
+}
+
+void MmcifFixResidues(               //
+    const MmcifLayout& layout,       //
+    absl::Span<PyObject*> comp_id,   //
+    absl::Span<PyObject*> atom_id,   //
+    absl::Span<const float> atom_x,  //
+    absl::Span<const float> atom_y,  //
+    absl::Span<const float> atom_z,  //
+    bool fix_arginine                //
+) {
+  std::optional<FixArginine> arginine;
+  std::size_t num_atoms = layout.num_atoms();
+  if (comp_id.size() != num_atoms || atom_id.size() != num_atoms ||
+      atom_x.size() != num_atoms || atom_y.size() != num_atoms ||
+      atom_z.size() != num_atoms) {
+    throw py::value_error(
+        absl::StrCat("Sizes must match. ",                   //
+                     "num_atoms=", num_atoms, ", ",          //
+                     "len(comp_id)=", comp_id.size(), ", ",  //
+                     "len(atom_id)=", atom_id.size(), ", ",  //
+                     "len(atom_x)=", atom_x.size(), ", ",    //
+                     "len(atom_y)=", atom_y.size(), ", ",    //
+                     "len(atom_z)=", atom_z.size()));
+  }
+
+  if (fix_arginine) {
+    arginine.emplace();
+  }
+  if (!arginine.has_value()) {
+    return;
+  }
+
+  for (std::size_t res_index = 0; res_index < layout.num_residues();
+       ++res_index) {
+    auto [atom_start, atom_end] = layout.atom_range(res_index);
+    std::size_t atom_count = atom_end - atom_start;
+    PyObject* resname = comp_id[atom_start];
+    if (arginine.has_value() && arginine->IsResidue(resname)) {
+      arginine->Fix(atom_id.subspan(atom_start, atom_count),
+                    atom_x.subspan(atom_start, atom_count),
+                    atom_y.subspan(atom_start, atom_count),
+                    atom_z.subspan(atom_start, atom_count));
+    }
+  }
+}
+
+std::vector<bool> SelectedPolymerResidueMask(
+    const MmcifLayout& layout,
+    const std::vector<absl::string_view>& atom_site_label_asym_ids,  //
+    const std::vector<absl::string_view>& atom_site_label_seq_ids,   //
+    const std::vector<absl::string_view>& atom_site_label_comp_ids,  //
+    const std::vector<absl::string_view>& poly_seq_asym_ids,         //
+    const std::vector<absl::string_view>& poly_seq_seq_ids,          //
+    const std::vector<absl::string_view>& poly_seq_mon_ids           //
+) {
+  absl::flat_hash_map<std::pair<absl::string_view, absl::string_view>,
+                      absl::string_view>
+      selected;
+  selected.reserve(layout.num_residues());
+  // layout.residues() is O(1) while layout.residue_starts() is O(num_res).
+  const std::vector<std::size_t>& residue_starts = layout.residue_starts();
+  for (int i = 0; i < layout.residues().size(); ++i) {
+    std::size_t res_start = residue_starts[i];
+    std::size_t res_end = layout.residues()[i];
+    if (res_start == res_end) {
+      continue;  // Skip empty residues (containing no atoms).
+    }
+
+    absl::string_view label_seq_id = atom_site_label_seq_ids[i];
+    if (label_seq_id == ".") {
+      continue;  // Skip non-polymers.
+    }
+
+    absl::string_view label_asym_id = atom_site_label_asym_ids[i];
+    absl::string_view label_comp_id = atom_site_label_comp_ids[i];
+    selected[std::make_pair(label_asym_id, label_seq_id)] = label_comp_id;
+  }
+
+  std::vector<bool> mask;
+  mask.reserve(poly_seq_mon_ids.size());
+  for (int i = 0; i < poly_seq_mon_ids.size(); ++i) {
+    absl::string_view poly_seq_asym_id = poly_seq_asym_ids[i];
+    absl::string_view poly_seq_seq_id = poly_seq_seq_ids[i];
+    absl::string_view poly_seq_mon_id = poly_seq_mon_ids[i];
+
+    auto it = selected.find(std::make_pair(poly_seq_asym_id, poly_seq_seq_id));
+    if (it != selected.end()) {
+      mask.push_back(it->second == poly_seq_mon_id);
+    } else {
+      mask.push_back(true);  // Missing residues are not heterogeneous.
+    }
+  }
+  return mask;
+}
+
+std::pair<std::vector<bool>, std::vector<bool>> SelectedLigandResidueMask(
+    const MmcifLayout& layout,                                           //
+    const std::vector<absl::string_view>& atom_site_label_asym_ids,      //
+    const std::vector<absl::string_view>& atom_site_label_seq_ids,       //
+    const std::vector<absl::string_view>& atom_site_auth_seq_ids,        //
+    const std::vector<absl::string_view>& atom_site_label_comp_ids,      //
+    const std::vector<absl::string_view>& atom_site_pdbx_pdb_ins_codes,  //
+    const std::vector<absl::string_view>& nonpoly_asym_ids,              //
+    const std::vector<absl::string_view>& nonpoly_auth_seq_ids,          //
+    const std::vector<absl::string_view>& nonpoly_pdb_ins_codes,         //
+    const std::vector<absl::string_view>& nonpoly_mon_ids,               //
+    const std::vector<absl::string_view>& branch_asym_ids,               //
+    const std::vector<absl::string_view>& branch_auth_seq_ids,           //
+    const std::vector<absl::string_view>& branch_pdb_ins_codes,          //
+    const std::vector<absl::string_view>& branch_mon_ids) {
+  absl::flat_hash_map<
+      std::tuple<absl::string_view, absl::string_view, absl::string_view>,
+      absl::string_view>
+      selected;
+  selected.reserve(layout.num_residues());
+  // layout.residues() is O(1) while layout.residue_starts() is O(num_res).
+  const std::vector<std::size_t>& residue_starts = layout.residue_starts();
+  for (int i = 0; i < layout.residues().size(); ++i) {
+    std::size_t res_start = residue_starts[i];
+    std::size_t res_end = layout.residues()[i];
+    if (res_start == res_end) {
+      continue;  // Skip empty residues (containing no atoms).
+    }
+
+    absl::string_view label_seq_id = atom_site_label_seq_ids[i];
+    if (label_seq_id != ".") {
+      continue;  // Skip polymers.
+    }
+
+    absl::string_view label_asym_id = atom_site_label_asym_ids[i];
+    absl::string_view auth_seq_id = atom_site_auth_seq_ids[i];
+    absl::string_view ins_code = atom_site_pdbx_pdb_ins_codes[i];
+    ins_code = ins_code == "?" ? "." : ins_code;  // Remap unknown to unset.
+    absl::string_view label_comp_id = atom_site_label_comp_ids[i];
+    selected[std::make_tuple(label_asym_id, auth_seq_id, ins_code)] =
+        label_comp_id;
+  }
+
+  std::vector<bool> nonpoly_mask;
+  nonpoly_mask.reserve(nonpoly_asym_ids.size());
+  for (int i = 0; i < nonpoly_asym_ids.size(); ++i) {
+    absl::string_view nonpoly_asym_id = nonpoly_asym_ids[i];
+    absl::string_view nonpoly_auth_seq_id = nonpoly_auth_seq_ids[i];
+    absl::string_view nonpoly_ins_code = nonpoly_pdb_ins_codes[i];
+    // Remap unknown to unset.
+    nonpoly_ins_code = nonpoly_ins_code == "?" ? "." : nonpoly_ins_code;
+    absl::string_view nonpoly_mon_id = nonpoly_mon_ids[i];
+
+    auto it = selected.find(std::make_tuple(
+        nonpoly_asym_id, nonpoly_auth_seq_id, nonpoly_ins_code));
+    if (it != selected.end()) {
+      nonpoly_mask.push_back(it->second == nonpoly_mon_id);
+    } else {
+      nonpoly_mask.push_back(true);  // Missing residues are not heterogeneous.
+    }
+  }
+
+  std::vector<bool> branch_mask;
+  branch_mask.reserve(branch_asym_ids.size());
+  for (int i = 0; i < branch_asym_ids.size(); ++i) {
+    absl::string_view branch_asym_id = branch_asym_ids[i];
+    absl::string_view branch_auth_seq_id = branch_auth_seq_ids[i];
+
+    // Insertion codes in _pdbx_branch_scheme are not required and can be
+    // missing. Default to unset ('.') in such case.
+    absl::string_view branch_ins_code;
+    if (i < branch_pdb_ins_codes.size()) {
+      branch_ins_code = branch_pdb_ins_codes[i];
+      // Remap unknown to unset.
+      branch_ins_code = branch_ins_code == "?" ? "." : branch_ins_code;
+    } else {
+      branch_ins_code = ".";
+    }
+
+    absl::string_view branch_mon_id = branch_mon_ids[i];
+
+    auto it = selected.find(
+        std::make_tuple(branch_asym_id, branch_auth_seq_id, branch_ins_code));
+    if (it != selected.end()) {
+      branch_mask.push_back(it->second == branch_mon_id);
+    } else {
+      branch_mask.push_back(true);  // Missing residues are not heterogeneous.
+    }
+  }
+
+  return std::make_pair(nonpoly_mask, branch_mask);
+}
+
+constexpr char kReadMmcifLayout[] = R"(
+Returns the layout of the cif_dict.
+
+Args:
+  mmcif: mmCIF to calculate the layout for.
+  model_id: If non-empty the layout of the given model is returned
+    otherwise the layout of all models are returned.
+Raises:
+  ValueError: if the mmCIF is malformed or the number of atoms in each
+    model are inconsistent.
+)";
+
+constexpr char kMmcifFilter[] = R"(
+Returns NumpyArray of selected rows in `_atom_site` and new layout.
+
+Args:
+  mmcif: mmCIF to filter.
+  include_nucleotides: Whether to include polymer entities of type:
+    "polypeptide(L)\", "polydeoxyribonucleotide", "polyribonucleotide".
+    Otherwise only "polypeptide(L)\". ("polypeptide(D)\" is never included.)
+  include_ligands: Whether to include non-polymer entities of type:
+    "non-polymer", "branched".
+  include_water: Whether to include entities of type water.
+  include_other: Whether to include other (non-standard) entity types
+    that are not covered by any of the above parameters.
+  model_id: If non-empty the model with given name is selected otherwise
+    all models are selected.
+
+Returns:
+  A tuple containing a numpy array with a shape (num_models, num_atoms)
+  with the atom_site indices selected and the new layout.
+
+Raises:
+  ValueError error if mmCIF dict does not have all required fields.
+)";
+
+constexpr char kMmcifFixResidues[] = R"(
+Fixes residue columns in-place.
+
+Args:
+  layout: layout from filter command.
+  comp_id: '_atom_site.label_comp_id' of first model.
+  group: '_atom_site.group_PDB' of first model.
+  atom_id: '_atom_site.label_atom_id' of first model.
+  type_symbol: '_atom_site.type_symbol' of first model.
+  atom_x: '_atom_site.Cartn_x' of first model.
+  atom_y: '_atom_site.Cartn_y' of first model.
+  atom_z: '_atom_site.Cartn_z' of first model.
+  fix_mse: Whether to convert MSE residues into MET residues.
+  fix_arg: Whether to ensure the atoms in ARG are in the correct order.
+  fix_unknown_dna: Whether to convert DNA residues from N to DN.
+  dna_mask: Which atoms are from DNA chains.
+
+Raises:
+  ValueError: If shapes are invalid.
+)";
+
+constexpr char kSelectedPolymerResidueMask[] = R"(
+Returns a _pdbx_poly_seq_scheme mask for selected hetero residues.
+
+Should be called after filtering the layout using mmcif_utils.filter.
+
+Args:
+  layout: Layout defining the _atom_site residue selection.
+  atom_site_label_asym_ids: Internal (label) chain ID, per selected residue.
+  atom_site_label_seq_ids: Internal (label) residue ID, per selected residue.
+  atom_site_label_comp_ids: Residue name, per selected residue.
+  poly_seq_asym_ids: Internal (label) chain ID, per residue.
+  poly_seq_seq_ids: Internal (label) residue ID, per residue.
+  poly_seq_mon_ids: Residue name, per residue.
+
+Returns:
+  A mask for the _pdbx_poly_seq_scheme table. If residues are selected
+  using this mask, they will have consistent heterogeneous residue
+  selection with the _atom_site table.
+)";
+
+constexpr char kSelectedLigandResidueMask[] = R"(
+Returns masks for selected ligand hetero residues.
+
+Should be called after filtering the layout using mmcif_utils.filter.
+
+Args:
+  layout: Layout defining the _atom_site residue selection.
+  atom_site_label_asym_ids: Internal (label) chain ID, per selected residue.
+  atom_site_label_seq_ids: Internal (author) residue ID, per selected residue.
+  atom_site_auth_seq_ids: External (author) residue ID, per selected residue.
+  atom_site_label_comp_ids: Residue name, per selected residue.
+  atom_site_pdbx_pdb_ins_codes: Insertion code, per selected residue.
+  nonpoly_asym_ids: Internal (label) chain ID, per residue from
+   _pdbx_nonpoly_scheme.
+  nonpoly_auth_seq_ids: External (author) residue ID, per residue from
+   _pdbx_nonpoly_scheme.
+  nonpoly_pdb_ins_codes: Residue name, per residue from
+   _pdbx_nonpoly_scheme.
+  nonpoly_mon_ids: Insertion code, per residue from _pdbx_nonpoly_scheme.
+  branch_asym_ids: Internal (label) chain ID, per residue from
+   _pdbx_branch_scheme.
+  branch_auth_seq_ids: External (author) residue ID, per residue from
+   _pdbx_branch_scheme.
+  branch_pdb_ins_codes: Residue name, per residue from _pdbx_branch_scheme.
+  branch_mon_ids: Insertion code, per residue from _pdbx_branch_scheme.
+
+Returns:
+  A tuple with masks for _pdbx_nonpoly_scheme and _pdbx_branch_scheme. If
+  residues are selected using these masks, they will have consistent
+  heterogeneous residue selection with the _atom_site table.
+)";
+
+}  // namespace
+
+void RegisterModuleMmcifUtils(pybind11::module m) {
+  m.def("read_layout", ReadMmcifLayout,
+        py::arg("mmcif"),              //
+        py::arg("model_id") = "",      //
+        py::doc(kReadMmcifLayout + 1)  //
+  );
+
+  m.def("filter", MmcifFilter,               //
+        py::arg("mmcif"),                    //
+        py::arg("include_nucleotides"),      //
+        py::arg("include_ligands") = false,  //
+        py::arg("include_water") = false,    //
+        py::arg("include_other") = false,    //
+        py::arg("model_id") = "",            //
+        py::doc(kMmcifFilter + 1)            //
+  );
+
+  m.def("fix_residues", MmcifFixResidues,
+        py::arg("layout"),              //
+        py::arg("comp_id"),             //
+        py::arg("atom_id"),             //
+        py::arg("atom_x"),              //
+        py::arg("atom_y"),              //
+        py::arg("atom_z"),              //
+        py::arg("fix_arg") = false,     //
+        py::doc(kMmcifFixResidues + 1)  //
+  );
+
+  m.def("selected_polymer_residue_mask", SelectedPolymerResidueMask,
+        py::arg("layout"),                         //
+        py::arg("atom_site_label_asym_ids"),       //
+        py::arg("atom_site_label_seq_ids"),        //
+        py::arg("atom_site_label_comp_ids"),       //
+        py::arg("poly_seq_asym_ids"),              //
+        py::arg("poly_seq_seq_ids"),               //
+        py::arg("poly_seq_mon_ids"),               //
+        py::call_guard<py::gil_scoped_release>(),  //
+        py::doc(kSelectedPolymerResidueMask + 1)   //
+  );
+
+  m.def("selected_ligand_residue_mask", SelectedLigandResidueMask,
+        py::arg("layout"),                         //
+        py::arg("atom_site_label_asym_ids"),       //
+        py::arg("atom_site_label_seq_ids"),        //
+        py::arg("atom_site_auth_seq_ids"),         //
+        py::arg("atom_site_label_comp_ids"),       //
+        py::arg("atom_site_pdbx_pdb_ins_codes"),   //
+        py::arg("nonpoly_asym_ids"),               //
+        py::arg("nonpoly_auth_seq_ids"),           //
+        py::arg("nonpoly_pdb_ins_codes"),          //
+        py::arg("nonpoly_mon_ids"),                //
+        py::arg("branch_asym_ids"),                //
+        py::arg("branch_auth_seq_ids"),            //
+        py::arg("branch_pdb_ins_codes"),           //
+        py::arg("branch_mon_ids"),                 //
+        py::call_guard<py::gil_scoped_release>(),  //
+        py::doc(kSelectedLigandResidueMask + 1)    //
+  );
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/mmcif_utils_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_UTILS_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_UTILS_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleMmcifUtils(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_MMCIF_UTILS_PYBIND_H_

flax_model/alphafold3/structure/cpp/string_array.pyi

@@ -0,0 +1,42 @@
+
+
+from collections.abc import Sequence
+from typing import Any, overload
+
+import numpy as np
+
+
+def format_float_array(
+    values: Sequence[float], num_decimal_places: int
+) -> list[str]: ...
+
+
+def isin(
+    array: np.ndarray[object],
+    test_elements: set[str | bytes],
+    *,
+    invert: bool = ...,
+) -> np.ndarray[bool]: ...
+
+
+@overload
+def remap(
+    array: np.ndarray[object],
+    mapping: dict[str, str],
+    default_value: str,
+    inplace: bool = ...,
+) -> np.ndarray[object]: ...
+
+
+@overload
+def remap(
+    array: np.ndarray[object],
+    mapping: dict[str, str],
+    inplace: bool = ...,
+) -> np.ndarray[object]: ...
+
+
+def remap_multiple(
+    arrays: Sequence[np.ndarray[object]],
+    mapping: dict[tuple[Any], int],
+) -> np.ndarray[int]: ...

flax_model/alphafold3/structure/cpp/string_array_pybind.cc

@@ -0,0 +1,329 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// if received directly from Google. Use is subject to terms of use available at
+// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+#include <Python.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <iterator>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "numpy/arrayobject.h"
+#include "numpy/ndarrayobject.h"
+#include "numpy/ndarraytypes.h"
+#include "numpy/npy_common.h"
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "pybind11/cast.h"
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11_abseil/absl_casters.h"
+
+namespace {
+
+namespace py = pybind11;
+
+PyObject* RemapNumpyArrayObjects(PyObject* array, PyObject* mapping,
+                                 bool inplace, PyObject* default_value) {
+  import_array();
+  if (!PyArray_Check(array)) {
+    PyErr_SetString(PyExc_TypeError, "'array' must be a np.ndarray.");
+    return nullptr;
+  }
+  if (!PyDict_Check(mapping)) {
+    PyErr_SetString(PyExc_TypeError, "'mapping' must be a Python dict.");
+    return nullptr;
+  }
+
+  PyArrayObject* array_obj = reinterpret_cast<PyArrayObject*>(array);
+  if (PyArray_TYPE(array_obj) != NPY_OBJECT) {
+    PyErr_SetString(PyExc_TypeError, "`array` must be an array of objects.");
+    return nullptr;
+  }
+
+  if (inplace) {
+    // We are returning original array so we need to increase the ref count.
+    Py_INCREF(array);
+  } else {
+    // We are returning a fresh copy.
+    array = PyArray_NewCopy(array_obj, NPY_CORDER);
+    if (array == nullptr) {
+      PyErr_SetString(PyExc_MemoryError, "Out of memory!");
+      return nullptr;
+    }
+    array_obj = reinterpret_cast<PyArrayObject*>(array);
+  }
+
+  if (PyArray_SIZE(array_obj) == 0) {
+    return array;
+  }
+
+  if (default_value == nullptr && PyDict_Size(mapping) == 0) {
+    return array;
+  }
+
+  NpyIter* iter = NpyIter_New(
+      array_obj, NPY_ITER_READWRITE | NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK,
+      NPY_KEEPORDER, NPY_NO_CASTING, nullptr);
+  if (iter == nullptr) {
+    PyErr_SetString(PyExc_MemoryError, "Out of memory!");
+    Py_XDECREF(array);
+    return nullptr;
+  }
+
+  NpyIter_IterNextFunc* iter_next = NpyIter_GetIterNext(iter, nullptr);
+  if (iter_next == nullptr) {
+    NpyIter_Deallocate(iter);
+    Py_XDECREF(array);
+    PyErr_SetString(PyExc_MemoryError, "Out of memory!");
+    return nullptr;
+  }
+
+  // Iterating arrays taken from:
+  // https://numpy.org/doc/stable/reference/c-api/iterator.html
+  char** data_pointer = NpyIter_GetDataPtrArray(iter);
+  npy_intp* stride_pointer = NpyIter_GetInnerStrideArray(iter);
+  npy_intp* inner_size_pointer = NpyIter_GetInnerLoopSizePtr(iter);
+  do {
+    char* data = *data_pointer;
+    npy_intp stride = *stride_pointer;
+    npy_intp count = *inner_size_pointer;
+    for (size_t i = 0; i < count; ++i) {
+      PyObject* entry;
+      std::memcpy(&entry, data, sizeof(PyObject*));
+      PyObject* result = PyDict_GetItem(mapping, entry);
+      if (result != nullptr) {
+        // Replace entry.
+        Py_INCREF(result);
+        Py_XDECREF(entry);
+        std::memcpy(data, &result, sizeof(PyObject*));
+      } else if (default_value != nullptr) {
+        // Replace entry with a default value.
+        Py_INCREF(default_value);
+        Py_XDECREF(entry);
+        std::memcpy(data, &default_value, sizeof(PyObject*));
+      }
+      data += stride;
+    }
+  } while (iter_next(iter));
+
+  NpyIter_Deallocate(iter);
+  return array;
+}
+
+// Convert 1D Numpy float array to a list of strings where each string has fixed
+// number of decimal points. This is faster than Python list comprehension.
+std::vector<std::string> FormatFloatArray(absl::Span<const float> values,
+                                          int num_decimal_places) {
+  std::vector<std::string> output;
+  output.reserve(values.size());
+
+  absl::c_transform(values, std::back_inserter(output),
+                    [num_decimal_places](float value) {
+                      return absl::StrFormat("%.*f", num_decimal_places, value);
+                    });
+  return output;
+}
+
+py::array_t<bool> IsIn(
+    const py::array_t<PyObject*, py::array::c_style>& array,
+    const absl::flat_hash_set<absl::string_view>& test_elements, bool invert) {
+  const size_t num_elements = array.size();
+  py::array_t<bool> output(num_elements);
+  std::fill(output.mutable_data(), output.mutable_data() + output.size(),
+            invert);
+
+  // Shortcut: The output will be trivially always false if test_elements empty.
+  if (test_elements.empty()) {
+    return output;
+  }
+
+  for (size_t i = 0; i < num_elements; ++i) {
+    // Compare the string values instead of comparing just object pointers.
+    py::handle handle = array.data()[i];
+    if (!PyUnicode_Check(handle.ptr()) && !PyBytes_Check(handle.ptr())) {
+      continue;
+    }
+    if (test_elements.contains(py::cast<absl::string_view>(handle))) {
+      output.mutable_data()[i] = !invert;
+    }
+  }
+  if (array.ndim() > 1) {
+    auto shape =
+        std::vector<ptrdiff_t>(array.shape(), array.shape() + array.ndim());
+    return output.reshape(shape);
+  }
+  return output;
+}
+
+py::array RemapMultipleArrays(
+    const std::vector<py::array_t<PyObject*, py::array::c_style>>& arrays,
+    const py::dict& mapping) {
+  size_t array_size = arrays[0].size();
+  for (const auto& array : arrays) {
+    if (array.size() != array_size) {
+      throw py::value_error("All arrays must have the same length.");
+    }
+  }
+
+  // Create a result buffer.
+  auto result = py::array_t<int64_t>(array_size);
+  absl::Span<int64_t> result_buffer(result.mutable_data(), array_size);
+  PyObject* entry = PyTuple_New(arrays.size());
+  if (entry == nullptr) {
+    throw py::error_already_set();
+  }
+  std::vector<absl::Span<PyObject* const>> array_spans;
+  array_spans.reserve(arrays.size());
+  for (const auto& array : arrays) {
+    array_spans.emplace_back(array.data(), array.size());
+  }
+
+  // Iterate over arrays and look up elements in the `py_dict`.
+  bool fail = false;
+  for (size_t i = 0; i < array_size; ++i) {
+    for (size_t j = 0; j < array_spans.size(); ++j) {
+      PyTuple_SET_ITEM(entry, j, array_spans[j][i]);
+    }
+    PyObject* result = PyDict_GetItem(mapping.ptr(), entry);
+    if (result != nullptr) {
+      int64_t result_value = PyLong_AsLongLong(result);
+      if (result_value == -1 && PyErr_Occurred()) {
+        fail = true;
+        break;
+      }
+      if (result_value > std::numeric_limits<int64_t>::max() ||
+          result_value < std::numeric_limits<int64_t>::lowest()) {
+        PyErr_SetString(PyExc_OverflowError, "Result value too large.");
+        fail = true;
+        break;
+      }
+      result_buffer[i] = result_value;
+    } else {
+      PyErr_Format(PyExc_KeyError, "%R", entry);
+      fail = true;
+      break;
+    }
+  }
+
+  for (size_t j = 0; j < array_spans.size(); ++j) {
+    PyTuple_SET_ITEM(entry, j, nullptr);
+  }
+  Py_XDECREF(entry);
+  if (fail) {
+    throw py::error_already_set();
+  }
+  return result;
+}
+
+constexpr char kRemapNumpyArrayObjects[] = R"(
+Replace objects in NumPy array of objects using mapping.
+
+Args:
+  array: NumPy array with dtype=object.
+  mapping: Dict mapping old values to new values.
+  inplace: Bool (default False) whether to replace values inplace or to
+    create a new array.
+  default_value: If given, what value to map to if the mapping is missing
+    for that particular item. If not given, such items are left unchanged.
+
+Returns
+  NumPy array of dtype object with values replaced according to mapping.
+  If inplace is True the original array is modified inplace otherwise a
+  new array is returned.
+)";
+
+constexpr char kFormatFloatArrayDoc[] = R"(
+Converts float -> string array with given number of decimal places.
+)";
+
+constexpr char kIsInDoc[] = R"(
+Computes whether each element is in test_elements.
+
+Same use as np.isin, but much faster. If len(array) = n, len(test_elements) = m:
+* This function has complexity O(n).
+* np.isin with arrays of objects has complexity O(m*log(m) + n * log(m)).
+
+Args:
+  array: Input NumPy array with dtype=object.
+  test_elements: The values against which to test each value of array.
+  invert: If True, the values in the returned array are inverted, as if
+    calculating `element not in test_elements`. Default is False.
+    `isin(a, b, invert=True)` is equivalent to but faster than `~isin(a, b)`.
+
+Returns
+  A boolean array of the same shape as the input array. Each value `val` is:
+  * `val in test_elements` if `invert=False`,
+  * `val not in test_elements` if `invert=True`.
+)";
+
+constexpr char kRemapMultipleDoc[] = R"(
+Maps keys from multiple aligned arrays to a single array.
+
+Args:
+  arrays: Numpy arrays of the same length. The tuple of aligned entries is used
+    as key for the mapping.
+  mapping: Dict mapping from tuples to integer values.
+
+Returns
+  NumPy array of dtype `int` with values looked up in mapping according to the
+  tuple of aligned array entries as keys.
+)";
+
+}  // namespace
+
+namespace alphafold3 {
+
+void RegisterModuleStringArray(pybind11::module m) {
+  m.def(
+      "remap",
+      [](py::object array, py::object mapping, bool inplace,
+         py::object default_value) -> py::object {
+        PyObject* result = RemapNumpyArrayObjects(array.ptr(), mapping.ptr(),
+                                                  inplace, default_value.ptr());
+        if (result == nullptr) {
+          throw py::error_already_set();
+        }
+        return py::reinterpret_steal<py::object>(result);
+      },
+      py::return_value_policy::take_ownership, py::arg("array"),
+      py::arg("mapping"), py::arg("inplace") = false, py::arg("default_value"),
+      py::doc(kRemapNumpyArrayObjects + 1));
+  m.def(
+      "remap",
+      [](py::object array, py::object mapping, bool inplace) -> py::object {
+        PyObject* result = RemapNumpyArrayObjects(array.ptr(), mapping.ptr(),
+                                                  inplace, nullptr);
+        if (result == nullptr) {
+          throw py::error_already_set();
+        }
+        return py::reinterpret_steal<py::object>(result);
+      },
+      py::return_value_policy::take_ownership, py::arg("array"),
+      py::arg("mapping"), py::arg("inplace") = false,
+      py::doc(kRemapNumpyArrayObjects + 1));
+  m.def("format_float_array", &FormatFloatArray, py::arg("values"),
+        py::arg("num_decimal_places"), py::doc(kFormatFloatArrayDoc + 1),
+        py::call_guard<py::gil_scoped_release>());
+  m.def("isin", &IsIn, py::arg("array"), py::arg("test_elements"),
+        py::kw_only(), py::arg("invert") = false, py::doc(kIsInDoc + 1));
+  m.def("remap_multiple", &RemapMultipleArrays, py::arg("arrays"),
+        py::arg("mapping"), py::doc(kRemapMultipleDoc + 1));
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/string_array_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * if received directly from Google. Use is subject to terms of use available at
+ * https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_STRING_ARRAY_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_STRING_ARRAY_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleStringArray(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_STRING_ARRAY_PYBIND_H_

flax_model/alphafold3/structure/mmcif.py

@@ -0,0 +1,321 @@
+
+
+"""Low level mmCIF parsing operations and wrappers for nicer C++/Py errors.
+
+Note that the cif_dict.CifDict class has many useful methods to help with data
+extraction which are not shown in this file. You can find them in cif_dict.clif
+together with docstrings. The cif_dict.CifDict class behaves like an immutable
+Python dictionary (some methods are not implemented though).
+"""
+from collections.abc import Callable, Mapping, Sequence
+import functools
+import itertools
+import re
+from typing import ParamSpec, TypeAlias, TypeVar
+
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.cpp import cif_dict
+from flax_model.alphafold3.cpp import mmcif_atom_site
+from flax_model.alphafold3.cpp import mmcif_struct_conn
+from flax_model.alphafold3.cpp import string_array
+import numpy as np
+
+Mmcif = cif_dict.CifDict
+
+
+_P = ParamSpec('_P')
+_T = TypeVar('_T')
+_WappedFn: TypeAlias = Callable[_P, _T]
+
+
+@functools.lru_cache(maxsize=256)
+def int_id_to_str_id(num: int) -> str:
+  """Encodes a number as a string, using reverse spreadsheet style naming.
+
+  Args:
+    num: A positive integer.
+
+  Returns:
+    A string that encodes the positive integer using reverse spreadsheet style,
+    naming e.g. 1 = A, 2 = B, ..., 27 = AA, 28 = BA, 29 = CA, ... This is the
+    usual way to encode chain IDs in mmCIF files.
+  """
+  if num <= 0:
+    raise ValueError(f'Only positive integers allowed, got {num}.')
+
+  num = num - 1  # 1-based indexing.
+  output = []
+  while num >= 0:
+    output.append(chr(num % 26 + ord('A')))
+    num = num // 26 - 1
+  return ''.join(output)
+
+
+@functools.lru_cache(maxsize=256)
+def str_id_to_int_id(str_id: str) -> int:
+  """Encodes an mmCIF-style string chain ID as an integer.
+
+  The integer IDs are one based so this function is the inverse of
+  int_id_to_str_id.
+
+  Args:
+    str_id: A string chain ID consisting only of upper case letters A-Z.
+
+  Returns:
+    An integer that can be used to order mmCIF chain IDs in the standard
+    (reverse spreadsheet style) ordering.
+  """
+  if not re.match('^[A-Z]+$', str_id):
+    raise ValueError(f'String ID must be upper case letters, got {str_id}.')
+
+  offset = ord('A') - 1
+  output = 0
+  for i, c in enumerate(str_id):
+    output += (ord(c) - offset) * int(26**i)
+  return output
+
+
+def from_string(mmcif_string: str | bytes) -> Mmcif:
+  return cif_dict.from_string(mmcif_string)
+
+
+def parse_multi_data_cif(cif_string: str) -> dict[str, Mmcif]:
+  """Parses a CIF string with multiple data records.
+
+  For instance, the CIF string:
+
+  ```
+  data_001
+  _foo bar
+  #
+  data_002
+  _foo baz
+  ```
+
+  is parsed as:
+
+  ```
+  {'001': Mmcif({'_foo': ['bar']}), '002': Mmcif({'_foo': ['baz']})}
+  ```
+
+  Args:
+    cif_string: The multi-data CIF string to be parsed.
+
+  Returns:
+    A dictionary mapping record names to Mmcif objects with data.
+  """
+  return cif_dict.parse_multi_data_cif(cif_string)
+
+
+# def tokenize(mmcif_string: str) -> list[str]:
+#   return cif_dict.tokenize(mmcif_string)
+
+
+# def split_line(line: str) -> list[str]:
+#   return cif_dict.split_line(line)
+
+
+class BondParsingError(Exception):
+  """Exception raised by errors when getting bond atom indices."""
+
+
+def get_bond_atom_indices(
+    mmcif: Mmcif,
+    model_id: str = '1',
+) -> tuple[Sequence[int], Sequence[int]]:
+  """Extracts the indices of the atoms that participate in bonds.
+
+  Args:
+    mmcif: The mmCIF object to process.
+    model_id: The ID of the model that the returned atoms will belong to. This
+      should be a value in the mmCIF's _atom_site.pdbx_PDB_model_num column.
+
+  Returns:
+    Two lists of atom indices, `from_atoms` and `to_atoms`, each one having
+    length num_bonds (as defined by _struct_conn, the bonds table). The bond
+    i, defined by the i'th row in _struct_conn, is a bond from atom at index
+    from_atoms[i], to the atom at index to_atoms[i]. The indices are simple
+    0-based indexes into the columns of the _atom_site table in the input
+    mmCIF, and do not necessarily correspond to the values in _atom_site.id,
+    or any other column.
+
+  Raises:
+    BondParsingError: If any of the required tables or columns are not present
+    in
+      the mmCIF, or if the _struct_conn table refers to atoms that cannot
+      be found in the _atom_site table.
+  """
+  try:
+    return mmcif_struct_conn.get_bond_atom_indices(mmcif, model_id)
+  except ValueError as e:
+    raise BondParsingError(str(e)) from e
+
+
+def get_or_infer_type_symbol(
+    mmcif: Mmcif, ccd: chemical_components.Ccd | None = None
+) -> Sequence[str]:
+  """Returns the type symbol (element) for all of the atoms.
+
+  Args:
+    mmcif: A parsed mmCIF file in the Mmcif format.
+    ccd: The chemical component dictionary. If not provided, defaults to the
+      cached CCD.
+
+  If present, returns the _atom_site.type_symbol. If not, infers it using
+  _atom_site.label_comp_id (residue name), _atom_site.label_atom_id (atom name)
+  and the CCD.
+  """
+  ccd = ccd or chemical_components.Ccd()
+  type_symbol_fn = lambda res_name, atom_name: chemical_components.type_symbol(
+      ccd, res_name, atom_name
+  )
+  return mmcif_atom_site.get_or_infer_type_symbol(mmcif, type_symbol_fn)
+
+
+def get_chain_type_by_entity_id(mmcif: Mmcif) -> Mapping[str, str]:
+  """Returns mapping from entity ID to its type or polymer type if available.
+
+  If the entity is in the _entity_poly table, returns its polymer chain type.
+  If not, returns the type as specified in the _entity table.
+
+  Args:
+    mmcif: CifDict holding the mmCIF.
+  """
+  poly_entity_id = mmcif.get('_entity_poly.entity_id', [])
+  poly_type = mmcif.get('_entity_poly.type', [])
+  poly_type_by_entity_id = dict(zip(poly_entity_id, poly_type, strict=True))
+
+  chain_type_by_entity_id = {}
+  for entity_id, entity_type in zip(
+      mmcif.get('_entity.id', []), mmcif.get('_entity.type', []), strict=True
+  ):
+    chain_type = poly_type_by_entity_id.get(entity_id) or entity_type
+    chain_type_by_entity_id[entity_id] = chain_type
+
+  return chain_type_by_entity_id
+
+
+def get_internal_to_author_chain_id_map(mmcif: Mmcif) -> Mapping[str, str]:
+  """Returns a mapping from internal chain ID to the author chain ID.
+
+  Note that this is not a bijection. One author chain ID can map to multiple
+  internal chain IDs. For example, a protein chain and a ligand bound to it will
+  share the same author chain ID, but they will each have a unique internal
+  chain ID).
+
+  Args:
+    mmcif: CifDict holding the mmCIF.
+  """
+  return mmcif_atom_site.get_internal_to_author_chain_id_map(mmcif)
+
+
+def get_experimental_method(mmcif: Mmcif) -> str | None:
+  field = '_exptl.method'
+  return ','.join(mmcif[field]).lower() if field in mmcif else None
+
+
+def get_release_date(mmcif: Mmcif) -> str | None:
+  """Returns the oldest revision date."""
+  if '_pdbx_audit_revision_history.revision_date' not in mmcif:
+    return None
+
+  # Release dates are ISO-8601, hence sort well.
+  return min(mmcif['_pdbx_audit_revision_history.revision_date'])
+
+
+def get_resolution(mmcif: Mmcif) -> float | None:
+  """Returns the resolution of the structure.
+
+  More than one resolution can be reported in an mmCIF. This function returns
+  the first one (in the order _refine.ls_d_res_high,
+  _em_3d_reconstruction.resolution, _reflns.d_resolution_high) that appears
+  in the mmCIF as is parseable as a float.
+
+  Args:
+    mmcif: An `Mmcif` object.
+
+  Returns:
+    The resolution as reported in the mmCIF.
+  """
+  for res_key in ('_refine.ls_d_res_high',
+                  '_em_3d_reconstruction.resolution',
+                  '_reflns.d_resolution_high'):
+    if res_key in mmcif:
+      try:
+        raw_resolution = mmcif[res_key][0]
+        return float(raw_resolution)
+      except ValueError:
+        continue
+  return None
+
+
+def parse_oper_expr(oper_expression: str) -> list[tuple[str, ...]]:
+  """Determines which transforms to apply based on an MMCIF oper_expression str.
+
+  Args:
+    oper_expression: the field oper_expression from MMCIF format data.
+      Transform ids may be either numbers or single letters. Hyphens are used to
+      denote a numeric range of transforms to apply, and commas are used to
+      delimit a sequence of transforms. Where two sets of parentheses are
+      adjacent without a comma, the two sets of transforms should be combined as
+      a cartesian product, i.e. all possible pairs.
+      example 1,2,3 -> generate 3 copies of each chain by applying 1, 2 or 3.
+      example (1-3) -> generate 3 copies of each chain by applying 1, 2 or 3.
+      example (1-3)(4-6) -> generate 9 copies of each chain by applying one of
+        [(1,4), (1,5), (1,6),
+         (2,4), (2,5), (2,6),
+         (3,4), (3,5), (3,6)]
+      example (P) -> apply transform with id P.
+
+  Raises:
+    ValueError: Failure to parse oper_expression.
+
+  Returns:
+    A list with one element for each chain copy that should be generated.
+    Each element is a list of transform ids to apply.
+  """
+  # Expand ranges, e.g. 1-4 -> 1,2,3,4.
+  def range_expander(match):
+    return ','.join(
+        [str(i) for i in range(int(match.group(1)),
+                               int(match.group(2)) + 1)])
+
+  ranges_expanded = re.sub(r'\b(\d+)-(\d+)', range_expander, oper_expression)
+
+  if re.fullmatch(r'(\w+,)*\w+', ranges_expanded):
+    # No brackets, just a single range, e.g. "1,2,3".
+    return [(t,) for t in ranges_expanded.split(',')]
+  elif re.fullmatch(r'\((\w+,)*\w+\)', ranges_expanded):
+    # Single range in brackets, e.g. "(1,2,3)".
+    return [(t,) for t in ranges_expanded[1:-1].split(',')]
+  elif re.fullmatch(r'\((\w+,)*\w+\)\((\w+,)*\w+\)', ranges_expanded):
+    # Cartesian product of two ranges, e.g. "(1,2,3)(4,5)".
+    part1, part2 = ranges_expanded[1:-1].split(')(')
+    return list(itertools.product(part1.split(','), part2.split(',')))
+  else:
+    raise ValueError(f'Unsupported oper_expression format: {oper_expression}')
+
+
+def format_float_array(
+    values: np.ndarray, num_decimal_places: int) -> Sequence[str]:
+  """Converts 1D array to a list of strings with the given number of decimals.
+
+  This function is faster than converting via Python list comprehension, e.g.:
+  atoms_x = ['%.3f' % x for x in atoms_x]
+
+  Args:
+    values: A numpy array with values to convert. This array is casted to
+      float32 before doing the conversion.
+    num_decimal_places: The number of decimal points to keep, including trailing
+      zeros. E.g. for 1.07 and num_decimal_places=1: 1.1,
+      num_decimal_places=2: 1.07, num_decimal_places=3: 1.070.
+
+  Returns:
+    A list of formatted strings.
+  """
+  if values.ndim != 1:
+    raise ValueError(f'The given array must be 1D, got {values.ndim}D')
+
+  return string_array.format_float_array(
+      values=values.astype(np.float32), num_decimal_places=num_decimal_places
+  )

flax_model/alphafold3/structure/parsing.py

@@ -0,0 +1,1801 @@
+
+
+"""Module for parsing various data sources and producing Structures."""
+
+from collections.abc import Collection, Mapping, MutableMapping, Sequence
+import dataclasses
+import datetime
+import enum
+import functools
+import itertools
+from typing import TypeAlias
+
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import mmcif_utils
+from flax_model.alphafold3.cpp import string_array
+from flax_model.alphafold3.structure import bioassemblies
+from flax_model.alphafold3.structure import bonds
+from flax_model.alphafold3.structure import chemical_components as struc_chem_comps
+from flax_model.alphafold3.structure import mmcif
+from flax_model.alphafold3.structure import structure
+from flax_model.alphafold3.structure import structure_tables
+import numpy as np
+
+
+ChainIndex: TypeAlias = int
+ResIndex: TypeAlias = int
+AtomName: TypeAlias = str
+BondAtomId: TypeAlias = tuple[ChainIndex, ResIndex, AtomName]
+
+_INSERTION_CODE_REMAP: Mapping[str, str] = {'.': '?'}
+
+
+class NoAtomsError(Exception):
+  """Raise when the mmCIF does not have any atoms."""
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class BondIndices:
+  from_indices: list[int]
+  dest_indices: list[int]
+
+
+@enum.unique
+class ModelID(enum.Enum):
+  """Values for specifying model IDs when parsing."""
+
+  FIRST = 1  # The first model in the file.
+  ALL = 2  # All models in the file.
+
+
+@enum.unique
+class SequenceFormat(enum.Enum):
+  """The possible formats for an input sequence."""
+
+  FASTA = 'fasta'  # One-letter code used in FASTA.
+  CCD_CODES = 'ccd_codes'  # Multiple-letter chemical components dictionary ids.
+  LIGAND_SMILES = 'ligand_smiles'  # SMILES string defining a molecule.
+
+
+def _create_bond_lookup(
+    bonded_atom_pairs: Sequence[tuple[BondAtomId, BondAtomId]],
+) -> Mapping[tuple[ChainIndex, ResIndex], Mapping[AtomName, BondIndices]]:
+  """Creates maps to help find bonds during a loop over residues."""
+  bond_lookup = {}
+  for bond_i, (from_atom_id, dest_atom_id) in enumerate(bonded_atom_pairs):
+    from_chain_i, from_res_i, from_atom_name = from_atom_id
+    dest_chain_i, dest_res_i, dest_atom_name = dest_atom_id
+    bonds_by_from_atom_name = bond_lookup.setdefault(
+        (from_chain_i, from_res_i), {}
+    )
+    bonds_by_dest_atom_name = bond_lookup.setdefault(
+        (dest_chain_i, dest_res_i), {}
+    )
+    bonds_by_from_atom_name.setdefault(
+        from_atom_name, BondIndices(from_indices=[], dest_indices=[])
+    ).from_indices.append(bond_i)
+    bonds_by_dest_atom_name.setdefault(
+        dest_atom_name, BondIndices(from_indices=[], dest_indices=[])
+    ).dest_indices.append(bond_i)
+  return bond_lookup
+
+
+def _get_atom_element(
+    ccd: chemical_components.Ccd, res_name: str, atom_name: str
+) -> str:
+  type_symbol = chemical_components.type_symbol(
+      ccd, res_name=res_name, atom_name=atom_name
+  )
+  return type_symbol or '?'
+
+
+def _get_representative_atom(
+    ccd: chemical_components.Ccd,
+    res_name: str,
+    chain_type: str,
+    sequence_format: SequenceFormat,
+) -> tuple[str, str]:
+  match sequence_format:
+    case SequenceFormat.CCD_CODES:
+      atom_name = _get_first_non_leaving_atom(ccd=ccd, res_name=res_name)
+      atom_element = _get_atom_element(
+          ccd=ccd, res_name=res_name, atom_name=atom_name
+      )
+      return atom_name, atom_element
+    case SequenceFormat.LIGAND_SMILES:
+      return '', '?'
+    case SequenceFormat.FASTA:
+      if chain_type in mmcif_names.PEPTIDE_CHAIN_TYPES:
+        return 'CA', 'C'
+      if chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES:
+        return "C1'", 'C'
+      else:
+        raise ValueError(chain_type)
+    case _:
+      raise ValueError(sequence_format)
+
+
+@functools.lru_cache(maxsize=128)
+def _get_first_non_leaving_atom(
+    ccd: chemical_components.Ccd, res_name: str
+) -> str:
+  """Returns first definitely non-leaving atom if exists, as a stand-in."""
+  all_atoms = struc_chem_comps.get_all_atoms_in_entry(ccd, res_name=res_name)[
+      '_chem_comp_atom.atom_id'
+  ]
+  representative_atom = all_atoms[0]
+  if representative_atom == 'O1' and len(all_atoms) > 1:
+    representative_atom = all_atoms[1]
+  return representative_atom
+
+
+def _add_ligand_to_chem_comp(
+    chem_comp: MutableMapping[str, struc_chem_comps.ChemCompEntry],
+    ligand_id: str,
+    ligand_smiles: str,
+):
+  """Adds a ligand to chemical components. Raises ValueError on mismatch."""
+  new_entry = struc_chem_comps.ChemCompEntry(
+      type='non-polymer', pdbx_smiles=ligand_smiles
+  )
+
+  existing_entry = chem_comp.get(ligand_id)
+  if existing_entry is None:
+    chem_comp[ligand_id] = new_entry
+  elif existing_entry != new_entry:
+    raise ValueError(
+        f'Mismatching data for ligand {ligand_id}: '
+        f'{new_entry} != {existing_entry}'
+    )
+
+
+def _get_first_model_id(cif: mmcif.Mmcif) -> str:
+  """Returns cheaply the first model ID from the mmCIF."""
+  return cif.get_array(
+      '_atom_site.pdbx_PDB_model_num', dtype=object, gather=slice(1)
+  )[0]
+
+
+def _get_str_model_id(
+    cif: mmcif.Mmcif,
+    model_id: ModelID | int,
+) -> str:
+  """Converts a user-specified model_id argument into a string."""
+  match model_id:
+    case int():
+      str_model_id = str(model_id)
+    case enum.Enum():
+      # We compare the enum's value attribute since regular enum comparison
+      # breaks when adhoc importing.
+      match model_id.value:
+        case ModelID.FIRST.value:
+          try:
+            str_model_id = _get_first_model_id(cif)
+          except IndexError as e:
+            raise NoAtomsError(
+                'The mmCIF does not have any atoms or'
+                ' _atom_site.pdbx_PDB_model_num is missing.'
+            ) from e
+        case ModelID.ALL.value:
+          str_model_id = ''
+        case _:
+          raise ValueError(
+              f'Model ID {model_id} with value {model_id.value} not recognized.'
+          )
+    case _:
+      raise ValueError(
+          f'Model ID {model_id} with type {type(model_id)} not recognized.'
+      )
+  return str_model_id
+
+
+def _parse_bonds(
+    cif: mmcif.Mmcif,
+    atom_key: np.ndarray,
+    model_id: str,
+) -> bonds.Bonds:
+  """Returns the bonds table extracted from the mmCIF.
+
+  Args:
+    cif: The raw mmCIF to extract the bond information from.
+    atom_key: A numpy array defining atom key for each atom in _atom_site. Note
+      that the atom key must be computed before resolving alt-locs since this
+      function operates on the raw mmCIF!
+    model_id: The ID of the model to get bonds for.
+  """
+  if '_struct_conn.id' not in cif:
+    # This is the category key item for the _struct_conn table, therefore
+    # we use it to determine whether to parse bond info.
+    return bonds.Bonds.make_empty()
+  from_atom, dest_atom = mmcif.get_bond_atom_indices(cif, model_id)
+  from_atom = np.array(from_atom, dtype=np.int64)
+  dest_atom = np.array(dest_atom, dtype=np.int64)
+  num_bonds = from_atom.shape[0]
+  bond_key = np.arange(num_bonds, dtype=np.int64)
+  bond_type = cif.get_array('_struct_conn.conn_type_id', dtype=object)
+  if '_struct_conn.pdbx_role' in cif:  # This column isn't always present.
+    bond_role = cif.get_array('_struct_conn.pdbx_role', dtype=object)
+  else:
+    bond_role = np.full((num_bonds,), '?', dtype=object)
+
+  bonds_mask = np.ones((num_bonds,), dtype=bool)
+  # Symmetries other than 1_555 imply the atom is not part of the asymmetric
+  # unit, and therefore this is a bond that only exists in the expanded
+  # bioassembly.
+  # We do not currently support parsing these types of bonds.
+  if '_struct_conn.ptnr1_symmetry' in cif:
+    ptnr1_symmetry = cif.get_array('_struct_conn.ptnr1_symmetry', dtype=object)
+    np.logical_and(bonds_mask, ptnr1_symmetry == '1_555', out=bonds_mask)
+  if '_struct_conn.ptnr2_symmetry' in cif:
+    ptnr2_symmetry = cif.get_array('_struct_conn.ptnr2_symmetry', dtype=object)
+    np.logical_and(bonds_mask, ptnr2_symmetry == '1_555', out=bonds_mask)
+  # Remove bonds that involve atoms that are not part of the structure,
+  # e.g. waters if include_water=False. In a rare case this also removes invalid
+  # bonds that are indicated by a key that is set to _atom_site size.
+  np.logical_and(bonds_mask, np.isin(from_atom, atom_key), out=bonds_mask)
+  np.logical_and(bonds_mask, np.isin(dest_atom, atom_key), out=bonds_mask)
+  return bonds.Bonds(
+      key=bond_key[bonds_mask],
+      type=bond_type[bonds_mask],
+      role=bond_role[bonds_mask],
+      from_atom_key=from_atom[bonds_mask],
+      dest_atom_key=dest_atom[bonds_mask],
+  )
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class _MmcifHeader:
+  name: str
+  resolution: float | None
+  release_date: datetime.date | None
+  structure_method: str | None
+  bioassembly_data: bioassemblies.BioassemblyData | None
+  chemical_components_data: struc_chem_comps.ChemicalComponentsData | None
+
+
+def _get_mmcif_header(
+    cif: mmcif.Mmcif,
+    fix_mse: bool,
+    fix_unknown_dna: bool,
+) -> _MmcifHeader:
+  """Extract header fields from an mmCIF object."""
+  entry_id = cif.get('_entry.id')
+  name = entry_id[0] if entry_id else cif.get_data_name()
+  resolution = mmcif.get_resolution(cif)
+
+  release_date = mmcif.get_release_date(cif)
+  if release_date is not None:
+    release_date = datetime.date.fromisoformat(release_date)
+
+  experiments = cif.get('_exptl.method')
+  structure_method = ','.join(experiments) if experiments else None
+
+  try:
+    bioassembly_data = bioassemblies.BioassemblyData.from_mmcif(cif)
+  except bioassemblies.MissingBioassemblyDataError:
+    bioassembly_data = None
+
+  try:
+    chemical_components_data = (
+        struc_chem_comps.ChemicalComponentsData.from_mmcif(
+            cif, fix_mse=fix_mse, fix_unknown_dna=fix_unknown_dna
+        )
+    )
+  except struc_chem_comps.MissingChemicalComponentsDataError:
+    chemical_components_data = None
+
+  return _MmcifHeader(
+      name=name,
+      resolution=resolution,
+      release_date=release_date,
+      structure_method=structure_method,
+      bioassembly_data=bioassembly_data,
+      chemical_components_data=chemical_components_data,
+  )
+
+
+def from_parsed_mmcif(
+    mmcif_object: mmcif.Mmcif,
+    *,
+    name: str | None = None,
+    fix_mse_residues: bool = False,
+    fix_arginines: bool = False,
+    fix_unknown_dna: bool = False,
+    include_water: bool = False,
+    include_other: bool = False,
+    include_bonds: bool = False,
+    model_id: int | ModelID = ModelID.FIRST,
+) -> structure.Structure:
+  """Construct a Structure from a parsed mmCIF object.
+
+  This function is called by `from_mmcif` but can be useful when an mmCIF has
+  already been parsed e.g. to extract extra information from the header before
+  then converting to Structure for further manipulation.
+
+  Args:
+    mmcif_object: A parsed mmcif.Mmcif object.
+    name: Optional name for the structure. If not provided, the name will be
+      taken from the mmCIF data_ field.
+    fix_mse_residues: If True, selenium atom sites (SE) in selenomethionine
+      (MSE) residues will be changed to sulphur atom sites (SD). This is because
+      methionine (MET) residues are often replaced with MSE to aid X-Ray
+      crystallography. If False, the SE MSE atom sites won't be modified.
+    fix_arginines: If True, NH1 and NH2 in arginine will be swapped if needed so
+      that NH1 is always closer to CD than NH2. If False, no atom sites in
+      arginine will be touched. Note that HH11, HH12, HH21, HH22 are fixed too.
+    fix_unknown_dna: If True, residues with name N in DNA chains will have their
+      res_name replaced with DN. Atoms are not changed.
+    include_water: If True, water (HOH) molecules will be parsed. Water
+      molecules may be grouped into chains, where number of residues > 1. Water
+      molecules are usually grouped into chains but do not necessarily all share
+      the same chain ID.
+    include_other: If True, all other atoms that are not included by any of the
+      above parameters will be included. This covers e.g. "polypeptide(D)" and
+      "macrolide" entities, as well as all other non-standard types.
+    include_bonds: If True, bond information will be parsed from the mmCIF and
+      stored in the Structure.
+    model_id: Either the integer model ID to parse, or one of ModelID.FIRST to
+      parse the first model, or ModelID.ALL to parse all models.
+
+  Returns:
+    A Structure representation of the mmCIF object.
+  """
+  str_model_id = _get_str_model_id(cif=mmcif_object, model_id=model_id)
+  header = _get_mmcif_header(
+      mmcif_object, fix_mse=fix_mse_residues, fix_unknown_dna=fix_unknown_dna
+  )
+
+  chains, residues, atoms = get_tables(
+      cif=mmcif_object,
+      fix_mse_residues=fix_mse_residues,
+      fix_arginines=fix_arginines,
+      fix_unknown_dna=fix_unknown_dna,
+      include_water=include_water,
+      include_other=include_other,
+      model_id=str_model_id,
+  )
+
+  if include_bonds and atoms.size > 0:
+    # NB: parsing the atom table before the bonds table allows for a more
+    # informative error message when dealing with bad multi-model mmCIFs.
+    # Also always use a specific model ID, even when parsing all models.
+    if str_model_id == '':  # pylint: disable=g-explicit-bool-comparison
+      bonds_model_id = _get_first_model_id(mmcif_object)
+    else:
+      bonds_model_id = str_model_id
+
+    bonds_table = _parse_bonds(
+        mmcif_object,
+        atom_key=atoms.key,
+        model_id=bonds_model_id,
+    )
+  else:
+    bonds_table = bonds.Bonds.make_empty()
+
+  return structure.Structure(
+      name=name if name is not None else header.name,
+      resolution=header.resolution,
+      release_date=header.release_date,
+      structure_method=header.structure_method,
+      bioassembly_data=header.bioassembly_data,
+      chemical_components_data=header.chemical_components_data,
+      bonds=bonds_table,
+      chains=chains,
+      residues=residues,
+      atoms=atoms,
+  )
+
+
+def from_mmcif(
+    mmcif_string: str | bytes,
+    *,
+    name: str | None = None,
+    fix_mse_residues: bool = False,
+    fix_arginines: bool = False,
+    fix_unknown_dna: bool = False,
+    include_water: bool = False,
+    include_other: bool = False,
+    include_bonds: bool = False,
+    model_id: int | ModelID = ModelID.FIRST,
+) -> structure.Structure:
+  """Construct a Structure from a mmCIF string.
+
+  Args:
+    mmcif_string: The string contents of an mmCIF file.
+    name: Optional name for the structure. If not provided, the name will be
+      taken from the mmCIF data_ field.
+    fix_mse_residues: If True, selenium atom sites (SE) in selenomethionine
+      (MSE) residues will be changed to sulphur atom sites (SD). This is because
+      methionine (MET) residues are often replaced with MSE to aid X-Ray
+      crystallography. If False, the SE MSE atom sites won't be modified.
+    fix_arginines: If True, NH1 and NH2 in arginine will be swapped if needed so
+      that NH1 is always closer to CD than NH2. If False, no atom sites in
+      arginine will be touched. Note that HH11, HH12, HH21, HH22 are fixed too.
+    fix_unknown_dna: If True, residues with name N in DNA chains will have their
+      res_name replaced with DN. Atoms are not changed.
+    include_water: If True, water (HOH) molecules will be parsed. Water
+      molecules may be grouped into chains, where number of residues > 1. Water
+      molecules are usually grouped into chains but do not necessarily all share
+      the same chain ID.
+    include_other: If True, all other atoms that are not included by any of the
+      above parameters will be included. This covers e.g. "polypeptide(D)" and
+      "macrolide" entities, as well as all other non-standard types.
+    include_bonds: If True, bond information will be parsed from the mmCIF and
+      stored in the Structure.
+    model_id: Either the integer model ID to parse, or one of ModelID.FIRST to
+      parse the first model, or ModelID.ALL to parse all models.
+
+  Returns:
+    A Structure representation of the mmCIF string.
+  """
+  mmcif_object = mmcif.from_string(mmcif_string)
+
+  return from_parsed_mmcif(
+      mmcif_object,
+      name=name,
+      fix_mse_residues=fix_mse_residues,
+      fix_arginines=fix_arginines,
+      fix_unknown_dna=fix_unknown_dna,
+      include_water=include_water,
+      include_other=include_other,
+      include_bonds=include_bonds,
+      model_id=model_id,
+  )
+
+
+def from_res_arrays(atom_mask: np.ndarray, **kwargs) -> structure.Structure:
+  """Returns Structure created from from arrays with a residue dimension.
+
+  All unset fields are filled with defaults (e.g. 1.0 for occupancy) or
+  unset/unknown values (e.g. UNK for residue type, or '.' for atom element).
+
+  Args:
+    atom_mask: A array with shape (num_res, num_atom). This is used to decide
+      which atoms in the atom dimension are present in a given residue. Present
+      atoms should have a nonzero value, e.g. 1.0 or True.
+    **kwargs: A mapping from field name to values. For all array-valued fields
+      these arrays must have a dimension of length num_res. Chain and residue
+      fields should have this as their only dimension and atom fields should be
+      shaped (num_res, num_atom). Coordinate fields may also have arbitrary
+      leading dimensions (they must be the same across all coordinate fields).
+      See structure.{CHAIN,RESIDUE,ATOM}_FIELDS for a list of allowed fields.
+  """
+  num_res, num_atom = atom_mask.shape
+  included_indices = np.flatnonzero(atom_mask)
+
+  array_fields = (
+      structure.CHAIN_FIELDS.keys()
+      | structure.RESIDUE_FIELDS.keys()
+      | structure.ATOM_FIELDS.keys()
+  )
+  initializer_kwargs = {}
+  fields = {}
+  for k, val in kwargs.items():
+    if k not in array_fields:
+      # The kwarg key isn't an array field name. Such kwargs are forwarded as-is
+      # to the constructor. They are expected to be global fields (e.g. name).
+      # Other values will raise an error when the constructor is called.
+      if k in structure.TABLE_FIELDS:
+        raise ValueError(f'Table fields must not be set. Got {k}.')
+      initializer_kwargs[k] = val
+      continue
+    elif val is None:
+      raise ValueError(f'{k} must be non-None.')
+
+    if not isinstance(val, np.ndarray):
+      raise TypeError(f'Value for {k} must be a NumPy array. Got {type(val)}.')
+    if k in structure.CHAIN_FIELDS or k in structure.RESIDUE_FIELDS:
+      if val.shape != (num_res,):
+        raise ValueError(
+            f'{k} must have shape ({num_res=},). Got {val.shape=}.'
+        )
+      # Do not reshape the chain/residue arrays, they have the shape we need.
+      fields[k] = val
+    else:
+      assert k in structure.ATOM_FIELDS
+      if val.shape[-2:] != (num_res, num_atom):
+        raise ValueError(
+            f'{k} must have final two dimensions of length '
+            f'{(num_res, num_atom)=}. Got {val.shape=}.'
+        )
+      leading_dims = val.shape[:-2]
+      flat_val = val.reshape(leading_dims + (-1,), order='C')
+      masked_val = flat_val[..., included_indices]
+      fields[k] = masked_val
+
+  # Get chain IDs or assume this is a single-chain structure.
+  chain_id = kwargs.get('chain_id', np.array(['A'] * num_res, dtype=object))
+  # Find chain starts in res-sized arrays, use these to make chain-sized arrays.
+  chain_start = np.concatenate(
+      ([0], np.where(chain_id[1:] != chain_id[:-1])[0] + 1)
+  )
+  if len(set(chain_id)) != len(chain_start):
+    raise ValueError(f'Chain IDs must be contiguous, but got {chain_id}')
+
+  chain_lengths = np.diff(chain_start, append=len(chain_id))
+  chain_key = np.repeat(np.arange(len(chain_start)), chain_lengths)
+
+  chain_entity_id = fields.get('chain_entity_id')
+  if chain_entity_id is not None:
+    entity_id = chain_entity_id[chain_start]
+  else:
+    entity_id = np.array(
+        [str(mmcif.str_id_to_int_id(cid)) for cid in chain_id[chain_start]],
+        dtype=object,
+    )
+  chain_str_empty = np.full((num_res,), '.', dtype=object)
+  chains_table = structure_tables.Chains(
+      key=chain_key[chain_start],
+      id=chain_id[chain_start],
+      type=fields.get('chain_type', chain_str_empty)[chain_start],
+      auth_asym_id=fields.get('chain_auth_asym_id', chain_id)[chain_start],
+      entity_id=entity_id,
+      entity_desc=fields.get('chain_entity_desc', chain_str_empty)[chain_start],
+  )
+
+  # Since all arrays are residue-shaped, we can use them directly.
+  res_key = np.arange(num_res, dtype=np.int64)
+  res_id = fields.get('res_id', res_key + 1).astype(np.int32)
+  residues_table = structure_tables.Residues(
+      key=res_key,
+      chain_key=chain_key,
+      id=res_id,
+      name=fields.get('res_name', np.full(num_res, 'UNK', dtype=object)),
+      auth_seq_id=fields.get(
+          'res_auth_seq_id', np.char.mod('%d', res_id).astype(object)
+      ),
+      insertion_code=fields.get(
+          'res_insertion_code', np.full(num_res, '?', dtype=object)
+      ),
+  )
+
+  # The atom-sized arrays have already been masked and reshaped.
+  num_atoms_per_res = np.sum(atom_mask, axis=1, dtype=np.int32)
+  num_atoms_total = np.sum(num_atoms_per_res, dtype=np.int32)
+  # Structure is immutable, so use the same array multiple times to save RAM.
+  atom_str_empty = np.full(num_atoms_total, '.', dtype=object)
+  atom_float32_zeros = np.zeros(num_atoms_total, dtype=np.float32)
+  atom_float32_ones = np.ones(num_atoms_total, dtype=np.float32)
+  atoms_table = structure_tables.Atoms(
+      key=np.arange(num_atoms_total, dtype=np.int64),
+      chain_key=np.repeat(chain_key, num_atoms_per_res),
+      res_key=np.repeat(res_key, num_atoms_per_res),
+      name=fields.get('atom_name', atom_str_empty),
+      element=fields.get('atom_element', atom_str_empty),
+      x=fields.get('atom_x', atom_float32_zeros),
+      y=fields.get('atom_y', atom_float32_zeros),
+      z=fields.get('atom_z', atom_float32_zeros),
+      b_factor=fields.get('atom_b_factor', atom_float32_zeros),
+      occupancy=fields.get('atom_occupancy', atom_float32_ones),
+  )
+
+  return structure.Structure(
+      chains=chains_table,
+      residues=residues_table,
+      atoms=atoms_table,
+      bonds=structure_tables.Bonds.make_empty(),  # Currently not set.
+      **initializer_kwargs,
+  )
+
+
+def expand_sequence(
+    sequence: str, chain_type: str, sequence_format: SequenceFormat
+) -> Sequence[str]:
+  """Returns full residue names based on a sequence string.
+
+  Args:
+    sequence: A string representing the sequence.
+    chain_type: The chain type of the sequence.
+    sequence_format: The format of the sequence argument.
+  """
+  match sequence_format:
+    case SequenceFormat.FASTA:
+      if not all(c.isalpha() for c in sequence):
+        raise ValueError(f'Sequence "{sequence}" has non-alphabetic characters')
+      match chain_type:
+        case mmcif_names.PROTEIN_CHAIN:
+          res_name_map = residue_names.PROTEIN_COMMON_ONE_TO_THREE
+          default_res_name = residue_names.UNK
+        case mmcif_names.RNA_CHAIN:
+          res_name_map = {r: r for r in residue_names.RNA_TYPES}
+          default_res_name = residue_names.UNK_RNA
+        case mmcif_names.DNA_CHAIN:
+          res_name_map = residue_names.DNA_COMMON_ONE_TO_TWO
+          default_res_name = residue_names.UNK_DNA
+        case _:
+          raise ValueError(f'{chain_type=} not supported for FASTA format.')
+      return [
+          res_name_map.get(one_letter_res, default_res_name)
+          for one_letter_res in sequence
+      ]
+    case SequenceFormat.CCD_CODES:
+      return sequence.strip('()').split(')(')
+    case SequenceFormat.LIGAND_SMILES:
+      ligand_id, _ = sequence.split(':', maxsplit=1)
+      return [ligand_id]
+
+
+def from_sequences_and_bonds(
+    *,
+    sequences: Sequence[str],
+    chain_types: Sequence[str],
+    sequence_formats: Sequence[SequenceFormat],
+    bonded_atom_pairs: Sequence[tuple[BondAtomId, BondAtomId]] | None,
+    ccd: chemical_components.Ccd,
+    chain_ids: Sequence[str] | None = None,
+    name: str = 'from_sequences_and_bonds',
+    bond_type: str | None = None,
+    **constructor_args,
+) -> structure.Structure:
+  """Returns a minimal structure for the input sequences and bonds.
+
+  The returned structure will have at least one atom per residue. If the
+  residue has any bonded atoms, according to `bonded_atom_pairs`, then
+  all (and only) those atoms will be present for that residue. If the residue
+  is not involved in any bond then an arbitrary atom will be created.
+
+  Args:
+    sequences: A sequence of strings, each one representing a single chain.
+    chain_types: The types of each chain, e.g. polypeptide(L). The n-th element
+      describes the n-th sequence in `sequences`.
+    sequence_formats: The format of each sequence. The n-th element describes
+      the n-th sequence in `sequences`.
+    bonded_atom_pairs: A sequence of bonded atom pairs. Each atom is described
+      as a tuple of (chain_index, res_index, atom_name), where the first two
+      values are 0-based indices. The chain_index is the index of the chain in
+      the `sequences` argument, and the res_index is the index of the residue in
+      that sequence. The atom_name is the name of the atom in the residue, e.g.
+      CA. If the atom is not found in the standard atoms for that residue
+      (according to the CCD) then an error is raised.
+    ccd: The chemical components dictionary.
+    chain_ids: A sequence of chain IDs, one for each chain in `sequences`. If
+      not provided, then the chain IDs will be generated automatically based on
+      sequence indices.
+    name: A name for the returned structure.
+    bond_type: This type will be used for all bonds in the structure, where type
+      follows PDB scheme, e.g. unknown (?), hydrog, metalc, covale, disulf.
+    **constructor_args: These arguments are passed directly to the
+      structure.Structure constructor.
+  """
+  chain_id = []
+  chain_type = []
+  chain_res_count = []
+  res_id = []
+  res_name = []
+  res_atom_count = []
+  atom_name = []
+  atom_element = []
+  chem_comp = {}
+
+  num_bonds = len(bonded_atom_pairs or ())
+  from_atom_key = np.full((num_bonds,), -1, dtype=np.int64)
+  dest_atom_key = np.full((num_bonds,), -1, dtype=np.int64)
+
+  # Create map (chain_i, res_i) -> {atom_name -> (from_idxs dest_idxs)}.
+  # This allows quick lookup of whether a residue has any bonded atoms, and
+  # which bonds those atoms participate in.
+  bond_lookup = _create_bond_lookup(bonded_atom_pairs or ())
+
+  current_atom_key = 0
+  for chain_i, (sequence, curr_chain_type, sequence_format) in enumerate(
+      zip(sequences, chain_types, sequence_formats, strict=True)
+  ):
+    if chain_ids is not None:
+      current_chain_id = chain_ids[chain_i]
+    else:
+      current_chain_id = mmcif.int_id_to_str_id(chain_i + 1)
+    num_chain_residues = 0
+    for res_i, full_res_name in enumerate(
+        expand_sequence(sequence, curr_chain_type, sequence_format)
+    ):
+      current_res_id = res_i + 1
+      num_res_atoms = 0
+
+      # Look for bonded atoms in the bond lookup and if any are found, add
+      # their atom keys to the bond atom_key columns.
+      if bond_indices_by_atom_name := bond_lookup.get((chain_i, res_i)):
+        comp_atoms = None
+        if sequence_format != SequenceFormat.LIGAND_SMILES:
+          comp_atoms = set(ccd.get(full_res_name)['_chem_comp_atom.atom_id'])
+        for bond_atom_name, bond_indices in bond_indices_by_atom_name.items():
+          if comp_atoms is not None and bond_atom_name not in comp_atoms:
+            raise ValueError(
+                f'Bonded atom "{bond_atom_name}" was not found in the list of'
+                f' atoms of the chemical component {full_res_name}. Valid atom'
+                f' names for {full_res_name} are: {sorted(comp_atoms)}.'
+                ' This is likely caused by an invalid atom name in the bonded'
+                f' atom (chain_id={current_chain_id}, res_id={current_res_id},'
+                f' atom_name={bond_atom_name}) specified in `bondedAtomPairs`'
+                ' in the input JSON.'
+            )
+          atom_name.append(bond_atom_name)
+          atom_element.append(
+              _get_atom_element(
+                  ccd=ccd, res_name=full_res_name, atom_name=bond_atom_name
+              )
+          )
+          for from_bond_i in bond_indices.from_indices:
+            from_atom_key[from_bond_i] = current_atom_key
+          for dest_bond_i in bond_indices.dest_indices:
+            dest_atom_key[dest_bond_i] = current_atom_key
+          current_atom_key += 1
+          num_res_atoms += 1
+      else:
+        # If this residue has no bonded atoms then we need to add one atom
+        # like in from_sequences.
+        assert num_res_atoms == 0
+        rep_atom_name, rep_atom_element = _get_representative_atom(
+            ccd=ccd,
+            res_name=full_res_name,
+            chain_type=curr_chain_type,
+            sequence_format=sequence_format,
+        )
+        atom_name.append(rep_atom_name)
+        atom_element.append(rep_atom_element)
+        num_res_atoms += 1
+        current_atom_key += 1
+
+      if sequence_format == SequenceFormat.LIGAND_SMILES:
+        # Sequence expect to be in the format <ligand_id>:<ligand_smiles>,
+        # which always corresponds to a single-residue chain.
+        ligand_id, ligand_smiles = sequence.split(':', maxsplit=1)
+        if ccd.get(ligand_id) is not None:
+          raise ValueError(
+              f'Ligand name {ligand_id} is in CCD - it is not supported to give'
+              ' ligands created from SMILES the same name as CCD components.'
+          )
+        # We need to provide additional chemical components metadata for
+        # ligands specified via SMILES strings since they might not be in CCD.
+        _add_ligand_to_chem_comp(chem_comp, ligand_id, ligand_smiles)
+
+      assert num_res_atoms >= 1
+      res_atom_count.append(num_res_atoms)
+      num_chain_residues += 1
+      res_id.append(current_res_id)
+      res_name.append(full_res_name)
+
+    chain_id.append(current_chain_id)
+    chain_type.append(curr_chain_type)
+    chain_res_count.append(num_chain_residues)
+
+  chem_comp_data = struc_chem_comps.ChemicalComponentsData(chem_comp)
+  chem_comp_data = struc_chem_comps.populate_missing_ccd_data(
+      ccd=ccd,
+      chemical_components_data=chem_comp_data,
+      chemical_component_ids=set(res_name),
+  )
+
+  if bonded_atom_pairs is not None:
+    unknown_bond_col = np.full((num_bonds,), '?', dtype=object)
+    if bond_type is None:
+      bond_type_col = unknown_bond_col
+    else:
+      bond_type_col = np.full((num_bonds,), bond_type, dtype=object)
+    bonds_table = bonds.Bonds(
+        key=np.arange(num_bonds, dtype=np.int64),
+        type=bond_type_col,
+        role=unknown_bond_col,
+        from_atom_key=from_atom_key,
+        dest_atom_key=dest_atom_key,
+    )
+  else:
+    bonds_table = structure_tables.Bonds.make_empty()
+
+  chain_key = np.arange(len(sequences), dtype=np.int64)  # 1 chain per sequence.
+  chain_id = np.array(chain_id, dtype=object)
+  chains_table = structure_tables.Chains(
+      key=chain_key,
+      id=chain_id,
+      type=np.array(chain_type, dtype=object),
+      auth_asym_id=chain_id,
+      entity_id=np.char.mod('%d', chain_key + 1).astype(object),
+      entity_desc=np.array(['.'] * len(chain_key), dtype=object),
+  )
+
+  res_key = np.arange(len(res_name), dtype=np.int64)
+  res_chain_key = np.repeat(chain_key, chain_res_count)
+  residues_table = structure_tables.Residues(
+      key=res_key,
+      chain_key=res_chain_key,
+      id=np.array(res_id, dtype=np.int32),
+      name=np.array(res_name, dtype=object),
+      auth_seq_id=np.char.mod('%d', res_id).astype(object),
+      insertion_code=np.full(len(res_name), '?', dtype=object),
+  )
+
+  num_atoms = current_atom_key
+  atom_float32_zeros = np.zeros(num_atoms, dtype=np.float32)
+  atoms_table = structure_tables.Atoms(
+      key=np.arange(num_atoms, dtype=np.int64),
+      chain_key=np.repeat(res_chain_key, res_atom_count),
+      res_key=np.repeat(res_key, res_atom_count),
+      name=np.array(atom_name, dtype=object),
+      element=np.array(atom_element, dtype=object),
+      x=atom_float32_zeros,
+      y=atom_float32_zeros,
+      z=atom_float32_zeros,
+      b_factor=atom_float32_zeros,
+      occupancy=np.ones(num_atoms, np.float32),
+  )
+
+  return structure.Structure(
+      name=name,
+      atoms=atoms_table,
+      residues=residues_table,
+      chains=chains_table,
+      bonds=bonds_table,
+      chemical_components_data=chem_comp_data,
+      **constructor_args,
+  )
+
+
+class _ChainResBuilder:
+  """Class for incrementally building chain and residue tables."""
+
+  def __init__(
+      self,
+      *,
+      chain_key_by_chain_id: Mapping[str, int],
+      entity_id_by_chain_id: Mapping[str, str],
+      chain_type_by_entity_id: Mapping[str, str],
+      entity_desc_by_entity_id: Mapping[str, str],
+      fix_mse_residues: bool,
+      fix_unknown_dna: bool,
+  ):
+    # Len: num_chains.
+    self.chain_key = []
+    self.chain_id = []
+    self.chain_type = []
+    self.chain_auth_asym_id = []
+    self.chain_entity_id = []
+    self.chain_entity_desc = []
+
+    # Len: num_residues.
+    self.res_key = []
+    self.res_chain_key = []
+    self.res_id = []
+    self.res_name = []
+    self.res_auth_seq_id = []
+    self.res_insertion_code = []
+
+    self.chain_key_by_chain_id = chain_key_by_chain_id
+    self.entity_id_by_chain_id = entity_id_by_chain_id
+    self.chain_type_by_entity_id = chain_type_by_entity_id
+    self.entity_desc_by_entity_id = entity_desc_by_entity_id
+    self.key_for_res: dict[tuple[str, str, str, str], int] = {}
+
+    self._fix_mse_residues = fix_mse_residues
+    self._fix_unknown_dna = fix_unknown_dna
+
+  def add_residues(
+      self,
+      *,
+      chain_ids: np.ndarray,
+      chain_auth_asym_ids: np.ndarray,
+      res_ids: np.ndarray,
+      res_names: np.ndarray,
+      res_auth_seq_ids: np.ndarray,
+      res_ins_codes: np.ndarray,
+  ):
+    """Adds a residue (and its chain) to the tables."""
+    # Create chain table data.
+    if chain_ids.size == 0:
+      return
+
+    chain_ids_with_prev = np.concatenate(
+        (([self.chain_id[-1] if self.chain_id else None], chain_ids))
+    )
+    chain_change_mask = chain_ids_with_prev[:-1] != chain_ids_with_prev[1:]
+    chain_change_ids = chain_ids[chain_change_mask]
+    chain_keys = string_array.remap(
+        chain_change_ids, self.chain_key_by_chain_id, inplace=False
+    )
+    self.chain_key.extend(chain_keys)
+    self.chain_id.extend(chain_change_ids)
+    self.chain_auth_asym_id.extend(chain_auth_asym_ids[chain_change_mask])
+    chain_entity_id = string_array.remap(
+        chain_change_ids, self.entity_id_by_chain_id, inplace=False
+    )
+    self.chain_entity_id.extend(chain_entity_id)
+    chain_type = string_array.remap(
+        chain_entity_id, self.chain_type_by_entity_id, inplace=False
+    )
+    self.chain_type.extend(chain_type)
+    chain_entity_desc = string_array.remap(
+        chain_entity_id, self.entity_desc_by_entity_id, inplace=False
+    )
+    self.chain_entity_desc.extend(chain_entity_desc)
+
+    # Create residue table data.
+    num_prev_res = len(self.res_id)
+    res_keys = np.arange(num_prev_res, num_prev_res + len(res_ids))
+    res_iter = zip(
+        chain_ids,
+        res_auth_seq_ids,
+        res_names,
+        res_ins_codes,
+        strict=True,
+    )
+    key_for_res_update = {
+        res_unique_id: res_key
+        for res_key, res_unique_id in enumerate(res_iter, num_prev_res)
+    }
+    self.key_for_res.update(key_for_res_update)
+    self.res_key.extend(res_keys)
+    self.res_chain_key.extend(
+        string_array.remap(chain_ids, self.chain_key_by_chain_id, inplace=False)
+    )
+    self.res_id.extend(res_ids)
+    self.res_name.extend(res_names)
+    self.res_auth_seq_id.extend(res_auth_seq_ids)
+    self.res_insertion_code.extend(res_ins_codes)
+
+  def make_chains_table(self) -> structure_tables.Chains:
+    """Returns the Structure chains table."""
+    chain_key = np.array(self.chain_key, dtype=np.int64)
+    if not np.all(chain_key[:-1] <= chain_key[1:]):
+      # If the order is inconsistent with the atoms table, sort so that it is.
+      order = np.argsort(self.chain_key, kind='stable')
+      return structure_tables.Chains(
+          key=chain_key[order],
+          id=np.array(self.chain_id, dtype=object)[order],
+          type=np.array(self.chain_type, dtype=object)[order],
+          auth_asym_id=np.array(self.chain_auth_asym_id, dtype=object)[order],
+          entity_id=np.array(self.chain_entity_id, dtype=object)[order],
+          entity_desc=np.array(self.chain_entity_desc, dtype=object)[order],
+      )
+    return structure_tables.Chains(
+        key=chain_key,
+        id=np.array(self.chain_id, dtype=object),
+        type=np.array(self.chain_type, dtype=object),
+        auth_asym_id=np.array(self.chain_auth_asym_id, dtype=object),
+        entity_id=np.array(self.chain_entity_id, dtype=object),
+        entity_desc=np.array(self.chain_entity_desc, dtype=object),
+    )
+
+  def make_residues_table(self) -> structure_tables.Residues:
+    """Returns the Structure residues table."""
+    res_name = np.array(self.res_name, dtype=object)
+    res_chain_key = np.array(self.res_chain_key, dtype=np.int64)
+
+    if self._fix_mse_residues:
+      string_array.remap(res_name, mapping={'MSE': 'MET'}, inplace=True)
+
+    if self._fix_unknown_dna:
+      # Remap residues from N -> DN in DNA chains only.
+      dna_chain_mask = (
+          np.array(self.chain_type, dtype=object) == mmcif_names.DNA_CHAIN
+      )
+      dna_chain_key = np.array(self.chain_key, dtype=object)[dna_chain_mask]
+      res_name[(res_name == 'N') & np.isin(res_chain_key, dna_chain_key)] = 'DN'
+
+    if not np.all(res_chain_key[:-1] <= res_chain_key[1:]):
+      # If the order is inconsistent with the atoms table, sort so that it is.
+      order = np.argsort(res_chain_key, kind='stable')
+      return structure_tables.Residues(
+          key=np.array(self.res_key, dtype=np.int64)[order],
+          chain_key=res_chain_key[order],
+          id=np.array(self.res_id, dtype=np.int32)[order],
+          name=res_name[order],
+          auth_seq_id=np.array(self.res_auth_seq_id, dtype=object)[order],
+          insertion_code=np.array(self.res_insertion_code, dtype=object)[order],
+      )
+    return structure_tables.Residues(
+        key=np.array(self.res_key, dtype=np.int64),
+        chain_key=res_chain_key,
+        id=np.array(self.res_id, dtype=np.int32),
+        name=res_name,
+        auth_seq_id=np.array(self.res_auth_seq_id, dtype=object),
+        insertion_code=np.array(self.res_insertion_code, dtype=object),
+    )
+
+
+def _get_string_array_default(cif: mmcif.Mmcif, key: str, default: list[str]):
+  try:
+    return cif.get_array(key, dtype=object)
+  except KeyError:
+    return default
+
+
+def _generate_required_tables_if_missing(
+    cif: mmcif.Mmcif,
+) -> Mapping[str, Sequence[str]]:
+  """Generates all required tables and columns if missing."""
+  update = {}
+
+  atom_site_entities = _get_string_array_default(
+      cif, '_atom_site.label_entity_id', []
+  )
+
+  # OpenMM produces files that don't have any of the tables and also have
+  # _atom_site.label_entity_id set to '?' for all atoms. We infer the entities
+  # based on the _atom_site.label_asym_id column. We start with cheaper O(1)
+  # checks to prevent running the expensive O(n) check on most files.
+  if (
+      len(atom_site_entities) > 0  # pylint: disable=g-explicit-length-test
+      and '_entity.id' not in cif  # Ignore if the _entity table exists.
+      and atom_site_entities[0] == '?'  # Cheap check.
+      and set(atom_site_entities) == {'?'}  # Expensive check.
+  ):
+    label_asym_ids = cif.get_array('_atom_site.label_asym_id', dtype=object)
+    atom_site_entities = [
+        str(mmcif.str_id_to_int_id(cid)) for cid in label_asym_ids
+    ]
+    # Update _atom_site.label_entity_id to be consistent with the new tables.
+    update['_atom_site.label_entity_id'] = atom_site_entities
+
+  # Check table existence by checking the presence of its primary key.
+  if '_struct_asym.id' not in cif:
+    # Infer the _struct_asym table using the _atom_site table.
+    asym_ids = _get_string_array_default(cif, '_atom_site.label_asym_id', [])
+
+    if len(atom_site_entities) == 0 or len(asym_ids) == 0:  # pylint: disable=g-explicit-length-test
+      raise ValueError(
+          'Could not parse an mmCIF with no _struct_asym table and also no '
+          '_atom_site.label_entity_id or _atom_site.label_asym_id columns.'
+      )
+
+    # Deduplicate, but keep the order intact - dict.fromkeys maintains order.
+    entity_id_chain_id_pairs = list(
+        dict.fromkeys(zip(atom_site_entities, asym_ids, strict=True))
+    )
+    update['_struct_asym.entity_id'] = [e for e, _ in entity_id_chain_id_pairs]
+    update['_struct_asym.id'] = [c for _, c in entity_id_chain_id_pairs]
+
+  if '_entity.id' not in cif:
+    # Infer the _entity_poly and _entity tables using the _atom_site table.
+    residues = _get_string_array_default(cif, '_atom_site.label_comp_id', [])
+    group_pdb = _get_string_array_default(cif, '_atom_site.group_PDB', [])
+    if '_atom_site.label_entity_id' in cif:
+      entities = atom_site_entities
+    else:
+      # If _atom_site.label_entity_id not set, use the asym_id -> entity_id map.
+      asym_to_entity = dict(
+          zip(
+              cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True
+          )
+      )
+      entities = string_array.remap(
+          cif.get_array('_atom_site.label_asym_id', dtype=object),
+          mapping=asym_to_entity,
+      )
+
+    entity_ids = []
+    entity_types = []
+    entity_poly_entity_ids = []
+    entity_poly_types = []
+    entity_poly_table_missing = '_entity_poly.entity_id' not in cif
+    for entity_id, group in itertools.groupby(
+        zip(entities, residues, group_pdb, strict=True), key=lambda e: e[0]
+    ):
+      _, entity_residues, entity_group_pdb = zip(*group, strict=True)
+      entity_type = _guess_entity_type(
+          chain_residues=entity_residues, atom_types=entity_group_pdb
+      )
+      entity_ids.append(entity_id)
+      entity_types.append(entity_type)
+
+      if entity_poly_table_missing and entity_type == mmcif_names.POLYMER_CHAIN:
+        polymer_type = mmcif_names.guess_polymer_type(entity_residues)
+        entity_poly_entity_ids.append(entity_id)
+        entity_poly_types.append(polymer_type)
+
+    update['_entity.id'] = entity_ids
+    update['_entity.type'] = entity_types
+    if entity_poly_table_missing:
+      update['_entity_poly.entity_id'] = entity_poly_entity_ids
+      update['_entity_poly.type'] = entity_poly_types
+
+  if '_atom_site.type_symbol' not in cif:
+    update['_atom_site.type_symbol'] = mmcif.get_or_infer_type_symbol(cif)
+
+  return update
+
+
+def _maybe_add_missing_scheme_tables(
+    cif: mmcif.Mmcif,
+    res_starts: Sequence[int],
+    label_asym_ids: np.ndarray,
+    label_seq_ids: np.ndarray,
+    label_comp_ids: np.ndarray,
+    auth_seq_ids: np.ndarray,
+    pdb_ins_codes: np.ndarray,
+) -> Mapping[str, Sequence[str]]:
+  """If missing, infers the scheme tables from the _atom_site table."""
+  update = {}
+
+  required_poly_seq_scheme_cols = (
+      '_pdbx_poly_seq_scheme.asym_id',
+      '_pdbx_poly_seq_scheme.pdb_seq_num',
+      '_pdbx_poly_seq_scheme.pdb_ins_code',
+      '_pdbx_poly_seq_scheme.seq_id',
+      '_pdbx_poly_seq_scheme.mon_id',
+      '_pdbx_poly_seq_scheme.pdb_strand_id',
+  )
+  if not all(col in cif for col in required_poly_seq_scheme_cols):
+    # Create a mask for atoms where each polymer residue start.
+    entity_id_by_chain_id = dict(
+        zip(cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True)
+    )
+    chain_type_by_entity_id = dict(
+        zip(cif['_entity.id'], cif['_entity.type'], strict=True)
+    )
+    # Remap asym ID -> entity ID.
+    label_entity_id = string_array.remap(
+        label_asym_ids, mapping=entity_id_by_chain_id, inplace=False
+    )
+    # Remap entity ID -> chain type.
+    chain_type = string_array.remap(
+        label_entity_id, mapping=chain_type_by_entity_id, inplace=False
+    )
+    res_mask = np.zeros_like(label_seq_ids, dtype=bool)
+    res_mask[res_starts] = True
+    res_mask &= chain_type == mmcif_names.POLYMER_CHAIN
+
+    entity_poly_seq_cols = (
+        '_entity_poly_seq.entity_id',
+        '_entity_poly_seq.num',
+        '_entity_poly_seq.mon_id',
+    )
+    if all(col in cif for col in entity_poly_seq_cols):
+      # Use _entity_poly_seq if available.
+      poly_seq_num = cif.get_array('_entity_poly_seq.num', dtype=object)
+      poly_seq_mon_id = cif.get_array('_entity_poly_seq.mon_id', dtype=object)
+      poly_seq_entity_id = cif.get_array(
+          '_entity_poly_seq.entity_id', dtype=object
+      )
+      # We have to add the entity ID to the residue ID because multiple residues
+      # can share the same ID. This also allows using string_array.remap.
+      label_seq_id_to_auth_seq_id = dict(
+          zip(
+              np.char.add(label_entity_id[res_mask], label_seq_ids[res_mask]),
+              auth_seq_ids[res_mask],
+              strict=True,
+          )
+      )
+      scheme_pdb_seq_num = string_array.remap(
+          np.char.add(poly_seq_entity_id, poly_seq_num),
+          mapping=label_seq_id_to_auth_seq_id,
+          default_value='.',
+      )
+      label_seq_id_to_ins_code = dict(
+          zip(label_seq_ids[res_mask], pdb_ins_codes[res_mask], strict=True)
+      )
+      scheme_pdb_ins_code = string_array.remap(
+          poly_seq_num, mapping=label_seq_id_to_ins_code, default_value='.'
+      )
+
+      # The _entity_poly_seq table is entity-based, while _pdbx_poly_seq_scheme
+      # is chain-based. A single entity could mean multiple chains (asym_ids),
+      # we therefore need to replicate each entity for all of the chains.
+      scheme_asym_id = []
+      select = []
+      indices = np.arange(len(poly_seq_entity_id), dtype=np.int32)
+      for asym_id, entity_id in zip(
+          cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True
+      ):
+        entity_mask = poly_seq_entity_id == entity_id
+        select.extend(indices[entity_mask])
+        scheme_asym_id.extend([asym_id] * sum(entity_mask))
+
+      scheme_pdb_strand_id = string_array.remap(
+          np.array(scheme_asym_id, dtype=object),
+          mapping=mmcif.get_internal_to_author_chain_id_map(cif),
+          inplace=False,
+      )
+
+      update['_pdbx_poly_seq_scheme.asym_id'] = scheme_asym_id
+      update['_pdbx_poly_seq_scheme.pdb_strand_id'] = scheme_pdb_strand_id
+      update['_pdbx_poly_seq_scheme.pdb_seq_num'] = scheme_pdb_seq_num[select]
+      update['_pdbx_poly_seq_scheme.pdb_ins_code'] = scheme_pdb_ins_code[select]
+      update['_pdbx_poly_seq_scheme.seq_id'] = poly_seq_num[select]
+      update['_pdbx_poly_seq_scheme.mon_id'] = poly_seq_mon_id[select]
+    else:
+      # _entity_poly_seq not available, fallback to _atom_site.
+      res_asym_ids = label_asym_ids[res_mask]
+      res_strand_ids = string_array.remap(
+          array=res_asym_ids,
+          mapping=mmcif.get_internal_to_author_chain_id_map(cif),
+          inplace=False,
+      )
+      update['_pdbx_poly_seq_scheme.asym_id'] = res_asym_ids
+      update['_pdbx_poly_seq_scheme.pdb_strand_id'] = res_strand_ids
+      update['_pdbx_poly_seq_scheme.pdb_seq_num'] = auth_seq_ids[res_mask]
+      update['_pdbx_poly_seq_scheme.pdb_ins_code'] = pdb_ins_codes[res_mask]
+      update['_pdbx_poly_seq_scheme.seq_id'] = label_seq_ids[res_mask]
+      update['_pdbx_poly_seq_scheme.mon_id'] = label_comp_ids[res_mask]
+
+  required_nonpoly_scheme_cols = (
+      '_pdbx_nonpoly_scheme.mon_id',
+      '_pdbx_nonpoly_scheme.asym_id',
+      '_pdbx_nonpoly_scheme.pdb_seq_num',
+      '_pdbx_nonpoly_scheme.pdb_ins_code',
+  )
+  required_branch_scheme_cols = (
+      '_pdbx_branch_scheme.mon_id',
+      '_pdbx_branch_scheme.asym_id',
+      '_pdbx_branch_scheme.pdb_seq_num',
+  )
+
+  # Generate _pdbx_nonpoly_scheme only if both tables are missing.
+  if not (
+      all(col in cif for col in required_nonpoly_scheme_cols)
+      or all(col in cif for col in required_branch_scheme_cols)
+  ):
+    # To be strictly semantically correct, multi-residue ligands should be
+    # written in _pdbx_branch_scheme. However, Structure parsing handles
+    # correctly multi-residue ligands in _pdbx_nonpoly_scheme and the tables
+    # constructed here live only while parsing, hence this is unnecessary.
+    entity_id_by_chain_id = dict(
+        zip(cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True)
+    )
+    chain_type_by_entity_id = dict(
+        zip(cif['_entity.id'], cif['_entity.type'], strict=True)
+    )
+    # Remap asym ID -> entity ID.
+    chain_type = string_array.remap(
+        label_asym_ids, mapping=entity_id_by_chain_id, inplace=False
+    )
+    # Remap entity ID -> chain type.
+    string_array.remap(
+        chain_type, mapping=chain_type_by_entity_id, inplace=True
+    )
+    res_mask = np.zeros_like(label_seq_ids, dtype=bool)
+    res_mask[res_starts] = True
+    res_mask &= chain_type != mmcif_names.POLYMER_CHAIN
+
+    if not np.any(res_mask):
+      return update  # Shortcut: no non-polymer residues.
+
+    ins_codes = string_array.remap(
+        pdb_ins_codes[res_mask], mapping={'?': '.'}, inplace=False
+    )
+
+    update['_pdbx_nonpoly_scheme.asym_id'] = label_asym_ids[res_mask]
+    update['_pdbx_nonpoly_scheme.pdb_seq_num'] = auth_seq_ids[res_mask]
+    update['_pdbx_nonpoly_scheme.pdb_ins_code'] = ins_codes
+    update['_pdbx_nonpoly_scheme.mon_id'] = label_comp_ids[res_mask]
+
+  return update
+
+
+def _get_chain_key_by_chain_id(
+    resolved_chain_ids: np.ndarray, struct_asym_chain_ids: np.ndarray
+) -> Mapping[str, int]:
+  """Returns chain key for each chain ID respecting resolved chain ordering."""
+  # Check that all chain IDs found in the (potentially filtered) _atom_site
+  # table are present in the _struct_asym table.
+  unique_resolved_chain_ids = set(resolved_chain_ids)
+  if not unique_resolved_chain_ids.issubset(set(struct_asym_chain_ids)):
+    unique_resolved_chain_ids = sorted(unique_resolved_chain_ids)
+    unique_struct_asym_chain_ids = sorted(set(struct_asym_chain_ids))
+    raise ValueError(
+        'Bad mmCIF: chain IDs in _atom_site.label_asym_id '
+        f'{unique_resolved_chain_ids} is not a subset of chain IDs in '
+        f'_struct_asym.id {unique_struct_asym_chain_ids}.'
+    )
+
+  resolved_mask = string_array.isin(
+      struct_asym_chain_ids, unique_resolved_chain_ids
+  )
+  # For all resolved chains, use the _atom_site order they appear in. E.g.
+  # resolved_chain_ids     = [B A   E D F]
+  # struct_asym_chain_ids  = [A B C D E F]
+  # consistent_chain_order = [B A C E D F]
+  # chain_keys             = [0 1 2 3 4 5]
+  consistent_chain_order = struct_asym_chain_ids.copy()
+  consistent_chain_order[resolved_mask] = resolved_chain_ids
+  return dict(zip(consistent_chain_order, range(len(struct_asym_chain_ids))))
+
+
+def get_tables(
+    cif: mmcif.Mmcif,
+    fix_mse_residues: bool,
+    fix_arginines: bool,
+    fix_unknown_dna: bool,
+    include_water: bool,
+    include_other: bool,
+    model_id: str,
+) -> tuple[
+    structure_tables.Chains, structure_tables.Residues, structure_tables.Atoms
+]:
+  """Returns chain, residue, and atom tables from a parsed mmcif.
+
+  Args:
+    cif: A parsed mmcif.Mmcif.
+    fix_mse_residues: See from_mmcif.
+    fix_arginines: See from_mmcif.
+    fix_unknown_dna: See from_mmcif.
+    include_water: See from_mmcif.
+    include_other: See from_mmcif.
+    model_id: A string defining which model ID to use. If set, only coordinates,
+      b-factors and occupancies for the given model are returned. If empty,
+      coordinates, b-factors and occupanciesall for models are returned with a
+      leading dimension of num_models. Note that the model_id argument in
+      from_mmcif is an integer and has slightly different use (see from_mmcif).
+  """
+  # Add any missing tables and columns we require for parsing.
+  if cif_update := _generate_required_tables_if_missing(cif):
+    cif = cif.copy_and_update(cif_update)
+
+  # Resolve alt-locs, selecting only a single option for each residue. Also
+  # computes the layout, which defines where chain and residue boundaries are.
+  atom_site_all_models, layout = mmcif_utils.filter(
+      cif,
+      include_nucleotides=True,
+      include_ligands=True,
+      include_water=include_water,
+      include_other=include_other,
+      model_id=model_id,
+  )
+  atom_site_first_model = atom_site_all_models[0]
+
+  # Get atom information from the _atom_site table.
+  def _first_model_string_array(col: str) -> np.ndarray:
+    return cif.get_array(col, dtype=object, gather=atom_site_first_model)
+
+  def _requested_models_float_array(col: str) -> np.ndarray:
+    if not model_id:
+      # Return data for all models with a leading dimension of num_models.
+      return cif.get_array(col, dtype=np.float32, gather=atom_site_all_models)
+    else:
+      # Return data only for the single requested model.
+      return cif.get_array(col, dtype=np.float32, gather=atom_site_first_model)
+
+  # These columns are the same for all models, fetch them just for the 1st one.
+  label_comp_ids = _first_model_string_array('_atom_site.label_comp_id')
+  label_asym_ids = _first_model_string_array('_atom_site.label_asym_id')
+  label_seq_ids = _first_model_string_array('_atom_site.label_seq_id')
+  label_atom_ids = _first_model_string_array('_atom_site.label_atom_id')
+  if '_atom_site.auth_seq_id' in cif:
+    auth_seq_ids = _first_model_string_array('_atom_site.auth_seq_id')
+  else:
+    auth_seq_ids = label_seq_ids  # auth_seq_id unset, fallback to label_seq_id.
+  type_symbols = _first_model_string_array('_atom_site.type_symbol')
+  pdbx_pdb_ins_codes = _first_model_string_array('_atom_site.pdbx_PDB_ins_code')
+
+  # These columns are different for all models, fetch them as requested.
+  atom_x = _requested_models_float_array('_atom_site.Cartn_x')
+  atom_y = _requested_models_float_array('_atom_site.Cartn_y')
+  atom_z = _requested_models_float_array('_atom_site.Cartn_z')
+  atom_b_factor = _requested_models_float_array('_atom_site.B_iso_or_equiv')
+  atom_occupancy = _requested_models_float_array('_atom_site.occupancy')
+
+  # Make sure the scheme (residue) tables exist in case they are not present.
+  if cif_update := _maybe_add_missing_scheme_tables(
+      cif,
+      res_starts=layout.residue_starts(),
+      label_asym_ids=label_asym_ids,
+      label_seq_ids=label_seq_ids,
+      label_comp_ids=label_comp_ids,
+      auth_seq_ids=auth_seq_ids,
+      pdb_ins_codes=pdbx_pdb_ins_codes,
+  ):
+    cif = cif.copy_and_update(cif_update)
+
+  # Fix common issues found in mmCIF files, like swapped arginine NH atoms.
+  mmcif_utils.fix_residues(
+      layout,
+      comp_id=label_comp_ids,
+      atom_id=label_atom_ids,
+      atom_x=atom_x[0] if not model_id else atom_x,
+      atom_y=atom_y[0] if not model_id else atom_y,
+      atom_z=atom_z[0] if not model_id else atom_z,
+      fix_arg=fix_arginines,
+  )
+
+  # Get keys for chains in the order they appear in _atom_site while also
+  # dealing with empty chains.
+  resolved_chain_ids = label_asym_ids[layout.chain_starts()]
+  struct_asym_chain_ids = cif.get_array('_struct_asym.id', dtype=object)
+
+  chain_key_by_chain_id = _get_chain_key_by_chain_id(
+      resolved_chain_ids=resolved_chain_ids,
+      struct_asym_chain_ids=struct_asym_chain_ids,
+  )
+  entity_id_by_chain_id = dict(
+      zip(struct_asym_chain_ids, cif['_struct_asym.entity_id'], strict=True)
+  )
+  entity_description = cif.get(
+      '_entity.pdbx_description', ['?'] * len(cif['_entity.id'])
+  )
+  entity_desc_by_entity_id = dict(
+      zip(cif['_entity.id'], entity_description, strict=True)
+  )
+  chain_type_by_entity_id = mmcif.get_chain_type_by_entity_id(cif)
+  auth_asym_id_by_chain_id = mmcif.get_internal_to_author_chain_id_map(cif)
+
+  chain_res_builder = _ChainResBuilder(
+      chain_key_by_chain_id=chain_key_by_chain_id,
+      entity_id_by_chain_id=entity_id_by_chain_id,
+      chain_type_by_entity_id=chain_type_by_entity_id,
+      entity_desc_by_entity_id=entity_desc_by_entity_id,
+      fix_mse_residues=fix_mse_residues,
+      fix_unknown_dna=fix_unknown_dna,
+  )
+
+  # Collect data for polymer chain and residue tables. _pdbx_poly_seq_scheme is
+  # guaranteed to be present thanks to _maybe_add_missing_scheme_tables.
+  def _get_poly_seq_scheme_col(col: str) -> np.ndarray:
+    return cif.get_array(key=f'_pdbx_poly_seq_scheme.{col}', dtype=object)
+
+  poly_seq_asym_ids = _get_poly_seq_scheme_col('asym_id')
+  poly_seq_pdb_seq_nums = _get_poly_seq_scheme_col('pdb_seq_num')
+  poly_seq_seq_ids = _get_poly_seq_scheme_col('seq_id')
+  poly_seq_mon_ids = _get_poly_seq_scheme_col('mon_id')
+  poly_seq_pdb_strand_ids = _get_poly_seq_scheme_col('pdb_strand_id')
+  poly_seq_pdb_ins_codes = _get_poly_seq_scheme_col('pdb_ins_code')
+  string_array.remap(
+      poly_seq_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+
+  # We resolved alt-locs earlier for the atoms table. In cases of heterogeneous
+  # residues (a residue with an alt-loc that is of different residue type), we
+  # need to also do the same resolution in the residues table. Compute a mask
+  # for the residues that were selected in the atoms table.
+  poly_seq_mask = mmcif_utils.selected_polymer_residue_mask(
+      layout=layout,
+      atom_site_label_asym_ids=label_asym_ids[layout.residue_starts()],
+      atom_site_label_seq_ids=label_seq_ids[layout.residue_starts()],
+      atom_site_label_comp_ids=label_comp_ids[layout.residue_starts()],
+      poly_seq_asym_ids=poly_seq_asym_ids,
+      poly_seq_seq_ids=poly_seq_seq_ids,
+      poly_seq_mon_ids=poly_seq_mon_ids,
+  )
+
+  if not include_other and poly_seq_mask:
+    # Mask filtered-out residues so that they are not treated as missing.
+    # Instead, we don't want them included in the chains/residues tables at all.
+    keep_mask = string_array.remap(
+        poly_seq_asym_ids,
+        mapping={cid: True for cid in resolved_chain_ids},
+        default_value=False,
+        inplace=False,
+    ).astype(bool)
+    poly_seq_mask &= keep_mask
+
+  chain_res_builder.add_residues(
+      chain_ids=poly_seq_asym_ids[poly_seq_mask],
+      chain_auth_asym_ids=poly_seq_pdb_strand_ids[poly_seq_mask],
+      res_ids=poly_seq_seq_ids[poly_seq_mask].astype(np.int32),
+      res_names=poly_seq_mon_ids[poly_seq_mask],
+      res_auth_seq_ids=poly_seq_pdb_seq_nums[poly_seq_mask],
+      res_ins_codes=poly_seq_pdb_ins_codes[poly_seq_mask],
+  )
+
+  # Collect data for ligand chain and residue tables. _pdbx_nonpoly_scheme
+  # could be empty/unset if there are only branched ligands.
+  def _get_nonpoly_scheme_col(col: str) -> np.ndarray:
+    key = f'_pdbx_nonpoly_scheme.{col}'
+    if f'_pdbx_nonpoly_scheme.{col}' in cif:
+      return cif.get_array(key=key, dtype=object)
+    else:
+      return np.array([], dtype=object)
+
+  nonpoly_asym_ids = _get_nonpoly_scheme_col('asym_id')
+  nonpoly_auth_seq_ids = _get_nonpoly_scheme_col('pdb_seq_num')
+  nonpoly_pdb_ins_codes = _get_nonpoly_scheme_col('pdb_ins_code')
+  nonpoly_mon_ids = _get_nonpoly_scheme_col('mon_id')
+  nonpoly_auth_asym_id = string_array.remap(
+      nonpoly_asym_ids, mapping=auth_asym_id_by_chain_id, inplace=False
+  )
+
+  def _get_branch_scheme_col(col: str) -> np.ndarray:
+    key = f'_pdbx_branch_scheme.{col}'
+    if f'_pdbx_branch_scheme.{col}' in cif:
+      return cif.get_array(key=key, dtype=object)
+    else:
+      return np.array([], dtype=object)
+
+  branch_asym_ids = _get_branch_scheme_col('asym_id')
+  branch_auth_seq_ids = _get_branch_scheme_col('pdb_seq_num')
+  branch_pdb_ins_codes = _get_branch_scheme_col('pdb_ins_code')
+  branch_mon_ids = _get_branch_scheme_col('mon_id')
+  branch_auth_asym_id = string_array.remap(
+      branch_asym_ids, mapping=auth_asym_id_by_chain_id, inplace=False
+  )
+
+  if branch_asym_ids.size > 0 and branch_pdb_ins_codes.size == 0:
+    branch_pdb_ins_codes = np.array(['.'] * branch_asym_ids.size, dtype=object)
+
+  # Compute the heterogeneous residue masks as above, this time for ligands.
+  nonpoly_mask, branch_mask = mmcif_utils.selected_ligand_residue_mask(
+      layout=layout,
+      atom_site_label_asym_ids=label_asym_ids[layout.residue_starts()],
+      atom_site_label_seq_ids=label_seq_ids[layout.residue_starts()],
+      atom_site_auth_seq_ids=auth_seq_ids[layout.residue_starts()],
+      atom_site_label_comp_ids=label_comp_ids[layout.residue_starts()],
+      atom_site_pdbx_pdb_ins_codes=pdbx_pdb_ins_codes[layout.residue_starts()],
+      nonpoly_asym_ids=nonpoly_asym_ids,
+      nonpoly_auth_seq_ids=nonpoly_auth_seq_ids,
+      nonpoly_pdb_ins_codes=nonpoly_pdb_ins_codes,
+      nonpoly_mon_ids=nonpoly_mon_ids,
+      branch_asym_ids=branch_asym_ids,
+      branch_auth_seq_ids=branch_auth_seq_ids,
+      branch_pdb_ins_codes=branch_pdb_ins_codes,
+      branch_mon_ids=branch_mon_ids,
+  )
+
+  if not include_water:
+    if nonpoly_mask:
+      nonpoly_mask &= (nonpoly_mon_ids != 'HOH') & (nonpoly_mon_ids != 'DOD')
+    if branch_mask:
+      # Fix for bad mmCIFs that have water in the branch scheme table.
+      branch_mask &= (branch_mon_ids != 'HOH') & (branch_mon_ids != 'DOD')
+
+  string_array.remap(
+      pdbx_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+  string_array.remap(
+      nonpoly_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+  string_array.remap(
+      branch_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+
+  def _ligand_residue_ids(chain_ids: np.ndarray) -> np.ndarray:
+    """Computes internal residue ID for ligand residues that don't have it."""
+
+    # E.g. chain_ids=[A, A, A, B, C, C, D, D, D] -> [1, 2, 3, 1, 1, 2, 1, 2, 3].
+    indices = np.arange(chain_ids.size, dtype=np.int32)
+    return (indices + 1) - np.maximum.accumulate(
+        indices * (chain_ids != np.roll(chain_ids, 1))
+    )
+
+  branch_residue_ids = _ligand_residue_ids(branch_asym_ids[branch_mask])
+  nonpoly_residue_ids = _ligand_residue_ids(nonpoly_asym_ids[nonpoly_mask])
+
+  chain_res_builder.add_residues(
+      chain_ids=branch_asym_ids[branch_mask],
+      chain_auth_asym_ids=branch_auth_asym_id[branch_mask],
+      res_ids=branch_residue_ids,
+      res_names=branch_mon_ids[branch_mask],
+      res_auth_seq_ids=branch_auth_seq_ids[branch_mask],
+      res_ins_codes=branch_pdb_ins_codes[branch_mask],
+  )
+
+  chain_res_builder.add_residues(
+      chain_ids=nonpoly_asym_ids[nonpoly_mask],
+      chain_auth_asym_ids=nonpoly_auth_asym_id[nonpoly_mask],
+      res_ids=nonpoly_residue_ids,
+      res_names=nonpoly_mon_ids[nonpoly_mask],
+      res_auth_seq_ids=nonpoly_auth_seq_ids[nonpoly_mask],
+      res_ins_codes=nonpoly_pdb_ins_codes[nonpoly_mask],
+  )
+
+  chains = chain_res_builder.make_chains_table()
+  residues = chain_res_builder.make_residues_table()
+
+  # Construct foreign residue keys for the atoms table.
+  res_ends = np.array(layout.residues(), dtype=np.int32)
+  res_starts = np.array(layout.residue_starts(), dtype=np.int32)
+  res_lengths = res_ends - res_starts
+
+  # Check just for HOH, DOD can be part e.g. of hydroxycysteine.
+  if include_water:
+    res_chain_types = chains.apply_array_to_column(
+        column_name='type', arr=residues.chain_key
+    )
+    water_mask = res_chain_types != mmcif_names.WATER
+    if 'HOH' in set(residues.name[water_mask]):
+      raise ValueError('Bad mmCIF file: non-water entity has water molecules.')
+  else:
+    # Include resolved and unresolved residues.
+    if 'HOH' in set(residues.name) | set(label_comp_ids[res_starts]):
+      raise ValueError('Bad mmCIF file: non-water entity has water molecules.')
+
+  atom_chain_key = string_array.remap(
+      label_asym_ids, mapping=chain_res_builder.chain_key_by_chain_id
+  ).astype(int)
+
+  # If any of the residue lookups failed, the mmCIF is corrupted.
+  try:
+    atom_res_key_per_res = string_array.remap_multiple(
+        (
+            label_asym_ids[res_starts],
+            auth_seq_ids[res_starts],
+            label_comp_ids[res_starts],
+            pdbx_pdb_ins_codes[res_starts],
+        ),
+        mapping=chain_res_builder.key_for_res,
+    )
+  except KeyError as e:
+    raise ValueError(
+        'Lookup for the following atom from the _atom_site table failed: '
+        f'(label_asym_id, auth_seq_id, res_name, ins_code)={e}. This typically '
+        'indicates that the _pdbx_poly_seq_scheme, _pdbx_nonpoly_scheme, or '
+        '_pdbx_branch_scheme tables do not have data for all residues present '
+        'in the _atom_site table. It could also be due to a known issue with '
+        'a small number of multi-model mmCIFs.'
+    ) from e
+
+  # The residue ID will be shared for all atoms within that residue.
+  atom_res_key = np.repeat(atom_res_key_per_res, repeats=res_lengths)
+
+  if fix_mse_residues:
+    met_residues_mask = (residues.name == 'MET')[atom_res_key]
+    unfixed_mse_selenium_mask = met_residues_mask & (label_atom_ids == 'SE')
+    label_atom_ids[unfixed_mse_selenium_mask] = 'SD'
+    type_symbols[unfixed_mse_selenium_mask] = 'S'
+
+  atoms = structure_tables.Atoms(
+      key=atom_site_first_model,
+      chain_key=atom_chain_key,
+      res_key=atom_res_key,
+      name=label_atom_ids,
+      element=type_symbols,
+      x=atom_x,
+      y=atom_y,
+      z=atom_z,
+      b_factor=atom_b_factor,
+      occupancy=atom_occupancy,
+  )
+
+  return chains, residues, atoms
+
+
+def from_atom_arrays(
+    *,
+    res_id: np.ndarray,
+    name: str = 'unset',
+    release_date: datetime.date | None = None,
+    resolution: float | None = None,
+    structure_method: str | None = None,
+    all_residues: Mapping[str, Sequence[tuple[str, int]]] | None = None,
+    bioassembly_data: bioassemblies.BioassemblyData | None = None,
+    chemical_components_data: (
+        struc_chem_comps.ChemicalComponentsData | None
+    ) = None,
+    bond_table: structure_tables.Bonds | None = None,
+    chain_id: np.ndarray | None = None,
+    chain_type: np.ndarray | None = None,
+    res_name: np.ndarray | None = None,
+    atom_key: np.ndarray | None = None,
+    atom_name: np.ndarray | None = None,
+    atom_element: np.ndarray | None = None,
+    atom_x: np.ndarray | None = None,
+    atom_y: np.ndarray | None = None,
+    atom_z: np.ndarray | None = None,
+    atom_b_factor: np.ndarray | None = None,
+    atom_occupancy: np.ndarray | None = None,
+) -> structure.Structure:
+  """Returns a Structure constructed from atom array level data.
+
+  All fields except name and, res_id are optional, all array fields consist of a
+  value for each atom in the structure - so residue and chain values should hold
+  the same value for each atom in the chain or residue. Fields which are not
+  defined are filled with default values.
+
+  Validation is performed by the Structure constructor where possible - but
+  author_naming scheme and all_residues must be checked in this function.
+
+  It is not possible to construct structures with chains that do not contain
+  any resolved residues using this function. If this is necessary, use the
+  structure.Structure constructor directly.
+
+  Args:
+    res_id: Integer array of shape [num_atom]. The unique residue identifier for
+      each residue. mmCIF field - _atom_site.label_seq_id.
+    name: The name of the structure. E.g. a PDB ID.
+    release_date: The release date of the structure as a `datetime.date`.
+    resolution: The resolution of the structure in Angstroms.
+    structure_method: The method used to solve this structure's coordinates.
+    all_residues: An optional mapping from each chain ID (i.e. label_asym_id) to
+      a sequence of (label_comp_id, label_seq_id) tuples, one per residue. This
+      can contain residues that aren't present in the atom arrays. This is
+      common in experimental data where some residues are not resolved but are
+      known to be present.
+    bioassembly_data: An optional instance of bioassembly.BioassemblyData. If
+      present then a new Structure representing a specific bioassembly can be
+      extracted using `Structure.generate_bioassembly(assembly_id)`.
+    chemical_components_data: An optional instance of ChemicalComponentsData.
+      Its content will be used for providing metadata about chemical components
+      in this Structure instance. If not specified information will be retrieved
+      from the standard chemical component dictionary (CCD, for more details see
+      https://www.wwpdb.org/data/ccd).
+    bond_table: A table representing manually-specified bonds. This corresponds
+      to the _struct_conn table in an mmCIF. Atoms are identified by their key,
+      as specified by the atom_key column. If this table is provided then the
+      atom_key column must also be defined.
+    chain_id: String array of shape [num_atom] of unique chain identifiers.
+      mmCIF field - _atom_site.label_asym_id.
+    chain_type: String array of shape [num_atom]. The molecular type of the
+      current chain (e.g. polyribonucleotide). mmCIF field - _entity_poly.type
+      OR _entity.type (for non-polymers).
+    res_name: String array of shape [num_atom].. The name of each residue,
+      typically a 3 letter string for polypeptides or 1-2 letter strings for
+      polynucleotides. mmCIF field - _atom_site.label_comp_id.
+    atom_key: A unique sorted integer array, used only by the bonds table to
+      identify the atoms participating in each bond. If the bonds table is
+      specified then this column must be non-None.
+    atom_name: String array of shape [num_atom]. The name of each atom (e.g CA,
+      O2', etc.). mmCIF field - _atom_site.label_atom_id.
+    atom_element: String array of shape [num_atom]. The element type of each
+      atom (e.g. C, O, N, etc.). mmCIF field - _atom_site.type_symbol.
+    atom_x: Float array of shape [..., num_atom] of atom x coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_y: Float array of shape [..., num_atom] of atom y coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_z: Float array of shape [..., num_atom] of atom z coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_b_factor: Float array of shape [..., num_atom] or [num_atom] of atom
+      b-factors or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+    atom_occupancy: Float array of shape [..., num_atom] or [num_atom] of atom
+      occupancies or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+  """
+
+  atoms, residues, chains = structure_tables.tables_from_atom_arrays(
+      res_id=res_id,
+      all_residues=all_residues,
+      chain_id=chain_id,
+      chain_type=chain_type,
+      res_name=res_name,
+      atom_key=atom_key,
+      atom_name=atom_name,
+      atom_element=atom_element,
+      atom_x=atom_x,
+      atom_y=atom_y,
+      atom_z=atom_z,
+      atom_b_factor=atom_b_factor,
+      atom_occupancy=atom_occupancy,
+  )
+
+  return structure.Structure(
+      name=name,
+      release_date=release_date,
+      resolution=resolution,
+      structure_method=structure_method,
+      bioassembly_data=bioassembly_data,
+      chemical_components_data=chemical_components_data,
+      atoms=atoms,
+      chains=chains,
+      residues=residues,
+      bonds=bond_table or structure_tables.Bonds.make_empty(),
+  )
+
+
+def _guess_entity_type(
+    chain_residues: Collection[str], atom_types: Collection[str]
+) -> str:
+  """Guess the entity type (polymer/non-polymer/water) based on residues/atoms.
+
+  We treat both arguments as unordered collections since we care only whether
+  all elements satisfy come conditions. The chain_residues can be either
+  grouped by residue (length num_res), or it can be raw (length num_atoms).
+  Atom type is unique for each atom in a residue, so don't group atom_types.
+
+  Args:
+    chain_residues: A sequence of full residue name (1-letter for DNA, 2-letters
+      for RNA, 3 for protein). The _atom_site.label_comp_id column in mmCIF.
+    atom_types: Atom type: ATOM or HETATM. The _atom_site.group_PDB column in
+      mmCIF.
+
+  Returns:
+    One of polymer/non-polymer/water based on the following criteria:
+    * If all atoms are HETATMs and all residues are water -> water.
+    * If all atoms are HETATMs and not all residues are water -> non-polymer.
+    * Otherwise -> polymer.
+  """
+  if not chain_residues or not atom_types:
+    raise ValueError(
+        f'chain_residues (len {len(chain_residues)}) and atom_types (len '
+        f'{len(atom_types)}) must be both non-empty. Got: {chain_residues=} '
+        f'and {atom_types=}'
+    )
+
+  if all(a == 'HETATM' for a in atom_types):
+    if all(c in residue_names.WATER_TYPES for c in chain_residues):
+      return mmcif_names.WATER
+    return mmcif_names.NON_POLYMER_CHAIN
+  return mmcif_names.POLYMER_CHAIN

flax_model/alphafold3/structure/sterics.py

@@ -0,0 +1,132 @@
+
+
+"""Functions relating to spatial locations of atoms within a structure."""
+
+from collections.abc import Collection, Sequence
+
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.structure import mmcif
+import numpy as np
+import scipy
+
+
+def _make_atom_has_clash_mask(
+    kd_query_result: np.ndarray,
+    struc: structure.Structure,
+    ignore_chains: Collection[str],
+) -> np.ndarray:
+  """Returns a boolean NumPy array representing whether each atom has a clash.
+
+  Args:
+    kd_query_result: NumPy array containing N-atoms arrays, each array
+      containing indices to atoms that clash with the N'th atom.
+    struc: Structure over which clashes were detected.
+    ignore_chains: Collection of chains that should not be considered clashing.
+      A boolean NumPy array of length N atoms.
+  """
+  atom_is_clashing = np.zeros((struc.num_atoms,), dtype=bool)
+  for atom_index, clashes in enumerate(kd_query_result):
+    chain_i = struc.chain_id[atom_index]
+    if chain_i in ignore_chains:
+      continue
+    islig_i = struc.is_ligand_mask[atom_index]
+    for clashing_atom_index in clashes:
+      chain_c = struc.chain_id[clashing_atom_index]
+      if chain_c in ignore_chains:
+        continue
+      islig_c = struc.is_ligand_mask[clashing_atom_index]
+      if (
+          clashing_atom_index == atom_index
+          or chain_i == chain_c
+          or islig_i != islig_c
+      ):
+        # Ignore clashes within chain or between ligand and polymer.
+        continue
+      atom_is_clashing[atom_index] = True
+  return atom_is_clashing
+
+
+def find_clashing_chains(
+    struc: structure.Structure,
+    clash_thresh_angstrom: float = 1.7,
+    clash_thresh_fraction: float = 0.3,
+) -> Sequence[str]:
+  """Finds chains that clash with others.
+
+  Clashes are defined by polymer backbone atoms and all ligand atoms.
+  Ligand-polymer clashes are not dropped.
+
+  Will not find clashes if all coordinates are 0. Coordinates are all 0s if
+  the structure is generated from sequences only, as done for inference in
+  dendro for example.
+
+  Args:
+    struc: The structure defining the chains and atom positions.
+    clash_thresh_angstrom: Below this distance, atoms are considered clashing.
+    clash_thresh_fraction: Chains with more than this fraction of their atoms
+      considered clashing will be dropped. This value should be in the range (0,
+      1].
+
+  Returns:
+    A sequence of chain ids for chains that clash.
+
+  Raises:
+    ValueError: If `clash_thresh_fraction` is not in range (0,1].
+  """
+  if not 0 < clash_thresh_fraction <= 1:
+    raise ValueError('clash_thresh_fraction must be in range (0,1]')
+
+  struc_backbone = struc.filter_polymers_to_single_atom_per_res()
+  if struc_backbone.num_chains == 0:
+    return []
+
+  # If the coordinates are all 0, do not search for clashes.
+  if not np.any(struc_backbone.coords):
+    return []
+
+  coord_kdtree = scipy.spatial.cKDTree(struc_backbone.coords)
+
+  # For each atom coordinate, find all atoms within the clash thresh radius.
+  clashing_per_atom = coord_kdtree.query_ball_point(
+      struc_backbone.coords, r=clash_thresh_angstrom
+  )
+  chain_ids = struc_backbone.chains
+  if struc_backbone.atom_occupancy is not None:
+    chain_occupancy = np.array([
+        np.mean(struc_backbone.atom_occupancy[start:end])
+        for start, end in struc_backbone.iter_chain_ranges()
+    ])
+  else:
+    chain_occupancy = None
+
+  # Remove chains until no more significant clashing.
+  chains_to_remove = set()
+  for _ in range(len(chain_ids)):
+    # Calculate maximally clashing.
+    atom_has_clash = _make_atom_has_clash_mask(
+        clashing_per_atom, struc_backbone, chains_to_remove
+    )
+    clashes_per_chain = np.array([
+        atom_has_clash[start:end].mean()
+        for start, end in struc_backbone.iter_chain_ranges()
+    ])
+    max_clash = np.max(clashes_per_chain)
+    if max_clash <= clash_thresh_fraction:
+      # None of the remaining chains exceed the clash fraction threshold, so
+      # we can exit.
+      break
+
+    # Greedily remove worst with the lowest occupancy.
+    most_clashes = np.nonzero(clashes_per_chain == max_clash)[0]
+    if chain_occupancy is not None:
+      occupancy_clashing = chain_occupancy[most_clashes]
+      last_lowest_occupancy = (
+          len(occupancy_clashing) - np.argmin(occupancy_clashing[::-1]) - 1
+      )
+      worst_and_last = most_clashes[last_lowest_occupancy]
+    else:
+      worst_and_last = most_clashes[-1]
+
+    chains_to_remove.add(chain_ids[worst_and_last])
+
+  return sorted(chains_to_remove, key=mmcif.str_id_to_int_id)

flax_model/alphafold3/structure/structure_tables.py

@@ -0,0 +1,811 @@
+
+
+"""Table implementations for the Structure class."""
+
+import collections
+from collections.abc import Mapping, Sequence
+import dataclasses
+import functools
+import itertools
+import typing
+from typing import Any, ClassVar, Self
+
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import aggregation
+from flax_model.alphafold3.cpp import string_array
+from flax_model.alphafold3.structure import bonds as bonds_module
+from flax_model.alphafold3.structure import mmcif
+from flax_model.alphafold3.structure import table
+import numpy as np
+
+
+Bonds = bonds_module.Bonds
+
+
+def _residue_name_to_record_name(
+    residue_name: np.ndarray,
+    polymer_mask: np.ndarray,
+) -> np.ndarray:
+  """Returns record names (ATOM/HETATM) given residue names and polymer mask."""
+  record_name = np.array(['HETATM'] * len(residue_name), dtype=object)
+  record_name[polymer_mask] = string_array.remap(
+      residue_name[polymer_mask],
+      mapping={r: 'ATOM' for r in residue_names.STANDARD_POLYMER_TYPES},
+      default_value='HETATM',
+  )
+  return record_name
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class AuthorNamingScheme:
+  """A mapping from internal values to author values in a mmCIF.
+
+  Fields:
+    auth_asym_id: A mapping from label_asym_id to auth_asym_id.
+    auth_seq_id: A mapping from label_asym_id to a mapping from
+      label_seq_id to auth_seq_id.
+    insertion_code: A mapping from label_asym_id to a mapping from
+      label_seq_id to insertion codes.
+    entity_id: A mapping from label_asym_id to _entity.id.
+    entity_desc: A mapping from _entity.id to _entity.pdbx_description.
+  """
+
+  auth_asym_id: Mapping[str, str]
+  auth_seq_id: Mapping[str, Mapping[int, str]]
+  insertion_code: Mapping[str, Mapping[int, str | None]]
+  entity_id: Mapping[str, str]
+  entity_desc: Mapping[str, str]
+
+
+def _default(
+    candidate_value: np.ndarray | None, default_value: Sequence[Any], dtype: Any
+) -> np.ndarray:
+  if candidate_value is None:
+    return np.array(default_value, dtype=dtype)
+  return np.array(candidate_value, dtype=dtype)
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Atoms(table.Table):
+  """Table of atoms in a Structure."""
+
+  chain_key: np.ndarray
+  res_key: np.ndarray
+  name: np.ndarray
+  element: np.ndarray
+  x: np.ndarray
+  y: np.ndarray
+  z: np.ndarray
+  b_factor: np.ndarray
+  occupancy: np.ndarray
+  multimodel_cols: ClassVar[tuple[str, ...]] = (
+      'x',
+      'y',
+      'z',
+      'b_factor',
+      'occupancy',
+  )
+
+  def __post_init__(self):
+    # Validates that the atom coordinates, b-factors and occupancies are finite.
+    for column_name in ('x', 'y', 'z', 'b_factor', 'occupancy'):
+      column = self.get_column(column_name)
+      if not np.isfinite(column).all():
+        raise ValueError(
+            f'Column {column_name} must not contain NaN/inf values.'
+        )
+    # super().__post_init__() can't be used as that causes the following error:
+    # TypeError: super(type, obj): obj must be an instance or subtype of type
+    super(Atoms, self).__post_init__()
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        chain_key=np.array([], dtype=np.int64),
+        res_key=np.array([], dtype=np.int64),
+        name=np.array([], dtype=object),
+        element=np.array([], dtype=object),
+        x=np.array([], dtype=np.float32),
+        y=np.array([], dtype=np.float32),
+        z=np.array([], dtype=np.float32),
+        b_factor=np.array([], dtype=np.float32),
+        occupancy=np.array([], dtype=np.float32),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      chain_key: np.ndarray,
+      res_key: np.ndarray,
+      key: np.ndarray | None = None,
+      name: np.ndarray | None = None,
+      element: np.ndarray | None = None,
+      x: np.ndarray | None = None,
+      y: np.ndarray | None = None,
+      z: np.ndarray | None = None,
+      b_factor: np.ndarray | None = None,
+      occupancy: np.ndarray | None = None,
+  ) -> Self:
+    """Create an Atoms table with minimal user inputs."""
+    num_atoms = len(chain_key)
+    if not num_atoms:
+      return cls.make_empty()
+    return Atoms(
+        chain_key=chain_key,
+        res_key=res_key,
+        key=_default(key, np.arange(num_atoms), np.int64),
+        name=_default(name, ['?'] * num_atoms, object),
+        element=_default(element, ['?'] * num_atoms, object),
+        x=_default(x, [0.0] * num_atoms, np.float32),
+        y=_default(y, [0.0] * num_atoms, np.float32),
+        z=_default(z, [0.0] * num_atoms, np.float32),
+        b_factor=_default(b_factor, [0.0] * num_atoms, np.float32),
+        occupancy=_default(occupancy, [1.0] * num_atoms, np.float32),
+    )
+
+  def get_value_by_index(
+      self, column_name: str, index: int
+  ) -> table.TableEntry | np.ndarray:
+    if column_name in self.multimodel_cols:
+      return self.get_column(column_name)[..., index]
+    else:
+      return self.get_column(column_name)[index]
+
+  def copy_and_update_coords(self, coords: np.ndarray) -> Self:
+    """Returns a copy with the x, y and z columns updated."""
+    if coords.shape[-1] != 3:
+      raise ValueError(
+          f'Expecting 3-dimensional coordinates, got {coords.shape}'
+      )
+    return typing.cast(
+        Atoms,
+        self.copy_and_update(
+            x=coords[..., 0], y=coords[..., 1], z=coords[..., 2]
+        ),
+    )
+
+  @property
+  def shape(self) -> tuple[int, ...]:
+    return self.x.shape
+
+  @property
+  def ndim(self) -> int:
+    return len(self.shape)
+
+  @functools.cached_property
+  def num_models(self) -> int:
+    """The number of models of this Structure."""
+    leading_dims = self.shape[:-1]
+    match leading_dims:
+      case ():
+        return 1
+      case (single_leading_dim_size,):
+        return single_leading_dim_size
+      case _:
+        raise ValueError(
+            'num_models not defined for atom tables with more than one '
+            'leading dimension.'
+        )
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Residues(table.Table):
+  """Table of residues in a Structure."""
+
+  chain_key: np.ndarray
+  id: np.ndarray
+  name: np.ndarray
+  auth_seq_id: np.ndarray
+  insertion_code: np.ndarray
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        chain_key=np.array([], dtype=np.int64),
+        id=np.array([], dtype=np.int32),
+        name=np.array([], dtype=object),
+        auth_seq_id=np.array([], dtype=object),
+        insertion_code=np.array([], dtype=object),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      id: np.ndarray,  # pylint:disable=redefined-builtin
+      chain_key: np.ndarray,
+      key: np.ndarray | None = None,
+      name: np.ndarray | None = None,
+      auth_seq_id: np.ndarray | None = None,
+      insertion_code: np.ndarray | None = None,
+  ) -> Self:
+    """Create a Residues table with minimal user inputs."""
+    num_res = len(id)
+    if not num_res:
+      return cls.make_empty()
+    return Residues(
+        key=_default(key, np.arange(num_res), np.int64),
+        id=id,
+        chain_key=chain_key,
+        name=_default(name, ['UNK'] * num_res, object),
+        auth_seq_id=_default(auth_seq_id, id.astype(str), object),
+        insertion_code=_default(insertion_code, ['?'] * num_res, object),
+    )
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Chains(table.Table):
+  """Table of chains in a Structure."""
+
+  id: np.ndarray
+  type: np.ndarray
+  auth_asym_id: np.ndarray
+  entity_id: np.ndarray
+  entity_desc: np.ndarray
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        id=np.array([], dtype=object),
+        type=np.array([], dtype=object),
+        auth_asym_id=np.array([], dtype=object),
+        entity_id=np.array([], dtype=object),
+        entity_desc=np.array([], dtype=object),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      id: np.ndarray,  # pylint:disable=redefined-builtin
+      key: np.ndarray | None = None,
+      type: np.ndarray | None = None,  # pylint:disable=redefined-builtin
+      auth_asym_id: np.ndarray | None = None,
+      entity_id: np.ndarray | None = None,
+      entity_desc: np.ndarray | None = None,
+  ) -> Self:
+    """Create a Chains table with minimal user inputs."""
+    num_chains = len(id)
+    if not num_chains:
+      return cls.make_empty()
+
+    return Chains(
+        key=_default(key, np.arange(num_chains), np.int64),
+        id=id,
+        type=_default(type, [mmcif_names.PROTEIN_CHAIN] * num_chains, object),
+        auth_asym_id=_default(auth_asym_id, id, object),
+        entity_id=_default(
+            entity_id, np.arange(1, num_chains + 1).astype(str), object
+        ),
+        entity_desc=_default(entity_desc, ['.'] * num_chains, object),
+    )
+
+
+def to_mmcif_sequence_and_entity_tables(
+    chains: Chains,
+    residues: Residues,
+    atom_res_key: np.ndarray,
+) -> Mapping[str, Sequence[str]]:
+  """Returns raw sequence and entity mmCIF tables."""
+  raw_mmcif = collections.defaultdict(list)
+  chains_by_entity_id = {}
+  written_entity_poly_seq_ids = set()
+  present_res_keys = set(atom_res_key)
+
+  # Performance optimisation: Find residue indices for each chain in advance, so
+  # that we don't have to do redunant masking work for each chain.
+  res_indices_for_chain = aggregation.indices_grouped_by_value(
+      residues.chain_key
+  )
+
+  for chain in chains.iterrows():
+    # Add all chain information to the _struct_asym table.
+    chain_id = chain['id']  # Saves multiple dict lookups.
+    auth_asym_id = chain['auth_asym_id']
+    entity_id = chain['entity_id']
+    chains_by_entity_id.setdefault(entity_id, []).append(chain)
+    raw_mmcif['_struct_asym.id'].append(chain_id)
+    raw_mmcif['_struct_asym.entity_id'].append(entity_id)
+
+    res_chain_indices = res_indices_for_chain[chain['key']]
+    chain_type = chain['type']
+    is_polymer = chain_type in mmcif_names.POLYMER_CHAIN_TYPES
+    is_water = chain_type == mmcif_names.WATER
+    is_branched = len(res_chain_indices) > 1 and not is_polymer and not is_water
+    write_entity_poly_seq = entity_id not in written_entity_poly_seq_ids
+
+    # Iterate over the individual masked residue table columns, as that doesn't
+    # create a copy (only a view), while residues[res_chain_indices] does.
+    for res_key, res_name, res_id, pdb_seq_num, res_ins_code in zip(
+        residues.key[res_chain_indices],
+        residues.name[res_chain_indices],
+        residues.id[res_chain_indices],
+        residues.auth_seq_id[res_chain_indices],
+        residues.insertion_code[res_chain_indices],
+        strict=True,
+    ):
+      is_missing = res_key not in present_res_keys
+      str_res_id = str(res_id)
+      # While atom_site uses "?" for insertion codes, scheme tables use ".".
+      ins_code = (res_ins_code or '.').replace('?', '.')
+      auth_seq_num = '?' if is_missing else pdb_seq_num
+
+      if is_polymer:
+        raw_mmcif['_pdbx_poly_seq_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.seq_id'].append(str_res_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_poly_seq_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_strand_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_ins_code'].append(ins_code)
+        # Structure doesn't support heterogeneous sequences.
+        raw_mmcif['_pdbx_poly_seq_scheme.hetero'].append('n')
+        if write_entity_poly_seq:
+          raw_mmcif['_entity_poly_seq.entity_id'].append(entity_id)
+          raw_mmcif['_entity_poly_seq.num'].append(str_res_id)
+          raw_mmcif['_entity_poly_seq.mon_id'].append(res_name)
+          # Structure doesn't support heterogeneous sequences.
+          raw_mmcif['_entity_poly_seq.hetero'].append('n')
+          written_entity_poly_seq_ids.add(entity_id)
+      elif is_branched:
+        raw_mmcif['_pdbx_branch_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_branch_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_branch_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_branch_scheme.num'].append(str_res_id)
+        raw_mmcif['_pdbx_branch_scheme.pdb_asym_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_branch_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_branch_scheme.auth_asym_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_branch_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_branch_scheme.pdb_ins_code'].append(ins_code)
+        # Structure doesn't support heterogeneous sequences.
+        raw_mmcif['_pdbx_branch_scheme.hetero'].append('n')
+      else:
+        raw_mmcif['_pdbx_nonpoly_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_nonpoly_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_strand_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_ins_code'].append(ins_code)
+
+  # Add _entity and _entity_poly tables.
+  for entity_id, chains in chains_by_entity_id.items():
+    # chains should always be a non-empty list because of how we constructed
+    # chains_by_entity_id.
+    assert chains
+    # All chains for a given entity should have the same type and sequence
+    # so we can pick the first one without losing information.
+    key_chain = chains[0]
+    raw_mmcif['_entity.id'].append(entity_id)
+    raw_mmcif['_entity.pdbx_description'].append(key_chain['entity_desc'])
+    entity_type = key_chain['type']
+    if entity_type not in mmcif_names.POLYMER_CHAIN_TYPES:
+      raw_mmcif['_entity.type'].append(entity_type)
+    else:
+      raw_mmcif['_entity.type'].append('polymer')
+      raw_mmcif['_entity_poly.entity_id'].append(entity_id)
+      raw_mmcif['_entity_poly.type'].append(entity_type)
+
+      # _entity_poly.pdbx_strand_id is a comma-separated list of
+      # auth_asym_ids that are part of the entity.
+      raw_mmcif['_entity_poly.pdbx_strand_id'].append(
+          ','.join(chain['auth_asym_id'] for chain in chains)
+      )
+  return raw_mmcif
+
+
+def to_mmcif_atom_site_and_bonds_table(
+    *,
+    chains: Chains,
+    residues: Residues,
+    atoms: Atoms,
+    bonds: Bonds,
+    coords_decimal_places: int,
+) -> Mapping[str, Sequence[str]]:
+  """Returns raw _atom_site and _struct_conn mmCIF tables."""
+  raw_mmcif = collections.defaultdict(list)
+  # Use [value] * num wherever possible since it is about 10x faster than list
+  # comprehension in such cases. Also use f-strings instead of str() - faster.
+  total_atoms = atoms.size * atoms.num_models
+  raw_mmcif['_atom_site.id'] = [f'{i}' for i in range(1, total_atoms + 1)]
+  raw_mmcif['_atom_site.label_alt_id'] = ['.'] * total_atoms
+  # Use format_float_array instead of list comprehension for performance.
+  raw_mmcif['_atom_site.Cartn_x'] = mmcif.format_float_array(
+      values=atoms.x.ravel(), num_decimal_places=coords_decimal_places
+  )
+  raw_mmcif['_atom_site.Cartn_y'] = mmcif.format_float_array(
+      values=atoms.y.ravel(), num_decimal_places=coords_decimal_places
+  )
+  raw_mmcif['_atom_site.Cartn_z'] = mmcif.format_float_array(
+      values=atoms.z.ravel(), num_decimal_places=coords_decimal_places
+  )
+
+  # atoms.b_factor or atoms.occupancy can be flat even when the coordinates have
+  # leading dimensions. In this case we tile it to match.
+  if atoms.b_factor.ndim == 1:
+    atom_b_factor = np.tile(atoms.b_factor, atoms.num_models)
+  else:
+    atom_b_factor = atoms.b_factor.ravel()
+  raw_mmcif['_atom_site.B_iso_or_equiv'] = mmcif.format_float_array(
+      values=atom_b_factor, num_decimal_places=2
+  )
+
+  if atoms.occupancy.ndim == 1:
+    atom_occupancy = np.tile(atoms.occupancy, atoms.num_models)
+  else:
+    atom_occupancy = atoms.occupancy.ravel()
+  raw_mmcif['_atom_site.occupancy'] = mmcif.format_float_array(
+      values=atom_occupancy.ravel(), num_decimal_places=2
+  )
+
+  label_atom_id = atoms.name
+  type_symbol = atoms.element
+  label_comp_id = residues.apply_array_to_column('name', atoms.res_key)
+  label_asym_id = chains.apply_array_to_column('id', atoms.chain_key)
+  label_entity_id = chains.apply_array_to_column('entity_id', atoms.chain_key)
+  # Performance optimisation: Do the int->str conversion on num_residue-sized,
+  # array, then select instead of selecting and then converting.
+  label_seq_id = residues.id.astype('str').astype(object)[
+      ..., residues.index_by_key[atoms.res_key]
+  ]
+
+  # _atom_site.label_seq_id is '.' for non-polymers.
+  non_polymer_chain_mask = string_array.isin(
+      chains.type, mmcif_names.POLYMER_CHAIN_TYPES, invert=True
+  )
+  non_polymer_chain_keys = chains.key[non_polymer_chain_mask]
+  non_polymer_atom_mask = np.isin(atoms.chain_key, non_polymer_chain_keys)
+  label_seq_id[non_polymer_atom_mask] = '.'
+
+  auth_asym_id = chains.apply_array_to_column('auth_asym_id', atoms.chain_key)
+  auth_seq_id = residues.apply_array_to_column('auth_seq_id', atoms.res_key)
+  pdbx_pdb_ins_code = residues.apply_array_to_column(
+      'insertion_code', atoms.res_key
+  )
+  string_array.remap(pdbx_pdb_ins_code, mapping={None: '?'}, inplace=True)
+
+  group_pdb = _residue_name_to_record_name(
+      residue_name=label_comp_id, polymer_mask=~non_polymer_atom_mask
+  )
+
+  def tile_for_models(arr: np.ndarray) -> list[str]:
+    if atoms.num_models == 1:
+      return arr.tolist()  # Memory optimisation: np.tile(arr, 1) does a copy.
+    return np.tile(arr, atoms.num_models).tolist()
+
+  raw_mmcif['_atom_site.group_PDB'] = tile_for_models(group_pdb)
+  raw_mmcif['_atom_site.label_atom_id'] = tile_for_models(label_atom_id)
+  raw_mmcif['_atom_site.type_symbol'] = tile_for_models(type_symbol)
+  raw_mmcif['_atom_site.label_comp_id'] = tile_for_models(label_comp_id)
+  raw_mmcif['_atom_site.label_asym_id'] = tile_for_models(label_asym_id)
+  raw_mmcif['_atom_site.label_entity_id'] = tile_for_models(label_entity_id)
+  raw_mmcif['_atom_site.label_seq_id'] = tile_for_models(label_seq_id)
+  raw_mmcif['_atom_site.auth_asym_id'] = tile_for_models(auth_asym_id)
+  raw_mmcif['_atom_site.auth_seq_id'] = tile_for_models(auth_seq_id)
+  raw_mmcif['_atom_site.pdbx_PDB_ins_code'] = tile_for_models(pdbx_pdb_ins_code)
+  model_id = np.array(
+      [str(i + 1) for i in range(atoms.num_models)], dtype=object
+  )
+  raw_mmcif['_atom_site.pdbx_PDB_model_num'] = np.repeat(
+      model_id, [atoms.size] * atoms.num_models
+  ).tolist()
+
+  if bonds.key.size > 0:
+    raw_mmcif.update(
+        bonds.to_mmcif_dict_from_atom_arrays(
+            atom_key=atoms.key,
+            chain_id=label_asym_id,
+            res_id=label_seq_id,
+            res_name=label_comp_id,
+            atom_name=label_atom_id,
+            auth_asym_id=auth_asym_id,
+            auth_seq_id=auth_seq_id,
+            insertion_code=np.array(pdbx_pdb_ins_code),
+        )
+    )
+  return raw_mmcif
+
+
+def _flatten_author_naming_scheme_table(
+    res_table: Mapping[str, Mapping[int, str]],
+    chain_ids: np.ndarray,
+    res_chain_ids: np.ndarray,
+    res_ids: np.ndarray,
+    default_if_missing: str,
+    table_name: str,
+) -> np.ndarray:
+  """Flattens an author naming scheme table consistently with res_ids."""
+  if not set(chain_ids).issubset(res_table):
+    raise ValueError(
+        f'Chain IDs in the chain_id array must be a subset of {table_name} in '
+        'author naming scheme:\n'
+        f'chain_ids: {sorted(chain_ids)}\n'
+        f'{table_name} keys: {sorted(res_table.keys())}'
+    )
+
+  chain_change_mask = res_chain_ids[1:] != res_chain_ids[:-1]
+  res_chain_boundaries = np.concatenate(
+      ([0], np.where(chain_change_mask)[0] + 1, [len(res_chain_ids)])
+  )
+
+  flat_vals = np.empty(len(res_ids), dtype=object)
+  for chain_start, chain_end in itertools.pairwise(res_chain_boundaries):
+    chain_id = res_chain_ids[chain_start]
+    chain_res_ids = res_ids[chain_start:chain_end]
+    chain_mapping = res_table[chain_id]
+    flat_vals[chain_start:chain_end] = [
+        chain_mapping.get(r, default_if_missing) for r in chain_res_ids
+    ]
+
+  return flat_vals
+
+
+def tables_from_atom_arrays(
+    *,
+    res_id: np.ndarray,
+    author_naming_scheme: AuthorNamingScheme | None = None,
+    all_residues: Mapping[str, Sequence[tuple[str, int]]] | None = None,
+    chain_id: np.ndarray | None = None,
+    chain_type: np.ndarray | None = None,
+    res_name: np.ndarray | None = None,
+    atom_key: np.ndarray | None = None,
+    atom_name: np.ndarray | None = None,
+    atom_element: np.ndarray | None = None,
+    atom_x: np.ndarray | None = None,
+    atom_y: np.ndarray | None = None,
+    atom_z: np.ndarray | None = None,
+    atom_b_factor: np.ndarray | None = None,
+    atom_occupancy: np.ndarray | None = None,
+) -> tuple[Atoms, Residues, Chains]:
+  """Returns Structure tables constructed from atom array level data.
+
+  All fields except name and, res_id are optional, all array fields consist of a
+  value for each atom in the structure - so residue and chain values should hold
+  the same value for each atom in the chain or residue. Fields which are not
+  defined are filled with default values.
+
+  Validation is performed by the Structure constructor where possible - but
+  author_naming scheme and all_residues must be checked in this function.
+
+  It is not possible to construct structures with chains that do not contain
+  any resolved residues using this function. If this is necessary, use the
+  structure.Structure constructor directly.
+
+  Args:
+    res_id: Integer array of shape [num_atom]. The unique residue identifier for
+      each residue. mmCIF field - _atom_site.label_seq_id.
+    author_naming_scheme: An optional instance of AuthorNamingScheme to use when
+      converting this structure to mmCIF.
+    all_residues: An optional mapping from each chain ID (i.e. label_asym_id) to
+      a sequence of (label_comp_id, label_seq_id) tuples, one per residue. This
+      can contain residues that aren't present in the atom arrays. This is
+      common in experimental data where some residues are not resolved but are
+      known to be present.
+    chain_id: String array of shape [num_atom] of unique chain identifiers.
+      mmCIF field - _atom_site.label_asym_id.
+    chain_type: String array of shape [num_atom]. The molecular type of the
+      current chain (e.g. polyribonucleotide). mmCIF field - _entity_poly.type
+      OR _entity.type (for non-polymers).
+    res_name: String array of shape [num_atom].. The name of each residue,
+      typically a 3 letter string for polypeptides or 1-2 letter strings for
+      polynucleotides. mmCIF field - _atom_site.label_comp_id.
+    atom_key: A unique sorted integer array, used only by the bonds table to
+      identify the atoms participating in each bond. If the bonds table is
+      specified then this column must be non-None.
+    atom_name: String array of shape [num_atom]. The name of each atom (e.g CA,
+      O2', etc.). mmCIF field - _atom_site.label_atom_id.
+    atom_element: String array of shape [num_atom]. The element type of each
+      atom (e.g. C, O, N, etc.). mmCIF field - _atom_site.type_symbol.
+    atom_x: Float array of shape [..., num_atom] of atom x coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_y: Float array of shape [..., num_atom] of atom y coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_z: Float array of shape [..., num_atom] of atom z coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_b_factor: Float array of shape [..., num_atom] or [num_atom] of atom
+      b-factors or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+    atom_occupancy: Float array of shape [..., num_atom] or [num_atom] of atom
+      occupancies or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+  """
+  num_atoms = len(res_id)
+
+  for arr_name, array, dtype in (
+      ('chain_id', chain_id, object),
+      ('chain_type', chain_type, object),
+      ('res_id', res_id, np.int32),
+      ('res_name', res_name, object),
+      ('atom_key', atom_key, np.int64),
+      ('atom_name', atom_name, object),
+      ('atom_element', atom_element, object),
+  ):
+    if array is not None and array.shape != (num_atoms,):
+      raise ValueError(f'{arr_name} shape {array.shape} != ({num_atoms},)')
+    if array is not None and array.dtype != dtype:
+      raise ValueError(f'{arr_name} dtype {array.dtype} != {dtype}')
+
+  for arr_name, array in (
+      ('atom_x', atom_x),
+      ('atom_y', atom_y),
+      ('atom_z', atom_z),
+      ('atom_b_factor', atom_b_factor),
+      ('atom_occupancy', atom_occupancy),
+  ):
+    if array is not None and array.shape[-1] != num_atoms:
+      raise ValueError(f'{arr_name} last dim {array.shape[-1]} != {num_atoms=}')
+    if (
+        array is not None
+        and array.dtype != np.float32
+        and array.dtype != np.float64
+    ):
+      raise ValueError(
+          f'{arr_name} must be np.float32 or np.float64, got {array.dtype=}'
+      )
+
+  if all_residues is not None and (res_name is None or res_id is None):
+    raise ValueError(
+        'If all_residues != None, res_name and res_id must not be None either.'
+    )
+
+  if num_atoms == 0:
+    return Atoms.make_empty(), Residues.make_empty(), Chains.make_empty()
+
+  if chain_id is None:
+    chain_id = np.full(shape=num_atoms, fill_value='A', dtype=object)
+  if res_name is None:
+    res_name = np.full(shape=num_atoms, fill_value='UNK', dtype=object)
+
+  chain_change_mask = chain_id[1:] != chain_id[:-1]
+  chain_start = np.concatenate(([0], np.where(chain_change_mask)[0] + 1))
+  res_start = np.concatenate(
+      ([0], np.where((res_id[1:] != res_id[:-1]) | chain_change_mask)[0] + 1)
+  )
+
+  if len(set(chain_id)) != len(chain_start):
+    raise ValueError(f'Chain IDs must be contiguous, but got {chain_id}')
+
+  # We do not support chains with unresolved residues-only in this function.
+  chain_ids = chain_id[chain_start]
+  if all_residues and set(all_residues.keys()) != set(chain_ids):
+    raise ValueError(
+        'all_residues must contain the same set of chain IDs as the chain_id '
+        f'array:\nall_residues keys: {sorted(all_residues.keys())}\n'
+        f'chain_ids: {sorted(chain_ids)}.'
+    )
+  # Make sure all_residue ordering is consistent with chain_id.
+  if all_residues and np.any(list(all_residues.keys()) != chain_ids):
+    all_residues = {cid: all_residues[cid] for cid in chain_ids}
+
+  # Create the chains table.
+  num_chains = len(chain_ids)
+  chain_keys = np.arange(num_chains, dtype=np.int64)
+  chain_key_by_chain_id = dict(zip(chain_ids, chain_keys, strict=True))
+
+  if chain_type is not None:
+    chain_types = chain_type[chain_start]
+  else:
+    chain_types = np.full(num_chains, mmcif_names.PROTEIN_CHAIN, dtype=object)
+
+  if author_naming_scheme is not None:
+    auth_asym_id = string_array.remap(
+        chain_ids, author_naming_scheme.auth_asym_id
+    )
+    entity_id = string_array.remap(
+        chain_ids, author_naming_scheme.entity_id, default_value='.'
+    )
+    entity_desc = string_array.remap(
+        entity_id, author_naming_scheme.entity_desc, default_value='.'
+    )
+  else:
+    auth_asym_id = chain_ids
+    entity_id = (chain_keys + 1).astype(str).astype(object)
+    entity_desc = np.full(num_chains, '.', dtype=object)
+
+  chains = Chains(
+      key=chain_keys,
+      id=chain_ids,
+      type=chain_types,
+      auth_asym_id=auth_asym_id,
+      entity_id=entity_id,
+      entity_desc=entity_desc,
+  )
+
+  # Create the residues table.
+  if all_residues is not None:
+    residue_order = []
+    for cid, residues in all_residues.items():
+      residue_order.extend((cid, rname, int(rid)) for (rname, rid) in residues)
+    res_chain_ids, res_names, res_ids = zip(*residue_order)
+    res_chain_ids = np.array(res_chain_ids, dtype=object)
+    res_ids = np.array(res_ids, dtype=np.int32)
+    res_names = np.array(res_names, dtype=object)
+  else:
+    res_chain_ids = chain_id[res_start]
+    res_ids = res_id[res_start]
+    res_names = res_name[res_start]
+    residue_order = list(zip(res_chain_ids, res_names, res_ids))
+
+  if author_naming_scheme is not None and author_naming_scheme.auth_seq_id:
+    auth_seq_id = _flatten_author_naming_scheme_table(
+        author_naming_scheme.auth_seq_id,
+        chain_ids=chain_ids,
+        res_chain_ids=res_chain_ids,
+        res_ids=res_ids,
+        default_if_missing='.',
+        table_name='auth_seq_id',
+    )
+  else:
+    auth_seq_id = res_ids.astype(str).astype(object)
+
+  if author_naming_scheme is not None and author_naming_scheme.insertion_code:
+    insertion_code = _flatten_author_naming_scheme_table(
+        author_naming_scheme.insertion_code,
+        chain_ids=chain_ids,
+        res_chain_ids=res_chain_ids,
+        res_ids=res_ids,
+        default_if_missing='?',
+        table_name='insertion_code',
+    )
+    # Make sure insertion code of None is mapped to '.'.
+    insertion_code = string_array.remap(insertion_code, {None: '?'})
+  else:
+    insertion_code = np.full(shape=len(res_ids), fill_value='?', dtype=object)
+
+  res_key_by_res = {res: i for i, res in enumerate(residue_order)}
+  res_keys = np.arange(len(residue_order), dtype=np.int64)
+  res_chain_keys = string_array.remap(
+      res_chain_ids, chain_key_by_chain_id
+  ).astype(np.int64)
+  residues = Residues(
+      chain_key=res_chain_keys,
+      key=res_keys,
+      id=res_ids,
+      name=res_names,
+      auth_seq_id=auth_seq_id,
+      insertion_code=insertion_code,
+  )
+
+  if atom_key is None:
+    atom_key = np.arange(num_atoms, dtype=np.int64)
+
+  atom_chain_keys = string_array.remap(chain_id, chain_key_by_chain_id).astype(
+      np.int64
+  )
+
+  try:
+    atom_res_keys = [res_key_by_res[r] for r in zip(chain_id, res_name, res_id)]
+  except KeyError as e:
+    missing_chain_id, missing_res_name, missing_res_id = e.args[0]
+    raise ValueError(
+        'Inconsistent res_name, res_id and all_residues. Could not find '
+        f'residue with chain_id={missing_chain_id}, '
+        f'res_name={missing_res_name}, res_id={missing_res_id} in all_residues.'
+    ) from e
+
+  atoms = Atoms(
+      key=atom_key,
+      chain_key=atom_chain_keys,
+      res_key=np.array(atom_res_keys, dtype=np.int64),
+      name=_default(atom_name, ['?'] * num_atoms, object),
+      element=_default(atom_element, ['?'] * num_atoms, object),
+      x=_default(atom_x, [0.0] * num_atoms, np.float32),
+      y=_default(atom_y, [0.0] * num_atoms, np.float32),
+      z=_default(atom_z, [0.0] * num_atoms, np.float32),
+      b_factor=_default(atom_b_factor, [0.0] * num_atoms, np.float32),
+      occupancy=_default(atom_occupancy, [1.0] * num_atoms, np.float32),
+  )
+  return atoms, residues, chains

flax_model/alphafold3/structure/table.py

@@ -0,0 +1,555 @@
+
+
+"""Table module for atom/residue/chain tables in Structure.
+
+Tables are intended to be lightweight collections of columns, loosely based
+on a pandas dataframe, for use in the Structure class.
+"""
+
+import abc
+from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence
+import dataclasses
+import functools
+import graphlib
+import typing
+from typing import Any, Protocol, Self, TypeAlias, TypeVar, overload
+
+from flax_model.alphafold3.cpp import string_array
+import numpy as np
+
+
+TableEntry: TypeAlias = str | int | float | None
+FilterPredicate: TypeAlias = (
+    TableEntry
+    | Iterable[Any]  # Workaround for b/326384670. Tighten once fixed.
+    | Callable[[Any], bool]  # Workaround for b/326384670. Tighten once fixed.
+    | Callable[[np.ndarray], bool]
+)
+
+
+class RowLookup(Protocol):
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> Mapping[str, Any]:
+    ...
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Table:
+  """Parent class for structure tables.
+
+  A table is a collection of columns of equal length, where one column is the
+  key. The key uniquely identifies each row in the table.
+
+  A table can refer to other tables by including a foreign key column, whose
+  values are key values from the other table's key column. These column can have
+  arbitrary names and are treated like any other integer-valued column.
+
+  See the `Database` class in this module for utilities for handing sets of
+  tables that are related via foreign keys.
+
+  NB: This does not correspond to an mmCIF table.
+  """
+
+  key: np.ndarray
+
+  def __post_init__(self):
+    for col_name in self.columns:
+      if (col_len := self.get_column(col_name).shape[-1]) != self.size:
+        raise ValueError(
+            f'All columns should have length {self.size} but got "{col_name}"'
+            f' with length {col_len}.'
+        )
+      self.get_column(col_name).flags.writeable = False  # Make col immutable.
+    if self.key.size and self.key.min() < 0:
+      raise ValueError(
+          'Key values must be non-negative. Got negative values:'
+          f' {set(self.key[self.key < 0])}'
+      )
+    self.key.flags.writeable = False  # Make key immutable.
+
+  def __getstate__(self) -> dict[str, Any]:
+    """Returns members with cached properties removed for pickling."""
+    cached_props = {
+        k
+        for k, v in self.__class__.__dict__.items()
+        if isinstance(v, functools.cached_property)
+    }
+    return {k: v for k, v in self.__dict__.items() if k not in cached_props}
+
+  @functools.cached_property
+  def index_by_key(self) -> np.ndarray:
+    """Mapping from key values to their index in the column arrays.
+
+    i.e.: self.key[index_by_key[k]] == k
+    """
+    if not self.key.size:
+      return np.array([], dtype=np.int64)
+    else:
+      index_by_key = np.zeros(np.max(self.key) + 1, dtype=np.int64)
+      index_by_key[self.key] = np.arange(self.size)
+      return index_by_key
+
+  @functools.cached_property
+  def columns(self) -> tuple[str, ...]:
+    """The names of the columns in the table, including the key column."""
+    return tuple(field.name for field in dataclasses.fields(self))
+
+  @functools.cached_property
+  def items(self) -> Mapping[str, np.ndarray]:
+    """Returns the mapping from column names to column values."""
+    return {col: getattr(self, col) for col in self.columns}
+
+  @functools.cached_property
+  def size(self) -> int:
+    """The number of rows in the table."""
+    return self.key.shape[-1]
+
+  def __len__(self) -> int:
+    return self.size
+
+  def get_column(self, column_name: str) -> np.ndarray:
+    """Gets a column by name."""
+    # Performance optimisation: use the cached columns, instead of getattr.
+    return self.items[column_name]
+
+  def apply_array(self, arr: np.ndarray) -> Self:
+    """Returns a sliced table using a key (!= index) array or a boolean mask."""
+    if arr.dtype == bool and np.all(arr):
+      return self  # Shortcut: No-op, so just return.
+
+    return self.copy_and_update(**{
+        column_name: self.apply_array_to_column(column_name, arr)
+        for column_name in self.columns
+    })
+
+  def apply_index(self, index_arr: np.ndarray) -> Self:
+    """Returns a sliced table using an index (!= key) array."""
+    if index_arr.dtype == bool:
+      raise ValueError('The index array must not be a boolean mask.')
+
+    return self.copy_and_update(
+        **{col: self.get_column(col)[..., index_arr] for col in self.columns}
+    )
+
+  def apply_array_to_column(
+      self,
+      column_name: str,
+      arr: np.ndarray,
+  ) -> np.ndarray:
+    """Returns a sliced column array using a key array or a boolean mask."""
+    if arr.dtype == bool:
+      return self.get_column(column_name)[..., arr]
+    else:
+      return self.get_column(column_name)[..., self.index_by_key[arr]]
+
+  def get_value_by_index(self, column_name: str, index: int) -> Any:
+    return self.get_column(column_name)[index]
+
+  def get_value_by_key(
+      self,
+      column_name: str,
+      key: int | np.integer,
+  ) -> TableEntry:
+    """Gets the value of a column at the row with specified key value."""
+    return self.get_value_by_index(column_name, self.index_by_key[key])
+
+  @overload
+  def __getitem__(self, key: str) -> np.ndarray:
+    ...
+
+  @overload
+  def __getitem__(self, key: np.ndarray) -> 'Table':
+    ...
+
+  @overload
+  def __getitem__(self, key: tuple[str, int | np.integer]) -> TableEntry:
+    ...
+
+  @overload
+  def __getitem__(self, key: tuple[str, np.ndarray]) -> np.ndarray:
+    ...
+
+  def __getitem__(self, key):
+    match key:
+      case str():
+        return self.get_column(key)
+      case np.ndarray() as key_arr_or_mask:
+        return self.apply_array(key_arr_or_mask)
+      case str() as col, int() | np.integer() as key_val:
+        return self.get_value_by_key(col, key_val)
+      case str() as col, np.ndarray() as key_arr_or_mask:
+        return self.apply_array_to_column(col, key_arr_or_mask)
+      case _:
+        if isinstance(key, tuple):
+          err_msg = f'{key}, type: tuple({[type(v) for v in key]})'
+        else:
+          err_msg = f'{key}, type: {type(key)}'
+        raise KeyError(err_msg)
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> dict[str, Any]:
+    """Gets the row with specified key value."""
+    return self.get_row_by_index(
+        self.index_by_key[key], column_name_map=column_name_map
+    )
+
+  def get_row_by_index(
+      self,
+      index: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> dict[str, Any]:
+    """Gets the row at the specified index."""
+    if column_name_map is not None:
+      return {
+          renamed_col: self.get_value_by_index(col, index)
+          for renamed_col, col in column_name_map.items()
+      }
+    else:
+      return {col: self.get_value_by_index(col, index) for col in self.columns}
+
+  def iterrows(
+      self,
+      *,
+      row_keys: np.ndarray | None = None,
+      column_name_map: Mapping[str, str] | None = None,
+      **table_by_foreign_key_col: RowLookup,
+  ) -> Iterator[Mapping[str, Any]]:
+    """Yields rows from the table.
+
+    This can be used to easily convert a table to a Pandas dataframe:
+
+    ```py
+    df = pd.DataFrame(table.iterrows())
+    ```
+
+    Args:
+      row_keys: An optional array of keys of rows to yield. If None, all rows
+        will be yielded.
+      column_name_map: An optional mapping from desired keys in the row dicts to
+        the names of the columns they correspond to.
+      **table_by_foreign_key_col: An optional mapping from column names in this
+        table, which are expected to be columns of foreign keys, to the table
+        that the foreign keys point into. If provided, then the yielded rows
+        will include data from the foreign tables at the appropriate key.
+    """
+    if row_keys is not None:
+      row_indices = self.index_by_key[row_keys]
+    else:
+      row_indices = range(self.size)
+    for i in row_indices:
+      row = self.get_row_by_index(i, column_name_map=column_name_map)
+      for key_col, table in table_by_foreign_key_col.items():
+        foreign_key = self[key_col][i]
+        foreign_row = table.get_row_by_key(foreign_key)
+        row.update(foreign_row)
+      yield row
+
+  def with_column_names(
+      self, column_name_map: Mapping[str, str]
+  ) -> 'RenamedTableView':
+    """Returns a view of this table with mapped column names."""
+    return RenamedTableView(self, column_name_map=column_name_map)
+
+  def make_filter_mask(
+      self,
+      mask: np.ndarray | None = None,
+      *,
+      apply_per_element: bool = False,
+      **predicate_by_col: FilterPredicate,
+  ) -> np.ndarray | None:
+    """Returns a boolean array of rows to keep, or None if all can be kept.
+
+    Args:
+      mask: See `Table.filter`.
+      apply_per_element: See `Table.filter`.
+      **predicate_by_col: See `Table.filter`.
+
+    Returns:
+      Either a boolean NumPy array of length `(self.size,)` denoting which rows
+      should be kept according to the input mask and predicates, or None. None
+      implies there is no filtering required, and is used where possible
+      instead of an all-True array to save time and space.
+    """
+    if mask is None:
+      if not predicate_by_col:
+        return None
+      else:
+        mask = np.ones((self.size,), dtype=bool)
+    else:
+      if mask.shape != (self.size,):
+        raise ValueError(
+            f'mask must have shape ({self.size},). Got: {mask.shape}.'
+        )
+      if mask.dtype != bool:
+        raise ValueError(f'mask must have dtype bool. Got: {mask.dtype}.')
+
+    for col, predicate in predicate_by_col.items():
+      if self[col].ndim > 1:
+        raise ValueError(
+            f'Cannot filter by column {col} with more than 1 dimension.'
+        )
+
+      callable_predicates = []
+      if not callable(predicate):
+        if isinstance(predicate, Iterable) and not isinstance(predicate, str):
+          target_vals = predicate
+        else:
+          target_vals = [predicate]
+        for target_val in target_vals:
+          callable_predicates.append(lambda x, target=target_val: x == target)
+      else:
+        callable_predicates.append(predicate)
+
+      field_mask = np.zeros_like(mask)
+      for callable_predicate in callable_predicates:
+        if not apply_per_element:
+          callable_predicate = typing.cast(
+              Callable[[np.ndarray], bool], callable_predicate
+          )
+          predicate_result = callable_predicate(self.get_column(col))
+        else:
+          predicate_result = np.array(
+              [callable_predicate(elem) for elem in self.get_column(col)]
+          )
+        np.logical_or(field_mask, predicate_result, out=field_mask)
+      np.logical_and(mask, field_mask, out=mask)  # Update in-place.
+    return mask
+
+  def filter(
+      self,
+      mask: np.ndarray | None = None,
+      *,
+      apply_per_element: bool = False,
+      invert: bool = False,
+      **predicate_by_col: FilterPredicate,
+  ) -> Self:
+    """Filters the table using mask and/or predicates and returns a new table.
+
+    Predicates can be either:
+      1. A constant value, e.g. `'CA'`. In this case then only rows that match
+        this value for the given column are retained.
+      2. A (non-string) iterable e.g. `('A', 'B')`. In this
+        case then rows are retained if they match any of the provided values for
+        the given column.
+      3. A boolean function e.g. `lambda b_fac: b_fac < 100.0`.
+        In this case then only rows that evaluate to `True` are retained. By
+        default this function's parameter is expected to be an array, unless
+        `apply_per_element=True`.
+
+    Args:
+      mask: An optional boolean NumPy array with length equal to the table size.
+        If provided then this will be combined with the other predicates so that
+        a row is included if it is masked-in *and* matches all the predicates.
+      apply_per_element: Whether apply predicates to each element in the column
+        individually, or to pass the whole column array to the predicate.
+      invert: If True then the returned table will contain exactly those rows
+        that would be removed if this was `False`.
+      **predicate_by_col: A mapping from column name to a predicate. Filtered
+        columns must be 1D arrays. If multiple columns are provided as keyword
+        arguments then each predicate is applied and the results are combined
+        using a boolean AND operation, so an atom is only retained if it passes
+        all predicates.
+
+    Returns:
+      A new table with the desired rows retained (or filtered out if
+      `invert=True`).
+
+    Raises:
+      ValueError: If mask is provided and is not a bool array with shape
+        `(num_atoms,)`.
+    """
+    filter_mask = self.make_filter_mask(
+        mask, apply_per_element=apply_per_element, **predicate_by_col
+    )
+    if filter_mask is None:
+      # No mask or predicate was specified, so we can return early.
+      if not invert:
+        return self
+      else:
+        return self[np.array((), dtype=np.int64)]
+    else:
+      return self[~filter_mask if invert else filter_mask]
+
+  def _validate_keys_are_column_names(self, keys: Collection[str]) -> None:
+    """Raises an error if any of the keys are not column names."""
+    if mismatches := set(keys) - set(self.columns):
+      raise ValueError(f'Invalid column names: {sorted(mismatches)}.')
+
+  def copy_and_update(self, **new_column_by_column_name: np.ndarray) -> Self:
+    """Returns a copy of this table with the specified changes applied.
+
+    Args:
+      **new_column_by_column_name: New values for the specified columns.
+
+    Raises:
+      ValueError: If a specified column name is not a column in this table.
+    """
+    self._validate_keys_are_column_names(new_column_by_column_name)
+    return dataclasses.replace(self, **new_column_by_column_name)
+
+  def copy_and_remap(
+      self, **mapping_by_col: Mapping[TableEntry, TableEntry]
+  ) -> Self:
+    """Returns a copy of the table with the specified columns remapped.
+
+    Args:
+      **mapping_by_col: Each kwarg key should be the name of one of this table's
+        columns, and each value should be a mapping. The values in the column
+        will be looked up in the mapping and replaced with the result if one is
+        found.
+
+    Raises:
+      ValueError: If a specified column name is not a column in this table.
+    """
+    self._validate_keys_are_column_names(mapping_by_col)
+    if not self.size:
+      return self
+    remapped_cols = {}
+    for column_name, mapping in mapping_by_col.items():
+      col_arr = self.get_column(column_name)
+      if col_arr.dtype == object:
+        remapped = string_array.remap(col_arr, mapping)
+      else:
+        remapped = np.vectorize(lambda x: mapping.get(x, x))(col_arr)  # pylint: disable=cell-var-from-loop
+      remapped_cols[column_name] = remapped
+    return self.copy_and_update(**remapped_cols)
+
+
+class RenamedTableView:
+  """View of a table with renamed column names."""
+
+  def __init__(self, table: Table, column_name_map: Mapping[str, str]):
+    self._table = table
+    self._column_name_map = column_name_map
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> Mapping[str, Any]:
+    del column_name_map
+    return self._table.get_row_by_key(
+        key, column_name_map=self._column_name_map
+    )
+
+
+_DatabaseT = TypeVar('_DatabaseT', bound='Database')
+
+
+class Database(abc.ABC):
+  """Relational database base class."""
+
+  @property
+  @abc.abstractmethod
+  def tables(self) -> Collection[str]:
+    """The names of the tables in this database."""
+
+  @abc.abstractmethod
+  def get_table(self, table_name: str) -> Table:
+    """Gets the table with the given name."""
+
+  @property
+  @abc.abstractmethod
+  def foreign_keys(self) -> Mapping[str, Collection[tuple[str, str]]]:
+    """Describes the relationship between keys in the database.
+
+    Returns:
+      A map from table names to pairs of `(column_name, foreign_table_name)`
+      where `column_name` is a column containing foreign keys in the table named
+      by the key, and the `foreign_table_name` is the name of the table that
+      those foreign keys refer to.
+    """
+
+  @abc.abstractmethod
+  def copy_and_update(
+      self: _DatabaseT,
+      **new_field_by_field_name: ...,
+  ) -> _DatabaseT:
+    """Returns a copy of this database with the specified changes applied."""
+
+
+def table_dependency_order(db: Database) -> Iterable[str]:
+  """Yields the names of the tables in the database in dependency order.
+
+  This order guarantees that a table appears after all other tables that
+  it refers to using foreign keys. Specifically A < B implies that A contains
+  no column that refers to B.key as a foreign key.
+
+  Args:
+    db: The database that defines the table names and foreign keys.
+  """
+  connections: dict[str, set[str]] = {}
+  for table_name in db.tables:
+    connection_set = set()
+    for _, foreign_table in db.foreign_keys.get(table_name, ()):
+      connection_set.add(foreign_table)
+    connections[table_name] = connection_set
+  yield from graphlib.TopologicalSorter(connections).static_order()
+
+
+def concat_databases(dbs: Sequence[_DatabaseT]) -> _DatabaseT:
+  """Concatenates the tables across a sequence of databases.
+
+  Args:
+    dbs: A non-empty sequence of database instances of the same type.
+
+  Returns:
+    A new database containing the concatenated tables from the input databases.
+
+  Raises:
+    ValueError: If `dbs` is empty or `dbs` contains different Database
+      types.
+  """
+  if not dbs:
+    raise ValueError('Need at least one value to concatenate.')
+  distinct_db_types = {type(db) for db in dbs}
+  if len(distinct_db_types) > 1:
+    raise ValueError(
+        f'All `dbs` must be of the same type, got: {distinct_db_types}'
+    )
+
+  first_db, *other_dbs = dbs
+  concatted_tables: dict[str, Table] = {}
+  key_offsets: dict[str, list[int]] = {}
+  for table_name in table_dependency_order(first_db):
+    first_table = first_db.get_table(table_name)
+    columns: dict[str, list[np.ndarray]] = {
+        column_name: [first_table.get_column(column_name)]
+        for column_name in first_table.columns
+    }
+    key_offsets[table_name] = [
+        first_table.key.max() + 1 if first_table.size else 0
+    ]
+
+    for prev_index, db in enumerate(other_dbs):
+      table = db.get_table(table_name)
+      for col_name in table.columns:
+        columns[col_name].append(table.get_column(col_name))
+      key_offset = key_offsets[table_name][prev_index]
+      offset_key = table.key + key_offset
+      columns['key'][-1] = offset_key
+      if table.size:
+        key_offsets[table_name].append(offset_key.max() + 1)
+      else:
+        key_offsets[table_name].append(key_offsets[table_name][prev_index])
+      for fkey_col_name, foreign_table_name in first_db.foreign_keys.get(
+          table_name, []
+      ):
+        fkey_columns = columns[fkey_col_name]
+        fkey_columns[-1] = (
+            fkey_columns[-1] + key_offsets[foreign_table_name][prev_index]
+        )
+
+    concatted_columns = {
+        column_name: np.concatenate(values, axis=-1)
+        for column_name, values in columns.items()
+    }
+    concatted_tables[table_name] = (type(first_table))(**concatted_columns)
+  return first_db.copy_and_update(**concatted_tables)

flax_model/alphafold3/structure/test_utils.py

@@ -0,0 +1,161 @@
+
+
+"""Utilities for structure module testing."""
+
+import dataclasses
+
+from absl.testing import parameterized
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.common.testing import data
+import numpy as np
+import tree
+
+
+class StructureTestCase(parameterized.TestCase):
+  """Testing utilities for working with structure.Structure."""
+
+  def assertAuthorNamingSchemeEqual(self, ans1, ans2):  # pylint: disable=invalid-name
+    """Walks naming scheme, making sure all elements are equal."""
+    if ans1 is None or ans2 is None:
+      self.assertIsNone(ans1)
+      self.assertIsNone(ans2)
+      return
+    flat_ans1 = dict(tree.flatten_with_path(dataclasses.asdict(ans1)))
+    flat_ans2 = dict(tree.flatten_with_path(dataclasses.asdict(ans2)))
+    for k, v in flat_ans1.items():
+      self.assertEqual(v, flat_ans2[k], msg=str(k))
+    for k, v in flat_ans2.items():
+      self.assertEqual(v, flat_ans1[k], msg=str(k))
+
+  def assertAllResiduesEqual(self, all_res1, all_res2):  # pylint: disable=invalid-name
+    """Walks all residues, making sure alll elements are equal."""
+    if all_res1 is None or all_res2 is None:
+      self.assertIsNone(all_res1)
+      self.assertIsNone(all_res2)
+      return
+    self.assertSameElements(all_res1.keys(), all_res2.keys())
+    for chain_id, chain_res in all_res1.items():
+      self.assertSequenceEqual(chain_res, all_res2[chain_id], msg=chain_id)
+
+  def assertBioassemblyDataEqual(self, data1, data2):  # pylint: disable=invalid-name
+    if data1 is None or data2 is None:
+      self.assertIsNone(data1)
+      self.assertIsNone(data2)
+      return
+    self.assertDictEqual(data1.to_mmcif_dict(), data2.to_mmcif_dict())
+
+  def assertChemicalComponentsDataEqual(  # pylint: disable=invalid-name
+      self,
+      data1,
+      data2,
+      allow_chem_comp_data_extension,
+  ):
+    """Checks whether two ChemicalComponentData objects are considered equal."""
+    if data1 is None or data2 is None:
+      self.assertIsNone(data1)
+      self.assertIsNone(data2)
+      return
+    if (not allow_chem_comp_data_extension) or (
+        data1.chem_comp.keys() ^ data2.chem_comp.keys()
+    ):
+      self.assertDictEqual(data1.chem_comp, data2.chem_comp)
+    else:
+      mismatching_values = []
+      for component_id in data1.chem_comp:
+        found = data1.chem_comp[component_id]
+        expected = data2.chem_comp[component_id]
+        if not found.extends(expected):
+          mismatching_values.append((component_id, expected, found))
+
+      if mismatching_values:
+        mismatch_err_msgs = '\n'.join(
+            f'{component_id}: {expected} or its extension expected,'
+            f' but {found} found.'
+            for component_id, expected, found in mismatching_values
+        )
+        self.fail(
+            f'Mismatching values for `_chem_comp` table: {mismatch_err_msgs}',
+        )
+
+  def assertBondsEqual(self, bonds1, bonds2, atom_key1, atom_key2):  # pylint: disable=invalid-name
+    """Checks whether two Bonds objects are considered equal."""
+    # An empty bonds table is functionally equivalent to an empty bonds table.
+    # NB: this can only ever be None in structure v1.
+    if bonds1 is None or not bonds1.size or bonds2 is None or not bonds2.size:
+      self.assertTrue(bonds1 is None or not bonds1.size, msg=f'{bonds1=}')
+      self.assertTrue(bonds2 is None or not bonds2.size, msg=f'{bonds2=}')
+      return
+
+    ptnr1_indices1, ptnr2_indices1 = bonds1.get_atom_indices(atom_key1)
+    ptnr1_indices2, ptnr2_indices2 = bonds2.get_atom_indices(atom_key2)
+    np.testing.assert_array_equal(ptnr1_indices1, ptnr1_indices2)
+    np.testing.assert_array_equal(ptnr2_indices1, ptnr2_indices2)
+    np.testing.assert_array_equal(bonds1.type, bonds2.type)
+    np.testing.assert_array_equal(bonds1.role, bonds2.role)
+
+  def assertStructuresEqual(  # pylint: disable=invalid-name
+      self,
+      struc1,
+      struc2,
+      *,
+      ignore_fields=None,
+      allow_chem_comp_data_extension=False,
+      atol=0,
+  ):
+    """Checks whether two Structure objects could be considered equal.
+
+    Args:
+      struc1: First Structure object.
+      struc2: Second Structure object.
+      ignore_fields: Fields not taken into account during comparison.
+      allow_chem_comp_data_extension: Whether to allow data of `_chem_comp`
+        table to differ if `struc2` is missing some fields, but `struc1` has
+        specific values for them.
+      atol: Absolute tolerance for floating point comparisons (in
+        np.testing.assert_allclose).
+    """
+    for field in sorted(structure.GLOBAL_FIELDS):
+      if ignore_fields and field in ignore_fields:
+        continue
+      if field == 'author_naming_scheme':
+        self.assertAuthorNamingSchemeEqual(struc1[field], struc2[field])
+      elif field == 'all_residues':
+        self.assertAllResiduesEqual(struc1[field], struc2[field])
+      elif field == 'bioassembly_data':
+        self.assertBioassemblyDataEqual(struc1[field], struc2[field])
+      elif field == 'chemical_components_data':
+        self.assertChemicalComponentsDataEqual(
+            struc1[field], struc2[field], allow_chem_comp_data_extension
+        )
+      elif field == 'bonds':
+        self.assertBondsEqual(
+            struc1.bonds, struc2.bonds, struc1.atom_key, struc2.atom_key
+        )
+      else:
+        self.assertEqual(struc1[field], struc2[field], msg=field)
+
+    # The chain order within a structure is arbitrary so in order to
+    # directly compare arrays we first align struc1 to struc2 and check that
+    # the number of atoms doesn't change.
+    num_atoms = struc1.num_atoms
+    self.assertEqual(struc2.num_atoms, num_atoms)
+    struc1 = struc1.order_and_drop_atoms_to_match(struc2)
+    self.assertEqual(struc1.num_atoms, num_atoms)
+
+    for field in sorted(structure.ARRAY_FIELDS):
+      if field == 'atom_key':
+        # atom_key has no external meaning, so it doesn't matter whether it
+        # differs between two structures.
+        continue
+      if ignore_fields and field in ignore_fields:
+        continue
+      self.assertEqual(struc1[field] is None, struc2[field] is None, msg=field)
+
+      if np.issubdtype(struc1[field].dtype, np.inexact):
+        np.testing.assert_allclose(
+            struc1[field], struc2[field], err_msg=field, atol=atol
+        )
+      else:
+        np.testing.assert_array_equal(
+            struc1[field], struc2[field], err_msg=field
+        )

flax_model/alphafold3/test_data/featurised_example.json

@@ -0,0 +1,68 @@
+[
+  {
+    "aatype": "46551530df4e79371352989100ae6af82ccd939009517b788173bdbb323ae40a",
+    "asym_id": "42c714885d17f6c0fcfebebd171b5ed7d7e77a56cf9c4983d9ff9ff90995bd1c",
+    "cleaned_struc": "<<structure>>",
+    "deletion_matrix": "54f980b5b3be8ce80cb6490c527e38d681deade50f239f2cb7d23cf9d0108c37",
+    "deletion_mean": "655a3ef0465a9f30fddf25f4dde0c19a05c6f9069b83961800c1944165955273",
+    "empty_output_struc": "<<structure>>",
+    "entity_id": "42c714885d17f6c0fcfebebd171b5ed7d7e77a56cf9c4983d9ff9ff90995bd1c",
+    "flat_output_layout": "<<atom-layout>>",
+    "frames_mask": "e8a48532fba32c318f7d809fae3ea0999baa01376b29d94ef3c27d82f12f6a90",
+    "is_dna": "1d83518b897b14e2943990eff655838246cc0207a7c95a5f3dfccc2e395f8bbf",
+    "is_ligand": "cfa4e080dc2885de94224db5075c2feb0805567d31e8898abdebaebc6646ce12",
+    "is_nonstandard_polymer_chain": "1d83518b897b14e2943990eff655838246cc0207a7c95a5f3dfccc2e395f8bbf",
+    "is_protein": "569b0e93b5f77c1a9d430b643f29e600c7db6f44ded72b96ab1dd25351d26075",
+    "is_rna": "1d83518b897b14e2943990eff655838246cc0207a7c95a5f3dfccc2e395f8bbf",
+    "is_water": "1d83518b897b14e2943990eff655838246cc0207a7c95a5f3dfccc2e395f8bbf",
+    "ligand_ligand_bonds": "<<None>>",
+    "msa": "871c93669d8a7dd046f8ba6a1ea7b9fc52e1a2e040c8d9bdc4ae73a77f95d3f7",
+    "msa_mask": "13a46a8a2c5305d2a07bb9587ad092d97749d2b2a2d10ad5d64474a4f0f4c803",
+    "num_alignments": "26b25d457597a7b0463f9620f666dd10aa2c4373a505967c7c8d70922a2d6ece",
+    "polymer_ligand_bonds": "<<None>>",
+    "pred_dense_atom_mask": "1e388e40167a4d31291903d74c98544868d0dc574706cdfd259ddd65f56d2993",
+    "profile": "6ec7da33d5f9801357ca6c82725416ef59025c7a8810bfd67fed2262845b480b",
+    "queries_to_keys:gather_idxs": "e0a2fd867fbae573b5613a98ee037eeb575cc5aa3285db2fe5158cda0e8c49cd",
+    "queries_to_keys:gather_mask": "7ddcbfda2904f013d9ab73c55517ac933ecb3dbfbfab9bb351ee43b5fd66e54e",
+    "queries_to_keys:input_shape": "cf801b14d3da833c486e3d73e355c3c2068523012b42c2b17decc45dbd8b9972",
+    "queries_to_token_atoms:gather_idxs": "3c2d0028564167c93b2f33e877974d45756c1e2e8934270c024fc5f1e85fcf14",
+    "queries_to_token_atoms:gather_mask": "1e388e40167a4d31291903d74c98544868d0dc574706cdfd259ddd65f56d2993",
+    "queries_to_token_atoms:input_shape": "cf801b14d3da833c486e3d73e355c3c2068523012b42c2b17decc45dbd8b9972",
+    "ref_atom_name_chars": "043f3242ff4dff04fd36699ccefdff718c57c325f72c71f33f4c4774e75a64dc",
+    "ref_charge": "27898ee238963105942db19f27a980c1fe4f7aeb396c1e5ea4e82da3876deef4",
+    "ref_element": "cc437747adeea094daf8a0e551be02f406c7e7d9b0a9328a84a59d907c5492b5",
+    "ref_mask": "1e388e40167a4d31291903d74c98544868d0dc574706cdfd259ddd65f56d2993",
+    "ref_space_uid": "97afd7d9d3ba835731457ef090ba4939d707bb2a270d48fde56369bde8e58e9f",
+    "residue_center_index": "dfef3704f73a7ee9a3be86b74e25ecd43e056c76574bc9bb65b6b446eeed87a7",
+    "residue_index": "6bf4ea803d0fcd31446e07f6ba9e012f794ced1fe5ad60b7181f66bb92601638",
+    "seq_length": "5a667734987b65ef5c4d07bf176148504bcf8378a5a1a0e23a7fa2c1121cc143",
+    "seq_mask": "a795c048f28bc4307f8599dc5140fe6e971bce6a34a0349b6c18d6cbb197f75f",
+    "sym_id": "5f170600b904a3464c72817993120957c1ce2eec8c2421af02d4912b5e8d8372",
+    "template_aatype": "6c5c3e7a5e23f064f1e8791f61bde8fb59ccf837095144b29e6c8e1607081398",
+    "template_atom_mask": "8a282b7476e13f6de5afbb52c14ea7baa3434281ed77b542bb790287e8c92e82",
+    "template_atom_positions": "b032ac0c813748a81b8d50ce94487c66510a177ad25e84581c7ea5f4b524987f",
+    "token_atoms_layout": "<<atom-layout>>",
+    "token_atoms_to_polymer_ligand_bonds:gather_idxs": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2",
+    "token_atoms_to_polymer_ligand_bonds:gather_mask": "d13d4a8b3b8add19b5970157f09d00c12cbda4fed4d74d8493156523f7069b66",
+    "token_atoms_to_polymer_ligand_bonds:input_shape": "6ed5c8b9711af5b9c7c3d98ba048bed4dfe5974be758f1a632abd746b8278ed2",
+    "token_atoms_to_pseudo_beta:gather_idxs": "3fd8de4dd02450c6e249f33de9c6573694f7991e6a36528d6f32b6c6334a4d4f",
+    "token_atoms_to_pseudo_beta:gather_mask": "a795c048f28bc4307f8599dc5140fe6e971bce6a34a0349b6c18d6cbb197f75f",
+    "token_atoms_to_pseudo_beta:input_shape": "6ed5c8b9711af5b9c7c3d98ba048bed4dfe5974be758f1a632abd746b8278ed2",
+    "token_atoms_to_queries:gather_idxs": "50a1708678bb3d848be6fb49816196b13fcac36891817dac585f4a7556a8310e",
+    "token_atoms_to_queries:gather_mask": "3f4ac4f311cea0c7f6b8464e225815b791a767c9a8d15f9c7fb488c49f674fa9",
+    "token_atoms_to_queries:input_shape": "6ed5c8b9711af5b9c7c3d98ba048bed4dfe5974be758f1a632abd746b8278ed2",
+    "token_index": "77627a5d3856b90a27ae1f5bf5aa14754743c03c00a97bab68f4e7961a405454",
+    "tokens_to_keys:gather_idxs": "bc16a2a2b81bdbb3fb3ff180d984704b201866a7cdc373fb345a9be417aa246e",
+    "tokens_to_keys:gather_mask": "7ddcbfda2904f013d9ab73c55517ac933ecb3dbfbfab9bb351ee43b5fd66e54e",
+    "tokens_to_keys:input_shape": "5545412ee9a21a04712910d53fc1447c73ff94a29ddd1268131b4ed30d1dbe07",
+    "tokens_to_ligand_ligand_bonds:gather_idxs": "bac3003591bd8ef44d85b82cd4d23191d48496c1992e61d3fca2a37e6ddd34ea",
+    "tokens_to_ligand_ligand_bonds:gather_mask": "4d33ec9784aaf54079b4445118c1be26b4a5555e1d6034ccc4182defd86f2068",
+    "tokens_to_ligand_ligand_bonds:input_shape": "5545412ee9a21a04712910d53fc1447c73ff94a29ddd1268131b4ed30d1dbe07",
+    "tokens_to_polymer_ligand_bonds:gather_idxs": "a0ee989ed2a0a2e3626520afa4032e06144865c8c8f6357293c9f4cd2069eaf2",
+    "tokens_to_polymer_ligand_bonds:gather_mask": "d13d4a8b3b8add19b5970157f09d00c12cbda4fed4d74d8493156523f7069b66",
+    "tokens_to_polymer_ligand_bonds:input_shape": "5545412ee9a21a04712910d53fc1447c73ff94a29ddd1268131b4ed30d1dbe07",
+    "tokens_to_queries:gather_idxs": "f82b61cc7f007a744fc9186107647ff333b598642a681b046d61aef4994931e1",
+    "tokens_to_queries:gather_mask": "3f4ac4f311cea0c7f6b8464e225815b791a767c9a8d15f9c7fb488c49f674fa9",
+    "tokens_to_queries:input_shape": "5545412ee9a21a04712910d53fc1447c73ff94a29ddd1268131b4ed30d1dbe07"
+  }
+]

flax_model/alphafold3/test_data/miniature_databases/pdb_mmcif/5y2e.cif

@@ -0,0 +1,562 @@
+data_5Y2E
+#
+_entry.id 5Y2E
+#
+loop_
+_chem_comp.formula
+_chem_comp.formula_weight
+_chem_comp.id
+_chem_comp.mon_nstd_flag
+_chem_comp.name
+_chem_comp.pdbx_synonyms
+_chem_comp.type
+"C3 H7 N O2"     89.093  ALA y ALANINE         ? "L-peptide linking" 
+"C6 H15 N4 O2 1" 175.209 ARG y ARGININE        ? "L-peptide linking" 
+"C4 H7 N O4"     133.103 ASP y "ASPARTIC ACID" ? "L-peptide linking" 
+"C5 H10 N2 O3"   146.144 GLN y GLUTAMINE       ? "L-peptide linking" 
+"C5 H9 N O4"     147.129 GLU y "GLUTAMIC ACID" ? "L-peptide linking" 
+"C2 H5 N O2"     75.067  GLY y GLYCINE         ? "peptide linking"   
+"C6 H10 N3 O2 1" 156.162 HIS y HISTIDINE       ? "L-peptide linking" 
+"H2 O"           18.015  HOH . WATER           ? non-polymer         
+"C6 H13 N O2"    131.173 ILE y ISOLEUCINE      ? "L-peptide linking" 
+"C6 H13 N O2"    131.173 LEU y LEUCINE         ? "L-peptide linking" 
+"C6 H15 N2 O2 1" 147.195 LYS y LYSINE          ? "L-peptide linking" 
+"C5 H11 N O2 S"  149.211 MET y METHIONINE      ? "L-peptide linking" 
+"C4 H9 N O3"     119.119 THR y THREONINE       ? "L-peptide linking" 
+"C5 H11 N O2"    117.146 VAL y VALINE          ? "L-peptide linking" 
+#
+_entity.id               1
+_entity.pdbx_description "Non-structural glycoprotein 4"
+_entity.type             polymer
+#
+_entity_poly.entity_id      1
+_entity_poly.pdbx_strand_id A
+_entity_poly.type           polypeptide(L)
+#
+loop_
+_entity_poly_seq.entity_id
+_entity_poly_seq.hetero
+_entity_poly_seq.mon_id
+_entity_poly_seq.num
+1 n MET 1  
+1 n ILE 2  
+1 n GLU 3  
+1 n LYS 4  
+1 n GLN 5  
+1 n MET 6  
+1 n ASP 7  
+1 n ARG 8  
+1 n VAL 9  
+1 n VAL 10 
+1 n LYS 11 
+1 n GLU 12 
+1 n MET 13 
+1 n ARG 14 
+1 n ARG 15 
+1 n GLN 16 
+1 n LEU 17 
+1 n GLU 18 
+1 n MET 19 
+1 n ILE 20 
+1 n ASP 21 
+1 n LYS 22 
+1 n LEU 23 
+1 n THR 24 
+1 n THR 25 
+1 n ARG 26 
+1 n GLY 27 
+1 n ILE 28 
+1 n GLU 29 
+1 n GLN 30 
+1 n VAL 31 
+1 n GLU 32 
+1 n LEU 33 
+1 n LEU 34 
+1 n LYS 35 
+1 n ARG 36 
+1 n ILE 37 
+1 n HIS 38 
+1 n ASP 39 
+1 n LYS 40 
+1 n LEU 41 
+1 n MET 42 
+1 n ILE 43 
+1 n ARG 44 
+1 n ALA 45 
+1 n VAL 46 
+1 n ASP 47 
+#
+_exptl.method "X-RAY DIFFRACTION"
+#
+_pdbx_audit_revision_history.revision_date 2018-03-14
+#
+_pdbx_database_status.recvd_initial_deposition_date 2018-03-14
+#
+loop_
+_pdbx_poly_seq_scheme.asym_id
+_pdbx_poly_seq_scheme.auth_seq_num
+_pdbx_poly_seq_scheme.entity_id
+_pdbx_poly_seq_scheme.hetero
+_pdbx_poly_seq_scheme.mon_id
+_pdbx_poly_seq_scheme.pdb_ins_code
+_pdbx_poly_seq_scheme.pdb_seq_num
+_pdbx_poly_seq_scheme.pdb_strand_id
+_pdbx_poly_seq_scheme.seq_id
+B 94  1 n MET . 94  A 1  
+B 95  1 n ILE . 95  A 2  
+B 96  1 n GLU . 96  A 3  
+B 97  1 n LYS . 97  A 4  
+B 98  1 n GLN . 98  A 5  
+B 99  1 n MET . 99  A 6  
+B 100 1 n ASP . 100 A 7  
+B 101 1 n ARG . 101 A 8  
+B 102 1 n VAL . 102 A 9  
+B 103 1 n VAL . 103 A 10 
+B 104 1 n LYS . 104 A 11 
+B 105 1 n GLU . 105 A 12 
+B 106 1 n MET . 106 A 13 
+B 107 1 n ARG . 107 A 14 
+B 108 1 n ARG . 108 A 15 
+B 109 1 n GLN . 109 A 16 
+B 110 1 n LEU . 110 A 17 
+B 111 1 n GLU . 111 A 18 
+B 112 1 n MET . 112 A 19 
+B 113 1 n ILE . 113 A 20 
+B 114 1 n ASP . 114 A 21 
+B 115 1 n LYS . 115 A 22 
+B 116 1 n LEU . 116 A 23 
+B 117 1 n THR . 117 A 24 
+B 118 1 n THR . 118 A 25 
+B 119 1 n ARG . 119 A 26 
+B 120 1 n GLY . 120 A 27 
+B 121 1 n ILE . 121 A 28 
+B 122 1 n GLU . 122 A 29 
+B 123 1 n GLN . 123 A 30 
+B 124 1 n VAL . 124 A 31 
+B 125 1 n GLU . 125 A 32 
+B 126 1 n LEU . 126 A 33 
+B 127 1 n LEU . 127 A 34 
+B 128 1 n LYS . 128 A 35 
+B 129 1 n ARG . 129 A 36 
+B 130 1 n ILE . 130 A 37 
+B 131 1 n HIS . 131 A 38 
+B 132 1 n ASP . 132 A 39 
+B 133 1 n LYS . 133 A 40 
+B 134 1 n LEU . 134 A 41 
+B 135 1 n MET . 135 A 42 
+B 136 1 n ILE . 136 A 43 
+B 137 1 n ARG . 137 A 44 
+B 138 1 n ALA . 138 A 45 
+B ?   1 n VAL . 139 A 46 
+B ?   1 n ASP . 140 A 47 
+#
+_pdbx_struct_assembly.details            author_and_software_defined_assembly
+_pdbx_struct_assembly.id                 1
+_pdbx_struct_assembly.method_details     PISA
+_pdbx_struct_assembly.oligomeric_count   4
+_pdbx_struct_assembly.oligomeric_details tetrameric
+#
+_pdbx_struct_assembly_gen.assembly_id     1
+_pdbx_struct_assembly_gen.asym_id_list    A,B,C,D,E,F,G
+_pdbx_struct_assembly_gen.oper_expression 1
+#
+_pdbx_struct_oper_list.id                 1
+_pdbx_struct_oper_list.matrix[1][1]       1.0000000000
+_pdbx_struct_oper_list.matrix[1][2]       0.0000000000
+_pdbx_struct_oper_list.matrix[1][3]       0.0000000000
+_pdbx_struct_oper_list.matrix[2][1]       0.0000000000
+_pdbx_struct_oper_list.matrix[2][2]       1.0000000000
+_pdbx_struct_oper_list.matrix[2][3]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][1]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][2]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][3]       1.0000000000
+_pdbx_struct_oper_list.name               1_555
+_pdbx_struct_oper_list.symmetry_operation x,y,z
+_pdbx_struct_oper_list.type               "identity operation"
+_pdbx_struct_oper_list.vector[1]          0.0000000000
+_pdbx_struct_oper_list.vector[2]          0.0000000000
+_pdbx_struct_oper_list.vector[3]          0.0000000000
+#
+_refine.ls_d_res_high 2.70
+#
+_software.classification other
+_software.name           "DeepMind Structure Class"
+_software.pdbx_ordinal   1
+_software.version        2.0.0
+#
+_struct_asym.entity_id 1
+_struct_asym.id        B
+#
+loop_
+_atom_site.group_PDB
+_atom_site.id
+_atom_site.type_symbol
+_atom_site.label_atom_id
+_atom_site.label_alt_id
+_atom_site.label_comp_id
+_atom_site.label_asym_id
+_atom_site.label_entity_id
+_atom_site.label_seq_id
+_atom_site.pdbx_PDB_ins_code
+_atom_site.Cartn_x
+_atom_site.Cartn_y
+_atom_site.Cartn_z
+_atom_site.occupancy
+_atom_site.B_iso_or_equiv
+_atom_site.auth_seq_id
+_atom_site.auth_asym_id
+_atom_site.pdbx_PDB_model_num
+ATOM 1   N N   . MET B 1 1  ? 2.644   6.839  -46.028 1.00 83.71  94  A 1 
+ATOM 2   C CA  . MET B 1 1  ? 1.400   7.612  -45.825 1.00 86.32  94  A 1 
+ATOM 3   C C   . MET B 1 1  ? 1.210   7.971  -44.344 1.00 87.15  94  A 1 
+ATOM 4   O O   . MET B 1 1  ? 1.479   7.188  -43.421 1.00 80.48  94  A 1 
+ATOM 5   C CB  . MET B 1 1  ? 0.170   6.890  -46.382 1.00 88.07  94  A 1 
+ATOM 6   C CG  . MET B 1 1  ? -0.124  7.275  -47.840 1.00 91.84  94  A 1 
+ATOM 7   S SD  . MET B 1 1  ? -1.789  6.807  -48.325 1.00 103.69 94  A 1 
+ATOM 8   C CE  . MET B 1 1  ? -2.264  8.027  -49.569 1.00 97.89  94  A 1 
+ATOM 9   N N   . ILE B 1 2  ? 0.696   9.166  -44.185 1.00 87.93  95  A 1 
+ATOM 10  C CA  . ILE B 1 2  ? 0.364   9.779  -42.948 1.00 84.67  95  A 1 
+ATOM 11  C C   . ILE B 1 2  ? -0.678  8.946  -42.258 1.00 75.98  95  A 1 
+ATOM 12  O O   . ILE B 1 2  ? -0.634  8.714  -41.080 1.00 67.98  95  A 1 
+ATOM 13  C CB  . ILE B 1 2  ? -0.321  11.082 -43.309 1.00 89.99  95  A 1 
+ATOM 14  C CG1 . ILE B 1 2  ? 0.696   12.146 -43.577 1.00 94.59  95  A 1 
+ATOM 15  C CG2 . ILE B 1 2  ? -1.197  11.546 -42.182 1.00 92.79  95  A 1 
+ATOM 16  C CD1 . ILE B 1 2  ? 1.132   12.780 -42.294 1.00 93.36  95  A 1 
+ATOM 17  N N   . GLU B 1 3  ? -1.625  8.499  -43.046 1.00 71.66  96  A 1 
+ATOM 18  C CA  . GLU B 1 3  ? -2.745  7.716  -42.585 1.00 68.44  96  A 1 
+ATOM 19  C C   . GLU B 1 3  ? -2.273  6.463  -41.883 1.00 69.98  96  A 1 
+ATOM 20  O O   . GLU B 1 3  ? -2.890  6.024  -40.946 1.00 67.88  96  A 1 
+ATOM 21  C CB  . GLU B 1 3  ? -3.613  7.332  -43.745 1.00 73.44  96  A 1 
+ATOM 22  C CG  . GLU B 1 3  ? -4.382  8.485  -44.320 1.00 87.42  96  A 1 
+ATOM 23  C CD  . GLU B 1 3  ? -3.735  9.039  -45.551 1.00 89.38  96  A 1 
+ATOM 24  O OE1 . GLU B 1 3  ? -4.454  9.413  -46.477 1.00 97.81  96  A 1 
+ATOM 25  O OE2 . GLU B 1 3  ? -2.511  9.105  -45.593 1.00 88.14  96  A 1 
+ATOM 26  N N   . LYS B 1 4  ? -1.212  5.844  -42.364 1.00 71.45  97  A 1 
+ATOM 27  C CA  . LYS B 1 4  ? -0.676  4.669  -41.671 1.00 72.76  97  A 1 
+ATOM 28  C C   . LYS B 1 4  ? 0.097   5.062  -40.408 1.00 74.11  97  A 1 
+ATOM 29  O O   . LYS B 1 4  ? 0.209   4.268  -39.462 1.00 77.84  97  A 1 
+ATOM 30  C CB  . LYS B 1 4  ? 0.197   3.807  -42.578 1.00 66.63  97  A 1 
+ATOM 31  N N   . GLN B 1 5  ? 0.624   6.270  -40.364 1.00 72.03  98  A 1 
+ATOM 32  C CA  . GLN B 1 5  ? 1.295   6.749  -39.158 1.00 65.60  98  A 1 
+ATOM 33  C C   . GLN B 1 5  ? 0.255   7.097  -38.099 1.00 58.64  98  A 1 
+ATOM 34  O O   . GLN B 1 5  ? 0.462   6.888  -36.942 1.00 52.26  98  A 1 
+ATOM 35  C CB  . GLN B 1 5  ? 2.143   7.976  -39.487 1.00 63.76  98  A 1 
+ATOM 36  N N   . MET B 1 6  ? -0.861  7.648  -38.542 1.00 49.04  99  A 1 
+ATOM 37  C CA  . MET B 1 6  ? -1.953  7.954  -37.679 1.00 45.66  99  A 1 
+ATOM 38  C C   . MET B 1 6  ? -2.486  6.679  -37.100 1.00 50.63  99  A 1 
+ATOM 39  O O   . MET B 1 6  ? -2.882  6.643  -35.948 1.00 50.11  99  A 1 
+ATOM 40  C CB  . MET B 1 6  ? -3.050  8.640  -38.449 1.00 42.97  99  A 1 
+ATOM 41  C CG  . MET B 1 6  ? -4.395  8.703  -37.753 1.00 42.67  99  A 1 
+ATOM 42  S SD  . MET B 1 6  ? -5.430  9.916  -38.524 1.00 47.21  99  A 1 
+ATOM 43  C CE  . MET B 1 6  ? -6.803  9.938  -37.397 1.00 48.85  99  A 1 
+ATOM 44  N N   . ASP B 1 7  ? -2.531  5.622  -37.894 1.00 56.67  100 A 1 
+ATOM 45  C CA  . ASP B 1 7  ? -3.070  4.383  -37.401 1.00 55.17  100 A 1 
+ATOM 46  C C   . ASP B 1 7  ? -2.223  3.846  -36.257 1.00 53.54  100 A 1 
+ATOM 47  O O   . ASP B 1 7  ? -2.785  3.365  -35.256 1.00 47.69  100 A 1 
+ATOM 48  C CB  . ASP B 1 7  ? -3.067  3.354  -38.520 1.00 57.94  100 A 1 
+ATOM 49  C CG  . ASP B 1 7  ? -3.936  2.164  -38.236 1.00 62.42  100 A 1 
+ATOM 50  O OD1 . ASP B 1 7  ? -5.111  2.416  -37.949 1.00 67.59  100 A 1 
+ATOM 51  O OD2 . ASP B 1 7  ? -3.401  1.020  -38.206 1.00 61.64  100 A 1 
+ATOM 52  N N   . ARG B 1 8  ? -0.900  3.909  -36.428 1.00 54.78  101 A 1 
+ATOM 53  C CA  . ARG B 1 8  ? -0.002  3.384  -35.432 1.00 61.23  101 A 1 
+ATOM 54  C C   . ARG B 1 8  ? -0.168  4.177  -34.127 1.00 63.20  101 A 1 
+ATOM 55  O O   . ARG B 1 8  ? -0.017  3.620  -33.077 1.00 67.94  101 A 1 
+ATOM 56  C CB  . ARG B 1 8  ? 1.484   3.327  -35.943 1.00 64.73  101 A 1 
+ATOM 57  C CG  . ARG B 1 8  ? 2.403   4.541  -35.635 1.00 73.16  101 A 1 
+ATOM 58  N N   . VAL B 1 9  ? -0.406  5.483  -34.222 1.00 63.14  102 A 1 
+ATOM 59  C CA  . VAL B 1 9  ? -0.546  6.334  -33.061 1.00 59.81  102 A 1 
+ATOM 60  C C   . VAL B 1 9  ? -1.853  6.121  -32.347 1.00 55.82  102 A 1 
+ATOM 61  O O   . VAL B 1 9  ? -1.880  6.066  -31.150 1.00 48.46  102 A 1 
+ATOM 62  C CB  . VAL B 1 9  ? -0.421  7.813  -33.454 1.00 60.39  102 A 1 
+ATOM 63  C CG1 . VAL B 1 9  ? -0.731  8.703  -32.272 1.00 60.03  102 A 1 
+ATOM 64  C CG2 . VAL B 1 9  ? 0.994   8.091  -33.943 1.00 61.68  102 A 1 
+ATOM 65  N N   . VAL B 1 10 ? -2.941  6.014  -33.079 1.00 56.18  103 A 1 
+ATOM 66  C CA  . VAL B 1 10 ? -4.231  5.743  -32.481 1.00 53.53  103 A 1 
+ATOM 67  C C   . VAL B 1 10 ? -4.236  4.419  -31.704 1.00 52.90  103 A 1 
+ATOM 68  O O   . VAL B 1 10 ? -4.772  4.355  -30.607 1.00 49.99  103 A 1 
+ATOM 69  C CB  . VAL B 1 10 ? -5.317  5.742  -33.553 1.00 51.75  103 A 1 
+ATOM 70  C CG1 . VAL B 1 10 ? -6.606  5.137  -33.045 1.00 52.18  103 A 1 
+ATOM 71  C CG2 . VAL B 1 10 ? -5.610  7.153  -33.983 1.00 52.67  103 A 1 
+ATOM 72  N N   . LYS B 1 11 ? -3.601  3.402  -32.243 1.00 55.55  104 A 1 
+ATOM 73  C CA  . LYS B 1 11 ? -3.521  2.104  -31.626 1.00 56.87  104 A 1 
+ATOM 74  C C   . LYS B 1 11 ? -2.709  2.164  -30.350 1.00 60.76  104 A 1 
+ATOM 75  O O   . LYS B 1 11 ? -2.986  1.464  -29.399 1.00 58.80  104 A 1 
+ATOM 76  C CB  . LYS B 1 11 ? -2.837  1.155  -32.582 1.00 60.72  104 A 1 
+ATOM 77  C CG  . LYS B 1 11 ? -3.729  0.487  -33.601 1.00 72.80  104 A 1 
+ATOM 78  C CD  . LYS B 1 11 ? -2.915  -0.472 -34.451 1.00 79.85  104 A 1 
+ATOM 79  C CE  . LYS B 1 11 ? -3.664  -0.963 -35.666 1.00 84.90  104 A 1 
+ATOM 80  N NZ  . LYS B 1 11 ? -5.110  -0.693 -35.513 1.00 94.50  104 A 1 
+ATOM 81  N N   . GLU B 1 12 ? -1.674  2.982  -30.369 1.00 56.70  105 A 1 
+ATOM 82  C CA  . GLU B 1 12 ? -0.798  3.204  -29.249 1.00 54.71  105 A 1 
+ATOM 83  C C   . GLU B 1 12 ? -1.507  3.946  -28.155 1.00 52.45  105 A 1 
+ATOM 84  O O   . GLU B 1 12 ? -1.387  3.608  -27.007 1.00 51.88  105 A 1 
+ATOM 85  C CB  . GLU B 1 12 ? 0.404   4.036  -29.667 1.00 60.63  105 A 1 
+ATOM 86  C CG  . GLU B 1 12 ? 1.418   4.261  -28.569 1.00 71.75  105 A 1 
+ATOM 87  C CD  . GLU B 1 12 ? 2.100   2.985  -28.105 1.00 83.66  105 A 1 
+ATOM 88  O OE1 . GLU B 1 12 ? 1.973   1.951  -28.793 1.00 97.71  105 A 1 
+ATOM 89  O OE2 . GLU B 1 12 ? 2.729   3.008  -27.039 1.00 91.62  105 A 1 
+ATOM 90  N N   . MET B 1 13 ? -2.250  4.972  -28.495 1.00 50.99  106 A 1 
+ATOM 91  C CA  . MET B 1 13 ? -2.993  5.703  -27.508 1.00 53.07  106 A 1 
+ATOM 92  C C   . MET B 1 13 ? -4.045  4.874  -26.849 1.00 51.54  106 A 1 
+ATOM 93  O O   . MET B 1 13 ? -4.220  4.974  -25.651 1.00 48.88  106 A 1 
+ATOM 94  C CB  . MET B 1 13 ? -3.640  6.915  -28.107 1.00 61.60  106 A 1 
+ATOM 95  C CG  . MET B 1 13 ? -2.651  8.032  -28.402 1.00 65.35  106 A 1 
+ATOM 96  S SD  . MET B 1 13 ? -3.551  9.590  -28.498 1.00 76.88  106 A 1 
+ATOM 97  C CE  . MET B 1 13 ? -3.620  9.931  -26.737 1.00 76.03  106 A 1 
+ATOM 98  N N   . ARG B 1 14 ? -4.714  4.016  -27.571 1.00 53.52  107 A 1 
+ATOM 99  C CA  . ARG B 1 14 ? -5.703  3.187  -26.952 1.00 50.53  107 A 1 
+ATOM 100 C C   . ARG B 1 14 ? -5.040  2.268  -25.975 1.00 51.01  107 A 1 
+ATOM 101 O O   . ARG B 1 14 ? -5.527  2.040  -24.903 1.00 48.89  107 A 1 
+ATOM 102 C CB  . ARG B 1 14 ? -6.358  2.367  -28.011 1.00 54.79  107 A 1 
+ATOM 103 C CG  . ARG B 1 14 ? -7.580  1.648  -27.566 1.00 60.71  107 A 1 
+ATOM 104 C CD  . ARG B 1 14 ? -7.895  0.586  -28.570 1.00 70.83  107 A 1 
+ATOM 105 N NE  . ARG B 1 14 ? -8.689  1.116  -29.658 1.00 78.79  107 A 1 
+ATOM 106 C CZ  . ARG B 1 14 ? -8.278  1.182  -30.911 1.00 81.77  107 A 1 
+ATOM 107 N NH1 . ARG B 1 14 ? -7.086  0.746  -31.233 1.00 91.09  107 A 1 
+ATOM 108 N NH2 . ARG B 1 14 ? -9.071  1.675  -31.832 1.00 79.36  107 A 1 
+ATOM 109 N N   . ARG B 1 15 ? -3.896  1.754  -26.355 1.00 46.73  108 A 1 
+ATOM 110 C CA  . ARG B 1 15 ? -3.174  0.827  -25.534 1.00 45.05  108 A 1 
+ATOM 111 C C   . ARG B 1 15 ? -2.681  1.503  -24.248 1.00 47.85  108 A 1 
+ATOM 112 O O   . ARG B 1 15 ? -2.752  0.930  -23.189 1.00 50.40  108 A 1 
+ATOM 113 C CB  . ARG B 1 15 ? -2.063  0.286  -26.404 1.00 50.23  108 A 1 
+ATOM 114 C CG  . ARG B 1 15 ? -1.165  -0.750 -25.876 1.00 59.06  108 A 1 
+ATOM 115 C CD  . ARG B 1 15 ? -0.358  -1.317 -27.046 1.00 59.59  108 A 1 
+ATOM 116 N NE  . ARG B 1 15 ? 0.861   -1.991 -26.583 1.00 65.92  108 A 1 
+ATOM 117 N N   . GLN B 1 16 ? -2.221  2.739  -24.349 1.00 49.74  109 A 1 
+ATOM 118 C CA  . GLN B 1 16 ? -1.810  3.511  -23.193 1.00 47.85  109 A 1 
+ATOM 119 C C   . GLN B 1 16 ? -2.993  3.774  -22.279 1.00 48.41  109 A 1 
+ATOM 120 O O   . GLN B 1 16 ? -2.898  3.560  -21.091 1.00 46.33  109 A 1 
+ATOM 121 C CB  . GLN B 1 16 ? -1.209  4.809  -23.611 1.00 49.18  109 A 1 
+ATOM 122 C CG  . GLN B 1 16 ? 0.208   4.681  -24.112 1.00 53.75  109 A 1 
+ATOM 123 C CD  . GLN B 1 16 ? 0.618   5.857  -24.974 1.00 61.24  109 A 1 
+ATOM 124 O OE1 . GLN B 1 16 ? -0.145  6.791  -25.159 1.00 69.70  109 A 1 
+ATOM 125 N NE2 . GLN B 1 16 ? 1.786   5.782  -25.546 1.00 63.00  109 A 1 
+ATOM 126 N N   . LEU B 1 17 ? -4.098  4.266  -22.817 1.00 45.61  110 A 1 
+ATOM 127 C CA  . LEU B 1 17 ? -5.278  4.506  -22.010 1.00 41.80  110 A 1 
+ATOM 128 C C   . LEU B 1 17 ? -5.848  3.244  -21.329 1.00 44.01  110 A 1 
+ATOM 129 O O   . LEU B 1 17 ? -6.378  3.314  -20.232 1.00 49.63  110 A 1 
+ATOM 130 C CB  . LEU B 1 17 ? -6.380  5.133  -22.818 1.00 41.83  110 A 1 
+ATOM 131 C CG  . LEU B 1 17 ? -6.545  6.610  -23.075 1.00 44.57  110 A 1 
+ATOM 132 C CD1 . LEU B 1 17 ? -6.481  7.448  -21.815 1.00 46.53  110 A 1 
+ATOM 133 C CD2 . LEU B 1 17 ? -5.534  7.076  -24.054 1.00 46.70  110 A 1 
+ATOM 134 N N   . GLU B 1 18 ? -5.724  2.085  -21.974 1.00 48.92  111 A 1 
+ATOM 135 C CA  . GLU B 1 18 ? -6.176  0.854  -21.354 1.00 50.45  111 A 1 
+ATOM 136 C C   . GLU B 1 18 ? -5.278  0.446  -20.173 1.00 50.00  111 A 1 
+ATOM 137 O O   . GLU B 1 18 ? -5.748  -0.088 -19.155 1.00 52.97  111 A 1 
+ATOM 138 C CB  . GLU B 1 18 ? -6.361  -0.236 -22.405 1.00 54.05  111 A 1 
+ATOM 139 C CG  . GLU B 1 18 ? -7.693  -0.095 -23.112 1.00 59.57  111 A 1 
+ATOM 140 C CD  . GLU B 1 18 ? -7.804  -0.858 -24.438 1.00 68.03  111 A 1 
+ATOM 141 O OE1 . GLU B 1 18 ? -6.834  -1.521 -24.893 1.00 73.96  111 A 1 
+ATOM 142 O OE2 . GLU B 1 18 ? -8.896  -0.808 -25.069 1.00 73.24  111 A 1 
+ATOM 143 N N   . MET B 1 19 ? -4.000  0.760  -20.284 1.00 48.79  112 A 1 
+ATOM 144 C CA  . MET B 1 19 ? -3.094  0.512  -19.186 1.00 52.15  112 A 1 
+ATOM 145 C C   . MET B 1 19 ? -3.371  1.430  -18.007 1.00 47.91  112 A 1 
+ATOM 146 O O   . MET B 1 19 ? -3.429  0.986  -16.872 1.00 45.58  112 A 1 
+ATOM 147 C CB  . MET B 1 19 ? -1.653  0.667  -19.617 1.00 54.11  112 A 1 
+ATOM 148 C CG  . MET B 1 19 ? -1.053  -0.654 -20.052 1.00 58.96  112 A 1 
+ATOM 149 S SD  . MET B 1 19 ? 0.471   -0.323 -20.923 1.00 65.09  112 A 1 
+ATOM 150 C CE  . MET B 1 19 ? 1.501   -0.237 -19.466 1.00 67.56  112 A 1 
+ATOM 151 N N   . ILE B 1 20 ? -3.541  2.711  -18.297 1.00 43.25  113 A 1 
+ATOM 152 C CA  . ILE B 1 20 ? -3.900  3.677  -17.287 1.00 40.09  113 A 1 
+ATOM 153 C C   . ILE B 1 20 ? -5.165  3.240  -16.596 1.00 40.28  113 A 1 
+ATOM 154 O O   . ILE B 1 20 ? -5.297  3.383  -15.395 1.00 41.67  113 A 1 
+ATOM 155 C CB  . ILE B 1 20 ? -4.079  5.077  -17.885 1.00 41.26  113 A 1 
+ATOM 156 C CG1 . ILE B 1 20 ? -2.700  5.633  -18.275 1.00 39.56  113 A 1 
+ATOM 157 C CG2 . ILE B 1 20 ? -4.715  6.029  -16.880 1.00 44.72  113 A 1 
+ATOM 158 C CD1 . ILE B 1 20 ? -2.722  6.963  -19.009 1.00 37.40  113 A 1 
+ATOM 159 N N   . ASP B 1 21 ? -6.096  2.680  -17.346 1.00 41.39  114 A 1 
+ATOM 160 C CA  . ASP B 1 21 ? -7.302  2.201  -16.728 1.00 46.35  114 A 1 
+ATOM 161 C C   . ASP B 1 21 ? -7.074  1.044  -15.762 1.00 45.76  114 A 1 
+ATOM 162 O O   . ASP B 1 21 ? -7.662  1.035  -14.682 1.00 43.56  114 A 1 
+ATOM 163 C CB  . ASP B 1 21 ? -8.272  1.734  -17.759 1.00 51.71  114 A 1 
+ATOM 164 C CG  . ASP B 1 21 ? -9.647  1.640  -17.208 1.00 60.74  114 A 1 
+ATOM 165 O OD1 . ASP B 1 21 ? -10.153 2.682  -16.742 1.00 68.50  114 A 1 
+ATOM 166 O OD2 . ASP B 1 21 ? -10.199 0.528  -17.210 1.00 73.36  114 A 1 
+ATOM 167 N N   . LYS B 1 22 ? -6.258  0.070  -16.178 1.00 47.51  115 A 1 
+ATOM 168 C CA  . LYS B 1 22 ? -5.942  -1.062 -15.321 1.00 50.59  115 A 1 
+ATOM 169 C C   . LYS B 1 22 ? -5.255  -0.563 -14.028 1.00 52.11  115 A 1 
+ATOM 170 O O   . LYS B 1 22 ? -5.648  -0.941 -12.919 1.00 55.27  115 A 1 
+ATOM 171 C CB  . LYS B 1 22 ? -5.085  -2.111 -16.025 1.00 52.86  115 A 1 
+ATOM 172 C CG  . LYS B 1 22 ? -5.829  -3.010 -17.008 1.00 57.80  115 A 1 
+ATOM 173 N N   . LEU B 1 23 ? -4.266  0.312  -14.176 1.00 48.62  116 A 1 
+ATOM 174 C CA  . LEU B 1 23 ? -3.594  0.880  -13.039 1.00 44.80  116 A 1 
+ATOM 175 C C   . LEU B 1 23 ? -4.511  1.665  -12.106 1.00 44.04  116 A 1 
+ATOM 176 O O   . LEU B 1 23 ? -4.412  1.565  -10.897 1.00 41.22  116 A 1 
+ATOM 177 C CB  . LEU B 1 23 ? -2.492  1.781  -13.511 1.00 43.26  116 A 1 
+ATOM 178 C CG  . LEU B 1 23 ? -1.353  1.039  -14.177 1.00 45.18  116 A 1 
+ATOM 179 C CD1 . LEU B 1 23 ? -0.449  2.064  -14.839 1.00 45.01  116 A 1 
+ATOM 180 C CD2 . LEU B 1 23 ? -0.582  0.210  -13.156 1.00 45.90  116 A 1 
+ATOM 181 N N   . THR B 1 24 ? -5.402  2.445  -12.674 1.00 49.17  117 A 1 
+ATOM 182 C CA  . THR B 1 24 ? -6.382  3.155  -11.886 1.00 49.60  117 A 1 
+ATOM 183 C C   . THR B 1 24 ? -7.250  2.198  -11.085 1.00 51.60  117 A 1 
+ATOM 184 O O   . THR B 1 24 ? -7.503  2.441  -9.918  1.00 50.82  117 A 1 
+ATOM 185 C CB  . THR B 1 24 ? -7.248  4.024  -12.811 1.00 54.67  117 A 1 
+ATOM 186 O OG1 . THR B 1 24 ? -6.463  5.126  -13.253 1.00 58.01  117 A 1 
+ATOM 187 C CG2 . THR B 1 24 ? -8.464  4.559  -12.141 1.00 57.77  117 A 1 
+ATOM 188 N N   . THR B 1 25 ? -7.718  1.116  -11.704 1.00 53.09  118 A 1 
+ATOM 189 C CA  . THR B 1 25 ? -8.495  0.117  -10.981 1.00 50.55  118 A 1 
+ATOM 190 C C   . THR B 1 25 ? -7.709  -0.503 -9.820  1.00 49.50  118 A 1 
+ATOM 191 O O   . THR B 1 25 ? -8.267  -0.695 -8.720  1.00 50.25  118 A 1 
+ATOM 192 C CB  . THR B 1 25 ? -8.933  -0.983 -11.947 1.00 53.52  118 A 1 
+ATOM 193 O OG1 . THR B 1 25 ? -9.796  -0.396 -12.916 1.00 60.17  118 A 1 
+ATOM 194 C CG2 . THR B 1 25 ? -9.619  -2.127 -11.260 1.00 58.18  118 A 1 
+ATOM 195 N N   . ARG B 1 26 ? -6.447  -0.762 -10.054 1.00 48.62  119 A 1 
+ATOM 196 C CA  . ARG B 1 26 ? -5.575  -1.319 -9.076  1.00 52.29  119 A 1 
+ATOM 197 C C   . ARG B 1 26 ? -5.362  -0.349 -7.936  1.00 56.15  119 A 1 
+ATOM 198 O O   . ARG B 1 26 ? -5.203  -0.758 -6.818  1.00 60.08  119 A 1 
+ATOM 199 C CB  . ARG B 1 26 ? -4.262  -1.609 -9.738  1.00 57.25  119 A 1 
+ATOM 200 C CG  . ARG B 1 26 ? -4.341  -2.608 -10.845 1.00 67.09  119 A 1 
+ATOM 201 C CD  . ARG B 1 26 ? -4.671  -3.974 -10.307 1.00 78.77  119 A 1 
+ATOM 202 N NE  . ARG B 1 26 ? -4.743  -4.923 -11.390 1.00 95.96  119 A 1 
+ATOM 203 C CZ  . ARG B 1 26 ? -3.703  -5.558 -11.900 1.00 111.57 119 A 1 
+ATOM 204 N NH1 . ARG B 1 26 ? -3.881  -6.412 -12.892 1.00 120.54 119 A 1 
+ATOM 205 N NH2 . ARG B 1 26 ? -2.494  -5.350 -11.417 1.00 114.95 119 A 1 
+ATOM 206 N N   . GLY B 1 27 ? -5.312  0.940  -8.241  1.00 57.29  120 A 1 
+ATOM 207 C CA  . GLY B 1 27 ? -5.176  2.005  -7.265  1.00 55.03  120 A 1 
+ATOM 208 C C   . GLY B 1 27 ? -6.391  2.105  -6.381  1.00 51.16  120 A 1 
+ATOM 209 O O   . GLY B 1 27 ? -6.273  2.252  -5.174  1.00 48.31  120 A 1 
+ATOM 210 N N   . ILE B 1 28 ? -7.570  1.983  -6.973  1.00 48.29  121 A 1 
+ATOM 211 C CA  . ILE B 1 28 ? -8.807  2.040  -6.198  1.00 51.59  121 A 1 
+ATOM 212 C C   . ILE B 1 28 ? -8.893  0.889  -5.227  1.00 54.18  121 A 1 
+ATOM 213 O O   . ILE B 1 28 ? -9.376  1.068  -4.118  1.00 53.53  121 A 1 
+ATOM 214 C CB  . ILE B 1 28 ? -10.053 2.141  -7.095  1.00 55.27  121 A 1 
+ATOM 215 C CG1 . ILE B 1 28 ? -10.187 3.601  -7.621  1.00 55.95  121 A 1 
+ATOM 216 C CG2 . ILE B 1 28 ? -11.306 1.807  -6.320  1.00 61.76  121 A 1 
+ATOM 217 C CD1 . ILE B 1 28 ? -11.126 3.822  -8.793  1.00 55.74  121 A 1 
+ATOM 218 N N   . GLU B 1 29 ? -8.420  -0.273 -5.627  1.00 61.84  122 A 1 
+ATOM 219 C CA  . GLU B 1 29 ? -8.319  -1.378 -4.700  1.00 67.68  122 A 1 
+ATOM 220 C C   . GLU B 1 29 ? -7.415  -1.069 -3.537  1.00 65.03  122 A 1 
+ATOM 221 O O   . GLU B 1 29 ? -7.791  -1.331 -2.393  1.00 64.70  122 A 1 
+ATOM 222 C CB  . GLU B 1 29 ? -7.868  -2.634 -5.431  1.00 78.78  122 A 1 
+ATOM 223 C CG  . GLU B 1 29 ? -9.069  -3.428 -5.962  1.00 86.76  122 A 1 
+ATOM 224 C CD  . GLU B 1 29 ? -8.878  -4.019 -7.353  1.00 87.49  122 A 1 
+ATOM 225 O OE1 . GLU B 1 29 ? -7.768  -4.558 -7.636  1.00 84.73  122 A 1 
+ATOM 226 O OE2 . GLU B 1 29 ? -9.870  -3.959 -8.121  1.00 77.87  122 A 1 
+ATOM 227 N N   . GLN B 1 30 ? -6.256  -0.484 -3.829  1.00 59.57  123 A 1 
+ATOM 228 C CA  . GLN B 1 30 ? -5.315  -0.084 -2.773  1.00 55.23  123 A 1 
+ATOM 229 C C   . GLN B 1 30 ? -5.910  0.955  -1.841  1.00 51.86  123 A 1 
+ATOM 230 O O   . GLN B 1 30 ? -5.662  0.942  -0.653  1.00 44.12  123 A 1 
+ATOM 231 C CB  . GLN B 1 30 ? -4.037  0.514  -3.362  1.00 55.92  123 A 1 
+ATOM 232 C CG  . GLN B 1 30 ? -3.221  -0.371 -4.262  1.00 61.37  123 A 1 
+ATOM 233 C CD  . GLN B 1 30 ? -2.863  -1.664 -3.614  1.00 62.35  123 A 1 
+ATOM 234 O OE1 . GLN B 1 30 ? -2.326  -1.663 -2.512  1.00 63.09  123 A 1 
+ATOM 235 N NE2 . GLN B 1 30 ? -3.188  -2.787 -4.276  1.00 70.63  123 A 1 
+ATOM 236 N N   . VAL B 1 31 ? -6.703  1.858  -2.386  1.00 57.05  124 A 1 
+ATOM 237 C CA  . VAL B 1 31 ? -7.323  2.883  -1.569  1.00 56.37  124 A 1 
+ATOM 238 C C   . VAL B 1 31 ? -8.333  2.246  -0.629  1.00 61.68  124 A 1 
+ATOM 239 O O   . VAL B 1 31 ? -8.461  2.699  0.470   1.00 58.51  124 A 1 
+ATOM 240 C CB  . VAL B 1 31 ? -7.969  3.950  -2.456  1.00 57.06  124 A 1 
+ATOM 241 C CG1 . VAL B 1 31 ? -9.007  4.756  -1.702  1.00 58.79  124 A 1 
+ATOM 242 C CG2 . VAL B 1 31 ? -6.871  4.864  -2.974  1.00 58.78  124 A 1 
+ATOM 243 N N   . GLU B 1 32 ? -9.044  1.205  -1.068  1.00 67.95  125 A 1 
+ATOM 244 C CA  . GLU B 1 32 ? -9.940  0.481  -0.175  1.00 64.07  125 A 1 
+ATOM 245 C C   . GLU B 1 32 ? -9.140  -0.246 0.887   1.00 61.43  125 A 1 
+ATOM 246 O O   . GLU B 1 32 ? -9.521  -0.248 2.049   1.00 58.85  125 A 1 
+ATOM 247 C CB  . GLU B 1 32 ? -10.821 -0.511 -0.911  1.00 70.57  125 A 1 
+ATOM 248 C CG  . GLU B 1 32 ? -11.844 0.047  -1.875  1.00 81.60  125 A 1 
+ATOM 249 C CD  . GLU B 1 32 ? -12.772 1.105  -1.306  1.00 92.01  125 A 1 
+ATOM 250 O OE1 . GLU B 1 32 ? -13.833 0.773  -0.713  1.00 101.96 125 A 1 
+ATOM 251 O OE2 . GLU B 1 32 ? -12.392 2.286  -1.482  1.00 95.64  125 A 1 
+ATOM 252 N N   . LEU B 1 33 ? -8.042  -0.849 0.489   1.00 55.72  126 A 1 
+ATOM 253 C CA  . LEU B 1 33 ? -7.201  -1.483 1.447   1.00 58.91  126 A 1 
+ATOM 254 C C   . LEU B 1 33 ? -6.628  -0.487 2.476   1.00 56.58  126 A 1 
+ATOM 255 O O   . LEU B 1 33 ? -6.271  -0.829 3.592   1.00 59.20  126 A 1 
+ATOM 256 C CB  . LEU B 1 33 ? -6.089  -2.149 0.676   1.00 63.68  126 A 1 
+ATOM 257 C CG  . LEU B 1 33 ? -5.181  -3.155 1.283   1.00 71.88  126 A 1 
+ATOM 258 C CD1 . LEU B 1 33 ? -4.453  -3.972 0.195   1.00 73.91  126 A 1 
+ATOM 259 C CD2 . LEU B 1 33 ? -4.179  -2.480 2.169   1.00 75.78  126 A 1 
+ATOM 260 N N   . LEU B 1 34 ? -6.482  0.743  2.064   1.00 51.59  127 A 1 
+ATOM 261 C CA  . LEU B 1 34 ? -6.054  1.821  2.952   1.00 47.43  127 A 1 
+ATOM 262 C C   . LEU B 1 34 ? -7.162  2.293  3.867   1.00 49.18  127 A 1 
+ATOM 263 O O   . LEU B 1 34 ? -6.917  2.748  4.956   1.00 48.29  127 A 1 
+ATOM 264 C CB  . LEU B 1 34 ? -5.537  2.996  2.165   1.00 43.17  127 A 1 
+ATOM 265 C CG  . LEU B 1 34 ? -4.031  2.968  1.890   1.00 42.73  127 A 1 
+ATOM 266 C CD1 . LEU B 1 34 ? -3.630  4.313  1.305   1.00 42.46  127 A 1 
+ATOM 267 C CD2 . LEU B 1 34 ? -3.088  2.610  3.005   1.00 43.31  127 A 1 
+ATOM 268 N N   . LYS B 1 35 ? -8.399  2.214  3.408   1.00 52.51  128 A 1 
+ATOM 269 C CA  . LYS B 1 35 ? -9.553  2.535  4.233   1.00 52.46  128 A 1 
+ATOM 270 C C   . LYS B 1 35 ? -9.710  1.492  5.330   1.00 53.78  128 A 1 
+ATOM 271 O O   . LYS B 1 35 ? -10.040 1.823  6.442   1.00 47.30  128 A 1 
+ATOM 272 C CB  . LYS B 1 35 ? -10.830 2.641  3.381   1.00 54.62  128 A 1 
+ATOM 273 C CG  . LYS B 1 35 ? -12.138 2.251  4.078   1.00 59.22  128 A 1 
+ATOM 274 N N   . ARG B 1 36 ? -9.454  0.236  4.997   1.00 56.56  129 A 1 
+ATOM 275 C CA  . ARG B 1 36 ? -9.471  -0.853 5.953   1.00 60.75  129 A 1 
+ATOM 276 C C   . ARG B 1 36 ? -8.392  -0.650 7.014   1.00 59.48  129 A 1 
+ATOM 277 O O   . ARG B 1 36 ? -8.620  -0.927 8.185   1.00 63.52  129 A 1 
+ATOM 278 C CB  . ARG B 1 36 ? -9.292  -2.210 5.247   1.00 65.76  129 A 1 
+ATOM 279 C CG  . ARG B 1 36 ? -10.535 -2.769 4.547   1.00 68.71  129 A 1 
+ATOM 280 C CD  . ARG B 1 36 ? -10.167 -3.835 3.488   1.00 67.67  129 A 1 
+ATOM 281 N N   . ILE B 1 37 ? -7.221  -0.161 6.611   1.00 54.56  130 A 1 
+ATOM 282 C CA  . ILE B 1 37 ? -6.140  0.129  7.539   1.00 49.62  130 A 1 
+ATOM 283 C C   . ILE B 1 37 ? -6.541  1.250  8.466   1.00 49.00  130 A 1 
+ATOM 284 O O   . ILE B 1 37 ? -6.319  1.155  9.647   1.00 47.20  130 A 1 
+ATOM 285 C CB  . ILE B 1 37 ? -4.839  0.472  6.806   1.00 51.01  130 A 1 
+ATOM 286 C CG1 . ILE B 1 37 ? -4.243  -0.830 6.231   1.00 57.41  130 A 1 
+ATOM 287 C CG2 . ILE B 1 37 ? -3.851  1.181  7.755   1.00 49.06  130 A 1 
+ATOM 288 C CD1 . ILE B 1 37 ? -2.922  -0.695 5.533   1.00 59.59  130 A 1 
+ATOM 289 N N   . HIS B 1 38 ? -7.179  2.276  7.957   1.00 50.21  131 A 1 
+ATOM 290 C CA  . HIS B 1 38 ? -7.618  3.356  8.797   1.00 53.53  131 A 1 
+ATOM 291 C C   . HIS B 1 38 ? -8.670  2.891  9.780   1.00 62.35  131 A 1 
+ATOM 292 O O   . HIS B 1 38 ? -8.744  3.378  10.880  1.00 61.76  131 A 1 
+ATOM 293 C CB  . HIS B 1 38 ? -8.175  4.452  7.947   1.00 51.42  131 A 1 
+ATOM 294 C CG  . HIS B 1 38 ? -9.034  5.412  8.688   1.00 54.14  131 A 1 
+ATOM 295 N ND1 . HIS B 1 38 ? -10.380 5.221  8.850   1.00 59.92  131 A 1 
+ATOM 296 C CD2 . HIS B 1 38 ? -8.749  6.586  9.283   1.00 55.36  131 A 1 
+ATOM 297 C CE1 . HIS B 1 38 ? -10.887 6.229  9.522   1.00 62.46  131 A 1 
+ATOM 298 N NE2 . HIS B 1 38 ? -9.918  7.073  9.795   1.00 61.48  131 A 1 
+ATOM 299 N N   . ASP B 1 39 ? -9.476  1.928  9.394   1.00 63.84  132 A 1 
+ATOM 300 C CA  . ASP B 1 39 ? -10.480 1.421  10.286  1.00 65.95  132 A 1 
+ATOM 301 C C   . ASP B 1 39 ? -9.837  0.648  11.420  1.00 61.61  132 A 1 
+ATOM 302 O O   . ASP B 1 39 ? -10.164 0.840  12.551  1.00 56.16  132 A 1 
+ATOM 303 C CB  . ASP B 1 39 ? -11.456 0.577  9.509   1.00 73.60  132 A 1 
+ATOM 304 C CG  . ASP B 1 39 ? -12.373 1.404  8.658   1.00 76.60  132 A 1 
+ATOM 305 O OD1 . ASP B 1 39 ? -12.277 2.631  8.703   1.00 70.60  132 A 1 
+ATOM 306 O OD2 . ASP B 1 39 ? -13.195 0.833  7.937   1.00 81.44  132 A 1 
+ATOM 307 N N   . LYS B 1 40 ? -8.926  -0.245 11.106  1.00 59.87  133 A 1 
+ATOM 308 C CA  . LYS B 1 40 ? -8.191  -0.985 12.129  1.00 59.13  133 A 1 
+ATOM 309 C C   . LYS B 1 40 ? -7.573  -0.021 13.157  1.00 58.88  133 A 1 
+ATOM 310 O O   . LYS B 1 40 ? -7.579  -0.267 14.361  1.00 63.83  133 A 1 
+ATOM 311 C CB  . LYS B 1 40 ? -7.093  -1.864 11.503  1.00 58.69  133 A 1 
+ATOM 312 C CG  . LYS B 1 40 ? -7.577  -3.018 10.646  1.00 62.42  133 A 1 
+ATOM 313 C CD  . LYS B 1 40 ? -6.463  -3.844 9.987   1.00 62.93  133 A 1 
+ATOM 314 N N   . LEU B 1 41 ? -7.038  1.086  12.674  1.00 60.19  134 A 1 
+ATOM 315 C CA  . LEU B 1 41 ? -6.454  2.092  13.561  1.00 62.93  134 A 1 
+ATOM 316 C C   . LEU B 1 41 ? -7.487  2.707  14.452  1.00 63.64  134 A 1 
+ATOM 317 O O   . LEU B 1 41 ? -7.205  2.982  15.605  1.00 73.06  134 A 1 
+ATOM 318 C CB  . LEU B 1 41 ? -5.765  3.197  12.764  1.00 64.98  134 A 1 
+ATOM 319 C CG  . LEU B 1 41 ? -4.531  2.611  12.036  1.00 68.88  134 A 1 
+ATOM 320 C CD1 . LEU B 1 41 ? -4.332  3.201  10.638  1.00 78.49  134 A 1 
+ATOM 321 C CD2 . LEU B 1 41 ? -3.301  2.685  12.900  1.00 67.40  134 A 1 
+ATOM 322 N N   . MET B 1 42 ? -8.669  2.955  13.913  1.00 63.76  135 A 1 
+ATOM 323 C CA  . MET B 1 42 ? -9.752  3.512  14.704  1.00 67.57  135 A 1 
+ATOM 324 C C   . MET B 1 42 ? -10.186 2.561  15.817  1.00 72.44  135 A 1 
+ATOM 325 O O   . MET B 1 42 ? -10.413 2.985  16.975  1.00 66.48  135 A 1 
+ATOM 326 C CB  . MET B 1 42 ? -10.932 3.853  13.816  1.00 69.82  135 A 1 
+ATOM 327 C CG  . MET B 1 42 ? -10.710 5.112  13.020  1.00 74.54  135 A 1 
+ATOM 328 S SD  . MET B 1 42 ? -10.062 6.547  13.958  1.00 82.27  135 A 1 
+ATOM 329 C CE  . MET B 1 42 ? -11.073 6.663  15.446  1.00 83.33  135 A 1 
+ATOM 330 N N   . ILE B 1 43 ? -10.271 1.274  15.477  1.00 81.93  136 A 1 
+ATOM 331 C CA  . ILE B 1 43 ? -10.594 0.243  16.446  1.00 79.15  136 A 1 
+ATOM 332 C C   . ILE B 1 43 ? -9.526  0.183  17.546  1.00 79.65  136 A 1 
+ATOM 333 O O   . ILE B 1 43 ? -9.854  0.128  18.719  1.00 82.01  136 A 1 
+ATOM 334 C CB  . ILE B 1 43 ? -10.781 -1.129 15.738  1.00 80.39  136 A 1 
+ATOM 335 C CG1 . ILE B 1 43 ? -12.176 -1.237 15.090  1.00 78.84  136 A 1 
+ATOM 336 C CG2 . ILE B 1 43 ? -10.586 -2.294 16.708  1.00 84.53  136 A 1 
+ATOM 337 C CD1 . ILE B 1 43 ? -12.480 -0.642 13.720  1.00 74.74  136 A 1 
+ATOM 338 N N   . ARG B 1 44 ? -8.250  0.200  17.158  1.00 85.15  137 A 1 
+ATOM 339 C CA  . ARG B 1 44 ? -7.131  0.216  18.111  1.00 90.78  137 A 1 
+ATOM 340 C C   . ARG B 1 44 ? -7.103  1.498  18.957  1.00 91.00  137 A 1 
+ATOM 341 O O   . ARG B 1 44 ? -6.632  1.453  20.068  1.00 96.79  137 A 1 
+ATOM 342 C CB  . ARG B 1 44 ? -5.803  0.033  17.373  1.00 95.82  137 A 1 
+ATOM 343 C CG  . ARG B 1 44 ? -4.883  -1.005 17.902  1.00 101.09 137 A 1 
+ATOM 344 C CD  . ARG B 1 44 ? -3.772  -1.099 16.880  1.00 107.64 137 A 1 
+ATOM 345 N NE  . ARG B 1 44 ? -2.978  -2.299 17.060  1.00 124.71 137 A 1 
+ATOM 346 C CZ  . ARG B 1 44 ? -2.237  -2.882 16.112  1.00 137.30 137 A 1 
+ATOM 347 N NH1 . ARG B 1 44 ? -2.186  -2.387 14.871  1.00 128.35 137 A 1 
+ATOM 348 N NH2 . ARG B 1 44 ? -1.533  -3.990 16.400  1.00 147.87 137 A 1 
+ATOM 349 N N   . ALA B 1 45 ? -7.688  2.601  18.506  1.00 92.41  138 A 1 
+ATOM 350 C CA  . ALA B 1 45 ? -7.783  3.810  19.381  1.00 96.29  138 A 1 
+ATOM 351 C C   . ALA B 1 45 ? -8.773  4.837  18.851  1.00 95.46  138 A 1 
+ATOM 352 O O   . ALA B 1 45 ? -8.386  5.847  18.235  1.00 81.53  138 A 1 
+ATOM 353 C CB  . ALA B 1 45 ? -6.408  4.475  19.603  1.00 97.91  138 A 1 
+#

flax_model/alphafold3/test_data/miniature_databases/pdb_mmcif/7rye.cif

@@ -0,0 +1,893 @@
+data_7RYE
+#
+_entry.id 7RYE
+#
+loop_
+_chem_comp.formula
+_chem_comp.formula_weight
+_chem_comp.id
+_chem_comp.mon_nstd_flag
+_chem_comp.name
+_chem_comp.pdbx_synonyms
+_chem_comp.type
+"C3 H7 N O2"     89.093  ALA y ALANINE         ? "L-peptide linking" 
+"C6 H15 N4 O2 1" 175.209 ARG y ARGININE        ? "L-peptide linking" 
+"C4 H8 N2 O3"    132.118 ASN y ASPARAGINE      ? "L-peptide linking" 
+"C4 H7 N O4"     133.103 ASP y "ASPARTIC ACID" ? "L-peptide linking" 
+"C3 H7 N O2 S"   121.158 CYS y CYSTEINE        ? "L-peptide linking" 
+"C5 H10 N2 O3"   146.144 GLN y GLUTAMINE       ? "L-peptide linking" 
+"C5 H9 N O4"     147.129 GLU y "GLUTAMIC ACID" ? "L-peptide linking" 
+"C2 H5 N O2"     75.067  GLY y GLYCINE         ? "peptide linking"   
+"C6 H10 N3 O2 1" 156.162 HIS y HISTIDINE       ? "L-peptide linking" 
+"C6 H13 N O2"    131.173 ILE y ISOLEUCINE      ? "L-peptide linking" 
+"C6 H13 N O2"    131.173 LEU y LEUCINE         ? "L-peptide linking" 
+"C6 H15 N2 O2 1" 147.195 LYS y LYSINE          ? "L-peptide linking" 
+"C5 H11 N O2 S"  149.211 MET y METHIONINE      ? "L-peptide linking" 
+"C9 H11 N O2"    165.189 PHE y PHENYLALANINE   ? "L-peptide linking" 
+"C5 H9 N O2"     115.130 PRO y PROLINE         ? "L-peptide linking" 
+"C3 H7 N O3"     105.093 SER y SERINE          ? "L-peptide linking" 
+"C4 H9 N O3"     119.119 THR y THREONINE       ? "L-peptide linking" 
+"C11 H12 N2 O2"  204.225 TRP y TRYPTOPHAN      ? "L-peptide linking" 
+"C9 H11 N O3"    181.189 TYR y TYROSINE        ? "L-peptide linking" 
+"C5 H11 N O2"    117.146 VAL y VALINE          ? "L-peptide linking" 
+#
+_entity.id               1
+_entity.pdbx_description "Protein PrgI"
+_entity.type             polymer
+#
+_entity_poly.entity_id      1
+_entity_poly.pdbx_strand_id J
+_entity_poly.type           polypeptide(L)
+#
+loop_
+_entity_poly_seq.entity_id
+_entity_poly_seq.hetero
+_entity_poly_seq.mon_id
+_entity_poly_seq.num
+1 n MET 1  
+1 n ALA 2  
+1 n THR 3  
+1 n PRO 4  
+1 n TRP 5  
+1 n SER 6  
+1 n GLY 7  
+1 n TYR 8  
+1 n LEU 9  
+1 n ASP 10 
+1 n ASP 11 
+1 n VAL 12 
+1 n SER 13 
+1 n ALA 14 
+1 n LYS 15 
+1 n PHE 16 
+1 n ASP 17 
+1 n THR 18 
+1 n GLY 19 
+1 n VAL 20 
+1 n ASP 21 
+1 n ASN 22 
+1 n LEU 23 
+1 n GLN 24 
+1 n THR 25 
+1 n GLN 26 
+1 n VAL 27 
+1 n THR 28 
+1 n GLU 29 
+1 n ALA 30 
+1 n LEU 31 
+1 n ASP 32 
+1 n LYS 33 
+1 n LEU 34 
+1 n ALA 35 
+1 n ALA 36 
+1 n LYS 37 
+1 n PRO 38 
+1 n SER 39 
+1 n ASP 40 
+1 n PRO 41 
+1 n ALA 42 
+1 n LEU 43 
+1 n LEU 44 
+1 n ALA 45 
+1 n ALA 46 
+1 n TYR 47 
+1 n GLN 48 
+1 n SER 49 
+1 n LYS 50 
+1 n LEU 51 
+1 n SER 52 
+1 n GLU 53 
+1 n TYR 54 
+1 n ASN 55 
+1 n LEU 56 
+1 n TYR 57 
+1 n ARG 58 
+1 n ASN 59 
+1 n ALA 60 
+1 n GLN 61 
+1 n SER 62 
+1 n ASN 63 
+1 n THR 64 
+1 n VAL 65 
+1 n LYS 66 
+1 n VAL 67 
+1 n PHE 68 
+1 n LYS 69 
+1 n ASP 70 
+1 n ILE 71 
+1 n ASP 72 
+1 n ALA 73 
+1 n ALA 74 
+1 n ILE 75 
+1 n ILE 76 
+1 n GLN 77 
+1 n ASN 78 
+1 n PHE 79 
+1 n ARG 80 
+#
+_exptl.method "ELECTRON MICROSCOPY"
+#
+_pdbx_audit_revision_history.revision_date 2021-11-10
+#
+_pdbx_database_status.recvd_initial_deposition_date 2021-11-10
+#
+loop_
+_pdbx_poly_seq_scheme.asym_id
+_pdbx_poly_seq_scheme.auth_seq_num
+_pdbx_poly_seq_scheme.entity_id
+_pdbx_poly_seq_scheme.hetero
+_pdbx_poly_seq_scheme.mon_id
+_pdbx_poly_seq_scheme.pdb_ins_code
+_pdbx_poly_seq_scheme.pdb_seq_num
+_pdbx_poly_seq_scheme.pdb_strand_id
+_pdbx_poly_seq_scheme.seq_id
+J ?  1 n MET . 1  J 1  
+J ?  1 n ALA . 2  J 2  
+J 3  1 n THR . 3  J 3  
+J 4  1 n PRO . 4  J 4  
+J 5  1 n TRP . 5  J 5  
+J 6  1 n SER . 6  J 6  
+J 7  1 n GLY . 7  J 7  
+J 8  1 n TYR . 8  J 8  
+J 9  1 n LEU . 9  J 9  
+J 10 1 n ASP . 10 J 10 
+J 11 1 n ASP . 11 J 11 
+J 12 1 n VAL . 12 J 12 
+J 13 1 n SER . 13 J 13 
+J 14 1 n ALA . 14 J 14 
+J 15 1 n LYS . 15 J 15 
+J 16 1 n PHE . 16 J 16 
+J 17 1 n ASP . 17 J 17 
+J 18 1 n THR . 18 J 18 
+J 19 1 n GLY . 19 J 19 
+J 20 1 n VAL . 20 J 20 
+J 21 1 n ASP . 21 J 21 
+J 22 1 n ASN . 22 J 22 
+J 23 1 n LEU . 23 J 23 
+J 24 1 n GLN . 24 J 24 
+J 25 1 n THR . 25 J 25 
+J 26 1 n GLN . 26 J 26 
+J 27 1 n VAL . 27 J 27 
+J 28 1 n THR . 28 J 28 
+J 29 1 n GLU . 29 J 29 
+J 30 1 n ALA . 30 J 30 
+J 31 1 n LEU . 31 J 31 
+J 32 1 n ASP . 32 J 32 
+J 33 1 n LYS . 33 J 33 
+J 34 1 n LEU . 34 J 34 
+J 35 1 n ALA . 35 J 35 
+J 36 1 n ALA . 36 J 36 
+J 37 1 n LYS . 37 J 37 
+J 38 1 n PRO . 38 J 38 
+J 39 1 n SER . 39 J 39 
+J 40 1 n ASP . 40 J 40 
+J 41 1 n PRO . 41 J 41 
+J 42 1 n ALA . 42 J 42 
+J 43 1 n LEU . 43 J 43 
+J 44 1 n LEU . 44 J 44 
+J 45 1 n ALA . 45 J 45 
+J 46 1 n ALA . 46 J 46 
+J 47 1 n TYR . 47 J 47 
+J 48 1 n GLN . 48 J 48 
+J 49 1 n SER . 49 J 49 
+J 50 1 n LYS . 50 J 50 
+J 51 1 n LEU . 51 J 51 
+J 52 1 n SER . 52 J 52 
+J 53 1 n GLU . 53 J 53 
+J 54 1 n TYR . 54 J 54 
+J 55 1 n ASN . 55 J 55 
+J 56 1 n LEU . 56 J 56 
+J 57 1 n TYR . 57 J 57 
+J 58 1 n ARG . 58 J 58 
+J 59 1 n ASN . 59 J 59 
+J 60 1 n ALA . 60 J 60 
+J 61 1 n GLN . 61 J 61 
+J 62 1 n SER . 62 J 62 
+J 63 1 n ASN . 63 J 63 
+J 64 1 n THR . 64 J 64 
+J 65 1 n VAL . 65 J 65 
+J 66 1 n LYS . 66 J 66 
+J 67 1 n VAL . 67 J 67 
+J 68 1 n PHE . 68 J 68 
+J 69 1 n LYS . 69 J 69 
+J 70 1 n ASP . 70 J 70 
+J 71 1 n ILE . 71 J 71 
+J 72 1 n ASP . 72 J 72 
+J 73 1 n ALA . 73 J 73 
+J 74 1 n ALA . 74 J 74 
+J 75 1 n ILE . 75 J 75 
+J 76 1 n ILE . 76 J 76 
+J 77 1 n GLN . 77 J 77 
+J 78 1 n ASN . 78 J 78 
+J 79 1 n PHE . 79 J 79 
+J 80 1 n ARG . 80 J 80 
+#
+_pdbx_struct_assembly.details            author_defined_assembly
+_pdbx_struct_assembly.id                 1
+_pdbx_struct_assembly.method_details     ?
+_pdbx_struct_assembly.oligomeric_count   24
+_pdbx_struct_assembly.oligomeric_details 24-meric
+#
+_pdbx_struct_assembly_gen.assembly_id     1
+_pdbx_struct_assembly_gen.asym_id_list    A,B,C,D,E,F,G,H,I,J,K,L,M,N,O,P,Q,R,S,T,U,V,W,X
+_pdbx_struct_assembly_gen.oper_expression 1
+#
+_pdbx_struct_oper_list.id                 1
+_pdbx_struct_oper_list.matrix[1][1]       1.0000000000
+_pdbx_struct_oper_list.matrix[1][2]       0.0000000000
+_pdbx_struct_oper_list.matrix[1][3]       0.0000000000
+_pdbx_struct_oper_list.matrix[2][1]       0.0000000000
+_pdbx_struct_oper_list.matrix[2][2]       1.0000000000
+_pdbx_struct_oper_list.matrix[2][3]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][1]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][2]       0.0000000000
+_pdbx_struct_oper_list.matrix[3][3]       1.0000000000
+_pdbx_struct_oper_list.name               1_555
+_pdbx_struct_oper_list.symmetry_operation ?
+_pdbx_struct_oper_list.type               "identity operation"
+_pdbx_struct_oper_list.vector[1]          0.0000000000
+_pdbx_struct_oper_list.vector[2]          0.0000000000
+_pdbx_struct_oper_list.vector[3]          0.0000000000
+#
+_refine.ls_d_res_high 3.90
+#
+_software.classification other
+_software.name           "DeepMind Structure Class"
+_software.pdbx_ordinal   1
+_software.version        2.0.0
+#
+_struct_asym.entity_id 1
+_struct_asym.id        J
+#
+loop_
+_atom_site.group_PDB
+_atom_site.id
+_atom_site.type_symbol
+_atom_site.label_atom_id
+_atom_site.label_alt_id
+_atom_site.label_comp_id
+_atom_site.label_asym_id
+_atom_site.label_entity_id
+_atom_site.label_seq_id
+_atom_site.pdbx_PDB_ins_code
+_atom_site.Cartn_x
+_atom_site.Cartn_y
+_atom_site.Cartn_z
+_atom_site.occupancy
+_atom_site.B_iso_or_equiv
+_atom_site.auth_seq_id
+_atom_site.auth_asym_id
+_atom_site.pdbx_PDB_model_num
+ATOM 1   N N   . THR J 1 3  ? 180.798 156.981 49.506 1.00 100.87 3  J 1 
+ATOM 2   C CA  . THR J 1 3  ? 180.336 158.053 50.380 1.00 100.87 3  J 1 
+ATOM 3   C C   . THR J 1 3  ? 178.888 157.811 50.811 1.00 100.87 3  J 1 
+ATOM 4   O O   . THR J 1 3  ? 178.076 157.312 50.030 1.00 100.87 3  J 1 
+ATOM 5   C CB  . THR J 1 3  ? 180.468 159.436 49.691 1.00 100.87 3  J 1 
+ATOM 6   O OG1 . THR J 1 3  ? 180.169 160.478 50.627 1.00 100.87 3  J 1 
+ATOM 7   C CG2 . THR J 1 3  ? 179.541 159.546 48.488 1.00 100.87 3  J 1 
+ATOM 8   N N   . PRO J 1 4  ? 178.569 158.139 52.065 1.00 96.05  4  J 1 
+ATOM 9   C CA  . PRO J 1 4  ? 177.201 157.918 52.554 1.00 96.05  4  J 1 
+ATOM 10  C C   . PRO J 1 4  ? 176.203 158.873 51.924 1.00 96.05  4  J 1 
+ATOM 11  O O   . PRO J 1 4  ? 176.293 160.090 52.112 1.00 96.05  4  J 1 
+ATOM 12  C CB  . PRO J 1 4  ? 177.326 158.152 54.065 1.00 96.05  4  J 1 
+ATOM 13  C CG  . PRO J 1 4  ? 178.504 159.047 54.212 1.00 96.05  4  J 1 
+ATOM 14  C CD  . PRO J 1 4  ? 179.457 158.664 53.117 1.00 96.05  4  J 1 
+ATOM 15  N N   . TRP J 1 5  ? 175.251 158.334 51.169 1.00 82.05  5  J 1 
+ATOM 16  C CA  . TRP J 1 5  ? 174.242 159.172 50.543 1.00 82.05  5  J 1 
+ATOM 17  C C   . TRP J 1 5  ? 173.275 159.710 51.591 1.00 82.05  5  J 1 
+ATOM 18  O O   . TRP J 1 5  ? 173.068 159.105 52.647 1.00 82.05  5  J 1 
+ATOM 19  C CB  . TRP J 1 5  ? 173.467 158.390 49.486 1.00 82.05  5  J 1 
+ATOM 20  C CG  . TRP J 1 5  ? 172.575 159.261 48.675 1.00 82.05  5  J 1 
+ATOM 21  C CD1 . TRP J 1 5  ? 172.953 160.130 47.697 1.00 82.05  5  J 1 
+ATOM 22  C CD2 . TRP J 1 5  ? 171.158 159.402 48.811 1.00 82.05  5  J 1 
+ATOM 23  N NE1 . TRP J 1 5  ? 171.856 160.778 47.189 1.00 82.05  5  J 1 
+ATOM 24  C CE2 . TRP J 1 5  ? 170.741 160.351 47.859 1.00 82.05  5  J 1 
+ATOM 25  C CE3 . TRP J 1 5  ? 170.199 158.807 49.634 1.00 82.05  5  J 1 
+ATOM 26  C CZ2 . TRP J 1 5  ? 169.408 160.720 47.709 1.00 82.05  5  J 1 
+ATOM 27  C CZ3 . TRP J 1 5  ? 168.877 159.176 49.482 1.00 82.05  5  J 1 
+ATOM 28  C CH2 . TRP J 1 5  ? 168.493 160.120 48.527 1.00 82.05  5  J 1 
+ATOM 29  N N   . SER J 1 6  ? 172.676 160.860 51.291 1.00 78.94  6  J 1 
+ATOM 30  C CA  . SER J 1 6  ? 171.730 161.484 52.213 1.00 78.94  6  J 1 
+ATOM 31  C C   . SER J 1 6  ? 170.799 162.391 51.422 1.00 78.94  6  J 1 
+ATOM 32  O O   . SER J 1 6  ? 171.235 163.417 50.892 1.00 78.94  6  J 1 
+ATOM 33  C CB  . SER J 1 6  ? 172.462 162.265 53.300 1.00 78.94  6  J 1 
+ATOM 34  O OG  . SER J 1 6  ? 173.149 161.389 54.176 1.00 78.94  6  J 1 
+ATOM 35  N N   . GLY J 1 7  ? 169.526 162.014 51.347 1.00 80.09  7  J 1 
+ATOM 36  C CA  . GLY J 1 7  ? 168.512 162.855 50.750 1.00 80.09  7  J 1 
+ATOM 37  C C   . GLY J 1 7  ? 168.082 163.958 51.698 1.00 80.09  7  J 1 
+ATOM 38  O O   . GLY J 1 7  ? 168.539 164.059 52.838 1.00 80.09  7  J 1 
+ATOM 39  N N   . TYR J 1 8  ? 167.174 164.804 51.207 1.00 78.37  8  J 1 
+ATOM 40  C CA  . TYR J 1 8  ? 166.760 165.964 51.991 1.00 78.37  8  J 1 
+ATOM 41  C C   . TYR J 1 8  ? 166.084 165.523 53.284 1.00 78.37  8  J 1 
+ATOM 42  O O   . TYR J 1 8  ? 166.332 166.093 54.351 1.00 78.37  8  J 1 
+ATOM 43  C CB  . TYR J 1 8  ? 165.830 166.858 51.167 1.00 78.37  8  J 1 
+ATOM 44  C CG  . TYR J 1 8  ? 165.412 168.151 51.846 1.00 78.37  8  J 1 
+ATOM 45  C CD1 . TYR J 1 8  ? 164.409 168.939 51.311 1.00 78.37  8  J 1 
+ATOM 46  C CD2 . TYR J 1 8  ? 166.073 168.622 52.974 1.00 78.37  8  J 1 
+ATOM 47  C CE1 . TYR J 1 8  ? 164.036 170.122 51.913 1.00 78.37  8  J 1 
+ATOM 48  C CE2 . TYR J 1 8  ? 165.705 169.801 53.581 1.00 78.37  8  J 1 
+ATOM 49  C CZ  . TYR J 1 8  ? 164.688 170.550 53.044 1.00 78.37  8  J 1 
+ATOM 50  O OH  . TYR J 1 8  ? 164.318 171.735 53.638 1.00 78.37  8  J 1 
+ATOM 51  N N   . LEU J 1 9  ? 165.226 164.506 53.213 1.00 73.64  9  J 1 
+ATOM 52  C CA  . LEU J 1 9  ? 164.645 163.967 54.439 1.00 73.64  9  J 1 
+ATOM 53  C C   . LEU J 1 9  ? 165.730 163.390 55.340 1.00 73.64  9  J 1 
+ATOM 54  O O   . LEU J 1 9  ? 165.719 163.604 56.560 1.00 73.64  9  J 1 
+ATOM 55  C CB  . LEU J 1 9  ? 163.595 162.906 54.109 1.00 73.64  9  J 1 
+ATOM 56  C CG  . LEU J 1 9  ? 162.192 163.335 53.661 1.00 73.64  9  J 1 
+ATOM 57  C CD1 . LEU J 1 9  ? 162.211 164.303 52.488 1.00 73.64  9  J 1 
+ATOM 58  C CD2 . LEU J 1 9  ? 161.355 162.112 53.321 1.00 73.64  9  J 1 
+ATOM 59  N N   . ASP J 1 10 ? 166.683 162.668 54.750 1.00 75.23  10 J 1 
+ATOM 60  C CA  . ASP J 1 10 ? 167.828 162.189 55.511 1.00 75.23  10 J 1 
+ATOM 61  C C   . ASP J 1 10 ? 168.631 163.360 56.062 1.00 75.23  10 J 1 
+ATOM 62  O O   . ASP J 1 10 ? 169.140 163.297 57.186 1.00 75.23  10 J 1 
+ATOM 63  C CB  . ASP J 1 10 ? 168.695 161.293 54.624 1.00 75.23  10 J 1 
+ATOM 64  C CG  . ASP J 1 10 ? 169.793 160.577 55.393 1.00 75.23  10 J 1 
+ATOM 65  O OD1 . ASP J 1 10 ? 169.921 160.785 56.617 1.00 75.23  10 J 1 
+ATOM 66  O OD2 . ASP J 1 10 ? 170.537 159.796 54.764 1.00 75.23  10 J 1 
+ATOM 67  N N   . ASP J 1 11 ? 168.742 164.442 55.289 1.00 76.21  11 J 1 
+ATOM 68  C CA  . ASP J 1 11 ? 169.449 165.624 55.770 1.00 76.21  11 J 1 
+ATOM 69  C C   . ASP J 1 11 ? 168.755 166.229 56.984 1.00 76.21  11 J 1 
+ATOM 70  O O   . ASP J 1 11 ? 169.416 166.658 57.932 1.00 76.21  11 J 1 
+ATOM 71  C CB  . ASP J 1 11 ? 169.573 166.658 54.651 1.00 76.21  11 J 1 
+ATOM 72  C CG  . ASP J 1 11 ? 170.510 166.210 53.546 1.00 76.21  11 J 1 
+ATOM 73  O OD1 . ASP J 1 11 ? 171.452 165.444 53.836 1.00 76.21  11 J 1 
+ATOM 74  O OD2 . ASP J 1 11 ? 170.303 166.623 52.386 1.00 76.21  11 J 1 
+ATOM 75  N N   . VAL J 1 12 ? 167.423 166.274 56.973 1.00 71.19  12 J 1 
+ATOM 76  C CA  . VAL J 1 12 ? 166.679 166.814 58.107 1.00 71.19  12 J 1 
+ATOM 77  C C   . VAL J 1 12 ? 166.846 165.921 59.330 1.00 71.19  12 J 1 
+ATOM 78  O O   . VAL J 1 12 ? 167.032 166.405 60.457 1.00 71.19  12 J 1 
+ATOM 79  C CB  . VAL J 1 12 ? 165.197 166.989 57.729 1.00 71.19  12 J 1 
+ATOM 80  C CG1 . VAL J 1 12 ? 164.358 167.240 58.967 1.00 71.19  12 J 1 
+ATOM 81  C CG2 . VAL J 1 12 ? 165.039 168.124 56.733 1.00 71.19  12 J 1 
+ATOM 82  N N   . SER J 1 13 ? 166.779 164.604 59.130 1.00 71.31  13 J 1 
+ATOM 83  C CA  . SER J 1 13 ? 167.004 163.691 60.244 1.00 71.31  13 J 1 
+ATOM 84  C C   . SER J 1 13 ? 168.397 163.882 60.831 1.00 71.31  13 J 1 
+ATOM 85  O O   . SER J 1 13 ? 168.555 163.965 62.054 1.00 71.31  13 J 1 
+ATOM 86  C CB  . SER J 1 13 ? 166.801 162.246 59.794 1.00 71.31  13 J 1 
+ATOM 87  O OG  . SER J 1 13 ? 165.459 162.023 59.402 1.00 71.31  13 J 1 
+ATOM 88  N N   . ALA J 1 14 ? 169.413 163.985 59.972 1.00 73.94  14 J 1 
+ATOM 89  C CA  . ALA J 1 14 ? 170.773 164.209 60.451 1.00 73.94  14 J 1 
+ATOM 90  C C   . ALA J 1 14 ? 170.904 165.566 61.130 1.00 73.94  14 J 1 
+ATOM 91  O O   . ALA J 1 14 ? 171.700 165.728 62.061 1.00 73.94  14 J 1 
+ATOM 92  C CB  . ALA J 1 14 ? 171.765 164.087 59.295 1.00 73.94  14 J 1 
+ATOM 93  N N   . LYS J 1 15 ? 170.140 166.557 60.668 1.00 71.21  15 J 1 
+ATOM 94  C CA  . LYS J 1 15 ? 170.140 167.860 61.321 1.00 71.21  15 J 1 
+ATOM 95  C C   . LYS J 1 15 ? 169.612 167.753 62.742 1.00 71.21  15 J 1 
+ATOM 96  O O   . LYS J 1 15 ? 170.146 168.383 63.660 1.00 71.21  15 J 1 
+ATOM 97  C CB  . LYS J 1 15 ? 169.311 168.855 60.511 1.00 71.21  15 J 1 
+ATOM 98  C CG  . LYS J 1 15 ? 170.142 169.773 59.635 1.00 71.21  15 J 1 
+ATOM 99  C CD  . LYS J 1 15 ? 169.313 170.920 59.089 1.00 71.21  15 J 1 
+ATOM 100 C CE  . LYS J 1 15 ? 170.176 171.888 58.300 1.00 71.21  15 J 1 
+ATOM 101 N NZ  . LYS J 1 15 ? 169.369 172.970 57.676 1.00 71.21  15 J 1 
+ATOM 102 N N   . PHE J 1 16 ? 168.559 166.963 62.943 1.00 68.61  16 J 1 
+ATOM 103 C CA  . PHE J 1 16 ? 168.102 166.713 64.308 1.00 68.61  16 J 1 
+ATOM 104 C C   . PHE J 1 16 ? 169.128 165.927 65.120 1.00 68.61  16 J 1 
+ATOM 105 O O   . PHE J 1 16 ? 169.342 166.226 66.300 1.00 68.61  16 J 1 
+ATOM 106 C CB  . PHE J 1 16 ? 166.749 166.005 64.309 1.00 68.61  16 J 1 
+ATOM 107 C CG  . PHE J 1 16 ? 165.587 166.929 64.088 1.00 68.61  16 J 1 
+ATOM 108 C CD1 . PHE J 1 16 ? 165.736 168.293 64.256 1.00 68.61  16 J 1 
+ATOM 109 C CD2 . PHE J 1 16 ? 164.345 166.437 63.736 1.00 68.61  16 J 1 
+ATOM 110 C CE1 . PHE J 1 16 ? 164.670 169.147 64.068 1.00 68.61  16 J 1 
+ATOM 111 C CE2 . PHE J 1 16 ? 163.277 167.286 63.544 1.00 68.61  16 J 1 
+ATOM 112 C CZ  . PHE J 1 16 ? 163.439 168.643 63.710 1.00 68.61  16 J 1 
+ATOM 113 N N   . ASP J 1 17 ? 169.768 164.918 64.521 1.00 77.06  17 J 1 
+ATOM 114 C CA  . ASP J 1 17 ? 170.752 164.149 65.286 1.00 77.06  17 J 1 
+ATOM 115 C C   . ASP J 1 17 ? 171.927 165.017 65.719 1.00 77.06  17 J 1 
+ATOM 116 O O   . ASP J 1 17 ? 172.413 164.888 66.848 1.00 77.06  17 J 1 
+ATOM 117 C CB  . ASP J 1 17 ? 171.254 162.929 64.509 1.00 77.06  17 J 1 
+ATOM 118 C CG  . ASP J 1 17 ? 170.136 162.119 63.888 1.00 77.06  17 J 1 
+ATOM 119 O OD1 . ASP J 1 17 ? 170.103 162.001 62.647 1.00 77.06  17 J 1 
+ATOM 120 O OD2 . ASP J 1 17 ? 169.286 161.602 64.642 1.00 77.06  17 J 1 
+ATOM 121 N N   . THR J 1 18 ? 172.399 165.901 64.845 1.00 76.90  18 J 1 
+ATOM 122 C CA  . THR J 1 18 ? 173.485 166.801 65.205 1.00 76.90  18 J 1 
+ATOM 123 C C   . THR J 1 18 ? 173.016 167.991 66.028 1.00 76.90  18 J 1 
+ATOM 124 O O   . THR J 1 18 ? 173.851 168.678 66.626 1.00 76.90  18 J 1 
+ATOM 125 C CB  . THR J 1 18 ? 174.203 167.303 63.949 1.00 76.90  18 J 1 
+ATOM 126 O OG1 . THR J 1 18 ? 175.476 167.849 64.314 1.00 76.90  18 J 1 
+ATOM 127 C CG2 . THR J 1 18 ? 173.381 168.373 63.253 1.00 76.90  18 J 1 
+ATOM 128 N N   . GLY J 1 19 ? 171.709 168.255 66.067 1.00 78.19  19 J 1 
+ATOM 129 C CA  . GLY J 1 19 ? 171.207 169.320 66.916 1.00 78.19  19 J 1 
+ATOM 130 C C   . GLY J 1 19 ? 171.381 169.012 68.389 1.00 78.19  19 J 1 
+ATOM 131 O O   . GLY J 1 19 ? 171.777 169.879 69.172 1.00 78.19  19 J 1 
+ATOM 132 N N   . VAL J 1 20 ? 171.089 167.780 68.785 1.00 76.12  20 J 1 
+ATOM 133 C CA  . VAL J 1 20 ? 171.317 167.329 70.150 1.00 76.12  20 J 1 
+ATOM 134 C C   . VAL J 1 20 ? 172.694 166.684 70.220 1.00 76.12  20 J 1 
+ATOM 135 O O   . VAL J 1 20 ? 173.054 165.862 69.368 1.00 76.12  20 J 1 
+ATOM 136 C CB  . VAL J 1 20 ? 170.212 166.362 70.606 1.00 76.12  20 J 1 
+ATOM 137 C CG1 . VAL J 1 20 ? 168.869 167.069 70.596 1.00 76.12  20 J 1 
+ATOM 138 C CG2 . VAL J 1 20 ? 170.158 165.132 69.713 1.00 76.12  20 J 1 
+ATOM 139 N N   . ASP J 1 21 ? 173.487 167.090 71.210 1.00 78.41  21 J 1 
+ATOM 140 C CA  . ASP J 1 21 ? 174.863 166.617 71.290 1.00 78.41  21 J 1 
+ATOM 141 C C   . ASP J 1 21 ? 174.916 165.113 71.524 1.00 78.41  21 J 1 
+ATOM 142 O O   . ASP J 1 21 ? 175.347 164.349 70.653 1.00 78.41  21 J 1 
+ATOM 143 C CB  . ASP J 1 21 ? 175.608 167.363 72.396 1.00 78.41  21 J 1 
+ATOM 144 C CG  . ASP J 1 21 ? 175.835 168.824 72.061 1.00 78.41  21 J 1 
+ATOM 145 O OD1 . ASP J 1 21 ? 175.946 169.149 70.861 1.00 78.41  21 J 1 
+ATOM 146 O OD2 . ASP J 1 21 ? 175.903 169.647 72.998 1.00 78.41  21 J 1 
+ATOM 147 N N   . ASN J 1 22 ? 174.471 164.668 72.693 1.00 74.91  22 J 1 
+ATOM 148 C CA  . ASN J 1 22 ? 174.565 163.267 73.084 1.00 74.91  22 J 1 
+ATOM 149 C C   . ASN J 1 22 ? 173.277 162.819 73.765 1.00 74.91  22 J 1 
+ATOM 150 O O   . ASN J 1 22 ? 173.294 162.151 74.799 1.00 74.91  22 J 1 
+ATOM 151 C CB  . ASN J 1 22 ? 175.777 163.050 73.987 1.00 74.91  22 J 1 
+ATOM 152 C CG  . ASN J 1 22 ? 175.825 164.032 75.138 1.00 74.91  22 J 1 
+ATOM 153 O OD1 . ASN J 1 22 ? 174.925 164.855 75.304 1.00 74.91  22 J 1 
+ATOM 154 N ND2 . ASN J 1 22 ? 176.881 163.956 75.938 1.00 74.91  22 J 1 
+ATOM 155 N N   . LEU J 1 23 ? 172.134 163.193 73.182 1.00 72.21  23 J 1 
+ATOM 156 C CA  . LEU J 1 23 ? 170.852 162.915 73.824 1.00 72.21  23 J 1 
+ATOM 157 C C   . LEU J 1 23 ? 170.648 161.425 74.053 1.00 72.21  23 J 1 
+ATOM 158 O O   . LEU J 1 23 ? 170.248 161.010 75.146 1.00 72.21  23 J 1 
+ATOM 159 C CB  . LEU J 1 23 ? 169.708 163.483 72.989 1.00 72.21  23 J 1 
+ATOM 160 C CG  . LEU J 1 23 ? 168.328 163.271 73.612 1.00 72.21  23 J 1 
+ATOM 161 C CD1 . LEU J 1 23 ? 168.229 163.983 74.946 1.00 72.21  23 J 1 
+ATOM 162 C CD2 . LEU J 1 23 ? 167.237 163.743 72.669 1.00 72.21  23 J 1 
+ATOM 163 N N   . GLN J 1 24 ? 170.920 160.601 73.039 1.00 70.63  24 J 1 
+ATOM 164 C CA  . GLN J 1 24 ? 170.840 159.157 73.230 1.00 70.63  24 J 1 
+ATOM 165 C C   . GLN J 1 24 ? 171.802 158.707 74.320 1.00 70.63  24 J 1 
+ATOM 166 O O   . GLN J 1 24 ? 171.425 157.974 75.242 1.00 70.63  24 J 1 
+ATOM 167 C CB  . GLN J 1 24 ? 171.143 158.435 71.917 1.00 70.63  24 J 1 
+ATOM 168 C CG  . GLN J 1 24 ? 170.834 156.949 71.947 1.00 70.63  24 J 1 
+ATOM 169 C CD  . GLN J 1 24 ? 169.405 156.644 71.556 1.00 70.63  24 J 1 
+ATOM 170 O OE1 . GLN J 1 24 ? 168.697 157.505 71.042 1.00 70.63  24 J 1 
+ATOM 171 N NE2 . GLN J 1 24 ? 168.972 155.412 71.801 1.00 70.63  24 J 1 
+ATOM 172 N N   . THR J 1 25 ? 173.052 159.167 74.242 1.00 72.05  25 J 1 
+ATOM 173 C CA  . THR J 1 25 ? 174.029 158.829 75.267 1.00 72.05  25 J 1 
+ATOM 174 C C   . THR J 1 25 ? 173.628 159.382 76.627 1.00 72.05  25 J 1 
+ATOM 175 O O   . THR J 1 25 ? 173.752 158.681 77.636 1.00 72.05  25 J 1 
+ATOM 176 C CB  . THR J 1 25 ? 175.407 159.349 74.865 1.00 72.05  25 J 1 
+ATOM 177 O OG1 . THR J 1 25 ? 175.777 158.785 73.601 1.00 72.05  25 J 1 
+ATOM 178 C CG2 . THR J 1 25 ? 176.447 158.961 75.898 1.00 72.05  25 J 1 
+ATOM 179 N N   . GLN J 1 26 ? 173.147 160.627 76.676 1.00 73.20  26 J 1 
+ATOM 180 C CA  . GLN J 1 26 ? 172.756 161.219 77.951 1.00 73.20  26 J 1 
+ATOM 181 C C   . GLN J 1 26 ? 171.656 160.400 78.610 1.00 73.20  26 J 1 
+ATOM 182 O O   . GLN J 1 26 ? 171.709 160.118 79.814 1.00 73.20  26 J 1 
+ATOM 183 C CB  . GLN J 1 26 ? 172.301 162.663 77.739 1.00 73.20  26 J 1 
+ATOM 184 C CG  . GLN J 1 26 ? 173.397 163.693 77.939 1.00 73.20  26 J 1 
+ATOM 185 C CD  . GLN J 1 26 ? 172.920 165.105 77.674 1.00 73.20  26 J 1 
+ATOM 186 O OE1 . GLN J 1 26 ? 171.782 165.454 77.979 1.00 73.20  26 J 1 
+ATOM 187 N NE2 . GLN J 1 26 ? 173.791 165.926 77.101 1.00 73.20  26 J 1 
+ATOM 188 N N   . VAL J 1 27 ? 170.660 159.989 77.824 1.00 70.38  27 J 1 
+ATOM 189 C CA  . VAL J 1 27 ? 169.615 159.113 78.339 1.00 70.38  27 J 1 
+ATOM 190 C C   . VAL J 1 27 ? 170.211 157.789 78.796 1.00 70.38  27 J 1 
+ATOM 191 O O   . VAL J 1 27 ? 169.787 157.217 79.808 1.00 70.38  27 J 1 
+ATOM 192 C CB  . VAL J 1 27 ? 168.521 158.911 77.274 1.00 70.38  27 J 1 
+ATOM 193 C CG1 . VAL J 1 27 ? 167.570 157.803 77.685 1.00 70.38  27 J 1 
+ATOM 194 C CG2 . VAL J 1 27 ? 167.761 160.204 77.059 1.00 70.38  27 J 1 
+ATOM 195 N N   . THR J 1 28 ? 171.208 157.285 78.067 1.00 66.97  28 J 1 
+ATOM 196 C CA  . THR J 1 28 ? 171.814 156.011 78.437 1.00 66.97  28 J 1 
+ATOM 197 C C   . THR J 1 28 ? 172.453 156.080 79.821 1.00 66.97  28 J 1 
+ATOM 198 O O   . THR J 1 28 ? 172.163 155.246 80.688 1.00 66.97  28 J 1 
+ATOM 199 C CB  . THR J 1 28 ? 172.841 155.591 77.388 1.00 66.97  28 J 1 
+ATOM 200 O OG1 . THR J 1 28 ? 172.192 155.449 76.118 1.00 66.97  28 J 1 
+ATOM 201 C CG2 . THR J 1 28 ? 173.475 154.265 77.774 1.00 66.97  28 J 1 
+ATOM 202 N N   . GLU J 1 29 ? 173.322 157.072 80.055 1.00 72.66  29 J 1 
+ATOM 203 C CA  . GLU J 1 29 ? 173.906 157.156 81.396 1.00 72.66  29 J 1 
+ATOM 204 C C   . GLU J 1 29 ? 172.874 157.531 82.451 1.00 72.66  29 J 1 
+ATOM 205 O O   . GLU J 1 29 ? 172.996 157.088 83.594 1.00 72.66  29 J 1 
+ATOM 206 C CB  . GLU J 1 29 ? 175.107 158.111 81.526 1.00 72.66  29 J 1 
+ATOM 207 C CG  . GLU J 1 29 ? 176.428 157.677 80.874 1.00 72.66  29 J 1 
+ATOM 208 C CD  . GLU J 1 29 ? 176.594 158.116 79.452 1.00 72.66  29 J 1 
+ATOM 209 O OE1 . GLU J 1 29 ? 175.740 158.874 78.981 1.00 72.66  29 J 1 
+ATOM 210 O OE2 . GLU J 1 29 ? 177.582 157.707 78.806 1.00 72.66  29 J 1 
+ATOM 211 N N   . ALA J 1 30 ? 171.858 158.331 82.113 1.00 71.55  30 J 1 
+ATOM 212 C CA  . ALA J 1 30 ? 170.840 158.641 83.113 1.00 71.55  30 J 1 
+ATOM 213 C C   . ALA J 1 30 ? 170.124 157.377 83.571 1.00 71.55  30 J 1 
+ATOM 214 O O   . ALA J 1 30 ? 169.974 157.134 84.776 1.00 71.55  30 J 1 
+ATOM 215 C CB  . ALA J 1 30 ? 169.842 159.654 82.555 1.00 71.55  30 J 1 
+ATOM 216 N N   . LEU J 1 31 ? 169.710 156.541 82.619 1.00 67.57  31 J 1 
+ATOM 217 C CA  . LEU J 1 31 ? 169.026 155.299 82.949 1.00 67.57  31 J 1 
+ATOM 218 C C   . LEU J 1 31 ? 169.947 154.332 83.681 1.00 67.57  31 J 1 
+ATOM 219 O O   . LEU J 1 31 ? 169.513 153.645 84.613 1.00 67.57  31 J 1 
+ATOM 220 C CB  . LEU J 1 31 ? 168.470 154.664 81.675 1.00 67.57  31 J 1 
+ATOM 221 C CG  . LEU J 1 31 ? 167.565 153.444 81.827 1.00 67.57  31 J 1 
+ATOM 222 C CD1 . LEU J 1 31 ? 166.434 153.521 80.822 1.00 67.57  31 J 1 
+ATOM 223 C CD2 . LEU J 1 31 ? 168.355 152.161 81.641 1.00 67.57  31 J 1 
+ATOM 224 N N   . ASP J 1 32 ? 171.217 154.256 83.275 1.00 72.67  32 J 1 
+ATOM 225 C CA  . ASP J 1 32 ? 172.149 153.367 83.960 1.00 72.67  32 J 1 
+ATOM 226 C C   . ASP J 1 32 ? 172.379 153.821 85.398 1.00 72.67  32 J 1 
+ATOM 227 O O   . ASP J 1 32 ? 172.414 152.997 86.319 1.00 72.67  32 J 1 
+ATOM 228 C CB  . ASP J 1 32 ? 173.467 153.305 83.188 1.00 72.67  32 J 1 
+ATOM 229 C CG  . ASP J 1 32 ? 174.413 152.240 83.718 1.00 72.67  32 J 1 
+ATOM 230 O OD1 . ASP J 1 32 ? 174.073 151.558 84.706 1.00 72.67  32 J 1 
+ATOM 231 O OD2 . ASP J 1 32 ? 175.507 152.083 83.137 1.00 72.67  32 J 1 
+ATOM 232 N N   . LYS J 1 33 ? 172.532 155.130 85.610 1.00 71.55  33 J 1 
+ATOM 233 C CA  . LYS J 1 33 ? 172.715 155.655 86.958 1.00 71.55  33 J 1 
+ATOM 234 C C   . LYS J 1 33 ? 171.496 155.382 87.824 1.00 71.55  33 J 1 
+ATOM 235 O O   . LYS J 1 33 ? 171.628 155.006 88.994 1.00 71.55  33 J 1 
+ATOM 236 C CB  . LYS J 1 33 ? 173.001 157.155 86.900 1.00 71.55  33 J 1 
+ATOM 237 C CG  . LYS J 1 33 ? 174.385 157.517 86.392 1.00 71.55  33 J 1 
+ATOM 238 C CD  . LYS J 1 33 ? 174.482 159.007 86.104 1.00 71.55  33 J 1 
+ATOM 239 C CE  . LYS J 1 33 ? 174.396 159.829 87.379 1.00 71.55  33 J 1 
+ATOM 240 N NZ  . LYS J 1 33 ? 174.607 161.279 87.119 1.00 71.55  33 J 1 
+ATOM 241 N N   . LEU J 1 34 ? 170.296 155.566 87.269 1.00 70.34  34 J 1 
+ATOM 242 C CA  . LEU J 1 34 ? 169.095 155.291 88.045 1.00 70.34  34 J 1 
+ATOM 243 C C   . LEU J 1 34 ? 168.944 153.807 88.358 1.00 70.34  34 J 1 
+ATOM 244 O O   . LEU J 1 34 ? 168.263 153.458 89.327 1.00 70.34  34 J 1 
+ATOM 245 C CB  . LEU J 1 34 ? 167.865 155.824 87.305 1.00 70.34  34 J 1 
+ATOM 246 C CG  . LEU J 1 34 ? 166.493 155.819 87.987 1.00 70.34  34 J 1 
+ATOM 247 C CD1 . LEU J 1 34 ? 165.645 156.923 87.403 1.00 70.34  34 J 1 
+ATOM 248 C CD2 . LEU J 1 34 ? 165.773 154.496 87.808 1.00 70.34  34 J 1 
+ATOM 249 N N   . ALA J 1 35 ? 169.587 152.930 87.582 1.00 66.69  35 J 1 
+ATOM 250 C CA  . ALA J 1 35 ? 169.434 151.494 87.785 1.00 66.69  35 J 1 
+ATOM 251 C C   . ALA J 1 35 ? 169.942 151.034 89.144 1.00 66.69  35 J 1 
+ATOM 252 O O   . ALA J 1 35 ? 169.492 149.995 89.640 1.00 66.69  35 J 1 
+ATOM 253 C CB  . ALA J 1 35 ? 170.153 150.726 86.678 1.00 66.69  35 J 1 
+ATOM 254 N N   . ALA J 1 36 ? 170.865 151.775 89.759 1.00 67.09  36 J 1 
+ATOM 255 C CA  . ALA J 1 36 ? 171.402 151.403 91.059 1.00 67.09  36 J 1 
+ATOM 256 C C   . ALA J 1 36 ? 171.016 152.357 92.179 1.00 67.09  36 J 1 
+ATOM 257 O O   . ALA J 1 36 ? 171.066 151.959 93.347 1.00 67.09  36 J 1 
+ATOM 258 C CB  . ALA J 1 36 ? 172.933 151.307 90.996 1.00 67.09  36 J 1 
+ATOM 259 N N   . LYS J 1 37 ? 170.632 153.590 91.860 1.00 66.94  37 J 1 
+ATOM 260 C CA  . LYS J 1 37 ? 170.209 154.579 92.850 1.00 66.94  37 J 1 
+ATOM 261 C C   . LYS J 1 37 ? 168.843 155.099 92.431 1.00 66.94  37 J 1 
+ATOM 262 O O   . LYS J 1 37 ? 168.736 156.111 91.726 1.00 66.94  37 J 1 
+ATOM 263 C CB  . LYS J 1 37 ? 171.229 155.709 92.979 1.00 66.94  37 J 1 
+ATOM 264 C CG  . LYS J 1 37 ? 172.582 155.242 93.477 1.00 66.94  37 J 1 
+ATOM 265 C CD  . LYS J 1 37 ? 172.450 154.552 94.823 1.00 66.94  37 J 1 
+ATOM 266 C CE  . LYS J 1 37 ? 173.808 154.194 95.400 1.00 66.94  37 J 1 
+ATOM 267 N NZ  . LYS J 1 37 ? 173.680 153.478 96.700 1.00 66.94  37 J 1 
+ATOM 268 N N   . PRO J 1 38 ? 167.768 154.426 92.850 1.00 68.68  38 J 1 
+ATOM 269 C CA  . PRO J 1 38 ? 166.442 154.759 92.315 1.00 68.68  38 J 1 
+ATOM 270 C C   . PRO J 1 38 ? 165.804 155.974 92.952 1.00 68.68  38 J 1 
+ATOM 271 O O   . PRO J 1 38 ? 165.061 156.684 92.265 1.00 68.68  38 J 1 
+ATOM 272 C CB  . PRO J 1 38 ? 165.630 153.501 92.616 1.00 68.68  38 J 1 
+ATOM 273 C CG  . PRO J 1 38 ? 166.222 153.021 93.902 1.00 68.68  38 J 1 
+ATOM 274 C CD  . PRO J 1 38 ? 167.701 153.332 93.834 1.00 68.68  38 J 1 
+ATOM 275 N N   . SER J 1 39 ? 166.056 156.240 94.230 1.00 70.54  39 J 1 
+ATOM 276 C CA  . SER J 1 39 ? 165.383 157.320 94.933 1.00 70.54  39 J 1 
+ATOM 277 C C   . SER J 1 39 ? 166.204 158.596 94.992 1.00 70.54  39 J 1 
+ATOM 278 O O   . SER J 1 39 ? 165.705 159.612 95.488 1.00 70.54  39 J 1 
+ATOM 279 C CB  . SER J 1 39 ? 165.027 156.887 96.360 1.00 70.54  39 J 1 
+ATOM 280 O OG  . SER J 1 39 ? 166.196 156.601 97.104 1.00 70.54  39 J 1 
+ATOM 281 N N   . ASP J 1 40 ? 167.425 158.579 94.500 1.00 71.51  40 J 1 
+ATOM 282 C CA  . ASP J 1 40 ? 168.310 159.702 94.730 1.00 71.51  40 J 1 
+ATOM 283 C C   . ASP J 1 40 ? 167.868 160.912 93.909 1.00 71.51  40 J 1 
+ATOM 284 O O   . ASP J 1 40 ? 167.348 160.766 92.800 1.00 71.51  40 J 1 
+ATOM 285 C CB  . ASP J 1 40 ? 169.741 159.324 94.374 1.00 71.51  40 J 1 
+ATOM 286 C CG  . ASP J 1 40 ? 170.519 158.829 95.569 1.00 71.51  40 J 1 
+ATOM 287 O OD1 . ASP J 1 40 ? 170.038 159.024 96.701 1.00 71.51  40 J 1 
+ATOM 288 O OD2 . ASP J 1 40 ? 171.606 158.245 95.379 1.00 71.51  40 J 1 
+ATOM 289 N N   . PRO J 1 41 ? 168.040 162.115 94.445 1.00 72.15  41 J 1 
+ATOM 290 C CA  . PRO J 1 41 ? 167.725 163.323 93.677 1.00 72.15  41 J 1 
+ATOM 291 C C   . PRO J 1 41 ? 168.792 163.594 92.630 1.00 72.15  41 J 1 
+ATOM 292 O O   . PRO J 1 41 ? 169.806 162.902 92.531 1.00 72.15  41 J 1 
+ATOM 293 C CB  . PRO J 1 41 ? 167.709 164.427 94.737 1.00 72.15  41 J 1 
+ATOM 294 C CG  . PRO J 1 41 ? 167.558 163.713 96.031 1.00 72.15  41 J 1 
+ATOM 295 C CD  . PRO J 1 41 ? 168.278 162.420 95.862 1.00 72.15  41 J 1 
+ATOM 296 N N   . ALA J 1 42 ? 168.537 164.627 91.831 1.00 68.90  42 J 1 
+ATOM 297 C CA  . ALA J 1 42 ? 169.449 165.183 90.837 1.00 68.90  42 J 1 
+ATOM 298 C C   . ALA J 1 42 ? 169.710 164.237 89.674 1.00 68.90  42 J 1 
+ATOM 299 O O   . ALA J 1 42 ? 170.388 164.632 88.718 1.00 68.90  42 J 1 
+ATOM 300 C CB  . ALA J 1 42 ? 170.795 165.605 91.445 1.00 68.90  42 J 1 
+ATOM 301 N N   . LEU J 1 43 ? 169.199 163.008 89.712 1.00 66.88  43 J 1 
+ATOM 302 C CA  . LEU J 1 43 ? 169.214 162.168 88.523 1.00 66.88  43 J 1 
+ATOM 303 C C   . LEU J 1 43 ? 167.793 161.722 88.216 1.00 66.88  43 J 1 
+ATOM 304 O O   . LEU J 1 43 ? 167.436 161.561 87.048 1.00 66.88  43 J 1 
+ATOM 305 C CB  . LEU J 1 43 ? 170.136 160.956 88.678 1.00 66.88  43 J 1 
+ATOM 306 C CG  . LEU J 1 43 ? 169.723 159.783 89.565 1.00 66.88  43 J 1 
+ATOM 307 C CD1 . LEU J 1 43 ? 170.578 158.575 89.251 1.00 66.88  43 J 1 
+ATOM 308 C CD2 . LEU J 1 43 ? 169.875 160.147 91.011 1.00 66.88  43 J 1 
+ATOM 309 N N   . LEU J 1 44 ? 166.976 161.526 89.255 1.00 62.81  44 J 1 
+ATOM 310 C CA  . LEU J 1 44 ? 165.564 161.231 89.034 1.00 62.81  44 J 1 
+ATOM 311 C C   . LEU J 1 44 ? 164.891 162.349 88.255 1.00 62.81  44 J 1 
+ATOM 312 O O   . LEU J 1 44 ? 164.229 162.103 87.240 1.00 62.81  44 J 1 
+ATOM 313 C CB  . LEU J 1 44 ? 164.850 161.025 90.366 1.00 62.81  44 J 1 
+ATOM 314 C CG  . LEU J 1 44 ? 164.662 159.603 90.876 1.00 62.81  44 J 1 
+ATOM 315 C CD1 . LEU J 1 44 ? 164.144 159.655 92.300 1.00 62.81  44 J 1 
+ATOM 316 C CD2 . LEU J 1 44 ? 163.706 158.844 89.980 1.00 62.81  44 J 1 
+ATOM 317 N N   . ALA J 1 45 ? 165.051 163.589 88.717 1.00 63.51  45 J 1 
+ATOM 318 C CA  . ALA J 1 45 ? 164.565 164.727 87.951 1.00 63.51  45 J 1 
+ATOM 319 C C   . ALA J 1 45 ? 165.278 164.805 86.611 1.00 63.51  45 J 1 
+ATOM 320 O O   . ALA J 1 45 ? 164.660 165.085 85.577 1.00 63.51  45 J 1 
+ATOM 321 C CB  . ALA J 1 45 ? 164.759 166.015 88.748 1.00 63.51  45 J 1 
+ATOM 322 N N   . ALA J 1 46 ? 166.587 164.547 86.611 1.00 63.88  46 J 1 
+ATOM 323 C CA  . ALA J 1 46 ? 167.322 164.470 85.356 1.00 63.88  46 J 1 
+ATOM 324 C C   . ALA J 1 46 ? 166.799 163.335 84.489 1.00 63.88  46 J 1 
+ATOM 325 O O   . ALA J 1 46 ? 166.675 163.486 83.270 1.00 63.88  46 J 1 
+ATOM 326 C CB  . ALA J 1 46 ? 168.815 164.297 85.627 1.00 63.88  46 J 1 
+ATOM 327 N N   . TYR J 1 47 ? 166.483 162.190 85.100 1.00 60.60  47 J 1 
+ATOM 328 C CA  . TYR J 1 47 ? 165.886 161.089 84.350 1.00 60.60  47 J 1 
+ATOM 329 C C   . TYR J 1 47 ? 164.618 161.528 83.639 1.00 60.60  47 J 1 
+ATOM 330 O O   . TYR J 1 47 ? 164.476 161.336 82.427 1.00 60.60  47 J 1 
+ATOM 331 C CB  . TYR J 1 47 ? 165.579 159.921 85.282 1.00 60.60  47 J 1 
+ATOM 332 C CG  . TYR J 1 47 ? 164.985 158.730 84.578 1.00 60.60  47 J 1 
+ATOM 333 C CD1 . TYR J 1 47 ? 165.771 157.885 83.811 1.00 60.60  47 J 1 
+ATOM 334 C CD2 . TYR J 1 47 ? 163.629 158.457 84.675 1.00 60.60  47 J 1 
+ATOM 335 C CE1 . TYR J 1 47 ? 165.223 156.798 83.168 1.00 60.60  47 J 1 
+ATOM 336 C CE2 . TYR J 1 47 ? 163.073 157.376 84.037 1.00 60.60  47 J 1 
+ATOM 337 C CZ  . TYR J 1 47 ? 163.872 156.550 83.285 1.00 60.60  47 J 1 
+ATOM 338 O OH  . TYR J 1 47 ? 163.317 155.471 82.646 1.00 60.60  47 J 1 
+ATOM 339 N N   . GLN J 1 48 ? 163.681 162.118 84.380 1.00 61.22  48 J 1 
+ATOM 340 C CA  . GLN J 1 48 ? 162.418 162.506 83.769 1.00 61.22  48 J 1 
+ATOM 341 C C   . GLN J 1 48 ? 162.618 163.572 82.703 1.00 61.22  48 J 1 
+ATOM 342 O O   . GLN J 1 48 ? 162.021 163.485 81.623 1.00 61.22  48 J 1 
+ATOM 343 C CB  . GLN J 1 48 ? 161.441 163.000 84.832 1.00 61.22  48 J 1 
+ATOM 344 C CG  . GLN J 1 48 ? 160.073 163.333 84.276 1.00 61.22  48 J 1 
+ATOM 345 C CD  . GLN J 1 48 ? 159.096 163.765 85.344 1.00 61.22  48 J 1 
+ATOM 346 O OE1 . GLN J 1 48 ? 159.445 163.860 86.519 1.00 61.22  48 J 1 
+ATOM 347 N NE2 . GLN J 1 48 ? 157.858 164.028 84.941 1.00 61.22  48 J 1 
+ATOM 348 N N   . SER J 1 49 ? 163.461 164.571 82.975 1.00 66.01  49 J 1 
+ATOM 349 C CA  . SER J 1 49 ? 163.666 165.643 82.010 1.00 66.01  49 J 1 
+ATOM 350 C C   . SER J 1 49 ? 164.298 165.117 80.729 1.00 66.01  49 J 1 
+ATOM 351 O O   . SER J 1 49 ? 163.843 165.437 79.628 1.00 66.01  49 J 1 
+ATOM 352 C CB  . SER J 1 49 ? 164.529 166.744 82.623 1.00 66.01  49 J 1 
+ATOM 353 O OG  . SER J 1 49 ? 164.839 167.732 81.658 1.00 66.01  49 J 1 
+ATOM 354 N N   . LYS J 1 50 ? 165.348 164.301 80.851 1.00 62.23  50 J 1 
+ATOM 355 C CA  . LYS J 1 50 ? 165.992 163.755 79.662 1.00 62.23  50 J 1 
+ATOM 356 C C   . LYS J 1 50 ? 165.074 162.806 78.908 1.00 62.23  50 J 1 
+ATOM 357 O O   . LYS J 1 50 ? 165.080 162.796 77.674 1.00 62.23  50 J 1 
+ATOM 358 C CB  . LYS J 1 50 ? 167.292 163.045 80.036 1.00 62.23  50 J 1 
+ATOM 359 C CG  . LYS J 1 50 ? 168.320 163.948 80.688 1.00 62.23  50 J 1 
+ATOM 360 C CD  . LYS J 1 50 ? 168.804 165.014 79.725 1.00 62.23  50 J 1 
+ATOM 361 C CE  . LYS J 1 50 ? 169.815 165.927 80.391 1.00 62.23  50 J 1 
+ATOM 362 N NZ  . LYS J 1 50 ? 171.093 165.218 80.672 1.00 62.23  50 J 1 
+ATOM 363 N N   . LEU J 1 51 ? 164.282 162.004 79.620 1.00 55.13  51 J 1 
+ATOM 364 C CA  . LEU J 1 51 ? 163.386 161.078 78.942 1.00 55.13  51 J 1 
+ATOM 365 C C   . LEU J 1 51 ? 162.306 161.830 78.175 1.00 55.13  51 J 1 
+ATOM 366 O O   . LEU J 1 51 ? 161.984 161.482 77.034 1.00 55.13  51 J 1 
+ATOM 367 C CB  . LEU J 1 51 ? 162.773 160.118 79.955 1.00 55.13  51 J 1 
+ATOM 368 C CG  . LEU J 1 51 ? 161.865 159.044 79.373 1.00 55.13  51 J 1 
+ATOM 369 C CD1 . LEU J 1 51 ? 162.675 158.158 78.450 1.00 55.13  51 J 1 
+ATOM 370 C CD2 . LEU J 1 51 ? 161.246 158.227 80.487 1.00 55.13  51 J 1 
+ATOM 371 N N   . SER J 1 52 ? 161.739 162.876 78.783 1.00 58.66  52 J 1 
+ATOM 372 C CA  . SER J 1 52 ? 160.769 163.700 78.069 1.00 58.66  52 J 1 
+ATOM 373 C C   . SER J 1 52 ? 161.418 164.423 76.896 1.00 58.66  52 J 1 
+ATOM 374 O O   . SER J 1 52 ? 160.815 164.553 75.825 1.00 58.66  52 J 1 
+ATOM 375 C CB  . SER J 1 52 ? 160.119 164.699 79.023 1.00 58.66  52 J 1 
+ATOM 376 O OG  . SER J 1 52 ? 159.266 164.041 79.940 1.00 58.66  52 J 1 
+ATOM 377 N N   . GLU J 1 53 ? 162.646 164.910 77.081 1.00 65.81  53 J 1 
+ATOM 378 C CA  . GLU J 1 53 ? 163.360 165.549 75.985 1.00 65.81  53 J 1 
+ATOM 379 C C   . GLU J 1 53 ? 163.519 164.593 74.816 1.00 65.81  53 J 1 
+ATOM 380 O O   . GLU J 1 53 ? 163.285 164.966 73.663 1.00 65.81  53 J 1 
+ATOM 381 C CB  . GLU J 1 53 ? 164.729 166.027 76.468 1.00 65.81  53 J 1 
+ATOM 382 C CG  . GLU J 1 53 ? 164.823 167.501 76.834 1.00 65.81  53 J 1 
+ATOM 383 C CD  . GLU J 1 53 ? 164.613 168.437 75.659 1.00 65.81  53 J 1 
+ATOM 384 O OE1 . GLU J 1 53 ? 164.395 167.965 74.524 1.00 65.81  53 J 1 
+ATOM 385 O OE2 . GLU J 1 53 ? 164.659 169.667 75.875 1.00 65.81  53 J 1 
+ATOM 386 N N   . TYR J 1 54 ? 163.901 163.351 75.104 1.00 59.32  54 J 1 
+ATOM 387 C CA  . TYR J 1 54 ? 164.070 162.346 74.065 1.00 59.32  54 J 1 
+ATOM 388 C C   . TYR J 1 54 ? 162.742 162.013 73.396 1.00 59.32  54 J 1 
+ATOM 389 O O   . TYR J 1 54 ? 162.683 161.828 72.175 1.00 59.32  54 J 1 
+ATOM 390 C CB  . TYR J 1 54 ? 164.710 161.104 74.676 1.00 59.32  54 J 1 
+ATOM 391 C CG  . TYR J 1 54 ? 164.986 159.991 73.707 1.00 59.32  54 J 1 
+ATOM 392 C CD1 . TYR J 1 54 ? 166.049 160.070 72.828 1.00 59.32  54 J 1 
+ATOM 393 C CD2 . TYR J 1 54 ? 164.208 158.845 73.697 1.00 59.32  54 J 1 
+ATOM 394 C CE1 . TYR J 1 54 ? 166.316 159.055 71.945 1.00 59.32  54 J 1 
+ATOM 395 C CE2 . TYR J 1 54 ? 164.470 157.819 72.819 1.00 59.32  54 J 1 
+ATOM 396 C CZ  . TYR J 1 54 ? 165.527 157.929 71.944 1.00 59.32  54 J 1 
+ATOM 397 O OH  . TYR J 1 54 ? 165.800 156.912 71.061 1.00 59.32  54 J 1 
+ATOM 398 N N   . ASN J 1 55 ? 161.668 161.931 74.183 1.00 55.92  55 J 1 
+ATOM 399 C CA  . ASN J 1 55 ? 160.344 161.659 73.629 1.00 55.92  55 J 1 
+ATOM 400 C C   . ASN J 1 55 ? 159.925 162.748 72.647 1.00 55.92  55 J 1 
+ATOM 401 O O   . ASN J 1 55 ? 159.524 162.458 71.511 1.00 55.92  55 J 1 
+ATOM 402 C CB  . ASN J 1 55 ? 159.340 161.529 74.777 1.00 55.92  55 J 1 
+ATOM 403 C CG  . ASN J 1 55 ? 157.922 161.278 74.306 1.00 55.92  55 J 1 
+ATOM 404 O OD1 . ASN J 1 55 ? 157.659 161.123 73.117 1.00 55.92  55 J 1 
+ATOM 405 N ND2 . ASN J 1 55 ? 156.998 161.213 75.253 1.00 55.92  55 J 1 
+ATOM 406 N N   . LEU J 1 56 ? 160.012 164.010 73.071 1.00 57.85  56 J 1 
+ATOM 407 C CA  . LEU J 1 56 ? 159.659 165.115 72.183 1.00 57.85  56 J 1 
+ATOM 408 C C   . LEU J 1 56 ? 160.587 165.185 70.980 1.00 57.85  56 J 1 
+ATOM 409 O O   . LEU J 1 56 ? 160.148 165.503 69.871 1.00 57.85  56 J 1 
+ATOM 410 C CB  . LEU J 1 56 ? 159.665 166.437 72.945 1.00 57.85  56 J 1 
+ATOM 411 C CG  . LEU J 1 56 ? 158.428 166.781 73.772 1.00 57.85  56 J 1 
+ATOM 412 C CD1 . LEU J 1 56 ? 158.257 165.920 74.998 1.00 57.85  56 J 1 
+ATOM 413 C CD2 . LEU J 1 56 ? 158.539 168.224 74.179 1.00 57.85  56 J 1 
+ATOM 414 N N   . TYR J 1 57 ? 161.872 164.895 71.176 1.00 60.05  57 J 1 
+ATOM 415 C CA  . TYR J 1 57 ? 162.800 164.844 70.054 1.00 60.05  57 J 1 
+ATOM 416 C C   . TYR J 1 57 ? 162.362 163.824 69.013 1.00 60.05  57 J 1 
+ATOM 417 O O   . TYR J 1 57 ? 162.299 164.134 67.817 1.00 60.05  57 J 1 
+ATOM 418 C CB  . TYR J 1 57 ? 164.203 164.534 70.578 1.00 60.05  57 J 1 
+ATOM 419 C CG  . TYR J 1 57 ? 165.239 164.289 69.513 1.00 60.05  57 J 1 
+ATOM 420 C CD1 . TYR J 1 57 ? 165.328 163.058 68.881 1.00 60.05  57 J 1 
+ATOM 421 C CD2 . TYR J 1 57 ? 166.114 165.288 69.127 1.00 60.05  57 J 1 
+ATOM 422 C CE1 . TYR J 1 57 ? 166.257 162.825 67.909 1.00 60.05  57 J 1 
+ATOM 423 C CE2 . TYR J 1 57 ? 167.052 165.063 68.149 1.00 60.05  57 J 1 
+ATOM 424 C CZ  . TYR J 1 57 ? 167.118 163.828 67.545 1.00 60.05  57 J 1 
+ATOM 425 O OH  . TYR J 1 57 ? 168.044 163.579 66.567 1.00 60.05  57 J 1 
+ATOM 426 N N   . ARG J 1 58 ? 162.069 162.598 69.446 1.00 58.38  58 J 1 
+ATOM 427 C CA  . ARG J 1 58 ? 161.680 161.562 68.498 1.00 58.38  58 J 1 
+ATOM 428 C C   . ARG J 1 58 ? 160.373 161.916 67.804 1.00 58.38  58 J 1 
+ATOM 429 O O   . ARG J 1 58 ? 160.238 161.726 66.587 1.00 58.38  58 J 1 
+ATOM 430 C CB  . ARG J 1 58 ? 161.566 160.217 69.208 1.00 58.38  58 J 1 
+ATOM 431 C CG  . ARG J 1 58 ? 162.900 159.553 69.491 1.00 58.38  58 J 1 
+ATOM 432 C CD  . ARG J 1 58 ? 163.871 159.707 68.335 1.00 58.38  58 J 1 
+ATOM 433 N NE  . ARG J 1 58 ? 165.241 159.437 68.753 1.00 58.38  58 J 1 
+ATOM 434 C CZ  . ARG J 1 58 ? 166.293 159.472 67.948 1.00 58.38  58 J 1 
+ATOM 435 N NH1 . ARG J 1 58 ? 166.171 159.760 66.663 1.00 58.38  58 J 1 
+ATOM 436 N NH2 . ARG J 1 58 ? 167.498 159.214 68.445 1.00 58.38  58 J 1 
+ATOM 437 N N   . ASN J 1 59 ? 159.399 162.432 68.558 1.00 58.16  59 J 1 
+ATOM 438 C CA  . ASN J 1 59 ? 158.133 162.817 67.945 1.00 58.16  59 J 1 
+ATOM 439 C C   . ASN J 1 59 ? 158.327 163.928 66.924 1.00 58.16  59 J 1 
+ATOM 440 O O   . ASN J 1 59 ? 157.735 163.893 65.839 1.00 58.16  59 J 1 
+ATOM 441 C CB  . ASN J 1 59 ? 157.133 163.240 69.017 1.00 58.16  59 J 1 
+ATOM 442 C CG  . ASN J 1 59 ? 156.401 162.065 69.619 1.00 58.16  59 J 1 
+ATOM 443 O OD1 . ASN J 1 59 ? 156.905 161.403 70.521 1.00 58.16  59 J 1 
+ATOM 444 N ND2 . ASN J 1 59 ? 155.203 161.796 69.117 1.00 58.16  59 J 1 
+ATOM 445 N N   . ALA J 1 60 ? 159.154 164.925 67.250 1.00 62.06  60 J 1 
+ATOM 446 C CA  . ALA J 1 60 ? 159.407 166.012 66.314 1.00 62.06  60 J 1 
+ATOM 447 C C   . ALA J 1 60 ? 160.081 165.504 65.052 1.00 62.06  60 J 1 
+ATOM 448 O O   . ALA J 1 60 ? 159.709 165.896 63.940 1.00 62.06  60 J 1 
+ATOM 449 C CB  . ALA J 1 60 ? 160.257 167.092 66.980 1.00 62.06  60 J 1 
+ATOM 450 N N   . GLN J 1 61 ? 161.067 164.621 65.203 1.00 65.39  61 J 1 
+ATOM 451 C CA  . GLN J 1 61 ? 161.748 164.072 64.037 1.00 65.39  61 J 1 
+ATOM 452 C C   . GLN J 1 61 ? 160.774 163.310 63.148 1.00 65.39  61 J 1 
+ATOM 453 O O   . GLN J 1 61 ? 160.735 163.512 61.925 1.00 65.39  61 J 1 
+ATOM 454 C CB  . GLN J 1 61 ? 162.895 163.171 64.494 1.00 65.39  61 J 1 
+ATOM 455 C CG  . GLN J 1 61 ? 164.044 163.067 63.519 1.00 65.39  61 J 1 
+ATOM 456 C CD  . GLN J 1 61 ? 165.192 162.262 64.080 1.00 65.39  61 J 1 
+ATOM 457 O OE1 . GLN J 1 61 ? 166.240 162.807 64.420 1.00 65.39  61 J 1 
+ATOM 458 N NE2 . GLN J 1 61 ? 164.998 160.957 64.189 1.00 65.39  61 J 1 
+ATOM 459 N N   . SER J 1 62 ? 159.959 162.442 63.753 1.00 64.39  62 J 1 
+ATOM 460 C CA  . SER J 1 62 ? 159.025 161.644 62.969 1.00 64.39  62 J 1 
+ATOM 461 C C   . SER J 1 62 ? 157.998 162.522 62.267 1.00 64.39  62 J 1 
+ATOM 462 O O   . SER J 1 62 ? 157.716 162.331 61.078 1.00 64.39  62 J 1 
+ATOM 463 C CB  . SER J 1 62 ? 158.333 160.619 63.865 1.00 64.39  62 J 1 
+ATOM 464 O OG  . SER J 1 62 ? 157.410 159.843 63.125 1.00 64.39  62 J 1 
+ATOM 465 N N   . ASN J 1 63 ? 157.438 163.504 62.976 1.00 65.16  63 J 1 
+ATOM 466 C CA  . ASN J 1 63 ? 156.437 164.361 62.357 1.00 65.16  63 J 1 
+ATOM 467 C C   . ASN J 1 63 ? 157.033 165.219 61.252 1.00 65.16  63 J 1 
+ATOM 468 O O   . ASN J 1 63 ? 156.389 165.421 60.216 1.00 65.16  63 J 1 
+ATOM 469 C CB  . ASN J 1 63 ? 155.767 165.233 63.413 1.00 65.16  63 J 1 
+ATOM 470 C CG  . ASN J 1 63 ? 154.853 164.439 64.316 1.00 65.16  63 J 1 
+ATOM 471 O OD1 . ASN J 1 63 ? 154.525 163.290 64.026 1.00 65.16  63 J 1 
+ATOM 472 N ND2 . ASN J 1 63 ? 154.434 165.045 65.416 1.00 65.16  63 J 1 
+ATOM 473 N N   . THR J 1 64 ? 158.252 165.729 61.444 1.00 66.60  64 J 1 
+ATOM 474 C CA  . THR J 1 64 ? 158.875 166.543 60.408 1.00 66.60  64 J 1 
+ATOM 475 C C   . THR J 1 64 ? 159.147 165.725 59.153 1.00 66.60  64 J 1 
+ATOM 476 O O   . THR J 1 64 ? 158.846 166.169 58.035 1.00 66.60  64 J 1 
+ATOM 477 C CB  . THR J 1 64 ? 160.169 167.158 60.935 1.00 66.60  64 J 1 
+ATOM 478 O OG1 . THR J 1 64 ? 159.877 167.977 62.072 1.00 66.60  64 J 1 
+ATOM 479 C CG2 . THR J 1 64 ? 160.819 168.011 59.865 1.00 66.60  64 J 1 
+ATOM 480 N N   . VAL J 1 65 ? 159.707 164.523 59.310 1.00 64.00  65 J 1 
+ATOM 481 C CA  . VAL J 1 65 ? 159.978 163.718 58.124 1.00 64.00  65 J 1 
+ATOM 482 C C   . VAL J 1 65 ? 158.671 163.316 57.452 1.00 64.00  65 J 1 
+ATOM 483 O O   . VAL J 1 65 ? 158.580 163.285 56.220 1.00 64.00  65 J 1 
+ATOM 484 C CB  . VAL J 1 65 ? 160.858 162.499 58.464 1.00 64.00  65 J 1 
+ATOM 485 C CG1 . VAL J 1 65 ? 162.203 162.957 58.999 1.00 64.00  65 J 1 
+ATOM 486 C CG2 . VAL J 1 65 ? 160.176 161.571 59.442 1.00 64.00  65 J 1 
+ATOM 487 N N   . LYS J 1 66 ? 157.626 163.047 58.241 1.00 65.12  66 J 1 
+ATOM 488 C CA  . LYS J 1 66 ? 156.341 162.691 57.649 1.00 65.12  66 J 1 
+ATOM 489 C C   . LYS J 1 66 ? 155.745 163.847 56.855 1.00 65.12  66 J 1 
+ATOM 490 O O   . LYS J 1 66 ? 155.240 163.645 55.745 1.00 65.12  66 J 1 
+ATOM 491 C CB  . LYS J 1 66 ? 155.366 162.235 58.733 1.00 65.12  66 J 1 
+ATOM 492 C CG  . LYS J 1 66 ? 153.992 161.855 58.202 1.00 65.12  66 J 1 
+ATOM 493 C CD  . LYS J 1 66 ? 154.088 160.891 57.032 1.00 65.12  66 J 1 
+ATOM 494 C CE  . LYS J 1 66 ? 152.711 160.542 56.496 1.00 65.12  66 J 1 
+ATOM 495 N NZ  . LYS J 1 66 ? 151.754 160.214 57.587 1.00 65.12  66 J 1 
+ATOM 496 N N   . VAL J 1 67 ? 155.779 165.064 57.402 1.00 65.92  67 J 1 
+ATOM 497 C CA  . VAL J 1 67 ? 155.163 166.185 56.696 1.00 65.92  67 J 1 
+ATOM 498 C C   . VAL J 1 67 ? 155.958 166.528 55.442 1.00 65.92  67 J 1 
+ATOM 499 O O   . VAL J 1 67 ? 155.381 166.863 54.397 1.00 65.92  67 J 1 
+ATOM 500 C CB  . VAL J 1 67 ? 154.993 167.401 57.629 1.00 65.92  67 J 1 
+ATOM 501 C CG1 . VAL J 1 67 ? 156.328 167.937 58.097 1.00 65.92  67 J 1 
+ATOM 502 C CG2 . VAL J 1 67 ? 154.200 168.492 56.929 1.00 65.92  67 J 1 
+ATOM 503 N N   . PHE J 1 68 ? 157.288 166.435 55.509 1.00 66.56  68 J 1 
+ATOM 504 C CA  . PHE J 1 68 ? 158.076 166.682 54.308 1.00 66.56  68 J 1 
+ATOM 505 C C   . PHE J 1 68 ? 157.826 165.606 53.260 1.00 66.56  68 J 1 
+ATOM 506 O O   . PHE J 1 68 ? 157.772 165.899 52.059 1.00 66.56  68 J 1 
+ATOM 507 C CB  . PHE J 1 68 ? 159.558 166.785 54.657 1.00 66.56  68 J 1 
+ATOM 508 C CG  . PHE J 1 68 ? 159.991 168.177 55.012 1.00 66.56  68 J 1 
+ATOM 509 C CD1 . PHE J 1 68 ? 160.099 168.568 56.333 1.00 66.56  68 J 1 
+ATOM 510 C CD2 . PHE J 1 68 ? 160.268 169.101 54.022 1.00 66.56  68 J 1 
+ATOM 511 C CE1 . PHE J 1 68 ? 160.490 169.850 56.660 1.00 66.56  68 J 1 
+ATOM 512 C CE2 . PHE J 1 68 ? 160.658 170.384 54.343 1.00 66.56  68 J 1 
+ATOM 513 C CZ  . PHE J 1 68 ? 160.769 170.758 55.664 1.00 66.56  68 J 1 
+ATOM 514 N N   . LYS J 1 69 ? 157.646 164.356 53.694 1.00 70.04  69 J 1 
+ATOM 515 C CA  . LYS J 1 69 ? 157.283 163.302 52.755 1.00 70.04  69 J 1 
+ATOM 516 C C   . LYS J 1 69 ? 155.929 163.584 52.123 1.00 70.04  69 J 1 
+ATOM 517 O O   . LYS J 1 69 ? 155.731 163.326 50.934 1.00 70.04  69 J 1 
+ATOM 518 C CB  . LYS J 1 69 ? 157.291 161.944 53.465 1.00 70.04  69 J 1 
+ATOM 519 C CG  . LYS J 1 69 ? 157.067 160.718 52.570 1.00 70.04  69 J 1 
+ATOM 520 C CD  . LYS J 1 69 ? 155.593 160.436 52.289 1.00 70.04  69 J 1 
+ATOM 521 C CE  . LYS J 1 69 ? 155.424 159.353 51.239 1.00 70.04  69 J 1 
+ATOM 522 N NZ  . LYS J 1 69 ? 153.999 158.953 51.085 1.00 70.04  69 J 1 
+ATOM 523 N N   . ASP J 1 70 ? 154.979 164.093 52.908 1.00 70.15  70 J 1 
+ATOM 524 C CA  . ASP J 1 70 ? 153.667 164.418 52.356 1.00 70.15  70 J 1 
+ATOM 525 C C   . ASP J 1 70 ? 153.770 165.503 51.293 1.00 70.15  70 J 1 
+ATOM 526 O O   . ASP J 1 70 ? 153.172 165.386 50.216 1.00 70.15  70 J 1 
+ATOM 527 C CB  . ASP J 1 70 ? 152.719 164.850 53.472 1.00 70.15  70 J 1 
+ATOM 528 C CG  . ASP J 1 70 ? 152.353 163.710 54.396 1.00 70.15  70 J 1 
+ATOM 529 O OD1 . ASP J 1 70 ? 152.337 162.552 53.930 1.00 70.15  70 J 1 
+ATOM 530 O OD2 . ASP J 1 70 ? 152.082 163.970 55.587 1.00 70.15  70 J 1 
+ATOM 531 N N   . ILE J 1 71 ? 154.538 166.558 51.574 1.00 64.31  71 J 1 
+ATOM 532 C CA  . ILE J 1 71 ? 154.721 167.619 50.588 1.00 64.31  71 J 1 
+ATOM 533 C C   . ILE J 1 71 ? 155.375 167.070 49.326 1.00 64.31  71 J 1 
+ATOM 534 O O   . ILE J 1 71 ? 154.941 167.362 48.203 1.00 64.31  71 J 1 
+ATOM 535 C CB  . ILE J 1 71 ? 155.545 168.772 51.188 1.00 64.31  71 J 1 
+ATOM 536 C CG1 . ILE J 1 71 ? 154.860 169.327 52.435 1.00 64.31  71 J 1 
+ATOM 537 C CG2 . ILE J 1 71 ? 155.751 169.866 50.159 1.00 64.31  71 J 1 
+ATOM 538 C CD1 . ILE J 1 71 ? 155.778 170.127 53.325 1.00 64.31  71 J 1 
+ATOM 539 N N   . ASP J 1 72 ? 156.423 166.261 49.492 1.00 68.85  72 J 1 
+ATOM 540 C CA  . ASP J 1 72 ? 157.144 165.727 48.343 1.00 68.85  72 J 1 
+ATOM 541 C C   . ASP J 1 72 ? 156.266 164.795 47.517 1.00 68.85  72 J 1 
+ATOM 542 O O   . ASP J 1 72 ? 156.316 164.814 46.282 1.00 68.85  72 J 1 
+ATOM 543 C CB  . ASP J 1 72 ? 158.398 165.005 48.825 1.00 68.85  72 J 1 
+ATOM 544 C CG  . ASP J 1 72 ? 159.551 165.947 49.062 1.00 68.85  72 J 1 
+ATOM 545 O OD1 . ASP J 1 72 ? 160.691 165.462 49.183 1.00 68.85  72 J 1 
+ATOM 546 O OD2 . ASP J 1 72 ? 159.318 167.171 49.125 1.00 68.85  72 J 1 
+ATOM 547 N N   . ALA J 1 73 ? 155.464 163.963 48.181 1.00 65.97  73 J 1 
+ATOM 548 C CA  . ALA J 1 73 ? 154.553 163.081 47.466 1.00 65.97  73 J 1 
+ATOM 549 C C   . ALA J 1 73 ? 153.501 163.878 46.715 1.00 65.97  73 J 1 
+ATOM 550 O O   . ALA J 1 73 ? 153.128 163.520 45.593 1.00 65.97  73 J 1 
+ATOM 551 C CB  . ALA J 1 73 ? 153.894 162.103 48.437 1.00 65.97  73 J 1 
+ATOM 552 N N   . ALA J 1 74 ? 153.004 164.960 47.319 1.00 64.99  74 J 1 
+ATOM 553 C CA  . ALA J 1 74 ? 152.063 165.819 46.610 1.00 64.99  74 J 1 
+ATOM 554 C C   . ALA J 1 74 ? 152.700 166.405 45.358 1.00 64.99  74 J 1 
+ATOM 555 O O   . ALA J 1 74 ? 152.080 166.428 44.288 1.00 64.99  74 J 1 
+ATOM 556 C CB  . ALA J 1 74 ? 151.567 166.931 47.533 1.00 64.99  74 J 1 
+ATOM 557 N N   . ILE J 1 75 ? 153.947 166.864 45.469 1.00 67.26  75 J 1 
+ATOM 558 C CA  . ILE J 1 75 ? 154.634 167.425 44.307 1.00 67.26  75 J 1 
+ATOM 559 C C   . ILE J 1 75 ? 154.797 166.367 43.222 1.00 67.26  75 J 1 
+ATOM 560 O O   . ILE J 1 75 ? 154.522 166.615 42.042 1.00 67.26  75 J 1 
+ATOM 561 C CB  . ILE J 1 75 ? 155.993 168.017 44.717 1.00 67.26  75 J 1 
+ATOM 562 C CG1 . ILE J 1 75 ? 155.807 169.113 45.765 1.00 67.26  75 J 1 
+ATOM 563 C CG2 . ILE J 1 75 ? 156.721 168.564 43.502 1.00 67.26  75 J 1 
+ATOM 564 C CD1 . ILE J 1 75 ? 154.983 170.274 45.286 1.00 67.26  75 J 1 
+ATOM 565 N N   . ILE J 1 76 ? 155.244 165.170 43.607 1.00 71.16  76 J 1 
+ATOM 566 C CA  . ILE J 1 76 ? 155.503 164.119 42.627 1.00 71.16  76 J 1 
+ATOM 567 C C   . ILE J 1 76 ? 154.216 163.694 41.935 1.00 71.16  76 J 1 
+ATOM 568 O O   . ILE J 1 76 ? 154.164 163.579 40.705 1.00 71.16  76 J 1 
+ATOM 569 C CB  . ILE J 1 76 ? 156.205 162.925 43.299 1.00 71.16  76 J 1 
+ATOM 570 C CG1 . ILE J 1 76 ? 157.602 163.325 43.761 1.00 71.16  76 J 1 
+ATOM 571 C CG2 . ILE J 1 76 ? 156.302 161.755 42.342 1.00 71.16  76 J 1 
+ATOM 572 C CD1 . ILE J 1 76 ? 158.489 163.764 42.631 1.00 71.16  76 J 1 
+ATOM 573 N N   . GLN J 1 77 ? 153.158 163.457 42.708 1.00 74.01  77 J 1 
+ATOM 574 C CA  . GLN J 1 77 ? 151.870 163.093 42.139 1.00 74.01  77 J 1 
+ATOM 575 C C   . GLN J 1 77 ? 151.276 164.203 41.291 1.00 74.01  77 J 1 
+ATOM 576 O O   . GLN J 1 77 ? 150.522 163.917 40.355 1.00 74.01  77 J 1 
+ATOM 577 C CB  . GLN J 1 77 ? 150.903 162.718 43.266 1.00 74.01  77 J 1 
+ATOM 578 C CG  . GLN J 1 77 ? 149.608 162.063 42.818 1.00 74.01  77 J 1 
+ATOM 579 C CD  . GLN J 1 77 ? 148.547 163.075 42.436 1.00 74.01  77 J 1 
+ATOM 580 O OE1 . GLN J 1 77 ? 148.119 163.881 43.261 1.00 74.01  77 J 1 
+ATOM 581 N NE2 . GLN J 1 77 ? 148.126 163.046 41.178 1.00 74.01  77 J 1 
+ATOM 582 N N   . ASN J 1 78 ? 151.610 165.453 41.583 1.00 75.47  78 J 1 
+ATOM 583 C CA  . ASN J 1 78 ? 151.007 166.584 40.903 1.00 75.47  78 J 1 
+ATOM 584 C C   . ASN J 1 78 ? 151.656 166.793 39.532 1.00 75.47  78 J 1 
+ATOM 585 O O   . ASN J 1 78 ? 151.100 167.503 38.692 1.00 75.47  78 J 1 
+ATOM 586 C CB  . ASN J 1 78 ? 151.109 167.797 41.862 1.00 75.47  78 J 1 
+ATOM 587 C CG  . ASN J 1 78 ? 150.672 169.138 41.268 1.00 75.47  78 J 1 
+ATOM 588 O OD1 . ASN J 1 78 ? 150.393 169.290 40.093 1.00 75.47  78 J 1 
+ATOM 589 N ND2 . ASN J 1 78 ? 150.555 170.120 42.149 1.00 75.47  78 J 1 
+ATOM 590 N N   . PHE J 1 79 ? 152.755 166.084 39.244 1.00 76.77  79 J 1 
+ATOM 591 C CA  . PHE J 1 79 ? 153.341 166.125 37.905 1.00 76.77  79 J 1 
+ATOM 592 C C   . PHE J 1 79 ? 152.308 165.775 36.843 1.00 76.77  79 J 1 
+ATOM 593 O O   . PHE J 1 79 ? 152.188 166.465 35.824 1.00 76.77  79 J 1 
+ATOM 594 C CB  . PHE J 1 79 ? 154.513 165.146 37.794 1.00 76.77  79 J 1 
+ATOM 595 C CG  . PHE J 1 79 ? 155.770 165.599 38.471 1.00 76.77  79 J 1 
+ATOM 596 C CD1 . PHE J 1 79 ? 155.885 166.866 39.000 1.00 76.77  79 J 1 
+ATOM 597 C CD2 . PHE J 1 79 ? 156.855 164.744 38.553 1.00 76.77  79 J 1 
+ATOM 598 C CE1 . PHE J 1 79 ? 157.054 167.261 39.614 1.00 76.77  79 J 1 
+ATOM 599 C CE2 . PHE J 1 79 ? 158.021 165.138 39.162 1.00 76.77  79 J 1 
+ATOM 600 C CZ  . PHE J 1 79 ? 158.121 166.398 39.692 1.00 76.77  79 J 1 
+ATOM 601 N N   . ARG J 1 80 ? 151.557 164.702 37.066 1.00 84.01  80 J 1 
+ATOM 602 C CA  . ARG J 1 80 ? 150.552 164.246 36.118 1.00 84.01  80 J 1 
+ATOM 603 C C   . ARG J 1 80 ? 149.236 164.973 36.361 1.00 84.01  80 J 1 
+ATOM 604 O O   . ARG J 1 80 ? 148.856 165.220 37.506 1.00 84.01  80 J 1 
+ATOM 605 C CB  . ARG J 1 80 ? 150.365 162.730 36.232 1.00 84.01  80 J 1 
+ATOM 606 C CG  . ARG J 1 80 ? 149.526 162.108 35.127 1.00 84.01  80 J 1 
+ATOM 607 C CD  . ARG J 1 80 ? 148.094 161.873 35.581 1.00 84.01  80 J 1 
+ATOM 608 N NE  . ARG J 1 80 ? 148.029 160.982 36.733 1.00 84.01  80 J 1 
+ATOM 609 C CZ  . ARG J 1 80 ? 146.964 160.838 37.510 1.00 84.01  80 J 1 
+ATOM 610 N NH1 . ARG J 1 80 ? 145.848 161.512 37.286 1.00 84.01  80 J 1 
+ATOM 611 N NH2 . ARG J 1 80 ? 147.022 159.999 38.541 1.00 84.01  80 J 1 
+ATOM 612 O OXT . ARG J 1 80 ? 148.527 165.332 35.421 1.00 84.01  80 J 1 
+#

flax_model/alphafold3/test_data/model_config.json

@@ -0,0 +1,231 @@
+{
+  "evoformer": {
+    "max_relative_chain": 2,
+    "max_relative_idx": 32,
+    "msa_channel": 64,
+    "msa_stack": {
+      "msa_attention": {
+        "num_head": 8
+      },
+      "msa_transition": {
+        "num_intermediate_factor": 4,
+        "use_glu_kernel": true
+      },
+      "num_layer": 4,
+      "outer_product_mean": {
+        "chunk_size": 128,
+        "num_outer_channel": 32
+      },
+      "pair_attention": {
+        "num_head": 4
+      },
+      "pair_transition": {
+        "num_intermediate_factor": 4,
+        "use_glu_kernel": true
+      },
+      "shard_transition_blocks": true,
+      "triangle_multiplication_incoming": {
+        "equation": "kjc,kic->ijc",
+        "use_glu_kernel": true
+      },
+      "triangle_multiplication_outgoing": {
+        "equation": "ikc,jkc->ijc",
+        "use_glu_kernel": true
+      }
+    },
+    "num_msa": 1024,
+    "pair_channel": 128,
+    "pairformer": {
+      "block_remat": false,
+      "num_layer": 48,
+      "pair_attention": {
+        "num_head": 4
+      },
+      "pair_transition": {
+        "num_intermediate_factor": 4,
+        "use_glu_kernel": true
+      },
+      "remat_block_size": 8,
+      "shard_transition_blocks": true,
+      "single_attention": {
+        "key_dim": null,
+        "num_head": 16,
+        "value_dim": null
+      },
+      "single_transition": {
+        "num_intermediate_factor": 4,
+        "use_glu_kernel": true
+      },
+      "triangle_multiplication_incoming": {
+        "equation": "kjc,kic->ijc",
+        "use_glu_kernel": true
+      },
+      "triangle_multiplication_outgoing": {
+        "equation": "ikc,jkc->ijc",
+        "use_glu_kernel": true
+      }
+    },
+    "per_atom_conditioning": {
+      "atom_transformer": {
+        "attention": {
+          "key_dim": 128,
+          "num_head": 4,
+          "value_dim": 128
+        },
+        "num_blocks": 3,
+        "num_intermediate_factor": 2
+      },
+      "per_atom_channels": 128,
+      "per_atom_pair_channels": 16,
+      "per_token_channels": 384
+    },
+    "seq_channel": 384,
+    "template": {
+      "dgram_features": {
+        "max_bin": 50.75,
+        "min_bin": 3.25,
+        "num_bins": 39
+      },
+      "num_channels": 64,
+      "template_stack": {
+        "num_layer": 2,
+        "pair_attention": {
+          "num_head": 4
+        },
+        "pair_transition": {
+          "num_intermediate_factor": 2,
+          "use_glu_kernel": true
+        },
+        "shard_transition_blocks": true,
+        "single_attention": null,
+        "single_transition": null,
+        "triangle_multiplication_incoming": {
+          "equation": "kjc,kic->ijc",
+          "use_glu_kernel": true
+        },
+        "triangle_multiplication_outgoing": {
+          "equation": "ikc,jkc->ijc",
+          "use_glu_kernel": true
+        }
+      }
+    }
+  },
+  "global_config": {
+    "bfloat16": "all",
+    "final_init": "zeros",
+    "flash_attention_implementation": "triton",
+    "pair_attention_chunk_size": [
+      [
+        1536,
+        128
+      ],
+      [
+        null,
+        32
+      ]
+    ],
+    "pair_transition_shard_spec": [
+      [
+        2048,
+        null
+      ],
+      [
+        null,
+        1024
+      ]
+    ]
+  },
+  "heads": {
+    "confidence": {
+      "dgram_features": {
+        "max_bin": 50.75,
+        "min_bin": 3.25,
+        "num_bins": 39
+      },
+      "max_error_bin": 31.0,
+      "no_embedding_prob": 0.2,
+      "num_bins": 64,
+      "num_plddt_bins": 50,
+      "pae": {
+        "max_error_bin": 31.0,
+        "num_bins": 64
+      },
+      "pairformer": {
+        "num_layer": 4,
+        "pair_attention": {
+          "num_head": 4
+        },
+        "pair_transition": {
+          "num_intermediate_factor": 4,
+          "use_glu_kernel": true
+        },
+        "shard_transition_blocks": true,
+        "single_attention": {
+          "key_dim": null,
+          "num_head": 16,
+          "value_dim": null
+        },
+        "single_transition": {
+          "num_intermediate_factor": 4,
+          "use_glu_kernel": true
+        },
+        "triangle_multiplication_incoming": {
+          "equation": "kjc,kic->ijc",
+          "use_glu_kernel": true
+        },
+        "triangle_multiplication_outgoing": {
+          "equation": "ikc,jkc->ijc",
+          "use_glu_kernel": true
+        }
+      }
+    },
+    "diffusion": {
+      "atom_transformer": {
+        "attention": {
+          "key_dim": 128,
+          "num_head": 4,
+          "value_dim": 128
+        },
+        "num_blocks": 3,
+        "num_intermediate_factor": 2
+      },
+      "conditioning": {
+        "pair_channel": 128,
+        "prob": 0.8,
+        "seq_channel": 384
+      },
+      "eval": {
+        "gamma_0": 0.8,
+        "gamma_min": 1.0,
+        "noise_scale": 1.003,
+        "num_samples": 5,
+        "step_scale": 1.5,
+        "steps": 200
+      },
+      "eval_batch_dim_shard_size": 5,
+      "eval_batch_size": 5,
+      "per_atom_channels": 128,
+      "per_atom_pair_channels": 16,
+      "per_token_channels": 768,
+      "transformer": {
+        "attention": {
+          "key_dim": null,
+          "num_head": 16,
+          "value_dim": null
+        },
+        "block_remat": false,
+        "num_blocks": 24,
+        "num_intermediate_factor": 2,
+        "super_block_size": 4
+      }
+    },
+    "distogram": {
+      "first_break": 2.3125,
+      "last_break": 21.6875,
+      "num_bins": 64
+    }
+  },
+  "num_recycles": 10,
+  "return_distogram": false,
+  "return_embeddings": false
+}

flax_model/alphafold3/version.py

@@ -0,0 +1,13 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+#
+# To request access to the AlphaFold 3 model parameters, follow the process set
+# out at https://github.com/google-deepmind/alphafold3. You may only use these
+# if received directly from Google. Use is subject to terms of use available at
+# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
+
+"""Single source of truth for the AlphaFold version."""
+
+__version__ = '3.0.1' 

inputs/t1119_search.json

@@ -0,0 +1,19 @@
+{
+  "dialect": "alphafold3",
+  "version": 1,
+  "name": "T1119",
+  "sequences": [
+    {
+      "protein": {
+        "id": "A",
+        "sequence": "QLEDSEVEAVAKGLEEMYANGVTEDNFKNYVKNNFAQQEISSVEEELNVNISDSCVANKIKDEFFAMISISAIVKAAQKKAWKELAVTVLRFAKANGLKTNAIIVAGQLALWAVQCG"
+      }
+    }
+  ],
+  "modelSeeds": [
+    100
+  ],
+  "bondedAtomPairs": null,
+  "userCCD": null
+}
+

scripts/infer.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+
+
+export TF_CPP_MIN_LOG_LEVEL="${TF_CPP_MIN_LOG_LEVEL:-0}"
+export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.95}"
+export TRITON_ENABLE_GLOBAL_TO_LOCAL="${TRITON_ENABLE_GLOBAL_TO_LOCAL:-1}"
+export TRITON_USE_MAKE_BLOCK_PTR="${TRITON_USE_MAKE_BLOCK_PTR:-1}"
+export TRITON_DEFAULT_ENABLE_NUM_VGPRS512="${TRITON_DEFAULT_ENABLE_NUM_VGPRS512:-1}"
+
+MODEL_ROOT="${ONESCIENCE_MODELS_DIR:-${PROJECT_ROOT}/weight}"
+MODEL_DIR="${ALPHAFOLD3_MODEL_DIR:-${MODEL_ROOT}}"
+JSON_PATH="${ALPHAFOLD3_JSON_PATH:-${PROJECT_ROOT}/inputs/7r6r_data.json}"
+OUTPUT_DIR="${ALPHAFOLD3_OUTPUT_DIR:-${PROJECT_ROOT}/outputs}"
+FLASH_ATTENTION="${ALPHAFOLD3_FLASH_ATTENTION:-triton}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+python "${PROJECT_ROOT}/scripts/run_alphafold.py" \
+  --json_path="${JSON_PATH}" \
+  --model_dir="${MODEL_DIR}" \
+  --output_dir="${OUTPUT_DIR}" \
+  --run_data_pipeline=false \
+  --flash_attention_implementation="${FLASH_ATTENTION}"

scripts/infer_jackhmmer.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+
+export PATH="${PROJECT_ROOT}/flax_model/alphafold3/_tools/hmmer/bin:${PATH}"
+export TF_CPP_MIN_LOG_LEVEL="${TF_CPP_MIN_LOG_LEVEL:-0}"
+export JAX_TRACEBACK_FILTERING="${JAX_TRACEBACK_FILTERING:-off}"
+export XLA_CLIENT_MEM_FRACTION="${XLA_CLIENT_MEM_FRACTION:-0.95}"
+export TRITON_ENABLE_GLOBAL_TO_LOCAL="${TRITON_ENABLE_GLOBAL_TO_LOCAL:-1}"
+export TRITON_USE_MAKE_BLOCK_PTR="${TRITON_USE_MAKE_BLOCK_PTR:-1}"
+export TRITON_DEFAULT_ENABLE_NUM_VGPRS512="${TRITON_DEFAULT_ENABLE_NUM_VGPRS512:-1}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+
+MODEL_ROOT="${ONESCIENCE_MODELS_DIR:-${PROJECT_ROOT}/weight}"
+MODEL_DIR="${ALPHAFOLD3_MODEL_DIR:-${MODEL_ROOT}}"
+DATASET_ROOT="${ALPHAFOLD3_DATASET_ROOT:-${ONESCIENCE_DATASETS_DIR:-${PROJECT_ROOT}/data}/alphafold3}"
+DB_DIRS="${ALPHAFOLD3_DB_DIR:-${DATASET_ROOT}/public_databases}"
+MMSEQS_DB_DIRS="${ALPHAFOLD3_MMSEQS_DB_DIR:-${DATASET_ROOT}/mmseqsDB}"
+JSON_PATH="${ALPHAFOLD3_JSON_PATH:-${PROJECT_ROOT}/inputs/t1119_search.json}"
+OUTPUT_DIR="${ALPHAFOLD3_OUTPUT_DIR:-${PROJECT_ROOT}/outputs}"
+RUN_INFERENCE="${ALPHAFOLD3_RUN_INFERENCE:-false}"
+FLASH_ATTENTION="${ALPHAFOLD3_FLASH_ATTENTION:-cutlass}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+python "${PROJECT_ROOT}/scripts/run_alphafold.py" \
+  --json_path="${JSON_PATH}" \
+  --model_dir="${MODEL_DIR}" \
+  --output_dir="${OUTPUT_DIR}" \
+  --run_data_pipeline=true \
+  --run_inference="${RUN_INFERENCE}" \
+  --flash_attention_implementation="${FLASH_ATTENTION}" \
+  --use_mmseqs=false \
+  --use_mmseqs_gpu=false \
+  --db_dir="${DB_DIRS}" \
+  --mmseqs_db_dir="${MMSEQS_DB_DIRS}" \
+  --small_bfd_database_path="${DATASET_ROOT}/jackhmmer_split/bfd-first_non_consensus_sequences.fasta@64" \
+  --small_bfd_z_value=65928866 \
+  --mgnify_database_path="${DATASET_ROOT}/jackhmmer_split/mgy_clusters_2022_05.fa@512" \
+  --mgnify_z_value=623796864 \
+  --uniprot_cluster_annot_database_path="${DATASET_ROOT}/jackhmmer_split/uniprot_cluster_annot_2021_04.fa@256" \
+  --uniprot_cluster_annot_z_value=225619586 \
+  --uniref90_database_path="${DATASET_ROOT}/jackhmmer_split/uniref90_2022_05.fa@128" \
+  --uniref90_z_value=153742194 \
+  --jackhmmer_n_cpu="${ALPHAFOLD3_JACKHMMER_N_CPU:-2}" \
+  --jackhmmer_max_parallel_shards="${ALPHAFOLD3_JACKHMMER_MAX_PARALLEL_SHARDS:-1}" \
+  --jackhmmer_max_threads="${ALPHAFOLD3_JACKHMMER_MAX_THREADS:-8}" \
+  --nhmmer_n_cpu="${ALPHAFOLD3_NHMMER_N_CPU:-2}" \
+  --nhmmer_max_parallel_shards="${ALPHAFOLD3_NHMMER_MAX_PARALLEL_SHARDS:-1}" \
+  --nhmmer_max_threads="${ALPHAFOLD3_NHMMER_MAX_THREADS:-8}"

scripts/infer_mmseqs.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+MODEL_ROOT="${ONESCIENCE_MODELS_DIR:-${PROJECT_ROOT}/weight}"
+MODEL_DIR="${ALPHAFOLD3_MODEL_DIR:-${MODEL_ROOT}}"
+MMSEQS_HOME="${ALPHAFOLD3_MMSEQS_HOME:-${MODEL_DIR}/mmseqs}"
+export PATH="${MMSEQS_HOME}/bin:${PATH}"
+export LD_LIBRARY_PATH="${MMSEQS_HOME}/lib:${LD_LIBRARY_PATH:-}"
+export TF_CPP_MIN_LOG_LEVEL="${TF_CPP_MIN_LOG_LEVEL:-0}"
+export JAX_TRACEBACK_FILTERING="${JAX_TRACEBACK_FILTERING:-off}"
+export XLA_PYTHON_CLIENT_ALLOCATOR="${XLA_PYTHON_CLIENT_ALLOCATOR:-platform}"
+export TRITON_ENABLE_GLOBAL_TO_LOCAL="${TRITON_ENABLE_GLOBAL_TO_LOCAL:-1}"
+export TRITON_USE_MAKE_BLOCK_PTR="${TRITON_USE_MAKE_BLOCK_PTR:-1}"
+export TRITON_DEFAULT_ENABLE_NUM_VGPRS512="${TRITON_DEFAULT_ENABLE_NUM_VGPRS512:-1}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+
+DATASET_ROOT="${ALPHAFOLD3_DATASET_ROOT:-${ONESCIENCE_DATASETS_DIR:-${PROJECT_ROOT}/data}/alphafold3}"
+DB_DIRS="${ALPHAFOLD3_DB_DIR:-${DATASET_ROOT}/public_databases}"
+MMSEQS_DB_DIRS="${ALPHAFOLD3_MMSEQS_DB_DIR:-${DATASET_ROOT}/mmseqsDB}"
+JSON_PATH="${ALPHAFOLD3_JSON_PATH:-${PROJECT_ROOT}/inputs/t1119_search.json}"
+OUTPUT_DIR="${ALPHAFOLD3_OUTPUT_DIR:-${PROJECT_ROOT}/outputs}"
+RUN_INFERENCE="${ALPHAFOLD3_RUN_INFERENCE:-false}"
+FLASH_ATTENTION="${ALPHAFOLD3_FLASH_ATTENTION:-triton}"
+MMSEQS_OPTIONS="${ALPHAFOLD3_MMSEQS_OPTIONS:---num-iterations 1 --db-load-mode 2 -a --max-seqs 10000 --prefilter-mode 3}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+python "${PROJECT_ROOT}/scripts/run_alphafold.py" \
+  --json_path="${JSON_PATH}" \
+  --model_dir="${MODEL_DIR}" \
+  --output_dir="${OUTPUT_DIR}" \
+  --run_data_pipeline=true \
+  --run_inference="${RUN_INFERENCE}" \
+  --flash_attention_implementation="${FLASH_ATTENTION}" \
+  --db_dir="${DB_DIRS}" \
+  --mmseqs_db_dir="${MMSEQS_DB_DIRS}" \
+  --use_mmseqs=true \
+  --use_mmseqs_gpu=true \
+  --mmseqs_options="${MMSEQS_OPTIONS}"