midas-nf-preprocess 0.1.0__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 (48) hide show
  1. midas_nf_preprocess-0.1.0/PKG-INFO +96 -0
  2. midas_nf_preprocess-0.1.0/README.md +67 -0
  3. midas_nf_preprocess-0.1.0/midas_nf_preprocess/__init__.py +45 -0
  4. midas_nf_preprocess-0.1.0/midas_nf_preprocess/cli.py +62 -0
  5. midas_nf_preprocess-0.1.0/midas_nf_preprocess/device.py +69 -0
  6. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/__init__.py +83 -0
  7. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/binary_io.py +118 -0
  8. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/cli.py +70 -0
  9. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/geometry.py +193 -0
  10. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/hkls.py +118 -0
  11. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/orientations.py +176 -0
  12. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/params.py +115 -0
  13. midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/pipeline.py +301 -0
  14. midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/__init__.py +30 -0
  15. midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/cli.py +58 -0
  16. midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/grid.py +237 -0
  17. midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/io.py +64 -0
  18. midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/params.py +62 -0
  19. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/__init__.py +40 -0
  20. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/cli.py +79 -0
  21. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/io.py +92 -0
  22. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/log_filter.py +130 -0
  23. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/median.py +85 -0
  24. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/params.py +133 -0
  25. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/peaks.py +355 -0
  26. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/pipeline.py +272 -0
  27. midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/spots_io.py +276 -0
  28. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/__init__.py +86 -0
  29. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/cli.py +150 -0
  30. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/crystal.py +136 -0
  31. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/dispatch.py +99 -0
  32. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_cache.py +141 -0
  33. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_grains.py +100 -0
  34. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_scratch.py +207 -0
  35. midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/io.py +77 -0
  36. midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/__init__.py +36 -0
  37. midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/cli.py +66 -0
  38. midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/filter.py +180 -0
  39. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/PKG-INFO +96 -0
  40. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/SOURCES.txt +46 -0
  41. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/dependency_links.txt +1 -0
  42. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/entry_points.txt +2 -0
  43. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/requires.txt +13 -0
  44. midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/top_level.txt +1 -0
  45. midas_nf_preprocess-0.1.0/pyproject.toml +57 -0
  46. midas_nf_preprocess-0.1.0/setup.cfg +4 -0
  47. midas_nf_preprocess-0.1.0/tests/test_cli_umbrella.py +86 -0
  48. midas_nf_preprocess-0.1.0/tests/test_device.py +104 -0
