diff --git a/cluster/ICON_fromImodFiles.slurm b/cluster/ICON_fromImodFiles.slurm index ae07e4225ab140685ed905d5755efd0b17e1b007..19eb0b47eab032a006588832c46badfe398b40d5 100755 --- a/cluster/ICON_fromImodFiles.slurm +++ b/cluster/ICON_fromImodFiles.slurm @@ -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} diff --git a/processing/mc2FromMDOC.py b/processing/mc2FromMDOC.py index 7b4644cd58c521b740b013570eb5dce70951e11c..1246370ab1cb0581e0565e224cf5fff39ba23af2 100755 --- a/processing/mc2FromMDOC.py +++ b/processing/mc2FromMDOC.py @@ -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 )