diff --git a/src/grid/molgrid.py b/src/grid/molgrid.py index 4ad59532..2a746dc4 100644 --- a/src/grid/molgrid.py +++ b/src/grid/molgrid.py @@ -611,6 +611,54 @@ 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 calculates how much + of that property belongs to each atom in the molecule. + + Parameters + ---------- + property_values : np.ndarray + Must have the same length as the number of points in the molecular grid. + + Returns + ------- + np.ndarray + One-dimensional array of length equal to the number of atoms, containing + the property value associated with each atom. + """ + + # 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)}" + ) + + property_values = np.asarray(property_values).reshape(-1) + if property_values.size != self.size: + raise ValueError( + "property_values must be one-dimensional with the same length as the grid.\n" + f"property_values.shape: {property_values.shape}, expected: ({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.weights[start_index:end_index] + ) + + return atomic_property + def _generate_default_rgrid(atnum: int): r""" diff --git a/src/grid/tests/test_molgrid.py b/src/grid/tests/test_molgrid.py index 0ca52b3e..95474618 100644 --- a/src/grid/tests/test_molgrid.py +++ b/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 must be one-dimensional"): + 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."""