Cryo-EM reconstruction of continuous heterogeneity by Laplacian spectral volumes

TL;DR

This paper introduces a novel cryo-EM reconstruction method that captures continuous molecular heterogeneity by analyzing conformational manifolds with spectral volumes derived from Laplacian eigenvectors.

Contribution

The paper presents a new approach combining graph Laplacian analysis and spectral volumes for high-resolution reconstruction of flexible macromolecules.

Findings

Successfully reconstructs continuous heterogeneity in simulated data

Provides high-resolution visualization of molecular dynamics

Theoretically analyzes the spectral volume method

Abstract

Single-particle electron cryomicroscopy is an essential tool for high-resolution 3D reconstruction of proteins and other biological macromolecules. An important challenge in cryo-EM is the reconstruction of non-rigid molecules with parts that move and deform. Traditional reconstruction methods fail in these cases, resulting in smeared reconstructions of the moving parts. This poses a major obstacle for structural biologists, who need high-resolution reconstructions of entire macromolecules, moving parts included. To address this challenge, we present a new method for the reconstruction of macromolecules exhibiting continuous heterogeneity. The proposed method uses projection images from multiple viewing directions to construct a graph Laplacian through which the manifold of three-dimensional conformations is analyzed. The 3D molecular structures are then expanded in a basis of Laplacian eigenvectors, using a novel generalized tomographic reconstruction algorithm to compute the expansion coefficients. These coefficients, which we name spectral volumes, provide a high-resolution visualization of the molecular dynamics. We provide a theoretical analysis and evaluate the method empirically on several simulated data sets.

Tables

Table 1. Table 1: List of symbols. Scalars are denoted by italics, vectors by boldface letters, matrices by non-italicized capitals, estimators are decorated with a hat.

Name	Domain	Description
$n$	$\mathbb{N}$	Number of images and underlying molecular volumes
$s$	$1,\dots ,n$	Index to molecular image/volume
$N$	$\mathbb{N}$	Image/volume size
$\check{N}$	$\mathbb{N}$	Downsampled image/volume size
$𝐱,{𝐱}_{\_}s$	${ℝ}^{N^3}$	Molecular volume
${\widehat{𝐱}}_{\_}s$	${ℝ}^{N^3}$	Our high-resolution molecular volume estimate
$𝐮$	${\{1,\dots ,N\}}^3$	Voxel index
$𝐲,{𝐲}_{\_}s$	${ℝ}^{N^2}$	Molecular image
$𝐡,{𝐡}_{\_}s$	${ℝ}^{N^2}$	Contrast transfer function (CTF)
$R,R_{\_}s$	$\mathrm{SO}\left(3\right)$	3D viewing orientation
$P,P_{\_}s$	${ℝ}^{N^2\times N^3}$	Imaging matrix (rotation, projection, and CTF)
${ℱ}_{\_}d$		The $d$-dimensional discrete Fourier transform
$𝝁$	${ℝ}^{N^3}$ or ${ℝ}^{{\check{N}}^3}$	Mean volume (high-res or low-res)
$\Sigma$	${ℝ}^{{\check{N}}^3\times {\check{N}}^3}$	Covariance matrix of downsampled molecular volumes
$q$	$\mathbb{N}$	Number of PCA eigenvolumes
${\hat{V}}_q$	${ℝ}^{{\check{N}}^3\times q}$	Eigenvolumes of the estimated covariance matrix
$𝜷\left(𝐱\right),{𝜷}_{\_}s$	${ℝ}^q$	PCA coordinates of a molecular volume
$𝐁$	$\subseteq {ℝ}^q$	The domain of PCA coordinates
$\nu \left(𝐁\right)$		Measure of volumes in PCA coordinate representation
$W$	${ℝ}^{n\times n}$	Edge weights matrix
$L$	${ℝ}^{n\times n}$	Graph Laplacian matrix
$ℳ$	$\subset {ℝ}^{N^3}$	Riemannian submanifold of molecular volumes
${\varphi }^{\left(\ell \right)}$	$𝐁\to ℝ$	Laplace–Beltrami eigenfunction of the $\ell$th smallest eigenvalue
$\varphi {\hat{}}^{\left(\ell \right)}$	${ℝ}^n$	Laplacian eigenvector of the $\ell$th smallest eigenvalue
$r$	$\mathbb{N}$	Number of spectral volumes
$K$	${ℝ}^{r{N}^3\times r{N}^3}$	Matrix of weighted projection-backprojections
$𝐛$	${ℝ}^{r{N}^3}$	Concatenation of weighted back-projection images
${𝜶}^{\left(\ell \right)}$	${ℝ}^{N^3}$	Spectral volumes

Table 2. Table 2: Runtimes for the main steps of our method on the ChannelSpin dataset, with n = 10 , 000 𝑛 10 000 n=10,000 images of 108 × 108 108 108 108\times 108 pixels.

Procedure	Running time (sec)
Calculation of $\widehat{𝝁}$	624.6
Calculation of $\hat{\Sigma }$	5044.7
Calculation ${\hat{V}}_q$	0.8
Calculation of {${\widehat{𝜷}}_{\_}s\}$	2084.8
Calculation of $\left\{\varphi {\hat{}}_{\_}s\right\}$	531.5
Calculation of $K$	12378.0
Calculation of $𝐛$	4014.1
Estimation of ${\left\{{\widehat{𝜶}}^{\left(\ell \right)}\right\}}_{\_}\ell =0^{15}$	1769.9

Equations

$${\hat{\mathbf{x}}}_{\_}s:=\sum _{\_}{\ell =0}^{r-1}{\hat{\mathbf{\alpha }}}^{(\ell )}{\hat{\varphi }}_s^{(\ell )},$$

$${\mathbf{y}}_{\_}s={\mathrm{P}}_{\_}s{\mathbf{x}}_{\_}s+{\mathbf{\epsilon }}_{\_}s\forall s=1,2,\dots ,n,$$

$$({\mathcal{F}}_{\_}d\mathbf{s})(\mathbf{k}):=\sum _{\_}\mathbf{u}\in M_{\_}{N}^d{e}^{-2\pi i⟨\mathbf{k},\mathbf{u}⟩}\mathbf{s}[\mathbf{u}]\forall \mathbf{k}\in {\mathbb{R}}^d$$

$$({\mathcal{F}}_{\_}2{\mathrm{P}}_{\_}s{\mathbf{x}}_{\_}s)([k_{\_}1,k_{\_}2{]}^{\mathrm{T}})=({\mathcal{F}}_{\_}3{\mathbf{x}}_{\_}s)(R_{\_}{s}^{-1}[k_{\_}1,k_{\_}2,0{]}^{\mathrm{T}})\cdot ({\mathcal{F}}_{\_}2{\mathbf{h}}_{\_}s)([k_{\_}1,k_{\_}2{]}^{\mathrm{T}}).$$

$${\mathbf{y}}_{\_}s={\mathrm{P}}_{\_}s\mathbf{\mu }+{\mathbf{\epsilon }}_{\_}s\forall s=1,2,\dots ,n,$$

$$\hat{\mathbf{\mu }}=\underset{\mathbf{\mu }\in {\mathbb{R}}^{N^3}}{\mathrm{argmin}}\sum _{\_}{s=1}^n{\Vert {\mathbf{y}}_{\_}s-{\mathrm{P}}_{\_}s\mathbf{\mu }\Vert }^2\text{mbox}.$$

$$\mathbf{x}[\mathbf{u}]\approx \sum _{\_}{\ell =0}^{r-1}{\alpha }_{\_}{\mathbf{u}}^{(\ell )}{\varphi }^{(\ell )}(\mathbf{\beta }(\mathbf{x})),$$

$$\mathbf{x}\approx \sum _{\_}{\ell =0}^{r-1}{\mathbf{\alpha }}^{(\ell )}{\varphi }^{(\ell )}(\mathbf{\beta }(\mathbf{x})),\forall \mathbf{\beta }\in \mathbf{B}.$$

$$\mathbf{\varphi }\hat{}{}^{(0)},\dots ,\mathbf{\varphi }\hat{}{}^{(r-1)}\in {\mathbb{R}}^n.$$

$${\hat{\varphi }}_s^{(\ell )}\approx \frac{1}{\sqrt{n}}{\varphi }^{(\ell )}({\mathbf{\beta }}_{\_}s)\forall s=1,2,\dots ,n,$$

$$\sum _{\_}{s=1}^n{\left({\hat{\varphi }}_s^{(\ell )}\right)}^2=1.$$

$${\mathbf{x}}_{\_}s\approx \sqrt{n}\sum _{\_}{\ell =0}^{r-1}{\mathbf{\alpha }}^{(\ell )}{\hat{\varphi }}_s^{(\ell )}\forall s=1,2,\dots ,n.$$

