Skip to content

Commit

Permalink
Weights rework for MDAnalysis#2429 (MDAnalysis#2610)
Browse files Browse the repository at this point in the history 
Fixes MDAnalysis#2429 

## Work done in this PR
weight_groupselections option added to allow selection of weights (*mass*, *None* or 1D float array) for each groupselection in RMSD. Weights for select now gets masses from mobile atoms.

## Files changed
 - rms.py
 - test_rms.py
 - CHANGELOG
  • Loading branch information
@yuxuanzhuang
yuxuanzhuang authored Mar 20, 2020
1 parent 89646b4 commit 4102481
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 50 deletions.
2 changes: 2 additions & 0 deletions package/CHANGELOG
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,8 @@ Enhancements
* Added ParmEdParser, ParmEdReader and ParmEdConverter to
convert between a parmed.Structure and MDAnalysis Universe (PR #2404)
* Improve the distance search in water bridge analysis with capped_distance (PR #2480)
* Added weights_groupselections option in RMSD for mapping custom weights
to groupselections (PR #2610, Issue #2429)

Changes
* Removed `details` from `ClusteringMethod`s (Issue #2575, PR #2620)
Expand Down
113 changes: 72 additions & 41 deletions package/MDAnalysis/analysis/rms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,6 @@ def rmsd(a, b, weights=None, center=False, superposition=False):
a = np.asarray(a, dtype=np.float64)
b = np.asarray(b, dtype=np.float64)
N = b.shape[0]

if a.shape != b.shape:
raise ValueError('a and b must have same shape')

Expand Down Expand Up @@ -336,8 +335,8 @@ class RMSD(AnalysisBase):
"""
def __init__(self, atomgroup, reference=None, select='all',
groupselections=None, weights=None, tol_mass=0.1,
ref_frame=0, **kwargs):
groupselections=None, weights=None, weights_groupselections=False,
tol_mass=0.1, ref_frame=0, **kwargs):
r"""Parameters
----------
atomgroup : AtomGroup or Universe
Expand Down Expand Up @@ -381,11 +380,22 @@ def __init__(self, atomgroup, reference=None, select='all',
.. Note:: Experimental feature. Only limited error checking
implemented.
weights : {"mass", ``None``} or array_like (optional)
choose weights. With ``"mass"`` uses masses as weights; with ``None``
weigh each atom equally. If a float array of the same length as
`atomgroup` is provided, use each element of the `array_like` as a
weight for the corresponding atom in `atomgroup`.
1. "mass" will use masses as weights for both `select` and `groupselections`.
2. ``None`` will weigh each atom equally for both `select` and `groupselections`.
3. If 1D float array of the same length as `atomgroup` is provided,
use each element of the `array_like` as a weight for the
corresponding atom in `select`, and assumes ``None`` for `groupselections`.
weights_groupselections : False or list of {"mass", ``None`` or array_like} (optional)
1. ``False`` will apply imposed weights to `groupselections` from ``weights`` option.
2. A list of {"mass", ``None`` or array_like} with the length of `groupselections`
will apply the weights to `groupselections` correspondingly.
tol_mass : float (optional)
Reject match if the atomic masses for matched atoms differ by more
than `tol_mass`.
Expand All @@ -400,17 +410,15 @@ def __init__(self, atomgroup, reference=None, select='all',
SelectionError
If the selections from `atomgroup` and `reference` do not match.
TypeError
If `weights` is not of the appropriate type; see also
:func:`MDAnalysis.lib.util.get_weights`
If `weights` or `weights_groupselections` is not of the appropriate type;
see also :func:`MDAnalysis.lib.util.get_weights`
ValueError
If `weights` are not compatible with `atomgroup` (not the same
length) or if it is not a 1D array (see
:func:`MDAnalysis.lib.util.get_weights`).
A :exc:`ValueError` is also raised if `weights` are not compatible
with `groupselections`: only equal weights (``weights=None``) or
mass-weighted (``weights="mass"``) are supported for additional
`groupselections`.
A :exc:`ValueError` is also raised if the length of `weights_groupselections`
are not compatible with `groupselections`.
Notes
-----
Expand Down Expand Up @@ -477,7 +485,7 @@ def __init__(self, atomgroup, reference=None, select='all',
self.weights = weights
self.tol_mass = tol_mass
self.ref_frame = ref_frame

self.weights_groupselections = weights_groupselections
self.ref_atoms = self.reference.select_atoms(*select['reference'])
self.mobile_atoms = self.atomgroup.select_atoms(*select['mobile'])

Expand Down Expand Up @@ -535,29 +543,52 @@ def __init__(self, atomgroup, reference=None, select='all',
igroup, sel['reference'], sel['mobile'],
len(atoms['reference']), len(atoms['mobile'])))

# Explicitly check for "mass" because this option CAN
# be used with groupselection. (get_weights() returns the mass array
# for "mass")
if not iterable(self.weights) and self.weights == "mass":
pass
else:
self.weights = get_weights(self.mobile_atoms, self.weights)
# check weights type
if iterable(self.weights) and (np.array(weights).dtype
not in (np.dtype('float64'),np.dtype('int64'))):
raise TypeError("weight should only be be 'mass', None or 1D float array."
"For weights on groupselections, use **weight_groupselections** ")
if iterable(self.weights) or self.weights != "mass":
get_weights(self.mobile_atoms, self.weights)

if self.weights_groupselections:
if len(self.weights_groupselections) != len(self.groupselections):
raise ValueError("Length of weights_groupselections is not equal to "
"length of groupselections ")
for weights, atoms, selection in zip(self.weights_groupselections,
self._groupselections_atoms,
self.groupselections):
try:
if iterable(weights) or weights != "mass":
get_weights(atoms['mobile'], weights)
except Exception as e:
raise type(e)(str(e) + ' happens in selection %s' % selection['mobile'])

# cannot use arbitrary weight array (for superposition) with
# groupselections because arrays will not match
if (len(self.groupselections) > 0 and (
iterable(self.weights) or self.weights not in ("mass", None))):
raise ValueError("groupselections can only be combined with "
"weights=None or weights='mass', not a weight "
"array.")

def _prepare(self):
self._n_atoms = self.mobile_atoms.n_atoms

if not iterable(self.weights) and self.weights == 'mass':
self.weights = self.ref_atoms.masses
if self.weights is not None:
self.weights = np.asarray(self.weights, dtype=np.float64) / np.mean(self.weights)
if not self.weights_groupselections:
if not iterable(self.weights): # apply 'mass' or 'None' to groupselections
self.weights_groupselections = [self.weights] * len(self.groupselections)
else:
self.weights_groupselections = [None] * len(self.groupselections)

for igroup, (weights, atoms) in enumerate(zip(self.weights_groupselections,
self._groupselections_atoms)):
if str(weights) == 'mass':
self.weights_groupselections[igroup] = atoms['mobile'].masses
if weights is not None:
self.weights_groupselections[igroup] = np.asarray(self.weights_groupselections[igroup],
dtype=np.float64) / \
np.mean(self.weights_groupselections[igroup])
# add the array of weights to weights_select
self.weights_select = get_weights(self.mobile_atoms, self.weights)
self.weights_ref = get_weights(self.ref_atoms, self.weights)
if self.weights_select is not None:
self.weights_select = np.asarray(self.weights_select, dtype=np.float64) / \
np.mean(self.weights_select)
self.weights_ref = np.asarray(self.weights_ref, dtype=np.float64) / \
np.mean(self.weights_ref)

current_frame = self.reference.universe.trajectory.ts.frame

Expand All @@ -566,14 +597,14 @@ def _prepare(self):
# (coordinates MUST be stored in case the ref traj is advanced
# elsewhere or if ref == mobile universe)
self.reference.universe.trajectory[self.ref_frame]
self._ref_com = self.ref_atoms.center(self.weights)
self._ref_com = self.ref_atoms.center(self.weights_ref)
# makes a copy
self._ref_coordinates = self.ref_atoms.positions - self._ref_com
if self._groupselections_atoms:
self._groupselections_ref_coords64 = [(self.reference.
select_atoms(*s['reference']).
positions.astype(np.float64)) for s in
self.groupselections]
select_atoms(*s['reference']).
positions.astype(np.float64)) for s in
self.groupselections]
finally:
# Move back to the original frame
self.reference.universe.trajectory[current_frame]
Expand All @@ -599,7 +630,7 @@ def _prepare(self):
self._mobile_coordinates64 = self.mobile_atoms.positions.copy().astype(np.float64)

def _single_frame(self):
mobile_com = self.mobile_atoms.center(self.weights).astype(np.float64)
mobile_com = self.mobile_atoms.center(self.weights_select).astype(np.float64)
self._mobile_coordinates64[:] = self.mobile_atoms.positions
self._mobile_coordinates64 -= mobile_com

Expand All @@ -614,7 +645,7 @@ def _single_frame(self):

self.rmsd[self._frame_index, 2] = qcp.CalcRMSDRotationalMatrix(
self._ref_coordinates64, self._mobile_coordinates64,
self._n_atoms, self._rot, self.weights)
self._n_atoms, self._rot, self.weights_select)

self._R[:, :] = self._rot.reshape(3, 3)
# Transform each atom in the trajectory (use inplace ops to
Expand All @@ -633,14 +664,14 @@ def _single_frame(self):
self._groupselections_atoms), 3):
self.rmsd[self._frame_index, igroup] = rmsd(
refpos, atoms['mobile'].positions,
weights=self.weights,
weights=self.weights_groupselections[igroup-3],
center=False, superposition=False)
else:
# only calculate RMSD by setting the Rmatrix to None (no need
# to carry out the rotation as we already get the optimum RMSD)
self.rmsd[self._frame_index, 2] = qcp.CalcRMSDRotationalMatrix(
self._ref_coordinates64, self._mobile_coordinates64,
self._n_atoms, None, self.weights)
self._n_atoms, None, self.weights_select)

