Skip to content

fix: point_source priors aligned to simulator truth + refresh constants (#514) #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
May 16, 2026
Merged

fix: point_source priors aligned to simulator truth + refresh constants (#514) #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 21 additions & 6 deletions likelihood/point_source/image_plane.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,13 +175,22 @@ def jit_profile(func, label, *args, n_repeats=10):
with timer.section("model_build"):
# GaussianPrior(mean=truth, sigma=small) centres prior-median at the
# simulator truth while keeping params free so gradient diagnostics
# have dimensionality.
# have dimensionality. Prior means MUST match the simulator's truth
# values exactly, otherwise the PointSolver finds fewer image-plane
# positions than the dataset contains and chi² explodes.
#
# Simulator truth (see autolens_workspace_developer/jax_profiling/
# dataset_setup/point_source.py):
# Isothermal at centre=(0, 0), einstein_radius=1.6,
# ell_comps = al.convert.ell_comps_from(axis_ratio=0.9, angle=45°)
# ≈ (0.0526316, 0.0)
# source point_0.centre = (0.07, 0.07)
mass = af.Model(al.mp.Isothermal)
mass.centre.centre_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.centre.centre_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.01, sigma=0.01)
mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.01, sigma=0.01)
mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.05263158, sigma=0.01)
mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.01)
lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)

point_0 = af.Model(al.ps.PointFlux)
Expand Down Expand Up @@ -443,7 +452,13 @@ def full_pipeline_from_params(params_tree):
# Simulator truth parameters + seeded noise (noise_seed=1 in
# simulators/point_source.py) make the image-plane log-likelihood
# deterministic. Eager, JIT, and vmap all agree to float64.
EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE = 0.07475703623045682
# Constant refreshed 2026-05-16 alongside the prior-truth-alignment fix
# above. The previous value (0.07475703623045682) was set on 2026-04-24
# against an earlier dataset+priors combination that has since been
# regenerated; the new value reflects the current truth-aligned
# evaluation against the dataset committed in
# autolens_workspace_developer@f8a5cef.
EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE = 7.196577317761017

np.testing.assert_allclose(
log_likelihood_ref,
Expand Down
32 changes: 26 additions & 6 deletions likelihood/point_source/source_plane.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,13 +158,23 @@ def jit_profile(func, label, *args, n_repeats=10):
with timer.section("model_build"):
# GaussianPrior(mean=truth, sigma=small) centres prior-median at the
# simulator truth while keeping params free so gradient vectors and
# finite-difference diagnostics have dimensionality.
# finite-difference diagnostics have dimensionality. Prior means MUST
# match the simulator's truth values exactly, otherwise the
# ray-traced source-plane positions cluster around the wrong centre
# and chi² explodes.
#
# Simulator truth (see autolens_workspace_developer/jax_profiling/
# dataset_setup/point_source.py):
# Isothermal at centre=(0, 0), einstein_radius=1.6,
# ell_comps = al.convert.ell_comps_from(axis_ratio=0.9, angle=45°)
# ≈ (0.0526316, 0.0)
# source point_0.centre = (0.07, 0.07)
mass = af.Model(al.mp.Isothermal)
mass.centre.centre_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.centre.centre_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.05263158, sigma=0.005)
mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)

point_0 = af.Model(al.ps.PointFlux)
Expand Down Expand Up @@ -495,7 +505,17 @@ def ray_trace_to_source_plane(params_tree, positions_raw):
# ell_comps=(0.01,0.01), source centre=(0.07,0.07)) + seeded noise
# (noise_seed=1 in simulators/point_source.py) make the log-likelihood
# deterministic. Eager numpy and full-pipeline JIT agree to float64.
EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE = -294.1401881258811
# Constant refreshed 2026-05-16 alongside the prior-truth-alignment fix.
# Previous value (-294.1401881258811) was set against an earlier
# dataset+priors combination. The source-plane chi² is more sensitive to
# small parameter changes than image-plane because the chi² formula
# weights residuals by magnifications² at the data positions — for a
# quad-image lens near a caustic configuration, magnifications can swing
# by 10-100x with small lens-parameter perturbations, dominating the
# log-likelihood. The refreshed value reflects the current truth-aligned
# evaluation against the dataset committed in
# autolens_workspace_developer@f8a5cef.
EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE = -33788.35731127962

np.testing.assert_allclose(
log_likelihood_ref,
Expand Down
Loading