$${\mathbf{y}}_{\_}s\approx \sqrt{n}\sum _{\_}{\ell =0}^{r-1}\left({\mathrm{P}}_{\_}s{\mathbf{\alpha }}^{(\ell )}\right){\hat{\varphi }}_s^{(\ell )}\forall s=1,2,\dots ,n.$$

$$\left({\hat{\mathbf{\alpha }}}^{(0)},\dots ,{\hat{\mathbf{\alpha }}}^{(r-1)}\right):=\underset{\left\{{\mathbf{\alpha }}^{(0)},\dots ,{\mathbf{\alpha }}^{(r-1)}\right\}}{\mathrm{argmin}}\sum _{\_}{s=1}^n{\Vert {\mathbf{y}}_{\_}s-\sqrt{n}\sum _{\_}{\ell =0}^{r-1}\left({\mathrm{P}}_{\_}s{\mathbf{\alpha }}^{(\ell )}\right){\hat{\varphi }}_s^{(\ell )}\Vert }^2.$$

$${\hat{\mathbf{x}}}_{\_}s=\sqrt{n}\sum _{\_}{\ell =0}^{r-1}{\hat{\varphi }}_s^{(\ell )}{\hat{\mathbf{\alpha }}}^{(\ell )}\forall s=1,2,\dots ,n\text{mbox}.$$

$$\frac{1}{n}\left(\sum _{\_}{s=1}^n{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathrm{P}}_{\_}s\right)\hat{\mathbf{\mu }}=\frac{1}{n}\sum _{\_}{s=1}^n{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathbf{y}}_{\_}s\text{mbox}.$$

$$\hat{\Sigma }=\underset{\Sigma \in {\mathbb{R}}^{N^3\times N^3}}{\mathrm{argmin}}\sum _{\_}{s=1}^n{\Vert ({\mathrm{P}}_{\_}s\Sigma {\mathrm{P}}_{\_}{s}^{\mathrm{T}}+{\sigma }^2{\mathrm{I}}_{\_}{N}^2)-({\mathbf{y}}_{\_}s-{\mathrm{P}}_{\_}s\hat{\mathbf{\mu }})({\mathbf{y}}_{\_}s-{\mathrm{P}}_{\_}s\hat{\mathbf{\mu }}{)}^{\mathrm{T}}\Vert }_{\_}^2\mathrm{F}\text{mbox}.$$

$${\hat{\mathbf{\beta }}}_{\_}s:=\underset{\mathbf{\beta }\in {\mathbb{R}}^q}{\mathrm{argmin}}\frac{1}{{\sigma }^2}{\Vert {\mathbf{y}}_{\_}s-{\mathrm{P}}_{\_}s\left(\hat{\mathbf{\mu }}+{\hat{\mathrm{V}}}_q\mathbf{\beta }\right)\Vert }^2+{\Vert {\Lambda }_{\_}{q}^{-1/2}\mathbf{\beta }\Vert }^2\text{mbox},$$

$$\mathrm{L}:={\mathrm{D}}^{-1/2}(\mathrm{D}-\mathrm{W}){\mathrm{D}}^{-1/2}=\mathrm{I}-{\mathrm{D}}^{-1/2}\mathrm{W}{\mathrm{D}}^{-1/2},$$

$$\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathbf{y}}_{\_}s=\sum _{\_}{m=0}^{r-1}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\hat{\varphi }}_s^{(m)}{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathrm{P}}_{\_}s{\mathbf{\alpha }}^{(m)}\forall \ell =0,1,\dots ,r-1.$$

$${\mathbf{b}}^{(\ell )}=\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathbf{y}}_{\_}s,$$

$${\mathrm{K}}^{(\ell ,m)}=\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\hat{\varphi }}_s^{(m)}{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathrm{P}}_{\_}s.$$

$$\mathbf{b}=\mathrm{K}\mathbf{\alpha }.$$

$${\mathrm{K}}^{(\ell ,\ell )}=\frac{1}{n}\sum _{\_}{s=1}^n{\mathrm{P}}_{\_}{s}^{\mathrm{T}}{\mathrm{P}}_{\_}s.$$

$$\sum _{\_}{s=1}^n{\Vert {\mathbf{z}}_{\_}s-\sqrt{n}\sum _{\_}{\ell =0}^{r-1}{\hat{\varphi }}_s^{(\ell )}{\mathbf{\alpha }}^{(\ell )}\Vert }^2\text{mbox}.$$

$${\hat{\alpha }}^{(\ell )}[\mathbf{u}]=\underset{{\alpha }^{(\ell )}[\mathbf{u}]}{\mathrm{argmin}}\sum _{\_}{s=1}^n{\left|z_{\_}s[\mathbf{u}]-\sqrt{n}\sum _{\_}{\ell =0}^{r-1}{\hat{\varphi }}_s^{(\ell )}{\alpha }^{(\ell )}[\mathbf{u}]\right|}^2\text{mbox}.$$

$${\hat{\alpha }}^{(\ell )}[\mathbf{u}]=\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{z}_{\_}s[\mathbf{u}]=\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}(x_{\_}s[\mathbf{u}]+{\epsilon }_{\_}s[\mathbf{u}]),$$

$${\hat{\mathbf{\alpha }}}^{(\ell )}=\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}\left({\mathbf{x}}_{\_}s+{\mathbf{\epsilon }}_{\_}s\right)=\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\mathbf{x}}_{\_}s+\mathcal{N}\left(0,\frac{{\sigma }^2}{n}{\mathrm{I}}_{\_}{N}^2\right).$$

$$\underset{\_}{\max }\ell \in \{0,\dots ,r-1\}\sum _{\_}{s=1}^n{\left({\hat{\varphi }}_s^{(\ell )}\right)}^4\to 0.$$

$${\hat{\mathbf{\alpha }}}^{(\ell )}=\mathbb{E}\left[\frac{1}{\sqrt{n}}\sum _{\_}{s=1}^n{\hat{\varphi }}_s^{(\ell )}{\mathbf{x}}_{\_}s\right]+O_{\mathrm{P}}\left(\frac{1}{\sqrt{n}}\right),$$

Full text

Keywords: single particle electron cryomicroscopy, heterogeneity, tomographic reconstruction, molecular conformation space, manifold learning, Laplacian eigenmaps, diffusion maps

1 Introduction

The function of biological macromolecules is determined not only by their chemical composition but also by their 3D configuration. Hence, accurately estimating these configurations is of great importance to the field of structural biology. Macromolecules may deform their structure, resulting in a continuum of possible configurations, known as conformations. It is not always possible to isolate different conformations and study each separately. As a result, practitioners often image a heterogeneous sample which is then “purified” computationally.

While X-ray crystallography and nuclear magnetic resonance (NMR) spectroscopy have been very successful in imaging homogeneous molecular structures to high resolution, they rely on aggregate measurements over an entire sample and are therefore ill-suited for imaging heterogeneous molecules. Single-particle electron cryomicroscopy (cryo-EM), on the other hand, produces a separate image for each individual molecule, opening the door to accurate determination of structural variability.

The cryo-EM process consists of rapidly freezing a solution containing the molecular sample and then imaging it using a transmission electron microscope. The electron detector captures a movie where each frame records the electron counts for every pixel. To reduce the effects of ionization damage—which destroys the fine structure of the molecules—the electron dose is kept low, resulting in exceptionally noisy images. See bottom row in Figure 1 for examples. Since each particle is randomly oriented with respect to the imaging plane, the resulting image contains projections of molecules from many random viewing directions. Each projection image is typically modeled as the line integral of the particle’s electric potential along the viewing direction, followed by convolution with a point spread function and the addition of noise [1, 2]. The goal of cryo-EM reconstruction is to invert this process and obtain an estimate of the molecular volume from its projection images. In recent years, better sample preparation techniques and improved detectors have led to reconstructions at a near-atomic resolution [3, 4, 5, 6].

The standard assumption in 3D reconstruction by cryo-EM is that of a homogeneous sample with no structural variability. Many methods for homogeneous 3D reconstruction have been proposed [1, 7, 8, 9, 10]. The prevalent methods are based on a Bayesian approach [11] which starts from some initial guess for the volume and then performs expectation-maximization (EM), alternating between estimating a latent distribution of viewing directions for every image and estimating the volume given these distributions. As discussed above, however, the homogeneous assumption does not hold in general. Resolving molecular structures with variability is known as the heterogeneity problem in single-particle cryo-EM. Two types of heterogeneity are typically considered: discrete and continuous.

1.1 Discrete heterogeneity

This is perhaps the simplest model for heterogeneity. In this model, it is assumed that the particles in the sample can be approximated by a finite number of fixed volumes. An example of a molecule that may be effectively modeled in this way is ATP synthase, an enzyme that acts as a molecular stepper motor and spends most of its time in one of three rotation angles [12].

Several software packages support reconstruction with discrete heterogeneity, also known as 3D classification in the cryo-EM community. These include RELION [13], cryoSPARC [14], FREALIGN/cisTEM [15, 16], and EMAN2 [17]. To recover $K$ distinct volumes, these methods assign, for each particle image, a latent distribution over the set $\{1,\ldots,K\}$. This is incorporated into the EM algorithm, similar to the latent distribution over the viewing directions.

