Skip to content
Snippets Groups Projects
run.py 6.11 KiB
# -*- coding: utf-8 -*-
import json
import cadnano
from cadnano.document import Document
import numpy as np

import glob,re

from beadModel import beadModel
from beadModelTwist import beadModelTwist
from atomicModel import atomicModel

import sys

# twistLp = 1e-6
twistLp = 90
def simulateDesign(inFile, outPrefix, gpu=0, shiftEulerZ=0):
    # ----------------- #
    # Read cadnano file #
    # ----------------- #

    try:
        with open(inFile) as ch:
            data = json.load(ch)
    except:
        with open(inFile) as ch:
            ## TODO: han
            content = ""
            for l in ch:
                l = re.sub(r"'", r'"', l)
                # https://stackoverflow.com/questions/4033633/handling-lazy-json-in-python-expecting-property-name
                # l = re.sub(r"{\s*(\w)", r'{"\1', l)
                # l = re.sub(r",\s*(\w)", r',"\1', l)
                # l = re.sub(r"(\w):", r'\1":', l)
                content += l+"\n"
            data = json.loads(content)

    try:
        doc = Document()
        cadnano.fileio.v3decode.decode(doc, data)
    except:
        doc = Document()
        cadnano.fileio.v2decode.decode(doc, data)


    def readArbdCoords(fname):
        coords = []
        with open(fname) as fh:
            for line in fh:
                coords.append([float(x) for x in line.split()[1:]])
        return np.array(coords)


    parts = [p for p in doc.getParts()]
    if len(parts) != 1:
        raise Exception("Only documents containing a single 'part' are implemented at this time.")
    part = parts[0]

    if shiftEulerZ != 0:
        ## Manually fix eulerZ, which is improperly set by cadnano2.5 for some structures
        props, origins = part.helixPropertiesAndOrigins()
        for t,i in zip( props['eulerZ'], range(part.getIdNumMax()+1) ):
            part.setVirtualHelixProperties(i,'eulerZ',t+shiftEulerZ)


    # ------------- #
    # Create models #
    # ------------- #

    coarsestModel = beadModel(part, twistPersistenceLength=twistLp, 
                            maxBpsPerDNode=7, maxNtsPerSNode=4)
    # return coarsestModel._writeArbdDihedralFile(outPrefix+"-1")
    coarseModel = beadModel(part, twistPersistenceLength=twistLp, 
                            maxBpsPerDNode=4, maxNtsPerSNode=2)
    fineModel = beadModel(part, twistPersistenceLength=twistLp, 
                          maxBpsPerDNode=1, maxNtsPerSNode=1)
    finerModel = beadModelTwist(part, twistPersistenceLength=twistLp, 
                          maxBpsPerDNode=1, maxNtsPerSNode=1)

    finestModel = atomicModel(part, ntScale=1.0)


    models = [coarsestModel, coarseModel, fineModel, finerModel, finestModel]
    # for m in models:
    #     m._removeIntrahelicalConnectionsBeyond(800)
    # for m in models[:3]:
    #     m._removeCrossoversBeyond(800)
    
    ## Custom adjust the z coordinates
    # com1 = np.array([-36.7, 42.2, -431.5])
    # M1 = np.array(((0,-1,0),(1,0,0),(0,0,1)))

    # com2 = np.array([-36.7, 177.0, -429.0])
    # M2 = np.array(((0,1,0),(-1,0,0),(0,0,1)))
    
    # for model in models[:-1]:
    #     beads = [b for h in model for i,b in h[1].nodes.items()]
    #     for h in model:
    #         for i,b in h[1].nodes.items():
    #             if np.abs(b.position[1]-50) < 80:
    #                 r = b.position
    #                 r = np.inner(M1,r-com1) + com1
    #                 r[0] -= 200
    #                 b.position = r
    #                 b.initialPosition = r
    #             elif b.position[1] > 120:
    #                 r = b.position
    #                 r = np.inner(M2,r-com2) + com2
    #                 b.position = r
    #                 b.initialPosition = r
                    

    # # coarsestModel, coarseModel, fineModel = [moveBeads(m) for m in models[:3]]
    # # finestModel = models[3]

    # for seg in finestModel.segments:
    #     for nt in seg:
    #         if nt.position[1] < 0:
    #             nt.position[0] -= 100

    # ------------------------------- #
    # Simulate and backmap each model #
    # ------------------------------- #

    out = outPrefix + "-1"
    coarsestModel.simulate(out, outputDirectory='output', numSteps=15000000, timestep=200e-6,
                       arbd="/home/cmaffeo2/development/cuda/arbd.dev/src/arbd", gpu=gpu)
    coarsestModelCoords = readArbdCoords( "output/%s.0.restart" % out )

    out = outPrefix + "-2"
    coarseModel.backmap( coarsestModel, coarsestModelCoords, 
                         dsDnaHelixNeighborDist=100, dsDnaAllNeighborDist=50,
                         ssDnaHelixNeighborDist=50)
    coarseModel.simulate( out, outputDirectory='output', numSteps=1000000,
                        arbd="/home/cmaffeo2/development/cuda/arbd.dev/src/arbd", gpu=gpu)
    coarseModelCoords = readArbdCoords( "output/%s.0.restart" % out )

    out = outPrefix + "-3"
    fineModel.backmap( coarseModel, coarseModelCoords,
                       dsDnaHelixNeighborDist=100, dsDnaAllNeighborDist=50,
                       ssDnaHelixNeighborDist=50)
    fineModel.simulate( out,  outputDirectory='output', numSteps=1000000,
                        arbd="/home/cmaffeo2/development/cuda/arbd.dev/src/arbd", gpu=gpu)
    fineModelCoords = readArbdCoords( "output/%s.0.restart" % out )

    out = outPrefix + "-4"
    finerModel.backmap( fineModel, fineModelCoords,
                        dsDnaHelixNeighborDist=100, dsDnaAllNeighborDist=50,
                        ssDnaHelixNeighborDist=50)
    finerModel.simulate( out,  outputDirectory='output', numSteps=1000000,
                        arbd="/home/cmaffeo2/development/cuda/arbd.dev/src/arbd", gpu=gpu)
    finerModelCoords = readArbdCoords( "output/%s.0.restart" % out )

    out = outPrefix + "-5"
    finestModel.randomizeUnassignedNtSeqs()
    finestModel.writePdbPsf( out + ".ideal" )
    finestModel.backmap( finerModel, finerModelCoords )
    finestModel.writePdbPsf( out )


# def simulateDesign(f,out,gpu=0,shiftEulerZ=180):
#     print(out)


if len(sys.argv) > 1:
    files = sys.argv[1:]
else:
    files = sorted(glob.glob('json/*.json'))

print(files)

gpu=0
for f in files[:]:
    # f = files
    print(f)
    out = re.match('json/(.*).json',f).group(1)
    simulateDesign(f,out,gpu=gpu)