run.py 7.96 KB
Newer Older
cmaffeo2's avatar
tmp  
cmaffeo2 committed
1
2
3
4
5
6
7
8
9
# -*- coding: utf-8 -*-
import json
import cadnano
from cadnano.document import Document
import numpy as np

import glob,re

from beadModel import beadModel
cmaffeo2's avatar
cmaffeo2 committed
10
from beadModelTwist import beadModelTwist
cmaffeo2's avatar
tmp  
cmaffeo2 committed
11
12
from atomicModel import atomicModel

13
import mdtraj as md
14
from coords import minimizeRmsd # mdtraj.superpose not working
15

16
17
import multiprocessing as mp

18
import pynvml as pynvml
cmaffeo2's avatar
cmaffeo2 committed
19
import sys
cmaffeo2's avatar
tmp  
cmaffeo2 committed
20
21

# twistLp = 1e-6
22
twistLp = 90
23

24
25
# arbd="/home/cmaffeo2/development/cuda/arbd.dev/src/arbd"
arbd="/home/cmaffeo2/development/cuda/arbd.dbg/src/arbd" # reduced the mem footprint cause vmd
26
namd="/home/cmaffeo2/development/namd-bin/NAMD_Git-2017-07-06_Linux-x86_64-multicore-CUDA/namd2"
27
# namd="/home/wilsja/namd/NAMD_CVS-2016-11-09_Linux-x86_64-multicore-CUDA/namd2"
28
# namd="/home/cmaffeo2/development/namd-bin/NAMD_2.12_Linux-x86_64-multicore-CUDA/namd2"
cmaffeo2's avatar
cmaffeo2 committed
29

30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def readSequenceFile(sequenceFile='resources/cadnano2pdb.seq.m13mp18.dat'):
    seq = []
    with open(sequenceFile) as ch:
        for l in ch:
            l = l.strip().replace(" ", "")
            if l[0] in (";","#"): continue
            seq.extend([c.upper() for c in l])
    return seq

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)

46
def readAvgArbdCoords(psf,pdb,dcd,rmsdThreshold=3.5):
47
    ## align trajectory to pdb
48
49
    ref = md.load(pdb, top=pdb)
    sel = md.load(dcd, top=pdb)
50
    # ref = sel[-1]
51
52
    ids = ref.topology.select("name =~ 'd.*'")
    assert(len(ids) > 3)
53

54
55
    # r0 = ref.xyz[0,ids,:]
    r0 = sel.xyz[-1,ids,:]
56
    t = -1                      # in case dcd_frames < 3
57
58
59
60
61
62
63
64
65
    for t in range(len(sel.xyz)-2,-1,-1):
        # print(t)
        R,comA,comB = minimizeRmsd(sel.xyz[t,ids,:],r0)
        sel.xyz[t,:,:] = np.array( [(r-comA).dot(R)+comB for r in sel.xyz[t]] )
        rmsd = np.mean( (sel.xyz[t,ids,:]-r0)**2 )
        if rmsd > (0.1*rmsdThreshold)**2:
            break
    t0=t+1
    print( "Averaging coordinates in %s after frame %d" % (dcd, t0) )
66

67
68
    pos = sel.xyz[t0:,:,:]
    pos = np.mean(pos, axis=0)
69
70
    return 10*pos               # convert to Angstroms

71
72
def simulateDesign(inFile, outPrefix, coarseSteps=1e7, fineSteps=1e7,
                   gpu=0, shiftEulerZ=0,
73
                   sequenceFile='resources/cadnano2pdb.seq.m13mp18.dat'):
74
75
    
    """ Read cadnano file """
cmaffeo2's avatar
cmaffeo2 committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
    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)

cmaffeo2's avatar
tmp  
cmaffeo2 committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
    try:
        doc = Document()
        cadnano.fileio.v3decode.decode(doc, data)
    except:
        doc = Document()
        cadnano.fileio.v2decode.decode(doc, data)

    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) ):
cmaffeo2's avatar
cmaffeo2 committed
108
            part.setVirtualHelixProperties(i,'eulerZ',t+shiftEulerZ)
cmaffeo2's avatar
tmp  
cmaffeo2 committed
109

110
    """ Create models """
111
112
    coarsestModelNoLongBonds = beadModel(part, twistPersistenceLength=twistLp, 
                                         maxBpsPerDNode=7, maxNtsPerSNode=4)
cmaffeo2's avatar
tmp  
cmaffeo2 committed
113
114
    coarsestModel = beadModel(part, twistPersistenceLength=twistLp, 
                            maxBpsPerDNode=7, maxNtsPerSNode=4)
cmaffeo2's avatar
cmaffeo2 committed
115
116
    finerModel = beadModelTwist(part, twistPersistenceLength=twistLp, 
                          maxBpsPerDNode=1, maxNtsPerSNode=1)
