diff --git a/AUTHORS b/AUTHORS index e6000c344..8361e2772 100644 --- a/AUTHORS +++ b/AUTHORS @@ -1,6 +1,7 @@ This program is developed in the group of Sjors Scheres at the MRC Laboratory of Molecular Biology, with contributions from the following people (in alphabetical order): - Tom Burnley (from the CCP-EM team at STFC) +- Alister Burt (from David Barford's group at the MRC-LMB) - Liyi Dong - Bjoern Forsberg (from the Lindahl group at SciLifeLabs) - Shaoda He @@ -10,6 +11,7 @@ This program is developed in the group of Sjors Scheres at the MRC Laboratory of - Takanori Nakane - Joaquin Oton (from the Briggs group at MRC-LMB) - Colin Palmer (from the CCP-EM team at STFC) +- Euan Pyle (from Giulia Zanetti's group at Birkbeck) - Sjors Scheres - Jasenko Zivanov diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..9e823fdfc --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,10 @@ +[build-system] +requires = [ + "setuptools>=42", + "wheel", + "setuptools_scm[toml]>=3.4" +] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] +write_to = "src/tomography_python_programs/_version.py" \ No newline at end of file diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 000000000..4e71cb2b8 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,100 @@ + +[metadata] +name = tomography_python_programs +url = https://github.com/3dem/relion +author = RELION team +author_email = aburt@mrc-lmb.cam.ac.uk +description = package description. +long_description = file: README.md +long_description_content_type = text/markdown +license = BSD license +classifiers = + Development Status :: 2 - Pre-Alpha + License :: OSI Approved :: BSD License + Natural Language :: English + Programming Language :: Python :: 3 + Programming Language :: Python :: 3.8 + Programming Language :: Python :: 3.9 + Programming Language :: Python :: 3.10 + +project_urls = + Source Code =https://github.com/3dem/relion + +[options] +zip_safe = False +package_dir = + =src +packages = find: +python_requires = >=3.8 +setup_requires = + setuptools_scm +install_requires = + numpy + pandas + makefun + starfile + mrcfile + mdocfile + typer + rich + einops + lil_aretomo + yet-another-imod-wrapper + +[options.extras_require] +testing = + pytest +dev = + ipython + jedi<0.18.0 + black + flake8 + flake8-docstrings + isort + mypy + pre-commit + pydocstyle + pytest + jupyter-book + +[options.entry_points] +console_scripts = + relion_tomo_import = tomography_python_programs.import_tilt_series:cli + relion_tomo_align_tilt_series = tomography_python_programs.tilt_series_alignment:cli + relion_tomo_denoise = tomography_python_programs.denoising:cli + +[bdist_wheel] +universal = 1 + +[flake8] +exclude = docs,_version.py,.eggs,examples +max-line-length = 88 +docstring-convention = numpy +ignore = D100, D213, D401, D413, D107, W503 + +[isort] +profile = black +src_paths = tomography_preprocessing + +[pydocstyle] +match_dir = tomography_preprocessing +convention = numpy +add_select = D402,D415,D417 +ignore = D100, D213, D401, D413, D107 + +[tool:pytest] +addopts = -W error + +[mypy] +files = tomography_preprocessing +warn_unused_configs = True +warn_unused_ignores = True +check_untyped_defs = True +implicit_reexport = False +# this is strict! +# disallow_untyped_defs = True +show_column_numbers = True +show_error_codes = True +ignore_missing_imports = True + + diff --git a/src/tomography_python_programs/__init__.py b/src/tomography_python_programs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/tomography_python_programs/denoising/__init__.py b/src/tomography_python_programs/denoising/__init__.py new file mode 100644 index 000000000..f34faeea2 --- /dev/null +++ b/src/tomography_python_programs/denoising/__init__.py @@ -0,0 +1,4 @@ +from .cryoCARE import cryoCARE_train as _cryoCARE_train +from .cryoCARE import cryoCARE_predict as _cryoCARE_predict + +from ._cli import cli diff --git a/src/tomography_python_programs/denoising/_cli.py b/src/tomography_python_programs/denoising/_cli.py new file mode 100644 index 000000000..338ec3ee4 --- /dev/null +++ b/src/tomography_python_programs/denoising/_cli.py @@ -0,0 +1,4 @@ +import typer + +CLI_NAME = 'relion_tomo_denoise' +cli = typer.Typer(name=CLI_NAME, add_completion=False) diff --git a/src/tomography_python_programs/denoising/cryoCARE/__init__.py b/src/tomography_python_programs/denoising/cryoCARE/__init__.py new file mode 100644 index 000000000..db8c7eed2 --- /dev/null +++ b/src/tomography_python_programs/denoising/cryoCARE/__init__.py @@ -0,0 +1,2 @@ +from .cryoCARE_train import cryoCARE_train +from .cryoCARE_predict import cryoCARE_predict diff --git a/src/tomography_python_programs/denoising/cryoCARE/_utils.py b/src/tomography_python_programs/denoising/cryoCARE/_utils.py new file mode 100644 index 000000000..ed959036d --- /dev/null +++ b/src/tomography_python_programs/denoising/cryoCARE/_utils.py @@ -0,0 +1,154 @@ +from pathlib import Path +from typing import Tuple, List, Dict + +import pandas as pd +import starfile +import typer +import json +import shutil + +def create_denoising_directory_structure( + output_directory: Path, + training_job: bool, +) -> Tuple[Path, Path, Path]: + """ + Creates directory structure for denoising jobs. Doe not create tomogram directory if the job is for training a + denoising model as no tomograms are generated in this step. + """ + training_dir = output_directory / 'external' / 'training' + training_dir.mkdir(parents=True, exist_ok=True) + tomogram_dir = output_directory / 'tomograms' + if not training_job: + tomogram_dir.mkdir(parents=True, exist_ok=True) + tilt_series_dir = output_directory / 'tilt_series' + tilt_series_dir.mkdir(parents=True, exist_ok=True) + return training_dir, tomogram_dir, tilt_series_dir + +def parse_training_tomograms( + training_tomograms: str +) -> List: + """ + Reads the string given to the CLI to ascertain which tomograms to train on. String should + be a list of : separated tomograms (name from rlnTomoName) + """ + training_tomograms = training_tomograms.strip().split(':') + return training_tomograms + +def generate_training_tomograms_star( + global_star: pd.DataFrame, + training_tomograms: List, +) -> pd.DataFrame: + """ + Generates a pandas dataframe of the tomograms the user has selected for training in global star format + """ + training_tomograms_idx = pd.DataFrame(global_star.rlnTomoName.tolist()).isin(training_tomograms).values + if not any(training_tomograms_idx): + e = f"Could not user specified training tomograms ({', '.join(str(x) for x in training_tomograms)}) in tilt series star file" + console.log(f'ERROR: {e}') + raise RuntimeError(e) + training_tomograms_star = global_star[training_tomograms_idx] + return training_tomograms_star + +def find_tomogram_halves( + training_tomograms_star: pd.DataFrame, +) -> Tuple[List, List]: + """ + Returns lists (even and odd) of the location of the the tomograms the user wishes to train on. + """ + return training_tomograms_star['rlnTomoReconstructedTomogramHalf1'].values.tolist(), training_tomograms_star['rlnTomoReconstructedTomogramHalf2'].values.tolist() + +def generate_train_data_config_json( + even_tomos: List, + odd_tomos: List, + training_dir: Path, + number_training_subvolumes: int, + subvolume_dimensions: int, +) -> Dict: + """ + Creates a Dict which can be saved as a json file for train_data_config.json file + """ + number_normalisation_subvolumes = round(number_training_subvolumes * 0.1) + train_data_config_json = json.loads(f'{{"even": {json.dumps(even_tomos)}, "odd": {json.dumps(odd_tomos)}, "patch_shape": [{subvolume_dimensions}, {subvolume_dimensions}, {subvolume_dimensions}], \ + "num_slices": {number_training_subvolumes}, "split": 0.9, "tilt_axis": "Y", "n_normalization_samples": {number_normalisation_subvolumes}, "path": "{training_dir}"}}') + return train_data_config_json + +def generate_train_config_json( + training_dir: Path, + output_directory: Path, + model_name: str, +) -> Dict: + """ + Creates a Dict which can be saved as a json file for train_config.json file + """ + train_config_json = json.loads(f'{{"train_data": "{training_dir}", "epochs": 100, "steps_per_epoch": 200, "batch_size": 16, "unet_kern_size": 3, \ + "unet_n_depth": 3, "unet_n_first": 16, "learning_rate": 0.0004, "model_name": "{model_name}", "path": "{output_directory}"}}') + return train_config_json + +def generate_predict_json( + even_tomos: List, + odd_tomos: List, + training_dir: Path, + model_name: Path, + output_directory: Path, + n_tiles: Tuple[int,int,int], +) -> Dict: + """ + Creates a Dict which can be saved as a json file for predict_config.json file + """ + predict_json = json.loads(f'{{"path": "{model_name}", "even": {json.dumps(even_tomos)}, \ + "odd": {json.dumps(odd_tomos)}, "n_tiles": {list(n_tiles)}, "output": "{output_directory / "tomograms"}"}}') + return predict_json + +def save_json( + training_dir: Path, + output_json: Dict, + json_prefix: str, +): + """ + Saves json file in output directory with desired file name (prefix). + """ + with open(f'{training_dir}/{json_prefix}.json', 'w') as outfile: + json.dump(output_json, outfile, indent=4) + +def save_tilt_series_stars( + global_star: pd.DataFrame, + tilt_series_dir: Path, +): + """ + Saves tilt series star files in output directory. + """ + for idx,row in global_star.iterrows(): + shutil.copyfile(f"{row['rlnTomoTiltSeriesStarFile']}", f'{tilt_series_dir}/{row["rlnTomoName"]}.star') + global_star['rlnTomoTiltSeriesStarFile'] = global_star.apply(lambda x: f'{tilt_series_dir}/{x["rlnTomoName"]}.star', axis=1) + +def add_denoised_tomo_to_global_star( + global_star: pd.DataFrame, + tomogram_dir: Path, + output_directory: Path, +): + """ + Adds location of the denoising tomogram to the global star file. + """ + global_star['rlnTomoReconstructedTomogramDenoised'] = global_star.apply(lambda x: f'{tomogram_dir}/rec_{x["rlnTomoName"]}.mrc', axis=1) + return global_star + +def save_global_star( + global_star: pd.DataFrame, + output_directory: Path, +): + """ + Saves global star file (tomograms.star) in output directory. + """ + starfile.write({'global': global_star}, f'{output_directory}/tomograms.star') + +def rename_predicted_tomograms( + even_tomos: List, + tomogram_dir: Path, + even_suffix: str, +): + """ + Gives denoised tomograms as cryoCARE likes to name them after the even tomograms. + """ + even_tomos = [Path(tomo) for tomo in even_tomos] + even_tomos = [Path(f"{tomogram_dir}/{tomo.name}") for tomo in even_tomos] + [tomo.rename(Path(f"{tomogram_dir}/{tomo.stem.replace(even_suffix,'')}{tomo.suffix}")) for tomo in even_tomos] diff --git a/src/tomography_python_programs/denoising/cryoCARE/constants.py b/src/tomography_python_programs/denoising/cryoCARE/constants.py new file mode 100644 index 000000000..a6966da90 --- /dev/null +++ b/src/tomography_python_programs/denoising/cryoCARE/constants.py @@ -0,0 +1,13 @@ +EVEN_SUFFIX = '_half1' #The suffix given to even tomograms during split tomogram generation. This will be removed from the names of the output tomograms + +PREDICT_CONFIG_PREFIX = 'predict_config' #Name (minus the suffix) of the predict_config.json file + +TRAIN_DATA_CONFIG_PREFIX = 'train_data_config' #Name (minus the suffix) of the train_data_config.json file + +MODEL_NAME = 'denoising_model' #Name of the model to be trained + +TRAIN_CONFIG_PREFIX = 'train_config' #Name (minus the suffix) of the train_config.json file + + + + diff --git a/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_predict.py b/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_predict.py new file mode 100644 index 000000000..3eddbce0e --- /dev/null +++ b/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_predict.py @@ -0,0 +1,130 @@ +from pathlib import Path +from typing import Optional, Tuple + +import pandas as pd +import starfile +import rich +import typer +import subprocess +from rich.progress import track + +from ._utils import ( + create_denoising_directory_structure, + find_tomogram_halves, + generate_predict_json, + save_json, + rename_predicted_tomograms, + save_tilt_series_stars, + add_denoised_tomo_to_global_star, + save_global_star, +) +from .constants import PREDICT_CONFIG_PREFIX, EVEN_SUFFIX +from .._cli import cli +from ...utils.relion import relion_pipeline_job + +console = rich.console.Console(record=True) + +@cli.command(name='cryoCARE:predict') +@relion_pipeline_job +def cryoCARE_predict( + tilt_series_star_file: Path = typer.Option(...), + output_directory: Path = typer.Option(...), + model_name: Path = typer.Option(...), + n_tiles: Optional[Tuple[int,int,int]] = typer.Option((1,1,1)) + +): + """Generates denoised tomograms using cryoCARE from a previously trained denoising model (.tar.gz) + (Euan Pyle version, https://github.com/EuanPyle/cryoCARE_mpido) branched from Thorsten Wagner + version, https://github.com/thorstenwagner/cryoCARE_mpido) + + Requires that two tomograms have been generated using the same sample. These can be generated via taking odd/even + frames during Motion Correction (optimal) or by taking odd/even tilts during tomogram reconstruction. + The location of these tomograms should be specified in the global star file for all tilt series with the headers: + + + rlnTomoReconstructedTomogramHalf1 + + rlnTomoReconstructedTomogramHalf2 + + Parameters + ---------- + + tilt_series_star_file: RELION tilt-series STAR file. + + output_directory: directory in which results will be stored. + + model_name: user should provide the path to the model.tar.gz produced by a previous cryoCARE denoising job. + + n-tiles (optional): Initial tiles per dimension during prediction step. Should get increased + if the tiles do not fit on the GPU. However, sometimes this parameter is a source of errors causing out of memory + problems and tomogram dimensions to be wrong. We have found other useful values for this include 4,4,2 and 2,4,4 + (default = 1,1,1). Enter as: `--n-tiles 4 4 2` + + Returns + ------- + Denoised tomograms. + """ + if not tilt_series_star_file.exists(): + e = 'Could not find tilt series star file' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + + global_star = starfile.read(tilt_series_star_file, always_dict=True)['global'] + + if not 'rlnTomoReconstructedTomogramHalf1' in global_star.columns: + e = 'Could not find rlnTomoReconstructedTomogramHalf1 in tilt series star file.' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + training_dir, tomogram_dir, tilt_series_dir = \ + create_denoising_directory_structure( + output_directory=output_directory, + training_job=False, + ) + + even_tomos, odd_tomos = find_tomogram_halves(global_star) + + predict_json = generate_predict_json( + even_tomos=even_tomos, + odd_tomos=odd_tomos, + training_dir=training_dir, + model_name=model_name, + output_directory=output_directory, + n_tiles=n_tiles, + ) + + save_json( + training_dir=training_dir, + output_json=predict_json, + json_prefix=PREDICT_CONFIG_PREFIX, + ) + + console.log('Generating denoised tomograms') + cmd = f"cryoCARE_predict.py --conf {training_dir}/{PREDICT_CONFIG_PREFIX}.json" + subprocess.run(cmd, shell=True) + + rename_predicted_tomograms( + even_tomos=even_tomos, + tomogram_dir=tomogram_dir, + even_suffix=EVEN_SUFFIX, + ) + + console.log('Denoised tomograms successfully generated, finalising metadata') + + save_tilt_series_stars( + global_star=global_star, + tilt_series_dir=tilt_series_dir, + ) + + add_denoised_tomo_to_global_star( + global_star=global_star, + tomogram_dir=tomogram_dir, + output_directory=output_directory, + ) + + save_global_star( + global_star=global_star, + output_directory=output_directory, + ) + + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) diff --git a/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_train.py b/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_train.py new file mode 100644 index 000000000..a8f41c674 --- /dev/null +++ b/src/tomography_python_programs/denoising/cryoCARE/cryoCARE_train.py @@ -0,0 +1,146 @@ +from pathlib import Path +from typing import Optional, Tuple, List + +import pandas as pd +import starfile +import rich +import typer +import subprocess +from rich.progress import track + +from ._utils import ( + create_denoising_directory_structure, + parse_training_tomograms, + generate_training_tomograms_star, + find_tomogram_halves, + generate_train_config_json, + generate_train_data_config_json, + save_json, + save_tilt_series_stars, + save_global_star, +) +from .constants import TRAIN_CONFIG_PREFIX, TRAIN_DATA_CONFIG_PREFIX, MODEL_NAME +from .._cli import cli +from ...utils.relion import relion_pipeline_job + +console = rich.console.Console(record=True) + +@cli.command(name='cryoCARE:train') +@relion_pipeline_job +def cryoCARE_train( + tilt_series_star_file: Path = typer.Option(...), + output_directory: Path = typer.Option(...), + training_tomograms: str = typer.Option(None), + number_training_subvolumes: Optional[int] = typer.Option(1200), + subvolume_dimensions: Optional[int] = typer.Option(72), +): + """Trains a denoising model using cryoCARE (Euan Pyle version, https://github.com/EuanPyle/cryoCARE_mpido) + branched from Thorsten Wagner version, https://github.com/thorstenwagner/cryoCARE_mpido) + + Requires that two tomograms have been generated using the same sample. These can be generated via taking odd/even + frames during Motion Correction (optimal) or by taking odd/even tilts during tomogram reconstruction. + The location of these tomograms should be specified in the global star file for all tilt series with the headers: + + + rlnTomoReconstructedTomogramHalf1 + + rlnTomoReconstructedTomogramHalf2 + + Parameters + ---------- + + tilt_series_star_file: RELION tilt-series STAR file. + + output_directory: directory in which results will be stored. + + training_tomograms: tomograms which are to be used for denoising neural network training. + User should enter tomogram names as rlnTomoName, separated by colons, e.g. TS_1:TS_2 + + number_training_subvolumes: Number of sub-volumes to be extracted per training tomogram + Corresponds to num_slices in cryoCARE_extract_train_data.py. Default is 1200. + Number of normalisation samples will be 10% of this value. + + subvolume_dimensions: Dimensions (for XYZ, in pixels) of the subvolumes extracted for training + Default is 72. This number should not be lower than 64. Corresponds to patch_shape in cryoCARE_extract_train_data.py + + Returns + ------- + A denoising model (denoising_model.tar.gz) in the external/training/ subdirectory of the output directory. + """ + if not tilt_series_star_file.exists(): + e = 'Could not find tilt series star file' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + + global_star = starfile.read(tilt_series_star_file, always_dict=True)['global'] + + if not 'rlnTomoReconstructedTomogramHalf1' in global_star.columns: + e = 'Could not find rlnTomoReconstructedTomogramHalf1 in tilt series star file.' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + + training_dir, tomogram_dir, tilt_series_dir = \ + create_denoising_directory_structure( + output_directory=output_directory, + training_job=True, + ) + + training_tomograms = parse_training_tomograms(training_tomograms) + + training_tomograms_star = generate_training_tomograms_star( + global_star=global_star, + training_tomograms=training_tomograms, + ) + + even_tomos, odd_tomos = find_tomogram_halves(training_tomograms_star) + + console.log('Beginning to train denoising model.') + + train_data_config_json = generate_train_data_config_json( + even_tomos=even_tomos, + odd_tomos=odd_tomos, + training_dir=training_dir, + number_training_subvolumes=number_training_subvolumes, + subvolume_dimensions=subvolume_dimensions, + ) + + save_json( + training_dir=training_dir, + output_json=train_data_config_json, + json_prefix=TRAIN_DATA_CONFIG_PREFIX, + ) + + cmd = f"cryoCARE_extract_train_data.py --conf {training_dir}/{TRAIN_DATA_CONFIG_PREFIX}.json" + subprocess.run(cmd, shell=True) + + train_config_json = generate_train_config_json( + training_dir=training_dir, + output_directory=output_directory, + model_name=MODEL_NAME, + ) + + save_json( + training_dir=training_dir, + output_json=train_config_json, + json_prefix=TRAIN_CONFIG_PREFIX, + ) + + cmd = f"cryoCARE_train.py --conf {training_dir}/{TRAIN_CONFIG_PREFIX}.json" + subprocess.run(cmd, shell=True) + + console.log(f'Finished training denoising model.') + + save_tilt_series_stars( + global_star=global_star, + tilt_series_dir=tilt_series_dir, + ) + + save_global_star( + global_star=global_star, + output_directory=output_directory, + ) + + console.log(f'Denoising model can be found in {output_directory}/{MODEL_NAME}.tar.gz') + + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) diff --git a/src/tomography_python_programs/import_tilt_series/__init__.py b/src/tomography_python_programs/import_tilt_series/__init__.py new file mode 100644 index 000000000..53476eb2a --- /dev/null +++ b/src/tomography_python_programs/import_tilt_series/__init__.py @@ -0,0 +1,3 @@ +from .serialem import import_tilt_series_from_serial_em +from .stub_alternative import stub +from ._cli import cli diff --git a/src/tomography_python_programs/import_tilt_series/_cli.py b/src/tomography_python_programs/import_tilt_series/_cli.py new file mode 100644 index 000000000..2c3aaa916 --- /dev/null +++ b/src/tomography_python_programs/import_tilt_series/_cli.py @@ -0,0 +1,4 @@ +import typer + +CLI_NAME = 'relion_tomo_import' +cli = typer.Typer(name=CLI_NAME, add_completion=False) diff --git a/src/tomography_python_programs/import_tilt_series/serialem.py b/src/tomography_python_programs/import_tilt_series/serialem.py new file mode 100644 index 000000000..070a27f9b --- /dev/null +++ b/src/tomography_python_programs/import_tilt_series/serialem.py @@ -0,0 +1,163 @@ +from pathlib import Path +from typing import Optional, List + +import mdocfile +import numpy as np +import pandas as pd +import rich +import starfile +import typer +from rich.progress import track + +from ._cli import cli +from .. import utils +from ..utils.relion import relion_pipeline_job + +console = rich.console.Console(record=True) + + +@cli.command(name='SerialEM') +@relion_pipeline_job +def import_tilt_series_from_serial_em( + tilt_image_movie_pattern: str = typer.Option(...), + mdoc_file_pattern: str = typer.Option(...), + output_directory: Path = typer.Option(...), + nominal_tilt_axis_angle: float = typer.Option(...), + nominal_pixel_size: float = typer.Option(...), + voltage: float = typer.Option(...), + spherical_aberration: float = typer.Option(...), + amplitude_contrast: float = typer.Option(...), + invert_defocus_handedness: Optional[bool] = typer.Option(None), + dose_per_tilt_image: Optional[float] = None, + prefix: str = '', + mtf_file: Optional[Path] = None, +) -> Path: + """Import tilt-series data using SerialEM metadata. + + Parameters + ---------- + tilt_image_movie_pattern : Filename pattern with wildcard characters + for tilt image movie files. + mdoc_file_pattern : Filename pattern wildcard characters for mdoc files + containing tilt-series metadata. + output_directory : Directory in which output files are written. + nominal_tilt_axis_angle : Nominal tilt axis angle in the images. + nominal_pixel_size : Nominal pixel spacing of the tilt-images in angstroms. + voltage : Acceleration voltage (keV) + spherical_aberration : Spherical aberration (mm) + amplitude_contrast : Amplitude contrast fraction (e.g. 0.1) + invert_defocus_handedess: Set this to flip the handedness of the defocus geometry (default=1). + The value of this parameter is either +1 or -1, and it describes whether the focus + increases or decreases as a function of Z distance. It has to be determined experimentally. + dose_per_tilt_image : dose in electrons per square angstrom in each tilt image. + If set, this will override the values from the mdoc file + prefix : a prefix which will be added to the tilt-series name. + This avoids name collisions when processing data from multiple collection + sessions. + mtf_file: a text file containing the MTF of the detector. + + Returns + ------- + tilt_series_star_file : a text file pointing to STAR files containing per tilt-series metadata + """ + console.log('Started import_tilt_series job.') + + # Create output directory structure + tilt_series_directory = Path(output_directory) / 'tilt_series' + console.log(f'Creating output directory at {tilt_series_directory}') + tilt_series_directory.mkdir(parents=True, exist_ok=True) + + # Get lists of necessary files + tilt_image_files = list(Path().glob(tilt_image_movie_pattern)) + mdoc_files = sorted(list(Path().glob(mdoc_file_pattern))) + console.log(f'Found {len(mdoc_files)} mdoc files and {len(tilt_image_files)} image files.') + + # Raise error if no tilt images or mdocs found + if len(tilt_image_files) == 0: + e = 'Could not find any image files' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + if len(mdoc_files) == 0: + e = 'Could not find any mdoc files' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + + # Get tomogram ids and construct paths for per-tilt-series STAR files + tomogram_ids = [ + utils.mdoc.construct_tomogram_id(mdoc_file, prefix) + for mdoc_file in mdoc_files + ] + tilt_series_star_files = [ + tilt_series_directory / f"{tomogram_id}.star" + for tomogram_id in tomogram_ids + ] + + # Construct a global dataframe pointing at per-tilt-series STAR files + global_star_file = Path(output_directory) / 'tilt_series.star' + global_df = pd.DataFrame({ + 'rlnTomoName': tomogram_ids, + 'rlnTomoTiltSeriesStarFile': tilt_series_star_files, + }) + global_df['rlnVoltage'] = voltage + global_df['rlnSphericalAberration'] = spherical_aberration + global_df['rlnAmplitudeContrast'] = amplitude_contrast + global_df['rlnMicrographOriginalPixelSize'] = nominal_pixel_size + global_df['rlnTomoHand'] = -1 if invert_defocus_handedness else 1 + if mtf_file is not None: + global_df['rlnMtfFileName'] = mtf_file + + starfile.write( + data={'global': global_df}, + filename=global_star_file, + overwrite=True + ) + console.log(f'Wrote tilt-series data STAR file {global_star_file}') + + # Write out per tilt-series STAR files + console.log('writing per tilt-series STAR files...') + for tomogram_id, mdoc_file, output_filename in \ + track(list(zip(tomogram_ids, mdoc_files, tilt_series_star_files))): + tilt_image_df = _generate_tilt_image_dataframe( + mdoc_file=mdoc_file, + tilt_image_files=tilt_image_files, + dose_per_tilt_image=dose_per_tilt_image, + prefix=prefix, + nominal_tilt_axis_angle=nominal_tilt_axis_angle, + ) + starfile.write( + data={f'{tomogram_id}': tilt_image_df}, + filename=output_filename, + overwrite=True + ) + console.log(f'Wrote STAR files for {len(tilt_series_star_files)} tilt-series.') + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) + return global_star_file + + +def _generate_tilt_image_dataframe( + mdoc_file: Path, + tilt_image_files: List[Path], + prefix: str, + nominal_tilt_axis_angle: float, + dose_per_tilt_image: Optional[float], +) -> pd.DataFrame: + """Generate a dataframe containing data about images in a tilt-series.""" + df = mdocfile.read(mdoc_file, camel_to_snake=True) + df = df.sort_values(by="z_value", ascending=True) + df = utils.mdoc.add_pre_exposure_dose(mdoc_df=df, dose_per_tilt=dose_per_tilt_image) + df = df.sort_values(by="tilt_angle", ascending=True) + df = utils.mdoc.add_tilt_image_files(mdoc_df=df, tilt_image_files=tilt_image_files) + df['tilt_series_id'] = utils.mdoc.construct_tomogram_id(mdoc_file, prefix) + df['nominal_tilt_axis_angle'] = nominal_tilt_axis_angle + + output_df = pd.DataFrame({ + 'rlnMicrographMovieName': df['tilt_image_file'], + 'rlnTomoTiltMovieFrameCount': df['num_sub_frames'], + 'rlnTomoTiltMovieIndex': np.array(df['z_value']) + 1, # 0 -> 1 indexing + 'rlnTomoNominalStageTiltAngle': df['tilt_angle'], + 'rlnTomoNominalTiltAxisAngle': df['nominal_tilt_axis_angle'], + 'rlnTomoNominalDefocus': df['target_defocus'], + 'rlnMicrographPreExposure': df['pre_exposure_dose'], + }) + return output_df diff --git a/src/tomography_python_programs/import_tilt_series/stub_alternative.py b/src/tomography_python_programs/import_tilt_series/stub_alternative.py new file mode 100644 index 000000000..bf385e584 --- /dev/null +++ b/src/tomography_python_programs/import_tilt_series/stub_alternative.py @@ -0,0 +1,9 @@ +"""Stub alternative CLI command for import_tilt_series job.""" + +from ._cli import cli + + +@cli.command(name='alternative', hidden=True) +def stub(): + """Stub for an alternative to SerialEM tomo metadata import_tilt_series.""" + pass diff --git a/src/tomography_python_programs/tilt_series_alignment/__init__.py b/src/tomography_python_programs/tilt_series_alignment/__init__.py new file mode 100644 index 000000000..e68524351 --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/__init__.py @@ -0,0 +1,8 @@ +# Import decorated batch alignment functions first +from .imod import batch_fiducials as _batch_imod_fiducials +from .imod import batch_patch_tracking as _batch_imod_patch_tracking +from .aretomo import batch_aretomo as _batch_aretomo +from ._job_utils import write_aligned_tilt_series_star_file as _write_global_output + +# Then import the cli, it will be decorated with all programs +from ._cli import cli diff --git a/src/tomography_python_programs/tilt_series_alignment/_cli.py b/src/tomography_python_programs/tilt_series_alignment/_cli.py new file mode 100644 index 000000000..470e9a81c --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/_cli.py @@ -0,0 +1,5 @@ +import typer + +cli = typer.Typer(name='relion_tomo_align_tilt_series', add_completion=False) + + diff --git a/src/tomography_python_programs/tilt_series_alignment/_job_utils.py b/src/tomography_python_programs/tilt_series_alignment/_job_utils.py new file mode 100644 index 000000000..2e5b512c7 --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/_job_utils.py @@ -0,0 +1,114 @@ +from pathlib import Path +from typing import Tuple, Optional + +import numpy as np +import pandas as pd +import starfile +import typer + +from .imod._utils import get_tilt_series_alignment_data, calculate_specimen_shifts, \ + get_xyz_extrinsic_euler_angles +from ..utils.transformations import S, Rx, Ry, Rz +from ._cli import cli + + +def create_alignment_job_directory_structure(output_directory: Path) -> Tuple[Path, Path, Path]: + """Create directory structure for a tilt-series alignment job.""" + stacks_directory = output_directory / 'stacks' + stacks_directory.mkdir(parents=True, exist_ok=True) + + external_directory = output_directory / 'external' + external_directory.mkdir(parents=True, exist_ok=True) + + metadata_directory = output_directory / 'tilt_series' + metadata_directory.mkdir(parents=True, exist_ok=True) + return stacks_directory, external_directory, metadata_directory + + +def tilt_series_alignment_parameters_to_relion_projection_matrices( + specimen_shifts: pd.DataFrame, + euler_angles: pd.DataFrame, + tilt_image_dimensions: np.ndarray, + tomogram_dimensions: np.ndarray, +): + """Generate affine matrices transforming points in 3D to 2D in tilt-images. + + Projection model: + 3D specimen is rotated about its center then translated such that the projection + of points onto the XY-plane gives their position in a tilt-image. + + More specifically + - 3D specimen is rotated about its center by + - shifting the origin to the specimen center + - rotated extrinsically about the Y-axis by the tilt angle + - rotated extrinsically about the Z-axis by the in plane rotation angle + - 3D specimen is translated to align coordinate system with tilt-image + - move center-of-rotation of specimen to center of tilt-image + - move center-of-rotation of specimen to rotation center in tilt-image + + Parameters + ---------- + specimen_shifts: XY-shifts which align the projected specimen with tilt-images + euler_angles: YZX intrinsic Euler angles which transform the specimen + tilt_image_dimensions: XY-dimensions of tilt-series. + tomogram_dimensions: size of tomogram in XYZ + """ + tilt_image_center = tilt_image_dimensions / 2 + specimen_center = tomogram_dimensions / 2 + + # Transformations, defined in order of application + s0 = S(-specimen_center) # put specimen center-of-rotation at the origin + r0 = Rx(euler_angles['rlnTomoXTilt']) # rotate specimen around X-axis + r1 = Ry(euler_angles['rlnTomoYTilt']) # rotate specimen around Y-axis + r2 = Rz(euler_angles['rlnTomoZRot']) # rotate specimen around Z-axis + s1 = S(specimen_shifts) # shift projected specimen in xy (camera) plane + s2 = S(tilt_image_center) # move specimen back into tilt-image coordinate system + + # compose matrices + transformations = s2 @ s1 @ r2 @ r1 @ r0 @ s0 + return np.squeeze(transformations) + + +def write_single_tilt_series_alignment_output( + tilt_image_df: pd.DataFrame, + tilt_series_id: str, + pixel_size: float, + alignment_directory: Path, + output_star_file: Path, +): + """Write metadata from a tilt-series alignment experiment.""" + xf, tlt = get_tilt_series_alignment_data(alignment_directory, tilt_series_id) + + shifts_px = calculate_specimen_shifts(xf) + tilt_image_df[['rlnTomoXShiftAngst', 'rlnTomoYShiftAngst']] = shifts_px * pixel_size + + euler_angles = get_xyz_extrinsic_euler_angles(xf, tlt) + tilt_image_df[['rlnTomoXTilt', 'rlnTomoYTilt', 'rlnTomoZRot']] = euler_angles + + starfile.write({tilt_series_id: tilt_image_df}, output_star_file) + + +@cli.command(name='write-global-output') +def write_aligned_tilt_series_star_file( + original_tilt_series_star_file: Path = typer.Option(...), + job_directory: Path = typer.Option(...) +): + """Write output from a batch of tilt-series alignment experiments.""" + df = starfile.read(original_tilt_series_star_file, always_dict=True)['global'] + + # update individual tilt series star files + df['rlnTomoTiltSeriesStarFile'] = [ + job_directory / 'tilt_series' / f'{tilt_series_id}.star' + for tilt_series_id in df['rlnTomoName'] + ] + df['rlnEtomoDirectiveFile'] = [ + job_directory / 'external' / tilt_series_id / f'{tilt_series_id}.edf' + for tilt_series_id in df['rlnTomoName'] + ] + etomo_directives_exist = any(df['rlnEtomoDirectiveFile'].apply(lambda x: Path(x).exists())) + if etomo_directives_exist is False: + df = df.drop(columns=['rlnEtomoDirectiveFile']) + + # check which output files were succesfully generated, take only those + df = df[df['rlnTomoTiltSeriesStarFile'].apply(lambda x: x.exists())] + starfile.write({'global': df}, job_directory / 'aligned_tilt_series.star') diff --git a/src/tomography_python_programs/tilt_series_alignment/aretomo/__init__.py b/src/tomography_python_programs/tilt_series_alignment/aretomo/__init__.py new file mode 100644 index 000000000..9ad96caba --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/aretomo/__init__.py @@ -0,0 +1 @@ +from .batch import batch_aretomo diff --git a/src/tomography_python_programs/tilt_series_alignment/aretomo/_utils.py b/src/tomography_python_programs/tilt_series_alignment/aretomo/_utils.py new file mode 100644 index 000000000..5b2724d75 --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/aretomo/_utils.py @@ -0,0 +1,7 @@ +from typing import Tuple + +def coerce_gpu_ids( + gpu_ids: str + ) -> Tuple: + """Convert GPU ID input from colon spaced (1:2:3) to no space inbetween GPU IDs (123)""" + return tuple(map(int, gpu_ids.strip().replace(':',''))) diff --git a/src/tomography_python_programs/tilt_series_alignment/aretomo/align_tilt_series.py b/src/tomography_python_programs/tilt_series_alignment/aretomo/align_tilt_series.py new file mode 100644 index 000000000..4130b3281 --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/aretomo/align_tilt_series.py @@ -0,0 +1,93 @@ +from pathlib import Path + +import pandas as pd +from lil_aretomo import run_aretomo_alignment +from rich.console import Console +from typing import Optional, Tuple + +from ._utils import coerce_gpu_ids + +from .._job_utils import ( + create_alignment_job_directory_structure, + write_single_tilt_series_alignment_output +) +from ... import utils + + +def align_single_tilt_series( + tilt_series_id: str, + tilt_series_df: pd.DataFrame, + tilt_image_df: pd.DataFrame, + do_local_alignments: bool, + alignment_resolution: float, + n_patches_xy: Tuple[int, int], + alignment_thickness_px: float, + tilt_angle_offset_correction: bool, + gpu_ids: Optional[str], + job_directory: Path, +): + """Align a single tilt-series in AreTomo using RELION tilt-series metadata. + + Parameters + ---------- + tilt_series_id: 'rlnTomoName' in RELION tilt-series metadata. + tilt_series_df: master file for tilt-series metadata. + tilt_image_df: file containing information for images in a single tilt-series. + do_local_alignments: flag to enable local alignments. + alignment_resolution: resolution for alignments in angstroms. + n_patches_xy: number of patches in x and y for local alignments + alignment_thickness_px: thickness of intermediate reconstruction during alignments. + tilt_angle_offset_correction: flag to enable/disable stage tilt offset correction (-TiltCor) in AreTomo + gpu_ids: string to specify GPUs. GPU identifiers should be separated by colons e.g. 0:1:2:3 + job_directory: directory in which results will be stored. + """ + console = Console(record=True) + + # Create output directory structure + stack_directory, external_directory, metadata_directory = \ + create_alignment_job_directory_structure(job_directory) + aretomo_directory = external_directory / tilt_series_id + aretomo_directory.mkdir(parents=True, exist_ok=True) + + # Establish filenames + tilt_series_filename = f'{tilt_series_id}.mrc' + tilt_image_metadata_filename = f'{tilt_series_id}.star' + + # Order is important in IMOD, sort by tilt angle + tilt_image_df = tilt_image_df.sort_values(by='rlnTomoNominalStageTiltAngle', ascending=True) + + # Create tilt-series stack and align using IMOD + # implicit assumption - one tilt-axis angle per tilt-series + console.log('Creating tilt series stack') + utils.image.stack_image_files( + image_files=tilt_image_df['rlnMicrographName'], + output_image_file=stack_directory / tilt_series_filename + ) + + if gpu_ids is not None: + gpu_ids = coerce_gpu_ids( + gpu_ids=gpu_ids + ) + + console.log('Running AreTomo') + run_aretomo_alignment( + tilt_series_file=stack_directory / tilt_series_filename, + tilt_angles=tilt_image_df['rlnTomoNominalStageTiltAngle'], + pixel_size=tilt_series_df['rlnTomoTiltSeriesPixelSize'], + nominal_rotation_angle=tilt_image_df['rlnTomoNominalTiltAxisAngle'][0], + output_directory=aretomo_directory, + local_align=do_local_alignments, + target_pixel_size=alignment_resolution / 2, + n_patches_xy=n_patches_xy, + correct_tilt_angle_offset=tilt_angle_offset_correction, + thickness_for_alignment=alignment_thickness_px, + gpu_ids=gpu_ids + ) + console.log('Writing STAR file for aligned tilt-series') + write_single_tilt_series_alignment_output( + tilt_image_df=tilt_image_df, + tilt_series_id=tilt_series_id, + pixel_size=tilt_series_df['rlnTomoTiltSeriesPixelSize'], + alignment_directory=aretomo_directory, + output_star_file=metadata_directory / tilt_image_metadata_filename, + ) diff --git a/src/tomography_python_programs/tilt_series_alignment/aretomo/batch.py b/src/tomography_python_programs/tilt_series_alignment/aretomo/batch.py new file mode 100644 index 000000000..e27a5d8ff --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/aretomo/batch.py @@ -0,0 +1,77 @@ +from pathlib import Path +from typing import Optional, Tuple + +import typer +from rich.console import Console + +from .align_tilt_series import align_single_tilt_series +from .._cli import cli +from .._job_utils import create_alignment_job_directory_structure, write_aligned_tilt_series_star_file +from ... import utils +from ...utils.relion import relion_pipeline_job + +console = Console(record=True) + + +@cli.command(name='AreTomo') +@relion_pipeline_job +def batch_aretomo( + tilt_series_star_file: Path = typer.Option(...), + output_directory: Path = typer.Option(...), + do_local_alignments: Optional[bool] = typer.Option(False), + n_patches_xy: Optional[Tuple[int, int]] = typer.Option((5,4)), + alignment_resolution: Optional[float] = typer.Option(10), + alignment_thickness: Optional[float] = typer.Option(800), + tomogram_name: Optional[str] = typer.Option(None), + tilt_angle_offset_correction: Optional[bool] = typer.Option(False), + gpu_ids: Optional[str] = typer.Option(None) +): + """Align one or multiple tilt-series in AreTomo using RELION tilt-series metadata. + + Parameters + ---------- + tilt_series_star_file: RELION tilt-series STAR file. + output_directory: directory in which results will be stored. + do_local_alignments: flag to enable/disable local alignments in AreTomo. + n_patches_xy: number of patches in x and y used in local alignments. + alignment_resolution: resolution for intermediate alignments. + alignment_thickness: thickness of intermediate reconstructions during alignments in px. + tomogram_name: 'rlnTomoName' for a specific tilt-series. + tilt_angle_offset_correction: flag to enable/disable stage tilt offset correction (-TiltCor) in AreTomo + gpu_ids: string to specify GPUs. GPU identifiers should be separated by colons e.g. 0:1:2:3 + + Returns + ------- + + """ + if not tilt_series_star_file.exists(): + e = 'Could not find tilt series star file' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + console.log('Extracting metadata for tilt series.') + tilt_series_metadata = list(utils.star.iterate_tilt_series_metadata( + tilt_series_star_file=tilt_series_star_file, + tilt_series_id=tomogram_name + )) + for tilt_series_id, tilt_series_df, tilt_image_df in tilt_series_metadata: + console.log(f'Aligning {tilt_series_id}...') + align_single_tilt_series( + tilt_series_id=tilt_series_id, + tilt_series_df=tilt_series_df, + tilt_image_df=tilt_image_df, + do_local_alignments=do_local_alignments, + alignment_resolution=alignment_resolution, + n_patches_xy=n_patches_xy, + alignment_thickness_px=alignment_thickness, + tilt_angle_offset_correction=tilt_angle_offset_correction, + gpu_ids=gpu_ids, + job_directory=output_directory, + ) + if tomogram_name is None: # write out STAR file for set of tilt-series + console.log('Writing aligned_tilt_series.star') + write_aligned_tilt_series_star_file( + original_tilt_series_star_file=tilt_series_star_file, + job_directory=output_directory + ) + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/__init__.py b/src/tomography_python_programs/tilt_series_alignment/imod/__init__.py new file mode 100644 index 000000000..dac768730 --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/__init__.py @@ -0,0 +1,2 @@ +from .batch_fiducials import batch_fiducials +from .batch_patch_tracking import batch_patch_tracking diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/_utils.py b/src/tomography_python_programs/tilt_series_alignment/imod/_utils.py new file mode 100644 index 000000000..1c2da1adc --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/_utils.py @@ -0,0 +1,131 @@ +import os +from functools import lru_cache +from pathlib import Path +from typing import Tuple + +import numpy as np + + +def read_xf(file: os.PathLike) -> np.ndarray: + """Read an IMOD xf file into an (n, 6) numpy array. + + The xf file with alignment transforms contains one + line per view, each with a linear transformation specified by six numbers: + A11 A12 A21 A22 DX DY + where the coordinate (X, Y) is transformed to (X', Y') by: + X' = A11 * X + A12 * Y + DX + Y' = A21 * X + A22 * Y + DY + """ + return np.loadtxt(fname=file, dtype=float).reshape((-1, 6)) + + +def read_tlt(file: os.PathLike) -> np.ndarray: + """Read an IMOD tlt file into an (n, ) numpy array.""" + return np.loadtxt(fname=file, dtype=float).reshape(-1) + + +def get_xf_shifts(xf: np.ndarray) -> np.ndarray: + """Extract XY shifts from an IMOD xf file. + + Output is an (n, 2) numpy array of XY shifts. Shifts in an xf file are + applied after rotations. IMOD xf files contain linear transformations. In + the context of tilt-series alignment they contain transformations which are + applied to 'align' a tilt-series such that images represent a fixed body rotating + around the Y-axis. + """ + return xf[:, -2:] + + +def get_xf_transformation_matrices(xf: np.ndarray) -> np.ndarray: + """Extract the 2D transformation matrix from IMOD xf data. + + Output is an (n, 2, 2) numpy array of matrices. + """ + return xf[:, :4].reshape((-1, 2, 2)) + + +def get_in_plane_rotations(xf: np.ndarray) -> np.ndarray: + """Extract the in plane rotation angle from IMOD xf data. + + Output is an (n, ) numpy array of angles in degrees. This function assumes + that the transformation in the xf file is a simple 2D rotation. + """ + transformation_matrices = get_xf_transformation_matrices(xf) + cos_theta = transformation_matrices[:, 0, 0] + return np.rad2deg(np.arccos(cos_theta)) + + +def calculate_specimen_shifts(xf: np.ndarray) -> np.ndarray: + """Extract specimen shifts from IMOD xf data. + + Output is an (n, 2) numpy array of XY shifts. Specimen shifts are shifts in + the camera plane applied to the projected image of a specimen to align it + with a tilt-image. + + This function relies on the fact that: + RSP = S'RP where S' = R_invS + + R is a rotation + R_inv is the inverse rotation + S is a shift vector + P is a point + S' is a rotated shift vector + + In words: Shifting then rotating is equivalent to rotating then shifting by a + rotated shift-vector. + + In IMOD, the rotation center is at (N - 1) / 2. In RELION, the rotation center + is at N / 2. An extra half-pixel shift is applied to account for these + differences. + """ + rotation_matrices = get_xf_transformation_matrices(xf) + inverse_rotation_matrices = rotation_matrices.transpose((0, 2, 1)) + + image_shifts = xf[:, -2:].reshape((-1, 2, 1)) + specimen_shifts = inverse_rotation_matrices @ -image_shifts + return specimen_shifts.reshape((-1, 2)) - 0.5 + + +@lru_cache(maxsize=100) +def get_tilt_series_id(alignment_directory: Path) -> str: + """Get the tilt-series stack basename from an Etomo directory.""" + edf_files = list(alignment_directory.glob('*.edf')) + if len(edf_files) != 1: + raise RuntimeError('Multiple Etomo directive files found') + return edf_files[0].stem + + +def get_edf_file(imod_directory: Path, tilt_series_id: str) -> Path: + """Get the Etomo directive file from an IMOD directory.""" + return imod_directory / f'{tilt_series_id}.edf' + + +def get_xf_file(alignment_directory: Path, tilt_series_id: str) -> Path: + """Get the xf file containing image transforms from an IMOD directory.""" + return alignment_directory / f'{tilt_series_id}.xf' + + +def get_tlt_file(alignment_directory: Path, tilt_series_id: str) -> Path: + """Get the tlt file containing tilt angles from an IMOD directory.""" + return alignment_directory / f'{tilt_series_id}.tlt' + + +def get_tilt_series_alignment_data( + alignment_directory: Path, tilt_series_id: str +) -> Tuple[np.ndarray, np.ndarray]: + """Get IMOD tilt-series alignment data.""" + xf = read_xf(get_xf_file(alignment_directory, tilt_series_id)) + tlt = read_tlt(get_tlt_file(alignment_directory, tilt_series_id)) + return xf, tlt + + +def get_xyz_extrinsic_euler_angles( + imod_xf_data: np.ndarray, imod_tlt_data: np.ndarray +) -> np.ndarray: + """Get XYZ extrinsic Euler angles which rotate the specimen.""" + euler_angles = np.zeros(shape=(len(imod_xf_data), 3)) + euler_angles[:, 1] = imod_tlt_data + euler_angles[:, 2] = get_in_plane_rotations(imod_xf_data) + return euler_angles + + diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/align_tilt_series.py b/src/tomography_python_programs/tilt_series_alignment/imod/align_tilt_series.py new file mode 100644 index 000000000..a8b1b01ab --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/align_tilt_series.py @@ -0,0 +1,66 @@ +from pathlib import Path +from typing import Callable, Dict, Any + +import mrcfile +import numpy as np +import pandas as pd +from rich.console import Console + +from .._job_utils import ( + create_alignment_job_directory_structure, + write_single_tilt_series_alignment_output +) + + +def align_single_tilt_series( + tilt_series_id: str, + tilt_series_df: pd.DataFrame, + tilt_image_df: pd.DataFrame, + alignment_function: Callable, + alignment_function_kwargs: Dict[str, Any], + output_directory: Path, +): + """Align a single tilt-series in IMOD using RELION tilt-series metadata. + + Parameters + ---------- + tilt_series_id: 'rlnTomoName' in RELION tilt-series metadata. + tilt_series_df: master file for tilt-series metadata. + tilt_image_df: file containing information for images in a single tilt-series. + alignment_function: alignment function from yet_another_imod_wrapper. + alignment_function_kwargs: keyword arguments specific to the alignment function. + output_directory: directory in which results will be stored. + """ + console = Console(record=True) + + # Create output directory structure + stack_directory, external_directory, metadata_directory = \ + create_alignment_job_directory_structure(output_directory) + imod_directory = external_directory / tilt_series_id + imod_directory.mkdir(parents=True, exist_ok=True) + tilt_image_metadata_filename = f'{tilt_series_id}.star' + + # Order is important in IMOD, sort by tilt angle + tilt_image_df = tilt_image_df.sort_values(by='rlnTomoNominalStageTiltAngle', ascending=True) + + # Align tilt-series using IMOD + # implicit assumption - one tilt-axis angle per tilt-series + console.log('Running IMOD alignment') + imod_output = alignment_function( + tilt_series=np.stack([mrcfile.read(f) for f in tilt_image_df['rlnMicrographName']]), + tilt_angles=tilt_image_df['rlnTomoNominalStageTiltAngle'], + pixel_size=tilt_series_df['rlnTomoTiltSeriesPixelSize'], + nominal_rotation_angle=tilt_image_df['rlnTomoNominalTiltAxisAngle'][0], + basename=tilt_series_id, + output_directory=imod_directory, + **alignment_function_kwargs, + ) + if imod_output.contains_alignment_results: + console.log('Writing STAR file for aligned tilt-series') + write_single_tilt_series_alignment_output( + tilt_image_df=tilt_image_df, + tilt_series_id=tilt_series_id, + pixel_size=tilt_series_df['rlnTomoTiltSeriesPixelSize'], + alignment_directory=imod_directory, + output_star_file=metadata_directory / tilt_image_metadata_filename, + ) diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/batch_fiducials.py b/src/tomography_python_programs/tilt_series_alignment/imod/batch_fiducials.py new file mode 100644 index 000000000..ab258180b --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/batch_fiducials.py @@ -0,0 +1,62 @@ +from pathlib import Path +from typing import Optional + +import typer +from yet_another_imod_wrapper import align_tilt_series_using_fiducials +from rich.console import Console + +from .._job_utils import write_aligned_tilt_series_star_file +from .align_tilt_series import align_single_tilt_series +from .._cli import cli +from ... import utils +from ...utils.relion import relion_pipeline_job + +console = Console(record=True) + + +@cli.command(name='IMOD:fiducials') +@relion_pipeline_job +def batch_fiducials( + tilt_series_star_file: Path = typer.Option(...), + output_directory: Path = typer.Option(...), + tomogram_name: Optional[str] = typer.Option(None), + nominal_fiducial_diameter_nanometres: float = typer.Option(...), +): + """Align one or multiple tilt-series with fiducials in IMOD. + + Parameters + ---------- + tilt_series_star_file: RELION tilt-series STAR file. + output_directory: directory in which results will be stored. + tomogram_name: 'rlnTomoName' in RELION tilt-series metadata. + nominal_fiducial_diameter_nanometres: nominal fiducial diameter in nanometers. + """ + if not tilt_series_star_file.exists(): + e = 'Could not find tilt series star file' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + console.log('Extracting metadata for tilt series.') + tilt_series_metadata = utils.star.iterate_tilt_series_metadata( + tilt_series_star_file=tilt_series_star_file, + tilt_series_id=tomogram_name + ) + for tilt_series_id, tilt_series_df, tilt_image_df in tilt_series_metadata: + console.log(f'Aligning {tilt_series_id}...') + align_single_tilt_series( + tilt_series_id=tilt_series_id, + tilt_series_df=tilt_series_df, + tilt_image_df=tilt_image_df, + alignment_function=align_tilt_series_using_fiducials, + alignment_function_kwargs={ + 'fiducial_size': nominal_fiducial_diameter_nanometres + }, + output_directory=output_directory, + ) + if tomogram_name is None: # write out STAR file for set of tilt-series + console.log('Writing aligned_tilt_series.star') + write_aligned_tilt_series_star_file( + original_tilt_series_star_file=tilt_series_star_file, + job_directory=output_directory + ) + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/batch_patch_tracking.py b/src/tomography_python_programs/tilt_series_alignment/imod/batch_patch_tracking.py new file mode 100644 index 000000000..c1217fa6a --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/batch_patch_tracking.py @@ -0,0 +1,65 @@ +from pathlib import Path +from typing import Optional + +import typer +from yet_another_imod_wrapper import align_tilt_series_using_patch_tracking +from rich.console import Console + +from .align_tilt_series import align_single_tilt_series +from .._cli import cli +from .._job_utils import write_aligned_tilt_series_star_file +from ... import utils +from ...utils.relion import relion_pipeline_job + +console = Console(record=True) + + +@cli.command(name='IMOD:patch-tracking') +@relion_pipeline_job +def batch_patch_tracking( + tilt_series_star_file: Path = typer.Option(...), + output_directory: Path = typer.Option(...), + tomogram_name: Optional[str] = typer.Option(None), + patch_size_angstroms: float = typer.Option(...), + patch_overlap_percentage: float = typer.Option(...), +): + """Align one or multiple tilt-series with patch-tracking in IMOD. + + Parameters + ---------- + tilt_series_star_file: RELION tilt-series STAR file. + output_directory: directory in which to store results. + tomogram_name: 'rlnTomoName' in tilt-series STAR file. + patch_size_angstroms: size of 2D patches used for alignment. + patch_overlap_percentage: percentage of overlap between tracked patches. + """ + if not tilt_series_star_file.exists(): + e = 'Could not find tilt series star file' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + console.log('Extracting metadata for tilt series.') + tilt_series_metadata = utils.star.iterate_tilt_series_metadata( + tilt_series_star_file=tilt_series_star_file, + tilt_series_id=tomogram_name + ) + for tilt_series_id, tilt_series_df, tilt_image_df in tilt_series_metadata: + console.log(f'Aligning {tilt_series_id}...') + align_single_tilt_series( + tilt_series_id=tilt_series_id, + tilt_series_df=tilt_series_df, + tilt_image_df=tilt_image_df, + alignment_function=align_tilt_series_using_patch_tracking, + alignment_function_kwargs={ + 'patch_size': patch_size_angstroms, + 'patch_overlap_percentage': patch_overlap_percentage, + }, + output_directory=output_directory, + ) + if tomogram_name is None: # write out STAR file for set of tilt-series + console.log('Writing aligned_tilt_series.star') + write_aligned_tilt_series_star_file( + original_tilt_series_star_file=tilt_series_star_file, + job_directory=output_directory + ) + console.save_html(str(output_directory / 'log.html'), clear=False) + console.save_text(str(output_directory / 'log.txt'), clear=False) diff --git a/src/tomography_python_programs/tilt_series_alignment/imod/etomo_launcher.py b/src/tomography_python_programs/tilt_series_alignment/imod/etomo_launcher.py new file mode 100644 index 000000000..cb0dc22ad --- /dev/null +++ b/src/tomography_python_programs/tilt_series_alignment/imod/etomo_launcher.py @@ -0,0 +1,27 @@ +from pathlib import Path + +import starfile + +from .._cli import cli +from .._job_utils import create_alignment_job_directory_structure +from ...utils.relion import relion_pipeline_job + + +@cli.command(name='IMOD:etomo-launcher') +@relion_pipeline_job +def etomo_launcher( + tilt_series_star_file: Path, + output_directory: Path, +): + tilt_series_directory, imod_alignments_directory = \ + create_alignment_job_directory_structure(output_directory) + tilt_series_df = starfile.read(tilt_series_star_file, always_dict=True)['global'] + + +class EtomoLauncher: + def __init__( + self, tilt_series_star_file: Path, output_directory: Path + ): + self.tilt_series_directory, self.imod_alignments_directory = \ + create_alignment_job_directory_structure(output_directory) + self.tilt_series_df = starfile.read(tilt_series_star_file, always_dict=True)['global'] diff --git a/src/tomography_python_programs/utils/__init__.py b/src/tomography_python_programs/utils/__init__.py new file mode 100644 index 000000000..9d3fab8e3 --- /dev/null +++ b/src/tomography_python_programs/utils/__init__.py @@ -0,0 +1 @@ +from . import mdoc, file, image, star, mrc, relion, transformations diff --git a/src/tomography_python_programs/utils/file.py b/src/tomography_python_programs/utils/file.py new file mode 100644 index 000000000..a07b8ee0d --- /dev/null +++ b/src/tomography_python_programs/utils/file.py @@ -0,0 +1,17 @@ +import os +from pathlib import Path +from typing import List + + +def basename(path: Path) -> str: + while path.stem != str(path): + path = Path(path.stem) + return path + + +def glob(pattern: str) -> List[Path]: + return list(Path().glob(pattern)) + + +def write_empty_file(filename: os.PathLike) -> None: + open(filename, mode='w').close() diff --git a/src/tomography_python_programs/utils/image.py b/src/tomography_python_programs/utils/image.py new file mode 100644 index 000000000..c3b529349 --- /dev/null +++ b/src/tomography_python_programs/utils/image.py @@ -0,0 +1,35 @@ +import os +from typing import List, Tuple + +import mrcfile +import numpy as np + +from .mrc import read_mrc + + +def stack_image_files(image_files: List[os.PathLike], output_image_file: os.PathLike): + n_images = len(image_files) + _, h, w = get_image_shape(image_files[0]) + stack_shape = (n_images, h, w) + mrc = mrcfile.new_mmap( + output_image_file, shape=stack_shape, mrc_mode=2, overwrite=True + ) + for idx, image_file in enumerate(image_files): + mrc.data[idx] = read_mrc(image_file).astype(np.float32) + mrc.reset_header_stats() + mrc.update_header_from_data() + pixel_size = get_pixel_size(image_files[0]) + mrc.voxel_size = (pixel_size, pixel_size, 0) + mrc.close() + + +def get_pixel_size(filename: os.PathLike) -> float: + with mrcfile.open(filename, header_only=True) as mrc: + pixel_size = mrc.voxel_size.x + return pixel_size + + +def get_image_shape(filename: os.PathLike) -> Tuple[int, int, int]: + with mrcfile.open(filename, header_only=True) as mrc: + nz, ny, nx = mrc.header.nz, mrc.header.ny, mrc.header.nx + return int(nz), int(ny), int(nx) \ No newline at end of file diff --git a/src/tomography_python_programs/utils/mdoc.py b/src/tomography_python_programs/utils/mdoc.py new file mode 100644 index 000000000..3c3ed641a --- /dev/null +++ b/src/tomography_python_programs/utils/mdoc.py @@ -0,0 +1,51 @@ +from os import PathLike +from pathlib import Path +from typing import Optional, List + +import numpy as np +import pandas as pd + +from .file import basename + + +def add_pre_exposure_dose( + mdoc_df: pd.DataFrame, dose_per_tilt: Optional[float] = None +) -> pd.DataFrame: + if dose_per_tilt is not None: + pre_exposure_dose = dose_per_tilt * np.arange(len(mdoc_df)) + else: # all zeros if exposure dose values not present in mdoc + pre_exposure_dose = [0] * len(mdoc_df) + + if "exposure_dose" not in mdoc_df.columns or dose_per_tilt is not None: + mdoc_df["pre_exposure_dose"] = pre_exposure_dose + else: + mdoc_df["pre_exposure_dose"] = np.cumsum(mdoc_df["exposure_dose"].to_numpy()) + return mdoc_df + + +def add_tilt_image_files( + mdoc_df: pd.DataFrame, tilt_image_files: List[PathLike] +) -> pd.DataFrame: + basename_filename_map = {Path(f).stem: f for f in tilt_image_files} + mdoc_tilt_image_basenames = get_tilt_image_basenames(mdoc_df["sub_frame_path"]) + mdoc_df["tilt_image_file"] = [ + basename_filename_map.get(basename, None) + for basename in mdoc_tilt_image_basenames + ] + mdoc_df = mdoc_df[mdoc_df["tilt_image_file"] != None] + return mdoc_df + + +def get_tilt_image_basenames(mdoc_tilt_image_files: pd.Series) -> pd.Series: + tilt_image_basenames = mdoc_tilt_image_files.apply( + lambda x: Path(str(x).split("\\")[-1]).stem + ) + return tilt_image_basenames + + +def construct_tomogram_id(mdoc_file: Path, prefix: str) -> str: + """Construct a unique tomogram ID.""" + if prefix == '': + return f'{utils.file.basename(mdoc_file)}' + else: + return f'{prefix}_{basename(mdoc_file)}' diff --git a/src/tomography_python_programs/utils/mrc.py b/src/tomography_python_programs/utils/mrc.py new file mode 100644 index 000000000..77c982493 --- /dev/null +++ b/src/tomography_python_programs/utils/mrc.py @@ -0,0 +1,16 @@ +import os + +import mrcfile +import numpy as np + + +def read_mrc(filename: os.PathLike) -> np.ndarray: + with mrcfile.open(filename, permissive=True) as mrc: + data = mrc.data + return data + + +def get_image_dimensions(filename: os.PathLike) -> np.ndarray: + """Get array of image dimensions (xyz) from an MRC file header.""" + with mrcfile.open(filename, header_only=True) as mrc: + return np.array([mrc.header.nx, mrc.header.ny, mrc.header.nz]) diff --git a/src/tomography_python_programs/utils/relion.py b/src/tomography_python_programs/utils/relion.py new file mode 100644 index 000000000..3392741ed --- /dev/null +++ b/src/tomography_python_programs/utils/relion.py @@ -0,0 +1,81 @@ +import inspect +from pathlib import Path +from typing import Callable + +import makefun +import typer + +from .file import write_empty_file + +JOB_SUCCESS_FILENAME = 'RELION_JOB_EXIT_SUCCESS' +JOB_FAILURE_FILENAME = 'RELION_JOB_EXIT_FAILURE' +ABORT_JOB_NOW_FILENAME = 'RELION_JOB_ABORT_NOW' +JOB_ABORTED_FILENAME = 'RELION_JOB_ABORTED' + + +def write_job_success_file(job_directory: Path) -> None: + """Write a file indicating job success.""" + output_file = job_directory / JOB_SUCCESS_FILENAME + write_empty_file(output_file) + + +def write_job_failure_file(job_directory: Path) -> None: + """Write a file indicating job failure.""" + output_file = job_directory / JOB_FAILURE_FILENAME + write_empty_file(output_file) + + +def _check_for_abort_job_now_file(job_directory: Path) -> bool: + """Check for the presence of a file indicating job termination.""" + output_file = job_directory / ABORT_JOB_NOW_FILENAME + return output_file.exists() + + +def _write_job_aborted_file(job_directory: Path) -> None: + """Write a file indicating that the job was succesfully terminated.""" + output_file = job_directory / JOB_ABORTED_FILENAME + write_empty_file(output_file) + + +def abort_job_if_necessary(job_directory: Path) -> None: + """Abort a job if file indicating job should be terminated is found.""" + if _check_for_abort_job_now_file(job_directory) is True: + _write_job_aborted_file(job_directory) + + +PIPELINE_CONTROL_KEYWORD_ARGUMENT = inspect.Parameter( + 'pipeline_control', + kind=inspect.Parameter.KEYWORD_ONLY, + default=typer.Option(None, '--pipeline_control'), + annotation=Path, +) + + +def relion_pipeline_job(func: Callable) -> Callable: + """Decorator which turns a function into a RELION pipeline-aware job. + + Specifically + - a file indicating job success will be written upon completion + - a file indicating job failure will be written upon failure + - a 'pipeline_control: Path' keyword argument will be added to the function + resulting in an autogenerated CLI option '--pipeline_control' + """ + function_signature = inspect.signature(func) + new_signature = makefun.add_signature_parameters( + function_signature, last=[PIPELINE_CONTROL_KEYWORD_ARGUMENT] + ) + + @makefun.wraps(func, new_sig=new_signature) + def pipeline_job(*args, **kwargs): + job_directory = kwargs.get('output_directory', None) + if job_directory is not None: + job_directory.mkdir(parents=True, exist_ok=True) + try: + kwargs.pop(PIPELINE_CONTROL_KEYWORD_ARGUMENT.name) + func(*args, **kwargs) + if job_directory is not None: + write_job_success_file(job_directory) + except: + if job_directory is not None: + write_job_failure_file(job_directory) + return pipeline_job diff --git a/src/tomography_python_programs/utils/star.py b/src/tomography_python_programs/utils/star.py new file mode 100644 index 000000000..c76267402 --- /dev/null +++ b/src/tomography_python_programs/utils/star.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple, Optional +from rich.console import Console + +import pandas as pd +import starfile + +console = Console(record=True) + +if TYPE_CHECKING: + import os + from typing import Iterable + +# (tilt_series_id, tilt_series_df, tilt_image_df) +TiltSeriesMetadata = Tuple[str, pd.DataFrame, pd.DataFrame] + + +def _iterate_all_tilt_series_metadata( + tilt_series_star_file: os.PathLike +) -> Iterable[TiltSeriesMetadata]: + """Yield all metadata from a tilt-series data STAR file.""" + star = starfile.read(tilt_series_star_file, always_dict=True) + for _, tilt_series_df in star['global'].iterrows(): + tilt_series_id = tilt_series_df['rlnTomoName'] + tilt_image_df = starfile.read( + tilt_series_df['rlnTomoTiltSeriesStarFile'], always_dict=True + )[tilt_series_id] + yield tilt_series_id, tilt_series_df, tilt_image_df + + +def _extract_single_tilt_series_metadata( + tilt_series_star_file: os.PathLike, + tilt_series_id: str +) -> TiltSeriesMetadata: + """Get metadata for a specific tilt-series from a tilt-series data STAR file.""" + star = starfile.read(tilt_series_star_file, always_dict=True) + # Check if Tomogram Name provided is actually found in the star file + if not star['global']['rlnTomoName'].str.contains(tilt_series_id).any(): + e = 'Specified Tomogram provided in Tomogram-Name is not found in rlnTomoName in the given star file.' + console.log(f'ERROR: {e}') + raise RuntimeError(e) + tilt_series_df = star['global'].set_index('rlnTomoName').loc[tilt_series_id, :] + tilt_image_df = starfile.read( + tilt_series_df['rlnTomoTiltSeriesStarFile'], always_dict=True + )[tilt_series_id] + return tilt_series_id, tilt_series_df, tilt_image_df + + +def iterate_tilt_series_metadata( + tilt_series_star_file: os.PathLike, + tilt_series_id: Optional[str] = None, +) -> Iterable[TiltSeriesMetadata]: + """Yield metadata from a tilt-series data STAR file.""" + if tilt_series_id is None: # align all tilt-series + all_tilt_series_metadata = _iterate_all_tilt_series_metadata( + tilt_series_star_file) + else: # do single tilt-series alignment + all_tilt_series_metadata = [ + _extract_single_tilt_series_metadata(tilt_series_star_file, tilt_series_id) + ] + for tilt_series_metadata in all_tilt_series_metadata: + yield tilt_series_metadata diff --git a/src/tomography_python_programs/utils/tests/test_transformations.py b/src/tomography_python_programs/utils/tests/test_transformations.py new file mode 100644 index 000000000..3d7fb8d73 --- /dev/null +++ b/src/tomography_python_programs/utils/tests/test_transformations.py @@ -0,0 +1,98 @@ +import numpy as np + +from tomography_preprocessing.utils.transformations import Rx, Ry, Rz, S + + +def test_single_x_rotation(): + """Rotation around X should be right handed. + """ + xyzw = np.array([1, 1, 1, 1]).reshape((4, 1)) + result = Rx(90) @ xyzw + expected = np.array([1, -1, 1, 1]).reshape((4, 1)) + assert result.shape == (4, 1) + assert np.allclose(result, expected) + + +def test_single_y_rotation(): + """Rotation around Y should be right handed. + """ + xyzw = np.array([1, 1, 1, 1]).reshape((4, 1)) + result = Ry(90) @ xyzw + expected = np.array([1, 1, -1, 1]).reshape((4, 1)) + assert result.shape == (4, 1) + assert np.allclose(result, expected) + + +def test_single_z_rotation(): + """Rotation around Z should be right handed. + """ + xyzw = np.array([1, 1, 1, 1]).reshape((4, 1)) + result = Rz(90) @ xyzw + expected = np.array([-1, 1, 1, 1]).reshape((4, 1)) + assert result.shape == (4, 1) + assert np.allclose(result, expected) + + +def test_multiple_x_rotations(): + """Should be able to generate multiple matrices from an array of angles.""" + angles = np.linspace(0, 90, 10) + matrices = Rx(angles) + assert matrices.shape == (10, 4, 4) + assert np.allclose(matrices[0], np.eye(4)) + assert np.allclose(matrices[-1], Rx(90)) + + +def test_multiple_y_rotations(): + """Should be able to generate multiple matrices from an array of angles.""" + angles = np.linspace(0, 90, 10) + matrices = Ry(angles) + assert matrices.shape == (10, 4, 4) + assert np.allclose(matrices[0], np.eye(4)) + assert np.allclose(matrices[-1], Ry(90)) + + +def test_multiple_z_rotations(): + """Should be able to generate multiple matrices from an array of angles.""" + angles = np.linspace(0, 90, 10) + matrices = Rz(angles) + assert matrices.shape == (10, 4, 4) + assert np.allclose(matrices[0], np.eye(4)) + assert np.allclose(matrices[-1], Rz(90)) + + +def test_single_shift_matrix_3d(): + """Generate a single (4, 4) affine matrix from a 3D shift vector.""" + shifts = [1, 2, 3] + matrix = S(shifts) + assert matrix.shape == (4, 4) + assert np.allclose(matrix[:3, 3], shifts) + assert np.allclose(matrix[:3, :3], np.eye(3)) + + +def test_single_shift_matrix_2d(): + """Generate a single (4, 4) affine matrix from a 2D shift vector.""" + shifts = [1, 2] + matrix = S(shifts) + assert matrix.shape == (4, 4) + assert np.allclose(matrix[:2, 3], shifts) + assert matrix[2, 3] == 0 + assert np.allclose(matrix[:3, :3], np.eye(3)) + + +def test_multiple_shift_matrices_3d(): + """Generate multiple (4, 4) affine matrices from an array of 3D shift vectors.""" + shifts = np.arange(18).reshape((6, 3)) + matrices = S(shifts) + assert matrices.shape == (6, 4, 4) + assert np.allclose(matrices[:, :3, 3], shifts) + assert np.allclose(matrices[:, :3, :3], np.eye(3)) + + +def test_multiple_shift_matrices_3d(): + """Generate multiple (4, 4) affine matrices from an array of 2D shift vectors.""" + shifts = np.arange(18).reshape((9, 2)) + matrices = S(shifts) + assert matrices.shape == (9, 4, 4) + assert np.allclose(matrices[:, :2, 3], shifts) + assert np.allclose(matrices[:, 2, 3], 0) + assert np.allclose(matrices[:, :3, :3], np.eye(3)) diff --git a/src/tomography_python_programs/utils/transformations.py b/src/tomography_python_programs/utils/transformations.py new file mode 100644 index 000000000..f03d10df5 --- /dev/null +++ b/src/tomography_python_programs/utils/transformations.py @@ -0,0 +1,68 @@ +import einops +import numpy as np + + +def Rx(angles_degrees: np.ndarray) -> np.ndarray: + """Affine matrix for a rotation around the X-axis.""" + angles_degrees = np.asarray(angles_degrees).reshape(-1) + c = np.cos(np.deg2rad(angles_degrees)) + s = np.sin(np.deg2rad(angles_degrees)) + matrices = einops.repeat( + np.eye(4), 'i j -> n i j', n=len(angles_degrees) + ) + matrices[:, 1, 1] = c + matrices[:, 1, 2] = -s + matrices[:, 2, 1] = s + matrices[:, 2, 2] = c + return np.squeeze(matrices) + + +def Ry(angles_degrees: np.ndarray) -> np.ndarray: + """Affine matrix for a rotation around the Y-axis.""" + angles_degrees = np.asarray(angles_degrees).reshape(-1) + c = np.cos(np.deg2rad(angles_degrees)) + s = np.sin(np.deg2rad(angles_degrees)) + matrices = einops.repeat( + np.eye(4), 'i j -> n i j', n=len(angles_degrees) + ) + matrices[:, 0, 0] = c + matrices[:, 0, 2] = s + matrices[:, 2, 0] = -s + matrices[:, 2, 2] = c + return np.squeeze(matrices) + + +def Rz(angles_degrees: float) -> np.ndarray: + """Affine matrix for a rotation around the Z-axis.""" + angle_degrees = np.asarray(angles_degrees).reshape(-1) + c = np.cos(np.deg2rad(angle_degrees)) + s = np.sin(np.deg2rad(angle_degrees)) + matrices = einops.repeat( + np.eye(4), 'i j -> n i j', n=len(angle_degrees) + ) + matrices[:, 0, 0] = c + matrices[:, 0, 1] = -s + matrices[:, 1, 0] = s + matrices[:, 1, 1] = c + return np.squeeze(matrices) + + +def S(shifts: np.ndarray) -> np.ndarray: + """Affine matrices for shifts. + + Shifts supplied can be 2D or 3D. + """ + shifts = np.asarray(shifts, dtype=float) + if shifts.shape[-1] == 2: + shifts = promote_2d_to_3d(shifts) + shifts = np.array(shifts).reshape((-1, 3)) + matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=shifts.shape[0]) + matrices[:, 0:3, 3] = shifts + return np.squeeze(matrices) + + +def promote_2d_to_3d(shifts: np.ndarray) -> np.ndarray: + """Promote 2D vectors to 3D with zeros in the last dimension.""" + shifts = np.asarray(shifts).reshape(-1, 2) + shifts = np.c_[shifts, np.zeros(shifts.shape[0])] + return np.squeeze(shifts) \ No newline at end of file