self._pm.rmsd = self.rmsd[self._frame_index, 2]

Expand Down
70 changes: 61 additions & 9 deletions testsuite/MDAnalysisTests/analysis/test_rms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,17 +172,20 @@ def correct_values(self):
def correct_values_mass(self):
return [[0, 1, 0], [49, 50, 4.74920]]

@pytest.fixture()
def correct_values_mass_add_ten(self):
return [[0, 1, 0.0632], [49, 50, 4.7710]]

@pytest.fixture()
def correct_values_group(self):
return [[0, 1, 0, 0, 0],
[49, 50, 4.7857, 4.7048, 4.6924]]
[49, 50, 4.7857, 4.7048, 4.6924]]

@pytest.fixture()
def correct_values_backbone_group(self):
return [[0, 1, 0, 0, 0],
[49, 50, 4.6997, 1.9154, 2.7139]]


def test_progress_meter(self, capsys, universe):
RMSD = MDAnalysis.analysis.rms.RMSD(universe, verbose=True)
RMSD.run()
Expand All @@ -202,8 +205,8 @@ def test_rmsd_unicode_selection(self, universe, correct_values):
RMSD = MDAnalysis.analysis.rms.RMSD(universe, select=u'name CA')
RMSD.run(step=49)
assert_almost_equal(RMSD.rmsd, correct_values, 4,
err_msg="error: rmsd profile should match" +
"test values")
err_msg="error: rmsd profile should match" +
"test values")

