fix: better init of spergel HLR and tabulate xValue

[beckermr]

This PR has a new way to initialize the spergel HLR and a better way to compute autodiff using implicit differentiation.

[codspeed-hq]

Merging this PR will regress 2 benchmarks

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 5 improved benchmarks
❌ 2 regressed benchmarks
✅ 29 untouched benchmarks
🆕 2 new benchmarks

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	WallTime	test_benchmark_spergel_init[run]	6.9 µs	41.5 µs	-83.39%
❌	Simulation	test_benchmark_spergel_init[run]	103.5 µs	266.2 µs	-61.13%
⚡	Simulation	test_benchmark_spergel_xvalue[run]	34,867 ms	382.8 ms	×91
⚡	WallTime	test_benchmark_spergel_xvalue[run]	1,665.5 ms	41 ms	×41
⚡	Simulation	test_benchmark_spergel_calcfluxrad[run]	1,091.7 µs	627.8 µs	+73.89%
⚡	WallTime	test_benchmark_spergel_calcfluxrad[run]	183.7 µs	117.7 µs	+55.98%
⚡	Simulation	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	1,152.4 µs	873.5 µs	+31.94%
🆕	Simulation	test_benchmark_spergel_init_float[run]	N/A	104.3 µs	N/A
🆕	WallTime	test_benchmark_spergel_init_float[run]	N/A	6.8 µs	N/A

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing test-spergel (db604d6) with main (f08967b)}

[ismael-mendoza]

hello - i started looking into what you requested in slack.

Regarding the image artifacts with slen=101 for main and this PR respectively (the plots show the residual between galsim - jax_galsim):

but using slen=201 (again main branch and then this PR):

I used the same code that's here to test. I don't understand why the residual changes so much in this PR after doubling the slen because GalSim reports that the good size is 52. Do you have any insight? (although in both cases it's smaller than what we had before so maybe it's ok anyways? )

[ismael-mendoza]

re the regression when nu is not constant - I could not find any. Perhaps I misunderstood but here is the code I was using to test this in my laptop.

import numpy as np
import jax_galsim as jgs

def _draw_spergel_jgs(
    params: dict, nu: float | Array, *, psf: jgs.GSObject, fft_size: int, slen: int
):
    gsparams = jgs.GSParams(minimum_fft_size=fft_size, maximum_fft_size=fft_size)
    prof = jgs.Spergel(nu=nu, flux=params["flux"], half_light_radius=params["hlr"])
    prof = prof.shear(q=params["q"], beta=params["beta"] * jgs.degrees)
    gal = jgs.Convolve([prof, psf]).withGSParams(gsparams)
    return gal.drawImage(nx=slen, ny=slen, scale=0.2, dtype=jnp.float64).array

# let's average the time over 1000 runs each while varying the parameters
# but fixing the PSF


times1 = []
times2 = []

xpsf = jgs.Gaussian(half_light_radius=0.7, flux=1.0)

draw_jax_nu_float = jit(
    partial(_draw_spergel_jgs, nu=-0.6, psf=xpsf, fft_size=256, slen=201)
)
draw_jax_nu_array = jit(partial(_draw_spergel_jgs, psf=xpsf, fft_size=256, slen=201))


for ii in range(1000):
    params = {
        "flux": np.random.uniform(1, 3, size=()).item(),
        "hlr": np.random.uniform(0.4, 0.5, size=()).item(),
        "q": 1.0,
        "beta": 0.0,
    }
    _nu = np.random.uniform(-0.55, -0.65, size=()).item()
    _nu = device_put(jnp.array(_nu))
    params_jax = device_put(params)

    # compilation
    if ii == 0:
        draw_jax_nu_float(params_jax)
        draw_jax_nu_array(params_jax, _nu)

    # timing
    t1 = time.time()
    with transfer_guard("disallow"):
        _ = block_until_ready(draw_jax_nu_float(params_jax))
    t2 = time.time()
    times1.append(t2 - t1)

    t1 = time.time()
    with transfer_guard("disallow"):
        _ = block_until_ready(draw_jax_nu_array(params_jax, _nu))
    t2 = time.time()
    times2.append(t2 - t1)


t_nu_float = np.mean(times1)
t_nu_array = np.mean(times2)
print(f"Average time nu float: {t_nu_float:.4f}")
print(f"Average time nu array: {t_nu_array:.4f}")

Did I miss something?

[beckermr]

This PR is ready for review @ismael-mendoza!

[ismael-mendoza]

Thanks Matt! I am not able to follow/check some of the detailed math calculations, but I did have some high level questions