Adding HSM module, demo1.py and demo2.py

[b-remy]

This PR adds a new HSM module that calls the standard GalSim implementation. The code is almost identical to hsm.py from GalSim, with minimal adaptations to accept JAX-GalSim objects as inputs.

I also added Image, Bounds, and Position methods that return the underlying GalSim C++ objects, which are required for calls into the _galsim C++ layer.

In addition, this PR introduces demo1.py and demo2.py, which are direct copies of the corresponding GalSim demos, modified only to import jax_galsim.

The logs produced by GalSim and JAX-GalSim differ slightly. This may be related to differences in how images are drawn in JAX-GalSim rather than the HSM computation itself. I’d be curious to hear your thoughts on this @beckermr.

demo1.py

GalSim

Starting demo script 1 using:
    - circular Gaussian galaxy (flux = 1.0e+05, sigma = 2.0),
    - circular Gaussian PSF (sigma = 1.0),
    - pixel scale = 0.20,
    - Gaussian noise (sigma = 30.00).
Wrote image to 'output/demo1.fits'
HSM reports that the image has observed shape and size:
    e1 = -0.025, e2 = -0.009, sigma = 11.060 (pixels)
Expected values in the limit that pixel response and noise are negligible:
    e1 = 0.000, e2 = 0.000, sigma = 11.180

JAX-GalSim

Starting demo script 1 using:
    - circular Gaussian galaxy (flux = 1.0e+05, sigma = 2.0),
    - circular Gaussian PSF (sigma = 1.0),
    - pixel scale = 0.20,
    - Gaussian noise (sigma = 30.00).
Wrote image to 'output/demo1.fits'
HSM reports that the image has observed shape and size:
    e1 = 0.008, e2 = 0.043, sigma = 11.010 (pixels)
Expected values in the limit that pixel response and noise are negligible:
    e1 = 0.000, e2 = 0.000, sigma = 11.180

demo2.py

GalSim

Starting demo script 2 using:
    - sheared (0.10,0.20) exponential galaxy (flux = 1.0e+05, scale radius = 2.70),
    - circular Moffat PSF (beta = 5.0, re = 1.00),
    - pixel scale = 0.20,
    - Poisson noise (sky level = 2.5e+03).
Wrote image to 'output/demo2.fits'
Wrote effective PSF image to 'output/demo2_epsf.fits'
HSM reports that the image has observed shape and size:
    e1 = 0.161, e2 = 0.343, sigma = 16.714 (pixels)
When carrying out Regaussianization PSF correction, HSM reports distortions
    e1, e2 = 0.176, 0.376
Expected values in the limit that noise and non-Gaussianity are negligible:
    g1, g2 = 0.190, 0.381

JAX-GalSim

Starting demo script 2 using:
    - sheared (0.10,0.20) exponential galaxy (flux = 1.0e+05, scale radius = 2.70),
    - circular Moffat PSF (beta = 5.0, re = 1.00),
    - pixel scale = 0.20,
    - Poisson noise (sky level = 2.5e+03).
/home/ben/Documents/repos/probabilistic-shear/.venv/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:188: UserWarning: Explicitly requested dtype float64 requested in asarray is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return asarray(x, dtype=self.dtype)
Wrote image to 'output/demo2.fits'
Wrote effective PSF image to 'output/demo2_epsf.fits'
HSM reports that the image has observed shape and size:
    e1 = 0.172, e2 = 0.352, sigma = 17.127 (pixels)
When carrying out Regaussianization PSF correction, HSM reports distortions
    e1, e2 = 0.188, 0.385
Expected values in the limit that noise and non-Gaussianity are negligible:
    g1, g2 = 0.190, 0.381

[beckermr]

I wonder if we do ourselves a disservice by making this a pattern we promote.

[b-remy]

I agree with your concern and think using import jax_galsim directly is clearer, so I've switched to this. Let me know if that works for you.

[beckermr]

I want to think carefully about emulating the C++ objects by returning pure galsim objects. It works for this case, but that as an API is odd to me. I don't think we'd do this everywhere.

[codspeed-hq]

Merging this PR will degrade performance by 38.28%

❌ 2 regressed benchmarks
✅ 29 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
❌	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	721.9 µs	1,141.1 µs	-36.74%
❌	test_benchmarks_lanczos_interp[xval-no_conserve_dc-run]	677.3 µs	1,097.3 µs	-38.28%

---

_{Comparing hsm (80eb06a) with main (40845dd)}