Skip to content
Snippets Groups Projects
beadModelTwist.py 55.6 KiB
Newer Older
cmaffeo2's avatar
cmaffeo2 committed
# -*- coding: utf-8 -*-
from datetime import datetime
from cadnano.cnenum import PointType
from math import pi,sqrt,exp,floor
import numpy as np
from scipy.special import erf
import scipy.optimize as opt
import os, sys, subprocess

import nbPot
from coords import minimizeRmsd, quaternionToMatrix3, rotationAboutAxis


class HarmonicPotential:
    def __init__(self, k, r0, rRange=(0,50), resolution=0.1, maxForce=None, maxPotential=None):
cmaffeo2's avatar
cmaffeo2 committed
        self.k = k
        self.r0 = r0
        self.maxForce = maxForce
        self.maxPotential = maxPotential
cmaffeo2's avatar
cmaffeo2 committed
        self.rRange = rRange
        self.resolution = resolution
cmaffeo2's avatar
cmaffeo2 committed
        self.periodic = False
        self.type = "None"
        self._kscale = None

    def filename(self, prefix='potentials/'):
cmaffeo2's avatar
cmaffeo2 committed
        # raise NotImplementedError("Not implemented")
        return "%s%s-%.3f-%.3f.dat" % (prefix, self.type,
cmaffeo2's avatar
cmaffeo2 committed
                                       self.k*self._kscale, self.r0)

    def write_file(self, prefix='potentials/'):
cmaffeo2's avatar
cmaffeo2 committed
        r = np.arange( self.rRange[0], 
                       self.rRange[1]+self.resolution, 
                       self.resolution )
        dr = r-self.r0

        if self.periodic == True:
            rSpan = self.rRange[1]-self.rRange[0]
            assert(rSpan > 0)
            dr = np.mod( dr+0.5*rSpan, rSpan) - 0.5*rSpan 

        u = 0.5*self.k*dr**2

        if self.maxForce is not None:
            assert(self.maxForce > 0)
            f = np.diff(u)/np.diff(r)
            f[f >  self.maxForce] =  self.maxForce
            f[f < -self.maxForce] = -self.maxForce            
cmaffeo2's avatar
cmaffeo2 committed
            u[0] = 0
            u[1:] = np.cumsum(f*np.diff(r))

        u = u - np.min(u)
        
        if self.maxPotential is not None:
            f = np.diff(u)/np.diff(r)
            ids = np.where( 0.5*(u[1:]+u[:-1]) > self.maxPotential )[0]

            w = np.sqrt(2*self.maxPotential/self.k)
            drAvg = 0.5*(np.abs(dr[ids]) + np.abs(dr[ids+1]))

            f[ids] = f[ids] * np.exp(-(drAvg-w)/(w))
            u[0] = 0
            u[1:] = np.cumsum(f*np.diff(r))

        u = u - np.min(u)

        np.savetxt( self.filename(prefix), np.array([r, u]).T, fmt="%f" )
cmaffeo2's avatar
cmaffeo2 committed

    def __hash__(self):
        assert(self.type != "None")
        return hash((self.type, self.k, self.r0, self.rRange, self.resolution, self.maxForce, self.maxPotential, self.periodic))
cmaffeo2's avatar
cmaffeo2 committed

    def __eq__(self, other):
        for a in ("type", "k", "r0", "rRange", "resolution", "maxForce", "maxPotential", "periodic"):
cmaffeo2's avatar
cmaffeo2 committed
            if self.__dict__[a] != other.__dict__[a]:
                return False
        return True

class NonBonded(HarmonicPotential):
    def __init__(self, k, r0, rRange=(0,50), resolution=0.1, maxForce=None, maxPotential=None):
        super().__init__(k,r0,rRange,resolution,maxForce,maxPotential)
cmaffeo2's avatar
cmaffeo2 committed
        self.type = "nonbonded"
        self._kscale = 1.0

