pytme 0.2.0b0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (85) hide show
  1. {pytme-0.2.0b0 → pytme-0.2.2}/MANIFEST.in +1 -2
  2. pytme-0.2.2/PKG-INFO +91 -0
  3. pytme-0.2.2/README.md +51 -0
  4. {pytme-0.2.0b0 → pytme-0.2.2}/pyproject.toml +15 -4
  5. {pytme-0.2.0b0 → pytme-0.2.2}/pytme.egg-info/SOURCES.txt +11 -3
  6. pytme-0.2.2/scripts/extract_candidates.py +257 -0
  7. pytme-0.2.0b0/scripts/match_template_filters.py → pytme-0.2.2/scripts/match_template.py +557 -222
  8. pytme-0.2.2/scripts/match_template_filters.py +1200 -0
  9. {pytme-0.2.0b0 → pytme-0.2.2}/scripts/postprocess.py +170 -71
  10. {pytme-0.2.0b0 → pytme-0.2.2}/scripts/preprocessor_gui.py +179 -86
  11. pytme-0.2.2/scripts/refine_matches.py +626 -0
  12. {pytme-0.2.0b0 → pytme-0.2.2}/tme/__init__.py +0 -1
  13. pytme-0.2.2/tme/__version__.py +1 -0
  14. pytme-0.2.2/tme/analyzer.py +1423 -0
  15. {pytme-0.2.0b0 → pytme-0.2.2}/tme/backends/__init__.py +41 -11
  16. pytme-0.2.2/tme/backends/_jax_utils.py +185 -0
  17. pytme-0.2.2/tme/backends/cupy_backend.py +273 -0
  18. pytme-0.2.2/tme/backends/jax_backend.py +282 -0
  19. pytme-0.2.2/tme/backends/matching_backend.py +1232 -0
  20. {pytme-0.2.0b0 → pytme-0.2.2}/tme/backends/mlx_backend.py +45 -68
  21. pytme-0.2.2/tme/backends/npfftw_backend.py +531 -0
  22. {pytme-0.2.0b0 → pytme-0.2.2}/tme/backends/pytorch_backend.py +41 -154
  23. {pytme-0.2.0b0 → pytme-0.2.2}/tme/density.py +312 -421
  24. pytme-0.2.2/tme/matching_data.py +735 -0
  25. pytme-0.2.2/tme/matching_exhaustive.py +507 -0
  26. {pytme-0.2.0b0 → pytme-0.2.2}/tme/matching_optimization.py +234 -129
  27. pytme-0.2.2/tme/matching_scores.py +884 -0
  28. {pytme-0.2.0b0 → pytme-0.2.2}/tme/matching_utils.py +281 -387
  29. pytme-0.2.2/tme/memory.py +377 -0
  30. {pytme-0.2.0b0 → pytme-0.2.2}/tme/orientations.py +226 -66
  31. {pytme-0.2.0b0 → pytme-0.2.2}/tme/parser.py +3 -4
  32. pytme-0.2.2/tme/preprocessing/__init__.py +2 -0
  33. pytme-0.2.2/tme/preprocessing/_utils.py +217 -0
  34. pytme-0.2.2/tme/preprocessing/composable_filter.py +31 -0
  35. pytme-0.2.2/tme/preprocessing/compose.py +55 -0
  36. pytme-0.2.2/tme/preprocessing/frequency_filters.py +388 -0
  37. pytme-0.2.2/tme/preprocessing/tilt_series.py +1011 -0
  38. {pytme-0.2.0b0 → pytme-0.2.2}/tme/preprocessor.py +574 -530
  39. {pytme-0.2.0b0 → pytme-0.2.2}/tme/structure.py +495 -189
  40. {pytme-0.2.0b0 → pytme-0.2.2}/tme/types.py +5 -3
  41. pytme-0.2.0b0/PKG-INFO +0 -73
  42. pytme-0.2.0b0/README.md +0 -46
  43. pytme-0.2.0b0/scripts/extract_candidates.py +0 -218
  44. pytme-0.2.0b0/scripts/match_template.py +0 -800
  45. pytme-0.2.0b0/scripts/refine_matches.py +0 -218
  46. pytme-0.2.0b0/tme/__version__.py +0 -1
  47. pytme-0.2.0b0/tme/analyzer.py +0 -1651
  48. pytme-0.2.0b0/tme/backends/cupy_backend.py +0 -378
  49. pytme-0.2.0b0/tme/backends/matching_backend.py +0 -1156
  50. pytme-0.2.0b0/tme/backends/npfftw_backend.py +0 -789
  51. pytme-0.2.0b0/tme/helpers.py +0 -881
  52. pytme-0.2.0b0/tme/matching_constrained.py +0 -195
  53. pytme-0.2.0b0/tme/matching_data.py +0 -672
  54. pytme-0.2.0b0/tme/matching_exhaustive.py +0 -1749
  55. {pytme-0.2.0b0 → pytme-0.2.2}/LICENSE +0 -0
  56. {pytme-0.2.0b0 → pytme-0.2.2}/scripts/__init__.py +0 -0
  57. {pytme-0.2.0b0 → pytme-0.2.2}/scripts/estimate_ram_usage.py +0 -0
  58. {pytme-0.2.0b0 → pytme-0.2.2}/scripts/preprocess.py +0 -0
  59. {pytme-0.2.0b0 → pytme-0.2.2}/setup.cfg +0 -0
  60. {pytme-0.2.0b0 → pytme-0.2.2}/setup.py +0 -0
  61. {pytme-0.2.0b0 → pytme-0.2.2}/src/extensions.cpp +0 -0
  62. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/__init__.py +0 -0
  63. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48n309.npy +0 -0
  64. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48n527.npy +0 -0
  65. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48n9.npy +0 -0
  66. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u1.npy +0 -0
  67. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u1153.npy +0 -0
  68. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u1201.npy +0 -0
  69. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u1641.npy +0 -0
  70. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u181.npy +0 -0
  71. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u2219.npy +0 -0
  72. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u27.npy +0 -0
  73. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u2947.npy +0 -0
  74. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u3733.npy +0 -0
  75. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u4749.npy +0 -0
  76. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u5879.npy +0 -0
  77. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u7111.npy +0 -0
  78. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u815.npy +0 -0
  79. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u83.npy +0 -0
  80. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c48u8649.npy +0 -0
  81. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c600v.npy +0 -0
  82. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/c600vc.npy +0 -0
  83. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/metadata.yaml +0 -0
  84. {pytme-0.2.0b0 → pytme-0.2.2}/tme/data/quat_to_numpy.py +0 -0
  85. {pytme-0.2.0b0 → pytme-0.2.2}/tme/matching_memory.py +0 -0
