Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 113 additions & 0 deletions config/forms/emw-mc-tomo.json
Original file line number Diff line number Diff line change
@@ -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."
}
]
}
]
}
114 changes: 95 additions & 19 deletions emwrap/motioncor/mcpipeline_tomo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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

Expand Down Expand Up @@ -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['']
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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:
Expand All @@ -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...")
Expand All @@ -216,6 +293,5 @@ def prerun(self):

self.addProcessor(outputQueue, self._output)


if __name__ == '__main__':
McPipelineTomo.main()
Loading