Volume.project and Image.backproject are not adjoint operators

[remy-abergel]

Hi,

First of all, thank you very much for sharing and maintaining this library.

Bug description:

I am posting here because some numerical experiments have led me to check whether Volume.project and Image.backproject are adjoint operations or not.

To make short, I found out that the answer is no. Actually axis permutation is needed after calling Image.backproject in order to perform the adjoint operation of the Volume.project operation.

EMproj = lambda vol,rots : vol.project(rots) # direct operator (volume projection)
EMbackproj = lambda proj,rots : proj.backproject(rots.matrices) # this is not the adjoint of EMproj
EMbackproj2 = lambda proj,rots : aspire.volume.Volume(np.moveaxis(proj.backproject(rots.matrices).asnumpy()[0],(0,1,2),(2,1,0))) # adjoint of EMproj

I provide below two different ways to check this statement. My apologies for the long post, I tried to be as precise and complete as possible.

To Reproduce:

In the following script, I considered a small 3D volume (vol) with shape (10,10,10) and I simulated 15 projections from this volume, leading to a measurement vector (y) with shape (15,10,10). Because of the small dimensions of the system, we are able to numerically evaluate a matrix representation (M) with shape (1500,1000) of the direct operator. Then, we can perform the back-projection using two different methods :

1. using y.backproject;

2. performing the matrix-vector multiplication of M.transpose() with a flattened representation of y (with shape (1500,)).

Comparing those two outputs yields significant differences.

#####################
# loading libraries #
#####################
import numpy as np
import aspire
import matplotlib.pyplot as plt
from aspire.image import Image
from aspire.utils import Rotation
from aspire.volume import Volume
from aspire.utils import (
    vol_to_vec,
    vec_to_vol,
    vec_to_im,
    im_to_vec
)
plt.ion()

##################################
# retrieve a 3D volume from EMDB #
##################################
from aspire.downloader import emdb_2660
huge = emdb_2660()

##################################################################
# Subsample the 3D volume to reduce the computational complexity #
##################################################################
L = 10
vol = huge.downsample(L)

#########################################################
# define direct & adjoint operators as lambda functions #
#########################################################
n_img = 15 # number of projections to simulate
rots = Rotation.generate_random_rotations(n=n_img, seed=31415) # sample random rotations

# retrieve projection and backprojection operators
EMproj = lambda vol,rots : vol.project(rots) # vol is an aspire volume and rots is an aspire rotation
EMbackproj = lambda proj,rots : proj.backproject(rots.matrices) # proj is an aspire image and rots is an aspire rotation

# same as above using ndarrays inputs and fixed rotations
A = lambda x_ndarr : EMproj(aspire.volume.Volume(x_ndarr),rots).asnumpy()
adjA = lambda y_ndarr : EMbackproj(aspire.image.Image(y_ndarr),rots).asnumpy()[0]

#####################################################################
# simulate noise free measurements (by projection of the 3D volume) #
#####################################################################
y = EMproj(vol,rots) # project the volume to generate the measurements y = (y_i)_{0 <= i < n_img}
y.show() # display the projected images

######################################
# compute the matrix associated to A #
######################################

# set matrix dimensions (m = number of rows, n = number of columns)
m = n_img*L**2
n = L**3

# reshape functions to switch from vector with shape (m,) to aspire.image.Image-like ndarray
imarray_to_vec = lambda im_ndarr : im_to_vec(np.moveaxis(im_ndarr,0,-1)).reshape((m,),order='F')
vec_to_imarray = lambda vec : np.moveaxis(vec_to_im(vec.reshape([L**2,n_img],order='F')),-1,0)

# compute the matrix entries
M = np.zeros((m,n),dtype=vol.dtype)
v = np.zeros((n,),dtype=vol.dtype)
for id in range(n):
    v[id-1] = 0
    v[id] = 1
    M[:,id] = imarray_to_vec(A(vec_to_vol(v)))
    print("computing column %d/%d : done" % (id+1,n))

Mt = np.transpose(M)

# projection and backprojection from matrix-vector multiplication
A2 = lambda x_ndarr : vec_to_imarray(M@vol_to_vec(x_ndarr))
adjA2 = lambda y_ndarr : vec_to_vol(Mt@imarray_to_vec(y_ndarr))

############################################
# recompute y using A2 (consistency check) #
############################################
ynew_ndarr = A2(vol.asnumpy()[0])
ynew = aspire.image.Image(ynew_ndarr)
y.show()
ynew.show()

rel = np.sqrt(np.sum((y.asnumpy()-ynew.asnumpy())**2)/np.sum(y.asnumpy()**2))
print("relative error between y and ynew : ||y-ynew|| / ||y|| =  %.3e " % rel) # should be close to machine epsilon --> OK

###################################
# compare adjoint implementations #
###################################
v1 = adjA(ynew.asnumpy()) # uses image.backproject
v2 = adjA2(ynew.asnumpy()) # uses matrix-vector multiplication

plt.clf()
plt.subplot(1,2,1)
plt.imshow(v1[:,:,np.int32(L/2)],cmap='gray')
plt.title("v1 (computed using image.backproject)")
plt.subplot(1,2,2)
plt.imshow(v2[:,:,np.int32(L/2)],cmap='gray')
plt.title("v2 (computed using matrix-vector multiplication)")

