File size: 3,224 Bytes
3e62986 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | // 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
|