Started parallel mapping issue #10

pycgtool/mapping.py

@@ -7,6 +7,8 @@ Each list corresponds to a single molecule.
 
 import numpy as np
 import warnings
+import itertools
+import multiprocessing
 
 from .frame import Atom, Residue, Frame
 from .parsers.cfg import CFG
@@ -83,6 +85,8 @@ class Mapping:
 
         self._map_center = options.map_center
 
+        # self._pool = multiprocessing.Pool(processes=4)
+
         with CFG(filename) as cfg:
             self._manual_charges = {}
             for mol in cfg:
@@ -179,7 +183,13 @@ class Mapping:
         cgframe.number = frame.number
         cgframe.box = frame.box
 
-        for aares, cgres in zip(aa_residues, cgframe):
+        with multiprocessing.Pool(processes=1) as pool:
+            cgframe.residues = list(pool.map(self.map_one, zip(aa_residues, cgframe, itertools.repeat(cgframe.box))))
+
+        return cgframe
+
+    def map_one(self, tup):
+        aares, cgres, box = tup
         molmap = self._mappings[aares.name]
 
         for i, (bead, bmap) in enumerate(zip(cgres, molmap)):
@@ -187,7 +197,7 @@ class Mapping:
             coords = np.array([aares[atom].coords for atom in bmap], dtype=np.float32)
 
             if self._map_center == "geom":
-                    bead.coords = calc_coords(ref_coords, coords, cgframe.box)
+                bead.coords = calc_coords(ref_coords, coords, box)
             else:
                 try:
                     weights = bmap.weights[self._map_center]
@@ -197,10 +207,9 @@ class Mapping:
                     else:
                         e.args = ("Error with mapping type '{0}', unknown mapping type.".format(e.args[0]),)
                     raise
-                    bead.coords = calc_coords_weight(ref_coords, coords, cgframe.box, weights)
-
-        return cgframe
+                bead.coords = calc_coords_weight(ref_coords, coords, box, weights)
 
+        return cgres
 
 @numba.jit
 def calc_coords_weight(ref_coords, coords, box, weights):