File size: 2,720 Bytes
35cdf53 | 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 | // 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 <vector>
#include "absl/container/flat_hash_set.h"
#include "pybind11/cast.h"
#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"
#include "pybind11_abseil/absl_casters.h"
namespace {
namespace py = pybind11;
py::array_t<bool> IsIn(const py::array_t<int64_t, py::array::c_style>& array,
const absl::flat_hash_set<int64_t>& 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) {
if (test_elements.contains(array.data()[i])) {
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;
}
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 kind='sort' has complexity O(m*log(m) + n * log(m)).
Args:
array: Input NumPy array with dtype=np.int64.
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`.
)";
} // namespace
namespace alphafold3 {
void RegisterModuleMembership(pybind11::module m) {
m.def("isin", &IsIn, py::arg("array"), py::arg("test_elements"),
py::kw_only(), py::arg("invert") = false, py::doc(kIsInDoc + 1));
}
} // namespace alphafold3
|