feat: port normalization to NeuNorm 2.0 (route a, #412)

conda.recipe/meta.yaml

@@ -23,7 +23,6 @@ build:
     - {{ PYTHON }} -m pip install . --no-deps --ignore-installed -vvv
     - {{ PYTHON }} -m pip install git+https://github.com/ornlneutronimaging/BraggEdge.git@2.0.6 --no-deps --no-build-isolation -vvv
     - {{ PYTHON }} -m pip install git+https://github.com/ruipgil/changepy.git@bf53e3a4af182b38b9c08fc56a0c6e68c2d33bda --no-deps --no-build-isolation -vvv
-    - {{ PYTHON }} -m pip install git+https://github.com/ornlneutronimaging/NeuNorm.git@v1.6.12 --no-deps --no-build-isolation -vvv
 
 requirements:
 
@@ -57,7 +56,10 @@ requirements:
     - tqdm
     - tomli
     - pydantic
-    # - neutronimaging::neunorm  # switch to pip install until we sorted out NeuNorm conda package version issue
+    - scipp
+    # 2.0 publishes a clean conda package to the neutronimaging channel,
+    # so the pip-vendoring workaround used for 1.x is gone (#412)
+    - neunorm >=2.0.0,<3
 
 about:
   home: {{ url }}

environment.yml

@@ -22,8 +22,9 @@ dependencies:
   # compute
   - astropy
   - lmfit
-  # NeuNorm 2.0 (on the neutronimaging channel) is a breaking rewrite — see #412
-  - neunorm>=1.6.12,<2
+  # NeuNorm 2.0 port landed with #412 — scipp comes with it
+  - neunorm>=2.0.0,<3
+  - scipp
   # storage
   - h5py
   - sqlite

pyproject.toml

@@ -42,10 +42,12 @@ dependencies = [
   "QtPy",
   "PyQt5",
   "pyqtgraph",
+  # scipp carries the DataArrays the NeuNorm 2.0 normalization runs on
+  "scipp",
   # These are the dependencies from our neutron imaging team
   "neutronbraggedge>=2.0.6,<3",
   "changepy>=0.4.0,<0.5",
-  "neunorm>=1.6.12,<2",
+  "neunorm>=2.0.0,<3",
 ]
 
 [project.urls]
@@ -136,6 +138,7 @@ loguru = "*"
 tqdm = "*"
 tomli = "*"
 pydantic = "*"
+scipp = "*"
 pyqtgraph = "*"
 qt = "*"
 qtpy = "*"

src/ibeatles/about/about_launcher.py

@@ -43,25 +43,11 @@ def __init__(self, parent=None):
             list_version.append(f"pandas: {pandas.__version__}")
 
         try:
-            import NeuNorm
+            import neunorm
         except ImportError:
-            list_version.append("NeuNorm: unknown")
+            list_version.append("neunorm: unknown")
         else:
-            list_version.append(f"NeuNorm: {NeuNorm.__version__}")
-
-        try:
-            import qtpy
-        except ImportError:
-            list_version.append("qtpy: unknown")
-        else:
-            list_version.append(f"qtpy: {qtpy.__version__}")
-
-        try:
-            import NeuNorm
-        except ImportError:
-            list_version.append("NeuNorm: unknown")
-        else:
-            list_version.append(f"NeuNorm: {NeuNorm.__version__}")
+            list_version.append(f"neunorm: {neunorm.__version__}")
 
         try:
             import qtpy

src/ibeatles/core/io/tiff.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+"""Float32 TIFF writing for analysis exports.
+
+This is the on-disk format every iBeatles export site produced through
+NeuNorm 1.x (PIL ``Image.fromarray`` on float32 frames): plain
+single-page TIFFs that step3's folder re-import and downstream tools
+read back. Keep the writer and dtype stable unless those readers move
+with it.
+"""
+
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+from PIL import Image
+
+
+def write_float32_tiff(data: np.ndarray, file_path: Union[str, Path]) -> None:
+    """Write a 2D array as a single-page float32 TIFF.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        2D image to write. Cast to float32, matching the NeuNorm 1.x
+        writer this replaces.
+    file_path : str or Path
+        Full output path, extension included.
+    """
+    frame = np.asarray(data, dtype=np.float32)
+    if frame.ndim != 2:
+        raise ValueError(f"Expected a 2D image, got shape {frame.shape}")
+    Image.fromarray(frame).save(str(file_path))

src/ibeatles/core/processing/normalization.py

@@ -1,18 +1,44 @@
 #!/usr/bin/env python
-"""Normalization functions for the TOF imaging data."""
+"""Normalization functions for the TOF imaging data.
+
+NeuNorm 2.0 port (route (a), issue #412): the transmission division runs
+through ``neunorm.processing.normalizer.normalize_transmission`` on scipp
+DataArrays. The 1.x behaviors that 2.0 dropped are bridged locally and
+pinned by tests/unit/ibeatles/core/processing/test_normalization_contract.py:
+
+- in-memory stacks (2.0 loaders are path-based), float32 working dtype
+- pooled multi-ROI background ("air") correction in the 1.x
+  ratio-of-means form, with INCLUSIVE extents — ROI (x0, y0, w, h)
+  covers (w+1) x (h+1) pixels
+- per-frame 1:1 OB pairing when stack counts match, nanmedian(OB)
+  fallback when they don't
+- zero-OB pixels (NaN/inf after division) zeroed in the output; the
+  sample-only path does NOT zero, matching 1.x
+- per-frame ``normalized_image_NNNN.tif`` export (1-based), the on-disk
+  layout the GUI's step3 folder re-import depends on
+
+No variances are attached to the DataArrays: iBeatles stacks may be
+pre-smoothed (moving average) or contain negative pixels, so Poisson
+(counts == variance) statistics would be wrong.
+"""
 
 import logging
 import shutil
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import numpy as np
-from NeuNorm.normalization import Normalization as NeuNormNormalization
-from NeuNorm.roi import ROI
+import scipp as sc
+from neunorm.processing.normalizer import normalize_transmission
 
 from ibeatles.core.config import IBeatlesUserConfig
+from ibeatles.core.io.tiff import write_float32_tiff
 from ibeatles.core.processing.moving_average import moving_average
 
+# (x0, y0, width, height) with 1.x inclusive extents: the region spans
+# columns x0..x0+width and rows y0..y0+height, boundaries included.
+BackgroundROI = Tuple[int, int, int, int]
+
 
 def normalize_data(
     sample_data: np.ndarray,
@@ -44,64 +70,191 @@ def normalize_data(
     """
     logging.info("Starting normalization process")
 
-    # Create NeuNorm object
-    o_norm = NeuNormNormalization()
-    o_norm.load(data=sample_data)
-    if ob_data is not None:
-        o_norm.load(data=ob_data, data_type="ob")
+    sample_stack = np.asarray(sample_data, dtype=np.float32)
+    ob_stack = np.asarray(ob_data, dtype=np.float32) if ob_data is not None else None
 
     # Apply moving average if configured
     if config.normalization.moving_average.active:
-        logging.info("Applying moving average")
         if config.normalization.processing_order == "Moving average, Normalization":
-            o_norm.data["sample"]["data"] = moving_average(
-                np.array(o_norm.data["sample"]["data"]),  # NeuNorm is return a list of arrays
-                config.normalization.moving_average,
-            )
-            if ob_data is not None:
-                o_norm.data["ob"]["data"] = moving_average(
-                    np.array(o_norm.data["ob"]["data"]),
-                    config.normalization.moving_average,
-                )
+            logging.info("Applying moving average")
+            sample_stack = moving_average(sample_stack, config.normalization.moving_average)
+            if ob_stack is not None:
+                ob_stack = moving_average(ob_stack, config.normalization.moving_average)
 
     # Perform normalization
     logging.info("Performing normalization")
-    background_roi = _get_background_roi(config)
-
-    if ob_data is None:
-        o_norm.normalization(roi=background_roi, use_only_sample=True)
-    else:
-        o_norm.normalization(roi=background_roi if background_roi else None)
+    normalized = normalize_stacks(
+        sample_stack,
+        ob_stack=ob_stack,
+        background_rois=_get_background_rois(config),
+    )
 
     # Apply moving average after normalization if configured
     if (
         config.normalization.moving_average.active
         and config.normalization.processing_order == "Normalization, Moving Average"
     ):
         logging.info("Applying moving average after normalization")
-        normalized_data = moving_average(
-            np.array(o_norm.get_normalized_data()),
-            config.normalization.moving_average,
-        )
-        # manual replace the normalized data
-        o_norm.data["normalized"] = normalized_data
+        normalized = moving_average(np.array(normalized), config.normalization.moving_average)
 
     # Export normalized data
     full_output_folder = _create_output_folder(output_folder, config)
-    o_norm.export(folder=full_output_folder)
+    export_normalized_stack(normalized, full_output_folder)
 
     # Move time spectra file
     _copy_time_spectra(time_spectra, full_output_folder)
 
-    return o_norm.get_normalized_data(), full_output_folder
+    return normalized, full_output_folder
+
+
+def normalize_stacks(
+    sample_stack: np.ndarray,
+    ob_stack: Optional[np.ndarray] = None,
+    background_rois: Optional[Sequence[BackgroundROI]] = None,
+) -> np.ndarray:
+    """Normalize an in-memory sample stack, with 1.x NeuNorm semantics.
+
+    This is the single math path shared by the headless pipeline
+    (``normalize_data``) and the GUI (step2). Inputs are cast to float32
+    (the 1.x working dtype), the division runs through NeuNorm 2.0 on
+    float64 scipp DataArrays, and the result is cast back to float32.
+
+    Parameters
+    ----------
+    sample_stack : np.ndarray
+        Sample images, shape (n_frames, height, width).
+    ob_stack : np.ndarray, optional
+        Open beam images, shape (n_ob, height, width). Frame shape must
+        match the sample. When n_ob != n_frames, the nanmedian of the
+        (ROI-corrected) OB stack is used as a single open beam.
+    background_rois : sequence of (x0, y0, width, height), optional
+        Background regions with inclusive 1.x extents. With an OB, each
+        sample and OB frame is divided by its own pooled mean over all
+        ROIs before the transmission division. Without an OB, at least
+        one ROI is required and each sample frame is divided by its
+        pooled ROI mean.
+
+    Returns
+    -------
+    np.ndarray
+        float32 stack, shape (n_frames, height, width).
+
+    Raises
+    ------
+    ValueError
+        If no OB and no background ROI is provided, if a ROI does not
+        fit inside the image, or if sample and OB frame shapes differ.
+    """
+    sample_stack = _as_float32_stack(sample_stack)
+    rois = [tuple(int(value) for value in roi) for roi in (background_rois or [])]
+    _validate_rois(rois, frame_shape=sample_stack.shape[1:])
+
+    sample = _make_data_array(sample_stack)
+
+    if ob_stack is None:
+        if not rois:
+            raise ValueError("You need to provide at least 1 background ROI to normalize without open beam data!")
+        corrected = sample / _pooled_roi_means(sample, rois)
+        # 1.x does not zero NaN/inf in the sample-only path
+        return corrected.values.astype(np.float32)
+
+    ob_stack = _as_float32_stack(ob_stack)
+    if ob_stack.shape[1:] != sample_stack.shape[1:]:
+        raise ValueError("Sample and open beam frames do not have the same shape!")
+
+    ob = _make_data_array(ob_stack)
+    if rois:
+        sample = sample / _pooled_roi_means(sample, rois)
+        ob = ob / _pooled_roi_means(ob, rois)
+
+    if ob_stack.shape[0] != sample_stack.shape[0]:
+        # 1.x fallback: collapse the (ROI-corrected) OB stack to its
+        # per-pixel nanmedian and divide every sample frame by it
+        ob = _collapse_to_nanmedian(ob)
+
+    transmission = normalize_transmission(sample, ob)
+
+    # 1.x zeroes the NaN/inf that zero-OB pixels produce
+    normalized = np.nan_to_num(transmission.values, nan=0.0, posinf=0.0, neginf=0.0)
+    return normalized.astype(np.float32)
+
+
+def export_normalized_stack(normalized: np.ndarray, folder: str) -> List[str]:
+    """Export one float32 TIFF per frame, in the 1.x on-disk layout.
+
+    Files are named ``normalized_image_NNNN.tif`` with a 1-based index —
+    the layout step3's folder re-import and the existing tests expect.
+
+    Returns the list of file paths written.
+    """
+    paths = []
+    for index, frame in enumerate(normalized, start=1):
+        path = Path(folder) / f"normalized_image_{index:04d}.tif"
+        write_float32_tiff(frame, path)
+        paths.append(str(path))
+    return paths
+
+
+def _as_float32_stack(data: np.ndarray) -> np.ndarray:
+    """Cast to the 1.x float32 working dtype and promote 2D to a 1-frame stack."""
+    stack = np.asarray(data, dtype=np.float32)
+    if stack.ndim == 2:
+        stack = stack[np.newaxis, :, :]
+    if stack.ndim != 3:
+        raise ValueError(f"Expected a 2D image or 3D stack, got shape {stack.shape}")
+    return stack
+
+
+def _make_data_array(stack: np.ndarray) -> sc.DataArray:
+    """Build a scipp DataArray from an in-memory stack (no variances)."""
+    values = np.ascontiguousarray(stack, dtype=np.float64)
+    _, n_y, n_x = values.shape
+    return sc.DataArray(
+        data=sc.array(dims=["N_image", "y", "x"], values=values, unit=sc.units.counts),
+        coords={
+            "y": sc.arange("y", n_y, unit=None),
+            "x": sc.arange("x", n_x, unit=None),
+        },
+    )
+
+
+def _pooled_roi_means(images: sc.DataArray, rois: Sequence[BackgroundROI]) -> sc.Variable:
+    """Per-frame pooled background mean, in the 1.x form.
+
+    Total counts over ALL ROIs divided by total ROI pixels, per frame,
+    with inclusive extents (hence the +1 on both slice stops).
+    """
+    total_counts = None
+    total_pixels = 0
+    for x0, y0, width, height in rois:
+        region = images["y", y0 : y0 + height + 1]["x", x0 : x0 + width + 1]
+        roi_counts = sc.sum(region.data, dim=["y", "x"])
+        total_counts = roi_counts if total_counts is None else total_counts + roi_counts
+        total_pixels += (height + 1) * (width + 1)
+    return total_counts / total_pixels
+
+
+def _collapse_to_nanmedian(images: sc.DataArray) -> sc.DataArray:
+    """Collapse a stack to its per-pixel nanmedian (the 1.x mismatched-counts OB)."""
+    median = np.nanmedian(images.values, axis=0)
+    return sc.DataArray(
+        data=sc.array(dims=["y", "x"], values=median, unit=images.unit),
+        coords={"y": images.coords["y"], "x": images.coords["x"]},
+    )
+
+
+def _validate_rois(rois: Sequence[BackgroundROI], frame_shape: Tuple[int, int]) -> None:
+    """Reject ROIs that do not fit, with the 1.x inclusive-bounds rule."""
+    n_y, n_x = frame_shape
+    for x0, y0, width, height in rois:
+        if x0 < 0 or y0 < 0 or x0 + width >= n_x or y0 + height >= n_y:
+            raise ValueError("roi does not fit into sample image!")
 
 
-def _get_background_roi(config: IBeatlesUserConfig) -> Optional[List[ROI]]:
-    """Extract background ROI from configuration."""
+def _get_background_rois(config: IBeatlesUserConfig) -> Optional[List[BackgroundROI]]:
+    """Extract background ROIs from configuration."""
     if config.normalization.sample_background:
-        return [
-            ROI(x0=bg.x0, y0=bg.y0, width=bg.width, height=bg.height) for bg in config.normalization.sample_background
-        ]
+        return [(bg.x0, bg.y0, bg.width, bg.height) for bg in config.normalization.sample_background]
     return None