Implement RELION dataloaders

mjo22 commented 8 months ago

This PR is a modified version of #161. It will close #8 (see discussion there).

The PR includes

cryojax conventions compliant with RELION and cisTEM. This entailed making sure that volume axis ordering is (z, y, x) and image axis ordering is (y, x). This fixed a bug in volume coordinate system generation.
A new submodule in cryojax called cryojax.data. This includes abstractions of a particle stack called an AbstractParticleStack and of a dataset called AbstractDataset. Concrete implementations are the RelionParticleStack and RelionDataset. An AbstractDataset is modeled after the pytorch API and can easily be wrapped into a pytorch dataset (see docstring in the AbstractDataset for details).
An example notebook that shows how to read in an image stack!
A slightly modified pattern for parameter type hinting to support vmapping over arbitrary pytree leaves.

mjo22 commented 8 months ago

Also adding a helical relion dataset loader. This simply grabs a particle stack all with the same filament index.

mjo22 commented 8 months ago

Also, a tutorial for reading a particle stack is now in the cryojax/docs/examples folder.

mjo22 / cryojax