class Bond(HarmonicPotential):
    def __init__(self, k, r0, rRange=(0,800), resolution=0.1, maxForce=5, maxPotential=None):
        super().__init__(k,r0,rRange,resolution,maxForce,maxPotential)
cmaffeo2's avatar
cmaffeo2 committed
        self.type = "bond"
        self._kscale = 1.0

class Angle(HarmonicPotential):
    def __init__(self, k, r0, rRange=(0,180), resolution=0.5, maxForce=None, maxPotential=None):
        super().__init__(k,r0,rRange,resolution,maxForce,maxPotential)
cmaffeo2's avatar
cmaffeo2 committed
        self.type = "angle"
        self._kscale = (180.0/pi)**2

class Dihedral(HarmonicPotential):
    def __init__(self, k, r0, rRange=(-180,180), resolution=1, maxForce=None, maxPotential=None):
        super().__init__(k,r0,rRange,resolution,maxForce,maxPotential)
cmaffeo2's avatar
cmaffeo2 committed
        self.periodic = True
        self.type = "dihedral"
        self._kscale = (180.0/pi)**2

class Node():
    def __init__(self, helix, pos, type="dsDNA"):
        self.helix    = helix
        self.position = np.array(pos)
        self.initialPosition = np.array(pos)
        self.type     = type
        self.nodeAbove = None
        self.nodeBelow = None
        self.xovers = []
        self.ssXovers = []

        self.orientationNode = None
        self.parentNode = None

        self.idx = helix.model.numParticles
        helix.model.numParticles += 1

    def addNodeAbove(self, node, separation):
        assert(self.nodeAbove is None)
        self.nodeAbove = node
        self.nodeAboveSep = separation # bp
    def addNodeBelow(self, node, separation):
        assert(self.nodeBelow is None)
        self.nodeBelow = node
        self.nodeBelowSep = separation # bp

    def addXover(self, node, fwds, double=False):
        ## TODO: what is meant by polarity?
        self.xovers.append( (node,fwds,double) )

    def addSsXover(self, node, fwds):
        self.ssXovers.append( (node,fwds) )
    
    def getNodesAbove(self,numNodes,inclusive=False):
        assert( type(numNodes) is int and numNodes > 0 )

        nodeList,sepList = [[],[]]
        n = self
        if inclusive:
            nodeList.append(n)

        for i in range(numNodes):
            if n.nodeAbove is None: break
            n = n.nodeAbove
            nodeList.append(n)
            sepList.append(n.nodeBelowSep)
            
        return nodeList,sepList

    def getNodesBelow(self,numNodes,inclusive=False):
        assert( type(numNodes) is int and numNodes > 0 )

        nodeList,sepList = [[],[]]
        n = self
        if inclusive:
            nodeList.append(n)

        for i in range(numNodes):
            if n.nodeBelow is None: break
            n = n.nodeBelow
            nodeList.append(n)
            sepList.append(n.nodeBelowSep)
            
        return nodeList,sepList

    def addOrientationNode(self, node):
        assert(self.nodeBelow is None)
        self.orientationNode = node
        node.parentNode = self

class helix():
    def __init__(self, model, part, hid):
        self.model = model
        self.props = part.getModelProperties().copy() # TODO: maybe move this out of here
        self.nodes = dict()
        self.orientationNodes = dict()

        self.hid = hid

        if self.props.get('point_type') == PointType.ARBITRARY:
            # TODO add code to encode Parts with ARBITRARY point configurations
            raise NotImplementedError("Not implemented")
        else:
            vh_props, origins = part.helixPropertiesAndOrigins()
            for x in vh_props:
                self.props[x] = vh_props[x][hid]

            self.origin = origins[hid]
            x,y = self.origin
            self.zIdxToPos = lambda idx: (x*10,y*10,-3.4*idx)

            ## get twizt
            keys = ['bases_per_repeat',
                    'turns_per_repeat',
                    'eulerZ','z']
            bpr,tpr,eulerZ,z = [vh_props[k][hid] for k in keys] 
            twist_per_base = tpr*360./bpr
Loading
Loading full blame...