__init__.py

# -*- coding: utf-8 -*-
from pathlib import Path

from .version import get_version
__version__ = get_version() 

import numpy as np
from copy import copy, deepcopy
from inspect import ismethod
import os, sys, subprocess

_RESOURCE_DIR = Path(__file__).parent / 'resources'
def get_resource_path(relative_path):
    return _RESOURCE_DIR / relative_path

## Abstract classes
class Transformable():
    def __init__(self, position, orientation=None):
        self.position = np.array(position)
        if orientation is not None:
            orientation = np.array(orientation)
        self.orientation = orientation

    def transform(self, R = ((1,0,0),(0,1,0),(0,0,1)),
                  center = (0,0,0), offset = (0,0,0)):

        R,center,offset = [np.array(x) for x in (R,center,offset)]

        self.position = R.dot(self.position-center)+center+offset
                
        if self.orientation is not None:
            ## TODO: what if self.orientation is taken from parent?!
            self.orientation = self.orientation.dot(R)
        ...        

    def collapsedPosition(self):
        # print("collapsedPosition called", type(self), self.name)
        if isinstance(self, Child):
            # print(self.parent, isinstance(self.parent,Transformable))
            if isinstance(self.parent, Transformable):
                return self.applyOrientation(self.position) + self.parent.collapsedPosition()
            
                # if self.parent.orientation is not None:
                #     return self.parent.collapsedOrientation().dot(self.position) + self.parent.collapsedPosition()
        return np.array(self.position) # return a copy
                
    def applyOrientation(self,obj):
        # print("applyOrientation called", self.name, obj)
        if isinstance(self, Child):
            # print( self.orientation, self.orientation is not None, None is not None )
            # if self.orientation is not None:
            #     # print("applyOrientation applying", self, self.name, self.orientation)
            #     obj = self.orientation.dot(obj)
            if isinstance(self.parent, Transformable):
                if self.parent.orientation is not None:
                    obj = self.parent.orientation.dot(obj)
                obj = self.parent.applyOrientation(obj)
        # print("applyOrientation returning", self.name, obj)
        return obj

