{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# NUFFT Overview\n",
    "\n",
    "This example illustrates how to use Nonuniform FFT (NUFFT)\n",
    "for image reconstruction in MRI\n",
    "using the Julia language."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Some MRI scans use non-Cartesian sampling patterns\n",
    "(like radial and spiral k-space trajectories, among others),\n",
    "often in the interest of acquisition speed.\n",
    "\n",
    "Image reconstruction from fully sampled Cartesian k-space data\n",
    "typically uses simple inverse FFT operations,\n",
    "whereas non-Cartesian sampling requires more complicated methods.\n",
    "\n",
    "The examples here illustrate both non-iterative (aka \"gridding\")\n",
    "and iterative methods for non-Cartesian MRI reconstruction.\n",
    "For simplicity the examples consider the case of single-coil data\n",
    "and ignore the effects of B0 field inhomogeneity.\n",
    "\n",
    "\n",
    "First we add the Julia packages that are need for these examples.\n",
    "Change `false` to `true` in the following code block\n",
    "if you are using any of the following packages for the first time."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "if false\n",
    "    import Pkg\n",
    "    Pkg.add([\n",
    "        \"ImagePhantoms\"\n",
    "        \"Unitful\"\n",
    "        \"Plots\"\n",
    "        \"LaTeXStrings\"\n",
    "        \"MIRTjim\"\n",
    "        \"MIRT\"\n",
    "        \"InteractiveUtils\"\n",
    "    ])\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now tell this Julia session to use the following packages.\n",
    "Run `Pkg.add()` in the preceding code block first, if needed."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using ImagePhantoms: shepp_logan, SheppLoganEmis, spectrum, phantom #, Gauss2\n",
    "using Unitful: mm # Allows use of physical units (mm here)\n",
    "using Plots; default(label=\"\", markerstrokecolor=:auto)\n",
    "using LaTeXStrings # for LaTeX in plot labels, e.g., L\"\\alpha_n\"\n",
    "using MIRTjim: jim, prompt # jiffy image display\n",
    "using MIRT: Anufft, diffl_map, ncg\n",
    "using InteractiveUtils: versioninfo"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following line is helpful when running this file as a script;\n",
    "this way it will prompt user to hit a key after each image is displayed."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && jim(:prompt, true);"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Radial k-space sampling\n",
    "\n",
    "We focus on radial sampling as a simple representative non-Cartesian case.\n",
    "Consider imaging a 256mm × 256mm field of FOV\n",
    "with the goal of reconstructing a 128 × 128 pixel image.\n",
    "The following radial and angular k-space sampling is reasonable."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "N = 128\n",
    "FOV = 256mm # physical units!\n",
    "Δx = FOV / N # pixel size\n",
    "kmax = 1 / 2Δx\n",
    "kr = ((-N÷2):(N÷2)) / (N÷2) * kmax # radial sampling in k-space\n",
    "Nϕ = 3N÷2 # theoretically should be about π/2 * N\n",
    "kϕ = (0:Nϕ-1)/Nϕ * π # angular samples\n",
    "νx = kr * cos.(kϕ)' # N × Nϕ k-space sampling in cycles/mm\n",
    "νy = kr * sin.(kϕ)'\n",
    "plot(νx, νy,\n",
    "    xlabel=L\"\\nu_x\", ylabel=L\"\\nu_y\",\n",
    "    aspect_ratio = 1,\n",
    "    title = \"Radial k-space sampling\",\n",
    ")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "For the NUFFT routines considered here,\n",
    "the sampling arguments must be \"Ω\" values:\n",
    "digital frequencies have pseudo-units of radians / pixel."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Ωx = (2π * Δx) * νx # N × Nϕ grid of k-space sample locations\n",
    "Ωy = (2π * Δx) * νy # in pseudo-units of radians / sample\n",
    "\n",
    "scatter(Ωx, Ωy,\n",
    "    xlabel=L\"\\Omega_x\", ylabel=L\"\\Omega_y\",\n",
    "    xticks=((-1:1)*π, [\"-π\", \"0\", \"π\"]),\n",
    "    yticks=((-1:1)*π, [\"-π\", \"0\", \"π\"]),\n",
    "    xlims=(-π,π) .* 1.1,\n",
    "    ylims=(-π,π) .* 1.1,\n",
    "    aspect_ratio = 1, markersize = 0.5,\n",
    "    title = \"Radial k-space sampling\",\n",
    ")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Radial k-space data for Shepp-Logan phantom\n",
    "\n",
    "Get the ellipse parameters for a MRI-suitable version of the Shepp-Logan phantom\n",
    "and calculate (analytically) the radial k-space data.\n",
    "Then display in polar coordinates."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "object = shepp_logan(SheppLoganEmis(); fovs=(FOV,FOV))\n",
    "#object = [Gauss2(18mm, 0mm, 100mm, 70mm, 0, 1)] # useful for validating DCF\n",
    "data = spectrum(object).(νx,νy)\n",
    "data = data / oneunit(eltype(data)) # abandon units at this point\n",
    "jim(kr, kϕ, abs.(data), title=\"k-space data magnitude\",\n",
    "    xlabel=L\"k_r\",\n",
    "    ylabel=L\"k_{\\phi}\",\n",
    "    xticks = (-1:1) .* maximum(abs, kr),\n",
    "    yticks = (0,π),\n",
    "    ylims = (0,π),\n",
    "    aspect_ratio = :none,\n",
    ")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Non-iterative gridding image reconstruction\n",
    "\n",
    "It would be impossible for a radiologist\n",
    "to diagnose a patient\n",
    "from the k-space data in polar coordinates,\n",
    "so image reconstruction is needed.\n",
    "\n",
    "The simplest approach\n",
    "is to (nearest-neighbor) interpolate the k-space data\n",
    "onto a Cartesian grid,\n",
    "and then reconstruction with an inverse FFT.\n",
    "\n",
    "One way to do that interpolation\n",
    "is to use a `Histogram` method\n",
    "in Julia's statistics package."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using StatsBase: fit, Histogram, weights"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following function is a work-around because `weights` in StatsBase is\n",
    "[limited to Real data](https://github.com/JuliaStats/StatsBase.jl/issues/745),\n",
    "so here we bin the real and imaginary k-space data separately,\n",
    "and handle the units."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function histogram(coord, vals::AbstractArray{<:Number}, edges)\n",
    "    u = oneunit(eltype(vals))\n",
    "    wr = weights(real(vec(vals / u)))\n",
    "    wi = weights(imag(vec(vals / u)))\n",
    "    tmp1 = fit(Histogram, coord, wr, edges)\n",
    "    tmp2 = fit(Histogram, coord, wi, edges)\n",
    "    return u * complex.(tmp1.weights, tmp2.weights)\n",
    "end\n",
    "\n",
    "kx = N * Δx * νx # N × Nϕ k-space sampling in cycles/mm\n",
    "ky = N * Δx * νy\n",
    "bin = (-(N÷2):(N÷2)) .- 0.5 # (N+1,) histogram bin edges\n",
    "gridded1 = histogram((vec(kx), vec(ky)), data, (bin,bin))\n",
    "\n",
    "using FFTW: ifft, fftshift\n",
    "tmp = fftshift(ifft(fftshift(gridded1)))\n",
    "x = (-(N÷2):(N÷2-1)) * Δx\n",
    "y = (-(N÷2):(N÷2-1)) * Δx\n",
    "jim(x, y, tmp, title=\"Elementary gridding reconstruction\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "That crummy gridding method does not work well.\n",
    "Here's what the phantom should look like:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ideal = phantom(x, y, object, 2)\n",
    "clim = (0, 9)\n",
    "p0 = jim(x, y, ideal, title=\"True Shepp-Logan phantom\"; clim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## NUFFT approach to gridding\n",
    "\n",
    "Basic nearest-neighbor gridding\n",
    "does not provide acceptable image quality in MRI,\n",
    "so now we turn to using the NUFFT.\n",
    "For MRI purposes,\n",
    "the NUFFT is a function that maps Cartesian spaced image data\n",
    "into non-Cartesian k-space data.\n",
    "The\n",
    "[NFFT.jl](https://github.com/tknopp/NFFT.jl)\n",
    "package has functions\n",
    "for computing the NUFFT and its adjoint.\n",
    "These are linear mappings (generalizations of matrices),\n",
    "so instead of calling those functions directly,\n",
    "here we use the NUFFT linear map object\n",
    "defined in MIRT\n",
    "that provides a non-Cartesian Fourier encoding \"matrix\"."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "A = Anufft([vec(Ωx) vec(Ωy)], (N,N); n_shift = [N/2,N/2])"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This linear map is constructed to map a N × N image into `length(Ωx)` k-space samples.\n",
    "So its\n",
    "[adjoint](https://en.wikipedia.org/wiki/Adjoint)\n",
    "goes the other direction.\n",
    "However, an adjoint is *not* an inverse!"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "gridded2 = A' * vec(data)\n",
    "jim(x, y, gridded2, title=\"NUFFT gridding without DCF\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Density compensation\n",
    "\n",
    "To get a decent image with NUFFT-based gridding of non-Cartesian data,\n",
    "one must compensate for the k-space sampling density.\n",
    "See\n",
    "[this book chapter](https://web.eecs.umich.edu/~fessler/book/c-four.pdf)\n",
    "for details.\n",
    "\n",
    "Because this example uses radial sampling,\n",
    "we can borrow ideas from tomography,\n",
    "especially the\n",
    "[ramp filter](https://en.wikipedia.org/wiki/Radon_transform#Radon_inversion_formula),\n",
    "to define a reasonable density compensation function (DCF).\n",
    "\n",
    "Here is a basic DCF version that uses the ramp filter in a simple way,\n",
    "corresponding to the areas of annular segments\n",
    "(Voronoi cells in polar coordinates).\n",
    "The `dcf .* data` line uses Julia's\n",
    "[broadcast](https://docs.julialang.org/en/v1/manual/arrays/#Broadcasting)\n",
    "feature\n",
    "to apply the 1D DCF to each radial spoke."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "dν = 1/FOV # k-space radial sample spacing\n",
    "dcf = pi / Nϕ * dν * abs.(kr) # (N+1) weights along k-space polar coordinate\n",
    "dcf = dcf / oneunit(eltype(dcf)) # kludge: units not working with LinearMap now\n",
    "gridded3 = A' * vec(dcf .* data)\n",
    "p3 = jim(x, y, gridded3, title=\"NUFFT gridding with simple ramp-filter DCF\"; clim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The image is more reasonable than without any DCF,\n",
    "but we can do better (quantitatively) using the correction of\n",
    "[Lauzon&Rutt, 1996](https://doi.org/10.1002/mrm.1910360617)\n",
    "and\n",
    "[Joseph, 1998](https://doi.org/10.1002/mrm.1910400317)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "dcf = pi / Nϕ * dν * abs.(kr) # see lauzon:96:eop, joseph:98:sei\n",
    "dcf[kr .== 0/mm] .= pi * (dν/2)^2 / Nϕ # area of center disk\n",
    "dcf = dcf / oneunit(eltype(dcf)) # kludge: units not working with LinearMap now\n",
    "gridded4 = A' * vec(dcf .* data)\n",
    "p4 = jim(x, y, gridded4, title=\"NUFFT gridding with better ramp-filter DCF\"; clim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "A profile helps illustrate the improvement."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "pp = plot(x, real(gridded4[:,N÷2]), label=\"Modified ramp DCF\")\n",
    "plot!(x, real(gridded3[:,N÷2]), label=\"Basic ramp DCF\")\n",
    "plot!(x, real(ideal[:,N÷2]), label=\"Ideal\", xlabel=L\"x\", ylabel=\"middle profile\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Finally we have made a NUFFT gridded image with DCF\n",
    "that has the appropriate range of values,\n",
    "but it still looks less than ideal.\n",
    "So next we try an iterative approach.\n",
    "\n",
    "\n",
    "## Iterative MR image reconstruction using NUFFT\n",
    "\n",
    "As an initial iterative approach,\n",
    "let's apply a few iterations of conjugate gradient (CG)\n",
    "to seek the minimizer of the least-squares cost function:\n",
    "$$\n",
    "\\arg \\min_{x} \\frac{1}{2} \\| A x - y \\|_2^2.\n",
    "$$\n",
    "CG is well-suited to minimizing quadratic cost functions,\n",
    "but we do not expect the image to be great quality\n",
    "because radial sampling omits the \"corners\" of k-space\n",
    "so the NUFFT operator $A$ is badly conditioned."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "dx = FOV / N # pixel size\n",
    "dx = dx / oneunit(dx) # abandon units for now\n",
    "gradf = u -> u - vec(data/dx^2) # gradient of f(u) = 1/2 \\| u - y \\|^2\n",
    "curvf = u -> 1 # curvature of f(u)\n",
    "x0 = gridded4 # initial guess: best gridding reconstruction\n",
    "xls, _ = ncg([A], [gradf], [curvf], x0; niter = 20)\n",
    "p5 = jim(x, y, xls, \"LS-CG reconstruction\"; clim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "To improve the results, we include regularization.\n",
    "Here we would like to reconstruct an image\n",
    "by finding the minimizer of a regularized LS cost function\n",
    "such as the following:\n",
    "$$\n",
    "\\arg \\min_{x} \\frac{1}{2} \\| A x - y \\|_2^2 + \\beta R(x)\n",
    ", \\qquad\n",
    "R(x) = 1' \\psi.(T x).\n",
    "$$\n",
    "We focus here on the case of edge-preserving regularization\n",
    "where $T$ is a 2D finite-differencing operator\n",
    "and $\\psi$ is a potential function.\n",
    "This operator maps a N×N image into a N×N×2 array\n",
    "with the horizontal and vertical finite differences."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "T = diffl_map((N,N), [1,2] ; T = ComplexF32)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Applying this operator to the ideal image illustrated its action:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "jim(x, y, T * ideal; ncol=1, title=\"Horizontal and vertical finite differences\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Edge-preserving regularization\n",
    "\n",
    "We use the Fair potential function:\n",
    "a rounded corner version of absolute value,\n",
    "an approximation to anisotropic total variation (TV)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "β = 2^13 # regularization parameter\n",
    "δ = 0.05 # edge-preserving parameter\n",
    "wpot = z -> 1 / (1 + abs(z)/δ) # weighting function"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Nonlinear CG algorithm\n",
    "\n",
    "We apply a (nonlinear) CG algorithm\n",
    "to seek the minimizer of the cost function.\n",
    "Nonlinear CG is well suited to convex problems\n",
    "that are locally quadratic\n",
    "like the regularized cost function considered here.\n",
    "See\n",
    "[this survey paper](https://doi.org/10.1109/MSP.2019.2943645)\n",
    "for an overview of optimization methods for MRI."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "B = [A, T] # see MIRT.ncg\n",
    "gradf = [u -> u - vec(data/dx^2), # data-term gradient, correct for pixel area\n",
    "         u -> β * (u .* wpot.(u))] # regularizer gradient\n",
    "curvf = [u -> 1, u -> β] # curvature of quadratic majorizers\n",
    "x0 = gridded4 # initial guess is best gridding reconstruction\n",
    "xhat, _ = ncg(B, gradf, curvf, x0; niter = 90)\n",
    "p6 = jim(x, y, xhat, \"Iterative reconstruction\"; clim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Here is a comparison of the profiles."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "plot!(pp, x, real(xls[:,N÷2]), label=\"LS-CG\")\n",
    "plot!(pp, x, real(xhat[:,N÷2]), label=\"Iterative edge-preserving\", color=:black)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "In this case, iterative image reconstruction provides the best looking image.\n",
    "One might argue this simulation was doomed to succeed,\n",
    "because the phantom is piece-wise constant,\n",
    "which is the best case for edge-preserving regularization.\n",
    "On the other hand,\n",
    "this was not an\n",
    "[inverse crime](https://doi.org/10.1016/j.cam.2005.09.027)\n",
    "([see also here](https://arxiv.org/abs/2109.08237))\n",
    "because the k-space data came from the analytical spectrum of ellipses,\n",
    "rather than from a discrete image.\n",
    "\n",
    "\n",
    "### Caveats\n",
    "\n",
    "* The phantom used here was real-valued, which is unrealistic\n",
    "  (although the reconstruction methods did not \"know\" it was real).\n",
    "* This simulation is for a single-coil scan\n",
    "  whereas modern MRI scanners generally use multiple receive coils.\n",
    "* There was no statistical noise in this simulation."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.1"
  },
  "kernelspec": {
   "name": "julia-1.9",
   "display_name": "Julia 1.9.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}