midas-nf-preprocess 0.1.0__py3-none-any.whl

midas_nf_preprocess/__init__.py

@@ -0,0 +1,45 @@
+"""midas-nf-preprocess: Differentiable PyTorch port of NF-HEDM preprocessing.
+
+Bundles the preprocessing stages that sit between raw experimental inputs
+and the orientation-fitting step (``FitOrientationOMP``):
+
+  - ``hex_grid``         : voxel grid generation (port of MakeHexGrid)
+  - ``tomo_filter``      : grid masking from a tomography image
+                           (port of filterGridfromTomo)
+  - ``seed_orientations``: NF seed-orientation generation (cache, from-scratch,
+                           or from an FF Grains.csv)
+  - ``diffr_spots``      : per-orientation diffraction-spot prediction
+                           (differentiable port of MakeDiffrSpots; wraps
+                           ``midas_diffract.HEDMForwardModel.calc_bragg_geometry``)
+  - ``process_images``   : raw TIFF -> SpotsInfo.bin
+                           (differentiable port of ProcessImagesCombined)
+
+The shared ``device`` module provides device/dtype resolution
+(env vars ``MIDAS_NF_PREPROCESS_DEVICE`` / ``MIDAS_NF_PREPROCESS_DTYPE``)
+matching the convention used by ``midas-transforms``.
+"""
+
+__version__ = "0.1.0"
+
+from . import (
+    device,
+    diffr_spots,
+    hex_grid,
+    process_images,
+    seed_orientations,
+    tomo_filter,
+)
+from .device import resolve_device, resolve_dtype, apply_cpu_threads
+
+__all__ = [
+    "__version__",
+    "device",
+    "diffr_spots",
+    "hex_grid",
+    "process_images",
+    "seed_orientations",
+    "tomo_filter",
+    "resolve_device",
+    "resolve_dtype",
+    "apply_cpu_threads",
+]

midas_nf_preprocess/cli.py

@@ -0,0 +1,62 @@
+"""Umbrella CLI: ``midas-nf-preprocess <subcommand> ...``.
+
+Subcommands:
+
+  - ``hex-grid``         : generate the voxel grid (port of MakeHexGrid)
+  - ``tomo-filter``      : mask the grid by a tomography image / bbox
+  - ``diffr-spots``      : forward-simulate diffraction spots (port of MakeDiffrSpots)
+  - ``process-images``   : raw TIFF -> SpotsInfo.bin (port of ProcessImagesCombined)
+
+Each subcommand mirrors its standalone ``python -m`` invocation:
+
+  python -m midas_nf_preprocess.hex_grid <args>          ==
+  midas-nf-preprocess hex-grid <args>
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import __version__
+from .diffr_spots import cli as diffr_spots_cli
+from .hex_grid import cli as hex_grid_cli
+from .process_images import cli as process_images_cli
+from .seed_orientations import cli as seed_orientations_cli
+from .tomo_filter import cli as tomo_filter_cli
+
+
+_SUBCOMMANDS = {
+    "hex-grid": hex_grid_cli,
+    "tomo-filter": tomo_filter_cli,
+    "diffr-spots": diffr_spots_cli,
+    "process-images": process_images_cli,
+    "seed-orientations": seed_orientations_cli,
+}
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="midas-nf-preprocess",
+        description="Differentiable PyTorch port of NF-HEDM preprocessing.",
+    )
+    parser.add_argument(
+        "--version", action="version", version=f"midas-nf-preprocess {__version__}"
+    )
+    subparsers = parser.add_subparsers(
+        dest="subcommand", required=True, metavar="<subcommand>"
+    )
+    for name, mod in _SUBCOMMANDS.items():
+        sp = subparsers.add_parser(name, help=mod.__doc__.splitlines()[0] if mod.__doc__ else None)
+        mod.add_arguments(sp)
+        sp.set_defaults(_run=mod.run)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    return args._run(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

midas_nf_preprocess/device.py

@@ -0,0 +1,69 @@
+"""Device / dtype resolution.
+
+Same contract as ``midas-transforms``:
+  - device default: cuda > mps > cpu (auto-detect)
+  - dtype default: float64 on cpu (matches the C binaries),
+                   float32 on cuda/mps (matches the GPU codepath idioms).
+  - Both float32 and float64 are supported on every device.
+  - Env vars MIDAS_NF_PREPROCESS_DEVICE / MIDAS_NF_PREPROCESS_DTYPE
+    override defaults.
+  - Library ``device=``/``dtype=`` kwargs override env vars.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional, Union
+
+import torch
+
+_DTYPE_ALIAS = {
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "float64": torch.float64,
+    "fp64": torch.float64,
+    "double": torch.float64,
+    "single": torch.float32,
+}
+
+ENV_DEVICE = "MIDAS_NF_PREPROCESS_DEVICE"
+ENV_DTYPE = "MIDAS_NF_PREPROCESS_DTYPE"
+
+
+def resolve_device(user: Optional[Union[str, torch.device]]) -> torch.device:
+    """Pick a torch.device. Precedence: arg > env var > auto-detect."""
+    if isinstance(user, torch.device):
+        return user
+    if user is None:
+        user = os.environ.get(ENV_DEVICE)
+    if user is None:
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+    return torch.device(user)
+
+
+def resolve_dtype(
+    device: torch.device, user: Optional[Union[str, torch.dtype]]
+) -> torch.dtype:
+    """Pick a torch.dtype. Precedence: arg > env var > per-device default."""
+    if isinstance(user, torch.dtype):
+        return user
+    if user is None:
+        user = os.environ.get(ENV_DTYPE)
+    if user is not None:
+        try:
+            return _DTYPE_ALIAS[user.lower()]
+        except KeyError as e:
+            raise ValueError(
+                f"Unsupported dtype '{user}'. Use one of {sorted(_DTYPE_ALIAS)}."
+            ) from e
+    return torch.float64 if device.type == "cpu" else torch.float32
+
+
+def apply_cpu_threads(num_procs: int, device: torch.device) -> None:
+    """Honor a legacy ``numProcs`` argv on CPU only (no-op on GPU)."""
+    if device.type == "cpu" and num_procs > 0:
+        torch.set_num_threads(num_procs)

