Testing fast implementations of A* A where A is the projection operator

[09pouchol]

Dear Aspire Team,

to start off, let me say how grateful I am for all your work in developing and maintaining Aspire!

Problem description

I am currently testing built-in implementations for fast computations (via convolution) of the operator A* A where A is the projection operator, as given in the class mean.MeanEstimator with the reconstruction.mean module. When comparing such an implementation with the successive application of A and then the adjoint A*, I consistently find relative $L^2$ errors of the order of 1.

Note that the operators A and A* are checked to be adjoint of one another, as I have corrected for the bug uncovered thanks to a recent bug report (made by a colleague of mine). Hence the problem could come

• either from me incorrectly using the class MeanEstimator,
• or an actual bug within the corresponding code.

To Reproduce:

The code below defines A, A*, as well as A* A as obtained by the aforementioned method, and compares applying A* A or A and then A*, to randomly chosen vectors.

import numpy as np
import os, aspire
import matplotlib.pyplot as plt
from scipy.ndimage import correlate
from aspire.image import Image
from aspire.utils import Rotation
from aspire.volume import Volume
from aspire.operators import CTFFilter, RadialCTFFilter, FunctionFilter, DualFilter
from aspire.basis import FBBasis3D, FFBBasis3D
from aspire.reconstruction import MeanEstimator
from aspire.source.simulation import Simulation

print("****************\n* Check the identity (A^ast A) u = A^ast (A u)\n*****************")


#########################################################
# define direct & adjoint operators as lambda functions #
#########################################################
L = 16 # Image size
n_img = 1 # number of projections to simulate
dtype = np.float32
def_val= 1e4 # defocus value
rots = Rotation.generate_random_rotations(n=n_img, seed=31415) # random rotations
angles = rots.angles # Retrieve corresponding angles
basis = FFBBasis3D(L, ell_max = None) # Choose a basis to represent volumes (required for mean_kernel function),

# Compute A and A^ast 
ctf = RadialCTFFilter(defocus=def_val) # Ctf function
dual_ctf = DualFilter(ctf) # Define adjoint of ctf, i.e., x \mapsto ctf(-x) #
EMproj = lambda vol,rots : vol.project(rots).filter(ctf) # vol is an aspire volume and rots is an aspire rotation
EMbackproj = lambda proj,rots : (proj.filter(dual_ctf)).backproject(rots.matrices) # proj is an aspire im and rots is an aspire rot

A = lambda x_arr : EMproj(aspire.volume.Volume(x_arr),rots).asnumpy()
adjA = lambda y_arr : EMbackproj(aspire.image.Image(y_arr),rots).asnumpy()[0]

# Set simulation and estimators, necessary to compute B = A^ast A

sim = Simulation(
   L=L,
   n=n_img,
   unique_filters=[RadialCTFFilter(defocus=d) for d in np.linspace(def_val, def_val, 1)],
   offsets = 0,
   dtype=dtype,
   angles = angles
)

mean_estimator = MeanEstimator(sim, basis=basis)
mean_kernel = mean_estimator.kernel
B = lambda x_arr :  n_img*mean_kernel.convolve_volume(x_arr).asnumpy()[0]    


# Lauch unitary test
Nsimu = 100
relmax = -np.inf


for id in range(Nsimu):
   
   # sample random signals 
   u = np.random.rand(L,L,L).astype(np.float32)
   
   # compute A^ast (A u) 
   test1 =  adjA(A(u))
   
   # compute (A^\ast A) u
   test2 = B(u)
   
   
   # compute the L^2 relative error between A^ast (A u) and (A^\ast A) u
   rel = np.sqrt(np.sum((test1-test2)**2)/(np.sum(test1**2)))
   relmax = np.max([rel,relmax])
       
   # display result
   print("simulation %d/%d : relative error = %.3e, max. relative error = %.3e" % (id+1,Nsimu,rel,relmax))

Further discussion

Of course, there are many parameters involved here (such as the volume size, the number of images, the basis chosen to represent volumes, etc), but the problem seems to be quite robust with respect to changing these.

This leads me to a related question: does Aspire host implementations of A* A that do not rely on Fourier-Bessel bases? I could not find one, but I might be wrong.

Thank you in advance for your help,
Camille Pouchol

[garrettwrong]

Hi, thanks for reporting. Sorry for the delay, we were participating in an event last week.

I just ran the script on develop which includes the proj/backproj changes and found the following:

