Heterogeneous multireference alignment for images with application to 2-D classification in single particle reconstruction

Abstract

Motivated by the task of 2-D classification in single particle reconstruction by cryo-electron microscopy (cryo-EM), we consider the problem of heterogeneous multireference alignment of images. In this problem, the goal is to estimate a (typically small) set of target images from a (typically large) collection of observations. Each observation is a rotated, noisy version of one of the target images. For each individual observation, neither the rotation nor which target image has been rotated are known. As the noise level in cryo-EM data is high, clustering the observations and estimating individual rotations is challenging. We propose a framework to estimate the target images directly from the observations, completely bypassing the need to cluster or register the images. The framework consists of two steps. First, we estimate rotation-invariant features of the images, such as the bispectrum. These features can be estimated to any desired accuracy, at any noise level, provided sufficiently many observations are collected. Then, we estimate the images from the invariant features. Numerical experiments on synthetic cryo-EM datasets demonstrate the effectiveness of the method. Ultimately, we outline future developments required to apply this method to experimental data.

I. Introduction

Single particle reconstruction using cryo-EM is a high-resolution imaging technique used in structural biology to image 3-D structures of macromolecules [1], [2]. In a cryo-EM experiment, multiple samples of a particle are frozen in a thin layer of vitreous ice. Within the ice, the samples are randomly oriented and positioned. The electron microscope produces a tomographic image of the ice and the embedded samples, called a micrograph. The goal is then to estimate the 3-D structure of the particle from the micrograph. Importantly, the signal to noise ratio (SNR) of the micrograph is usually low because of the limited electron dose that can be applied without causing excessive radiation damage.

The first stage of existing cryo-EM algorithmic pipelines is called particle picking. In this stage, one aims to detect the projections of the samples within the micrograph and extract them. We refer to these extracted images as projection images. Throughout this paper, we assume perfect particle picking, that is, we obtain a large number of projection images, each containing a centered particle projection corresponding to an unknown viewing direction.

An important intermediate stage in the 3-D reconstruction procedure of cryo-EM is called 2-D classification. The goal of this stage is to produce $K$ 2-D images—called class averages—with higher SNR. Each one of the $K$ images should represent a subset of the projections taken from a similar viewing direction. The 2-D class averages can be used as templates for particle picking, to construct ab initio 3-D structures [3]–[6], to provide a quick assessment of the particles, to remove picked particles which are associated with non-informative classes, and for symmetry detection [7]–[9].

Different solutions were proposed for the 2-D classification problem. One approach is the reference-free alignment (RFA) technique [10]. RFA tries to align all projection images globally by estimating all individual rotation parameters. However, when the images arise from many different viewing directions, RFA tends to produce large errors as no assignment of in-plane rotational angles can align all images simultaneously. Methods based on expectation-maximization (EM)—an iterative algorithm that aims to find the marginalized maximum likelihood—are also popular. In the context of cryo-EM, the method is usually referred to as Maximum Likelihood 2-D classification (ML2D). The method was first proposed in [11], and is implemented in the popular software package RELION [12], [13]. Nevertheless, the EM framework lacks theoretical analysis and may be computationally expensive. In addition, EM suffers from an intrinsic resolution-computational load trade-off, since the sampling of the in-plane rotation angles is Nyquist sampling. In Section IV we present some numerical results of EM and discuss more of its properties.

A different 2-D classification technique is based on multireference alignment (MRA) [14]–[16]. In MRA, the images are clustered into $K$ classes and the images within each class are averaged to suppress the noise. The averaged images are the class averages. As projection images can be similar up to rotation, the clustering is based on either rotationally aligning the images within each class, or on features of the images which are invariant under rotations, such as autocorrelation [17] or bispectrum [18]. MRA and invariant features play a key role in this paper and are discussed in detail later. Notably, our proposed method avoids the clustering stage, which may be inaccurate at low SNR. Instead, we aim to estimate the $K$ class averages directly from the projection images, with no intermediate clustering stage.

In this paper, we propose to model the 2-D classification problem as an instance of the heterogeneous multireference alignment (hMRA) problem, for the case of 2-D images [19], [20]. In the hMRA problem, we observe $N$ projection images. Each observed image is an in-plane rotated, noisy version of one of the $K$ underlying images—the class averages, corresponding to $K$ viewing directions. For each observation, the specific underlying image and the in-plane rotation are unknown. Since in cryo-EM all in-plane rotations are equally likely to appear, we assume uniform distribution of rotations. The goal is to estimate the $K$ class averages, as well as the distribution among the class averages. Crucially, for each observed particle, which class average it came from (its label) and which in-plane rotation was applied to it are treated as nuisance variables; that is: they are unknown, but we do not seek to estimate them. A detailed mathematical model of the hMRA problem is provided in Section II.

In the high SNR regime, the nuisance variables could be estimated accurately, at least in principle. Given an accurate estimate of these variables, the problem becomes trivial: one can cluster the observations into the $K$ class averages, undo the rotations, and average within each class to suppress the noise. However, in the low SNR regime, estimating the labels and rotations becomes challenging, and indeed impossible as the SNR drops to zero; see for instance [21] for analysis in a related model. Notwithstanding, it was shown in a series of papers that in many MRA setups the underlying signal (or signals in our case) can be estimated at any noise level, provided sufficiently many observations are recorded [19], [22]–[27]. Remarkably, it was shown that in the low SNR regime, the method of moments achieves the optimal sample complexity under rather moderate conditions [24], [28].

Consequently, targeting the low SNR regime, we propose to apply the method of moments of MRA to the 2-D classification problem in cryo-EM. Our work builds upon the notion of bispectrum, first proposed by Tukey [29], and currently used in signal processing [30], [31]. The bispectrum is invariant under rotations; that is, the bispectrum of an image remains unchanged after an arbitrary in-plane rotation. This property enables us to bypass estimation of individual rotations associated with each one of the observations. Under the assumption of uniform distribution of rotations, the bispectrum is equivalent to the third-order moment of the image. Inspired by the seminal work of Kam [32], previous works [20], [23] studied the MRA and hMRA problems for 1-D signals using the bispectrum as a simplified model for the 3-D reconstruction problem in cryo-EM. In this paper, we study the more involved hMRA model for 2-D images as a model for 2-D classification.

In a nutshell, our proposed approach for 2-D classification consists of the following stages. First, we expand each image in a steerable basis. In this paper, we use the Fourier-Bessel basis, but alternative bases, such as the prolate spheroidal wave functions, can be used alternatively [33]. As explained in Section II, in such a basis all in-plane rotations of an image admit the same expansion coefficients, up to complex phase modulations. This property is called steerability. As a result, specific monomials in these coefficients are invariant under in-plane rotations. We refer to these monomials as invariant features. Specifically, we make use of monomials of the first-, second- and third-order called the mean, power spectrum and bispectrum, respectively.

In practice, instead of working directly on the expansion coefficients, we employ a dimensionality-reduction and denoising technique called steerable principal component analysis (sPCA). This technique is similar to the standard PCA, while boosting the SNR by accounting for all in-plane rotations of the data in an efficient way [34]. In addition, the sPCA coefficients preserve the steerability property. Therefore, the invariant monomials can be computed in the sPCA space.

After computing the invariant features of each image, we average over all images. These averages are consistent estimators (up to bias terms that can be removed easily) of the mixed invariant features of the $K$ class averages. Ultimately, a nonconvex least squares (LS) optimization problem is designed to recover the sPCA coefficients of the individual class averages from these mixed invariant features. All the ingredients of this algorithmic pipeline are provided in Sections II and III. Figure 1 illustrates the flowchart of the procedure. Numerical results and comparison with EM are provided in Section IV. Section V concludes this work, discusses its limitations and potential future extensions.

Fig. 1.

Flow chart of Algorithm 1 for 2-D classification using rotationally invariant features.

II. Statistical model and invariant features

In this section, we first describe in detail the hMRA model. Then, we introduce our framework based on computing features that are invariant under rotations in a steerable basis.

A. Statistical model

Let $\left\{I_1,I_2,\dots ,I_K\right\}$ be a set of $K$ images of size $L\times L$, with $L$ odd¹ and pixel values in [0, 1]: these are the class averages, our target parameters. The pixels in an image are indexed by a pair of integers $(x,y)$ with $-(L-1)/2\le x,y\le (L-1)/2$. The support of the images is assumed to lie in the disk $x^2+y^2\le {(L-1)}^2/4$; as a result, any of their rotations have the same property. Let $R_{\theta }$ be a rotation operator which rotates an image counter-clockwise by angle $\theta$, and let $\xi$ be a random variable following a uniform distribution on [$0$, $2\pi$). In addition, let $\pi$ be a random variable on the set $\{1,2,\dots ,K\}$ with distribution $\left({\pi }_1,\dots ,{\pi }_K\right)$:

$${\pi }_k:=\text{ℙ}(\pi =k)>0,k\in \{1,2,\dots ,K\}.$$

Then, our observations are i.i.d. random samples from the model

$$Y=R_{\xi }{I}_{\pi }+\epsilon ,$$ (II.1)

where $\epsilon =\left({\epsilon }_{ij}\right)\in {ℝ}^{L\times L}$ is a noise matrix of i.i.d. Gaussian variables with zero mean and variance ${\sigma }^2$; the random variables $\xi$, $\pi$, $\epsilon$ are independent. Indeed, it was observed that the background noise in cryo-EM experiments can be treated as Gaussian [35].

Suppose we collect $N$ independent observations from the generative model works when L is even. (II.1),

$$𝓨=\left\{Y_1,Y_2,\dots ,Y_N\right\},$$

so that

$$Y_i=R_{{\xi }_i}{I}_{{\pi }_i}+{\epsilon }_i,i=1,\dots ,N.$$

From the observed data $𝓨$, we seek to estimate the target images $\left\{I_1,I_2,\dots ,I_K\right\}$ (class averages) and, possibly, the distribution $\pi$, without estimating the in-plane rotations of individual observations ${\xi }_i$ or the labels ${\pi }_i$. The model (II.1) suffers from unavoidable ambiguities of rotations (of each class average) and permutation (across the $K$ images). Therefore, naturally, a solution is defined up to these symmetries. In Section IV we define a suitable error metric.

B. Steerable basis

As mentioned above, we aim to bypass estimating the nuisance variables by computing features that are invariant under rotation. To this end, we first expand the images with respect to a steerable basis. In polar coordinates, steerable basis functions take the form

$$u^{k,q}(r,\theta )=f^{k,q}(r)e^{\iota k\theta },$$ (II.2)

where $\iota :=\sqrt{-1}$. Notice the separation of variables: If we expand an image in $u^{k,q}$,

$$I(r,\theta )=\sum _{k,a}{a}_{k,q}{u}^{k,q}(r,\theta ),$$ (II.3)

then the expansion of the rotated image follows from:

$$\begin{gathered} \left(R_{\alpha }I\right)(r,\theta )=I(r,\theta -\alpha )=\sum _{k,a}{a}_{k,q}{u}^{k,q}(r,\theta -\alpha ) \\ =\sum _{k,q}{a}_{k,q}{e}^{-\iota k\alpha }{u}^{k,q}(r,\theta ). \end{gathered}$$ (II.4)

Since our images are real, the coefficients satisfy a conjugacy symmetry: $a_{k,q}=\overline{a_{-k,q}}$. Therefore, coefficients with $k\ge 0$ suffice to represent the images.

Examples of steerable bases include the Fourier-Bessel basis and prolate spheroidal wavefunctions. See [34] [36] [33] [37] for efficient expansion algorithms. In this paper, we work with the Fourier-Bessel basis on a disk with radius $c$ defined as:

$$u^{k,q}(r,\theta )=\{\begin{array}{cc}{N}_{k,q}{J}_k\left(R_k\frac{r}{c}\right)e^{\iota k\theta },& r\le c,\\ 0,& r>c,\end{array}$$ (II.5)

where $J_k$ is the Bessel function of the first kind, $R_{k,q}$ is the $q^{th}$ root of $J_k$ and $N_{k,q}={\left(c\sqrt{\pi }\left|J_{k+1}\left(R_{k,q}\right)\right|\right)}^{-1}$ is a normalization factor. We take $c$ to be $(L-1)/2$, in accordance with the assumed support of the images.

To reduce the dimensionality of the representation and denoise the image, we perform sPCA after expanding the images in a steerable basis [34]. The sPCA results in a new, data driven basis to represent the images. Importantly, this new basis preserves the steerability property and consequently the rotation property (II.4) holds true. Section III-A introduces the sPCA technique in more details. With some abuse of notation, in what follows the coefficients in a sPCA basis are also denoted by $\left\{a_{k,q}\right\}$.

C. Invariant features

The steerability property (II.4) enables us to determine features of images which are invariant under rotation. Specifically, features that are invariant to an action of $SO$ (2): the special orthogonal group in 2-D. We assume that the images are “band-limited” in the sense that their expansion in a steerable basis is finite.

From (II.4), it is clear that coefficients $a_{k,q}$ corresponding to $k=0$ are not affected by rotation. Hence, the first-order invariants are just the mean values, or the “DC components”:

$$m_q=a_{0,q},$$ (II.6)

for all $q$. The second-order invariants, which form the power spectrum, are given by

$$p_{k,q_1,q_2}=a_{k,q_1}\overline{a_{k,q_2},}$$ (II.7)

for all $k$, $q_1$, $q_2$. The power spectrum coefficients are invariant to rotation since for all $\alpha$:

$$\left(a_{k,q_1}{e}^{-\iota k\alpha }\right)\left(\overline{a_{k,q_2}{e}^{-\iota k\alpha }}\right)=a_{k,q_1}\overline{a_{k,q_2}}.$$

Unfortunately, the power spectrum does not determine the image uniquely: a multiplication of the expansion coefficients by $e^{\iota h[k]}$ for an arbitrary function $h$ does not change the power spectrum, yet it does change the image.

The third-order invariant, the bispectrum, is defined by

$$b_{k_1,k_2,q_1,q_2,q_3}=a_{k_1,q_1}{a}_{k_2,q_2}\overline{a_{k_1+k_2,q_3}},$$ (II.8)

for all $k_1$, $k_2$, $q_1$, $q_2$, $q_3$. Using equation (II.4), one can verify that indeed:

$$\begin{gathered} \left(a_{k_1,q_1}{e}^{-\iota k_1\alpha }\right)\left(a_{k_2,q_2}{e}^{-\iota k_2\alpha }\right)\left(\overline{a_{k_1+k_2,q_3}{e}^{-\iota \left(k_1+k_2\right)\alpha }}\right) \\ =\left(a_{k_1,q_1}{a}_{k_2,q_2}\overline{a_{k_1+k_2,q_3}}\right)\left(e^{-\iota k_1\alpha }{e}^{-\iota k_2\alpha }\overline{e^{-\iota \left(k_1+k_2\right)\alpha }}\right) \\ =a_{k_1,q_1}{a}_{k_2,q_2}\overline{a_{k_1+k_2,q_3}}. \end{gathered}$$

The combined power spectrum and bispectrum do determine an image uniquely, up to global rotation:

Theorem II.1.

Consider two images with steerable basis coefficients $a_{k,q}$ and $a_{k,q}^{\prime }$, respectively, in the range $-k_{max}\le k\le k_{max}$. Assume that for all $-k_{max}\le k\le k_{max}$ there exists at least one $q$ such that $a_{k,q}\ne 0$. If for all indices

$$a_{k,q_1}\overline{a_{k,q_2}}=a_{k,q_1}^{\prime }\overline{a_{k,q_2}^{\prime }},$$ (II.9)

$$a_{k_1,q_1}{a}_{k_2,q_2}\overline{a_{k_1+k_2,q_3}}=a_{k_1,q_1}^{\prime }{a}_{k_2,q_2}^{\prime }\overline{a_{k_1+k_2,q_3}^{\prime }},$$ (II.10)

then there exists $\theta \in [0,2\pi )$ such that

$$a_{k,q}^{\prime }=a_{k,q}{e}^{-\iota k\theta }.$$ (II.11)

for all $k$, $q$. That is, the two images only differ by a rotation.

Proof: Set $q_1=q_2$ in (II.9), we have $\left|a_{k,q}\right|=\left|a_{k,q}^{\prime }\right|$ for any $k$ and $q$. Hence, $a_{k,q}^{\prime }\ne 0$ if and only if $a_{k,q}\ne 0$, and there exists ${\theta }_{k,q}\in [0,2\pi )$ such that $a_{k,q}^{\prime }=a_{k,q}{e}^{-\iota {\theta }_{k,q}}$. Then, still by (II.9), we have

$$a_{k,q_1}^{\prime }\overline{a_{k,q_2}^{\prime }}=a_{k,q_1}\overline{a_{k,q_2}}{e}^{-\iota \left({\theta }_{k,q_1}-{\theta }_{k,q_2}\right)}.$$

This means that, for fixed $k$, ${\theta }_{k,q}$ take a same value (in [$0$, $2\pi$)) for all $q$ satisfying $a_{k,q}\ne 0$. Hence, for each $k$, there exists a single ${\theta }_k\in [0,2\pi )$ such that

$$a_{k,q}^{\prime }=a_{k,q}{e}^{-\iota {\theta }_k}.$$

Next, by (II.10), we have for all $k_1$, $k_2$, $q_1$, $q_2$, and $q_3$,

$$\begin{gathered} a_{k_1,q_1}^{\prime }{a}_{k_2,q_2}^{\prime }\overline{a_{k_1+k_2,q_3}^{\prime }}=a_{k_1,q_1}{a}_{k_2,q_2}\overline{a_{k_1+k_2,q_3}} \\ \times e^{-\iota \left({\theta }_{k_1}+{\theta }_{k_2}-{\theta }_{k_1+k_2}\right)}. \end{gathered}$$

By assumption, we can always choose $q_1$, $q_2$ and $q_3$ such that $a_{k_1,q_1},a_{k_2,q_2}$ and $a_{k_1+k_2,q_3}$ are all nonzero. Then, we have

$${\theta }_{k_1}+{\theta }_{k_2}={\theta }_{k_1+k_2},$$

for all $-k_{max}\le k_1,k_2\le k_{max}$. This in turn implies

$${\theta }_k=k\theta ,$$

for some constant $\theta$. Thus we conclude that

$$a_{k,q}^{\prime }=a_{k,q}{e}^{-\iota k\theta }.$$

■

III. 2-D classification using invariant features

In this section, we introduce the whole pipeline of our algorithm in detail. We start by expanding the observed images in Fourier-Bessel basis. We then perform sPCA on the resulting coefficients. Next, we estimate the mixed invariants of the $K$ true images (class averages) from the observations’ sPCA coefficients. Ultimately, we estimate the sPCA coefficients of the true images, and thus the images themselves, from the mixed invariants via a nonconvex LS optimization problem. Figure 1 shows a schematic flow chart of our algorithm and Algorithm 1 describes our algorithm step by step. Next, we elaborate on each of the steps.

Algorithm 1:

2-D classification by invariant features

1	Input: Observations $Y_1,\dots ,Y_N$; noise variance ${\sigma }^2$.
2	Expand the observations in the Fourier-Bessel basis, and perform sPCA on the expansion coefficients using the method described in [34].
3	Estimate the mixed invariants of the true images using the sPCA coefficients of the data by (III.5), (III.6) and (III.7).
4	Estimate the sPCA coefficients of the true images and the distribution $\pi$ by solving the optimization problem (III.13).
5	Recover the images ${\widehat{I}}_1,\dots ,{\widehat{I}}_K$ from the sPCA coefficients by (III.16).
6	Output: images ${\widehat{I}}_1,\dots ,{\widehat{I}}_K$ (up to permutation and rotations); distribution $\widehat{\pi }$.

A. Fourier-Bessel sPCA

We perform sPCA on the images after they were expanded in a Fourier-Bessel steerable basis, introduced in Section II-B. Like in a standard PCA, the first step is to subtract the mean observed image from each observation to center the data. The mean image is added back at the last step of the PCA. To ease exposition, we assume (only in this subsection) that the images have zero mean.

For a regular PCA, one would construct the data matrix $X$ such that each column holds the expansion coefficients of one image. Then, PCA would extract the dominant eigenvectors of $X{X}^{\ast }$, where $X^{\ast }$ is the conjugate transpose of $X$. Note that since $X{X}^{\ast }$ is Hermitian, it has real eigenvalues. Since in our model (and in cryo-EM) any observed images could have been observed after arbitrary in-plane rotation with the same probability, we wish to include all such rotated versions of all images in the PCA procedure. This can be done efficiently owing to steerability, as described in [38]. The dominant eigenimages obtained through sPCA form an orthonormal basis for a lower dimensional subspace, where we now project all observations. Crucially, this eigenbasis is also steerable (since each eigenimage is a linear combination of steerable basis functions.) We still get the two usual benefits of PCA—dimensionality reduction and denoising—with the added benefit that we exploited all of the available information. The sPCA has been proven to be an effective denoising tool for cryo-EM reconstruction [18], [33], [39], [40]. The procedure is actually faster than standard PCA because the data covariance $X{X}^{\ast }$ is block-diagonal upon factoring in all rotated images. In the next subsection, we use the expansion coefficients in the sPCA basis to compute invariant features.

B. Estimating the invariant features of the class average images

After performing sPCA, we get a steerable eigenbasis (a collection of eigenimages) and the expansion coefficients of the observed images in that basis (after projection to the corresponding subspace). Then, we can compute the invariant features using these coefficients. The invariants can be computed according to equations (II.6), (II.7) and (II.8). Next, we estimate the mixed invariants of the underlying class average images using the invariants of the noisy data, which we now explain.

Let $\left\{a_{k,q}^i\right\}$ be the sPCA coefficients of the $i\text{th}$ target image $I_i$. As per our model (II.1), in the absence of noise, the coefficients of an observation $Y$ are given by $a_{k,q}^{\pi }{e}^{-\iota k\xi }$. Let $b_{k_1,k_2,q_1,q_2,q_3}^{\pi ,\xi }$ be the bispectrum computed from the latter. By construction, this is independent of $\xi$: this is simply the bispectrum of the target image $I_{\pi }$. Marginalizing over the remaining nuisance variable $\pi$, we find

$${\mathbb{E}}_{\pi }{b}_{k_1,k_2,q_1,q_2,q_3}^{\pi ,\xi }={\mathbb{E}}_{\pi }{a}_{k_1,q_1}^{\pi }{a}_{k_2,q_2}^{\pi }\overline{a_{k_1+k_2,q_3}^{\pi }},$$ (III.1)

where ${\mathbb{E}}_{\pi }$ represents expectation taken against $\pi$: a sum over $\pi =1,\dots ,K$ weighted by $\left({\pi }_1,\dots ,{\pi }_K\right)$.

This relation implies that by averaging over all the bispectra of the observations we can estimate the mixed bispectra of the $K$ class averages. Estimation of mixed mean and power spectra can be similarly obtained. Crucially, we approximate the mixed invariants of the true images without estimating the rotations $\xi$ or the labels $\pi$ of individual observations.

The same method can be applied in the presence of noise. Now, the coefficients of the observations are given by

$$a_{k,q}^{\pi }{e}^{-\iota k\xi }+{\epsilon }_{k,q}^c{e}^{-\iota k\xi },$$ (III.2)

where ${\epsilon }^c$ denotes the complex Gaussian noise in the coefficients, satisfying $\mathbb{E}{\epsilon }^c=0$ and $\mathbb{E}{\epsilon }^c{\left({\epsilon }^c\right)}^{\ast }={\sigma }^{2}I$, where $I$ is the identity matrix. The noise terms induce bias in the power spectrum and bispectrum estimation. Particularly, the noisy power spectrum of $Y$ satisfies

$${\mathbb{E}}_{\pi ,\epsilon }{p}_{k,q_1,q_2}^{\pi ,\xi ,\epsilon }={\mathbb{E}}_{\pi }{a}_{k,q_1}^{\pi }\overline{a_{k,q_2}^{\pi }}+{\sigma }^2{\delta }_{q_1,q_2},$$ (III.3)

where ${\delta }_{q_1,q_2}$ is the Kronecker delta function. Hence, we get a bias term which depends solely on $\sigma$, which is usually estimated in the cryo-EM algorithmic pipeline. Similarly, expectation over the noisy bispectrum results in

$$\begin{gathered} {\mathbb{E}}_{\pi ,\epsilon }{b}_{k_1,k_2,q_1,q_2,q_3}^{\pi ,\xi ,\epsilon }={\mathbb{E}}_{\pi }{a}_{k_1,q_1}^{\pi }{a}_{k_2,q_2}^{\pi }\overline{a_{k_1+k_2,q_3}^{\pi }} \\ +{\sigma }^2{\mathbb{E}}_{\pi }{A}^{\pi }, \end{gathered}$$ (III.4)

where

$$A^{\pi }:={\delta }_{q_2,q_3}{\delta }_{k_1,0}{a}_{0,q_1}^{\pi }+{\delta }_{q_1,q_3}{\delta }_{k_2,0}{a}_{0,q_2}^{\pi }+{\delta }_{q_1,q_2}{\delta }_{k_1+k_2,0}{a}_{0,q_3}^{\pi }.$$

Here, the bias term depends on both ${\sigma }^2$ and the coefficients $a_{0,q}^{\pi }$. Noise does not introduce bias in estimates of the mean.

Equipped with (III.3) and (III.4), estimating the mixed invariants can be executed by averaging over the invariants of the observations and removing the bias terms. Specifically, let $m_q^{Y_i},p_{k,q_1,q_2}^{Y_i},b_{k_1,k_2,q_1,q_2,q_3}^{Y_i}$ be, respectively, the mean, power spectrum and bispectrum of the noisy observation $Y_i$. Then, our estimators of the mixed invariants are easily computed as:

$${\widehat{m}}_q=\frac{1}{N}\sum _{i=1}^N{m}_q^{Y_i},$$ (III.5)

$${\widehat{p}}_{k,q_1,q_2}=\frac{1}{N}\sum _{i=1}^N{p}_{k,q_1,q_2}^{Y_i}-{\sigma }^2{\delta }_{q_1,q_2},$$ (III.6)

$${\widehat{b}}_{k_1,k_2,q_1,q_2,q_3}=\frac{1}{N}\sum _{i=1}^N{b}_{k_1,k_2,q_1,q_2,q_3}^{Y_i}-{\sigma }^2{E}_{\pi }{A}^{\pi }.$$ (III.7)

In (III.7), the bias term $E_{\pi }{A}^{\pi }$ can be estimated by ${\widehat{m}}_q$.

C. Estimating the coefficients of the class averages

In the last section, we showed how the mixed invariants of the class averages can be estimated from the data. Now, we turn our attention to estimating the sPCA coefficients of the class average images from their mixed invariants by a LS optimization problem.

Our optimization problem consists of two types of variables. The first type represents the sPCA coefficients of the target images:

$${\widehat{a}}^i=\left\{{\widehat{a}}_{k,q}^i\right\},i=1,2,\dots ,K.$$ (III.8)

The second type represents the distribution from which the observations are sampled:

$$\widehat{\pi }=\left({\widehat{\pi }}_1,\dots {\widehat{\pi }}_K\right).$$ (III.9)

In practice, as long as we have sufficiently many observations so that the empirical estimates of the invariant features are accurate, the following identities hold approximately:

$$\sum _{i=1}^K{\pi }_i{a}_{0,q}^i\approx {\widehat{m}}_{0,q},$$ (III.10)

$$\sum _{i=1}^K{\pi }_i{a}_{k,q_1}^i\overline{a_{k,q_2}^i}\approx {\widehat{p}}_{k,q_1,q_2},$$ (III.11)

$$\sum _{i=1}^K{\pi }_i{a}_{k_1,q_2}^i{a}_{k_2,q_2}^i\overline{a_{k_1+k_2,q_3}^i}\approx {\widehat{b}}_{k_1,k_2,q_1,q_2,q_3}.$$ (III.12)

In an ideal case, we want ${\widehat{a}}^i$ to be coefficients of $I_i$ (or an in-plane rotated version thereof), and ${\widehat{\pi }}_i={\pi }_i$. Hence, we design an LS problem to minimize the difference between left- and right-hand sides of equations (III.10), (III.11) and (III.12). Let $M_q(\widehat{a},\widehat{\pi }),P_{k,q_1,q_2}(\widehat{a},\widehat{\pi }),B_{k_1,k_2,q_1,q_2,q_3}(\widehat{a},\widehat{\pi })$ capture the left-hand sides of the three equations above, respectively, with the estimators $(\widehat{a},\widehat{\pi })$ rather the unknown, underlying parameters $(a,\pi )$. Our objective function reads

$$\begin{gathered} F(\widehat{a},\widehat{\pi })=\sum _q{\left|M_q(\widehat{a},\widehat{\pi })-{\widehat{m}}_{0,q}\right|}^2 \\ +\frac{1}{1+{\sigma }^2}\sum _{k,q_1,q_2}{\left|P_{k,q_1,q_2}(\widehat{a},\widehat{\pi })-{\widehat{p}}_{k,q_1,q_2}\right|}^2 \\ +\frac{1}{1+{\sigma }^2+{\sigma }^4}\sum _{k_1,k_2,q_1,q_2,q_3}{\left|B_{k_1,k_2,q_1,q_2,q_3}(\widehat{a},\widehat{\pi })-{\widehat{b}}_{k_1,k_2,q_1,q_2,q_3}\right|}^2. \end{gathered}$$ (III.13)

Here we use $1$, $1+{\sigma }^2$ and $1+{\sigma }^2+{\sigma }^4$ as rough estimates for the variances of corresponding terms; see also [20].

Some constraints need to be imposed. First, as variables $\widehat{\pi }$ are used to represent a distribution, they should lie on the simplex, that is, ${\widehat{\pi }}_i\ge 0$ and ${\widehat{\pi }}_1+\cdots +{\widehat{\pi }}_K=1$. Second, as our images are real, and Fourier-Bessel sPCA basis functions with $k=0$ are real, we have $\Im \left(a_{0,q}^i\right)=0$ (where $\Im$ extracts imaginary part). Consequently, we can force ${\widehat{a}}_{0,q}^i$ to be real. Similarly, coefficients with the same $q$ but opposite $k$ are conjugate,

$${\widehat{a}}_{k,q}^i=\overline{{\widehat{a}}_{-k,q}^i}.$$ (III.14)

In the optimization problem, we can just consider those coefficients with nonnegative $k$.

To conclude, we aim to solve the following constrained LS optimization problem:

$$\begin{gathered} \underset{{\widehat{a}}_{k,q}^i\in ℂ,{\widehat{\pi }}_i\in ℝ}{\min }F(\widehat{a},\widehat{\pi }), \\ \text{subject to}\sum _{i=1}^K{\widehat{\pi }}_i=1,{\widehat{\pi }}_i\ge 0, \\ \Im \left({\widehat{a}}_{0,q}^i\right)=0. \end{gathered}$$ (III.15)

While the LS is nonconvex, we find that we can solve it satisfactorily in practice—see Section IV. This is in line with recent related work [20], [23], [26], [41].

We attempt to solve the optimization problem (III.15) by a trust-regions algorithm or conjugate gradient method using Manopt [42]: a toolbox for optimization on manifolds². In our problem, the variable $\widehat{a}$ lies in Euclidean space, while $\widehat{\pi }$ lies on the simplex, whose relative interior is endowed with a Riemannian geometry in the toolbox [43]. When using Manopt, we only provide the gradient on Euclidean space. The gradient on the manifold is computed from the gradient on Euclidean space together with the representation of the manifold. The Hessian is approximated automatically by finite differences of the gradient.

D. Recovery of the images

After solving the optimization problem, we obtain a collection of coefficients ${\widehat{a}}_{k,q}^i$, which are believed to approximate the sPCA coefficients of the target images, $a_{k,q}^i$. To recover the images themselves, up to rotation, we simply compute a linear combination of the basis images given by the sPCA, using coefficients ${\widehat{a}}^i$. We then add the mean image that was subtracted from all observations during the sPCA preprocessing, denoted by $I_m^i$—see Section III-A. Specifically, letting $\text{Φ}$ be the sPCA basis, we recover the images by

$${\widehat{I}}_i=\text{Φ}{\widehat{a}}^i+I_m^i.$$ (III.16)

E. Computational complexity

In this section we discuss the computational complexity of each step of Algorithm 1. By [34], the computational cost of the sPCA step is $𝓞\left(N{L}^3+L^4\right)$, where $N$ is the number of observations and $L$ is the side length of the images. For each image, assume the sPCA provides $M$ components and the maximum angular frequency is $k_{max}$. Then, we obtain $𝓞\left(\frac{M^3}{k_{max}}\right)$ invariants in total [18]. Hence, $𝓞\left(\frac{N{M}^3}{k_{max}}\right)$ computations are required to compute the invariants of all observations and estimate the mixed invariants for groundtruth images. Next, in the optimization part, computing the gradient requires going through all terms in the objective function, and each term contributes $𝓞(K)$ elements of the gradient. Hence, if $T$ iterations are performed, the computational complexity of the optimization step is $𝓞\left(\frac{TK{M}^3}{k_{max}}\right)$. Finally, building the recovered images just involves $K$ linear combinations of the principal components with computational cost $𝓞\left(KM{L}^2\right)$.

IV. Numerical experiments

In this section, we show results of numerical experiments using our algorithm. First, we use random projections of the E. coli 70S ribosome volume [44] as the groundtruth images to explore the performance of our algorithm under different noise levels and distributions. The volume is available in the software ASPIRE package³. The size of each image is 129² (i.e., $L=129$). Figure 2 shows some examples of the class averages and noisy input images. Later, we apply our algorithm on two other molecules with projections of larger size. Code for our algorithm and all experiments is available at https://github.com/chaom1026/2DhMRA. The experiments presented below are conducted by MATLAB on a machine with 4× E7–8880 v3 CPUs, and 750GB of RAM.

Fig. 2.

Examples of class averages (the right column) and rotated noisy images (the first three columns) for the E. coli 70S ribosome volume. $\text{SNR}=1/50$.

Following [20], we define error metrics suitable for the inherent symmetries of our problem. For two images $I$ and $\widehat{I}$, we define a rotationally invariant distance as

$$\text{dist}(I,\widehat{I})=\underset{\theta \in [0,2\pi ]}{\min }{\Vert R_{\theta }I-\widehat{I}\Vert }_{\text{F}}.$$ (IV.1)

This distance measures the Frobenius norm between all rotational alignments of $I$ and $\widehat{I}$. Let $ℐ=\left\{I_1,\dots ,I_K\right\}$ be a set of underlying images, and let $\widehat{ℐ}=\left\{{\widehat{I}}_1,\dots ,{\widehat{I}}_K\right\}$ be the estimates. To be invariant under both rotations and permutations, we use the following definition:

$$\text{dist}{(ℐ,\widehat{ℐ})}^2=\underset{p\in S_K}{\min }\sum _{i=1}^K\text{dist}{\left(I_i,{\widehat{I}}_{p(i)}\right)}^2,$$ (IV.2)

where $S_K$ is the set of all permutations of $\{1,2,\dots ,K\}$. The relative error between $ℐ$ and $\widehat{ℐ}$ is defined by

$${\text{dist}}_r(ℐ,\widehat{ℐ})=\frac{\text{dist}(ℐ,\widehat{ℐ})}{\sqrt{{\sum }_{i=1}^K{\Vert I_i\Vert }_{\text{F}}^2}}.$$ (IV.3)

Note that when computing errors we only consider the disk area with diameter $L$ and ignore the corners. We measure the error between $\widehat{\pi }$ and the true distribution $\pi$ by the total variation (TV) distance which takes values in [0,1]:

$${\text{dist}}_{\text{TV}}(\widehat{\pi },\pi )=\frac{1}{2}\sum _{i=1}^K\left|{\widehat{\pi }}_i-{\pi }_i\right|.$$ (IV.4)

Here we assume that a permutation given by equation (IV.2) has been applied to $\widehat{\pi }$. In what follows, we define the SNR as

$$\text{SNR}=\frac{\mathbb{E}\left({\text{Signal}}^2\right)}{\mathbb{E}\left({\text{Noise}}^2\right)}:=\frac{{\sum }_{i=1}^K{\Vert I_i\Vert }_{\text{F}}^2}{K{L}^2{\sigma }^2}.$$ (IV.5)

During sPCA, we use the method introduced in [34] to choose the eigenimages automatically, based on properties of the Marchenko-Pastur distribution. Specifically, for each frequency $k$, we take those eigenimages with eigenvalues satisfying

$${\lambda }^{(k)}>1.005{\sigma }^2\left(1+\sqrt{{\gamma }_k}\right),$$ (IV.6)

where ${\sigma }^2$ is the variance of the noise, ${\gamma }_0=\frac{p_0}{N}$, and ${\gamma }_k=\frac{p_k}{2N}$ for $k\ne 0$. Here, $p_k$ is the number of eigenimages for frequency $k$ and $N$ is the total number of observations. The factor 1.005 in (IV.6) is chosen heuristically to control the number of sPCA coefficients.

In our first experiment we consider a uniform distribution $\pi$, and assume that the algorithm knows that $\pi$ is uniform. Hence, in the optimization problem, variables ${\widehat{\pi }}_i$ are fixed to be $1/K$. We choose $K=10$, $\text{SNR}=1/50$, and take $N={10}^4$ observations in total. Examples of noisy observations are shown in the third column of Figure 3. With this much noise, it would be challenging to rotationally align and cluster the observations. Nevertheless, the proposed algorithm gets estimates of the images without (even implicitly) doing either alignment or clustering. During sPCA, 83 coefficients are automatically chosen to represent each image, according to (IV.6). Hence, in total we have 830 variables. Figure 3 shows some examples of original images (before and after sPCA), noisy observations and the recovered images by our algorithm. We can see that our algorithm produces accurate recovery of the groundtruth images from noisy samples. We split the measured error into two terms: the error caused by sPCA (the error between groundtruth images before and after sPCA) and estimation error in the sPCA space caused by the optimization problem (the error between recovered images and groundtruth images after sPCA). We refer to these errors as sPCA error and estimation error, respectively. For this experiment, the relative sPCA error compared to the groundtruth images is about 19.6%, while the relative estimation error compared to groundtruth images after sPCA is about 5.2%. Table I shows the CPU time of each step of the algorithm. In this experiment, and all the experiments on uniform distribution in the following, conjugate gradient method is used to solve the optimization problem.

Fig. 3.

First column (left to right): Groundtruth images before sPCA. Second column: Groundtruth images after sPCA. Third column: examples of noisy observations. Fourth column: recovered images by our algorithm, rotated and permuted to align with groundtruth images.

TABLE I.

CPU time cost (in seconds) of different steps for experiments on E. coli 70s ribosome with uniform and non-uniform distribution (Corresponding to Figure 3 and 5).

Step╲Distribution	Uniform	Non-uniform
Computing sPCA	252.2s	259.7s
Computing mixed invariants	16.3s	17.5s
Optimization	428.2s	7157.3s

We conducted experiments to study how the recovery error increases with the noise level. As before, we set $K=10$, $\pi$ is the (known) uniform distribution, and the number of observations is 10⁴. Figure 4 shows the recovery results and relative errors for different SNRs. As can be seen, both the relative sPCA error and estimation error are, more or less, inversely proportional to the SNR. When the noise is larger, the sPCA gives less coefficients and results in larger sPCA errors. In addition, the estimation of the invariants, and thus the coefficients of groundtruth images, are less accurate under larger noise. Of course, when the noise is larger, we need more observations to average out the noise.

Fig. 4.

Recovery results for different SNR. The first to fourth rows of images are groundtruth images, noisy images, groundtruth images after sPCA, and recovered images, respectively. The first to third rows of numbers are SNR, relative sPCA error compared to groundtruth images (the first row), and relative estimation error compared to groundtruth images after sPCA (the third row), respectively. From left to right, the numbers of coefficients chosen by sPCA are 172, 88, 60, 45, 41, 33, 20 and 13.

The next experiment aims to examine our algorithm when optimizing over $\widehat{\pi }$ and the images simultaneously. As before, we fix $K=10$ and $\text{SNR}=1/50$. We take 500 observations for each class for the first 5 classes, and 1500 observations for the other 5 classes, so that

$${\pi }_i=0.05,i=1,\dots ,5;{\pi }_i=0.15,i=6,\dots ,10.$$ (IV.7)

After applying our algorithm, the TV distance between $\widehat{\pi }$ and $\pi$ turns out to be 0.0086. The relative estimation error of all 10 images compared to groundtruth images after sPCA is 6.05%. The relative estimation error of the 5 images with ${\pi }_i=0.05$ is 7.39%, while the relative estimation error of the other 5 images with ${\pi }_i=0.15$ is 3.94%. Empirically, a non-uniform distribution does not influence much the overall quality of the recovery, though it seems underrepresented images suffer more. Figure 5 shows some results of this experiment with nonuniform distribution. CPU time cost is shown in table I. For the experiments with non-uniform distribution, trust-regions algorithm is used to solve the optimization problem. Usually, trust-regions algorithm provide more accurate estimates than conjugate gradient at the cost of running times.

Fig. 5.

First column: Groundtruth images before sPCA. Second column: Groundtruth images after sPCA. Third column: examples of noisy observations. Fourth column: recovered images by our algorithm, rotated to align with groundtruth images. The rows: The first two rows are images with $\pi =0.15$ and the last two rows are images with ${\pi }_i=0.05$.

Next, we study the influence of $\pi$ on the quality of the recovery. In this experiment, we take $K=2$, $\text{SNR}=1/50$, and $N={10}^4$. Table II shows the relative estimation errors on recovered images and the TV error on the recovered distribution as the distribution $\left({\pi }_1,{\pi }_2\right)$ shifts away from uniform. From the table, we can see that in all cases our estimated distributions are close to the true distribution with TV distance less than 0.01. When the number of observations for different images are not equal, the image with more observations tends to have lower estimation error.

TABLE II.

Relative estimation error compared to groundtruth images after sPCA and total variation error of experimental results with varying distribution $\pi =\left({\pi }_1,{\pi }_2\right)$; ${\text{error}}_i$ represents the relative error on image $i=1,2$.

${\pi }_1$	0.1	0.2	0.3	0.4	0.5
${\text{dist}}_{\text{TV}}$	0.0025	0.0007	0.0022	0.0009	0.0014
${\text{error}}_1$	3.59%	5.01%	1.59%	3.60%	3.95%
${\text{error}}_2$	2.01%	4.66%	0.91%	4.08%	4.63%

Usually the projection images are not perfectly centered because the particle picking is not ideal. We conducted experiments to examine the robustness of our algorithm to small shifts of the input images, albeit our model does not take these shifts into account. This time, we consider a problem with only one class $(K=1)$ of size 129 × 129, 5 × 10³ noisy observations and $\text{SNR}=1/50$. A random shift is applied to each observation. The shift is generated by a 2-D uniform distribution on all the shifts within a circle of radius s. When s ranges from 0 to 5, the relative estimation errors of the recovered images are 3.30%, 4.86%, 6.58%, 10.28%, 11.64% and 18.30%, respectively. While the error increases with the size of the shifts, it does so at a reasonable pace: when the shifts are small, the recovery errors are too.

Comparison of our algorithm with the EM method was made. We implement a vanilla version of the EM algorithm, which is applied on sPCA coefficients rather than the images themselves, and considers only a finite set of in-plane rotations. Our EM is different from the EM-based algorithms implemented in cryo-EM software packages (such as RELION), which are more sophisticated and include many heuristics to improve running time and accuracy. Here we aim to underscore the resolution-computational load trade-off of EM In the first experiment, we take $K=5$ classes and 10³ observations per class with $\text{SNR}=1/50$. Figure 6 compares the relative estimation error compared to groundtruth images after sPCA for each class and the running time with different number of in-plane rotations considered by EM. From the figure we can see that EM performs better when the number of rotations is large (≥ 32 in this experiments). However, at the same time EM becomes time-costing, taking nearly 10 times more CPU-time than our algorithm. In the right plot of Figure 6, the time costs of computing invariant features are shown, which are already included in the time cost of our algorithm (the red line). Figure 7 shows the results of another experiment for nonuniform distribution and larger noise. We take the distribution (IV.7) and $\text{SNR}=1/100$, with $K=10$ and $N={10}^4$. We can observe similar phenomena as the last experiment.

Fig. 6.

Comparison of our algorithm (MRA) with EM in terms of relative estimation error and computation time. Images are generated from 5 classes with uniformly random in-plane rotations. Contrary to MRA, EM needs to assume the rotations are selected from a discrete set. Here, we see the accuracy/computation time trade-off of EM for $\text{SNR}=1/50$. Left panel: for each discretization value (i.e., number of uniformly sampled angles) and algorithm, each point represents the relative estimation error compared to groundtruth images after sPCA for one of the classes. Right panel: integers indicate the number of EM iterations. The CPU time is measured in seconds.

Fig. 7.

Comparison of our algorithm (MRA) with EM in terms of relative estimation error and computation time. The 10 classes are distributed nonuniformly according to (IV.7) and $\text{SNR}=1/100$. Left panel: for each discretization value and algorithm, each point represents the relative estimation error compared to groundtruth images after sPCA for one of the classes. Right panel: integers indicate the number of EM iterations. The CPU time is measured in seconds.

To demonstrate that our algorithm applies to other data sets as well, we considered two additional molecules: the transient receptor potential cation channel subfamily V member 1 (TrpV1) [45] and the yeast mitochondrial ribosome [46]. The volumes are downloaded from The Electronic Microscopy Data Bank⁴. Similarly to the experiments before, groundtruth images are projections randomly generated from the volumes. The size of the projections is 181 × 181 and we set $K=10$, $\text{SNR}=1/50$, and $N={10}^4$. Figure 8 shows some results of TrpV1 data. In this experiment, the distribution is taken to be uniform and the variable $\widehat{\pi }$ is fixed. The relative estimation error compared to groundtruth images after sPCA is 4.29%. Figure 9 shows part of the results from the yeast mitochondrial ribosome data. In this experiment, we use the distribution given by (IV.7). The total variation error of $\widehat{\pi }$ is 0.0141. The relative estimation error of all the images compared to groundtruth images after sPCA is 6.03%, and 3.83% and 7.96% for images with ${\pi }_i=0.15$ and ${\pi }_i=0.05$, respectively.

Fig. 8.

Recovery results of the TrpV1 data. First column: Groundtruth images before sPCA. Second column: Groundtruth images after sPCA. Third column: examples of noisy observations. Fourth column: recovered images by our algorithm, rotated to align with groundtruth images.

Fig. 9.

Recovery results of the yeast mitochondrial ribosome data. First column: Groundtruth images before sPCA. Second column: Groundtruth images after sPCA. Third column: examples of noisy observations. Fourth column: recovered images by our algorithm, rotated to align with groundtruth images. The rows: The first two rows are images with $\pi =0.15$. The last two rows are images with ${\pi }_i=0.05$.

V. Conclusion

In this paper, we studied the problem of heterogeneous MRA for 2-D images and proposed a new algorithmic framework for 2-D classification for SPR. Experimental results show that our algorithm can provide high-quality recovery of the groundtruth images (class averages), even when the noise level is high. The algorithm requires only one pass over the data and thus suits for large experimental data sets.

In practice, the projection images in cryo-EM suffer from small random shifts. Hence, a more accurate generative model reads

$$Y=T_s{R}_{\xi }{I}_{\pi }+\epsilon ,$$

where $T_s$ is a small random shift by $s$; compare with (II.1). In future work we intend to extend our framework to take shifts into account. A recent study [26] shows that non-uniform distributions of translations makes MRA easier. We may take this issue into account in the future. Meanwhile, we have shown that our algorithm is robust against small shifts.

More importantly, our algorithm considers a discrete set of viewing directions. Yet, more realistically, cryo-EM micrographs contain projections sampled from a continuous distribution of viewing directions. We hope to extend our algorithm to the continuous case in the future. To apply the proposed techniques to experimental data, it is necessary to handle effects of the contrast transfer functions (CTF) and of colored noise as well.

In [20], [47], it was shown that the number of classes that can be demixed in 1-D hMRA is, approximately, $\sqrt{L}$, where $L$ is the length of the signals. Our experiments indicate that we can demix 40 − 50 classes. How this number depends on the size of the image or number of sPCA coefficients is left for future study.

Appendix A Gradient of the Objective Function

In this section we give the gradient of our objective function (III.13) in Euclidean space. For complex variables $\widehat{a}$, we treat them as a matrix, and define the gradient $\partial F/\partial \widehat{a}$ as the only matrix $g=g(\widehat{a},\widehat{\pi })$ satisfying

$$\Re \left\{\text{Tr}\left(g^{\ast }y\right)\right\}=DF(\widehat{a},\widehat{\pi })[y],$$ (A.1)

where $y$ is a matrix with the same size as $\widehat{a}$, $\Re \left\{\text{Tr}\left(g^{\ast }y\right)\right\}$ is an inner product, and $DF(\widehat{a},\widehat{\pi })[y]$ is the directional derivative of $F$ at $\widehat{a}$ along $y$. As the objective function is a summation of least squares, the gradient of the objective function is the summation of the gradient of least squares terms. Hence, we only need to compute gradients for the following 3 groups of terms:

$${\left|\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{0,q}^i-{\widehat{m}}_{0,q}\right|}^2,$$ (A.2)

$${\left|\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{k,q_1}^i\overline{{\widehat{a}}_{k,q_2}^i}-{\widehat{p}}_{k,q_1,q_2}\right|}^2,$$ (A.3)

$${\left|\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{k_1,q_2}^i{\widehat{a}}_{k_2,q_2}^i\overline{{\widehat{a}}_{k_1+k_2,q_3}^i}-{\widehat{b}}_{k_1,k_2,q_1,q_2,q_3}\right|}^2.$$ (A.4)

We call (A.2), (A.3) and (A.4) the first-, second- and third-order terms according to the order of moments they contain.

A. First-order terms

Let

$${ℳ}_q(\widehat{a},\widehat{\pi })=\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{0,q}^i-{\widehat{m}}_{0,q}.$$

Then, for real variables ${\widehat{\pi }}_i$ we can easily get

$$\frac{\partial {\left|{ℳ}_q(\widehat{a},\widehat{\pi })\right|}^2}{\partial {\widehat{\pi }}_i}=2\Re \left\{{ℳ}_q(\widehat{a},\widehat{\pi }){\widehat{a}}_{0,q}^i\right\}.$$ (A.5)

For complex variables ${\widehat{a}}_{0,q}^i$, we have

$$\frac{\partial {\left|{ℳ}_q(\widehat{a},\widehat{\pi })\right|}^2}{\partial {\widehat{a}}_{0,q}^i}=2{ℳ}_q(\widehat{a},\widehat{\pi }){\widehat{\pi }}_i.$$ (A.6)

For ${\widehat{a}}_{k^{\prime },q^{\prime }}^i$ with $k^{\prime }\ne 0$ or $q^{\prime }\ne q$ we always have

$$\frac{\partial {\left|{ℳ}_q(\widehat{a},\widehat{\pi })\right|}^2}{\partial {\widehat{a}}_{k^{\prime },q^{\prime }}^i}=0,$$ (A.7)

as they do not appear in the term. In the following subsections, we ignore the gradient with respect to such variables.

B. Second-order terms

Let

$$𝓟(\widehat{a},\widehat{\pi }):={𝓟}_{k,q_1,q_2}(\widehat{a},\widehat{\pi })=\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{k,q_1}^i\overline{{\widehat{a}}_{k,q_2}^i}-{\widehat{p}}_{k,q_1,q_2}.$$

Then, for ${\widehat{\pi }}_i$, similar to the last subsection, we have

$$\frac{\partial |𝓟(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{\pi }}_i}=2\Re \left\{𝓟(\widehat{a},\widehat{\pi }){\widehat{a}}_{k,q_1}^i\overline{{\widehat{a}}_{k,q_2}^i}\right\}.$$ (A.8)

For ${\widehat{a}}_{k,q}^i$, there are two cases. If $q_1=q_2$, then

$$\frac{\partial |𝓟(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k,q_1}^i}=4\Re \{𝓟(\widehat{a},\widehat{\pi })\}{\widehat{a}}_{k,q_1}^i{\widehat{\pi }}_i.$$ (A.9)

If $q_1\ne q_2$ then we have

$$\frac{\partial |𝓟(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k,q_1}^i}=2𝓟(\widehat{a},\widehat{\pi }){\widehat{a}}_{k,q_2}^i{\widehat{\pi }}_i,$$ (A.10)

and

$$\frac{\partial |𝓟(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k,q_2}^i}=2\overline{𝓟(\widehat{a},\widehat{\pi })}{\widehat{a}}_{k,q_1}^i{\widehat{\pi }}_i,$$ (A.11)

C. Third-order terms

Let

$$\begin{gathered} ℬ(\widehat{a},\widehat{\pi }):={ℬ}_{k_1,k_2,q_1,q_2,q_3}(\widehat{a},\widehat{\pi }) \\ =\sum _{i=1}^K{\widehat{\pi }}_i{\widehat{a}}_{k_1,q_2}^i{\widehat{a}}_{k_2,q_2}^i\overline{{\widehat{a}}_{k_1+k_2,q_3}^i}-{\widehat{b}}_{k_1,k_2,q_1,q_2,q_3}. \end{gathered}$$

Then, firstly we have

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{\pi }}_i}=2\Re \left\{ℬ(\widehat{a},\widehat{\pi }){\widehat{a}}_{k_1,q_2}^i{\widehat{a}}_{k_2,q_2}^i\overline{{\widehat{a}}_{k_1+k_2,q_3}^i}\right\}.$$ (A.12)

Next, again we consider two cases. If $k_1=k_2$ and $q_1=q_2$, then

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k_1,q_1}^i}=4ℬ(\widehat{a},\widehat{\pi })\overline{{\widehat{a}}_{k_1,q_1}^i}{\widehat{a}}_{k_1+k_2,q_3}^i{\widehat{\pi }}_i,$$ (A.13)

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k_1+k_2,q_3}^i}=2\overline{ℬ(\widehat{a},\widehat{\pi })}{\left({\widehat{a}}_{k_1,q_1}^i\right)}^2{\widehat{\pi }}_i.$$ (A.14)

Otherwise, we have

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k_1,q_1}^i}=2ℬ(\widehat{a},\widehat{\pi })\overline{{\widehat{a}}_{k_2,q_2}^i}{\widehat{a}}_{k_1+k_2,q_3}^i{\widehat{\pi }}_i,$$ (A.15)

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k_2,q_2}^i}=2ℬ(\widehat{a},\widehat{\pi })\overline{{\widehat{a}}_{k_1,q_1}^i}{\widehat{a}}_{k_1+k_2,q_3}^i{\widehat{\pi }}_i,$$ (A.16)

$$\frac{\partial |ℬ(\widehat{a},\widehat{\pi }){|}^2}{\partial {\widehat{a}}_{k_1+k_2,q_3}^i}=2\overline{ℬ(\widehat{a},\widehat{\pi })}{\widehat{a}}_{k_1,q_1}^i{\widehat{a}}_{k_2,q_2}^i{\widehat{\pi }}_i.$$ (A.17)