Updated LADMAP implementation with JAX JIT

peng-lab / BaSiCPy

MIT License

62 stars 20 forks source link

Updated LADMAP implementation with JAX JIT #58

Closed yfukai closed 2 years ago

yfukai commented 2 years ago

This PR attempts to add

The original approximate algorithm
The updated LADMAP algorithm both accelerated by JAX. (still WIP with some experimental code)

yfukai commented 2 years ago

Still not super-ready, but I think the main logic is working fine. I'd be happy to hear feedback.

yfukai commented 2 years ago

Please note that experimental_code will be deleted at the merge.

yfukai commented 2 years ago

Dropping Python 3.7 as @jit of a class fails.

yfukai commented 2 years ago

Wrongly closed, reopening!

yfukai commented 2 years ago

It seems the tests are not working with the class property changes, but all other parts including baseline calculation are working!

yfukai commented 2 years ago

Still not perfect, but I believe the approximate implementation reproduces the result of the original impl. Please check this https://github.com/yfukai/BaSiCpy/blob/jax_imprementation/misc_notebooks/compare_with_reference_impl.ipynb for the details! @tying84 @Nicholas-Schaub @tdmorello

yfukai commented 2 years ago

I changed the transform test a bit to get it working. Not sure I'm going in the correct direction, so please let me know if you would have any comments! @Nicholas-Schaub

yfukai commented 2 years ago

This will close #4.

yfukai commented 2 years ago

I could add Python 3.7 with setting frozen=True in pydantic config. Ready for review for now.

yfukai commented 2 years ago

Closing as we merge PRs to #62.