"""Align ragged per-spectrum arrays onto a single shared x grid.
The ``AutoBK`` and ``Xftf`` tasks produce one (x, y) spectrum per input row.
When the inputs differ slightly (e.g. different ``e0``), Larch returns k/R
grids of different lengths, which cannot be stacked into a single Orange
``Table`` (one shared x-axis per table). This module resamples such ragged
spectra onto one common grid, either by linear interpolation (default) or by
rebinning.
"""
from __future__ import annotations
from enum import Enum
from typing import TYPE_CHECKING
import numpy as np
if TYPE_CHECKING:
from collections.abc import Sequence
from numpy.typing import NDArray
[docs]
class Alignment(str, Enum):
"""How to align spectra of differing length onto a common x grid."""
INTERPOLATION = "interpolation"
REBINNING = "rebinning"
def _common_step(x_list: Sequence[NDArray[np.float64]]) -> float:
"""The shared sampling step: the median spacing of the longest spectrum.
Larch k/R grids share a fixed step, so every spectrum is an aligned subset
of one underlying grid; the longest spectrum is the most reliable estimate.
"""
longest = max(x_list, key=len)
return float(np.median(np.abs(np.diff(longest))))
def _grid(
x_min: float, x_max: float, step: float, *, inclusive: bool
) -> NDArray[np.float64]:
"""Build an evenly spaced grid from ``x_min`` with the given ``step``.
``inclusive=False`` keeps every point at or below ``x_max`` (used for the
interpolation overlap, so no point ever falls outside a spectrum's range).
``inclusive=True`` rounds the endpoint so the grid spans the full union.
"""
span = (x_max - x_min) / step
n = (int(np.round(span)) if inclusive else int(np.floor(span))) + 1
return x_min + step * np.arange(max(n, 1))
[docs]
def interpolate_spectrum(
target_x: NDArray[np.float64],
x: NDArray[np.float64],
y: NDArray[np.float64],
) -> NDArray[np.float64]:
"""Linearly interpolate the spectrum ``(x, y)`` onto ``target_x``.
``x`` need not be sorted: ``np.interp`` requires increasing sample points,
so ``x`` and ``y`` are sorted together first.
"""
order = np.argsort(x)
return np.interp(target_x, x[order], y[order])
def _interpolate(
x_list: Sequence[NDArray[np.float64]],
y_list: Sequence[NDArray[np.float64]],
step: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""Linearly interpolate each spectrum onto the shared overlap grid."""
x_min = max(float(x[0]) for x in x_list)
x_max = min(float(x[-1]) for x in x_list)
common_x = _grid(x_min, x_max, step, inclusive=False)
rows = [
interpolate_spectrum(common_x, x, y)
for x, y in zip(x_list, y_list, strict=True)
]
return common_x, np.asarray(rows, dtype=np.float64)
def _rebin(
x_list: Sequence[NDArray[np.float64]],
y_list: Sequence[NDArray[np.float64]],
step: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""Bin each spectrum onto the shared union grid; empty bins become NaN."""
x_min = min(float(x[0]) for x in x_list)
x_max = max(float(x[-1]) for x in x_list)
common_x = _grid(x_min, x_max, step, inclusive=True)
n = common_x.size
rows = []
for x, y in zip(x_list, y_list, strict=True):
idx = np.rint((x - x_min) / step).astype(int)
valid = (idx >= 0) & (idx < n)
idx, y_valid = idx[valid], y[valid]
sums = np.zeros(n, dtype=np.float64)
counts = np.zeros(n, dtype=np.float64)
np.add.at(sums, idx, y_valid)
np.add.at(counts, idx, 1.0)
row = np.full(n, np.nan, dtype=np.float64)
filled = counts > 0
row[filled] = sums[filled] / counts[filled]
rows.append(row)
return common_x, np.asarray(rows, dtype=np.float64)
[docs]
def resample_to_common_grid(
x_list: Sequence[NDArray[np.float64]],
y_list: Sequence[NDArray[np.float64]],
mode: Alignment | str = Alignment.INTERPOLATION,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
"""Resample ragged ``(x, y)`` spectra onto one shared x grid.
Args:
x_list: Per-spectrum 1D x arrays (may differ in length).
y_list: Per-spectrum 1D y arrays, each paired with ``x_list``.
mode: :class:`Alignment.INTERPOLATION` (default) linearly interpolates
each spectrum onto a grid spanning the *overlap* of all ranges, so
no extrapolation and no NaN. :class:`Alignment.REBINNING` bins each
spectrum onto a grid spanning the *union* of all ranges, leaving
bins with no data for a given spectrum as ``NaN``.
Returns:
``(common_x, Y)`` with ``Y`` of shape ``(len(x_list), common_x.size)``.
Raises:
ValueError: If ``x_list`` is empty or lengths of the two lists differ.
"""
if not x_list:
raise ValueError("x_list must contain at least one spectrum.")
if len(x_list) != len(y_list):
raise ValueError("x_list and y_list must have the same length.")
mode = Alignment(mode)
# Fast path: identical grids need no resampling and stay byte-identical.
first = x_list[0]
if all(x.shape == first.shape and np.array_equal(x, first) for x in x_list):
return np.asarray(first, dtype=np.float64), np.asarray(y_list, dtype=np.float64)
step = _common_step(x_list)
if mode is Alignment.REBINNING:
return _rebin(x_list, y_list, step)
return _interpolate(x_list, y_list, step)