Skip to content

Commit 1abc6ae

Browse files
Jammy2211Jammy2211
authored
Merge pull request #123 from PyAutoLabs/feature/array2d-native-jit-safety
feat: restore @jax.jit roundtrip in simulator parity tests
2 parents 87fa08a + cbfb812 commit 1abc6ae

2 files changed

Lines changed: 231 additions & 0 deletions

File tree

‎scripts/imaging/simulator_use_jax_parity.py‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
"""
2+
SimulatorImaging use_jax parity test
3+
=====================================
4+
5+
Runs ``al.SimulatorImaging(use_jax=False)`` and ``al.SimulatorImaging(use_jax=True)``
6+
against the same tracer + grid and asserts the noise-free simulated image agrees
7+
to machine precision. Cross-xp numerical validation for ``SimulatorImaging.use_jax=True``
8+
— library unit tests stay NumPy-only.
9+
10+
Noise is disabled on both paths because JAX and NumPy use different RNG algorithms;
11+
the same seed produces different draws. The simulator's xp-threaded noise path is
12+
tested separately (``test_autoarray/dataset/imaging/test_simulator_use_jax.py``
13+
covers the constructor wiring; this script covers numerical agreement of the
14+
deterministic image-generation path).
15+
16+
Also tests the @jax.jit roundtrip: the simulator under jit produces a dataset
17+
whose data array agrees with the eager path. This exercises the
18+
``register_tracer_classes(tracer)`` registration walker shipped in PR 1
19+
(PyAutoLens#538) — without prior registration the @jax.jit decoration would fail
20+
to flatten the Tracer at the JIT boundary.
21+
"""
22+
23+
from autoconf import jax_wrapper # Sets JAX float64 before other imports
24+
25+
import jax
26+
import numpy as np
27+
28+
import autolens as al
29+
30+
31+
grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)
32+
33+
lens_galaxy = al.Galaxy(
34+
redshift=0.5,
35+
light=al.lp.Sersic(
36+
centre=(0.0, 0.0),
37+
ell_comps=(0.01, 0.01),
38+
intensity=1.0,
39+
effective_radius=0.5,
40+
sersic_index=4.0,
41+
),
42+
mass=al.mp.Isothermal(
43+
centre=(0.0, 0.0),
44+
ell_comps=(0.01, 0.01),
45+
einstein_radius=1.6,
46+
),
47+
)
48+
source_galaxy = al.Galaxy(
49+
redshift=1.0,
50+
light=al.lp.Sersic(
51+
centre=(0.1, 0.1),
52+
ell_comps=(0.01, 0.01),
53+
intensity=0.5,
54+
effective_radius=0.2,
55+
sersic_index=1.0,
56+
),
57+
)
58+
tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
59+
60+
psf = al.Convolver.from_gaussian(
61+
shape_native=(11, 11), sigma=0.05, pixel_scales=grid.pixel_scales
62+
)
63+
64+
# Noise-free configuration: both backends must produce identical deterministic images.
65+
common_kwargs = dict(
66+
exposure_time=300.0,
67+
psf=psf,
68+
background_sky_level=0.1,
69+
add_poisson_noise_to_data=False,
70+
include_poisson_noise_in_noise_map=False,
71+
)
72+
73+
# NumPy path.
74+
simulator_np = al.SimulatorImaging(**common_kwargs)
75+
dataset_np = simulator_np.via_tracer_from(tracer=tracer, grid=grid)
76+
data_np = np.asarray(dataset_np.data.array)
77+
78+
# JAX path (eager, not yet under @jax.jit).
79+
simulator_jax = al.SimulatorImaging(use_jax=True, **common_kwargs)
80+
dataset_jax = simulator_jax.via_tracer_from(tracer=tracer, grid=grid)
81+
data_jax = np.asarray(dataset_jax.data.array)
82+
83+
assert data_np.shape == data_jax.shape, (
84+
f"Shape mismatch: numpy={data_np.shape} vs jax={data_jax.shape}"
85+
)
86+
np.testing.assert_allclose(
87+
data_np,
88+
data_jax,
89+
atol=1e-8,
90+
err_msg="SimulatorImaging(use_jax=True) data differs from use_jax=False",
91+
)
92+
print(
93+
f"PASS: SimulatorImaging(use_jax=True) data matches use_jax=False "
94+
f"to atol=1e-8 (shape {data_np.shape})."
95+
)
96+
97+
# @jax.jit roundtrip: register tracer classes so JAX can flatten/unflatten
98+
# the Tracer pytree at the JIT boundary, then run the simulator under jit
99+
# and verify the output matches the eager JAX path.
100+
al.util.register_tracer_classes(tracer)
101+
102+
103+
@jax.jit
104+
def simulate_jit(tracer):
105+
dataset = simulator_jax.via_tracer_from(tracer=tracer, grid=grid)
106+
return dataset.data._array
107+
108+
109+
data_jit = np.asarray(simulate_jit(tracer))
110+
111+
np.testing.assert_allclose(
112+
data_jax,
113+
data_jit,
114+
atol=1e-8,
115+
err_msg="SimulatorImaging @jax.jit data differs from eager JAX path",
116+
)
117+
print(
118+
f"PASS: SimulatorImaging @jax.jit roundtrip matches eager JAX "
119+
f"to atol=1e-8 (shape {data_jit.shape})."
120+
)

