Adds get_atomic_property method (issue #216)

src/grid/molgrid.py

@@ -611,6 +611,49 @@ def __getitem__(self, index: int):
             )
         return self._atgrids[index]
 
+    def get_atomic_property(self, property_values):
+        """Convert property at grid points to property per atom.
+
+        This method takes a property evaluated at all grid points and calculated how much
+        of that propery belongs to each atom in the molecule.
+
+        Parameters
+        ----------
+        property_values : np.narray
+            Must have the same length as the number of points in the molecular grid.
+        """
+
+        # check if the input is valid
+        if not isinstance(property_values, np.ndarray):
+            raise TypeError(
+                f"property_values must be of numpy array type, got {type(property_values)}"
+            )
+
+        if property_values.size != self.size:
+            raise ValueError(
+                "property_values is not the same size as grid. \n"
+                f"property_values.size: {property_values.size}, self.size: {self.size}"
+            )
+
+        # initialize output array
+        num_atoms = len(self.atcoords)
+        atomic_property = np.zeros(num_atoms)
+
+        # calculate property for each atom
+        for i in range(num_atoms):
+            # get the range of grid points for each atom
+            start_index = self.indices[i]
+            end_index = self.indices[i + 1]
+
+            # finding property of each atom
+            atomic_property[i] = np.sum(
+                property_values[start_index:end_index]
+                * self.aim_weights[start_index:end_index]
+                * self.atweights[start_index:end_index]
+            )
+
+        return atomic_property
+
 
 def _generate_default_rgrid(atnum: int):
     r"""

src/grid/tests/test_molgrid.py

@@ -810,6 +810,55 @@ def test_integrate_hirshfeld_weights_pair_1s(self):
         occupation = mg.integrate(fn)
         assert_almost_equal(occupation, 2.5, decimal=5)
 
+    def test_get_atomic_property_constant(self):
+        """Test with constant property values"""
+        coords = np.array([[0, 0, -0.5], [0, 0, 0.5]])
+        mg = MolGrid.from_preset(
+            atnums=np.array([1, 1]),
+            atcoords=coords,
+            preset="coarse",
+            aim_weights=BeckeWeights(order=3),
+        )
+
+        property_values = np.ones(mg.size)
+        atomic_properties = mg.get_atomic_property(property_values)
+
+        assert atomic_properties.size == 2
+        assert_allclose(np.sum(atomic_properties), mg.integrate(property_values))
+
+    def test_get_atomic_property_wrong_size(self):
+        """Test error handling for wrong array size"""
+        coords = np.array([[0, 0, -0.5], [0, 0, 0.5]])
+        mg = MolGrid.from_preset(
+            atnums=np.array([1, 1]),
+            atcoords=coords,
+            preset="coarse",
+            aim_weights=BeckeWeights(order=3),
+        )
+
+        wrong_size = np.ones(mg.size + 1)
+        with pytest.raises(ValueError, match="property_values is not the same size as grid"):
+            mg.get_atomic_property(wrong_size)
+
+    def test_get_atomic_property_nonuniform(self):
+        """Test with non-uniform property(distance from origin)"""
+        coords = np.array([[0, 0, -0.5], [0, 0, 0.5]])
+        mg = MolGrid.from_preset(
+            atnums=np.array([1, 1]),
+            atcoords=coords,
+            preset="coarse",
+            aim_weights=BeckeWeights(order=3),
+        )
+
+        property_values = np.linalg.norm(mg.points, axis=1)
+        atomic_properties = mg.get_atomic_property(property_values)
+
+        # Due to symmetry, both atoms should have equal properties
+        assert atomic_properties.size == 2
+
+        assert_allclose(atomic_properties[0], atomic_properties[1], rtol=1e-3)
+        assert_allclose(np.sum(atomic_properties), mg.integrate(property_values), rtol=1e-3)
+
 
 def test_interpolation_with_gaussian_center():
     r"""Test interpolation with molecular grid of sum of two Gaussian examples."""