Feat/jax dct - Githubissues

peng-lab / BaSiCPy

MIT License

62 stars 20 forks source link

Closed tdmorello closed 2 years ago

tdmorello commented 2 years ago

This pull request adds a custom implementation of the inverse discrete cosine transform (a.k.a. dct type 3).

The function (.tools._jax_idct.idct) was verified against scipy.fft.idct, is n-dimensional, and contains an ortho=norm option for scaling
tools.dct2d_tools was updated to include this custom function
Tests are included

This function should eventually be replaced by jax.scipy.idct whenever that becomes available.

Closes #17

tdmorello commented 2 years ago

Whoops, I see that adding gpu support is a part of #17. I will add that.

tdmorello commented 2 years ago

Note that the JAX dct and idct need to be somewhere inside a @jax.jit decoration to be fast.

Nicholas-Schaub commented 2 years ago

Just need to fix merge conflicts and it's good to go