diff --git a/scripts/jax_likelihood_functions/weak/__init__.py b/scripts/jax_likelihood_functions/weak/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/scripts/jax_likelihood_functions/weak/shear.py b/scripts/jax_likelihood_functions/weak/shear.py new file mode 100644 index 00000000..92c714c6 --- /dev/null +++ b/scripts/jax_likelihood_functions/weak/shear.py @@ -0,0 +1,154 @@ +""" +Func Grad: Weak Lensing Shear Chi-Squared +========================================= + +Test that JAX can compute the log-likelihood of a ``WeakDataset`` via ``al.AnalysisWeak`` +with ``use_jax=True``: the fitness ``_vmap`` path (which registers the model pytrees and +forces tracer propagation) must match the eager NumPy likelihood exactly. + +Covers the ``FitWeak`` xp-threaded statistics and the ``AnalysisWeak._register_fit_weak_pytrees`` +registration added by PyAutoLens feature/weak-sigma-crit-jax (issue #590). The redshift +scale factors are concrete per-dataset constants, so a dataset carrying per-galaxy redshifts +exercises the eager-scaling + traced-statistics combination too. +""" + +# %matplotlib inline +# from pyprojroot import here +# workspace_path = str(here()) +# %cd $workspace_path +# print(f"Working Directory has been set to `{workspace_path}`") + +import numpy as np +import jax.numpy as jnp +from pathlib import Path + +import autoarray as aa +import autofit as af +import autolens as al + +""" +__Dataset__ +""" +dataset_path = Path("dataset") / "weak" / "simple" + +""" +__Dataset Auto-Simulation__ +""" +if al.util.dataset.should_simulate(str(dataset_path)): + import subprocess + import sys + + subprocess.run( + [sys.executable, "scripts/jax_likelihood_functions/weak/simulator.py"], + check=True, + ) + +dataset = al.from_json(file_path=dataset_path / "dataset.json") + +""" +__Model__ +""" +mass = af.Model(al.mp.Isothermal) + +mass.centre.centre_0 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.centre.centre_1 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.ell_comps.ell_comps_0 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.ell_comps.ell_comps_1 = af.UniformPrior(lower_limit=0.0, upper_limit=0.02) +mass.einstein_radius = af.UniformPrior(lower_limit=1.5, upper_limit=1.8) + +lens = af.Model(al.Galaxy, redshift=0.5, mass=mass) + +source = af.Model(al.Galaxy, redshift=1.0) + +model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + +print(model.info) + +""" +__NumPy reference__ +""" +analysis_numpy = al.AnalysisWeak(dataset=dataset, use_jax=False) + +instance = model.instance_from_prior_medians() + +log_likelihood_numpy = analysis_numpy.log_likelihood_function(instance=instance) + +print(f"NumPy log likelihood : {log_likelihood_numpy}") + +""" +__JAX vmap parity__ +""" +from autofit.non_linear.fitness import Fitness +import time + +analysis = al.AnalysisWeak(dataset=dataset, use_jax=True) + +batch_size = 2 + +fitness = Fitness( + model=model, + analysis=analysis, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) + +parameters = np.zeros((batch_size, model.total_free_parameters)) +for i in range(batch_size): + parameters[i, :] = model.physical_values_from_prior_medians +parameters = jnp.array(parameters) + +start = time.time() +result = fitness._vmap(parameters) +print(f"JAX vmap log likelihood : {np.array(result)}") +print("JAX Time To VMAP + JIT Function", time.time() - start) + +np.testing.assert_allclose( + np.array(result), + log_likelihood_numpy, + rtol=1e-6, + err_msg="weak/shear: JAX vmap likelihood mismatch vs eager NumPy", +) + +""" +__Redshift-scaled variant__ + +Per-galaxy sigma_crit scaling: the scale factors are concrete constants computed before any +trace, so the scaled likelihood must also match between NumPy and vmap paths. +""" +redshifts = list(np.random.default_rng(2).uniform(0.6, 2.0, dataset.n_galaxies)) +dataset_scaled = al.WeakDataset( + shear_yx=dataset.shear_yx, + noise_map=dataset.noise_map, + name="simple_scaled", + redshifts=redshifts, +) + +analysis_scaled_numpy = al.AnalysisWeak(dataset=dataset_scaled, use_jax=False) +log_likelihood_scaled_numpy = analysis_scaled_numpy.log_likelihood_function( + instance=instance +) + +analysis_scaled = al.AnalysisWeak(dataset=dataset_scaled, use_jax=True) +fitness_scaled = Fitness( + model=model, + analysis=analysis_scaled, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, +) +result_scaled = fitness_scaled._vmap(parameters) + +print(f"NumPy scaled log likelihood : {log_likelihood_scaled_numpy}") +print(f"JAX scaled log likelihood : {np.array(result_scaled)}") + +np.testing.assert_allclose( + np.array(result_scaled), + log_likelihood_scaled_numpy, + rtol=1e-6, + err_msg="weak/shear: scaled JAX vmap likelihood mismatch vs eager NumPy", +) + +assert not np.isclose(log_likelihood_scaled_numpy, log_likelihood_numpy), ( + "scaling changed nothing — redshift scale factors not applied" +) + +print("WEAK JAX PARITY PASSED") diff --git a/scripts/jax_likelihood_functions/weak/simulator.py b/scripts/jax_likelihood_functions/weak/simulator.py new file mode 100644 index 00000000..ffadff6e --- /dev/null +++ b/scripts/jax_likelihood_functions/weak/simulator.py @@ -0,0 +1,40 @@ +""" +Simulator for the weak-lensing JAX parity scripts. + +Writes a seeded, noise-controlled `WeakDataset` to `dataset/weak/simple/dataset.json` so the +`shear.py` vmap-parity regression constant is stable across runs. +""" +from pathlib import Path + +import autoarray as aa +import autolens as al + +import numpy as np + +dataset_path = Path("dataset") / "weak" / "simple" + +grid = aa.Grid2DIrregular(values=np.random.default_rng(1).uniform(-3.0, 3.0, (100, 2))) + +tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=(0.0, 0.05), + einstein_radius=1.6, + ), + ), + al.Galaxy(redshift=1.0), + ] +) + +dataset = al.SimulatorShearYX(noise_sigma=0.3, seed=1).via_tracer_from( + tracer=tracer, grid=grid, name="simple" +) + +dataset_path.mkdir(parents=True, exist_ok=True) +al.output_to_json(obj=dataset, file_path=dataset_path / "dataset.json") +al.output_to_json(obj=tracer, file_path=dataset_path / "tracer.json") + +print(f"Wrote weak parity dataset to {dataset_path}")