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.
- midas_nf_preprocess-0.1.0/PKG-INFO +96 -0
- midas_nf_preprocess-0.1.0/README.md +67 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/__init__.py +45 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/cli.py +62 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/device.py +69 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/__init__.py +83 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/binary_io.py +118 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/cli.py +70 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/geometry.py +193 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/hkls.py +118 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/orientations.py +176 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/params.py +115 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/diffr_spots/pipeline.py +301 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/__init__.py +30 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/cli.py +58 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/grid.py +237 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/io.py +64 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/hex_grid/params.py +62 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/__init__.py +40 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/cli.py +79 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/io.py +92 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/log_filter.py +130 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/median.py +85 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/params.py +133 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/peaks.py +355 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/pipeline.py +272 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/process_images/spots_io.py +276 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/__init__.py +86 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/cli.py +150 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/crystal.py +136 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/dispatch.py +99 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_cache.py +141 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_grains.py +100 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/from_scratch.py +207 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/seed_orientations/io.py +77 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/__init__.py +36 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/cli.py +66 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess/tomo_filter/filter.py +180 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/PKG-INFO +96 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/SOURCES.txt +46 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/dependency_links.txt +1 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/entry_points.txt +2 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/requires.txt +13 -0
- midas_nf_preprocess-0.1.0/midas_nf_preprocess.egg-info/top_level.txt +1 -0
- midas_nf_preprocess-0.1.0/pyproject.toml +57 -0
- midas_nf_preprocess-0.1.0/setup.cfg +4 -0
- midas_nf_preprocess-0.1.0/tests/test_cli_umbrella.py +86 -0
- 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())
|