{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 2D dynamic: temporal TV\n",
    "\n",
    "This example illustrates 2D dynamic MRI image reconstruction\n",
    "from golden angle (GA) radial sampled k-space data\n",
    "collected with multiple coils\n",
    "(parallel MRI),\n",
    "with temporal \"TV\" regularizer (corner-rounded)\n",
    "using the Julia language."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using ImageGeoms: ImageGeom\n",
    "using ImagePhantoms: shepp_logan, SouthPark, phantom, ellipse, spectrum\n",
    "using MIRTjim: jim, prompt\n",
    "using MIRT: Anufft, diffl_map, ncg\n",
    "using MIRT: ir_mri_sensemap_sim, ir_mri_kspace_ga_radial\n",
    "using Plots: gui, plot, scatter, default; default(markerstrokecolor=:auto, label=\"\")\n",
    "using Plots: gif, @animate, Plots\n",
    "using LinearAlgebra: norm, dot, Diagonal\n",
    "using LinearMapsAA: LinearMapAA, block_diag\n",
    "using Random: seed!\n",
    "jim(:abswarn, false); # suppress warnings about display of |complex| images"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following line is helpful when running this jl-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": [
    "## Define dynamic image sequence"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "N = (60,64)\n",
    "fov = 220\n",
    "nt = 8 # frames\n",
    "ig = ImageGeom(; dims = N, deltas=(fov,fov) ./ N)\n",
    "\n",
    "object0 = shepp_logan(SouthPark(); fovs = (fov,fov))\n",
    "objects = Array{typeof(object0)}(undef, nt)\n",
    "xtrue = Array{ComplexF32}(undef, N..., nt)\n",
    "for it in 1:nt\n",
    "    tmp = copy(object0)\n",
    "    width2 = 15 + 5 * sin(2π*it/nt) # mouth open/close\n",
    "    mouth = tmp[2]\n",
    "    tmp[2] = ellipse(mouth.center, (mouth.width[1], width2), mouth.angle[1], mouth.value)\n",
    "    objects[it] = tmp\n",
    "    xtrue[:,:,it] = phantom(axes(ig)..., tmp, 4)\n",
    "end\n",
    "jimxy = (args...; aspect_ratio=1, kwargs...) ->\n",
    "    jim(axes(ig)..., args...; aspect_ratio, kwargs...)\n",
    "jimxy(xtrue, \"True images\"; nrow=2)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Animate true image:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "anim1 = @animate for it in 1:nt\n",
    "    jimxy(xtrue[:,:,it], title=\"Frame $it\")\n",
    "end\n",
    "gif(anim1; fps = 6)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Plot one time course to see temporal change:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ix,iy = 30,14\n",
    "plot(1:nt, abs.(xtrue[ix,iy,:]), label=\"ix=$ix, iy=$iy\",\n",
    "    marker=:o, xlabel=\"frame\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define k-space sampling"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "accelerate = 3\n",
    "nspf = round(Int, maximum(N)/accelerate) # spokes per frame\n",
    "\n",
    "Nro = maximum(N)\n",
    "Nspoke = nspf * nt\n",
    "kspace = ir_mri_kspace_ga_radial(; Nro, Nspoke)\n",
    "fovs = (fov, fov)\n",
    "kspace[:,:,1] ./= fovs[1]\n",
    "kspace[:,:,2] ./= fovs[2]\n",
    "kspace = reshape(kspace, Nro, nspf, nt, 2)\n",
    "(size(kspace), extrema(kspace))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Plot sampling (in units of cycles/pixel):"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ps = Array{Any}(undef, nt)\n",
    "for it in 1:nt\n",
    "    ps[it] = scatter(kspace[:,:,it,1] * fovs[1], kspace[:,:,it,2] * fovs[2],\n",
    "        xtick=(-1:1)*0.5, ytick=(-1:1)*0.5, xlim=(-1,1).*0.5, ylim=(-1,1).*0.5,\n",
    "        aspect_ratio=1, markersize=1, title=\"Frame $it\", widen=true)\n",
    "end\n",
    "plot(ps...; layout=(2,4))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();\n",
    "\n",
    "anim2 = @animate for it in 1:nt\n",
    "    plot(ps[it])\n",
    "end\n",
    "gif(anim2; fps = 6)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Make sensitivity maps, normalized so SSoS = 1:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ncoil = 2\n",
    "smap_raw = ir_mri_sensemap_sim(dims=N, ncoil=ncoil, orbit_start=[90])\n",
    "p1 = jim(smap_raw, \"Sensitivity maps raw\");\n",
    "\n",
    "sum_last = (f, x) -> selectdim(sum(f, x; dims=ndims(x)), ndims(x), 1)\n",
    "ssos_fun = smap -> sqrt.(sum_last(abs2, smap)) # SSoS\n",
    "ssos_raw = ssos_fun(smap_raw)\n",
    "p2 = jim(ssos_raw, \"SSoS for ncoil=$ncoil\");\n",
    "\n",
    "smap = smap_raw ./ ssos_raw\n",
    "p3 = jim(smap, \"Sensitivity maps\");\n",
    "\n",
    "ssos = ssos_fun(smap) # SSoS\n",
    "@assert all(≈(1), ssos)\n",
    "jim(p1, p2, p3)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## System matrix\n",
    "Make system matrix for dynamic non-Cartesian parallel MRI:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Fs = Array{Any}(undef, nt)\n",
    "for it in 1:nt # a NUFFT object for each frame\n",
    "    Ω = [kspace[:,:,it,1][:] kspace[:,:,it,2][:]] * fov * 2pi\n",
    "    Fs[it] = Anufft(Ω, N, n_shift = [N...]/2)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Block diagonal system matrix, with one NUFFT per frame"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "S = [Diagonal(vec(selectdim(smap, ndims(smap), ic))) for ic in 1:ncoil]\n",
    "SO = s -> LinearMapAA(s ; idim=N, odim=N) # LinearMapAO for coil maps\n",
    "AS1 = F -> vcat([F * SO(s) for s in S]...); # [A1*S1; ... ; A1*Sncoil]"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Linear operator input is `(N..., nt)`, output is `(nspf*Nro, Ncoil, nt)`"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "A = block_diag([AS1(F) for F in Fs]...) # todo: refine show()\n",
    "(size(A), A._odim, A._idim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Simulate k-space data via an inverse crime (todo):"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ytrue = A * xtrue\n",
    "snr2sigma = (db, yb) -> # compute noise sigma from SNR (no sqrt(2) needed)\n",
    "    10^(-db/20) * norm(yb) / sqrt(length(yb))\n",
    "sig = Float32(snr2sigma(50, ytrue))\n",
    "seed!(0)\n",
    "y = ytrue + sig * randn(ComplexF32, size(ytrue))\n",
    "ysnr = 20*log10(norm(ytrue) / norm(y - ytrue)) # verify SNR"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Initial image\n",
    "\n",
    "Initial image via zero-fill and scaling:\n",
    "todo: should use density compensation, perhaps via\n",
    "[VoronoiDelaunay.jl](https://github.com/JuliaGeometry/VoronoiDelaunay.jl)"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "x0 = A' * y # zero-filled recon (for each frame)\n",
    "tmp = A * x0 # (Nkspace, Ncoil, Nframe)\n",
    "x0 = (dot(tmp,y) / norm(tmp)^2) * x0 # scale sensibly\n",
    "jimxy(x0, \"initial image\"; nrow=2)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Temporal finite differences:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Dt = diffl_map((N..., nt), length(N)+1 ; T=eltype(A))\n",
    "tmp = Dt' * (Dt * xtrue)\n",
    "jimxy(tmp, \"Time differences are sparse\"; nrow=2)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Regularized CG reconstruction\n",
    "Run nonlinear CG on \"temporal edge-preserving\" regularized LS cost function"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "niter = 90\n",
    "delta = Float32(0.1) # small relative to temporal differences\n",
    "reg = Float32(2^20) # trial and error here\n",
    "ffair = (t,d) -> d^2 * (abs(t)/d - log(1 + abs(t)/d))\n",
    "pot = z -> ffair(z, delta)\n",
    "dpot = z -> z / (Float32(1) + abs(z/delta))\n",
    "cost = x -> 0.5 * norm(A*x - y)^2 + reg * sum(pot.(Dt * x))\n",
    "fun = (x,iter) -> cost(x)\n",
    "gradf = [v -> v - y, u -> reg * dpot.(u)]\n",
    "curvf = [v -> Float32(1), u -> reg]\n",
    "(xh, out) = ncg([A, Dt], gradf, curvf, x0 ; niter, fun)\n",
    "costs = [out[i+1][1] for i=0:niter];"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Show results"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "pr = plot(\n",
    "    jimxy(xtrue, \"|xtrue|\"; nrow=2),\n",
    "    jimxy(xh, \"|recon|\"; nrow=2),\n",
    "    jimxy(xh-xtrue, \"|error|\"; nrow=2),\n",
    "    scatter(0:niter, log.(costs), label=\"cost\", xlabel=\"iteration\"),\n",
    ")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() && prompt();"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Animate true, recon, error"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "anim3 = @animate for it in 1:nt\n",
    "    plot(\n",
    "        jimxy(xtrue[:,:,it], clim=(0,120), title=\"True\"),\n",
    "        jimxy(xh[:,:,it], clim=(0,120), title=\"|Recon|\"),\n",
    "        jimxy(xh[:,:,it] - xtrue[:,:,it], clim=(0,30), title=\"|Error|\"),\n",
    "        plot_title = \"Frame $it\",\n",
    "        layout = (1,3),\n",
    "    )\n",
    "end\n",
    "gif(anim3; fps = 6)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "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
}