Merge pull request #149 from marrink-lab/cog

Allow do_average_bead to use different centerings

vermouth/data/force_fields/elnedyn/general.ff

@@ -0,0 +1,16 @@
+; Copyright 2018 University of Groningen
+;
+; Licensed under the Apache License, Version 2.0 (the "License");
+; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;    http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+
+[ variables ]
+center_weight "mass"

vermouth/data/force_fields/elnedyn22/general.ff

@@ -0,0 +1,16 @@
+; Copyright 2018 University of Groningen
+;
+; Licensed under the Apache License, Version 2.0 (the "License");
+; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;    http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+
+[ variables ]
+center_weight "mass"

vermouth/data/force_fields/elnedyn22p/general.ff

@@ -0,0 +1,16 @@
+; Copyright 2018 University of Groningen
+;
+; Licensed under the Apache License, Version 2.0 (the "License");
+; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;    http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+
+[ variables ]
+center_weight "mass"

vermouth/data/force_fields/martini22/general.ff

@@ -0,0 +1,16 @@
+; Copyright 2018 University of Groningen
+;
+; Licensed under the Apache License, Version 2.0 (the "License");
+; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;    http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+
+[ variables ]
+center_weight "mass"

vermouth/data/force_fields/martini22p/general.ff

@@ -0,0 +1,16 @@
+; Copyright 2018 University of Groningen
+;
+; Licensed under the Apache License, Version 2.0 (the "License");
+; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;    http://www.apache.org/licenses/LICENSE-2.0
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS,
+; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+; See the License for the specific language governing permissions and
+; limitations under the License.
+
+[ variables ]
+center_weight "mass"

vermouth/processors/average_beads.py

@@ -21,7 +21,7 @@
 from .processor import Processor
 
 
-def do_average_bead(molecule, ignore_missing_graphs=False):
+def do_average_bead(molecule, ignore_missing_graphs=False, weight=None):
     """
     Set the position of the particles to the mean of the underlying atoms.
 
@@ -36,6 +36,13 @@ def do_average_bead(molecule, ignore_missing_graphs=False):
     atomname from the underlying graph as keys, and the weights as values.
     Atoms without a weight set use a default weight of 1.
 
+    The average can also be weighted using an arbitrary node attribute by
+    giving the attribute name with the `weight` keyword argument. This can be
+    used to get the center of mass for instance; assuming the mass of the
+    underlying atoms is stored under the "mass" attribute, setting `weight` to
+    "mass" will place the bead at the center of mass. By default, `weight` is
+    set to `None` and the center of geometry is used.
+
     The atoms in the underlying graph must have a position. If they do not,
     they are ignored from the average.
 
@@ -48,12 +55,22 @@ def do_average_bead(molecule, ignore_missing_graphs=False):
     ignore_missing_graphs: bool
         If `True`, skip the atoms that do not have a `graph` attribute; else
         fail if not all the atoms in the molecule have a `graph` attribute.
+    weight: collections.abc.Hashable
+        The name of the attribute used to weight the position of the node. The
+        attribute is read from the underlying atoms.
     """
     # Make sure the molecule fullfill the requirements.
     missing = []
     for node in molecule.nodes.values():
         if 'graph' not in node:
             missing.append(node)
+        elif weight is not None:
+            have_all_weights =  all(
+                weight in subnode for subnode in node['graph'].nodes.values()
+            )
+            if not have_all_weights:
+                raise KeyError('Not all underlying atoms have an attribute {}.'
+                               .format(weight))
     if missing and not ignore_missing_graphs:
         raise ValueError('{} particles are missing the graph attribute'
                          .format(len(missing)))
@@ -66,7 +83,7 @@ def do_average_bead(molecule, ignore_missing_graphs=False):
                 if subnode.get('position') is not None
             ])
             weights = np.array([
-                node.get('mapping_weights', {}).get(subnode_key, 1) * subnode['mass']
+                node.get('mapping_weights', {}).get(subnode_key, 1) * subnode.get(weight, 1)
                 for subnode_key, subnode in node['graph'].nodes.items()
                 if subnode.get('position') is not None
             ])
@@ -76,9 +93,16 @@ def do_average_bead(molecule, ignore_missing_graphs=False):
 
 
 class DoAverageBead(Processor):
-    def __init__(self, ignore_missing_graphs=False):
+    def __init__(self, ignore_missing_graphs=False, weight=None):
         super().__init__()
         self.ignore_missing_graphs = ignore_missing_graphs
+        self.weight = weight
 
     def run_molecule(self, molecule):
-        return do_average_bead(molecule, self.ignore_missing_graphs)
+        if self.weight is None:
+            weight = molecule.force_field.variables.get('center_weight', None)
+        elif self.weight is False:
+            weight = None
+        else:
+            weight = self.weight
+        return do_average_bead(molecule, self.ignore_missing_graphs, weight=weight)

vermouth/tests/test_average_beads.py