class Parent():
    def __init__(self, children=None, remove_duplicate_bonded_terms=False):
        self.children = []
        if children is not None:
            for x in children:
                self.add(x)
        
        self.remove_duplicate_bonded_terms = remove_duplicate_bonded_terms
        self.bonds = []
        self.angles = []
        self.dihedrals = []
        self.impropers = []
        self.exclusions = []
        self.bond_angles = []

        self.rigid = False

        ## TODO: self.cacheInvalid = True # What will be in the cache?

    def add(self,x):
        ## TODO: check the parent-child tree to make sure there are no cycles
        if not isinstance(x,Child):
            raise Exception('Attempted to add an object to a group that does not inherit from the "Child" type')

        if x.parent is not None and x.parent is not self:
            raise Exception("Child {} already belongs to some group".format(x))
        x.parent = self
        self.children.append(x)

    def insert(self,idx,x):
        ## TODO: check the parent-child tree to make sure there are no cycles
        if not isinstance(x,Child):
            raise Exception('Attempted to add an object to a group that does not inherit from the "Child" type')

        if x.parent is not None and x.parent is not self:
            raise Exception("Child {} already belongs to some group".format(x))
        x.parent = self
        self.children.insert(idx,x)

    def index(self, x):
        return self.children.index(x)

    def clear_all(self, keep_children=False):
        if keep_children == False:
            for x in self.children:
                x.parent = None
            self.children = []
        self.bonds = []
        self.angles = []
        self.dihedrals = []
        self.impropers = []
        self.exclusions = []

    def remove(self,x):
        if x in self.children:
            self.children.remove(x)
            if x.parent is self:
                x.parent = None

    def add_bond(self, i,j, bond, exclude=False):
        assert( i is not j )
        ## TODO: how to handle duplicating and cloning bonds
        # beads = [b for b in self]
        # for b in (i,j): assert(b in beads)
        self.bonds.append( (i,j, bond, exclude) )

    def add_angle(self, i,j,k, angle):
        assert( len(set((i,j,k))) == 3 )
        # beads = [b for b in self]
        # for b in (i,j,k): assert(b in beads)
        self.angles.append( (i,j,k, angle) )

    def add_dihedral(self, i,j,k,l, dihedral):
        assert( len(set((i,j,k,l))) == 4 )

        # beads = [b for b in self]
        # for b in (i,j,k,l): assert(b in beads)
        self.dihedrals.append( (i,j,k,l, dihedral) )

    def add_improper(self, i,j,k,l, dihedral):
        # beads = [b for b in self]
        # for b in (i,j,k,l): assert(b in beads)
        self.impropers.append( (i,j,k,l, dihedral) )

    def add_exclusion(self, i,j):
        ## TODO: how to handle duplicating and cloning bonds
        ## TODO: perform following check elsewhere
        # beads = [b for b in self]
        # for b in (i,j): assert(b in beads)
        self.exclusions.append( (i,j) )

    def add_bond_angle(self, i,j,k, bond_angle, exclude=False):
        assert( len(set((i,j,k))) == 3 )
        ## TODO: how to handle duplicating and cloning bonds
        # beads = [b for b in self]
        # for b in (i,j): assert(b in beads)
        self.bond_angles.append( (i,j,k, bond_angle) )

    def get_restraints(self):
        ret = []
        for c in self.children:
            ret.extend( c.get_restraints() )
        return ret

    def get_bonds(self):
        ret = self.bonds
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_bonds() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret


    def get_angles(self):
        ret = self.angles
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_angles() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret

    def get_dihedrals(self):
        ret = self.dihedrals
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_dihedrals() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret

    def get_impropers(self):
        ret = self.impropers
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_impropers() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret

    def get_exclusions(self):
        ret = self.exclusions
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_exclusions() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret

    def get_bond_angles(self):
        ret = self.bond_angles
        for c in self.children:
            if isinstance(c,Parent): ret.extend( c.get_bond_angles() )
        if self.remove_duplicate_bonded_terms:
            return list(set(ret))
        else:
            return ret

    def _get_bond_potentials(self):
        bonds =  [b for i,j,b,ex in self.get_bonds()]
        bondangles1 = [b[1] for i,j,k,b in self.get_bond_angles()]
        bondangles2 = [b[2] for i,j,k,b in self.get_bond_angles()]
        return list(set( bonds+bondangles1+bondangles2 ))

    def _get_angle_potentials(self):
        angles =  [b for i,j,k,b in self.get_angles()]
        bondangles = [b[0] for i,j,k,b in self.get_bond_angles()]
        return list(set( angles+bondangles ))


    ## Removed because prohibitively slow
    # def remove_duplicate_terms(self):
    #     for key in "bonds angles dihedrals impropers exclusions".split():
    #         self.remove_duplicate_item(key)

    # def remove_duplicate_item(self, dict_key, existing=None):
    #     if existing is None: existing = []
    #     ret = [i for i in list(set(self.__dict__[dict_key])) if i not in existing]
    #     self.__dict__[dict_key] = ret
    #     existing.extend(ret)
    #     for c in self.children:
    #         if isinstance(c,Parent): 
    #             ret = ret + c.remove_duplicate_item(dict_key, existing)
    #     return ret
        

    def __iter__(self):
        ## TODO: decide if this is the nicest way to do it!
        """Depth-first iteration through tree"""
        for x in self.children:
            if isinstance(x,Parent):
                if isinstance(x,Clone) and not isinstance(x.get_original_recursively(),Parent):
                    yield x
                else:
                    for y in x:
                        yield y
            else:
                yield x    

    def __len__(self):
        l = 0
        for x in self.children:
            if isinstance(x,Parent):
                l += len(x)
            else:
                l += 1
        return l
        
    def __getitem__(self, i):
        return self.children[i]
    
    def __setitem__(self, i, val):
        x = self.children[i]
        x.parent = None
        val.parent = self
        self.children[i] = val
        
