<commit_before>from __future__ import print_function, division

__all__ = ['load', 'clone', 'Compound', 'Particle']

import collections
from collections import OrderedDict, defaultdict
from copy import deepcopy
import itertools
import os
import sys
import tempfile
from warnings import warn

import mdtraj as md
import numpy as np
from oset import oset as OrderedSet
import parmed as pmd
from parmed.periodic_table import AtomicNum, element_by_name, Mass
from six import integer_types, string_types

from mbuild.bond_graph import BondGraph
from mbuild.box import Box
from mbuild.coordinate_transform import translate
from mbuild.exceptions import MBuildError
from mbuild.formats.hoomdxml import write_hoomdxml
from mbuild.formats.lammpsdata import write_lammpsdata
from mbuild.formats.gsdwriter import write_gsd
from mbuild.periodic_kdtree import PeriodicCKDTree
from mbuild.utils.io import run_from_ipython, import_
from mbuild.coordinate_transform import _translate, _rotate

def load(filename, relative_to_module=None, compound=None, coords_only=False,
         rigid=False, **kwargs):
    """Load a file into an mbuild compound.

    Files are read using the MDTraj package. Please refer to http://mdtraj.org/
    1.8.0/load_functions.html for supported formats.

    Parameters
    ----------
    filename : str
        Name of the file from which to load atom and bond information.
    relative_to_module : str, optional, default=None
        Instead of looking in the current working directory, look for the file
        where this module is defined. This is typically used in Compound classes
        that will be instantiated from a different directory (such as the
        Compounds located in mbuild.lib).
    compound : mb.Compound, optional, default=None
        Existing compound to load atom and bond information into.
    coords_only : bool, optional, default=False
        Only load the coordinates into an existing compoint.
    rigid : bool, optional
        Treat the compound as a rigid body

    Returns
    -------
    compound : mb.Compound

    """
    # Handle mbuild *.py files containing a class that wraps a structure file
    # in its own folder. E.g., you build a system from ~/foo.py and it imports
    # from ~/bar/baz.py where baz.py loads ~/bar/baz.pdb.
    if relative_to_module:
        script_path = os.path.realpath(sys.modules[relative_to_module].__file__)
        file_dir = os.path.dirname(script_path)
        filename = os.path.join(file_dir, filename)

    if compound is None:
        compound = Compound()

    traj = md.load(filename, **kwargs)
    compound.from_trajectory(traj, frame=-1, coords_only=coords_only)
    if rigid:
        compound.label_rigid_bodies()
    return compound