1.2 Continuous heterogeneity

In this model, the molecular volumes in the sample vary continuously subject to the many constraints due to molecular bonds. If the number of degrees of freedom associated with the flexible motion is small then the space of molecular volumes forms a low-dimensional manifold (up to thermal vibrations). Figure 1 shows a simple molecular model with continuous heterogeneity that we use in our simulations. Here, the continuous motion is the free rotation of the top part around the vertical axis. In this case, the manifold of molecular volumes is diffeomorphic to the unit circle $S^{1}$.

One approach for analyzing structural heterogeneity is to perform principal component analysis (PCA) of the 3D molecular structures represented as densities on an $N\times N\times N$ voxel grid. This idea goes back to [18] and was further developed by [19, 20, 21, 22, 23]. These methods estimate the $N^{3}\times N^{3}$ covariance matrix of the 3D volumes and compute its leading eigenvectors, known as eigenvolumes. One variant relies on a consistent least-squares estimator for the covariance [24, 25, 26]. These methods may capture continuous heterogeneity—as illustrated by [26]—but are ill-suited for high-resolution reconstruction, as we discuss in Section 3.3. A notable exception is the method proposed in [27] that attempts to directly compute the leading eigenvectors, at high resolution, without estimating the entire covariance matrix.

A different approach is taken in [28, 29, 30, 31] and is based on diffusion maps, a non-linear dimensionality reduction method that is well-suited for recovering low-dimensional manifold structure [32, 33]. These methods first cluster the projection images by their viewing direction and then compute a separate low-dimensional embedding for each cluster. All of these different embeddings are then aligned, yielding a global embedding of conformations. Sets of close points in the global embedding may then be used to reconstruct a 3D volume corresponding to a particular conformation. This approach faces two important challenges: first, unsupervised global registration of the embeddings is by itself a very challenging problem [34, 35]; second, each individual embedding uses only a small subset of images from a particular viewing direction, which may be insufficient for accurate manifold recovery.

The RELION software package has also been recently extended to include multi-body refinement [36]. This method takes a segmentation of a 3D molecular reconstruction and attempts to refine each part separately from a static base model, with independent viewing directions and shift parameters for each part. Multi-body refinement, however, is limited to rigid variability and may fail to accurately reconstruct the interface between moving parts.

Other methods have been proposed based on normal mode analysis of the molecular structure reconstruction [37, 38]. However, the underlying harmonic oscillator model used in these methods may be too simple to describe sophisticated continuous variability such as structural deformations. See [39] for a survey of methods for studying continuous heterogeneity using cryo-EM.

1.3 Our contribution

We present a new method for recovering continuous variability based on manifold learning. In contrast to the viewing-direction specific manifold estimates of [28], our method directly approximates the global manifold of conformations from all projection images, regardless of their viewing direction.

Throughout this paper, we identify molecular volumes with their electric potential sampled on a 3D voxel grid of dimension $N^{3}$. Under the continuous heterogeneity model, a single molecule corresponds to an embedded submanifold of ${\mathbb{R}}^{N^{3}}$. This manifold is the range of a smooth function that maps a set of conformation parameters to a volume. A standard technique for approximating smooth functions on manifolds is by series expansion in Laplacian eigenfunctions. This technique generalizes the familiar Fourier series expansion in Euclidean space. However, to apply it we need to have the Laplacian eigenfunctions. This is a “chicken and egg” problem: The computation of the Laplacian eigenfunctions requires the distribution of 3D volumes, which is the very thing we would like to estimate. To resolve this problem, we use the covariance-based approach [26] to obtain low-resolution estimates of the 3D volumes. These reconstructions are then used to form an empirical graph Laplacian whose $r$ eigenvectors with lowest eigenvalues $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)},\ldots,\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(r-1)}\in{\mathbb{R}}^{n}$ are used in lieu of the unknown Laplacian eigenfunctions. Then we compute a set of expansion coefficient vectors ${\hat{\boldsymbol{\alpha}}}^{(0)},\ldots,{\hat{\boldsymbol{\alpha}}}^{(r-1)}\in{\mathbb{R}}^{N^{3}}$, which we refer to as spectral volumes. Together, they define a high-resolution 3D reconstruction $\hat{{\bf x}}_{\_}s$ for each projection image:

[TABLE]

To compute the expansion coefficients $\{{\hat{\boldsymbol{\alpha}}}^{(\ell)}\}_{\_}{\ell=0}^{r-1}$, we formulate a novel generalized tomographic reconstruction problem posed as a 3D deconvolution, similar to [40]. The convolution kernel is computed efficiently using a non-uniform fast Fourier transform (NUFFT) [41, 42] and the solution is computed using the conjugate gradient method, leveraging the fast Fourier transform (FFT) for the application of the convolution. These computational details are key for scaling up to high-resolution. See Figure 2 for a diagram of the main steps which constitute our method.

Remark 1**.**

The eigenvectors of the Laplacian can be used not only for function representation but also for non-linear dimensionality reduction (e.g. [43, 32]). In our case, they define an embedding of the low-resolution reconstructions that is useful for visualizing the underlying manifold of conformations. In Figure 5 we show the two-dimensional embedding of the ChannelSpin dataset using the second and third eigenvectors. Despite the high noise levels, the underlying circular manifold of motions is recovered.

Remark 2**.**

The spectral volumes have the same dimensionality as the high-resolution volumes that we reconstruct. They may therefore be visualized as 3D molecular volumes, albeit with negative values as well as positive. These visualizations provide insight regarding the range of motions of the molecule. See Figure 6 for examples and Section 5 for an asymptotic analysis of the spectral volumes.

Section 2 defines the forward model and formulates the inverse problem for continuous heterogeneity in cryo-EM. We describe our method in Section 3, including the generalized tomographic reconstruction from noisy projection images. Section 4 outlines the algorithms used and their computational complexity. In Section 5 we prove the convergence of the spectral volumes and high-resolution reconstructions under the manifold assumption. Finally, we present results on synthetic datasets in Section 6.

2 Problem formulation

We begin by describing the forward model for cryo-EM and then define the inverse problem that we wish to solve, first by considering the simpler case without heterogeneity and then by generalizing to the case of continuous heterogeneity.

2.1 Forward model

A sample of many identical molecules is prepared in a solution and then rapidly frozen, forming a thin sheet of vitreous ice which is then imaged using a transmission electron microscope. The resulting image is a measurement of the electrostatic potential of this thin sheet, integrated along the direction perpendicular to the imaging plane. The individual molecules, known as “particles” in the cryo-EM literature, are all captured in different orientations.

For every molecule in a particular 3D conformation, there is a corresponding real-valued electrostatic density map which we simply refer to as the volume and discretize it on an $N\times N\times N$ grid of voxels. We now describe the data generation model. First, the volumes ${\bf x}_{\_}1,\ldots,{\bf x}_{\_}n$ are drawn i.i.d. from some distribution on ${\mathbb{R}}^{N^{3}}$ which describes the structural variability of the molecule. Then, linear imaging operators $\mathrm{P}_{\_}1,\ldots,\mathrm{P}_{\_}n\in{\mathbb{R}}^{N^{2}\times N^{3}}$ are drawn i.i.d from some distribution. These operators are the composition of a volume rotation operator $R_{\_}s,$ tomographic projection onto the imaging plane, and convolution with a point spread function. The individual particle images ${\bf y}_{\_}1,\ldots,{\bf y}_{\_}n\in{\mathbb{R}}^{N^{2}}$ are formed by

[TABLE]

where $\boldsymbol{\varepsilon}_{\_}s$ are noise terms. For simplicity, we assume that $\boldsymbol{\varepsilon}_{\_}s\sim\mathcal{N}(0,\sigma^{2}\mathrm{I}_{\_}{N\times N})$. The cryo-EM forward operator also includes an in-plane shift after the projection and filtering. In our pipeline, this is estimated and corrected for during the classical reconstruction stage (Figure 2, step (i)).

We consider the volumes ${\bf x}_{\_}s\in{\mathbb{R}}^{N^{3}}$ as functions ${\bf x}_{\_}s:M_{\_}N^{3}\to{\mathbb{R}}$, where $M_{\_}{N}:=[-1,-1+2/N,\ldots,1-2/N]$, is the grid for even values of $N$ (a similar grid may be defined for odd $N$). Similarly, the images ${\bf y}_{\_}s\in{\mathbb{R}}^{N^{2}}$ are functions ${\bf y}_{\_}s:M_{\_}{N}^{2}\to{\mathbb{R}}$. To define the imaging operators $\mathrm{P}_{\_}s$ we must define the tomographic projection operation. One approach to this is in terms of line integrals perpendicular to the projection plane but since the volumes lie on a discrete grid one must incorporate an interpolation scheme. An alternative is to express tomographic projection in the Fourier domain. Let $\bf s$ be a $d$-dimensional signal on $M_{\_}{N}^{d}$, its discrete Fourier transform (DFT) is given by

