Centre the oversampled model grid on the detector pixels

docs/source/defaults.rst

@@ -416,7 +416,7 @@ The ``other:`` section includes options specifying the types of data that are be
 
 - ``ang_res_unit`` is the number of pixels in an angular resolution unit for ARTS
 
-- ``lam_res_unit`` is the number of pixels in an specular resolution unit for ARTS
+- ``lam_res_unit`` is the number of pixels in an specular resolution unit for ARTS. If it does not divide the number of wavelength pixels, the leftover pixels form a narrower final unit when the electron distribution function is 2D with resonance quadrature enabled, and are dropped otherwise
 
 
 Plotting

test_angular_2v.py

@@ -58,7 +58,7 @@ def test_arts2d_forward_pass():
                 config["data"]["fit_rng"]["forward_iaw_start"],
                 config["data"]["fit_rng"]["forward_iaw_end"],
             ]
-            config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+            config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
             sas = get_scattering_angles(config)
 
             [axisxE, _, _, _, _, _] = get_calibrations(

test_arts2v_random.py

@@ -162,7 +162,7 @@ def test_arts2d_inverse(config_path: str = "tests/configs/arts2d_test_defaults.y
                 config["data"]["fit_rng"]["forward_iaw_start"],
                 config["data"]["fit_rng"]["forward_iaw_end"],
             ]
-            config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+            config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
             sas = get_scattering_angles(config)
 
             sas["angAxis"], _, _, _, _, _ = get_calibrations(

tests/test_forward/test_1d.py

@@ -48,7 +48,17 @@ def test_1d_forward_pass():
             config["data"]["fit_rng"]["forward_iaw_start"],
             config["data"]["fit_rng"]["forward_iaw_end"],
         ]
-        config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+        config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
+        # Accounts for the pixel-centred model grid fix: ThryE-1d.npy was generated on the old grid,
+        # linspace(start, end, npts). Shrink the ranges so the new grid lands on those same points.
+        # Remove this block when the reference is regenerated on the new grid.
+        ppp = config["other"]["points_per_pixel"]
+        n_pixels = config["other"]["npts"] // ppp
+        for key in ("lamrangE", "lamrangI"):
+            start, end = config["other"][key]
+            pixel_width = (end - start) / (n_pixels - 1 + (1 - 1 / ppp))
+            overhang = 0.5 * pixel_width * (1 - 1 / ppp)
+            config["other"][key] = [start + overhang, end - overhang]
         sas = get_scattering_angles(config)
 
         dummy_batch = {

tests/test_forward/test_angular_1v.py

@@ -48,7 +48,17 @@ def test_arts1d_forward_pass():
             config["data"]["fit_rng"]["forward_iaw_start"],
             config["data"]["fit_rng"]["forward_iaw_end"],
         ]
-        config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+        config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
+        # Accounts for the pixel-centred model grid fix: ThryE-arts1v.npy was generated on the old grid,
+        # linspace(start, end, npts). Shrink the ranges so the new grid lands on those same points.
+        # Remove this block when the reference is regenerated on the new grid.
+        ppp = config["other"]["points_per_pixel"]
+        n_pixels = config["other"]["npts"] // ppp
+        for key in ("lamrangE", "lamrangI"):
+            start, end = config["other"][key]
+            pixel_width = (end - start) / (n_pixels - 1 + (1 - 1 / ppp))
+            overhang = 0.5 * pixel_width * (1 - 1 / ppp)
+            config["other"][key] = [start + overhang, end - overhang]
         sas = get_scattering_angles(config)
 
         [axisxE, _, _, _, _, _] = get_calibrations(

tests/test_forward/test_angular_2v.py

@@ -58,7 +58,7 @@ def test_arts2d_forward_pass():
                 config["data"]["fit_rng"]["forward_iaw_start"],
                 config["data"]["fit_rng"]["forward_iaw_end"],
             ]
-            config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+            config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
             sas = get_scattering_angles(config)
 
             [axisxE, _, _, _, _, _] = get_calibrations(

tests/test_forward/test_irf_area.py

@@ -18,17 +18,14 @@
     _gaussian_bin_response,
     add_ATS_IRF,
 )
+from tsadar.utils.detector_grid import detector_edges_from_centers, grouped_detector_edges, oversampled_axis
 from tsadar.core import thomson_diagnostic
 from tsadar.core.thomson_diagnostic import (
     ThomsonScatteringDiagnostic,
     _ensure_angular_detector_edges,
 )
 from tsadar.data import prepare
-from tsadar.data.calibration import (
-    detector_edges_from_centers,
-    get_calibrations,
-    grouped_detector_edges,
-)
+from tsadar.data.calibration import get_calibrations
 from tsadar.utils.misc import merge_defaults_and_inputs
 
 
@@ -69,7 +66,38 @@ def test_grouped_detector_edges_reject_invalid_group_size(bad_group):
         grouped_detector_edges(np.arange(4.0), bad_group)
 
 
-def test_prepare_data_retains_exact_arts_spectral_edges(monkeypatch):
+@pytest.mark.parametrize("points_per_pixel", [1, 2, 5])
+def test_oversampled_axis_groups_are_centered_on_the_detector_pixels(points_per_pixel):
+    centers = np.linspace(400.0, 700.0, 7)
+
+    fine = oversampled_axis(centers, points_per_pixel)
+
+    assert fine.shape == (centers.size * points_per_pixel,)
+    np.testing.assert_allclose(np.diff(fine), 50.0 / points_per_pixel, rtol=0, atol=1e-10)
+    np.testing.assert_allclose(fine.reshape(centers.size, -1).mean(axis=1), centers, rtol=0, atol=1e-10)
+
+
+def test_oversampled_axis_stays_inside_each_pixel_of_a_nonuniform_axis():
+    centers = np.array([400.0, 400.5, 401.0, 402.0, 404.0])
+    edges = detector_edges_from_centers(centers)
+
+    fine = oversampled_axis(centers, 5).reshape(centers.size, -1)
+
+    assert np.all(np.diff(fine.ravel()) > 0)
+    assert np.all(fine > edges[:-1, None]) and np.all(fine < edges[1:, None])
+    np.testing.assert_allclose(fine.mean(axis=1), 0.5 * (edges[:-1] + edges[1:]), rtol=0, atol=1e-10)
+
+
+@pytest.mark.parametrize("bad_points_per_pixel", [0, -1, 1.5, True])
+def test_oversampled_axis_rejects_invalid_points_per_pixel(bad_points_per_pixel):
+    with pytest.raises(ValueError):
+        oversampled_axis(np.arange(4.0), bad_points_per_pixel)
+
+
+@pytest.mark.parametrize("fe_dim, detector_binned", [(2, True), (1, False)])
+def test_prepare_data_keeps_a_ragged_final_arts_spectral_unit_only_for_detector_bins(
+    monkeypatch, fe_dim, detector_binned
+):
     electron_data = np.arange(20.0).reshape(5, 4)
     ion_data = np.zeros_like(electron_data)
     wavelength_centers = np.arange(5.0) + 0.5
@@ -122,6 +150,7 @@ def test_prepare_data_retains_exact_arts_spectral_edges(monkeypatch):
         },
         "feature_detector": {"estimate_lineouts_epw": False, "estimate_lineouts_iaw": False},
         "optimizer": {"batch_size": 1},