def clone(existing_compound, 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.

    Parameters
    ----------
    existing_compound : mb.Compound
        Existing Compound that will be copied

    Other Parameters
    ----------------
    clone_of : dict, optional
    root_container : mb.Compound, optional

    """
    if clone_of is None:
        clone_of = dict()

    newone = existing_compound._clone(clone_of=clone_of,
                                      root_container=root_container)
    existing_compound._clone_bonds(clone_of=clone_of)
    return newone


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; <NAME>; <NAME>; <NAME> (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 : mb.Compound or list of mb.Compound, optional, default=None
        One or more compounds to be added to self.
    name : str, optional, default=self.__class__.__name__
        The type of Compound.
    pos : np.ndarray, shape=(3,), dtype=float, optional, default=[0, 0, 0]
        The position of the Compound in Cartestian space
    charge : float, optional, default=0.0
        Currently not used. Likely removed in next release.
    periodicity : np.ndarray, shape=(3,), dtype=float, optional, default=[0, 0, 0]
        The periodic lengths of the Compound in the x, y and z directions.
        Defaults to zeros which is treated as non-periodic.
    port_particle : bool, optional, default=False
        Whether or not this Compound is part of a Port

    Attributes
    ----------
    bond_graph : mb.BondGraph
        Graph-like object that stores bond information for this Compound
    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.
    rigid_id : int, default=None
        The ID of the rigid body that this Compound belongs to.  Only Particles
        (the bottom of the containment hierarchy) can have integer values for
        `rigid_id`. Compounds containing rigid particles will always have 
        `rigid_id == None`. See also `contains_rigid`.
    boundingbox
    center
    contains_rigid
    max_rigid_id
    n_particles
    n_bonds
    root
    xyz
    xyz_with_ports

    """

    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 periodicity 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._rigid_id = None
        self._contains_rigid = False
        self._check_if_contains_rigid_bodies = False

        # self.add() must be called after labels and children are initialized.
        if subcompounds:
            self.add(subcompounds)

    def particles(self, include_ports=False):
        """Return all Particles of the Compound.
        
        Parameters
        ----------
        include_ports : bool, optional, default=False
            Include port particles

        Yields
        -------
        mb.Compound
            The next Particle in the Compound

        """
        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.

        Yields
        -------
        mb.Compound
            The next Particle 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):
        """Return the number of Particles in the Compound.

        Returns
        -------
        int
            The number of Particles in the Compound

        """
        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.

        Yields
        ------
        mb.Compound
            The next Compound above self in the hierarchy

        """
        if self.parent is not None:
            yield self.parent
            for ancestor in self.parent.ancestors():
                yield ancestor

    @property
    def root(self):
        """The Compound at the top of self's hierarchy.

        Returns
        -------
        mb.Compound
            The Compound at the top of self's hierarchy

        """
        parent = None
        for parent in self.ancestors():
            pass
        if parent is None:
            return self
        return parent

    def particles_by_name(self, name):
        """Return all Particles of the Compound with a specific name

        Parameters
        ----------
        name : str
            Only particles with this name are returned

        Yields
        ------
        mb.Compound
            The next Particle in the Compound with the user-specified name

        """
        for particle in self.particles():
            if particle.name == name:
                yield particle

    @property
    def rigid_id(self):
        return self._rigid_id

    @rigid_id.setter
    def rigid_id(self, value):
        if self._contains_only_ports():
            self._rigid_id = value
            for ancestor in self.ancestors():
                ancestor._check_if_contains_rigid_bodies = True
        else:
            raise AttributeError("rigid_id is immutable for Compounds that are "
                                 "not at the bottom of the containment hierarchy.")

    @property
    def contains_rigid(self):
        """Returns True if the Compound contains rigid bodies

        If the Compound contains any particle with a rigid_id != None 
        then contains_rigid will return True. If the Compound has no
        children (i.e. the Compound resides at the bottom of the containment
        hierarchy) then contains_rigid will return False.

        Returns
        -------
        bool
            True if the Compound contains any particle with a rigid_id != None

        Notes
        -----
        The private variable '_check_if_contains_rigid_bodies' is used to help
        cache the status of 'contains_rigid'. If '_check_if_contains_rigid_bodies'
        is False, then the rigid body containment of the Compound has not changed,
        and the particle tree is not traversed, boosting performance.

        """
        if self._check_if_contains_rigid_bodies:
            self._check_if_contains_rigid_bodies = False
            if any(particle.rigid_id is not None for particle in self._particles()):
                self._contains_rigid = True
            else:
                self._contains_rigid = False
        return self._contains_rigid

    @property
    def max_rigid_id(self):
        """Returns the maximum rigid body ID contained in the Compound.

        This is usually used by compound.root to determine the maximum
        rigid_id in the containment hierarchy.

        Returns
        -------
        int or None
            The maximum rigid body ID contained in the Compound. If no
            rigid body IDs are found, None is returned

        """
        try:
            return max([particle.rigid_id for particle in self.particles()
                        if particle.rigid_id is not None])
        except ValueError:
            return

    def rigid_particles(self, rigid_id=None):
        """Generate all particles in rigid bodies.

        If a rigid_id is specified, then this function will only yield particles
        with a matching rigid_id.

        Parameters
        ----------
        rigid_id : int, optional
            Include only particles with this rigid body ID

        Yields
        ------
        mb.Compound
            The next particle with a rigid_id that is not None, or the next
            particle with a matching rigid_id if specified

        """
        for particle in self.particles():
            if rigid_id is not None:
                if particle.rigid_id == rigid_id:
                    yield particle
            else:
                if particle.rigid_id is not None:
                    yield particle

    def label_rigid_bodies(self, discrete_bodies=None, rigid_particles=None):
        """Designate which Compounds should be treated as rigid bodies

        If no arguments are provided, this function will treat the compound
        as a single rigid body by providing all particles in `self` with the
        same rigid_id. If `discrete_bodies` is not None, each instance of 
        a Compound with a name found in `discrete_bodies` will be treated as a 
        unique rigid body. If `rigid_particles` is not None, only Particles 
        (Compounds at the bottom of the containment hierarchy) matching this name 
        will be considered part of the rigid body.

        Parameters
        ----------
        discrete_bodies : str or list of str, optional, default=None
            Name(s) of Compound instances to be treated as unique rigid bodies.
            Compound instances matching this (these) name(s) will be provided 
            with unique rigid_ids
        rigid_particles : str or list of str, optional, default=None
            Name(s) of Compound instances at the bottom of the containment
            hierarchy (Particles) to be included in rigid bodies. Only Particles
            matching this (these) name(s) will have their rigid_ids altered to 
            match the rigid body number.

        Examples
        --------
        Creating a rigid benzene

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.label_rigid_bodies()

        Creating a semi-rigid benzene, where only the carbons are treated as
        a rigid body

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.label_rigid_bodies(rigid_particles='C')

        Create a box of rigid benzenes, where each benzene has a unique rigid
        body ID.

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.name = 'Benzene'
        >>> filled = mb.fill_box(benzene,
        ...                      n_compounds=10,
        ...                      box=[0, 0, 0, 4, 4, 4])
        >>> filled.label_rigid_bodies(distinct_bodies='Benzene')

        Create a box of semi-rigid benzenes, where each benzene has a unique
        rigid body ID and only the carbon portion is treated as rigid.

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.name = 'Benzene'
        >>> filled = mb.fill_box(benzene,
        ...                      n_compounds=10,
        ...                      box=[0, 0, 0, 4, 4, 4])
        >>> filled.label_rigid_bodies(distinct_bodies='Benzene',
        ...                           rigid_particles='C')

        """
        if discrete_bodies is not None:
            if isinstance(discrete_bodies, string_types):
                discrete_bodies = [discrete_bodies]
        if rigid_particles is not None:
            if isinstance(rigid_particles, string_types):
                rigid_particles = [rigid_particles]

        if self.root.max_rigid_id is not None:
            rigid_id = self.root.max_rigid_id + 1
            warn("{} rigid bodies already exist.  Incrementing 'rigid_id'"
                 "starting from {}.".format(rigid_id, rigid_id))
        else:
            rigid_id = 0

        for successor in self.successors():
            if discrete_bodies and successor.name not in discrete_bodies:
                continue
            for particle in successor.particles():
                if rigid_particles and particle.name not in rigid_particles:
                    continue
                particle.rigid_id = rigid_id
            if discrete_bodies:
                rigid_id += 1

    def unlabel_rigid_bodies(self):
        """Remove all rigid body labels from the Compound """
        self._check_if_contains_rigid_bodies = True
        for child in self.children:
            child._check_if_contains_rigid_bodies = True
        for particle in self.particles():
            particle.rigid_id = None

    def _increment_rigid_ids(self, increment):
        """Increment the rigid_id of all rigid Particles in a Compound

        Adds `increment` to the rigid_id of all Particles in `self` that
        already have an integer rigid_id.
        """
        for particle in self.particles():
            if particle.rigid_id is not None:
                particle.rigid_id += increment

    def _reorder_rigid_ids(self):
        """Reorder rigid body IDs ensuring consecutiveness.

        Primarily used internally to ensure consecutive rigid_ids following 
        removal of a Compound.

        """
        max_rigid = self.max_rigid_id
        unique_rigid_ids = sorted(set([p.rigid_id for p in self.rigid_particles()]))
        n_unique_rigid = len(unique_rigid_ids)
        if max_rigid and n_unique_rigid != max_rigid + 1:
            missing_rigid_id = (unique_rigid_ids[-1] * (unique_rigid_ids[-1] + 1))/2 - sum(unique_rigid_ids)
            for successor in self.successors():
                if successor.rigid_id is not None:
                    if successor.rigid_id > missing_rigid_id:
                        successor.rigid_id -= 1
            if self.rigid_id:
                if self.rigid_id > missing_rigid_id:
                    self.rigid_id -= 1

    def add(self, new_child, label=None, containment=True, replace=False,
            inherit_periodicity=True, reset_rigid_ids=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.
        inherit_periodicity : bool, optional, default=True
            Replace the periodicity of self with the periodicity of the
            Compound being added
        reset_rigid_ids : bool, optional, default=True
            If the Compound to be added contains rigid bodies, reset the 
            rigid_ids such that values remain distinct from rigid_ids
            already present in `self`. Can be set to False if attempting
            to add Compounds to an existing rigid body.

        """
        # 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, reset_rigid_ids=reset_rigid_ids)
            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))

        if new_child.contains_rigid or new_child.rigid_id is not None:
            if self.contains_rigid and reset_rigid_ids:
                new_child._increment_rigid_ids(increment=self.max_rigid_id + 1)
            self._check_if_contains_rigid_bodies = True
        if self.rigid_id is not None:
            self.rigid_id = None

        # 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.

        Parameters
        ----------
        objs_to_remove : mb.Compound or list of mb.Compound
            The Compound(s) to be removed from self

        """
        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

        remove_from_here = objs_to_remove.intersection(self.children)
        self.children -= remove_from_here
        yet_to_remove = objs_to_remove - remove_from_here

        for removed in remove_from_here:
            for child in removed.children:
                removed.remove(child)

        for removed_part in remove_from_here:
            if removed_part.rigid_id is not None:
                for ancestor in removed_part.ancestors():
                    ancestor._check_if_contains_rigid_bodies = True
            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.
        for child in self.children:
            child.remove(yet_to_remove)
            if child.contains_rigid:
                self.root._reorder_rigid_ids()

    def _remove_references(self, 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 list(removed_part.labels.items()):
                if not isinstance(part, Compound):
                    for p in part:
                        self._remove_references(p)
                elif removed_part not in part.ancestors():
                    try:
                        part.referrers.discard(removed_part)
                    except KeyError:
                        pass
                    else:
                        labels_to_delete.append(label)
        for label in labels_to_delete:
            removed_part.labels.pop(label, None)

    def referenced_ports(self):
        """Return all Ports referenced by this Compound.

        Returns
        -------
        list of mb.Compound
            A list of all ports referenced by the Compound

        """
        from mbuild.port import Port
        return [port for port in self.labels.values()
                if isinstance(port, Port)]

    def all_ports(self):
        """Return all Ports referenced by this Compound and its successors

        Returns
        -------
        list of mb.Compound
            A list of all Ports referenced by this Compound and its successors

        """
        from mbuild.port import Port
        return [successor for successor in self.successors()
                if isinstance(successor, Port)]

    def available_ports(self):
        """Return all unoccupied Ports referenced by this Compound.

        Returns
        -------
        list of mb.Compound
            A list of all unoccupied ports referenced by the 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):
        """Return all bonds in the Compound and sub-Compounds.

        Yields
        -------
        tuple of mb.Compound
            The next bond in the Compound

        See Also
        --------
        bond_graph.edges_iter : Iterates over all edges in a BondGraph

        """
        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.

        Returns
        -------
        int
            The number of bonds in the Compound

        """
        return sum(1 for _ in self.bonds())

    def add_bond(self, particle_pair):
        """Add a bond between two Particles.

        Parameters
        ----------
        particle_pair : indexable object, length=2, dtype=mb.Compound
            The pair of Particles to add a bond between

        """
        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].

        Parameters
        ----------
        name_a : str
            The name of one of the Particles to be in each bond
        name_b : str
            The name of the other Particle to be in each bond
        dmin : float
            The minimum distance between Particles for considering a bond
        dmax : float
            The maximum distance between Particles for considering a bond

        """
        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
                min_dist = self.min_periodic_distance(p2.pos, p1.pos)
                if (p2.name == name_b) and (dmin <= min_dist <= dmax):
                    self.add_bond((p1, p2))
                    added_bonds.append(bond_tuple)

    def remove_bond(self, particle_pair):
        """Deletes a bond between a pair of Particles

        Parameters
        ----------
        particle_pair : indexable object, length=2, dtype=mb.Compound
            The pair of Particles to remove the bond between

        """
        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), dtype=float
            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.

        Returns
        -------
        pos : np.ndarray, shape=(n, 3), dtype=float
            Array with the positions of all particles and 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

    @xyz.setter
    def xyz(self, arrnx3):
        """Set the positions of the particles in the Compound, excluding the Ports.

        This function does not set the position of the ports.

        Parameters
        ----------
        arrnx3 : np.ndarray, shape=(n,3), dtype=float
            The new particle positions

        """
        if not self.children:
            if not arrnx3.shape[0] == 1:
                raise ValueError('Trying to set position of {} with more than one'
                                 'coordinate: {}'.format(self, arrnx3))
            self.pos = np.squeeze(arrnx3)
        else:
            for atom, coords in zip(self._particles(include_ports=False), arrnx3):
                atom.pos = coords

    @xyz_with_ports.setter
    def xyz_with_ports(self, arrnx3):
        """Set the positions of the particles in the Compound, including the Ports.

        Parameters
        ----------
        arrnx3 : np.ndarray, shape=(n,3), dtype=float
            The new particle positions

        """
        if not self.children:
            if not arrnx3.shape[0] == 1:
                raise ValueError('Trying to set position of {} with more than one'
                                 'coordinate: {}'.format(self, arrnx3))
            self.pos = np.squeeze(arrnx3)
        else:
            for atom, coords in zip(self._particles(include_ports=True), arrnx3):
                atom.pos = coords

    @property
    def center(self):
        """The cartesian center of the Compound based on its Particles.

        Returns
        -------
        np.ndarray, shape=(3,), dtype=float
            The cartesian center of the Compound based on its Particles

        """
        if self.xyz.any():
            return np.mean(self.xyz, axis=0)

    @property
    def boundingbox(self):
        """Compute the bounding box of the compound.

        Returns
        -------
        mb.Box
            The bounding box for this 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.

        Parameters
        ----------
        xyz0 : np.ndarray, shape=(3,), dtype=float
            Coordinates of first point
        xyz1 : np.ndarray, shape=(3,), dtype=float
            Coordinates of second point

        Returns
        -------
        float
            Vectorized distance between the two points following minimum
            image convention

        """
        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.

        Parameters
        ----------
        compound : mb.Compound
            Reference particle to find other particles in range of
        dmax : float
            Maximum distance from 'compound' to look for Particles
        max_particles : int, optional, default=20
            Maximum number of Particles to return
        particle_kdtree : mb.PeriodicCKDTree, optional
            KD-tree for looking up nearest neighbors. If not provided, a KD-
            tree will be generated from all Particles in self
        particle_array : np.ndarray, shape=(n,), dtype=mb.Compound, optional
            Array of possible particles to consider for return. If not
            provided, this defaults to all Particles in self

        Returns
        -------
        np.ndarray, shape=(n,), dtype=mb.Compound
            Particles in range of compound according to user-defined limits

        See Also
        --------
        periodic_kdtree.PerioidicCKDTree : mBuild implementation of kd-trees
        scipy.spatial.ckdtree : Further details on kd-trees

        """
        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 visualize(self, show_ports=False):
        """Visualize the Compound using nglview.

        Allows for visualization of a Compound within a Jupyter Notebook.

        Parameters
        ----------
        show_ports : bool, optional, default=False
            Visualize Ports in addition to Particles

        """
        nglview = import_('nglview')
        if run_from_ipython():
            structure = self.to_trajectory(show_ports)
            return nglview.show_mdtraj(structure)
        else:
            raise RuntimeError('Visualization is only supported in Jupyter '
                               'Notebooks.')

    def update_coordinates(self, filename, update_port_locations=True):
        """Update the coordinates of this Compound from a file.

        Parameters
        ----------
        filename : str
            Name of file from which to load coordinates. Supported file types
            are the same as those supported by load()
        update_port_locations : bool, optional, default=True
            Update the locations of Ports so that they are shifted along with
            their anchor particles.  Note: This conserves the location of
            Ports with respect to the anchor Particle, but does not conserve
            the orientation of Ports with respect to the molecule as a whole.

        See Also
        --------
        load : Load coordinates from a file

        """
        if update_port_locations:
            compound_init = clone(self)
            load(filename, compound=self, coords_only=True)
            self._update_port_locations(compound_init)
        else:
            load(filename, compound=self, coords_only=True)

    def _update_port_locations(self, initial_compound):
        """Adjust port locations after particles have moved
