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

.. only:: html

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

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

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

.. _sphx_glr_auto_examples_demo_fista.py:


Demo FISTA.
===================================================

.. GENERATED FROM PYTHON SOURCE LINES 8-10

Select Working Directory and Device
-----------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 10-38

.. code-block:: Python

    import os

    from torch.utils import data

    os.chdir(os.path.dirname(os.getcwd()))
    print("Current Working Directory ", os.getcwd())

    import sys

    sys.path.append(os.path.join(os.getcwd()))

    # General imports
    import matplotlib.pyplot as plt
    import torch
    import os

    # Set random seed for reproducibility
    torch.manual_seed(0)

    manual_device = "cpu"
    # Check GPU support
    print("GPU support: ", torch.cuda.is_available())

    if manual_device:
        device = manual_device
    else:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")





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

 .. code-block:: none

    Current Working Directory  /home/runner/work/pycolibri/pycolibri
    GPU support:  False




.. GENERATED FROM PYTHON SOURCE LINES 39-41

Load dataset
-----------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 41-53

.. code-block:: Python

    from colibri.data.datasets import CustomDataset

    name = "cifar10"
    path = "."
    batch_size = 1


    dataset = CustomDataset(name, path)


    acquisition_name = 'spc'  # ['spc', 'cassi']





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

 .. code-block:: none


      0%|          | 0.00/170M [00:00<?, ?B/s]
      1%|          | 1.15M/170M [00:00<00:14, 11.4MB/s]
      2%|▏         | 3.93M/170M [00:00<00:08, 20.8MB/s]
      4%|▍         | 7.67M/170M [00:00<00:05, 28.2MB/s]
      6%|▌         | 10.6M/170M [00:00<00:05, 28.6MB/s]
      8%|▊         | 13.5M/170M [00:00<00:05, 28.0MB/s]
     10%|▉         | 16.3M/170M [00:00<00:05, 27.9MB/s]
     11%|█         | 19.1M/170M [00:00<00:05, 28.0MB/s]
     13%|█▎        | 22.0M/170M [00:00<00:05, 27.7MB/s]
     15%|█▍        | 24.8M/170M [00:00<00:05, 27.8MB/s]
     16%|█▌        | 27.6M/170M [00:01<00:05, 27.7MB/s]
     18%|█▊        | 30.3M/170M [00:01<00:05, 27.3MB/s]
     19%|█▉        | 33.2M/170M [00:01<00:05, 27.5MB/s]
     22%|██▏       | 37.6M/170M [00:01<00:04, 32.3MB/s]
     26%|██▋       | 44.9M/170M [00:01<00:02, 44.5MB/s]
     31%|███▏      | 53.5M/170M [00:01<00:02, 57.0MB/s]
     36%|███▌      | 61.7M/170M [00:01<00:01, 64.3MB/s]
     41%|████      | 70.3M/170M [00:01<00:01, 70.7MB/s]
     47%|████▋     | 79.5M/170M [00:01<00:01, 77.0MB/s]
     52%|█████▏    | 88.4M/170M [00:01<00:01, 80.8MB/s]
     57%|█████▋    | 97.6M/170M [00:02<00:00, 83.9MB/s]
     63%|██████▎   | 107M/170M [00:02<00:00, 85.7MB/s] 
     68%|██████▊   | 116M/170M [00:02<00:00, 87.4MB/s]
     73%|███████▎  | 125M/170M [00:02<00:00, 85.3MB/s]
     78%|███████▊  | 133M/170M [00:02<00:00, 86.4MB/s]
     84%|████████▎ | 143M/170M [00:02<00:00, 87.7MB/s]
     89%|████████▉ | 151M/170M [00:02<00:00, 87.8MB/s]
     94%|█████████▍| 161M/170M [00:02<00:00, 89.0MB/s]
     99%|█████████▉| 170M/170M [00:02<00:00, 89.4MB/s]
    100%|██████████| 170M/170M [00:02<00:00, 59.9MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 54-56

Visualize dataset
-----------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 56-61

.. code-block:: Python

    from torchvision.utils import make_grid

    sample = dataset[0]["input"]
    sample = sample.unsqueeze(0).to(device)








.. GENERATED FROM PYTHON SOURCE LINES 62-63

Optics forward model

.. GENERATED FROM PYTHON SOURCE LINES 63-145

.. code-block:: Python


    import math
    from colibri.optics import SPC, SD_CASSI, DD_CASSI, C_CASSI


    img_size = sample.shape[1:]

    acquisition_config = dict(
        input_shape=img_size,
    )

    if acquisition_name == "spc":
        n_measurements = 25**2
        n_measurements_sqrt = int(math.sqrt(n_measurements))
        acquisition_config["n_measurements"] = n_measurements

    acquisition_model = {"spc": SPC, "sd_cassi": SD_CASSI, "dd_cassi": DD_CASSI, "c_cassi": C_CASSI}[
        acquisition_name
    ]

    acquisition_model = acquisition_model(**acquisition_config)

    y = acquisition_model(sample)

    # Reconstruct image
    from colibri.recovery.fista import Fista
    from colibri.recovery.terms.prior import Sparsity
    from colibri.recovery.terms.fidelity import L2
    from colibri.recovery.terms.transforms import DCT2D


    algo_params = {
        "max_iters": 2000,
        "alpha": 1e-4,
        "_lambda": 0.01,
    }

    fidelity = L2()
    prior = Sparsity(basis="dct")

    fista = Fista(acquisition_model, fidelity, prior, **algo_params)

    x0 = acquisition_model.forward(y, type_calculation="backward")
    x_hat = fista(y, x0=x0)

    basis = DCT2D()

    theta = basis.forward(x_hat).detach()

    normalize = lambda x: (x - torch.min(x)) / (torch.max(x) - torch.min(x))

    plt.figure(figsize=(10, 10))

    plt.subplot(1, 4, 1)
    plt.title("Reference")
    plt.imshow(sample[0, :, :].permute(1, 2, 0), cmap="gray")
    plt.xticks([])
    plt.yticks([])

    plt.subplot(1, 4, 2)
    plt.title("Sparse Representation")
    plt.imshow(abs(normalize(theta[0, :, :])).permute(1, 2, 0), cmap="gray")
    plt.xticks([])
    plt.yticks([])

    if acquisition_name == "spc":
        y = y.reshape(y.shape[0], -1, n_measurements_sqrt, n_measurements_sqrt)

    plt.subplot(1, 4, 3)
    plt.title("Measurement")
    plt.imshow(normalize(y[0, :, :]).permute(1, 2, 0), cmap="gray")
    plt.xticks([])
    plt.yticks([])

    plt.subplot(1, 4, 4)
    plt.title("Reconstruction")
    plt.imshow(normalize(x_hat[0, :, :]).permute(1, 2, 0).detach().cpu().numpy(), cmap="gray")
    plt.xticks([])
    plt.yticks([])

    plt.show()




.. image-sg:: /auto_examples/images/sphx_glr_demo_fista_001.png
   :alt: Reference, Sparse Representation, Measurement, Reconstruction
   :srcset: /auto_examples/images/sphx_glr_demo_fista_001.png
   :class: sphx-glr-single-img






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

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


.. _sphx_glr_download_auto_examples_demo_fista.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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