simulation 100/100 : relative error = 8.106e-01, max. relative error = 8.106e-01, is this consistent with your run? I see running the same before the patch is 1-2 magnitudes higher.

Regarding non FB 3D basis, no, not at this time.

[09pouchol]

Thank you for the quick answer.
I just ran the script to be sure, and I get

simulation 100/100 : relative error = 2.092e+00, max. relative error = 2.392e+00

which is not what you get, but roughly has the same order of magnitude. Even in your case, this seems very high since we are dealing with $L^2$ relative errors.

[garrettwrong]

Okay, I suspect you don't have all the patches I put in #1110. I had similar values when I ran on the main branch to compare.

If you are using git, you could try git pull origin develop.

Regardless, I will continue to look into this, thanks.

[janden]

This leads me to a related question: does Aspire host implementations of A* A that do not rely on Fourier-Bessel bases? I could not find one, but I might be wrong.

Just wanted to drop in here to note that the implementation of A*A does not rely on Fourier–Bessel bases. These bases are mostly used to improve the conditioning of the problem, but are not strictly necessary. The implementation works on 3D data on a uniform voxel grid.

[09pouchol]

@garrettwrong all right, thanks again.

@janden thanks for dropping in! I understand that Fourier-Bessel bases are only meant to ease computations. In using the classMeanEstimator, though, I could not find how to not transit through one of these bases. I might be missing that there is an easy way to force the basis to be that associated to a 3D uniform voxel grid?

[janden]

There used to be a Basis class called DiracBasis that would do this, but seems to have been removed in recent versions. I don't recall the reason for doing this exactly, At any rate, the code is here:

import logging

import numpy as np

from aspire.basis import Basis
from aspire.utils import roll_dim, unroll_dim
from aspire.utils.matlab_compat import m_flatten, m_reshape

logger = logging.getLogger(__name__)

class DiracBasis(Basis):
    """
    Define a derived class for Dirac basis
    """

    def __init__(self, size, mask=None, dtype=np.float32):
        """
        Initialize an object for Dirac basis
        :param size: The shape of the vectors for which to define the basis.
        May be a 2-tuple or an integer, in which case, a square basis is assumed.
        :param mask: A boolean _mask of size sz indicating which coordinates
            to include in the basis (default np.full(sz, True)).
        """
        if isinstance(size, int):
            size = (size, size)
        if mask is None:
            mask = np.full(size, True)
        self._mask = m_flatten(mask)

        super().__init__(size, dtype=dtype)

    def _build(self):
        """
        Build the internal data structure to Dirac basis
        """
        logger.info("Expanding object in a Dirac basis.")
        self.count = np.sum(self._mask)
        self._sz_prod = self.nres**self.ndim

    def evaluate(self, v):
        """
        Evaluate coefficients in standard coordinate basis from those in Dirac basis

        :param v: A coefficient vector (or an array of coefficient vectors) to
            be evaluated. The first dimension must equal `self.count`.
        :return: The evaluation of the coefficient vector(s) `v` for this basis.
            This is an array whose first dimensions equal `self.sz` and the remaining
            dimensions correspond to dimensions two and higher of `v`.
        """
        v, sz_roll = unroll_dim(v, 2)
        x = np.zeros(shape=(self._sz_prod,) + v.shape[1:], dtype=self.dtype)
        x[self._mask, ...] = v
        x = m_reshape(x, self.sz + x.shape[1:])
        x = roll_dim(x, sz_roll)

        return x

    def evaluate_t(self, x):
        """
        Evaluate coefficient in Dirac basis from those in standard coordinate basis

        :param x: The coefficient array to be evaluated. The first dimensions
            must equal `self.sz`.
        :return: The evaluation of the coefficient array `v` in the dual basis
            of `basis`. This is an array of vectors whose first dimension equals
             `self.count` and whose remaining dimensions correspond to
             higher dimensions of `v`.
        """
        x, sz_roll = unroll_dim(x, self.ndim + 1)
        x = m_reshape(x, new_shape=(self._sz_prod,) + x.shape[self.ndim :])
        v = np.zeros(shape=(self.count,) + x.shape[1:], dtype=self.dtype)
        v = x[self._mask, ...]
        v = roll_dim(v, sz_roll)

        return v

    def expand(self, x):
        return self.evaluate_t(x)

I haven't tested it, but it should work (perhaps with some small modifications).

[09pouchol]

This will come in very handy. Thanks!

[janden]

No worries. Let me know if you run into any issue.

[garrettwrong]

Added basis (without the m_* functions.) Closing.