[TABLE]

where $\bf k$ is a wave vector that corresponds to a particular directional frequency. By the Fourier slice theorem, a tomographic projection along the $z$ axis in the spatial domain is equivalent to a restriction to the $x$-$y$ plane in the Fourier domain [44]. We use this fact to express the projection image $\mathrm{P}_{\_}s{\bf x}_{\_}s$ in the Fourier domain as follows:

[TABLE]

where $[k_{\_}1,k_{\_}2]$ is a wave vector in the resulting 2D projection image, $R_{\_}s\in{\mathbb{R}}^{3\times 3}$ is the rotation of particle number $s$ and ${\bf h}_{\_}s$ is the point-spread function whose Fourier transform $\mathcal{F}_{\_}2{\bf h}_{\_}s$ is known as the contrast transfer function (CTF). See Section 2 of [26] for more details on the forward model.

2.2 Inverse problem

Homogeneous case.  The traditional inverse problem in single-particle cryo-EM assumes that all of the molecular volumes in the sample are identical. Thus, the forward model (2) simplifies to

[TABLE]

where $\boldsymbol{\mu}$ is a mean volume. Suppose the orientations and CTFs are known so that we have the imaging operators $\mathrm{P}_{\_}1,\ldots,\mathrm{P}_{\_}n$. Furthermore, suppose that the images are centered (i.e. in-plane shifts have been accounted for). Then for a white Gaussian noise model, the maximum-likelihood estimate of $\boldsymbol{\mu}$ is the solution to the following least-squares problem:

[TABLE]

This problem and regularized variants of it are not well-posed in general, with the condition number depending on the distribution of the viewing angles, the CTFs, and the desired resolution of the reconstruction. Nevertheless, high accuracy solutions are routinely obtained using cryo-EM software packages. [13, 14, 16, 17].

Continuous heterogeneity. Our main goal when analyzing a heterogeneous sample is to estimate the density of volumes ${\bf x}\in{\mathbb{R}}^{N^{3}}$ associated with a given molecule. We approach this problem by performing reconstructions of the individual volumes ${\bf x}_{\_}1,\ldots,{\bf x}_{\_}n$. Clearly, estimating $nN^{3}$ voxel values from merely $nN^{2}$ noisy measurements is an ill-posed problem and much harder than the homogeneous problem, where only a single volume of $N^{3}$ voxels needs to be estimated. In this paper we make two main assumptions: The first is that the molecular volumes in the sample lie near a low-dimensional manifold. This model is natural since many heterogeneous macromolecules only have a few degrees of freedom that describe their range of motions [28, 29, 30, 31]. Varying these degrees of freedom traces out a smooth, low-dimensional manifold $\mathcal{M}\subset{\mathbb{R}}^{N^{3}}$. The second assumption is that the imaging operators $\mathrm{P}_{\_}s$ can be accurately estimated using standard cryo-EM reconstruction tools. This is the case when the molecule contains a large fixed component and a smaller heterogeneous part. A good indication that this is indeed the case for a particular dataset is when the reconstruction of the mean volume has a high resolution in some regions and lower resolution in others.

In the next section, we explain how we combine these assumptions with spectral techniques for function approximation on low-dimensional spaces to reconstruct all of the volumes in a heterogeneous molecular sample.

3 Methods

In this section, we describe our spectral approach to the reconstruction of molecular samples with continuous heterogeneity. Our approach is based on the representation and approximation of molecular volumes using an orthogonal basis expansion of eigenfunctions. By expanding the molecular volumes in this basis and imposing the projection constraints we obtain a generalized spectral formulation of the cryo-EM reconstruction problem.

3.1 Manifold spectral representation

Our method builds on the output of a low-resolution reconstruction method [26] that we describe in Section 3.3. In this method, each reconstructed volume is a linear combination of $q$ PCA eigenvolumes, hence it defines some mapping $({\bf y}_{\_}s,\mathrm{P}_{\_}s)\mapsto\boldsymbol{\beta}_{\_}s$ where $\boldsymbol{\beta}_{\_}s\in{\mathbb{R}}^{q}$ is the vector of eigenvolume coefficients corresponding to a low-dimensional representation of ${\bf x}_{\_}s$. In what follows, we ignore potential ambiguities due to the projection and consider the low-resolution reconstruction as a linear dimensionality reduction of the underlying volume ${\bf x}_{\_}s\mapsto\boldsymbol{\beta}_{\_}s.$ Since we assumed the underlying manifold of volumes is $d$-dimensional, then if $d<q$ the image of this mapping is some compact domain $\mathbf{B}\subseteq{\mathbb{R}}^{q}$ that is a $d$-dimensional immersed manifold.

In what follows we consider the approximation of smooth functions on general domains $\mathbf{B}$ via eigenfunctions of the Laplacian operator. We briefly review some relevant facts [45]. The Laplacian has a set of real eigenfunctions $\phi^{(\ell)}:\mathbf{B}\to{\mathbb{R}}$ that form a complete orthonormal basis of $L^{2}(\mathbf{B})$ with corresponding non-negative eigenvalues $0=\lambda_{\_}0\leq\lambda_{\_}1\leq\ldots\rightarrow\infty$. The smoothness of $\phi^{(\ell)}$ is controlled by $\lambda_{\_}\ell$, which corresponds to the spatial frequency of $\phi^{(\ell)}$. Consequently, the eigenfunctions with lowest eigenvalues form a natural basis for approximating smooth functions on $\mathbf{B}$. In fact, this basis is optimal for the approximation of smooth functions with $L^{2}$ bounded gradient magnitudes [46]. The idea of using Laplacian eigenfunctions for approximation and regression over arbitrary domains is a generalization of the classical approach for signal representation by Fourier series expansion [47].

Let us therefore consider the basis formed by the first $r$ eigenfunctions $\phi^{(0)},\ldots,\phi^{(r-1)}$. Fix a voxel ${\bf u}\in N^{3}$ and consider its associated restriction function ${\bf x}[{\bf u}]$. We may approximate this function using low-frequency eigenfunctions

[TABLE]

where $\boldsymbol{\beta}({\bf x})\in\mathbf{B}$ is the image of ${\bf x}$ in PCA coordinates. This can be written more succinctly by aggregating the coefficients for all voxels into a single volume, yielding

[TABLE]

We call the coefficient vectors ${\boldsymbol{{\alpha}}}^{(0)},\ldots,{\boldsymbol{{\alpha}}}^{(r-1)}\in{\mathbb{R}}^{N^{3}}$ spectral volumes. Note that the above construction does not rely on a voxel-wise representation of the volumes as the same type of expansion can be done for volumes represented in any spatial basis.

The eigenfunctions are unknown, so we employ a widely used technique from the field of manifold learning, replacing them with estimates given by eigenvectors of a data-driven graph Laplacian. More specifically, we build a weighted undirected graph, where the vertices correspond to the projection images ${\bf y}_{\_}1,\ldots,{\bf y}_{\_}n$ and the edge weights are estimates of the affinity between the underlying molecular conformations. In our case, the affinities are computed from the low-resolution reconstruction coordinate $\boldsymbol{\hat{\beta}}_{\_}s$ described in Section 3.3. We then form the symmetric normalized graph Laplacian and compute its $r$ eigenvectors with the lowest eigenvalues,

[TABLE]

See Section 4.1 for the specific algorithms used for forming the graph and computing these eigenvectors. As we explain in Section 5.3, we may assume that these estimates converge to the eigenfunctions in the sense that

[TABLE]

where the $\sqrt{n}$ factor is needed for proper normalization, so that

[TABLE]

We can now write a data-driven variant of the spectral expansion in (8),

[TABLE]

In the next section we explain how we estimate the coefficients of this expansion.

3.2 Generalized tomographic reconstruction

We assume that the molecular orientations can be accurately estimated using standard methods for homogeneous cryo-EM reconstruction [22, 48], so that the projection operators $\mathrm{P}_{\_}s$ are estimated to high accuracy. By applying the imaging matrix $\mathrm{P}_{\_}s$ to both sides of (12) and plugging in the forward model (2), we obtain

[TABLE]

We seek spectral volumes that minimize the squared error

[TABLE]

The minimizer can be calculated efficiently by forming the normal equations and solving them using the conjugate gradient method. See Section 4.2 for more details on the numerical solution of this minimization problem. Note that in contrast to the low-resolution PCA eigenvolumes, the spectral volumes are at the full resolution $N$. Our high-resolution reconstructions of the molecular volumes are now given by

[TABLE]

This estimator generalizes the least-squares estimator (6) for a single mean volume to multiple volumes ${\hat{\boldsymbol{\alpha}}}^{(0)},\ldots,{\hat{\boldsymbol{\alpha}}}^{(r-1)}$ whose contribution to the reconstructed volumes is given by the Laplacian eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)},\ldots,\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(r-1)}$ defined in Eq. (9).

