feat: FitWeak per-galaxy sigma_crit scaling + JAX support (weak series step 10)#591
Merged
Conversation
- When WeakDataset.redshifts is present, FitWeak scales the model signal per galaxy by beta_i/beta_ref via LensingCosmology.scaling_factor_between_ redshifts_from (unity at the tracer's source plane, zero at/below the lens plane); reduced datasets scale both gamma and kappa: g_i = s*g/(1-s*k). No redshifts -> bit-identical to the previous single-plane behaviour. Factors are concrete NumPy constants (plane/catalogue redshifts are not sampled parameters), keeping them outside any JAX trace. - xp threaded through the fit statistics (LensCalc hessian methods already take xp); model_shear returns a raw array on the jnp path (LensCalc guard pattern); AnalysisWeak._register_fit_weak_pytrees mirrors AnalysisPoint (no_flatten: dataset, _xp, _redshift_scale_factors). use_jax default stays False. - 4 new NumPy-only tests (factor endpoints, per-galaxy scaling, unchanged no-redshift behaviour, reduced+scaled composition); JAX validated by the companion autolens_workspace_test weak vmap-parity script (plain + scaled variants match eager NumPy at rtol 1e-6). Weak series step 10 (#590). 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.
Summary
The two library follow-ups from the completed weak series, on one branch (both rework
FitWeakinternals). Σ_crit scaling: whendataset.redshiftsis present, the model signal is scaled per galaxy by β_i/β_ref (LensingCosmology.scaling_factor_between_redshifts_from; unity at the tracer's source plane, zero at/below the lens plane); reduced datasets scale both γ and κ (g_i = s·γ/(1−s·κ)); no redshifts → bit-identical previous behaviour. JAX:xpthreaded through the fit statistics,AnalysisWeak._register_fit_weak_pytreesmirroringAnalysisPoint,use_jax=Truesupported (default staysFalse). Scale factors are concrete constants (redshifts are not sampled parameters), documented. Closes #590 together with the companion workspace_test PR.API Changes
Behaviour addition, no signature breaks:
FitWeak(dataset, tracer, xp=np)gains the optionalxp; datasets carryingredshifts(stored-but-unused since #589, merged earlier today) now have them applied.See full details on #590.
Test Plan
test_autolens/: 373 passed (52 weak; 4 new scaling tests — factor endpoints, per-galaxy scaling ratio, unchanged no-redshift behaviour, reduced+scaled composition).fitness._vmap== eager NumPy at rtol 1e-6 for plain AND redshift-scaled datasets.Validation checklist (--auto run)
🤖 Generated with Claude Code