Downsample bug

[itaysason]

Seems like downsample in src/aspire/image/__init__.py doesn't work as the matlab code.
I can put replace the crop and downsample functions in src/aspire/utils/preprocess.py with my code here line 88 and 104.

The implementation is a lot simpler, but for now it doesn't support up sampling (only crop needs to be revised to do so).

[junchaoxia]

I will not say it is a bug. I believe that src/aspire/image/__init__.py is from the Cov3d Matlab code which performs the downsampling in a different way with the ASPIRE matlab code that has been reimplemented partially in src/aspire/utils/preprocess.py. Of cause you can add more functionality and make them callable in Image class. The unit tests need to be added or updated too once you finish the implementation. Please see also the general discussion in issue #122.

[yoelsh]

I don't understand how is it possible that we have two downsample functions that return different results. I am certain that cryo_downsample in MATLAB returns the correct results, so the results of any Python implementation should be compatible with that one.

[junchaoxia]

Depending on the different choices, the implementation details could be different such as scaling the density and cropping. Both of them can be reasonable but reproducing exact result can be difficult. That is the reason to create the issue #122. Theoretical we should only use one copy or at least we should create uniform options for different implementations.

[yoelsh]

We have a reference implementation that we know is correct. Any port of this code must give identical results, for otherwise we won't be able to identify bugs or to test new code.
Also, it makes no sense that different implementations of the same functionality will give different results.

[junchaoxia]

Yes. Make sure to code or update the unit tests for your parts if you think it is right. Otherwise it will be very time consuming for other developers to chasing the problem.

[chris-langfield]

Did a bit of work on this issue today. Came away with two perhaps unsurprising conclusions

1. As a baseline I confirmed Image.downsample() did not change between March 2020 and now. It produces the exact same result.
2. The MATLAB cryo_downsample.m code does not produce the same output as Image.downsample(). Since this method is also used by ImageSource.downsample(), the downsampling done to the sources in general is not the same as the MATLAB.

There is another piece of downsample code in the old commit that I had checked out. This is in src/aspire/aspire/preprocessor.py. I could not get it to work at all so far, but I wonder if this code is a more direct port of the MATLAB. If so, we could revise it and replace our current downsample method with it.

The next step is, of course, to continue trying to use this second downsampling method and compare it to MATLAB. Also to check out @itaysason's branch and see if we can maybe use that code.

Create projections for comparison

src_data = np.load("tests/saved_test_data/clean70SRibosome_vol.npy")
sim = Simulation(n=100, L=65, vols = Volume(src_data))
im = sim.images(0,100)
im.save("projections_70S.mrcs")

Load, downsample and save in Python

src_data = mrcfile.open("projections_70S.mrcs").data
im = Image(src_data)
ds = im.downsample(33)
ds.save("downsample_70S_current.mrcs", overwrite=True)

Load, downsample and save in MATLAB

im = ReadMRC("/Users/langfield/ASPIRE-Python/projections_70S.mrcs")
ds = cryo_downsample(im, 33, true)
WriteMRC(ds, 1.0, "downsample_70S_matlab.mrcs")

Compare results using np.allclose and plot differences

current = mrcfile.open("downsample_70S_current.mrcs").data
matlab = mrcfile.open("/Users/langfield/aspire/io/mrc/downsample_70S_matlab.mrcs").data
if not np.allclose(current, matlab):
    print("XXXXXXXXXXXXX MATLAB downsample does not match current!")
    diff = Image(np.array([current[i] - matlab[i] for i in range(20)]))
    diff.show()

Below are some example deltas between Python and MATLAB downsamples. To my eye the pairs of images looked identical except with slight shifts.

[chris-langfield]

Adding the original Python vs MATLAB downsampled images (resolution is 33)
The images above are the element-wise difference between the Python and MATLAB images

Python:

MATLAB:

[chris-langfield]

The second method from the old code in fact produces the exact same output as the MATLAB! This is in aspire.aspire.preprocessor