@@ -0,0 +1,96 @@
1
+ Metadata-Version: 2.4
2
+ Name: midas-nf-preprocess
3
+ Version: 0.1.0
4
+ Summary: Differentiable PyTorch port of NF-HEDM preprocessing (CPU/CUDA/MPS): hex grid, tomo filter, diffraction-spot prediction, image processing
5
+ Author-email: Hemant Sharma <hsharma@anl.gov>
6
+ License-Expression: BSD-3-Clause
7
+ Project-URL: Homepage, https://github.com/marinerhemant/MIDAS
8
+ Project-URL: Documentation, https://github.com/marinerhemant/MIDAS/tree/master/packages/midas_nf_preprocess
9
+ Project-URL: Issues, https://github.com/marinerhemant/MIDAS/issues
10
+ Keywords: MIDAS,NF-HEDM,diffraction,PyTorch,differentiable,image processing,peak detection,near-field,preprocessing,hex grid,tomography
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Topic :: Scientific/Engineering :: Physics
15
+ Classifier: Topic :: Scientific/Engineering :: Image Processing
16
+ Requires-Python: >=3.9
17
+ Description-Content-Type: text/markdown
18
+ Requires-Dist: numpy>=1.22
19
+ Requires-Dist: torch>=2.1
20
+ Requires-Dist: tifffile>=2022.0
21
+ Requires-Dist: midas-diffract>=0.1.0
22
+ Requires-Dist: midas-stress>=0.6.0
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest>=7.0; extra == "dev"
25
+ Requires-Dist: pytest-cov; extra == "dev"
26
+ Requires-Dist: scipy>=1.9; extra == "dev"
27
+ Provides-Extra: parity
28
+ Requires-Dist: scipy>=1.9; extra == "parity"
29
+
30
+ # midas-nf-process-images
31
+
32
+ Differentiable PyTorch port of `NF_HEDM/src/ProcessImagesCombined.c` for CPU, CUDA, and MPS.
33
+
34
+ ## Pipeline
35
+
36
+ For one layer of a near-field HEDM scan:
37
+
38
+ 1. **Load** `NrFilesPerLayer` raw TIFF frames into a `[N, Z, Y]` tensor.
39
+ 2. **Temporal median** across frames → per-pixel background `[Z, Y]`.
40
+ 3. **Per-frame**:
41
+ - Subtract median + `BlanketSubtraction`, clamp at 0.
42
+ - Spatial median (3x3 or 5x5).
43
+ - LoG response at one or more scales.
44
+ - Threshold and label connected components → integer spot labels (detached).
45
+ - Sigmoid surrogate → continuous spot-probability map (autograd path).
46
+ 4. **Accumulate** spot pixels into a bit-packed `SpotsInfo.bin` mmap (drop-in compatible
47
+ with `FitOrientationOMP` / `simulateNF`).
48
+
49
+ Phases 1, 2, and 3 (filtered, log_response, soft spot probability) are differentiable
50
+ end-to-end. Connected components and the binary `SpotsInfo.bin` output are computed on
51
+ detached tensors and do not participate in autograd.
52
+
53
+ ## Quickstart
54
+
55
+ ```python
56
+ from midas_nf_process_images import ProcessParams, ProcessImagesPipeline
57
+
58
+ params = ProcessParams.from_paramfile("ps.txt")
59
+ pipe = ProcessImagesPipeline(params, device="cuda") # or "mps", "cpu"
60
+ spots = pipe.process_layer(layer_nr=1)
61
+ spots.write("/path/to/SpotsInfo.bin")
62
+
63
+ # Or all layers in one call:
64
+ spots = pipe.process_all(layers=range(1, params.n_distances + 1))
65
+ ```
66
+
67
+ ## CLI
68
+
69
+ ```bash
70
+ midas-nf-process-images <ParameterFile> <LayerNr> [--device cuda] [--dtype fp32] [--n-cpus 8]
71
+ ```
72
+
73
+ Behaves like the C `ProcessImagesCombined` executable.
74
+
75
+ ## Backend selection
76
+
77
+ Same contract as `midas-transforms`:
78
+
79
+ - Default device: `cuda` -> `mps` -> `cpu`.
80
+ - Default dtype: `float64` on CPU, `float32` on CUDA/MPS.
81
+ - Override with `device=` / `dtype=` kwargs, or env vars
82
+ `MIDAS_NF_PROCESS_IMAGES_DEVICE` / `MIDAS_NF_PROCESS_IMAGES_DTYPE`.
83
+
84
+ ## Differentiability
85
+
86
+ The end-to-end pipeline produces three tensors:
87
+
88
+ - `filtered`: median-subtracted, spatial-median-filtered image (autograd).
89
+ - `log_response`: LoG convolution output (autograd).
90
+ - `spot_prob`: soft, continuous spot-probability map via `sigmoid(L / temperature)` (autograd).
91
+
92
+ Plus one detached output:
93
+
94
+ - `labels`: integer connected-component IDs from hard-thresholded LoG (no gradient).
95
+
96
+ Optimization through the discrete spot mask uses the soft `spot_prob` surrogate.
@@ -0,0 +1,67 @@
1
+ # midas-nf-process-images
2
+
3
+ Differentiable PyTorch port of `NF_HEDM/src/ProcessImagesCombined.c` for CPU, CUDA, and MPS.
4
+
5
+ ## Pipeline
6
+
7
+ For one layer of a near-field HEDM scan:
8
+
9
+ 1. **Load** `NrFilesPerLayer` raw TIFF frames into a `[N, Z, Y]` tensor.
10
+ 2. **Temporal median** across frames → per-pixel background `[Z, Y]`.
11
+ 3. **Per-frame**:
12
+ - Subtract median + `BlanketSubtraction`, clamp at 0.
13
+ - Spatial median (3x3 or 5x5).
14
+ - LoG response at one or more scales.
15
+ - Threshold and label connected components → integer spot labels (detached).
16
+ - Sigmoid surrogate → continuous spot-probability map (autograd path).
17
+ 4. **Accumulate** spot pixels into a bit-packed `SpotsInfo.bin` mmap (drop-in compatible
18
+ with `FitOrientationOMP` / `simulateNF`).
19
+
20
+ Phases 1, 2, and 3 (filtered, log_response, soft spot probability) are differentiable
21
+ end-to-end. Connected components and the binary `SpotsInfo.bin` output are computed on
22
+ detached tensors and do not participate in autograd.
23
+
24
+ ## Quickstart
25
+
26
+ ```python
27
+ from midas_nf_process_images import ProcessParams, ProcessImagesPipeline
28
+
29
+ params = ProcessParams.from_paramfile("ps.txt")
30
+ pipe = ProcessImagesPipeline(params, device="cuda") # or "mps", "cpu"
31
+ spots = pipe.process_layer(layer_nr=1)
32
+ spots.write("/path/to/SpotsInfo.bin")
33
+
34
+ # Or all layers in one call:
35
+ spots = pipe.process_all(layers=range(1, params.n_distances + 1))
36
+ ```
37
+
38
+ ## CLI
39
+
40
+ ```bash
41
+ midas-nf-process-images <ParameterFile> <LayerNr> [--device cuda] [--dtype fp32] [--n-cpus 8]
42
+ ```
43
+
44
+ Behaves like the C `ProcessImagesCombined` executable.
45
+
46
+ ## Backend selection
47
+
48
+ Same contract as `midas-transforms`:
49
+
50
+ - Default device: `cuda` -> `mps` -> `cpu`.
51
+ - Default dtype: `float64` on CPU, `float32` on CUDA/MPS.
52
+ - Override with `device=` / `dtype=` kwargs, or env vars
53
+ `MIDAS_NF_PROCESS_IMAGES_DEVICE` / `MIDAS_NF_PROCESS_IMAGES_DTYPE`.
54
+
55
+ ## Differentiability
56
+
57
+ The end-to-end pipeline produces three tensors:
58
+
59
+ - `filtered`: median-subtracted, spatial-median-filtered image (autograd).
60
+ - `log_response`: LoG convolution output (autograd).
61
+ - `spot_prob`: soft, continuous spot-probability map via `sigmoid(L / temperature)` (autograd).
62
+
63
+ Plus one detached output:
64
+
65
+ - `labels`: integer connected-component IDs from hard-thresholded LoG (no gradient).
66
+
67
+ Optimization through the discrete spot mask uses the soft `spot_prob` surrogate.
@@ -0,0 +1,45 @@
1
+ """midas-nf-preprocess: Differentiable PyTorch port of NF-HEDM preprocessing.
2
+
3
+ Bundles the preprocessing stages that sit between raw experimental inputs
4
+ and the orientation-fitting step (``FitOrientationOMP``):
5
+
6
+ - ``hex_grid`` : voxel grid generation (port of MakeHexGrid)
7
+ - ``tomo_filter`` : grid masking from a tomography image
8
+ (port of filterGridfromTomo)
9
+ - ``seed_orientations``: NF seed-orientation generation (cache, from-scratch,
10
+ or from an FF Grains.csv)
11
+ - ``diffr_spots`` : per-orientation diffraction-spot prediction
12
+ (differentiable port of MakeDiffrSpots; wraps
13
+ ``midas_diffract.HEDMForwardModel.calc_bragg_geometry``)
14
+ - ``process_images`` : raw TIFF -> SpotsInfo.bin
15
+ (differentiable port of ProcessImagesCombined)
16
+
17
+ The shared ``device`` module provides device/dtype resolution
18
+ (env vars ``MIDAS_NF_PREPROCESS_DEVICE`` / ``MIDAS_NF_PREPROCESS_DTYPE``)
19
+ matching the convention used by ``midas-transforms``.
20
+ """
21
+
22
+ __version__ = "0.1.0"
23
+
24
+ from . import (
25
+ device,
26
+ diffr_spots,
27
+ hex_grid,
28
+ process_images,
29
+ seed_orientations,
30
+ tomo_filter,
31
+ )
32
+ from .device import resolve_device, resolve_dtype, apply_cpu_threads
33
+
34
+ __all__ = [
35
+ "__version__",
36
+ "device",
37
+ "diffr_spots",
38
+ "hex_grid",
39
+ "process_images",
40
+ "seed_orientations",
41
+ "tomo_filter",
42
+ "resolve_device",
43
+ "resolve_dtype",
44
+ "apply_cpu_threads",
45
+ ]
@@ -0,0 +1,62 @@
1
+ """Umbrella CLI: ``midas-nf-preprocess <subcommand> ...``.
2
+
3
+ Subcommands:
4
+
5
+ - ``hex-grid`` : generate the voxel grid (port of MakeHexGrid)
6
+ - ``tomo-filter`` : mask the grid by a tomography image / bbox
7
+ - ``diffr-spots`` : forward-simulate diffraction spots (port of MakeDiffrSpots)
8
+ - ``process-images`` : raw TIFF -> SpotsInfo.bin (port of ProcessImagesCombined)
9
+
10
+ Each subcommand mirrors its standalone ``python -m`` invocation:
11
+
12
+ python -m midas_nf_preprocess.hex_grid <args> ==
13
+ midas-nf-preprocess hex-grid <args>
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import argparse
19
+ import sys
20
+
21
+ from . import __version__
22
+ from .diffr_spots import cli as diffr_spots_cli
23
+ from .hex_grid import cli as hex_grid_cli
24
+ from .process_images import cli as process_images_cli
25
+ from .seed_orientations import cli as seed_orientations_cli
26
+ from .tomo_filter import cli as tomo_filter_cli
27
+
28
+
29
+ _SUBCOMMANDS = {
30
+ "hex-grid": hex_grid_cli,
31
+ "tomo-filter": tomo_filter_cli,
32
+ "diffr-spots": diffr_spots_cli,
33
+ "process-images": process_images_cli,
34
+ "seed-orientations": seed_orientations_cli,
35
+ }
36
+
37
+
38
+ def _build_parser() -> argparse.ArgumentParser:
39
+ parser = argparse.ArgumentParser(
40
+ prog="midas-nf-preprocess",
41
+ description="Differentiable PyTorch port of NF-HEDM preprocessing.",
42
+ )
43
+ parser.add_argument(
44
+ "--version", action="version", version=f"midas-nf-preprocess {__version__}"
45
+ )
46
+ subparsers = parser.add_subparsers(
47
+ dest="subcommand", required=True, metavar="<subcommand>"
48
+ )
49
+ for name, mod in _SUBCOMMANDS.items():
50
+ sp = subparsers.add_parser(name, help=mod.__doc__.splitlines()[0] if mod.__doc__ else None)
51
+ mod.add_arguments(sp)
52
+ sp.set_defaults(_run=mod.run)
53
+ return parser
54
+
55
+
56
+ def main(argv: list[str] | None = None) -> int:
57
+ args = _build_parser().parse_args(argv)
58
+ return args._run(args)
59
+
60
+
61
+ if __name__ == "__main__":
62
+ sys.exit(main())
@@ -0,0 +1,69 @@
1
+ """Device / dtype resolution.
2
+
3
+ Same contract as ``midas-transforms``:
4
+ - device default: cuda > mps > cpu (auto-detect)
5
+ - dtype default: float64 on cpu (matches the C binaries),
6
+ float32 on cuda/mps (matches the GPU codepath idioms).
7
+ - Both float32 and float64 are supported on every device.
8
+ - Env vars MIDAS_NF_PREPROCESS_DEVICE / MIDAS_NF_PREPROCESS_DTYPE
9
+ override defaults.
10
+ - Library ``device=``/``dtype=`` kwargs override env vars.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import os
16
+ from typing import Optional, Union
17
+
18
+ import torch
19
+
20
+ _DTYPE_ALIAS = {
21
+ "float32": torch.float32,
22
+ "fp32": torch.float32,
23
+ "float64": torch.float64,
24
+ "fp64": torch.float64,
25
+ "double": torch.float64,
26
+ "single": torch.float32,
27
+ }
28
+
29
+ ENV_DEVICE = "MIDAS_NF_PREPROCESS_DEVICE"
30
+ ENV_DTYPE = "MIDAS_NF_PREPROCESS_DTYPE"
31
+
32
+
33
+ def resolve_device(user: Optional[Union[str, torch.device]]) -> torch.device:
34
+ """Pick a torch.device. Precedence: arg > env var > auto-detect."""
35
+ if isinstance(user, torch.device):
36
+ return user
37
+ if user is None:
38
+ user = os.environ.get(ENV_DEVICE)
39
+ if user is None:
40
+ if torch.cuda.is_available():
41
+ return torch.device("cuda")
42
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
43
+ return torch.device("mps")
44
+ return torch.device("cpu")
45
+ return torch.device(user)
46
+
47
+
48
+ def resolve_dtype(
49
+ device: torch.device, user: Optional[Union[str, torch.dtype]]
50
+ ) -> torch.dtype:
51
+ """Pick a torch.dtype. Precedence: arg > env var > per-device default."""
52
+ if isinstance(user, torch.dtype):
53
+ return user
54
+ if user is None:
55
+ user = os.environ.get(ENV_DTYPE)
56
+ if user is not None:
57
+ try:
58
+ return _DTYPE_ALIAS[user.lower()]
59
+ except KeyError as e:
60
+ raise ValueError(
61
+ f"Unsupported dtype '{user}'. Use one of {sorted(_DTYPE_ALIAS)}."
62
+ ) from e
63
+ return torch.float64 if device.type == "cpu" else torch.float32
64
+
65
+
66
+ def apply_cpu_threads(num_procs: int, device: torch.device) -> None:
67
+ """Honor a legacy ``numProcs`` argv on CPU only (no-op on GPU)."""
68
+ if device.type == "cpu" and num_procs > 0:
69
+ torch.set_num_threads(num_procs)
@@ -0,0 +1,83 @@
1
+ """diffr_spots: per-orientation diffraction-spot prediction (port of MakeDiffrSpots).
2
+
3
+ Differentiable PyTorch port of ``NF_HEDM/src/MakeDiffrSpots.c``. For each seed
4
+ orientation and each HKL, solve the Bragg condition:
5
+
6
+ - ``omega(s)`` (1 or 2 solutions): rotation angles where the reciprocal-space
7
+ G-vector intersects the Ewald sphere
8
+ - ``eta``: azimuthal angle on the detector
9
+ - ``(yl, zl)``: lab-frame detector position at the primary distance
10
+
11
+ then filter by ``OmegaRange``, ``BoxSize``, and ``ExcludePoleAngle``. Output is
12
+ either a structured tensor bundle or the same three binary files the C
13
+ executable produces:
14
+
15
+ - ``DiffractionSpots.bin`` : ``[TotalSpots, 3]`` float64 ``(yl, zl, omega)``
16
+ - ``Key.bin`` : ``[N_orient, 2]`` int32 ``(count_i, offset_i)``
17
+ - ``OrientMat.bin`` : ``[N_orient, 3, 3]`` float64 row-major
18
+
19
+ Differentiability:
20
+
21
+ - ``omega``, ``eta``, ``yl``, ``zl`` are differentiable in the input
22
+ quaternions / Euler angles / orientation matrices and in the lattice
23
+ parameters that produced ``hkls``.
24
+ - The hard ``OmegaRange`` / ``BoxSize`` / ``ExcludePoleAngle`` filters return
25
+ a boolean ``valid`` mask (no gradient), matching the convention used by
26
+ ``midas_diffract.SpotDescriptors``.
27
+
28
+ The geometry primitives (``calc_omega``, ``calc_eta``, ``calc_spot_position``)
29
+ match the C implementations in ``MakeDiffrSpots.c`` L99-L204 line-by-line and
30
+ extend them to batched, vectorized PyTorch operations.
31
+ """
32
+
33
+ from .params import DiffrSpotsParams
34
+ from .orientations import (
35
+ quat_to_orient_matrix,
36
+ orient_matrix_to_quat,
37
+ euler_to_orient_matrix,
38
+ )
39
+ from .geometry import (
40
+ bragg_omega_eta,
41
+ calc_eta_deg,
42
+ calc_spot_position,
43
+ rotate_around_z,
44
+ )
45
+ from .pipeline import (
46
+ DiffrSpotsPipeline,
47
+ DiffrSpotsResult,
48
+ predict_spots,
49
+ )
50
+ from .binary_io import (
51
+ write_diffr_spots_bin,
52
+ write_key_bin,
53
+ write_orient_mat_bin,
54
+ read_diffr_spots_bin,
55
+ read_key_bin,
56
+ read_orient_mat_bin,
57
+ write_all,
58
+ )
59
+ from .hkls import HKLEntry, read_hkls_csv, read_seed_orientations
60
+
61
+ __all__ = [
62
+ "DiffrSpotsParams",
63
+ "quat_to_orient_matrix",
64
+ "orient_matrix_to_quat",
65
+ "euler_to_orient_matrix",
66
+ "bragg_omega_eta",
67
+ "calc_eta_deg",
68
+ "calc_spot_position",
69
+ "rotate_around_z",
70
+ "DiffrSpotsPipeline",
71
+ "DiffrSpotsResult",
72
+ "predict_spots",
73
+ "write_diffr_spots_bin",
74
+ "write_key_bin",
75
+ "write_orient_mat_bin",
76
+ "read_diffr_spots_bin",
77
+ "read_key_bin",
78
+ "read_orient_mat_bin",
79
+ "write_all",
80
+ "HKLEntry",
81
+ "read_hkls_csv",
82
+ "read_seed_orientations",
83
+ ]
@@ -0,0 +1,118 @@
1
+ """Binary I/O for the three MakeDiffrSpots output files.
2
+
3
+ File formats (from MakeDiffrSpots.c L595-L626):
4
+
5
+ - ``Key.bin`` : ``int32 [N_orient * 2]`` little-endian, with
6
+ entries ``(count_i, offset_i)`` for each orientation.
7
+ - ``DiffractionSpots.bin`` : ``float64 [TotalSpots * 3]`` little-endian, with
8
+ row layout ``(yl, zl, omega)``.
9
+ - ``OrientMat.bin`` : ``float64 [N_orient * 9]`` little-endian, with
10
+ row-major flattened 3x3 matrices.
11
+
12
+ Total spots = sum of per-orientation counts; ``offset_i`` is the cumulative
13
+ prefix sum (excluding ``count_i`` itself) so that orientation i's spots span
14
+ ``DiffractionSpots[offset_i : offset_i + count_i]``.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ from pathlib import Path
20
+ from typing import Union
21
+
22
+ import numpy as np
23
+ import torch
24
+
25
+
26
+ # -----------------------------------------------------------------------------
27
+ # Writers
28
+ # -----------------------------------------------------------------------------
29
+
30
+
31
+ def write_key_bin(
32
+ counts: torch.Tensor, path: Union[str, Path]
33
+ ) -> None:
34
+ """Write ``Key.bin`` from a ``(N,)`` int tensor of per-orientation spot counts."""
35
+ if counts.ndim != 1:
36
+ raise ValueError(f"Expected 1D counts tensor, got shape {tuple(counts.shape)}")
37
+ counts_np = counts.detach().cpu().numpy().astype(np.int32)
38
+ offsets = np.zeros_like(counts_np)
39
+ if counts_np.size > 0:
40
+ offsets[1:] = np.cumsum(counts_np[:-1])
41
+ key = np.stack([counts_np, offsets], axis=1) # (N, 2)
42
+ path = Path(path)
43
+ path.parent.mkdir(parents=True, exist_ok=True)
44
+ key.tofile(path)
45
+
46
+
47
+ def write_diffr_spots_bin(
48
+ spots: torch.Tensor, path: Union[str, Path]
49
+ ) -> None:
50
+ """Write ``DiffractionSpots.bin`` from a ``(TotalSpots, 3)`` float tensor."""
51
+ if spots.ndim != 2 or spots.shape[1] != 3:
52
+ raise ValueError(
53
+ f"Expected (TotalSpots, 3) tensor, got shape {tuple(spots.shape)}"
54
+ )
55
+ arr = spots.detach().cpu().numpy().astype(np.float64)
56
+ path = Path(path)
57
+ path.parent.mkdir(parents=True, exist_ok=True)
58
+ arr.tofile(path)
59
+
60
+
61
+ def write_orient_mat_bin(
62
+ orient_mats: torch.Tensor, path: Union[str, Path]
63
+ ) -> None:
64
+ """Write ``OrientMat.bin`` from a ``(N, 3, 3)`` float tensor."""
65
+ if orient_mats.ndim != 3 or orient_mats.shape[1:] != (3, 3):
66
+ raise ValueError(
67
+ f"Expected (N, 3, 3) tensor, got shape {tuple(orient_mats.shape)}"
68
+ )
69
+ arr = orient_mats.detach().cpu().numpy().astype(np.float64)
70
+ path = Path(path)
71
+ path.parent.mkdir(parents=True, exist_ok=True)
72
+ arr.tofile(path)
73
+
74
+
75
+ def write_all(
76
+ output_dir: Union[str, Path],
77
+ counts: torch.Tensor,
78
+ spots: torch.Tensor,
79
+ orient_mats: torch.Tensor,
80
+ ) -> dict[str, Path]:
81
+ """Convenience wrapper that writes all three files into ``output_dir``."""
82
+ output_dir = Path(output_dir)
83
+ paths = {
84
+ "Key.bin": output_dir / "Key.bin",
85
+ "DiffractionSpots.bin": output_dir / "DiffractionSpots.bin",
86
+ "OrientMat.bin": output_dir / "OrientMat.bin",
87
+ }
88
+ write_key_bin(counts, paths["Key.bin"])
89
+ write_diffr_spots_bin(spots, paths["DiffractionSpots.bin"])
90
+ write_orient_mat_bin(orient_mats, paths["OrientMat.bin"])
91
+ return paths
92
+
93
+
94
+ # -----------------------------------------------------------------------------
95
+ # Readers (for parity tests / round-trip)
96
+ # -----------------------------------------------------------------------------
97
+
98
+
99
+ def read_key_bin(path: Union[str, Path], n_orient: int) -> torch.Tensor:
100
+ """Read ``Key.bin`` and return a ``(N, 2)`` int64 tensor ``(count_i, offset_i)``."""
101
+ arr = np.fromfile(path, dtype=np.int32, count=n_orient * 2).reshape(n_orient, 2)
102
+ return torch.from_numpy(arr.astype(np.int64))
103
+
104
+
105
+ def read_diffr_spots_bin(
106
+ path: Union[str, Path], n_spots: int
107
+ ) -> torch.Tensor:
108
+ """Read ``DiffractionSpots.bin`` and return a ``(N_spots, 3)`` float64 tensor."""
109
+ arr = np.fromfile(path, dtype=np.float64, count=n_spots * 3).reshape(n_spots, 3)
110
+ return torch.from_numpy(arr)
111
+
112
+
113
+ def read_orient_mat_bin(
114
+ path: Union[str, Path], n_orient: int
115
+ ) -> torch.Tensor:
116
+ """Read ``OrientMat.bin`` and return a ``(N, 3, 3)`` float64 tensor."""
117
+ arr = np.fromfile(path, dtype=np.float64, count=n_orient * 9).reshape(n_orient, 3, 3)
118
+ return torch.from_numpy(arr)
@@ -0,0 +1,70 @@
1
+ """``midas-nf-preprocess diffr-spots`` subcommand."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import sys
7
+ from pathlib import Path
8
+
9
+ from .params import DiffrSpotsParams
10
+ from .pipeline import DiffrSpotsPipeline
11
+
12
+
13
+ def add_arguments(parser: argparse.ArgumentParser) -> None:
14
+ parser.add_argument(
15
+ "parameter_file", help="MIDAS parameter file (Lsd, SeedOrientations, etc.)"
16
+ )
17
+ parser.add_argument(
18
+ "--device", default=None, help="cpu | cuda | mps (default: auto-detect)."
19
+ )
20
+ parser.add_argument(
21
+ "--dtype", default=None, help="fp32 | fp64 (default: fp64 on cpu)."
22
+ )
23
+ parser.add_argument(
24
+ "--hkls-csv",
25
+ default=None,
26
+ help="Override hkls.csv path (default: <DataDirectory>/hkls.csv).",
27
+ )
28
+ parser.add_argument(
29
+ "--seeds",
30
+ default=None,
31
+ help="Override SeedOrientations path (default: from parameter file).",
32
+ )
33
+ parser.add_argument(
34
+ "--output-dir",
35
+ default=None,
36
+ help="Override output directory for Key.bin / DiffractionSpots.bin / OrientMat.bin.",
37
+ )
38
+
39
+
40
+ def run(args: argparse.Namespace) -> int:
41
+ params = DiffrSpotsParams.from_paramfile(args.parameter_file)
42
+ pipe = DiffrSpotsPipeline(
43
+ params,
44
+ device=args.device,
45
+ dtype=args.dtype,
46
+ hkls_csv=args.hkls_csv,
47
+ seed_orientations_csv=args.seeds,
48
+ )
49
+ result, paths = pipe.run(output_dir=args.output_dir)
50
+ total = int(result.counts.sum().item())
51
+ print(
52
+ f"Wrote {len(paths)} files to {Path(paths['Key.bin']).parent} "
53
+ f"({pipe.n_orientations} orientations, {total} total spots)"
54
+ )
55
+ for name, p in paths.items():
56
+ print(f" {name}: {p}")
57
+ return 0
58
+
59
+
60
+ def main(argv: list[str] | None = None) -> int:
61
+ parser = argparse.ArgumentParser(
62
+ prog="midas-nf-preprocess diffr-spots",
63
+ description="Differentiable diffraction-spot prediction (port of MakeDiffrSpots).",
64
+ )
65
+ add_arguments(parser)
66
+ return run(parser.parse_args(argv))
67
+
68
+
69
+ if __name__ == "__main__":
70
+ sys.exit(main())