diff --git a/config/forms/emw-mc-tomo.json b/config/forms/emw-mc-tomo.json new file mode 100644 index 0000000..0d47ece --- /dev/null +++ b/config/forms/emw-mc-tomo.json @@ -0,0 +1,113 @@ +{ + "name": "emw-mc-tomo", + "help": "Wrapper to run MotionCor motion correction.", + "sections": [ + { + "label": "Motion", + "params": [ + { + "name": "input_star_mics", + "label": "Input movies STAR file:", + "pattern": "STAR Files (*.star)", + "help": "A STAR file with all micrographs to run MotionCor on", + "paramClass": "PathParam", + "pointerClass": "FrameSeries" + }, + { + "paramClass": "Line", + "label": "Patches:", + "params": [ + { + "name": "patch_x", + "label": "x", + "default": 5, + "paramClass": "IntParam", + "help": "Number of patches in x" + }, + { + "name": "patch_y", + "label": "y", + "default": 5, + "paramClass": "IntParam", + "help": "Number of patches in y" + } + ] + }, + { + "name": "do_split_sum", + "label": "Split into odd+even frames?", + "default": false, + "paramClass": "BooleanParam", + "help": "If set to True, frames will be split into even & odd half sets (useful for CryoCARE, etc.)." + }, + { + "name": "do_dose_weighting", + "label": "Perform dose-weighting?", + "default": false, + "paramClass": "BooleanParam", + "help": "If set to True, the averaged micrographs will be dose-weighted. Dose per frame required if selected." + }, + { + "name": "dose_per_frame", + "label": "Dose per frame (e/A2):", + "paramClass": "FloatParam", + "help": "Dose per movie frame (in electrons per squared Angstrom). Required if performing dose-weighting or if using EER movies." + }, + { + "name": "gpu_ids", + "label": "Which GPUs to use:", + "help": "Provide a list of which GPUs (0,1,2,3, etc) to use.\nNote that multiple MotionCor processes should not share a GPU; otherwise, it can lead to crash or broken outputs (e.g. blank images) ." + }, + { + "name": "other_motioncor_args", + "label": "Other MOTIONCOR arguments:", + "help": "Additional arguments that need to be passed to MotionCor." + } + ] + }, + { + "label": "Advanced", + "params": [ + { + "name": "min_fraction_dose", + "label": "Minimum dose per fraction (e/A2):", + "default": 0.15, + "paramClass": "FloatParam", + "help": "Hardware frames are grouped into one 'fraction' such that the dose per fraction is at least this value. Too little dose per fraction, and there will be not enough signal for accurate motion correction. Too much dose dose per fraction, and motion will be sampled more poorly. See https://www3.mrc-lmb.cam.ac.uk/relion/index.php/Image_compression#Falcon4_EER for detailed guidance on EER processing." + }, + { + "name": "err_tolerance", + "label": "Error tolerance:", + "default": 0.10, + "paramClass": "FloatParam", + "help": "Iterative alignment once alignment error is less than this number of pixels. " + }, + { + "name": "reference_frame", + "label": "Reference frame:", + "default": -1, + "paramClass": "IntParam", + "help": "Frame to use for relative shifts (default: middle frame). " + }, + { + "name": "gain_rot", + "label": "Gain rotation:", + "default": 0, + "help": "Rotate the gain reference by this number times 90 degrees clockwise in relion_display. This is the same as -RotGain in MotionCor2. Note that MotionCor2 uses a different convention for rotation so it says 'counter-clockwise'. Valid values are 0, 1, 2 and 3." + }, + { + "name": "gain_flip", + "label": "Gain flip:", + "default": 0, + "help": "Flip the gain reference after rotation. This is the same as -FlipGain in MotionCor2. 0 means do nothing, 1 means flip Y (upside down) and 2 means flip X (left to right)." + }, + { + "name": "fn_defect", + "label": "Defect file:", + "paramClass": "PathParam", + "help": "Location of a UCSF MotionCor2-style defect text file or a defect map that describe the defect pixels on the detector. Each line of a defect text file should contain four numbers specifying x, y, width and height of a defect region. A defect map is an image (MRC or TIFF), where 0 means good and 1 means bad pixels. The coordinate system is the same as the input movie before application of binning, rotation and/or flipping.\nNote that the format of the defect text is DIFFERENT from the defect text produced by SerialEM! One can convert a SerialEM-style defect file into a defect map using IMOD utilities e.g. \"clip defect -D defect.txt -f tif movie.mrc defect_map.tif\". See explanations in the SerialEM manual.\n\nLeave empty if you don't have any defects, or don't want to correct for defects on your detector." + } + ] + } + ] +} diff --git a/emwrap/motioncor/mcpipeline_tomo.py b/emwrap/motioncor/mcpipeline_tomo.py index cec9bb7..d08e911 100644 --- a/emwrap/motioncor/mcpipeline_tomo.py +++ b/emwrap/motioncor/mcpipeline_tomo.py @@ -17,9 +17,11 @@ import os import shutil import json -import argparse import time import sys +import subprocess +import re +import shlex from emtools.utils import Color, FolderManager, Path from emtools.metadata import Table, Column, StarFile, RelionStar, Acquisition @@ -34,23 +36,67 @@ class McPipelineTomo(ProcessingPipeline): """ Pipeline specific to Motioncor tilt-series processing. """ name = 'emw-mc-tomo' - def __init__(self, input_args): - ProcessingPipeline.__init__(self, input_args) + def __init__(self, input_args, output): + super().__init__(input_args, output) args = self._args - self.gpuList = args['gpu'].split() + self.get_gpu_list(args['gpu_ids']) self.outputMicDir = self.join('Micrographs') self.inputLen = 0 + self.micsPerTs = 0 + self.movieDims = () self.acq = self.loadAcquisition() self.inputGain = self.acq.get('gain', None) self.outputTsDir = 'TS' self._DEBUG_only_output = 'DEBUG_only_output' in args - extra = self._args['motioncor']['extra_args'] - self.bin = float(extra.get('-FtBin', 1.0)) + self.get_extras() + ###print(f"\n{os.path.basename(__file__)}:50: self._args='{self._args}'") + + def get_gpu_list(self, gpu_field): + """ + If GPU list not provided, then uses all. + If GPU list provided, then parses into list. + """ + + # Trap for double double quotes + if gpu_field.startswith('"') and gpu_field.endswith('"'): + gpu_field = gpu_field[1:-1] + + gpu_field = gpu_field.strip() + + # Use all GPUs + if not gpu_field: + gpuResult = subprocess.run( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, text=True, check=True) + self.gpuList = [int(line.strip()) for line in gpuResult.stdout.splitlines()] + + else: + parts = re.split(r"[,\s]+", gpu_field) + self.gpuList = [int(p) for p in parts if p] + + def get_extras(self): + """ + Split other_motioncor_args + Get FtBin. + """ + + extra = self._args.get('other_motioncor_args', '') + + # Turn other_motioncor2_args into a dictionary + if extra: + tokens = shlex.split(extra) + extra_dict = dict(zip(tokens[::2], tokens[1::2])) + else: + extra_dict = {} + + self.bin = float(extra_dict.get('-FtBin', 1.0)) + self._args['extra_args'] = extra_dict + del self._args['other_motioncor_args'] def get_motioncor_proc(self, gpu): def _motioncor(batch): - # In this pipeline, batch are not created until now, when we are - # processing each one. + # In this pipeline, batches are not created until now, + # when we are processing each one. # We also need to create the links to movie files items = batch['items'] batch.create() @@ -71,7 +117,7 @@ def _absfn(item): if self.inputGain: acq['gain'] = batch.link(self.inputGain) - mc = Motioncor(acq, **self._args['motioncor']) + mc = Motioncor(acq, **self._args) mc.process_batch(batch, gpu=gpu) return batch @@ -120,8 +166,9 @@ def _output(self, batch): name = f'{baseName}{suffix}.mrc' src = batch.join('output', f'micrograph-{name}') dst = batchFolder.join(name) - #self.log(f"Moving {src} -> {dst}") - shutil.move(src, dst) + if os.path.exists(src): + #self.log(f"Moving {src} -> {dst}") + shutil.move(src, dst) files[suffix] = dst micFile = files[''] @@ -151,7 +198,6 @@ def _output(self, batch): sfOut2.writeTable('global_shift', sf.getTable('global_shift')) sfOut.writeRowValues(values) - self._writeCorrectedTS() batch.info['tsName'] = batch['tsName'] # Store tsName in the info.json @@ -165,12 +211,31 @@ def _output(self, batch): return batch def _getInputTsTable(self): - """ Read input star file and return the 'global' table. """ - inputStar = self._args['input_tiltseries'] + """ + Read input star file and return the 'global' table. + Also stores: + inputLen : number of tilt series + micsPerTs : number of micrographs per tilt series + movieDims : movie dimensions (x, y, num_frames) + + Adapted from warp.WarpBaseTsAlign._getInfo() + """ + + inputStar = self._args['input_star_mics'] with StarFile(inputStar) as sf: - t = sf.getTable('global') - self.inputLen = len(t) # Let's update the inputLen property - return t + tsAllTable = sf.getTable('global') + ###sf.printTable(tsAllTable) + self.inputLen = len(tsAllTable) # Let's update the inputLen property + + # Get number of frame from first movie in first tilt series + first = tsAllTable[0] + tsTable = StarFile.getTableFromFile(first.rlnTomoName, first.rlnTomoTiltSeriesStarFile) + self.micsPerTs = len(tsTable) + movieFn = tsTable[0].rlnMicrographMovieName + self._args['movieDims'] = self.movieDims = Image.get_dimensions(movieFn) + # (What happens if it isn't a movie? (only 2 dimensions will be returned)) + + return tsAllTable return None def _getOutputTsFolder(self, tsName): @@ -181,8 +246,9 @@ def _writeCorrectedTS(self): cols = inputTs.getColumnNames() outTs = Table(cols + ['rlnTomoTiltSeriesPixelSize']) newPixelSize = self.acq.pixel_size * self.bin + newTsStarFile = self.join('corrected_tilt_series.star') - with StarFile(self.join('corrected_tilt_series.star'), 'w') as sfOut: + with StarFile(newTsStarFile, 'w') as sfOut: sfOut.writeTimeStamp() sfOut.writeHeader('global', outTs) for row in inputTs: @@ -195,6 +261,17 @@ def _writeCorrectedTS(self): rlnTomoTiltSeriesPixelSize=newPixelSize) sfOut.writeRowValues(values) + self.outputs = { + 'TiltSeries': { + 'label': 'Tilt Series', + 'type': 'TiltSeries', + 'info': f"{len(inputTs)} items, {self.movieDims[0]} x {self.movieDims[1]} x {self.micsPerTs}, {newPixelSize:0.3f} Å/px", + 'files': [ + [newTsStarFile, 'TomogramGroupMetadata.star.relion.tomo.import'] + ] + } + } + def prerun(self): if self._DEBUG_only_output: print("DEBUG: Only generating output...") @@ -216,6 +293,5 @@ def prerun(self): self.addProcessor(outputQueue, self._output) - if __name__ == '__main__': McPipelineTomo.main() diff --git a/emwrap/motioncor/motioncor.py b/emwrap/motioncor/motioncor.py index a0a2489..729299a 100644 --- a/emwrap/motioncor/motioncor.py +++ b/emwrap/motioncor/motioncor.py @@ -21,7 +21,6 @@ from emtools.metadata import Table, StarFile, TextFile, Acquisition from emtools.image import Image - class Motioncor: """ Motioncor wrapper to run in a batch folder. """ def __init__(self, acq, **kwargs): @@ -31,9 +30,16 @@ def __init__(self, acq, **kwargs): else: self.path, self.version = Motioncor.__get_environ() self.ctf = kwargs.get('ctf', False) + self.framesPerMovie = kwargs['movieDims'][2] + self.minFractionDose = float(kwargs['min_fraction_dose']) + + # Get from project acquisition settings: PixSize, kV, Cs, AmpCont, Gain self.acq = Acquisition(acq) self.args = self.argsFromAcq(acq) - self.args.update(kwargs.get('extra_args', {})) + + # Get from protocol GUI + self.frameDose = 0.0 + self.argsFromGui(kwargs) self.outputPrefix = "output/aligned_" @property @@ -63,6 +69,7 @@ def process_batch(self, batch, **kwargs): elif extLower.startswith('.eer'): inArg = '-InEer' kwargs['-EerSampling'] = 1 + self.make_frame_integration(batch) else: raise Exception(f"Unsupported movie format: {ext}") @@ -136,6 +143,15 @@ def _rename(output, outputPrefix): 'mc_output': total }) + def make_frame_integration(self, batch): + """Create frame integration file for EER files.""" + + self.frame_file = os.path.join(batch.path, 'motioncor-frame.txt') + numFramesToCombine = max( int(self.minFractionDose/self.frameDose), 1) + with open(self.frame_file, 'w') as f: + f.write(f"{self.framesPerMovie} {numFramesToCombine} {self.frameDose:0.3f}") + self.args['–FmIntFile'] = self.frame_file + def __expect(self, fileName): if not os.path.exists(fileName): raise Exception(f"Missing expected output: {fileName}") @@ -186,14 +202,14 @@ def __get_environ(): if program := os.environ.get(varPath, None): if not os.path.exists(program): - raise Exception(f"Motioncor path ({varPath}={program}) does not exists.") + raise Exception(f"MotionCor path ({varPath}={program}) does not exists.") else: - raise Exception(f"Motioncor path variable {varPath} is not defined.") + raise Exception(f"MotionCor path variable {varPath} is not defined.") if version := int(os.getenv(varVersion, 3)): pass else: - raise Exception(f"Motioncor version variable {varVersion} is not defined.") + raise Exception(f"MotionCor version variable {varVersion} is not defined.") return program, version @@ -207,7 +223,15 @@ def __parse_dimensions(logFile): return None def argsFromAcq(self, acq): - """ Define arguments from a given acquisition """ + """ + Add arguments from acquisition parameters: + PixSize + kV + Cs + AmpCont + Gain (can be overridden from protocol GUI) + """ + args = Args({ '-PixSize': acq.pixel_size, '-kV': acq.voltage, @@ -218,3 +242,31 @@ def argsFromAcq(self, acq): args['-Gain'] = gain return args + + def argsFromGui(self, kwargs): + """Add arguments from GUI entry.""" + + ###print(f"\n{os.path.basename(__file__)}:248: self.args({type(self.args)})='{self.args}'") + self.args.update(kwargs.get('extra_args', {})) + ###print(f"\n{os.path.basename(__file__)}:251: kwargs({type(kwargs)})='{kwargs}'") + + self.args['-Patch'] = f"{kwargs['patch_x']} {kwargs['patch_y']}" + if kwargs['gain_rot'] != '0' : self.args['-RotGain'] = kwargs['gain_rot'] + if kwargs['gain_flip'] != '0' : self.args['-FlipGain'] = kwargs['gain_flip'] + if kwargs['fn_defect'] : self.args['-DefectFile'] = kwargs['fn_defect'] + if int(kwargs['reference_frame']) > 0 : self.args['-FmRef'] = kwargs['reference_frame'] + if kwargs['do_split_sum'] : self.args['-SplitSum'] = 1 + if 'num_iters' in kwargs: + if kwargs['num_iters'] : self.args['-Iter'] = kwargs['num_iters'] + if kwargs['err_tolerance'] : self.args['-Tol'] = kwargs['err_tolerance'] + + # If Kv, PixSize, and FmDose are provided, then dose-weighted sums are generated + if kwargs['do_dose_weighting'] : self.args['-FmDose'] = kwargs['dose_per_frame'] + if 'dose_per_frame' in kwargs : + if kwargs['dose_per_frame'] : self.frameDose = float(kwargs['dose_per_frame']) + + # Override gain reference if provided in protocol GUI + if 'fn_gain_ref' in kwargs: + if kwargs['fn_gain_ref'] : self.args['-Gain'] = kwargs['fn_gain_ref'] + + return