diff --git a/pycgtool/frame.py b/pycgtool/frame.py index dc4907a7adbba6003e3f5e858bd40f76fd3b6d77..f7d0c831473e45e413b97010ecc5904d011f8e24 100644 --- a/pycgtool/frame.py +++ b/pycgtool/frame.py @@ -136,6 +136,7 @@ class Frame: self.residues = [] self.number = frame_start - 1 self.numframes = 0 + self.natoms = 0 self.box = np.zeros(3, dtype=np.float32) self._xtc_buffer = None diff --git a/pycgtool/mapping.py b/pycgtool/mapping.py index 701490369ed05fbe710704c9f8c45a8653a9ab4d..da12cebe78779c2bd1c4f9ae1e7caac065b3d12f 100644 --- a/pycgtool/mapping.py +++ b/pycgtool/mapping.py @@ -116,80 +116,78 @@ class Mapping: def __iter__(self): return iter(self._mappings) - def apply(self, frame, cgframe=None, exclude=None): + def _cg_frame_setup(self, aa_residues, name=None): """ - Apply the AA->CG mapping to an atomistic Frame. - - :param frame: Frame to which mapping will be applied - :param cgframe: CG Frame to remap - optional - :param exclude: Set of molecule names to exclude from mapping - e.g. solvent - :return: A new Frame instance containing the CG frame + Create a new CG Frame and populate beads + :param aa_residues: Iterable of atomistic residues to map from + :param name: Name of Frame + :return: New CG Frame instance """ - if cgframe is None: - cgframe = Frame() - cgframe.name = frame.name - - cgframe.natoms = 0 - cgframe.residues = [] + cgframe = Frame() + cgframe.name = name - cgframe.number = frame.number - cgframe.box = frame.box + for aares in aa_residues: + molmap = self._mappings[aares.name] + cgres = Residue(name=aares.name, num=aares.num) + cgres.atoms = [Atom(name=bead.name, type=bead.type, charge=bead.charge, mass=bead.mass) for bead in molmap] - for aares in frame: - if aares.name not in self._mappings: - continue - if exclude is not None and aares.name in exclude: - continue + for i, (bead, bmap) in enumerate(zip(cgres, molmap)): + cgres.name_to_num[bead.name] = i + bead.charge = bmap.charge + bead.mass = bmap.mass - res = self._apply_res_pbc(aares, frame.box) + cgframe.add_residue(cgres) + cgframe.natoms += len(cgres) - cgframe.natoms += len(res) - cgframe.residues.append(res) return cgframe - def _apply_res_pbc(self, aares, box): + def apply(self, frame, cgframe=None, exclude=None): """ - Apply mapping transformation to a single residue to allow multithreading. + Apply the AA->CG mapping to an atomistic Frame. - :param aares: Atomistic residue to apply mapping - :param box: Cubic periodic box vectors - :return: A single coarse grained residue + :param frame: Frame to which mapping will be applied + :param cgframe: CG Frame to remap - optional + :param exclude: Set of molecule names to exclude from mapping - e.g. solvent + :return: Frame instance containing the CG frame """ + select_predicate = lambda res: res.name in self._mappings and not (exclude is not None and res.name in exclude) + aa_residues = (aares for aares in frame if select_predicate(aares)) - molmap = self._mappings[aares.name] - res = Residue(name=aares.name, num=aares.num) - res.atoms = [Atom(name=bead.name, type=bead.type, charge=bead.charge, mass=bead.mass) for bead in molmap] - - # Perform mapping - for i, (bead, bmap) in enumerate(zip(res, molmap)): - res.name_to_num[bead.name] = i - bead.charge = bmap.charge - bead.mass = bmap.mass + if cgframe is None: + aa_residues = list(aa_residues) + cgframe = self._cg_frame_setup(aa_residues, frame.name) - ref_coords = aares[bmap[0]].coords - coords = np.array([aares[atom].coords for atom in bmap], dtype=np.float32) + cgframe.number = frame.number + cgframe.box = frame.box - 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, unknown mapping type '{0}'".format(e.args[0]),) - raise - bead.coords = calc_coords_weight(ref_coords, coords, box, weights) + 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, unknown mapping type '{0}'".format(e.args[0]),) + raise + bead.coords = calc_coords_weight(ref_coords, coords, cgframe.box, weights) - return res + return cgframe @jit def calc_coords_weight(ref_coords, coords, box, weights): coords = dist_with_pbc(ref_coords, coords, box) coords = np.sum(weights * coords, axis=0) - coords /= np.sum(weights) + coords /= sum(weights) coords += ref_coords return coords