<commit_msg>Port shifting in update_coordinates is no longer default.  Using as a test only<commit_after>from __future__ import print_function, division

__all__ = ['load', 'clone', 'Compound', 'Particle']

import collections
from collections import OrderedDict, defaultdict
from copy import deepcopy
import itertools
import os
import sys
import tempfile
from warnings import warn

import mdtraj as md
import numpy as np
from oset import oset as OrderedSet
import parmed as pmd
from parmed.periodic_table import AtomicNum, element_by_name, Mass
from six import integer_types, string_types

from mbuild.bond_graph import BondGraph
from mbuild.box import Box
from mbuild.coordinate_transform import translate
from mbuild.exceptions import MBuildError
from mbuild.formats.hoomdxml import write_hoomdxml
from mbuild.formats.lammpsdata import write_lammpsdata
from mbuild.formats.gsdwriter import write_gsd
from mbuild.periodic_kdtree import PeriodicCKDTree
from mbuild.utils.io import run_from_ipython, import_
from mbuild.coordinate_transform import _translate, _rotate

def load(filename, relative_to_module=None, compound=None, coords_only=False,
         rigid=False, **kwargs):
    """Load a file into an mbuild compound.

    Files are read using the MDTraj package. Please refer to http://mdtraj.org/
    1.8.0/load_functions.html for supported formats.

    Parameters
    ----------
    filename : str
        Name of the file from which to load atom and bond information.
    relative_to_module : str, optional, default=None
        Instead of looking in the current working directory, look for the file
        where this module is defined. This is typically used in Compound classes
        that will be instantiated from a different directory (such as the
        Compounds located in mbuild.lib).
    compound : mb.Compound, optional, default=None
        Existing compound to load atom and bond information into.
    coords_only : bool, optional, default=False
        Only load the coordinates into an existing compoint.
    rigid : bool, optional
        Treat the compound as a rigid body

    Returns
    -------
    compound : mb.Compound

    """
    # Handle mbuild *.py files containing a class that wraps a structure file
    # in its own folder. E.g., you build a system from ~/foo.py and it imports
    # from ~/bar/baz.py where baz.py loads ~/bar/baz.pdb.
    if relative_to_module:
        script_path = os.path.realpath(sys.modules[relative_to_module].__file__)
        file_dir = os.path.dirname(script_path)
        filename = os.path.join(file_dir, filename)

    if compound is None:
        compound = Compound()

    traj = md.load(filename, **kwargs)
    compound.from_trajectory(traj, frame=-1, coords_only=coords_only)
    if rigid:
        compound.label_rigid_bodies()
    return compound


