diff --git a/pycgtool/mapping.py b/pycgtool/mapping.py index e93a0101e20b8ec3a7a096dd2203765282643247..a84be223fc8912bf5bb09364a7dff8fccff39340 100644 --- a/pycgtool/mapping.py +++ b/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,28 +183,33 @@ class Mapping: cgframe.number = frame.number cgframe.box = frame.box - for aares, cgres in zip(aa_residues, cgframe): - molmap = self._mappings[aares.name] - - for i, (bead, bmap) in enumerate(zip(cgres, molmap)): - ref_coords = aares[bmap[0]].coords - 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) - else: - try: - weights = bmap.weights[self._map_center] - except KeyError as e: - if self._map_center == "mass": - e.args = ("Error with mapping type 'mass', did you provide an itp file?",) - 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) + 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)): + ref_coords = aares[bmap[0]].coords + 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, box) + else: + try: + weights = bmap.weights[self._map_center] + except KeyError as e: + if self._map_center == "mass": + e.args = ("Error with mapping type 'mass', did you provide an itp file?",) + else: + e.args = ("Error with mapping type '{0}', unknown mapping type.".format(e.args[0]),) + raise + bead.coords = calc_coords_weight(ref_coords, coords, box, weights) + + return cgres @numba.jit def calc_coords_weight(ref_coords, coords, box, weights):