+        "parameters": {"electron": {"fe": {"dim": fe_dim}}},
         "other": {
             "CCDsize": [5, 4],
             "ang_res_unit": 2,
@@ -134,16 +163,32 @@ def test_prepare_data_retains_exact_arts_spectral_edges(monkeypatch):
 
     all_data, _, all_axes = prepare.prepare_data(cfg, 1)
 
-    np.testing.assert_array_equal(all_axes["epw_y"].squeeze(), np.array([1.0, 3.0, 4.5]))
-    np.testing.assert_array_equal(
-        cfg["other"]["detector_specs"]["electron_wavelength_edges"], np.array([0.0, 2.0, 4.0, 5.0])
-    )
-    np.testing.assert_array_equal(
-        cfg["other"]["detector_specs"]["electron_wavelength_centers"],
-        np.array([1.0, 3.0, 4.5]),
-    )
-    assert all_data["e_data"].shape == (2, 3)
-    assert cfg["other"]["CCDsize"] == (2, 3)
+    # 5 wavelength pixels in units of 2. The detector-bin quadrature keeps the fifth pixel as a
+    # half-width unit; the uniform model grid drops it.
+    if detector_binned:
+        centers, edges = np.array([1.0, 3.0, 4.5]), np.array([0.0, 2.0, 4.0, 5.0])
+    else:
+        centers, edges = np.array([1.0, 3.0]), np.array([0.0, 2.0, 4.0])
+    np.testing.assert_array_equal(all_axes["epw_y"].squeeze(), centers)
+    np.testing.assert_array_equal(cfg["other"]["detector_specs"]["electron_wavelength_edges"], edges)
+    np.testing.assert_array_equal(cfg["other"]["detector_specs"]["electron_wavelength_centers"], centers)
+    assert all_data["e_data"].shape == (2, centers.size)
+    assert cfg["other"]["CCDsize"] == (2, centers.size)
+    assert cfg["other"]["lamrangE"] == [centers[0], centers[-1]]
+    assert cfg["other"]["npts"] == centers.size
+
+
+def test_angular_irf_can_be_jit_compiled():
+    import jax
+
+    irf = AngularIRF(spect_stddev=0.5, ang_stddev=1.0, ang_axis=np.linspace(50.0, 70.0, 6), normalize=0)
+    wavelengths = jnp.linspace(400.0, 500.0, 16)
+    model = jnp.arange(96.0).reshape(6, 16)
+
+    expected = add_ATS_IRF(irf, wavelengths, model, _instrument_params())[1]
+    compiled = jax.jit(lambda values: add_ATS_IRF(irf, wavelengths, values, _instrument_params())[1])(model)
+
+    np.testing.assert_allclose(compiled, expected, rtol=1e-12)
 
 
 def test_forward_arts_bounds_remain_first_and_last_detector_centers():

tests/test_inverse/test_1d_random.py

@@ -89,7 +89,7 @@ def test_1d_inverse():
             config["data"]["fit_rng"]["forward_iaw_start"],
             config["data"]["fit_rng"]["forward_iaw_end"],
         ]
-        config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+        config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
         sas = get_scattering_angles(config)
 
         dummy_batch = {

tests/test_inverse/test_angular_1v.py

@@ -135,7 +135,7 @@ def test_arts1d_inverse(arbitrary_distribution: bool):
                 config["data"]["fit_rng"]["forward_iaw_start"],
                 config["data"]["fit_rng"]["forward_iaw_end"],
             ]
-            config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+            config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
             sas = get_scattering_angles(config)
 
             sas["angAxis"], _, _, _, _, _ = get_calibrations(

tests/test_inverse/test_angular_2v.py

@@ -153,7 +153,7 @@ def test_arts2d_inverse(dist_type: bool):
                 config["data"]["fit_rng"]["forward_iaw_start"],
                 config["data"]["fit_rng"]["forward_iaw_end"],
             ]
-            config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"])
+            config["other"]["npts"] = int(config["other"]["CCDsize"][0] * config["other"]["points_per_pixel"])
             sas = get_scattering_angles(config)
 
             sas["angAxis"], _, _, _, _, _ = get_calibrations(

tests/test_inverse/test_mcmc.py

@@ -776,13 +776,29 @@ def test_covar_throughput_correction_does_not_depend_on_points_per_pixel(fitted_
     cfg = copy.deepcopy(fitted_fixture["config"])
     cfg["optimizer"]["loss_method"] = "covar"
     cfg["data"]["background"]["bg_subtract"] = False
-    # oversample the model: its binned wavelength axis then no longer coincides with the pixel centers
+    # oversample the model: its binned wavelength axis is still the detector pixel centers
     cfg["other"]["points_per_pixel"] = 2
-    cfg["other"]["npts"] = int(cfg["other"]["CCDsize"][1] * 2)
+    cfg["other"]["npts"] = int(cfg["other"]["CCDsize"][0] * 2)
     loss_fn = _build_loss_fn_for_draw(cfg, fitted_fixture["sa"], fitted_fixture["all_data"], cfg["optimizer"]["batch_size"])
 
     # the correction is the one applied to the data, on the detector pixel centers
     expected = throughput_correction(
         cfg["other"]["extraoptions"]["spectype"], np.asarray(fitted_fixture["all_axes"]["epw_y"]), cfg["data"]["shotnum"]
     )
     np.testing.assert_allclose(np.asarray(loss_fn.covar_throughput_e), expected, rtol=1e-10)
+
+
+@pytest.mark.parametrize("points_per_pixel", [1, 2, 5])
+def test_binned_model_axis_matches_the_detector_pixel_centers(fitted_fixture, points_per_pixel):
+    from tsadar.inverse.loss_function import LossFunction
+
+    cfg = copy.deepcopy(fitted_fixture["config"])
+    batch_size = cfg["optimizer"]["batch_size"]
+    cfg["other"]["points_per_pixel"] = points_per_pixel
+    cfg["other"]["npts"] = int(cfg["other"]["CCDsize"][0] * points_per_pixel)
+    sample = {k: v[:batch_size] for k, v in fitted_fixture["all_data"].items()}
+    sample = {"noise_e": sample["noiseE"], "noise_i": sample["noiseI"]} | sample
+    loss_fn = LossFunction(cfg, fitted_fixture["sa"], sample)
+
+    _, _, lamAxisE, _ = loss_fn.ts_diag(fitted_fixture["fitted_weights"][0], fitted_fixture["batch"])
+    np.testing.assert_allclose(np.asarray(lamAxisE)[0], np.asarray(fitted_fixture["all_axes"]["epw_y"]), rtol=0, atol=1e-9)

tests/test_physics/test_arts_recovery.py

@@ -29,6 +29,7 @@ def _arts_config(*, refined=False):
     }
     cfg["other"].update(
         iawfilter=[False, 0, 0, 0],
+        points_per_pixel=1,
         npts=64,
         CCDsize=[64, 3],
         lamrangE=[440.0, 515.0],

tests/test_physics/test_production_recovery.py

@@ -41,6 +41,8 @@ def _recovery_config(feature, *, refined=False):
     # peak normalization even when its wavelengths are excluded from the loss.
     cfg["other"].update(
         CCDsize=[128, 128],
+        # 128 detector pixels, refined by sampling more points inside each pixel
+        points_per_pixel=16 if refined else 8,
         npts=2048 if refined else 1024,
         lamrangE=[440.0, 630.0],
         lamrangI=[524.5, 528.5],

tsadar/core/instrument/irf.py

@@ -21,21 +21,7 @@
 from jax.scipy.special import ndtr
 from jax.scipy.signal import fftconvolve
 
-
-def _bin_edges_from_centers(axis):
-    """Infer finite physical-coordinate cells around ordered sample centers."""
-
-    axis = jnp.ravel(jnp.asarray(axis))
-    if axis.size < 2:
-        raise ValueError("At least two coordinate centers are required for a Gaussian IRF")
-    spacing = jnp.diff(axis)
-    return jnp.concatenate(
-        (
-            axis[:1] - 0.5 * spacing[:1],
-            axis[:-1] + 0.5 * spacing,
-            axis[-1:] + 0.5 * spacing[-1:],
-        )
-    )
+from ...utils.detector_grid import detector_edges_from_centers
 
 
 def _normal_cdf_antiderivative(z):
@@ -65,11 +51,12 @@ def _gaussian_bin_response(axis, stddev):
     if stddev <= 0:
         raise ValueError(f"Gaussian standard deviation must be positive, got {stddev!r}.")
 
-    axis = jnp.ravel(jnp.asarray(axis))
+    # the axis is a static detector coordinate, so the edges are built in numpy, before any traced operation
+    axis = np.ravel(np.asarray(axis))
     if axis.size == 1:
         return jnp.ones((1, 1), dtype=axis.dtype)
 
-    edges = _bin_edges_from_centers(axis)
+    edges = jnp.asarray(detector_edges_from_centers(axis))
     output_lower = edges[:-1, None]
     output_upper = edges[1:, None]
     source_lower = edges[None, :-1]
@@ -119,8 +106,9 @@ def _uniform_gaussian_bin_kernel(axis, stddev):
 def _apply_uniform_gaussian_bin_response(axis, values, stddev, *, response_axis):
     """Apply the exact uniform-bin Gaussian response without a dense matrix.
 
-    Spectral axes in the forward model are generated by ``linspace`` and are therefore
-    uniform. For modest grids direct convolution avoids FFT overhead; fine legacy 1D
+    Spectral axes in the forward model are built by ``oversampled_axis`` from evenly spaced
+    pixel centers and are therefore uniform; a non-uniform axis is not supported here. For
+    modest grids direct convolution avoids FFT overhead; fine legacy 1D
     grids use FFT convolution, reducing the response storage from quadratic to linear.
     """
 
@@ -274,8 +262,8 @@ def add_ion_IRF(irf: SpectrometerIRF, lamAxisI, modlI, amps, TSins) -> Tuple[jnp
         #print(f"ThryI max {jnp.max(ThryI)}")
         #print(f"amps max {jnp.max(amps)}")
 
+        lamAxisI = jnp.average(lamAxisI.reshape(irf.n_spectral_pixels, -1), axis=1)
         if irf.normalize == 0:
-            lamAxisI = jnp.average(lamAxisI.reshape(irf.n_spectral_pixels, -1), axis=1)
             ThryI = TSins["general"]["amp3"] * amps * ThryI / _reference_amplitude(lamAxisI, ThryI, irf.amp_ranges)
             # lamAxisE = jnp.average(lamAxisE.reshape(irf.n_spectral_pixels, -1), axis=1)
     else:
@@ -319,8 +307,8 @@ def add_electron_IRF(irf: SpectrometerIRF, lamAxisE, modlE, amps, TSins) -> Tupl
         )
 
     ThryE = jnp.average(ThryE.reshape(irf.n_spectral_pixels, -1), axis=1)
+    lamAxisE = jnp.average(lamAxisE.reshape(irf.n_spectral_pixels, -1), axis=1)
     if irf.normalize == 0:
-        lamAxisE = jnp.average(lamAxisE.reshape(irf.n_spectral_pixels, -1), axis=1)
         ThryE = amps * ThryE / _reference_amplitude(lamAxisE, ThryE, irf.amp_ranges)
         ThryE = jnp.where(
             lamAxisE < TSins["general"]["lam"], TSins["general"]["amp1"] * ThryE, TSins["general"]["amp2"] * ThryE

tsadar/core/physics/form_factor.py

@@ -15,6 +15,7 @@
 from jax import checkpoint
 
 from . import ratintn
+from ...utils.detector_grid import oversampled_axis
 from .interpolation import interp_uniform
 from ...utils.vector_tools import vsub, vdot, vdiv
 
@@ -119,8 +120,8 @@ class FormFactor:
     scattering structure factor or spectral density function, supporting both 1D and 2D electron distribution
     functions (EDFs), multiple plasma conditions, and scattering angles.
     Args:
-        lambda_range (list): Starting and ending wavelengths over which to calculate the spectrum.
-        npts (int): Number of wavelength points to use in the calculation.
+        lambda_range (list): Centers of the first and last detector pixels the spectrum is calculated for.
+        npts (int): Number of wavelength points to use in the calculation, points_per_pixel per pixel.
         lam_shift (float): Wavelength shift to apply.
         scattering_angles (dict): Dictionary containing scattering angles (in degrees).
         num_grad_points (int): Number of gradient points for plasma parameter profiles.
@@ -129,6 +130,7 @@ class FormFactor:
         n_beta (int): Number of angles at which the 2D EDF's projection is tabulated. Trades accuracy for
             speed in the 2D calculation; 0 disables the tabulation and rotates the EDF exactly at every
             evaluation point instead.
+        points_per_pixel (int): Number of wavelength points inside each detector pixel.
     Attributes:
         C (float): Speed of light in cm/s.
         Me (float): Electron mass in keV/C^2.
@@ -206,6 +208,7 @@ def __init__(
         va_ang,
         calc_gain,
         n_beta=DEFAULT_N_BETA,
+        points_per_pixel=1,
     ):
 
         # basic quantities
@@ -217,7 +220,15 @@ def __init__(
         self.h = 0.01
         minmax = 8.2
         h1 = 1024  # 1024
-        lamAxis = jnp.linspace(lambda_range[0], lambda_range[1], npts)
+        # lambda_range is the first and last pixel centers; each pixel holds points_per_pixel grid points
+        if npts % points_per_pixel:
+            raise ValueError(f"npts={npts} must be a multiple of points_per_pixel={points_per_pixel}.")
+        n_pixels = npts // points_per_pixel
+        if n_pixels == 1:
+            lamAxis = jnp.linspace(lambda_range[0], lambda_range[1], npts)
+        else:
+            pixel_centers = np.linspace(lambda_range[0], lambda_range[1], n_pixels)
+            lamAxis = jnp.asarray(oversampled_axis(pixel_centers, points_per_pixel))
         self.lambda_axis_nm = lamAxis
         self.omgL_num = 2 * jnp.pi * 1e7 * self.C
         omgs = 2e7 * jnp.pi * self.C / lamAxis  # Scattered frequency axis(1 / sec)