Repository navigation
Rename gain config keys and nest LPI gain options under other.LPI_gain - #162
Merged
Merged
Conversation
… the electron suseptibility
…end the code will fit from right to left instead of throwing an error
…tandalone postprocessor runner Moves saving of fit results from postprocessing into the main fitter, unifies the previously-duplicated batch-creation logic between fitter and postprocess, and makes the postprocessor optional so a run can be restarted or replotted without redoing the fit. Adds run_postprocessor.py so plots can be remade or run separately after a fit completes, with regression coverage for refit_bad_fits and the new runner. Also updates a test_manifest.py mock to match the new process_data signature, and fixes an unrelated pre-existing Windows-only test failure (mlflow artifact_location needs a file:// URI, not a bare Windows path).
…actor - refit_bad_fits passed the full sample_indices array to one_d_loop with num_batches=1, so it always refit lineout 0 regardless of which lineout was actually flagged as bad. Now passes just the flagged index. - postprocess_runner's replay never applied the ang_res_unit lineout start/end conversion multirun_angular_optax applies during the original fit, so it postprocessed a different (unconverted) lineout range for angular fits. That conversion, and the nvx/window-length growth from multi-pass angular refinement (config["optimizer"]["num_mins"] > 1), are now factored out of loops.py (apply_ang_res_unit, advance_refinement_shape) and replayed before reconstructing the ThomsonParams skeleton, so a refined checkpoint can actually be deserialized and the correct batch is rebuilt. - run_postprocess_local/_load_merged_config and run_postprocess_remote hardcoded defaults.yaml + inputs.yaml, excluding app-originated runs (runner.run_for_app), which only ever log a single config.yaml. Both now check for config.yaml first and fall back to the defaults/inputs pair. Adds regression coverage for all four: tests/test_inverse/test_loops.py for the two extracted helpers, and new cases in test_postprocess_runner.py for the config.yaml artifact layout (both local and remote). Full test suite: 133 passed, 5 skipped (up from 125 -- 8 new tests). Note: separately investigated a report that plain "angular" (non-"range" lineout type) fitting looks non-functional -- build_angular_batch assumes a full CCD frame that only "angular_full" actually has, and there is no one_d_loop-style batched path for it. This predates the work in this PR (traced back to at least the 2024-08 optax-based fitter, unrelated to the postprocessor refactor) and is left unfixed here as a separate, tracked issue.
…n shots The unknown-shot-number fallback in get_calibrations had a raise Warning(...) inside a try/except KeyError, which doesn't actually catch Warning (it isn't a KeyError), so it crashed instead of warning. Switched to warnings.warn(...), and surfaced EPWtcc/IAWtcc from detector_specs for the imaging-calibration fallback path (previously only the non-imaging path set fallback values). Docs updated to describe the detector_specs fallback fields. Also replaces sa_lookup's nine near-identical elif beam == branches with a data-driven dict lookup (verified exact match on all beams and the error path).
launch_data_visualizer opened several figures per shot without ever closing them; over a long run (or many shots) these accumulate and can exhaust memory. Adds plt.close(fig) after each figure is saved, and saves the previously-unsaved temporal-comparison figure. Also switches the lineouts-with-background plot to gate on background:report_background instead of hardcoding background:type in ["fit", "pixel"].
…strings, static typing Correctness fixes: - calc_series.py: fixed a real NameError (undefined fe_val/velocity) and a wrong-argument-count call to plot_dist in the non-angular 2D-distribution plotting branch. - loss_function.py: loss() called self.unravel_pytree, an attribute that was never assigned (only unravel_weights was ever set); loss_functionals() raised an opaque UnboundLocalError for an unrecognized loss_method instead of a clear ValueError. - warpcorr.py: added a clear NotImplementedError for non-EPW instruments instead of a silent crash. - distribution_functions/base.py: calc_moment now raises a clear error for unsupported distribution-function shapes instead of returning an unbound variable. - load_ts_data.py: fixed two invalid \m escape sequences in axis-label strings. Efficiency improvements: - loss_function.py: removed a duplicate forward-model call in calc_loss, vectorized calculate_covariance_matrix with vmap instead of an unrolled Python loop. - form_factor.py: cached the Z-prime table disk reads instead of reloading them on every FormFactor construction. - postprocess.py / calc_series.py: hoisted ThomsonScatteringDiagnostic construction out of per-batch/per-iteration loops. - warpcorr.py: vectorized the pixel-by-pixel dewarping loop (up to ~1M iterations per shot) using np.add.at-based scatter-accumulation; verified byte-exact against the original loop-based implementation on real shot data. - correct_throughput.py: replaced a tile+transpose allocation with plain broadcasting, and a cell-by-cell xlrd read loop with vectorized col_values calls; also drops the deprecated numpy.matlib.repmat import. Both verified exact-match against the originals. Redundancy / simplification: - generate_spectra.py: consolidated ion_spectrum/electron_spectrum and their _detailed counterparts around shared core methods; verified bit-exact against the pristine pre-refactor implementation on both the fast and detailed paths. - plotters.py: factored detailed_lineouts's duplicated worst/best plotting blocks into a shared helper, fixing a real bug uncovered in the process (the "worst" block never set the ion subplot's y-limit due to a copy-paste slip); factored the duplicated white-transition colormap construction. - example_plot.py: deleted -- never functional as written, and nothing referenced it beyond a stale auto-generated docs stub. - edf_movie.py: replaced a hardcoded personal file path with a documented, editable constant. Docstrings: added missing public function/method docstrings across the touched modules, and module-level docstrings to every file in the package that lacked one. Documentation: docs/source/getting_started.rst now documents the standalone postprocessor runner and cluster job submission (queue_tsadar.py, the machine config field, and the required Slurm environment variables), neither of which were previously documented. Full test suite (pytest tests/) passes: 125 passed, 5 skipped, throughout.
…model options Add "brem" as a bg_alg option under background type "fit", fitting only an overall scale and offset with Z/Te/ne fixed at the deck's initial conditions (they are fully degenerate with the scale, which is why the original 5-parameter curve_fit never converged). Also add a "brem_model" background type that instead adds the bremsstrahlung emission directly into the forward model, evaluated at each fit iteration's own Z/Te/ne so they're constrained by the full signal+background spectral shape instead of a separately pre-fit, under-determined background. - tsadar/data/evaluate_background.py, lineouts.py: add the "brem" bg_alg and the "brem_model" no-op background type - tsadar/core/physics/bremsstrahlung.py: shared bremsstrahlung emission model - tsadar/core/modules/ts_params.py: optional brem_amp/brem_c fit parameters - tsadar/core/thomson_diagnostic.py: add the brem background into ThryE/modlE when background type is "brem_model" - tsadar/data/data_visualizer.py: fix two background-plotting assumptions (single pixel location, static per-pixel array) that didn't hold for brem_model - configs/1d: switch the example deck's background to bg_alg "brem"
… support
Introduces a standalone Metropolis-Hastings MCMC postprocessor as an
alternative to the existing Hessian/Laplace-approximation uncertainty,
giving a fuller (non-Gaussian-capable) picture of per-lineout parameter
uncertainty and covariance. It can also fold in uncertainty on instrument
calibration values (gain, IRF widths, dispersion, offset) by pooling
independent chains run under sampled calibration realizations, since none
of those quantities are cheap to resample at MCMC-step granularity without
an invasive rewrite of the core physics forward model.
- New tsadar/inverse/postprocess/ package: the existing postprocessor moves
to .laplace unchanged, re-exported from __init__.py so no call sites
change; .mcmc (sampler core), .mcmc_calibration (calibration draws), and
.mcmc_postprocess (entry point) are new.
- LossFunction gains neg_log_likelihood (shared by the Hessian and MCMC
uncertainty methods) and fixes a pre-existing crash in the Hessian path's
return-value unpacking, uncovered while wiring the sampler's own Laplace
step-size seed.
- New standalone runner (tsadar/mcmc_postprocess_runner.py,
run_mcmc_postprocessor.py) mirrors postprocess_runner.py, sharing its
fit-state reconstruction via a newly extracted _reconstruct_fit_state.
- New config schema under other.mcmc / other.calibration_uncertainty, fully
optional and backward-compatible (an unconfigured deck runs unchanged).
- Known limitation: sampling the electron distribution function ("fe"
active) isn't supported yet -- its parameters are stored per-lineout as a
list of objects rather than a single batched array, which this sampler's
vectorized kernel doesn't handle; raises NotImplementedError rather than
silently mishandling it.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
MCMC postprocessor: - Fix a broken `_reconstruct_fit_state` import that crashed the standalone postprocessor entry point (run_mcmc_postprocessor.py) immediately. - Fix a production bug where Laplace-seeded step-scale calibration silently failed under `vmap` (a `bool()` call on a batching tracer), always falling back to a badly miscalibrated flat default and producing near-frozen chains. - Add tqdm progress reporting (burn-in/sampling step counters, per-chain acceptance-rate logging) so long runs aren't silent. - Add corner plots of the MCMC posterior for a capped, evenly-spaced subset of lineouts; full posterior samples remain saved (`mcmc_samples.nc`) so any other lineout can be corner-plotted after the fact. Adds `corner` as a dependency. - Add `--overrides <path.yaml>` to run_postprocessor.py and run_mcmc_postprocessor.py: a small stub deck merged onto a completed fit's saved config in memory, so postprocessing-only settings (plotting ranges, MCMC settings) can be tweaked without mutating the original saved deck. Adds configs/postprocessor/postprocessor_stub.yaml as a starting template. - Unify `calibration_uncertainty.num_draws` and a new `mcmc.init_dispersion_factor` into one "number of independent chains" concept: chains may differ by calibration, by a perturbed starting point, by both, or by neither (still legitimate, differing only by their own MH noise). `num_draws > 1` with all `*_sigma` at 0 no longer collapses to a single chain. Report Gelman-Rubin R-hat automatically whenever multiple chains are run. - Document all of the above in docs/source/mcmc.rst and defaults.rst. Other staged fixes: - Fix the cosine-decay learning-rate schedule for non-LBFGS optimizers (tsadar/inverse/loops.py). - Fix a relative-import bug in tsadar/inverse/postprocess/laplace.py. - Normalize reversed lineout ranges (start > end) in the data visualizer. - Enable `jax_debug_nans` in run_tsadar.py. - Update configs/1d/inputs.yaml to a different example shot configuration.
…m job queueing script, and improve sampling efficiency with thinning
_run_window was called bare (never wrapped in jit), so every burn-in window and sampling chunk retraced the forward-model jaxpr from scratch even though the compiled XLA kernel underneath was reused. Benchmarked at ~800-900ms of pure Python retracing per call, ~92% of total wall time on a small test fit -- and identically so in the production path (run_mcmc_for_fit_batches' eqx.filter_vmap still calls this bare function fresh per window). filter_jit gives repeat calls with the same shapes a cache hit, ~15x faster with no change to the algorithm or outputs.
…ion. Also updated the R-hat label to note it is a target
Pulls in PR #142 (ARTS2D susceptibility geometry and flow-frame fixes, including the 2D-path Im(chi_e) sign fix and per-ion-species flow handling via _charge_weighted_flow/_electron_resonance) and the ratintn principal-value quadrature fix, along with related config schema and test additions. Manual conflict resolutions: - tsadar/inverse/loss_function.py: kept feature/mcmc's neg_log_likelihood (shared Laplace/MCMC likelihood, per_lineout support) and folded in origin/main's angular-data special case (routes through calc_loss/_angular_data_objective, whose noise model is purpose-built for ARTS2D rather than the generic |data| approximation used elsewhere); per_lineout=True now raises for angular data too, mirroring the existing covar-loss_method guard. - tsadar/utils/manifest.py: kept both branches' independent new manifest-role entries (MCMC uncertainty artifacts, angular objective diagnostics). - tsadar/postprocess_runner.py: kept feature/mcmc's _reconstruct_fit_state helper and ported over origin/main's checkpoint_refinements config override, which the helper was missing. Also fixed tests/test_form_factor/test_arts2d_consistency.py's new per-ion-flow schema test to skip config files without a top-level "parameters" key (e.g. configs/postprocessor/postprocessor_stub.yaml, a feature/mcmc-only partial-override stub for the MCMC postprocessor, never a full parameter deck). Verified: tests/test_form_factor, tests/test_forward (120 passed, 5 skipped), tests/test_inverse/test_angular_objective.py (13 passed), tests/test_inverse/test_mcmc.py + tests/test_mcmc_postprocess_runner.py + tests/test_postprocess_runner.py (26 passed).
The calc_sigmas/compare_to_laplace Hessian was taken over the entire ThomsonParams pytree instead of just the active fit parameters, risking a multi-hundred-GB allocation on an ordinary fit, and get_sigmas indexed the result with dict-style brackets that could never work against its actual (non-dict) shape; the angular path also called a get_plasma_parameters method that didn't exist anywhere in the codebase. Restrict the Hessian to diff_params (mirroring the same restriction mcmc.py's Laplace-seeded step scale already used) and rewrite get_sigmas to match the resulting structure and the column order plotting code expects. Add regression tests covering the fix: no silent fallback, a cross-check of get_sigmas' output against an independent recomputation of the column ordering, and the fe-active guard. Also wire up adapt_shape in the 1D and postprocessor-stub config decks and document it in mcmc.rst: burn-in now also corrects each lineout's relative per-parameter step-scale balance, which the Robbins-Monro magnitude update alone can't fix. Unset JAX_PLATFORMS in the NERSC SLURM launch scripts, and drop a stale dated comment in tsadar/data/calibration.py.
… fix the compile time
… Hessian-diagonal validation test Adds timestamped progress prints around run_mcmc_for_fit_batches' step-scale seeding loop so a slow/stalled Perlmutter run can be distinguished as one-time JIT compile vs. genuinely stuck on a specific fit-batch, and adds test_seed_step_scale_from_laplace_matches_full_hessian_diagonal to verify the grad+jvp diagonal-Hessian trick against a ground-truth full Hessian.
run_mcmc_postprocessor.py (and other postprocess-replay entry points) had several minutes of dead time before any output on NERSC, because tsadar/__init__.py, laplace.py, fitter.py, and mcmc_postprocess.py all eagerly imported tsadar.utils.plotting.plotters (which pulls in plotly/corner/matplotlib) at module top, even though every use of it is confined to a couple of functions that a postprocess-only invocation may not even call. On Lustre, each such import chain does many slow stat()/ open() calls, adding roughly 140s before the process shows any sign of life. tsadar/__init__.py now uses PEP 562 module __getattr__/__dir__ to lazily resolve run_for_app/ThomsonScatteringDiagnostic/ThomsonParams/ get_scattering_angles on first access, and the plotters/postprocess imports in laplace.py, fitter.py, and mcmc_postprocess.py are deferred into the specific functions that actually use them. Pure import reordering, no behavior change. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…, not a flat number _seed_step_scale_from_laplace and _seed_step_scale_default both used one flat init_step_scale for every leaf/lineout whenever curvature was unusable (h_ii <= 0 or use_laplace_seed off). A flat logit-space step maps to wildly different physical steps depending on where a parameter's value sits in [lb, ub] (physical_step = (ub-lb) * sigmoid'(x) * logit_step), so the same flat number that's fine for one leaf can badly mis-scale another -- and since _mh_accept makes one joint accept/reject decision per lineout across every active leaf, a single badly-scaled fallback leaf (brem_amp/brem_c, documented as degenerate with Z/Te/ne) can drag down the whole lineout's acceptance, not just that parameter's own. Replace it with fallback_step_frac: a dimensionless fraction of a parameter's own range, inverted through sigmoid'(x) at that leaf's current position so the resulting logit-space step always corresponds to the same physical-range fraction regardless of where the parameter sits (see _fallback_step_scale). This makes init_step_scale dead, so it's dropped from both config decks that overrode it. Added a regression test that reproduces a genuine non-positive diagonal Hessian entry for brem_c (confirmed h_ii < 0 at the fitted point) and checks the seeded fallback matches the position-aware formula rather than a flat constant -- which fails against the old flat pass-through, confirmed directly. Note: with the generous burn-in this codebase's other MCMC tests use, Robbins-Monro/adapt_shape adaptation converges brem_c's post-burn-in acceptance rate to about the same value whether the fallback is flat or position-aware on this small fixture, so that metric wasn't a reliable regression signal here; the seeded value itself is. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
postprocess_runner.py, mcmc_postprocess_runner.py, and queue_mcmc_postprocessor.py each downloaded a run's 2-5 config/weights files one at a time via mlflow.artifacts.download_artifacts(run_id=..., artifact_path=...). Each such call builds its own RunsArtifactRepository, whose __init__ does a fresh store.get_run(run_id) REST round-trip -- so fetching N files cost N+ redundant round-trips just to keep re-resolving the same run. Adds postprocess_runner._resolve_artifact_uri (one mlflow.get_run() call, returning the run's absolute artifact_uri) and postprocess_runner._download_run_artifact (downloads one file via download_artifacts(artifact_uri=posixpath.join(base, fname), ...), bypassing RunsArtifactRepository's per-file re-resolution; posixpath because mlflow artifact paths are always POSIX-style regardless of the local platform). All three call sites now resolve the run once and reuse that URI for every file, with their existing per-file try/except/fallback logic (which files are required vs. optional, what happens if one is missing) left untouched -- only the fetch mechanism changed. One intentional behavior change: postprocess_runner.py's config.yaml probe broadened from `except MlflowException` to `except Exception` to match the other two sites, since a missing-file-within-a-valid-run error can now surface through a different exception type once the resolved artifact_uri path is used instead of RunsArtifactRepository. Also: a genuinely nonexistent run_id now fails immediately with a clear "run not found" error from mlflow.get_run(), instead of the old confusing sequence (config.yaml probe fails -> falls back to defaults.yaml/ inputs.yaml -> that fails too). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
run_mcmc_postprocessor.py, run_postprocessor.py, and run_tsadar.py had no persistent JAX compilation cache configured, so every process invocation repaid the full XLA compile cost from scratch -- a real cost on a shared cluster running many similar postprocessing jobs. Defaults to a directory under $PSCRATCH on NERSC (falling back to ~/.cache locally), overridable via JAX_COMPILATION_CACHE_DIR. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… dense over the whole batch Different lineouts share no parameters and don't interact in the forward model, so every cross-lineout entry of the full batch Hessian is structurally zero -- but h_loss_wrt_params computed and discarded that dense (batch_size, batch_size) block per parameter pair anyway, which has been observed to attempt an 18+GiB allocation on a real production shot's compare_to_laplace=True path. LossFunction.h_loss_wrt_params_per_lineout gets the same per-lineout (active-parameter x active-parameter) covariance via one extra forward-mode jvp per active parameter, without ever materializing the cross-lineout terms, cutting memory from O((n_active * batch_size)^2) to O(n_active^2 * batch_size). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ulted mcmc config fields; save postprocessor overrides as an artifact The position-aware fallback_step_frac formula (introduced to make a degenerate parameter's fallback proposal scale position-aware rather than a flat logit-space number) silently broke production runs whose deck still set the old init_step_scale key: since it renamed the field under the same config['other']['mcmc'] namespace with different units, the deck's tuned value was dropped unnoticed and the sampler fell back to the new default -- roughly 2-3 orders of magnitude larger at typical positions -- collapsing acceptance to ~0. Revert to the original flat semantics (default 0.1, matching pre-refactor), keeping the unrelated low-memory jax.lax.map/jvp Hessian-diagonal extraction mechanism from later commits intact. To make this failure mode visible instead of silent next time: _mcmc_cfg now warns whenever any other.mcmc field isn't explicitly set in the deck and falls back to a built-in default. And since part of what made the original regression hard to diagnose was not knowing which override stub actually launched the run, run_mcmc_postprocess now saves the --overrides deck as its own overrides.yaml artifact at run start, mirroring runner.load_and_make_folders' save-the-deck-at-run-start pattern for ordinary fits. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ectly Diagnosing a still-elevated R-hat after the init_step_scale fix required a one-off script to split the pooled posterior back into its K per-chain blocks and recolor a corner plot manually, which confirmed a handful of chains were getting stuck away from the consensus region and never mixing back in within the burn-in budget -- invisible in the existing single-color corner plot, since chains sharing a real degenerate direction (e.g. Te-ne-amp1) can look like one filled ridge even when badly unmixed. plot_corner now accepts an optional num_chains (default 1, exactly reproducing today's single-color plot) that recolors each chain's own points separately via corner.corner's plot_datapoints=False plus a manual per-chain scatter overlay, leaving the pooled density contours/marginal histograms themselves untouched. mcmc_postprocess.py threads the already-known chain count through to its plot_corner calls. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ted via Robust Adaptive Metropolis (Vihola 2012) Real production fits showed genuine multi-parameter degeneracies (a near-exact lam-Va ridge, a Ti-Z coupling with Ti's own diagonal curvature negative at the reported best fit) that a diagonal-only Laplace-seeded proposal cannot move along efficiently, causing poor mixing/convergence. This replaces it with a full per-lineout proposal covariance: - Laplace seeding now uses the full per-lineout Hessian (not just its diagonal) via LossFunction.h_loss_wrt_params_per_lineout, regularized by eigenvalue-clipping in a normalized (correlation-scaled) basis so an indefinite Hessian (a saddle direction) still yields a valid positive-definite proposal covariance while preserving which *combinations* of parameters form each near-degenerate direction. - step_scale is now the (batch_size, n_active, n_active) Cholesky factor of that covariance rather than a per-leaf scalar; _propose draws correlated joint steps accordingly. - Burn-in adaptation is replaced with Vihola's (2012) Robust Adaptive Metropolis: a per-step rank-one Cholesky update driven by the continuous MH acceptance probability, positive-definite by construction, with no eigenvalue floor or separate magnitude/shape decomposition needed. This replaces an earlier windowed Robbins-Monro + determinant-1-normalized shape re-estimation scheme, which had two production-confirmed failure modes: a positive-feedback runaway when one direction's regularization floor inflated every other direction to compensate, and (independently) a NaN-producing Cholesky failure when a single under-sampled window's near-zero empirical variance forced the same kind of compensating blowup. RAM has no such cross-direction coupling -- confirmed via a full 20-draw production comparison run landing every lineout's acceptance rate in a tight band near the configured target, instead of pinned near 0 or 1. Also fixes an unrelated column-alignment bug in save_sigmas_params: the per-species column index was reset to 0 for every species instead of running globally across all species' combined columns, silently misattributing sigma values for any config with more than one active species.
…ostprocessor Chains for ill-conditioned lineouts (e.g. those needing ne/Ti deactivation) were failing to converge under a single full-covariance RAM proposal. Detect problem parameters from the per-lineout Hessian's normalized eigenvalues, union them shot-wide into a fixed (well-conditioned, problem) partition, and sample the two blocks with independently RAM-adapted proposals (Turek et al. 2015). This lets ne/Ti stay active without manual deactivation, matching or beating the old workaround's quality. Also replace classical Gelman-Rubin R-hat with the rank-normalized, folded, split R-hat of Vehtari et al. (2021), which is far less biased by heavy-tailed/high-variance chains. Chain-selection now votes per parameter: a parameter that's structurally unreliable across many chains no longer vetoes the whole lineout's summary statistics, it's just excluded from voting (own mean/std still reported). On real saved run6 chains this cut unreliable lineouts from 83/84 to 10/84 and dropped chains from 1395 to 201. Update tsadar_math.tex with the above plus citations, and note the geodesic/Riemannian-manifold-aware sampling alternative (Byrne & Girolami 2013, Betancourt 2013) as a documented option for a future session. Also fixes a column-ordering bug in save_sigmas_params surfaced while extending plot_mcmc_diagnostics to the new per-parameter diagnostics. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…egative-theory sqrt NaN guard - calculate_covariance_matrix's sig_s now built from the forward-model prediction (ThryI too, matching what ThryE already did) instead of raw data, removing the Neyman/Pearson bias and the near-zero-data variance blowup at its source. - neg_log_likelihood(per_lineout=True) now supported for loss_method="covar" (previously raised NotImplementedError, which is what blocked MCMC from ever sampling against this likelihood). Returns the raw per-lineout quadratic form rather than the point-estimate loss's dof-normalized scalar. Solve upgraded from a naive jnp.linalg.solve to a Cholesky factorization. - Floored sig_s's sqrt() input at 1e-10 (not 0): real shot 116773 data has forward-model pixels that are a hair negative from floating-point noise, which NaNs sqrt() and, via the PSF convolution, the entire lineout's covariance matrix -- and a floor of exactly 0 fixes the value but leaves the gradient NaN (d/dx sqrt(x) diverges at x=0), stalling any gradient-based refit. - postprocessor_stub.yaml: added optimizer.loss_method=covar, read only by the standalone MCMC postprocessor's reconstructed loss_fn so this changes what MCMC samples against without touching the point-estimate optimizer that already produced fitted_weights.eqx. Validated via 25-draw synthetic Monte Carlo recovery (within ~20% of analytic Fisher-information uncertainty for 4/5 parameters) and a real-data spot check on shot 116773 lineout 1825 (the near-zero-pixel blowup's top-10-pixel share drops from 46% to 3.8%; point estimate vs. covar optimum gap shrinks from a 7,000-34,000 unit runaway to a 28% NLL improvement). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… pixels calculate_covariance_matrix was factorizing a (1024 x 1024) matrix per lineout per call regardless of how few pixels (~150-200) are actually in the fit range, making loss_method="covar" ~3.2x slower than l2 per gradient call -- confirmed to be the dominant cost behind an MCMC run only reaching ~7 steps/s on GPU instead of the expected ~50. Precomputes, once per LossFunction (in _init_covar_pixel_indices), a padded pixel-index window per feature (IAW/EPW-blue/EPW-red): the true in-range span padded by the PSF kernel's own half-width on each side, so the covariance convolution still sees each in-range pixel's true neighbors. Naively gathering to ONLY the true in-range pixels (no padding) was tried first and shifted the NLL by ~0.6% on real data by dropping that cross-boundary correlation; with the padding, results match the original full-array computation to floating-point precision. Safe to compute once (not per call) because the experimental wavelength axis ts_diag returns doesn't depend on the fitted weights (confirmed: bit-identical across widely different "lam" values) -- verified directly. Net effect: covar's per-call gradient cost drops from ~137ms to ~43ms, matching l2 (~43ms). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- loss_function.py: _loss_for_hess_fn_ and _h_loss_wrt_params_per_lineout_impl were differentiating neg_log_likelihood(per_lineout=False), which for loss_method="covar" is dof-normalized (divided by finite_pixels*batch_size - free_params, order of a few hundred to a couple thousand) -- correct for the point-estimate loss's reduced-chi2-style scale, but wrong for a Hessian: it silently divided the curvature by that same factor. Confirmed via finite-difference cross-check on real data: AD/FD ratio was ~0.001-0.003 (a ~350-1400x systematic understatement) before the fix, ~0.98-1.0 for every parameter except Ti after it. This inflated the Laplace-seeded MCMC proposal step scale by ~sqrt(norm), throwing chains into catastrophically unphysical starting points (per-lineout NLL blowups of 5-6 orders of magnitude) regardless of block-Gibbs. Fixed by summing neg_log_likelihood(per_lineout=True) instead -- identical to the old behavior for l2 (already an unnormalized sum) and every other non-angular method; angular is untouched (per_lineout=True isn't supported there, and its own per_lineout=False objective was never dof-normalized to begin with). - mcmc.py/mcmc_postprocess.py: each calibration draw's raw (samples, diagnostics) is now pickled and uploaded to the active mlflow run (draw_checkpoints/draw_NN.pkl) as soon as that draw finishes, rather than only once at the very end alongside the final pooled summary/plots -- so a run that dies partway through (walltime, OOM, an unrelated crash) still leaves every completed draw's results recoverable. static_params is deliberately excluded from the checkpoint (contains an electron distribution-function closure plain pickle can't serialize; fully reconstructible from fitted_weights.eqx + get_filter_spec anyway). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…data fixture fitted_fixture's shared data was prepared with load_ion_spec/fit_IAW off (EPW-only), so Z -- which only enters the forward model through the IAW spectrum -- has zero effect on the loss there: its diagonal Hessian entry sits at ~0 (a genuinely flat, unidentifiable direction, not merely "poorly conditioned"), producing the large negative normalized eigenvalue this test failed on. Re-prepares a private copy of the data with IAW loaded and fit so Z is actually constrained, instead of mutating the module-scoped fixture shared by every other test in this file. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… MLflow export - Cap the gain exponent (other.gain_cap) so an unresolved resonance cannot overflow exp() - Force data.background.bg_subtract=false when loss_method is covar - Realign within-chain R-hat columns to active_keys order in mcmc_postprocess - Skip fe/f/flm in save_sigmas_params (not counted in num_params) and add a regression test - Stash the --overrides stub at queue time in queue_mcmc_postprocessor.py - Log per-epoch metrics asynchronously - Remove export_run/upload_dir_to_s3 and the mlflow_export_import dependency - Update default IAW calibration/windows and the working input deck Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…mutter settings - Shorten docstrings and comments across the MCMC, Laplace, loss-function, runner, plotting and test modules; code is unchanged apart from the wording of the covar/bg_subtract warning - docs/tsadar_math.tex: add a work-in-progress note; document the gain cap, the model-based noise covariance and its pixel window, the shared likelihood, the per-lineout Hessian, the Laplace-seeded proposal, block-Gibbs thresholds, cross-chain outlier detection, and implementation notes - docs/source/math_full.rst: rendered version of the math document, generated by docs/make_math_page.py, with pre-rendered TikZ figures; linked from math.rst - docs/source/mcmc.rst, defaults.rst, inputs.rst: describe the current sampler (RAM, block Gibbs, convergence checks, covar likelihood) and its options; add gain_cap - docs/source/generated: regenerate the autosummary stubs - calibration.py: update the spectral IRF widths - configs/1d decks and batch scripts: take the versions in use on Perlmutter Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- configs/1d: restore the decks from main; add only the new gain_cap field - configs/postprocessor/postprocessor_stub.yaml: list every MCMC option and use placeholder mlflow names - refit_bad_fits: warm-start the refit from the previous lineout's ThomsonParams instead of rebuilding it through the config, which failed for parameters the deck does not define (brem_amp, brem_c); write the refit result back to the current lineout's own batch - tests: run test_refit with lineouts in one batch and in separate batches; compare the Laplace-seeded proposal against the regularized inverse Hessian for any conditioning - run_tsadar.py: leave jax_debug_nans disabled, as on main - Bump version to 0.3.1 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…and MCMC diagnostics - loss_function: build covar pixel windows only for features that are fit; raise a clear error when an enabled fit range selects no pixels - postprocess_runner: load fitted_weights.eqx saved before the bremsstrahlung parameters existed through a legacy-layout fallback - mcmc_postprocess: treat chains that never moved or have a non-finite R-hat as failed; match Laplace comparison sigmas to parameters by name - mcmc: use the Blom denominator N + 1/4 in rank normalization; run the burn-in remainder so exactly burn_in steps execute - mcmc_calibration: leave an unperturbed IRF width (including zero) unchanged - tests for each of the above Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…and gain; bump to 0.4.0 Statistical model - MCMC target is now log-likelihood + log-prior + log-Jacobian of the logit transform; the prior is selected by name (other.mcmc.prior, default uniform within lb/ub) from a registry - covar likelihood for MCMC and the Hessian includes log(det K), evaluated on the in-range pixels - covar fit statistic is differentiated with K held fixed, as in Swadling et al. (2022) - blue/red EPW errors are averaged only for mean reductions; sum-based paths (covar, the -2 log-likelihood) add them - Laplace sigmas are reported in physical units, J (2 H^-1) J^T; the MCMC proposal seed uses the Hessian of the negative log-posterior Data handling - lineouts sum 2*dpixel+1 rows centered on the lineout position (was 2*dpixel) - new data.spectral_smoothing option (default: 2*dpixel+1; 1: none); covar and the MCMC postprocessor force 1, calc_sigmas warns when the data are smoothed - configs without data.spectral_smoothing warn that they predate 0.4.0 - a single camera gain: the noise model uses other.gain (default now 108, Ghosh et al. 2004) and works in photoelectrons; covar warns when the gain is not 108 - model amplitude is referenced to its maximum inside the fit ranges, for the EPW and IAW, so an unresolved peak outside them cannot set the normalization Docs and tests - math document: statistical model, prior and Jacobian, likelihood vs fitting statistic, source of the detector constants, units of the Laplace sigmas; rendered page regenerated - defaults reference and MCMC page updated; tests for each change; test decks use gain 108 - version 0.4.0 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Resolves the conflicts with #156 (async MLflow logging): keep the mlflow>=3.11.1 floor from main, keep this branch's removal of mlflow_export_import and MLFLOW_EXPORT, and keep the short comments. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ility flags - mcmc: the logit Jacobian and its curvature run over every active leaf (the set from get_filter_spec) instead of leaves selected by name, so an active ion fraction is included - mcmc_postprocess: write binary/mcmc_reliability.nc (param_unreliable per lineout and parameter, lineout_unreliable and n_chains_dropped per lineout) alongside the sigmas; registered in the manifest - docs: reliability artifact on the MCMC page; note on the ion-fraction prior in the math document - tests: active-fraction Jacobian test; runner test checks the reliability artifact Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- `other.gain` -> `other.detector_gain` - `other.calibration_uncertainty.gain_sigma` -> `detector_gain_sigma` - `other.include_gains` -> `other.LPI_gain.include_LPI_gains` - `other.gain_cap` -> `other.LPI_gain.LPI_gain_cap` - `other.Ipump_14`, `other.beam_diam_um` -> moved under `other.LPI_gain` Updates all decks, the code that reads these keys, tests, and docs. No backward-compatibility shim: decks using the old names must be updated. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
main squash-merged feature/mcmc (#158), so every file the two histories share conflicted. Resolved by taking main's content throughout and re-applying the gain key rename on top of it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…eprecation warning Temporary helper `_rename_deprecated_gain_keys_` moves the pre-rename keys of `other` to their new names and emits a FutureWarning listing them. To be removed in a future version, most likely 0.5.0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Moves the temporary rename helper to tsadar.utils.misc so both config validation and forward_pass call it, making the key rename backward compatible on every entry point. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![ergodic-reviewer[bot]](https://avatars.githubusercontent.com/u/74554135?s=60&v=4){kind=small}
There was a problem hiding this comment.
Requesting changes for a configuration-precedence regression: the compatibility helper can replace explicit new-style overrides with legacy defaults or saved values, causing MCMC to use the wrong detector calibration. See the inline finding.
Postprocessing: existing forward plots/NetCDF, fit checkpoints/parameter CSVs, and MCMC diagnostics/artifacts remain implemented and reachable. The finding affects the configuration used when replaying legacy fits, not a missing output path.
Static review of the complete diff and relevant callers/tests; I did not execute tests or postprocessing.
| if old_path[-1] not in parent: | ||
| continue | ||
| target = other | ||
| for key in new_path[:-1]: |
There was a problem hiding this comment.
[P2] Preserve input/override precedence across the migration
The new key is not necessarily a default: both postprocessor runners merge --overrides into the saved configuration before _reconstruct_fit_state() calls _validate_inputs_(). For a saved fit with other.gain: 1, the updated postprocessor stub's other.detector_gain: 108 survives the merge as a separate leaf, then this assignment replaces it with 1. Before the rename, the stub's gain: 108 correctly replaced the saved value. Consequently, the documented MCMC override workflow cannot correct a legacy fit's detector gain, and get_lineouts() and LossFunction.G use the wrong calibration for photoelectron conversion and the covariance likelihood. New-style LPI overrides and new inputs layered over old defaults have the same precedence problem.
Normalize the individual configuration layers before merging them (including a saved config.yaml before applying overrides), or otherwise retain their precedence rather than unconditionally preferring the old spelling. Add regressions for old defaults plus new inputs and a legacy saved fit plus new-style postprocessor overrides.
…gain argument - Note in defaults.rst and inputs.rst listing the renamed keys and that old decks are converted with a warning until about 0.5.0 - Math docs point at the deck key that enables the LPI gain, not only the internal calc_gain name - FormFactor.__init__ docstring entry for calc_gain Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
The previously reported P2 configuration-precedence defect remains at this head. rename_deprecated_gain_keys (tsadar/utils/misc.py:297) overwrites explicit new-style overrides with legacy values after the postprocessor runners merge the configuration. A saved other.gain: 1 plus the updated stub's other.detector_gain: 108 therefore still uses 1 for photoelectron conversion and the covariance likelihood. Normalize configuration layers before merging so inputs/overrides retain precedence, and cover legacy saved configurations plus new-style overrides in regression tests. The existing inline thread contains the full finding; I have not duplicated it.
Postprocessing: forward plots/NetCDF, fit checkpoints/parameter CSVs, and standalone Laplace/MCMC diagnostics and artifacts remain implemented and reachable. This rename does not require duplicate postprocessing; the blocker is the calibration used when replaying legacy fits.
Reviewed the complete diff and relevant implementation, guidance, callers, and tests at this commit. Tests and postprocessing were not executed.
When both spellings are present the deprecated key is now discarded instead of overwriting the new key, so an explicit new-style override (e.g. the postprocessor stub's detector_gain) is no longer replaced by a legacy saved value. The warning names each ignored key with both values. A deprecated key is still converted when its new key is absent. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
No material findings at this head. The previous detector-calibration override defect is fixed: rename_deprecated_gain_keys preserves an explicit detector_gain: 108 alongside a legacy saved gain: 1. The documented new-key-wins policy also means new-style defaults can override legacy input values; the helper warns with the discarded and selected values.
Postprocessing: existing forward plots/NetCDF, fit checkpoints/parameter CSVs, and standalone Laplace/MCMC diagnostics and saved artifacts remain reachable through the migrated configuration paths. This rename does not require additional postprocessing.
Reviewed the complete diff and relevant implementation, callers, guidance, and tests. Tests and postprocessing were not executed. CPU CI is still running; GPU CI is skipped. Enabled-LPI and legacy-forward migration coverage remain validation gaps, not demonstrated defects.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The
othersection had four unrelated keys all called some variant of "gain". This renames them so the detector gain and the LPI (SRS/SBS) amplification are distinguishable, and groups the LPI options together.other.gainother.detector_gainother.calibration_uncertainty.gain_sigmaother.calibration_uncertainty.detector_gain_sigmaother.include_gainsother.LPI_gain.include_LPI_gainsother.gain_capother.LPI_gain.LPI_gain_capother.Ipump_14other.LPI_gain.Ipump_14other.beam_diam_umother.LPI_gain.beam_diam_umUpdates all decks, the code that reads these keys, tests, and docs.
optimizer.angular_objective.gainis a different quantity and is unchanged, as are the form factor's internalcalc_gaindict keys.Backward compatibility
Decks using the old names still run, with a warning. A temporary helper,
tsadar.utils.misc.rename_deprecated_gain_keys, runs first in_validate_inputs_(fits,calc_vs_data, the postprocessor) and inforward_pass. For each old key it finds:Either way the old key is removed and a single
FutureWarninglists every key, saying which were ignored (with both values) and which were converted. Example:Consequence: an old-style inputs deck run against the current defaults has most of its old gain values ignored, because the shipped defaults already set the new keys. That is intended; the warning tells the user to correct the deck rather than keep old keys working indefinitely.
This addresses the review finding: a new-style override such as the postprocessor stub's
detector_gain: 108now takes precedence over a legacy savedgain: 1.The helper is to be removed in a future version, most likely 0.5.0.
Tests
No new regression tests were added for the compatibility helper. It is a temporary shim in pre-1.0 code and will be deleted, so the layer-precedence tests requested in review were deliberately left out. The one existing helper test was updated to the new-key-wins behaviour.
Not included
.zipbundles underdocs/source/examples/still contain decks with the old key names.include_LPI_gainson, and theforward_passcall to the helper has no test.Validation
pytest -m "not slow" tests/pytest -m slow tests/(note who ran it or link the result)pytest -q tests/test_form_factor/test_ratintn_gpu.pyor link the result)Only targeted tests were run locally:
tests/test_inverse/test_validate_inputs.py(12 passed on the final code), plus the 1D forward pass and the MCMC calibration-draw and covar tests before the helper was added. No numerical code changed.🤖 Generated with Claude Code