Skip to content

Commit

Permalink
python: add __array__() where we had buffer protocol previously
Browse files Browse the repository at this point in the history
  • Loading branch information
wojdyr committed Oct 7, 2024
1 parent 3467d24 commit 729151e
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 20 deletions.
30 changes: 17 additions & 13 deletions python/mtz.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,19 @@ auto make_new_column(const Mtz& mtz, int dataset, Func func) {
return numpy_arr;
}

auto mtz_to_array(Mtz& self) {
size_t nrow = self.has_data() ? (size_t) self.nreflections : 0;
size_t ncol = self.columns.size();
return nb::ndarray<nb::numpy, float, nb::ndim<2>>(self.data.data(), {nrow, ncol}, nb::handle());
}

auto column_to_array(Mtz::Column& self) {
return nb::ndarray<nb::numpy, float, nb::ndim<1>>(self.parent->data.data() + self.idx,
{(size_t)self.size()},
nb::handle(),
{(int64_t) self.stride()});
}

} // anonymous namespace

void add_mtz(nb::module_& m) {
Expand All @@ -78,12 +91,8 @@ void add_mtz(nb::module_& m) {

mtz
.def(nb::init<bool>(), nb::arg("with_base")=false)
.def_prop_ro("array", [](Mtz& self) {
size_t nrow = self.has_data() ? (size_t) self.nreflections : 0;
size_t ncol = self.columns.size();
return nb::ndarray<nb::numpy, float, nb::ndim<2>>(
self.data.data(), {nrow, ncol}, nb::handle());
}, nb::rv_policy::reference_internal)
.def_prop_ro("array", &mtz_to_array, nb::rv_policy::reference_internal)
.def("__array__", &mtz_to_array, nb::rv_policy::reference_internal)
.def_rw("title", &Mtz::title)
.def_rw("nreflections", &Mtz::nreflections)
.def_rw("sort_order", &Mtz::sort_order)
Expand Down Expand Up @@ -292,13 +301,8 @@ void add_mtz(nb::module_& m) {
})
;
pyMtzColumn
.def_prop_ro("array", [](const Mtz::Column& self) {
return nb::ndarray<nb::numpy, float, nb::ndim<1>>(
self.parent->data.data() + self.idx,
{(size_t)self.size()},
nb::handle(),
{(int64_t) self.stride()});
}, nb::rv_policy::reference_internal)
.def_prop_ro("array", &column_to_array, nb::rv_policy::reference_internal)
.def("__array__", &column_to_array, nb::rv_policy::reference_internal)
.def_prop_ro("dataset",
(Mtz::Dataset& (Mtz::Column::*)()) &Mtz::Column::dataset)
.def_rw("dataset_id", &Mtz::Column::dataset_id)
Expand Down
21 changes: 14 additions & 7 deletions python/unitcell.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,23 +17,32 @@

using namespace gemmi;

static std::string triple(double x, double y, double z) {
namespace {

std::string triple(double x, double y, double z) {
char buf[128];
auto r = [](double d) { return std::fabs(d) < 1e-15 ? 0 : d; };
snprintf_z(buf, 128, "%g, %g, %g", r(x), r(y), r(z));
return std::string(buf);
}

static void mat33_from_list(Mat33& self, std::array<std::array<double,3>,3>& m) {
void mat33_from_list(Mat33& self, std::array<std::array<double,3>,3>& m) {
for (int i = 0; i < 3; ++i)
for (int j = 0; j < 3; ++j)
self.a[i][j] = m[i][j];
}

static nb::tuple make_six_tuple(const std::array<double,6>& v) {
nb::tuple make_six_tuple(const std::array<double,6>& v) {
return nb::make_tuple(v[0], v[1], v[2], v[3], v[4], v[5]);
}

auto mat33_to_array(Mat33& self) {
return nb::ndarray<nb::numpy, double, nb::shape<3,3>, nb::c_contig>(
&self.a[0][0], {3, 3}, nb::handle());
}

} // anonymous namespace

template<typename T> void add_smat33(nb::module_& m, const char* name) {
using M = SMat33<T>;
nb::class_<M>(m, name)
Expand Down Expand Up @@ -135,10 +144,8 @@ void add_unitcell(nb::module_& m) {
new(mat) Mat33();
mat33_from_list(*mat, arr);
})
.def_prop_ro("array", [](Mat33& self) {
return nb::ndarray<nb::numpy, double, nb::shape<3,3>, nb::c_contig>(
&self.a[0][0], {3, 3}, nb::handle());
}, nb::rv_policy::reference_internal)
.def_prop_ro("array", &mat33_to_array, nb::rv_policy::reference_internal)
.def("__array__", &mat33_to_array, nb::rv_policy::reference_internal)
.def("row_copy", &Mat33::row_copy)
.def("column_copy", &Mat33::column_copy)
.def(nb::self + nb::self)
Expand Down

0 comments on commit 729151e

Please sign in to comment.