From 229c828733568471fb9bae4280cf4a911254f13e Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Mon, 1 Jun 2026 15:23:44 +0100 Subject: [PATCH] Vectorize angular (ATS) reductions to fix XLA compile blowup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The angular forward model wrapped two functions whose Python list comprehensions were fully unrolled into the jit-traced graph, once per output bin, and then re-traced through value_and_grad: - thomson_diagnostic.reduce_ATS_to_resunit: ~1024 + 1024 slice+mean+stack ops to bin the wavelength and angular axes. - irf.add_ATS_IRF: ~2048 + 1024 jnp.convolve ops for the separable 2D instrument-response smoothing. At angular shapes (npts=2048, 1024 angles) this is ~5000 unrolled iterations baked into a single graph. On CPU the angular value_and_grad failed to finish compiling in 15+ minutes. Replace both with vectorized equivalents: - _bin_average: reshape (+ NaN-pad for a ragged final window) and mean, a single op; bit-for-bit equal to the loop including the ragged case. - add_ATS_IRF: vmap each 1D convolution over the batch axis. After the change the same angular value_and_grad compiles in ~4.6s on CPU. Outputs are unchanged (1D fits recover identical parameters). Add tests/test_forward/test_ats_vectorization.py pinning both vectorized forms to the original loop semantics — these run on CPU, unlike the angular integration tests which skip without a GPU. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_forward/test_ats_vectorization.py | 78 ++++++++++++++++++++ tsadar/core/physics/irf.py | 12 +-- tsadar/core/thomson_diagnostic.py | 30 ++++++-- 3 files changed, 110 insertions(+), 10 deletions(-) create mode 100644 tests/test_forward/test_ats_vectorization.py diff --git a/tests/test_forward/test_ats_vectorization.py b/tests/test_forward/test_ats_vectorization.py new file mode 100644 index 00000000..29cb0576 --- /dev/null +++ b/tests/test_forward/test_ats_vectorization.py @@ -0,0 +1,78 @@ +"""Regression tests locking the vectorized ATS reductions to their original semantics. + +The angular (ATS) integration tests skip on CPU, so these pure-function checks guard +the vectorized replacements of the former unrolled Python loops in +``thomson_diagnostic._bin_average`` and ``irf.add_ATS_IRF`` against the original +list-comprehension implementations. +""" +import numpy as np +import pytest +from jax import config + +config.update("jax_enable_x64", True) +config.update("jax_platform_name", "cpu") + +from jax import numpy as jnp, vmap + +from tsadar.core.thomson_diagnostic import _bin_average + + +def _old_reduce(ThryE, lamAxisE, lam_step, ang_step): + ThryE = jnp.array([jnp.average(ThryE[:, i : i + lam_step], axis=1) for i in range(0, ThryE.shape[1], lam_step)]) + ThryE = jnp.array([jnp.average(ThryE[:, i : i + ang_step], axis=1) for i in range(0, ThryE.shape[1], ang_step)]) + lamAxisE = jnp.array([jnp.average(lamAxisE[i : i + lam_step], axis=0) for i in range(0, lamAxisE.shape[0], lam_step)]) + return ThryE, lamAxisE + + +def _new_reduce(ThryE, lamAxisE, lam_step, ang_step): + ThryE = _bin_average(ThryE, lam_step, axis=1) + ThryE = _bin_average(ThryE, ang_step, axis=0) + lamAxisE = _bin_average(lamAxisE, lam_step, axis=0) + return ThryE, lamAxisE + + +def _old_conv(modlE, inst_func_ang, inst_func_lam): + ThryE = jnp.array([jnp.convolve(modlE[:, i], inst_func_ang, "same") for i in range(modlE.shape[1])]) + ThryE = jnp.array([jnp.convolve(ThryE[:, i], inst_func_lam, "same") for i in range(ThryE.shape[1])]) + return ThryE + + +def _new_conv(modlE, inst_func_ang, inst_func_lam): + ThryE = vmap(lambda col: jnp.convolve(col, inst_func_ang, "same"), in_axes=1, out_axes=1)(modlE) + ThryE = vmap(lambda row: jnp.convolve(row, inst_func_lam, "same"), in_axes=0, out_axes=0)(ThryE) + return ThryE + + +@pytest.mark.parametrize( + "nang, npts, lam_step, ang_step", + [ + (1024, 2048, 2, 1), # angular shapes (divisible) + (200, 1000, 5, 4), # divisible + (130, 1003, 7, 3), # ragged final window on both axes + ], +) +def test_bin_average_matches_loop(nang, npts, lam_step, ang_step): + rng = np.random.default_rng(0) + ThryE = jnp.array(rng.random((nang, npts))) + lamAxisE = jnp.array(rng.random((npts,))) + + oT, oL = _old_reduce(ThryE, lamAxisE, lam_step, ang_step) + nT, nL = _new_reduce(ThryE, lamAxisE, lam_step, ang_step) + + assert oT.shape == nT.shape and oL.shape == nL.shape + np.testing.assert_allclose(np.asarray(nT), np.asarray(oT), rtol=0, atol=1e-12) + np.testing.assert_allclose(np.asarray(nL), np.asarray(oL), rtol=0, atol=1e-12) + + +@pytest.mark.parametrize("nang, npts", [(256, 512), (300, 200), (64, 333)]) +def test_ats_irf_conv_matches_loop(nang, npts): + rng = np.random.default_rng(0) + modlE = jnp.array(rng.random((nang, npts))) + inst_func_ang = jnp.array(rng.random((nang,))) + inst_func_lam = jnp.array(rng.random((npts,))) + + old = _old_conv(modlE, inst_func_ang, inst_func_lam) + new = _new_conv(modlE, inst_func_ang, inst_func_lam) + + assert old.shape == new.shape + np.testing.assert_allclose(np.asarray(new), np.asarray(old), rtol=0, atol=1e-12) diff --git a/tsadar/core/physics/irf.py b/tsadar/core/physics/irf.py index 4886e090..48158c39 100644 --- a/tsadar/core/physics/irf.py +++ b/tsadar/core/physics/irf.py @@ -1,5 +1,5 @@ from typing import Tuple -from jax import numpy as jnp +from jax import numpy as jnp, vmap def add_ATS_IRF(config, sas, lamAxisE, modlE, amps, TSins) -> Tuple[jnp.ndarray, jnp.ndarray]: @@ -31,10 +31,12 @@ def add_ATS_IRF(config, sas, lamAxisE, modlE, amps, TSins) -> Tuple[jnp.ndarray, (1.0 / (stddev_ang * jnp.sqrt(2.0 * jnp.pi))) * jnp.exp(-((sas["angAxis"] - origin_ang) ** 2.0) / (2.0 * (stddev_ang) ** 2.0)) ) # Gaussian - ThryE = jnp.array([jnp.convolve(modlE[:, i], inst_func_ang, "same") for i in range(modlE.shape[1])]) - # ThryE = jnp.array([fftconvolve(modlE[:, i], inst_func_ang, "same") for i in range(modlE.shape[1])]) - ThryE = jnp.array([jnp.convolve(ThryE[:, i], inst_func_lam, "same") for i in range(ThryE.shape[1])]) - # ThryE = jnp.array([fftconvolve(ThryE[:, i], inst_func_lam, "same") for i in range(ThryE.shape[1])]) + # Separable 2D convolution: smooth along the angular axis (axis 0) for every + # wavelength column, then along the wavelength axis (axis 1) for every angle row. + # vmap batches each 1D convolution into a single op instead of unrolling a Python + # loop over thousands of columns/rows into the traced graph (huge XLA compile cost). + ThryE = vmap(lambda col: jnp.convolve(col, inst_func_ang, "same"), in_axes=1, out_axes=1)(modlE) + ThryE = vmap(lambda row: jnp.convolve(row, inst_func_lam, "same"), in_axes=0, out_axes=0)(ThryE) ThryE = jnp.amax(modlE, axis=1, keepdims=True) / jnp.amax(ThryE, axis=1, keepdims=True) * ThryE diff --git a/tsadar/core/thomson_diagnostic.py b/tsadar/core/thomson_diagnostic.py index 14e1ad42..14dd3b1e 100644 --- a/tsadar/core/thomson_diagnostic.py +++ b/tsadar/core/thomson_diagnostic.py @@ -7,6 +7,28 @@ from .physics.generate_spectra import FitModel +def _bin_average(arr, step, axis): + """Average non-overlapping windows of length ``step`` along ``axis``. + + Vectorized replacement for ``jnp.array([jnp.average(arr[..., i:i+step], axis) for i + in range(0, n, step)])``. A ragged final window (when ``n`` is not a multiple of + ``step``) is averaged over its real elements only, matching the original + list-comprehension semantics. ``step`` and the array shape are static at trace time, + so this compiles to a single reshape+mean instead of unrolling the loop into the graph. + """ + n = arr.shape[axis] + n_bins = -(-n // step) # ceil division -> matches len(range(0, n, step)) + pad = n_bins * step - n + if pad: + pad_width = [(0, 0)] * arr.ndim + pad_width[axis] = (0, pad) + arr = jnp.pad(arr, pad_width, constant_values=jnp.nan) + arr = jnp.moveaxis(arr, axis, 0) + arr = arr.reshape(n_bins, step, *arr.shape[1:]) + arr = jnp.nanmean(arr, axis=1) + return jnp.moveaxis(arr, 0, axis) + + class ThomsonScatteringDiagnostic: """ The SpectrumCalculator class wraps the FitModel class adding instrumental effects to the calculated spectrum so it @@ -94,12 +116,10 @@ def reduce_ATS_to_resunit(self, ThryE, lamAxisE, TSins, batch): lam_step = round(ThryE.shape[1] / batch["e_data"].shape[1]) ang_step = round(ThryE.shape[0] / self.cfg["other"]["CCDsize"][0]) - ThryE = jnp.array([jnp.average(ThryE[:, i : i + lam_step], axis=1) for i in range(0, ThryE.shape[1], lam_step)]) - ThryE = jnp.array([jnp.average(ThryE[:, i : i + ang_step], axis=1) for i in range(0, ThryE.shape[1], ang_step)]) + ThryE = _bin_average(ThryE, lam_step, axis=1) # bin the wavelength axis + ThryE = _bin_average(ThryE, ang_step, axis=0) # bin the angular axis - lamAxisE = jnp.array( - [jnp.average(lamAxisE[i : i + lam_step], axis=0) for i in range(0, lamAxisE.shape[0], lam_step)] - ) + lamAxisE = _bin_average(lamAxisE, lam_step, axis=0) ThryE = ThryE[self.cfg["data"]["lineouts"]["start"] : self.cfg["data"]["lineouts"]["end"], :] ThryE = batch["e_amps"] * ThryE / jnp.amax(ThryE, axis=1, keepdims=True) ThryE = jnp.where(