mjo22 / cryojax

Cryo electron microscopy image simulation and analysis built on JAX.
https://mjo22.github.io/cryojax/
GNU Lesser General Public License v2.1
29 stars 9 forks source link

Add a grid search API in cryojax.inference #207

Closed mjo22 closed 4 months ago

mjo22 commented 5 months ago

Addresses #201. The API is very generic and assumes nothing about cryo-EM. It could be released to the JAX community, operating at a similar level of abstraction as something like optax. I think it will be really useful for everyone to have this as a part of a maintained package, but my dilemma is that I am a little hesitant right now to maintain something for a more general community outside of cryo-EM. So I'm not sure yet if its appropriate to add this into cryojax and keep it there, make it an independent package, or add it now and later make it an independent package.

Aside from this, this PR also includes important packaging updates. The submodules in cryojax were getting a little superfluous. For example, I've eliminated cryojax.core and absorbed cryojax.io into cryojax.data.

mjo22 commented 5 months ago

Whether or not this is added to cryojax, to finish this PR I would like to add an inference tutorial notebook into the documentation for a cross-correlation based grid search.

mjo22 commented 4 months ago

Finished the tutorial for the grid search and ready to merge. This PR now also includes the new API, so merging will mark a big version release.