Merge remote-tracking branch 'origin/nanobind-environment-bod' into n…

…anobind-environment

freud/environment.py

@@ -254,10 +254,12 @@ class BondOrder(_SpatialHistogram):
             (Default value = :code:`'bod'`).
     """
 
-    known_modes = {'bod': freud._environment.bod,
-                   'lbod': freud._environment.lbod,
-                   'obcd': freud._environment.obcd,
-                   'oocd': freud._environment.oocd}
+    known_modes = {
+        "bod": freud._environment.bod,
+        "lbod": freud._environment.lbod,
+        "obcd": freud._environment.obcd,
+        "oocd": freud._environment.oocd,
+    }
 
     def __init__(self, bins, mode="bod"):
         if isinstance(bins, collections.abc.Sequence):
@@ -273,7 +275,6 @@ def __init__(self, bins, mode="bod"):
 
         self._cpp_obj = freud._environment.BondOrder(n_bins_theta, n_bins_phi, l_mode)
 
-
     @property
     def default_query_args(self):
         """No default query arguments."""
@@ -282,8 +283,15 @@ def default_query_args(self):
             NO_DEFAULT_QUERY_ARGS_MESSAGE.format(type(self).__name__)
         )
 
-    def compute(self, system, orientations=None, query_points=None,
-                query_orientations=None, neighbors=None, reset=True):
+    def compute(
+        self,
+        system,
+        orientations=None,
+        query_points=None,
+        query_orientations=None,
+        neighbors=None,
+        reset=True,
+    ):
         r"""Calculates the correlation function and adds to the current
         histogram.
 
@@ -317,57 +325,57 @@ def compute(self, system, orientations=None, query_points=None,
         if reset:
             self._reset()
 
-        nq, nlist, qargs, l_query_points, num_query_points = \
-            self._preprocess_arguments(system, query_points, neighbors)
+        nq, nlist, qargs, l_query_points, num_query_points = self._preprocess_arguments(
+            system, query_points, neighbors
+        )
         if orientations is None:
             orientations = np.array([[1, 0, 0, 0]] * nq.points.shape[0])
         if query_orientations is None:
             query_orientations = orientations
         if query_points is None:
             query_points = nq.points
 
-        orientations = freud.util._convert_array(orientations, shape=(nq.points.shape[0], 4))
-        query_orientations = freud.util._convert_array(query_orientations, shape=(num_query_points, 4))
-        query_points = freud.util._convert_array(query_points, shape=(num_query_points, 3))
-
+        orientations = freud.util._convert_array(
+            orientations, shape=(nq.points.shape[0], 4)
+        )
+        query_orientations = freud.util._convert_array(
+            query_orientations, shape=(num_query_points, 4)
+        )
+        query_points = freud.util._convert_array(
+            query_points, shape=(num_query_points, 3)
+        )
 
         self._cpp_obj.accumulate(
             nq._cpp_obj,
             orientations,
             query_points,
             query_orientations,
             nlist._cpp_obj,
-            qargs._cpp_obj
+            qargs._cpp_obj,
         )
         return self
 
     @_Compute._computed_property
     def bond_order(self):
         """:math:`\\left(N_{\\phi}, N_{\\theta} \\right)` :class:`numpy.ndarray`: Bond order."""  # noqa: E501
-        # return freud.util.make_managed_numpy_array(
-        #     &self.thisptr.getBondOrder(),
-        #     freud.util.arr_type_t.FLOAT)
-        return self._cpp_obj.getBondOrder()#.toNumpyArray()
+        return self._cpp_obj.getBondOrder().toNumpyArray()
 
     @_Compute._computed_property
     def box(self):
         """:class:`freud.box.Box`: Box used in the calculation."""
         return freud.box.BoxFromCPP(self._cpp_obj.getBox())
 
     def __repr__(self):
-        return ("freud.environment.{cls}(bins=({bins}), mode='{mode}')".format(
+        return "freud.environment.{cls}(bins=({bins}), mode='{mode}')".format(
             cls=type(self).__name__,
-            bins=', '.join([str(b) for b in self.nbins]),
-            mode=self.mode))
+            bins=", ".join([str(b) for b in self.nbins]),
+            mode=self.mode,
+        )
 
     @property
     def mode(self):
         """str: Bond order mode."""
-        mode = self._cpp_obj.getMode()
-        for key, value in self.known_modes.items():
-            if value == mode:
-                return key
-        return mode
+        return self._cpp_obj.getMode().name
 
 
 class LocalDescriptors(_PairCompute):

freud/environment/BondOrder.cc

@@ -57,37 +57,41 @@ BondOrder::BondOrder(unsigned int n_bins_theta, unsigned int n_bins_phi, BondOrd
             (*m_sa_array)(i, j) = sa;
         }
     }
