Merge branch 'dev'

mrdna/arbdmodel/LICENSE

+../LICENSE
\ No newline at end of file

mrdna/arbdmodel/README.md

+../README.md
\ No newline at end of file

mrdna/arbdmodel/__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.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 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
+
+    ## 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 """
+
+        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)))