midas_nf_preprocess/diffr_spots/__init__.py

@@ -0,0 +1,83 @@
+"""diffr_spots: per-orientation diffraction-spot prediction (port of MakeDiffrSpots).
+
+Differentiable PyTorch port of ``NF_HEDM/src/MakeDiffrSpots.c``. For each seed
+orientation and each HKL, solve the Bragg condition:
+
+  - ``omega(s)`` (1 or 2 solutions): rotation angles where the reciprocal-space
+    G-vector intersects the Ewald sphere
+  - ``eta``: azimuthal angle on the detector
+  - ``(yl, zl)``: lab-frame detector position at the primary distance
+
+then filter by ``OmegaRange``, ``BoxSize``, and ``ExcludePoleAngle``. Output is
+either a structured tensor bundle or the same three binary files the C
+executable produces:
+
+  - ``DiffractionSpots.bin`` : ``[TotalSpots, 3]`` float64 ``(yl, zl, omega)``
+  - ``Key.bin``              : ``[N_orient, 2]`` int32 ``(count_i, offset_i)``
+  - ``OrientMat.bin``        : ``[N_orient, 3, 3]`` float64 row-major
+
+Differentiability:
+
+  - ``omega``, ``eta``, ``yl``, ``zl`` are differentiable in the input
+    quaternions / Euler angles / orientation matrices and in the lattice
+    parameters that produced ``hkls``.
+  - The hard ``OmegaRange`` / ``BoxSize`` / ``ExcludePoleAngle`` filters return
+    a boolean ``valid`` mask (no gradient), matching the convention used by
+    ``midas_diffract.SpotDescriptors``.
+
+The geometry primitives (``calc_omega``, ``calc_eta``, ``calc_spot_position``)
+match the C implementations in ``MakeDiffrSpots.c`` L99-L204 line-by-line and
+extend them to batched, vectorized PyTorch operations.
+"""
+
+from .params import DiffrSpotsParams
+from .orientations import (
+    quat_to_orient_matrix,
+    orient_matrix_to_quat,
+    euler_to_orient_matrix,
+)
+from .geometry import (
+    bragg_omega_eta,
+    calc_eta_deg,
+    calc_spot_position,
+    rotate_around_z,
+)
+from .pipeline import (
+    DiffrSpotsPipeline,
+    DiffrSpotsResult,
+    predict_spots,
+)
+from .binary_io import (
+    write_diffr_spots_bin,
+    write_key_bin,
+    write_orient_mat_bin,
+    read_diffr_spots_bin,
+    read_key_bin,
+    read_orient_mat_bin,
+    write_all,
+)
+from .hkls import HKLEntry, read_hkls_csv, read_seed_orientations
+
+__all__ = [
+    "DiffrSpotsParams",
+    "quat_to_orient_matrix",
+    "orient_matrix_to_quat",
+    "euler_to_orient_matrix",
+    "bragg_omega_eta",
+    "calc_eta_deg",
+    "calc_spot_position",
+    "rotate_around_z",
+    "DiffrSpotsPipeline",
+    "DiffrSpotsResult",
+    "predict_spots",
+    "write_diffr_spots_bin",
+    "write_key_bin",
+    "write_orient_mat_bin",
+    "read_diffr_spots_bin",
+    "read_key_bin",
+    "read_orient_mat_bin",
+    "write_all",
+    "HKLEntry",
+    "read_hkls_csv",
+    "read_seed_orientations",
+]

