FLE2D updates

[garrettwrong]

Last week I was informed that there have been updates and bug fixes to the FLE 2D basis which significantly improve the accuracy and should fix the odd/even issue we saw when originally ported.

We should i) port these updates to our implementation ii) confirm the accuracy for our use iii) compare speed/accuracy with FFB2D.

[garrettwrong]

Working on this now. It is unfortunately slow going. The changes the fle_2d authors made are in fact very minor (yay), but our implementation and their implementation are different (woe), particularly in the basis function area. It basically requires the same level of effort as re-porting to understand where those few lines of changes go. Bummer.

The two main fixes so far include a change to the conventions used in the real-complex conversions and a change in their padding approach. There is a third change converting a method (for radial convolution) from a singleton to vectorized (tensor) calls, but we may already have this ability (will check later).

The algorithm for odd images is ultimately still being addressed via a pad-crop to the next even size, just in a different way (lpad) than before (rpad). This wasn't what I was hoping for, and probably also requires that sign fix, but if it works well, sure.

In hind sight it would have been wiser not to port it when it wasn't working well enough or to keep it more/less the same until that time so parallel development could be migrated effectively. Because it was not working well enough to use in practice we didn't gain anything by porting it early.

[garrettwrong]

I was able to implement the padding change today and maintain our tests. I'm moving on to the sign convention changes.

[garrettwrong]

After the work to backport the odd crop-pad change, it doesn't seem to improve anything for us and may make some results less accurate compared to our code. I think the grid and norm fixes that I put in our code might be better :D. I want to double check later, and saved the work if I need to recall it.

I am refreshing myself on the basis function discrepancy. It looks like the issue cited in #738 has actually been resolved in our (develop) code and can probably be closed after this. That's nice. I (develop code) am getting good correspondence between FB2D, FLE2D, and FLE2D using the dense matrix application for all one-hot coef vectors with both odds and evens.

Currently I am going to attempt backporting the FLE2D author's sign convention changes in isolation from the odd patches while keeping the correspondence we have now. If I can do that, I see no reason not to take those changes even though it probably doesn't change anything in practice (IE, we don't use the slow dense method outside unit tests).

[garrettwrong]

Began working through sign changes, but they appear they would only be superficial for our version of the code and not impact the accuracy or speed performance. Development discussion/review decided to pass on sign convention changes.

Instead we will use this time to measure the roundtrip accuracy and speed for a batch of images across various resolution and compare to the current default of FFB2D. Hopefully that will aide in determining next steps.

[garrettwrong]

Okay, first results are in, basically as I expected from working with them. Our implementation of FFB2D is faster, by about a factor of 2 for experimental data sized images. Accuracy is virtually the same using the basis test below. I do not understand the jump around 256 since it comes down for larger sizes, but I've run it again to confirm.

Stripping out the timing and looping code... Generate a batch of masked images, round trip them through the basis, mask off again then compute a relative norm of the differences.

        vol = emdb_2660().astype(np.float64).downsample(sz).astype(dt)
        src = Simulation(vols=vol, offsets=0, amplitudes=1, dtype=dt, n=512)
        mask = grid_2d(sz, dtype=np.float64, normalized=True)['r'] < 0.99
        imgs = src.downsample(sz).images[:]*mask

        ....
            coefs = basis.evaluate_t(imgs)
            ....
            rt_imgs = basis.evaluate(coefs)

            # Zero values outside disk with the mask
            diff = (rt_imgs - imgs)*mask
            rel_err = np.mean(np.linalg.norm(diff, axis=(1,2)) / np.linalg.norm(imgs*mask, axis=(1,2)))
            ....

I'll try to do something similar over rotations next.

[janden]

Thank you for running these tests. I think this provides enough support to keep the FFB basis as a default for now.

[garrettwrong]

No problem. The L_inf norm was a little less dramatic for everything. I was thinking maybe to rerun in bg with expand and a prayer to see if things were better overall. We can talk today if this is inline with expectations, or if I can change things to yield a more fair evaluation. Fortunately, they all do look pretty good these days even if the measurement is off a bit.

The rotation comparison is still running.

[garrettwrong]

Closing this out. For both basis they had error round tripping and rotating one hot coef vectors using evaluate_t. On average the rotations of cryo images seem serviceable for our application, but it is a stretch to call them accurate in general. Using expand faired much better, but the current implementation is remarkably (unusable in experimental sized problems) slow even at lowest possible tolerances. We can consider that as by design or a separate issue.

[oscarmickelin]

Hi! I just saw this and thought I'd chime in.

I think something has gone a little wrong when FLE2D was re-implemented into Aspire. On my machine, with fresh installs of aspire and fle_2d (from the github https://github.com/nmarshallf/fle_2d - the pip install version is a few bug fixes behind, sorry!), I ran the accuracy and timing comparison in the attached code on an M4 Macbook Air, with the following results:

for a 256x256 image, using expand with FLE2D takes roughly 1/6 of the computation time, compared to using expand with FFBbasis as implemented in aspire, and both get comparable errors. The FLE2D precomputation step is also around 20x faster than the FFB precomputation step.

for a 512x512 image, FLE2D and FFB get comparable reconstruction errrors, but expand in FLE2D is roughly 10x faster and it has roughly 50x faster precomputation than FFB.

The only relevant part of the FLE2D github that has really changed since work on the aspire port started is the cropping for odd image sizes and a sign-convention. Neither of these matter for the tests: signs don't matter when doing evaluate(evaluate_t(...)) and the test mentioned here are for even-size images. It could be that the aspire port changed some tolerances/tuned constants in the code or something else, which then changed the performance. Alternatively, maybe the FFBbasis has some more tricks with the vectorization for > 1 image, but that should be fixable if FLE2D is this much faster for a single image.

We will keep maintaining the FLE2D github and add more functionality (there are many fun things that can be done with these basis functions!), so it would be nice if we could find a less painful way to keep the versions in sync. We do want to keep it as a standalone package, since the basis expansions have applications also outside cryo. In particular, I think this is important for all the method-of-moment projects in the group, since these use the FLE2D code to get faster covariance estimation.