Train progress bar#84
Open
maja-jablonska wants to merge 17 commits into
Open
Conversation
Rewrite the training/test steps, vectorizer, continuum and model code to use jax.numpy and jax-based autodiff/optimization (jax.vmap, jax.jacfwd, jaxopt) in place of the numpy/scipy + multiprocessing implementation. - Enable 64-bit JAX globally and shim jax.tree_map for jaxopt 0.8.3 - Replace hand-written label-vector derivative with jax.jacfwd - numpy 2.0 compat (RankWarning, collections.abc.Iterable) - Add jax/jaxlib/jaxopt deps; pytest as a test extra - Add parity tests against frozen golden numpy/scipy outputs Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fit pixels in fixed-size, pad-aligned batches instead of one fused vmap. This drives a tqdm progress bar (the single fused call is opaque), compiles the optimizer once and reuses it across batches, and bounds peak memory. Pixels are independent so results are identical. Keep theta/s2 as JAX arrays end-to-end (block_until_ready instead of np.asarray, jnp.concatenate, single tolist() for metadata). Round-trip now preserves them as JAX arrays; update the test contract accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The unregularized, unbounded pixel objective is a convex quadratic, so its minimum is the normal-equations solution -- running L-BFGS on it is wasted work (verified to land on the same point to ~1e-13). Add make_pixel_closed_form() and route the default reg=0/no-bounds case through it, ~3x faster at P=4200 with identical results. Censored coefficients are pinned to zero via a unit diagonal on the Gram matrix. Also JIT the L-BFGS init computation (one fused kernel, built only when the optimizer runs) and block per batch only when the progress bar is shown so batches can overlap otherwise. Add a closed-form-vs-L-BFGS parity test covering censoring. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
scripts/sweep_cannon.py: k-fold cross-validation over a grid of label sets, polynomial orders, regularization strengths and censoring schemes. Reports per-label bias/scatter, median r_chi_sq and convex-hull coverage, writes a tidy CSV, and (optionally) logs one offline W&B run per grid point -- config + scalar metrics + residual histograms + a per-run one-to-one spread figure -- plus a sweep-level summary run with a results table and a spread-vs-parameters bar chart. scripts/sweep.pbs / sync_wandb.pbs: run the sweep offline on a compute node, then spawn a copyq (data-mover) job that `wandb sync`s the offline runs to the cloud. Interpreter is selected via PYTHON_BIN. .gitignore: ignore wandb/ run data and sweep_*results.csv. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
scripts/train_cannon.py loads an APOGEE parquet table, assembles the dispersion/flux/ivar arrays, continuum-normalizes, splits into train/validation (dropping non-finite labels), fits a polynomial CannonModel, runs the test step, and writes a one-to-one figure, a predictions CSV, and optionally the trained model. Parametrized via a CLI with a --demo mode on the bundled golden data. Device is left to JAX_PLATFORMS rather than forced to cpu. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
scripts/run_sweep.py loads + continuum-normalizes the spectra via the train_cannon helpers, drops non-finite-label stars across the union of swept labels, then cross-validates a grid of label sets / orders / regularizations through sweep_cannon.sweep (offline W&B per grid point). CLI with nested-label-set defaults and a --demo mode. sweep.pbs now runs this driver instead of the bare sweep_cannon demo. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
run_sweep now constructs the sweep grid from a fixed core (--base, default
teff/logg/fe_h/mg_h), one or more age columns (--age-cols, default
age_Dnu,age_L -- each a base variant so the two ages can be compared),
an optional mass column, and a list of abundances, combined via
--label-set-mode {one-at-a-time,cumulative,minimal}. Label sets that
reference a column missing from the table are skipped with a warning, so
e.g. an absent age_Dnu no longer crashes the run.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- continuum.normalize: replace in-place masked assignment with jnp.where so JAX-array inputs no longer raise on immutable item assignment - model.train/test: drive progress with jax-tqdm via on-device lax.scan over batches (vmap within each batch); train keeps its prior batching, test gains a progressbar arg. Numerically identical to the previous vmap path. - continuum.normalize: add a plain tqdm bar over the host-side per-star loop - add end-to-end JAX test covering normalize -> build -> train -> test on a tiny synthetic sample - declare jax-tqdm and tqdm in install_requires Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the host-side per-star loop in sines_and_cosines with a single jax.jit + jax.vmap solve (_continuum_amplitudes) that fits every star at once per region; the design matrices are shared across stars so this is a clean batch (recompiled once per region shape). Aggregate the out-of-region pixel warning into one message and move the optional tqdm bar to iterate over regions. Continuum values match the previous implementation to 1e-10. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Training now optionally accounts for uncertainties on the training-set labels. Passing `training_set_label_err` (per-label 1-sigma values) makes `train()` propagate the label errors into the per-pixel weights and refine them with iteratively reweighted least squares (`n_irls` passes), so stars with uncertain labels are down-weighted. Omitting it reproduces the exact -label fit bit-for-bit. - fitting.py: add make_pixel_closed_form_eiv / make_pixel_fitter_eiv and a shared _label_variance_term (first-order propagation g^T Sigma g folded into ivar_eff = ivar / (1 + ivar * v_label)). - model.py: store/validate training_set_label_err (serialized; old models read back as None), build the label Jacobian + scaled variances once in train(), and select the EIV fitters. - model.py: make the default batch_size order-aware -- cap it to a memory budget so higher polynomial orders shrink the batch automatically instead of blowing up peak memory. - tests: add test_label_errors.py (no-regression, finite, regularized, down-weighting mechanism, round-trip, validation). - notebooks/start_with_label_errors.ipynb: copy of start.ipynb showing how to assemble per-label errors and train with them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eport chi2 A driver that adds the data knobs (S/N cutoff and arbitrary pandas-query row filters) on top of the existing label-set/order/regularization sweep, and for each grid point trains, validates, then: - reports the recovered spread (per-label bias/scatter/RMSE + mean_scatter) and writes a one-to-one spread figure; - reconstructs the validation spectra by forward-modelling the recovered (and reference) labels through the trained model; - plots a few observed-vs-reconstructed spectra spanning the chi2 range; - reports reconstruction reduced chi2 at both recovered and reference labels alongside the test-step chi2. Results go to experiment_results.csv plus per-config figures. Reuses the load/normalize/split helpers and spread_figure; optional --label-err trains with per-label uncertainties. Includes a self-contained --demo. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The test step was already batched over stars, but two things made it OOM the device on large models: - the default batch_size ignored the model size, so a model with many pixels / labels used a batch far larger than the device could hold. Cap it to a memory budget (batch_size * per_star, per_star ~ (P, L) Jacobian + (P, T) design), shrinking it automatically; an explicit batch_size still wins. - the per-star model flux (P,) returned by the spectrum fitter is not part of test()'s output, yet lax.scan accumulated it across batches -- wasting device memory of order the whole validation set (S, P). Drop it inside the scan body. Results are unchanged (stars are fit independently; verified batch-size independent). Full suite green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ipts - model.py: drop the default test working-set budget from 512 to 256 MiB so the memory-aware test batch_size is more conservative out of the box. - train_cannon.py / experiment_cannon.py: add a --test-batch-size knob threaded into model.test(...), so the test batch can be hand-lowered when the device OOMs (overrides the auto default). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…y on GPU Previously every train()/test() call built fresh @jax.jit closures with the design matrix and trained theta baked in as compile-time constants, so a hyper-parameter sweep recompiled the XLA programs for every fold of every grid point -- on a GPU the compiles dominate wall-clock for a sequential sweep. - Hoist the train/test lax.scan runners to module level, memoize them per (fitter, progress bar), and pass everything that varies between calls (design matrix, theta/s2/fiducials/scales, data batches) as arguments, so same-shape calls hit JAX's jit cache: all folds and all regularization strengths of a grid point now share one compiled program. - Memoize the per-pixel fitter factories and key the per-spectrum core on the vectorizer's terms so repeated factory calls return the same function object. - Move the initial-theta stack to a module-level jitted function. - cross_validate: one vectorizer for all folds and no per-fold tqdm (the host callbacks would be baked into the scan and force recompiles). - run_sweep: enable JAX's persistent on-disk compilation cache (programs now contain no array constants, so identical shapes hit it across processes and restarted jobs); add --jax-cache-dir. - Add scripts/sweep_gpu.pbs: single-GPU gpuvolta job that runs the sweep sequentially (each grid point already saturates the device via vmap), with a backend assertion, no-prealloc, and the shared compilation cache. Demo sweep: scan compiles 48 -> 24 (regs now share programs), identical results; cross-process warm start halves a small train step on CPU. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Pass --mass-col= (empty) so build_label_sets produces no base+mass variants; the swept label sets are base(+age) plus abundance additions only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
load_spectra now derives a <x>_fe = raw_<x>_h - raw_fe_h column for every raw [X/H] abundance except iron itself (never overwriting existing columns), and the sweep defaults use them: the base set carries mg_fe and the swept abundances are ce/ca/si/ni/mn/al/c/n over Fe. raw_fe_h itself stays [Fe/H]. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add quality_mask: stars with spectrum_flags != 0 or any warn_* column set to True are dropped. run_sweep applies it right after loading -- before continuum normalization and training -- and aborts if the cuts reject every star. Missing columns skip the corresponding cut with a warning; NaN warn values count as not-set, NaN spectrum_flags as not-clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
No description provided.