Commit 9bca8552 authored by Adrian Nievergelt's avatar Adrian Nievergelt
Browse files

MotionCor2 processing script update to allow for wildcard mdoc input and...

MotionCor2 processing script update to allow for wildcard mdoc input and switch for line correction script.
parent 8c1392e8
......@@ -40,5 +40,5 @@ thickness=$((numSlicesFull/binFactor))
mkdir -p $baseDir
${iconPreproc} -i ${inputAli} -t ${tiltFile} -th ${specimenThickness} -o ${baseDir}_ipp.mrc
${iconGPU} -i ${baseDir}_ipp.mrc -t $2 -iter 80,160,80 -thr 0 -o ${baseDir} -s 0,${numSlices} -d 1
${iconGPU} -i ${baseDir}_ipp.mrc -t ${tiltFile} -iter 80,160,80 -thr 0 -o ${baseDir} -s 0,${numSlices} -d 1
${iconMask} -i ${baseDir}/reconstruction -t ${tiltFile} -th ${thickness} -cf ${baseDir}/crossValidation/crossV.frc -ff ${baseDir}/crossValidation/fullRec.frc -o ${baseDir}_icon.rec -s 0,${numSlices} -z ${zShift}
......@@ -22,6 +22,7 @@ rawDir = 'raw'
correctedDir = 'corrected'
alignedDir = 'aligned'
tiltStackDir = 'tiltStacks'
lineCorrectDir = 'lineCorrected'
subframesDir = None
correctAlignSLURM = "/projects/project-nieverge/scripts/cluster/correctAndAlign.slurm"
......@@ -29,7 +30,7 @@ correctAlignSLURM = "/projects/project-nieverge/scripts/cluster/correctAndAlign.
parser = argparse.ArgumentParser(description = "Generates batch processing scripts for aligning frames with MotionCor2 from MDOC files.")
# parser.add_argument("-f", "--file", help="create overlay for a single file",
parser.add_argument('-i', '--input', nargs='?', default='.', help='Input MDOC file to process.')
parser.add_argument('-i', '--input', nargs='+', default='.', help='Input MDOC file to process.')
# parser.add_argument('-d', '--outDir', nargs='?', default='.', help='Optional output directory to write scripts to.')
# parser.add_argument('-dp', '--dataPath', nargs='?', default='.', help='Optional override to use if the data is mounted in a different network base path.')
# parser.add_argument('-cp', '--clusterPath', nargs='?', default='/projects/project-nieverge/', help='Destination path of the HPC filesystem data directory.')
......@@ -38,6 +39,7 @@ parser.add_argument('-ps', '--patchSize', nargs='?', type=int, default=3, help='
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('-cl', '--correctLines', action='store_true', help='Use correctK2Lines.py script for correcting lines due to drifty camera sensor.')
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.')
......@@ -52,172 +54,188 @@ args = parser.parse_args()
def FakeSecHead(fp): yield '[general]\n'; yield from fp
if args.input == '.':
print("Missing input file");
print(parser.print_help())
sys.exit()
for inputFile in args.input:
# process all files passed by the input parameter sequentially
parser = configparser.ConfigParser()
if inputFile == '.':
print("Missing input file");
print(parser.print_help())
sys.exit()
with open(args.input) as mdocFile:
parser.readfp(FakeSecHead(mdocFile))
parser = configparser.ConfigParser()
# getFilename = re.compile()
with open(inputFile) as mdocFile:
parser.readfp(FakeSecHead(mdocFile))
stackParams = []
# getFilename = re.compile()
stackParams = []
for section in parser.sections():
if not section.startswith("ZValue"):
continue
tempDict = {}
for option in parser.options(section):
tempDict[option] = parser.get(section, option)
for section in parser.sections():
if not section.startswith("ZValue"):
continue
tempDict = {}
for option in parser.options(section):
tempDict[option] = parser.get(section, option)
try:
tempDict['filename'] = re.search("\\\\([\w.-]+)$", tempDict['subframepath']).group(1)
except:
print("Section \"%s\" has no subframes, skipping." % (section))
continue
for imageParam in ['countsperelectron', 'defocus', 'exposuredose', 'exposuretime', 'intensity', 'pixelspacing', 'priorrecorddose', 'rotationangle', 'stagez', 'targetdefocus', 'tiltangle']:
try:
tempDict[imageParam] = float(tempDict[imageParam])
except:
tempDict[imageParam] = 0
for imageParam in ['binning', 'cameraindex', 'dividedby2', 'magindex', 'magnification', 'numsubframes', 'spotsize']:
try:
tempDict[imageParam] = int(tempDict[imageParam])
tempDict['filename'] = re.search("\\\\([\w.-]+)$", tempDict['subframepath']).group(1)
except:
tempDict[imageParam] = 0
tempDict['frameDose'] = float(tempDict['framedosesandnumber'].split()[0])
stackParams.append(tempDict)
# Data transfer to cluster
inPath = os.path.splitext(args.input)
# print(inPath)
# pprint(stackParams[0:2])
# find subframes directory
mdocDir = os.path.dirname(args.input)
if mdocDir == "":
mdocDir = "."
for (dirpath, dirnames, filenames) in os.walk(mdocDir):
for filename in filenames:
if filename == stackParams[0]['filename'] and not dirpath.endswith(alignedDir):
subframesDir = os.path.relpath(dirpath, mdocDir)
print("Found subframes in directory \"%s\"" % (subframesDir))
break
if not subframesDir:
print("No subframes found in current directory tree, terminating!")
sys.exit()
stackParams = sorted(stackParams, key = lambda i: i['tiltangle'])
try:
inBaseName = re.search("([^\/]+)\.\w+.mdoc$", args.input).group(1)
except:
print("Not a .mdoc file, skipping.")
sys.exit()
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")
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, 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:
tiltFile.write("%f\n" % (imageParams['tiltangle']) )
if args.verbose:
verboseString = ""
elif args.noLog:
verboseString = " &>/dev/null"
else:
verboseString = ">> %s 2>&1" % (inBaseName + "_alignment.log")
scriptFilePath = Path(os.path.join(mdocDir, inBaseName + "_alignFramesMC2.sh"))
with open(scriptFilePath, "w+") as scriptFile:
scriptFile.write("#!/bin/bash\n\n")
scriptFile.write("# Script file for motioncor2 alignment of tilt series created by mc2FromMDOC.py on %s by %s\n\n" % (datetime.today().strftime('%Y-%m-%d %H:%M'), getpass.getuser()) )
scriptFile.write('cd "$(dirname "$0")"\n')
scriptFile.write("mkdir -p %s\n" % (alignedDir))
scriptFile.write("rm -f %s\n\n" % (inBaseName + "_alignment.log"))
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("source /sw/apps/imod/current/imodenv\n")
print("Section \"%s\" has no subframes, skipping." % (section))
continue
for imageParam in ['countsperelectron', 'defocus', 'exposuredose', 'exposuretime', 'intensity', 'pixelspacing', 'priorrecorddose', 'rotationangle', 'stagez', 'targetdefocus', 'tiltangle']:
try:
tempDict[imageParam] = float(tempDict[imageParam])
except:
tempDict[imageParam] = 0
for imageParam in ['binning', 'cameraindex', 'dividedby2', 'magindex', 'magnification', 'numsubframes', 'spotsize']:
try:
tempDict[imageParam] = int(tempDict[imageParam])
except:
tempDict[imageParam] = 0
tempDict['frameDose'] = float(tempDict['framedosesandnumber'].split()[0])
stackParams.append(tempDict)
# Data transfer to cluster
inPath = os.path.splitext(inputFile)
# print(inPath)
# pprint(stackParams[0:2])
# find subframes directory
mdocDir = os.path.dirname(inputFile)
if mdocDir == "":
mdocDir = "."
for (dirpath, dirnames, filenames) in os.walk(mdocDir):
for filename in filenames:
if filename == stackParams[0]['filename'] and not dirpath.endswith(alignedDir) and not dirpath.endswith(lineCorrectDir):
subframesDir = os.path.relpath(dirpath, mdocDir)
print("Found subframes in directory \"%s\"" % (subframesDir))
break
if not subframesDir:
print("No subframes found in current directory tree, terminating!")
sys.exit()
stackParams = sorted(stackParams, key = lambda i: i['tiltangle'])
try:
inBaseName = re.search("([^\/]+)\.\w+.mdoc$", inputFile).group(1)
except:
print("Not a .mdoc file, skipping.")
sys.exit()
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")
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*'
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, 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:
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))
tiltFile.write("%f\n" % (imageParams['tiltangle']) )
if args.verbose:
verboseString = ""
elif args.noLog:
verboseString = " &>/dev/null"
else:
verboseString = ">> %s 2>&1" % (inBaseName + "_alignment.log")
scriptFilePath = Path(os.path.join(mdocDir, inBaseName + "_alignFramesMC2.sh"))
scriptFilePath.chmod(scriptFilePath.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP )
with open(scriptFilePath, "w+") as scriptFile:
scriptFile.write("#!/bin/bash\n\n")
scriptFile.write("# Script file for motioncor2 alignment of tilt series created by mc2FromMDOC.py on %s by %s\n\n" % (datetime.today().strftime('%Y-%m-%d %H:%M'), getpass.getuser()) )
scriptFile.write('cd "$(dirname "$0")"\n')
scriptFile.write("mkdir -p %s\n" % (alignedDir))
scriptFile.write("rm -f %s\n\n" % (inBaseName + "_alignment.log"))
if args.correctLines:
scriptFile.write("mkdir -p %s\n" % (lineCorrectDir))
for imageParams in stackParams:
if args.correctLines:
scriptFile.write(
"echo Correcting lines in frame stack %s \n" % (os.path.join(subframesDir, imageParams['filename'])))
scriptFile.write(
"correctK2Lines.py -i %s -d %s -c 0.7 \n" % (
os.path.join(subframesDir, imageParams['filename']),
lineCorrectDir))
sourceDir = lineCorrectDir
else:
sourceDir = subframesDir
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(sourceDir, 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(sourceDir, 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")
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:
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 )
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment