Skip to content

Refactor postprocessing into a single batch-handling path; add standalone postprocessor runner - #128

Open
almilder wants to merge 11 commits into
mainfrom
pr-1-postprocessor
Open

Refactor postprocessing into a single batch-handling path; add standalone postprocessor runner#128
almilder wants to merge 11 commits into
mainfrom
pr-1-postprocessor

Conversation

@almilder

Copy link
Copy Markdown
Collaborator

Summary

  • Moves saving of fit results from postprocessing into the main fitter, and unifies the previously-duplicated batch-creation logic between fitter and postprocess (new build_batch/build_angular_batch/unbatch_fitted_params helpers in loops.py).
  • Makes the postprocessor optional (other:run_postprocess) so a run can be restarted or replotted without redoing the fit.
  • Adds run_postprocessor.py / tsadar/postprocess_runner.py, a standalone tool to replay postprocess() on an already-completed fit, either from a local run directory or a remote MLflow run id/URL.
  • Deletes angular_optax, an orphaned duplicate of angular_multiple_optax's training loop with no remaining callers.
  • Fixes an existing test (test_manifest.py) whose process_data mock didn't match the new signature, and a pre-existing Windows-only MLflow artifact-URI bug unrelated to this change.

Test plan

  • pytest tests/ — 125 passed, 5 skipped, run on this branch alone (based on current main)
  • New regression tests: tests/test_inverse/test_refit.py (guards a real historical crash in the refit path), tests/test_postprocess_runner.py (exercises both local and remote replay entry points end-to-end)

Avi added 9 commits May 30, 2025 15:56
…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).

@joglekara joglekara left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requesting changes for these correctness issues:

  • Angular replay skips the ang_res_unit conversion used during fitting, so it postprocesses a different lineout range.
  • refit_bad_fits passes every index to one_d_loop with num_batches=1, so it optimizes lineout 0 instead of the flagged lineout.
  • Replay rebuilds checkpoints from the original config, but angular refinement mutates fe.nvx; refined checkpoints therefore cannot be deserialized.
  • Remote replay also assumes defaults.yaml/inputs.yaml, excluding app runs that only have config.yaml.

CPU/GPU checks are green, but these paths need regression coverage.

…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.
@almilder

Copy link
Copy Markdown
Collaborator Author

Thanks for the thorough review — all four are fixed in the latest commit (978a6b99).

Angular replay skips the ang_res_unit conversion. Confirmed: multirun_angular_optax divides lineouts.start/end by ang_res_unit before batching during the original fit, but postprocess_runner was replaying from a freshly-loaded config that never went through that mutation. Factored the conversion out into apply_ang_res_unit() (loops.py), called from both multirun_angular_optax (unchanged behavior) and postprocess_runner.run_postprocess before it touches all_data. Also worth noting for the record: this same value feeds ThomsonScatteringDiagnostic.reduce_ATS_to_resunit's slice on the theory side too, gated on spectype == "angular_full" specifically — so the fix needs to run before LossFunction is constructed, which it now does.

refit_bad_fits always refits lineout 0. Confirmed: it passed the entire sample_indices array into one_d_loop with num_batches=1; with batch_size forced to 1, that reshapes to (N, 1) and the loop only ever touches batch index 0 — lineout 0, never the flagged lineout i. Now passes np.array([i]) so the one batch actually contains the flagged lineout.

Refined checkpoints can't be deserialized. Confirmed: multirun_angular_optax's num_mins refinement loop grows nvx (and the smoothing window length) between passes, but replay was building the ThomsonParams skeleton from the un-refined config, so eqx.tree_deserialise_leaves hit a shape mismatch whenever num_mins > 1. Factored the shape progression (not the array interpolation, which isn't needed to build a skeleton) into advance_refinement_shape(), and run_postprocess now replays it num_mins - 1 times before constructing the skeleton.

Remote replay assumes defaults.yaml/inputs.yaml. Confirmed: app-originated runs (runner.run_for_app) only ever log a single config.yaml, never defaults.yaml/inputs.yaml. Both _load_merged_config (local path) and run_postprocess_remote (remote path) now check for config.yaml first and only fall back to the defaults/inputs pair if it's absent.

Added regression coverage for all four in tests/test_inverse/test_loops.py and new cases in tests/test_postprocess_runner.py. Full suite: 133 passed, 5 skipped (up from 125).

Separately, while chasing the angular-replay fix I found that plain "angular" (non-"range" lineout type) fitting looks structurally broken independent of this PR — build_angular_batch assumes a full CCD frame that only "angular_full" actually produces, and there's no batched per-lineout path for it. Traced the dispatch pattern back to the 2024-08 introduction of the optax-based fitter and it's been this way throughout, so it predates this PR and isn't something introduced here. Leaving it out of scope for this PR — tracking separately.

@almilder
almilder requested a review from joglekara August 25, 2026 14:30
…batch_size

run_postprocess rebuilt its normalization sample from the first batch_size
raw rows instead of the fitted lineout range, and never reset batch_size to
1 for angular fits, so replayed sigmas/losses/plots could silently diverge
from (or crash relative to) the original fit. Also avoid downloading
config.yaml twice in run_postprocess_remote.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants