float32 consistency

[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 improve performance by 63.41%

⚠️ 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

⚡ 12 improved benchmarks
✅ 26 untouched benchmarks

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	WallTime	test_benchmark_moffat_init[run]	97 µs	42.1 µs	×2.3
⚡	WallTime	test_benchmark_spergel_init[run]	157 µs	88 µs	+78.44%
⚡	WallTime	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	154.7 µs	87.6 µs	+76.66%
⚡	WallTime	test_benchmark_spergel_calcfluxrad[run]	301.1 µs	172 µs	+75.03%
⚡	WallTime	test_benchmarks_lanczos_interp[kval-no_conserve_dc-run]	70.9 µs	41.4 µs	+71.21%
⚡	WallTime	test_benchmarks_lanczos_interp[kval-conserve_dc-run]	69.6 µs	41.5 µs	+67.76%
⚡	WallTime	test_benchmarks_lanczos_interp[xval-no_conserve_dc-run]	130.6 µs	78.3 µs	+66.77%
⚡	WallTime	test_benchmark_invert_ab_noraise[run]	138.1 µs	91.8 µs	+50.46%
⚡	WallTime	test_benchmark_interpimage_flux_frac[run]	489.9 µs	338.2 µs	+44.84%
⚡	WallTime	test_benchmark_spergel_init_float[run]	7.7 µs	5.5 µs	+40.36%
⚡	WallTime	test_benchmark_rng_discard[run]	120.4 µs	86.1 µs	+39.91%
⚡	WallTime	test_benchmark_gaussian_init[run]	7.6 µs	5.4 µs	+39.13%

Tip

Curious why performance improved? Comment @codspeedbot explain why performance improved on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing b-biswas:dtype-consistency (4113589) with main (5d90e4e)}

[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!