3.3 Low-resolution reconstruction

While the approach outlined above provides a recipe for computing the eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)},\ldots,\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(r-1)}$ and using them to obtain high-resolution volume estimates, a crucial ingredient is missing still: the graph weights $W_{\_}{ij}$. We would like them to approximate an affinity of the underlying molecular volumes.

Several approaches have been proposed for computing affinities between projection images of heterogeneous ensembles. One of the earliest was to compute affinities using a common-line distance [49], without estimating the relative orientations. This procedure finds the best common-line correspondence out of all candidate common lines, resulting in very noisy affinity estimates. To reduce the noise one can first estimate the orientations of the projection images and then compute the common line distance based on the relative orientation. This was proposed in [50], however, the resulting affinity measure is still very noisy, so the authors first performed 2D class averaging within each set of projection images from the same orientation. However, this may average different conformations together.

We define the affinity $W_{\_}{ij}$ to be the Euclidean distance between the low-resolution reconstructions, obtained using the covariance estimation method [26]. This approach achieves robustness to noise without averaging different conformations together. We now briefly describe their method. The first step is to estimate the mean $\boldsymbol{\mu}=\mathbb{E}[{\bf x}]$ of the distribution of molecular volumes. This is done by taking the derivative of Equation (6) with respect to $\boldsymbol{\mu}$ and setting it equal to zero. This yields the normal equations

[TABLE]

This formulation corresponds to the maximum-likelihood estimator of $\mathbb{E}[{\bf x}]$ in the setting of Gaussian white noise. As a consequence, $\hat{\boldsymbol{\mu}}$ is a consistent estimator [24]. A similar estimator for the covariance matrix $\mathrm{Cov}[{\bf x}]:=\mathbb{E}[({\bf x}-\mathbb{E}[{\bf x}])({\bf x}-\mathbb{E}[{\bf x}])^{\mathrm{T}}]$ is given by

[TABLE]

While not a maximum-likelihood estimator, it is consistent under mild conditions [24]. Computing its normal equations yields a linear system in $O(N^{6})$ variables. Fortunately, this linear system can be reformulated as a deconvolution problem in six dimensions. Precalculating the convolution kernel requires $O(N^{6}\log N+nN^{4})$ operations, but it can then be applied with complexity $O(N^{6}\log N)$. The equations can now be solved using the preconditioned conjugate gradient method. Empirically, it takes around $50$ iterations to converge [26].

While more efficient than a naive approach, the algorithm outlined above still scales poorly in image size $N$. As a result, this covariance estimation method is not currently practical for $N>25$. Furthermore, from a simple dimensionality argument, to estimate the $O(N^{6})$ elements of $\mathrm{Cov}[{\bf x}]$ from $n$ images of size $N\times N$, we need at least $n=O(N^{6}/N^{2})=O(N^{4})$ images. So to apply the algorithm to experimental data, we must first downsample the images from $N\times N$ to $\check{N}\times\check{N}$. It is possible to gain insight on the structural variability using this approach, but the resulting reconstructions are of low-resolution.

After obtaining the mean and covariance estimates $\hat{\boldsymbol{\mu}}$ and $\hat{\Sigma}$, the volumes ${\bf x}_{\_}1,\ldots,{\bf x}_{\_}n$ can be reconstructed by the PCA method introduced in [22]. First, the $q$ eigenvectors, or eigenvolumes, of $\hat{\Sigma}$ are extracted and arranged as columns in a $\check{N}^{3}\times q$ matrix $\hat{\mathrm{V}}_{q}$. They represent the principal directions of molecular volume variability in ${\mathbb{R}}^{\check{N}^{3}}$. Together with the estimated mean, they define an affine $q$-dimensional subspace of ${\mathbb{R}}^{\check{N}^{3}}$ of the form $\hat{\boldsymbol{\mu}}+\hat{\mathrm{V}}_{q}\boldsymbol{\beta}$, where $\boldsymbol{\beta}\in\mathbf{B}\subseteq{\mathbb{R}}^{q}$ is a coordinates vector. Each image ${\bf y}_{\_}s$ may then be associated with a volume in the affine subspace through [26]

[TABLE]

where $\Lambda_{\_}q=\hat{\mathrm{V}}_{q}^{\mathrm{T}}\hat{\Sigma}\hat{\mathrm{V}}_{q}$ is the diagonal matrix of the leading $q$ eigenvalues of $\hat{\Sigma}$. The above estimator is the maximum a posteriori (MAP) estimator of the coordinates of ${\bf x}_{\_}s$ for Gaussian distributions of ${\bf x}_{\_}s$ and $\boldsymbol{\varepsilon}_{\_}s$. It is also equal to the Wiener filter estimator and the linear minimum mean squared error estimator of the coordinates [51, 52].

Given the solutions to (18), we have a low-resolution estimate of each volume ${\bf x}_{\_}s$ given by $\hat{\boldsymbol{\mu}}+\hat{\mathrm{V}}_{q}\boldsymbol{\hat{\beta}}_{\_}s$. We assume that the manifold structure of $\mathcal{M}$ is not destroyed by the mapping of projection images to coordinate vectors in ${\mathbb{R}}^{q}$, hence that it is possible to invert this process and associate a unique molecular conformation with every low-dimensional reconstruction. If the intrinsic dimensionality of the conformation space is low and the volumes vary smoothly along this space then the inverse map $\mathbf{B}\to{\mathbb{R}}^{N^{3}}$ can be approximated by a small number of spectral volumes.

4 Algorithms and computational complexity

In this section, we provide the technical details of our reconstruction method. In Section 4.1 we describe the precise methods used to form the graph Laplacian and compute its eigenvectors, and in Section 4.2 we describe the deconvolution-based solution of the generalized tomographic reconstruction problem (14).

4.1 Graph computations

To compute the PCA eigenvolumes, we begin by downsampling the input images to size $\check{N}\times\check{N}$, where $\check{N}$ is typically about 16. These images are then fed into the mean and covariance estimation pipeline described in [26]. It has computational complexity $O(n\check{N}^{4}+\sqrt{\kappa^{\prime}}\check{N}^{6}\log\check{N})$. The condition number $\kappa^{\prime}$ is of the order of $100$. The top $q$ eigenvectors of the estimated covariance $\hat{\Sigma}$ are computed and the $q$-dimensional coordinates $\boldsymbol{\hat{\beta}}_{\_}s$ of each image are obtained via (18). This step has computational complexity $O(q\check{N}^{3}\log\check{N}+nq^{2}\check{N}^{2})$, following the algorithm described in [26]. A weighted undirected graph is then constructed with vertices that correspond to the images ${\bf y}_{\_}1,\ldots,{\bf y}_{\_}n$ and edge weights calculated from the PCA coordinates $\boldsymbol{\hat{\beta}}_{\_}1,\ldots,\boldsymbol{\hat{\beta}}_{\_}n$. We tested two kinds of weight matrices:

Gaussian kernel weights $W_{\_}{ij}=\mathrm{e}^{-\|\boldsymbol{\hat{\beta}}_{\_}i-\boldsymbol{\hat{\beta}}_{\_}j\|^{2}/2\sigma^{2}}$. 2. 2.

Binary symmetric KNN matrices, whereby $W_{\_}{ij}=1$ if and only if $\boldsymbol{\hat{\beta}}_{\_}i$ is one of the $k$ nearest neighbors of $\boldsymbol{\hat{\beta}}_{\_}j$ or vice versa, and $W_{\_}{ij}=0$ otherwise.

In our preliminary experiments we obtained similar results with both choices. For our final results, we chose to use the symmetric KNN graph since it is sparse, which reduces the memory and computational costs. For the Laplacian matrix, we use the symmetric normalized graph Laplacian

[TABLE]

where $\mathrm{D}$ is a diagonal matrix that satisfies $D_{\_}{ii}=\sum_{\_}jW_{\_}{ij}$. The symmetry of $\mathrm{L}$ permits the use of specialized algorithms for eigenvector calculation and guarantees that the resulting eigenvectors are orthogonal. See the tutorial by [53] for other common choices of weight and Laplacian matrices.

We build the KNN weights matrix $\mathrm{W}$ using MATLAB’s knnsearch function which for low dimensions is based on a KDTree [54]. The running time of this part is $O(qn\log n)$ where $q$ is the dimension of the PCA coordinates $\boldsymbol{\hat{\beta}}_{\_}s$ used in the low-resolution reconstruction. We then form the Laplacian matrix $\mathrm{L}$ and compute its $r$ eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)},\ldots,\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(r-1)}$ with lowest eigenvalues using MATLAB’s eigs function. This function implements the Krylov–Schur algorithm [55]. The matrices $\mathrm{W}$ and $\mathrm{L}$ are stored as sparse matrices of average degree $O(k)$, hence their memory usage is $O(nk)$. There exist newer methods of computing eigenvectors, such as the algebraic multigrid preconditioner used by the megaman manifold learning package [56, 57]. We did not incorporate such methods in the current work, as the eigenvector calculation step was not a bottleneck in our implementation.

