Commit f2f91983 authored by Adrian Nievergelt's avatar Adrian Nievergelt

added script to create local motioncor2 series alignment from mdoc files

parent 38467e1a
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 15 09:22:40 2019
@author: adrian
"""
import configparser
import os, sys
from pprint import pprint
import re
import argparse
from datetime import datetime
import getpass
# clusterBaseDir = '/projects/project-nieverge/'
procDir = 'tomoProc'
rawDir = 'raw'
correctedDir = 'corrected'
alignedDir = 'aligned'
tiltStackDir = 'tiltStacks'
subframesDir = 'subframes'
correctAlignSLURM = "/projects/project-nieverge/scripts/cluster/correctAndAlign.slurm"
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('-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.')
parser.add_argument('-cv', '--cudaVer', nargs='?', type=int, default=92, help='Override CUDA version. Default 9.2')
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('-v', '--verbose', 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.')
# parser.add_argument('-c', '--cutoff', nargs='?', type=float, default='1', help='The threshold of how many standard deviations above or below the a line must deviate before it is corrected')
# parser.add_argument('-g', '--graph', action='store_true', help='Store a graph of the corrected and uncorrected line intensities for debug purposes.')
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()
parser = configparser.ConfigParser()
with open(args.input) as mdocFile:
parser.readfp(FakeSecHead(mdocFile))
# 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)
tempDict['filename'] = re.search("\\\\([\w.-]+)$", tempDict['subframepath']).group(1)
for imageParam in ['countsperelectron', 'defocus', 'exposuredose', 'exposuretime', 'intensity', 'pixelspacing', 'priorrecorddose', 'rotationangle', 'stagez', 'targetdefocus', 'tiltangle']:
tempDict[imageParam] = float(tempDict[imageParam])
for imageParam in ['binning', 'cameraindex', 'dividedby2', 'magindex', 'magnification', 'numsubframes', 'spotsize']:
tempDict[imageParam] = int(tempDict[imageParam])
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)
for (dirpath, dirnames, filenames) in os.walk(mdocDir):
for filename in filenames:
if filename == stackParams[0]['filename']:
subframesDir = os.path.relpath(dirpath, mdocDir)
print("Found subframes in directory \"%s\"" % (subframesDir))
break
stackParams = sorted(stackParams, key = lambda i: i['tiltangle'])
inBaseName = re.search("(\w+)\..*", args.input).group(1)
stackFileName = os.path.join(mdocDir, inBaseName + ".stackFiles")
tiltFileName = os.path.join(mdocDir, inBaseName + ".rawTlt")
with open(stackFileName, "w+") as stackFile:
for imageParams in stackParams:
stackFile.write("%s\n" % (os.path.join(alignedDir, imageParams['filename']) ))
with open(tiltFileName, "w+") as tiltFile:
for imageParams in stackParams:
tiltFile.write("%f\n" % (imageParams['tiltangle']) )
with open(os.path.join(mdocDir, inBaseName + "_alignFramesMC2.sh"), "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\n" % (alignedDir))
for imageParams in stackParams:
if args.doseWeight:
doseWeightParams = " -Kv 300 -FmDose %f -InitDose %f" % (imageParams['frameDose'], imageParams['priorrecorddose'])
else:
doseWeightParams = ""
if args.verbose:
verboseString = ""
else:
verboseString = " &>/dev/null"
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 1 -Iter 30 -Tilt %f -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.binning, imageParams['pixelspacing'], doseWeightParams, verboseString) )
scriptFile.write("echo Creating combined tilt series stack %s\n" % (inBaseName + "_mc2.mrc"))
scriptFile.write("newstack -fileinlist %s -output %s -tilt %s\n\n" % ( os.path.basename(stackFileName), inBaseName + "_mc2.mrc", os.path.basename(tiltFileName) ))
if not args.keepFrames:
scriptFile.write("echo Removing temporary aligned files.\n")
for imageParams in stackParams:
scriptFile.write("rm %s\n" % (os.path.join(alignedDir, imageParams['filename'])))
scriptFile.write("rmdir %s\n" % (alignedDir))
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