rel = np.sqrt(np.sum((v1-v2)**2)/np.sum(v2**2))
print("relative error between v1 and v2 : ||v1-v2|| / ||v2|| =  %.3e " % rel) # should be close to machine epsilon --> NOT OK

##########################################################
# exchange axes after backprojection: (0,1,2) -> (2,1,0) #
##########################################################
adjA3 = lambda y_ndarr : np.moveaxis(adjA(y_ndarr),(0,1,2),(2,1,0))
v3 = adjA3(ynew.asnumpy()) # uses image.backproject and axis permutation

plt.clf()
plt.subplot(1,2,1)
plt.imshow(v3[:,:,np.int32(L/2)],cmap='gray')
plt.title("v3 (image.backproject + axes permutation)")
plt.subplot(1,2,2)
plt.imshow(v2[:,:,np.int32(L/2)],cmap='gray')
plt.title("v2 (computed using matrix-vector multiplication)")

rel = np.sqrt(np.sum((v3-v2)**2)/np.sum(v2**2))
print("relative error between v3 and v2 : ||v3-v2|| / ||v2|| =  %.3e " % rel) # now we are close to the machine epsilon

Expected behavior:

If the Volume.project and Image.backproject are adjoint operators, then v1 should be the same as v2 (up to the machine epsilon), which is not the case here (we display below one slice of the back-projected volumes computed using the two methods 1) and 2) mentioned above) :

On my laptop, I get 59% of relative difference between v1 and v2.

When reverting the axes order of the back-projected volume v1 (axis order (0,1,2) is changed into (2,1,0)) we obtain a back-projected volume v3 identical (up to machine epsilon) to v2:

On my laptop, the relative difference between v2 and v3 is only $2.7\cdot 10^{-7}$ (and thus, close to the float32 machine epsilon).

Numerical validation of the adjoint relation:

With the script below, we can check that the identity

$$\langle A(u) , v\rangle = \langle u, A^*(v)\rangle$$

is satisfied up to the float32 machine epsilon when axis permutation is performed after the back-projection operation (denoting by $A$ the direct projection operator and by $A^*$ its adjoint).

dtype = 'float32'
Nsimu = 1000
DIROP = lambda u_ndarr,rots : EMproj(aspire.volume.Volume(u_ndarr),rots).asnumpy()
ADJOP = lambda y_ndarr,rots : np.moveaxis(EMbackproj(aspire.image.Image(y_ndarr),rots).asnumpy()[0],(0,1,2),(2,1,0))
relmax = -np.inf
for id in range(Nsimu):
    
    # compute random dimensions
    _L = np.random.randint(10,20)
    _n_img = np.random.randint(1,20)
    
    # compute random rotation matrices
    _rots = Rotation.generate_random_rotations(n=_n_img)
    
    # compute random signals (u,v)
    u = np.array(np.random.rand(_L,_L,_L),dtype=dtype)
    v = np.array(np.random.rand(_n_img,_L,_L),dtype=dtype)
    
    # compute A(u) & adjA(v)
    Au = DIROP(u,_rots)
    adjAv = ADJOP(v,_rots)
    
    # compute inner products
    inprod1 = np.sum(Au * v)
    inprod2 = np.sum(u * adjAv)
    
    # compute & display error
    rel = np.abs(1-inprod2/inprod1)
    relmax = np.max([rel,relmax])
    print("simulation %d/%d, relative error = %.3e, max. relative error = %.3e" % (id+1,Nsimu,rel,relmax))

The above script yields relative errors close to the float32 machine epsilon. When replacing above ADJOP by

ADJOP = lambda y_ndarr,rots : EMbackproj(aspire.image.Image(y_ndarr),rots).asnumpy()[0]

we can observe much larger relative errors (close to 1%).

Environment (please complete the following information):

• OS: Ubuntu 22.04
• Python: 3.10.12
• ASPIRE Version: 0.12.1 dev
• CPU Type: Intel® Core™ i7-7920HQ CPU @ 3.10GHz × 8
• GPU Extensions: yes

[garrettwrong]

Hi, thank you for your outstanding bug report!

ASPIRE recently underwent several changes intended to better match Relion. These were specifically in our projection/backprojection/rotation conventions and it is very possible that we overlooked this while attempting to match that test case.

I think both your cases are resolved by this line in backproject:

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 81589b14..b0f45543 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -548,7 +548,7 @@ class Image:
             )
             pts_rot = pts_rot.reshape((3, -1))
 
-            vol += anufft(im_f, pts_rot[::-1], (L, L, L), real=True)
+            vol += anufft(im_f, pts_rot, (L, L, L), real=True)
 
         vol /= L

However, I'm going to pass this along to my colleague Josh who recently made the Relion convention updates. He may prefer that this change occurs in a different method or have preference as to how we should update the unit tests. We will follow up more next week.

Thank you again for your report; this was one of the best bug reports I've ever seen 🙏 .

[remy-abergel]

Thank you for your reply and for these kind words.

From my side, fixing this bug solved many problems of non-converging optimization schemes (proximal algorithms, conjugate gradient,...).

I may have found another bug related to Image.backproject (probably less critical though). I am going to report it in another thread.

[garrettwrong]

Closed by #1110