4.2 Spectral volume estimation

Recall that the spectral volumes are defined in (14) as minimizers of the generalized tomographic reconstruction equation, (14). To find this minimum, we compute the gradient with respect to $\{{\boldsymbol{{\alpha}}}^{(\ell)}\}_{\_}{\ell=0}^{r-1}$ and set it to zero, obtaining the normal equations

[TABLE]

We can rewrite the equation in vector notation by defining the vectors ${\bf b}^{(0)},\ldots,{\bf b}^{(r-1)}\in\mathbb{R}^{N^{3}}$ to be weighted backprojected images

[TABLE]

and $\mathrm{K}\in\mathbb{R}^{rN^{3}\times rN^{3}}$ to be an $r\times r$ block matrix, with blocks of size $N^{3}\times N^{3}$. Each block is a weighted sum of projection-backprojection matrices, with its $(\ell,m)$ block given by

[TABLE]

By defining the vector ${\bf b}\in\mathbb{R}^{rN^{3}}$ to be the concatenation of ${\bf b}^{(0)},\ldots,{\bf b}^{(r-1)}$ and ${\boldsymbol{{\alpha}}}\in\mathbb{R}^{rN^{3}}$ to be the concatenation of ${\boldsymbol{{\alpha}}}^{(0)},\ldots,{\boldsymbol{{\alpha}}}^{(r-1)}$ we can rewrite (20) as

[TABLE]

Since $\mathrm{K}$ is of size $N^{3}r\times N^{3}r$, it would be very expensive to directly solve this equation using standard direct inversion algorithms such as those based on LU or Cholesky decomposition, since this would require $O(N^{9}r^{3})$ operations. Even merely storing the matrix $\mathrm{K}$ in RAM may be prohibitive. However, if we use an iterative solver such as the conjugate gradient method, we do not need to explicitly store the matrix $\mathrm{K}$ so long as we have an efficient method to apply it. To this end, we draw on the work of [40] and note that applying $\mathrm{P}_{\_}s^{\mathrm{T}}\mathrm{P}_{\_}s$ to a volume is equivalent to convolving that volume with a kernel calculated from $R_{\_}s$ and ${\bf h}_{\_}s$. A complication arises from the fact that the points $R_{\_}s^{-1}[k_{\_}1,k_{\_}2,0]^{\mathrm{T}}$ in (4) do not lie on a regular grid, hence to evaluate the expression $(\mathcal{F}_{\_}3{\bf x}_{\_}s)(R_{\_}s^{-1}[k_{\_}1,k_{\_}2,0]^{\mathrm{T}})$ we need to compute Fourier amplitudes on a non-regular grid which cannot be achieved through the standard FFT. Instead, we use the FINUFFT non-uniform fast Fourier transform software package [58]. It has computational complexity $O(N^{3}\log N+S)$ where $S$ is the number of points at which the transform is computed. Here, $S=N^{2}n$, as both ${\bf h}_{\_}s$ and ${\bf y}_{\_}s$ are of size $N\times N$, and we consider $n$ instances of projection images. We must compute the convolution kernel that corresponds to $\mathrm{K}^{(\ell,m)}$ for each of the $r^{2}$ $(\ell,m)$-pairs, and ${\bf b}^{(\ell)}$ for each $\ell$. Thus, the total time to calculate the convolution kernels of all the blocks of $\mathrm{K}$ is $O(r^{2}N^{3}\log N+r^{2}nN^{2})$. The backprojected images vector ${\bf b}$ is also calculated from $R_{\_}s$, ${\bf h}_{\_}s$, and ${\bf y}_{\_}s$ using a non-uniform FFT at a total computational cost of $O(rN^{3}\log N+rnN^{2})$.

Each step of the conjugate gradient method involves applying the forward operator $\mathrm{K}$ as well as performing several vector dot products and additions. Applying the forward operator is done using $r^{2}$ FFT operations of size $N\times N\times N$, which has a total complexity of $O(r^{2}N^{3}\log N)$. The complexity of the conjugate gradient method is thus $O(\sqrt{\kappa}r^{2}N^{3}\log N)$, where $\kappa$ is the condition number of $\mathrm{K}$, since the conjugate gradient method converges in $O(\sqrt{\kappa})$ steps [59, 60]. In conclusion, the total runtime for solving the normal equations (20) is $O(r^{2}nN^{2}+\sqrt{\kappa}r^{2}N^{3}\log N)$. For our synthetic data sets ChannelSpin and ChannelStretch, using $r=15$ spectral volumes we found that $\kappa$ is of the order of $10$–$30$. See Section 6.3 for empirical runtimes on these data sets.

Remark 3**.**

The running time may be reduced by computing an approximation to $\mathrm{K}$. In the proof of Theorem 1 we show that $\mathrm{K}^{(\ell,m)}\to\delta_{\_}{\ell,m}\mathbb{E}[\mathrm{P}^{\mathrm{T}}\mathrm{P}]$ in probability. We can thus approximate $\mathrm{K}$ by setting the off-diagonal blocks to zero and setting the diagonal blocks to the empirical estimate of $\mathbb{E}[\mathrm{P}^{\mathrm{T}}\mathrm{P}]$

[TABLE]

With this approximation, the time to approximate $\mathrm{K}$ reduces to $O(N^{3}\log N+nN^{2})$ and is now dominated by the computation of ${\bf b}$. The time to multiply vectors by $\mathrm{K}$ is $O(rN^{3}\log N)$, so the total runtime drops by a factor of $r$ to $O(rnN^{2}+\sqrt{\kappa}rN^{3}\log N)$.

5 Theory

In this section, we analyze the solution to the generalized tomographic reconstruction as defined in (14), starting with a simplified special case.

5.1 Warmup: Spectral volumes without projections

We first analyze the solution in an easy setting where the imaging operators $\mathrm{P}_{\_}1,\ldots,\mathrm{P}_{\_}n$ are all equal to the identity matrix. That is, we have direct, albeit noisy, measurements ${\bf z}_{\_}s={\bf x}_{\_}s+\boldsymbol{\varepsilon}_{\_}s$ without projections and point spread function. This case is directly applicable for reconstructing a manifold of 2D images, as we later demonstrate in Section 6.1. In this setting, the spectral volumes ${\hat{\boldsymbol{\alpha}}}^{(0)},\ldots,{\hat{\boldsymbol{\alpha}}}^{(r-1)}$ minimize

[TABLE]

In this sum, each voxel $\bf u$ can be considered separately, giving

[TABLE]

For a symmetric graph Laplacian $\mathrm{L}$, the eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}_{\_}0,\ldots,\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}_{\_}{r-1}$ form an orthonormal set. Hence the coefficient ${\hat{\alpha}}^{(\ell)}[{\bf u}]$ is given by an orthonormal projection of $z[{\bf u}]$ onto $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(\ell)}$

[TABLE]

or, in vector form,

[TABLE]

The last equality stems from the fact that the noise terms satisfy $\boldsymbol{\varepsilon}_{\_}s\sim\mathcal{N}(0,\sigma^{2}\mathrm{I}_{\_}{N^{2}})$. Consequently, the spectral volumes in this simplified model are, up to a noise term, orthogonal projections of the true volumes ${\bf x}_{\_}1,\ldots,{\bf x}_{\_}n$ onto the basis of Laplacian eigenvectors. In the next subsection, we show that this is also the case when tomographic projections are incorporated into the model.

5.2 Spectral volumes with projections

We now consider the full forward model with non-trivial imaging operators $\mathrm{P}_{\_}1,\ldots,\mathrm{P}_{\_}n$. First note that in our model, the images ${\bf y}_{\_}1,\ldots,{\bf y}_{\_}n$ and the imaging operators are random vectors, therefore the Laplacian eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)}\ldots\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(r-1)}\in{\mathbb{R}}^{n}$ are also random vectors. For our analysis, we make the following two assumptions:

Assumption 1**.**

Let ${\bf y}_{\_}s$ be an image, drawn according to the forward model (2). Then its Laplacian eigenvector coordinates $\hat{\phi}_{s}^{(0)},\ldots,\hat{\phi}_{s}^{(r-1)}$ are independent of $\mathrm{P}_{\_}s$.

In other words, the Laplacian eigenmap (or diffusion map) coordinates are independent of the viewing direction and CTF of the particle. We can justify this assumption by assuming that the covariance-based method of [26] performs accurate low-resolution reconstruction, regardless of the viewing angle.

Assumption 2**.**

For any $r>0$, the following sum converges in probability:

[TABLE]

