[b-biswas]

While trying to run JAX metadetection in float32, I noticed that the FFT kernels were still using float64 internally due to JAX-Galsim. Some small data type fixes seem to solve this issue.
Note: Running on float32 is significantly faster!

TODO:

• Write an explicit test to check if things are working as expected
• Check differences wrt original gasim

[codspeed]

Merging this PR will regress 7 benchmarks

4 warnings

⚠️ Walltime results may be unreliable

Some benchmarks used the Walltime instrument outside CodSpeed Macro Runners, which can make the results inconsistent.

Learn more about Macro Runners

See affected benchmarks

⚠️ 12 benchmarks are dominated by system calls

System calls cannot be simulated consistently and are left out of the measure. The reported value understates the real cost. Switch to the Walltime instrument to measure them.

Learn more about system calls

See affected benchmarks

⚠️ 12 benchmarks with a significant change ran on different hardware

Their base and head results were measured on machines with a different CPU, core count or memory, so the changes might not come from the code. Run base and head on the same runner type.

Learn more about Macro Runners

See affected benchmarks

⚠️ 12 benchmarks with a significant change ran in a different software environment

Their base and head results were measured with a different OS or toolchain, such as compiler, runtime or library versions, so the changes might come from this difference rather than from the code. Open an affected benchmark to see what changed.

See affected benchmarks

⚡ 5 improved benchmarks
❌ 7 regressed benchmarks
✅ 26 untouched benchmarks

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	Walltime	test_benchmarks_lanczos_interp[kval-no_conserve_dc-run]	31.7 µs	87.6 µs	-63.82%
❌	Walltime	test_benchmarks_lanczos_interp[kval-conserve_dc-run]	32 µs	86.8 µs	-63.1%
❌	Walltime	test_benchmark_spergel_conv[run]	112.5 ms	276.1 ms	-59.25%
❌	Walltime	test_benchmark_moffat_conv[run]	152.7 ms	302.8 ms	-49.58%
❌	Walltime	test_benchmark_moffat_conv_grad[run]	236.8 ms	443.2 ms	-46.57%
❌	Walltime	test_benchmarks_lanczos_interp[xval-no_conserve_dc-run]	81.9 µs	134.1 µs	-38.94%
❌	Walltime	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	89.5 µs	143.7 µs	-37.69%
⚡	Walltime	test_benchmark_spergel_init[run]	90.7 µs	13.9 µs	×6.5
⚡	Walltime	test_benchmark_spergel_calcfluxrad[run]	194 µs	39.2 µs	×4.9
⚡	CPU Simulation	test_benchmark_spergel_calcfluxrad[run]	2,621.6 µs	687.5 µs	×3.8
⚡	CPU Simulation	test_benchmark_spergel_init[run]	1,367 µs	396.1 µs	×3.5
⚡	CPU Simulation	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	2.8 ms	2.2 ms	+31.86%

Tip

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

---

_{Comparing b-biswas:dtype-consistency (03376b3) with main (eac0cdc)}

[ismael-mendoza]

Thanks @b-biswas for this PR! I haven't appreciated that the float32 does not propagate correctly in this case. Let me know if a review would be heplful whenever

[b-biswas]

@ismael-mendoza let's talk about it in the next meeting... I will write some tests meanwhile!

[b-biswas]

Some tests are failing with the Interpolated images because galsim.InterpolatedImage converts image arrays to np.float64 even if the input image is float32 and now jax_galsim.InterpolatedImage keeps them as float32 leading to numerical instabilities ~ $10^7$

[beckermr]

We'll need to make sure we run some tests in 64 and have others where we run in 32 and test at a reduced precision.

[beckermr]

This won't work. You'll need to use a separate python job for a float32 test.

[ismael-mendoza]

apologies, I accidentally pushed an old branch without rebasing Biswajit's changes. Working on a fix...