-    const auto axes = util::Axes {
-        std::make_shared<util::RegularAxis>(n_bins_theta, 0, constants::TWO_PI),
-        std::make_shared<util::RegularAxis>(n_bins_phi, 0, M_PI)
-    };
-    // m_histogram = BondHistogram(axes);
+    const auto axes = util::Axes {std::make_shared<util::RegularAxis>(n_bins_theta, 0, constants::TWO_PI),
+                                  std::make_shared<util::RegularAxis>(n_bins_phi, 0, M_PI)};
+    m_histogram = BondHistogram(axes);
+    m_bo_array = std::make_shared<util::ManagedArray<float>>(std::vector<size_t> {m_histogram.shape()});
 
     m_local_histograms = BondHistogram::ThreadLocalHistogram(m_histogram);
 }
 
-void BondOrder::reduce()
+void BondOrder::reset()
 {
-    // TODO: previously, we could prepare the histogram: but not anymore?
-    // m_histogram.prepare(m_histogram.shape());
-    // m_histogram = std::make_shared<util::ManagedArray<unsigned int>>(std::vector<size_t> {m_histogram.shape()});
-
+    BondHistogramCompute::reset();
     m_bo_array = std::make_shared<util::ManagedArray<float>>(std::vector<size_t> {m_histogram.shape()});
+}
 
-
+void BondOrder::reduce()
+{
     m_histogram.reduceOverThreadsPerBin(m_local_histograms, [&](size_t i) {
         (*m_bo_array)[i] = m_histogram[i] / (*m_sa_array)[i] / static_cast<float>(m_frame_counter);
     });
 }
 
+std::vector<std::vector<float>> BondOrder::getBinCenters()
+{
+    return m_histogram.getBinCenters();
+}
+
 const std::shared_ptr<util::ManagedArray<float>> BondOrder::getBondOrder()
 {
     return reduceAndReturn(m_bo_array);
 }
 
-void BondOrder::accumulate(const std::shared_ptr<locality::NeighborQuery>& neighbor_query, quat<float>* orientations,
-                           vec3<float>* query_points, quat<float>* query_orientations,
-                           unsigned int n_query_points, const std::shared_ptr<freud::locality::NeighborList>& nlist,
+void BondOrder::accumulate(const std::shared_ptr<locality::NeighborQuery>& neighbor_query,
+                           quat<float>* orientations, vec3<float>* query_points,
+                           quat<float>* query_orientations, unsigned int n_query_points,
+                           const std::shared_ptr<freud::locality::NeighborList>& nlist,
                            freud::locality::QueryArgs qargs)
 {
     accumulateGeneral(neighbor_query, query_points, n_query_points, nlist, qargs,

freud/environment/BondOrder.h

@@ -42,9 +42,12 @@ class BondOrder : public locality::BondHistogramCompute
     //! Accumulate the bond order
     void accumulate(const std::shared_ptr<locality::NeighborQuery>& neighbor_query, quat<float>* orientations,
                     vec3<float>* query_points, quat<float>* query_orientations, unsigned int n_query_points,
-                    const std::shared_ptr<freud::locality::NeighborList>& nlist, freud::locality::QueryArgs qargs);
+                    const std::shared_ptr<freud::locality::NeighborList>& nlist,
+                    freud::locality::QueryArgs qargs);
 
     void reduce() override;
+    void reset() override;
+    std::vector<std::vector<float>> getBinCenters();
 
     //! Get a shared_ptr to the last computed bond order
     const std::shared_ptr<util::ManagedArray<float>> getBondOrder();
@@ -57,7 +60,7 @@ class BondOrder : public locality::BondHistogramCompute
 private:
     std::shared_ptr<util::ManagedArray<float>> m_bo_array; //!< bond order array computed
     std::shared_ptr<util::ManagedArray<float>> m_sa_array; //!< surface area array computed
-    BondOrderMode m_mode;                 //!< The mode to calculate with.
+    BondOrderMode m_mode;                                  //!< The mode to calculate with.
 };
 
 }; }; // end namespace freud::environment

freud/environment/export-BondOrder.cc

@@ -1,7 +1,6 @@
 // Copyright (c) 2010-2024 The Regents of the University of Michigan
 // This file is from the freud project, released under the BSD 3-Clause License.
 
-#include <iostream>
 #include <memory>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