@@ -60,7 +60,6 @@ global-exclude .DS_Store
60
60
  global-exclude .gitconfig
61
61
  global-exclude *.pickle
62
62
  prune tme/temp
63
- prune tme/preprocessing
64
63
  global-exclude doc/compileApiReference.sh
65
64
  prune public
66
65
  global-exclude tme/matching_backend.py
@@ -70,4 +69,4 @@ global-exclude tme/scoring.py
70
69
  global-exclude tme/package.py
71
70
  global-exclude tme/transforms.py
72
71
  global-exclude tme/tests/test_packaging.py
73
- global-exclude scripts/match_template_devel.py
72
+ global-exclude scripts/match_template_devel.py
pytme-0.2.2/PKG-INFO ADDED
@@ -0,0 +1,91 @@
1
+ Metadata-Version: 2.1
2
+ Name: pytme
3
+ Version: 0.2.2
4
+ Summary: Python Template Matching Engine
5
+ Author: Valentin Maurer
6
+ Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
7
+ License: Proprietary
8
+ Project-URL: Homepage, https://github.com/KosinskiLab/pyTME
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.11
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+ Requires-Dist: mrcfile>=1.4.3
15
+ Requires-Dist: numpy>=1.22.2
16
+ Requires-Dist: pyfftw>=0.13.1
17
+ Requires-Dist: pytest>=6.2.5
18
+ Requires-Dist: PyYAML>=6.0
19
+ Requires-Dist: scikit-image>=0.19.0
20
+ Requires-Dist: scikit_learn>=1.2.1
21
+ Requires-Dist: scipy>=1.9.1
22
+ Requires-Dist: pybind11
23
+ Requires-Dist: psutil
24
+ Requires-Dist: tifffile
25
+ Requires-Dist: h5py
26
+ Provides-Extra: cupy
27
+ Requires-Dist: cupy; extra == "cupy"
28
+ Requires-Dist: voltools==0.6.0; extra == "cupy"
29
+ Provides-Extra: pytorch
30
+ Requires-Dist: torch; extra == "pytorch"
31
+ Provides-Extra: jax
32
+ Requires-Dist: jax; extra == "jax"
33
+ Requires-Dist: jaxlib; extra == "jax"
34
+ Provides-Extra: all
35
+ Requires-Dist: cupy; extra == "all"
36
+ Requires-Dist: voltools==0.6.0; extra == "all"
37
+ Requires-Dist: torch; extra == "all"
38
+ Requires-Dist: jax; extra == "all"
39
+ Requires-Dist: jaxlib; extra == "all"
40
+
41
+ # Python Template Matching Engine (PyTME)
42
+
43
+ [![Build Status](https://img.shields.io/github/actions/workflow/status/KosinskiLab/pyTME/main.yml?label=CI)](https://github.com/KosinskiLab/pyTME/actions)
44
+ [![PyPI](https://img.shields.io/pypi/v/pytme.svg)](https://pypi.org/project/pytme/)
45
+
46
+ **[Documentation](https://kosinskilab.github.io/pyTME/)** | **[Installation](https://kosinskilab.github.io/pyTME/quickstart/installation.html)** | **[API](https://kosinskilab.github.io/pyTME/reference/index.html)**
47
+
48
+ PyTME is a Python library for data-intensive n-dimensional template matching using CPUs and GPUs.
49
+
50
+ With its [backend-agnostic design](https://kosinskilab.github.io/pyTME/reference/backends.html), the same code can be run on diverse hardware platforms using a best-of-breed approach. The underyling abstract backend specification allows for adding new backends to benefit from gains in performance and capabilities without modifying the library's core routines. The implementation of template matching scores is modular and provides developers with a flexible framework for rapid prototyping. Furthermore, pyTME supports a unique callback capability through [analyzers](https://kosinskilab.github.io/pyTME/reference/analyzer.html), which allows for injection of custom code, enabling real-time processing and manipulation of results
51
+
52
+ PyTME includes a [graphical user interface](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html#practical-example) that provides simplified mask creation, interactive filter exploration, result visualization, and manual refinement capabilities. This GUI serves as an accessible entry point to the library's core functionalities, allowing users to efficiently interact with and analyze their data.
53
+
54
+ Finally, pyTME offers specialized tools for cryogenic electron microscopy data, such as wedge masks, CTF correction, as well as [means for handling structural data](https://kosinskilab.github.io/pyTME/reference/data_structure.html). Through dedicated [integrations](https://kosinskilab.github.io/pyTME/quickstart/integrations.html), the output of pyTME seamlessly integrates with commonly used cryogenic electron microscopy software such as RELION, Dynamo and IMOD.
55
+
56
+ Running into bugs or missing a feature? Help us improve the project by opening an [issue](https://github.com/KosinskiLab/pyTME/issues).
57
+
58
+ ## Installation
59
+
60
+ We recommend installation using one of the following methods
61
+
62
+ | Method | Command |
63
+ |----------|---------------------------------------------------------|
64
+ | PyPi | `pip install pytme` |
65
+ | Source | `pip install git+https://github.com/KosinskiLab/pyTME` |
66
+ | Docker | `docker build -t pytme -f docker/Dockerfile_GPU .` |
67
+
68
+ You can find alternative installation methods in the [documentation](https://kosinskilab.github.io/pyTME/quickstart/installation.html).
69
+
70
+
71
+ ## Quickstart
72
+
73
+ Learn how to get started with
74
+
75
+ - [Preprocessing:](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html) Picking the right mask and filter for template matching.
76
+ - [Template matching:](https://kosinskilab.github.io/pyTME/quickstart/match_template.html) Find your template of interest.
77
+ - [Postprocessing](https://kosinskilab.github.io/pyTME/quickstart/postprocessing.html) Analyze template matching results and downstream integrations.
78
+
79
+ ## How to Cite
80
+
81
+ If PyTME contributed significantly to your research, please cite the corresponding publication on [SoftwareX](https://www.sciencedirect.com/science/article/pii/S2352711024000074).
82
+
83
+ ```bibtex
84
+ @article{Maurer:2024aa,
85
+ author = {Maurer, Valentin J. and Siggel, Marc and Kosinski, Jan},
86
+ journal = {SoftwareX},
87
+ pages = {101636},
88
+ title = {PyTME (Python Template Matching Engine): A fast, flexible, and multi-purpose template matching library for cryogenic electron microscopy data},
89
+ volume = {25},
90
+ year = {2024}}
91
+ ```
pytme-0.2.2/README.md ADDED
@@ -0,0 +1,51 @@
1
+ # Python Template Matching Engine (PyTME)
2
+
3
+ [![Build Status](https://img.shields.io/github/actions/workflow/status/KosinskiLab/pyTME/main.yml?label=CI)](https://github.com/KosinskiLab/pyTME/actions)
4
+ [![PyPI](https://img.shields.io/pypi/v/pytme.svg)](https://pypi.org/project/pytme/)
5
+
6
+ **[Documentation](https://kosinskilab.github.io/pyTME/)** | **[Installation](https://kosinskilab.github.io/pyTME/quickstart/installation.html)** | **[API](https://kosinskilab.github.io/pyTME/reference/index.html)**
7
+
8
+ PyTME is a Python library for data-intensive n-dimensional template matching using CPUs and GPUs.
9
+
10
+ With its [backend-agnostic design](https://kosinskilab.github.io/pyTME/reference/backends.html), the same code can be run on diverse hardware platforms using a best-of-breed approach. The underyling abstract backend specification allows for adding new backends to benefit from gains in performance and capabilities without modifying the library's core routines. The implementation of template matching scores is modular and provides developers with a flexible framework for rapid prototyping. Furthermore, pyTME supports a unique callback capability through [analyzers](https://kosinskilab.github.io/pyTME/reference/analyzer.html), which allows for injection of custom code, enabling real-time processing and manipulation of results
11
+
12
+ PyTME includes a [graphical user interface](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html#practical-example) that provides simplified mask creation, interactive filter exploration, result visualization, and manual refinement capabilities. This GUI serves as an accessible entry point to the library's core functionalities, allowing users to efficiently interact with and analyze their data.
13
+
14
+ Finally, pyTME offers specialized tools for cryogenic electron microscopy data, such as wedge masks, CTF correction, as well as [means for handling structural data](https://kosinskilab.github.io/pyTME/reference/data_structure.html). Through dedicated [integrations](https://kosinskilab.github.io/pyTME/quickstart/integrations.html), the output of pyTME seamlessly integrates with commonly used cryogenic electron microscopy software such as RELION, Dynamo and IMOD.
15
+
16
+ Running into bugs or missing a feature? Help us improve the project by opening an [issue](https://github.com/KosinskiLab/pyTME/issues).
17
+
18
+ ## Installation
19
+
20
+ We recommend installation using one of the following methods
21
+
22
+ | Method | Command |
23
+ |----------|---------------------------------------------------------|
24
+ | PyPi | `pip install pytme` |
25
+ | Source | `pip install git+https://github.com/KosinskiLab/pyTME` |
26
+ | Docker | `docker build -t pytme -f docker/Dockerfile_GPU .` |
27
+
28
+ You can find alternative installation methods in the [documentation](https://kosinskilab.github.io/pyTME/quickstart/installation.html).
29
+
30
+
31
+ ## Quickstart
32
+
33
+ Learn how to get started with
34
+
35
+ - [Preprocessing:](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html) Picking the right mask and filter for template matching.
36
+ - [Template matching:](https://kosinskilab.github.io/pyTME/quickstart/match_template.html) Find your template of interest.
37
+ - [Postprocessing](https://kosinskilab.github.io/pyTME/quickstart/postprocessing.html) Analyze template matching results and downstream integrations.
38
+
39
+ ## How to Cite
40
+
41
+ If PyTME contributed significantly to your research, please cite the corresponding publication on [SoftwareX](https://www.sciencedirect.com/science/article/pii/S2352711024000074).
42
+
43
+ ```bibtex
44
+ @article{Maurer:2024aa,
45
+ author = {Maurer, Valentin J. and Siggel, Marc and Kosinski, Jan},
46
+ journal = {SoftwareX},
47
+ pages = {101636},
48
+ title = {PyTME (Python Template Matching Engine): A fast, flexible, and multi-purpose template matching library for cryogenic electron microscopy data},
49
+ volume = {25},
50
+ year = {2024}}
51
+ ```
@@ -7,7 +7,7 @@ name="pytme"
7
7
  authors = [
8
8
  { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
9
9
  ]
10
- version="0.2.0b"
10
+ version="0.2.2"
11
11
  description="Python Template Matching Engine"
12
12
  readme="README.md"
13
13
  requires-python = ">=3.11"
@@ -16,15 +16,14 @@ dependencies=[
16
16
  "numpy>=1.22.2",
17
17
  "pyfftw>=0.13.1",
18
18
  "pytest>=6.2.5",
19
- "PyWavelets>=1.2.0",
20
19
  "PyYAML>=6.0",
21
20
  "scikit-image>=0.19.0",
22
21
  "scikit_learn>=1.2.1",
23
22
  "scipy>=1.9.1",
24
23
  "pybind11",
25
24
  "psutil",
26
- "tifffile[all]",
27
- "h5py"
25
+ "tifffile",
26
+ "h5py",
28
27
  ]
29
28
  license = {text = "Proprietary"}
30
29
  classifiers = [
@@ -32,6 +31,18 @@ classifiers = [
32
31
  "Operating System :: OS Independent",
33
32
  ]
34
33
 
34
+ [project.optional-dependencies]
35
+ cupy = ["cupy", "voltools==0.6.0"]
36
+ pytorch = ["torch"]
37
+ jax = ["jax", "jaxlib"]
38
+ all = [
39
+ "cupy",
40
+ "voltools==0.6.0",
41
+ "torch",
42
+ "jax",
43
+ "jaxlib"
44
+ ]
45
+
35
46
  [project.scripts]
36
47
  match_template = "scripts:match_template.main"
37
48
  estimate_ram_usage = "scripts:estimate_ram_usage.main"
@@ -17,20 +17,22 @@ tme/__init__.py
17
17
  tme/__version__.py
18
18
  tme/analyzer.py
19
19
  tme/density.py
20
- tme/helpers.py
21
- tme/matching_constrained.py
22
20
  tme/matching_data.py
23
21
  tme/matching_exhaustive.py
24
22
  tme/matching_memory.py
25
23
  tme/matching_optimization.py
24
+ tme/matching_scores.py
26
25
  tme/matching_utils.py
26
+ tme/memory.py
27
27
  tme/orientations.py
28
28
  tme/parser.py
29
29
  tme/preprocessor.py
30
30
  tme/structure.py
31
31
  tme/types.py
32
32
  tme/backends/__init__.py
33
+ tme/backends/_jax_utils.py
33
34
  tme/backends/cupy_backend.py
35
+ tme/backends/jax_backend.py
34
36
  tme/backends/matching_backend.py
35
37
  tme/backends/mlx_backend.py
36
38
  tme/backends/npfftw_backend.py
@@ -57,4 +59,10 @@ tme/data/c48u8649.npy
57
59
  tme/data/c600v.npy
58
60
  tme/data/c600vc.npy
59
61
  tme/data/metadata.yaml
60
- tme/data/quat_to_numpy.py
62
+ tme/data/quat_to_numpy.py
63
+ tme/preprocessing/__init__.py
64
+ tme/preprocessing/_utils.py
65
+ tme/preprocessing/composable_filter.py
66
+ tme/preprocessing/compose.py
67
+ tme/preprocessing/frequency_filters.py
68
+ tme/preprocessing/tilt_series.py
@@ -0,0 +1,257 @@
1
+ #!python3
2
+ """ Prepare orientations stack for refinement.
3
+
4
+ Copyright (c) 2023 European Molecular Biology Laboratory
5
+
6
+ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
7
+ """
8
+ import argparse
9
+ from os.path import splitext
10
+
11
+ import numpy as np
12
+
13
+ from tme import Density, Orientations
14
+ from tme.matching_utils import (
15
+ load_pickle,
16
+ generate_tempfile_name,
17
+ rotation_aligning_vectors,
18
+ euler_from_rotationmatrix,
19
+ euler_to_rotationmatrix,
20
+ )
21
+
22
+
23
+ class ProgressBar:
24
+ """
25
+ ASCII progress bar.
26
+ """
27
+
28
+ def __init__(self, message : str, nchars : int, total : int):
29
+ self._size = nchars - len(message)
30
+ self._message = message
31
+ self._total = total
32
+
33
+ def update(self, cur):
34
+ x = int(cur * self._size / self._total)
35
+ print(
36
+ "%s[%s%s] %i/%i\r"
37
+ % (self._message, "#" * x, "." * (self._size - x), cur, self._total),
38
+ end=''
39
+ )
40
+
41
+
42
+ def parse_args():
43
+ parser = argparse.ArgumentParser(
44
+ description="Extract matching candidates for further refinement."
45
+ )
46
+
47
+ io_group = parser.add_argument_group("Input / Output")
48
+ io_group.add_argument(
49
+ "--input_file",
50
+ required=False,
51
+ type=str,
52
+ help="Path to the output of match_template.py.",
53
+ )
54
+ io_group.add_argument(
55
+ "--orientations",
56
+ required=True,
57
+ type=str,
58
+ help="Path to file generated by postprocess.py using output_format orientations.",
59
+ )
60
+ io_group.add_argument(
61
+ "--target",
62
+ required=False,
63
+ type=str,
64
+ help="Extract candidates from this target, can be at different sampling rate.",
65
+ )
66
+ io_group.add_argument(
67
+ "--template",
68
+ required=False,
69
+ type=str,
70
+ help="Extract candidates from this target, can be at different sampling rate.",
71
+ )
72
+ io_group.add_argument(
73
+ "-o",
74
+ "--output_file",
75
+ required=True,
76
+ type=str,
77
+ help="Path to output HDF5 file.",
78
+ )
79
+
80
+ alignment_group = parser.add_argument_group("Alignment")
81
+ alignment_group.add_argument(
82
+ "--align_orientations",
83
+ action="store_true",
84
+ required=False,
85
+ help="Whether to align extracted orientations based on their angles. Allows "
86
+ "for efficient subsequent sampling of cone angles.",
87
+ )
88
+ alignment_group.add_argument(
89
+ "--angles_are_vector",
90
+ action="store_true",
91
+ required=False,
92
+ help="Considers euler_z euler_y, euler_x as vector that will be rotated to align "
93
+ "with the z-axis (1,0,0). Only considered when --align_orientations is set."
94
+ )
95
+ alignment_group.add_argument(
96
+ "--interpolation_order",
97
+ dest="interpolation_order",
98
+ required=False,
99
+ type=int,
100
+ default=1,
101
+ help="Interpolation order for alignment, less than zero is no interpolation."
102
+ )
103
+
104
+ extraction_group = parser.add_argument_group("Extraction")
105
+ extraction_group.add_argument(
106
+ "--box_size",
107
+ required=False,
108
+ type=int,
109
+ help="Box size for extraction, defaults to two times the template.",
110
+ )
111
+ extraction_group.add_argument(
112
+ "--translation_uncertainty",
113
+ required=False,
114
+ type=int,
115
+ help="Sets box size for extraction to template box plus this value.",
116
+ )
117
+ extraction_group.add_argument(
118
+ "--keep_out_of_box",
119
+ action="store_true",
120
+ required=False,
121
+ help="Whether to keep orientations that fall outside the box. If the "
122
+ "orientations are sensible, it is safe to pass this flag.",
123
+ )
124
+
125
+ args = parser.parse_args()
126
+
127
+ data_present = args.target is not None and args.template is not None
128
+ if args.input_file is None and not data_present:
129
+ raise ValueError(
130
+ "Either --input_file or --target and --template need to be specified."
131
+ )
132
+ elif args.input_file is not None and data_present:
133
+ raise ValueError(
134
+ "Please specific either --input_file or --target and --template."
135
+ )
136
+
137
+ return args
138
+
139
+
140
+ def main():
141
+ args = parse_args()
142
+ orientations = Orientations.from_file(args.orientations)
143
+
144
+ if args.input_file is not None:
145
+ data = load_pickle(args.input_file)
146
+ target_origin, _, sampling_rate, cli_args = data[-1]
147
+ args.target, args.template = cli_args.target, cli_args.template
148
+
149
+ target = Density.from_file(args.target, use_memmap=True)
150
+
151
+ try:
152
+ template = Density.from_file(args.template)
153
+ except Exception:
154
+ template = Density.from_structure(args.template, sampling_rate = target.sampling_rate)
155
+
156
+ box_size = args.box_size
157
+ if box_size is None:
158
+ box_size = np.multiply(template.shape, 2)
159
+ if args.translation_uncertainty is not None:
160
+ box_size = np.add(template.shape, args.translation_uncertainty)
161
+
162
+ box_size = np.array(box_size)
163
+ box_size = np.repeat(box_size, template.data.ndim // box_size.size).astype(int)
164
+
165
+ extraction_shape = np.copy(box_size)
166
+ if args.align_orientations:
167
+ extraction_shape[:] = int(np.linalg.norm(box_size) + 1)
168
+
169
+ orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
170
+ target_shape=target.shape,
171
+ extraction_shape=extraction_shape,
172
+ drop_out_of_box=not args.keep_out_of_box,
173
+ return_orientations=True,
174
+ )
175
+
176
+ if args.align_orientations:
177
+ orientations.rotations = orientations.rotations.astype(np.float32)
178
+ for index in range(orientations.rotations.shape[0]):
179
+ rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
180
+ rotation_matrix = np.linalg.inv(rotation_matrix)
181
+ if args.angles_are_vector:
182
+ rotation_matrix = rotation_aligning_vectors(
183
+ orientations.rotations[index], target_vector=(1,0,0)
184
+ )
185
+ orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)
186
+
187
+
188
+ filename = generate_tempfile_name()
189
+ output_dtype = target.data.dtype
190
+ if args.align_orientations is not None:
191
+ output_dtype = np.float32
192
+
193
+ target.data = target.data.astype(output_dtype)
194
+
195
+ dens = Density(
196
+ np.memmap(
197
+ filename,
198
+ mode="w+",
199
+ shape=(len(obs_slices), *box_size),
200
+ dtype=output_dtype,
201
+ ),
202
+ sampling_rate=(1, *target.sampling_rate),
203
+ origin=(0, *target.origin),
204
+ )
205
+ dens.data[:] = target.metadata["mean"]
206
+
207
+ data_subset = np.zeros(extraction_shape, dtype = target.data.dtype)
208
+ pbar = ProgressBar(message = "Orientation ", nchars = 80, total = len(obs_slices))
209
+ for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
210
+ pbar.update(index + 1)
211
+
212
+ data_subset.fill(0)
213
+ data_subset[cand_slice] = target.data[obs_slice]
214
+ target_subset = Density(
215
+ data_subset,
216
+ sampling_rate=target.sampling_rate,
217
+ origin=target.origin,
218
+ )
219
+
220
+ if args.align_orientations:
221
+ rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
222
+ target_subset = target_subset.rigid_transform(
223
+ rotation_matrix=rotation_matrix,
224
+ use_geometric_center=True,
225
+ order=args.interpolation_order,
226
+ )
227
+ target_subset.pad(box_size, center=True)
228
+
229
+ # target_value = target.data[tuple(orientations.translations[index].astype(int))]
230
+ # center = np.divide(target_subset.data.shape, 2).astype(int ) + np.mod(target_subset.shape, 2)
231
+ # print(np.where(target_subset.data == target_value), center)
232
+ # print(target_subset.data[tuple(center.astype(int))],
233
+ # target_value,
234
+ # target_subset.data[tuple(center.astype(int))] == target_value
235
+ # )
236
+
237
+ dens.data[index] = target_subset.data
238
+ print("")
239
+
240
+ target_meta = {
241
+ k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
242
+ }
243
+ dens.metadata.update(target_meta)
244
+ dens.metadata["batch_dimension"] = (0, )
245
+
246
+ dens.to_file(args.output_file)
247
+ orientations.to_file(
248
+ f"{splitext(args.output_file)[0]}_aligned.tsv",
249
+ file_format="text"
250
+ )
251
+ # orientations.to_file(
252
+ # f"{splitext(args.output_file)[0]}_aligned.star",
253
+ # file_format="relion"
254
+ # )
255
+
256
+ if __name__ == "__main__":
257
+ main()