midas_nf_preprocess/diffr_spots/binary_io.py

@@ -0,0 +1,118 @@
+"""Binary I/O for the three MakeDiffrSpots output files.
+
+File formats (from MakeDiffrSpots.c L595-L626):
+
+  - ``Key.bin``              : ``int32 [N_orient * 2]`` little-endian, with
+                               entries ``(count_i, offset_i)`` for each orientation.
+  - ``DiffractionSpots.bin`` : ``float64 [TotalSpots * 3]`` little-endian, with
+                               row layout ``(yl, zl, omega)``.
+  - ``OrientMat.bin``        : ``float64 [N_orient * 9]`` little-endian, with
+                               row-major flattened 3x3 matrices.
+
+Total spots = sum of per-orientation counts; ``offset_i`` is the cumulative
+prefix sum (excluding ``count_i`` itself) so that orientation i's spots span
+``DiffractionSpots[offset_i : offset_i + count_i]``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import torch
+
+
+# -----------------------------------------------------------------------------
+# Writers
+# -----------------------------------------------------------------------------
+
+
+def write_key_bin(
+    counts: torch.Tensor, path: Union[str, Path]
+) -> None:
+    """Write ``Key.bin`` from a ``(N,)`` int tensor of per-orientation spot counts."""
+    if counts.ndim != 1:
+        raise ValueError(f"Expected 1D counts tensor, got shape {tuple(counts.shape)}")
+    counts_np = counts.detach().cpu().numpy().astype(np.int32)
+    offsets = np.zeros_like(counts_np)
+    if counts_np.size > 0:
+        offsets[1:] = np.cumsum(counts_np[:-1])
+    key = np.stack([counts_np, offsets], axis=1)  # (N, 2)
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    key.tofile(path)
+
+
+def write_diffr_spots_bin(
+    spots: torch.Tensor, path: Union[str, Path]
+) -> None:
+    """Write ``DiffractionSpots.bin`` from a ``(TotalSpots, 3)`` float tensor."""
+    if spots.ndim != 2 or spots.shape[1] != 3:
+        raise ValueError(
+            f"Expected (TotalSpots, 3) tensor, got shape {tuple(spots.shape)}"
+        )
+    arr = spots.detach().cpu().numpy().astype(np.float64)
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr.tofile(path)
+
+
+def write_orient_mat_bin(
+    orient_mats: torch.Tensor, path: Union[str, Path]
+) -> None:
+    """Write ``OrientMat.bin`` from a ``(N, 3, 3)`` float tensor."""
+    if orient_mats.ndim != 3 or orient_mats.shape[1:] != (3, 3):
+        raise ValueError(
+            f"Expected (N, 3, 3) tensor, got shape {tuple(orient_mats.shape)}"
+        )
+    arr = orient_mats.detach().cpu().numpy().astype(np.float64)
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr.tofile(path)
+
+
+def write_all(
+    output_dir: Union[str, Path],
+    counts: torch.Tensor,
+    spots: torch.Tensor,
+    orient_mats: torch.Tensor,
+) -> dict[str, Path]:
+    """Convenience wrapper that writes all three files into ``output_dir``."""
+    output_dir = Path(output_dir)
+    paths = {
+        "Key.bin": output_dir / "Key.bin",
+        "DiffractionSpots.bin": output_dir / "DiffractionSpots.bin",
+        "OrientMat.bin": output_dir / "OrientMat.bin",
+    }
+    write_key_bin(counts, paths["Key.bin"])
+    write_diffr_spots_bin(spots, paths["DiffractionSpots.bin"])
+    write_orient_mat_bin(orient_mats, paths["OrientMat.bin"])
+    return paths
+
+
+# -----------------------------------------------------------------------------
+# Readers (for parity tests / round-trip)
+# -----------------------------------------------------------------------------
+
+
+def read_key_bin(path: Union[str, Path], n_orient: int) -> torch.Tensor:
+    """Read ``Key.bin`` and return a ``(N, 2)`` int64 tensor ``(count_i, offset_i)``."""
+    arr = np.fromfile(path, dtype=np.int32, count=n_orient * 2).reshape(n_orient, 2)
+    return torch.from_numpy(arr.astype(np.int64))
+
+
+def read_diffr_spots_bin(
+    path: Union[str, Path], n_spots: int
+) -> torch.Tensor:
+    """Read ``DiffractionSpots.bin`` and return a ``(N_spots, 3)`` float64 tensor."""
+    arr = np.fromfile(path, dtype=np.float64, count=n_spots * 3).reshape(n_spots, 3)
+    return torch.from_numpy(arr)
+
+
+def read_orient_mat_bin(
+    path: Union[str, Path], n_orient: int
+) -> torch.Tensor:
+    """Read ``OrientMat.bin`` and return a ``(N, 3, 3)`` float64 tensor."""
+    arr = np.fromfile(path, dtype=np.float64, count=n_orient * 9).reshape(n_orient, 3, 3)
+    return torch.from_numpy(arr)

midas_nf_preprocess/diffr_spots/cli.py

@@ -0,0 +1,70 @@
+"""``midas-nf-preprocess diffr-spots`` subcommand."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from .params import DiffrSpotsParams
+from .pipeline import DiffrSpotsPipeline
+
+
+def add_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "parameter_file", help="MIDAS parameter file (Lsd, SeedOrientations, etc.)"
+    )
+    parser.add_argument(
+        "--device", default=None, help="cpu | cuda | mps (default: auto-detect)."
+    )
+    parser.add_argument(
+        "--dtype", default=None, help="fp32 | fp64 (default: fp64 on cpu)."
+    )
+    parser.add_argument(
+        "--hkls-csv",
+        default=None,
+        help="Override hkls.csv path (default: <DataDirectory>/hkls.csv).",
+    )
+    parser.add_argument(
+        "--seeds",
+        default=None,
+        help="Override SeedOrientations path (default: from parameter file).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default=None,
+        help="Override output directory for Key.bin / DiffractionSpots.bin / OrientMat.bin.",
+    )
+
+
+def run(args: argparse.Namespace) -> int:
+    params = DiffrSpotsParams.from_paramfile(args.parameter_file)
+    pipe = DiffrSpotsPipeline(
+        params,
+        device=args.device,
+        dtype=args.dtype,
+        hkls_csv=args.hkls_csv,
+        seed_orientations_csv=args.seeds,
+    )
+    result, paths = pipe.run(output_dir=args.output_dir)
+    total = int(result.counts.sum().item())
+    print(
+        f"Wrote {len(paths)} files to {Path(paths['Key.bin']).parent} "
+        f"({pipe.n_orientations} orientations, {total} total spots)"
+    )
+    for name, p in paths.items():
+        print(f"  {name}: {p}")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="midas-nf-preprocess diffr-spots",
+        description="Differentiable diffraction-spot prediction (port of MakeDiffrSpots).",
+    )
+    add_arguments(parser)
+    return run(parser.parse_args(argv))
+
+
+if __name__ == "__main__":
+    sys.exit(main())

