From 3350ba6cc9f09c68aea0a4791ac80bf5a09d240b Mon Sep 17 00:00:00 2001 From: lbgee Date: Thu, 16 Jul 2026 23:41:36 -0700 Subject: [PATCH] Add common_mode_correction and subtract_polynomial_background steps common_mode_correction subtracts per-row, per-column, or per-bank electronic baseline estimated from a signal-free reference band, on 3D detector data before shot reduction, preserving shape. subtract_polynomial_background fits a low-order polynomial to signal-free regions along the spatial axis and subtracts it, writing a non-destructive _bkgsub. peak_mask accepts multiple ranges so several emission lines dispersed on one detector can be masked at once. Per-row weighting excludes NaNs from the fit. Reuses the vectorized weighted-polyfit approach from patch_pixels. Adds unit tests for both steps and documents them in the YAML guide and README step lists. Closes #102 --- README.md | 4 +- XSpect/analysis/spectroscopy.py | 197 +++++++++++++++++++++++++++++++ docs/YAML_PIPELINE_GUIDE.md | 24 ++++ tests/test_spectroscopy_steps.py | 159 +++++++++++++++++++++++++ 4 files changed, 382 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 67d7bb1..8951000 100644 --- a/README.md +++ b/README.md @@ -88,9 +88,9 @@ The same pipeline machinery covers the common LCLS spectroscopy modes. Example c Full list and arguments in the [YAML pipeline guide](docs/YAML_PIPELINE_GUIDE.md). - Shot filtering: `filter_shots`, `filter_detector_adu`, `filter_detector_variance` (sklearn VarianceThreshold, data-driven alternative to ADU/keV thresholds) -- Detector prep: `patch_pixels` (manual or auto-detected bad columns), `rotate_detector`, `find_rotation_angle`, `droplet_reconstruction` +- Detector prep: `patch_pixels` (manual or auto-detected bad columns), `rotate_detector`, `find_rotation_angle`, `common_mode_correction` (per-row/column/bank baseline from a dark band), `droplet_reconstruction` - Binning and reduction: `reduce_detector_temporal`, `combine_runs` -- Spectra: `make_energy_axis`, `normalize_xes` +- Spectra: `make_energy_axis`, `normalize_xes`, `subtract_polynomial_background` (polynomial baseline from the spatial axis, multi-peak masking) ## Skills diff --git a/XSpect/analysis/spectroscopy.py b/XSpect/analysis/spectroscopy.py index 5cf672f..aac0a7a 100644 --- a/XSpect/analysis/spectroscopy.py +++ b/XSpect/analysis/spectroscopy.py @@ -169,6 +169,203 @@ def filter_detector_variance(run, **kwargs): ) +@register_step("common_mode_correction") +def common_mode_correction(run, **kwargs): + """Subtract per-row, per-column, or per-bank common-mode offsets. + + Detector electronics add a slowly varying baseline that is shared across + all pixels in a readout unit (a row, a column, or an ePix100 bank). The + offset is estimated from a reference region that carries no signal (a dark + band) and subtracted from the whole unit. Runs on 3D detector data + (shots x rows x cols) before shot reduction and preserves the input shape. + + Parameters from YAML: + on: detector key (3D: shots x rows x cols, or 2D: rows x cols) + axis: "row" (default), "column", or "bank". "row" removes a per-row + offset (shared across columns), "column" a per-column offset, + "bank" a per-bank offset across fixed-width column blocks. + method: "median" (default, robust to outliers) or "mean". + reference: [start, end] pixel range of the signal-free band used to + estimate the offset. Indexed along the axis orthogonal to the + correction: for axis="row" it is a column range, for axis="column" + (and "bank") it is a row range. Default: full extent. + bank_size: int, column width of an ePix100 bank (default 128). Only + used when axis="bank". + """ + detector_key = kwargs.get("on") + axis = kwargs.get("axis", "row") + method = kwargs.get("method", "median") + reference = kwargs.get("reference", None) + bank_size = int(kwargs.get("bank_size", 128)) + if detector_key is None: + return + + data = getattr(run, detector_key, None) + if data is None: + run.update_status(f"common_mode_correction: {detector_key} not found") + return + + data = np.asarray(data, dtype=np.float64) + if data.ndim == 2: + # promote to (1, rows, cols) so one code path handles both + working = data[np.newaxis, ...] + squeezed = True + elif data.ndim == 3: + working = data + squeezed = False + else: + run.update_status( + f"common_mode_correction: {detector_key} must be 2D or 3D, got {data.ndim}D" + ) + return + + reducer = np.nanmedian if method == "median" else np.nanmean + + # working is (shots, rows, cols). The reference range slices the axis + # orthogonal to the correction direction so the offset is estimated only + # from the dark band. + if axis == "row": + # per-row offset shared across columns; reference is a column range + if reference is not None: + band = working[:, :, reference[0] : reference[1]] + else: + band = working + offset = reducer(band, axis=2, keepdims=True) # (shots, rows, 1) + corrected = working - offset + elif axis == "column": + # per-column offset shared across rows; reference is a row range + if reference is not None: + band = working[:, reference[0] : reference[1], :] + else: + band = working + offset = reducer(band, axis=1, keepdims=True) # (shots, 1, cols) + corrected = working - offset + elif axis == "bank": + # per-bank offset: split columns into fixed-width blocks, estimate one + # offset per block from the reference row band, subtract per block. + if reference is not None: + band = working[:, reference[0] : reference[1], :] + else: + band = working + n_cols = working.shape[2] + corrected = working.copy() + for start in range(0, n_cols, bank_size): + end = min(start + bank_size, n_cols) + offset = reducer(band[:, :, start:end], axis=(1, 2), keepdims=True) + corrected[:, :, start:end] = working[:, :, start:end] - offset + else: + run.update_status( + f"common_mode_correction: unknown axis '{axis}' (use row|column|bank)" + ) + return + + if squeezed: + corrected = corrected[0] + setattr(run, detector_key, corrected) + run.update_status( + f"Common-mode corrected {detector_key} (axis={axis}, method={method})" + ) + + +@register_step("subtract_polynomial_background") +def subtract_polynomial_background(run, **kwargs): + """Subtract a polynomial baseline fit along the spatial axis. + + Fits a low-order polynomial to signal-free regions along one axis and + subtracts it, removing smooth scattering/fluorescence background while + preserving peak area. The peak region is excluded from the fit either by + naming the background regions explicitly (``background``) or by masking + the peak (``peak_mask``). Non-destructive: writes ``_bkgsub``. + + Works on a 1D spectrum or a 2D array (bins x pixels). The fit reuses the + vectorized weighted-polynomial projection from patch_pixels: the offset + vector depends only on the sample positions and weights, so it is built + once and applied to every row with a single matrix multiply. + + Parameters from YAML: + on: spectrum key (1D pixels, or 2D bins x pixels). + axis: spatial axis to fit along (default: last axis). + order: polynomial degree (default 2). + background: list of [start, end] pixel ranges to fit (signal-free). + If given, only these ranges anchor the fit. + peak_mask: pixel range(s) to EXCLUDE from the fit (the emission peaks). + Accepts a single [start, end] or a list of [start, end] ranges, so + several dispersed lines (e.g. Kalpha and Kbeta on one detector) can + all be masked at once. Used when naming the peaks is easier than the + background. Ignored if ``background`` is given. + """ + detector_key = kwargs.get("on") + axis = kwargs.get("axis", None) + order = int(kwargs.get("order", 2)) + background = kwargs.get("background", None) + peak_mask = kwargs.get("peak_mask", None) + if detector_key is None: + return + + data = getattr(run, detector_key, None) + if data is None: + run.update_status(f"subtract_polynomial_background: {detector_key} not found") + return + + data = np.asarray(data, dtype=np.float64) + if axis is None: + axis = data.ndim - 1 + + n_pixels = data.shape[axis] + x = np.arange(n_pixels, dtype=np.float64) + + # weights select which pixels anchor the fit: 1 for background, 0 for peak. + weights = np.ones(n_pixels, dtype=np.float64) + if background is not None: + weights[:] = 0.0 + for rng in background: + weights[rng[0] : rng[1]] = 1.0 + elif peak_mask is not None: + # accept a single [start, end] or a list of ranges (multiple lines) + mask_ranges = peak_mask + if len(peak_mask) == 2 and np.isscalar(peak_mask[0]): + mask_ranges = [peak_mask] + for rng in mask_ranges: + weights[rng[0] : rng[1]] = 0.0 + + # Bring the fit axis to front so every other axis is a batch dimension. + moved = np.moveaxis(data, axis, 0) # (n_pixels, ...) + flat = moved.reshape(n_pixels, -1).copy() # (n_pixels, N_rows) + nan_mask = np.isnan(flat) + flat[nan_mask] = 0.0 + + if np.sum(weights > 0.5) < order + 1: + run.update_status( + f"subtract_polynomial_background: too few background pixels " + f"({int(np.sum(weights > 0.5))}) for order {order}; skipped" + ) + return + + # Per-row weights: the base background/peak mask, zeroed wherever a row has + # a NaN so those points never enter that row's fit. NaN positions differ + # per row, so the normal equations are solved per row (batched). The base + # weight is squared to match numpy.polyfit's 1/sigma convention + # (minimise sum(w**2 * r**2)). + V = np.vander(x, order + 1) # (n_pixels, order+1) + w_full = (weights**2)[:, np.newaxis] * (~nan_mask) # (n_pixels, N_rows) + + # A[r] = V^T diag(w_r) V ; b[r] = V^T diag(w_r) flat_r, batched over rows r. + A = np.einsum("pk,pr,pl->rkl", V, w_full, V) # (N_rows, order+1, order+1) + b = np.einsum("pk,pr->rk", V, w_full * flat) # (N_rows, order+1) + coeffs = np.linalg.solve(A, b) # (N_rows, order+1) + baseline = (V @ coeffs.T) # (n_pixels, N_rows) + + subtracted = flat - baseline + subtracted[nan_mask] = np.nan # keep original NaN positions + result = np.moveaxis(subtracted.reshape(moved.shape), 0, axis) + + setattr(run, f"{detector_key}_bkgsub", result) + run.update_status( + f"Polynomial background subtracted {detector_key} -> " + f"{detector_key}_bkgsub (order={order}, axis={axis})" + ) + + @register_step("find_rotation_angle") def find_rotation_angle(run, **kwargs): """Auto-detect the tilt angle of a dispersed spectral signal. diff --git a/docs/YAML_PIPELINE_GUIDE.md b/docs/YAML_PIPELINE_GUIDE.md index 671fece..ca9090e 100644 --- a/docs/YAML_PIPELINE_GUIDE.md +++ b/docs/YAML_PIPELINE_GUIDE.md @@ -139,6 +139,7 @@ reduction: | `filter_shots` | `on`, `filter_key`, `threshold` | Zero out shots below threshold (or outside range if threshold is [min, max]) | | `filter_detector_adu` | `on`, `adu_threshold` | Zero detector pixels below ADU threshold | | `filter_detector_variance` | `on`, `variance_threshold` | Zero low-variance detector pixels using sklearn `VarianceThreshold`. Data-driven alternative to `filter_detector_adu` — no ADU cutoff to hand-tune | +| `common_mode_correction` | `on`, `axis`, `method`, `reference`, `bank_size` | Subtract per-row, per-column, or per-bank electronic baseline estimated from a signal-free reference band. Runs on 3D detector data before shot reduction; preserves shape | `filter_detector_variance` computes each pixel's variance across shots and zeros the pixels that barely change (dead pixels, static hot pixels, constant background). Signal-bearing pixels vary shot to shot and are kept. It writes back to the same detector key and stores the retained boolean mask as `_variance_mask` for inspection. `variance_threshold` defaults to `0.0` (removes only constant pixels); raise it to drop low- but nonzero-variance pixels. @@ -149,6 +150,17 @@ reduction: variance_threshold: 0.01 ``` +`common_mode_correction` removes the slowly varying baseline shared across a readout unit. The offset is estimated per shot from a `reference` range that carries no signal (a dark band) and subtracted from the whole unit. `axis: row` subtracts a per-row offset (reference is a column range), `axis: column` a per-column offset (reference is a row range), `axis: bank` a per-bank offset over `bank_size`-wide column blocks (ePix100 default 128). `method` is `median` (default, robust) or `mean`. + +```yaml +# Per-row common mode from a dark column band +- step: common_mode_correction + on: epix + axis: row + method: median + reference: [700, 768] # signal-free columns +``` + ### Shot Combination | Step | Parameters | Description | @@ -176,9 +188,21 @@ reduction: | Step | Parameters | Description | |------|-----------|-------------| | `normalize_xes` | `on`, `pixel_range` | Normalize each row by its sum. Produces `{on}_normalized` | +| `subtract_polynomial_background` | `on`, `axis`, `order`, `background`, `peak_mask` | Fit a polynomial baseline to signal-free regions along the spatial axis and subtract it. Non-destructive: produces `{on}_bkgsub` | | `make_energy_axis` | `detector_key`, `n_pixels`, `crystal_detector_distance`, `crystal_radius`, `d_spacing`, `mm_per_pixel`, `name` | Convert pixels to energy via vonHamos geometry. Produces `{name}_energy` | | `patch_pixels` | `on`, `pixels`, `mode` | Repair bad pixels by interpolation or zeroing | +`subtract_polynomial_background` removes smooth scattering/fluorescence background while preserving peak area. It fits a low-order polynomial (`order`, default 2) along `axis` (default last) using only signal-free pixels, then subtracts the fitted baseline. Name the fit regions with `background` (a list of `[start, end]` ranges), or name the peaks to exclude with `peak_mask`. `peak_mask` accepts a single `[start, end]` or a list of ranges, so several emission lines dispersed on one detector can all be masked at once. NaN pixels are excluded from the fit and left as NaN in the output. + +```yaml +# Two emission lines on one detector: mask both, fit the rest +- step: subtract_polynomial_background + on: epix_ROI_1_time_binned + axis: 1 + order: 2 + peak_mask: [[30, 70], [95, 125]] +``` + ### XAS-Specific | Step | Parameters | Description | diff --git a/tests/test_spectroscopy_steps.py b/tests/test_spectroscopy_steps.py index a5fd248..29c5ca1 100644 --- a/tests/test_spectroscopy_steps.py +++ b/tests/test_spectroscopy_steps.py @@ -184,3 +184,162 @@ def test_purge_keys(self, mock_run): step = get_step("purge_keys") step(mock_run, keys=["epix_ROI_1"]) assert mock_run.epix_ROI_1 is None + + +class TestCommonModeCorrection: + def test_per_row_offset_removed(self): + # 3 shots, 4 rows, 6 cols; each row carries a known constant offset + run = MockRun() + offsets = np.array([1.0, 2.0, 3.0, 4.0]) + data = np.zeros((3, 4, 6)) + offsets[np.newaxis, :, np.newaxis] + run.det = data.copy() + get_step("common_mode_correction")(run, on="det", axis="row") + np.testing.assert_allclose(run.det, 0.0, atol=1e-12) + + def test_flat_frame_unchanged(self): + run = MockRun() + run.det = np.full((2, 5, 5), 7.0) + get_step("common_mode_correction")(run, on="det", axis="row", method="mean") + np.testing.assert_allclose(run.det, 0.0, atol=1e-12) + + def test_reference_band_isolates_offset(self): + # signal in cols 0-2, dark band in cols 3-5 carrying offset 5 + run = MockRun() + data = np.zeros((1, 3, 6)) + data[0, :, 0:3] = 100.0 + data[0, :, 3:6] = 5.0 + run.det = data.copy() + get_step("common_mode_correction")(run, on="det", axis="row", reference=[3, 6]) + np.testing.assert_allclose(run.det[0, :, 0:3], 95.0, atol=1e-12) + np.testing.assert_allclose(run.det[0, :, 3:6], 0.0, atol=1e-12) + + def test_column_axis(self): + run = MockRun() + col_offsets = np.array([1.0, 2.0, 3.0]) + run.det = np.zeros((2, 4, 3)) + col_offsets[np.newaxis, np.newaxis, :] + get_step("common_mode_correction")(run, on="det", axis="column") + np.testing.assert_allclose(run.det, 0.0, atol=1e-12) + + def test_bank_axis(self): + # 8 cols, bank_size 4 -> two banks with distinct offsets + run = MockRun() + data = np.zeros((1, 3, 8)) + data[0, :, 0:4] = 10.0 + data[0, :, 4:8] = 20.0 + run.det = data.copy() + get_step("common_mode_correction")(run, on="det", axis="bank", bank_size=4) + np.testing.assert_allclose(run.det, 0.0, atol=1e-12) + + def test_shape_preserved_2d(self): + run = MockRun() + run.det = np.random.rand(6, 6) + get_step("common_mode_correction")(run, on="det", axis="row") + assert run.det.shape == (6, 6) + + def test_missing_key_noop(self): + run = MockRun() + get_step("common_mode_correction")(run, on="nope", axis="row") + assert any("not found" in s for s in run.status) + + +class TestSubtractPolynomialBackground: + def test_gaussian_on_linear_background(self): + # narrow peak, generous mask so tails don't contaminate the fit region + n = 100 + x = np.arange(n, dtype=float) + bkg = 2.0 + 0.05 * x + peak = 500.0 * np.exp(-((x - 50) ** 2) / (2 * 3.0**2)) + run = MockRun() + run.spec = bkg + peak + get_step("subtract_polynomial_background")( + run, on="spec", order=1, peak_mask=[30, 70] + ) + result = run.spec_bkgsub + np.testing.assert_allclose(result[0:25], 0.0, atol=1e-6) + np.testing.assert_allclose(result[75:], 0.0, atol=1e-6) + np.testing.assert_allclose(result.sum(), peak.sum(), rtol=1e-4) + + def test_pure_background_yields_zero(self): + n = 80 + x = np.arange(n, dtype=float) + run = MockRun() + run.spec = 3.0 - 0.02 * x + 0.001 * x**2 + get_step("subtract_polynomial_background")(run, on="spec", order=2) + np.testing.assert_allclose(run.spec_bkgsub, 0.0, atol=1e-8) + + def test_background_ranges(self): + n = 100 + x = np.arange(n, dtype=float) + peak = 300.0 * np.exp(-((x - 50) ** 2) / (2 * 3.0**2)) + run = MockRun() + run.spec = 10.0 + 0.1 * x + peak + get_step("subtract_polynomial_background")( + run, on="spec", order=1, background=[[0, 30], [70, 100]] + ) + np.testing.assert_allclose(run.spec_bkgsub[0:28], 0.0, atol=1e-6) + + def test_multiple_peak_masks(self): + # two dispersed emission lines on one axis; mask both + n = 150 + x = np.arange(n, dtype=float) + bkg = 5.0 + 0.02 * x + line1 = 400.0 * np.exp(-((x - 40) ** 2) / (2 * 3.0**2)) + line2 = 300.0 * np.exp(-((x - 100) ** 2) / (2 * 3.0**2)) + run = MockRun() + run.spec = bkg + line1 + line2 + get_step("subtract_polynomial_background")( + run, on="spec", order=1, peak_mask=[[25, 55], [85, 115]] + ) + result = run.spec_bkgsub + # background between and around the two lines returns to zero + # (atol allows for negligible Gaussian-tail leakage at the mask edges) + np.testing.assert_allclose(result[0:20], 0.0, atol=1e-3) + np.testing.assert_allclose(result[65:80], 0.0, atol=1e-3) + np.testing.assert_allclose(result[125:], 0.0, atol=1e-3) + # combined line area preserved + np.testing.assert_allclose( + result.sum(), (line1 + line2).sum(), rtol=1e-3 + ) + + def test_2d_rows_fit_independently(self): + n_bins, n_pix = 5, 100 + x = np.arange(n_pix, dtype=float) + data = np.zeros((n_bins, n_pix)) + for i in range(n_bins): + data[i] = (1.0 + i) + 0.03 * x + data[i] += 200.0 * np.exp(-((x - 50) ** 2) / (2 * 3.0**2)) + run = MockRun() + run.spec = data + get_step("subtract_polynomial_background")( + run, on="spec", order=1, peak_mask=[30, 70] + ) + result = run.spec_bkgsub + assert result.shape == (n_bins, n_pix) + np.testing.assert_allclose(result[:, 0:25], 0.0, atol=1e-6) + + def test_nans_do_not_propagate(self): + n = 80 + x = np.arange(n, dtype=float) + spec = 5.0 + 0.1 * x + spec[10] = np.nan + spec[20] = np.nan + run = MockRun() + run.spec = spec + get_step("subtract_polynomial_background")(run, on="spec", order=1) + result = run.spec_bkgsub + assert np.isnan(result[10]) and np.isnan(result[20]) + finite = result[np.isfinite(result)] + np.testing.assert_allclose(finite, 0.0, atol=1e-6) + + def test_non_destructive(self): + run = MockRun() + original = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + run.spec = original.copy() + get_step("subtract_polynomial_background")(run, on="spec", order=1) + np.testing.assert_array_equal(run.spec, original) + assert hasattr(run, "spec_bkgsub") + + def test_missing_key_noop(self): + run = MockRun() + get_step("subtract_polynomial_background")(run, on="nope") + assert any("not found" in s for s in run.status)