probixi 0.1.0__py3-none-any.whl

probixi/__init__.py

@@ -0,0 +1 @@
+from .probixi import Probixi as Probixi

probixi/cli.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from typing import Literal, Optional, cast
+
+import click
+import torch
+
+from probixi.io import DataOffloader, PeakOffloader
+from probixi.probixi import Probixi
+
+
+@click.command()
+@click.option(
+    "-i",
+    "--input",
+    "list_file",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="CrystFEL list file (.lst) of HDF5 inputs.",
+)
+@click.option(
+    "-g",
+    "--geometry",
+    "geometry_file",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="CrystFEL geometry file (.geom).",
+)
+@click.option(
+    "-p",
+    "--cell",
+    "cell_file",
+    default=None,
+    type=click.Path(exists=True, dir_okay=False),
+    help="CrystFEL unit-cell file (.cell). Required unless --peaks-only.",
+)
+@click.option(
+    "-o",
+    "--output",
+    "output",
+    required=True,
+    type=click.Path(dir_okay=False, writable=True),
+    help="Output CrystFEL-style .stream file.",
+)
+@click.option(
+    "--peaks-only",
+    is_flag=True,
+    help="Only run peak finding and export a peaks-only stream "
+    "Does not require a unit cell.",
+)
+@click.option(
+    "--gif",
+    "gif",
+    default=None,
+    type=click.Path(dir_okay=False, writable=True),
+    help="Also write a noise-model diagnostic GIF over the seed frames.",
+)
+@click.option("--start", type=int, default=None, help="First frame index (inclusive).")
+@click.option("--stop", type=int, default=None, help="Stop frame index (exclusive).")
+@click.option(
+    "--batch-size",
+    type=int,
+    default=8,
+    show_default=True,
+    help="Frames per batched refinement pass.",
+)
+@click.option("--device", default=None, help="Torch device (default: auto).")
+@click.option(
+    "--noise-mode",
+    type=click.Choice(["online", "per_frame"]),
+    default="online",
+    show_default=True,
+    help="Noise-model update mode.",
+)
+@click.option("--warmup-frames", type=int, default=16, show_default=True)
+@click.option(
+    "--seed-frames",
+    type=int,
+    default=32,
+    show_default=True,
+    help="Frames used to calibrate the noise model and detection threshold.",
+)
+@click.option(
+    "--target-noise-peaks",
+    type=float,
+    default=5.0,
+    show_default=True,
+    help="Calibrate the detection threshold so a signal-free frame "
+    "yields at most this many noise blobs.",
+)
+@click.option(
+    "--panel",
+    default="0",
+    show_default=True,
+    help="Fallback panel name for peaks outside all geometry panels.",
+)
+@click.option("-q", "--quiet", is_flag=True, help="Suppress per-frame progress.")
+def main(
+    list_file: str,
+    geometry_file: str,
+    cell_file: Optional[str],
+    output: str,
+    peaks_only: bool,
+    gif: Optional[str],
+    start: Optional[int],
+    stop: Optional[int],
+    batch_size: int,
+    device: Optional[str],
+    noise_mode: str,
+    warmup_frames: int,
+    seed_frames: int,
+    target_noise_peaks: float,
+    panel: str,
+    quiet: bool,
+) -> None:
+    """Run the probixi pipeline and write indexed frames to a CrystFEL stream.
+
+    With --peaks-only, peak finding runs but indexing does not, and a
+    peaks-only stream is written instead.
+    """
+    if not peaks_only and cell_file is None:
+        raise click.UsageError(
+            "a unit cell (-p/--cell) is required unless --peaks-only"
+        )
+    dev = torch.device(device) if device else None
+    probixi = Probixi(
+        list_file=list_file,
+        geometry_file=geometry_file,
+        cell_file=cell_file,
+        noise_mode=cast("Literal['online', 'per_frame']", noise_mode),
+        warmup_frames=warmup_frames,
+        device=dev,
+    )
+
+    meta = probixi.metadata
+    if not quiet:
+        click.echo(f"Loaded {meta.n_frames} frames from {meta.n_files} file(s).")
+
+    if gif:
+        probixi.noise_diagnostics(
+            gif, stop=seed_frames, batch_size=max(2, seed_frames // 8)
+        )
+        if not quiet:
+            click.echo(f"Wrote noise diagnostic GIF to {gif}")
+
+    cal = probixi.calibrate(n_seed=seed_frames, target_noise_peaks=target_noise_peaks)
+    if not quiet:
+        tc = probixi.threshold_calibration
+        msg = (
+            f"Calibrated on {seed_frames} frames: kappa={cal.kappa:.2f} "
+            f"prior_peak={cal.prior_peak:.4f} var_scale={cal.var_scale:.3f}"
+        )
+        if tc is not None:
+            msg += f" mf_threshold={tc.threshold:.2f}"
+        click.echo(msg)
+
+    frames = probixi.frames(start=start, stop=stop)
+
+    if peaks_only:
+        peaks = probixi.peak_stream(
+            frames, start_index=start or 0, estimate_scale=False
+        )
+        with PeakOffloader(
+            output,
+            geometry=probixi.geometry,
+            geometry_file=geometry_file,
+            files=meta.files,
+            panel=panel,
+        ) as off:
+            n = 0
+            for result in peaks:
+                if len(result) == 0:
+                    continue
+                off.write(result)
+                n += 1
+                if not quiet:
+                    click.echo(f"  frame {result.frame_index}: {len(result)} peaks")
+        click.echo(f"Wrote peaks for {n} frame(s) to {output}")
+        return
+
+    stream = probixi.index_stream(frames, batch_size=batch_size, start_index=start or 0)
+
+    with DataOffloader(
+        output,
+        geometry=probixi.geometry,
+        cell=probixi.target_cell,
+        geometry_file=geometry_file,
+        files=meta.files,
+        panel=panel,
+    ) as off:
+        n = 0
+        for result in stream:
+            off.write(result)
+            n += 1
+            if not quiet:
+                click.echo(
+                    f"  frame {result.frame_index}: "
+                    f"{result.n_indexed}/{result.n_peaks} peaks indexed "
+                    f"(rmsd {result.rmsd:.4f})"
+                )
+
+    click.echo(f"Wrote {n} indexed frame(s) to {output}")
+
+
+if __name__ == "__main__":
+    main()

probixi/indexer/__init__.py

@@ -0,0 +1,17 @@
+from .indexer import (
+    CellMatchConfig,
+    Indexer,
+    IndexStream,
+    IntegrateConfig,
+    RefineConfig,
+    SeedConfig,
+)
+
+__all__ = [
+    "Indexer",
+    "IndexStream",
+    "CellMatchConfig",
+    "IntegrateConfig",
+    "RefineConfig",
+    "SeedConfig",
+]

probixi/indexer/forward.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def _geometry_constants(
+    geometry: dict,
+    device: Optional[torch.device] = None,
+    dtype: torch.dtype = torch.float64,
+) -> dict:
+    bc = geometry["beam_center"]
+    return {
+        "bc_row": torch.tensor(float(bc[0]), dtype=dtype, device=device),
+        "bc_col": torch.tensor(float(bc[1]), dtype=dtype, device=device),
+        "clen_A": torch.tensor(
+            float(geometry["clen"]) * 1e10, dtype=dtype, device=device
+        ),
+        "pix_A": torch.tensor(
+            float(geometry["pixel_size"]) * 1e10, dtype=dtype, device=device
+        ),
+        "wavelength_A": torch.tensor(
+            float(geometry["wavelength"]), dtype=dtype, device=device
+        ),
+    }
+
+
+def detector_to_q(
+    positions: Tensor,
+    geometry: dict,
+    frame_rotation: Optional[Tensor] = None,
+    dtype: torch.dtype = torch.float64,
+) -> Tensor:
+    """Map detector pixel positions to reciprocal-space q-vectors.
+
+    Elastic (Ewald) construction with the beam along +z: a pixel at lab position
+    ``(x, y, clen)`` defines a scattered unit direction ``s_hat`` and
+    ``q = (s_hat - s0_hat) / lambda`` with ``s0_hat = (0, 0, 1)``.
+
+    Parameters
+    ----------
+    positions : Tensor
+        (N, 2) detector ``(row, col)`` pixel positions.
+    geometry : dict
+        Geometry dict (``beam_center``, ``clen``, ``pixel_size``, ``wavelength``).
+    frame_rotation : Tensor, optional
+        (3, 3) rotation taking q from lab into the crystal frame.
+    dtype : torch.dtype, optional
+        Working dtype (default ``torch.float64``).
+
+    Returns
+    -------
+    Tensor
+        (N, 3) reciprocal-space q-vectors (A^-1).
+    """
+    if positions.ndim != 2 or positions.shape[-1] != 2:
+        raise ValueError("positions must be (N, 2)")
+    device = positions.device
+    g = _geometry_constants(geometry, device=device, dtype=dtype)
+    pos = positions.to(device=device, dtype=dtype)
+    rows, cols = pos[:, 0], pos[:, 1]
+    # pixel offset from beam center translation to lab-frame coordinates \AA
+    x_lab = (cols - g["bc_col"]) * g["pix_A"]
+    y_lab = (rows - g["bc_row"]) * g["pix_A"]
+    z_lab = g["clen_A"].expand_as(x_lab)
+    r = torch.sqrt(x_lab * x_lab + y_lab * y_lab + z_lab * z_lab)
+    inv_lambda = 1.0 / g["wavelength_A"]
+    qx = (x_lab / r) * inv_lambda
+    qy = (y_lab / r) * inv_lambda
+    qz = (z_lab / r - 1.0) * inv_lambda
+    q_lab = torch.stack([qx, qy, qz], dim=-1)
+    if frame_rotation is None:
+        return q_lab
+    R = frame_rotation.to(device=device, dtype=dtype)
+    if R.shape != (3, 3):
+        raise ValueError("frame_rotation must be (3, 3)")
+    return q_lab @ R
+
+
+def q_to_detector(
+    q: Tensor,
+    geometry: dict,
+    frame_rotation: Optional[Tensor] = None,
+    dtype: torch.dtype = torch.float64,
+) -> Tensor:
+    # S = q*lambda + s0
+    # Intersect the ray with z = clen.
+    if q.ndim != 2 or q.shape[-1] != 3:
+        raise ValueError("q must be (N, 3)")
+    device = q.device
+    g = _geometry_constants(geometry, device=device, dtype=dtype)
+    qv = q.to(device=device, dtype=dtype)
+    if frame_rotation is not None:
+        R = frame_rotation.to(device=device, dtype=dtype)
+        if R.shape != (3, 3):
+            raise ValueError("frame_rotation must be (3, 3)")
+        qv = qv @ R.transpose(-1, -2)
+    inv_lambda = 1.0 / g["wavelength_A"]
+    Sx = qv[:, 0] / inv_lambda
+    Sy = qv[:, 1] / inv_lambda
+    Sz = qv[:, 2] / inv_lambda + 1.0
+    scale = g["clen_A"] / Sz.clamp_min(1e-30)
+    col = (Sx * scale) / g["pix_A"] + g["bc_col"]
+    row = (Sy * scale) / g["pix_A"] + g["bc_row"]
+    return torch.stack([row, col], dim=-1)