{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class='alert alert-warning'>\n\n# Need GPU warning\n\nRunning this example requires a GPU, and hence is NOT possible on binder currently We request you to kindly run this notebook on Google Colab by clicking the link below. Additionally, please make sure to set the runtime on Colab to use a GPU and install the below libraries before running.\n</div>\n<div class=\"colab-button\">\n            <a href=\"https://colab.research.google.com/https://github.com/paquiteau/snake-fmri/blob/gh-pages/examples/auto_examples/anatomical/example_gpu_anat_spirals_slice_single.ipynb\" target=\"_blank\">\n                <img src=\"https://colab.research.google.com/assets/colab-badge.svg\"\n                alt=\"Open In Colab\"/>\n            </a>\n        </div>\n        "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Compare Fourier Model and T2* Model for 2D Stack of Spirals trajectory\n\nThis examples walks through the elementary components of SNAKE.\n\nHere we proceed step by step and use the Python interface. A more integrated\nalternative is to use the CLI ``snake-main``\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Install libraries\n\n!pip installsnake-fmri mri-nufft[finufft,cufinufft]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Imports\nimport matplotlib.pyplot as plt\n\nfrom snake.core.simulation import SimConfig, default_hardware, GreConfig\nfrom snake.core.phantom import Phantom\nfrom snake.core.smaps import get_smaps\nfrom snake.core.sampling import StackOfSpiralSampler\n\nfrom mrinufft import get_operator\n\n\n# For faster computation, try to use the GPU\n\nNUFFT_BACKEND = \"stacked-gpunufft\"\nCOMPUTE_BACKEND = \"cupy\"\n\ntry:\n    import cupy as cp\n\n    if not cp.cupy.cuda.runtime.getDeviceCount():\n        raise ValueError(\"No CUDA Device found\")\n\n    get_operator(\"gpunufft\")\nexcept Exception:\n    try:\n        get_operator(\"finufft\")\n    except ValueError as e:\n        raise ValueError(\"No NUFFT backend available\") from e\n\n    NUFFT_BACKEND = \"finufft\"\n    COMPUTE_BACKEND = \"numpy\"\n\nprint(\n    f\"Using NUFFT backend: {NUFFT_BACKEND}\", f\"Using Compute backend: {COMPUTE_BACKEND}\"\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sim_conf = SimConfig(\n    max_sim_time=3,\n    seq=GreConfig(TR=50, TE=22, FA=12),\n    hardware=default_hardware,\n)\nsim_conf.hardware.n_coils = 1  # Update to get multi coil results.\nsim_conf.hardware.field_strength = 7\nphantom = Phantom.from_brainweb(sub_id=4, sim_conf=sim_conf, tissue_file=\"tissue_7T\")\n\n\n# Create a FOV of 192x192x3mm with a resolution of 3x3x3 mm\nsim_conf.fov.size = (192, 192, 3)\nsim_conf.fov.res_mm = (3, 3, 3)\nsim_conf.fov.offset = (-90, -110, 0)  # Set the z-offset to 0 in MNIspace"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setting up Acquisition Pattern and Initializing Result file.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# The next piece of simulation is the acquisition trajectory.\n# Here nothing fancy, we are using a stack of spiral, that samples a 3D\n# k-space, with an acceleration factor AF=4 on the z-axis.\n\nsampler = StackOfSpiralSampler(\n    accelz=1,\n    acsz=0.1,\n    orderz=\"top-down\",\n    nb_revolutions=12,\n    obs_time_ms=30,\n    constant=True,\n)\n\nsmaps = None\nif sim_conf.hardware.n_coils > 1:\n    smaps = get_smaps(sim_conf.shape, n_coils=sim_conf.hardware.n_coils)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The acquisition trajectory looks like this\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "traj = sampler.get_next_frame(sim_conf)\nfrom mrinufft.trajectories.display import display_3D_trajectory\n\ndisplay_3D_trajectory(traj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Adding noise in Image\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from snake.core.handlers.noise import NoiseHandler\n\nnoise_handler = NoiseHandler(variance=0.01)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Acquisition with Cartesian Engine\n\nThe generated file ``example_EPI.mrd`` does not contains any k-space data for\nnow, only the sampling trajectory. let's put some in. In order to do so, we\nneed to setup the **acquisition engine** that models the MR physics, and get\nsampled at the specified k-space trajectory.\n\nSNAKE comes with two models for the MR Physics:\n\n- model=\"simple\" :: Each k-space shot acquires a constant signal, which is the\n  image contrast at TE.\n- model=\"T2s\" :: Each k-space shot is degraded by the T2* decay induced by\n  each tissue.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Here we will use the \"simple\" model, which is faster.\n#\n# SNAKE's Engine are capable of simulating the data in parallel, by distributing\n# the shots to be acquired to a set of processes. To do so , we need to specify\n# the number of jobs that will run in parallel, as well as the size of a job.\n# Setting the job size and the number of jobs can have a great impact on total\n# runtime and memory consumption.\n#\n# Here, we have a single frame to acquire with 60 frames (one EPI per slice), so\n# a single worker will do.\n\nfrom snake.core.engine import NufftAcquisitionEngine\n\n# engine = NufftAcquisitionEngine(model=\"simple\", snr=30000, slice_2d=True)\n\n# engine(\n#     \"example_spiral_2D.mrd\",\n#     sampler,\n#     phantom,\n#     sim_conf,\n#     handlers=[noise_handler],\n#     smaps=smaps,\n#     worker_chunk_size=60,\n#     n_workers=1,\n#     nufft_backend=NUFFT_BACKEND,\n# )\nengine_t2s = NufftAcquisitionEngine(model=\"T2s\", snr=30000, slice_2d=True)\n\nengine_t2s(\n    \"example_spiral_t2s_2D.mrd\",\n    sampler,\n    phantom,\n    sim_conf,\n    handlers=[noise_handler],\n    worker_chunk_size=60,\n    n_workers=1,\n    nufft_backend=NUFFT_BACKEND,\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from snake.mrd_utils import NonCartesianFrameDataLoader\n\nwith NonCartesianFrameDataLoader(\"example_spiral_t2s_2D.mrd\") as data_loader:\n    traj, kspace_data = data_loader.get_kspace_frame(0, shot_dim=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kspace_data = kspace_data.squeeze(1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kspace_data.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "shot = traj[0].copy()\nnufft = get_operator(NUFFT_BACKEND)(\n    samples=shot[:, :2],\n    shape=data_loader.shape[:-1],\n    density=None,\n    n_batchs=len(kspace_data),\n)\nnufft.samples = shot[:, :2]\nimage = nufft.adj_op(kspace_data)\n\nfig, ax = plt.subplots()\n# axis3dcut(abs(image), None, cuts=(40, 40, 40), ax=ax)\nplt.imshow(\n    abs(image).T,\n)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}