mc2FromMDOC.py 10.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
#!/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
16 17
import stat
from pathlib import Path
18 19 20 21 22 23 24

# clusterBaseDir = '/projects/project-nieverge/'
procDir = 'tomoProc'
rawDir = 'raw'
correctedDir = 'corrected'
alignedDir = 'aligned'
tiltStackDir = 'tiltStacks'
25
subframesDir = None
26 27 28 29 30 31 32
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.')
33 34 35
# 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.')
Adrian Nievergelt's avatar
Adrian Nievergelt committed
36
parser.add_argument('-cv', '--cudaVer', nargs='?', type=int, default=92, help='Override CUDA version. Default 92. Reduce to a lower value if MotionCor2 does not run properly with the installed graphics card.')
37 38 39
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.')
40
parser.add_argument('-ss', '--splitSum', action='store_true', help='Split sum saving for cryoCARE.')
41
parser.add_argument('-v', '--verbose', action='store_true', help='Allow verbose output from motioncor2 during script run.')
42
parser.add_argument('-nl', '--noLog', action='store_true', help='Allow verbose output from motioncor2 during script run.')
43
parser.add_argument('-kf', '--keepFrames', action='store_true', help='Keep temporary individual aligned frames.')
Adrian Nievergelt's avatar
Adrian Nievergelt committed
44
parser.add_argument('-psf', '--pixelSizeFactor', nargs='?', type=float, default=1, help='Correction factor for pixel size, set to 0.5 for super resolution with binned mdoc.')
45
parser.add_argument('-g', '--group', nargs='?', type=int, default=2, help='Group parameter to motioncor2. Default 2, set to 3-4 on very low signal input.')
46
#parser.add_argument('-g', '--group', nargs='?', type=int, default=1, help='Group parameter to motioncor2. Default 1, set to 2 or 3 on very low signal input.')
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76

# 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)
        
77 78 79 80 81
    try:
        tempDict['filename'] = re.search("\\\\([\w.-]+)$", tempDict['subframepath']).group(1)
    except:
        print("Section \"%s\" has no subframes, skipping." % (section))
        continue
82
    for imageParam in ['countsperelectron', 'defocus', 'exposuredose', 'exposuretime', 'intensity', 'pixelspacing', 'priorrecorddose', 'rotationangle', 'stagez', 'targetdefocus', 'tiltangle']:
83 84 85 86
        try:
            tempDict[imageParam] = float(tempDict[imageParam])
        except:
            tempDict[imageParam] = 0
87
    for imageParam in ['binning', 'cameraindex', 'dividedby2', 'magindex', 'magnification', 'numsubframes', 'spotsize']:
88 89 90 91 92
        try:
            tempDict[imageParam] = int(tempDict[imageParam])
        except:
            tempDict[imageParam] = 0
            
93 94 95 96 97 98 99 100 101 102 103 104
    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)
105 106
if mdocDir == "":
    mdocDir = "."
107 108
for (dirpath, dirnames, filenames) in os.walk(mdocDir):
    for filename in filenames:
109
        if filename == stackParams[0]['filename'] and not dirpath.endswith(alignedDir):
110 111 112 113
            subframesDir = os.path.relpath(dirpath, mdocDir)
            print("Found subframes in directory \"%s\"" % (subframesDir))
            break

114 115 116 117
if not subframesDir:
    print("No subframes found in current directory tree, terminating!")
    sys.exit()

Adrian Nievergelt's avatar
Adrian Nievergelt committed
118 119 120 121 122 123
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()
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143

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")
144 145
tiltFileName = os.path.join(mdocDir, inBaseName + ".rawTlt")

146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
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'])) ))

163 164 165 166 167

with open(tiltFileName, "w+") as tiltFile:
    for imageParams in stackParams:
        tiltFile.write("%f\n" % (imageParams['tiltangle']) )

Adrian Nievergelt's avatar
Adrian Nievergelt committed
168 169
if args.verbose:
    verboseString = ""
170
elif args.noLog:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
171
    verboseString = " &>/dev/null"
172
else:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
173
    verboseString = ">> %s 2>&1" % (inBaseName + "_alignment.log")
Adrian Nievergelt's avatar
Adrian Nievergelt committed
174

175 176 177
scriptFilePath = Path(os.path.join(mdocDir, inBaseName + "_alignFramesMC2.sh"))

with open(scriptFilePath, "w+") as scriptFile:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
178
    scriptFile.write("#!/bin/bash\n\n")
179 180 181
    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')                 
Adrian Nievergelt's avatar
Adrian Nievergelt committed
182 183
    scriptFile.write("mkdir -p %s\n" % (alignedDir))
    scriptFile.write("rm -f %s\n\n" % (inBaseName + "_alignment.log"))
184 185 186 187
    
    for imageParams in stackParams:
        if args.doseWeight:
            doseWeightParams = " -Kv 300 -FmDose %f -InitDose %f" % (imageParams['frameDose'], imageParams['priorrecorddose'])
188 189
            if args.splitSum:
                doseWeightParams += " -SplitSum 1"
190 191
        else:
            doseWeightParams = ""
192

193
        scriptFile.write("echo Aligning frame stack %s \n" % (os.path.join(subframesDir, imageParams['filename']) ))
194 195 196 197 198 199 200 201 202 203 204 205
        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))

206
    scriptFile.write("source /sw/apps/imod/current/imodenv\n")    
207 208 209 210 211 212 213

    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))

214 215 216
    
    if not args.keepFrames:
        scriptFile.write("echo Removing temporary aligned files.\n")
217
        modeSuffixes['doseWeight'] = '_DW*'
218
        for imageParams in stackParams:
219 220 221
            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))
222 223

scriptFilePath.chmod(scriptFilePath.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP )