diff --git a/NOTES.md b/NOTES.md index fa6a35735..eb747e687 100644 --- a/NOTES.md +++ b/NOTES.md @@ -40,3 +40,46 @@ Append-only record of numerical-physics investigations and the decisions needed - Across five scan-grid phases, full root-aware quadrature had 0.0393–0.0399% detector L1 error and 0.0569–0.1035% loss-gradient L2 error. The first-order local epsilon/numerator residue correction had 0.0138–0.3162% detector error but 15.5–210.5% gradient error. - Corrected CPU medians (100 interleaved blocking calls): full/local warm values were 3.721/3.722 ms and value-plus-gradient calls were 3.805/2.622 ms. Cold value calls were 0.602/0.726 s; cold value-plus-gradient calls were 2.290/2.467 s. There is no useful forward speedup, and the faster local backward pass is numerically wrong. - Consequence: implement only the full root-aware quadrature in production. Record the local-correction prototype as a rejected benchmark rather than exposing an inference-unsafe method through the API. + +## 2026-08-30 — RESULT: detector-edge and IRF pipeline audit for #134 + +- ARTS data preparation averages raw wavelength pixels into resolution-unit mean densities. The calibrated wavelength arrays retained by plotting and loss code are bin centers; exact grouped outer edges are now stored separately as `other.detector_specs.electron_wavelength_edges`, including the true width of a ragged final group. +- The legacy spectral and angular Gaussian samples are now normalized to discrete unit sum, and the hidden post-convolution peak restoration was removed. An unresolved contained delta-line test preserves area to `1e-12` in float64, while its peak is broadened rather than restored. +- The ATS IRF can now apply angular blur without a second spectral blur when the root-aware quadrature has already returned spectral detector-bin means. Nineteen focused detector/IRF/vectorization tests passed on CPU. +- Remaining limitation observed during the audit: the calibrated ARTS angle axis is nonuniform (adjacent spacing 0.0771–0.1598 degrees), whereas the legacy angular `jnp.convolve` assumes translation in sample index. Discrete normalization fixes grid-dependent amplitude scaling but does not by itself make angular blur exactly coordinate-aware. + +## 2026-08-30 — RESULT: production root-aware detector quadrature for #134 + +- The production kernel finds sign-changing roots of `Re(epsilon)` on a fixed wavelength scan, differentiates each bisection root implicitly, tan-maps Gauss--Legendre nodes through the local complex-linearized resonance, and integrates exact Gaussian-CDF probability mass into nonuniform detector bins. No collision term or fixed dielectric floor was added. +- Every gradient/scattering-angle geometry is integrated independently before aperture weighting. A memory-bounded `lax.map` avoids materializing the complete angle-by-node-by-detector tensor. Downstream ARTS processing applies only the angular IRF and angular resolution reduction; it does not convolve or bin wavelength a second time. +- Fatal numerical conditions (root-capacity overflow, zero local width, non-finite evaluation, or invalid edges/IRF/source bounds/scan phase) poison detector bins with NaN and are also exposed through fixed-shape diagnostics. +- In the physical 22-degree CPU regression, the 256-panel result varied by at most `2.10e-6` in detector relative L1 and `2.06e-5` in drift-gradient relative L2 across three scan phases. Refining 128 to 256 panels changed values by at most `1.17e-5` and gradients by at most `1.06e-4`, well inside the declared 1% tolerance. +- Synthetic root benchmarks preserve a constant density and a contained unresolved line's area, support nonuniform detector edges, and match transformed reference values and implicit gradients. Focused production-plumbing tests verify geometry ordering, aperture weighting, calibrated centers/edges, and that no second spectral IRF or spectral reduction occurs. + +## 2026-08-30 — RESULT: physical-coordinate spectral and angular IRFs for #134 + +- This entry supersedes the nonuniform-angular-axis limitation recorded above. On the real 1024-pixel ARTS angular calibration, the legacy index-space convolution gave a centered unresolved line physical area of `1.04699` and shifted its centroid by about `0.084` degrees; its effective physical width also varied with the local pixel spacing. Even-length uniform spectral grids had an additional half-sample centering error. +- Spectral and angular blur now use an analytic bin-to-bin Gaussian response matrix. Each element exactly integrates the unit-area Gaussian over a finite source cell and output cell on the calibrated physical coordinate, then divides by the output-cell width to return a bin-mean density. This supports even-length and nonuniform axes without peak restoration or an index-spacing approximation. +- CPU regressions on the real nonuniform angular calibration preserve a centered line's physical area to `2e-11` and centroid to `1e-3` degrees. Uniform 400- and 801-point spectral grids preserve line area to `2e-12` and centroid to `2e-11`; constant densities remain unity away from finite-detector boundaries. + +## 2026-08-30 — RESULT: full-range root-scan and refinement validation for #134 + +- The representative 449–670 nm, 22-degree ARTS2D spectrum has eight sign-changing roots of `Re(epsilon)` at `[474.201875844894, 514.081647830358, 526.276288254028, 526.451996163684, 526.548008175237, 526.723802255372, 539.209067615960, 591.042039725023]` nm. The closest pair is separated by `0.0960120` nm; a root capacity of four was therefore both insufficient and capable of returning a poisoned model over the real detector range. +- Across 39 scan phases from `-0.95` through `0.95`, 2048 root-scan panels missed a close pair in six phases. Both 4096 and 8192 panels found all eight roots at every phase; their maximum solved-root phase spreads were `1.36e-12` and `6.82e-13` nm. Warm single-geometry medians were `23.5`, `32.2`, and `55.8` ms for 2048, 4096, and 8192 panels. The correctness-first default is therefore 4096 root-scan panels, independent of a 256-panel regular integration grid, with static capacity for 16 roots. +- A physical narrow-line regression comparing 64 integration panels against 128 and 256 panels at three phases had maximum detector relative-L1 changes of `4.94e-5` and `3.96e-5`; maximum drift-gradient relative changes were `1.86e-4` and `2.64e-4`. The 256-panel phase spreads were `2.10e-6` for values and `2.06e-5` for gradients, all well below the one-percent acceptance threshold. +- The full-range regression includes phases at which 2048 panels failed, requires exactly the eight roots above, and verifies finite nonzero reverse-mode drift gradients. Its first version exposed NaN cotangents from statically evaluated inactive tan-map candidates; replacing only those inactive evaluations with finite full-panel sentinel mappings fixed the VJP without adding an epsilon floor or changing any active resonance segment. +- Data-fitting mode carries exact calibrated detector edges, including ragged resolution groups. Forward-only decks have no edge calibration, so `forward_epw_start`/`forward_epw_end` retain their documented legacy role as the first and last wavelength centers; finite outer edges are inferred by half-spacing extrapolation rather than silently reinterpreting those inputs as detector boundaries. +- The analytic response is stored densely only for the modest nonuniform angular detector axis. Uniform spectral grids use the identical Toeplitz bin-integrated Gaussian kernel, with FFT convolution above 2048 points; this keeps the legacy 10,240-point 1D path memory-linear instead of materializing an approximately 800 MiB float64 response plus intermediates. Dense-versus-Toeplitz regressions agree to `5e-11`, and the 10,240-point 1D forward snapshot remains green. + +## 2026-08-30 — RESULT: final #134 validation + +- The local-correction comparison and root-scan timings were measured on Apple arm64 with Python 3.14 and JAX 0.9.1. The production conclusion is unchanged: local correction offered no forward speedup and had unacceptable loss-gradient errors, while the full root-aware method remained below the one-percent value/gradient target. +- The stored ARTS1D golden spectrum was regenerated because it encoded the removed index-space convolution. Replaying the legacy response matched the old snapshot to `5.72e-4` maximum absolute error; the physical-coordinate result differed by `0.104201` at a sharp EPW point. Independent delta-response checks showed that the legacy angular kernel shifted that point by `0.1138` degrees and broadened a requested `0.42466`-degree sigma to `0.5848` degrees, while the analytic response shifted it by `0.0020` degrees and produced `0.4276` degrees after finite-bin integration. +- Final local CPU validation, with the optional HDF4 dependency installed and MLflow redirected to a file store, completed with `186 passed, 5 skipped` in `196.46` seconds. The five skips are GPU-only tests; the only warning is MLflow's filesystem-backend deprecation. `git diff --check` and Python byte-compilation also passed. + +## 2026-08-30 — RESULT: #134 GPU gradient-test audit + +- The first manually dispatched stacked-branch workflow exposed two independent runner outcomes: the hosted CPU runner was recycled at 64% after 122 passing tests and one skip (exit 143, no assertion or traceback), while the GPU runner completed the suite and reported that the new scalar-drift detector Jacobian regressions attempted 12.5 GB and 16.0 GB allocations. +- The GPU allocations came from using reverse-mode `jacrev` for one scalar input and hundreds of detector-bin outputs. This orientation repeats the reverse pass for every output and does not represent the scalar-loss gradient used by inference. +- The regression now forms the complete scalar-to-detector Jacobian with one forward-mode `jacfwd` tangent evaluation, and independently checks a weighted scalar detector loss with reverse-mode `grad` against the forward-Jacobian contraction. This preserves elementwise tangent coverage and production-style reverse-mode coverage without the artificial quadratic memory cost. +- The focused CPU file remained green (`3 passed`) and fell from `81.72` to `63.87` seconds; the full-range eight-root test fell from `24.46` to `12.87` seconds. diff --git a/configs/arts-2d/defaults.yaml b/configs/arts-2d/defaults.yaml index 1f9b245e6..937592e74 100644 --- a/configs/arts-2d/defaults.yaml +++ b/configs/arts-2d/defaults.yaml @@ -83,7 +83,7 @@ other: calc_sigmas: false crop_window: 1 expandedions: false - extraoptions: + extraoptions: {} flatbg: 0 gain: 1 iawfilter: @@ -99,6 +99,27 @@ other: # ~1e-2 in relative L2 on the gradient w.r.t. the EDF -- the gradient is the binding # constraint, so raise this if a 2D fit struggles to converge. n_beta: 1024 + # Integrate the collisionless ARTS2D spectrum into finite detector bins before + # angular postprocessing. A fine fixed scan finds dielectric roots independently of + # the coarser integration grid; root neighborhoods use a tan mapping so resonances much + # narrower than a wavelength pixel retain their area and differentiable position. The + # spectral Gaussian IRF is included analytically through bin-edge CDF differences and + # must not be applied again downstream. Enabled iawfilter edges become exact source-side + # integration breakpoints before that IRF is applied. + resonance_quadrature: + enabled: true + root_scan_panels: 4096 + integration_panels: 256 + regular_order: 8 + root_order: 32 + max_roots: 16 + neighbor_panels: 1 + bisection_iterations: 48 + tail_sigma: 6.0 + scan_phase: 0.0 + # Number of gradient/angle spectra evaluated concurrently. One bounds memory most + # tightly; raise cautiously on a device with enough memory. + map_batch_size: 1 points_per_pixel: 1 refit: true refit_thresh: 0.25 diff --git a/configs/arts-2d/inputs.yaml b/configs/arts-2d/inputs.yaml index 18a2df8d8..4145c46b7 100644 --- a/configs/arts-2d/inputs.yaml +++ b/configs/arts-2d/inputs.yaml @@ -21,7 +21,7 @@ mlflow: experiment: inverse-thomson-scattering run: multiplex_arts_rot90_nosmoothtest_nonete_1 other: - PhysParams: + detector_specs: widIRF: ang_FWHM_ele: 1.0 spect_FWHM_ele: 0.9 diff --git a/docs/source/defaults.rst b/docs/source/defaults.rst index 737d7d437..77cd64beb 100644 --- a/docs/source/defaults.rst +++ b/docs/source/defaults.rst @@ -214,9 +214,9 @@ The ``data:`` section contains the specifics on which shot and what region of th - ``iaw_cf_max`` ending wavelength for a central feature in the IAW that is to be excluded from analysis in nm, must be larger than ``iaw_cf_min`` - - ``forward_epw_start`` starting wavelength in nm for the EPW calculation for forward model only - - - ``forward_epw_end`` ending wavelength in nm for the EPW calculation for forward model only + - ``forward_epw_start`` first wavelength center in nm for the EPW calculation in forward mode. For detector-integrated ARTS2D, the outer detector edge is inferred by half-spacing extrapolation. + + - ``forward_epw_end`` last wavelength center in nm for the EPW calculation in forward mode. For detector-integrated ARTS2D, the outer detector edge is inferred by half-spacing extrapolation. - ``forward_iaw_start`` starting wavelength in nm for the IAW calculation for forward model only @@ -317,7 +317,9 @@ The ``other:`` section includes options specifying the types of data that are be - ``gain`` CCD counts per photo-electron; the standard OMEGA ROSS has a gain of 144. Gain must be accurate for appropriate use of Poisson statistics but the gain is generaly not important for the fitting process as the data is normalized by default. -- ``points_per_pixel`` number of wavelength points computed in the spectrum per pixel in the data being analyzed, for most cases 1 is sufficient but if the peaks in the data are very narrow then it may be necessary to use a value larger than 1 to ensure the peaks are well resolved in the computed spectrum, due to the scaling behaviour its not recommended to use a value larger than 10 here. +- ``points_per_pixel`` number of wavelength points computed per detector pixel by the legacy sampled-spectrum path. ARTS2D does not use this setting to resolve a narrow resonance: it integrates the continuous spectrum directly into the calibrated detector-bin edges. + +- ``resonance_quadrature`` controls the ARTS2D detector-bin integration and is enabled by default. ``root_scan_panels`` (4096 by default) is the fine grid used only to bracket sign-changing zeros of the real dielectric, while ``integration_panels`` (256 by default) controls the coarser regular quadrature grid. Separating them resolves closely spaced roots without paying the fine-grid cost in the detector response matrix. ``regular_order`` and the even ``root_order`` set the Gauss--Legendre rules away from and near a root; ``neighbor_panels`` expands the root-mapped neighborhood; ``max_roots`` is the fixed root capacity (16 by default); ``bisection_iterations`` controls the differentiable root solve; and ``tail_sigma`` extends the source integration domain beyond the detector by that many spectral-IRF standard deviations. When ``iawfilter`` is enabled, every filter edge strictly inside that source domain is inserted as an exact integration breakpoint: attenuation is applied to the continuous source spectrum before Gaussian spectral blur and detector integration, including when a filter cuts through a detector bin. ``iawoff`` remains a detector-space mask. ``scan_phase`` is intended for convergence tests and must lie strictly between -1 and 1. ``map_batch_size`` trades device parallelism for peak memory across scattering geometries. A detected root overflow, zero-width resonance, invalid breakpoint, invalid detector geometry, or non-finite evaluation produces a non-finite model rather than silently returning a partial integral. - ``ang_res_unit`` is the number of pixels in an angular resolution unit for ARTS diff --git a/docs/source/forward_pass.rst b/docs/source/forward_pass.rst index 7dc9ea5e2..a3e7e834e 100644 --- a/docs/source/forward_pass.rst +++ b/docs/source/forward_pass.rst @@ -13,13 +13,14 @@ Other - ``extraoptions`` - - ``spectype`` the type of spectrum to be computed. This field is self determined from the data when fitting. For a forward pass somthing has to be specified but it deos not effect the spectrum. Options are ``temporal`` or ``imaging`` In this context they produce the same spectrum. + - ``spectype`` the type of spectrum to be computed. This field is determined from the data when fitting, but must be supplied for a forward pass. Options include ``temporal``, ``imaging``, and ``angular``. - - ``detector_specs`` the subfields define instrumental properties +- ``detector_specs`` defines instrumental properties and is a sibling of ``extraoptions`` under ``other``. - - ``widIRF`` the subfields define the instrumental response functions, when fitting data it is determined from the calibration files but when running in forward mode these must be supplied. The subfields are: + - ``widIRF`` defines the instrumental response widths. They are read from calibration files when fitting data and must be supplied in forward mode. Its subfields are: - - ``spect_std_ion`` the standard deviation of the gaussian ion instrumental response function in nanometers + - ``spect_stddev_ion`` the standard deviation of the Gaussian ion spectral response in nanometers. - - ``spect_std_ele`` the standard deviation of the gaussian electron instrumental response function in nanometers + - ``spect_stddev_ele`` the standard deviation of the Gaussian electron spectral response in nanometers for non-ARTS spectra. + - ``spect_FWHM_ele`` and ``ang_FWHM_ele`` the spectral width in nanometers and angular width in degrees, respectively, for ARTS. These two values are full widths at half maximum. diff --git a/test_angular_2v.py b/test_angular_2v.py index 39f91e583..5b6118165 100644 --- a/test_angular_2v.py +++ b/test_angular_2v.py @@ -62,7 +62,11 @@ def test_arts2d_forward_pass(): sas = get_scattering_angles(config) [axisxE, _, _, _, _, _] = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" diff --git a/test_arts2v_random.py b/test_arts2v_random.py index 445781c18..4c59119f6 100644 --- a/test_arts2v_random.py +++ b/test_arts2v_random.py @@ -166,7 +166,11 @@ def test_arts2d_inverse(config_path: str = "tests/configs/arts2d_test_defaults.y sas = get_scattering_angles(config) sas["angAxis"], _, _, _, _, _ = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" diff --git a/tests/test_form_factor/test_arts2d_spectral_terms.py b/tests/test_form_factor/test_arts2d_spectral_terms.py new file mode 100644 index 000000000..f97e6bd9b --- /dev/null +++ b/tests/test_form_factor/test_arts2d_spectral_terms.py @@ -0,0 +1,152 @@ +"""Focused contracts for reusable ARTS2D wavelength-space spectral terms.""" + +import numpy as np + +from jax import config, grad, jit, numpy as jnp + +config.update("jax_enable_x64", True) + +from tsadar.core.physics.form_factor import FormFactor + + +def _form_factor(): + return FormFactor( + lambda_range=[516.0, 524.0], + npts=7, + lam_shift=0.0, + scattering_angles={ + "sa": np.array([38.0, 61.0]), + "weights": np.ones((1, 2)), + }, + num_grad_points=2, + ud_ang=17.0, + va_ang={"ion-1": -23.0}, + calc_gain={"calc": False}, + n_beta=8, + ) + + +def _params(ud=0.13): + vx = jnp.linspace(-7.0, 7.0, 33) + grid_x, grid_y = jnp.meshgrid(vx, vx) + fe = jnp.exp( + -0.5 * ((grid_x - 0.2) / 0.85) ** 2 + - 0.5 * ((grid_y + 0.1) / 1.15) ** 2 + ) + fe /= jnp.sum(fe) * (vx[1] - vx[0]) ** 2 + return { + "electron": {"ne": 0.24, "Te": 0.62, "fe": fe, "v": vx}, + "general": { + "ne_gradient": 4.0, + "Te_gradient": 6.0, + "lam": 526.5, + "ud": ud, + }, + "ion-1": { + "A": 1.0, + "Z": 1.0, + "Ti": 0.09, + "fract": 1.0, + "Va": 0.08, + }, + } + + +def test_spectral_terms_reconstruct_calc_in_2D_and_support_arbitrary_nodes(): + ff = _form_factor() + params = _params() + sinogram = ff.prepare_2D_sinogram(params) + + numerator, epsilon = ff.calc_2D_spectral_terms( + params, + ff.lambda_axis_nm, + sinogram=sinogram, + ) + spectrum, wavelengths_cm = ff.calc_in_2D(params) + + assert numerator.shape == (ff.npts, ff.num_grad_points, 2) + assert epsilon.shape == numerator.shape + np.testing.assert_allclose( + np.asarray(jnp.transpose(numerator / jnp.abs(epsilon) ** 2, (1, 0, 2))), + np.asarray(spectrum), + rtol=2e-14, + atol=0, + ) + np.testing.assert_allclose( + np.squeeze(np.asarray(wavelengths_cm)) * 1.0e7, + np.asarray(ff.lambda_axis_nm), + rtol=1e-15, + atol=0, + ) + + arbitrary_nodes = jnp.array([517.25, 519.875, 523.1]) + single_angle_terms = ff.calc_2D_spectral_terms( + params, + arbitrary_nodes, + sinogram=sinogram, + scattering_angles=47.0, + ) + assert single_angle_terms[0].shape == (3, ff.num_grad_points, 1) + assert bool(jnp.all(jnp.isfinite(single_angle_terms[0]))) + assert bool(jnp.all(jnp.isfinite(single_angle_terms[1]))) + + +def test_pointwise_nodes_match_common_grid_and_gradient_contract(): + ff = _form_factor() + nodes = jnp.array([517.4, 520.2, 522.7]) + + def reconstructed_spectrum(ud, pointwise): + params = _params(ud) + sinogram = ff.prepare_2D_sinogram(params) + if pointwise: + node_mesh = jnp.broadcast_to( + nodes[:, None, None], + (nodes.size, ff.num_grad_points, 2), + ) + numerator, epsilon = ff.calc_2D_spectral_terms_at_points( + params, + node_mesh, + sinogram=sinogram, + ) + else: + numerator, epsilon = ff.calc_2D_spectral_terms( + params, + nodes, + sinogram=sinogram, + ) + return jnp.sum(numerator / jnp.abs(epsilon) ** 2) + + common_value = reconstructed_spectrum(0.13, False) + pointwise_value = reconstructed_spectrum(0.13, True) + common_gradient = grad(reconstructed_spectrum, argnums=0)(0.13, False) + pointwise_gradient = grad(reconstructed_spectrum, argnums=0)(0.13, True) + + assert bool(jnp.all(jnp.isfinite(jnp.asarray([common_value, common_gradient])))) + np.testing.assert_allclose( + np.asarray(pointwise_value), np.asarray(common_value), rtol=0, atol=0 + ) + np.testing.assert_allclose( + np.asarray(pointwise_gradient), np.asarray(common_gradient), rtol=0, atol=0 + ) + + +def test_prepared_evaluator_builds_sinogram_only_once(monkeypatch): + ff = _form_factor() + params = _params() + build_sinogram = ff._build_sinogram + calls = 0 + + def counted_build(vx, fe): + nonlocal calls + calls += 1 + return build_sinogram(vx, fe) + + monkeypatch.setattr(ff, "_build_sinogram", counted_build) + evaluate = ff.prepare_2D_spectral_evaluator(params, scattering_angles=49.0) + assert calls == 1 + + compiled_evaluate = jit(evaluate) + numerator, epsilon = compiled_evaluate(jnp.array([518.0, 520.0])) + assert numerator.shape == epsilon.shape == (2, ff.num_grad_points, 1) + evaluate(jnp.array([519.0, 521.0, 523.0])) + assert calls == 1 diff --git a/tests/test_form_factor/test_resonance_quadrature.py b/tests/test_form_factor/test_resonance_quadrature.py new file mode 100644 index 000000000..dcd5b83ad --- /dev/null +++ b/tests/test_form_factor/test_resonance_quadrature.py @@ -0,0 +1,508 @@ +"""Focused tests for unresolved-resonance detector quadrature.""" + +import numpy as np +import pytest +from scipy.integrate import quad +from scipy.special import ndtr as scipy_ndtr + +from jax import config + +config.update("jax_enable_x64", True) + +import jax +import jax.numpy as jnp + +from tsadar.core.physics.resonance_quadrature import ( + gaussian_bin_probabilities, + integrate_detector_bins, + raise_for_diagnostics, +) + + +ROOT_NM = 474.199678891 +HWHM_NM = 3.49399e-4 +IRF_SIGMA_NM = 0.035 +DETECTOR_EDGES_NM = np.array([473.72, 474.02, 474.15, 474.225, 474.41, 474.76]) +SOURCE_BOUNDS_NM = np.array( + [DETECTOR_EDGES_NM[0] - 6 * IRF_SIGMA_NM, DETECTOR_EDGES_NM[-1] + 6 * IRF_SIGMA_NM] +) + + +def _lorentz_terms(wavelengths_nm, root_nm, half_width_nm=HWHM_NM): + numerator = jnp.full_like(wavelengths_nm, half_width_nm / jnp.pi) + epsilon = wavelengths_nm - root_nm + 1j * half_width_nm + return numerator, epsilon + + +def _cauchy_detector_reference(root_nm, half_width_nm=HWHM_NM): + """High-accuracy reference using the exact Cauchy tan transformation.""" + + t_lower = np.arctan((SOURCE_BOUNDS_NM[0] - root_nm) / half_width_nm) + t_upper = np.arctan((SOURCE_BOUNDS_NM[1] - root_nm) / half_width_nm) + integrals = [] + for edge_lower, edge_upper in zip(DETECTOR_EDGES_NM[:-1], DETECTOR_EDGES_NM[1:]): + + def transformed_integrand(t): + wavelength = root_nm + half_width_nm * np.tan(t) + probability = scipy_ndtr((edge_upper - wavelength) / IRF_SIGMA_NM) - scipy_ndtr( + (edge_lower - wavelength) / IRF_SIGMA_NM + ) + # (gamma / pi) / ((lambda-root)^2 + gamma^2) * d lambda / dt = 1 / pi. + return probability / np.pi + + integral = quad( + transformed_integrand, + t_lower, + t_upper, + epsabs=2e-13, + epsrel=2e-13, + limit=300, + )[0] + integrals.append(integral) + return np.asarray(integrals) / np.diff(DETECTOR_EDGES_NM) + + +def _unresolved_result(root_nm, *, phase=0.0, panels=256): + return integrate_detector_bins( + lambda wavelengths: _lorentz_terms(wavelengths, root_nm), + jnp.asarray(DETECTOR_EDGES_NM), + IRF_SIGMA_NM, + source_bounds_nm=jnp.asarray(SOURCE_BOUNDS_NM), + root_scan_panels=1024, + integration_panels=panels, + regular_order=8, + root_order=32, + max_roots=4, + neighbor_panels=1, + bisection_iterations=48, + scan_phase=phase, + ) + + +def _relative_detector_l1(actual, expected): + bin_widths = np.diff(DETECTOR_EDGES_NM) + return np.sum(np.abs(actual - expected) * bin_widths) / np.sum(np.abs(expected) * bin_widths) + + +def test_gaussian_bin_probabilities_are_exact_cdf_differences(): + wavelengths = np.array([473.9, 474.2, 474.7]) + actual = gaussian_bin_probabilities( + jnp.asarray(wavelengths), jnp.asarray(DETECTOR_EDGES_NM), IRF_SIGMA_NM + ) + expected = scipy_ndtr( + (DETECTOR_EDGES_NM[None, 1:] - wavelengths[:, None]) / IRF_SIGMA_NM + ) - scipy_ndtr((DETECTOR_EDGES_NM[None, :-1] - wavelengths[:, None]) / IRF_SIGMA_NM) + np.testing.assert_allclose(np.asarray(actual), expected, rtol=2e-14, atol=2e-14) + + +def test_bin_means_preserve_area_for_nonuniform_detector_edges(): + edges = jnp.asarray([0.0, 0.17, 0.83, 1.9, 2.4]) + + result = integrate_detector_bins( + lambda wavelengths: (jnp.ones_like(wavelengths), jnp.ones_like(wavelengths)), + edges, + 0.04, + root_scan_panels=96, + integration_panels=96, + ) + + raise_for_diagnostics(result) + # A constant source convolved with a normalized Gaussian remains constant; + # the default six-sigma truncation loses less than four parts in 1e9. + np.testing.assert_allclose(np.asarray(result.bin_mean), 1.0, rtol=0, atol=4e-9) + area = jnp.sum(result.bin_mean * jnp.diff(edges)) + np.testing.assert_allclose(np.asarray(area), 2.4, rtol=0, atol=5e-9) + + +def test_unresolved_root_is_phase_and_refinement_stable(): + reference = _cauchy_detector_reference(ROOT_NM) + phases = jnp.arange(8, dtype=jnp.float64) / 8 + + def evaluate(phase): + result = _unresolved_result(ROOT_NM, phase=phase, panels=256) + return result.bin_mean, result.diagnostics + + values, diagnostics = jax.jit(jax.vmap(evaluate))(phases) + values = np.asarray(values) + + assert np.all(np.asarray(diagnostics.root_count) == 1) + assert not np.any(np.asarray(diagnostics.root_overflow)) + assert not np.any(np.asarray(diagnostics.nonfinite)) + assert not np.any(np.asarray(diagnostics.zero_width)) + np.testing.assert_allclose( + np.asarray(diagnostics.roots_nm)[:, 0], ROOT_NM, rtol=0, atol=3e-12 + ) + np.testing.assert_allclose( + np.asarray(diagnostics.resonance_centers_nm)[:, 0], ROOT_NM, rtol=0, atol=3e-12 + ) + np.testing.assert_allclose( + np.asarray(diagnostics.resonance_half_widths_nm)[:, 0], HWHM_NM, rtol=3e-10, atol=2e-13 + ) + + errors = np.asarray([_relative_detector_l1(value, reference) for value in values]) + assert np.max(errors) < 1.0e-3 + assert np.max(np.ptp(values, axis=0) * np.diff(DETECTOR_EDGES_NM)) < 2.0e-5 + + coarse = np.asarray(_unresolved_result(ROOT_NM, phase=0.375, panels=128).bin_mean) + fine = np.asarray(_unresolved_result(ROOT_NM, phase=0.375, panels=256).bin_mean) + assert _relative_detector_l1(coarse, fine) < 5.0e-4 + + +def test_implicit_root_and_integral_gradients_match_reference(): + bin_weights = jnp.asarray([0.2, -0.4, 0.7, -0.1, 0.3]) * jnp.asarray( + np.diff(DETECTOR_EDGES_NM) + ) + + def loss(root_nm, phase): + result = _unresolved_result(root_nm, phase=phase, panels=256) + return jnp.sum(result.bin_mean * bin_weights) + + phases = jnp.asarray([0.0, 0.375, 0.75]) + gradients = np.asarray(jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0)))(ROOT_NM, phases)) + + delta = 2.0e-6 + + def reference_loss(root_nm): + return np.sum(_cauchy_detector_reference(root_nm) * np.asarray(bin_weights)) + + reference_gradient = (reference_loss(ROOT_NM + delta) - reference_loss(ROOT_NM - delta)) / ( + 2 * delta + ) + np.testing.assert_allclose(gradients, reference_gradient, rtol=3e-3, atol=2e-5) + assert np.ptp(gradients) < 3e-3 * abs(reference_gradient) + 2e-5 + + root_gradient = jax.grad(lambda root: _unresolved_result(root).diagnostics.roots_nm[0])( + ROOT_NM + ) + np.testing.assert_allclose(np.asarray(root_gradient), 1.0, rtol=2e-12, atol=2e-12) + + +def test_no_root_path_is_jittable_vmap_compatible_and_supports_components(): + edges = jnp.asarray([-1.0, -0.1, 0.25, 1.0]) + + def one(amplitude): + def terms(wavelengths): + shape = jnp.exp(-(wavelengths / 0.4) ** 2) + numerator = jnp.stack((amplitude * shape, 2 * amplitude * shape), axis=-1) + return numerator, jnp.ones_like(wavelengths) + 0.2j + + return integrate_detector_bins( + terms, + edges, + 0.06, + source_bounds_nm=jnp.asarray([-1.4, 1.4]), + root_scan_panels=64, + integration_panels=64, + ) + + result = jax.jit(jax.vmap(one))(jnp.asarray([0.5, 1.0, 1.5])) + assert result.bin_mean.shape == (3, 3, 2) + assert np.all(np.asarray(result.diagnostics.root_count) == 0) + assert not np.any(np.asarray(result.diagnostics.nonfinite)) + np.testing.assert_allclose( + np.asarray(result.bin_mean[..., 1]), 2 * np.asarray(result.bin_mean[..., 0]), rtol=2e-14, atol=2e-14 + ) + np.testing.assert_allclose( + np.asarray(result.bin_mean[1]), 2 * np.asarray(result.bin_mean[0]), rtol=2e-14, atol=2e-14 + ) + + +def test_root_on_scan_edge_is_counted_once(): + result = integrate_detector_bins( + lambda wavelengths: (jnp.ones_like(wavelengths), wavelengths - 0.5 + 0.01j), + jnp.asarray([0.0, 0.4, 0.8, 1.0]), + 0.02, + source_bounds_nm=jnp.asarray([0.0, 1.0]), + root_scan_panels=40, + integration_panels=10, + ) + assert int(result.diagnostics.root_count) == 1 + np.testing.assert_allclose(np.asarray(result.diagnostics.roots_nm[0]), 0.5, rtol=0, atol=1e-14) + + +def test_two_close_roots_share_one_coarse_panel_without_phase_error(): + """A fine root scan may assign multiple resonances to one coarse panel.""" + + edges = jnp.asarray([-0.25, 0.0, 0.02, 0.05, 0.25]) + source_bounds = jnp.asarray([-0.5, 0.5]) + root_separation = 0.022 + imaginary_epsilon = 4.0e-5 + phases = jnp.asarray([-0.6, 0.0, 0.6]) + + def one(shift, phase, integration_panels): + roots = jnp.asarray([0.012, 0.012 + root_separation]) + shift + + def terms(wavelengths): + real_epsilon = jnp.prod(wavelengths[:, None] - roots[None, :], axis=1) + numerator = jnp.full_like( + wavelengths, imaginary_epsilon * root_separation / jnp.pi + ) + return numerator, real_epsilon + 1j * imaginary_epsilon + + return integrate_detector_bins( + terms, + edges, + 0.018, + source_bounds_nm=source_bounds, + root_scan_panels=1024, + integration_panels=integration_panels, + regular_order=16, + root_order=48, + max_roots=4, + neighbor_panels=1, + scan_phase=phase, + ) + + def values_and_gradients(integration_panels): + def evaluate(phase): + result = one(0.0, phase, integration_panels) + gradient = jax.jacrev(lambda shift: one(shift, phase, integration_panels).bin_mean)( + 0.0 + ) + return result, gradient + + return jax.jit(jax.vmap(evaluate))(phases) + + coarse, coarse_gradient = values_and_gradients(16) + fine, fine_gradient = values_and_gradients(128) + coarse_values = np.asarray(coarse.bin_mean) + fine_values = np.asarray(fine.bin_mean) + coarse_gradient = np.asarray(coarse_gradient) + fine_gradient = np.asarray(fine_gradient) + + assert np.all(np.asarray(coarse.diagnostics.root_count) == 2) + assert np.all(np.asarray(fine.diagnostics.root_count) == 2) + # At phase zero both roots lie in the same 0.0625-nm coarse panel. + coarse_panel_ids = np.floor((np.array([0.012, 0.034]) + 0.5) / (1.0 / 16)).astype(int) + assert coarse_panel_ids[0] == coarse_panel_ids[1] + for phase_index in range(phases.size): + np.testing.assert_allclose( + coarse_values[phase_index], fine_values[phase_index], rtol=3e-3, atol=2e-6 + ) + np.testing.assert_allclose( + coarse_gradient[phase_index], fine_gradient[phase_index], rtol=5e-3, atol=2e-5 + ) + np.testing.assert_allclose( + coarse_values, + np.broadcast_to(coarse_values[1], coarse_values.shape), + rtol=3e-3, + atol=2e-6, + ) + np.testing.assert_allclose( + coarse_gradient, + np.broadcast_to(coarse_gradient[1], coarse_gradient.shape), + rtol=5e-3, + atol=2e-5, + ) + + +def test_integration_breakpoint_exactly_partitions_a_rectangular_transmission(): + edges = np.asarray([-0.6, -0.1, 0.3, 0.7]) + source_bounds = np.asarray([-1.0, 1.0]) + cutoff = 0.137 + sigma = 0.04 + + def terms(wavelengths): + transmission = jnp.where(wavelengths < cutoff, 1.0, 0.2) + return transmission, jnp.ones_like(wavelengths) + 0j + + def evaluate(phase): + return integrate_detector_bins( + terms, + jnp.asarray(edges), + sigma, + source_bounds_nm=jnp.asarray(source_bounds), + root_scan_panels=32, + integration_panels=8, + integration_breakpoints_nm=jnp.asarray([cutoff]), + regular_order=16, + scan_phase=phase, + ) + + phases = jnp.asarray([-0.55, 0.0, 0.55]) + results = jax.jit(jax.vmap(evaluate))(phases) + + reference_integrals = [] + for lower_edge, upper_edge in zip(edges[:-1], edges[1:]): + probability = lambda wavelength: scipy_ndtr( + (upper_edge - wavelength) / sigma + ) - scipy_ndtr((lower_edge - wavelength) / sigma) + integral = quad(probability, source_bounds[0], cutoff, epsabs=2e-13, epsrel=2e-13)[ + 0 + ] + integral += 0.2 * quad( + probability, cutoff, source_bounds[1], epsabs=2e-13, epsrel=2e-13 + )[0] + reference_integrals.append(integral) + reference = np.asarray(reference_integrals) / np.diff(edges) + + assert not np.any(np.asarray(results.diagnostics.invalid_integration_breakpoints)) + np.testing.assert_allclose( + np.asarray(results.bin_mean), + np.broadcast_to(reference, results.bin_mean.shape), + rtol=0, + atol=2e-9, + ) + + +def test_close_breakpoints_share_one_original_panel_exactly_across_phases(): + edges = np.asarray([-0.6, -0.1, 0.3, 0.7]) + source_bounds = np.asarray([-1.0, 1.0]) + filter_lower = 0.031 + filter_upper = 0.047 + attenuation = 0.15 + sigma = 0.04 + breakpoints = np.asarray([filter_lower, filter_upper]) + + # Both boundaries select the same nearest edge of the original eight-panel + # phase-zero grid. They still need two distinct exact integration boundaries. + original_interior_edges = np.linspace(*source_bounds, 9)[1:-1] + nearest_edges = np.argmin( + np.abs(breakpoints[:, None] - original_interior_edges[None, :]), axis=1 + ) + assert nearest_edges[0] == nearest_edges[1] + + def terms(wavelengths): + inside_filter = (wavelengths > filter_lower) & (wavelengths < filter_upper) + transmission = jnp.where(inside_filter, attenuation, 1.0) + return transmission, jnp.ones_like(wavelengths) + 0j + + def evaluate(phase): + return integrate_detector_bins( + terms, + jnp.asarray(edges), + sigma, + source_bounds_nm=jnp.asarray(source_bounds), + root_scan_panels=32, + integration_panels=8, + integration_breakpoints_nm=jnp.asarray(breakpoints), + regular_order=32, + scan_phase=phase, + ) + + phases = jnp.asarray([-0.55, 0.0, 0.55]) + results = jax.jit(jax.vmap(evaluate))(phases) + + reference_integrals = [] + for lower_edge, upper_edge in zip(edges[:-1], edges[1:]): + probability = lambda wavelength: scipy_ndtr( + (upper_edge - wavelength) / sigma + ) - scipy_ndtr((lower_edge - wavelength) / sigma) + integral = quad( + probability, + source_bounds[0], + filter_lower, + epsabs=2e-13, + epsrel=2e-13, + )[0] + integral += attenuation * quad( + probability, + filter_lower, + filter_upper, + epsabs=2e-13, + epsrel=2e-13, + )[0] + integral += quad( + probability, + filter_upper, + source_bounds[1], + epsabs=2e-13, + epsrel=2e-13, + )[0] + reference_integrals.append(integral) + reference = np.asarray(reference_integrals) / np.diff(edges) + + assert not np.any(np.asarray(results.diagnostics.invalid_integration_breakpoints)) + np.testing.assert_allclose( + np.asarray(results.bin_mean), + np.broadcast_to(reference, results.bin_mean.shape), + rtol=0, + atol=2e-9, + ) + + +def test_diagnostics_report_overflow_nonfinite_and_zero_width(): + roots = jnp.asarray([-0.8, -0.4, 0.0, 0.4, 0.8]) + + def too_many_terms(wavelengths): + real_epsilon = jnp.prod(wavelengths[:, None] - roots[None, :], axis=1) + return jnp.ones_like(wavelengths), real_epsilon + 0.02j + + overflow = integrate_detector_bins( + too_many_terms, + jnp.asarray([-0.9, 0.0, 0.9]), + 0.03, + source_bounds_nm=jnp.asarray([-1.0, 1.0]), + root_scan_panels=100, + integration_panels=100, + max_roots=2, + ) + assert int(overflow.diagnostics.root_count) == 5 + assert bool(overflow.diagnostics.root_overflow) + assert np.all(np.isnan(np.asarray(overflow.bin_mean))) + with pytest.raises(ValueError, match="detected 5 roots"): + raise_for_diagnostics(overflow) + + zero_width = integrate_detector_bins( + lambda wavelengths: (jnp.ones_like(wavelengths), wavelengths - 0.13 + 0j), + jnp.asarray([-0.5, 0.0, 0.5]), + 0.04, + source_bounds_nm=jnp.asarray([-0.7, 0.7]), + root_scan_panels=64, + integration_panels=64, + ) + assert bool(zero_width.diagnostics.zero_width) + assert np.all(np.isnan(np.asarray(zero_width.bin_mean))) + with pytest.raises(ValueError, match="zero local half-width"): + raise_for_diagnostics(zero_width) + + nonfinite = integrate_detector_bins( + lambda wavelengths: ( + jnp.ones_like(wavelengths), + jnp.where(wavelengths > 0, jnp.nan + 0j, 1.0 + 0j), + ), + jnp.asarray([-0.5, 0.0, 0.5]), + 0.04, + source_bounds_nm=jnp.asarray([-0.7, 0.7]), + root_scan_panels=32, + integration_panels=32, + ) + assert bool(nonfinite.diagnostics.nonfinite) + assert np.all(np.isnan(np.asarray(nonfinite.bin_mean))) + + invalid_edges = integrate_detector_bins( + lambda wavelengths: (jnp.ones_like(wavelengths), jnp.ones_like(wavelengths) + 0j), + jnp.asarray([0.0, 0.2, 0.2, 0.5]), + 0.0, + scan_phase=1.0, + ) + assert bool(invalid_edges.diagnostics.invalid_bin_width) + assert bool(invalid_edges.diagnostics.invalid_irf_sigma) + assert bool(invalid_edges.diagnostics.invalid_scan_phase) + assert np.all(np.isnan(np.asarray(invalid_edges.bin_mean))) + + @jax.jit + def invalid_breakpoints(breakpoints): + return integrate_detector_bins( + lambda wavelengths: (jnp.ones_like(wavelengths), jnp.ones_like(wavelengths) + 0j), + jnp.asarray([-0.5, 0.0, 0.5]), + 0.04, + source_bounds_nm=jnp.asarray([-0.7, 0.7]), + root_scan_panels=32, + integration_panels=8, + integration_breakpoints_nm=breakpoints, + ) + + bad_breakpoint_sets = ( + [0.2, 0.1], + [0.1, 0.1], + [-0.8, 0.1], + [0.1, np.nan], + ) + bad_breakpoints = None + for breakpoint_set in bad_breakpoint_sets: + bad_breakpoints = invalid_breakpoints(jnp.asarray(breakpoint_set)) + assert bool(bad_breakpoints.diagnostics.invalid_integration_breakpoints) + assert np.all(np.isnan(np.asarray(bad_breakpoints.bin_mean))) + with pytest.raises(ValueError, match="integration breakpoints"): + raise_for_diagnostics(bad_breakpoints) diff --git a/tests/test_forward/ThryE-arts1v.npy b/tests/test_forward/ThryE-arts1v.npy index c06d3c5cf..fe47482c6 100644 Binary files a/tests/test_forward/ThryE-arts1v.npy and b/tests/test_forward/ThryE-arts1v.npy differ diff --git a/tests/test_forward/test_angular_1v.py b/tests/test_forward/test_angular_1v.py index 15eb79de8..3e42206c5 100644 --- a/tests/test_forward/test_angular_1v.py +++ b/tests/test_forward/test_angular_1v.py @@ -95,7 +95,8 @@ def test_arts1d_forward_pass(): fig2.savefig(os.path.join(td, "SpectralDifference.png"), bbox_inches="tight") mlflow.log_artifacts(td) - # Testing criteria was loosened , comparsions showed absolute differences <0.0006 across the spectrum. While it is unclear why this occurred the change is smaller then expected changes from resolution changes. + # The reference was refreshed for analytic physical-coordinate IRF + # centering and removal of the legacy per-row peak restoration. np.testing.assert_allclose(ThryE, ground_truth, atol=1e-3) diff --git a/tests/test_forward/test_angular_2v.py b/tests/test_forward/test_angular_2v.py index ca81df04f..8a699c911 100644 --- a/tests/test_forward/test_angular_2v.py +++ b/tests/test_forward/test_angular_2v.py @@ -62,7 +62,11 @@ def test_arts2d_forward_pass(): sas = get_scattering_angles(config) [axisxE, _, _, _, _, _] = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" diff --git a/tests/test_forward/test_irf_area.py b/tests/test_forward/test_irf_area.py new file mode 100644 index 000000000..315316344 --- /dev/null +++ b/tests/test_forward/test_irf_area.py @@ -0,0 +1,417 @@ +"""Detector-coordinate and area-preserving IRF regressions for ARTS.""" + +from types import SimpleNamespace +from pathlib import Path + +import numpy as np +import pytest +import yaml +from jax import config + +config.update("jax_enable_x64", True) + +from jax import numpy as jnp + +from tsadar.core.instrument.irf import ( + AngularIRF, + _apply_uniform_gaussian_bin_response, + _gaussian_bin_response, + add_ATS_IRF, +) +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.utils.misc import merge_defaults_and_inputs + + +def _instrument_params(): + return {"general": {"lam": 0.0, "amp1": 1.0, "amp2": 1.0}} + + +def test_shipped_arts2d_deck_exposes_detector_irf_and_quadrature(): + repository = Path(__file__).resolve().parents[2] + defaults = yaml.safe_load((repository / "configs/arts-2d/defaults.yaml").read_text()) + inputs = yaml.safe_load((repository / "configs/arts-2d/inputs.yaml").read_text()) + + merged = merge_defaults_and_inputs(defaults, inputs) + + assert merged["other"]["extraoptions"]["spectype"] == "angular" + assert merged["other"]["detector_specs"]["widIRF"]["spect_FWHM_ele"] == 0.9 + assert merged["other"]["resonance_quadrature"]["enabled"] is True + + +def test_detector_edges_distinguish_centers_from_finite_pixel_support(): + centers = np.array([0.5, 1.5, 2.5, 3.5, 4.5]) + + np.testing.assert_array_equal(detector_edges_from_centers(centers), np.arange(6.0)) + # Two full two-pixel groups followed by a one-pixel ragged group. Re-inferring + # edges from the group centers [1, 3, 4.5] would incorrectly move the 4.0 edge. + np.testing.assert_array_equal(grouped_detector_edges(centers, 2), np.array([0.0, 2.0, 4.0, 5.0])) + + +@pytest.mark.parametrize("bad_centers", [[1.0], [1.0, 1.0], [2.0, 1.0], [1.0, np.nan]]) +def test_detector_edges_reject_invalid_center_axes(bad_centers): + with pytest.raises(ValueError): + detector_edges_from_centers(np.asarray(bad_centers)) + + +@pytest.mark.parametrize("bad_group", [0, -1, 1.5, True]) +def test_grouped_detector_edges_reject_invalid_group_size(bad_group): + with pytest.raises(ValueError): + grouped_detector_edges(np.arange(4.0), bad_group) + + +def test_prepare_data_retains_exact_arts_spectral_edges(monkeypatch): + electron_data = np.arange(20.0).reshape(5, 4) + ion_data = np.zeros_like(electron_data) + wavelength_centers = np.arange(5.0) + 0.5 + + monkeypatch.setattr( + prepare, + "loadData", + lambda *args, **kwargs: [electron_data, ion_data, "Angle", [0.0, 0.0], "angular", [5, 4]], + ) + monkeypatch.setattr( + prepare, + "get_scattering_angles", + lambda cfg: {"sa": np.arange(4.0), "weights": np.ones((4, 4))}, + ) + monkeypatch.setattr( + prepare, + "get_calibrations", + lambda *args, **kwargs: [ + np.arange(4.0), + np.arange(4.0), + wavelength_centers, + wavelength_centers, + 1.0, + { + "spect_stddev_ion": 1.0, + "spect_stddev_ele": 1.0, + "spect_FWHM_ele": 1.0, + "ang_FWHM_ele": 1.0, + }, + ], + ) + monkeypatch.setattr(prepare, "correctThroughput", lambda data, *args, **kwargs: data) + monkeypatch.setattr( + prepare, + "get_shot_bg", + lambda *args, **kwargs: [np.zeros_like(electron_data), np.zeros_like(ion_data)], + ) + + cfg = { + "data": { + "lineouts": {"type": "range", "val": [0], "start": 0, "end": 4}, + "load_ele_spec": True, + "load_ion_spec": False, + "fit_EPWb": True, + "fit_EPWr": True, + "fit_IAW": False, + "launch_data_visualizer": False, + "shotDay": False, + "bgscaleE": 0.0, + }, + "feature_detector": {"estimate_lineouts_epw": False, "estimate_lineouts_iaw": False}, + "optimizer": {"batch_size": 1}, + "other": { + "CCDsize": [5, 4], + "ang_res_unit": 2, + "lam_res_unit": 2, + "points_per_pixel": 1, + "detector_specs": {}, + "extraoptions": {"spectype": "angular"}, + }, + } + + 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) + + +def test_forward_arts_bounds_remain_first_and_last_detector_centers(): + cfg = { + "data": {"load_ele_spec": True}, + "other": { + "CCDsize": [5, 4], + "lamrangE": [400.0, 500.0], + "detector_specs": {}, + "extraoptions": {"spectype": "angular_full"}, + }, + } + + _ensure_angular_detector_edges(cfg) + + np.testing.assert_array_equal( + cfg["other"]["detector_specs"]["electron_wavelength_edges"], + np.array([387.5, 412.5, 437.5, 462.5, 487.5, 512.5]), + ) + np.testing.assert_array_equal( + cfg["other"]["detector_specs"]["electron_wavelength_centers"], + np.array([400.0, 425.0, 450.0, 475.0, 500.0]), + ) + + +def test_single_bin_forward_arts_uses_the_requested_finite_support(): + cfg = { + "data": {"load_ele_spec": True}, + "other": { + "CCDsize": [1, 4], + "lamrangE": [400.0, 500.0], + "detector_specs": {}, + "extraoptions": {"spectype": "angular_full"}, + }, + } + + _ensure_angular_detector_edges(cfg) + + np.testing.assert_array_equal( + cfg["other"]["detector_specs"]["electron_wavelength_edges"], + np.array([400.0, 500.0]), + ) + np.testing.assert_array_equal( + cfg["other"]["detector_specs"]["electron_wavelength_centers"], + np.array([450.0]), + ) + + +@pytest.mark.parametrize("number_of_points", [400, 801]) +def test_normalized_spectral_irf_preserves_unresolved_line_area_and_centroid( + number_of_points, +): + wavelengths = jnp.linspace(-10.0, 10.0, number_of_points) + edges = detector_edges_from_centers(np.asarray(wavelengths)) + widths = np.diff(edges) + source_index = number_of_points // 2 + unresolved_line = jnp.zeros((1, wavelengths.size)).at[0, source_index].set( + 1.0 / widths[source_index] + ) + irf = AngularIRF( + spect_stddev=1.0, + ang_stddev=1.0, + ang_axis=np.array([0.0]), + normalize=0, + ) + + _, blurred = add_ATS_IRF(irf, wavelengths, unresolved_line, _instrument_params()) + + blurred = np.asarray(blurred[0]) + area = np.sum(blurred * widths) + centroid = np.sum(blurred * widths * np.asarray(wavelengths)) / area + np.testing.assert_allclose(area, 1.0, rtol=0, atol=2e-12) + np.testing.assert_allclose( + centroid, + float(wavelengths[source_index]), + rtol=0, + atol=2e-11, + ) + assert float(jnp.max(blurred)) < float(jnp.max(unresolved_line)) + + +def test_normalized_ats_irf_keeps_constant_density_away_from_boundaries(): + angles = np.linspace(-5.0, 5.0, 101) + wavelengths = jnp.linspace(-10.0, 10.0, 201) + constant = jnp.ones((angles.size, wavelengths.size)) + irf = AngularIRF(spect_stddev=0.5, ang_stddev=0.3, ang_axis=angles, normalize=0) + + _, blurred = add_ATS_IRF(irf, wavelengths, constant, _instrument_params()) + + np.testing.assert_allclose(np.asarray(blurred[20:-20, 30:-30]), 1.0, rtol=0, atol=1e-9) + + +@pytest.mark.parametrize("number_of_points", [20, 21, 2050]) +def test_uniform_spectral_operator_matches_exact_dense_response(number_of_points): + wavelengths = jnp.linspace(-3.0, 4.0, number_of_points) + values = jnp.stack( + ( + jnp.exp(-((wavelengths + 0.7) / 0.4) ** 2), + 0.3 + jnp.exp(-((wavelengths - 1.2) / 0.8) ** 2), + ) + ) + + actual = _apply_uniform_gaussian_bin_response( + wavelengths, + values, + 0.35, + response_axis=1, + ) + exact = values @ _gaussian_bin_response(wavelengths, 0.35).T + + np.testing.assert_allclose( + np.asarray(actual), + np.asarray(exact), + rtol=5e-11, + atol=5e-11, + ) + + +def test_real_nonuniform_angular_irf_preserves_physical_area_and_center(): + angles, *_ = get_calibrations(104000, "angular", 0.0, [1024, 1024], {}) + edges = detector_edges_from_centers(angles) + widths = np.diff(edges) + source_index = angles.size // 2 + line = jnp.zeros((angles.size, 1)).at[source_index, 0].set( + 1.0 / widths[source_index] + ) + irf = AngularIRF( + spect_stddev=1.0, + ang_stddev=1.0 / 2.3548, + ang_axis=angles, + normalize=0, + ) + + _, blurred = add_ATS_IRF( + irf, + jnp.asarray([500.0]), + line, + _instrument_params(), + apply_spectral_blur=False, + ) + + blurred = np.asarray(blurred[:, 0]) + area = np.sum(blurred * widths) + centroid = np.sum(blurred * widths * angles) / area + np.testing.assert_allclose(area, 1.0, rtol=0, atol=2e-11) + np.testing.assert_allclose( + centroid, + angles[source_index], + rtol=0, + atol=1e-3, + ) + + constant = np.asarray( + add_ATS_IRF( + irf, + jnp.asarray([500.0]), + jnp.ones((angles.size, 1)), + _instrument_params(), + apply_spectral_blur=False, + )[1][:, 0] + ) + interior = (angles > angles[0] + 5 * irf.ang_stddev) & ( + angles < angles[-1] - 5 * irf.ang_stddev + ) + np.testing.assert_allclose(constant[interior], 1.0, rtol=0, atol=1e-6) + + +def test_ats_can_skip_spectral_blur_after_detector_bin_quadrature(): + angles = np.linspace(-5.0, 5.0, 101) + wavelengths = jnp.array([-1.0, 0.0, 1.0]) + detector_bin_means = jnp.tile(jnp.array([[0.0, 1.0, 0.0]]), (angles.size, 1)) + irf = AngularIRF(spect_stddev=0.5, ang_stddev=0.3, ang_axis=angles, normalize=0) + + _, angular_only = add_ATS_IRF( + irf, + wavelengths, + detector_bin_means, + _instrument_params(), + apply_spectral_blur=False, + ) + + np.testing.assert_allclose(np.asarray(angular_only[angles.size // 2]), np.array([0.0, 1.0, 0.0]), atol=1e-12) + + +def test_diagnostic_does_not_apply_a_second_spectral_irf(monkeypatch): + diagnostic = ThomsonScatteringDiagnostic.__new__(ThomsonScatteringDiagnostic) + diagnostic.cfg = { + "data": {"load_ele_spec": True, "load_ion_spec": False}, + "other": {"extraoptions": {"spectype": "angular_full"}}, + } + diagnostic.model = SimpleNamespace(electron_spectrum_is_detector_binned=True) + diagnostic.ats_irf = object() + called = {} + + def capture_ats_irf( + irf_description, + wavelength_axis, + model, + instrument_params, + *, + apply_spectral_blur, + ): + del irf_description, instrument_params + called["apply_spectral_blur"] = apply_spectral_blur + return wavelength_axis, model + + monkeypatch.setattr(thomson_diagnostic.irf, "add_ATS_IRF", capture_ats_irf) + wavelengths = jnp.asarray([1.0, 2.0, 3.0]) + model = jnp.ones((2, 3)) + output, _, output_axis, _ = diagnostic.postprocess_theory( + model, + 0, + wavelengths, + jnp.zeros(1), + {"e_amps": 1.0, "i_amps": 1.0}, + _instrument_params(), + ) + + assert called["apply_spectral_blur"] is False + np.testing.assert_array_equal(np.asarray(output), np.asarray(model)) + np.testing.assert_array_equal(np.asarray(output_axis), np.asarray(wavelengths)) + + +def test_detector_binned_arts_reduction_only_reduces_angle(): + diagnostic = ThomsonScatteringDiagnostic.__new__(ThomsonScatteringDiagnostic) + diagnostic.model = SimpleNamespace(electron_spectrum_is_detector_binned=True) + diagnostic.cfg = { + "other": {"CCDsize": [2, 3], "ang_res_unit": 2}, + "data": {"lineouts": {"start": 0, "end": 2}}, + } + wavelengths = jnp.asarray([1.0, 2.0, 4.0]) + model = jnp.asarray( + [ + [1.0, 2.0, 3.0], + [2.0, 4.0, 6.0], + [3.0, 6.0, 9.0], + [4.0, 8.0, 12.0], + ] + ) + batch = {"e_data": np.ones((2, 3)), "e_amps": jnp.ones((2, 1))} + instrument_params = { + "general": {"lam": 0.0, "amp1": 1.0, "amp2": 1.0} + } + + reduced, reduced_axis = diagnostic.reduce_ATS_to_resunit( + model, wavelengths, instrument_params, batch + ) + + assert reduced.shape == (2, 3) + np.testing.assert_array_equal(np.asarray(reduced_axis), np.asarray(wavelengths)) + + +def test_detector_binned_arts_reduction_keeps_a_ragged_angular_group(): + diagnostic = ThomsonScatteringDiagnostic.__new__(ThomsonScatteringDiagnostic) + diagnostic.model = SimpleNamespace(electron_spectrum_is_detector_binned=True) + diagnostic.cfg = { + "other": {"CCDsize": [2, 2], "ang_res_unit": 3}, + "data": {"lineouts": {"start": 0, "end": 2}}, + } + wavelengths = jnp.asarray([1.0, 2.0]) + model = jnp.arange(10.0).reshape(5, 2) + 1 + batch = {"e_data": np.ones((2, 2)), "e_amps": jnp.ones((2, 1))} + instrument_params = { + "general": {"lam": 0.0, "amp1": 1.0, "amp2": 1.0} + } + + reduced, _ = diagnostic.reduce_ATS_to_resunit( + model, wavelengths, instrument_params, batch + ) + + assert reduced.shape == (2, 2) diff --git a/tests/test_forward/test_resonance_notch_filter.py b/tests/test_forward/test_resonance_notch_filter.py new file mode 100644 index 000000000..d27aefcf5 --- /dev/null +++ b/tests/test_forward/test_resonance_notch_filter.py @@ -0,0 +1,169 @@ +"""Production wiring regression for source-side ARTS2D notch filtering.""" + +import numpy as np +from scipy.integrate import quad +from scipy.special import ndtr + +from jax import config + +config.update("jax_enable_x64", True) + +from jax import numpy as jnp + +from tsadar.core.physics.generate_spectra import FitModel + + +DETECTOR_EDGES_NM = np.asarray([500.0, 501.0, 502.0]) +IRF_SIGMA_NM = 0.1 +NOTCH_LOWER_NM = 500.25 +NOTCH_UPPER_NM = 500.75 +NOTCH_ATTENUATION = 0.1 +TAIL_SIGMA = 6.0 + + +def _config(integration_panels, scan_phase): + return { + "data": { + "load_ele_spec": True, + "load_ion_spec": False, + "ele_lam_shift": 0.0, + }, + "parameters": { + "general": { + "Te_gradient": {"num_grad_points": 1}, + "ne_gradient": {"num_grad_points": 1}, + "ud": {"angle": 0.0}, + }, + "electron": {"fe": {"dim": 2}}, + "ion-1": {"Va": {"angle": 0.0}}, + }, + "other": { + "detector_specs": { + "electron_wavelength_edges": DETECTOR_EDGES_NM, + "electron_wavelength_centers": 0.5 + * (DETECTOR_EDGES_NM[:-1] + DETECTOR_EDGES_NM[1:]), + "widIRF": {"spect_FWHM_ele": IRF_SIGMA_NM * 2.3548}, + }, + "extraoptions": {"spectype": "angular_full"}, + "resonance_quadrature": { + "enabled": True, + "root_scan_panels": 64, + "integration_panels": integration_panels, + "regular_order": 8, + "root_order": 16, + "max_roots": 16, + "tail_sigma": TAIL_SIGMA, + "scan_phase": scan_phase, + "map_batch_size": 1, + }, + "lamrangE": [500.0, 502.0], + "lamrangI": [500.0, 502.0], + "npts": 2, + "n_beta": 4, + # OD 1 over [500.25, 500.75] nm: the notch cuts through only part + # of the first [500, 501] detector bin. + "iawfilter": [1, 1.0, 0.5, 500.5], + "iawoff": 0, + }, + } + + +def _constant_spectrum_model(integration_panels, scan_phase): + model = FitModel( + _config(integration_panels, scan_phase), + {"sa": np.asarray([30.0]), "weights": np.ones((1, 1))}, + ) + model.electron_form_factor.prepare_2D_sinogram = lambda params: None + + def constant_terms(params, wavelengths_nm, sinogram=None, scattering_angles=None): + del params, sinogram, scattering_angles + shape = (wavelengths_nm.size, 1, 1) + return jnp.ones(shape), jnp.ones(shape, dtype=jnp.complex128) + + model.electron_form_factor.calc_2D_spectral_terms = constant_terms + return model + + +def _exact_source_side_bin_means(): + source_lower = DETECTOR_EDGES_NM[0] - TAIL_SIGMA * IRF_SIGMA_NM + source_upper = DETECTOR_EDGES_NM[-1] + TAIL_SIGMA * IRF_SIGMA_NM + + def transmission(source_nm): + if NOTCH_LOWER_NM < source_nm < NOTCH_UPPER_NM: + return NOTCH_ATTENUATION + return 1.0 + + def one_bin(bin_lower, bin_upper): + def integrand(source_nm): + probability = ndtr((bin_upper - source_nm) / IRF_SIGMA_NM) - ndtr( + (bin_lower - source_nm) / IRF_SIGMA_NM + ) + return transmission(source_nm) * probability + + integral = 0.0 + boundaries = [ + source_lower, + NOTCH_LOWER_NM, + NOTCH_UPPER_NM, + source_upper, + ] + for lower, upper in zip(boundaries[:-1], boundaries[1:]): + integral += quad(integrand, lower, upper, epsabs=1.0e-13, epsrel=1.0e-13)[0] + return integral / (bin_upper - bin_lower) + + return np.asarray( + [ + one_bin(lower, upper) + for lower, upper in zip( + DETECTOR_EDGES_NM[:-1], DETECTOR_EDGES_NM[1:] + ) + ] + ) + + +def test_partial_bin_notch_is_exact_source_side_and_phase_stable(): + """A notch boundary splits source integration, not a detector-space bin.""" + + expected = _exact_source_side_bin_means() + outputs = [] + + for integration_panels, scan_phase in [(32, 0.0), (32, 0.37), (64, -0.41)]: + model = _constant_spectrum_model(integration_panels, scan_phase) + np.testing.assert_array_equal( + np.asarray( + model.resonance_quadrature_options[ + "integration_breakpoints_nm" + ] + ), + np.asarray([NOTCH_LOWER_NM, NOTCH_UPPER_NM]), + ) + + _, detector_model, raw_model, diagnostics = ( + model.detector_integrated_electron_spectrum( + {"general": {"lam": 501.5}} + ) + ) + output = np.asarray(detector_model[0]) + outputs.append(output) + + np.testing.assert_allclose(output, expected, rtol=2.0e-10, atol=2.0e-10) + np.testing.assert_allclose( + np.asarray(raw_model[0, :, 0]), + expected, + rtol=2.0e-10, + atol=2.0e-10, + ) + assert not np.any(np.asarray(diagnostics.invalid_integration_breakpoints)) + assert np.all(np.asarray(diagnostics.root_count) == 0) + assert np.all(np.isfinite(output)) + + # A detector-space mask selected by the first bin's 500.5-nm center would attenuate + # the entire bin to roughly 0.1. Source-side filtering attenuates only its central + # half, leaving the correctly blurred/bin-integrated density near 0.55. + assert expected[0] > 0.5 + np.testing.assert_allclose( + outputs, + np.broadcast_to(expected, (3, 2)), + rtol=2.0e-10, + atol=2.0e-10, + ) diff --git a/tests/test_forward/test_unresolved_arts2d.py b/tests/test_forward/test_unresolved_arts2d.py new file mode 100644 index 000000000..b524621ce --- /dev/null +++ b/tests/test_forward/test_unresolved_arts2d.py @@ -0,0 +1,375 @@ +"""Regression for the unresolved small-angle ARTS2D EPW from #124/#134.""" + +import numpy as np + +from jax import config + +config.update("jax_enable_x64", True) + +import jax +from jax import numpy as jnp + +from tsadar.core.physics.form_factor import FormFactor +from tsadar.core.physics.generate_spectra import FitModel +from tsadar.core.physics.resonance_quadrature import integrate_detector_bins + + +DETECTOR_EDGES_NM = jnp.linspace(468.0, 476.0, 257) +FULL_DETECTOR_EDGES_NM = jnp.linspace(449.0, 670.0, 222) +IRF_SIGMA_NM = 0.9 / 2.3548 +PHASES = jnp.asarray([0.0, 0.31, 0.67]) +# Includes phases at which a 2048-panel root scan misses the closely spaced +# central pair, plus both extremes of the shifted scan grid. +FULL_RANGE_PHASES = jnp.asarray([-0.95, -0.10, 0.0, 0.90, 0.95]) +EXPECTED_FULL_RANGE_ROOTS_NM = np.asarray( + [ + 474.201875844894, + 514.081647830358, + 526.276288254028, + 526.451996163684, + 526.548008175237, + 526.723802255372, + 539.209067615960, + 591.042039725023, + ] +) +FINE_ROOT_SCAN_PANELS = 4096 +MAX_ROOTS = 16 + + +def _small_angle_problem(detector_edges_nm=DETECTOR_EDGES_NM): + velocity = jnp.linspace(-9.0, 9.0, 129) + projection = jnp.exp(-0.5 * velocity**2) / jnp.sqrt(2.0 * jnp.pi) + derivative = jnp.gradient(projection, velocity[1] - velocity[0]) + n_beta = 16 + sinogram = ( + jnp.broadcast_to(projection, (n_beta, velocity.size)), + jnp.broadcast_to(derivative, (n_beta, velocity.size)), + ) + + vx, vy = jnp.meshgrid(velocity, velocity) + distribution = jnp.exp(-0.5 * (vx**2 + vy**2)) / (2.0 * jnp.pi) + detector_edges_nm = jnp.asarray(detector_edges_nm) + form_factor = FormFactor( + lambda_range=[float(detector_edges_nm[0]), float(detector_edges_nm[-1])], + npts=32, + lam_shift=0.0, + scattering_angles={"sa": np.asarray([22.0]), "weights": np.ones((1, 1))}, + num_grad_points=1, + va_ang={"ion-1": 0.0}, + ud_ang=0.0, + calc_gain={"calc": False}, + n_beta=n_beta, + ) + + def parameters(electron_drift): + return { + "electron": { + "ne": 0.44, + "Te": 1.1, + "fe": distribution, + "v": velocity, + }, + "general": { + "ne_gradient": 0.0, + "Te_gradient": 0.0, + "lam": 526.5, + "ud": electron_drift, + }, + "ion-1": { + "A": 1.0, + "Z": 1.0, + "Ti": 0.03, + "fract": 1.0, + "Va": 0.0, + }, + } + + def integrate( + electron_drift, + phase, + integration_panels, + root_scan_panels=FINE_ROOT_SCAN_PANELS, + ): + params = parameters(electron_drift) + + def terms(wavelengths_nm): + numerator, epsilon = form_factor.calc_2D_spectral_terms( + params, + wavelengths_nm, + sinogram=sinogram, + ) + return numerator[:, 0, 0], epsilon[:, 0, 0] + + return integrate_detector_bins( + terms, + detector_edges_nm, + IRF_SIGMA_NM, + root_scan_panels=root_scan_panels, + integration_panels=integration_panels, + regular_order=8, + root_order=32, + max_roots=MAX_ROOTS, + neighbor_panels=1, + scan_phase=phase, + ) + + return integrate + + +def _evaluate_values_and_drift_gradients( + integrate, + integration_panels, + *, + root_scan_panels=FINE_ROOT_SCAN_PANELS, + phases=PHASES, +): + def one_phase(phase): + result = integrate(5.0, phase, integration_panels, root_scan_panels) + # The input is one scalar while the output has one entry per detector bin, + # so forward mode forms the complete Jacobian in one tangent evaluation. + gradient = jax.jacfwd( + lambda electron_drift: integrate( + electron_drift, + phase, + integration_panels, + root_scan_panels, + ).bin_mean + )(5.0) + return result, gradient + + return jax.jit(jax.vmap(one_phase))(phases) + + +def _relative_detector_l1(actual, expected, detector_edges_nm=DETECTOR_EDGES_NM): + widths = np.diff(np.asarray(detector_edges_nm)) + return np.sum(np.abs(actual - expected) * widths) / np.sum( + np.abs(expected) * widths + ) + + +def test_small_angle_unresolved_epw_values_and_gradients_converge(): + """The physical unresolved line is stable to phase and 2x/4x refinement. + + The acceptance tolerance is one percent. This regression retains substantial + margin while exercising the corrected signed susceptibility from #135, implicit + root motion with electron drift, exact spectral-IRF bin probabilities, and the + narrow resonance that the sampled wavelength grid missed in #124. + """ + + integrate = _small_angle_problem() + panel_counts = (64, 128, 256) + evaluations = [ + _evaluate_values_and_drift_gradients(integrate, panel_count) + for panel_count in panel_counts + ] + values = [np.asarray(results.bin_mean) for results, _ in evaluations] + gradients = [np.asarray(gradient) for _, gradient in evaluations] + + for results, gradient in evaluations: + assert np.all(np.asarray(results.diagnostics.root_count) == 1) + assert not np.any(np.asarray(results.diagnostics.root_overflow)) + assert np.all(np.isfinite(np.asarray(results.bin_mean))) + assert np.all(np.isfinite(np.asarray(gradient))) + assert np.linalg.norm(gradients[-1][0]) > 0 + + # Retain a focused reverse-mode check representative of inference: one scalar + # detector loss must have the same drift derivative as contracting the full + # forward-mode detector Jacobian with its loss weights. + loss_weights = jnp.linspace(-0.5, 0.75, DETECTOR_EDGES_NM.size - 1) + loss_weights *= jnp.diff(DETECTOR_EDGES_NM) + + def weighted_detector_loss(electron_drift): + return jnp.vdot( + loss_weights, + integrate( + electron_drift, + PHASES[0], + panel_counts[-1], + FINE_ROOT_SCAN_PANELS, + ).bin_mean, + ) + + reverse_loss_gradient = jax.jit(jax.grad(weighted_detector_loss))(5.0) + forward_loss_gradient = jnp.vdot(loss_weights, evaluations[-1][1][0]) + assert np.isfinite(float(reverse_loss_gradient)) + assert abs(float(reverse_loss_gradient)) > 0 + np.testing.assert_allclose( + np.asarray(reverse_loss_gradient), + np.asarray(forward_loss_gradient), + rtol=2e-10, + atol=2e-12, + ) + + # The corrected benchmark root lies near 474.2 nm and is roughly one thousand + # times narrower than the 0.9-nm FWHM instrument response. + roots = np.asarray(evaluations[-1][0].diagnostics.roots_nm)[:, 0] + widths = np.asarray( + evaluations[-1][0].diagnostics.resonance_half_widths_nm + )[:, 0] + np.testing.assert_allclose(roots, 474.22, rtol=0, atol=0.03) + assert np.all(widths < 1.0e-3) + + for phase_index in range(1, PHASES.size): + assert _relative_detector_l1(values[-1][phase_index], values[-1][0]) < 0.01 + assert ( + np.linalg.norm(gradients[-1][phase_index] - gradients[-1][0]) + / np.linalg.norm(gradients[-1][0]) + < 0.01 + ) + + # The same 64-panel baseline is compared with both 2x and 4x integration + # refinement while the fine root scan remains fixed. This isolates integration + # convergence from root-discovery topology. + for refined_index in (1, 2): + for phase_index in range(PHASES.size): + assert ( + _relative_detector_l1( + values[0][phase_index], values[refined_index][phase_index] + ) + < 0.01 + ) + assert ( + np.linalg.norm( + gradients[0][phase_index] - gradients[refined_index][phase_index] + ) + / np.linalg.norm(gradients[refined_index][phase_index]) + < 0.01 + ) + + +def test_full_detector_range_finds_every_physical_root_with_finite_gradients(): + """A fine root scan sees the complete small-angle dielectric topology. + + Across the representative 449--670 nm ARTS range the collisionless dielectric + has narrow blue/red EPW roots, broad crossings, and four closely spaced roots + around the probe wavelength. Root discovery must be independent of the coarser + detector-integration grid and its phase; missing the close pairs can silently + integrate the wrong spectrum even when a narrow-wing regression still passes. + """ + + integrate = _small_angle_problem(FULL_DETECTOR_EDGES_NM) + results, gradients = _evaluate_values_and_drift_gradients( + integrate, + 256, + root_scan_panels=FINE_ROOT_SCAN_PANELS, + phases=FULL_RANGE_PHASES, + ) + + root_counts = np.asarray(results.diagnostics.root_count) + roots = np.asarray(results.diagnostics.roots_nm) + root_masks = np.asarray(results.diagnostics.root_mask) + values = np.asarray(results.bin_mean) + gradients = np.asarray(gradients) + + np.testing.assert_array_equal(root_counts, np.full(FULL_RANGE_PHASES.shape, 8)) + assert not np.any(np.asarray(results.diagnostics.root_overflow)) + assert np.all(np.sum(root_masks, axis=1) == 8) + assert np.all(np.isfinite(values)) + assert np.all(np.isfinite(gradients)) + assert np.all(np.linalg.norm(gradients, axis=1) > 0) + + active_roots = np.stack( + [ + phase_roots[phase_mask] + for phase_roots, phase_mask in zip(roots, root_masks) + ] + ) + np.testing.assert_allclose( + active_roots, + np.broadcast_to(active_roots[0], active_roots.shape), + rtol=0, + atol=2e-8, + ) + np.testing.assert_allclose( + active_roots, + np.broadcast_to(EXPECTED_FULL_RANGE_ROOTS_NM, active_roots.shape), + rtol=0, + atol=2e-6, + ) + + +def test_fit_model_integrates_each_geometry_then_aperture_weights(monkeypatch): + """Production plumbing preserves gradient/angle topology without excess memory.""" + + detector_edges = np.asarray([0.0, 0.45, 0.9, 1.35, 2.0]) + detector_centers = 0.5 * (detector_edges[:-1] + detector_edges[1:]) + weights = np.asarray([[0.25, 0.75], [0.6, 0.4]]) + scattering_angles = {"sa": np.asarray([1.0, 2.0]), "weights": weights} + cfg = { + "data": { + "load_ele_spec": True, + "load_ion_spec": False, + "ele_lam_shift": 0.0, + }, + "parameters": { + "general": { + "Te_gradient": {"num_grad_points": 2}, + "ne_gradient": {"num_grad_points": 2}, + "ud": {"angle": 0.0}, + }, + "electron": {"fe": {"dim": 2}}, + "ion-1": {"Va": {"angle": 0.0}}, + }, + "other": { + "detector_specs": { + "electron_wavelength_edges": detector_edges, + "electron_wavelength_centers": detector_centers, + "widIRF": {"spect_FWHM_ele": 0.1}, + }, + "extraoptions": {"spectype": "angular_full"}, + "resonance_quadrature": { + "root_scan_panels": 64, + "integration_panels": 64, + "regular_order": 4, + "root_order": 16, + "max_roots": MAX_ROOTS, + "map_batch_size": 1, + }, + "lamrangE": [0.0, 2.0], + "lamrangI": [0.0, 2.0], + "npts": 8, + "n_beta": 4, + "iawoff": 0, + "iawfilter": [0, 0.0, 0.0, 0.0], + }, + } + model = FitModel(cfg, scattering_angles) + + monkeypatch.setattr( + model.electron_form_factor, + "prepare_2D_sinogram", + lambda params: None, + ) + + def synthetic_terms(params, wavelengths_nm, sinogram=None, scattering_angles=None): + del params, sinogram + angles = jnp.atleast_1d(jnp.asarray(scattering_angles)) + gradient_offsets = 0.03 * jnp.arange(2)[:, None] + roots = 0.65 + 0.10 * angles[None, :] + gradient_offsets + epsilon = wavelengths_nm[:, None, None] - roots[None, :, :] + 0.02j + numerator = jnp.full(epsilon.shape, 0.02 / jnp.pi) + return numerator, epsilon + + monkeypatch.setattr( + model.electron_form_factor, + "calc_2D_spectral_terms", + synthetic_terms, + ) + + axis, model_bins, raw_bins, diagnostics = model.detector_integrated_electron_spectrum( + {"general": {"lam": 1.0}} + ) + + assert model.electron_spectrum_is_detector_binned + assert raw_bins.shape == (2, detector_edges.size - 1, 2) + assert model_bins.shape == (2, detector_edges.size - 1) + assert np.all(np.asarray(diagnostics.root_count) == 1) + assert np.all(np.isfinite(np.asarray(raw_bins))) + np.testing.assert_array_equal(np.asarray(axis), detector_centers) + np.testing.assert_allclose( + np.asarray(model_bins), + weights @ np.mean(np.asarray(raw_bins), axis=0).T, + rtol=2e-14, + atol=2e-14, + ) diff --git a/tests/test_inverse/test_angular_1v.py b/tests/test_inverse/test_angular_1v.py index ec1911972..31c91268e 100644 --- a/tests/test_inverse/test_angular_1v.py +++ b/tests/test_inverse/test_angular_1v.py @@ -139,7 +139,11 @@ def test_arts1d_inverse(arbitrary_distribution: bool): sas = get_scattering_angles(config) sas["angAxis"], _, _, _, _, _ = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" diff --git a/tests/test_inverse/test_angular_2v.py b/tests/test_inverse/test_angular_2v.py index 69a6f630e..8b69958b5 100644 --- a/tests/test_inverse/test_angular_2v.py +++ b/tests/test_inverse/test_angular_2v.py @@ -157,7 +157,11 @@ def test_arts2d_inverse(dist_type: bool): sas = get_scattering_angles(config) sas["angAxis"], _, _, _, _, _ = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" diff --git a/tsadar/core/instrument/irf.py b/tsadar/core/instrument/irf.py index 73fa16da0..82d8f3ad9 100644 --- a/tsadar/core/instrument/irf.py +++ b/tsadar/core/instrument/irf.py @@ -1,4 +1,4 @@ -"""Instrument response function: spectral (and angular) blur, then binning onto pixels. +"""Physical-coordinate instrument responses for spectral and angular detector bins. This is device-dependent code, not physics. It holds both the value objects describing *this* detector's response and the routines that apply it, because those two belong to @@ -16,12 +16,134 @@ from typing import Tuple import numpy as np -from jax import numpy as jnp, vmap +from jax import vmap +from jax import numpy as jnp +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:], + ) + ) + + +def _normal_cdf_antiderivative(z): + """Return an antiderivative of the standard-normal CDF.""" + + density = jnp.exp(-0.5 * z**2) / jnp.sqrt(2.0 * jnp.pi) + return z * ndtr(z) + density + + +def _gaussian_bin_response(axis, stddev): + r"""Map physical-coordinate source-bin means to blurred output-bin means. + + For inferred source/output edges :math:`x_j` and :math:`y_i`, this returns + + .. math:: + + R_{ij} = \frac{1}{\Delta y_i} + \int_{y_i}^{y_{i+1}}\int_{x_j}^{x_{j+1}} + G_\sigma(y-x)\,dx\,dy. + + The double integral is evaluated analytically with the antiderivative of the + normal CDF. Unlike an index convolution, this is centered and unit-area on an + even or nonuniform calibrated axis. A single-bin axis has no resolvable blur and + maps identically. + """ + + if stddev <= 0: + raise ValueError(f"Gaussian standard deviation must be positive, got {stddev!r}.") + + axis = jnp.ravel(jnp.asarray(axis)) + if axis.size == 1: + return jnp.ones((1, 1), dtype=axis.dtype) + + edges = _bin_edges_from_centers(axis) + output_lower = edges[:-1, None] + output_upper = edges[1:, None] + source_lower = edges[None, :-1] + source_upper = edges[None, 1:] + sigma = jnp.asarray(stddev, dtype=axis.dtype) + + integrated_mass = sigma * ( + _normal_cdf_antiderivative((output_upper - source_lower) / sigma) + - _normal_cdf_antiderivative((output_upper - source_upper) / sigma) + - _normal_cdf_antiderivative((output_lower - source_lower) / sigma) + + _normal_cdf_antiderivative((output_lower - source_upper) / sigma) + ) + # Roundoff in far-tail cancellation can produce tiny negative entries. The matrix + # depends only on static detector coordinates, so clipping those does not alter any + # fitted-physics tangent. + integrated_mass = jnp.maximum(integrated_mass, 0) + return integrated_mass / jnp.diff(edges)[:, None] + + +def _uniform_gaussian_bin_kernel(axis, stddev): + r"""Return the exact Toeplitz response kernel for uniform physical bins. + + On a uniform axis, :func:`_gaussian_bin_response` depends only on the + output-minus-source bin lag. Storing those ``2*N - 1`` lags instead of the + dense ``N*N`` matrix keeps fine 1D spectral grids memory-linear while retaining + the same analytic double-bin integral and finite-detector boundary behavior. + """ + + if stddev <= 0: + raise ValueError(f"Gaussian standard deviation must be positive, got {stddev!r}.") + + axis = jnp.ravel(jnp.asarray(axis)) + if axis.size == 1: + return jnp.ones((1,), dtype=axis.dtype) + + cell_width = jnp.abs((axis[-1] - axis[0]) / (axis.size - 1)) + lags = jnp.arange(-(axis.size - 1), axis.size, dtype=axis.dtype) * cell_width + sigma = jnp.asarray(stddev, dtype=axis.dtype) + integrated_mass = sigma * ( + _normal_cdf_antiderivative((lags + cell_width) / sigma) + - 2 * _normal_cdf_antiderivative(lags / sigma) + + _normal_cdf_antiderivative((lags - cell_width) / sigma) + ) + return jnp.maximum(integrated_mass, 0) / cell_width + + +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 + grids use FFT convolution, reducing the response storage from quadratic to linear. + """ + + axis = jnp.ravel(jnp.asarray(axis)) + values = jnp.asarray(values) + if values.shape[response_axis] != axis.size: + raise ValueError( + "Gaussian-response axis length must match the selected value dimension, " + f"got {axis.size} and {values.shape[response_axis]}." + ) + + kernel = _uniform_gaussian_bin_kernel(axis, stddev=stddev) + moved = jnp.moveaxis(values, response_axis, -1) + flat = moved.reshape((-1, axis.size)) + convolve = fftconvolve if axis.size > 2048 else jnp.convolve + blurred = vmap(lambda row: convolve(row, kernel, mode="valid"))(flat) + blurred = blurred.reshape(moved.shape) + return jnp.moveaxis(blurred, -1, response_axis) @dataclass(frozen=True) class SpectrometerIRF: - """Response of a 1D spectrometer: spectral blur, then binning onto detector pixels. + """Response of a 1D spectrometer: physical spectral blur, then pixel reduction. Every field is required. There are no defaults on purpose: a default for ``n_spectral_pixels`` would be silently correct on OMEGA and silently wrong @@ -67,15 +189,22 @@ class AngularIRF: normalize: int -def add_ATS_IRF(irf: AngularIRF, lamAxisE, modlE, TSins) -> Tuple[jnp.ndarray, jnp.ndarray]: +def add_ATS_IRF( + irf: AngularIRF, lamAxisE, modlE, TSins, *, apply_spectral_blur: bool = True +) -> Tuple[jnp.ndarray, jnp.ndarray]: """ - Applies a 2D Gaussian smoothing to angular Thomson scattering data to account for the instrument response function (IRF) of the diagnostic. - This function convolves the synthetic spectra with Gaussian kernels along both the wavelength and angular axes, simulating the broadening effects introduced by the instrument. The resulting spectrum is optionally normalized according to the IRF description. + Applies a separable 2D Gaussian detector response to angular Thomson scattering data. + The response is integrated over physical wavelength and angular cells rather than + sampled in array-index space. The resulting spectrum is optionally normalized + according to the IRF description. Args: irf (AngularIRF): Description of the angular detector's response. lamAxisE (jnp.ndarray): Array of wavelengths (in nm) at which the spectrum is computed. modlE (jnp.ndarray): Synthetic spectra produced by the formfactor routine, shape (n_angles, n_wavelengths). TSins (dict): Dictionary of Thomson scattering instrument parameters and their values. + apply_spectral_blur: Whether to blur along wavelength. Set this to + ``False`` when an upstream detector-bin quadrature has already folded in + the spectral IRF; the angular IRF is still applied. Returns: lamAxisE (jnp.ndarray): Wavelength axis (in nm). ThryE (jnp.ndarray): Smoothed and optionally normalized synthetic spectra, shape (n_angles, n_wavelengths). @@ -83,25 +212,17 @@ def add_ATS_IRF(irf: AngularIRF, lamAxisE, modlE, TSins) -> Tuple[jnp.ndarray, j stddev_lam = irf.spect_stddev stddev_ang = irf.ang_stddev - # Conceptual_origin so the convolution donsn't shift the signal - origin_lam = (jnp.amax(lamAxisE) + jnp.amin(lamAxisE)) / 2.0 - origin_ang = (jnp.amax(irf.ang_axis) + jnp.amin(irf.ang_axis)) / 2.0 - inst_func_lam = jnp.squeeze( - (1.0 / (stddev_lam * jnp.sqrt(2.0 * jnp.pi))) - * jnp.exp(-((lamAxisE - origin_lam) ** 2.0) / (2.0 * (stddev_lam) ** 2.0)) - ) # Gaussian - inst_func_ang = jnp.squeeze( - (1.0 / (stddev_ang * jnp.sqrt(2.0 * jnp.pi))) - * jnp.exp(-((irf.ang_axis - origin_ang) ** 2.0) / (2.0 * (stddev_ang) ** 2.0)) - ) # Gaussian - # 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 + angular_response = _gaussian_bin_response(irf.ang_axis, stddev_ang) + # The response matrices act on detector-bin mean densities. Matrix multiplication + # applies every wavelength/angle column in one compiled operation. + ThryE = angular_response @ modlE + if apply_spectral_blur: + ThryE = _apply_uniform_gaussian_bin_response( + lamAxisE, + ThryE, + stddev_lam, + response_axis=1, + ) if irf.normalize > 0: ThryE = jnp.where( @@ -128,13 +249,12 @@ def add_ion_IRF(irf: SpectrometerIRF, lamAxisI, modlI, amps, TSins) -> Tuple[jnp stddevI = irf.spect_stddev if stddevI: - originI = (jnp.amax(lamAxisI) + jnp.amin(lamAxisI)) / 2.0 - inst_funcI = jnp.squeeze( - (1.0 / (stddevI * jnp.sqrt(2.0 * jnp.pi))) - * jnp.exp(-((lamAxisI - originI) ** 2.0) / (2.0 * (stddevI) ** 2.0)) - ) # Gaussian - ThryI = jnp.convolve(modlI, inst_funcI, "same") - ThryI = (jnp.amax(modlI) / jnp.amax(ThryI)) * ThryI + ThryI = _apply_uniform_gaussian_bin_response( + lamAxisI, + modlI, + stddevI, + response_axis=0, + ) ThryI = jnp.average(ThryI.reshape(irf.n_spectral_pixels, -1), axis=1) #print(f"modlI max {jnp.max(modlI)}") #print(f"ThryI max {jnp.max(ThryI)}") @@ -170,13 +290,12 @@ def add_electron_IRF(irf: SpectrometerIRF, lamAxisE, modlE, amps, TSins) -> Tupl """ stddevE = irf.spect_stddev - # Conceptual_origin so the convolution doesn't shift the signal - originE = (jnp.amax(lamAxisE) + jnp.amin(lamAxisE)) / 2.0 - inst_funcE = jnp.squeeze( - (1.0 / (stddevE * jnp.sqrt(2.0 * jnp.pi))) * jnp.exp(-((lamAxisE - originE) ** 2.0) / (2.0 * (stddevE) ** 2.0)) - ) # Gaussian - ThryE = jnp.convolve(modlE, inst_funcE, "same") - ThryE = (jnp.amax(modlE) / jnp.amax(ThryE)) * ThryE + ThryE = _apply_uniform_gaussian_bin_response( + lamAxisE, + modlE, + stddevE, + response_axis=0, + ) if irf.normalize > 0: ThryE = jnp.where( diff --git a/tsadar/core/physics/form_factor.py b/tsadar/core/physics/form_factor.py index 31cd85b6e..da77b9e4d 100644 --- a/tsadar/core/physics/form_factor.py +++ b/tsadar/core/physics/form_factor.py @@ -257,6 +257,7 @@ def __init__( minmax = 8.2 h1 = 1024 # 1024 lamAxis = jnp.linspace(lambda_range[0], lambda_range[1], npts) + 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) self.omgs = omgs[None, ..., None, None] # [1, npts, 1, 1] @@ -715,7 +716,7 @@ def calc_chi_vals(self, vx, sinogram, inputs): chiERrat = -1.0 / (klde_mag_at**2) * _principal_value_integral(df, vx, xi_at) return fe_vphi, chiEI, chiERrat - def _calc_all_chi_vals_(self, vx, DF, beta, xi, klde_mag): + def _calc_all_chi_vals_(self, vx, DF, beta, xi, klde_mag, sinogram=None): """ Calculate the susceptibility values for all the desired points xie @@ -734,26 +735,36 @@ def _calc_all_chi_vals_(self, vx, DF, beta, xi, klde_mag): chiERrat: real part of the electron susceptibility """ - calc_chi_vals = "batch_vmap" - - flattened_inputs = (beta.flatten(), xi.flatten(), klde_mag.flatten()) - # Tabulate the projection over angles once, rather than rotating the whole 2D # distribution function again at every one of the (many) evaluation points. When # `n_beta` is 0 the distribution function is passed through and each point does # its own exact rotation, which is the behaviour this replaced. - df_or_sinogram = self._build_sinogram(vx, jnp.squeeze(DF)) if self.n_beta else jnp.squeeze(DF) + if sinogram is None: + sinogram = self._build_sinogram(vx, jnp.squeeze(DF)) if self.n_beta else jnp.squeeze(DF) + + return self._calc_all_chi_vals_from_sinogram(vx, sinogram, beta, xi, klde_mag) + + def _calc_all_chi_vals_from_sinogram(self, vx, sinogram, beta, xi, klde_mag): + """Evaluate electron terms using an already prepared EDF projection. + + Keeping this separate from :meth:`_calc_all_chi_vals_` lets wavelength-space + quadrature reuse the expensive numerical-EDF sinogram across root searches and + repeated batches of quadrature nodes. + """ + + calc_chi_vals = "batch_vmap" + flattened_inputs = (beta.flatten(), xi.flatten(), klde_mag.flatten()) if calc_chi_vals == "scan": _, (fe_vphi, chiEI, chiERrat) = scan( - self.scan_calc_chi_vals, (vx, df_or_sinogram), flattened_inputs, unroll=1 + self.scan_calc_chi_vals, (vx, sinogram), flattened_inputs, unroll=1 ) elif calc_chi_vals == "vmap": - fe_vphi, chiEI, chiERrat = self.vmap_calc_chi_vals(vx, df_or_sinogram, flattened_inputs) + fe_vphi, chiEI, chiERrat = self.vmap_calc_chi_vals(vx, sinogram, flattened_inputs) elif calc_chi_vals == "batch_vmap": - batch_vmap_calc_chi_vals = partial(self.calc_chi_vals, vx, df_or_sinogram) + batch_vmap_calc_chi_vals = partial(self.calc_chi_vals, vx, sinogram) fe_vphi, chiEI, chiERrat = jmap(batch_vmap_calc_chi_vals, xs=flattened_inputs, batch_size=128) else: raise NotImplementedError @@ -764,7 +775,7 @@ def _calc_all_chi_vals_(self, vx, DF, beta, xi, klde_mag): return fe_vphi, chiEI, chiERrat - def parallel_calc_all_chi_vals(self, x, DF, beta, xi, klde_mag): + def parallel_calc_all_chi_vals(self, x, DF, beta, xi, klde_mag, sinogram=None): """ Multi-device counterpart to _calc_all_chi_vals_: flattens beta/xi/klde_mag, distributes them across devices via self.sharding (device_put), then delegates to _calc_all_chi_vals_ to compute the @@ -793,7 +804,14 @@ def parallel_calc_all_chi_vals(self, x, DF, beta, xi, klde_mag): flat_xi = device_put(f_xi, self.sharding) flat_klde_mag = device_put(f_klde_mag, self.sharding) - fe_vphi, chiEI, chiERrat = self._calc_all_chi_vals_(x, DF, flat_beta, flat_xi, flat_klde_mag) + fe_vphi, chiEI, chiERrat = self._calc_all_chi_vals_( + x, + DF, + flat_beta, + flat_xi, + flat_klde_mag, + sinogram=sinogram, + ) fe_vphi = fe_vphi.reshape(beta.shape) chiEI = chiEI.reshape(beta.shape) @@ -819,15 +837,118 @@ def _ion_flow_angles(self, ion_species): ) return angles - def calc_in_2D(self, params): - """Calculate the collisionless Thomson spectrum for a 2-D numerical EDF. + def prepare_2D_sinogram(self, params): + """Prepare the wavelength-independent numerical-EDF projection state. + + Building the sinogram is the expensive part of evaluating a 2-D numerical + distribution. Detector quadrature should call this once, then pass the returned + state to every :meth:`calc_2D_spectral_terms` call made for root searches and + quadrature nodes. When ``n_beta == 0`` the returned state is the exact 2-D EDF + itself, preserving the existing exact-rotation fallback. + """ + + vx = params["electron"]["v"] + fe = jnp.squeeze(params["electron"]["fe"]) + return self._build_sinogram(vx, fe) if self.n_beta else fe + + def prepare_2D_spectral_evaluator(self, params, scattering_angles=None): + """Return an arbitrary-wavelength evaluator backed by one cached sinogram. - Each ion species has its own flow vector. ``general.ud`` is the electron drift - relative to their charge-weighted bulk flow, so the electron lab-frame velocity - is ``sum(Z * fract * Va) / Zbar + ud``. The longitudinal EDF projection is fixed - by ``k_hat`` and is sampled at a signed resonance coordinate. + ``scattering_angles`` is an optional scalar or one-dimensional array in degrees. + The returned callable accepts a one-dimensional wavelength array in nm and + returns ``(numerator_lambda, epsilon)``; see + :meth:`calc_2D_spectral_terms` for the precise convention and shapes. """ + sinogram = self.prepare_2D_sinogram(params) + return partial( + self.calc_2D_spectral_terms, + params, + sinogram=sinogram, + scattering_angles=scattering_angles, + ) + + def calc_2D_spectral_terms( + self, + params, + wavelengths_nm, + sinogram=None, + scattering_angles=None, + ): + """Evaluate reusable 2-D spectral terms on arbitrary wavelength nodes. + + Args: + params: Runtime plasma and numerical-EDF parameters. + wavelengths_nm: One-dimensional physical scattered-wavelength nodes in nm. + sinogram: Optional state returned by :meth:`prepare_2D_sinogram`. Supplying + it prevents the numerical EDF from being projected again. + scattering_angles: Optional scalar or one-dimensional array of scattering + angles in degrees. By default, uses ``self.scattering_angles["sa"]``. + + Returns: + A tuple ``(numerator_lambda, epsilon)``. Both arrays have shape + ``[wavelength, gradient, angle]``. The physical wavelength-space spectrum is + exactly ``numerator_lambda / abs(epsilon)**2``. ``numerator_lambda`` includes + the electron and ion structure-factor numerators, ``1 / |k|``, the + laboratory-frequency factor ``1 + 2*omega/omega_L``, ``r_e**2 n_e``, and the + ``d omega / d lambda`` Jacobian. No dielectric denominator is included. + """ + + wavelengths_nm = jnp.asarray(wavelengths_nm) + if wavelengths_nm.ndim != 1: + raise ValueError( + "wavelengths_nm must be one-dimensional; use " + "calc_2D_spectral_terms_at_points for per-spectrum nodes" + ) + + angles = self.scattering_angles["sa"] if scattering_angles is None else scattering_angles + angles = jnp.atleast_1d(jnp.asarray(angles)) + wavelength_points = jnp.broadcast_to( + wavelengths_nm[:, None, None], + (wavelengths_nm.size, self.num_grad_points, angles.size), + ) + return self.calc_2D_spectral_terms_at_points( + params, + wavelength_points, + sinogram=sinogram, + scattering_angles=angles, + ) + + def calc_2D_spectral_terms_at_points( + self, + params, + wavelengths_nm, + sinogram=None, + scattering_angles=None, + ): + """Evaluate terms at a distinct wavelength mesh for every plasma spectrum. + + ``wavelengths_nm`` has shape ``sample_shape + [gradient, angle]`` and the two + returned arrays have the identical shape. This is the efficient interface for + root-mapped quadrature, where each gradient/angle spectrum needs its own nodes. + For a common one-dimensional wavelength grid, prefer + :meth:`calc_2D_spectral_terms`. + """ + + angles = self.scattering_angles["sa"] if scattering_angles is None else scattering_angles + angles = jnp.atleast_1d(jnp.asarray(angles)) + wavelengths_nm = jnp.asarray(wavelengths_nm) + expected_trailing_shape = (self.num_grad_points, angles.size) + if wavelengths_nm.ndim < 2 or wavelengths_nm.shape[-2:] != expected_trailing_shape: + raise ValueError( + "wavelengths_nm must have trailing [gradient, angle] dimensions " + f"{expected_trailing_shape}, got {wavelengths_nm.shape}" + ) + + sample_shape = wavelengths_nm.shape[:-2] + # Canonical internal axes are [gradient, flattened sample, angle]. Flattening + # only sample axes keeps all shapes static under jit while supporting arbitrary + # quadrature-node layouts. + wavelengths_nm_canonical = jnp.transpose( + wavelengths_nm.reshape((-1,) + expected_trailing_shape), + (1, 0, 2), + ) + ne = ( 1.0e20 * params["electron"]["ne"] @@ -870,11 +991,13 @@ def calc_in_2D(self, params): Esq = self.Me * self.C**2 * re constants = jnp.sqrt(4 * jnp.pi * Esq / self.Me) - # Keep the calculation axes explicit: [gradient, wavelength, angle], adding - # species only as the final axis for ion quantities. - sarad = self.scattering_angles["sa"][None, None, :] * jnp.pi / 180 + sarad = angles[None, None, :] * jnp.pi / 180 omgL = self.omgL_num / lam - omgs = self.omgs[..., 0] + omgs = 2 * jnp.pi * 1.0e7 * self.C / wavelengths_nm_canonical + # Derive the wavelength-space Jacobian from the same frequency values as the + # historical path. Besides making the physical conversion explicit, this keeps + # configured-grid evaluation numerically identical down to roundoff. + lams_cm = 2 * jnp.pi * self.C / omgs omgpe = constants * jnp.sqrt(ne) omg = omgs - omgL @@ -907,7 +1030,16 @@ def calc_in_2D(self, params): ZpiI = interp_uniform(xii, self.xi2, self.Zpi[1, :], left=0, right=0) chiI = jnp.sum(-0.5 / (kldi**2) * (ZpiR + 1j * ZpiI), axis=-1) - fe_vphi, chiEI, chiERrat = self.calc_all_chi_vals(vx, fe, beta, xi, klde_mag) + if sinogram is None: + sinogram = self.prepare_2D_sinogram(params) + fe_vphi, chiEI, chiERrat = self.calc_all_chi_vals( + vx, + fe, + beta, + xi, + klde_mag, + sinogram=sinogram, + ) chiE = chiERrat + 1j * chiEI epsilon = 1.0 + chiE + chiI @@ -920,19 +1052,48 @@ def calc_in_2D(self, params): ) ele_comp = jnp.abs(1.0 + chiI) ** 2 * fe_vphi / vTe - SKW_ion_omg = jnp.sum( - ion_comp / k_mag[..., None] / jnp.abs(epsilon[..., None]) ** 2, - axis=-1, + # Everything except the common dielectric denominator belongs in the reusable + # wavelength-space numerator. Keeping the lab-frequency factor here is important: + # it is part of the physical spectrum, not a detector-integration correction. + structure_numerator = ( + jnp.sum(ion_comp / k_mag[..., None], axis=-1) + + ele_comp / k_mag ) - SKW_ele_omg = ele_comp / k_mag / jnp.abs(epsilon) ** 2 - - PsOmg = ( - (SKW_ion_omg + SKW_ele_omg) + numerator_lambda = ( + structure_numerator * (1 + 2 * omg / omgL) * re**2 * ne + * 2 + * jnp.pi + * self.C + / lams_cm**2 ) - lams = 2 * jnp.pi * self.C / self.omgs - formfactor = PsOmg * 2 * jnp.pi * self.C / lams[..., 0] ** 2 + output_shape = sample_shape + expected_trailing_shape + + def restore_sample_axes(values): + return jnp.transpose(values, (1, 0, 2)).reshape(output_shape) + + return restore_sample_axes(numerator_lambda), restore_sample_axes(epsilon) + + def calc_in_2D(self, params): + """Calculate the collisionless Thomson spectrum for a 2-D numerical EDF. + + This compatibility wrapper evaluates the reusable numerator and dielectric on + the object's configured wavelength grid, divides by ``abs(epsilon)**2``, and + restores the historical output shape ``[gradient, wavelength, angle]``. + """ + + sinogram = self.prepare_2D_sinogram(params) + numerator_lambda, epsilon = self.calc_2D_spectral_terms( + params, + self.lambda_axis_nm, + sinogram=sinogram, + ) + formfactor = jnp.transpose( + numerator_lambda / jnp.abs(epsilon) ** 2, + (1, 0, 2), + ) + lams = 2 * jnp.pi * self.C / self.omgs return formfactor, lams diff --git a/tsadar/core/physics/generate_spectra.py b/tsadar/core/physics/generate_spectra.py index dc0b672b6..15aacb0a8 100644 --- a/tsadar/core/physics/generate_spectra.py +++ b/tsadar/core/physics/generate_spectra.py @@ -3,8 +3,10 @@ from typing import Dict from .form_factor import DEFAULT_N_BETA, FormFactor +from .resonance_quadrature import integrate_detector_bins -from jax import numpy as jnp +from jax import lax, numpy as jnp +from jax.tree_util import tree_map class FitModel: @@ -64,6 +66,94 @@ def __init__(self, config: Dict, scattering_angles: Dict): self.config = config self.scattering_angles = scattering_angles + quadrature_config = config["other"].get("resonance_quadrature", {}) + self.electron_spectrum_is_detector_binned = bool( + config["data"]["load_ele_spec"] + and config["parameters"]["electron"]["fe"]["dim"] == 2 + and config["other"]["extraoptions"]["spectype"] == "angular_full" + and quadrature_config.get("enabled", True) + ) + if self.electron_spectrum_is_detector_binned: + detector_specs = config["other"]["detector_specs"] + self.electron_detector_edges_nm = jnp.asarray( + detector_specs["electron_wavelength_edges"] + ) + self.electron_detector_centers_nm = jnp.asarray( + detector_specs.get( + "electron_wavelength_centers", + 0.5 + * ( + self.electron_detector_edges_nm[:-1] + + self.electron_detector_edges_nm[1:] + ), + ) + ) + self.electron_irf_sigma_nm = ( + detector_specs["widIRF"]["spect_FWHM_ele"] / 2.3548 + ) + tail_sigma = float(quadrature_config.get("tail_sigma", 6.0)) + self.resonance_quadrature_options = { + "root_scan_panels": int( + quadrature_config.get("root_scan_panels", 4096) + ), + "integration_panels": int( + quadrature_config.get("integration_panels", 256) + ), + "regular_order": int(quadrature_config.get("regular_order", 8)), + "root_order": int(quadrature_config.get("root_order", 32)), + "max_roots": int(quadrature_config.get("max_roots", 16)), + "neighbor_panels": int(quadrature_config.get("neighbor_panels", 1)), + "bisection_iterations": int( + quadrature_config.get("bisection_iterations", 48) + ), + "tail_sigma": tail_sigma, + "scan_phase": float(quadrature_config.get("scan_phase", 0.0)), + } + # ``lax.map`` keeps the expensive node-by-detector response matrix bounded + # in memory. A small explicit batch can recover device parallelism without + # ever materializing every scattering geometry at once. + self.resonance_quadrature_map_batch_size = int( + quadrature_config.get("map_batch_size", 1) + ) + if self.resonance_quadrature_map_batch_size < 1: + raise ValueError("resonance_quadrature.map_batch_size must be positive") + + self.electron_notch_filter = None + if config["other"]["iawfilter"][0]: + filter_center = float(config["other"]["iawfilter"][3]) + filter_width = float(config["other"]["iawfilter"][2]) + if filter_width <= 0: + raise ValueError("enabled iawfilter width must be positive") + self.electron_notch_filter = ( + filter_center - filter_width / 2, + filter_center + filter_width / 2, + 10 ** (-float(config["other"]["iawfilter"][1])), + ) + + # The rectangular filter is applied to the continuous source spectrum, + # before the Gaussian spectral response and detector-bin integration. + # Preserve every discontinuity that actually lies inside the source + # integration domain as an exact coarse-panel boundary. Boundaries on or + # outside the domain do not split an integration panel and must not be + # passed to the quadrature kernel, which requires strict interior points. + source_lower = ( + float(self.electron_detector_edges_nm[0]) + - tail_sigma * self.electron_irf_sigma_nm + ) + source_upper = ( + float(self.electron_detector_edges_nm[-1]) + + tail_sigma * self.electron_irf_sigma_nm + ) + integration_breakpoints = tuple( + boundary + for boundary in self.electron_notch_filter[:2] + if source_lower < boundary < source_upper + ) + if integration_breakpoints: + self.resonance_quadrature_options[ + "integration_breakpoints_nm" + ] = jnp.asarray(integration_breakpoints) + assert ( config["parameters"]["general"]["Te_gradient"]["num_grad_points"] == config["parameters"]["general"]["ne_gradient"]["num_grad_points"] @@ -197,6 +287,12 @@ def _electron_spectrum_core(self, all_params, want_thry): for the detailed/postprocessing output, so electron_spectrum passes False to skip it entirely rather than computing and discarding it on every forward pass. """ + if self.electron_spectrum_is_detector_binned: + lamAxisE, modlE, ThryE, _ = self.detector_integrated_electron_spectrum( + all_params + ) + return lamAxisE, modlE, ThryE if want_thry else 0 + if self.config["data"]["load_ele_spec"]: if self.config["parameters"]["electron"]["fe"]["dim"] == 1: ThryE, lamAxisE_orig = self.electron_form_factor(all_params) @@ -242,3 +338,93 @@ def _electron_spectrum_core(self, all_params, want_thry): raw_thry = 0 lamAxisE = [] return lamAxisE, modlE, raw_thry + + def detector_integrated_electron_spectrum(self, all_params): + """Integrate the ARTS2D electron spectrum directly into detector bins. + + Each gradient/scattering-angle spectrum gets an independent root search and + tan-mapped quadrature because its dielectric roots occur at different physical + wavelengths. The Gaussian spectral IRF is folded into the integral through exact + CDF differences, so callers must not apply a second spectral convolution or + wavelength reduction. An enabled rectangular ``iawfilter`` multiplies the + continuous source numerator, with its in-domain edges inserted as exact + integration breakpoints. ``iawoff`` remains a detector-bin mask. + + Returns: + ``(wavelength_centers_nm, aperture_weighted_bin_means, + per_geometry_bin_means, diagnostics)``. The first model array has shape + ``[calibrated angle, detector bin]``; the per-geometry array has shape + ``[gradient, detector bin, scattering angle]``. Every diagnostic field has + leading ``[gradient, scattering angle]`` axes. + """ + + if not self.electron_spectrum_is_detector_binned: + raise ValueError( + "detector-integrated electron spectra require enabled ARTS2D " + "resonance quadrature" + ) + + form_factor = self.electron_form_factor + sinogram = form_factor.prepare_2D_sinogram(all_params) + angles = jnp.asarray(self.scattering_angles["sa"]) + num_gradients = form_factor.num_grad_points + num_angles = angles.size + flat_indices = jnp.arange(num_gradients * num_angles, dtype=jnp.int32) + + def integrate_one(flat_index): + gradient_index = flat_index // num_angles + angle_index = flat_index % num_angles + angle = angles[angle_index] + + def terms_at(wavelengths_nm): + numerator, epsilon = form_factor.calc_2D_spectral_terms( + all_params, + wavelengths_nm, + sinogram=sinogram, + scattering_angles=angle, + ) + if self.electron_notch_filter is not None: + filter_lower, filter_upper, attenuation = self.electron_notch_filter + transmission = jnp.where( + (wavelengths_nm > filter_lower) + & (wavelengths_nm < filter_upper), + attenuation, + 1.0, + ) + numerator = numerator * transmission[:, None, None] + return numerator[:, gradient_index, 0], epsilon[:, gradient_index, 0] + + return integrate_detector_bins( + terms_at, + self.electron_detector_edges_nm, + self.electron_irf_sigma_nm, + **self.resonance_quadrature_options, + ) + + result = lax.map( + integrate_one, + flat_indices, + batch_size=self.resonance_quadrature_map_batch_size, + ) + per_geometry = result.bin_mean.reshape( + num_gradients, num_angles, self.electron_detector_edges_nm.size - 1 + ) + # Restore FormFactor's historical raw ordering [gradient, wavelength, angle]. + ThryE = jnp.transpose(per_geometry, (0, 2, 1)) + gradient_average = jnp.mean(per_geometry, axis=0) + modlE = jnp.matmul(self.scattering_angles["weights"], gradient_average) + lamAxisE = self.electron_detector_centers_nm + + lam = all_params["general"]["lam"] + if self.config["other"]["iawoff"]: + ion_feature = (lamAxisE > lam - 3.0) & (lamAxisE < lam + 3.0) + modlE = jnp.where(ion_feature[None, :], 0, modlE) + ThryE = jnp.where(ion_feature[None, :, None], 0, ThryE) + + diagnostics = tree_map( + lambda value: value.reshape( + (num_gradients, num_angles) + value.shape[1:] + ), + result.diagnostics, + ) + return lamAxisE, modlE, ThryE, diagnostics diff --git a/tsadar/core/physics/resonance_quadrature.py b/tsadar/core/physics/resonance_quadrature.py new file mode 100644 index 000000000..89eec6e3f --- /dev/null +++ b/tsadar/core/physics/resonance_quadrature.py @@ -0,0 +1,829 @@ +r"""Root-aware wavelength quadrature for unresolved collective resonances. + +The detector does not sample a point spectrum. A source wavelength contributes +to every detector bin according to the Gaussian instrument response, so the +quantity required by the forward model is + +.. math:: + + \bar S_j = \frac{1}{\Delta\lambda_j} + \int S(\lambda)\left[\Phi\left(\frac{e_{j+1}-\lambda}{\sigma}\right) + - \Phi\left(\frac{e_j-\lambda}{\sigma}\right)\right]d\lambda. + +For a collective electron-plasma-wave feature, ``S`` can be much narrower than +the detector grid. This module locates zeros of ``Re(epsilon)`` on a fixed scan, +then uses the local complex-linearized dielectric to tan-map Gauss--Legendre +nodes in coarse root panels and their neighbors. A separate, finer scan locates +roots, and coarse panels containing multiple roots are split at root midpoints. +All other panels use ordinary Gauss--Legendre quadrature. The topology (which +panels contain roots) is +necessarily nondifferentiable, while roots, mapped nodes, and mapped weights +remain differentiable. + +The public kernel handles one dielectric spectrum. It is deliberately written +so callers can :func:`jax.vmap` it over scattering geometries. ``terms_fn`` must +accept a one-dimensional wavelength array in nm and return +``(numerator, epsilon)``. ``epsilon`` has the same one-dimensional shape; +``numerator`` may additionally have trailing component axes. The intrinsic +spectrum is ``numerator / abs(epsilon)**2``. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Callable, NamedTuple + +import jax +from jax import lax +import jax.numpy as jnp +from jax.scipy.special import ndtr +import numpy as np + + +Array = jax.Array +SpectralTerms = Callable[[Array], tuple[Array, Array]] + + +class ResonanceQuadratureDiagnostics(NamedTuple): + """JAX-compatible status and root information for one integration. + + ``root_overflow``, ``nonfinite``, ``zero_width``, or any invalid-input flag + means the numerical value must not be used; ``bin_mean`` is poisoned with + ``nan`` in that case so compiled fit code cannot consume a partial result. + The fixed-size root arrays are padded with finite values; ``root_mask`` + identifies populated entries. + """ + + root_count: Array + used_root_count: Array + root_overflow: Array + nonfinite: Array + zero_width: Array + invalid_bin_width: Array + nonmonotonic_edges: Array + invalid_irf_sigma: Array + invalid_source_bounds: Array + invalid_scan_phase: Array + invalid_integration_breakpoints: Array + roots_nm: Array + resonance_centers_nm: Array + resonance_half_widths_nm: Array + root_mask: Array + + +class ResonanceQuadratureResult(NamedTuple): + """Detector-bin mean densities and their numerical diagnostics.""" + + bin_mean: Array + diagnostics: ResonanceQuadratureDiagnostics + + +@lru_cache(maxsize=None) +def _legendre_rule(order: int) -> tuple[np.ndarray, np.ndarray]: + nodes, weights = np.polynomial.legendre.leggauss(order) + return nodes, weights + + +def gaussian_bin_probabilities(wavelengths_nm, detector_edges_nm, irf_sigma_nm): + """Return exact Gaussian probabilities for source wavelengths and bins. + + The result has shape ``wavelengths_nm.shape + (number_of_bins,)``. This is + a probability *mass* in each detector bin, not a point-sampled Gaussian. + ``irf_sigma_nm`` must be strictly positive; the integration result exposes a + diagnostic flag instead of silently accepting a zero or negative value. + """ + + wavelengths_nm = jnp.asarray(wavelengths_nm) + detector_edges_nm = jnp.asarray(detector_edges_nm) + irf_sigma_nm = jnp.asarray(irf_sigma_nm, dtype=wavelengths_nm.dtype) + + z = (detector_edges_nm - wavelengths_nm[..., None]) / irf_sigma_nm + cdf = ndtr(z) + return cdf[..., 1:] - cdf[..., :-1] + + +def _epsilon_real(terms_fn: SpectralTerms, wavelength_nm): + """Evaluate scalar ``Re(epsilon)`` while retaining a vector callback API.""" + + _, epsilon = terms_fn(jnp.reshape(wavelength_nm, (1,))) + if epsilon.shape != (1,): + raise ValueError("terms_fn epsilon must have the same 1-D shape as its wavelength input") + return jnp.real(epsilon[0]) + + +def _fixed_bisection(function, lower, upper, iterations: int): + """Fixed-work scalar bisection used as the primal custom-root solve.""" + + f_lower = function(lower) + f_upper = function(upper) + + lower_is_root = f_lower == 0 + upper_is_root = f_upper == 0 + collapsed = jnp.where(lower_is_root, lower, jnp.where(upper_is_root, upper, lower)) + lower = jnp.where(lower_is_root | upper_is_root, collapsed, lower) + upper = jnp.where(lower_is_root | upper_is_root, collapsed, upper) + f_lower = jnp.where(lower_is_root | upper_is_root, 0, f_lower) + + def body(_, state): + lo, hi, f_lo = state + midpoint = lo + 0.5 * (hi - lo) + f_mid = function(midpoint) + crosses_left = (f_mid == 0) | (jnp.signbit(f_lo) != jnp.signbit(f_mid)) + next_lo = jnp.where(crosses_left, lo, midpoint) + next_hi = jnp.where(crosses_left, midpoint, hi) + next_f_lo = jnp.where(crosses_left, f_lo, f_mid) + return next_lo, next_hi, next_f_lo + + lower, upper, _ = lax.fori_loop(0, iterations, body, (lower, upper, f_lower)) + return lower + 0.5 * (upper - lower) + + +def _implicit_bisection(function, lower, upper, iterations: int): + """Find a root by bisection and differentiate the defining equation. + + Differentiating the bisection decisions would make a root appear locally + fixed inside its scan panel. ``lax.custom_root`` instead applies the + implicit-function tangent ``dx = -df_parameter / df_x``. Brackets and root + topology are stopped because changing topology is genuinely discrete. + """ + + lower = lax.stop_gradient(lower) + upper = lax.stop_gradient(upper) + initial = lower + 0.5 * (upper - lower) + + def solve(f, _): + return _fixed_bisection(f, lower, upper, iterations) + + def tangent_solve(linearized_function, right_hand_side): + slope = linearized_function(jnp.ones_like(right_hand_side)) + safe_slope = jnp.where(slope != 0, slope, jnp.ones_like(slope)) + return right_hand_side / safe_slope + + return lax.custom_root(function, initial, solve, tangent_solve) + + +def _evaluate_weighted_density( + terms_fn: SpectralTerms, + wavelengths_nm, + detector_edges_nm, + safe_irf_sigma_nm, +): + """Evaluate detector-weighted density and a finite flag at flat nodes.""" + + wavelengths_nm = jnp.ravel(wavelengths_nm) + numerator, epsilon = terms_fn(wavelengths_nm) + numerator = jnp.asarray(numerator) + epsilon = jnp.asarray(epsilon) + + if epsilon.shape != wavelengths_nm.shape: + raise ValueError("terms_fn epsilon must have the same 1-D shape as its wavelength input") + if numerator.ndim == 0 or numerator.shape[0] != wavelengths_nm.shape[0]: + raise ValueError("terms_fn numerator must have wavelength as its first axis") + + denominator = jnp.abs(epsilon) ** 2 + denominator = denominator.reshape(denominator.shape + (1,) * (numerator.ndim - 1)) + density = numerator / denominator + probabilities = gaussian_bin_probabilities(wavelengths_nm, detector_edges_nm, safe_irf_sigma_nm) + probabilities = probabilities.reshape( + probabilities.shape + (1,) * (numerator.ndim - 1) + ) + weighted = probabilities * density[:, None, ...] + + finite = jnp.isfinite(wavelengths_nm) & jnp.isfinite(epsilon) + finite = finite & jnp.all(jnp.isfinite(numerator).reshape((wavelengths_nm.size, -1)), axis=1) + finite = finite & jnp.all(jnp.isfinite(weighted).reshape((wavelengths_nm.size, -1)), axis=1) + safe_weighted = jnp.where(jnp.isfinite(weighted), weighted, jnp.zeros_like(weighted)) + return safe_weighted, finite + + +def _regular_panel_integrals( + terms_fn: SpectralTerms, + panel_edges_nm, + detector_edges_nm, + safe_irf_sigma_nm, + order: int, +): + """Integrate every fixed scan panel with ordinary Gauss--Legendre nodes.""" + + numpy_nodes, numpy_weights = _legendre_rule(order) + nodes = jnp.asarray(numpy_nodes, dtype=panel_edges_nm.dtype) + weights = jnp.asarray(numpy_weights, dtype=panel_edges_nm.dtype) + + lower = panel_edges_nm[:-1] + upper = panel_edges_nm[1:] + midpoint = 0.5 * (lower + upper) + half_width = 0.5 * (upper - lower) + wavelengths = midpoint[:, None] + half_width[:, None] * nodes[None, :] + weighted, all_finite = _evaluate_weighted_density( + terms_fn, wavelengths, detector_edges_nm, safe_irf_sigma_nm + ) + + output_shape = (wavelengths.shape[0], wavelengths.shape[1]) + weighted.shape[1:] + weighted = weighted.reshape(output_shape) + quadrature_weights = half_width[:, None] * weights[None, :] + quadrature_weights = quadrature_weights.reshape( + quadrature_weights.shape + (1,) * (weighted.ndim - 2) + ) + panel_integrals = jnp.sum(quadrature_weights * weighted, axis=1) + finite_panels = jnp.all(all_finite.reshape(wavelengths.shape), axis=1) + return panel_integrals, finite_panels + + +def _tan_mapped_panel_integrals( + terms_fn: SpectralTerms, + lower_nm, + upper_nm, + center_nm, + half_width_nm, + detector_edges_nm, + safe_irf_sigma_nm, + order: int, +): + """Integrate selected panels after ``lambda = c + gamma tan(t)``.""" + + numpy_nodes, numpy_weights = _legendre_rule(order) + nodes = jnp.asarray(numpy_nodes, dtype=lower_nm.dtype) + weights = jnp.asarray(numpy_weights, dtype=lower_nm.dtype) + + t_lower = jnp.arctan((lower_nm - center_nm) / half_width_nm) + t_upper = jnp.arctan((upper_nm - center_nm) / half_width_nm) + t_midpoint = 0.5 * (t_lower + t_upper) + t_half_width = 0.5 * (t_upper - t_lower) + t_nodes = t_midpoint[:, None] + t_half_width[:, None] * nodes[None, :] + + cos_t = jnp.cos(t_nodes) + wavelengths = center_nm[:, None] + half_width_nm[:, None] * jnp.tan(t_nodes) + jacobian = half_width_nm[:, None] / (cos_t * cos_t) + mapped_weights = t_half_width[:, None] * weights[None, :] * jacobian + + weighted, all_finite = _evaluate_weighted_density( + terms_fn, wavelengths, detector_edges_nm, safe_irf_sigma_nm + ) + output_shape = (wavelengths.shape[0], wavelengths.shape[1]) + weighted.shape[1:] + weighted = weighted.reshape(output_shape) + mapped_weights = mapped_weights.reshape( + mapped_weights.shape + (1,) * (weighted.ndim - 2) + ) + panel_integrals = jnp.sum(mapped_weights * weighted, axis=1) + finite_panels = jnp.all(all_finite.reshape(wavelengths.shape), axis=1) + return panel_integrals, finite_panels + + +def _root_crossings(real_epsilon): + """Mark each root once, including roots exactly on scan-panel boundaries.""" + + left = real_epsilon[:-1] + right = real_epsilon[1:] + finite = jnp.isfinite(left) & jnp.isfinite(right) + strict_crossing = (left != 0) & (right != 0) & (jnp.signbit(left) != jnp.signbit(right)) + + # A root on a shared scan edge belongs to the panel that starts there. The + # right endpoint is included only for the final panel, avoiding duplicates. + left_root = left == 0 + right_root_at_end = jnp.arange(left.shape[0]) == left.shape[0] - 1 + right_root_at_end = right_root_at_end & (right == 0) + return finite & (strict_crossing | left_root | right_root_at_end) + + +def _selected_root_panels(crossings, max_roots: int): + panel_ids = jnp.arange(crossings.shape[0], dtype=jnp.int32) + sentinel = jnp.asarray(crossings.shape[0], dtype=jnp.int32) + candidates = jnp.where(crossings, panel_ids, sentinel) + selected = jnp.sort(candidates)[:max_roots] + active = selected < sentinel + return jnp.minimum(selected, sentinel - 1), active + + +def _shifted_panel_edges(lower_nm, upper_nm, panel_count: int, phase): + """Return fixed-count panel edges with a common interior phase shift.""" + + uniform_edges_nm = jnp.linspace(lower_nm, upper_nm, panel_count + 1) + panel_step_nm = (upper_nm - lower_nm) / panel_count + return jnp.concatenate( + ( + uniform_edges_nm[:1], + uniform_edges_nm[1:-1] + phase * panel_step_nm, + uniform_edges_nm[-1:], + ) + ) + + +def _integration_panel_edges( + lower_nm, + upper_nm, + integration_panels: int, + phase, + integration_breakpoints_nm, +): + """Build coarse integration edges while retaining exact static breakpoints. + + Breakpoints replace distinct interior phased edges. Each breakpoint greedily + removes the nearest base edge that has not already been removed; the remaining + base edges and exact breakpoints are then sorted together. This keeps the number + of integration panels fixed (and therefore keeps JIT topology and memory fixed), + while allowing arbitrarily close ordered discontinuities to share what was one + base panel. Candidate edges that are not strictly monotone are reported as + invalid. + """ + + base_edges_nm = _shifted_panel_edges(lower_nm, upper_nm, integration_panels, phase) + if integration_breakpoints_nm is None: + return base_edges_nm, jnp.asarray(False) + + breakpoints_nm = jnp.asarray(integration_breakpoints_nm, dtype=base_edges_nm.dtype) + if breakpoints_nm.ndim != 1: + raise ValueError("integration_breakpoints_nm must be one-dimensional") + breakpoint_count = breakpoints_nm.shape[0] + if breakpoint_count == 0: + return base_edges_nm, jnp.asarray(False) + if breakpoint_count > integration_panels - 1: + raise ValueError( + "integration_breakpoints_nm has more entries than available interior " + "integration-panel edges" + ) + + finite = jnp.all(jnp.isfinite(breakpoints_nm)) + interior = jnp.all((breakpoints_nm > lower_nm) & (breakpoints_nm < upper_nm)) + ordered = jnp.all(jnp.diff(breakpoints_nm) > 0) + + # Keep all index calculations finite even when a traced breakpoint is invalid. + fallback_breakpoints = lower_nm + (upper_nm - lower_nm) * ( + jnp.arange(1, breakpoint_count + 1, dtype=base_edges_nm.dtype) + / (breakpoint_count + 1) + ) + safe_breakpoints = jnp.where(jnp.isfinite(breakpoints_nm), breakpoints_nm, fallback_breakpoints) + interior_edges_nm = base_edges_nm[1:-1] + + def remove_nearest_edge(breakpoint_index, removed): + distances = jnp.abs(interior_edges_nm - safe_breakpoints[breakpoint_index]) + distances = jnp.where(removed, jnp.asarray(jnp.inf, dtype=distances.dtype), distances) + removal_index = lax.stop_gradient(jnp.argmin(distances)) + return removed.at[removal_index].set(True) + + removed = lax.fori_loop( + 0, + breakpoint_count, + remove_nearest_edge, + jnp.zeros(interior_edges_nm.shape, dtype=bool), + ) + retained_or_infinite = jnp.where( + removed, + jnp.asarray(jnp.inf, dtype=interior_edges_nm.dtype), + interior_edges_nm, + ) + candidate_interior_edges_nm = jnp.sort( + jnp.concatenate((retained_or_infinite, safe_breakpoints)) + )[: integration_panels - 1] + candidate_edges_nm = jnp.concatenate( + (base_edges_nm[:1], candidate_interior_edges_nm, base_edges_nm[-1:]) + ) + candidate_monotone = jnp.all(jnp.diff(candidate_edges_nm) > 0) + invalid = ~(finite & interior & ordered & candidate_monotone) + return jnp.where(invalid, base_edges_nm, candidate_edges_nm), invalid + + +def _candidate_root_segments( + root_panels, + roots_nm, + root_mask, + valid_mapping, + panel_edges_nm, + neighbor_panels: int, +): + """Build disjoint tan-mapped segments for every root/coarse-panel pair. + + Multiple fine-scan roots may occupy one coarse integration panel. All roots + whose coarse neighborhoods include that panel are ordered by their already + sorted root ids, and the panel is split at pairwise root midpoints. Thus each + portion of a replaced panel is integrated exactly once without forcing the + smooth background onto the fine root-scan grid. + """ + + n_panels = panel_edges_nm.shape[0] - 1 + + offsets = jnp.arange(-neighbor_panels, neighbor_panels + 1, dtype=jnp.int32) + raw_panels = root_panels[:, None] + offsets[None, :] + in_bounds = (raw_panels >= 0) & (raw_panels < n_panels) + candidate_valid = root_mask[:, None] & valid_mapping[:, None] & in_bounds + + candidate_panels = jnp.clip(raw_panels, 0, n_panels - 1).reshape(-1) + candidate_valid = candidate_valid.reshape(-1) + repeats = 2 * neighbor_panels + 1 + candidate_root_ids = jnp.repeat(jnp.arange(root_panels.shape[0], dtype=jnp.int32), repeats) + candidate_roots_nm = jnp.repeat(roots_nm, repeats) + + same_panel = candidate_panels[:, None] == candidate_panels[None, :] + valid_other = candidate_valid[None, :] + previous_other = ( + same_panel + & valid_other + & (candidate_root_ids[None, :] < candidate_root_ids[:, None]) + ) + next_other = ( + same_panel + & valid_other + & (candidate_root_ids[None, :] > candidate_root_ids[:, None]) + ) + source_span_nm = panel_edges_nm[-1] - panel_edges_nm[0] + lower_sentinel_nm = panel_edges_nm[0] - source_span_nm + upper_sentinel_nm = panel_edges_nm[-1] + source_span_nm + previous_root_nm = jnp.max( + jnp.where(previous_other, candidate_roots_nm[None, :], lower_sentinel_nm), axis=1 + ) + next_root_nm = jnp.min( + jnp.where(next_other, candidate_roots_nm[None, :], upper_sentinel_nm), axis=1 + ) + has_previous = jnp.any(previous_other, axis=1) + has_next = jnp.any(next_other, axis=1) + safe_previous_root_nm = jnp.where(has_previous, previous_root_nm, candidate_roots_nm) + safe_next_root_nm = jnp.where(has_next, next_root_nm, candidate_roots_nm) + + panel_lower_nm = panel_edges_nm[candidate_panels] + panel_upper_nm = panel_edges_nm[candidate_panels + 1] + segment_lower_nm = jnp.where( + has_previous, + jnp.maximum(panel_lower_nm, 0.5 * (safe_previous_root_nm + candidate_roots_nm)), + panel_lower_nm, + ) + segment_upper_nm = jnp.where( + has_next, + jnp.minimum(panel_upper_nm, 0.5 * (candidate_roots_nm + safe_next_root_nm)), + panel_upper_nm, + ) + candidate_use = candidate_valid & (segment_upper_nm > segment_lower_nm) + return candidate_panels, segment_lower_nm, segment_upper_nm, candidate_use + + +def integrate_detector_bins( + terms_fn: SpectralTerms, + detector_edges_nm, + irf_sigma_nm, + *, + source_bounds_nm=None, + root_scan_panels: int = 4096, + integration_panels: int = 256, + integration_breakpoints_nm=None, + regular_order: int = 8, + root_order: int = 32, + max_roots: int = 16, + neighbor_panels: int = 1, + bisection_iterations: int = 48, + tail_sigma: float = 6.0, + scan_phase: float = 0.0, +) -> ResonanceQuadratureResult: + """Integrate one possibly unresolved spectrum into detector bins. + + Args: + terms_fn: Callable taking a 1-D wavelength array in nm and returning + ``(numerator, epsilon)``. The numerator's first axis and epsilon's + sole axis must match the wavelength array. The intrinsic spectral + density is ``numerator / abs(epsilon)**2``. + detector_edges_nm: Strictly increasing detector-bin edges. Nonuniform + bins are supported. + irf_sigma_nm: Positive Gaussian IRF standard deviation in nm. + source_bounds_nm: Optional two-element integration interval. By default + the detector range is extended by ``tail_sigma * irf_sigma_nm`` on + each side. + root_scan_panels: Number of fine fixed panels used only to bracket roots. + integration_panels: Number of coarse panels used to cover the source + interval. Root neighborhoods split these panels as needed, so more + than one root may occupy one integration panel. + integration_breakpoints_nm: Optional ordered one-dimensional array of + static source-wavelength discontinuities. Each value replaces a distinct + nearby coarse interior edge, making rectangular transmission-filter + boundaries exact without refining the root scan or regular coverage. + regular_order: Gauss--Legendre order outside root neighborhoods. + root_order: Even Gauss--Legendre order used on tan-mapped panels. An + even order avoids sampling the mapped resonance center exactly. + max_roots: Static root capacity. More roots set ``root_overflow``. + neighbor_panels: Number of coarse integration panels on either side of + each root panel to tan-map. + bisection_iterations: Fixed bisection work per root. + tail_sigma: Default source-domain extension in Gaussian sigmas. + scan_phase: Shift of every interior fine-scan and coarse-integration + boundary as a fraction of its respective panel width, while preserving + the source endpoints and exact integration breakpoints. Values must lie + strictly between -1 and 1. This is primarily useful for convergence + checks against arbitrary numerical-grid phase. + + Returns: + :class:`ResonanceQuadratureResult`. ``bin_mean`` is a spectral density, + so detector-bin area is ``sum(bin_mean * diff(detector_edges_nm))``. + + Notes: + Integer controls are compile-time topology. Close over them when using + :func:`jax.jit`. Root-panel selection is nondifferentiable, but root + locations use implicit differentiation and the tan-mapped nodes and + weights retain their full derivatives. + """ + + if root_scan_panels < 1: + raise ValueError("root_scan_panels must be positive") + if integration_panels < 1: + raise ValueError("integration_panels must be positive") + if regular_order < 1: + raise ValueError("regular_order must be positive") + if root_order < 2 or root_order % 2: + raise ValueError("root_order must be a positive even integer") + if max_roots < 1: + raise ValueError("max_roots must be positive") + if neighbor_panels < 0: + raise ValueError("neighbor_panels must be nonnegative") + if bisection_iterations < 1: + raise ValueError("bisection_iterations must be positive") + if tail_sigma < 0: + raise ValueError("tail_sigma must be nonnegative") + + detector_edges_nm = jnp.asarray(detector_edges_nm) + if detector_edges_nm.ndim != 1 or detector_edges_nm.shape[0] < 2: + raise ValueError("detector_edges_nm must be a 1-D array with at least two entries") + dtype = jnp.result_type(detector_edges_nm.dtype, jnp.asarray(irf_sigma_nm).dtype, jnp.float32) + detector_edges_nm = detector_edges_nm.astype(dtype) + irf_sigma_nm = jnp.asarray(irf_sigma_nm, dtype=dtype) + if irf_sigma_nm.ndim != 0: + raise ValueError("irf_sigma_nm must be scalar") + scan_phase = jnp.asarray(scan_phase, dtype=dtype) + if scan_phase.ndim != 0: + raise ValueError("scan_phase must be scalar") + + bin_widths = jnp.diff(detector_edges_nm) + invalid_bin_width = jnp.any(bin_widths == 0) + nonmonotonic_edges = jnp.any(bin_widths < 0) + invalid_irf_sigma = (~jnp.isfinite(irf_sigma_nm)) | (irf_sigma_nm <= 0) + invalid_scan_phase = (~jnp.isfinite(scan_phase)) | (jnp.abs(scan_phase) >= 1) + safe_irf_sigma_nm = jnp.where(invalid_irf_sigma, jnp.asarray(1, dtype=dtype), irf_sigma_nm) + safe_scan_phase = jnp.where(invalid_scan_phase, jnp.asarray(0, dtype=dtype), scan_phase) + + if source_bounds_nm is None: + source_lower = detector_edges_nm[0] - tail_sigma * safe_irf_sigma_nm + source_upper = detector_edges_nm[-1] + tail_sigma * safe_irf_sigma_nm + else: + source_bounds_nm = jnp.asarray(source_bounds_nm, dtype=dtype) + if source_bounds_nm.shape != (2,): + raise ValueError("source_bounds_nm must contain exactly two entries") + source_lower, source_upper = source_bounds_nm + + invalid_source_bounds = ( + (~jnp.isfinite(source_lower)) + | (~jnp.isfinite(source_upper)) + | (source_upper <= source_lower) + ) + safe_source_lower = jnp.where(invalid_source_bounds, detector_edges_nm[0], source_lower) + fallback_span = jnp.maximum(jnp.abs(detector_edges_nm[-1] - detector_edges_nm[0]), 1) + safe_source_upper = jnp.where( + invalid_source_bounds, safe_source_lower + fallback_span, source_upper + ) + root_scan_edges_nm = _shifted_panel_edges( + safe_source_lower, safe_source_upper, root_scan_panels, safe_scan_phase + ) + integration_edges_nm, invalid_integration_breakpoints = _integration_panel_edges( + safe_source_lower, + safe_source_upper, + integration_panels, + safe_scan_phase, + integration_breakpoints_nm, + ) + + _, epsilon_scan = terms_fn(root_scan_edges_nm) + epsilon_scan = jnp.asarray(epsilon_scan) + if epsilon_scan.shape != root_scan_edges_nm.shape: + raise ValueError("terms_fn epsilon must have the same 1-D shape as its wavelength input") + real_epsilon_scan = jnp.real(epsilon_scan) + crossings = _root_crossings(real_epsilon_scan) + root_count = jnp.sum(crossings, dtype=jnp.int32) + used_root_count = jnp.minimum(root_count, max_roots) + root_overflow = root_count > max_roots + scan_nonfinite = ~jnp.all(jnp.isfinite(epsilon_scan)) + + def no_root_branch(_): + regular, integration_finite = _regular_panel_integrals( + terms_fn, + integration_edges_nm, + detector_edges_nm, + safe_irf_sigma_nm, + regular_order, + ) + integral = jnp.sum(regular, axis=0) + roots = jnp.full((max_roots,), safe_source_lower, dtype=dtype) + centers = roots + widths = jnp.zeros((max_roots,), dtype=dtype) + mask = jnp.zeros((max_roots,), dtype=bool) + return integral, ~jnp.all(integration_finite), False, roots, centers, widths, mask + + def root_branch(_): + root_scan_panel_ids, root_mask = _selected_root_panels(crossings, max_roots) + lower = root_scan_edges_nm[root_scan_panel_ids] + upper = root_scan_edges_nm[root_scan_panel_ids + 1] + dummy = lower + 0.5 * (upper - lower) + + def solve_one(lo, hi, is_active, dummy_root): + def equation(wavelength_nm): + residual = _epsilon_real(terms_fn, wavelength_nm) + return jnp.where(is_active, residual, wavelength_nm - dummy_root) + + return _implicit_bisection(equation, lo, hi, bisection_iterations) + + roots = jax.vmap(solve_one)(lower, upper, root_mask, dummy) + + def epsilon_scalar(wavelength_nm): + _, epsilon = terms_fn(jnp.reshape(wavelength_nm, (1,))) + return epsilon[0] + + epsilon_at_root = jax.vmap(epsilon_scalar)(roots) + epsilon_slope = jax.vmap( + lambda root: jax.jvp(epsilon_scalar, (root,), (jnp.ones_like(root),))[1] + )(roots) + slope_norm = jnp.abs(epsilon_slope) ** 2 + safe_slope_norm = jnp.where(slope_norm > 0, slope_norm, jnp.ones_like(slope_norm)) + product = epsilon_at_root * jnp.conj(epsilon_slope) + centers = roots - jnp.real(product) / safe_slope_norm + widths = jnp.abs(jnp.imag(product)) / safe_slope_norm + + root_finite = ( + jnp.isfinite(roots) + & jnp.isfinite(epsilon_at_root) + & jnp.isfinite(epsilon_slope) + & jnp.isfinite(centers) + & jnp.isfinite(widths) + ) + zero_width_roots = root_mask & ((slope_norm == 0) | (widths <= 0)) + valid_mapping = root_mask & root_finite & ~zero_width_roots + + regular, regular_finite = _regular_panel_integrals( + terms_fn, + integration_edges_nm, + detector_edges_nm, + safe_irf_sigma_nm, + regular_order, + ) + # Root-to-coarse-panel assignment is discrete topology. The solved root and + # all mapped nodes/weights remain differentiable after that assignment. + integration_root_panels = jnp.searchsorted( + lax.stop_gradient(integration_edges_nm), + lax.stop_gradient(roots), + side="right", + ) - 1 + integration_root_panels = jnp.clip( + integration_root_panels, 0, integration_panels - 1 + ).astype(jnp.int32) + ( + candidate_panels, + candidate_lower, + candidate_upper, + candidate_use, + ) = _candidate_root_segments( + integration_root_panels, + roots, + root_mask, + valid_mapping, + integration_edges_nm, + neighbor_panels, + ) + repeats = 2 * neighbor_panels + 1 + candidate_centers = jnp.repeat(centers, repeats) + candidate_widths = jnp.repeat(widths, repeats) + + # Empty/inactive candidate segments are still evaluated under JAX's static + # control flow. Give them a well-conditioned full-panel mapping so masked + # reverse-mode cotangents cannot encounter 0/tiny intermediate derivatives. + fallback_lower = integration_edges_nm[candidate_panels] + fallback_upper = integration_edges_nm[candidate_panels + 1] + safe_candidate_lower = jnp.where(candidate_use, candidate_lower, fallback_lower) + safe_candidate_upper = jnp.where(candidate_use, candidate_upper, fallback_upper) + panel_midpoints = 0.5 * (safe_candidate_lower + safe_candidate_upper) + safe_centers = jnp.where(candidate_use, candidate_centers, panel_midpoints) + safe_widths = jnp.where( + candidate_use, + candidate_widths, + safe_candidate_upper - safe_candidate_lower, + ) + mapped, mapped_finite = _tan_mapped_panel_integrals( + terms_fn, + safe_candidate_lower, + safe_candidate_upper, + safe_centers, + safe_widths, + detector_edges_nm, + safe_irf_sigma_nm, + root_order, + ) + + component_dims = (1,) * (regular.ndim - 1) + candidate_mask = candidate_use.reshape(candidate_use.shape + component_dims) + mapped_by_panel = jnp.zeros_like(regular).at[candidate_panels].add( + jnp.where(candidate_mask, mapped, jnp.zeros_like(mapped)) + ) + replaced = ( + jnp.zeros((integration_panels,), dtype=jnp.int32) + .at[candidate_panels] + .add(candidate_use.astype(jnp.int32)) + > 0 + ) + selected_panels = jnp.where( + replaced.reshape(replaced.shape + component_dims), mapped_by_panel, regular + ) + integral = jnp.sum(selected_panels, axis=0) + + regular_nonfinite = jnp.any((~replaced) & ~regular_finite) + mapped_nonfinite = jnp.any(candidate_use & ~mapped_finite) + integration_nonfinite = regular_nonfinite | mapped_nonfinite + root_nonfinite = jnp.any(root_mask & ~root_finite) + return ( + integral, + integration_nonfinite | root_nonfinite, + jnp.any(zero_width_roots), + roots, + centers, + widths, + root_mask, + ) + + branch_result = lax.cond(root_count > 0, root_branch, no_root_branch, operand=None) + integral, branch_nonfinite, zero_width, roots, centers, widths, root_mask = branch_result + + safe_bin_widths = jnp.where(bin_widths > 0, bin_widths, jnp.ones_like(bin_widths)) + width_shape = safe_bin_widths.shape + (1,) * (integral.ndim - 1) + bin_mean = integral / safe_bin_widths.reshape(width_shape) + output_nonfinite = ~jnp.all(jnp.isfinite(bin_mean)) + nonfinite = ( + scan_nonfinite + | branch_nonfinite + | output_nonfinite + | ~jnp.all(jnp.isfinite(detector_edges_nm)) + ) + fatal = ( + root_overflow + | nonfinite + | zero_width + | invalid_bin_width + | nonmonotonic_edges + | invalid_irf_sigma + | invalid_source_bounds + | invalid_scan_phase + | invalid_integration_breakpoints + ) + bin_mean = jnp.where(fatal, jnp.full_like(bin_mean, jnp.nan), bin_mean) + + diagnostics = ResonanceQuadratureDiagnostics( + root_count=root_count, + used_root_count=used_root_count, + root_overflow=root_overflow, + nonfinite=nonfinite, + zero_width=zero_width, + invalid_bin_width=invalid_bin_width, + nonmonotonic_edges=nonmonotonic_edges, + invalid_irf_sigma=invalid_irf_sigma, + invalid_source_bounds=invalid_source_bounds, + invalid_scan_phase=invalid_scan_phase, + invalid_integration_breakpoints=invalid_integration_breakpoints, + roots_nm=roots, + resonance_centers_nm=centers, + resonance_half_widths_nm=widths, + root_mask=root_mask, + ) + return ResonanceQuadratureResult(bin_mean=bin_mean, diagnostics=diagnostics) + + +def raise_for_diagnostics(result: ResonanceQuadratureResult) -> None: + """Raise a host-side exception if a completed result is not trustworthy. + + This helper is intentionally separate from the JIT-compatible kernel. Call + it after transferring a result to the host when fail-fast behavior is useful. + """ + + diagnostics = jax.device_get(result.diagnostics) + failures = [] + if bool(diagnostics.root_overflow): + failures.append( + f"detected {int(diagnostics.root_count)} roots, capacity is {int(diagnostics.used_root_count)}" + ) + if bool(diagnostics.nonfinite): + failures.append("non-finite input, dielectric, integrand, or quadrature output") + if bool(diagnostics.zero_width): + failures.append("a detected resonance has zero local half-width") + if bool(diagnostics.invalid_bin_width): + failures.append("detector edges contain a zero-width bin") + if bool(diagnostics.nonmonotonic_edges): + failures.append("detector edges are not increasing") + if bool(diagnostics.invalid_irf_sigma): + failures.append("Gaussian IRF sigma must be finite and positive") + if bool(diagnostics.invalid_source_bounds): + failures.append("source bounds must be finite and increasing") + if bool(diagnostics.invalid_scan_phase): + failures.append("scan phase must be finite and strictly between -1 and 1") + if bool(diagnostics.invalid_integration_breakpoints): + failures.append( + "integration breakpoints must be finite, strictly ordered, interior, and " + "form strictly monotone integration-panel edges" + ) + if failures: + raise ValueError("invalid resonance quadrature: " + "; ".join(failures)) + + +__all__ = [ + "ResonanceQuadratureDiagnostics", + "ResonanceQuadratureResult", + "gaussian_bin_probabilities", + "integrate_detector_bins", + "raise_for_diagnostics", +] diff --git a/tsadar/core/thomson_diagnostic.py b/tsadar/core/thomson_diagnostic.py index 5afed0dc1..9a7568f15 100644 --- a/tsadar/core/thomson_diagnostic.py +++ b/tsadar/core/thomson_diagnostic.py @@ -1,5 +1,6 @@ """ThomsonScatteringDiagnostic: wraps FitModel with instrument-response effects (spectrometer/angular IRFs) to turn a theoretical Thomson spectrum into the synthetic-detector-space spectrum comparable to real data.""" +import numpy as np from jax import numpy as jnp, vmap from scipy.signal import find_peaks @@ -32,6 +33,50 @@ def _bin_average(arr, step, axis): return jnp.moveaxis(arr, 0, axis) +def _ensure_angular_detector_edges(cfg): + """Populate forward-only ARTS wavelength edges when no calibration supplied them. + + In data-fitting mode :func:`tsadar.data.prepare.prepare_data` stores exact calibrated + edges before this class is constructed. Forward decks have only ``lamrangE``; its + endpoints retain their documented legacy meaning as the first and last wavelength + centers in the calculation. Before data preparation, calibration defines ``CCDsize`` + as ``[wavelength, angle]``, so ``CCDsize[0]`` is the number of spectral bins. + Prepared data already carries explicit edges and returns above. + """ + + if ( + cfg["other"]["extraoptions"]["spectype"] != "angular_full" + or not cfg["data"]["load_ele_spec"] + ): + return + + detector_specs = cfg["other"]["detector_specs"] + if "electron_wavelength_edges" in detector_specs: + return + + lower, upper = (float(value) for value in cfg["other"]["lamrangE"]) + n_spectral_bins = int(cfg["other"]["CCDsize"][0]) + if not lower < upper: + raise ValueError(f"ARTS wavelength bounds must be increasing, got [{lower}, {upper}].") + if n_spectral_bins < 1: + raise ValueError(f"ARTS must have at least one spectral detector bin, got {n_spectral_bins}.") + if n_spectral_bins == 1: + centers = np.asarray([0.5 * (lower + upper)]) + edges = np.asarray([lower, upper]) + else: + centers = np.linspace(lower, upper, n_spectral_bins) + spacing = np.diff(centers) + edges = np.concatenate( + ( + centers[:1] - 0.5 * spacing[:1], + centers[:-1] + 0.5 * spacing, + centers[-1:] + 0.5 * spacing[-1:], + ) + ) + detector_specs["electron_wavelength_edges"] = edges + detector_specs["electron_wavelength_centers"] = centers + + def _irfs_from_config(cfg, scattering_angles): """Adapter from the input deck to device-agnostic IRF descriptions. @@ -109,6 +154,7 @@ def __init__(self, cfg, scattering_angles): super().__init__() self.cfg = cfg self.scattering_angles = scattering_angles + _ensure_angular_detector_edges(cfg) self.model = FitModel(cfg, scattering_angles) self.ele_irf, self.ion_irf, self.ats_irf = _irfs_from_config(cfg, scattering_angles) @@ -151,7 +197,13 @@ def postprocess_theory(self, modlE, modlI, lamAxisE, lamAxisI, amps, TSins): if self.cfg["data"]["load_ele_spec"]: if self.cfg["other"]["extraoptions"]["spectype"] == "angular_full": - lamAxisE, ThryE = irf.add_ATS_IRF(self.ats_irf, lamAxisE, modlE, TSins) + lamAxisE, ThryE = irf.add_ATS_IRF( + self.ats_irf, + lamAxisE, + modlE, + TSins, + apply_spectral_blur=not self.model.electron_spectrum_is_detector_binned, + ) else: lamAxisE, ThryE = irf.add_electron_IRF(self.ele_irf, lamAxisE, modlE, amps["e_amps"], TSins) else: @@ -174,13 +226,33 @@ def reduce_ATS_to_resunit(self, ThryE, lamAxisE, TSins, batch): lamAxisE: the input wavelength axis integrated over a wavelngth resolution unit and correspondingly downsized """ - lam_step = round(ThryE.shape[1] / batch["e_data"].shape[1]) - ang_step = round(ThryE.shape[0] / self.cfg["other"]["CCDsize"][0]) + target_angular_bins = batch["e_data"].shape[0] + if ThryE.shape[0] == target_angular_bins: + ang_step = 1 + else: + ang_step = int(self.cfg["other"]["ang_res_unit"]) + if ang_step < 1: + raise ValueError("ang_res_unit must be a positive integer") + expected_angular_bins = -(-ThryE.shape[0] // ang_step) + if expected_angular_bins != target_angular_bins: + raise ValueError( + "Cannot reduce ARTS angular pixels with the configured resolution " + f"unit: {ThryE.shape[0]} raw rows grouped by {ang_step} gives " + f"{expected_angular_bins}, but the data has {target_angular_bins}." + ) + + if self.model.electron_spectrum_is_detector_binned: + if ThryE.shape[1] != batch["e_data"].shape[1]: + raise ValueError( + "Root-aware ARTS spectrum already has one value per detector bin, " + f"but model/data have {ThryE.shape[1]}/{batch['e_data'].shape[1]} bins." + ) + else: + lam_step = round(ThryE.shape[1] / batch["e_data"].shape[1]) + ThryE = _bin_average(ThryE, lam_step, axis=1) + lamAxisE = _bin_average(lamAxisE, lam_step, axis=0) - 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 = _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( diff --git a/tsadar/data/calibration.py b/tsadar/data/calibration.py index e12a412e5..0124f6a66 100644 --- a/tsadar/data/calibration.py +++ b/tsadar/data/calibration.py @@ -9,6 +9,54 @@ BASE_FILES_PATH = os.path.join(os.path.dirname(__file__), "..", "external") +def detector_edges_from_centers(centers: np.ndarray) -> np.ndarray: + """Return pixel edges for a strictly increasing one-dimensional center axis. + + Interior edges are the midpoints between neighboring centers. The two outer + edges use the adjacent half-spacing. This keeps calibrated plotting coordinates + (centers) distinct from the finite detector support needed by bin quadrature. + """ + + centers = np.asarray(centers) + if centers.ndim != 1: + raise ValueError(f"Detector centers must be one-dimensional, got shape {centers.shape}.") + if centers.size < 2: + raise ValueError("At least two detector centers are required to infer outer pixel edges.") + if not np.all(np.isfinite(centers)): + raise ValueError("Detector centers must all be finite.") + + spacing = np.diff(centers) + if not np.all(spacing > 0): + raise ValueError("Detector centers must be strictly increasing.") + + interior = centers[:-1] + 0.5 * spacing + return np.concatenate( + ( + centers[:1] - 0.5 * spacing[:1], + interior, + centers[-1:] + 0.5 * spacing[-1:], + ) + ) + + +def grouped_detector_edges(centers: np.ndarray, group_size: int) -> np.ndarray: + """Return exact outer edges after averaging consecutive detector pixels. + + The last group may be shorter than ``group_size``. Selecting boundaries from + the original pixel-edge array, rather than re-inferring them from the averaged + centers, preserves that ragged group's true physical width. + """ + + if isinstance(group_size, bool) or not isinstance(group_size, (int, np.integer)) or group_size <= 0: + raise ValueError(f"group_size must be a positive integer, got {group_size!r}.") + + centers = np.asarray(centers) + edges = detector_edges_from_centers(centers) + boundary_indices = np.arange(0, centers.size, int(group_size), dtype=int) + boundary_indices = np.concatenate((boundary_indices, np.array([centers.size], dtype=int))) + return edges[boundary_indices] + + _SCATTERING_ANGLES = { # beam: (sa_start_deg, sa_end_deg, weights) -- all values precalculated for OMEGA TIM6 TS "P9": ( diff --git a/tsadar/data/prepare.py b/tsadar/data/prepare.py index 1825a43e6..e4042f44e 100644 --- a/tsadar/data/prepare.py +++ b/tsadar/data/prepare.py @@ -8,7 +8,12 @@ from .evaluate_background import get_shot_bg from .load_ts_data import loadData from .correct_throughput import correctThroughput -from .calibration import get_calibrations, get_scattering_angles +from .calibration import ( + detector_edges_from_centers, + get_calibrations, + get_scattering_angles, + grouped_detector_edges, +) from .lineouts import get_lineouts from .data_visualizer import launch_data_visualizer from .feature_detector import first_guess @@ -130,6 +135,8 @@ def prepare_data(config: Dict, shotNum: int) -> Dict: config["data"]["lineouts"]["val"] = config["data"]["lineouts"]["val"][: -(num_slices % batch_size)] print(f"final {num_slices % batch_size} lineouts have been removed") + electron_wavelength_edges = None + # extract ARTS section if (config["data"]["lineouts"]["type"] == "range") & (config["other"]["extraoptions"]["spectype"] == "angular"): config["other"]["extraoptions"]["spectype"] = "angular_full" @@ -139,6 +146,7 @@ def prepare_data(config: Dict, shotNum: int) -> Dict: # down sample image to resolution units by summation ang_res_unit = config["other"]["ang_res_unit"] # in pixels lam_res_unit = config["other"]["lam_res_unit"] # in pixels + electron_wavelength_edges = grouped_detector_edges(axisyE, lam_res_unit) data_res_unit = np.array( [np.average(elecData[i : i + lam_res_unit, :], axis=0) for i in range(0, elecData.shape[0], lam_res_unit)] @@ -198,6 +206,14 @@ def prepare_data(config: Dict, shotNum: int) -> Dict: launch_data_visualizer(elecData, ionData, all_data, all_axes, config) config["other"]["detector_specs"]["widIRF"] = stddev + if config["other"]["extraoptions"]["spectype"] == "angular_full": + # Keep the calibrated axis exposed to plotting/loss code as bin centers, but + # retain the finite support of every (possibly ragged) spectral resolution + # unit for detector-bin quadrature. + if electron_wavelength_edges is None: + electron_wavelength_edges = detector_edges_from_centers(np.ravel(axisyE)) + config["other"]["detector_specs"]["electron_wavelength_edges"] = electron_wavelength_edges + config["other"]["detector_specs"]["electron_wavelength_centers"] = np.ravel(axisyE) config["other"]["lamrangE"] = [axisyE[0], axisyE[-1]] config["other"]["lamrangI"] = [axisyI[0], axisyI[-1]] config["other"]["npts"] = int(config["other"]["CCDsize"][1] * config["other"]["points_per_pixel"]) diff --git a/tsadar/forward/calc_series.py b/tsadar/forward/calc_series.py index 6e46c33a3..e3021e0c8 100644 --- a/tsadar/forward/calc_series.py +++ b/tsadar/forward/calc_series.py @@ -51,7 +51,17 @@ def forward_pass(config): 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"]) + # Forward ARTS decks retain the raw calibration convention + # ``CCDsize == [wavelength, angle]`` until the synthetic detector is built. + # Non-angular forward decks retain the historical wavelength-second layout. + spectral_pixels = ( + config["other"]["CCDsize"][0] + if is_angular + else config["other"]["CCDsize"][1] + ) + config["other"]["npts"] = int( + spectral_pixels * config["other"]["points_per_pixel"] + ) sas = get_scattering_angles(config) if not is_angular: @@ -68,13 +78,21 @@ def forward_pass(config): if is_angular: [axisxE, _, _, _, _, _] = get_calibrations( - 104000, config["other"]["extraoptions"]["spectype"], 0.0, config["other"]["CCDsize"] + 104000, + config["other"]["extraoptions"]["spectype"], + 0.0, + config["other"]["CCDsize"], + config["other"]["detector_specs"], ) # shot number hardcoded to get calibration config["other"]["extraoptions"]["spectype"] = "angular_full" sas["angAxis"] = axisxE - dummy_batch["i_data"] = np.ones((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1])) - dummy_batch["e_data"] = np.ones((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1])) + # Raw calibration defines CCDsize as [wavelength, angle], while ARTS model/data + # arrays are [angle, wavelength]. Keep that orientation explicit for non-square + # detectors rather than relying on OMEGA's 1024x1024 shape. + angular_shape = (config["other"]["CCDsize"][1], config["other"]["CCDsize"][0]) + dummy_batch["i_data"] = np.ones(angular_shape) + dummy_batch["e_data"] = np.ones(angular_shape) series={} for species in config["parameters"].keys(): @@ -137,7 +155,11 @@ def forward_pass(config): savedata = plotters.plot_data_angular( config, {"ele": np.squeeze(ThryE)}, - {"e_data": np.zeros((config["other"]["CCDsize"][0], config["other"]["CCDsize"][1]))}, + { + "e_data": np.zeros( + (config["other"]["CCDsize"][1], config["other"]["CCDsize"][0]) + ) + }, {"epw_x": sas["angAxis"], "epw_y": lamAxisE, 'x_label': 'Angle'}, td, )