@@ -0,0 +1,133 @@
+# Copyright 2018 University of Groningen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for the :mod:`vermouth.processors.average_beads` module.
+"""
+
+import pytest
+import networkx as nx
+import numpy as np
+
+from vermouth.processors import average_beads
+
+
+@pytest.fixture
+def mol_with_subgraph():
+    """
+    Create a molecule with a subgraph under the "graph" attribute of each node.
+    """
+    mol = nx.OrderedGraph()
+
+    subgraph = nx.Graph()
+    subgraph.add_nodes_from((
+        (0, {'mass': 1.1, 'not mass': 2.1, 'position': np.array([1, 2, 3], dtype=float),}),
+        (1, {'mass': 1.2, 'not mass': 2.2, 'position': np.array([2, 3, 4], dtype=float),}),
+        (2, {'mass': 1.3, 'not mass': 2.3, 'position': np.array([3, 4, 5], dtype=float),}),
+    ))
+    mol.add_node(0, **{
+        "graph": subgraph,
+        "mapping_weights": {0: 1, 1: 2, 2: 3},
+        "target mass": np.array([2.378378378, 3.378378378, 4.378378378]),
+        "target not mass": np.array([2.358208955, 3.358208955, 4.358208955]),
+        "target None": np.array([2.333333333, 3.333333333, 4.333333333]),
+        "target False": np.array([2.333333333, 3.333333333, 4.333333333]),
+    })
+
+
+    subgraph = nx.Graph()
+    subgraph.add_nodes_from((
+        (0, {'mass': 1.2, 'not mass': 2.2, 'position': np.array([2, 3, 4], dtype=float),}),
+        (1, {'mass': 1.3, 'not mass': 2.3, 'position': np.array([3, 4, 5], dtype=float),}),
+    ))
+    mol.add_node(1, **{
+        "graph": subgraph,
+        "mapping_weights": {0: 2, 1: 3},
+        "target mass": np.array([2.619047619, 3.619047619, 4.619047619]),
+        "target not mass": np.array([2.610619469, 3.610619469, 4.610619469]),
+        "target None": np.array([2.6, 3.6, 4.6]),
+        "target False": np.array([2.6, 3.6, 4.6]),
+    })
+
+    return mol
+
+@pytest.fixture(params=(None, 'mass', 'not mass'))
+def mol_with_variable(request, mol_with_subgraph):
+    """
+    Build a mock molecule with a mock force field declaring 'center_weight'.
+    """
+    weight = request.param
+
+    class MockForceField(object):
+        pass
+
+    ff = MockForceField()
+    ff.variables = {'center_weight': weight}
+
+    mol_with_subgraph.force_field = ff
+    return mol_with_subgraph
+
+
+@pytest.mark.parametrize('weight', (None, 'mass', 'not mass'))
+def test_do_average_bead(mol_with_subgraph, weight):
+    """
+    Test normal operation of :func:`average_beads.do_average_bead`.
+    """
+    result_mol = average_beads.do_average_bead(
+        mol_with_subgraph, ignore_missing_graphs=False, weight=weight,
+    )
+    target_key = 'target {}'.format(weight)
+    target_positions = np.stack([node[target_key] for node in mol_with_subgraph.nodes.values()])
+    positions = np.stack([node['position'] for node in mol_with_subgraph.nodes.values()])
+    assert np.allclose(positions, target_positions)
+
+
+@pytest.mark.parametrize('weight', ('mass', 'not mass'))
+def test_shoot_weight(mol_with_subgraph, weight):
+    """
+    Test that :func:`average_beads.do_average_bead` fails if a weight is missing.
+    """
+    del mol_with_subgraph.nodes[0]['graph'].nodes[1][weight]
+    with pytest.raises(KeyError):
+        average_beads.do_average_bead(
+            mol_with_subgraph, ignore_missing_graphs=False, weight=weight,
+        )
+
+
+def test_shoot_graph(mol_with_subgraph):
+    """
+    Test that :func:`average_beads.do_average_bead` fails if a subgraph is missing.
+    """
+    del mol_with_subgraph.nodes[1]['graph']
+    with pytest.raises(ValueError):
+        average_beads.do_average_bead(mol_with_subgraph)
+
+
+def test_processor_variable(mol_with_variable):
+    processor = average_beads.DoAverageBead()
+    mol = processor.run_molecule(mol_with_variable)
+    weight = mol_with_variable.force_field.variables['center_weight']
+    target_key = 'target {}'.format(weight)
+    target_positions = np.stack([node[target_key] for node in mol_with_variable.nodes.values()])
+    positions = np.stack([node['position'] for node in mol_with_variable.nodes.values()])
+    assert np.allclose(positions, target_positions)
+
+
+@pytest.mark.parametrize('weight', (False, 'mass', 'not mass'))
+def test_processor_weight(mol_with_variable, weight):
+    processor = average_beads.DoAverageBead(weight=weight)
+    mol = processor.run_molecule(mol_with_variable)
+    target_key = 'target {}'.format(weight)
+    target_positions = np.stack([node[target_key] for node in mol_with_variable.nodes.values()])
+    positions = np.stack([node['position'] for node in mol_with_variable.nodes.values()])
+    assert np.allclose(positions, target_positions)