.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_tutorials/tutorials/cov3d_simulation.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_tutorials_tutorials_cov3d_simulation.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_tutorials_tutorials_cov3d_simulation.py:


Covariance 3D Estimation
========================

This script illustrates the example of Covariance 3D estimation using simulation images
generated from Gaussian blob volumes.

.. GENERATED FROM PYTHON SOURCE LINES 8-23

.. code-block:: Python


    import numpy as np
    from scipy.cluster.vq import kmeans2

    from aspire.basis import FBBasis3D
    from aspire.covariance import CovarianceEstimator
    from aspire.denoising import src_wiener_coords
    from aspire.noise import WhiteNoiseEstimator
    from aspire.operators import RadialCTFFilter
    from aspire.reconstruction import MeanEstimator
    from aspire.source.simulation import Simulation
    from aspire.utils import eigs
    from aspire.utils.random import Random
    from aspire.volume import LegacyVolume, Volume








.. GENERATED FROM PYTHON SOURCE LINES 24-26

Create Simulation Object
------------------------

.. GENERATED FROM PYTHON SOURCE LINES 26-63

.. code-block:: Python


    # Specify parameters
    img_size = 8  # image size in square
    num_imgs = 1024  # number of images
    num_eigs = 16  # number of eigen-vectors to keep
    dtype = np.float32

    # Generate a ``Volume`` object for use in the simulation. Here we use a ``LegacyVolume`` and
    # set C = 3 to generate 3 unique random volumes.
    vols = LegacyVolume(
        L=img_size,
        C=3,
        dtype=dtype,
    ).generate()

    # Create a simulation object with specified filters
    sim = Simulation(
        L=img_size,
        n=num_imgs,
        vols=vols,
        unique_filters=[
            RadialCTFFilter(pixel_size=10, defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)
        ],
        dtype=dtype,
    )

    # The Simulation object was created using 3 volumes.
    num_vols = sim.C

    # Specify the normal FB basis method for expending the 2D images
    basis = FBBasis3D(img_size)

    # Estimate the noise variance. This is needed for the covariance estimation step below.
    noise_estimator = WhiteNoiseEstimator(sim, batch_size=500)
    noise_variance = noise_estimator.estimate()
    print(f"Noise Variance = {noise_variance}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/3 [00:00<?, ?it/s]
    100%|██████████| 3/3 [00:00<00:00, 35.12it/s]
    Noise Variance = 0.0023168901493349162




.. GENERATED FROM PYTHON SOURCE LINES 64-71

Estimate Mean Volume and Covariance
-----------------------------------

Estimate the mean. This uses conjugate gradient on the normal equations for
the least-squares estimator of the mean volume. The mean volume is represented internally
using the basis object, but the output is in the form of an
L-by-L-by-L array.

.. GENERATED FROM PYTHON SOURCE LINES 71-81

.. code-block:: Python


    mean_estimator = MeanEstimator(sim, basis=basis)
    mean_est = mean_estimator.estimate()

    # Passing in a mean_kernel argument to the following constructor speeds up some calculations
    covar_estimator = CovarianceEstimator(
        sim, basis=basis, mean_kernel=mean_estimator.kernel
    )
    covar_est = covar_estimator.estimate(mean_est, noise_variance)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/2 [00:00<?, ?it/s]
     50%|█████     | 1/2 [00:02<00:02,  2.15s/it]
    100%|██████████| 2/2 [00:04<00:00,  2.16s/it]
    100%|██████████| 2/2 [00:04<00:00,  2.16s/it]




.. GENERATED FROM PYTHON SOURCE LINES 82-84

Use Top Eigenpairs to Form a Basis
----------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 84-116

.. code-block:: Python


    # Extract the top eigenvectors and eigenvalues of the covariance estimate.
    # Since we know the population covariance is low-rank, we are only interested
    # in the top eigenvectors.

    eigs_est, lambdas_est = eigs(covar_est, num_eigs)

    # Eigs returns column-major, so we transpose and construct a volume.
    eigs_est = Volume(np.transpose(eigs_est, (3, 0, 1, 2)))

    # Truncate the eigendecomposition. Since we know the true rank of the
    # covariance matrix, we enforce it here.

    eigs_est_trunc = eigs_est[: num_vols - 1]
    lambdas_est_trunc = lambdas_est[: num_vols - 1, : num_vols - 1]

    # Estimate the coordinates in the eigenbasis. Given the images, we find the
    # coordinates in the basis that minimize the mean squared error, given the
    # (estimated) covariances of the volumes and the noise process.
    coords_est = src_wiener_coords(
        sim, mean_est, eigs_est_trunc, lambdas_est_trunc, noise_variance
    )

    # Cluster the coordinates using k-means. Again, we know how many volumes
    # we expect, so we can use this parameter here. Typically, one would take
    # the number of clusters to be one plus the number of eigenvectors extracted.

    # Since kmeans2 relies on randomness for initialization, important to push random seed to context manager here.
    with Random(0):
        centers, vol_idx = kmeans2(coords_est.T, num_vols)
        centers = centers.squeeze()








.. GENERATED FROM PYTHON SOURCE LINES 117-119

Performance Evaluation
----------------------

.. GENERATED FROM PYTHON SOURCE LINES 119-143

.. code-block:: Python


    # Evaluate performance of mean estimation.

    mean_perf = sim.eval_mean(mean_est)


    # Evaluate performance of covariance estimation. We also evaluate the truncated
    # eigendecomposition. This is expected to be a closer approximation since it
    # imposes an additional low-rank condition on the estimate.

    covar_perf = sim.eval_covar(covar_est)
    eigs_perf = sim.eval_eigs(eigs_est_trunc, lambdas_est_trunc)

    # Evaluate clustering performance.

    clustering_accuracy = sim.eval_clustering(vol_idx)

    # Assign the cluster centroids to the different images. Since we expect a discrete distribution of volumes
    # (and therefore of coordinates), we assign the centroid coordinate to each image that belongs to that cluster.
    # Evaluate the coordinates estimated

    clustered_coords_est = centers[vol_idx]
    coords_perf = sim.eval_coords(mean_est, eigs_est_trunc, clustered_coords_est)








.. GENERATED FROM PYTHON SOURCE LINES 144-146

Results
-------

.. GENERATED FROM PYTHON SOURCE LINES 146-167

.. code-block:: Python


    # Output estimated covariance spectrum.

    print(f"Population Covariance Spectrum = {np.diag(lambdas_est)}")


    # Output performance results.

    print(f'Mean (rel. error) = {mean_perf["rel_err"]}')
    print(f'Mean (correlation) = {mean_perf["corr"]}')
    print(f'Covariance (rel. error) = {covar_perf["rel_err"]}')
    print(f'Covariance (correlation) = {covar_perf["corr"]}')
    print(f'Eigendecomposition (rel. error) = {eigs_perf["rel_err"]}')
    print(f"Clustering (accuracy) = {clustering_accuracy}")
    print(f'Coordinates (mean rel. error) = {coords_perf["rel_err"]}')
    print(f'Coordinates (mean correlation) = {np.mean(coords_perf["corr"])}')

    # Basic Check
    assert covar_perf["rel_err"] <= 0.80
    assert np.mean(coords_perf["corr"]) >= 0.97
    assert clustering_accuracy >= 0.99




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Population Covariance Spectrum = [11.67564     6.42179     1.759699    1.370867    1.0250142   0.90507495
      0.76535636  0.6299399   0.558662    0.52021676  0.4597896   0.417423
      0.35928726  0.3185292   0.2915414   0.26145974]
    Mean (rel. error) = 0.171009823679924
    Mean (correlation) = 0.9861127138137817
    Covariance (rel. error) = 0.7859532237052917
    Covariance (correlation) = 0.6227817535400391
    Eigendecomposition (rel. error) = 0.7407597303390503
    Clustering (accuracy) = 1.0
    Coordinates (mean rel. error) = [[0.15486291 0.15486291 0.2977213  ... 0.15122503 0.2977213  0.15122503]
     [0.15508513 0.15508513 0.28991634 ... 0.15925645 0.28991634 0.15925645]]
    Coordinates (mean correlation) = 0.9783612607948894





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 37.341 seconds)


.. _sphx_glr_download_auto_tutorials_tutorials_cov3d_simulation.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: cov3d_simulation.ipynb <cov3d_simulation.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: cov3d_simulation.py <cov3d_simulation.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: cov3d_simulation.zip <cov3d_simulation.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_