‎scripts/interferometer/simulator_use_jax_parity.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
"""
2+
SimulatorInterferometer use_jax parity test
3+
============================================
4+
5+
Runs ``al.SimulatorInterferometer(use_jax=False)`` and ``al.SimulatorInterferometer(use_jax=True)``
6+
against the same tracer + grid and asserts the noise-free simulated visibilities
7+
agree to machine precision. Cross-xp numerical validation for
8+
``SimulatorInterferometer.use_jax=True`` — library unit tests stay NumPy-only.
9+
10+
Noise is disabled (``noise_sigma=None``) because JAX and NumPy use different RNG
11+
algorithms; the same seed produces different draws. This test covers the
12+
deterministic Fourier-transform path.
13+
14+
Also tests the @jax.jit roundtrip for the interferometer simulator.
15+
"""
16+
17+
from autoconf import jax_wrapper # Sets JAX float64 before other imports
18+
19+
import jax
20+
import numpy as np
21+
22+
import autolens as al
23+
24+
25+
grid = al.Grid2D.uniform(shape_native=(64, 64), pixel_scales=0.1)
26+
27+
lens_galaxy = al.Galaxy(
28+
redshift=0.5,
29+
light=al.lp.Sersic(
30+
centre=(0.0, 0.0),
31+
ell_comps=(0.01, 0.01),
32+
intensity=1.0,
33+
effective_radius=0.5,
34+
sersic_index=4.0,
35+
),
36+
mass=al.mp.Isothermal(
37+
centre=(0.0, 0.0),
38+
ell_comps=(0.01, 0.01),
39+
einstein_radius=1.6,
40+
),
41+
)
42+
source_galaxy = al.Galaxy(
43+
redshift=1.0,
44+
light=al.lp.Sersic(
45+
centre=(0.1, 0.1),
46+
ell_comps=(0.01, 0.01),
47+
intensity=0.5,
48+
effective_radius=0.2,
49+
sersic_index=1.0,
50+
),
51+
)
52+
tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
53+
54+
uv_wavelengths = np.random.RandomState(seed=0).uniform(
55+
low=-1000.0, high=1000.0, size=(200, 2)
56+
)
57+
58+
# Noise-free configuration: both backends produce identical deterministic visibilities.
59+
common_kwargs = dict(
60+
uv_wavelengths=uv_wavelengths,
61+
exposure_time=300.0,
62+
noise_sigma=None, # disable noise for deterministic parity
63+
)
64+
65+
# NumPy path.
66+
simulator_np = al.SimulatorInterferometer(**common_kwargs)
67+
dataset_np = simulator_np.via_tracer_from(tracer=tracer, grid=grid)
68+
data_np = np.asarray(dataset_np.data)
69+
70+
# JAX path (eager).
71+
simulator_jax = al.SimulatorInterferometer(use_jax=True, **common_kwargs)
72+
dataset_jax = simulator_jax.via_tracer_from(tracer=tracer, grid=grid)
73+
data_jax = np.asarray(dataset_jax.data)
74+
75+
assert data_np.shape == data_jax.shape, (
76+
f"Shape mismatch: numpy={data_np.shape} vs jax={data_jax.shape}"
77+
)
78+
np.testing.assert_allclose(
79+
data_np,
80+
data_jax,
81+
atol=1e-8,
82+
err_msg="SimulatorInterferometer(use_jax=True) visibilities differ from use_jax=False",
83+
)
84+
85+
print(
86+
f"PASS: SimulatorInterferometer(use_jax=True) visibilities match use_jax=False "
87+
f"to atol=1e-8 ({data_np.shape[0]} visibilities)."
88+
)
89+
90+
# @jax.jit roundtrip.
91+
al.util.register_tracer_classes(tracer)
92+
93+
94+
@jax.jit
95+
def simulate_jit(tracer):
96+
dataset = simulator_jax.via_tracer_from(tracer=tracer, grid=grid)
97+
return dataset.data._array
98+
99+
100+
data_jit = np.asarray(simulate_jit(tracer))
101+
102+
np.testing.assert_allclose(
103+
data_jax,
104+
data_jit,
105+
atol=1e-8,
106+
err_msg="SimulatorInterferometer @jax.jit visibilities differ from eager JAX path",
107+
)
108+
print(
109+
f"PASS: SimulatorInterferometer @jax.jit roundtrip matches eager JAX "
110+
f"to atol=1e-8 ({data_jit.shape[0]} visibilities)."
111+
)

0 commit comments

Comments
 (0)