feat(moments): add density parameter to moment invariant calculation

geometricus/geometricus.py

@@ -421,6 +421,8 @@ class MomentInvariants(Structure):
     """Filled with moment invariant values for each structural fragment"""
     moment_types: List[MomentType] = None
     """Names of moments used"""
+    density: np.ndarray = None
+    """Density of each residue"""
 
     @classmethod
     def from_coordinates(
@@ -437,6 +439,7 @@ def from_coordinates(
             MomentType.O_5,
             MomentType.F,
         ),
+        density: np.ndarray = None
     ):
         """
         Construct MomentInvariants instance from a set of coordinates.
@@ -454,7 +457,10 @@ def from_coordinates(
             split_size=split_size,
             upsample_rate=upsample_rate,
             moment_types=moment_types,
+            density=density
         )
+        if shape.density is None:
+            shape.density = np.ones(shape.coordinates.shape[0])
         shape._split(split_type)
         return shape
 
@@ -473,6 +479,8 @@ def from_prody_atomgroup(
             MomentType.O_5,
             MomentType.F,
         ),
+        density: np.ndarray = None
+
     ):
         """
         Construct MomentInvariants instance from a ProDy AtomGroup object.
@@ -497,7 +505,10 @@ def from_prody_atomgroup(
             split_size=split_size,
             upsample_rate=upsample_rate,
             moment_types=moment_types,
+            density=density
         )
+        if shape.density is None:
+            shape.density = np.ones(shape.coordinates.shape[0])
         shape._split(split_type)
         return shape
 
@@ -533,6 +544,7 @@ def from_pdb_file(
             MomentType.O_5,
             MomentType.F,
         ),
+        density: np.ndarray = None
     ):
         """
         Construct MomentInvariants instance from a PDB file and optional chain.
@@ -555,6 +567,7 @@ def from_pdb_file(
             selection=selection,
             upsample_rate=upsample_rate,
             moment_types=moment_types,
+            density=density
         )
 
     @classmethod
@@ -572,6 +585,7 @@ def from_pdb_id(
             MomentType.O_5,
             MomentType.F,
         ),
+        density: np.ndarray = None
     ):
         """
         Construct MomentInvariants instance from a PDB ID and optional chain (downloads the PDB file from RCSB).
@@ -594,6 +608,7 @@ def from_pdb_id(
             selection=selection,
             upsample_rate=upsample_rate,
             moment_types=moment_types,
+            density=density
         )
 
     def _kmerize(self):
@@ -639,9 +654,9 @@ def _allmerize(self):
     def _get_moments(self, split_indices):
         moments = np.zeros((len(split_indices), len(self.moment_types)))
         for i, indices in enumerate(split_indices):
-            moments[i] = get_moments_from_coordinates(
-                self.coordinates[indices], self.moment_types
-            )
+            moments[i] = get_moments_from_coordinates(self.coordinates[indices],
+                                                      self.moment_types,
+                                                      self.density[indices])
         return split_indices, moments
 
     def _split_radius(self):
@@ -678,9 +693,7 @@ def _split_radius_upsample(self):
             split_indices.append(kd_tree.getIndices())
         moments = np.zeros((len(split_indices), len(self.moment_types)))
         for i, indices in enumerate(split_indices_upsample):
-            moments[i] = get_moments_from_coordinates(
-                coordinates_upsample[indices], self.moment_types
-            )
+            moments[i] = get_moments_from_coordinates(coordinates_upsample[indices], self.moment_types)
         return split_indices, moments
 
 

geometricus/moment_utility.py

@@ -23,14 +23,15 @@ class MomentInfo:
 
 
 @nb.njit(cache=False)
-def mu(p, q, r, coords, centroid):
+def mu(p, q, r, coords, density, centroid):
     """
     Central moment
     """
     return np.sum(
         ((coords[:, 0] - centroid[0]) ** p)
         * ((coords[:, 1] - centroid[1]) ** q)
         * ((coords[:, 2] - centroid[2]) ** r)
+        * density
     )
 
 
@@ -1294,6 +1295,7 @@ def get_moments_from_coordinates(
         MomentType.O_5,
         MomentType.F,
     ),
+    density: np.ndarray = None
 ) -> ty.List[float]:
     """
     Gets rotation-invariant moments for a set of coordinates
@@ -1304,17 +1306,22 @@ def get_moments_from_coordinates(
     moment_types
         Which moments to calculate
         Choose from ['O_3', 'O_4', 'O_5', 'F', 'phi_2', 'phi_3', 'phi_4', 'phi_5', 'phi_6', 'phi_7', 'phi_8', 'phi_9', 'phi_10', 'phi_11', 'phi_12', 'phi_13']
+    density
+        assign a density to each residue/coordinate. 1 by default
 
     Returns
     -------
     list of moments
     """
+    if density is None:
+        density = np.ones(coordinates.shape[0])
+    assert density.shape[0] == coordinates.shape[0]
     all_moment_mu_types: ty.Set[ty.Tuple[int, int, int]] = set(
         m for moment_type in moment_types for m in moment_type.value.mu_arguments
     )
     centroid = nb_mean_axis_0(coordinates)
     mus = {
-        (x, y, z): mu(float(x), float(y), float(z), coordinates, centroid)
+        (x, y, z): mu(float(x), float(y), float(z), coordinates, density, centroid)
         for (x, y, z) in all_moment_mu_types
     }
     moments = [