iterative ref: grid SO(3) with healpy

[geoffwoollard]

implement simple monte carlo method: random uniform samples.

use geomstats or pytorchgeometric libraries

[geoffwoollard]

Thanks to @fredericpoitevin for the suggestion for healpy: https://healpy.readthedocs.io/en/latest/

healpy is a Python package to handle pixelated data on the sphere. It is based on the Hierarchical Equal Area isoLatitude Pixelization (HEALPix) scheme and bundles the HEALPix C++ library.

Also see the 2012 Relion paper section "2.3. Increasing computational speed: adaptive expectation– maximization"

RELION parameterizes 3D orientations by three Euler angles, and approximates a uniform sampling of the first two Euler angles on the sphere using the HEALPix framework (Gorski et al., 2005). The HEALPix approach was originally proposed for the field of astronomy (where pixelized images of the sky are represented on a sphere), and it has two characteristics that are particularly useful for the adaptive expectation–maximization algorithm outlined above: (i) it yields a reasonable approximation to a uniform sampling of the sphere so that bias towards certain orientations may be minimized; and (ii) it generates discrete grids in a hierarchical way that allows fast calculation of neighbouring sampling points in grids with distinct angular sampling rates. In particular, each subsequent grid in the hierarchy contains four times more sampling points than the previous one, yielding an angular sampling rate that is approximately twice as high.

The implemented adaptive expectation maximization algorithm uses a given grid in the HEALPix hierarchy for the coarse sampling of the first two Euler angles, and the next one in the hierarchy for the fine sampling. In addition, it uses a two times finer, linear sampling of the third Euler angle and of both translations sampling points than the coarse sampling grid. Consequently, the in the fine grid. Thereby, the fine grid will have 25 ¼ 32 times more maximum speed-up of the adaptive approach will be close to 32 (i.e. if only one sampling point contributes to 99.9% of the probability mass on the coarse grid). In practice, the posterior distributions are typically relatively broad during the initial stages of refinement (where low-resolution models provide less information to distinguish different orientations), and these distributions become more ‘‘spiky’’ towards convergence. Therefore, more orientations will contribute significantly to the probability mass on the coarse grid during the first few iterations when speed-ups are typically less pronounced, while towards the end of the refinement speed-ups become much more important.

• Scheres, S. H. W. (2012). RELION: Implementation of a Bayesian approach to cryo-EM structure determination. Journal of Structural Biology, 180(3), 519–530. http://doi.org/10.1016/j.jsb.2012.09.006

[thisFreya]

Looked into healpy - it seems to be adjacent to our application, but doesn't fit our requirements exactly. Regardless it could be used with some awkwardness, but only if n_rotations is a power of 2. Otherwise healpy can't grid the sphere (as it uses the healpix formulation).

In the meantime I wrote a random rotation version.

[geoffwoollard]

Having n_rotations a power of 2 seems OK and useful. It just means that a grid search divides things by halves, instead of less drastically. That seems OK. If you can get some benchmarking for run time on generating samples (e..g. 2, 4, ..., millions) that would be helpful in knowing if/how to use this.

Other than that, we can leave it for now and then circle back when the whole iterative refinement is working.

[thisFreya]

@geoffwoollard I looked a bit more into implementing healpy for this - because we need rotation matrices, which I couldn't find a way for healpy to output, I think the extra computational overhead of converting angles or rotation vectors to matrices will be a detriment to healpy's performance compared to random rotations. For now I think we should stick to randomized.

[geoffwoollard]

@thisTyler pytorch3d has all sorts of conversions. https://pytorch3d.readthedocs.io/en/latest/modules/transforms.html

so if you can get Euler angles, you can easily get a rotation. I see in the healpy docs that only two rotations are output. probably the third could be assumed to be uniform, for uniform rotations.

I looked over the healpy docs. I'm confused about npix, nside, ipix... https://healpy.readthedocs.io/en/latest/healpy_pix.html#conversion-from-to-sky-coordinates

If you have figured that out and can paste in a code snippet that takes in some number for the granularity and returns sets (theta, phi), then I can show you what I mean.

[geoffwoollard]

Oh, actually I think I understand now... https://healpy.readthedocs.io/en/latest/tutorial.html#NSIDE-and-ordering

[thisFreya]

I'm familiar with the workflow you're referring to - I'll look into pytorch and get this working sometime soon then!

[geoffwoollard]

Ok it should be something like so: https://github.com/geoffwoollard/learn_cryoem_math/blob/master/nb/healpy.ipynb

That code snippet will be enough to benchmark for different NPIX.

For benchmarking it shouldn't matter what the actual Euler convention is. We can just use 'ZYZ' as a placeholder. You can check if 'ZYX' and other conventions would affect performance. I doubt it.

[geoffwoollard]

Here are some initial numbers. https://github.com/geoffwoollard/learn_cryoem_math/blob/master/nb/healpy.ipynb

NSIDE	NPIX	time
1	12	324 µs
2	48	1.13 ms
4	192	9.85 ms
8	768	188 ms
16	3072	3.01 s

If you could compare and contrast with sampling, that would be useful.

1. memory, not only compute speed
2. speeds up on GPU?
3. compute and memory time for random sampling
4. how uniform the sampling is for high / low n_rot. Maybe healpy is good for low n_rot, but eventually for high it is memory inefficient and random sampling is OK.