Commit 2e1b1332 authored by cmaffeo2's avatar cmaffeo2
Browse files

Merge branch 'dev'

parents bb71bbfb d3c116bf
../LICENSE
\ No newline at end of file
../README.md
\ No newline at end of file
# -*- 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