Once it ran, the code I had checked out (0eb59315) produced the following images that seemed to be reflected in an odd way

The problem was that the cropped FFT of the image was being cast to real prior to applying the IFFT. The fix was simply to move the float32 cast from before the IFFT to after.

In PreProcessor.downsample()

ASPIRE-Python/src/aspire/aspire/preprocessor.py

Lines 235 to 239 in 0eb5931

	elif ndim == 2:
	# return a 2D image scaled from the original image
	fx = cls.crop(fftshift(fft2(img)), side) * mask
	out = ifft2(ifftshift(fx)) * (np.prod(szout) / np.prod(szin))

ifft2(ifftshift(fx)) ---> ifft2(ifftshift(fx)).astype("float32")

In PreProcessor.crop()

ASPIRE-Python/src/aspire/aspire/preprocessor.py

Lines 97 to 100 in 0eb5931

	if start_x >= 0 and start_y >= 0: # cropping
	start_x, start_y = math.floor(start_x), math.floor(start_y)
	logger.debug(f'crop:cropping from {mat.shape} to {n}..')
	return mat[start_x: start_x + int(n), start_y: start_y + int(n)].astype('float32')

Remove astype("float32")

Python code used to generate the downsampled images (after above change was made). Most of the difficulty was just in converting between the changed conventions for image stacks. It seems that the old code wanted [image_index, L, L], whereas the current code assumes [L,L,image_index]
Again, checked out at 0eb59315f40d4ed7d7539ea0171dab36ade18d5c

from aspire.aspire.preprocessor import PreProcessor
from aspire.image import *
import mrcfile
import numpy as np

data_fp = "projections_70S.mrcs"
src_data = mrcfile.open(data_fp).data
data =  np.stack([src_data[i,:,:] for i in range(src_data.shape[0])],-1)
# Call on each 2D image from MRC because cannot get the "stack" functionality of this method to work
ds_method2 = np.array([PreProcessor.downsample(data[:,:,i], side=33) for i in range(data.shape[2])])
# count along axis 2, not axis 0, otherwise we cannot create an Image (behavior has changed since this stage of the code)
ds_method2 = Image(np.stack([ds_method2[i,:,:] for i in range(ds_method2.shape[0])],2))
# using our CURRENT save method, as the save method in this version saves in a different shape
with mrcfile.new("downsample_70S_march2020_method2.mrcs",overwrite=True) as mrc:
            mrc.set_data(ds_method2.data.astype(np.float32))

Checked out to current code for comparison

import mrcfile
import numpy as np

mar20_method2 = mrcfile.open("/Users/langfield/March2020/downsample_70S_march2020_method2.mrcs").data
# this convention has changed since the old code was written, moves axis 2 back to axis 0
mar20_method2 = np.stack([mar20_method2[:,:,i] for i in range(mar20_method2.shape[2])],0)

matlab = mrcfile.open("/Users/langfield/aspire/io/mrc/downsample_70S_matlab.mrcs").data

assert np.allclose(mar20_method2, matlab)

The main difference between the methods seems to be that in the old code, the FFT, FFTSHIFT, IFFT, etc. were from numpy. In the current code, the FFT method used is determined by the config (cuPy, SciPy, pyfftw). I wonder if the exact correspondence relies on a specific version of numpy?

[chris-langfield]

Implementing the old method from preprocessor.py in the current code replicates the MATLAB. In fact we gain approximately a 40% speed up in downsampling by doing this.

I have made the change on the downsample_matlab branch. It's not ready to go in as some tests that rely on comparisons to hardcoded sample data will fail. Another issue is that by using numpy.fft methods it goes outside the aspire.numeric framework that decides what package a call to fft will use.

Just a proof of concept, and we can decide how to proceed from there.

Cases checked:
✅ 65 -> 33
✅ 65 -> 32
❌ 64 -> 33
✅ 64 -> 32