cmaffeo2's avatar
tmp  
cmaffeo2 committed
117

118
    finestModel = atomicModel(part, ntScale=0.25)
cmaffeo2's avatar
cmaffeo2 committed
119

120
121
    coarsestModelNoLongBonds._removeIntrahelicalConnectionsBeyond(100)
    coarsestModelNoLongBonds._removeCrossoversBeyond(100)
cmaffeo2's avatar
cmaffeo2 committed
122

123
    models = [coarsestModelNoLongBonds, coarsestModel, finerModel, finestModel]
124

125
    """ Examples of how models can be transformed """
cmaffeo2's avatar
cmaffeo2 committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
    ## 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
                    
    # for seg in finestModel.segments:
    #     for nt in seg:
    #         if nt.position[1] < 0:
    #             nt.position[0] -= 100
cmaffeo2's avatar
tmp  
cmaffeo2 committed
153

154
155
156
    """ Simulate and backmap each model """
    numSteps = [int(0.05*coarseSteps), int(0.95*coarseSteps), fineSteps]
    timestep = [200e-6, 200e-6, 50e-6]
157
158
    lastModel = None

159
160
161
    i = 0
    for i in range(len(models)-1):
        model = models[i]
162
163
        out = "%s-%d" % (outPrefix,i)
        lastOut = "%s-%d" % (outPrefix,i-1)
164

165
166
167
168
169
170
171
        model.writePsf(out + '.ideal.psf')
        model.writePdb(out + '.ideal.pdb')
        
        if lastModel is not None:
            if i == 1:
                coords = readArbdCoords( "output/%s.0.restart" % lastOut )
            else: 
172
                coords = readAvgArbdCoords("%s.psf" % lastOut, "%s.ideal.pdb" % lastOut,
173
174
175
176
177
                                           "output/%s.0.dcd" % lastOut )
            model.backmap( lastModel, coords,
                           dsDnaHelixNeighborDist=100, dsDnaAllNeighborDist=50,
                           ssDnaHelixNeighborDist=50)
                    
178
179
180
181
182
        model.simulate(out, outputDirectory = 'output',
                       arbd=arbd, gpu=gpu,
                       numSteps = numSteps[i], 
                       timestep = timestep[i])

183
184
185
186
187
        lastModel = model

    coords = readAvgArbdCoords(
        "%s.psf" % out, "%s.pdb" % out,
        "output/%s.0.dcd" % out )
cmaffeo2's avatar
cmaffeo2 committed
188

189
    out = outPrefix + "-%d" % (i+1)
190
191
192
193
    try:
        finestModel.setScaffoldSeq( readSequenceFile(sequenceFile) )
    except Exception as err:
        print("WARNING:", err)
cmaffeo2's avatar
cmaffeo2 committed
194
    finestModel.randomizeUnassignedNtSeqs()
195
    # finestModel.setUnassignedNtSeqs("T")
cmaffeo2's avatar
cmaffeo2 committed
196
    finestModel.writePdbPsf( out + ".ideal" )
197
    finestModel.backmap( lastModel, coords )
198
    finestModel.simulate( out, numSteps=480000, gpus=gpu, numprocs=8, namd=namd )
cmaffeo2's avatar
tmp  
cmaffeo2 committed
199

cmaffeo2's avatar
cmaffeo2 committed
200
201
202
203
if len(sys.argv) > 1:
    files = sys.argv[1:]
else:
    files = sorted(glob.glob('json/*.json'))
cmaffeo2's avatar
cmaffeo2 committed
204

205
def run(f):
206
    print("Working on %s with gpu %d" % (f,gpu))
cmaffeo2's avatar
cmaffeo2 committed
207
    out = re.match('json/(.*).json',f).group(1)
208
209
210
211
212
213
    oldOut = sys.stdout
    oldErr = sys.stderr
    with open("%s.log" % out, "w") as fh:
        sys.stdout = fh
        sys.stderr = fh
        try:
214
215
            simulateDesign(f, out, gpu = gpu)

216
217
218
219
220
221
222
        except Exception as err:
            print("WARNING: %s: failed to simulate design" % f)
            print("ERROR:", err)
        sys.stdout.flush()

    sys.stdout = oldOut
    sys.stderr = oldErr
223

224
225
226
227
228
pynvml.nvmlInit()
gpus = range(len(pynvml.nvmlDeviceGetCount()))
pynvml.nvmlShutdown()
# gpus = [0,1,2]
print(gpus)
229
230
231
232
233
234
235
236
237
238
239
240
241

if __name__ == '__main__':

    def init(queue):
        global gpu
        gpu = queue.get()

    idQueue = mp.Manager().Queue()
    for gpu in gpus:
        idQueue.put(gpu)
        
    pool = mp.Pool(len(gpus), init, (idQueue,))
    pool.map(run, files)