diff --git a/emtools/__init__.py b/emtools/__init__.py index 26dc68e..50a8321 100644 --- a/emtools/__init__.py +++ b/emtools/__init__.py @@ -24,5 +24,5 @@ # * # ************************************************************************** -__version__ = '0.1.3' +__version__ = '0.2.0-rc260805' diff --git a/emtools/datatypes.py b/emtools/datatypes.py new file mode 100644 index 0000000..e153e85 --- /dev/null +++ b/emtools/datatypes.py @@ -0,0 +1,21 @@ +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + + +"""Shared data-type labels for EM file and workflow object metadata.""" + +VOLUME = 'Volume' +STACK_2D = '2D stack' diff --git a/emtools/image/__init__.py b/emtools/image/__init__.py index 619aec6..7341c21 100644 --- a/emtools/image/__init__.py +++ b/emtools/image/__init__.py @@ -1,8 +1,6 @@ # ************************************************************************** # * -# * Authors: J.M. De la Rosa Trevin (delarosatrevin@scilifelab.se) [1] -# * -# * [1] SciLifeLab, Stockholm University +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) # * # * This program is free software; you can redistribute it and/or modify # * it under the terms of the GNU General Public License as published by @@ -14,18 +12,11 @@ # * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # * GNU General Public License for more details. # * -# * You should have received a copy of the GNU General Public License -# * along with this program; if not, write to the Free Software -# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA -# * 02111-1307 USA -# * -# * All comments concerning this program package may be sent to the -# * e-mail address 'delarosatrevin@scilifelab.se' -# * # ************************************************************************** -from .thumbnail import Thumbnail +from .thumbnail import Thumbnail, Image +from emtools.datatypes import STACK_2D, VOLUME -__all__ = [Thumbnail] +__all__ = ["Thumbnail", "Image", "STACK_2D", "VOLUME"] diff --git a/emtools/image/__main__.py b/emtools/image/__main__.py new file mode 100644 index 0000000..5f6de66 --- /dev/null +++ b/emtools/image/__main__.py @@ -0,0 +1,97 @@ +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + +import argparse +import mrcfile +import numpy as np + +from .thumbnail import Image, Thumbnail + + + +def show2D(inputFile, **kwargs): + import matplotlib.pyplot as plt + size = kwargs.get('size', 512) + print("size: ", size) + thumb = Thumbnail(max_size=(size, size), contrast_factor=0.15, std_threshold=1) + array = Image.get_array(inputFile) + plt.imshow(thumb.from_array(array), cmap='gray') + plt.axis('off') # Optional: Turn off axis labels and ticks + plt.show() + + +def fourier_crop(inputFile, scale): + array = Image.get_array(inputFile) + return Image.fourier_crop(inputFile, scale) + + +def save_array(array, outputFile): + if outputFile.endswith('.mrc'): + with mrcfile.new(outputFile, overwrite=True) as mrc: + mrc.set_data(array) + elif outputFile.endswith('.png'): + thumb = Thumbnail(max_size=None, contrast_factor=0.15)#, std_threshold=1) + pil_img = thumb.from_array(array) + pil_img.save(outputFile) + else: + raise ValueError(f"Unsupported file type: {outputFile}") + return outputFile + + +def compare_images(inputFile1, inputFile2): + array1 = Image.get_array(inputFile1) + array2 = Image.get_array(inputFile2) + return np.array_equal(array1, array2) + + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('path', metavar="IMAGE_PATH", + help="Image path") + p.add_argument('--show', '-s', action='store_true', + help="Show the image") + p.add_argument('--max-size', '-m', type=int, default=512, + help="Size of the image") + p.add_argument('--bin', '-b', type=float, default=None, + help="Bin factor for the image, using Fourier cropping. Bin 1 means no cropping, bin 2 means half the original size, etc.") + p.add_argument('--output', '-o', type=str, default=None, + help="Output file name") + p.add_argument('--compare', '-c', + help="Compare the image with another image.") + args = p.parse_args() + + if args.show: + show2D(args.path, size=args.max_size) + elif args.bin: + if args.bin > 1: + scale = 1 / args.bin + array = fourier_crop(args.path, scale) + if args.output: + save_array(array, args.output) + else: + print("Bin factor must be greater than 1") + elif args.compare: + if compare_images(args.path, args.compare): + print("Images are the same") + else: + print("Images are different") + else: + print(Image.get_dimensions(args.path)) + + +if __name__ == '__main__': + main() diff --git a/emtools/image/thumbnail.py b/emtools/image/thumbnail.py index 0da7d7d..f20b036 100644 --- a/emtools/image/thumbnail.py +++ b/emtools/image/thumbnail.py @@ -1,8 +1,6 @@ # ************************************************************************** # * -# * Authors: J.M. De la Rosa Trevin (delarosatrevin@scilifelab.se) [1] -# * -# * [1] SciLifeLab, Stockholm University +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) # * # * This program is free software; you can redistribute it and/or modify # * it under the terms of the GNU General Public License as published by @@ -14,22 +12,21 @@ # * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # * GNU General Public License for more details. # * -# * You should have received a copy of the GNU General Public License -# * along with this program; if not, write to the Free Software -# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA -# * 02111-1307 USA -# * -# * All comments concerning this program package may be sent to the -# * e-mail address 'delarosatrevin@scilifelab.se' -# * # ************************************************************************** +from doctest import OutputChecker import io import numpy as np import base64 import mrcfile +import tifffile + +import PIL +from PIL import ImageFilter, ImageOps + +from emtools.datatypes import STACK_2D, VOLUME -from PIL import Image, ImageOps, ImageFilter +from emtools.utils import Path, Pretty class Thumbnail: @@ -49,7 +46,6 @@ def __init__(self, **kwargs): self.min_max = kwargs.get('min_max', None) self.std_threshold = kwargs.get('std_threshold', 0) - def __format(self, pil_img): format = self.output_format @@ -82,7 +78,8 @@ def from_pil(self, pil_img): pil_img = ImageOps.autocontrast(pil_img, cutoff=self.contrast_factor) if self.gaussian_radius is not None: - pil_img = pil_img.filter(ImageFilter.GaussianBlur(radius=self.gaussian_radius)) + pil_img = pil_img.filter( + ImageFilter.GaussianBlur(radius=self.gaussian_radius)) return self.__format(pil_img) @@ -90,7 +87,7 @@ def from_path(self, path): """ Read the image path as a PIL image and encode it as base64. """ try: - img = Image.open(path) + img = PIL.Image.open(path) encoded = self.from_pil(img) img.close() except: @@ -105,7 +102,7 @@ def from_array(self, imageArray): array = imageArray else: if self.std_threshold > 0: - array = np.array(imageArray) + array = np.asarray(imageArray, dtype=np.float64) imean = array.mean() isd = array.std() isdTh = self.std_threshold * isd @@ -121,7 +118,7 @@ def from_array(self, imageArray): im255 = ((array - iMin) / (iMax - iMin) * 255).astype(np.uint8) - pil_img = Image.fromarray(im255) + pil_img = PIL.Image.fromarray(im255) return self.from_pil(pil_img) @@ -143,8 +140,8 @@ def from_mrc(self, mrc_path): @staticmethod def Micrograph(**kwargs): - """ Shortcut method with presets for Micrograph thumbail. - All settings can be overwriten with kwargs. + """ Shortcut method with presets for Micrograph thumbnail. + All settings can be overwritten with kwargs. """ defaults = { 'output_format': 'base64', @@ -157,8 +154,8 @@ def Micrograph(**kwargs): @staticmethod def Psd(**kwargs): - """ Shortcut method with presets for PSD thumbails. - All settings can be overwriten with kwargs. + """ Shortcut method with presets for PSD thumbnails. + All settings can be overwritten with kwargs. """ defaults = { 'output_format': 'base64', @@ -168,3 +165,304 @@ def Psd(**kwargs): defaults.update(kwargs) return Thumbnail(**defaults) + @staticmethod + def _preview_stack_indices(n): + """Return 4 frame indices: first, two equally spaced, last.""" + if n <= 1: + return [0, 0, 0, 0] + if n == 2: + return [0, 0, 1, 1] + if n == 3: + return [0, 1, 1, 2] + return [ + 0, + int(round((n - 1) / 3)), + int(round(2 * (n - 1) / 3)), + n - 1, + ] + + @staticmethod + def _array_to_uint8(array): + i_min = array.min() + i_max = array.max() + if i_max == i_min: + return np.zeros(array.shape, dtype=np.uint8) + return ((array - i_min) / (i_max - i_min) * 255).astype(np.uint8) + + @staticmethod + def _montage_grid(pil_images, cols, pad=2, bg_color=(255, 255, 255)): + w, h = pil_images[0].size + rows = (len(pil_images) + cols - 1) // cols + canvas_w = cols * w + (cols + 1) * pad + canvas_h = rows * h + (rows + 1) * pad + montage = PIL.Image.new('RGB', (canvas_w, canvas_h), bg_color) + + for i, img in enumerate(pil_images): + row, col = divmod(i, cols) + x = pad + col * (w + pad) + y = pad + row * (h + pad) + if img.mode != 'RGB': + img = img.convert('RGB') + montage.paste(img, (x, y)) + + return montage + + @staticmethod + def _volume_slice_montage(data): + im255 = Thumbnail._array_to_uint8(data) + _, y, x = im255.shape + z = im255.shape[0] + + ximg = PIL.Image.fromarray(im255[:, :, x // 2]) + yimg = PIL.Image.fromarray(im255[:, y // 2, :]) + zimg = PIL.Image.fromarray(im255[z // 2, :, :]) + + xw, xh = ximg.size + yw, yh = yimg.size + pad = 2 + canvas_w = (xw + yw) + (3 * pad) + canvas_h = (xh + yh) + (3 * pad) + montage = PIL.Image.new('RGB', (canvas_w, canvas_h), (255, 255, 255)) + + montage.paste(ximg.convert('RGB'), (pad, pad)) + montage.paste(zimg.convert('RGB'), (pad, xh + 2 * pad)) + montage.paste(yimg.convert('RGB'), (xw + 2 * pad, xh + 2 * pad)) + + return montage + + @staticmethod + def _stack_frame_montage(data): + n = data.shape[0] + im255 = Thumbnail._array_to_uint8(data) + indices = Thumbnail._preview_stack_indices(n) + images = [PIL.Image.fromarray(im255[i, :, :]) for i in indices] + return Thumbnail._montage_grid(images, cols=2) + + @staticmethod + def _mrc_is_stack(mrc, image_lower): + if image_lower.endswith('.mrcs'): + return True + if len(mrc.data.shape) < 3: + return False + if mrc.is_volume(): + return False + if mrc.is_image_stack(): + return True + return True + + @staticmethod + def Preview(imagePath, **kwargs): + imageLower = imagePath.lower() + thumb = Thumbnail.Micrograph(max_size=(256, 256)) + + if not (Path.isImage(imagePath) or Path.isEmImage(imagePath)): + raise Exception("Can not generate preview for: %s" % imagePath) + + if Path.isImage(imagePath): + return thumb.from_path(imagePath) + + if imageLower.endswith('.mrc') or imageLower.endswith('.mrcs'): + with mrcfile.open(imagePath, permissive=True) as mrc: + data = mrc.data + if len(data.shape) == 2: + return thumb.from_array(data) + + if len(data.shape) != 3: + raise Exception("Invalid dimensions: %s" % (data.shape,)) + + preview_thumb = Thumbnail( + max_size=(256, 256), output_format='base64') + + if Thumbnail._mrc_is_stack(mrc, imageLower): + montage = Thumbnail._stack_frame_montage(data) + else: + montage = Thumbnail._volume_slice_montage(data) + + return preview_thumb.from_pil(montage) + + raise Exception("Can not generate preview for: %s" % imagePath) + + +class Image: + @staticmethod + def _mrc_data_type(mrc, image_lower): + if image_lower.endswith('.mrcs'): + return STACK_2D + if mrc.is_volume(): + return VOLUME + if mrc.is_image_stack(): + return STACK_2D + return None + + @staticmethod + def get_metadata(imagePath): + """Return structured metadata for EM image files, or None.""" + imageLower = imagePath.lower() + if imageLower.endswith('.mrc') or imageLower.endswith('.mrcs'): + with mrcfile.open(imagePath) as mrc: + dims = mrc.data.shape[::-1] + if len(dims) == 2: + x, y = dims + return { + 'info': f'{x} x {y}', + } + if len(dims) == 3: + x, y, third = dims + is_cube = x == y == third + data_type = VOLUME if is_cube else Image._mrc_data_type(mrc, imageLower) + if data_type == VOLUME: + return { + 'dataType': data_type, + 'info': f'{x} x {y} x {third}', + } + return { + 'dataType': STACK_2D, + 'info': f'{x} x {y} x {third}', + } + return None + + if (imageLower.endswith('.tif') or + imageLower.endswith('.tiff') or + imageLower.endswith('.eer') or + imageLower.endswith('.gain')): + dims = Image.get_dimensions(imagePath) + if len(dims) == 2: + x, y = dims + return { + 'info': f'{x} x {y}', + } + if len(dims) == 3: + x, y, n = dims + return { + 'dataType': STACK_2D, + 'info': f'{x} x {y} x {n}', + } + return None + + @staticmethod + def get_dimensions(imagePath): + imageLower = imagePath.lower() + if imageLower.endswith('.mrc') or imageLower.endswith('.mrcs'): + with mrcfile.open(imagePath) as mrc: + return mrc.data.shape[::-1] # in reverse order + elif (imageLower.endswith('.tif') or + imageLower.endswith('.tiff') or + imageLower.endswith('.eer') or + imageLower.endswith('.gain')): + with tifffile.TiffFile(imagePath) as tif: + n = len(tif.pages) + y, x = tif.pages[0].shape + return (x, y, n) if n > 1 else (x, y) + + @staticmethod + def get_array(imagePath): + imageLower = imagePath.lower() + if imageLower.endswith('.mrc') or imageLower.endswith('.mrcs'): + with mrcfile.open(imagePath) as mrc: + return mrc.data + elif (imageLower.endswith('.tif') or + imageLower.endswith('.tiff') or + imageLower.endswith('.eer') or + imageLower.endswith('.gain')): + with tifffile.TiffFile(imagePath) as tif: + return tif.asarray() + return None + + @staticmethod + def _fourier_output_size(size, scale): + """Return rounded output size for a given scale factor.""" + return max(1, int(round(size * scale))) + + @staticmethod + def _fourier_axis_slices(in_size, out_size): + """Return source/destination slices for DC-centered crop or pad. + + The DC component lives at in_size // 2 in the shifted spectrum, so the + crop/pad is centered there rather than on the array geometric center. + This keeps even and odd input/output sizes aligned correctly. + """ + if out_size <= in_size: + src_start = in_size // 2 - out_size // 2 + dst_start = 0 + length = out_size + else: + src_start = 0 + dst_start = out_size // 2 - in_size // 2 + length = in_size + return src_start, dst_start, length + + @staticmethod + def _is_volume(imagePath, array): + """Return True when a 3D array should be treated as a volume.""" + if array.ndim != 3: + return False + + if Path.exists(imagePath): + metadata = Image.get_metadata(imagePath) + if metadata and metadata.get('dataType') == VOLUME: + return True + + z, y, x = array.shape + return z == y == x + + @staticmethod + def _fourier_rescale(image, scale): + """Resize an n-D image by Fourier cropping or zero-padding.""" + shape = image.shape + new_shape = tuple(Image._fourier_output_size(s, scale) for s in shape) + + if new_shape == shape: + return np.array(image, copy=True) + + spectrum = np.fft.fftshift(np.fft.fftn(image)) + resized = np.zeros(new_shape, dtype=spectrum.dtype) + + src_slices = [] + dst_slices = [] + for in_size, out_size in zip(shape, new_shape): + src_start, dst_start, length = Image._fourier_axis_slices( + in_size, out_size) + src_slices.append(slice(src_start, src_start + length)) + dst_slices.append(slice(dst_start, dst_start + length)) + + resized[tuple(dst_slices)] = spectrum[tuple(src_slices)] + + # Preserve average intensity when changing the number of pixels. + amp = np.prod(new_shape) / np.prod(shape) + result = np.real(np.fft.ifftn(np.fft.ifftshift(resized * amp))) + + if np.issubdtype(image.dtype, np.floating): + return result.astype(image.dtype, copy=False) + return result + + @staticmethod + def fourier_crop(imagePath, scale): + """Resize an image by Fourier cropping or padding. + + Args: + imagePath: Path to a supported image file. + scale: Output-size multiplier per axis (e.g. 0.5 halves the size). + + Returns: + Rescaled numpy array with the same dimensionality as the input. + """ + array = Image.get_array(imagePath) + if array is None: + raise ValueError("Unsupported image format: %s" % imagePath) + if scale <= 0: + raise ValueError("Scale must be positive, got: %s" % scale) + + if array.ndim == 2: + return Image._fourier_rescale(array, scale) + + if array.ndim == 3: + if Image._is_volume(imagePath, array): + return Image._fourier_rescale(array, scale) + + return np.stack( + [Image._fourier_rescale(array[i], scale) + for i in range(array.shape[0])], + axis=0, + ) + + raise ValueError("Expected 2D or 3D image, got shape: %s" % (array.shape,)) diff --git a/emtools/jobs/__init__.py b/emtools/jobs/__init__.py index 672d307..c230e72 100644 --- a/emtools/jobs/__init__.py +++ b/emtools/jobs/__init__.py @@ -14,7 +14,11 @@ # * # ************************************************************************** -from .pipeline import Pipeline, ProcessingPipeline -from .batch_manager import BatchManager +from .pipeline import Pipeline +from .batch_manager import (Args, NumericList, Batch, Vars, + BatchManager, MdocBatchManager, TsStarBatchManager) +from .workflow import Workflow -__all__ = ["Pipeline", "BatchManager", "ProcessingPipeline"] \ No newline at end of file +__all__ = ["Pipeline", "Workflow", + "Args", "NumericList", "Vars", + "Batch", "BatchManager", "MdocBatchManager", "TsStarBatchManager"] diff --git a/emtools/jobs/batch_manager.py b/emtools/jobs/batch_manager.py index e4ebcd2..da2c5b4 100644 --- a/emtools/jobs/batch_manager.py +++ b/emtools/jobs/batch_manager.py @@ -15,10 +15,265 @@ # ************************************************************************** import os +import json +import subprocess +import traceback +import shlex +import time +import re +from glob import glob from uuid import uuid4 -from datetime import datetime +from datetime import datetime, timedelta +from contextlib import contextmanager -from emtools.utils import Process +from emtools.utils import Color, FolderManager, Timer, Pretty, Path +from emtools.metadata import Mdoc, StarFile + + +class Args(dict): + """ Subclass from dict with some utilities related to arguments. """ + + def toList(self): + args = [] + for k, v in self.items(): + args.append(str(k)) + if isinstance(v, list): + args.extend(str(e) for e in v) + elif v != '': + args.append((str(v))) + return args + + def toLine(self): + return ' '.join("%s %s" % (k, v) for k, v in self.items()) + + @staticmethod + def fromString(string): + return Args.fromList(shlex.split(string)) + + @staticmethod + def fromList(iterable): + r = re.compile(r"^-{1,2}[a-zA-Z][a-zA-Z0-9_-]+$") + def _is_arg(v): + return r.match(v) is not None + + args = Args() + for p in iterable: + if _is_arg(p): + last_key = p + args[p] = '' + else: + v = args[last_key] + + if v: + if isinstance(v, list): + v.append(p) + else: + v = [v, p] + else: + v = p + args[last_key] = v + + return args + + def subset(self, prefix, new_prefix='', filters=None, + inverted_booleans=None, possitive=None, multiple_values=None): + """Return a new Args object with a subset of the keys.""" + filters = filters or [] + inverted_booleans = inverted_booleans or [] + possitive = possitive or [] + multiple_values = multiple_values or [] + + full_prefix = f'{prefix}.' + result = Args() + + for k, v in self.items(): + if not k.startswith(full_prefix): + continue + + k_suffix = k.replace(full_prefix, '') + nk = k.replace(full_prefix, new_prefix) + + if isinstance(v, bool): + if 'binary_boolean' in filters: + result[nk] = '1' if v else '0' + elif 'remove_false' in filters: + add_boolean = not v if k_suffix in inverted_booleans else v + if add_boolean: + result[nk] = '' + else: + result[nk] = v + + continue + + if not v and 'remove_empty' in filters: + continue + + if k_suffix in possitive and float(v) <= 0: + continue + + if 'multiple_values' in filters and k_suffix in multiple_values: + tokens = str(v).split() + result[nk] = tokens if len(tokens) > 1 else v + else: + result[nk] = v + + return result + + +class NumericList(list): + """ List of integers parsed from a string, subclassing 'list' so the + result can be used directly wherever a list of ints is expected. + + Accepted syntax (used e.g. for GPU ids, tilt/frame indices, etc.): + - Individual values, separated by spaces and/or commas: + '1,2,3' or '1 2 3' or '1, 2 3' + - Inclusive ranges, written as 'start-end' (no spaces around the + dash), which get expanded to every value in between: + '10-12' -> 10, 11, 12 + '1, 5-7, 9' -> 1, 5, 6, 7, 9 + """ + + _RANGE = re.compile(r'^(?P\d+)-(?P\d+)$') + + @classmethod + def fromString(cls, string): + values = cls() + for token in re.split(r'[,\s]+', str(string).strip()): + if not token: + continue + m = cls._RANGE.match(token) + if m: + start, end = int(m.group('start')), int(m.group('end')) + if end < start: + raise ValueError( + f"Invalid range '{token}': end must be >= start.") + values.extend(range(start, end + 1)) + else: + try: + values.append(int(token)) + except ValueError: + raise ValueError( + f"Invalid numeric value or range '{token}' in '{string}'.") + return values + + +class Vars: + """ Handle variable definitions, either from input dict + or from os.environ. + """ + def __init__(self, vars={}): + self._vars = vars + + def get(self, key, is_path=False): + """ Get the var for that Key, raising exception if the var does not exist. + If is_path = True, validates that the path exists. + """ + value = self._vars.get(key, os.environ.get(key, None)) + + if value is None: + raise Exception(f"ERROR: Missing expected variable {key}.") + + if is_path and not os.path.exists(value): + raise Exception(f"ERROR: Variable {key}={value} does not exist.") + + return value + + +class Batch(dict, FolderManager): + """ Subclass from dict with some utilities related to Batch logic. """ + def __init__(self, *args, **kwargs): + dict.__init__(self, *args, **kwargs) + FolderManager.__init__(self, self['path']) + self._logId = f" {self.id}:" + self._timer = Timer() # Create a timer to monitor batch execution + self._timerPrefix = '' + + def clone(self): + return Batch(self) + + @property + def id(self): + return self['id'] + + @property + def index(self): + return self['index'] + + @property + def info(self): + if 'info' not in self: + self['info'] = {} + return self['info'] + + @property + def error(self): + return self.info.get('error', None) + + @error.setter + def error(self, value): + self.info['error'] = str(value) + + def dump_info(self): + self.dump(self.info, 'info.json') + + def dump_all(self, fn=None): + fileName = fn or 'batch.json' + self.dump(self, fileName) + + def load_all(self, fn=None): + filePath = fn or self.join('batch.json') + with open(filePath) as f: + self.update(json.load(f)) + + def call(self, program, kwargs, logfile=None, verbose=False, cwd=True): + """ + If cwd is True, call the program from the batch directory. + If cwd is a string, call the program from that directory. + If cwd is False, use the current working directory. + + Returns the exit code of the program. + """ + if isinstance(kwargs, dict): + args = Args(kwargs).toList() + elif isinstance(kwargs, list): + args = list(kwargs) + elif isinstance(kwargs, str): + args = shlex.split(kwargs) + else: + raise Exception("Expecting dict or list as arguments") + + args.insert(0, program) + logfile = logfile or self.join('batch.log') + + with open(logfile, 'a') as f: + cmd = self.log(f"{Color.green(args[0])} {Color.bold(' '.join(args[1:]))}") + f.write(f"\n{cmd}\n") + f.flush() + kwargs = {'stderr': f, 'stdout': f} + if cwd is not False: + kwargs['cwd'] = self.path if cwd is True else cwd + return subprocess.call(args, **kwargs) + + def tic(self, prefix=''): + self._timer.tic() + self._timerPrefix = prefix + + def toc(self): + self.info.update({ + f'{self._timerPrefix}_start': self._timer.getTic(), + f'{self._timerPrefix}_end': Pretty.now(), + f'{self._timerPrefix}_elapsed': str(self._timer.getElapsedTime()) + }) + + @contextmanager + def execute(self, prefix=''): + try: + self.tic(prefix=prefix) + yield self + except Exception as e: + self.error = traceback.format_exc() + finally: + self.toc() class BatchManager: @@ -30,7 +285,8 @@ class BatchManager: folder. """ def __init__(self, batchSize, inputItemsIterator, workingPath, - itemFileNameFunc=lambda item: item.getFileName()): + itemFileNameFunc=lambda item: item.getFileName(), + createBatch=True): """ Args: batchSize: Number of items that will be grouped into one batch @@ -39,49 +295,50 @@ def __init__(self, batchSize, inputItemsIterator, workingPath, itemFileNameFunc: function to extract a filename from each item (by default: lambda item: item.getFileName()) """ - self._items = inputItemsIterator + self._itemsIterator = inputItemsIterator self._batchSize = batchSize self._batchCount = 0 self._workingPath = workingPath self._itemFileNameFunc = itemFileNameFunc + self._create = createBatch def _createBatchId(self): # We will use batchCount, before the batch is created nowPrefix = datetime.now().strftime('%y%m%d-%H%M%S') - countStr = '%02d' % (self._batchCount + 1) + countStr = '%02d' % self._batchCount uuidSuffix = str(uuid4()).split('-')[0] return f"{nowPrefix}_{countStr}_{uuidSuffix}" - def _createBatch(self, items, inputFolder=None): + def _createBatch(self, items, inputFolder=None, **batchAttrs): + self._batchCount += 1 batch_id = self._createBatchId() batch_path = os.path.join(self._workingPath, batch_id) - print(f"Creating batch: {batch_path}") - Process.system(f"rm -rf '{batch_path}'") - Process.system(f"mkdir '{batch_path}'") + batch = Batch(id=batch_id, + index=self._batchCount, + path=batch_path, + items=items, + **batchAttrs) + if self._create: + batch.create() + self._createBatchLinks(batch, items, inputFolder=inputFolder) + return batch + + def _createBatchLinks(self, batch, items, inputFolder=None): if inputFolder is not None: - Process.system(f"mkdir '{batch_path}/{inputFolder}'") + batch.mkdir(inputFolder) for item in items: fn = self._itemFileNameFunc(item) baseName = os.path.basename(fn) if inputFolder is not None: baseName = os.path.join(inputFolder, baseName) - os.symlink(os.path.abspath(fn), - os.path.join(batch_path, baseName)) - - self._batchCount += 1 - return { - 'items': items, - 'id': batch_id, - 'path': batch_path, - 'index': self._batchCount - } + os.symlink(os.path.abspath(fn), batch.join(baseName)) def generate(self): """ Generate batches based on the input items. """ items = [] - for item in self._items: + for item in self._itemsIterator: items.append(item) if len(items) == self._batchSize: @@ -91,3 +348,133 @@ def generate(self): if items: yield self._createBatch(items) + +class MdocBatchManager(BatchManager): + """ Batch manager for Tilt-series. """ + + def __init__(self, mdocsPattern, workingPath, + moviesPath=None, **kwargs): + """ + Args: + mdocsPattern: input pattern of Mdocs files + workingPath: path where the batches folder will be created + moviesPath: path where the frames pointed by Mdocs are + + Kwargs: + wait: waiting time in seconds to check for new files + timeout: time in seconds to quit after no new files found + blacklist: container of tsName that have been processed or want + to be avoided + """ + if not glob(mdocsPattern): + raise Exception(f"No mdoc files were found with pattern: {mdocsPattern}") + + BatchManager.__init__(self, 0, self._iterMdocs(mdocsPattern), workingPath, + itemFileNameFunc=lambda item: item[1]['SubFramePath'], + createBatch=kwargs.get('createBatch', True)) + self._moviesPath = moviesPath + self._wait = kwargs.get('wait', 60) + self._timeout = timedelta(seconds=kwargs.get('timeout', 3600)) + self._blacklist = set(kwargs.get('blacklist', [])) + + def _iterMdocs(self, mdocsPattern): + """ Iterate over a provided Mdocs pattern. """ + one_min = timedelta(minutes=1) + + def _newMdoc(now, fn): + """ Return True if the file meets the following two conditions: + - It has not been processed (in blacklist) + - Modification time is more than 1 minute. + """ + tsName = self._tsName(fn) + if tsName not in self._blacklist: + s = os.stat(fn) + dt = datetime.fromtimestamp(s.st_mtime) + # Ignore also sessions that have not been updated for + # more than X days or that have not been modified since last check + if now - dt > one_min: + self._blacklist.add(tsName) + return True + return False + + last_found = datetime.now() + now = datetime.now() + + def _print(msg): + print(f"INPUT MDOCS: {Pretty.now()}: {msg}", flush=True) + + while now - last_found < self._timeout: + _print("Checking for new mdocs") + if new_mdocs := [fn for fn in glob(mdocsPattern) if _newMdoc(now, fn)]: + _print(f"New mdocs found: {str(new_mdocs)}") + for mdocFn in new_mdocs: + mdoc = Mdoc.parse(mdocFn) + mdoc['MdocFile'] = {'Path': mdocFn} + yield mdoc + last_found = now + else: + _print("No new Mdocs found, sleeping.") + + time.sleep(self._wait) + now = datetime.now() + + def _subframePath(self, mdocFn, section): + movieFolder = self._moviesPath or os.path.dirname(mdocFn) + return os.path.join(movieFolder, Mdoc.getSubFrameBase(section)) + + def _tsName(self, mdocFn): + # Remove all extensions, there are cases like .mrc.mdoc + name = mdocFn + while Path.getExt(name): + name = Path.removeBaseExt(name) + return name + + def generate(self): + """ Generate batches based on the input items. """ + for mdoc in self._itemsIterator: + mdocFn = mdoc['MdocFile']['Path'] + yield self._createBatch(mdoc.zvalues, mdoc=mdoc, tsName=self._tsName(mdocFn)) + + def _createBatchLinks(self, batch, items, inputFolder=None): + mdocFn = batch['mdoc']['MdocFile']['Path'] + + def _absfn(item): + return os.path.abspath(self._subframePath(mdocFn, item[1])) + + framesFolder = os.path.dirname(_absfn(items[0])) + os.symlink(framesFolder, batch.join('frames')) + + for item in items: + baseName = os.path.basename(_absfn(item)) + os.symlink(os.path.join('frames', baseName), batch.join(baseName)) + + +class TsStarBatchManager(BatchManager): + """ + Batch manager from a Relion tilt_series.star file. + (e.g. after the TS import job) + """ + + def __init__(self, tsIterator, workingPath): + """ + Args: + tsIterator: input tilt-series iterator + workingPath: path where the batches folder will be created + """ + BatchManager.__init__(self, 0, tsIterator, workingPath, + itemFileNameFunc=lambda item: item.rlnMicrographMovieName) + self._create = False # Do not create batch folder until processing + + def generate(self): + """ Generate batches based on the input items. """ + for tsRow in self._itemsIterator: + tsName = tsRow.rlnTomoName + tsMdoc = tsRow.rlnTomoMdocFile + with StarFile(tsRow.rlnTomoTiltSeriesStarFile) as sf: + # Force float typing for angle columns: guessing from the + # first row alone can lock the column to int and then break + # on later rows with decimal values (e.g. '3.00013'). + items = [row._asdict() for row in sf.iterTable( + tsName, types={'rlnTomoNominalStageTiltAngle': float, + 'rlnTomoNominalTiltAxisAngle': float})] + yield self._createBatch(items, tsName=tsName, tsMdoc=tsMdoc, rowDict=tsRow._asdict()) diff --git a/emtools/jobs/pipeline.py b/emtools/jobs/pipeline.py index 5f40420..53f91a6 100644 --- a/emtools/jobs/pipeline.py +++ b/emtools/jobs/pipeline.py @@ -14,12 +14,8 @@ # * # ************************************************************************** -import os -import sys from collections import OrderedDict import threading -import signal -import traceback class Pipeline: @@ -64,82 +60,86 @@ class TaskQueue: """ Queue of tasks where producers can deposit tasks and consumers can get it. """ - def __init__(self): + def __init__(self, maxsize=None): self._activeGenerators = 0 - self._condition = threading.Condition() self._tasks = [] + self._maxsize = maxsize + self._lock = threading.Lock() # Lock to access tasks + self._condEmpty = threading.Condition(self._lock) + self._condFull = threading.Condition(self._lock) def getTask(self, proc): """ This function should be called from a consumer of this output instance. """ - self._condition.acquire() - - proc._print("Inside condition lock, queue._activeGenerators: ", - self._activeGenerators) - doWait = True - task = None - - while doWait: - doWait = False - if self._tasks: - proc._print("There are tasks") - task = self._tasks.pop(0) - elif self._activeGenerators > 0: - proc._print("No tasks, but not Done, waiting...") - self._condition.wait() - doWait = True - else: - proc._print("No tasks and done, should return None task.") - - self._condition.release() + with self._lock: + proc._print("Inside condition lock, queue._activeGenerators: ", + self._activeGenerators) + doWait = True + task = None + + while doWait: + doWait = False + if self._tasks: + proc._print("There are tasks") + task = self._tasks.pop(0) + self._condFull.notify() + elif self._activeGenerators > 0: + proc._print("No tasks, but not Done, waiting...") + self._condEmpty.wait() + doWait = True + else: + proc._print("No tasks and done, should return None task.") # Return the task, either None if nothing else should be # done, or a task to be processed return task - def putTask(self, task): + def putTask(self, task, proc): """ This function should be used by subclasses of Output that produces items that will be used by consumers. """ - self._condition.acquire() - self._tasks.append(task) - self._condition.notify() - self._condition.release() + with self._lock: + if self._maxsize and len(self._tasks) == self._maxsize: + self._condFull.wait() + self._tasks.append(task) + self._condEmpty.notify() def notifyGeneratorStarts(self): """ When this queue is associated to a generator, this method should be used to notify that the generator has started to run. """ - self._condition.acquire() - self._activeGenerators += 1 - self._condition.release() + with self._lock: + self._activeGenerators += 1 def notifyGeneratorEnds(self): """ This function should be used by generators associated to this queue to notify that they are done and not more tasks will be produced. """ - self._condition.acquire() - self._activeGenerators -= 1 - if self._activeGenerators == 0: - self._condition.notifyAll() - self._condition.release() + with self._lock: + self._activeGenerators -= 1 + if self._activeGenerators == 0: + self._condEmpty.notifyAll() def isDone(self): - self._condition.acquire() - is_done = self._activeGenerators == 0 - self._condition.release() + with self._lock: + is_done = self._activeGenerators == 0 + return is_done class TaskGenerator(threading.Thread): def __init__(self, generator, outputQueue=None, - name='', debug=False): + name='', debug=False, queueMaxSize=None): """ Params: generator: function generating new tasks outputQueue: queue to put new tasks. If None, a new queue will be created + queueMaxSize: maximum number of task that can be in + output queue. After that, a call to putTask block + the generator. If outputQueue is not None, this + parameter is ignored. """ threading.Thread.__init__(self) self.id = None @@ -148,7 +148,7 @@ def __init__(self, generator, outputQueue=None, self._generator = generator if outputQueue is None: - self.outputQueue = TaskQueue() + self.outputQueue = TaskQueue(maxsize=queueMaxSize) else: self.outputQueue = outputQueue @@ -156,8 +156,11 @@ def run(self): self.outputQueue.notifyGeneratorStarts() self.id = threading.get_ident() + # self._print(">>>>>> Iterating generator tasks") for task in self._generator(): - self.outputQueue.putTask(task) + # self._print(">>>>>>>> Got task: ", task['id'], "...putting it queue.") + self.outputQueue.putTask(task, self) + # self._print(">>>>>>>> SENT task: ", task['id']) self.outputQueue.notifyGeneratorEnds() @@ -169,8 +172,10 @@ def _print(self, *args): class TaskProcessor(TaskGenerator): def __init__(self, inputQueue, processor, outputQueue=None, - name='', debug=False): - TaskGenerator.__init__(self, self._process, outputQueue, name, debug) + name='', debug=False, queueMaxSize=None): + TaskGenerator.__init__(self, self._process, + outputQueue=outputQueue, name=name, + debug=debug, queueMaxSize=queueMaxSize) self._processor = processor self._inputQueue = inputQueue @@ -186,71 +191,3 @@ def _process(self): self._print("Got task: None") - -class ProcessingPipeline(Pipeline): - """ Subclass of Pipeline that is commonly used to run programs. - - This class will define a workingDir (usually os.getcwd) - and an output dir where all output should be generated. - It will also add some helper functions to manipulate file - paths relative to the working dir. - """ - def __init__(self, workingDir, outputDir, **kwargs): - Pipeline.__init__(self, **kwargs) - self.workingDir = self.__validate(workingDir, 'working') - self.outputDir = self.__validate(outputDir, 'output') - - def __validate(self, path, key): - if not path: - raise Exception(f'Invalid {key} directory: {path}') - if not os.path.exists(path): - raise Exception(f'Non-existing {key} directory: {path}') - - return path - - def get_arg(self, argDict, key, envKey, default=None): - """ Get an argument from the argDict or from the environment. - - Args: - argDict: arguments dict from where to get the 'key' value - key: string key of the argument name in argDict - envKey: string key of the environment variable - default: default value if not found in argDict or environ - """ - return argDict.get(key, os.environ.get(envKey, default)) - - def join(self, *p): - return os.path.join(self.outputDir, *p) - - def relpath(self, p): - return os.path.relpath(p, self.workingDir) - - def prerun(self): - """ This method will be called before the run. """ - pass - - def postrun(self): - """ This method will be called after the run. """ - pass - - def __file(self, suffix): - with open(self.join(f'RELION_JOB_EXIT_{suffix}'), 'w'): - pass - - def __abort(self, signum, frame): - self.__file('ABORTED') - sys.exit(0) - - def run(self): - try: - signal.signal(signal.SIGINT, self.__abort) - signal.signal(signal.SIGTERM, self.__abort) - self.prerun() - Pipeline.run(self) - self.postrun() - self.__file('SUCCESS') - except Exception as e: - self.__file('FAILURE') - traceback.print_exc() - - diff --git a/emtools/jobs/workflow.py b/emtools/jobs/workflow.py new file mode 100644 index 0000000..5471e87 --- /dev/null +++ b/emtools/jobs/workflow.py @@ -0,0 +1,163 @@ +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + + +class Workflow: + """ + Simple implementation of a Workflow class for management of Jobs and + produced Data. The workflow is represented as a directed acyclic graph. + """ + + def __init__(self, **kwargs): + self._jobs = {} + self.data = {} + self.jobNextIndex = 1 + + def jobs(self): + """ Iterate over the jobs sorted by index. """ + return sorted(self._jobs.values(), key=lambda j: j.index) + + def root(self): + """ Iterator over nodes that does not have any input. """ + for j in self.jobs(): + if not j.inputs: + yield j + + def hasJob(self, jobId): + return jobId in self._jobs + + def getJob(self, jobId, default=None): + return self._jobs.get(jobId, default) + + def hasData(self, dataId): + return dataId in self.data + + def getData(self, dataId): + return self.data[dataId] + + def registerJob(self, jobId, inputs=None, **kwargs): + jobIndex = kwargs.get('jobindex', self.jobNextIndex) + job = Workflow.Job(self, jobId, jobIndex, + inputs=inputs, **kwargs) + self.jobNextIndex = jobIndex + 1 + self._jobs[jobId] = job + return job + + def deleteJob(self, job): + job.clearOutputs() + del self._jobs[job.id] + + def dot(self): + """ Print the workflow to the terminal. """ + dot = 'digraph G {\n compound=true;\n' + links = '' + + for j in self.jobs(): + dot += (f' subgraph cluster_{j.index} {{\n' + f' style=filled; color=lightgrey; \n' + f' node [style=filled,color=white];\n' + f' i{j.index} [color=lightgrey,fontcolor=lightgrey];\n' + f' label="{j.id}";\n') + for o in j.outputs: + oid = o.id.replace('/', '_').replace('.', '_') + dot += f' {oid}\n' + for c in o.childs: + links += f'{oid} -> i{c.index} [lhead=cluster_{c.index}];\n' + dot += ' }\n' + + dot += f'\n{links}\n}}\n' + return dot + + class Job(dict): + def __init__(self, wf, jobId, index, inputs=None, **kwargs): + dict.__init__(self, **kwargs) + self.wf = wf + self.id = jobId + self.index = index + self._inputs = {} + self._outputs = {} + self.addInputs(inputs) + + @property + def outputs(self): + return self._outputs.values() + + def registerOutput(self, dataId, **kwargs): + kwargs.setdefault('datatype', 'File') + data = Workflow.Data(self, dataId, **kwargs) + self.wf.data[dataId] = data + self._outputs[dataId] = data + return data + + def hasOutput(self, dataId): + return dataId in self._outputs + + def getOutput(self, dataId, default=None): + return self._outputs.get(dataId, default) + + def _validateInputs(self, inputs): + for i in inputs: + if not isinstance(i, Workflow.Data): + raise Exception(f"Input {i} is not of type Workflow.Data") + if i.id in self._inputs: + Exception(f'Input {i} was already added.') + # TODO validate cyclic dependencies + + @property + def inputs(self): + return self._inputs.values() + + def getInput(self, inputId, default=None): + return self._inputs.get(inputId, default) + + def addInputs(self, inputs): + if not inputs: + return + + self._validateInputs(inputs) + + for i in inputs: + self._inputs[i.id] = i + if self not in i.childs: + i.childs.append(self) + + def hasInput(self, inputId): + return inputId in self._inputs + + def clearInputs(self): + for data in self._inputs.values(): + if self in data.childs: + data.childs.remove(self) + self._inputs = {} + + def removeOutput(self, output_id): + if output_id in self._outputs: + if output_id in self.wf.data: + del self.wf.data[output_id] + del self._outputs[output_id] + + def clearOutputs(self): + for output_id in list(self._outputs.keys()): + self.removeOutput(output_id) + + class Data(dict): + def __init__(self, parent, dataId, **kwargs): + dict.__init__(self, **kwargs) + self.id = dataId + self.parent = parent + self.childs = [] + + diff --git a/emtools/metadata/__init__.py b/emtools/metadata/__init__.py index ece86c5..0618916 100644 --- a/emtools/metadata/__init__.py +++ b/emtools/metadata/__init__.py @@ -15,12 +15,14 @@ # ************************************************************************** from .table import Column, ColumnList, Table -from .starfile import StarFile, StarMonitor +from .starfile import StarFile, StarMonitor, RelionStar from .epu import EPU -from .misc import Bins, TsBins, DataFiles, MovieFiles, Mdoc, TextFile +from .misc import (Bins, TsBins, DataFiles, MovieFiles, + Mdoc, TextFile, Acquisition, WarpXml, WarpPopulation, Imod) from .sqlite import SqliteFile -__all__ = ["Column", "ColumnList", "Table", "StarFile", "StarMonitor", "EPU", +__all__ = ["Column", "ColumnList", "Table", + "StarFile", "StarMonitor", "RelionStar", "EPU", "Bins", "TsBins", "SqliteFile", "DataFiles", "MovieFiles", - "Mdoc", "TextFile"] + "Mdoc", "TextFile", "Acquisition", "WarpXml", "WarpPopulation", "Imod"] diff --git a/emtools/metadata/epu.py b/emtools/metadata/epu.py index 2190866..6a23f65 100644 --- a/emtools/metadata/epu.py +++ b/emtools/metadata/epu.py @@ -25,7 +25,8 @@ class EPU: - MOVIES_SUFFICES = ['_fractions.tiff', '_EER.eer'] + MOVIES_SUFFICES = ['_fractions.tiff', '_EER.eer', '_fractions.mrc'] + @staticmethod def get_acquisition(movieXmlFn): """ Parse acquisition parameters from EPU's xml movie file. """ @@ -40,7 +41,7 @@ def get_acquisition(movieXmlFn): def _pixelSize(k): if not pixelSize: return '' - ps = float(pixelSize[k]['numericValue']) * (10**10) + ps = float(pixelSize[k]['numericValue']) * (10 ** 10) return f'{ps:0.5f}' data = { @@ -56,7 +57,7 @@ def _pixelSize(k): 'ExposureTime': camera['ExposureTime'], 'ReadoutArea': {'height': camera['ReadoutArea']['a:height'], 'width': camera['ReadoutArea']['a:width']} - } + } } return data @@ -129,10 +130,20 @@ def get_movie_xml(fn): return fn.replace(s, '.xml') return '' + @staticmethod + def count_movies(folder): + m = 0 + for root, dirs, files in os.walk(folder): + for fn in files: + if EPU.is_movie_fn(fn) and not os.path.islink(os.path.join(root, fn)): + m += 1 + return m + class Data: """ Class to keep track of EPU files and associated metadata. The information can be read/write from/to a STAR file. """ + def __init__(self, rootFolder, epuStar): self._acq = None self._rootFolder = rootFolder @@ -229,6 +240,7 @@ class Session: Monitor EPU session files and allow to make a copy of GridSquares images and xml files. """ + def __init__(self, inputDir, outputStar=None, backupFolder=None, pl=None): """ Create a new EPU.Session instance. diff --git a/emtools/metadata/misc.py b/emtools/metadata/misc.py index 85286bd..bcb2324 100644 --- a/emtools/metadata/misc.py +++ b/emtools/metadata/misc.py @@ -14,11 +14,15 @@ # * # ************************************************************************** +import hashlib import os - +import pathlib from datetime import datetime, timedelta +from glob import glob +from readline import insert_text +import xmltodict -from emtools.utils import Path, Pretty, Process +from emtools.utils import Path, Pretty, Color, Timer class Bins: @@ -134,7 +138,8 @@ def print(self, name): f"\n\ttime: {last_dt}") if self.first and self.last_ts: - print(f"Duration: {(last_dt - first_dt).seconds / 3600:0.2f} hours") + #print(f"Duration: {(last_dt - first_dt).seconds / 3600:0.2f} hours") + print(f"Duration: {Pretty.delta(last_dt - first_dt)}") print(f"Total {name}s: {self.total}, size: {Pretty.size(self.total_size)}") @@ -160,13 +165,18 @@ def __init__(self, filters=[], root=None): def scan(self, folder): """ Scan a folder and register all files recursively. """ + t = Timer() + self.root = Path.addslash(folder) + self._total_dirs = 0 for root, dirs, files in os.walk(folder): for fn in files: self.register(os.path.join(root, fn)) self._total_dirs += len(dirs) + #t.toc("Scanned") + def register(self, filename, stat=None): """ Register a file, if stat is None it will be calculated. """ if stat or os.path.exists(filename): @@ -207,7 +217,7 @@ class MovieFiles(DataFiles): def __init__(self, **kwargs): DataFiles.__init__(self, filters=[self.is_movie], **kwargs) self._moviesSuffix = kwargs.get('moviesSuffix', - ['fractions.tiff', '.eer']) + ['fractions.tiff', '.eer', 'fractions.mrc']) def is_movie(self, fn): return any(fn.endswith(s) for s in self._moviesSuffix) @@ -274,9 +284,55 @@ def parse(mdocFn): return mdoc + @staticmethod + def glob(mdocPattern): + mdocs = [] + for mdocFn in glob(mdocPattern): + mdoc = Mdoc.parse(mdocFn) + mdoc['MdocFile'] = {'Path': mdocFn} + mdocs.append(mdoc) + + return mdocs + + @staticmethod + def getSubFrameBase(section): + """ Helper method to extract the subframe base filename. """ + subFramePath = section.get('SubFramePath', '') + return pathlib.PureWindowsPath(subFramePath).parts[-1] + + MDOC_DATE_FMTS = ('%d-%b-%Y %H:%M:%S', '%d-%b-%y %H:%M:%S') + + @staticmethod + def parseDate(dateStr): + """ Parse an mdoc DateTime field (e.g. '31-Jul-19 17:20:05'). + + SerialEM used dd-Mon-yy before 4.1; yyyy since 4.1 (July 2022). + """ + for fmt in Mdoc.MDOC_DATE_FMTS: + try: + return datetime.strptime(dateStr, fmt) + except ValueError: + continue + raise ValueError(f"Could not parse mdoc DateTime: {dateStr!r}") + @property def zvalues(self): - return [(k, v) for k, v in self.items() if k.startswith('ZValue')] + """ Get the Z values from the mdoc file. + Returns: + list[tuple[str, dict]]: list of Z values with the section data + """ + return list(self.zsections()) + + def zsections(self, sort=None): + """ Iterate over ZValue sections in the mdoc file. + Args: + sort: Use 'date' to sort by acquisition date (newest first). + """ + sections = [(k, v) for k, v in self.items() if k.startswith('ZValue')] + if sort == 'date': + sections.sort(key=lambda x: Mdoc.parseDate(x[1]['DateTime'])) + for k, v in sections: + yield k, v def write(self, path): with open(path, 'w') as f: @@ -296,3 +352,210 @@ def stripLines(fn, **kwargs): if line and not line.startswith('#'): yield line + +class WarpXml: + """ Helper class to read Warp's XML files. """ + def __init__(self, xmlPath): + with open(xmlPath) as f: + self._data = xmltodict.parse(f.read()) + + def getDict(self, *keys): + """ Navigate the provided keys and get a dict from Name=Value pairs. + """ + d = self._data + for k in keys: + d = d[k] + + return {e['@Name']: e['@Value'] for e in d} + + +class WarpSpecies(dict): + """ Helper class to read Warp's .species files. """ + def __init__(self, speciesFile): + with open(speciesFile) as f: + xmlDict = xmltodict.parse(f.read()) + self._data = xmlDict['Species'] + for item in self._data['Param']: + self[item['@Name']] = item['@Value'] + +class WarpPopulation: + """ Helper class to read Warp's .population files. """ + def __init__(self, populationFile): + self._filepath = populationFile + with open(populationFile) as f: + xmlDict = xmltodict.parse(f.read()) + self._data = xmlDict['Population'] + self.Name = self._data['Param']['@Value'] + self.LastRefinementOptions = {e['@Name']: e['@Value'] for e in self._data['LastRefinementOptions']['Param']} + self.Sources = self._parseList(self._data.get('Sources'), 'Source') + self.Species = self._parseList(self._data.get('Species'), 'Species') + + def __repr__(self): + r = f"Population: {self.Name}\n" + r += f" {Color.bold('Last Refinement Options:')}\n" + for k, v in self.LastRefinementOptions.items(): + r += f" {k:<30}: {v:<}\n" + r += f" {Color.green('Species:')}\n" + for s in self.Species: + r += f" {s['name']:<30}: {s['path']:<}\n" + r += f" {Color.cyan('Sources:')}\n" + for s in self.Sources: + r += f" {s['name']:<30}: {s['path']:<}\n" + return r + + def _parseList(self, data, key): + if not data: + return [] + + suffix = f".{key.lower()}" + + def _parseItem(item): + p = item['@Path'] + name = os.path.basename(p).replace(suffix, '') + return {'id': item['@GUID'], 'path': p, 'name': name} + + d = data.get(key) if isinstance(data, dict) else None + if not d: + return [] + + if isinstance(d, list): + return [_parseItem(item) for item in d] + return [_parseItem(d)] + + def digest(self): + """ Content digest of this population and the species it points to. + + Refinements rewrite the species files while leaving the .population + file itself untouched, so both are needed to tell two states apart. + """ + md5 = hashlib.md5() + folder = os.path.dirname(self._filepath) + entries = [('', self._filepath)] + entries += [(s['path'], os.path.join(folder, s['path'])) + for s in self.Species] + + for relpath, path in entries: + md5.update(relpath.encode()) + if os.path.isfile(path): + with open(path, 'rb') as f: + md5.update(f.read()) + + return md5.hexdigest() + + def getSource(self, nameOrIndex): + entry = None + if isinstance(nameOrIndex, int): + entry = self.Sources[nameOrIndex] + else: + for s in self.Sources: + if s['name'] == nameOrIndex: + entry = s + break + + if not entry: + raise Exception(f"Source {nameOrIndex} not found") + + folder = os.path.dirname(self._filepath) + source_file = os.path.join(folder, entry['path']) + if not os.path.isfile(source_file): + raise Exception(f"Source file not found: {source_file}") + + return source_file + + def getSpecies(self, nameOrIndex): + entry = None + if isinstance(nameOrIndex, int): + entry = self.Species[nameOrIndex] + else: + for s in self.Species: + if s['name'] == nameOrIndex: + entry = s + break + + if not entry: + raise Exception(f"Species {nameOrIndex} not found") + + folder = os.path.dirname(self._filepath) + species_file = os.path.join(folder, entry['path']) + if not os.path.isfile(species_file): + raise Exception(f"Species file not found: {species_file}") + + return WarpSpecies(species_file) + + +class Acquisition(dict): + """ Subclass from dict with some utilities related to Acquisition. """ + + @property + def pixel_size(self): + return float(self['pixel_size']) + + @pixel_size.setter + def pixel_size(self, value): + self['pixel_size'] = float(value) + + @property + def voltage(self): + return float(self['voltage']) + + @voltage.setter + def voltage(self, value): + self['voltage'] = float(value) + + @property + def cs(self): + return float(self['cs']) + + @cs.setter + def cs(self, value): + self['cs'] = float(value) + + @property + def amplitude_contrast(self): + return float(self.get('amplitude_contrast', 0.1)) + + @amplitude_contrast.setter + def amplitude_contrast(self, value): + self['amplitude_contrast'] = float(value) + + @property + def dose(self): + return float(self.get('dose', 0.0)) + + @dose.setter + def dose(self, value): + self['dose'] = float(value) + + @property + def total_dose(self): + return float(self.get('total_dose', 0.0)) + + @total_dose.setter + def total_dose(self, value): + self['total_dose'] = float(value) + + +class Imod: + @staticmethod + def get_angles_from_tlt(tltFile): + """ Read AreTomo3/IMOD file with tilt angles. + + Expected file: + TS_NAME_Imod/TS_NAME_st.tlt + Returns: + list[float]: list of tilt angles (as floats) in the same order as in the input file. + """ + return [float(line) for line in TextFile.stripLines(tltFile)] + + @staticmethod + def get_alignment_from_xf(xfFile): + """ Read IMOD XF transformation matrices from .xf file. + + Expected file: + TS_NAME_Imod/TS_NAME_st.xf + Each row contains: + A11 A12 A21 A22 DX DY + Returns: + list[list[float]] + """ + return [list(map(float, line.split())) for line in TextFile.stripLines(xfFile)] \ No newline at end of file diff --git a/emtools/metadata/starfile.py b/emtools/metadata/starfile.py index a2f2a96..f3a87c1 100644 --- a/emtools/metadata/starfile.py +++ b/emtools/metadata/starfile.py @@ -25,11 +25,15 @@ import sys import time import re +import math from contextlib import AbstractContextManager -from collections import OrderedDict from datetime import datetime, timedelta +import emtools +from emtools.utils import Pretty, Color, Path + from .table import ColumnList, Table +from .misc import Acquisition class StarFile(AbstractContextManager): @@ -46,9 +50,12 @@ class StarFile(AbstractContextManager): _splitRegex = re.compile('\"[^"]*\"|[^"\s]+') @staticmethod - def printTable(table, tableName=''): + def printTable(table, tableName='', + computeFormat=False, + timeStamp=False): w = StarFile(sys.stdout, closeFile=False) - w.writeTable(tableName, table, singleRow=len(table) <= 1) + w.writeTable(tableName, table, singleRow=len(table) <= 1, + computeFormat=computeFormat, timeStamp=timeStamp) def __init__(self, inputFile, mode='r', **kwargs): """ @@ -90,9 +97,9 @@ def getTableNames(self): # While searching for a data line, we will store the offsets # for any data_ line that we find if line.startswith('data_'): - tn = line.strip().replace('data_', '') - self._offsets[tn] = offset - self._names.append(tn) + ds = line.strip() + self._offsets[ds] = offset + self._names.append(ds.replace('data_', '', 1)) offset = f.tell() line = f.readline() @@ -120,7 +127,11 @@ def getTable(self, tableName, **kwargs): types=None, optional types dict with {columnName: columnType} pairs that allows to specify types for certain columns. """ - self.__createTable(tableName, **kwargs) + try: + self.__createTable(tableName, **kwargs) + except: + return None + if self._singleRow: self._table.addRow(self.__rowFromValues(self._values)) else: @@ -129,6 +140,22 @@ def getTable(self, tableName, **kwargs): return self._table + @staticmethod + def getTableFromFile(tableName, starFileName, **kwargs): + """ Shortcut to read a table from file. + **kwargs are the same expected by getTable function. + """ + with StarFile(starFileName) as sf: + return sf.getTable(tableName, **kwargs) + + @staticmethod + def getTablesDict(starFileName, **kwargs): + """ Shortcut to read all tables from file as a dictionary. + **kwargs are the same expected by getTable function. + """ + with StarFile(starFileName) as sf: + return {table: sf.getTable(table, **kwargs) for table in sf.getTableNames()} + def getTableSize(self, tableName): """ Return the number of elements in the given table without parsing @@ -274,6 +301,8 @@ def _findDataLine(self, dataName): break line = f.readline() # Start from the beginning and scann until complete the full loop + if initial_offset == 0: + break f.seek(0) offset = 0 line = f.readline() @@ -319,6 +348,15 @@ def writeLine(self, line): """ Write a line to the opened file. """ self._file.write(f"{line}\n") + def writeTimeStamp(self): + """ Write a comment line with current datetime and library version. """ + self.writeLine(f"\n# StarFile written on {Pretty.now()} " + f"by emtools ({emtools.__version__})\n") + + def writeVersion(self, version): + """ Write a RELION metadata-table version tag (e.g. 50001). """ + self.writeLine(f"# version {version}") + def _writeTableName(self, tableName): self._file.write("\ndata_%s\n\n" % (tableName or '')) @@ -341,13 +379,16 @@ def writeHeader(self, tableName, table): self._file.write("loop_\n") self._columns = table.getColumns() # Write column names - for col in self._columns: - self._file.write("_%s \n" % col.getName()) + for i, col in enumerate(self._columns, start=1): + self._file.write(f"_{col.getName()} #{i}\n") def writeRowValues(self, values): """ Write to file a line for these row values. Order should be ensured that is the same of the expected columns. """ + if isinstance(values, dict): + values = values.values() + if not self._format: self._computeLineFormat([values]) @@ -363,37 +404,57 @@ def writeRow(self, row): def _writeNewline(self): self._file.write('\n') - def _computeLineFormat(self, valuesList): + def _computeLineFormat(self, valuesList, computeFormat=False): """ Compute format base on row values width. """ # Take a hint for the columns width from the first row widths = [len(_formatValue(v)) for v in valuesList[0]] formats = [_getFormatStr(v) for v in valuesList[0]] n = len(valuesList) + a = '>' if n > 1: # Check middle and last row, just in case ;) - for index in [n // 2, -1]: + indexes = list(range(len(valuesList))) if computeFormat else [n // 2, -1] + if computeFormat == 'left': + a = '<' + for index in indexes: for i, v in enumerate(valuesList[index]): w = len(_formatValue(v)) if w > widths[i]: widths[i] = w - self._format = " ".join("{:>%d%s} " % (w + 1, f) + self._format = " ".join("{:%s%d%s} " % (a, w + 1, f) for w, f in zip(widths, formats)) + '\n' - def writeTable(self, tableName, table, singleRow=False): + def writeTable(self, tableName, table, + singleRow=False, + computeFormat=False, + timeStamp=False, + version=None): """ Write a Table in Star format to the given file. Args: tableName: The name of the table to write. table: Table that is going to be written singleRow: If True, don't write *loop\_*, just label/value pairs. + computeFormat: compute format based on widest first column, + just for aesthetics and not recommended for large tables. + Values can be 'left' or 'rigth' for alignment. + version: If set, write a RELION ``# version`` tag before the table. """ + if timeStamp: + self.writeTimeStamp() + if version is not None: + self.writeVersion(version) + if table.size(): if singleRow: self.writeSingleRow(tableName, table[0]) else: self.writeHeader(tableName, table) + if computeFormat: + valuesList = [row._asdict().values() for row in table] + self._computeLineFormat(valuesList, computeFormat=computeFormat) for row in table: self.writeRow(row) @@ -416,11 +477,10 @@ def __init__(self, fileName, tableName, rowKeyFunc, **kwargs): self.fileName = fileName self._tableName = tableName self._rowKeyFunc = rowKeyFunc - self._wait = kwargs.get('wait', 10) + self._wait = kwargs.get('wait', 30) self._timeout = timedelta(seconds=kwargs.get('timeout', 300)) self.lastCheck = None # Last timestamp when input was checked self.lastUpdate = None # Last timestamp when new items were found - self.inputCount = 0 # Count all input elements # Black list some items to not be monitored again # We are not interested in the items but just skip them from @@ -432,21 +492,23 @@ def __init__(self, fileName, tableName, rowKeyFunc, **kwargs): def update(self): newRows = [] - now = datetime.now() - mTime = datetime.fromtimestamp(os.path.getmtime(self.fileName)) - - if self.lastCheck is None or mTime > self.lastCheck: - with StarFile(self.fileName) as sf: - for row in sf.iterTable(self._tableName): - rowKey = self._rowKeyFunc(row) - if rowKey not in self._seenItems: - self.inputCount += 1 - self._seenItems.add(rowKey) - newRows.append(row) - - self.lastCheck = now - if newRows: - self.lastUpdate = now + + if os.path.exists(self.fileName): + now = datetime.now() + mTime = datetime.fromtimestamp(os.path.getmtime(self.fileName)) + + if self.lastCheck is None or mTime > self.lastCheck: + with StarFile(self.fileName) as sf: + for row in sf.iterTable(self._tableName): + rowKey = self._rowKeyFunc(row) + if rowKey not in self._seenItems: + self._seenItems.add(rowKey) + newRows.append(row) + + self.lastCheck = now + if newRows: + self.lastUpdate = now + return newRows def timedOut(self): @@ -455,7 +517,7 @@ def timedOut(self): if self.lastCheck is None or self.lastUpdate is None: return False else: - return self.lastCheck - self.lastUpdate > self._timeout + return (self.lastCheck - self.lastUpdate) > self._timeout def newItems(self, sleep=10): """ Yield new items since last update until the stream is closed. """ @@ -478,3 +540,647 @@ def _escapeStrValue(v): """ Escape string values by adding quotes if the string is empty or contains spaces. """ return '"%s"' % v if isinstance(v, str) and (not v or ' ' in v) else v + + +class RelionStar: + + # RELION 5 pipeline STAR tables; prevents re-running pre-5.0 label migration. + PIPELINE_VERSION = 50001 + + JOB_INDEX = re.compile('job(\d{3})') + TRUE_VALUES = ['Yes', 'True', 'true'] + FALSE_VALUES = ['No', 'False', 'false'] + + TOMO_FRAME_SERIES_COLUMNS = [ + 'rlnMicrographMovieName', + 'rlnTomoTiltMovieFrameCount', + 'rlnTomoNominalStageTiltAngle', + 'rlnTomoNominalTiltAxisAngle', + 'rlnMicrographPreExposure', + 'rlnTomoNominalDefocus' + ] + + TOMO_ALIGNMENT_COLUMNS = [ + "rlnTomoXTilt", + "rlnTomoYTilt", + "rlnTomoZRot", + "rlnTomoXShiftAngst", + "rlnTomoYShiftAngst" + ] + + TOMO_OPTIMISATION_SET_TABLE = 'optimisation_set' + TOMO_OPTIMISATION_SET_COLUMNS = [ + 'rlnTomoParticlesFile', + 'rlnTomoTomogramsFile' + ] + + TOMO_PARTICLES_TABLE = 'particles' + TOMO_PARTICLES_COLUMNS = [ + 'rlnTomoName', + ] + TOMO_PARTICLES_PIXEL_COORD_COLUMNS = [ + 'rlnCoordinateX', + 'rlnCoordinateY', + 'rlnCoordinateZ', + ] + TOMO_PARTICLES_CENTERED_COORD_COLUMNS = [ + 'rlnCenteredCoordinateXAngst', + 'rlnCenteredCoordinateYAngst', + 'rlnCenteredCoordinateZAngst', + ] + + @staticmethod + def hasTomoParticleCoordinates(table): + """Return True if the table has tomography particle coordinates.""" + return ( + table.hasAllColumns(RelionStar.TOMO_PARTICLES_PIXEL_COORD_COLUMNS) + or table.hasAllColumns(RelionStar.TOMO_PARTICLES_CENTERED_COORD_COLUMNS) + ) + + @staticmethod + def isTomoOptimisationSet(starFile): + """Return True if the STAR file has a compliant optimisation_set table.""" + try: + RelionStar.readTomoOptimisationSet(starFile) + return True + except (OSError, IOError, Exception): + return False + + @staticmethod + def isTomoParticles(starFile): + """Return True if the STAR file has a compliant tomography particles table.""" + if not starFile or not os.path.isfile(starFile): + return False + try: + with StarFile(starFile) as sf: + if RelionStar.TOMO_PARTICLES_TABLE not in sf.getTableNames(): + return False + if sf.getTableSize(RelionStar.TOMO_PARTICLES_TABLE) < 1: + return False + table = sf.getTableInfo(RelionStar.TOMO_PARTICLES_TABLE) + return ( + table is not None + and table.hasAllColumns(RelionStar.TOMO_PARTICLES_COLUMNS) + and RelionStar.hasTomoParticleCoordinates(table) + ) + except (OSError, IOError, Exception): + return False + + @staticmethod + def readTomoOptimisationSet(starFile): + """Read the optimisation_set table or raise ValueError.""" + if not starFile or not os.path.isfile(starFile): + raise ValueError(f"Invalid STAR file: {starFile}") + all_tables = StarFile.getTablesDict(starFile) + first_table = next(iter(all_tables.values())) + + if not first_table.hasAllColumns(RelionStar.TOMO_OPTIMISATION_SET_COLUMNS): + raise ValueError( + f"{starFile} is not a compliant tomography optimisation_set STAR file." + ) + return first_table + + @staticmethod + def readTomoParticles(starFile): + """Read the particles table or raise ValueError.""" + if not RelionStar.isTomoParticles(starFile): + raise ValueError( + f"{starFile} is not a compliant tomography particles STAR file." + ) + table = StarFile.getTableFromFile(RelionStar.TOMO_PARTICLES_TABLE, starFile) + if not table: + raise ValueError( + f"Could not read '{RelionStar.TOMO_PARTICLES_TABLE}' " + f"table from {starFile}." + ) + return table + + @staticmethod + def to_bool(strValue): + """ Convert Relion Yes/No to True/False. """ + if strValue == 'Yes': + return True + elif strValue == 'False': + return False + else: + raise Exception(f"Invalid Relion bool value: {strValue}") + + @staticmethod + def from_bool(boolValue): + """ Return Yes or No string from True/False. """ + if not isinstance(boolValue): + raise Exception("Expecting bool value for Yes/No conversion") + + return 'Yes' if boolValue else 'No' + + @staticmethod + def true_value(v): + return v in RelionStar.TRUE_VALUES + + @staticmethod + def false_value(v): + return v in RelionStar.FALSE_VALUES + + @staticmethod + def getTomoBinning(row): + return float(getattr(row, 'rlnTomoTomogramBinning', 1)) + + @staticmethod + def getTomoPixelSize(row): + """Compute the tomogram pixel size from TS pixel size and binning.""" + return (float(getattr(row, 'rlnTomoTiltSeriesPixelSize', 0)) + * RelionStar.getTomoBinning(row)) + + @staticmethod + def reconstructedTomoSize(row, axis): + """Return reconstructed tomogram size in pixels along X/Y/Z.""" + return float(getattr(row, axis)) / RelionStar.getTomoBinning(row) + + @staticmethod + def centeredAngstToPixel(centered_angst, row, axis): + """Convert Relion centered Angstrom coordinates to tomogram pixels.""" + return (float(centered_angst) / RelionStar.getTomoPixelSize(row) + + RelionStar.reconstructedTomoSize(row, axis) / 2) + + @staticmethod + def particleCoordsToPixel(particle_row, tomo_row): + """Return particle X/Y/Z coordinates in reconstructed tomogram pixels.""" + if hasattr(particle_row, 'rlnCoordinateX'): + return ( + float(particle_row.rlnCoordinateX), + float(particle_row.rlnCoordinateY), + float(particle_row.rlnCoordinateZ), + ) + return ( + RelionStar.centeredAngstToPixel( + particle_row.rlnCenteredCoordinateXAngst, tomo_row, 'rlnTomoSizeX'), + RelionStar.centeredAngstToPixel( + particle_row.rlnCenteredCoordinateYAngst, tomo_row, 'rlnTomoSizeY'), + RelionStar.centeredAngstToPixel( + particle_row.rlnCenteredCoordinateZAngst, tomo_row, 'rlnTomoSizeZ'), + ) + + @staticmethod + def getTomogram(row): + """Return tomogram path, trying from different columns.""" + cols = ['rlnTomoReconstructedTomogram', 'rlnTomoReconstructedTomogramDenoised'] + for col in cols: + if value := row.get(col): + return value + raise ValueError(f"No tomogram column ({', '.join(cols)}) found in row: {row}") + + @staticmethod + def read_jobstar(jobStarFile): + tValues = StarFile.getTableFromFile('joboptions_values', + jobStarFile, + guessType=False) + def _val(v): + if RelionStar.true_value(v): + return True + elif RelionStar.false_value(v): + return False + else: + return v + + return {row.rlnJobOptionVariable: _val(row.rlnJobOptionValue) for row in tValues} + + @staticmethod + def write_jobstar(jobType, values, jobStarFile, isTomo=0, isContinue=0): + """ Convert params dict to a Relion job.star file. """ + with StarFile(jobStarFile, 'w') as sfOut: + tJob = Table(['rlnJobTypeLabel', 'rlnJobIsContinue', 'rlnJobIsTomo']) + tJob.addRowValues(jobType, isContinue, isTomo) + sfOut.writeTimeStamp() + sfOut.writeTable('job', tJob, singleRow=True) + tValues = Table(['rlnJobOptionVariable', 'rlnJobOptionValue']) + for k, v in values.items(): + if isinstance(v, bool): + val = 'Yes' if v else 'No' + elif v is None: + val = '' + else: + # Job option values are always strings in STAR files. + val = str(v) + tValues.addRowValues(k, val) + sfOut.writeTable('joboptions_values', tValues, computeFormat='left') + + @staticmethod + def optics_table(acq, opticsGroup=1, opticsGroupName="opticsGroup1", + mtf=None, originalPixelSize=None): + origPs = originalPixelSize or acq['pixel_size'] + + values = { + 'rlnOpticsGroupName': opticsGroupName, + 'rlnOpticsGroup': opticsGroup, + 'rlnMicrographOriginalPixelSize': origPs, + 'rlnVoltage': acq['voltage'], + 'rlnSphericalAberration': acq['cs'], + 'rlnAmplitudeContrast': acq.get('amplitude_contrast', 0.1), + 'rlnMicrographPixelSize': acq['pixel_size'] + } + if mtf: + values['rlnMtfFileName'] = mtf + return Table.fromDict(values) + + @staticmethod + def movies_table(**kwargs): + extra_cols = kwargs.get('extra_cols', []) + return Table([ + 'rlnMicrographMovieName', + 'rlnOpticsGroup' + ] + extra_cols) + + @staticmethod + def micrograph_table(**kwargs): + cols = [] + if image_id := kwargs.get('image_id', None): + cols.append(image_id) + cols.extend([ + 'rlnMicrographName', + 'rlnOpticsGroup', + 'rlnCtfImage', + 'rlnDefocusU', + 'rlnDefocusV', + 'rlnCtfAstigmatism', + 'rlnDefocusAngle', + 'rlnCtfFigureOfMerit', + 'rlnCtfMaxResolution' + ]) + if extra_cols := kwargs.get('extra_cols', []): + cols.extend(extra_cols) + return Table(cols) + + @staticmethod + def coordinates_table(**kwargs): + return Table(['rlnMicrographName', 'rlnMicrographCoordinates']) + + @staticmethod + def tiltseries_table(mc=True, ctf=True, **kwargs): + cols = list(RelionStar.TOMO_FRAME_SERIES_COLUMNS) + cols.extend([ + 'rlnMicrographName', + 'rlnMicrographNameEven', + 'rlnMicrographNameOdd' + ]) + + if mc: + cols.extend([ + 'rlnMicrographMetadata', + 'rlnAccumMotionTotal', + 'rlnAccumMotionEarly', + 'rlnAccumMotionLate' + ]) + + if ctf: + cols.extend([ + 'rlnCtfImage', + 'rlnDefocusU', + 'rlnDefocusV', + 'rlnCtfAstigmatism', + 'rlnDefocusAngle', + 'rlnCtfFigureOfMerit', + 'rlnCtfMaxResolution', + 'rlnCtfIceRingDensity' + ]) + + cols.extend(kwargs.get('extra_cols', [])) + + return Table(cols) + + @staticmethod + def global_tiltseries_table(**kwargs): + cols = [ + 'rlnTomoName', + 'rlnTomoTiltSeriesStarFile', + 'rlnVoltage', + 'rlnSphericalAberration', + 'rlnAmplitudeContrast', + 'rlnMicrographOriginalPixelSize', + 'rlnTomoHand', + 'rlnOpticsGroupName', + 'rlnTomoTiltSeriesPixelSize' + ] + cols.extend(kwargs.get('extra_cols', [])) + + return Table(cols) + + @staticmethod + def _acquisition_from_row(row): + """ Build Acquisition from an optics or tomography global row. """ + if getattr(row, 'rlnTomoTiltSeriesPixelSize', None): + pixel_size = RelionStar.getTomoPixelSize(row) + else: + pixel_size = (getattr(row, 'rlnMicrographPixelSize', None) + or row.rlnMicrographOriginalPixelSize) + + acq = Acquisition( + pixel_size=pixel_size, + voltage=row.rlnVoltage, + cs=row.rlnSphericalAberration, + amplitude_contrast=getattr(row, 'rlnAmplitudeContrast', 0.1) + ) + if gain := getattr(row, 'rlnMicrographGainName', None): + acq['gain'] = gain + if dose := getattr(row, 'rlnMicrographDoseRate', None): + acq['total_dose'] = float(dose) + + return acq + + @staticmethod + def _resolve_linked_star(baseStarFile, linkedPath): + if not linkedPath: + return None + if os.path.isabs(linkedPath): + return linkedPath + + candidates = [ + os.path.normpath(os.path.join(os.path.dirname(baseStarFile), + linkedPath)), + os.path.normpath(os.path.join(os.getcwd(), linkedPath)), + ] + for candidate in candidates: + if os.path.exists(candidate): + return candidate + return candidates[0] + + @staticmethod + def getAcquisition(inputTableOrFile): + """ Load acquisition parameters from an optics/global table row, + or a given input STAR file (movies, tilt series, tomograms, etc.). + """ + if hasattr(inputTableOrFile, 'rlnVoltage'): + return RelionStar._acquisition_from_row(inputTableOrFile) + + if isinstance(inputTableOrFile, Table): + return RelionStar._acquisition_from_row(inputTableOrFile[0]) + + starFile = inputTableOrFile + if RelionStar.isTomoOptimisationSet(starFile): + row = RelionStar.readTomoOptimisationSet(starFile)[0] + if tomogramsStar := getattr(row, 'rlnTomoTomogramsFile', None): + return RelionStar.getAcquisition( + RelionStar._resolve_linked_star(starFile, tomogramsStar)) + if particlesStar := getattr(row, 'rlnTomoParticlesFile', None): + return RelionStar.getAcquisition( + RelionStar._resolve_linked_star(starFile, particlesStar)) + + with StarFile(starFile) as sf: + if t := sf.getTable('optics'): + return RelionStar._acquisition_from_row(t[0]) + if t := sf.getTable('global'): + return RelionStar._acquisition_from_row(t[0]) + + raise Exception(f"Could not read acquisition parameters from {starFile}") + + @staticmethod + def alignment_from_xf(xf_row, pixel_size): + """Convert one IMOD XF row into Relion alignment labels. + IMOD XF row: + A11 A12 A21 A22 DX DY + The translation should be taken from the inverse transform, then + converted from pixels to Angstroms. + """ + a11, a12, a21, a22, dx, dy = xf_row + + det = a11 * a22 - a12 * a21 + if abs(det) < 1e-12: + return { + 'rlnTomoZRot': '', + 'rlnTomoXShiftAngst': '', + 'rlnTomoYShiftAngst': '', + } + + z_rot = math.degrees(math.atan2(a12, a11)) + + # Inverse affine translation: + # inv(M) * -t + inv_dx = -((a22 * dx - a12 * dy) / det) + inv_dy = -((-a21 * dx + a11 * dy) / det) + + return { + 'rlnTomoZRot': z_rot, + 'rlnTomoXShiftAngst': inv_dx * pixel_size, + 'rlnTomoYShiftAngst': inv_dy * pixel_size, + } + + @staticmethod + def alignment_to_xf(alignment, pixel_size): + """Convert one RELION alignment row into an AreTomo/IMOD XF row. + + The AreTomo XF convention used here stores: + + A11 A12 A21 A22 DX DY + + Matrix coefficients are rounded to three decimals before computing + DX and DY. Translations are then rounded to two decimals. This is + required to reverse the existing XF-to-RELION conversion exactly + for AreTomo's rounded rigid transforms. + """ + if pixel_size <= 0: + raise ValueError( + f"Alignment pixel size must be greater than zero: {pixel_size}" + ) + + def _get_value(name): + if isinstance(alignment, dict): + value = alignment.get(name) + else: + value = getattr(alignment, name, None) + + if value in (None, ''): + raise ValueError(f"Missing RELION alignment value: {name}") + + return float(value) + + z_rot = _get_value('rlnTomoZRot') + + shift_x_pixels = ( + _get_value('rlnTomoXShiftAngst') / pixel_size + ) + shift_y_pixels = ( + _get_value('rlnTomoYShiftAngst') / pixel_size + ) + + angle = math.radians(z_rot) + + # Important: AreTomo writes the matrix coefficients with three + # decimals. Round them before reconstructing the translation. + a11 = float(f'{math.cos(angle):.3f}') + a12 = float(f'{math.sin(angle):.3f}') + a21 = float(f'{-math.sin(angle):.3f}') + a22 = float(f'{math.cos(angle):.3f}') + + # Existing forward conversion: + # relion_shift = -inverse(M) @ imod_translation + # Therefore: + # imod_translation = -M @ relion_shift + dx = -(a11 * shift_x_pixels + a12 * shift_y_pixels) + dy = -(a21 * shift_x_pixels + a22 * shift_y_pixels) + + # AreTomo writes translations with two decimals. + dx = float(f'{dx:.2f}') + dy = float(f'{dy:.2f}') + + # Avoid writing "-0.00". + if dx == 0: + dx = 0.0 + if dy == 0: + dy = 0.0 + + return [a11, a12, a21, a22, dx, dy] + + @staticmethod + def alignments_from_imod(tlt_angles, xf_alignments, pixel_size): + """ Read tilt angles (.tlt file) and IMOD transforms (.xf file) to compute Relion alignments. + Returns: + list[dict]: list of Relion alignments + """ + rln_alignments = [] + + for tilt, xf_row in zip(tlt_angles, xf_alignments): + xf_values = RelionStar.alignment_from_xf(xf_row, pixel_size) + ctf_scale = math.cos(math.radians(tilt)) + + rln_alignments.append({ + 'tilt': tilt, + 'rlnTomoXTilt': 0.0 if tilt != '' else '', + 'rlnTomoYTilt': tilt, + 'rlnTomoZRot': xf_values.get('rlnTomoZRot', ''), + 'rlnTomoXShiftAngst': xf_values.get('rlnTomoXShiftAngst', ''), + 'rlnTomoYShiftAngst': xf_values.get('rlnTomoYShiftAngst', ''), + 'rlnCtfScalefactor': ctf_scale, + }) + + return rln_alignments + + @staticmethod + def pipeline_tables(): + return { + 'processes': Table(['rlnPipeLineProcessName', + 'rlnPipeLineProcessAlias', + 'rlnPipeLineProcessTypeLabel', + 'rlnPipeLineProcessStatusLabel']), + 'nodes': Table(['rlnPipeLineNodeName', + 'rlnPipeLineNodeTypeLabel', + 'rlnPipeLineNodeTypeLabelDepth']), + 'output_edges': Table(['rlnPipeLineEdgeProcess', + 'rlnPipeLineEdgeToNode']), + 'input_edges': Table(['rlnPipeLineEdgeFromNode', + 'rlnPipeLineEdgeProcess']) + } + + @staticmethod + def write_pipeline(pipeline_star, jobCounter=1, tables=None): + version = RelionStar.PIPELINE_VERSION + with StarFile(pipeline_star, 'w') as sf: + sf.writeTimeStamp() + tGeneral = Table(['rlnPipeLineJobCounter']) + tGeneral.addRowValues(jobCounter) + sf.writeTable('pipeline_general', tGeneral, + singleRow=True, version=version) + + if tables: + for name, t in tables.items(): + if len(t): + sf.writeTable(f"pipeline_{name}", t, + computeFormat=True, version=version) + + @staticmethod + def job_index(jobId): + """ Return the integer job index from the name of the form Folder/jobXXX. """ + m = RelionStar.JOB_INDEX.search(jobId) + if m is None: + return None + else: + return int(m.groups()[0]) + + @staticmethod + def pipeline_to_workflow(pipelineStar): + """ Read the Relion pipeline star file and build the proper Workflow. """ + from emtools.jobs import Workflow # import here to avoid circular imports + + wf = Workflow() + with StarFile(pipelineStar) as sf: + tables = sf.getTableNames() + + def _table(name): + fullname = f"pipeline_{name}" + return sf.getTable(fullname) if fullname in tables else None + + if tGeneral := _table('general'): + wf.jobNextIndex = int(tGeneral[0].rlnPipeLineJobCounter) + else: + wf.jobNextIndex = 1 + + if tProc := _table('processes'): + for row in tProc: + jobId = Path.rmslash(row.rlnPipeLineProcessName) + wf.registerJob(jobId, + alias=row.rlnPipeLineProcessAlias, + status=row.rlnPipeLineProcessStatusLabel, + jobtype=row.rlnPipeLineProcessTypeLabel, + jobindex=RelionStar.job_index(jobId)) + + if tNodes := _table('nodes'): + nodes = {row.rlnPipeLineNodeName: row.rlnPipeLineNodeTypeLabel + for row in tNodes} + else: + nodes = {} + + if tOutput := _table('output_edges'): + for row in tOutput: + job = wf.getJob(Path.rmslash(row.rlnPipeLineEdgeProcess)) + nodeName = row.rlnPipeLineEdgeToNode + job.registerOutput(nodeName, datatype=nodes.get(nodeName, 'File')) + + if tInput := _table('input_edges'): + for row in tInput: + job = wf.getJob(Path.rmslash(row.rlnPipeLineEdgeProcess)) + if wf.hasData(row.rlnPipeLineEdgeFromNode): + job.addInputs([wf.getData(row.rlnPipeLineEdgeFromNode)]) + else: + print(f"WARNING: Missing input edge: {row.rlnPipeLineEdgeFromNode}") + + return wf + + @staticmethod + def workflow_to_pipeline(wf, pipelineStar): + """ Write the input workflow as the expected Relion pipeline STAR file. """ + tables = RelionStar.pipeline_tables() + tProc = tables['processes'] + tNodes = tables['nodes'] + tOutput = tables['output_edges'] + tInput = tables['input_edges'] + + # There are some job'status that are not supported by Relion, so we need to map them to the expected values + status_map = { + 'Launched': 'Scheduled', + 'Saved': 'Scheduled' + } + + for job in wf.jobs(): + status = status_map.get(job['status'], job['status']) + tProc.addRowValues( + rlnPipeLineProcessName=Path.addslash(job.id), + rlnPipeLineProcessAlias=job['alias'], + rlnPipeLineProcessStatusLabel=status, + rlnPipeLineProcessTypeLabel=job['jobtype'] + ) + for i in job.inputs: + tInput.addRowValues( + rlnPipeLineEdgeProcess=Path.addslash(job.id), + rlnPipeLineEdgeFromNode=i.id + ) + + for o in job.outputs: + tNodes.addRowValues( + rlnPipeLineNodeName=o.id, + rlnPipeLineNodeTypeLabel=o.get('datatype', 'File'), + rlnPipeLineNodeTypeLabelDepth=1 + ) + tOutput.addRowValues( + rlnPipeLineEdgeProcess=Path.addslash(job.id), + rlnPipeLineEdgeToNode=o.id + ) + + RelionStar.write_pipeline(pipelineStar, wf.jobNextIndex, tables) diff --git a/emtools/metadata/table.py b/emtools/metadata/table.py index 26126f3..e068168 100644 --- a/emtools/metadata/table.py +++ b/emtools/metadata/table.py @@ -23,6 +23,7 @@ from collections import OrderedDict, namedtuple +import re class Column: @@ -49,6 +50,9 @@ def getType(self): def setType(self, colType): self._type = colType + def clone(self): + return Column(self._name, type=self._type) + class ColumnList: def __init__(self, columns=None): @@ -111,6 +115,18 @@ def get(self, key, default=None): return Row + def cloneColumns(self, exclude=None): + """ Create a new Table that will have exactly the same columns + as this table. Optionally, some columns can be excluded. """ + excludeList = exclude or [] + newCols = [] + + for colName, col in self._columns.items(): + if colName not in excludeList: + newCols.append(col.clone()) + + return Table(newCols) + @staticmethod def createColumns(colNames, values, guessType=True, types=None): """ Return a list of Columns create from the names. @@ -143,6 +159,29 @@ def __init__(self, columns=None): self.Row = self.createRowClass() self._rows = [] + @staticmethod + def fromDict(valuesDict): + """ Create a Table from a dictionary of values or a list of dictionaries. + If it is a list, all dictionaries must have the same keys. + + Args: + valuesDict: a dictionary of values or a list of dictionaries + Returns: + Table: a Table object + """ + if isinstance(valuesDict, dict): + rows = [valuesDict] + elif isinstance(valuesDict, list): + rows = valuesDict + else: + raise ValueError(f"Invalid type {type(valuesDict)} for valuesDict") + + t = Table(list(rows[0].keys())) + for row in rows: + t.addRowValues(**row) + + return t + def clear(self): self.Row = None self._columns.clear() @@ -262,6 +301,51 @@ def keyFunc(r): return getattr(r, key) if isinstance(key, str) else key self._rows.sort(key=keyFunc, reverse=reverse) + def update(self, spec): + """Update column values from a comma-separated spec of col=expr assignments. + + Each assignment must contain exactly one '=' character. Expressions are + evaluated per row using column names as variables. + + Args: + spec: Comma-separated column=expression pairs, e.g. + 'rlnPixelSize=rlnOriginalPixelSize/2, rlnOpticsGroup=1' + + Returns: + self, to allow method chaining. + """ + updates = _parseUpdateSpec(spec) + colNames = self.getColumnNames() + for col, _ in updates: + if col not in colNames: + raise ValueError(f"Column '{col}' not found in table") + + newRows = [] + for row in self._rows: + values = row._asdict() + for col, expr in updates: + values[col] = _evalRowExpr(expr, row, values=values) + newRows.append(self.Row(**values)) + self._rows = newRows + return self + + def filter(self, spec): + """Keep rows where the expression evaluates to True. + + Args: + spec: Boolean expression evaluated per row using column names + as variables, e.g. 'rlnOpticsGroup == 1' + + Returns: + self, to allow method chaining. + """ + spec = spec.strip() + if not spec: + raise ValueError("FILTER requires a non-empty expression") + _validateFilterSpec(spec) + self._rows = [row for row in self._rows if _evalRowExpr(spec, row)] + return self + def print(self, formatStr=None): for row in self._rows: print(formatStr.format(**row._asdict())) @@ -280,6 +364,48 @@ def __setitem__(self, key, value): # --------- Helper functions ------------------------ +def _parseUpdateSpec(spec): + """Parse 'col1=expr1, col2=expr2' into a list of (column, expression) pairs.""" + spec = spec.strip() + if not spec: + raise ValueError("UPDATE requires a non-empty expression spec") + + updates = [] + for part in spec.split(','): + part = part.strip() + if not part: + continue + if part.count('=') != 1: + raise ValueError( + f"Invalid UPDATE spec (expected exactly one '=' per column): {part}") + col, expr = part.split('=', 1) + col = col.strip() + expr = expr.strip() + if not col: + raise ValueError(f"Invalid UPDATE spec (empty column name): {part}") + if not expr: + raise ValueError(f"Invalid UPDATE spec (empty expression): {part}") + updates.append((col, expr)) + + if not updates: + raise ValueError("UPDATE requires at least one column assignment") + return updates + + +def _validateFilterSpec(spec): + """Reject bare column names that often mean the shell ate a comparison.""" + if re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*', spec): + raise ValueError( + f"FILTER expression '{spec}' looks like a column name only. " + "Quote the expression in the shell, e.g. 'rlnDefocusAngle > 50'") + + +def _evalRowExpr(expr, row, values=None): + """Evaluate an expression using row column values as variables.""" + namespace = dict(values if values is not None else row._asdict()) + return eval(expr, {"__builtins__": {}}, namespace) + + def _str(s): """ Get the string value but stripping quotes if present. """ return s[1:-1] if s.startswith('"') and s.endswith('"') else s diff --git a/emtools/scripts/emt-scipion-otf.py b/emtools/scripts/emt-scipion-otf.py index b947d39..69d2185 100755 --- a/emtools/scripts/emt-scipion-otf.py +++ b/emtools/scripts/emt-scipion-otf.py @@ -24,6 +24,7 @@ from collections import OrderedDict import datetime as dt import re +from pprint import pprint from emtools.utils import Process, Color, System from emtools.metadata import EPU, SqliteFile, StarFile, Table @@ -314,7 +315,7 @@ def _path(*p): sphericalAberration=acq['cs'], doseInitial=0.0, dosePerFrame=acq['dose'], - gainFile=gain, + gainFile=os.path.abspath(gain), dataStreaming=True ) @@ -346,6 +347,7 @@ def _path(*p): 'motioncorr.protocols.ProtMotionCorrTasks', objLabel='motioncor', patchX=patchX, patchY=patchY, + gainFlip=1, # Fli numberOfThreads=1, streamingBatchSize=16, gpuList=' '.join(str(g) for g in params['mcGpus']) @@ -701,6 +703,180 @@ def fix_run_links(workingDir, srcRuns): logger.system(f"cd Runs && ln -s runs/{fn}") +class CryoSparc: + STATUS_FAILED = "failed" + STATUS_ABORTED = "aborted" + STATUS_COMPLETED = "completed" + STATUS_KILLED = "killed" + STATUS_RUNNING = "running" + STATUS_QUEUED = "queued" + STATUS_LAUNCHED = "launched" + STATUS_STARTED = "started" + STATUS_BUILDING = "building" + + STOP_STATUSES = [STATUS_ABORTED, STATUS_COMPLETED, STATUS_FAILED, STATUS_KILLED] + ACTIVE_STATUSES = [STATUS_QUEUED, STATUS_RUNNING, STATUS_STARTED, + STATUS_LAUNCHED, STATUS_BUILDING] + + def __init__(self, projId): + self.projId = projId + from cryosparc.tools import CryoSPARC, CommandClient + cs_config = os.environ.get('CRYOSPARC_CONFIG', None) + if cs_config is None: + raise Exception('Please define CRYOSPARC_CONFIG="LICENSE|URL|PORT"') + + license, url, port = cs_config.split('|') + print("\n>>> Using license: ", Color.green(license)) + print(">>> URL/port: ", Color.bold(f"{url}:{port}")) + self._cli = CommandClient(host=url, port=port, headers={"License-ID": license}) + projInfo = self.cli('get_project', projId) + print("\n", "=" * 20, Color.green(f"PROJECT: {projId}"), "=" * 20) + pprint(projInfo) + print("=" * 50, "\n") + self.userId = projInfo['owner_user_id'] + lanes = self.cli('get_scheduler_lanes') + pprint(lanes) + + def __call__(self, cmd, **kwargs): + p = Process(self.csm, 'cli', cmd) + lines = list(p.lines()) + + try: + for i, line in enumerate(lines): + print(Color.cyan(i), Color.bold(line)) + return lines[0] + except Exception as e: + print(Color.red(f"Error: running command {cmd}")) + print(e) + + def _argstr(self, args): + return json.dumps(args).replace('true', 'True') + + def cli(self, function, *args, **kwargs): + def _val(v): + return Color.bold(json.dumps(v)) + + argsStr = ','.join(_val(a) for a in args) + sepStr = ', ' if argsStr else '' + kwargsStr = ','.join("%s=%s" % (Color.cyan(k), _val(v)) for k, v in kwargs.items()) + print(f"\n{Color.green(function)}({argsStr}{sepStr}{kwargsStr})") + func = getattr(self._cli, function) + return func(*args, **kwargs) + + def job_status(self, jobId): + """ Return the job status. """ + status = self.cli('get_job_status', project_uid=self.projId, job_uid=jobId) + print(status) + return status + + def job_wait(self, jobId): + """ Wait for a job to complete (in any stop status). """ + while self.job_status(jobId) not in self.STOP_STATUSES: + time.sleep(10) + + def job_run(self, wsId, jobType, args, inputs={}, wait=True): + #cmd = (f'make_job("{jobType}", "{self.projId}", "{wsId}", "{self.userId}", None, None, None, ' + # f'{self._argstr(args)}, {self._argstr(inputs)})') + #jobId = self(cmd) + # jobId = self._cli.make_job(job_type=jobType, project_uid=self.projId, workspace_uid=wsId, + # user_id=self.userId, params=args, input_group_connects=inputs) + jobId = self.cli('make_job', + job_type=jobType, project_uid=self.projId, workspace_uid=wsId, + user_id=self.userId, params=args, input_group_connects=inputs) + #cmd = f'enqueue_job("{self.projId}", "{jobId}", "default", "{self.userId}")' + #self(cmd) + #self._cli.enqueue_job(project_uid=self.projId, user_id=self.userId, job_uid=jobId, lane='default') + self.cli('enqueue_job', + project_uid=self.projId, user_id=self.userId, job_uid=jobId) + if wait: + self.job_wait(jobId) + + return jobId + + +def cryosparc_prepare(): + if os.path.exists('CS'): + raise Exception("CS folder already exists. Remove it before running this command.") + + logger = Process.Logger(format="%(message)s", only_log=False)#True) + + for folder in ['Micrographs', 'Movies', 'XML']: + logger.mkdir(f'CS/{folder}') + + fn = 'micrographs_ctf.star' + + with StarFile(fn) as sf: + with StarFile('CS/particles.star', 'w') as sfOut: + ctfCols = ['rlnDefocusU', 'rlnDefocusV', 'rlnDefocusAngle', 'rlnCtfFigureOfMerit', 'rlnCtfMaxResolution'] + ctfCols = [] # CS is giving an error when using CTF + t = Table(['rlnMicrographName', 'rlnCoordinateX', 'rlnCoordinateY'] + ctfCols) + print("cols", len(t.getColumnNames()), t.getColumnNames()) + + sfOut.writeHeader('particles', t) + + for row in sf.iterTable('micrographs'): + micFn = row.rlnMicrographName + movFn = row.rlnMicrographMovieName.replace('Images-Disc1_', '') + xmlFn = movFn.replace('_EER.eer', '.xml') + base = os.path.basename(micFn) + micName = base.replace('_DW.mrc', '') + movName = micName.replace('mic_', 'mov_') + logger.system(f'ln -s ../../{micFn} CS/Micrographs/{micName}.mrc') + logger.system(f'ln -s ../../{movFn} CS/Movies/{movName}.eer') + logger.system(f'ln -s ../../{xmlFn} CS/XML/{movName}.xml') + coordsFn = f'Coordinates/{micName}_DW_coordinates.star' + ctfValues = [getattr(row, k) for k in ctfCols] + print(len(ctfValues)) + if os.path.exists(coordsFn): + with StarFile(coordsFn) as sfCoords: + for rowCoord in sfCoords.iterTable(''): + sfOut.writeRow(t.Row(f'{micName}.mrc', + rowCoord.rlnCoordinateX, + rowCoord.rlnCoordinateY, + *ctfValues)) + + +def cryosparc_import(projId, dataRoot): + + acq = { + "psize_A": 0.724, + "accel_kv": 300, + "cs_mm": 0.1, + } + + cs = CryoSparc(projId) + csRoot = os.path.join(dataRoot, 'CS') + + print(f">>> Importing data from: {Color.green(dataRoot)}") + + args = { + "blob_paths": f"{csRoot}/Micrographs/mic_*.mrc", + "total_dose_e_per_A2": 40, + "parse_xml_files": True, + "xml_paths": f"{csRoot}/XML/mov_*.xml", + "mov_cut_prefix_xml": 4, + "mov_cut_suffix_xml": 4, + "xml_cut_prefix_xml": 4, + "xml_cut_suffix_xml": 4 + } + args.update(acq) + micsImport = cs.job_run("W1", "import_micrographs", args) + time.sleep(5) # FIXME: wait for job completion + args = { + "ignore_blob": True, + "particle_meta_path": f"{csRoot}/particles.star", + "query_cut_suff": 4, + "remove_leading_uid": True, + "source_cut_suff": 4, + "enable_validation": True, + "location_exists": True, + "amp_contrast": 2.7, + } + args.update(acq) + ptsImport = cs.job_run("W1", "import_particles", args, + {'micrographs': f'{micsImport}.imported_micrographs'}) + + def main(): p = argparse.ArgumentParser(prog='scipion-otf') g = p.add_mutually_exclusive_group() @@ -728,6 +904,15 @@ def main(): "and the Cryolo picking for picking. One can pass a string" "with the protocol ids for ctfs and/or picking. For example:" "--write_starts 'ctfs=1524 picking=1711'") + g.add_argument('--cs_prepare', action='store_true', + help="Prepare a folder CS to be used to import movies, micrographs " + "and particles into CryoSparc. ") + g.add_argument('--cs_import', nargs='+', + metavar=('CRYOSPARC_PROJECT_ID', 'DATA_ROOT'), + help="Import data from CS into a running project. ") + #get_scheduler_lanes + g.add_argument('--cs_test', metavar='CRYOSPARC_PROJECT_ID', + help="Test connection to CryoSparc server. ") g.add_argument('--clone_project', nargs=2, metavar=('SRC', 'DST'), help="Clone an existing Scipion project") g.add_argument('--fix_run_links', metavar='RUNS_SRC', @@ -762,6 +947,14 @@ def main(): fix_run_links(cwd, args.fix_run_links) elif protId := args.print_protocol: print_protocol(cwd, protId) + elif cs := args.cs_prepare: + cryosparc_prepare() + elif projId := args.cs_test: + cs = CryoSparc(projId) + elif cs := args.cs_import: + projId = cs[0] + dataRoot = cs[1] + cryosparc_import(projId, dataRoot) else: # by default open the GUI from pyworkflow.gui.project import ProjectWindow ProjectWindow(cwd).show() diff --git a/emtools/scripts/emt_files.py b/emtools/scripts/emt_files.py index 2b8b99f..21bf616 100755 --- a/emtools/scripts/emt_files.py +++ b/emtools/scripts/emt_files.py @@ -18,6 +18,7 @@ import os import time import argparse +import json from glob import glob from datetime import datetime, timedelta from pprint import pprint @@ -27,15 +28,112 @@ from emtools.metadata import EPU, MovieFiles +def scan_folder(folder): + """Scan a folder; return (files_dict, dirs_set). + files_dict: relative_path -> {size, mtime} + dirs_set: set of relative directory paths (including '.' for the root). + """ + folder = os.path.abspath(os.path.expanduser(folder)) + if not os.path.isdir(folder): + raise SystemExit(f"ERROR: Not a directory: {folder}") + files_result = {} + dirs_set = set() + for root, _dirs, files in os.walk(folder): + rel_root = os.path.relpath(root, folder) + if rel_root == '.': + dirs_set.add('.') + else: + dirs_set.add(rel_root) + for fn in files: + path = os.path.join(root, fn) + try: + st = os.stat(path) + except OSError: + continue + rel = os.path.relpath(path, folder) + files_result[rel] = {'size': st.st_size, 'mtime': st.st_mtime} + return files_result, dirs_set + + +def scan_save(folder, output_path): + """Scan folder and write snapshot to a JSON file.""" + files_snapshot, dirs_set = scan_folder(folder) + folder_abs = os.path.abspath(os.path.expanduser(folder)) + data = { + 'folder': folder_abs, + 'scanned_at': datetime.now().isoformat(), + 'files': files_snapshot, + 'dirs': sorted(dirs_set), + } + output_path = os.path.abspath(os.path.expanduser(output_path)) + os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True) + with open(output_path, 'w') as f: + json.dump(data, f, indent=2) + print(f"Scan saved: {len(files_snapshot)} files, {len(dirs_set)} dirs -> {output_path}") + + +def scan_compare(folder, compare_path): + """Scan folder and compare to a previously saved JSON snapshot.""" + folder_abs = os.path.abspath(os.path.expanduser(folder)) + compare_path = os.path.abspath(os.path.expanduser(compare_path)) + if not os.path.isfile(compare_path): + raise SystemExit(f"ERROR: Compare file not found: {compare_path}") + + with open(compare_path) as f: + data = json.load(f) + previous_files = data.get('files', data) if 'files' in data else data + if isinstance(previous_files, dict) and not previous_files and 'files' in data: + previous_files = data['files'] + previous_dirs = set(data.get('dirs', [])) + + current_files, current_dirs = scan_folder(folder) + prev_file_keys = set(previous_files) + curr_file_keys = set(current_files) + + new_files = sorted(curr_file_keys - prev_file_keys) + deleted_files = sorted(prev_file_keys - curr_file_keys) + modified = [] + for k in sorted(prev_file_keys & curr_file_keys): + p, c = previous_files[k], current_files[k] + if p.get('size') != c.get('size') or p.get('mtime') != c.get('mtime'): + modified.append(k) + + new_dirs = sorted(current_dirs - previous_dirs) + deleted_dirs = sorted(previous_dirs - current_dirs) + + def _report(label, items, color_fn=Color.red): + if not items: + return + print(color_fn(f"\n{label} ({len(items)}):")) + for rel in items: + print(f" {rel}") + + print(f"Comparison: current scan vs {compare_path}") + print(f" Files: previous {len(prev_file_keys)} | current {len(curr_file_keys)}") + print(f" Dirs: previous {len(previous_dirs)} | current {len(current_dirs)}") + _report("New folders", new_dirs, Color.green) + _report("Deleted folders", deleted_dirs, Color.red) + _report("New files", new_files, Color.green) + _report("Deleted files", deleted_files, Color.red) + _report("Modified files", modified, Color.red if modified else lambda x: x) + + if not new_files and not deleted_files and not modified and not new_dirs and not deleted_dirs: + print(Color.green("\nNo changes detected.")) + + def statsDir(folder, sort): df = MovieFiles() df.scan(folder) df.print(sort=sort) - df.counters[1].print('movie') -def timeStats(pattern, bin, plot): - files = glob(pattern) +def timeStats(pattern, bin, plot, data): + files = [] + if os.path.isdir(pattern): + for root, dirs, dfiles in os.walk(pattern): + files.extend(os.path.join(root, fn) for fn in dfiles) + else: + files = glob(pattern) total_size = 0 filesDict = {} @@ -51,6 +149,8 @@ def timeStats(pattern, bin, plot): first = fs[0] last = fs[-1] + to_GB = 1 / (1024 ** 3) + if bin: bindelta = timedelta(minutes=bin) start = datetime.fromtimestamp(first[1]['ts']) @@ -60,7 +160,8 @@ def timeStats(pattern, bin, plot): end = last_bin['end'] ts = datetime.fromtimestamp(v['ts']) if ts <= end: - last_bin['count'] += 1 + value = 1 if not data else v['size'] * to_GB + last_bin['count'] += value else: bins.append({'start': end, 'end': end + bindelta, @@ -117,7 +218,8 @@ def _addDt(b, onlyTime=False): w = width * 0.9 ax.bar(x + w / 2, values, w, label='Men') # Add some text for labels, title and custom x-axis tick labels, etc. - ax.set_ylabel('Files') + ylabel = 'Files' if not data else 'Data (Gb)' + ax.set_ylabel(ylabel) ax.set_title(f'Files generated every {bin} minutes') ax.set_xticks(x) ax.set_xticklabels(labels) @@ -151,26 +253,46 @@ def main(): g = p.add_mutually_exclusive_group() g.add_argument('--stats', '-s', metavar='FOLDER', help="Statistics of the files in a given folder.") - g.add_argument('--timing', metavar='PATTERN', + g.add_argument('--timing', metavar='FOLDER_OR_PATTERN', help="Compute histogram from the timestamps of files " - "matching the pattern.") + "in the folder or matching the pattern.") + g.add_argument('--count_movies', '-m', nargs='+', + help="Count number of movies for each input folder") g.add_argument('--copy_dir', nargs=2, metavar=('SRC_DIR', 'NEW_DIR'), help='Copy directory with some delay') g.add_argument('--check_dirs', nargs=2, metavar=('DIR1', 'DIR2'), help='Check if the two directories are synchronized. ') - + g.add_argument('--rsync_dirs', nargs=2, metavar=('DIR1', 'DIR2'), + help='Rsync both directories and print the number of ' + 'transferred files. ') + g.add_argument('--scan', metavar='FOLDER', + help='Scan folder. Use with --output to save snapshot to JSON, ' + 'or with --compare to diff against a saved snapshot.') + g.add_argument('--relink', nargs=2, metavar=('OLD_PREFIX', 'NEW_PREFIX'), + help='Relink the symbolic links in the current directory, changing the prefix to the new one') + g.add_argument('--transfer', nargs=3, metavar=('FRAMES_DIR', 'RAW_DIR', 'EPU_DIR'), + help='REVIEW: Transfer files from FRAMES_DIR to RAW_DIR and EPU_DIR') + + p.add_argument('--output', '-o', metavar='FILE', + help='Save scan snapshot to this JSON file (with --scan)') + p.add_argument('--compare', '-c', metavar='FILE', + help='Compare current scan to this JSON snapshot (with --scan)') p.add_argument('--bin', '-b', type=int, default=6000, help="Create bins of the given time in minutes " "(with --timing)") p.add_argument('--plot', '-p', action='store_true', help="Plot the number of files per bin " - "(with --stats)") + "(with --timing)") + p.add_argument('--data', '-a', action='store_true', + help="Use file size for the timing plot") p.add_argument('--delay', '-d', type=float, default=0, help="Delay in seconds when copying files " "(with --copy_dir)") p.add_argument('--sort', choices=['count', 'size'], help="Sort results from --stats with a folder" "based on count or size (with --stats FOLDER)") + p.add_argument('--dry-run', action='store_true', + help="Dry run, without actually performing the operation") args = p.parse_args() @@ -214,8 +336,31 @@ def _mkdir(d): s = Color.green('in SYNC') if sync else Color.red('NOT in SYNC') print(f"Dirs are {s}") + elif dirs := args.count_movies: + maxlen = max(len(d) for d in dirs) + def _pad(s): + return (maxlen - len(s)) * ' ' + s + + for d in dirs: + print(f"{_pad(d)}: {EPU.count_movies(d):>8}") + + elif dirs := args.rsync_dirs: + n = Path.rsync(dirs[0], dirs[1], verbose=True) + print(f"Transferred files: {n}") + elif pattern := args.timing: - timeStats(pattern, args.bin, args.plot) + timeStats(pattern, args.bin, args.plot, args.data) + + elif folder := args.scan: + if args.output and args.compare: + p.error("--scan: use either --output or --compare, not both") + elif args.output: + scan_save(folder, args.output) + elif args.compare: + scan_compare(folder, args.compare) + else: + p.error("--scan requires either --output FILE (save snapshot) " + "or --compare FILE (compare to snapshot)") # TODO: check from here elif args.transfer: @@ -258,15 +403,31 @@ def _moveFile(srcFile, dstFile): pprint(epuData.info()) - elif args.parse: - ed = Path.ExtDict() - for root, dirs, files in os.walk(args.parse): - for f in files: - srcFn = os.path.join(root, f) - if os.path.isfile(srcFn): - ed.register(os.path.join(root, f)) - ed.print() - + # elif args.parse: + # ed = Path.ExtDict() + # for root, dirs, files in os.walk(args.parse): + # for f in files: + # srcFn = os.path.join(root, f) + # if os.path.isfile(srcFn): + # ed.register(os.path.join(root, f)) + # ed.print() + + elif args.relink: + old_prefix, new_prefix = args.relink + cwd = os.getcwd() + print(f"Relinking files in {cwd} from {old_prefix} to {new_prefix}") + for fn in os.listdir(cwd): + filepath = os.path.join(cwd, fn) + if os.path.islink(filepath): + target = os.readlink(filepath) + if target.startswith(old_prefix): + new_target = target.replace(old_prefix, new_prefix) + print(f"LINK: {Color.bold(filepath)}\n" + f" OLD: {Color.red(target)}\n" + f" NEW: {Color.green(new_target)}") + if not args.dry_run: + os.unlink(filepath) + os.symlink(new_target, filepath) if __name__ == '__main__': main() diff --git a/emtools/scripts/emt_ps.py b/emtools/scripts/emt_ps.py index bcc4bab..9ad65ef 100755 --- a/emtools/scripts/emt_ps.py +++ b/emtools/scripts/emt_ps.py @@ -55,44 +55,9 @@ def main(): print(System.hostname()) sys.exit(0) - v = args.verbose - - kill = args.kill folderPath = os.path.abspath(args.folder) if args.folder else args.folder print('path', folderPath) - processes = Process.ps(args.name, workingDir=folderPath, children=args.children) - - color = Color.red if kill else Color.bold - - for folder, procs in processes.items(): - print(Color.warn(f"{folder}")) - header = f" {'USER':<15} {'PPID/PID':<15} {color('PROGRAM'):<30}" - if v > 0: - header += f" {'CPU(%)':>10} {'MEMORY(%)':>10}" - if v > 1: - header += f" {'COMMAND LINE'}" - - print(Color.bold(header)) - - prefix = 'Killing' if kill else '' - for p in procs: - pidstr = f"{p.info['ppid']}/{p.pid}" - msg = f" {prefix} {p.info['username']:<15} {pidstr:<15} {color(p.info['name']):<30}" - if v > 0: - try: - cpu_percent = p.cpu_percent(interval=1) / cpus - except: - continue - - msg += f" {cpu_percent:>10,.2f} {p.info['memory_percent']:>10,.2f}" - if v > 1: - msg += f" {p.cmdline()}" - print(msg) - if kill: - try: - p.kill() - except: - pass + Process.checkChilds(args.name, folderPath, kill=args.kill, verbose=args.verbose) if __name__ == '__main__': diff --git a/emtools/scripts/emt_star.py b/emtools/scripts/emt_star.py new file mode 100755 index 0000000..0da65a2 --- /dev/null +++ b/emtools/scripts/emt_star.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + +import os +import sys +import time +import argparse +from glob import glob +from datetime import datetime, timedelta +from pprint import pprint +import numpy as np +from collections import defaultdict + +from emtools.utils import Process, Color, Path, Timer, Pretty +from emtools.metadata import StarFile, Table + + +ALL_TABLES = 'all' +OP_DROP = 'DROP' +OP_UPDATE = 'UPDATE' +OP_FILTER = 'FILTER' +VALID_OPERATIONS = {OP_DROP, OP_UPDATE, OP_FILTER} + + +def _normalizeOperation(op): + op = op.strip().upper() + if op not in VALID_OPERATIONS: + raise ValueError( + f"Unknown operation '{op}'. Valid values: {', '.join(sorted(VALID_OPERATIONS))}") + return op + + +def _resolveTableName(table, tableNames): + """Resolve table name; only the special 'all' token is case-insensitive.""" + table = table.strip() + if table.lower() == ALL_TABLES: + return ALL_TABLES + if table not in tableNames: + raise ValueError(f"Table '{table}' not found in input STAR file") + return table + + +def _parseOperateArgs(operateArgs): + """Flatten --operate argument groups into (TABLE, OPERATION, EXPR) tuples.""" + operations = [] + for group in operateArgs: + if len(group) % 3 != 0: + raise ValueError( + "--operate expects groups of TABLE OPERATION EXPR " + f"(multiple of 3 arguments), got {len(group)} in one --operate") + for i in range(0, len(group), 3): + operations.append((group[i], group[i + 1], group[i + 2])) + return operations + + +def _parseOperateActions(operations, tableNames): + """Build per-table action lists and the set of tables to drop.""" + explicitTables = set() + dropTables = set() + tableActions = defaultdict(list) + allActions = [] + allDrop = False + + for table, operation, expr in operations: + tableName = _resolveTableName(table, tableNames) + operation = _normalizeOperation(operation) + + if tableName == ALL_TABLES: + if operation == OP_DROP: + allDrop = True + else: + allActions.append((operation, expr)) + continue + + explicitTables.add(tableName) + + if operation == OP_DROP: + dropTables.add(tableName) + else: + tableActions[tableName].append((operation, expr)) + + if allDrop: + for tableName in tableNames: + if tableName not in explicitTables: + dropTables.add(tableName) + + return explicitTables, dropTables, tableActions, allActions + + +def _iterTableRows(sf, tableName, subset=None): + kwargs = {'limit': subset} if subset is not None else {} + return list(sf.iterTable(tableName, **kwargs)) + + +def _writeProcessedTable(sfOut, tableName, tableInfo, rows, singleRow): + if not rows: + sfOut.writeTable(tableName, tableInfo) + return + + if singleRow: + sfOut.writeSingleRow(tableName, rows[0]) + else: + result = tableInfo.cloneColumns() + for row in rows: + result.addRow(row) + sfOut.writeTable(tableName, result) + + +def _writeUnchangedTable(sfIn, sfOut, tableName, subset=None): + sfIn.getTableInfo(tableName) + singleRow = sfIn._singleRow + rows = _iterTableRows(sfIn, tableName, subset=subset) + tableInfo = sfIn.getTableInfo(tableName) + _writeProcessedTable(sfOut, tableName, tableInfo, rows, singleRow) + + +def _processTable(sf, tableName, actions, subset=None): + tableInfo = sf.getTableInfo(tableName) + singleRow = sf._singleRow + table = tableInfo.cloneColumns() + kwargs = {'limit': subset} if subset is not None else {} + for row in sf.iterTable(tableName, **kwargs): + table.addRow(row) + + for operation, expr in actions: + if operation == OP_UPDATE: + table.update(expr) + elif operation == OP_FILTER: + table.filter(expr) + + return tableInfo, list(table), singleRow + + +def mergeStarFiles(pattern, tableName, output=None): + files = sorted(glob(pattern)) + if not files: + raise FileNotFoundError(f"No STAR files match pattern: {pattern}") + + tableInfo = None + refColumnNames = None + singleRow = None + rows = [] + + for starPath in files: + with StarFile(starPath) as sf: + if tableName not in sf.getTableNames(): + raise ValueError( + f"Table '{tableName}' not found in {starPath}") + + info = sf.getTableInfo(tableName) + columnNames = info.getColumnNames() + fileSingleRow = sf._singleRow + + if tableInfo is None: + tableInfo = info + refColumnNames = columnNames + singleRow = fileSingleRow + else: + if len(columnNames) != len(refColumnNames): + raise ValueError( + f"Column count mismatch in {starPath}: " + f"expected {len(refColumnNames)}, got {len(columnNames)}") + if columnNames != refColumnNames: + raise ValueError( + f"Column names mismatch in {starPath}: " + f"expected {refColumnNames}, got {columnNames}") + if fileSingleRow != singleRow: + raise ValueError( + f"Table layout mismatch in {starPath}: " + f"expected {'single-row' if singleRow else 'loop'} " + f"format") + + for row in sf.iterTable(tableName): + rows.append(row) + + if len(rows) > 1: + singleRow = False + + closeOutput = output is not None + out = open(output, 'w') if closeOutput else sys.stdout + try: + with StarFile(out, closeFile=closeOutput) as sfOut: + sfOut.writeTimeStamp() + _writeProcessedTable(sfOut, tableName, tableInfo, rows, singleRow) + finally: + if closeOutput: + out.close() + + +def operateStarFile(inputStar, operations, subset=None, output=None): + if not operations: + raise ValueError("At least one --operate action is required") + + if not os.path.exists(inputStar): + raise FileNotFoundError(f"Input star file does not exist: {inputStar}") + + closeOutput = output is not None + out = open(output, 'w') if closeOutput else sys.stdout + + try: + with StarFile(inputStar) as sfIn: + tableNames = sfIn.getTableNames() + explicitTables, dropTables, tableActions, allActions = ( + _parseOperateActions(operations, tableNames)) + + with StarFile(out, closeFile=closeOutput) as sfOut: + sfOut.writeTimeStamp() + + for tableName in tableNames: + if tableName in dropTables: + continue + + if tableName in explicitTables: + actions = tableActions.get(tableName, []) + else: + actions = list(allActions) + + if not actions: + _writeUnchangedTable(sfIn, sfOut, tableName, subset=subset) + else: + tableInfo, rows, singleRow = _processTable( + sfIn, tableName, actions, subset=subset) + _writeProcessedTable(sfOut, tableName, tableInfo, rows, singleRow) + finally: + if closeOutput: + out.close() + + +def printStarInfo(starFile): + with StarFile(starFile) as sf: + tables = sf.getTableNames() + for t in tables: + cols = sf.getTableInfo(t).getColumnNames() + tSize = sf.getTableSize(t) + print(f">>> {Color.bold('Table')}: {Color.green(t)}" + f"\n - Columns: {Color.cyan(len(cols))} [{' '.join(c for c in cols)}]" + f"\n - Rows: {Color.cyan(tSize)}") + + +def groupBy(starFile, table, column): + group = defaultdict(lambda: 0) + + with StarFile(starFile) as sf: + for row in sf.iterTable(table): + group[row.get(column)] += 1 + + for k, v in group.items(): + print(k, v) + + +def checkDuplicates(inputStar, table, column): + items = set() + duplicates = [] + + with StarFile(inputStar) as sf: + for row in sf.iterTable(table): + value = row.get(column) + if value in items: + duplicates.append(value) + else: + items.add(value) + + print(f">>> Duplicates: {len(duplicates)}\n" + f" {duplicates}") + +def printColumns(inputStar, tableName=None, columns=None, subset=None): + if not os.path.exists(inputStar): + raise FileNotFoundError(f"Input star file does not exist: {inputStar}") + + with StarFile(inputStar) as sf: + existingTables = sf.getTableNames() + if tableName is None: + if not existingTables: + return + tableName = existingTables[0] + elif tableName not in existingTables: + raise ValueError(f"Table name does not exist: {tableName}") + + columnList = columns.split() if columns else sf.getTableInfo(tableName).getColumnNames() + table = _buildPrintTable(sf, tableName, columnList, subset=subset) + StarFile.printTable(table, tableName) + + +def printAllTables(inputStar, subset=None): + if not os.path.exists(inputStar): + raise FileNotFoundError(f"Input star file does not exist: {inputStar}") + + with StarFile(inputStar) as sf: + for tableName in sf.getTableNames(): + tableInfo = sf.getTableInfo(tableName) + table = _buildPrintTable(sf, tableName, tableInfo.getColumnNames(), + subset=subset) + StarFile.printTable(table, tableName) + + +def _buildPrintTable(sf, tableName, columnList, subset=None): + tableInfo = sf.getTableInfo(tableName) + cols = [col for col in tableInfo.getColumns() if col.getName() in columnList] + newTable = Table(columns=cols) + kwargs = {'limit': subset} if subset is not None else {} + for row in sf.iterTable(tableName, **kwargs): + values = {k: getattr(row, k) for k in columnList} + newTable.addRowValues(**values) + return newTable + + +def splitBy(starFile, column, minSize): + with StarFile(starFile) as sf: + tOptics = sf.getTable('optics') + tParticles = sf.getTableInfo('particles') + rows = [] + count = 0 + map = {} + + def _writeStar(minSize=0): + nonlocal count + nonlocal rows + + if len(rows) <= minSize: + return + + count += 1 + outStarFile = Path.replaceExt(starFile, f'_{count:03}.star') + with StarFile(outStarFile, 'w') as sfOut: + sfOut.writeTimeStamp() + sfOut.writeTable('optics', tOptics) + sfOut.writeHeader('particles', tParticles) + for row in rows: + sfOut.writeRow(row) + rows = [] + + lastValue = None + lastIndex = 0 + + for row in sf.iterTable('particles'): + value = getattr(row, column) + if lastValue is not None and lastValue != value: + _writeStar(int(minSize)) + rows.append(row) + lastValue = value + + if rows: + _writeStar(0) # Write all remaining + + +def main(): + p = argparse.ArgumentParser(prog='emt-star') + p.add_argument('input', + help="Input STAR file, or a glob pattern when using --merge.") + p.add_argument('--group_by', '-g', nargs=2, + metavar=('TABLE', 'COLUMN'), + help="Count rows grouped by a given label") + p.add_argument('--split_particles', '-s', nargs='+', metavar=('COLUMN', 'minsize'), + help="Split input particles by some column") + p.add_argument('--duplicates', '-d', nargs=2, + metavar=('TABLE', 'COLUMN'), + help="Check duplicates values for a given label") + outputMode = p.add_mutually_exclusive_group() + outputMode.add_argument('--print', '-p', nargs='*', + metavar=('COLUMNS', 'TABLE'), + help="Print columns from STAR file tables to stdout. " + "With no arguments, print all tables and all columns. " + "With COLUMNS only, print from the first table. " + "With COLUMNS and TABLE, print from the given table.") + outputMode.add_argument('--operate', '-e', action='append', nargs='+', + metavar='TRIPLET', + help="Apply one or more operations. Each operation is a triplet " + "TABLE OPERATION EXPR; multiple triplets can be passed in a " + "single --operate. TABLE is the exact table name from the " + "input STAR file, or 'all' (case insensitive) for all tables " + "not explicitly listed in other operations. OPERATION can be " + "UPDATE, FILTER, or DROP. For UPDATE, EXPR is comma-separated " + "column=expression assignments. For FILTER, EXPR is a boolean " + "expression per row. For DROP, EXPR is ignored.") + outputMode.add_argument('--merge', '-m', metavar='TABLE', + help="Merge TABLE from every STAR file matching the input " + "glob pattern into one table.") + p.add_argument('--subset', '-n', type=int, default=None, metavar='N', + help="Process at most N rows per table (for debugging)") + p.add_argument('--output', '-o', default=None, metavar='FILE', + help="Write output STAR file to this path (default: stdout)") + + args = p.parse_args() + inputStar = args.input + + if args.operate: + operateStarFile(inputStar, _parseOperateArgs(args.operate), + subset=args.subset, output=args.output) + elif args.merge: + mergeStarFiles(inputStar, args.merge, output=args.output) + elif args.group_by: + table, column = args.group_by + groupBy(inputStar, table, column) + elif split := args.split_particles: + column = split[0] + minSize = split[1] if len(split) > 1 else 0 + splitBy(inputStar, column, minSize) + elif args.duplicates: + table, column = args.duplicates + checkDuplicates(inputStar, table, column) + elif args.print is not None: + if len(args.print) == 0: + printAllTables(inputStar, subset=args.subset) + else: + tableName = None + cols = args.print[0] + if len(args.print) > 2: + raise ValueError("Only pass columns and optionally the table name") + elif len(args.print) > 1: + tableName = args.print[1] + + printColumns(inputStar, tableName, cols, subset=args.subset) + else: + printStarInfo(args.input) + + +if __name__ == '__main__': + main() diff --git a/emtools/scripts/emt_sysinfo.py b/emtools/scripts/emt_sysinfo.py new file mode 100755 index 0000000..5a5768e --- /dev/null +++ b/emtools/scripts/emt_sysinfo.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + +""" +Quick summary of this workstation: OS/distro, kernel, CPUs, RAM, disk and +GPUs. Mostly a thin, friendly wrapper around emtools.utils.System. +""" + +import json +import argparse +import platform + +from emtools.utils import Color, Pretty, System + + +def get_info(disk_path='/'): + """ Collect a dict with a quick summary of this workstation. + Reuses emtools.utils.System for all the underlying lookups. """ + return { + 'hostname': System.hostname(), + 'os': System.distro(), + 'kernel': System.kernel(), + 'arch': platform.machine(), + 'cpus': System.cpus(), + 'memory_gb': System.memory(), + 'disk': System.disk(disk_path), + 'disk_path': disk_path, + 'gpus': System.gpus(), + } + + +def print_info(info): + """ Print a nicely formatted report from the dict returned by get_info(). """ + width = 70 + title = f" SYSTEM INFO: {info['hostname']} " + print(Color.bold(title.center(width, '='))) + print(f" {'OS':<10}: {info['os']}") + print(f" {'Kernel':<10}: {info['kernel']} ({info['arch']})") + print(f" {'CPUs':<10}: {info['cpus']}") + print(f" {'Memory':<10}: {info['memory_gb']} GB") + + disk = info['disk'] + if disk: + pct = 100 * disk['used'] / disk['total'] if disk['total'] else 0 + print(f" {'Disk (' + info['disk_path'] + ')':<10}: " + f"{Pretty.size(disk['total'])} total, " + f"{Pretty.size(disk['free'])} free ({pct:.0f}% used)") + + gpus = info['gpus'] + print(f" {'GPUs':<10}: {len(gpus)}") + for g in gpus: + line = (f" [{g.get('index', '?')}] {g.get('name', 'Unknown GPU'):<28} " + f"{g.get('memory.total', '?'):>10} total " + f"{g.get('memory.used', '?'):>10} used " + f"driver {g.get('driver_version', '?')}") + print(Color.cyan(line)) + + print(Color.bold('=' * width)) + + +def main(): + p = argparse.ArgumentParser( + prog='emt-sysinfo', + description="Print a quick summary of this workstation's hardware " + "and OS: Linux distro/version, kernel, CPUs, RAM, disk " + "usage and GPUs.") + p.add_argument('--json', '-j', action='store_true', + help='Print the info as JSON instead of the formatted report.') + p.add_argument('--disk-path', default='/', + help='Path used to report disk usage (default: /).') + + args = p.parse_args() + info = get_info(disk_path=args.disk_path) + + if args.json: + print(json.dumps(info, indent=2)) + else: + print_info(info) + + +if __name__ == '__main__': + main() diff --git a/emtools/tests/test_image.py b/emtools/tests/test_image.py new file mode 100644 index 0000000..d5a9681 --- /dev/null +++ b/emtools/tests/test_image.py @@ -0,0 +1,58 @@ +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + +import os +import unittest +import tempfile +import random +import time +import threading +import tempfile +from pprint import pprint +from datetime import datetime + +from emtools.utils import Timer, Color, Pretty +from emtools.metadata import StarFile, SqliteFile, EPU, StarMonitor +from emtools.jobs import BatchManager +from emtools.tests import testpath +from emtools.image import Image + +from .star_pipeline_tester import StarPipelineTester + + +class TestImage(unittest.TestCase): + """ + Tests for Image class. + """ + + def test_dimensions(self): + """ + Read a star file with several blocks + """ + names = ['May08_03.05.02.bin.mrc', + 'gain.mrc', + '20170629_00021_frameImage.tiff'] + dims = [(1240, 1200, 50), + (3710, 3838), + (3710, 3838, 24)] + files = [testpath('movies', n) for n in names] + + if any(f is None for f in files): + return + + for f, d in zip(files, dims): + self.assertEqual(Image.get_dimensions(f), d) + diff --git a/emtools/tests/test_metadata.py b/emtools/tests/test_metadata.py index cdf9ab4..56797a4 100644 --- a/emtools/tests/test_metadata.py +++ b/emtools/tests/test_metadata.py @@ -24,7 +24,7 @@ from datetime import datetime from emtools.utils import Timer, Color, Pretty -from emtools.metadata import StarFile, SqliteFile, EPU, StarMonitor +from emtools.metadata import StarFile, SqliteFile, EPU, StarMonitor, RelionStar, Table from emtools.jobs import BatchManager from emtools.tests import testpath @@ -373,6 +373,216 @@ def _pipeline(monitor): #self.__test_star_streaming(_pipeline, inputStreaming=True) self.__test_star_streaming(_pipeline, inputStreaming=False) + def test_table_offsets_cache(self): + ftmp = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.star') + ftmp.write(""" +# version 50001 + +data_general + +_rlnTomoSubTomosAre2DStacks 1 + +# version 50001 + +data_optics + +loop_ +_rlnOpticsGroup #1 +_rlnOpticsGroupName #2 + 1 opticsGroup1 + 2 opticsGroup2 + +# version 50001 + +data_particles + +loop_ +_rlnTomoName #1 +_rlnTomoParticleId #2 +TS_01 1 +TS_01 2 +TS_02 3 +""") + ftmp.close() + + with StarFile(ftmp.name) as sf: + self.assertEqual(sf.getTableNames(), ['general', 'optics', 'particles']) + self.assertEqual(set(sf._offsets), + {'data_general', 'data_optics', 'data_particles'}) + + # Read out of order to exercise seeking from cached offsets + self.assertEqual(len(sf.getTable('particles')), 3) + self.assertEqual(sf.getTable('general')[0].rlnTomoSubTomosAre2DStacks, 1) + self.assertEqual(sf.getTableSize('optics'), 2) + self.assertEqual(sf.getTableInfo('particles').getColumnNames(), + ['rlnTomoName', 'rlnTomoParticleId']) + self.assertIsNone(sf.getTable('missing')) + + os.unlink(ftmp.name) + + +class TestTable(unittest.TestCase): + """Tests for Table.update and Table.filter.""" + + @staticmethod + def _sampleTable(): + t = Table(['rlnA', 'rlnB', 'rlnStatus']) + t.addRowValues(1, 10, 'Scheduled') + t.addRowValues(2, 20, 'Scheduled') + t.addRowValues(3, 30, 'Finished') + return t + + def test_update_constant(self): + t = self._sampleTable() + t.update('rlnB=100') + self.assertEqual(t.size(), 3) + for row in t: + self.assertEqual(row.rlnB, 100) + + def test_update_from_column_expression(self): + t = self._sampleTable() + t.update('rlnB = rlnA * 10') + self.assertEqual(t[0].rlnB, 10) + self.assertEqual(t[1].rlnB, 20) + self.assertEqual(t[2].rlnB, 30) + + def test_update_multiple_columns(self): + t = self._sampleTable() + t.update('rlnB = rlnA * 10, rlnStatus = "Updated"') + self.assertEqual(t[0].rlnB, 10) + self.assertEqual(t[0].rlnStatus, 'Updated') + self.assertEqual(t[2].rlnB, 30) + + def test_update_later_assignments_see_earlier_updates(self): + t = Table(['rlnA', 'rlnB', 'rlnC']) + t.addRowValues(2, 0, 0) + t.update('rlnB = rlnA * 10, rlnC = rlnB + 1') + self.assertEqual(t[0].rlnB, 20) + self.assertEqual(t[0].rlnC, 21) + + def test_update_whitespace_in_spec(self): + t = self._sampleTable() + t.update(' rlnB = rlnA * 10 , rlnStatus = "Updated" ') + self.assertEqual(t[1].rlnB, 20) + self.assertEqual(t[1].rlnStatus, 'Updated') + + def test_update_returns_self(self): + t = self._sampleTable() + self.assertIs(t.update('rlnB=0'), t) + + def test_update_raises_on_invalid_spec(self): + t = self._sampleTable() + for spec in ['', 'rlnB', 'rlnB=1=2', 'rlnB==1']: + with self.subTest(spec=spec): + with self.assertRaises(ValueError): + t.update(spec) + + def test_update_raises_on_unknown_column(self): + t = self._sampleTable() + with self.assertRaisesRegex(ValueError, "not found"): + t.update('missingCol=1') + + def test_filter_by_numeric_condition(self): + t = self._sampleTable() + t.filter('rlnA > 1') + self.assertEqual(t.size(), 2) + self.assertEqual([row.rlnA for row in t], [2, 3]) + + def test_filter_by_string_condition(self): + t = self._sampleTable() + t.filter('rlnStatus == "Scheduled"') + self.assertEqual(t.size(), 2) + self.assertTrue(all(row.rlnStatus == 'Scheduled' for row in t)) + + def test_filter_returns_self(self): + t = self._sampleTable() + self.assertIs(t.filter('rlnA > 0'), t) + + def test_filter_raises_on_empty_spec(self): + t = self._sampleTable() + with self.assertRaisesRegex(ValueError, "non-empty"): + t.filter(' ') + + def test_filter_raises_on_bare_column_name(self): + t = self._sampleTable() + with self.assertRaisesRegex(ValueError, "column name only"): + t.filter('rlnA') + + def test_update_then_filter_chaining(self): + t = self._sampleTable() + t.update('rlnB = rlnA * 10').filter('rlnB >= 20') + self.assertEqual(t.size(), 2) + self.assertEqual([row.rlnA for row in t], [2, 3]) + self.assertEqual([row.rlnB for row in t], [20, 30]) + + +class TestRelionStarTomo(unittest.TestCase): + """Tests for Relion tomography STAR file validation helpers.""" + + def _write_star(self, tables): + ftmp = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.star') + with StarFile(ftmp.name, 'w') as sf: + for table_name, table in tables.items(): + sf.writeTable(table_name, table, timeStamp=False, singleRow=len(table) <= 1) + ftmp.close() + return ftmp.name + + def test_isTomoOptimisationSet(self): + opt_star = self._write_star({ + 'optimisation_set': Table.fromDict({ + 'rlnTomoParticlesFile': 'particles.star', + 'rlnTomoTomogramsFile': 'tomograms.star', + }), + }) + wrong_table = self._write_star({ + 'global': Table.fromDict({'rlnTomoName': 'tomo1'}), + }) + missing_link = self._write_star({ + 'optimisation_set': Table.fromDict({'rlnTomoTomogramsFile': 'tomograms.star'}), + }) + + self.assertTrue(RelionStar.isTomoOptimisationSet(opt_star)) + self.assertFalse(RelionStar.isTomoOptimisationSet(wrong_table)) + self.assertFalse(RelionStar.isTomoOptimisationSet(missing_link)) + self.assertFalse(RelionStar.isTomoOptimisationSet(__file__)) + + row = RelionStar.readTomoOptimisationSet(opt_star)[0] + self.assertEqual(row.rlnTomoParticlesFile, 'particles.star') + + for fn in (opt_star, wrong_table, missing_link): + os.unlink(fn) + + def test_isTomoParticles(self): + particles_star = self._write_star({ + 'particles': Table.fromDict({ + 'rlnTomoName': 'tomo1', + 'rlnCoordinateX': 1.0, + 'rlnCoordinateY': 2.0, + 'rlnCoordinateZ': 3.0, + }), + }) + centered_star = self._write_star({ + 'particles': Table.fromDict({ + 'rlnTomoName': 'tomo1', + 'rlnCenteredCoordinateXAngst': 10.0, + 'rlnCenteredCoordinateYAngst': 20.0, + 'rlnCenteredCoordinateZAngst': 30.0, + }), + }) + missing_coords = self._write_star({ + 'particles': Table.fromDict({'rlnTomoName': 'tomo1'}), + }) + + self.assertTrue(RelionStar.isTomoParticles(particles_star)) + self.assertTrue(RelionStar.isTomoParticles(centered_star)) + self.assertFalse(RelionStar.isTomoParticles(missing_coords)) + + table = RelionStar.readTomoParticles(particles_star) + self.assertEqual(table[0].rlnTomoName, 'tomo1') + + for fn in (particles_star, centered_star, missing_coords): + os.unlink(fn) + class TestEPU(unittest.TestCase): """ Tests for EPU class. """ diff --git a/emtools/tests/test_pipeline.py b/emtools/tests/test_pipeline.py index 33687c5..d1bfdb1 100644 --- a/emtools/tests/test_pipeline.py +++ b/emtools/tests/test_pipeline.py @@ -18,9 +18,12 @@ import numpy as np import time + +from emtools.utils import Color from emtools.jobs import Pipeline + class TestThreading(unittest.TestCase): def test_threads_processors(self): @@ -65,4 +68,28 @@ def picking(mic): pipeline.run() + print("PROCESSING DONE!!!") + + def test_queueMaxSize(self): + def generate(): + n = 8 + for i in range(1, n+1): + batch = "batch_%03d" % i + print("Generated batch: %s" % Color.green(batch)) + yield batch + time.sleep(1) + + def process(batch): + print("Processing batch: %s" % Color.warn(batch)) + time.sleep(8) + return batch + + pipeline = Pipeline(debug=False) + + g = pipeline.addGenerator(generate, + name='GENERATOR', + queueMaxSize=2) + + pipeline.addProcessor(g.outputQueue, process, name='PROC') + pipeline.run() print("PROCESSING DONE!!!") \ No newline at end of file diff --git a/emtools/tests/test_workflow.py b/emtools/tests/test_workflow.py new file mode 100644 index 0000000..2c1282c --- /dev/null +++ b/emtools/tests/test_workflow.py @@ -0,0 +1,48 @@ +# ************************************************************************** +# * +# * Authors: J.M. de la Rosa Trevin (delarosatrevin@gmail.com) +# * +# * This program is free software; you can redistribute it and/or modify +# * it under the terms of the GNU General Public License as published by +# * the Free Software Foundation; either version 3 of the License, or +# * (at your option) any later version. +# * +# * This program is distributed in the hope that it will be useful, +# * but WITHOUT ANY WARRANTY; without even the implied warranty of +# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# * GNU General Public License for more details. +# * +# ************************************************************************** + +import unittest +import numpy as np +import time + + +from emtools.utils import Color +from emtools.jobs import Pipeline, Workflow + + +class TestWorkflow(unittest.TestCase): + def test_basic(self): + wf = Workflow() + + j1 = wf.registerJob('job01') + d1 = j1.registerOutput('d1') + j2 = wf.registerJob('job02', inputs=[d1]) + j3 = wf.registerJob('job03', inputs=[d1]) + d3a = j3.registerOutput('d3a') + d3b = j3.registerOutput('d3b') + j5 = wf.registerJob('job05') + d5 = j5.registerOutput('d5') + j6 = wf.registerJob('job06', inputs=[d3b, d5]) + + dot = wf.dot() + + def test_relion_pipeline(self): + pipelineStar = '/Users/jdela80/work/data/emwrap/testing/Relion5-Tutorial-emwrap/default_pipeline.star' + + wf = Workflow.fromRelionPipeline(pipelineStar) + + print("\n") + print(wf.dot()) diff --git a/emtools/utils/__init__.py b/emtools/utils/__init__.py index 01f256c..e3c1f3d 100644 --- a/emtools/utils/__init__.py +++ b/emtools/utils/__init__.py @@ -19,12 +19,12 @@ from .time import Timer from .process import Process -from .path import Path -from .system import System +from .path import Path, FolderManager +from .system import System, GpuMonitor from .server import JsonTCPServer, JsonTCPClient -__all__ = ["Color", "Pretty", "Timer", "Process", "Path", "System", - "JsonTCPServer", "JsonTCPClient"] +__all__ = ["Color", "Pretty", "Timer", "Process", "Path", "FolderManager", + "System", "JsonTCPServer", "JsonTCPClient", "GpuMonitor"] diff --git a/emtools/utils/path.py b/emtools/utils/path.py index e256170..243f7b7 100644 --- a/emtools/utils/path.py +++ b/emtools/utils/path.py @@ -15,12 +15,30 @@ # ************************************************************************** import os +import shutil import time +import tempfile +import json +import hashlib +from glob import glob from datetime import datetime as dt from collections import OrderedDict +from contextlib import contextmanager from .pretty import Pretty from .process import Process +from .color import Color + + +GLOB_CHARS = ['*', '?', '[', ']'] + +IMAGE_EXT = ['tiff', 'tif', 'png', 'jpg', 'jpeg'] +EM_EXT = ['mrc', 'mrcs', 'eer', 'gain'] +TEXT_EXT = ['txt', 'log', 'err', 'out', 'json', 'csv', + 'star', 'sh', 'out', 'err', 'bashrc', 'xml', + 'script', 'settings', 'job', 'tomostar', 'mdoc', + 'population', 'species', 'yaml', 'id', + 'aln', 'com', 'rawtlt', 'tlt', 'xf', 'xtilt'] class Path: @@ -79,7 +97,7 @@ def splitall(path): @staticmethod def addslash(path): - """ Add an slash (/) to the end of the path if not present. """ + """ Add a slash (/) to the end of the path if not present. """ return path if path.endswith('/') else path + '/' @staticmethod @@ -93,19 +111,50 @@ def inSync(dir1, dir2, verbose=False): Use rsync as a subprocess to check if the two directories are synchronized. Both directories must exist. """ + return Path.rsync(dir1, dir2, '--dry-run', verbose=verbose) == 0 + + @staticmethod + def rsync(dir1, dir2, *args, + verbose=False, + size=False): + """ Run rsync to synchronize dir1 and dir2 are synchronized (i.e. same content) + Use rsync as a subprocess to synchronize dir1 and dir2 and return + the number of files transferred. + Args: + dir1: source directory + dir2: destination directory + *args: extra arguments to rsync + verbose: If True, print the command to stdout + size: If True, a tuple is returned with transferred files and transferred data size + """ dir1 = Path.addslash(dir1) dir2 = Path.addslash(dir2) - p = Process('rsync', '--dry-run', '-a', '--stats', dir1, dir2) + cmd = ['rsync', '-a', '--stats'] + list(args) + [dir1, dir2] + p = Process(*cmd, doRaise=True) + if verbose: p.print(stdout=True) - transf = 1 + def _value(line): + # Get the value after the colon (:) + # and remove , that is used to separate thousands + v = line.split(':')[1].replace(',', '') + if ' ' in v: # MacOS have a different rsync output format + v = v.strip().split()[0] + return int(v) + + transf = 0 + transfSize = 0 + for line in p.lines(): - if 'files transferred:' in line: - transf = int(line.split(':')[1]) - break - return transf == 0 + if 'Number of regular files transferred:' in line: + transf = _value(line) + elif 'Total transferred file size:' in line: + transfSize = _value(line.replace('bytes', '')) + + return (transf, transfSize) if size else transf + @staticmethod def lastModified(folder): @@ -115,11 +164,15 @@ def lastModified(folder): for fn in files: f = os.path.join(folder, fn) - s = os.stat(f) - t = (f, s.st_mtime) - last = t if not last or s.st_mtime > last[1] else last + if os.path.exists(f): + s = os.stat(f) + t = (f, s.st_mtime) + last = t if not last or s.st_mtime > last[1] else last - return last[0], dt.fromtimestamp(last[1]) + if last: + return last[0], dt.fromtimestamp(last[1]) + else: + return None, None @staticmethod def copyFile(file1, file2, sleep=0): @@ -132,7 +185,6 @@ def copyFile(file1, file2, sleep=0): f2.write(rbytes) if sleep: time.sleep(sleep) - #Process.system(f'cp {file1} {file2}') @staticmethod def copyDir(dir1, dir2, copyFileFunc=None, pl=None, **kwargs): @@ -162,6 +214,31 @@ def _mkdir(d): for f in files: _copy(os.path.join(root, f), os.path.join(root2, f), **kwargs) + @staticmethod + @contextmanager + def tmpDir(**kwargs): + tmp = tempfile.mkdtemp(prefix=kwargs.get('prefix', '')) + + chdir = kwargs.get('chdir', False) + cwd = os.getcwd() + if chdir: + os.chdir(tmp) + + if kwargs.get('verbose', True): + print(f"Using temporary dir: {tmp}") + + yield tmp + + if chdir: + os.chdir(cwd) + + globalClean = int(os.environ.get('EMWRAP_CLEAN', 1)) + if kwargs.get('clean', globalClean): + shutil.rmtree(tmp) + else: + print(f"Temporary directory was not deleted, " + f"remove it with the following command: \n" + f"{Color.bold('rm -rf %s' % tmp)}") @staticmethod def replaceExt(filename, newExt): @@ -198,3 +275,137 @@ def exists(path): """ return path and os.path.exists(path) + @staticmethod + def isPattern(path): + return any(c in path for c in GLOB_CHARS) + + @staticmethod + def isImage(path): + return Path.getExt(path).lower()[1:] in IMAGE_EXT + + @staticmethod + def isText(path): + return Path.getExt(path).lower()[1:] in TEXT_EXT + + @staticmethod + def isEmImage(path): + return Path.getExt(path).lower()[1:] in EM_EXT + + @staticmethod + def computeHashDict(path, verbose=False): + """ Get the hash of a file. """ + import hashlib + + result = {} + + # Ensure the input path is absolute for consistent splitting + base_path = os.path.abspath(path) + + for root, dirs, files in os.walk(base_path): + # 1. Handle folder entries (directories) + for dir_name in dirs: + dir_full_path = os.path.join(root, dir_name) + # Calculate path relative to the input folder + rel_dir_path = os.path.relpath(dir_full_path, base_path) + result[rel_dir_path] = "" + + # 2. Handle file entries + for file_name in files: + file_full_path = os.path.join(root, file_name) + rel_file_path = os.path.relpath(file_full_path, base_path) + + # Calculate MD5 by reading the entire file into memory + try: + if verbose: + print(f"Computing hash for {rel_file_path}") + with open(file_full_path, "rb") as f: + file_bytes = f.read() # Loads the whole file into RAM + + # Hash the complete byte string at once + result[rel_file_path] = hashlib.md5(file_bytes).hexdigest() + except (PermissionError, FileNotFoundError): + result[rel_file_path] = "ERROR: Cannot read file" + + return result + + +class FolderManager: + """ Helper class with some path utilities from a given path. """ + def __init__(self, path): + self.__path = path + self._logId = "" + self.__extraLog = None + + def join(self, *p): + return os.path.join(self.__path, *p) + + def relpath(self, p): + return os.path.relpath(p, self.path) + + def mkdir(self, *p, **kwargs): + d = self.join(*p) + Process.system(f"mkdir -p '{d}'", **kwargs) + return d + + def exists(self, *p): + return os.path.exists(self.join(*p)) + + @property + def path(self): + return self.__path + + @path.setter + def path(self, value): + if not isinstance(value, str): + raise Exception(f"FolderManger: Path must be a string, got {type(value)}") + self.__path = value + + def clear(self): + """ Remove existing path. """ + Process.system(f"rm -rf '{self.path}'") + + def create(self, **kwargs): + """ Create batch folder. """ + self.log(f"Creating folder: {self.path}") + Process.system(f"rm -rf '{self.path}'", **kwargs) + Process.system(f"mkdir -p '{self.path}'", **kwargs) + + def log(self, msg, flush=False): + logMsg = f"{Pretty.now()}:{self._logId} {msg}" + print(logMsg, flush=flush) + if self.__extraLog: + self.__extraLog(logMsg, flush=flush) + return logMsg + + def setExtraLog(self, logFunc): + self.__extraLog = logFunc + + def listdir(self): + """ Return files relative to the path. """ + return os.listdir(self.path) + + def glob(self, pattern): + return glob(self.join(pattern)) + + def dump(self, obj, fn): + filePath = self.join(fn) + with open(filePath, 'w') as f: + json.dump(obj, f, indent=4) + + def rename(self, oldFn, newFn): + os.rename(self.join(oldFn), self.join(newFn)) + + def link(self, fn, absolute=False, name=None): + """ Link a file inside the folder and return the basename. + If name is None, the basename of the fn will be used. + """ + base = name or os.path.basename(fn) + src = os.path.abspath(fn) if absolute else self.relpath(fn) + os.symlink(src, self.join(base)) + return base + + def copy(self, *paths): + """ Copy one or many files into the path. """ + for p in paths: + shutil.copy(p, self.__path) + diff --git a/emtools/utils/pretty.py b/emtools/utils/pretty.py index 626f4c3..7ad7ec7 100644 --- a/emtools/utils/pretty.py +++ b/emtools/utils/pretty.py @@ -16,7 +16,7 @@ import math import os -from datetime import datetime +from datetime import datetime, timedelta class Pretty: @@ -69,6 +69,20 @@ def parse_datetime(dt_str, **kwargs): f = kwargs.get('format', Pretty.DATETIME_FORMAT) return datetime.strptime(dt_str, f) + @staticmethod + def parse_timedelta(td_str, **kwargs): + """Parse 'HH:MM:SS' or 'D days, HH:MM:SS' format""" + parts = td_str.split(', ') + days = 0 + if len(parts) == 2: + days = int(parts[0].split()[0]) + time_part = parts[1] + else: + time_part = parts[0] + + h, m, s = map(float, time_part.split(':')) + return timedelta(days=days, hours=h, minutes=m, seconds=s) + @staticmethod def modified(fn, **kwargs): if not os.path.exists(fn): @@ -80,7 +94,7 @@ def modified(fn, **kwargs): @staticmethod def elapsed(timestamp, now=None): """ - Get a datetime object or a int() Epoch timestamp and return a + Get a datetime object or an int() Epoch timestamp and return a pretty string like 'an hour ago', 'Yesterday', '3 months ago', 'just now', etc """ @@ -129,4 +143,9 @@ def _plural(div, noun): return _plural(365, 'year') + @staticmethod + def dprint(msg): + """ DEBUG print with timestamp and flush. """ + print(f"{Pretty.now()}: >>> DEBUG: {msg}", flush=True) + diff --git a/emtools/utils/process.py b/emtools/utils/process.py index 7f47c6c..5307c0e 100644 --- a/emtools/utils/process.py +++ b/emtools/utils/process.py @@ -22,6 +22,12 @@ import subprocess import logging +from .color import Color + + +def _print(*msgs): + print(*msgs) + class Process: def __init__(self, *args, **kwargs): @@ -29,7 +35,8 @@ def __init__(self, *args, **kwargs): self.args = args error = '' try: - self._p = subprocess.run(args, capture_output=True, text=True) + self._p = subprocess.run(args, capture_output=True, text=True, + input=kwargs.get('input', None)) self.stdout = self._p.stdout self.stderr = self._p.stderr self.returncode = self._p.returncode @@ -46,7 +53,7 @@ def __init__(self, *args, **kwargs): def lines(self): """ Iterate over the lines of the process output. """ - for line in self.stdout.split('\n'): + for line in self.stdout.splitlines(): yield line def print(self, args=True, stdout=False): @@ -56,7 +63,7 @@ def print(self, args=True, stdout=False): print(self.stdout) @staticmethod - def system(cmd, only_print=False, color=None, do_print=True): + def system(cmd, only_print=False, color=None, print=_print): """ Execute and print a command. Args: @@ -65,7 +72,7 @@ def system(cmd, only_print=False, color=None, do_print=True): not executed color: Optional color for the command """ - if do_print: + if print: printCmd = cmd if color is None else color(cmd) print(printCmd) if not only_print: @@ -91,10 +98,18 @@ def _addProc(f, proc): pids.add(proc.pid) attrs = ['pid', 'ppid', 'name', 'cwd', 'username', 'memory_percent', 'cpu_percent'] + def _filter_name(proc): + if program and program not in proc.info['name']: + cmdline = proc.cmdline() + if len(cmdline) == 0 or all(program not in cmd for cmd in cmdline): + return False + return True + for proc in psutil.process_iter(attrs): - if not program or program in proc.info['name']: + if _filter_name(proc): folder = proc.info['cwd'] if workingDir is None or folder == workingDir: + print(f"program: {program}, proc_info: {proc.info['name']}") _addProc(folder, proc) if children: for child in proc.children(recursive=True): @@ -103,6 +118,73 @@ def _addProc(f, proc): return processes + @staticmethod + def checkChilds(programName, folderPath, kill=False, verbose=0, pid=None): + from .system import System + specs = System.specs() + cpus = specs['CPUs'] + attrs = ['pid', 'ppid', 'name', 'cwd', 'username', + 'memory_percent', 'cpu_percent'] + + if pid is not None: + try: + root = psutil.Process(int(pid)) + except (psutil.NoSuchProcess, psutil.AccessDenied, ValueError): + return False + + procs = [] + seen = set() + + def _add(proc): + if proc.pid in seen: + return + proc.info = proc.as_dict(attrs) + procs.append(proc) + seen.add(proc.pid) + + _add(root) + for child in root.children(recursive=True): + _add(child) + folder = root.info.get('cwd') or folderPath or '' + processes = {folder: procs} + else: + processes = Process.ps(programName, workingDir=folderPath, + children=True) + + color = Color.red if kill else Color.bold + + for folder, procs in processes.items(): + print(Color.warn(f"{folder}")) + header = f" {'USER':<15} {'PPID/PID':<15} {color('PROGRAM'):<30}" + if verbose > 0: + header += f" {'CPU(%)':>10} {'MEMORY(%)':>10}" + if verbose > 1: + header += f" {'COMMAND LINE'}" + + print(Color.bold(header)) + + prefix = 'Killing' if kill else '' + for p in procs: + pidstr = f"{p.info['ppid']}/{p.pid}" + msg = f" {prefix} {p.info['username']:<15} {pidstr:<15} {color(p.info['name']):<30}" + if verbose > 0: + try: + cpu_percent = p.cpu_percent(interval=1) / cpus + except: + continue + + msg += f" {cpu_percent:>10,.2f} {p.info['memory_percent']:>10,.2f}" + if verbose > 1: + msg += f" {p.cmdline()}" + print(msg) + if kill: + try: + p.kill() + except: + pass + + return True if pid is not None else None + class Logger: """ Use a logger to log commands that are executed via os.system. """ def __init__(self, logger=None, only_log=False, @@ -119,6 +201,7 @@ def __init__(self, logger=None, only_log=False, # Shortcuts self.logger = logger self.info = logger.info + self.debug = logger.debug self.error = logger.error self.warning = logger.warning diff --git a/emtools/utils/system.py b/emtools/utils/system.py index 8f54c5d..4c5a7dd 100644 --- a/emtools/utils/system.py +++ b/emtools/utils/system.py @@ -21,8 +21,13 @@ """ import socket +import shutil import platform import psutil +import time +import json +import threading +from datetime import datetime from .process import Process @@ -97,3 +102,107 @@ def specs(): def hostname(): """ Return the hostname. """ return socket.gethostname() + + @staticmethod + def distro(): + """ Return a human-readable OS distribution name and version. + On Linux, it reads /etc/os-release (PRETTY_NAME); on macOS it uses + the product version; falls back to platform.platform() otherwise. + """ + system = platform.system() + + if system == 'Linux': + info = {} + try: + with open('/etc/os-release') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#') or '=' not in line: + continue + key, _, value = line.partition('=') + info[key] = value.strip().strip('"') + except (FileNotFoundError, OSError): + pass + + name = info.get('PRETTY_NAME') or info.get('NAME') + return name or f"Linux {platform.release()}" + + elif system == 'Darwin': + version, _, _ = platform.mac_ver() + return f"macOS {version}" if version else "macOS" + + return platform.platform() + + @staticmethod + def kernel(): + """ Return the kernel release (e.g. what 'uname -r' would print). """ + return platform.release() + + @staticmethod + def disk(path='/'): + """ Return disk usage (in bytes) for the filesystem containing path. + Keys: 'total', 'used', 'free'. Returns None if it can't be read. """ + try: + usage = shutil.disk_usage(path) + except OSError: + return None + return {'total': usage.total, 'used': usage.used, 'free': usage.free} + + +class GpuMonitor(threading.Thread): + """ Monitor GPU utilization. + Keeps an internal record of utilization data points, indexed by time. """ + + def __init__(self): + super().__init__() + self._stopEvent = threading.Event() + self._data = { + "sample": System.gpus(), + "columns": ["timestamp", ["temperature.gpu", + "utilization.gpu", + "utilization.memory"]], + "rows": [] + } + self.sleep = 1 + self.outputLog = 'gpu_monitor.json' + + def sample(self, verbose=False): + now = datetime.now() + gpus = System.gpus() + gpuLine = f'\r{now} ' + row = [str(now), []] + gpuEntries = {} + for gpuDict in sorted(gpus, key=lambda r: r['index']): + i = gpuDict['index'] + ugpu = gpuDict["utilization.gpu"].split()[0] # Remove % character + umem = gpuDict["utilization.memory"].split()[0] + gpuStr = f'{i}: gpu {ugpu}, mem {umem}' + gpuLine += f"{gpuStr:<30}" + gpuEntries[i] = [ugpu, umem] + if verbose: + print(gpuLine, end="") + self._data['rows'].append([str(now), gpuEntries]) + + def monitor(self, outputLog=None): + if outputLog: + self.outputLog = outputLog + c = 0 + while not self._stopEvent.is_set(): + self.sample() + c += 1 + if self.outputLog and c % 10 == 1: + with open(self.outputLog, 'w') as f: + json.dump(self._data, f) + c = 0 + + time.sleep(self.sleep) + + def run(self): + self.monitor() + + def stop(self): + """ Stop the current thread. """ + self._stopEvent.set() + self.join() + + diff --git a/emtools/utils/time.py b/emtools/utils/time.py index b7de988..4aa4570 100644 --- a/emtools/utils/time.py +++ b/emtools/utils/time.py @@ -14,7 +14,7 @@ # * # ************************************************************************** -from datetime import datetime +from datetime import datetime, timedelta from functools import wraps from .pretty import Pretty @@ -38,6 +38,9 @@ def getElapsedTime(self): def toc(self, message=None, pretty=False): print(self.getToc(message=message, pretty=pretty)) + def getTic(self): + return Pretty.datetime(self._dt) + def getToc(self, message=None, pretty=False): if message: self.message = message @@ -62,3 +65,8 @@ def wrap(*args, **kw): t.toc(f"Function {func.__name__} took: ") return result return wrap + + @staticmethod + def parse_timedelta(tdStr): + hours, minutes, seconds = tuple(map(float, tdStr.split(':'))) + return timedelta(hours=hours, minutes=minutes, seconds=seconds) diff --git a/requirements.txt b/requirements.txt index 9e549ea..7a15b9e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ mrcfile numpy Pillow>=9.0.1 xmltodict -psutil \ No newline at end of file +psutil +tifffile diff --git a/setup.py b/setup.py index 486211c..acec377 100644 --- a/setup.py +++ b/setup.py @@ -74,10 +74,14 @@ entry_points={ # Optional 'console_scripts': [ 'emt-ps = emtools.scripts.emt_ps:main', + 'emt-sysinfo = emtools.scripts.emt_sysinfo:main', 'emt-files = emtools.scripts.emt_files:main', 'emt-epu = emtools.scripts.emt_epu:main', 'emt-beamshifts = emtools.scripts.emt_beamshifts:main', - 'emt-angdist = emtools.scripts.emt_angdist:main' + 'emt-angdist = emtools.scripts.emt_angdist:main', + 'emt-star = emtools.scripts.emt_star:main', + 'emt-image = emtools.image.__main__:main' + ], },