Modify Context.getState to accept a set of indices in addition to a bitmask

wrappers/python/src/swig_doxygen/swig_lib/python/extend.i

@@ -55,25 +55,30 @@
          - getEnergy (bool=False) whether to store potential and kinetic energy in the State
          - getParameter (bool=False) whether to store context parameters in the State
          - enforcePeriodicBox (bool=False) if false, the position of each particle will be whatever position is stored in the Context, regardless of periodic boundary conditions.  If true, particle positions will be translated so the center of every molecule lies in the same periodic box.
-         - groups (int=-1) a set of bit flags for which force groups to include when computing forces and energies.  Group i will be included if (groups&(1<<i)) != 0.  The default value includes all groups.
+         - groups (set={0,1,2,...,31}) a set of indices for which force groups
+           to include when computing forces and energies. The default value
+           includes all groups. groups can also be passed as an unsigned integer
+           interpreted as a bitmask, in which case group i will be included if
+           (groups&(1<<i)) != 0.
         """
+        getP, getV, getF, getE, getPa, enforcePeriodic = map(bool,
+            (getPositions, getVelocities, getForces, getEnergy, getParameters,
+             enforcePeriodicBox))
 
-        if getPositions: getP=1
-        else: getP=0
-        if getVelocities: getV=1
-        else: getV=0
-        if getForces: getF=1
-        else: getF=0
-        if getEnergy: getE=1
-        else: getE=0
-        if getParameters: getPa=1
-        else: getPa=0
-        if enforcePeriodicBox: enforcePeriodic=1
-        else: enforcePeriodic=0
+        try:
+            # is the input integer-like?
+            groups_mask = int(groups)
+        except TypeError:
+            if isinstance(groups, set):
+                # nope, okay, then it should be an set
+                groups_mask = functools.reduce(operator.or_,
+                        ((1<<x) & 0xffffffff for x in groups))
+            else:
+                raise TypeError('%s is neither an int nor set' % groups)
 
         (simTime, periodicBoxVectorsList, energy, coordList, velList,
          forceList, paramMap) = \
-            self._getStateAsLists(getP, getV, getF, getE, getPa, enforcePeriodic, groups)
+            self._getStateAsLists(getP, getV, getF, getE, getPa, enforcePeriodic, groups_mask)
 
         state = State(simTime=simTime,
                       energy=energy,
@@ -207,25 +212,29 @@ Parameters:
          - getEnergy (bool=False) whether to store potential and kinetic energy in the State
          - getParameter (bool=False) whether to store context parameters in the State
          - enforcePeriodicBox (bool=False) if false, the position of each particle will be whatever position is stored in the Context, regardless of periodic boundary conditions.  If true, particle positions will be translated so the center of every molecule lies in the same periodic box.
-         - groups (int=-1) a set of bit flags for which force groups to include when computing forces and energies.  Group i will be included if (groups&(1<<i)) != 0.  The default value includes all groups.
+         - groups (set={0,1,2,...,31}) a set of indices for which force groups
+           to include when computing forces and energies. The default value
+           includes all groups. groups can also be passed as an unsigned integer
+           interpreted as a bitmask, in which case group i will be included if
+           (groups&(1<<i)) != 0.
         """
+        getP, getV, getF, getE, getPa, enforcePeriodic = map(bool,
+            (getPositions, getVelocities, getForces, getEnergy, getParameters,
+             enforcePeriodicBox))
 
-        if getPositions: getP=1
-        else: getP=0
-        if getVelocities: getV=1
-        else: getV=0
-        if getForces: getF=1
-        else: getF=0
-        if getEnergy: getE=1
-        else: getE=0
-        if getParameters: getPa=1
-        else: getPa=0
-        if enforcePeriodicBox: enforcePeriodic=1
-        else: enforcePeriodic=0
+        try:
+            # is the input integer-like?
+            groups_mask = int(groups)
+        except TypeError:
+            if isinstance(groups, set):
+                groups_mask = functools.reduce(operator.or_,
+                        ((1<<x) & 0xffffffff for x in groups))
+            else:
+                raise TypeError('%s is neither an int nor set' % groups)
 
         (simTime, periodicBoxVectorsList, energy, coordList, velList,
           forceList, paramMap) = \
-            self._getStateAsLists(copy, getP, getV, getF, getE, getPa, enforcePeriodic, groups)
+             self._getStateAsLists(copy, getP, getV, getF, getE, getPa, enforcePeriodic, groups_mask)
 
         state = State(simTime=simTime,
                       energy=energy,

wrappers/python/src/swig_doxygen/swig_lib/python/pythoncode.i

@@ -8,6 +8,8 @@ except ImportError:
 import copy
 import sys
 import math
+import functools
+import operator
 RMIN_PER_SIGMA=math.pow(2, 1/6.0)
 RVDW_PER_SIGMA=math.pow(2, 1/6.0)/2.0
 if sys.version_info[0] == 2:

wrappers/python/tests/TestForceGroups.py

+import unittest
+import itertools
+import simtk.openmm as mm
+
+
+class TestForceGroups(unittest.TestCase):
+    def setUp(self):
+        system = mm.System()
+        system.addParticle(1.0)
+
+        for i in range(32):
+            force = mm.CustomExternalForce(str(i))
+            force.addParticle(0, [])
+            force.setForceGroup(i)
+            system.addForce(force)
+
+        platform = mm.Platform.getPlatformByName('Reference')
+        context = mm.Context(system, mm.VerletIntegrator(0), platform)
+        context.setPositions([(0,0,0)])
+        self.context = context
+
+    def test1(self):
+        n = 31  # Should be 32, but github issue #1198
+        for (i,j) in itertools.combinations(range(n), 2):
+            groups = 1<<i | 1<<j
+            e_0 = self.context.getState(getEnergy=True, groups=groups).getPotentialEnergy()._value
+            e_1 = self.context.getState(getEnergy=True, groups={i,j}).getPotentialEnergy()._value
+            e_ref = i+j
+            self.assertEqual(e_0, e_ref)
+            self.assertEqual(e_1, e_ref)
+
+    def test2(self):
+        with self.assertRaises(TypeError):
+            # groups must be an int or set
+            self.context.getState(getEnergy=True, groups=(1, 2))
+
+if __name__ == '__main__':
+    unittest.main()
+
+