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/image/__init__.py b/emtools/image/__init__.py index 619aec6..4563128 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,10 @@ # * 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 -__all__ = [Thumbnail] +__all__ = ["Thumbnail", "Image"] diff --git a/emtools/image/__main__.py b/emtools/image/__main__.py new file mode 100644 index 0000000..5254c7f --- /dev/null +++ b/emtools/image/__main__.py @@ -0,0 +1,32 @@ +# ************************************************************************** +# * +# * 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 +from .thumbnail import Image + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('path', metavar="IMAGE_PATH", + help="Image path") + + args = p.parse_args() + + print(Image.get_dimensions(args.path)) + + +if __name__ == '__main__': + main() diff --git a/emtools/image/thumbnail.py b/emtools/image/thumbnail.py index 0da7d7d..14f5e86 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,19 @@ # * 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 -from PIL import Image, ImageOps, ImageFilter +import PIL +from PIL import Image + +from emtools.utils import Path, Pretty class Thumbnail: @@ -49,7 +44,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 @@ -79,10 +73,10 @@ def from_pil(self, pil_img): self.scale = scale if self.contrast_factor is not None: - pil_img = ImageOps.autocontrast(pil_img, cutoff=self.contrast_factor) + pil_img = PIL.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(PIL.ImageFilter.GaussianBlur(radius=self.gaussian_radius)) return self.__format(pil_img) @@ -90,7 +84,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: @@ -121,7 +115,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 +137,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 +151,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 +162,80 @@ def Psd(**kwargs): defaults.update(kwargs) return Thumbnail(**defaults) + @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'): + dims = Image.get_dimensions(imagePath) + mrc = mrcfile.open(imagePath, permissive=True) + thumb = Thumbnail.Micrograph() + if len(dims) == 2: + array = mrc.data + elif len(dims) == 3: + x, y, z = dims + if mrc.is_volume() or (x == y and y == z): + thumb = Thumbnail(max_size=(256, 256), output_format='base64') + iMax = mrc.data.max() # min(imean + 10 * isd, imageArray.max()) + iMin = mrc.data.min() # max(imean - 10 * isd, imageArray.min()) + im255 = ((mrc.data - iMin) / (iMax - iMin) * 255).astype(np.uint8) + + # 1. Setup + 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 + zw, zh = zimg.size + + pad = 2 # The thickness of the dark gray lines/borders + + # 2. Calculate canvas size for a full grid with outer borders + # Total Width = (2 * image width) + (3 * padding for left, middle, right) + canvas_w = (xw + yw) + (3 * pad) + canvas_h = (xh + zh) + (3 * pad) + + # Create canvas with a white background (matching your image) + bg_color = (256, 256, 256) + montage = PIL.Image.new('RGB', (canvas_w, canvas_h), bg_color) + + # Top-Left: x slice + montage.paste(ximg, (pad, pad)) + # Bottom-Left: z slice + montage.paste(zimg, (pad, xh + 2 * pad)) + # Bottom-Right: y slice + montage.paste(yimg, (xw + 2 * pad, xh + 2 * pad)) + + return thumb.from_pil(montage) + Pretty.dprint("Loading MRC volume: %s" % str(array.shape)) + else: + array = mrc.data[z//2, :, :] # FIXME + Pretty.dprint("Loading MRC 2D: %s" % str(array.shape)) + return thumb.from_array(array) + else: + raise Exception("Invalid dimensions: %s" % dims) + + +class Image: + @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) diff --git a/emtools/jobs/__init__.py b/emtools/jobs/__init__.py index 672d307..88b1afd 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, Batch, Vars, + BatchManager, MdocBatchManager, TsStarBatchManager) +from .workflow import Workflow -__all__ = ["Pipeline", "BatchManager", "ProcessingPipeline"] \ No newline at end of file +__all__ = ["Pipeline", "Workflow", + "Args", "Vars", + "Batch", "BatchManager", "MdocBatchManager", "TsStarBatchManager"] diff --git a/emtools/jobs/batch_manager.py b/emtools/jobs/batch_manager.py index e4ebcd2..59a9e3c 100644 --- a/emtools/jobs/batch_manager.py +++ b/emtools/jobs/batch_manager.py @@ -15,10 +15,224 @@ # ************************************************************************** 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 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 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: + kwargs['cwd'] = self.path + 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 +244,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 +254,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 +307,128 @@ 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: + items = [row._asdict() for row in sf.iterTable(tsName)] + 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..1908273 --- /dev/null +++ b/emtools/jobs/workflow.py @@ -0,0 +1,158 @@ +# ************************************************************************** +# * +# * 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): + 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 + i.childs.append(self) + + def hasInput(self, inputId): + return inputId in self._inputs + + def clearInputs(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..fcf449a 100644 --- a/emtools/metadata/misc.py +++ b/emtools/metadata/misc.py @@ -15,10 +15,13 @@ # ************************************************************************** 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 +137,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 +164,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 +216,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 +283,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 +351,164 @@ 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['Sources'], 'Source') + self.Species = self._parseList(self._data['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): + 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} + + from pprint import pprint + d = data[key] + + if isinstance(d, list): + return [_parseItem(item) for item in d] + else: + return [_parseItem(d)] + + 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) + + return WarpSpecies(os.path.join(folder, entry['path'])) + + +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..28c18b9 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): """ @@ -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,11 @@ 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 _writeTableName(self, tableName): self._file.write("\ndata_%s\n\n" % (tableName or '')) @@ -341,13 +375,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 +400,53 @@ 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): """ 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. """ + if timeStamp: + self.writeTimeStamp() + 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 +469,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 +484,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 +509,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 +532,469 @@ 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: + + 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" + ] + + @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 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) # FIXME check continue and isTomo + sfOut.writeTimeStamp() + sfOut.writeTable('job', tJob, singleRow=True) + tValues = Table(['rlnJobOptionVariable', 'rlnJobOptionValue']) + for k, v in values.items(): + val = ('Yes' if v else 'No') if isinstance(v, bool) else 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 starFile.endswith('optimisation_set.star'): + with StarFile(starFile) as sf: + tableNames = sf.getTableNames() + tableName = ('optimisation_set' if 'optimisation_set' in tableNames + else tableNames[0]) + t = sf.getTable(tableName) + row = t[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 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): + with StarFile(pipeline_star, 'w') as sf: + sf.writeTimeStamp() + tGeneral = Table(['rlnPipeLineJobCounter']) + tGeneral.addRowValues(jobCounter) + sf.writeTable('pipeline_general', tGeneral, singleRow=True) + + if tables: + for name, t in tables.items(): + if len(t): + sf.writeTable(f"pipeline_{name}", t, computeFormat=True) + + @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[nodeName]) + + 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['datatype'], + 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..9ebc63c 100644 --- a/emtools/metadata/table.py +++ b/emtools/metadata/table.py @@ -49,6 +49,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 +114,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 +158,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() 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..46cae0e --- /dev/null +++ b/emtools/scripts/emt_star.py @@ -0,0 +1,174 @@ +#!/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 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 + + +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, columns): + + if not os.path.exists(inputStar): + raise Exception(f"Input star file does not exist: {inputStar}") + + with StarFile(inputStar) as sf: + existingTables = sf.getTableNames() + if tableName is None: + tableName = existingTables[0] + else: + if not tableName in existingTables: + raise Exception(f"Table name does not exist: {tableName}") + + table = StarFile.getTableFromFile(tableName, inputStar) + columnList = columns.split() + newTable = Table(columns=[col for col in table.getColumns() if col.getName() in columnList]) + for row in table: + values = {k: getattr(row, k) for k in columnList} + newTable.addRowValues(**values) + + StarFile.printTable(newTable, tableName) + + +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. ") + 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") + p.add_argument('--print', '-p', nargs='+', + metavar=('COLUMNS', 'TABLE'), + help="Print some columns from the given table.") + + args = p.parse_args() + inputStar = args.input + + if 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: + tableName = None + n = len(args.print) + cols = args.print[0] + if n > 2: + raise Exception(f"Only pass columns and optionally the tableName") + elif n > 1: # n == 2 + tableName = args.print[1] + + printColumns(args.input, tableName, cols) + else: + printStarInfo(args.input) + + +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_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..806dc64 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', + '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..2c58979 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,45 @@ def _addProc(f, proc): return processes + @staticmethod + def checkChilds(programName, folderPath, kill=False, verbose=0): + from .system import System + specs = System.specs() + cpus = specs['CPUs'] + 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 + class Logger: """ Use a logger to log commands that are executed via os.system. """ def __init__(self, logger=None, only_log=False, @@ -119,6 +173,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..f907e0c 100644 --- a/emtools/utils/system.py +++ b/emtools/utils/system.py @@ -23,6 +23,10 @@ import socket import platform import psutil +import time +import json +import threading +from datetime import datetime from .process import Process @@ -97,3 +101,62 @@ def specs(): def hostname(): """ Return the hostname. """ return socket.gethostname() + + +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..35b91cf 100644 --- a/setup.py +++ b/setup.py @@ -77,7 +77,10 @@ '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' + ], },