@@ -15,20 +14,16 @@ namespace nb = nanobind;
 
 namespace freud { namespace environment {
 
-template<typename T, typename shape>
-using nb_array = nb::ndarray<T, shape, nb::device::cpu, nb::c_contig>;
+template<typename T, typename shape> using nb_array = nb::ndarray<T, shape, nb::device::cpu, nb::c_contig>;
 
 namespace wrap {
 
-
 void accumulateBondOrder(const std::shared_ptr<BondOrder>& self,
-    const std::shared_ptr<locality::NeighborQuery>& nq,
-    const nb_array<float, nb::shape<-1, 4>>& orientations,
-    const nb_array<float, nb::shape<-1, 3>>& query_points,
-    const nb_array<float, nb::shape<-1, 4>>& query_orientations,
-    std::shared_ptr<locality::NeighborList> nlist, 
-    const locality::QueryArgs& qargs
-  )
+                         const std::shared_ptr<locality::NeighborQuery>& nq,
+                         const nb_array<float, nb::shape<-1, 4>>& orientations,
+                         const nb_array<float, nb::shape<-1, 3>>& query_points,
+                         const nb_array<float, nb::shape<-1, 4>>& query_orientations,
+                         std::shared_ptr<locality::NeighborList> nlist, const locality::QueryArgs& qargs)
 {
     unsigned int const n_query_points = query_points.shape(0);
     // std::cout << n_query_points << std::endl;
@@ -39,15 +34,13 @@ void accumulateBondOrder(const std::shared_ptr<BondOrder>& self,
     // else {
     //   auto* query_points_data= reinterpret_cast<vec3<float>*>(query_points.data());
     // }
-    
+
     auto* orientations_data = reinterpret_cast<quat<float>*>(orientations.data());
-    auto* query_points_data= reinterpret_cast<vec3<float>*>(query_points.data());
+    auto* query_points_data = reinterpret_cast<vec3<float>*>(query_points.data());
     auto* query_orientations_data = reinterpret_cast<quat<float>*>(query_orientations.data());
 
-
-    self->accumulate(
-        nq, orientations_data, query_points_data, query_orientations_data, n_query_points, nlist, qargs
-    );
+    self->accumulate(nq, orientations_data, query_points_data, query_orientations_data, n_query_points, nlist,
+                     qargs);
 }
 
 }; // namespace wrap
@@ -56,33 +49,30 @@ namespace detail {
 
 void export_BondOrder(nb::module_& module)
 {
-
     nb::enum_<BondOrderMode>(module, "BondOrderMode")
         .value("bod", BondOrderMode::bod)
         .value("lbod", BondOrderMode::lbod)
         .value("obcd", BondOrderMode::obcd)
         .value("oocd", BondOrderMode::oocd)
         .export_values();
-  
+
     nb::class_<BondOrder>(module, "BondOrder")
         .def(nb::init<unsigned int, unsigned int, BondOrderMode>())
         .def("getBondOrder", &BondOrder::getBondOrder)
         .def("getBinCounts", &BondOrder::getBinCounts)
+        .def("getBinCenters", &BondOrder::getBinCenters)
+        .def("getBinEdges", &BondOrder::getBinEdges)
+        .def("getBox", &BondOrder::getBox)
         .def("getAxisSizes", &BondOrder::getAxisSizes)
         .def("getMode", &BondOrder::getMode)
-        .def("accumulate", &wrap::accumulateBondOrder,
-          nanobind::arg("nq").none(),
-          nanobind::arg("orientations"),
-          nanobind::arg("query_points"),
-          nanobind::arg("query_orientations"),
-          // nanobind::arg("n_query_points"),
-          nanobind::arg("nlist").none(),
-          nanobind::arg("qargs").none()
-        )
-        .def("reset", &BondOrder::reset)
-        ;
+        .def("accumulate", &wrap::accumulateBondOrder, nanobind::arg("nq").none(),
+             nanobind::arg("orientations"), nanobind::arg("query_points"),
+             nanobind::arg("query_orientations"),
+             // nanobind::arg("n_query_points"),
+             nanobind::arg("nlist").none(), nanobind::arg("qargs").none())
+        .def("reset", &BondOrder::reset);
 }
 
 }; // namespace detail
 
-}; }; // namespace freud::locality
+}; }; // namespace freud::environment

freud/locality.py

@@ -992,17 +992,18 @@ def box(self):
 
     @_Compute._computed_property
     def bin_counts(self):
-        """:class:`numpy.ndarray`: The bin counts in the histogram."""
+        """:math:`\\left(N_0, N_1 \\right)` :class:`numpy.ndarray`: The bin counts in the histogram."""  # noqa: E501
         return self._cpp_obj.getBinCounts().toNumpyArray()
 
-    @property
+    @property  # TODO: Actual shape is 2, (d1=1, d2+1) and type is list
     def bin_centers(self):
         """:class:`numpy.ndarray`: The centers of each bin in the histogram
         (has the same shape as the histogram itself)."""
         centers = self._cpp_obj.getBinCenters()
+        # TODO: update docs to list the correct dimensions and type
         return [np.array(c) for c in centers]
 
-    @property
+    @property  # TODO: Actual shape is 2, (d1=1, d2+1) and type is list
     def bin_edges(self):
         """:class:`numpy.ndarray`: The edges of each bin in the histogram (is
         one element larger in each dimension than the histogram because each

tests/test_environment_bond_order.py

@@ -40,6 +40,9 @@ def test_bond_order(self):
         # Test access
         bo.box
         bo.bond_order
+        bo.bin_counts
+        bo.bin_edges
+        bo.bin_centers
 
         # Test all the basic attributes.
         assert bo.nbins[0] == n_bins_theta
@@ -57,8 +60,7 @@ def test_bond_order(self):
             box, positions, positions, "nearest", r_max, num_neighbors, True
         )
         for nq, neighbors in test_set:
-            # Test that lbod gives identical results when orientations are the
-            # same.
+            # Test that lbod gives identical results when orientations are the same.
             # TODO: Find a way to test a rotated system to ensure that lbod gives
             # the desired results.
             bo = freud.environment.BondOrder(nbins, mode="lbod")