midas_nf_preprocess/diffr_spots/geometry.py

@@ -0,0 +1,193 @@
+"""Bragg geometry for the diffr_spots pipeline.
+
+We **delegate** the omega/eta calculation to
+``midas_diffract.HEDMForwardModel.calc_bragg_geometry`` (forward.py:680-857)
+to avoid maintaining two copies of the same quadratic-Bragg solver. That
+function ports ``CalcDiffractionSpots.c::CalcOmega`` (NF) and
+``ForwardSimulationCompressed.c`` (FF), and adds a wedge correction.
+
+Only the helpers that are MakeDiffrSpots-specific live here:
+
+  - ``calc_spot_position`` : ``(eta, ring_radius) -> (yl, zl)`` lab-frame
+    projection at a given detector distance (matches MakeDiffrSpots.c L115-L120).
+  - ``rotate_around_z``    : convenience helper used by tests.
+
+Returned omega/eta from calc_bragg_geometry are in radians; the rest of this
+package converts to degrees at the boundary so the binary output matches the
+C convention.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+
+_DEG2RAD = math.pi / 180.0
+_RAD2DEG = 180.0 / math.pi
+
+
+def rotate_around_z(v: torch.Tensor, alpha_deg: torch.Tensor) -> torch.Tensor:
+    """Apply a rotation about the +Z axis (active, right-handed).
+
+    Direct port of ``RotateAroundZ`` (MakeDiffrSpots.c L99-L107). Used in tests
+    to validate the geometry; the production path goes through midas_diffract.
+    """
+    if v.shape[-1] != 3:
+        raise ValueError(f"Expected last dim = 3, got shape {tuple(v.shape)}")
+    a = alpha_deg * _DEG2RAD
+    cosa = torch.cos(a)
+    sina = torch.sin(a)
+    x = v[..., 0]
+    y = v[..., 1]
+    z = v[..., 2]
+    return torch.stack(
+        [cosa * x - sina * y, sina * x + cosa * y, z], dim=-1
+    )
+
+
+def calc_eta_deg(y: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+    """Convert lab-frame (y, z) into eta in degrees.
+
+    Port of ``MakeDiffrSpots.c::CalcEtaAngle`` (L109-L113), used as a
+    cross-check against the radians version returned by midas_diffract.
+    """
+    r = torch.sqrt(y * y + z * z)
+    eps = torch.finfo(r.dtype).eps
+    cos_eta = torch.where(r > 0, z / torch.clamp(r, min=eps), torch.zeros_like(z))
+    cos_eta = torch.clamp(cos_eta, -1.0, 1.0)
+    alpha = torch.acos(cos_eta) * _RAD2DEG
+    return torch.where(y > 0, -alpha, alpha)
+
+
+def calc_spot_position(
+    ring_radius: torch.Tensor, eta_deg: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Lab-frame ``(yl, zl)`` from ring radius and eta (degrees).
+
+    Direct port of ``MakeDiffrSpots.c::CalcSpotPosition`` (L115-L120):
+
+        yl = -RingRadius * sin(eta)
+        zl =  RingRadius * cos(eta)
+
+    This is the bit MakeDiffrSpots needs but ``midas_diffract`` does not
+    expose -- the latter goes straight from eta to detector pixel coordinates,
+    skipping the lab-frame intermediate.
+    """
+    eta_rad = eta_deg * _DEG2RAD
+    yl = -ring_radius * torch.sin(eta_rad)
+    zl = ring_radius * torch.cos(eta_rad)
+    return yl, zl
+
+
+# -----------------------------------------------------------------------------
+# Lazy adapter to midas_diffract.HEDMForwardModel for the Bragg geometry call.
+# -----------------------------------------------------------------------------
+
+
+def _make_minimal_geometry(
+    distance_um: float, *, n_pixels: int = 2048, wedge_deg: float = 0.0
+):
+    """Build an HEDMGeometry whose pixel/beam fields are placeholders.
+
+    ``calc_bragg_geometry`` only reads ``self.wedge`` and ``self.epsilon``
+    (the Bragg quadratic itself doesn't touch pixels or beam centers), so any
+    consistent values for the unused fields work.
+    """
+    from midas_diffract.forward import HEDMGeometry
+
+    return HEDMGeometry(
+        Lsd=float(distance_um),
+        y_BC=0.0,
+        z_BC=0.0,
+        px=1.0,
+        omega_start=0.0,
+        omega_step=0.1,
+        n_frames=1,
+        n_pixels_y=n_pixels,
+        n_pixels_z=n_pixels,
+        min_eta=0.0,
+        wavelength=0.0,
+        wedge=float(wedge_deg),
+        flip_y=False,
+    )
+
+
+def bragg_omega_eta(
+    orient_mats: torch.Tensor,
+    hkls: torch.Tensor,
+    thetas_deg: torch.Tensor,
+    *,
+    distance_um: float,
+    wedge_deg: float = 0.0,
+    device: Optional[torch.device] = None,
+):
+    """Compute (omega_deg, eta_deg, valid) by delegating to midas_diffract.
+
+    Wraps ``HEDMForwardModel.calc_bragg_geometry``: the work happens there.
+    Output is converted from radians (midas_diffract) to degrees (the
+    convention used by MakeDiffrSpots and the rest of this package).
+
+    Parameters
+    ----------
+    orient_mats : Tensor of shape ``(N, 3, 3)``.
+    hkls        : Tensor of shape ``(M, 3)``.
+    thetas_deg  : Tensor of shape ``(M,)`` -- Bragg angles in degrees.
+
+    Returns
+    -------
+    omega_deg : Tensor of shape ``(N, 2, M)`` -- two solutions per HKL.
+    eta_deg   : Tensor of shape ``(N, 2, M)``.
+    valid     : Bool Tensor of shape ``(N, 2, M)``.
+    """
+    from midas_diffract.forward import HEDMForwardModel
+
+    if orient_mats.ndim != 3 or orient_mats.shape[-2:] != (3, 3):
+        raise ValueError(
+            f"Expected (N, 3, 3), got shape {tuple(orient_mats.shape)}"
+        )
+    if hkls.shape[-1] != 3:
+        raise ValueError(f"Expected hkls last dim = 3, got {tuple(hkls.shape)}")
+
+    device = device or orient_mats.device
+
+    # Construct a one-shot model. The constructor stores hkls/thetas as
+    # buffers, but calc_bragg_geometry will use the explicit args we pass.
+    geom = _make_minimal_geometry(distance_um, wedge_deg=wedge_deg)
+    model = HEDMForwardModel(
+        hkls=hkls.to(device),
+        thetas=thetas_deg.to(device) * _DEG2RAD,  # midas_diffract expects radians
+        geometry=geom,
+        device=device,
+    )
+
+    omega, eta, _two_theta, valid = model.calc_bragg_geometry(
+        orientation_matrices=orient_mats,
+        hkls_cart=hkls.to(device),
+        thetas=thetas_deg.to(device) * _DEG2RAD,
+    )
+    # midas_diffract layout: (..., 2N, M) -- the leading 2N is a *concatenation*
+    # of the two omega solutions across the N axis. For our (N, 1, 3, 3) input
+    # we get omega shape (1, 2*N, M); with N=N orientations this becomes
+    # (2*N, M) which we reshape into (N, 2, M).
+    # In our use we always pass orient_mats of shape (N, 3, 3) which makes
+    # midas_diffract interpret it as N "voxels" in a 1D batch. The concat
+    # along the K=2N axis stacks "omega_p (N entries)" then "omega_n (N
+    # entries)", so reshape with stride pattern needs care.
+    # Output of calc_bragg_geometry for our input shape:
+    #   omega: (2*N, M)  [no leading batch dim]
+    # The first N rows are omega_p (one per orientation); the next N rows
+    # are omega_n (same orientations, the other quadratic root). Reshape to
+    # (2, N, M) then permute to (N, 2, M).
+    n = orient_mats.shape[0]
+    omega_2nm = omega.reshape(2, n, -1).permute(1, 0, 2).contiguous()
+    eta_2nm = eta.reshape(2, n, -1).permute(1, 0, 2).contiguous()
+    valid_2nm = valid.reshape(2, n, -1).permute(1, 0, 2).bool().contiguous()
+
+    return (
+        omega_2nm * _RAD2DEG,
+        eta_2nm * _RAD2DEG,
+        valid_2nm,
+    )