diff --git a/geometricus/geometricus.py b/geometricus/geometricus.py index eab94cd..e5f4e76 100644 --- a/geometricus/geometricus.py +++ b/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 diff --git a/geometricus/moment_utility.py b/geometricus/moment_utility.py index 1e7b086..7eb4205 100644 --- a/geometricus/moment_utility.py +++ b/geometricus/moment_utility.py @@ -23,7 +23,7 @@ class MomentInfo: @nb.njit(cache=False) -def mu(p, q, r, coords, centroid): +def mu(p, q, r, coords, density, centroid): """ Central moment """ @@ -31,6 +31,7 @@ def mu(p, q, r, coords, centroid): ((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 = [