From 57151e254c6a08c9b99837078ac265116f4a6d08 Mon Sep 17 00:00:00 2001 From: gph82 Date: Wed, 20 Mar 2024 11:43:20 +0100 Subject: [PATCH] [utils.COM.geom2COMxyz] API BREAK change shape of returned array to only those residues involved --- mdciao/utils/COM.py | 19 ++++++++----------- tests/test_COM_utils.py | 4 ++-- 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/mdciao/utils/COM.py b/mdciao/utils/COM.py index 0587f2d3..953ac480 100644 --- a/mdciao/utils/COM.py +++ b/mdciao/utils/COM.py @@ -140,7 +140,7 @@ def geom2COMdist(geom, residue_pairs, subtract_max_radii=False, low_mem=True, unwrapped_residue_geom = geom # This would be worth migrating to mdanalysis # https://docs.mdanalysis.org/1.0.1/documentation_pages/core/groups.html#MDAnalysis.core.groups.ResidueGroup.center - COMs_xyz = geom2COMxyz(unwrapped_residue_geom, residue_idxs=residue_idxs_unique)[:, residue_idxs_unique] + COMs_xyz = geom2COMxyz(unwrapped_residue_geom, residue_idxs=residue_idxs_unique) COM_dists_t = _compute_distances_core(COMs_xyz, pair_map, @@ -191,13 +191,11 @@ def geom2COMxyz(igeom, residue_idxs=None): residue_idxs : iterable, default is None Residues for which the center of mass will be computed. Default - is to compute all residues. The excluded residues will appear - as np.nans in the returned value, which has the same shape - regardless of the input + is to compute all residues. Returns ------- - COMs : numpy.ndarray of shape (igeom.n_frames, igeom.n_residues,3) + COMs : numpy.ndarray of shape (igeom.n_frames, len(residue_idxs),3) """ @@ -205,11 +203,10 @@ def geom2COMxyz(igeom, residue_idxs=None): residue_idxs=_np.arange(igeom.top.n_residues) masses = [_np.hstack([aa.element.mass for aa in rr.atoms]) for rr in igeom.top.residues] - COMs_res_time_coords = _np.zeros((igeom.n_residues,igeom.n_frames,3)) - COMs_res_time_coords[:,:,:] = _np.nan - COMs_res_time_coords[residue_idxs] = [_np.average(igeom.xyz[:, [aa.index for aa in igeom.top.residue(index).atoms], :], axis=1, weights=masses[index]) - for index in residue_idxs] - COMs_time_res_coords = _np.swapaxes(_np.array(COMs_res_time_coords),0,1) + COMs_res_time_coords = _np.array([_np.average(igeom.xyz[:, [aa.index for aa in igeom.top.residue(index).atoms], :], + axis=1, weights=masses[index]) + for index in residue_idxs]) + COMs_time_res_coords = _np.swapaxes(_np.array(COMs_res_time_coords), 0, 1) return COMs_time_res_coords def geom2max_residue_radius(geom, residue_idxs=None, res_COMs=None) -> _np.ndarray: @@ -254,7 +251,7 @@ def geom2max_residue_radius(geom, residue_idxs=None, res_COMs=None) -> _np.ndarr raise ValueError("If 'residue_idxs' is None, then 'res_COMs' has to be None as well.") if res_COMs is None: - res_COMs = geom2COMxyz(geom, residue_idxs=residue_idxs)[:, residue_idxs] + res_COMs = geom2COMxyz(geom, residue_idxs=residue_idxs) else: assert res_COMs.shape[0] == geom.n_frames assert res_COMs.shape[1] == len(residue_idxs) diff --git a/tests/test_COM_utils.py b/tests/test_COM_utils.py index ec3952c3..aa37e7a3 100644 --- a/tests/test_COM_utils.py +++ b/tests/test_COM_utils.py @@ -35,8 +35,8 @@ def test_COMxyz_works_some_residues(self): residue_idxs = [1, 3, 5, 7] COMSs_mine = geom2COMxyz(self.traj_5_frames, residue_idxs=residue_idxs) - _np.testing.assert_allclose(COMSs_mine[:, [residue_idxs]], - self.COMS_mdtraj[:,[residue_idxs]]) + _np.testing.assert_allclose(COMSs_mine, + self.COMS_mdtraj[:,residue_idxs]) def test_COMdist_works(self): res_pairs = [[0,10], [10,20]]