midas-index 0.1.0__tar.gz

midas_index-0.1.0/PKG-INFO

@@ -0,0 +1,97 @@
+Metadata-Version: 2.4
+Name: midas-index
+Version: 0.1.0
+Summary: Pure-Python/PyTorch FF-HEDM indexer (drop-in replacement for IndexerOMP/IndexerGPU)
+Author-email: Hemant Sharma <hsharma@anl.gov>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/marinerhemant/MIDAS
+Project-URL: Documentation, https://github.com/marinerhemant/MIDAS/tree/master/packages/midas_index
+Project-URL: Issues, https://github.com/marinerhemant/MIDAS/issues
+Keywords: MIDAS,HEDM,indexing,PyTorch,far-field,diffraction,crystallography,polycrystal,grain
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.22
+Requires-Dist: torch>=2.0
+Requires-Dist: midas-diffract>=0.1.0
+Requires-Dist: midas-stress<1.0,>=0.5.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-benchmark; extra == "dev"
+
+# midas-index
+
+Pure-Python/PyTorch FF-HEDM indexer. Drop-in replacement for `IndexerOMP` /
+`IndexerGPU` from MIDAS, with seamless CPU / CUDA / MPS device switching.
+
+**Status:** v0.1.0 in development. Detailed design doc lives in
+`dev/implementation_plan.md` (gitignored).
+
+## Install
+
+```bash
+pip install midas-index
+```
+
+For local development:
+
+```bash
+cd packages/midas_index
+pip install -e .[dev]
+```
+
+## Quick start
+
+```bash
+# CLI — drop-in for IndexerOMP / IndexerGPU
+midas-index paramstest.txt 0 1 1000 8
+
+# Pin device / dtype via env vars (auto-detect: CUDA -> MPS -> CPU)
+MIDAS_INDEX_DEVICE=cuda MIDAS_INDEX_DTYPE=float32 \
+    midas-index paramstest.txt 0 1 1000 8
+```
+
+Library API:
+
+```python
+from midas_index import Indexer
+
+result = Indexer.from_param_file("paramstest.txt", device="cuda").run(
+    block_nr=0, n_blocks=1, n_spots_to_index=1000,
+)
+```
+
+## Drive from `ff_MIDAS.py`
+
+Pass `-useTorchIndexer 1` to switch the indexing stage from C `IndexerOMP` /
+`IndexerGPU` to this package:
+
+```bash
+python ff_MIDAS.py -paramFN paramstest.txt … -useTorchIndexer 1
+```
+
+## Architecture
+
+`midas-index` is a thin orchestration layer. Heavy lifting is delegated to:
+
+- [`midas-diffract`](../midas_diffract/) — forward simulation (HKL -> theoretical spots).
+- [`midas-stress`](../midas_stress/) — orientation conversions, symmetry, fundamental zone.
+
+This package itself owns: seed enumeration, orientation / position grid layout,
+binned matching, scoring, I/O, and the CLI / library API.
+
+## Benchmark
+
+A bundled benchmark drives the full per-seed pipeline end-to-end:
+
+```bash
+python -m midas_index.benchmarks.bench_seed --n-grains 5 --n-iter 3
+```
+
+## License
+
+BSD-3-Clause. Part of [MIDAS](https://github.com/marinerhemant/MIDAS).

midas_index-0.1.0/README.md

@@ -0,0 +1,72 @@
+# midas-index
+
+Pure-Python/PyTorch FF-HEDM indexer. Drop-in replacement for `IndexerOMP` /
+`IndexerGPU` from MIDAS, with seamless CPU / CUDA / MPS device switching.
+
+**Status:** v0.1.0 in development. Detailed design doc lives in
+`dev/implementation_plan.md` (gitignored).
+
+## Install
+
+```bash
+pip install midas-index
+```
+
+For local development:
+
+```bash
+cd packages/midas_index
+pip install -e .[dev]
+```
+
+## Quick start
+
+```bash
+# CLI — drop-in for IndexerOMP / IndexerGPU
+midas-index paramstest.txt 0 1 1000 8
+
+# Pin device / dtype via env vars (auto-detect: CUDA -> MPS -> CPU)
+MIDAS_INDEX_DEVICE=cuda MIDAS_INDEX_DTYPE=float32 \
+    midas-index paramstest.txt 0 1 1000 8
+```
+
+Library API:
+
+```python
+from midas_index import Indexer
+
+result = Indexer.from_param_file("paramstest.txt", device="cuda").run(
+    block_nr=0, n_blocks=1, n_spots_to_index=1000,
+)
+```
+
+## Drive from `ff_MIDAS.py`
+
+Pass `-useTorchIndexer 1` to switch the indexing stage from C `IndexerOMP` /
+`IndexerGPU` to this package:
+
+```bash
+python ff_MIDAS.py -paramFN paramstest.txt … -useTorchIndexer 1
+```
+
+## Architecture
+
+`midas-index` is a thin orchestration layer. Heavy lifting is delegated to:
+
+- [`midas-diffract`](../midas_diffract/) — forward simulation (HKL -> theoretical spots).
+- [`midas-stress`](../midas_stress/) — orientation conversions, symmetry, fundamental zone.
+
+This package itself owns: seed enumeration, orientation / position grid layout,
+binned matching, scoring, I/O, and the CLI / library API.
+
+## Benchmark
+
+A bundled benchmark drives the full per-seed pipeline end-to-end:
+
+```bash
+python -m midas_index.benchmarks.bench_seed --n-grains 5 --n-iter 3
+```
+
+## License
+
+BSD-3-Clause. Part of [MIDAS](https://github.com/marinerhemant/MIDAS).

midas_index-0.1.0/midas_index/__init__.py

@@ -0,0 +1,32 @@
+"""midas-index: Pure-Python/PyTorch FF-HEDM indexer.
+
+Drop-in replacement for the C binaries `IndexerOMP` and `IndexerGPU` from
+MIDAS, with seamless CPU/CUDA/MPS device switching.
+
+See `dev/implementation_plan.md` (gitignored) for design and roadmap.
+
+Quick start
+-----------
+    # CLI
+    midas-index paramstest.txt 0 1 1000 8
+
+    # Library
+    from midas_index import Indexer, IndexerParams
+    result = Indexer.from_param_file("paramstest.txt", device="cuda").run(
+        block_nr=0, n_blocks=1, n_spots_to_index=1000,
+    )
+"""
+
+__version__ = "0.1.0"
+
+from .params import IndexerParams
+from .result import IndexerResult, SeedResult
+from .indexer import Indexer
+
+__all__ = [
+    "Indexer",
+    "IndexerParams",
+    "IndexerResult",
+    "SeedResult",
+    "__version__",
+]

midas_index-0.1.0/midas_index/__main__.py

@@ -0,0 +1,6 @@
+"""Allow `python -m midas_index ...` to invoke the CLI."""
+
+from .cli import main
+
+if __name__ == "__main__":
+    main()

midas_index-0.1.0/midas_index/benchmarks/bench_seed.py

@@ -0,0 +1,183 @@
+"""Per-seed throughput benchmark for `process_seed`.
+
+Drives the full forward + match pipeline on a synthetic 5-grain dataset
+(generated in-memory via `IndexerForwardAdapter`) and reports seeds/sec.
+
+Useful for validating perf changes (e.g. the (y0,z0) cartesian-batch
+vectorization in pipeline.py).
+
+Usage:
+    python -m midas_index.benchmarks.bench_seed --n-grains 5 --n-iter 3
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+import time
+
+import numpy as np
+import torch
+
+from midas_index import IndexerParams
+from midas_index.compute.forward_adapter import IndexerForwardAdapter
+from midas_index.io import build_bin_index
+from midas_index.pipeline import IndexerContext, process_seed
+
+
+def _toy_params() -> IndexerParams:
+    p = IndexerParams()
+    p.Distance = 1_000_000.0
+    p.Wavelength = 0.172979
+    p.Rsample = 250.0
+    p.Hbeam = 200.0
+    p.px = 200.0
+    p.SpaceGroup = 225
+    p.LatticeConstant = (4.08, 4.08, 4.08, 90.0, 90.0, 90.0)
+    p.StepsizePos = 100.0
+    p.StepsizeOrient = 0.5
+    p.MarginOme = 0.5
+    p.MarginRad = 500.0
+    p.MarginRadial = 500.0
+    p.MarginEta = 500.0
+    p.EtaBinSize = 0.1
+    p.OmeBinSize = 0.1
+    p.ExcludePoleAngle = 6.0
+    p.MinMatchesToAcceptFrac = 0.1
+    p.RingNumbers = [1, 2, 3, 4]
+    p.RingRadii = {
+        1: 73582.31550724161,
+        2: 85023.04143552633,
+        3: 120567.4304605822,
+        4: 141666.90427076322,
+    }
+    p.OmegaRanges = [(-180.0, 180.0)]
+    p.BoxSizes = [(-2_000_000.0, 2_000_000.0, -2_000_000.0, 2_000_000.0)]
+    p.UseFriedelPairs = 0
+    p.OutputFolder = "."
+    return p
+
+
+def _toy_hkls():
+    a = 4.08
+    one_over_a = 1.0 / a
+    wl = 0.172979
+    hkls = []
+    int_rows = []
+    real_rows = []
+    for h, k, l, ring in [
+        (1, -1, -1, 1), (1, 1, 1, 1), (-1, 1, 1, 1), (1, -1, 1, 1),
+        (-1, -1, 1, 1), (-1, 1, -1, 1), (1, 1, -1, 1), (-1, -1, -1, 1),
+        (2, 0, 0, 2), (-2, 0, 0, 2), (0, 2, 0, 2), (0, -2, 0, 2),
+        (0, 0, 2, 2), (0, 0, -2, 2),
+    ]:
+        g = (h * one_over_a, k * one_over_a, l * one_over_a)
+        d = 1.0 / math.sqrt(sum(x * x for x in g))
+        sin_th = wl / (2.0 * d)
+        if not 0 < sin_th <= 1.0:
+            continue
+        th_rad = math.asin(sin_th)
+        radius = 73582.31550724161 if ring == 1 else 85023.04143552633
+        real_rows.append((g[0], g[1], g[2], float(ring), d, th_rad, radius))
+        int_rows.append((h, k, l, ring))
+    return (
+        np.asarray(real_rows, dtype=np.float64),
+        np.asarray(int_rows, dtype=np.int64),
+    )
+
+
+def _build_dataset(n_grains: int, seed: int):
+    rng = np.random.default_rng(seed)
+    q = rng.normal(size=(n_grains, 4))
+    q /= np.linalg.norm(q, axis=1, keepdims=True)
+    w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
+    R = np.empty((n_grains, 3, 3))
+    R[:, 0, 0] = 1 - 2 * (y * y + z * z)
+    R[:, 0, 1] = 2 * (x * y - w * z)
+    R[:, 0, 2] = 2 * (x * z + w * y)
+    R[:, 1, 0] = 2 * (x * y + w * z)
+    R[:, 1, 1] = 1 - 2 * (x * x + z * z)
+    R[:, 1, 2] = 2 * (y * z - w * x)
+    R[:, 2, 0] = 2 * (x * z - w * y)
+    R[:, 2, 1] = 2 * (y * z + w * x)
+    R[:, 2, 2] = 1 - 2 * (x * x + y * y)
+    return torch.as_tensor(R, dtype=torch.float64)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n-grains", type=int, default=5)
+    parser.add_argument("--n-iter", type=int, default=3,
+                        help="Iterations of process_seed to time")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    params = _toy_params()
+    hkls_real_np, hkls_int_np = _toy_hkls()
+    R = _build_dataset(args.n_grains, args.seed)
+
+    # Forward-sim obs via the adapter
+    adapter = IndexerForwardAdapter(
+        params=params,
+        hkls_real=torch.as_tensor(hkls_real_np, dtype=torch.float64),
+        hkls_int=torch.as_tensor(hkls_int_np, dtype=torch.long),
+        device=torch.device("cpu"),
+        dtype=torch.float64,
+    )
+    pos = torch.zeros(R.shape[0], 3, dtype=torch.float64)
+    theor, valid = adapter.simulate(R, pos)
+
+    rows = []
+    spot_id = 1
+    for g in range(theor.shape[0]):
+        for k in range(theor.shape[1]):
+            if not bool(valid[g, k]):
+                continue
+            y = float(theor[g, k, 10])
+            z = float(theor[g, k, 11])
+            rows.append([
+                y, z, float(theor[g, k, 6]), math.sqrt(y * y + z * z),
+                float(spot_id), float(theor[g, k, 9]),
+                float(theor[g, k, 7]), float(theor[g, k, 8]) * 2.0,
+                float(theor[g, k, 13]),
+            ])
+            spot_id += 1
+    obs = np.asarray(rows, dtype=np.float64)
+    bin_data, bin_ndata = build_bin_index(
+        obs, eta_bin_size=0.1, ome_bin_size=0.1, n_rings=4,
+        margin_eta=params.MarginEta, margin_ome=params.MarginOme,
+        stepsize_orient=params.StepsizeOrient,
+        ring_radii=params.RingRadii,
+    )
+
+    ctx = IndexerContext(
+        params=params, hkls_real=hkls_real_np, hkls_int=hkls_int_np,
+        obs=obs, bin_data=bin_data, bin_ndata=bin_ndata,
+        device=torch.device("cpu"), dtype=torch.float64,
+    )
+
+    spot_ids = obs[:, 4].astype(int)
+    print(f"benchmark: {args.n_grains} grains -> {len(obs)} obs spots")
+    print(f"           {len(spot_ids)} candidate seed spots")
+
+    # Warm up
+    process_seed(int(spot_ids[0]), ctx)
+
+    # Time process_seed across a sample
+    sample = spot_ids[: min(args.n_grains, len(spot_ids))]
+    timings = []
+    for _ in range(args.n_iter):
+        t0 = time.perf_counter()
+        for sid in sample:
+            process_seed(int(sid), ctx)
+        timings.append(time.perf_counter() - t0)
+    n_seeds = len(sample)
+    best = min(timings)
+    print(f"\n  best of {args.n_iter}: {best:.3f}s for {n_seeds} seeds")
+    print(f"  -> {best / n_seeds * 1000:.1f} ms/seed, {n_seeds / best:.1f} seeds/sec")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

midas_index-0.1.0/midas_index/cli.py

@@ -0,0 +1,99 @@
+"""CLI entry point for midas-index.
+
+Drop-in replacement for the IndexerOMP / IndexerGPU positional argv:
+
+    midas-index <param_file> <block_nr> <n_blocks> <n_spots_to_index> <num_procs>
+
+Optional flags (extension over C binaries):
+
+    --device {cpu,cuda,mps}       override auto-detection (env: MIDAS_INDEX_DEVICE)
+    --dtype  {float32,float64}    override per-device default (env: MIDAS_INDEX_DTYPE)
+    --version                     print version and exit
+
+Behaviour mirrors `FF_HEDM/src/IndexerOMP.c::main` exactly:
+  - Reads <param_file> for all 33 paramstest.txt keys.
+  - Computes [startRowNr, endRowNr] from (block_nr, n_blocks, n_spots_to_index).
+  - Writes BestPos_<block_nr>.csv plus consolidated binaries to OutputFolder.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import __version__
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="midas-index",
+        description="Pure-Python/PyTorch FF-HEDM indexer (drop-in for IndexerOMP/IndexerGPU).",
+    )
+    parser.add_argument("param_file", help="Path to paramstest.txt")
+    parser.add_argument("block_nr", type=int, help="Block index for sharded run")
+    parser.add_argument("n_blocks", type=int, help="Total number of blocks")
+    parser.add_argument(
+        "n_spots_to_index", type=int, help="Total number of seed spots to process"
+    )
+    parser.add_argument(
+        "num_procs",
+        type=int,
+        help="Threads on CPU (passed to torch.set_num_threads); ignored on GPU/MPS",
+    )
+    parser.add_argument("--device", choices=["cpu", "cuda", "mps"], default=None)
+    parser.add_argument("--dtype", choices=["float32", "float64"], default=None)
+    parser.add_argument("--version", action="version", version=f"midas-index {__version__}")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point. Returns process exit code."""
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    from .indexer import Indexer
+    from .io.output import close_output_files, open_output_files, write_seed_record
+
+    indexer = Indexer.from_param_file(
+        args.param_file, device=args.device, dtype=args.dtype,
+    )
+    indexer.load_observations()
+    obs = indexer._observations
+    assert obs is not None
+
+    spot_ids = obs["spot_ids"]
+    n_total = int(min(args.n_spots_to_index, len(spot_ids)))
+
+    result = indexer.run(
+        block_nr=args.block_nr,
+        n_blocks=args.n_blocks,
+        n_spots_to_index=n_total,
+        num_procs=args.num_procs,
+    )
+
+    # Build a spot_id -> offset map (offset is the row in SpotsToIndex.csv).
+    sid_to_offset: dict[int, int] = {
+        int(sid): i for i, sid in enumerate(spot_ids[:n_total].tolist())
+    }
+
+    output_folder = indexer.params.OutputFolder
+    fd_best, fd_full = open_output_files(output_folder, n_total, args.block_nr)
+    try:
+        for seed in result.seeds:
+            offset = sid_to_offset.get(int(seed.spot_id), -1)
+            if offset < 0:
+                continue
+            write_seed_record(fd_best, seed, offset)
+    finally:
+        close_output_files(fd_best, fd_full)
+
+    print(
+        f"midas-index {__version__}: block {args.block_nr}/{args.n_blocks} "
+        f"completed. {len(result.seeds)} seeds indexed -> {output_folder}/IndexBest.bin",
+        file=sys.stderr,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

midas_index-0.1.0/midas_index/compute/__init__.py

@@ -0,0 +1,32 @@
+"""Compute kernels for midas-index.
+
+Most heavy lifting is delegated:
+  - Forward simulation         -> midas-diffract.HEDMForwardModel (forward_adapter.py)
+  - Orientation conversions    -> midas-stress.orientation         (rotation.py shim)
+
+Owned here: seed enumeration, orientation/position grids, binned matching, scoring.
+"""
+
+from . import (
+    binning,
+    constants,
+    forward_adapter,
+    matching,
+    orientation_grid,
+    position_grid,
+    reduce,
+    rotation,
+    seeds,
+)
+
+__all__ = [
+    "binning",
+    "constants",
+    "forward_adapter",
+    "matching",
+    "orientation_grid",
+    "position_grid",
+    "reduce",
+    "rotation",
+    "seeds",
+]

midas_index-0.1.0/midas_index/compute/binning.py

@@ -0,0 +1,48 @@
+"""Binned-spot index lookups.
+
+Mirrors `GetBin` from `FF_HEDM/src/IndexerOMP.c:115`. Given (ring_nr, eta,
+omega) of a theoretical spot, computes a flat bin index and gathers the
+candidate observed-spot row IDs from `data`/`ndata` arrays.
+
+Vectorized: works on arbitrary leading-dim tensors of theoretical spots and
+returns (n_in_bin, data_offset) per element. The actual gather of candidate
+rows is left to `compare_spots` (it depends on the matching strategy —
+dense vs. jagged).
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def get_bin_indices(
+    ring_nr: torch.Tensor,        # int (..)
+    eta_deg: torch.Tensor,        # float (..)
+    omega_deg: torch.Tensor,      # float (..)
+    eta_bin_size: float,
+    ome_bin_size: float,
+    n_eta_bins: int,
+    n_ome_bins: int,
+) -> torch.Tensor:
+    """Compute flat bin indices `pos = (ring-1) * (n_eta * n_ome) + iEta * n_ome + iOme`.
+
+    Mirrors C: `iEta = floor((180 + eta) / EtaBinSize)`, `iOme` likewise.
+    """
+    i_ring = (ring_nr.to(torch.int64) - 1)
+    i_eta = torch.floor((180.0 + eta_deg) / eta_bin_size).to(torch.int64)
+    i_ome = torch.floor((180.0 + omega_deg) / ome_bin_size).to(torch.int64)
+    return i_ring * (n_eta_bins * n_ome_bins) + i_eta * n_ome_bins + i_ome
+
+
+def lookup_bin_counts(
+    pos: torch.Tensor,            # int64 (..,)
+    ndata: torch.Tensor,          # int32 (2*n_bins,)
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return (n_in_bin, data_offset) for each pos.
+
+    `ndata` layout is interleaved: `[count_0, offset_0, count_1, offset_1, ...]`.
+    """
+    pos = pos.to(torch.int64)
+    n_per = ndata[pos * 2].to(torch.int64)
+    offset = ndata[pos * 2 + 1].to(torch.int64)
+    return n_per, offset

midas_index-0.1.0/midas_index/compute/constants.py

@@ -0,0 +1,30 @@
+"""Compile-time constants from `FF_HEDM/src/IndexerOMP.c` / `IndexerGPU.cu`.
+
+In Python these are advisory caps + bookkeeping (column counts).
+Buffer sizes are determined dynamically at runtime, not from these.
+"""
+
+# --- Bookkeeping caps (matches C; advisory only) ---
+MAX_N_SPOTS = 100_000_000
+MAX_N_RINGS = 500
+MAX_N_HKLS = 5000
+MAX_N_STEPS = 2000
+MAX_N_OMEGARANGES = 2000
+MAX_N_OR_CPU = 36_000     # IndexerOMP.c
+MAX_N_OR_GPU = 7_200      # IndexerGPU.cu
+
+# --- Column counts of legacy flat layouts ---
+N_COL_THEORSPOTS = 14   # TheorSpots row layout
+N_COL_OBSSPOTS = 9      # Spots.bin row layout
+N_COL_GRAINSPOTS = 17
+N_COL_GRAINMATCHES = 16
+
+# --- Numerical ---
+EPS_F32 = 1e-9
+EPS_F64 = 1e-12
+
+# --- Conversions ---
+import math
+
+DEG2RAD = math.pi / 180.0
+RAD2DEG = 180.0 / math.pi