class Child():
    def __init__(self, parent=None):
        self.parent = parent
        if parent is not None:
            assert( isinstance(parent, Parent) )
            parent.children.append(self)

    def __getattr__(self, name):
        """
        Try to get attribute from the parent

        """
        # if self.parent is not None:
        if "parent" not in self.__dict__ or self.__dict__["parent"] is None or name is "children":
            raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name))


        excluded_attributes = ['parent','rigid']
        if name in excluded_attributes:
            raise AttributeError("'{}' object has no attribute '{}' and cannot look it up from the parent".format(type(self).__name__, name))

        ## TODO: determine if there is a way to avoid __getattr__ if a method is being looked up  
        try:
            ret = getattr(self.parent,name)
        except:
            raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name))
        if ismethod(ret):
            raise AttributeError("'{}' object has no method '{}'".format(type(self).__name__, name))
        return ret 

            
    # def __getstate__(self):
    #     print("Child getstate called", self)
    #     print(self.__dict__)
    #     return (self.__dict__,)

    # def __setstate__(self, state):
    #     self.__dict__, = state

class Clone(Transformable, Parent, Child):
    def __init__(self, original, parent=None,
                 position = None,
                 orientation = None):
        if position is None and original.position is not None:
            position = np.array( original.position )
        if orientation is None and original.orientation is not None:
            orientation = np.array( original.orientation )
        if parent is None:
            parent = original.parent
        self.original = original
        Child.__init__(self, parent)        
        Transformable.__init__(self, position, orientation)        

        ## TODO: keep own bond_list, etc, update when needed original changes

        if "children" in original.__dict__ and len(original.children) > 0:
            self.children = [Clone(c, parent = self) for c in original.children]
        else:
            self.children = []

    def get_original_recursively(self):
        if isinstance(self.original, Clone):
            return self.original.get_original_recursively()
        else:
            return self.original

    def __getattr__(self, name):
        """
        Try to get attribute from the original without descending the tree heirarchy, then look up parent

        TODO: handle PointParticle lookups into ParticleType
        """
        # print("Clone getattr",name)
        if name in self.original.__dict__:
            return self.original.__dict__[name]
        else:
            if "parent" not in self.__dict__ or self.__dict__["parent"] is None:
                raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name))
            return getattr(self.parent, name)
        

## Particle classes
class ParticleType():
    """Class that hold common attributes that particles can point to"""

    excludedAttributes = ("idx","type_",
                          "position",
                          "children",
                          "parent", "excludedAttributes","rigid"
    )

    def __init__(self, name, charge=0, parent=None, **kargs):
        """ Parent type is used to fall back on for nonbonded interactions if this type is not specifically referenced """

        if parent is not None:
            for k,v in parent.__dict__.items():
                if k not in ParticleType.excludedAttributes:
                    self.__dict__[k] = v

        self.name   = name
        self.charge = charge
        self.parent = parent

        for key in ParticleType.excludedAttributes:
            assert( key not in kargs )

        for key,val in kargs.items():
            self.__dict__[key] = val

    def is_same_type(self, other):
        assert(isinstance(other,ParticleType))
        if self == other:
            return True
        elif self.parent is not None and self.parent == other:
            return True
        else:
            return False

    def __hash_key(self):
        l = [self.name,self.charge]
        for keyval in sorted(self.__dict__.items()):
            if isinstance(keyval[1], list): keyval = (keyval[0],tuple(keyval[1]))
            l.extend(keyval)
        return tuple(l)

    def __hash__(self):
        return hash(self.__hash_key())
    
    def _equal_check(a,b):
        if a.name == b.name:
            if a.__hash_key() != b.__hash_key():
                raise Exception("Two different ParticleTypes have same 'name' attribute")

    def __eq__(a,b):
        a._equal_check(b)
        return a.name == b.name
    def __lt__(a,b):
        a._equal_check(b)
        return a.name < b.name
    def __le__(a,b):
        a._equal_check(b)
        return a.name <= b.name
    def __gt__(a,b):
        a._equal_check(b)
        return a.name > b.name
    def __ge__(a,b):
        a._equal_check(b)
        return a.name >= b.name

    def __repr__(self):
        return '<{} {}{}>'.format( type(self), self.name, '[{}]'.format(self.parent) if self.parent is not None else '' )


