Python wrapper now has a proper Vec3 class that supports math operators

wrappers/python/simtk/openmm/__init__.py

@@ -34,4 +34,5 @@ else:
 
 
 from simtk.openmm.openmm import *
+from simtk.openmm.vec3 import Vec3
 pluginLoadedLibNames = Platform.loadPluginsFromDirectory(Platform.getDefaultPluginsDirectory())
\ No newline at end of file

wrappers/python/simtk/openmm/vec3.py

+"""
+vec3.py: Defines the Vec3 class used by OpenMM
+"""
+__author__ = "Peter Eastman"
+__version__ = "1.0"
+
+import simtk.unit as unit
+
+class Vec3(tuple):
+    """Vec3 is a 3-element tuple that supports many math operations."""
+    
+    def __new__(cls, x, y, z):
+        """Create a new Vec3."""
+        return tuple.__new__(cls, (x, y, z))
+    
+    def __add__(self, other):
+        """Add two Vec3s."""
+        return Vec3(self[0]+other[0], self[1]+other[1], self[2]+other[2])
+    
+    def __radd__(self, other):
+        """Add two Vec3s."""
+        return Vec3(self[0]+other[0], self[1]+other[1], self[2]+other[2])
+    
+    def __sub__(self, other):
+        """Add two Vec3s."""
+        return Vec3(self[0]-other[0], self[1]-other[1], self[2]-other[2])
+    
+    def __rsub__(self, other):
+        """Add two Vec3s."""
+        return Vec3(other[0]-self[0], other[1]-self[1], other[2]-self[2])
+    
+    def __mul__(self, other):
+        """Multiply a Vec3 by a constant."""
+        if unit.is_unit(other):
+            return unit.Quantity(self, other)
+        return Vec3(other*self[0], other*self[1], other*self[2])
+    
+    def __rmul__(self, other):
+        """Multiply a Vec3 by a constant."""
+        if unit.is_unit(other):
+            return unit.Quantity(self, other)
+        return Vec3(other*self[0], other*self[1], other*self[2])
+    
+    def __div__(self, other):
+        """Divide a Vec3 by a constant."""
+        return Vec3(self[0]/other, self[1]/other, self[2]/other)

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

@@ -29,10 +29,18 @@
     OpenMM::Vec3 myVecB;
     OpenMM::Vec3 myVecC;
     state.getPeriodicBoxVectors(myVecA, myVecB, myVecC);
-    pPeriodicBoxVectorsList = Py_BuildValue("(d,d,d),(d,d,d), (d,d,d)",
-                                            myVecA[0], myVecA[1], myVecA[2],
-                                            myVecB[0], myVecB[1], myVecB[2],
-                                            myVecC[0], myVecC[1], myVecC[2]);
+    PyObject* mm = PyImport_AddModule("simtk.openmm");
+    PyObject* vec3 = PyObject_GetAttrString(mm, "Vec3");
+    PyObject* args1 = Py_BuildValue("(d,d,d)", myVecA[0], myVecA[1], myVecA[2]);
+    PyObject* args2 = Py_BuildValue("(d,d,d)", myVecB[0], myVecB[1], myVecB[2]);
+    PyObject* args3 = Py_BuildValue("(d,d,d)", myVecC[0], myVecC[1], myVecC[2]);
+    PyObject* pyVec1 = PyObject_CallObject(vec3, args1);
+    PyObject* pyVec2 = PyObject_CallObject(vec3, args2);
+    PyObject* pyVec3 = PyObject_CallObject(vec3, args3);
+    Py_DECREF(args1);
+    Py_DECREF(args2);
+    Py_DECREF(args3);
+    pPeriodicBoxVectorsList = Py_BuildValue("N,N,N", pyVec1, pyVec2, pyVec3);
 
     if (getPositions) {
       pPositions = copyVVec3ToList(state.getPositions());

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

@@ -4,17 +4,18 @@ namespace OpenMM {
 
 PyObject *copyVVec3ToList(std::vector<Vec3> vVec3) {
   int i, n;
-  PyObject *pyTuple;
   PyObject *pyList;
-  OpenMM::Vec3 vec3;
 
   n=vVec3.size(); 
   pyList=PyList_New(n);
+  PyObject* mm = PyImport_AddModule("simtk.openmm");
+  PyObject* vec3 = PyObject_GetAttrString(mm, "Vec3");
   for (i=0; i<n; i++) {
-    vec3=vVec3.at(i);
-    pyTuple=Py_BuildValue("(d,d,d)",
-                          vec3[0], vec3[1], vec3[2]);
-    PyList_SET_ITEM(pyList, i, pyTuple);
+    OpenMM::Vec3& v = vVec3.at(i);
+    PyObject* args = Py_BuildValue("(d,d,d)", v[0], v[1], v[2]);
+    PyObject* pyVec = PyObject_CallObject(vec3, args);
+    Py_DECREF(args);
+    PyList_SET_ITEM(pyList, i, pyVec);
   }
   return pyList;
 }

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

@@ -74,11 +74,19 @@
 /* Convert C++ (Vec3&, Vec3&, Vec3&) object to python tuple or tuples */
 %typemap(argout) (Vec3& a, Vec3& b, Vec3& c) {
   // %typemap(argout) (Vec3& a, Vec3& b, Vec3& c)
+  PyObject* mm = PyImport_AddModule("simtk.openmm");
+  PyObject* vec3 = PyObject_GetAttrString(mm, "Vec3");
+  PyObject* args1 = Py_BuildValue("(d,d,d)", (*$1)[0], (*$1)[1], (*$1)[2]);
+  PyObject* args2 = Py_BuildValue("(d,d,d)", (*$2)[0], (*$2)[1], (*$2)[2]);
+  PyObject* args3 = Py_BuildValue("(d,d,d)", (*$3)[0], (*$3)[1], (*$3)[2]);
+  PyObject* pyVec1 = PyObject_CallObject(vec3, args1);
+  PyObject* pyVec2 = PyObject_CallObject(vec3, args2);
+  PyObject* pyVec3 = PyObject_CallObject(vec3, args3);
+  Py_DECREF(args1);
+  Py_DECREF(args2);
+  Py_DECREF(args3);
   PyObject *o, *o2, *o3;
-  o = Py_BuildValue("[N, N, N]",
-                    Py_BuildValue("(d, d, d)", (*$1)[0], (*$1)[1], (*$1)[2]),
-                    Py_BuildValue("(d, d, d)", (*$2)[0], (*$2)[1], (*$2)[2]),
-                    Py_BuildValue("(d, d, d)", (*$3)[0], (*$3)[1], (*$3)[2]));
+  o = Py_BuildValue("[N, N, N]", pyVec1, pyVec2, pyVec3);
   if ((!$result) || ($result == Py_None)) {
     $result = o;
   } else {