def test_rmsd_atomgroup_selections(self, universe):
# see Issue #1684
Expand Down Expand Up @@ -255,6 +258,29 @@ def test_custom_weighted_list(self, universe, correct_values_mass):
err_msg="error: rmsd profile should match" +
"test values")

def test_custom_groupselection_weights_applied_1D_array(self, universe):
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
select='backbone',
groupselections=['name CA and resid 1-5', 'name CA and resid 1'],
weights=None,
weights_groupselections=[[1, 0, 0, 0, 0], None]).run(step=49)

assert_almost_equal(RMSD.rmsd.T[3], RMSD.rmsd.T[4], 4,
err_msg="error: rmsd profile should match "
"for applied weight array and selected resid")

def test_custom_groupselection_weights_applied_mass(self, universe, correct_values_mass):
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
select='backbone',
groupselections=['all', 'all'],
weights=None,
weights_groupselections=['mass',
universe.atoms.masses]).run(step=49)

assert_almost_equal(RMSD.rmsd.T[3], RMSD.rmsd.T[4], 4,
err_msg="error: rmsd profile should match "
"between applied mass and universe.atoms.masses")

def test_rmsd_scalar_weights_raises_TypeError(self, universe):
with pytest.raises(TypeError):
RMSD = MDAnalysis.analysis.rms.RMSD(
Expand All @@ -270,11 +296,25 @@ def test_rmsd_mismatched_weights_raises_ValueError(self, universe):
RMSD = MDAnalysis.analysis.rms.RMSD(
universe, weights=universe.atoms.masses[:-1])

def test_rmsd_group_selections_wrong_weights(self, universe):
def test_rmsd_misuse_weights_for_groupselection_raises_TypeError(self, universe):
with pytest.raises(TypeError):
RMSD = MDAnalysis.analysis.rms.RMSD(
universe, groupselections=['all'],
weights=[universe.atoms.masses, universe.atoms.masses[:-1]])

def test_rmsd_mismatched_weights_in_groupselection_raises_ValueError(self, universe):
with pytest.raises(ValueError):
RMSD = MDAnalysis.analysis.rms.RMSD(
universe, groupselections=['all'],
weights=universe.atoms.masses,
weights_groupselections = [universe.atoms.masses[:-1]])

def test_rmsd_list_of_weights_wrong_length(self, universe):
with pytest.raises(ValueError):
RMSD = MDAnalysis.analysis.rms.RMSD(
universe, groupselections=['backbone', 'name CA'],
weights=universe.atoms.masses)
weights='mass',
weights_groupselections=[None])

def test_rmsd_group_selections(self, universe, correct_values_group):
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
Expand Down Expand Up @@ -310,14 +350,26 @@ def test_mass_mismatches(self, universe):
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
reference=reference)

def test_ref_mobile_mass_mismapped(self, universe,correct_values_mass_add_ten):
reference = MDAnalysis.Universe(PSF, DCD)
universe.atoms.masses = universe.atoms.masses + 10
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
reference=reference,
select='all',
weights='mass',
tol_mass=100)
RMSD.run(step=49)
assert_almost_equal(RMSD.rmsd, correct_values_mass_add_ten, 4,
err_msg="error: rmsd profile should match "
"between true values and calculated values")

def test_group_selections_unequal_len(self, universe):
reference = MDAnalysis.Universe(PSF, DCD)
reference.atoms[0].residue.resname='NOTMET'
reference.atoms[0].residue.resname = 'NOTMET'
with pytest.raises(SelectionError):
RMSD = MDAnalysis.analysis.rms.RMSD(universe,
reference=reference,
groupselections=
['resname MET','type NH3'])
groupselections=['resname MET', 'type NH3'])


class TestRMSF(object):
Expand Down

0 comments on commit 4102481

Please sign in to comment.