class PointParticle(Transformable, Child):
    def __init__(self, type_, position, name="A", **kwargs):
        parent = None
        if 'parent' in kwargs:
            parent = kwargs['parent']
        Child.__init__(self, parent=parent)
        Transformable.__init__(self,position)

        self.type_    = type_                
        self.idx     = None
        self.name = name
        self.counter = 0
        self.restraints = []

        for key,val in kwargs.items():
            self.__dict__[key] = val
        
    def add_restraint(self, restraint):
        ## TODO: how to handle duplicating and cloning bonds
        self.restraints.append( restraint )

    def get_restraints(self):
        return [(self,r) for r in self.restraints]

    def duplicate(self):
        new = deepcopy(self)
        return new

    def __getattr__(self, name):
        """
        First try to get attribute from the parent, then type_
        
        Note that this data structure seems to be fragile, can result in stack overflow
        
        """
        # return Child.__getattr__(self,name)
        try:
            return Child.__getattr__(self,name)
        except Exception as e:
            if 'type_' in self.__dict__:
                return getattr(self.type_, name)
            else:
                raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name))

    def _get_psfpdb_dictionary(self):
        p = self
        try:
            segname = p.segname
        except:
            segname = "A"
        try:
            chain = p.chain
        except:
            chain = "A"
        try:
            resname = p.resname
        except:
            resname = p.name[:3]
        try:
            resid = p.resid
        except:
            resid = p.idx+1
        try:
            mass = p.mass
        except:
            mass = 1

        try:
            occ = p.occupancy
        except:
            occ = 0
        try:
            beta = p.beta
        except:
            beta = 0

        data = dict(segname = segname,
                    resname = resname,
                    name = str(p.name)[:4],
                    chain = chain[0],
                    resid = int(resid),
                    idx = p.idx+1,
                    type = p.type_.name[:7],
                    charge = p.charge,
                    mass = mass,
                    occupancy = occ,
                    beta = beta
                )
        return data


class Group(Transformable, Parent, Child):

    def __init__(self, name=None, children = None, parent=None, 
                 position = np.array((0,0,0)),
                 orientation = np.array(((1,0,0),(0,1,0),(0,0,1))),
                 remove_duplicate_bonded_terms = False,
                 **kwargs):

        Transformable.__init__(self, position, orientation)
        Child.__init__(self, parent) # Initialize Child first
        Parent.__init__(self, children, remove_duplicate_bonded_terms)
        self.name = name
        self.isClone = False

        for key,val in kwargs.items():
            self.__dict__[key] = val


    def clone(self):
        return Clone(self)
        g = copy(self)
        g.isClone = True        # TODO: use?
        g.children = [copy(c) for c in g.children]
        for c in g.children:
            c.parent = g
        return g
        g = Group(position = self.position,
                  orientation = self.orientation)
        g.children = self.children # lists point to the same object

    def duplicate(self):
        new = deepcopy(self)
        for c in new.children:
            c.parent = new
        return new
        # Group(position = self.position,
        #       orientation = self.orientation)
        # g.children = deepcopy self.children.deepcopy() # lists are the same object

    ## TODO override deepcopy so parent can be excluded from copying?
        
    # def __getstate__(self):
    #     return (self.children, self.parent, self.position, self.orientation)

    # def __setstate__(self, state):
    #     self.children, self.parent, self.position, self.orientation = state

        
