def write_gmx_topology(system, filename, header=""): """ Write out gmx topology file Parameters ---------- system : mb.Compound() filename : str header : str, optional Header string for include statements etc Notes ----- Molecule names are based on the one-level-up parents of the particles """ molecules = OrderedSet() for p in system.particles(): if p.parent is not None: molecules.add(p.parent) with open(filename, 'w') as f: f.write("{}\n\n".format(header)) f.write("[ system ]\n") f.write("mBuild Bilayer System\n\n") f.write("[ molecules ]\n") f.write("\n".join([ "{:<8s} {}".format(name, sum(1 for _ in g)) for name, g in groupby([c.name for c in molecules]) ]))
class UniqueQueue(Queue): def _init(self, maxsize): self.queue = OrderedSet() def _put(self, item): self.queue.add(item) def _get(self): return self.queue.pop() def __contains__(self, item): with self.mutex: return item in self.queue
def __init__(self, subcompounds=None, name=None, pos=None, charge=0.0, periodicity=None, port_particle=False): super(Compound, self).__init__() if name: if not isinstance(name, string_types): raise ValueError( 'Compound.name should be a string. You passed ' '{}'.format(name)) self.name = name else: self.name = self.__class__.__name__ # A periodocity of zero in any direction is treated as non-periodic. if periodicity is None: self._periodicity = np.array([0.0, 0.0, 0.0]) else: self._periodicity = np.asarray(periodicity) if pos is not None: self._pos = np.asarray(pos, dtype=float) else: self._pos = np.zeros(3) self.charge = charge self.parent = None self.children = OrderedSet() self.labels = OrderedDict() self.referrers = set() self.bond_graph = None self.port_particle = port_particle # self.add() must be called after labels and children are initialized. if subcompounds: self.add(subcompounds)
def __init__(self, subcompounds=None, name=None, pos=None, charge=0.0, periodicity=None, port_particle=False): super(Compound, self).__init__() if name: if not isinstance(name, string_types): raise ValueError('Compound.name should be a string. You passed' '{}'.format(name)) self.name = name else: self.name = self.__class__.__name__ # A periodocity of zero in any direction is treated as non-periodic. if periodicity is None: self._periodicity = np.array([0.0, 0.0, 0.0]) else: self._periodicity = np.asarray(periodicity) if pos is not None: self._pos = np.asarray(pos, dtype=float) else: self._pos = np.zeros(3) self.charge = charge self.parent = None self.children = OrderedSet() self.labels = OrderedDict() self.referrers = set() self.bond_graph = None self.port_particle = port_particle # self.add() must be called after labels and children are initialized. if subcompounds: self.add(subcompounds)
class Compound(object): """A building block in the mBuild hierarchy. Compound is the superclass of all composite building blocks in the mBuild hierarchy. That is, all composite building blocks must inherit from compound, either directly or indirectly. The design of Compound follows the Composite design pattern (Gamma, Erich; Richard Helm; Ralph Johnson; John M. Vlissides (1995). Design Patterns: Elements of Reusable Object-Oriented Software. Addison-Wesley. p. 395. ISBN 0-201-63361-2.), with Compound being the composite, and Particle playing the role of the primitive (leaf) part, where Particle is in fact simply an alias to the Compound class. Compound maintains a list of children (other Compounds contained within), and provides a means to tag the children with labels, so that the compounds can be easily looked up later. Labels may also point to objects outside the Compound's containment hierarchy. Compound has built-in support for copying and deepcopying Compound hierarchies, enumerating particles or bonds in the hierarchy, proximity based searches, visualization, I/O operations, and a number of other convenience methods. Parameters ---------- subcompounds : Compound, optional, default=None One or more compounds to be added to self. name : str, optional, default=self.__class__.__name__ The type of Compound. periodicity : np.ndarray, shape=(3,), dtype=float, optional The periodic lengths of the Compound in the x, y and z directions. Defaults to zeros which is treated as non-periodic. Attributes ---------- name : str, optional, default=self.__class__.__name__ The type of Compound. periodicity : np.ndarray, shape=(3,), dtype=float, optional The periodic lengths of the Compound in the x, y and z directions. Defaults to zeros which is treated as non-periodic. children : OrderedSet Contains all children (other Compounds). labels : OrderedDict Labels to Compound/Atom mappings. These do not necessarily need not be in self.children. parent : mb.Compound The parent Compound that contains this part. Can be None if this compound is the root of the containment hierarchy. referrers : set Other compounds that reference this part with labels. """ def __init__(self, subcompounds=None, name=None, pos=None, charge=0.0, periodicity=None, port_particle=False): super(Compound, self).__init__() if name: if not isinstance(name, string_types): raise ValueError('Compound.name should be a string. You passed' '{}'.format(name)) self.name = name else: self.name = self.__class__.__name__ # A periodocity of zero in any direction is treated as non-periodic. if periodicity is None: self._periodicity = np.array([0.0, 0.0, 0.0]) else: self._periodicity = np.asarray(periodicity) if pos is not None: self._pos = np.asarray(pos, dtype=float) else: self._pos = np.zeros(3) self.charge = charge self.parent = None self.children = OrderedSet() self.labels = OrderedDict() self.referrers = set() self.bond_graph = None self.port_particle = port_particle # self.add() must be called after labels and children are initialized. if subcompounds: self.add(subcompounds) def particles(self, include_ports=False): """ """ if not self.children: yield self else: for particle in self._particles(include_ports): yield particle def _particles(self, include_ports=False): """Return all Particles of the Compound. """ for child in self.successors(): if not child.children: if include_ports or not child.port_particle: yield child def successors(self): """Yield Compounds below self in the hierarchy. """ if not self.children: return for part in self.children: # Parts local to the current Compound. yield part # Parts further down the hierarchy. for subpart in part.successors(): yield subpart @property def n_particles(self): if not self.children: return 1 else: return self._n_particles(include_ports=False) def _n_particles(self, include_ports=False): """Return the number of Particles in the Compound. """ return sum(1 for _ in self._particles(include_ports)) def _contains_only_ports(self): for part in self.children: if not part.port_particle: return False return True def ancestors(self): """Generate all ancestors of the Compound recursively. """ if self.parent is not None: yield self.parent for ancestor in self.parent.ancestors(): yield ancestor @property def root(self): parent = None for parent in self.ancestors(): pass if parent is None: return self return parent def particles_by_name(self, name): for particle in self.particles(): if particle.name == name: yield particle def add(self, new_child, label=None, containment=True, replace=False, inherit_periodicity=True): """Add a part to the Compound. Note: This does not necessarily add the part to self.children but may instead be used to add a reference to the part to self.labels. See 'containment' argument. Parameters ---------- new_child : mb.Compound The object to be added to this Compound. label : str, optional A descriptive string for the part. containment : bool, optional, default=True Add the part to self.children. replace : bool, optional, default=True Replace the label if it already exists. """ # Support batch add via lists, tuples and sets. if not isinstance(new_child, Compound): for child in new_child: self.add(child) return # Create children and labels on the first add operation if self.children is None: self.children = OrderedSet() if self.labels is None: self.labels = OrderedDict() if not isinstance(new_child, Compound): raise ValueError('Only objects that inherit from mbuild.Compound ' 'can be added to Compounds. You tried to add ' '{}'.format(new_child)) if containment: if new_child.parent is not None: raise ValueError('Part {} already has a parent: {}'.format( new_child, new_child.parent)) self.children.add(new_child) new_child.parent = self if new_child.bond_graph is not None: if self.root.bond_graph is None: self.root.bond_graph = new_child.bond_graph else: self.root.bond_graph = nx.compose(self.root.bond_graph, new_child.bond_graph) new_child.bond_graph = None # Add new_part to labels. Does not currently support batch add. if label is None: label = '{0}[$]'.format(new_child.__class__.__name__) if label is not None: if label.endswith('[$]'): label = label[:-3] if label not in self.labels: self.labels[label] = [] label_pattern = label + '[{}]' count = len(self.labels[label]) self.labels[label].append(new_child) label = label_pattern.format(count) if not replace and label in self.labels: raise Exception('Label {0} already exists in {1}'.format(label, self)) else: self.labels[label] = new_child new_child.referrers.add(self) if (inherit_periodicity and isinstance(new_child, Compound) and new_child.periodicity.any()): self.periodicity = new_child.periodicity def remove(self, objs_to_remove): """Remove children from the Compound. """ if not self.children: return if not hasattr(objs_to_remove, '__iter__'): objs_to_remove = [objs_to_remove] objs_to_remove = set(objs_to_remove) if len(objs_to_remove) == 0: return intersection = objs_to_remove.intersection(self.children) self.children -= intersection objs_to_remove -= intersection for removed_part in intersection: if self.root.bond_graph and self.root.bond_graph.has_node(removed_part): self.root.bond_graph.remove_node(removed_part) self._remove_references(removed_part) # Remove the part recursively from sub-compounds. if self.children: for part in self.children: part.remove(objs_to_remove) @staticmethod def _remove_references(removed_part): """Remove labels pointing to this part and vice versa. """ removed_part.parent = None # Remove labels in the hierarchy pointing to this part. referrers_to_remove = set() for referrer in removed_part.referrers: if removed_part not in referrer.ancestors(): for label, referred_part in referrer.labels.items(): if referred_part is removed_part: del referrer.labels[label] referrers_to_remove.add(referrer) removed_part.referrers -= referrers_to_remove # Remove labels in this part pointing into the hierarchy. labels_to_delete = [] if isinstance(removed_part, Compound): for label, part in removed_part.labels.items(): if removed_part not in part.ancestors(): part.referrers.remove(removed_part) labels_to_delete.append(label) for label in labels_to_delete: del removed_part.labels[label] def referenced_ports(self): """Return all Ports referenced by this Compound. """ from mbuild.port import Port return [port for port in self.labels.values() if isinstance(port, Port)] def bonds(self): """A list of all Bonds in the Compound and sub-Compounds. """ if self.root.bond_graph: return self.root.bond_graph.subgraph(self.particles()).edges_iter() else: return iter(()) @property def n_bonds(self): """Return the number of Bonds in the Compound. """ return sum(1 for _ in self.bonds()) def add_bond(self, particle_pair): """""" if self.root.bond_graph is None: self.root.bond_graph = nx.Graph() self.root.bond_graph.add_edge(particle_pair[0], particle_pair[1]) def generate_bonds(self, name_a, name_b, dmin, dmax): """Add Bonds between all pairs of types a/b within [dmin, dmax]. """ particle_kdtree = PeriodicCKDTree(data=self.xyz, bounds=self.periodicity) particle_array = np.array(list(self.particles())) added_bonds = list() for p1 in self.particles_by_name(name_a): nearest = self.particles_in_range(p1, dmax, max_particles=20, particle_kdtree=particle_kdtree, particle_array=particle_array) for p2 in nearest: if p2 == p1: continue bond_tuple = (p1, p2) if id(p1) < id(p2) else (p2, p1) if bond_tuple in added_bonds: continue if (p2.name == name_b) and (dmin <= self.min_periodic_distance(p2.pos, p1.pos) <= dmax): self.add_bond((p1, p2)) added_bonds.append(bond_tuple) def remove_bond(self, particle_pair): if self.root.bond_graph is None or not self.root.bond_graph.has_edge(*particle_pair): warn("Bond between {} and {} doesn't exist!".format(*particle_pair)) return self.root.bond_graph.remove_edge(*particle_pair) for particle in particle_pair: if not self.root.bond_graph.neighbors(particle): self.root.bond_graph.remove_node(particle) @property def pos(self): if not self.children: return self._pos else: return self.center @pos.setter def pos(self, value): if not self.children: self._pos = value else: raise Exception("Cannot set position on a Compound that has" " children.") @property def periodicity(self): return self._periodicity @periodicity.setter def periodicity(self, periods): self._periodicity = np.array(periods) @property def xyz(self): """Return all particle coordinates in this compound. Returns ------- pos : np.ndarray, shape=(n, 3) Array with the positions of all particles. """ if not self.children: pos = np.expand_dims(self._pos, axis=0) else: arr = np.fromiter(itertools.chain.from_iterable( particle.pos for particle in self.particles()), dtype=float) pos = arr.reshape((-1, 3)) return pos @property def xyz_with_ports(self): """Return all particle coordinates in this compound including ports. """ if not self.children: pos = self._pos else: arr = np.fromiter(itertools.chain.from_iterable( particle.pos for particle in self.particles(include_ports=True)), dtype=float) pos = arr.reshape((-1, 3)) return pos @property def center(self): """The cartesian center of the Compound based on its Atoms. """ if self.xyz.any(): return np.mean(self.xyz, axis=0) @property def boundingbox(self): """Compute the bounding box of the compound. """ xyz = self.xyz return Box(mins=xyz.min(axis=0), maxs=xyz.max(axis=0)) def min_periodic_distance(self, xyz0, xyz1): """Vectorized distance calculation considering minimum image. """ d = np.abs(xyz0 - xyz1) d = np.where(d > 0.5 * self.periodicity, self.periodicity - d, d) return np.sqrt((d ** 2).sum(axis=-1)) def particles_in_range(self, compound, dmax, max_particles=20, particle_kdtree=None, particle_array=None): """Find particles ithin a specified range of another particle. """ if particle_kdtree is None: particle_kdtree = PeriodicCKDTree(data=self.xyz, bounds=self.periodicity) _, idxs = particle_kdtree.query(compound.pos, k=max_particles, distance_upper_bound=dmax) # TODO: why we are doing this idxs = idxs[idxs != self.n_particles] if particle_array is None: particle_array = np.array(list(self.particles())) return particle_array[idxs] def view_hierarchy(self, show_ports=False): """Visualize a compound hierarchy as a tree. A tree is constructed from the compound hierarchy with self as the root. The tree is then rendered in a web browser window using D3.js. Note ------ Portions of this code are adapted from https://gist.github.com/mbostock/4339083. """ raise NotImplementedError('To be replaced with igraph') def visualize(self, show_ports=False, shader='lambert', drawing_type='ball and stick', camera_type='perspective', element_properties=None): """Visualize the Compound using imolecule. """ json_mol = self._to_json(show_ports) imolecule.draw(json_mol, format='json', shader=shader, drawing_type=drawing_type, camera_type=camera_type, element_properties=element_properties) def _to_json(self, show_ports=False): atoms = list() for idx, particle in enumerate(self._particles(include_ports=show_ports)): particle.index = idx atoms.append({'element': particle.name, 'location': list(np.asarray(particle.pos, dtype=float) * 10)}) bonds = [{'atoms': [atom1.index, atom2.index], 'order': 1} for atom1, atom2 in self.bonds()] output = {'name': self.name, 'atoms': atoms, 'bonds': bonds} # Remove the index attribute on particles. for idx, particle in enumerate(self.particles()): if not show_ports and particle.port_particle: continue del particle.index return imolecule.json_formatter.compress(output) def update_coordinates(self, filename): """Update the coordinates of this Compound from a file. """ load(filename, compound=self, coords_only=True) def save(self, filename, show_ports=False, forcefield=None, **kwargs): """Save the Compound to a file. Parameters ---------- filename : str Filesystem path in which to save the trajectory. The extension or prefix will be parsed and will control the format. show_ports : bool, default=False Save ports contained within the compound. forcefield : str, default=None Apply a forcefield to the output file using the `foyer` package. Other Parameters ---------------- force_overwrite : bool """ extension = os.path.splitext(filename)[-1] savers = {'.hoomdxml': self.save_hoomdxml, '.gro': self.save_gromacs, '.top': self.save_gromacs, '.mol2': self.save_mol2, '.lammps': self.save_lammpsdata, '.lmp': self.save_lammpsdata} try: saver = savers[extension] except KeyError: # TODO: better reporting saver = None if (not saver or extension == '.mol2') and forcefield: ff_formats = ', '.join(set(savers.keys()) - set(['.mol2'])) raise ValueError('The only supported formats with forcefield' 'information are: {0}'.format(ff_formats)) if saver: # mBuild/InterMol supported saver. traj = self.to_trajectory(show_ports=show_ports, **kwargs) return saver(filename, traj, forcefield, **kwargs) else: # MDTraj supported saver. traj = self.to_trajectory(show_ports=show_ports, **kwargs) return traj.save(filename, **kwargs) def save_mol2(self, filename, traj, forcefield, **kwargs): """ """ write_mol2(filename, traj) def save_hoomdxml(self, filename, traj, forcefield, force_overwrite=False, **kwargs): """ """ raise NotImplementedError('Interface to InterMol missing') def save_gromacs(self, filename, traj, forcefield, force_overwrite=False, **kwargs): """ """ from foyer.forcefield import apply_forcefield # Create separate file paths for .gro and .top filepath, filename = os.path.split(filename) basename = os.path.splitext(filename)[0] top_filename = os.path.join(filepath, basename + '.top') gro_filename = os.path.join(filepath, basename + '.gro') structure = self.to_parmed() if forcefield: structure = apply_forcefield(structure, forcefield=forcefield) structure.save(top_filename, 'gromacs', **kwargs) structure.save(gro_filename, 'gro', **kwargs) def save_lammpsdata(self, filename, traj, forcefield, force_overwrite=False, **kwargs): """ """ from foyer.forcefield import apply_forcefield import intermol.lammps as lmp # Create separate file paths for .gro and .top filepath, filename = os.path.split(filename) basename = os.path.splitext(filename)[0] inp_filename = os.path.join(filepath, basename + '.input') intermol_system = self.to_intermol() if forcefield: apply_forcefield(intermol_system, forcefield=forcefield) lmp.save(inp_filename, intermol_system) # Interface to Trajectory for reading/writing .pdb and .mol2 files. # ----------------------------------------------------------------- def from_trajectory(self, traj, frame=-1, coords_only=False): """Extract atoms and bonds from a md.Trajectory. Will create sub-compounds for every chain if there is more than one and sub-sub-compounds for every residue. Parameters ---------- traj : md.Trajectory The trajectory to load. frame : int The frame to take coordinates from. """ if coords_only: if traj.n_atoms != self.n_particles: raise ValueError('Number of atoms in {traj} does not match {self}'.format(**locals())) for mdtraj_atom, particle in zip(traj.topology.atoms, self._particles(include_ports=False)): particle.pos = traj.xyz[frame, mdtraj_atom.index] return atom_mapping = dict() for chain in traj.topology.chains: if traj.topology.n_chains > 1: chain_compound = Compound() self.add(chain_compound, 'chain[$]') else: chain_compound = self for res in chain.residues: for atom in res.atoms: new_atom = Particle(name=str(atom.name), pos=traj.xyz[frame, atom.index]) chain_compound.add(new_atom, label='{0}[$]'.format(atom.name)) atom_mapping[atom] = new_atom for mdtraj_atom1, mdtraj_atom2 in traj.topology.bonds: atom1 = atom_mapping[mdtraj_atom1] atom2 = atom_mapping[mdtraj_atom2] self.add_bond((atom1, atom2)) if np.any(traj.unitcell_lengths) and np.any(traj.unitcell_lengths[0]): self.periodicity = traj.unitcell_lengths[0] else: self.periodicity = np.array([0., 0., 0.]) def to_trajectory(self, show_ports=False, chain_types=None, residue_types=None, **kwargs): """Convert to an md.Trajectory and flatten the compound. Parameters ---------- show_ports : bool, optional, default=False Include all port atoms when converting to trajectory. Returns ------- trajectory : md.Trajectory See also -------- _to_topology """ atom_list = [particle for particle in self.particles(show_ports)] top = self._to_topology(atom_list, chain_types, residue_types) # Coordinates. xyz = np.ndarray(shape=(1, top.n_atoms, 3), dtype='float') for idx, atom in enumerate(atom_list): xyz[0, idx] = atom.pos # Unitcell information. box = self.boundingbox unitcell_lengths = np.empty(3) for dim, val in enumerate(self.periodicity): if val: unitcell_lengths[dim] = val else: unitcell_lengths[dim] = box.lengths[dim] return md.Trajectory(xyz, top, unitcell_lengths=unitcell_lengths, unitcell_angles=np.array([90, 90, 90])) def _to_topology(self, atom_list, chain_types=None, residue_types=None): """Create a mdtraj.Topology from a Compound. Parameters ---------- atom_list : chain_types : residue_types : Returns ------- top : mtraj.Topology """ if isinstance(chain_types, Compound): chain_types = [Compound] if isinstance(chain_types, (list, set)): chain_types = tuple(chain_types) if isinstance(residue_types, Compound): residue_types = [Compound] if isinstance(residue_types, (list, set)): residue_types = tuple(residue_types) top = Topology() atom_mapping = {} default_chain = top.add_chain() default_residue = top.add_residue('RES', default_chain) last_residue_compound = None last_chain_compound = None last_residue = None last_chain = None for atom in atom_list: # Chains for parent in atom.ancestors(): if chain_types and isinstance(parent, chain_types): if parent != last_chain_compound: last_chain_compound = parent last_chain = top.add_chain() last_chain_default_residue = top.add_residue('RES', last_chain) last_chain.compound = last_chain_compound break else: last_chain = default_chain last_chain.compound = last_chain_compound # Residues for parent in atom.ancestors(): if residue_types and isinstance(parent, residue_types): if parent != last_residue_compound: last_residue_compound = parent last_residue = top.add_residue(parent.__class__.__name__, last_chain) last_residue.compound = last_residue_compound break else: if last_chain != default_chain: last_residue = last_chain_default_residue else: last_residue = default_residue last_residue.compound = last_residue_compound # Add the actual atoms try: elem = get_by_symbol(atom.name) except KeyError: elem = get_by_symbol("VS") at = top.add_atom(atom.name, elem, last_residue) at.charge = atom.charge atom_mapping[atom] = at # Remove empty default residues. chains_to_remove = [chain for chain in top.chains if chain.n_atoms == 0] residues_to_remove = [res for res in top.residues if res.n_atoms == 0] for chain in chains_to_remove: top._chains.remove(chain) for res in residues_to_remove: for chain in top.chains: try: chain._residues.remove(res) except ValueError: # Already gone. pass for atom1, atom2 in self.bonds(): # Ensure that both atoms are part of the compound. This becomes an # issue if you try to convert a sub-compound to a topology which is # bonded to a different subcompound. if all(a in atom_mapping.keys() for a in [atom1, atom2]): top.add_bond(atom_mapping[atom1], atom_mapping[atom2]) return top def to_parmed(self, title=''): """Create a ParmEd Structure from a Compound. """ structure = pmd.Structure() structure.title = title if title else self.name atom_mapping = {} # For creating bonds below for atom in self.particles(): try: atomic_number = AtomicNum[atom.name] except KeyError: element = element_by_name(atom.name) warn('Guessing that {} is element: {}'.format(atom, element)) else: element = atom.name atomic_number = AtomicNum[element] mass = Mass[element] pmd_atom = pmd.Atom(atomic_number=atomic_number, name=atom.name, mass=mass) pmd_atom.xx, pmd_atom.xy, pmd_atom.xz = atom.pos * 10 # Angstroms structure.add_atom(pmd_atom, resname='RES', resnum=1) atom_mapping[atom] = pmd_atom for atom1, atom2 in self.bonds(): bond = pmd.Bond(atom_mapping[atom1], atom_mapping[atom2]) structure.bonds.append(bond) box = self.boundingbox box_vector = np.empty(6) box_vector[3] = box_vector[4] = box_vector[5] = 90.0 for dim, val in enumerate(self.periodicity): if val: box_vector[dim] = val * 10 else: box_vector[dim] = box.lengths[dim] * 10 + 5 structure.box = box_vector return structure def to_intermol(self, molecule_types=None): """Create an InterMol system from a Compound. Parameters ---------- molecule_types : list or tuple of subclasses of Compound Returns ------- intermol_system : intermol.system.System """ from intermol.atom import Atom as InterMolAtom from intermol.molecule import Molecule from intermol.system import System import simtk.unit as u if isinstance(molecule_types, list): molecule_types = tuple(molecule_types) elif molecule_types is None: molecule_types = (type(self),) intermol_system = System() last_molecule_compound = None for atom_index, atom in enumerate(self.particles()): for parent in atom.ancestors(): # Don't want inheritance via isinstance(). if type(parent) in molecule_types: # Check if we have encountered this molecule type before. if parent.name not in intermol_system.molecule_types: self._add_intermol_molecule_type(intermol_system, parent) if parent != last_molecule_compound: last_molecule_compound = parent last_molecule = Molecule(name=parent.name) intermol_system.add_molecule(last_molecule) break else: # Should never happen if molecule_types only contains type(self) raise ValueError('Found an atom {} that is not part of any of ' 'the specified molecule types {}'.format(atom, molecule_types)) # Add the actual intermol atoms. intermol_atom = InterMolAtom(atom_index + 1, name=atom.name, residue_index=1, residue_name='RES') intermol_atom.position = atom.pos * u.nanometers last_molecule.add_atom(intermol_atom) return intermol_system @staticmethod def _add_intermol_molecule_type(intermol_system, parent): """Create a molecule type for the parent and add bonds. """ from intermol.moleculetype import MoleculeType from intermol.forces.bond import Bond as InterMolBond molecule_type = MoleculeType(name=parent.name) intermol_system.add_molecule_type(molecule_type) for index, parent_atom in enumerate(parent.particles()): parent_atom.index = index + 1 for atom1, atom2 in parent.bonds(): intermol_bond = InterMolBond(atom1.index, atom2.index) molecule_type.bonds.add(intermol_bond) def __getitem__(self, selection): if isinstance(selection, integer_types): return list(self.particles())[selection] if isinstance(selection, string_types): return self.labels.get(selection) def __repr__(self): descr = list('<') descr.append(self.name + ' ') if self.children: descr.append('{:d} particles, '.format(self.n_particles)) if any(self.periodicity): descr.append('periodicity: {}, '.format(self.periodicity)) else: descr.append('non-periodic, ') else: descr.append('pos=({: .4f},{: .4f},{: .4f}), '.format(self.pos[0], self.pos[1], self.pos[2])) descr.append('{:d} bonds, '.format(self.n_bonds)) descr.append('id: {}>'.format(id(self))) return ''.join(descr) def _clone(self, clone_of=None, root_container=None): """A faster alternative to deepcopying. Does not resolve circular dependencies. This should be safe provided you never try to add the top of a Compound hierarchy to a sub-Compound. Clones compound hierarchy only, not the bonds. """ if root_container is None: root_container = self if clone_of is None: clone_of = dict() # If this compound has already been cloned, return that. if self in clone_of: return clone_of[self] # Otherwise we make a new clone. cls = self.__class__ newone = cls.__new__(cls) # Remember that we're cloning the new one of self. clone_of[self] = newone newone.name = deepcopy(self.name) newone.periodicity = deepcopy(self.periodicity) newone._pos = deepcopy(self._pos) newone.charge = deepcopy(self.charge) newone.port_particle = deepcopy(self.port_particle) if hasattr(self, 'index'): newone.index = deepcopy(self.index) if self.children is None: newone.children = None else: newone.children = OrderedSet() # Parent should be None initially. newone.parent = None newone.labels = OrderedDict() newone.referrers = set() newone.bond_graph = None # Add children to clone. if self.children: for child in self.children: newchild = child._clone(clone_of, root_container) newone.children.add(newchild) newchild.parent = newone # Copy labels, except bonds with atoms outside the hierarchy. if self.labels: for label, compound in self.labels.items(): if not isinstance(compound, list): newone.labels[label] = compound._clone(clone_of, root_container) compound.referrers.add(clone_of[compound]) else: # compound is a list of compounds, so we create an empty # list, and add the clones of the original list elements. newone.labels[label] = [] for subpart in compound: newone.labels[label].append(subpart._clone(clone_of, root_container)) # Referrers must have been handled already, or the will be handled return newone def _clone_bonds(self, clone_of=None): newone = clone_of[self] for c1, c2 in self.bonds(): newone.add_bond((clone_of[c1], clone_of[c2]))
def add(self, new_child, label=None, containment=True, replace=False, inherit_periodicity=True): """Add a part to the Compound. Note: This does not necessarily add the part to self.children but may instead be used to add a reference to the part to self.labels. See 'containment' argument. Parameters ---------- new_child : mb.Compound The object to be added to this Compound. label : str, optional A descriptive string for the part. containment : bool, optional, default=True Add the part to self.children. replace : bool, optional, default=True Replace the label if it already exists. """ # Support batch add via lists, tuples and sets. if not isinstance(new_child, Compound): for child in new_child: self.add(child) return # Create children and labels on the first add operation if self.children is None: self.children = OrderedSet() if self.labels is None: self.labels = OrderedDict() if not isinstance(new_child, Compound): raise ValueError('Only objects that inherit from mbuild.Compound ' 'can be added to Compounds. You tried to add ' '{}'.format(new_child)) if containment: if new_child.parent is not None: raise ValueError('Part {} already has a parent: {}'.format( new_child, new_child.parent)) self.children.add(new_child) new_child.parent = self if new_child.bond_graph is not None: if self.root.bond_graph is None: self.root.bond_graph = new_child.bond_graph else: self.root.bond_graph = nx.compose(self.root.bond_graph, new_child.bond_graph) new_child.bond_graph = None # Add new_part to labels. Does not currently support batch add. if label is None: label = '{0}[$]'.format(new_child.__class__.__name__) if label is not None: if label.endswith('[$]'): label = label[:-3] if label not in self.labels: self.labels[label] = [] label_pattern = label + '[{}]' count = len(self.labels[label]) self.labels[label].append(new_child) label = label_pattern.format(count) if not replace and label in self.labels: raise Exception('Label {0} already exists in {1}'.format(label, self)) else: self.labels[label] = new_child new_child.referrers.add(self) if (inherit_periodicity and isinstance(new_child, Compound) and new_child.periodicity.any()): self.periodicity = new_child.periodicity
def _init(self, maxsize): self.queue = OrderedSet()
class Compound(object): """A building block in the mBuild hierarchy. Compound is the superclass of all composite building blocks in the mBuild hierarchy. That is, all composite building blocks must inherit from compound, either directly or indirectly. The design of Compound follows the Composite design pattern (Gamma, Erich; Richard Helm; Ralph Johnson; John M. Vlissides (1995). Design Patterns: Elements of Reusable Object-Oriented Software. Addison-Wesley. p. 395. ISBN 0-201-63361-2.), with Compound being the composite, and Particle playing the role of the primitive (leaf) part, where Particle is in fact simply an alias to the Compound class. Compound maintains a list of children (other Compounds contained within), and provides a means to tag the children with labels, so that the compounds can be easily looked up later. Labels may also point to objects outside the Compound's containment hierarchy. Compound has built-in support for copying and deepcopying Compound hierarchies, enumerating particles or bonds in the hierarchy, proximity based searches, visualization, I/O operations, and a number of other convenience methods. Parameters ---------- subcompounds : Compound, optional, default=None One or more compounds to be added to self. name : str, optional, default=self.__class__.__name__ The type of Compound. periodicity : np.ndarray, shape=(3,), dtype=float, optional The periodic lengths of the Compound in the x, y and z directions. Defaults to zeros which is treated as non-periodic. Attributes ---------- name : str, optional, default=self.__class__.__name__ The type of Compound. periodicity : np.ndarray, shape=(3,), dtype=float, optional The periodic lengths of the Compound in the x, y and z directions. Defaults to zeros which is treated as non-periodic. children : OrderedSet Contains all children (other Compounds). labels : OrderedDict Labels to Compound/Atom mappings. These do not necessarily need not be in self.children. parent : mb.Compound The parent Compound that contains this part. Can be None if this compound is the root of the containment hierarchy. referrers : set Other compounds that reference this part with labels. """ def __init__(self, subcompounds=None, name=None, pos=None, charge=0.0, periodicity=None, port_particle=False): super(Compound, self).__init__() if name: if not isinstance(name, string_types): raise ValueError( 'Compound.name should be a string. You passed ' '{}'.format(name)) self.name = name else: self.name = self.__class__.__name__ # A periodocity of zero in any direction is treated as non-periodic. if periodicity is None: self._periodicity = np.array([0.0, 0.0, 0.0]) else: self._periodicity = np.asarray(periodicity) if pos is not None: self._pos = np.asarray(pos, dtype=float) else: self._pos = np.zeros(3) self.charge = charge self.parent = None self.children = OrderedSet() self.labels = OrderedDict() self.referrers = set() self.bond_graph = None self.port_particle = port_particle # self.add() must be called after labels and children are initialized. if subcompounds: self.add(subcompounds) def particles(self, include_ports=False): """ """ if not self.children: yield self else: for particle in self._particles(include_ports): yield particle def _particles(self, include_ports=False): """Return all Particles of the Compound. """ for child in self.successors(): if not child.children: if include_ports or not child.port_particle: yield child def successors(self): """Yield Compounds below self in the hierarchy. """ if not self.children: return for part in self.children: # Parts local to the current Compound. yield part # Parts further down the hierarchy. for subpart in part.successors(): yield subpart @property def n_particles(self): if not self.children: return 1 else: return self._n_particles(include_ports=False) def _n_particles(self, include_ports=False): """Return the number of Particles in the Compound. """ return sum(1 for _ in self._particles(include_ports)) def _contains_only_ports(self): for part in self.children: if not part.port_particle: return False return True def ancestors(self): """Generate all ancestors of the Compound recursively. """ if self.parent is not None: yield self.parent for ancestor in self.parent.ancestors(): yield ancestor @property def root(self): parent = None for parent in self.ancestors(): pass if parent is None: return self return parent def particles_by_name(self, name): for particle in self.particles(): if particle.name == name: yield particle def add(self, new_child, label=None, containment=True, replace=False, inherit_periodicity=True): """Add a part to the Compound. Note: This does not necessarily add the part to self.children but may instead be used to add a reference to the part to self.labels. See 'containment' argument. Parameters ---------- new_child : mb.Compound or list-like of mb.Compound The object(s) to be added to this Compound. label : str, optional A descriptive string for the part. containment : bool, optional, default=True Add the part to self.children. replace : bool, optional, default=True Replace the label if it already exists. """ # Support batch add via lists, tuples and sets. if (isinstance(new_child, collections.Iterable) and not isinstance(new_child, string_types)): for child in new_child: self.add(child) return if not isinstance(new_child, Compound): raise ValueError('Only objects that inherit from mbuild.Compound ' 'can be added to Compounds. You tried to add ' '"{}".'.format(new_child)) # Create children and labels on the first add operation if self.children is None: self.children = OrderedSet() if self.labels is None: self.labels = OrderedDict() if containment: if new_child.parent is not None: raise MBuildError('Part {} already has a parent: {}'.format( new_child, new_child.parent)) self.children.add(new_child) new_child.parent = self if new_child.bond_graph is not None: if self.root.bond_graph is None: self.root.bond_graph = new_child.bond_graph else: self.root.bond_graph.compose(new_child.bond_graph) new_child.bond_graph = None # Add new_part to labels. Does not currently support batch add. if label is None: label = '{0}[$]'.format(new_child.__class__.__name__) if label.endswith('[$]'): label = label[:-3] if label not in self.labels: self.labels[label] = [] label_pattern = label + '[{}]' count = len(self.labels[label]) self.labels[label].append(new_child) label = label_pattern.format(count) if not replace and label in self.labels: raise MBuildError('Label "{0}" already exists in {1}.'.format( label, self)) else: self.labels[label] = new_child new_child.referrers.add(self) if (inherit_periodicity and isinstance(new_child, Compound) and new_child.periodicity.any()): self.periodicity = new_child.periodicity def remove(self, objs_to_remove): """Remove children from the Compound. """ if not self.children: return if not hasattr(objs_to_remove, '__iter__'): objs_to_remove = [objs_to_remove] objs_to_remove = set(objs_to_remove) if len(objs_to_remove) == 0: return intersection = objs_to_remove.intersection(self.children) self.children -= intersection objs_to_remove -= intersection for removed_part in intersection: if self.root.bond_graph and self.root.bond_graph.has_node( removed_part): self.root.bond_graph.remove_node(removed_part) self._remove_references(removed_part) # Remove the part recursively from sub-compounds. if self.children: for part in self.children: part.remove(objs_to_remove) @staticmethod def _remove_references(removed_part): """Remove labels pointing to this part and vice versa. """ removed_part.parent = None # Remove labels in the hierarchy pointing to this part. referrers_to_remove = set() for referrer in removed_part.referrers: if removed_part not in referrer.ancestors(): for label, referred_part in list(referrer.labels.items()): if referred_part is removed_part: del referrer.labels[label] referrers_to_remove.add(referrer) removed_part.referrers -= referrers_to_remove # Remove labels in this part pointing into the hierarchy. labels_to_delete = [] if isinstance(removed_part, Compound): for label, part in removed_part.labels.items(): if removed_part not in part.ancestors(): part.referrers.remove(removed_part) labels_to_delete.append(label) for label in labels_to_delete: del removed_part.labels[label] def referenced_ports(self): """Return all Ports referenced by this Compound. """ from mbuild.port import Port return [ port for port in self.labels.values() if isinstance(port, Port) ] def available_ports(self): """Return all unoccupied Ports referenced by this Compound. """ from mbuild.port import Port return [ port for port in self.labels.values() if isinstance(port, Port) and not port.used ] def bonds(self): """A list of all Bonds in the Compound and sub-Compounds. """ if self.root.bond_graph: if self.root == self: return self.root.bond_graph.edges_iter() else: return self.root.bond_graph.subgraph( self.particles()).edges_iter() else: return iter(()) @property def n_bonds(self): """Return the number of Bonds in the Compound. """ return sum(1 for _ in self.bonds()) def add_bond(self, particle_pair): """""" if self.root.bond_graph is None: self.root.bond_graph = BondGraph() self.root.bond_graph.add_edge(particle_pair[0], particle_pair[1]) def generate_bonds(self, name_a, name_b, dmin, dmax): """Add Bonds between all pairs of types a/b within [dmin, dmax]. """ particle_kdtree = PeriodicCKDTree(data=self.xyz, bounds=self.periodicity) particle_array = np.array(list(self.particles())) added_bonds = list() for p1 in self.particles_by_name(name_a): nearest = self.particles_in_range(p1, dmax, max_particles=20, particle_kdtree=particle_kdtree, particle_array=particle_array) for p2 in nearest: if p2 == p1: continue bond_tuple = (p1, p2) if id(p1) < id(p2) else (p2, p1) if bond_tuple in added_bonds: continue if (p2.name == name_b) and (dmin <= self.min_periodic_distance( p2.pos, p1.pos) <= dmax): self.add_bond((p1, p2)) added_bonds.append(bond_tuple) def remove_bond(self, particle_pair): if self.root.bond_graph is None or not self.root.bond_graph.has_edge( *particle_pair): warn( "Bond between {} and {} doesn't exist!".format(*particle_pair)) return self.root.bond_graph.remove_edge(*particle_pair) @property def pos(self): if not self.children: return self._pos else: return self.center @pos.setter def pos(self, value): if not self.children: self._pos = value else: raise MBuildError('Cannot set position on a Compound that has' ' children.') @property def periodicity(self): return self._periodicity @periodicity.setter def periodicity(self, periods): self._periodicity = np.array(periods) @property def xyz(self): """Return all particle coordinates in this compound. Returns ------- pos : np.ndarray, shape=(n, 3) Array with the positions of all particles. """ if not self.children: pos = np.expand_dims(self._pos, axis=0) else: arr = np.fromiter(itertools.chain.from_iterable( particle.pos for particle in self.particles()), dtype=float) pos = arr.reshape((-1, 3)) return pos @property def xyz_with_ports(self): """Return all particle coordinates in this compound including ports. """ if not self.children: pos = self._pos else: arr = np.fromiter(itertools.chain.from_iterable( particle.pos for particle in self.particles(include_ports=True)), dtype=float) pos = arr.reshape((-1, 3)) return pos @property def center(self): """The cartesian center of the Compound based on its Atoms. """ if self.xyz.any(): return np.mean(self.xyz, axis=0) @property def boundingbox(self): """Compute the bounding box of the compound. """ xyz = self.xyz return Box(mins=xyz.min(axis=0), maxs=xyz.max(axis=0)) def min_periodic_distance(self, xyz0, xyz1): """Vectorized distance calculation considering minimum image. """ d = np.abs(xyz0 - xyz1) d = np.where(d > 0.5 * self.periodicity, self.periodicity - d, d) return np.sqrt((d**2).sum(axis=-1)) def particles_in_range(self, compound, dmax, max_particles=20, particle_kdtree=None, particle_array=None): """Find particles within a specified range of another particle. """ if particle_kdtree is None: particle_kdtree = PeriodicCKDTree(data=self.xyz, bounds=self.periodicity) _, idxs = particle_kdtree.query(compound.pos, k=max_particles, distance_upper_bound=dmax) idxs = idxs[idxs != self.n_particles] if particle_array is None: particle_array = np.array(list(self.particles())) return particle_array[idxs] def view_hierarchy(self, show_ports=False): """Visualize a compound hierarchy as a tree. """ raise NotImplementedError('Coming soon!') def visualize(self, show_ports=False): """Visualize the Compound using nglview. """ nglview = import_('nglview') if run_from_ipython(): structure = self.to_trajectory(show_ports) return nglview.show_mdtraj(structure) else: try: """Visualize the Compound using imolecule. """ import imolecule json_mol = self._to_json(show_ports) imolecule.draw(json_mol, format='json', shader='lambert', drawing_type='ball and stick', camera_type='perspective', element_properties=None) except ImportError: raise RuntimeError( 'Visualization is only supported in Jupyter ' 'Notebooks.') def _to_json(self, show_ports=False): import imolecule atoms = list() for idx, particle in enumerate( self._particles(include_ports=show_ports)): particle.index = idx atoms.append({ 'element': particle.name, 'location': list(np.asarray(particle.pos, dtype=float) * 10) }) bonds = [{ 'atoms': [atom1.index, atom2.index], 'order': 1 } for atom1, atom2 in self.bonds()] output = {'name': self.name, 'atoms': atoms, 'bonds': bonds} # Remove the index attribute on particles. for idx, particle in enumerate(self.particles()): if not show_ports and particle.port_particle: continue del particle.index return imolecule.json_formatter.compress(output) def update_coordinates(self, filename): """Update the coordinates of this Compound from a file. """ load(filename, compound=self, coords_only=True) def save(self, filename, show_ports=False, forcefield_name=None, forcefield_files=None, box=None, **kwargs): """Save the Compound to a file. Parameters ---------- filename : str Filesystem path in which to save the trajectory. The extension or prefix will be parsed and will control the format. show_ports : bool, default=False Save ports contained within the compound. forcefield_name : str, default=None Apply a forcefield to the output file using a forcefield provided by the `foyer` package. forcefield_name : str, default=None Apply a forcefield to the output file using the `foyer` package and a specific forcefield.xml file. Other Parameters ---------------- force_overwrite : bool """ extension = os.path.splitext(filename)[-1] savers = { '.hoomdxml': self.save_hoomdxml, '.gsd': self.save_gsd, '.gro': self.save_gromacs, '.top': self.save_gromacs, '.lammps': self.save_lammpsdata, '.lmp': self.save_lammpsdata } try: saver = savers[extension] except KeyError: # TODO: better reporting saver = None structure = self.to_parmed(**kwargs) if saver: # mBuild/InterMol supported saver. saver(filename, structure, forcefield_name, forcefield_files, box, **kwargs) elif extension == '.xyz': traj = self.to_trajectory(show_ports=show_ports) traj.save(filename) else: # ParmEd supported saver. return structure.save(filename, **kwargs) def _apply_forcefield(self, structure, forcefield_files, forcefield_name): from foyer import Forcefield ff = Forcefield(forcefield_files=forcefield_files, name=forcefield_name) structure = ff.apply(structure) return structure def _gen_box(self): box = self.boundingbox for dim, val in enumerate(self.periodicity): if val: box.lengths[dim] = val box.maxs[dim] = val box.mins[dim] = 0.0 if not val: box.maxs[dim] += 0.25 box.mins[dim] -= 0.25 box.lengths[dim] += 0.5 return box def save_hoomdxml(self, filename, structure, forcefield_name, forcefield_files, box, **kwargs): """ """ forcefield = False if forcefield_name or forcefield_files: forcefield = True structure = self._apply_forcefield(structure, forcefield_files, forcefield_name) if box is None: box = self._gen_box() write_hoomdxml(structure, filename, forcefield, box, **kwargs) def save_gsd(self, filename, structure, forcefield_name, forcefield_files, box=None, **kwargs): """ """ from mbuild.formats.gsdwriter import write_gsd forcefield = False if forcefield_name or forcefield_files: forcefield = True structure = self._apply_forcefield(structure, forcefield_files, forcefield_name) if box is None: box = self._gen_box() write_gsd(structure, filename, forcefield, box, **kwargs) def save_gromacs(self, filename, structure, forcefield_name, forcefield_files, box, **kwargs): """ """ # Create separate file paths for .gro and .top filepath, filename = os.path.split(filename) basename = os.path.splitext(filename)[0] top_filename = os.path.join(filepath, basename + '.top') gro_filename = os.path.join(filepath, basename + '.gro') forcefield = False if forcefield_name or forcefield_files: forcefield = True structure = self._apply_forcefield(structure, forcefield_files, forcefield_name) if box is None: box = self._gen_box() structure.save(top_filename, 'gromacs', **kwargs) structure.save(gro_filename, 'gro', **kwargs) def save_lammpsdata(self, filename, structure, forcefield_name, forcefield_files, box, **kwargs): """ """ forcefield = False if forcefield_name or forcefield_files: forcefield = True structure = self._apply_forcefield(structure, forcefield_files, forcefield_name) if box is None: box = self._gen_box() write_lammpsdata(structure, filename, forcefield, box, **kwargs) # Interface to Trajectory for reading/writing .pdb and .mol2 files. # ----------------------------------------------------------------- def from_trajectory(self, traj, frame=-1, coords_only=False): """Extract atoms and bonds from a md.Trajectory. Will create sub-compounds for every chain if there is more than one and sub-sub-compounds for every residue. Parameters ---------- traj : md.Trajectory The trajectory to load. frame : int The frame to take coordinates from. """ if coords_only: if traj.n_atoms != self.n_particles: raise ValueError('Number of atoms in {traj} does not match' ' {self}'.format(**locals())) atoms_particles = zip(traj.topology.atoms, self._particles(include_ports=False)) for mdtraj_atom, particle in atoms_particles: particle.pos = traj.xyz[frame, mdtraj_atom.index] return atom_mapping = dict() for chain in traj.topology.chains: if traj.topology.n_chains > 1: chain_compound = Compound() self.add(chain_compound, 'chain[$]') else: chain_compound = self for res in chain.residues: for atom in res.atoms: new_atom = Particle(name=str(atom.name), pos=traj.xyz[frame, atom.index]) chain_compound.add(new_atom, label='{0}[$]'.format(atom.name)) atom_mapping[atom] = new_atom for mdtraj_atom1, mdtraj_atom2 in traj.topology.bonds: atom1 = atom_mapping[mdtraj_atom1] atom2 = atom_mapping[mdtraj_atom2] self.add_bond((atom1, atom2)) if np.any(traj.unitcell_lengths) and np.any(traj.unitcell_lengths[0]): self.periodicity = traj.unitcell_lengths[0] else: self.periodicity = np.array([0., 0., 0.]) def to_trajectory(self, show_ports=False, chain_types=None, residue_types=None, **kwargs): """Convert to an md.Trajectory and flatten the compound. Parameters ---------- show_ports : bool, optional, default=False Include all port atoms when converting to trajectory. Returns ------- trajectory : md.Trajectory See also -------- _to_topology """ import mdtraj as md atom_list = [particle for particle in self.particles(show_ports)] top = self._to_topology(atom_list, chain_types, residue_types) # Coordinates. xyz = np.ndarray(shape=(1, top.n_atoms, 3), dtype='float') for idx, atom in enumerate(atom_list): xyz[0, idx] = atom.pos # Unitcell information. box = self.boundingbox unitcell_lengths = np.empty(3) for dim, val in enumerate(self.periodicity): if val: unitcell_lengths[dim] = val else: unitcell_lengths[dim] = box.lengths[dim] return md.Trajectory(xyz, top, unitcell_lengths=unitcell_lengths, unitcell_angles=np.array([90, 90, 90])) def _to_topology(self, atom_list, chain_types=None, residue_types=None): """Create a mdtraj.Topology from a Compound. Parameters ---------- atom_list : chain_types : residue_types : Returns ------- top : mtraj.Topology """ from mdtraj.core.element import get_by_symbol from mdtraj.core.topology import Topology if isinstance(chain_types, Compound): chain_types = [Compound] if isinstance(chain_types, (list, set)): chain_types = tuple(chain_types) if isinstance(residue_types, Compound): residue_types = [Compound] if isinstance(residue_types, (list, set)): residue_types = tuple(residue_types) top = Topology() atom_mapping = {} default_chain = top.add_chain() default_residue = top.add_residue('RES', default_chain) last_residue_compound = None last_chain_compound = None last_residue = None last_chain = None for atom in atom_list: # Chains for parent in atom.ancestors(): if chain_types and isinstance(parent, chain_types): if parent != last_chain_compound: last_chain_compound = parent last_chain = top.add_chain() last_chain_default_residue = top.add_residue( 'RES', last_chain) last_chain.compound = last_chain_compound break else: last_chain = default_chain last_chain.compound = last_chain_compound # Residues for parent in atom.ancestors(): if residue_types and isinstance(parent, residue_types): if parent != last_residue_compound: last_residue_compound = parent last_residue = top.add_residue( parent.__class__.__name__, last_chain) last_residue.compound = last_residue_compound break else: if last_chain != default_chain: last_residue = last_chain_default_residue else: last_residue = default_residue last_residue.compound = last_residue_compound # Add the actual atoms try: elem = get_by_symbol(atom.name) except KeyError: elem = get_by_symbol("VS") at = top.add_atom(atom.name, elem, last_residue) at.charge = atom.charge atom_mapping[atom] = at # Remove empty default residues. chains_to_remove = [ chain for chain in top.chains if chain.n_atoms == 0 ] residues_to_remove = [res for res in top.residues if res.n_atoms == 0] for chain in chains_to_remove: top._chains.remove(chain) for res in residues_to_remove: for chain in top.chains: try: chain._residues.remove(res) except ValueError: # Already gone. pass for atom1, atom2 in self.bonds(): # Ensure that both atoms are part of the compound. This becomes an # issue if you try to convert a sub-compound to a topology which is # bonded to a different subcompound. if all(a in atom_mapping.keys() for a in [atom1, atom2]): top.add_bond(atom_mapping[atom1], atom_mapping[atom2]) return top def from_parmed(self, structure, coords_only=False): """Extract atoms and bonds from a pmd.Structure. Will create sub-compounds for every chain if there is more than one and sub-sub-compounds for every residue. Parameters ---------- structure : pmd.Structure The structure to load. coords_only : bool Set preexisting atoms in compound to coordinates given by structure. """ if coords_only: if len(structure.atoms) != self.n_particles: raise ValueError( 'Number of atoms in {structure} does not match' ' {self}'.format(**locals())) atoms_particles = zip(structure.atoms, self._particles(include_ports=False)) for parmed_atom, particle in atoms_particles: particle.pos = structure.coordinates[parmed_atom.idx] return atom_mapping = dict() chain_id = None chains = defaultdict(list) for residue in structure.residues: chains[residue.chain].append(residue) for chain, residues in chains.items(): if len(chains) > 1: chain_compound = Compound() self.add(chain_compound, chain_id) else: chain_compound = self for residue in residues: for atom in residue.atoms: new_atom = Particle(name=str(atom.name), pos=structure.coordinates[atom.idx]) chain_compound.add(new_atom, label='{0}[$]'.format(atom.name)) atom_mapping[atom] = new_atom for bond in structure.bonds: atom1 = atom_mapping[bond.atom1] atom2 = atom_mapping[bond.atom2] self.add_bond((atom1, atom2)) if structure.box is not None: self.periodicity = structure.box[0:3] else: self.periodicity = np.array([0., 0., 0.]) def to_parmed(self, title='', **kwargs): """Create a ParmEd Structure from a Compound. """ structure = pmd.Structure() structure.title = title if title else self.name atom_mapping = {} # For creating bonds below guessed_elements = set() for atom in self.particles(): atomic_number = None name = ''.join(char for char in atom.name if not char.isdigit()) try: atomic_number = AtomicNum[atom.name] except KeyError: element = element_by_name(atom.name) if name not in guessed_elements: warn('Guessing that "{}" is element: "{}"'.format( atom, element)) guessed_elements.add(name) else: element = atom.name atomic_number = atomic_number or AtomicNum[element] mass = Mass[element] pmd_atom = pmd.Atom(atomic_number=atomic_number, name=atom.name, mass=mass) pmd_atom.xx, pmd_atom.xy, pmd_atom.xz = atom.pos * 10 # Angstroms structure.add_atom(pmd_atom, resname='RES', resnum=1) atom_mapping[atom] = pmd_atom for atom1, atom2 in self.bonds(): bond = pmd.Bond(atom_mapping[atom1], atom_mapping[atom2]) structure.bonds.append(bond) box = self.boundingbox box_vector = np.empty(6) box_vector[3] = box_vector[4] = box_vector[5] = 90.0 for dim, val in enumerate(self.periodicity): if val: box_vector[dim] = val * 10 else: box_vector[dim] = box.lengths[dim] * 10 + 5 structure.box = box_vector return structure def to_intermol(self, molecule_types=None): """Create an InterMol system from a Compound. Parameters ---------- molecule_types : list or tuple of subclasses of Compound Returns ------- intermol_system : intermol.system.System """ from intermol.atom import Atom as InterMolAtom from intermol.molecule import Molecule from intermol.system import System import simtk.unit as u if isinstance(molecule_types, list): molecule_types = tuple(molecule_types) elif molecule_types is None: molecule_types = (type(self), ) intermol_system = System() last_molecule_compound = None for atom_index, atom in enumerate(self.particles()): for parent in atom.ancestors(): # Don't want inheritance via isinstance(). if type(parent) in molecule_types: # Check if we have encountered this molecule type before. if parent.name not in intermol_system.molecule_types: self._add_intermol_molecule_type( intermol_system, parent) if parent != last_molecule_compound: last_molecule_compound = parent last_molecule = Molecule(name=parent.name) intermol_system.add_molecule(last_molecule) break else: # Should never happen if molecule_types only contains type(self) raise ValueError('Found an atom {} that is not part of any of ' 'the specified molecule types {}'.format( atom, molecule_types)) # Add the actual intermol atoms. intermol_atom = InterMolAtom(atom_index + 1, name=atom.name, residue_index=1, residue_name='RES') intermol_atom.position = atom.pos * u.nanometers last_molecule.add_atom(intermol_atom) return intermol_system @staticmethod def _add_intermol_molecule_type(intermol_system, parent): """Create a molecule type for the parent and add bonds. """ from intermol.moleculetype import MoleculeType from intermol.forces.bond import Bond as InterMolBond molecule_type = MoleculeType(name=parent.name) intermol_system.add_molecule_type(molecule_type) for index, parent_atom in enumerate(parent.particles()): parent_atom.index = index + 1 for atom1, atom2 in parent.bonds(): intermol_bond = InterMolBond(atom1.index, atom2.index) molecule_type.bonds.add(intermol_bond) def __getitem__(self, selection): if isinstance(selection, integer_types): return list(self.particles())[selection] if isinstance(selection, string_types): return self.labels.get(selection) def __repr__(self): descr = list('<') descr.append(self.name + ' ') if self.children: descr.append('{:d} particles, '.format(self.n_particles)) if any(self.periodicity): descr.append('periodicity: {}, '.format(self.periodicity)) else: descr.append('non-periodic, ') else: descr.append('pos=({: .4f},{: .4f},{: .4f}), '.format( self.pos[0], self.pos[1], self.pos[2])) descr.append('{:d} bonds, '.format(self.n_bonds)) descr.append('id: {}>'.format(id(self))) return ''.join(descr) def _clone(self, clone_of=None, root_container=None): """A faster alternative to deepcopying. Does not resolve circular dependencies. This should be safe provided you never try to add the top of a Compound hierarchy to a sub-Compound. Clones compound hierarchy only, not the bonds. """ if root_container is None: root_container = self if clone_of is None: clone_of = dict() # If this compound has already been cloned, return that. if self in clone_of: return clone_of[self] # Otherwise we make a new clone. cls = self.__class__ newone = cls.__new__(cls) # Remember that we're cloning the new one of self. clone_of[self] = newone newone.name = deepcopy(self.name) newone.periodicity = deepcopy(self.periodicity) newone._pos = deepcopy(self._pos) newone.charge = deepcopy(self.charge) newone.port_particle = deepcopy(self.port_particle) if hasattr(self, 'index'): newone.index = deepcopy(self.index) if self.children is None: newone.children = None else: newone.children = OrderedSet() # Parent should be None initially. newone.parent = None newone.labels = OrderedDict() newone.referrers = set() newone.bond_graph = None # Add children to clone. if self.children: for child in self.children: newchild = child._clone(clone_of, root_container) newone.children.add(newchild) newchild.parent = newone # Copy labels, except bonds with atoms outside the hierarchy. if self.labels: for label, compound in self.labels.items(): if not isinstance(compound, list): newone.labels[label] = compound._clone( clone_of, root_container) compound.referrers.add(clone_of[compound]) else: # compound is a list of compounds, so we create an empty # list, and add the clones of the original list elements. newone.labels[label] = [] for subpart in compound: newone.labels[label].append( subpart._clone(clone_of, root_container)) # Referrers must have been handled already, or the will be handled return newone def _clone_bonds(self, clone_of=None): newone = clone_of[self] for c1, c2 in self.bonds(): newone.add_bond((clone_of[c1], clone_of[c2]))
def add(self, new_child, label=None, containment=True, replace=False, inherit_periodicity=True): """Add a part to the Compound. Note: This does not necessarily add the part to self.children but may instead be used to add a reference to the part to self.labels. See 'containment' argument. Parameters ---------- new_child : mb.Compound or list-like of mb.Compound The object(s) to be added to this Compound. label : str, optional A descriptive string for the part. containment : bool, optional, default=True Add the part to self.children. replace : bool, optional, default=True Replace the label if it already exists. """ # Support batch add via lists, tuples and sets. if (isinstance(new_child, collections.Iterable) and not isinstance(new_child, string_types)): for child in new_child: self.add(child) return if not isinstance(new_child, Compound): raise ValueError('Only objects that inherit from mbuild.Compound ' 'can be added to Compounds. You tried to add ' '"{}".'.format(new_child)) # Create children and labels on the first add operation if self.children is None: self.children = OrderedSet() if self.labels is None: self.labels = OrderedDict() if containment: if new_child.parent is not None: raise MBuildError('Part {} already has a parent: {}'.format( new_child, new_child.parent)) self.children.add(new_child) new_child.parent = self if new_child.bond_graph is not None: if self.root.bond_graph is None: self.root.bond_graph = new_child.bond_graph else: self.root.bond_graph.compose(new_child.bond_graph) new_child.bond_graph = None # Add new_part to labels. Does not currently support batch add. if label is None: label = '{0}[$]'.format(new_child.__class__.__name__) if label.endswith('[$]'): label = label[:-3] if label not in self.labels: self.labels[label] = [] label_pattern = label + '[{}]' count = len(self.labels[label]) self.labels[label].append(new_child) label = label_pattern.format(count) if not replace and label in self.labels: raise MBuildError('Label "{0}" already exists in {1}.'.format( label, self)) else: self.labels[label] = new_child new_child.referrers.add(self) if (inherit_periodicity and isinstance(new_child, Compound) and new_child.periodicity.any()): self.periodicity = new_child.periodicity