def clone(existing_compound, 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.

    Parameters
    ----------
    existing_compound : mb.Compound
        Existing Compound that will be copied

    Other Parameters
    ----------------
    clone_of : dict, optional
    root_container : mb.Compound, optional

    """
    if clone_of is None:
        clone_of = dict()

    newone = existing_compound._clone(clone_of=clone_of,
                                      root_container=root_container)
    existing_compound._clone_bonds(clone_of=clone_of)
    return newone


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; <NAME>; <NAME>; <NAME> (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 : mb.Compound or list of mb.Compound, optional, default=None
        One or more compounds to be added to self.
    name : str, optional, default=self.__class__.__name__
        The type of Compound.
    pos : np.ndarray, shape=(3,), dtype=float, optional, default=[0, 0, 0]
        The position of the Compound in Cartestian space
    charge : float, optional, default=0.0
        Currently not used. Likely removed in next release.
    periodicity : np.ndarray, shape=(3,), dtype=float, optional, default=[0, 0, 0]
        The periodic lengths of the Compound in the x, y and z directions.
        Defaults to zeros which is treated as non-periodic.
    port_particle : bool, optional, default=False
        Whether or not this Compound is part of a Port

    Attributes
    ----------
    bond_graph : mb.BondGraph
        Graph-like object that stores bond information for this Compound
    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.
    rigid_id : int, default=None
        The ID of the rigid body that this Compound belongs to.  Only Particles
        (the bottom of the containment hierarchy) can have integer values for
        `rigid_id`. Compounds containing rigid particles will always have 
        `rigid_id == None`. See also `contains_rigid`.
    boundingbox
    center
    contains_rigid
    max_rigid_id
    n_particles
    n_bonds
    root
    xyz
    xyz_with_ports

    """

    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 periodicity 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._rigid_id = None
        self._contains_rigid = False
        self._check_if_contains_rigid_bodies = False

        # self.add() must be called after labels and children are initialized.
        if subcompounds:
            self.add(subcompounds)

    def particles(self, include_ports=False):
        """Return all Particles of the Compound.
        
        Parameters
        ----------
        include_ports : bool, optional, default=False
            Include port particles

        Yields
        -------
        mb.Compound
            The next Particle in the Compound

        """
        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.

        Yields
        -------
        mb.Compound
            The next Particle 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):
        """Return the number of Particles in the Compound.

        Returns
        -------
        int
            The number of Particles in the Compound

        """
        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.

        Yields
        ------
        mb.Compound
            The next Compound above self in the hierarchy

        """
        if self.parent is not None:
            yield self.parent
            for ancestor in self.parent.ancestors():
                yield ancestor

    @property
    def root(self):
        """The Compound at the top of self's hierarchy.

        Returns
        -------
        mb.Compound
            The Compound at the top of self's hierarchy

        """
        parent = None
        for parent in self.ancestors():
            pass
        if parent is None:
            return self
        return parent

    def particles_by_name(self, name):
        """Return all Particles of the Compound with a specific name

        Parameters
        ----------
        name : str
            Only particles with this name are returned

        Yields
        ------
        mb.Compound
            The next Particle in the Compound with the user-specified name

        """
        for particle in self.particles():
            if particle.name == name:
                yield particle

    @property
    def rigid_id(self):
        return self._rigid_id

    @rigid_id.setter
    def rigid_id(self, value):
        if self._contains_only_ports():
            self._rigid_id = value
            for ancestor in self.ancestors():
                ancestor._check_if_contains_rigid_bodies = True
        else:
            raise AttributeError("rigid_id is immutable for Compounds that are "
                                 "not at the bottom of the containment hierarchy.")

    @property
    def contains_rigid(self):
        """Returns True if the Compound contains rigid bodies

        If the Compound contains any particle with a rigid_id != None 
        then contains_rigid will return True. If the Compound has no
        children (i.e. the Compound resides at the bottom of the containment
        hierarchy) then contains_rigid will return False.

        Returns
        -------
        bool
            True if the Compound contains any particle with a rigid_id != None

        Notes
        -----
        The private variable '_check_if_contains_rigid_bodies' is used to help
        cache the status of 'contains_rigid'. If '_check_if_contains_rigid_bodies'
        is False, then the rigid body containment of the Compound has not changed,
        and the particle tree is not traversed, boosting performance.

        """
        if self._check_if_contains_rigid_bodies:
            self._check_if_contains_rigid_bodies = False
            if any(particle.rigid_id is not None for particle in self._particles()):
                self._contains_rigid = True
            else:
                self._contains_rigid = False
        return self._contains_rigid

    @property
    def max_rigid_id(self):
        """Returns the maximum rigid body ID contained in the Compound.

        This is usually used by compound.root to determine the maximum
        rigid_id in the containment hierarchy.

        Returns
        -------
        int or None
            The maximum rigid body ID contained in the Compound. If no
            rigid body IDs are found, None is returned

        """
        try:
            return max([particle.rigid_id for particle in self.particles()
                        if particle.rigid_id is not None])
        except ValueError:
            return

    def rigid_particles(self, rigid_id=None):
        """Generate all particles in rigid bodies.

        If a rigid_id is specified, then this function will only yield particles
        with a matching rigid_id.

        Parameters
        ----------
        rigid_id : int, optional
            Include only particles with this rigid body ID

        Yields
        ------
        mb.Compound
            The next particle with a rigid_id that is not None, or the next
            particle with a matching rigid_id if specified

        """
        for particle in self.particles():
            if rigid_id is not None:
                if particle.rigid_id == rigid_id:
                    yield particle
            else:
                if particle.rigid_id is not None:
                    yield particle

    def label_rigid_bodies(self, discrete_bodies=None, rigid_particles=None):
        """Designate which Compounds should be treated as rigid bodies

        If no arguments are provided, this function will treat the compound
        as a single rigid body by providing all particles in `self` with the
        same rigid_id. If `discrete_bodies` is not None, each instance of 
        a Compound with a name found in `discrete_bodies` will be treated as a 
        unique rigid body. If `rigid_particles` is not None, only Particles 
        (Compounds at the bottom of the containment hierarchy) matching this name 
        will be considered part of the rigid body.

        Parameters
        ----------
        discrete_bodies : str or list of str, optional, default=None
            Name(s) of Compound instances to be treated as unique rigid bodies.
            Compound instances matching this (these) name(s) will be provided 
            with unique rigid_ids
        rigid_particles : str or list of str, optional, default=None
            Name(s) of Compound instances at the bottom of the containment
            hierarchy (Particles) to be included in rigid bodies. Only Particles
            matching this (these) name(s) will have their rigid_ids altered to 
            match the rigid body number.

        Examples
        --------
        Creating a rigid benzene

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.label_rigid_bodies()

        Creating a semi-rigid benzene, where only the carbons are treated as
        a rigid body

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.label_rigid_bodies(rigid_particles='C')

        Create a box of rigid benzenes, where each benzene has a unique rigid
        body ID.

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.name = 'Benzene'
        >>> filled = mb.fill_box(benzene,
        ...                      n_compounds=10,
        ...                      box=[0, 0, 0, 4, 4, 4])
        >>> filled.label_rigid_bodies(distinct_bodies='Benzene')

        Create a box of semi-rigid benzenes, where each benzene has a unique
        rigid body ID and only the carbon portion is treated as rigid.

        >>> import mbuild as mb
        >>> from mbuild.utils.io import get_fn
        >>> benzene = mb.load(get_fn('benzene.mol2'))
        >>> benzene.name = 'Benzene'
        >>> filled = mb.fill_box(benzene,
        ...                      n_compounds=10,
        ...                      box=[0, 0, 0, 4, 4, 4])
        >>> filled.label_rigid_bodies(distinct_bodies='Benzene',
        ...                           rigid_particles='C')

        """
        if discrete_bodies is not None:
            if isinstance(discrete_bodies, string_types):
                discrete_bodies = [discrete_bodies]
        if rigid_particles is not None:
            if isinstance(rigid_particles, string_types):
                rigid_particles = [rigid_particles]

        if self.root.max_rigid_id is not None:
            rigid_id = self.root.max_rigid_id + 1
            warn("{} rigid bodies already exist.  Incrementing 'rigid_id'"
                 "starting from {}.".format(rigid_id, rigid_id))
        else:
            rigid_id = 0

        for successor in self.successors():
            if discrete_bodies and successor.name not in discrete_bodies:
                continue
            for particle in successor.particles():
                if rigid_particles and particle.name not in rigid_particles:
                    continue
                particle.rigid_id = rigid_id
            if discrete_bodies:
                rigid_id += 1

    def unlabel_rigid_bodies(self):
        """Remove all rigid body labels from the Compound """
        self._check_if_contains_rigid_bodies = True
        for child in self.children:
            child._check_if_contains_rigid_bodies = True
        for particle in self.particles():
            particle.rigid_id = None

    def _increment_rigid_ids(self, increment):
        """Increment the rigid_id of all rigid Particles in a Compound

        Adds `increment` to the rigid_id of all Particles in `self` that
        already have an integer rigid_id.
        """
        for particle in self.particles():
            if particle.rigid_id is not None:
                particle.rigid_id += increment

    def _reorder_rigid_ids(self):
        """Reorder rigid body IDs ensuring consecutiveness.

        Primarily used internally to ensure consecutive rigid_ids following 
        removal of a Compound.

        """
        max_rigid = self.max_rigid_id
        unique_rigid_ids = sorted(set([p.rigid_id for p in self.rigid_particles()]))
        n_unique_rigid = len(unique_rigid_ids)
        if max_rigid and n_unique_rigid != max_rigid + 1:
            missing_rigid_id = (unique_rigid_ids[-1] * (unique_rigid_ids[-1] + 1))/2 - sum(unique_rigid_ids)
            for successor in self.successors():
                if successor.rigid_id is not None:
                    if successor.rigid_id > missing_rigid_id:
                        successor.rigid_id -= 1
            if self.rigid_id:
                if self.rigid_id > missing_rigid_id:
                    self.rigid_id -= 1

    def add(self, new_child, label=None, containment=True, replace=False,
            inherit_periodicity=True, reset_rigid_ids=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.
        inherit_periodicity : bool, optional, default=True
            Replace the periodicity of self with the periodicity of the
            Compound being added
        reset_rigid_ids : bool, optional, default=True
            If the Compound to be added contains rigid bodies, reset the 
            rigid_ids such that values remain distinct from rigid_ids
            already present in `self`. Can be set to False if attempting
            to add Compounds to an existing rigid body.

        """
        # 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, reset_rigid_ids=reset_rigid_ids)
            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))

        if new_child.contains_rigid or new_child.rigid_id is not None:
            if self.contains_rigid and reset_rigid_ids:
                new_child._increment_rigid_ids(increment=self.max_rigid_id + 1)
            self._check_if_contains_rigid_bodies = True
        if self.rigid_id is not None:
            self.rigid_id = None

        # 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.

        Parameters
        ----------
        objs_to_remove : mb.Compound or list of mb.Compound
            The Compound(s) to be removed from self

        """
        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

        remove_from_here = objs_to_remove.intersection(self.children)
        self.children -= remove_from_here
        yet_to_remove = objs_to_remove - remove_from_here

        for removed in remove_from_here:
            for child in removed.children:
                removed.remove(child)

        for removed_part in remove_from_here:
            if removed_part.rigid_id is not None:
                for ancestor in removed_part.ancestors():
                    ancestor._check_if_contains_rigid_bodies = True
            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.
        for child in self.children:
            child.remove(yet_to_remove)
            if child.contains_rigid:
                self.root._reorder_rigid_ids()

    def _remove_references(self, 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 list(removed_part.labels.items()):
                if not isinstance(part, Compound):
                    for p in part:
                        self._remove_references(p)
                elif removed_part not in part.ancestors():
                    try:
                        part.referrers.discard(removed_part)
                    except KeyError:
                        pass
                    else:
                        labels_to_delete.append(label)
        for label in labels_to_delete:
            removed_part.labels.pop(label, None)

    def referenced_ports(self):
        """Return all Ports referenced by this Compound.

        Returns
        -------
        list of mb.Compound
            A list of all ports referenced by the Compound

        """
        from mbuild.port import Port
        return [port for port in self.labels.values()
                if isinstance(port, Port)]

    def all_ports(self):
        """Return all Ports referenced by this Compound and its successors

        Returns
        -------
        list of mb.Compound
            A list of all Ports referenced by this Compound and its successors

        """
        from mbuild.port import Port
        return [successor for successor in self.successors()
                if isinstance(successor, Port)]

    def available_ports(self):
        """Return all unoccupied Ports referenced by this Compound.

        Returns
        -------
        list of mb.Compound
            A list of all unoccupied ports referenced by the 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):
        """Return all bonds in the Compound and sub-Compounds.

        Yields
        -------
        tuple of mb.Compound
            The next bond in the Compound

        See Also
        --------
        bond_graph.edges_iter : Iterates over all edges in a BondGraph

        """
        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.

        Returns
        -------
        int
            The number of bonds in the Compound

        """
        return sum(1 for _ in self.bonds())

    def add_bond(self, particle_pair):
        """Add a bond between two Particles.

        Parameters
        ----------
        particle_pair : indexable object, length=2, dtype=mb.Compound
            The pair of Particles to add a bond between

        """
        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].

        Parameters
        ----------
        name_a : str
            The name of one of the Particles to be in each bond
        name_b : str
            The name of the other Particle to be in each bond
        dmin : float
            The minimum distance between Particles for considering a bond
        dmax : float
            The maximum distance between Particles for considering a bond

        """
        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
                min_dist = self.min_periodic_distance(p2.pos, p1.pos)
                if (p2.name == name_b) and (dmin <= min_dist <= dmax):
                    self.add_bond((p1, p2))
                    added_bonds.append(bond_tuple)

    def remove_bond(self, particle_pair):
        """Deletes a bond between a pair of Particles

        Parameters
        ----------
        particle_pair : indexable object, length=2, dtype=mb.Compound
            The pair of Particles to remove the bond between

        """
        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), dtype=float
            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.

        Returns
        -------
        pos : np.ndarray, shape=(n, 3), dtype=float
            Array with the positions of all particles and 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

    @xyz.setter
    def xyz(self, arrnx3):
        """Set the positions of the particles in the Compound, excluding the Ports.

        This function does not set the position of the ports.

        Parameters
        ----------
        arrnx3 : np.ndarray, shape=(n,3), dtype=float
            The new particle positions

        """
        if not self.children:
            if not arrnx3.shape[0] == 1:
                raise ValueError('Trying to set position of {} with more than one'
                                 'coordinate: {}'.format(self, arrnx3))
            self.pos = np.squeeze(arrnx3)
        else:
            for atom, coords in zip(self._particles(include_ports=False), arrnx3):
                atom.pos = coords

    @xyz_with_ports.setter
    def xyz_with_ports(self, arrnx3):
        """Set the positions of the particles in the Compound, including the Ports.

        Parameters
        ----------
        arrnx3 : np.ndarray, shape=(n,3), dtype=float
            The new particle positions

        """
        if not self.children:
            if not arrnx3.shape[0] == 1:
                raise ValueError('Trying to set position of {} with more than one'
                                 'coordinate: {}'.format(self, arrnx3))
            self.pos = np.squeeze(arrnx3)
        else:
            for atom, coords in zip(self._particles(include_ports=True), arrnx3):
                atom.pos = coords

    @property
    def center(self):
        """The cartesian center of the Compound based on its Particles.

        Returns
        -------
        np.ndarray, shape=(3,), dtype=float
            The cartesian center of the Compound based on its Particles

        """
        if self.xyz.any():
            return np.mean(self.xyz, axis=0)

    @property
    def boundingbox(self):
        """Compute the bounding box of the compound.

        Returns
        -------
        mb.Box
            The bounding box for this 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.

        Parameters
        ----------
        xyz0 : np.ndarray, shape=(3,), dtype=float
            Coordinates of first point
        xyz1 : np.ndarray, shape=(3,), dtype=float
            Coordinates of second point

        Returns
        -------
        float
            Vectorized distance between the two points following minimum
            image convention

        """
        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.

        Parameters
        ----------
        compound : mb.Compound
            Reference particle to find other particles in range of
        dmax : float
            Maximum distance from 'compound' to look for Particles
        max_particles : int, optional, default=20
            Maximum number of Particles to return
        particle_kdtree : mb.PeriodicCKDTree, optional
            KD-tree for looking up nearest neighbors. If not provided, a KD-
            tree will be generated from all Particles in self
        particle_array : np.ndarray, shape=(n,), dtype=mb.Compound, optional
            Array of possible particles to consider for return. If not
            provided, this defaults to all Particles in self

        Returns
        -------
        np.ndarray, shape=(n,), dtype=mb.Compound
            Particles in range of compound according to user-defined limits

        See Also
        --------
        periodic_kdtree.PerioidicCKDTree : mBuild implementation of kd-trees
        scipy.spatial.ckdtree : Further details on kd-trees

        """
        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 visualize(self, show_ports=False):
        """Visualize the Compound using nglview.

        Allows for visualization of a Compound within a Jupyter Notebook.

        Parameters
        ----------
        show_ports : bool, optional, default=False
            Visualize Ports in addition to Particles

        """
        nglview = import_('nglview')
        if run_from_ipython():
            structure = self.to_trajectory(show_ports)
            return nglview.show_mdtraj(structure)
        else:
            raise RuntimeError('Visualization is only supported in Jupyter '
                               'Notebooks.')

    def update_coordinates(self, filename, update_port_locations=False):
        """Update the coordinates of this Compound from a file.

        Parameters
        ----------
        filename : str
            Name of file from which to load coordinates. Supported file types
            are the same as those supported by load()
        update_port_locations : bool, optional, default=True
            Update the locations of Ports so that they are shifted along with
            their anchor particles.  Note: This conserves the location of
            Ports with respect to the anchor Particle, but does not conserve
            the orientation of Ports with respect to the molecule as a whole.

        See Also
        --------
        load : Load coordinates from a file

        """
        if update_port_locations:
            compound_init = clone(self)
            load(filename, compound=self, coords_only=True)
            self._update_port_locations(compound_init)
        else:
            load(filename, compound=self, coords_only=True)

    def _update_port_locations(self, initial_compound):
        """Adjust port locations after particles have moved

{'fraction_non_alphanumeric': 0.07570990172085615, 'fraction_numerical': 0.0037809631097126726, 'mean_word_length': 1.9069898746565288, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 135, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2596453', 'n_tokens_mistral': 21048, 'n_tokens_neox': 19907, 'n_words': 6862}