That is, for any $\epsilon>0$, the probability that $\sum_{\_}{s=1}^{n}\left(\hat{\phi}_{s}^{(\ell)}\right)^{4}>\epsilon$ tends to zero as $n\to\infty$.

Note that from the normalization constraint $\sum_{\_}{s=1}^{n}(\hat{\phi}_{s}^{(\ell)})^{2}=1$, unless the energy of the eigenvectors is highly concentrated, we expect to have $\hat{\phi}_{s}^{(\ell)}\sim 1/\sqrt{n}$ and thus $\sum_{\_}{s=1}^{n}(\hat{\phi}_{s}^{(\ell)})^{4}\sim 1/n$, in which case Assumption 2 holds. Under standard assumptions the Laplacian eigenvectors converge to limiting eigenfunctions of some differential operator. As we show in Section 5.3, if the eigenfunctions are bounded and this spectral convergence holds then Assumption 2 follows.

Before stating our main result, we recall big-O in probability notation for stochastic boundedness: a sequence of random variables $\{X_{\_}n\}_{\_}{n=1}^{\infty}$ satisfies $X_{\_}n=O_{\mathrm{P}}(f(n))$ if for every $\epsilon>0$, there is some bound $M_{\_}\epsilon$ such that $\Pr[|X_{\_}n|/f(n)>M_{\_}\epsilon]<\epsilon$. We now state our main result which characterizes the estimated spectral volumes ${\hat{\boldsymbol{\alpha}}}^{(0)},\ldots,{\hat{\boldsymbol{\alpha}}}^{(r-1)}$ up to a stochastically bounded error.

Theorem 1**.**

(Spectral volume convergence) Let $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(\ell)}$ be an eigenvector of the symmetric graph Laplacian described in Section 4.1. Under Assumptions 1 and 2, the spectral volumes as defined in (14) satisfy

[TABLE]

where the expectation is taken with respect to the random draw of projection images as described in Section 2.1.

The proof is in A.

5.3 Convergence of the reconstructed volumes

Consider the graph Laplacian eigenvectors $\{\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(\ell)}\}_{\_}{\ell=0}^{n-1}$ computed from the low-resolution reconstruction coordinates $\boldsymbol{\hat{\beta}}_{\_}1,\ldots,\boldsymbol{\hat{\beta}}_{\_}n\in\mathbf{B}$. Several variants of the discrete graph Laplacian are known to converge to a continuous linear operator on $\mathbf{B}$. This convergence is not only pointwise, but also spectral, meaning that the eigenvectors of the graph Laplacian converge to the eigenfunctions $\phi^{(\ell)}$ of this operator [61, 62, 63]. In particular cases, the limiting operator is the continuous Laplacian, but more generally it is a weighted Laplacian operator, or Fokker–Planck operator, which has an additional drift term towards, or away from, high-density regions [33, 64, 65]. For our theoretical analysis we only need spectral convergence towards some set of eigenfunctions, not necessarily the Laplacian eigenfunctions. We formulate this requirement in the following assumption. For simplicity, we ignore possible eigenvalue multiplicities.

Assumption 3**.**

The domain $\mathbf{B}$ is compact and the eigenvectors $\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(0)},\boldsymbol{\phi}\kern-6.49994pt\hat{\phantom{\boldsymbol{\phi}}}\phantom{}^{(1)},\ldots\in{\mathbb{R}}^{n}$ of the graph Laplacian, ordered by their eigenvalues, converge in probability to a set of eigenfunctions $\phi^{(0)},\phi^{(1)},\ldots:\mathbf{B}\to{\mathbb{R}}$ of some continuous linear differential operator on $\mathbf{B}$, in the sense that

[TABLE]

Furthermore, $\{\phi^{(\ell)}\}$ form an orthonormal set with respect to the measure $\nu(\mathbf{B})$,

[TABLE]

Remark 4**.**

Under this assumption, the eigenfunctions $\{\phi^{(\ell)}\}$ have upper bounds, which we denote as $U_{\_}\ell$. This is due to the fact that they are continuous functions on a compact domain. It follows that,

[TABLE]

Thus, Assumption 2 follows from Assumption 3.

Remark 5**.**

The $\sqrt{n}$ term in (31) is necessary for the eigenvector normalization, since

[TABLE]

By Propositions 2.1 and 2.2 of [24], the low-resolution mean and covariance estimates are consistent. However, unlike these aggregate quantities, the PCA coordinates $\boldsymbol{\hat{\beta}}_{\_}s$ are computed from a single image, so they must contain an irreducible error term due to the finite noise level. We codify this in the following assumption.

Assumption 4**.**

The estimated PCA coordinates are correct up to some stochastically bounded noise term,

[TABLE]

We now show that, up to noise, the spectral volumes are merely voxel-wise orthogonal projections of the true volumes ${\bf x}_{\_}1,\ldots,{\bf x}_{\_}n$ onto a basis of eigenfunctions.

Corollary 1**.**

Under Assumptions 1, 3 and 4 it follows from Theorem 1 that

[TABLE]

where the expectation is with respect to the distribution of ${\bf x}\in\mathcal{M}$.

So far we have treated the convergence of the spectral volumes. We now turn to the convergence of the high-resolution reconstructions. As discussed in Section 3.1, we assume that the manifold of the molecular volumes can be well approximated by a small number of eigenfunctions. We define this notion precisely in the following assumption.

Assumption 5**.**

There is a set of spectral volumes ${\boldsymbol{{\alpha}}}^{(0)},\ldots,{\boldsymbol{{\alpha}}}^{(r-1)}$ and a non-negative function $h(r)$ that satisfies $h(r)\to 0$ such that $h(r)$ bounds the approximation of $\mathcal{M}$ by $r$ spectral volumes:

[TABLE]

In that case, we can prove that the true volumes are recovered up to noise.

Theorem 2**.**

Consider a sample from a manifold that conforms to (37), then it follows from Assumptions 1, 3, 4 and 5 that as $n\to\infty$ we have

[TABLE]

where $C_{\_}\ell$ is an upper bound on the norm of the gradient of $\phi^{(\ell)}$.

The proof of this theorem is in A. Note that the first error term contains an irreducible error from the finite level of noise in the PCA coordinate assignment.

6 Results

In this section, we apply our method to several synthetic datasets with a low-dimensional conformation space. We first consider clean images of a clock face with a rotating hand, then more realistic datasets of molecular volumes with one- and two-dimensional motions.

6.1 Clock dataset

We begin with a toy model of a 2D clock face with a single moving hand. Since the objects we wish to reconstruct are images rather than volumes, no projections are involved. This is the setting studied in Section 5.1. The simulated dataset is comprised of $n=10^{4}$ noisy images ${\bf z}_{\_}1,\ldots,{\bf z}_{\_}n\in{\mathbb{R}}^{N\times N}$ with $N=128$, where each image shows the clock hand at a random angle with additive Gaussian noise. The affinity matrix was chosen to be $W_{\_}{ij}=\mathrm{e}^{-\|{\bf z}_{\_}i-{\bf z}_{\_}j\|^{2}/N^{2}\sigma^{2}}$ where $\sigma$ is the standard deviation of the Gaussian noise. We then constructed a normalized graph Laplacian, extracted its eigenvectors, and computed $r=15$ spectral volumes by solving the least-squares problem of (25). Figure 3 shows representative input images and their corresponding reconstructions. Figure 4 shows the estimated spectral volumes (spectral images in this case). We need a large value of $r$ to get good reconstructions, since the clock hand has sharp discontinuities.

To analyze this example we note that the clock dataset has the manifold geometry of the unit circle $S^{1}$. Ignoring an arbitrary phase offset, the set of real eigenfunctions of the Laplace–Beltrami operator on $S^{1}$ are $\phi^{(0)}(\theta)=1/2\pi$ and for all integer $\ell\geq 1$, $\phi^{(2\ell-1)}(\theta)=\sqrt{1/\pi}\sin(\ell\theta)$ and $\phi^{(2\ell)}(\theta)=\sqrt{1/\pi}\cos(\ell\theta)$. It follows from (28) that

[TABLE]

Let ${\bf z}_{\_}\theta$ denote the image with the clock hand at angle $\theta$, we may rewrite the above as

[TABLE]

Rather than fixing a pixel and rotating the clock hand, we may fix the clock hand and rotate the pixel in the other direction. For pixels $[x,y]$ inside the disk of the clock face,

[TABLE]

We conclude that in the case of simple rotation heterogeneity, a pixel of the $\ell$th spectral volume in the rotating domain converges to the $\ell$th Fourier coefficient of the function $f(\theta)={\bf z}_{\_}0[R_{\_}{-\theta}[x,y]]$. Put differently, the coefficients ${\hat{\boldsymbol{\alpha}}}^{(0)}[x,y],{\hat{\boldsymbol{\alpha}}}^{(1)}[x,y],\ldots$ converge to the Fourier coefficients of the rotating domain, in polar representation.

