motioncor processing added options for doseweighted stacks including split sums for cryoCARE.

processing/mc2FromMDOC.py

@@ -37,6 +37,7 @@ parser.add_argument('-cv', '--cudaVer', nargs='?', type=int, default=92, help='O
 parser.add_argument('-ps', '--patchSize', nargs='?', type=int, default=3, help='Override local alignment patch count. Default 3.')
 parser.add_argument('-b', '--binning', nargs='?', type=int, default=1, help='Output binning. Default 1 (no binning).')
 parser.add_argument('-dw', '--doseWeight', action='store_true', help='Do dose-weighting on alignment.')
+parser.add_argument('-ss', '--splitSum', action='store_true', help='Split sum saving for cryoCARE.')
 parser.add_argument('-v', '--verbose', action='store_true', help='Allow verbose output from motioncor2 during script run.')
 parser.add_argument('-nl', '--noLog', action='store_true', help='Allow verbose output from motioncor2 during script run.')
 parser.add_argument('-kf', '--keepFrames', action='store_true', help='Keep temporary individual aligned frames.')
@@ -120,13 +121,45 @@ try:
 except:
     print("Not a .mdoc file, skipping.")
     sys.exit()
-stackFileName = os.path.join(mdocDir, inBaseName + ".stackFiles")
+
+modeSuffixes = {'default': '',
+                'doseWeight': '_DW',
+                'even': '_EVN',
+                'odd': '_ODD'}
+
+outputModes = ['default']
+if args.doseWeight:
+    outputModes += ['doseWeight']
+    if args.splitSum:
+        outputModes += ['even', 'odd']
+
+stackFileNames = {}
+for key in modeSuffixes:
+    stackFileNames[key] = os.path.join(mdocDir, inBaseName + modeSuffixes[key] + ".stackFiles")
+
+# stackFileName = os.path.join(mdocDir, inBaseName + ".stackFiles")
+# stackFileDwName = os.path.join(mdocDir, inBaseName + "_DW.stackFiles")
+# stackFileEvnName = os.path.join(mdocDir, inBaseName + "_EVN.stackFiles")
+# stackFileOddName = os.path.join(mdocDir, inBaseName + "_ODD.stackFiles")
 tiltFileName = os.path.join(mdocDir, inBaseName + ".rawTlt")
 
-with open(stackFileName, "w+") as stackFile:
-    stackFile.write("%i\n" % ( len(stackParams) ))
+for outputMode in outputModes:
+    with open(stackFileNames[outputMode], "w+") as stackFile:
+        stackFile.write("%i\n" % (len(stackParams)))
         for imageParams in stackParams:
-        stackFile.write("%s\n0\n" % (os.path.join(alignedDir, imageParams['filename']) ))
+            stackFile.write("%s\n0\n" % (os.path.join(alignedDir, re.sub(".mrc$", modeSuffixes[outputMode] + ".mrc", imageParams['filename'])) ))
+
+# with open(stackFileName, "w+") as stackFile:
+#     stackFile.write("%i\n" % ( len(stackParams) ))
+#     for imageParams in stackParams:
+#         stackFile.write("%s\n0\n" % (os.path.join(alignedDir, imageParams['filename']) ))
+#
+# if args.doseWeight:
+#     with open(stackFileDwName, "w+") as stackFile:
+#         stackFile.write("%i\n" % ( len(stackParams) ))
+#         for imageParams in stackParams:
+#             stackFile.write("%s\n0\n" % (os.path.join(alignedDir, re.sub(".mrc$", "_DW.mrc", imageParams['filename'])) ))
+
 
 with open(tiltFileName, "w+") as tiltFile:
     for imageParams in stackParams:
@@ -152,19 +185,39 @@ with open(scriptFilePath, "w+") as scriptFile:
     for imageParams in stackParams:
         if args.doseWeight:
             doseWeightParams = " -Kv 300 -FmDose %f -InitDose %f" % (imageParams['frameDose'], imageParams['priorrecorddose'])
+            if args.splitSum:
+                doseWeightParams += " -SplitSum 1"
         else:
             doseWeightParams = ""
 
         scriptFile.write("echo Aligning frame stack %s \n" % (os.path.join(subframesDir, imageParams['filename']) ))
-        scriptFile.write("MotionCor2-cuda%i -InMrc \"%s\" -OutMrc \"%s\" -Patch %i %i -InFmMotion 0 -Iter 30 -Tilt %f -Group %i -FtBin %i -PixSize %f%s%s\n\n" %(args.cudaVer, os.path.join(subframesDir, imageParams['filename']), os.path.join(alignedDir, imageParams['filename']), args.patchSize, args.patchSize,  imageParams['tiltangle'], args.group, args.binning,  imageParams['pixelspacing']*args.pixelSizeFactor, doseWeightParams, verboseString) )
-    scriptFile.write("echo Creating combined tilt series stack %s\n" % (inBaseName + "_mc2.mrc"))
+        scriptFile.write("MotionCor2-cuda%i -InMrc \"%s\" -OutMrc \"%s\" -Patch %i %i \
+         -InFmMotion 0 -Iter 30 -Tilt %f -Group %i -FtBin %i -PixSize %f%s %s\n\n" %(args.cudaVer,
+                                                                                      os.path.join(subframesDir, imageParams['filename']),
+                                                                                      os.path.join(alignedDir, imageParams['filename']),
+                                                                                      args.patchSize,
+                                                                                      args.patchSize,
+                                                                                      imageParams['tiltangle'],
+                                                                                      args.group,
+                                                                                      args.binning,
+                                                                                      imageParams['pixelspacing']*args.pixelSizeFactor,
+                                                                                      doseWeightParams, verboseString))
+
     scriptFile.write("source /sw/apps/imod/current/imodenv\n")    
-    scriptFile.write("newstack -fileinlist %s -output %s -tilt %s%s\n\n" % ( os.path.basename(stackFileName), inBaseName + "_mc2.mrc", os.path.basename(tiltFileName), verboseString ))
+
+    for outputMode in outputModes:
+        scriptFile.write("echo Creating combined tilt series stack %s\n" % (inBaseName + "_mc2" + modeSuffixes[outputMode] + ".mrc"))
+        scriptFile.write("newstack -fileinlist %s -output %s -tilt %s%s\n\n" % (os.path.basename(stackFileNames[outputMode]),
+                                                                                inBaseName + "_mc2" + modeSuffixes[outputMode] + ".mrc",
+                                                                                os.path.basename(tiltFileName), verboseString))
+
     
     if not args.keepFrames:
         scriptFile.write("echo Removing temporary aligned files.\n")
+        modeSuffixes['doseWeight'] = '_DW*'
         for imageParams in stackParams:
-            scriptFile.write("rm %s\n" % (os.path.join(alignedDir, imageParams['filename'])))
+            for outputMode in outputModes:
+                scriptFile.write("rm %s\n" % (os.path.join(alignedDir, re.sub(".mrc$", modeSuffixes[outputMode] + ".mrc", imageParams['filename']))))
         scriptFile.write("rmdir %s\n" % (alignedDir))
 
 scriptFilePath.chmod(scriptFilePath.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP )