pmelchior / scarlet2

Scarlet, all new and shiny
MIT License
15 stars 4 forks source link

`scene.set_spectra_to_match` taking to long #69

Open b-remy opened 5 months ago

b-remy commented 5 months ago

Hi @pmelchior , I ran a profiler on the scene_spectra_to_match function and as you can see, the interpolation function (in the resampling renderer) is taking most of the time (3 calls per rendering operation)

image 37 seconds overall

This can be reduced to 6.3 seconds by jitting the resample_obs function in interpolate.py

@partial(jax.jit, static_argnums=(1,2,3,4))
def resample_ops(kimage, shape_in, shape_out, res_in, res_out, interpolant=Quintic()):

image

Interpolating images is indeed expensive, so should we jit this function in the code?

pmelchior commented 4 months ago

I also noticed that this function takes long. Until now, we have only jitted the inner loops of the sampling and the fitting functions because jit usually works best when it sees large code blocks. But it might be a good idea to jit (some part of) the renderer. The most high-level function would be Observation.render Can you try if this helps?