We tested a similar clock dataset in 3D, using the same clock hand shape, this time with tomographic projections. The results we obtained are similar to the results of the 2D clock dataset, in accordance with Corollary 1. See B for details.

6.2 Simulated ion channel

We created two synthetic datasets based on a voltage-gated potassium channel, shown in Figure 1. The first dataset ChannelSpin demonstrates a rotational motion of the top part about the $z$ axis. The second dataset ChannelStretch demonstrates a nonrigid stretching of the bottom part. Specifically, each slice was displaced along the $x$-$y$ plane by an amount that is proportional to the squared distance of the slice from the center of the molecule. We used a spatial resolution of $N=108$ and generated $n=10,000$ volumes for each dataset. The angles of rotation in the ChannelSpin dataset were drawn uniformly, which gives a conformational manifold diffeomorphic to the circle. For the ChannelStretch dataset, we drew random displacements $\delta_{\_}x,\delta_{\_}y\in\{-16,-15,\ldots,16\}$ that parameterize a non-rigid shift of every $x$-$y$ slice in the bottom half of the molecule. Let ${\bf v}\in\mathbb{R}^{N\times N\times N}$ be the original (unstretched) ion channel, the stretched ion channel $\bf v^{\prime}$ is defined for every $0\leq z\leq N/2$ by

[TABLE]

Note that $z=N/2$ is the center of mass of the ion channel and $z_{\_}0=16$ is the bottom of the molecule.

In both datasets, we projected the molecules along random orientations, drawn uniformly from ${\mathrm{SO}}(3)$. We then applied a simulated point spread function with a defocus value chosen uniformly at random from $1.50$, $1.67$, $1.83$, $2.00$, $2.17$, $2.33$, or $2.50$ microns. Finally, we added white Gaussian noise, with a variance chosen such that the total energy of the noise was 30 times that of the total energy of each clean image. No in-plane shift was applied. See Figure 1 for example images.

Using the projection images, we ran the covariance estimation method with $q=4$ components to build the adjacency matrix for the ChannelSpin dataset and $q=8$ for the ChannelStretch dataset. We then reconstructed the volumes using $r=1,\ldots,15$ spectral volumes. We used the true orientations of the projection images for both the covariance and spectral volume estimation procedures.

Examining the Laplacian eigenmaps embedding

Figure 5 shows the embeddings of a random sample from the ChannelSpin and ChannelStretch datasets. The embedding of ChannelSpin clearly shows a circle whereas the embedding of the ChannelStretch dataset shows a 2-dimensional square in the $\hat{\phi}_{s}^{(1)}-\hat{\phi}_{s}^{(2)}$ plane that is shaped like a saddle. Both of these results are in accordance with the underlying motion manifold.

Examining the spectral volumes

Figure 6 shows the first few spectral volumes. For the ChannelSpin dataset, as expected ${\hat{\boldsymbol{\alpha}}}^{(0)}$ captures the mean over all rotations. Higher order spectral volumes have increasing angular frequency, capturing more and more detail. Note that ${\hat{\boldsymbol{\alpha}}}^{(1)}$ and ${\hat{\boldsymbol{\alpha}}}^{(2)}$ are a quarter period out of phase, while ${\hat{\boldsymbol{\alpha}}}^{(3)}$ has twice the angular frequency. Due to the $C_{\_}4$ symmetry of the ion channel, the lowest frequency of the ChannelSpin dataset has a period of 90 degrees.

For the ChannelStretch dataset, we see that ${\hat{\boldsymbol{\alpha}}}^{(0)}$ captures the fixed part of the molecule with high resolution and shows a “smeared” bottom portion. The first and second nontrivial spectral volumes each have a low spatial frequency along the $x$ and $y$ axes. Higher spectral volumes show higher spatial frequencies. ${\hat{\boldsymbol{\alpha}}}^{(3)}$ shows a mix of the directions of ${\hat{\boldsymbol{\alpha}}}^{(1)}$ and ${\hat{\boldsymbol{\alpha}}}^{(2)}$ whereas ${\hat{\boldsymbol{\alpha}}}^{(4)}$ is similar to ${\hat{\boldsymbol{\alpha}}}^{(1)}$ but with a double spatial frequency. Recall that by Corollary 1 up to noise the spectral volumes are $\mathbb{E}[\phi^{(\ell)}(\boldsymbol{\beta}({\bf x})){\bf x}]$. In this case the eigenfunctions of the Laplacian on the square are the 2D discrete cosine transform basis functions, which are up to scale $\phi^{(n_{\_}x,n_{\_}y)}=\cos(n_{\_}xx)\cos(n_{\_}yy)$ with eigenvalue $\propto n_{\_}x^{2}+n_{\_}y^{2}$, See [45, Section 3.1]. This agrees with our empirical observations.

Reconstruction accuracy

Figure 7 shows reconstructions with increasing numbers of spectral volumes alongside the original simulated volume. Note that the reconstructions for the ChannelSpin volume are of higher quality than for the ChannelStretch volume. This is expected, since the manifold of conformations of ChannelSpin is one-dimensional whereas ChannelStretch has two-dimensional motion. Hence more samples are needed to get a dense cover of the conformational manifold.

To quantify the accuracy of our reconstructions, we use the Fourier shell correlation (FSC), which is the standard evaluation criterion in the cryo-EM literature [66]. Given two volumes ${\bf x}_{\_}1,{\bf x}_{\_}2$, the FSC takes their Fourier transforms and computes the correlation between each frequency shell. Because we wish to estimate the quality of reconstructing the variable part of the molecule, instead of reporting the FSC between the reconstructed volumes and the original volumes, we report the mean-subtracted FSC

[TABLE]

Figure 8 shows the results for each simulated dataset. As expected, the reconstruction quality increases with the number of spectral volumes. The reconstruction of the high frequencies is less accurate than that of the low frequencies. In both cases, as $r$ increases, the FSC curves converge, suggesting a number after which more spectral volumes yield diminishing returns.

6.3 Runtime

Table 2 details the running time of our method for the ChannelSpin simulation with $N=108$, $r=15$, and $q=8$. The method is implemented in MATLAB 2017b and runs on 16 cores of a 2.3 GHz Intel Xeon CPU; memory usage was about 60 GB.

7 Conclusion

Today, rigid macromolecules are routinely reconstructed to near-atomic resolution using standard cryo-EM software tools. However, the high-resolution reconstruction of molecular samples with continuous heterogeneity remains one of the grand challenges of the field. This work describes a new method which addresses this challenge. It combines spectral graph theory with recent techniques for covariance-based low-resolution reconstruction. Our procedure computes conformation-dependent Laplacian eigenmap coordinates and then generates a set of spectral volumes that characterize the variability of the molecule under study. Together these define a high-resolution 3D reconstruction for every projection image.

In the context of machine learning, our method combines and extends two classical methods: (i) the low-resolution covariance-based reconstruction, which we use to form the affinity graph, may be viewed as a generalization of PCA, as it finds the principal volumes in the space of molecular conformations. Unlike PCA, the input is projection images rather than full observations. (ii) the construction of an affinity graph from the low-resolution reconstructions and the generalized tomographic reconstruction using the eigenvectors of the graph Laplacian. This can be viewed as an extension of standard approaches for nonparametric regression, semi-supervised learning and matrix completion on graphs and manifolds (see for example [67, 68, 63, 69, 70]). The key difference is that rather than partially-labeled or noisy observations, we reconstruct a smooth high-dimensional function from noisy tomographic measurement. We note that the combination of PCA and graph Laplacian representations has been used for dimensionality reduction and denoising, for example in [71, 72].

Similar to the hyper-molecules method proposed by [73, 74], our method expands the molecular volumes which generated the projection images using a small set of basis volumes. However, in the hyper-molecules model the basis volumes are obtained from a user-specified manifold. Similarly, the multi-body refinement of RELION 3 [36] requires that the user manually segment the molecule into components that exhibit motion relative to each other. In contrast, our method is data-driven and requires no such user input. It relies only on the assumption that the molecule deforms continuously in a manner that is determined by a small number of parameters.

To conclude, in this paper we have described a method for the reconstruction of molecules with continuous heterogeneity, studied it theoretically and demonstrated the high-resolution reconstruction of synthetic data with one-dimensional and two-dimensional motion manifolds. In future work, we will continue to scale up the method and apply it to the analysis of experimental datasets.

Software

Code for computing spectral volume reconstructions and producing the figures in this paper is available at http://github.com/PrincetonUniversity/specvols

We would like to thank Ronald R. Coifman and Mark Tygert for interesting discussions. AS, AM, and AH were partially supported by NIGMS Award Number R01GM090200, AFOSR FA9550-17-1-0291, ARO W911NF-17-1-0512, Simons Investigator Award, the Moore Foundation Data-Driven Discovery Investigator Award, and NSF BIGDATA Award IIS-1837992. The Flatiron Institute is a division of the Simons Foundation. 3D molecular graphics were rendered using UCSF Chimera [75].