class PdbModel(Transformable, Parent):

    def __init__(self, children=None, dimensions=None, remove_duplicate_bonded_terms=False):
        Transformable.__init__(self,(0,0,0))
        Parent.__init__(self, children, remove_duplicate_bonded_terms)
        self.dimensions = dimensions
        self.particles = [p for p in self]
        self.cacheInvalid = True

    def _updateParticleOrder(self):
        pass

    def writePdb(self, filename, beta_from_fixed=False):
        if self.cacheInvalid:
            self._updateParticleOrder()
        with open(filename,'w') as fh:
            ## Write header
            fh.write("CRYST1{:>9.3f}{:>9.3f}{:>9.3f}  90.00  90.00  90.00 P 1           1\n".format( *self.dimensions ))

            ## Write coordinates
            formatString = "ATOM {idx:>6.6s} {name:^4.4s} {resname:3.3s} {chain:1.1s}{resid:>5.5s}   {x:8.8s}{y:8.8s}{z:8.8s}{occupancy:6.2f}{beta:6.2f}  {charge:2d}{segname:>6s}\n"
            for p in self.particles:
                ## http://www.wwpdb.org/documentation/file-format-content/format33/sect9.html#ATOM
                data = p._get_psfpdb_dictionary()
                idx = data['idx']

                if np.log10(idx) >= 5:
                    idx = " *****"
                else:
                    idx = "{:>6d}".format(idx)
                data['idx'] = idx

                if beta_from_fixed:
                    data['beta'] = 1 if 'fixed' in p.__dict__ else 0

                pos = p.collapsedPosition()
                dig = [max(int(np.log10(np.abs(x)+1e-6)//1),0)+1 for x in pos]
                for d in dig: assert( d <= 7 )
                # assert( np.all(dig <= 7) )
                fs = ["{: %d.%df}" % (8,7-d) for d in dig]
                x,y,z = [f.format(x) for f,x in zip(fs,pos)] 
                data['x'] = x
                data['y'] = y
                data['z'] = z
                assert(data['resid'] < 1e5)
                data['charge'] = int(data['charge'])
                data['resid'] = "{:<4d}".format(data['resid'])
                fh.write( formatString.format(**data) )

        return
        
    def writePsf(self, filename):
        if self.cacheUpToDate == False:
            self._updateParticleOrder()
        with open(filename,'w') as fh:
            ## Write header
            fh.write("PSF NAMD\n\n") # create NAMD formatted psf
            fh.write("{:>8d} !NTITLE\n\n".format(0))
            
            ## ATOMS section
            fh.write("{:>8d} !NATOM\n".format(len(self.particles)))

            ## From vmd/plugins/molfile_plugin/src/psfplugin.c
            ## "%d %7s %10s %7s %7s %7s %f %f"
            formatString = "{idx:>8d} {segname:7.7s} {resid:<10.10s} {resname:7.7s}" + \
                           " {name:7.7s} {type:7.7s} {charge:f} {mass:f}\n"
            for p in self.particles:
                data = p._get_psfpdb_dictionary()
                data['resid'] = "%d%c%c" % (data['resid']," "," ") # TODO: work with large indices
                fh.write( formatString.format(**data) )
            fh.write("\n")

            ## Write out bonds
            bonds = self.get_bonds()
            fh.write("{:>8d} !NBOND\n".format(len(bonds)))
            counter = 0
            for p1,p2,b,ex in bonds:
                fh.write( "{:>8d}{:>8d}".format(p1.idx+1,p2.idx+1) )
                counter += 1
                if counter == 4:
                    fh.write("\n")
                    counter = 0
                else:
                    fh.write(" ")
            fh.write("\n" if counter == 0 else "\n\n")

            ## Write out angles
            angles = self.get_angles()
            fh.write("{:>8d} !NTHETA\n".format(len(angles)))
            counter = 0
            for p1,p2,p3,a in angles:
                fh.write( "{:>8d}{:>8d}{:>8d}".format(p1.idx+1,p2.idx+1,p3.idx+1) )
                counter += 1
                if counter == 3:
                    fh.write("\n")
                    counter = 0
                else:
                    fh.write(" ")
            fh.write("\n" if counter == 0 else "\n\n")

            ## Write out dihedrals
            dihedrals = self.get_dihedrals()
            fh.write("{:>8d} !NPHI\n".format(len(dihedrals)))
            counter = 0
            for p1,p2,p3,p4,a in dihedrals:
                fh.write( "{:>8d}{:>8d}{:>8d}{:>8d}".format(p1.idx+1,p2.idx+1,p3.idx+1,p4.idx+1) )
                counter += 1
                if counter == 2:
                    fh.write("\n")
                    counter = 0
                else:
                    fh.write(" ") 
            fh.write("\n" if counter == 0 else "\n\n")

            ## Write out impropers
            impropers = self.get_impropers()
            fh.write("{:>8d} !NIMPHI\n".format(len(impropers)))
            counter = 0
            for p1,p2,p3,p4,a in impropers:
                fh.write( "{:>8d}{:>8d}{:>8d}{:>8d}".format(p1.idx+1,p2.idx+1,p3.idx+1,p4.idx+1) )
                counter += 1
                if counter == 2:
                    fh.write("\n")
                    counter = 0
                else:
                    fh.write(" ")
            fh.write("\n" if counter == 0 else "\n\n")

            fh.write("\n{:>8d} !NDON: donors\n\n\n".format(0))
            fh.write("\n{:>8d} !NACC: acceptors\n\n\n".format(0))
            fh.write("\n       0 !NNB\n\n")
            natoms = len(self.particles)
            for i in range(natoms//8):
                fh.write("      0       0       0       0       0       0       0       0\n")
            for i in range(natoms-8*(natoms//8)):
                fh.write("      0")
            fh.write("\n\n       1       0 !NGRP\n\n")


class ArbdModel(PdbModel):
    def __init__(self, children, origin=None, dimensions=(1000,1000,1000), temperature=291, timestep=50e-6,
                 particle_integrator = 'Brown',
                 cutoff=50, decomp_period=1000, pairlist_distance=None, nonbonded_resolution=0.1,
                 remove_duplicate_bonded_terms=True, extra_bd_file_lines=""):

        PdbModel.__init__(self, children, dimensions, remove_duplicate_bonded_terms)
        self.origin = origin
        self.temperature = temperature

        self.timestep = timestep
        self.cutoff  =  cutoff

        self.particle_integrator = particle_integrator
        
        if pairlist_distance == None:
            pairlist_distance = cutoff+30
        
        self.decomp_period = decomp_period
        self.pairlist_distance = pairlist_distance

        self.extra_bd_file_lines = extra_bd_file_lines

        self.numParticles = 0
        self.particles = []
        self.type_counts = None

        self.nbSchemes = []
        self._nbParamFiles = [] # This could be made more robust
        self.nbResolution = 0.1

        self._written_bond_files = dict()        

        self.cacheUpToDate = False

    def clear_all(self, keep_children=False):
        Parent.clear_all(self, keep_children=keep_children)
        self.particles = []
        self.numParticles = 0
        self.type_counts = None
        self._nbParamFiles = []
        self._written_bond_files = dict()

    def _getNbScheme(self, typeA, typeB):
        scheme = None
        for s,A,B in self.nbSchemes:
            if A is None or B is None:
                if A is None and B is None:
                    return s
                elif A is None and typeB.is_same_type(B):
                    return s
                elif B is None and typeA.is_same_type(A):
                    return s
            elif typeA.is_same_type(A) and typeB.is_same_type(B):
                return s
        raise Exception("No nonbonded scheme found for %s and %s" % (typeA.name, typeB.name))

    def _countParticleTypes(self):
        ## TODO: check for modifications to particle that require
        ## automatic generation of new particle type
        type_counts = dict()
        for p in self:
            t = p.type_
            if t in type_counts:
                type_counts[t] += 1
            else:
                type_counts[t] = 1
        self.type_counts = type_counts

    def _updateParticleOrder(self):
        ## Create ordered list
        self.particles = [p for p in self]
        # self.particles = sorted(particles, key=lambda p: (p.type_, p.idx))
        
        ## Update particle indices
        for p,i in zip(self.particles,range(len(self.particles))):
            p.idx = i
            
        # self.initialCoords = np.array([p.initialPosition for p in self.particles])

    def useNonbondedScheme(self, nbScheme, typeA=None, typeB=None):
        self.add_nonbonded_scheme(nbScheme, typeA, typeB)

    def add_nonbonded_scheme(self, nonbonded_scheme, typeA=None, typeB=None):
        self.nbSchemes.append( (nonbonded_scheme, typeA, typeB) )
        if typeA != typeB:
            self.nbSchemes.append( (nonbonded_scheme, typeB, typeA) )

    def simulate(self, output_name, output_directory='output', num_steps=100000000, timestep=None, gpu=0, output_period=1e4, arbd=None, directory='.', restart_file=None, replicas=1, log_file=None, dry_run = False):
        assert(type(gpu) is int)
        num_steps = int(num_steps)

        d_orig = os.getcwd()
        try:
            if not os.path.exists(directory):
                os.makedirs(directory)
            os.chdir(directory)

            if timestep is not None:
                self.timestep = timestep

            if self.cacheUpToDate == False: # TODO: remove cache?
                self._countParticleTypes()
                self._updateParticleOrder()

            if output_directory == '': output_directory='.'

            if dry_run:
                if arbd is None: arbd='arbd'
            else:
                if arbd is None:
                    for path in os.environ["PATH"].split(os.pathsep):
                        path = path.strip('"')
                        fname = os.path.join(path, "arbd")
                        if os.path.isfile(fname) and os.access(fname, os.X_OK):
                            arbd = fname
                            break

                if arbd is None: raise Exception("ARBD was not found")

                if not os.path.exists(arbd):
                    raise Exception("ARBD was not found")
                if not os.path.isfile(arbd):
                    raise Exception("ARBD was not found")
                if not os.access(arbd, os.X_OK):
                    raise Exception("ARBD is not executable")


            if not os.path.exists(output_directory):
                os.makedirs(output_directory)
            elif not os.path.isdir(output_directory):
                raise Exception("output_directory '%s' is not a directory!" % output_directory)


            self.writePdb( output_name + ".pdb" )
            self.writePsf( output_name + ".psf" )
            self.writeArbdFiles( output_name, numSteps=num_steps, outputPeriod=output_period, restart_file=restart_file )
            # os.sync()

            ## http://stackoverflow.com/questions/18421757/live-output-from-subprocess-command

            cmd = [arbd, '-g', "%d" % gpu]
            if replicas > 1:
                cmd = cmd + ['-r',replicas]
            cmd = cmd + ["%s.bd" % output_name, "%s/%s" % (output_directory, output_name)]
            cmd = tuple(str(x) for x in cmd)

            if dry_run:
                print("Run with: %s" % " ".join(cmd))
            else:
                print("Running ARBD with: %s" % " ".join(cmd))
                if log_file is None or (hasattr(log_file,'write') and callable(log_file.write)):
                    fd = sys.stdout if log_file is None else log_file
                    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
                    for line in process.stdout:
                        try:
                            fd.write(line)
                        except:
                            print("WARNING: could not encode line; your locale might not be set correctly")
                        fd.flush()
                else:
                    with open(log_file,'w') as fd:
                        process = subprocess.Popen(cmd, stdout=log_file, universal_newlines=True)
                        process.communicate()

        except:
            raise
        finally:
            os.chdir(d_orig)


    # -------------------------- #
    # Methods for printing model #
    # -------------------------- #

    def writeArbdFiles(self, prefix, numSteps=100000000, outputPeriod=10000, restart_file=None):
        ## TODO: save and reference directories and prefixes using member data
        d = self.potential_directory = "potentials"
        if not os.path.exists(d):
            os.makedirs(d)
        self._restraint_filename = "%s/%s.restraint.txt" % (d, prefix)
        self._bond_filename = "%s/%s.bonds.txt" % (d, prefix)
        self._angle_filename = "%s/%s.angles.txt" % (d, prefix)
        self._dihedral_filename = "%s/%s.dihedrals.txt" % (d, prefix)
        self._exclusion_filename = "%s/%s.exculsions.txt" % (d, prefix)
        self._bond_angle_filename = "%s/%s.bondangles.txt" % (d, prefix)
        
        # self._writeArbdCoordFile( prefix + ".coord.txt" )
        self._writeArbdParticleFile( prefix + ".particles.txt" )
        self._writeArbdRestraintFile()
        self._writeArbdBondFile()
        self._writeArbdAngleFile()
        self._writeArbdDihedralFile()
        self._writeArbdExclusionFile()
        self._writeArbdBondAngleFile()
        self._writeArbdPotentialFiles( prefix, directory = d )
        self._writeArbdConf( prefix, numSteps=numSteps, outputPeriod=outputPeriod, restart_file=restart_file )
        
    # def _writeArbdCoordFile(self, filename):
    #     with open(filename,'w') as fh:
    #         for p in self.particles:
    #             fh.write("%f %f %f\n" % tuple(x for x in p.collapsedPosition()))

    def _writeArbdParticleFile(self, filename):
        with open(filename,'w') as fh:
            if self.particle_integrator == "Brown":
                for p in self.particles:
                    data = tuple([p.idx,p.type_.name] + [x for x in p.collapsedPosition()])
                    fh.write("ATOM %d %s %f %f %f\n" % data)
            else:
                for p in self.particles:
                    data = [p.idx,p.type_.name] + [x for x in p.collapsedPosition()]
                    try:
                        data = data + p.momentum
                    except:
                        try:
                            data = data + p.velocity*p.mass
                        except:
                            data = data + [0,0,0]
                    fh.write("ATOM %d %s %f %f %f %f %f %f\n" % tuple(data))
                
    def _write_rigid_group_file(self, filename, groups):
        with open(filename,'w') as fh:
            for g in groups:
                fh.write("#Group\n")
                try:
                    if len(g.trans_damping) != 3: raise
                    fh.write(" ".join(str(v) for v in g.trans_damping) + " ")
                except:
                    raise ValueError("Group {} lacks 3-value 'trans_damping' attribute")
                try:
                    if len(g.rot_damping) != 3: raise
                    fh.write(" ".join(str(v) for v in g.rot_damping) + " ")
                except:
                    raise ValueError("Group {} lacks 3-value 'rot_damping' attribute")
                fh.write("{}\n".format(len(g)))
                particles = [p for p in g]

                def chunks(l, n):
                    """Yield successive n-sized chunks from l."""
                    for i in range(0, len(l), n):
                        yield l[i:i + n]

                for c in chunks(particles,8):
                    fh.write(" ".join(str(p.idx) for p in c) + "\n")


    def _writeArbdConf(self, prefix, randomSeed=None, numSteps=100000000, outputPeriod=10000, restart_file=None):
        ## TODO: raise exception if _writeArbdPotentialFiles has not been called
        filename = "%s.bd" % prefix

        ## Prepare a dictionary to fill in placeholders in the configuration file
        params = self.__dict__.copy() # get parameters from System object

        if randomSeed is None:
            params['randomSeed']     = ""
        else:
            params['randomSeed'] = "seed %s" % randomSeed
        params['numSteps']       = int(numSteps)

        # params['coordinateFile'] = "%s.coord.txt" % prefix
        params['particleFile'] = "%s.particles.txt" % prefix
        if restart_file is None:
            params['restartCoordinates'] = ""
        else:
            params['restartCoordinates'] = "restartCoordinates %s" % restart_file
        params['outputPeriod'] = outputPeriod

        for k,v in zip('XYZ', self.dimensions):
            params['dim'+k] = v

        if self.origin is None:
            for k,v in zip('XYZ', self.dimensions):
                params['origin'+k] = -v*0.5
        else:
            for k,v in zip('XYZ', self.origin):
                params['origin'+k] = v
        
        params['pairlist_distance'] -= params['cutoff'] 

        """ Find rigid groups """
        rigid_groups = []
        def get_rigid_groups(parent):
            ret_list = []
            for c in parent.children:
                is_rigid = c.rigid if 'rigid' in c.__dict__ else False
                if is_rigid:
                    rigid_groups.append(c)
                elif isinstance(c,Group):
                    get_rigid_groups(c)
        get_rigid_groups(self)