Updated run script to remove unneeded simulation steps and make...

Updated run script to remove unneeded simulation steps and make fine-resolution model run longer; also added dependency on pynvml to query gpus

run.py

@@ -15,6 +15,7 @@ from coords import minimizeRmsd # mdtraj.superpose not working
 
 import multiprocessing as mp
 
+import pynvml as pynvml
 import sys
 
 # twistLp = 1e-6
@@ -67,12 +68,11 @@ def readAvgArbdCoords(psf,pdb,dcd,rmsdThreshold=3.5):
     pos = np.mean(pos, axis=0)
     return 10*pos               # convert to Angstroms
 
-def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEulerZ=0,
+def simulateDesign(inFile, outPrefix, coarseSteps=1e7, fineSteps=1e7,
+                   gpu=0, shiftEulerZ=0,
                    sequenceFile='resources/cadnano2pdb.seq.m13mp18.dat'):
-    # ----------------- #
-    # Read cadnano file #
-    # ----------------- #
-
+    
+    """ Read cadnano file """
     try:
         with open(inFile) as ch:
             data = json.load(ch)
@@ -96,7 +96,6 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
         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.")
@@ -108,19 +107,11 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
         for t,i in zip( props['eulerZ'], range(part.getIdNumMax()+1) ):
             part.setVirtualHelixProperties(i,'eulerZ',t+shiftEulerZ)
 
-
-    # ------------- #
-    # Create models #
-    # ------------- #
-
+    """ Create models """
     coarsestModelNoLongBonds = beadModel(part, twistPersistenceLength=twistLp, 
                                          maxBpsPerDNode=7, maxNtsPerSNode=4)
     coarsestModel = beadModel(part, twistPersistenceLength=twistLp, 
                             maxBpsPerDNode=7, maxNtsPerSNode=4)
-    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)
 
@@ -129,13 +120,9 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
     coarsestModelNoLongBonds._removeIntrahelicalConnectionsBeyond(100)
     coarsestModelNoLongBonds._removeCrossoversBeyond(100)
 
+    models = [coarsestModelNoLongBonds, coarsestModel, finerModel, finestModel]
 
-    models = [coarsestModelNoLongBonds, coarsestModel, coarseModel, fineModel, finerModel, finestModel]
-    # for m in models:
-    #     m._removeIntrahelicalConnectionsBeyond(800)
-    # for m in models[:3]:
-    #     m._removeCrossoversBeyond(800)
-    
+    """ Examples of how models can be transformed """
     ## 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)))
@@ -159,37 +146,22 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
     #                 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 #
-    # ------------------------------- #
-
-    stepId = [i for i in range(5)]
-    numSteps = [int(0.05*numCoarsestSteps), int(0.95*numCoarsestSteps)]
-    timestep = [200e-6, 200e-6]
+    """ Simulate and backmap each model """
+    numSteps = [int(0.05*coarseSteps), int(0.95*coarseSteps), fineSteps]
+    timestep = [200e-6, 200e-6, 50e-6]
     lastModel = None
 
-    for i in stepId:
-        args = dict(arbd=arbd, gpu=gpu, numSteps=10000,
-                    outputDirectory = 'output')
-
-        if i < len(numSteps): args['numSteps'] = numSteps[i]
-        if i < len(numSteps): args['timestep'] = timestep[i]
-        if i == 4: 
-            args['numSteps'] = 1000000
-            args['timestep'] = timestep=50e-6
-
+    i = 0
+    for i in range(len(models)-1):
+        model = models[i]
         out = "%s-%d" % (outPrefix,i)
         lastOut = "%s-%d" % (outPrefix,i-1)
-        model = models[i]
+
         model.writePsf(out + '.ideal.psf')
         model.writePdb(out + '.ideal.pdb')
         
@@ -203,14 +175,18 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
                            dsDnaHelixNeighborDist=100, dsDnaAllNeighborDist=50,
                            ssDnaHelixNeighborDist=50)
                     
-        model.simulate(out, **args)
+        model.simulate(out, outputDirectory = 'output',
+                       arbd=arbd, gpu=gpu,
+                       numSteps = numSteps[i], 
+                       timestep = timestep[i])
+
         lastModel = model
 
     coords = readAvgArbdCoords(
         "%s.psf" % out, "%s.pdb" % out,
         "output/%s.0.dcd" % out )
 
-    out = outPrefix + "-5"
+    out = outPrefix + "-%d" % (i+1)
     try:
         finestModel.setScaffoldSeq( readSequenceFile(sequenceFile) )
     except Exception as err:
@@ -219,7 +195,7 @@ def simulateDesign(inFile, outPrefix, numCoarsestSteps=5000000, gpu=0, shiftEule
     # finestModel.setUnassignedNtSeqs("T")
     finestModel.writePdbPsf( out + ".ideal" )
     finestModel.backmap( lastModel, coords )
-    finestModel.simulate( out, numSteps=480000, gpus=gpu, numprocs=4, namd=namd )
+    finestModel.simulate( out, numSteps=480000, gpus=gpu, numprocs=8, namd=namd )
 
 if len(sys.argv) > 1:
     files = sys.argv[1:]
@@ -235,7 +211,8 @@ def run(f):
         sys.stdout = fh
         sys.stderr = fh
         try:
-            simulateDesign(f, out, numCoarsestSteps=10000000, gpu=gpu)
+            simulateDesign(f, out, gpu = gpu)
+
         except Exception as err:
             print("WARNING: %s: failed to simulate design" % f)
             print("ERROR:", err)
@@ -243,10 +220,12 @@ def run(f):
 
     sys.stdout = oldOut
     sys.stderr = oldErr
-                    
 
-gpus = [0,1,2]
-# gpus = [2]
+pynvml.nvmlInit()
+gpus = range(len(pynvml.nvmlDeviceGetCount()))
+pynvml.nvmlShutdown()
+# gpus = [0,1,2]
+print(gpus)
 
 if __name__ == '__main__':