mc2FromMDOC.py 7.84 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#!/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'
23
subframesDir = None
24
25
26
27
28
29
30
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.')
31
32
33
# 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
34
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.')
35
36
37
38
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.')
39
parser.add_argument('-nl', '--noLog', action='store_true', help='Allow verbose output from motioncor2 during script run.')
40
parser.add_argument('-kf', '--keepFrames', action='store_true', help='Keep temporary individual aligned frames.')
Adrian Nievergelt's avatar
Adrian Nievergelt committed
41
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.')
42
43
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.')
#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.')
44
45
46
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

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

111
112
113
114
if not subframesDir:
    print("No subframes found in current directory tree, terminating!")
    sys.exit()

Adrian Nievergelt's avatar
Adrian Nievergelt committed
115
116
117
118
119
120
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()
121
122
123
124
stackFileName = os.path.join(mdocDir, inBaseName + ".stackFiles")
tiltFileName = os.path.join(mdocDir, inBaseName + ".rawTlt")

with open(stackFileName, "w+") as stackFile:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
125
    stackFile.write("%i\n" % ( len(stackParams) ))
126
    for imageParams in stackParams:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
127
        stackFile.write("%s\n0\n" % (os.path.join(alignedDir, imageParams['filename']) ))
128
129
130
131
132

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

Adrian Nievergelt's avatar
Adrian Nievergelt committed
133
134
if args.verbose:
    verboseString = ""
135
elif args.noLog:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
136
    verboseString = " &>/dev/null"
137
else:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
138
    verboseString = ">> %s 2>&1" % (inBaseName + "_alignment.log")
Adrian Nievergelt's avatar
Adrian Nievergelt committed
139

140
with open(os.path.join(mdocDir, inBaseName + "_alignFramesMC2.sh"), "w+") as scriptFile:
Adrian Nievergelt's avatar
Adrian Nievergelt committed
141
    scriptFile.write("#!/bin/bash\n\n")
142
143
144
    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
145
146
    scriptFile.write("mkdir -p %s\n" % (alignedDir))
    scriptFile.write("rm -f %s\n\n" % (inBaseName + "_alignment.log"))
147
148
149
150
151
152
153
154
    
    for imageParams in stackParams:
        if args.doseWeight:
            doseWeightParams = " -Kv 300 -FmDose %f -InitDose %f" % (imageParams['frameDose'], imageParams['priorrecorddose'])
        else:
            doseWeightParams = ""
            
        scriptFile.write("echo Aligning frame stack %s \n" % (os.path.join(subframesDir, imageParams['filename']) ))
155
        scriptFile.write("MotionCor2-cuda%i -InMrc \"%s\" -OutMrc \"%s\" -Patch %i %i -InFmMotion 1 -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) )
Adrian Nievergelt's avatar
Adrian Nievergelt committed
156
    scriptFile.write("echo Creating combined tilt series stack %s\n" % (inBaseName + "_mc2.mrc"))
157
    scriptFile.write("source /sw/apps/imod/current/imodenv\n")    
Adrian Nievergelt's avatar
Adrian Nievergelt committed
158
    scriptFile.write("newstack -fileinlist %s -output %s -tilt %s%s\n\n" % ( os.path.basename(stackFileName), inBaseName + "_mc2.mrc", os.path.basename(tiltFileName), verboseString ))
159
160
161
162
163
164
    
    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))