{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "29abc7d4",
   "metadata": {},
   "source": [
    "# Differentiable beamforming for ultrasound autofocusing\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28d4bf8",
   "metadata": {},
   "source": [
    "In this tutorial we will implement a basic differential beamformer. We will use a gradient descent method to minimize a pixelwise common midpoint phase error to estimate a speed of sound map. The algorithm is slightly simplified, loss is computed without patching. \n",
    "\n",
    "For more information we would like to refer you to the original research [project page](https://waltersimson.com/dbua/) of the differential beamformer for ultrasound autofocusing (DBUA) [paper](https://doi.org/10.1007/978-3-031-43999-5_41):\n",
    "- Simson, W., Zhuang, L., Sanabria, S.J., Antil, N., Dahl, J.J., Hyun, D. (2023). Differentiable Beamforming for Ultrasound Autofocusing. Medical Image Computing and Computer Assisted Intervention (MICCAI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b61732db",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tue-bmd/zea/blob/main/docs/source/notebooks/pipeline/dbua_example.ipynb)\n",
    " \n",
    "[![View on GitHub](https://img.shields.io/badge/GitHub-View%20Source-blue?logo=github)](https://github.com/tue-bmd/zea/blob/main/docs/source/notebooks/pipeline/dbua_example.ipynb)\n",
    " \n",
    "[![Hugging Face dataset](https://img.shields.io/badge/Hugging%20Face-Dataset-yellow?logo=huggingface)](https://huggingface.co/datasets/zeahub/simulations/blob/main/circular_inclusion_simulation.hdf5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231c7ee0",
   "metadata": {},
   "source": [
    "‼️ **Important:** This notebook is optimized for **GPU/TPU**. Code execution on a **CPU** may be very slow.\n",
    "\n",
    "If you are running in Colab, please enable a hardware accelerator via:\n",
    "\n",
    "**Runtime → Change runtime type → Hardware accelerator → GPU/TPU** 🚀."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "701a9138",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "87a10c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\"\n",
    "os.environ[\"ZEA_DISABLE_CACHE\"] = \"1\"\n",
    "os.environ[\"ZEA_LOG_LEVEL\"] = \"INFO\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f548c87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'jax'\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import keras\n",
    "from keras import ops\n",
    "from keras.utils import Progbar\n",
    "\n",
    "import zea\n",
    "from zea import File, init_device\n",
    "from zea.visualize import set_mpl_style\n",
    "from zea.io_lib import matplotlib_figure_to_numpy, save_to_gif\n",
    "from zea.backend.optimizer import adam\n",
    "from zea.backend.autograd import AutoGrad\n",
    "from zea.ops import (\n",
    "    TOFCorrection,\n",
    "    CommonMidpointPhaseError,\n",
    "    PatchedGrid,\n",
    "    EnvelopeDetect,\n",
    "    Normalize,\n",
    "    DelayAndSum,\n",
    "    LogCompress,\n",
    "    ReshapeGrid,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccdf3474",
   "metadata": {},
   "source": [
    "We will work with the GPU if available, and initialize using `init_device` to pick the best available device. Also, (optionally), we will set the matplotlib style for plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "62bf6ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76834576",
   "metadata": {},
   "source": [
    "## Load dataset \n",
    "First let's load a dataset. In this case a circular inclusion in an isoechoic medium, simulated in [k-Wave](http://www.k-wave.org/) and stored to [zea data format](../../data-acquisition.rst). It will automatically load the dataset from our [Hugging Face](https://huggingface.co/zeahub) repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9b8c0f96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf957f6db96541c88a5f71457fedadfe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "circular_inclusion_simulation.hdf5:   0%|          | 0.00/308M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "file_path = \"hf://zeahub/simulations/circular_inclusion_simulation.hdf5\"\n",
    "file = File(file_path)\n",
    "data_frame = file.load_data(data_type=\"raw_data\")\n",
    "probe = file.probe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60799899",
   "metadata": {},
   "source": [
    "## Setting up the pipelines"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b15c6f4f",
   "metadata": {},
   "source": [
    "We will set up two separate pipelines for processing the data in this notebook. This is a good example of how to use multiple different ultrasound pipelines for different purposes.\n",
    "- `loss_pipeline`: does time-of-flight correction and computes the pixelwise loss based on the common midpoint error.\n",
    "- `image_plot_pipeline`: does time-of-flight correction and delay-and-sum beamforming to produce a B-mode image for visualization.\n",
    "\n",
    "Lets start with defining a field of view and f-number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bbcf9856",
   "metadata": {},
   "outputs": [],
   "source": [
    "xlims = (-15e-3, 15e-3)\n",
    "zlims = (0e-3, 30e-3)\n",
    "f_number = 0.7\n",
    "\n",
    "x_min, x_max = xlims\n",
    "z_min, z_max = zlims"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "348adf0d",
   "metadata": {},
   "source": [
    "### Loss pipeline\n",
    "\n",
    "We begin with selecting a grid of points where we will be calculating the loss. This grid is more an indication of how many points we consider per step as we resample them every iteration within the bounds of the field of view. The main star of the loss is the [zea.ops.CommonMidpointPhaseError](../../_autosummary/zea.ops.rst#zea.ops.CommonMidpointPhaseError), which essentially computes the phase error between signals from two different subapertures that originate from the same point in the field of view. See [this paper](https://arxiv.org/pdf/2410.03008) for more info. We will use this as a pixelwise loss function to optimize the sound speed map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1d788ceb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No initial times provided, using zeros\n",
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No azimuth angles provided, using zeros\n",
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No transmit origins provided, using zeros\n"
     ]
    }
   ],
   "source": [
    "grid_size_x = grid_size_z = 25\n",
    "\n",
    "scan = file.scan(\n",
    "    xlims=xlims,\n",
    "    zlims=zlims,\n",
    "    grid_size_x=grid_size_x,\n",
    "    grid_size_z=grid_size_z,\n",
    "    f_number=f_number,\n",
    ")\n",
    "\n",
    "loss_pipeline = zea.Pipeline(\n",
    "    [\n",
    "        PatchedGrid(\n",
    "            [\n",
    "                TOFCorrection(),\n",
    "                CommonMidpointPhaseError(),\n",
    "            ],\n",
    "            num_patches=grid_size_z * grid_size_x,\n",
    "        ),\n",
    "        ReshapeGrid(),\n",
    "    ],\n",
    "    jit_options=\"pipeline\",\n",
    ")\n",
    "parameters = loss_pipeline.prepare_parameters(probe, scan)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "433727dd",
   "metadata": {},
   "source": [
    "### Image pipeline\n",
    "We can now construct a pipeline for the B-mode image for visualization of the B-mode while optimizing the speed of sound map. Note that we now use a denser grid for the beamforming grid to produce an image without aliasing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "805bf879",
   "metadata": {},
   "outputs": [],
   "source": [
    "width, height = xlims[1] - xlims[0], zlims[1] - zlims[0]\n",
    "wavelength = 1540 / probe.center_frequency\n",
    "\n",
    "grid_size_x = int(width / (0.5 * wavelength) / 4) + 1\n",
    "grid_size_z = int(height / (0.5 * wavelength) / 4) + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "61c16b91",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No initial times provided, using zeros\n",
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No azimuth angles provided, using zeros\n",
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No transmit origins provided, using zeros\n"
     ]
    }
   ],
   "source": [
    "scan_plot = file.scan(\n",
    "    xlims=xlims,\n",
    "    zlims=zlims,\n",
    "    grid_size_x=grid_size_x,\n",
    "    grid_size_z=grid_size_z,\n",
    "    f_number=f_number,\n",
    ")\n",
    "\n",
    "image_plot_pipeline = zea.Pipeline(\n",
    "    [\n",
    "        PatchedGrid(\n",
    "            [TOFCorrection(), DelayAndSum()],\n",
    "            num_patches=grid_size_x * grid_size_z,\n",
    "        ),\n",
    "        ReshapeGrid(),\n",
    "        EnvelopeDetect(),\n",
    "        Normalize(),\n",
    "        LogCompress(),\n",
    "    ],\n",
    "    jit_options=\"pipeline\",\n",
    ")\n",
    "\n",
    "parameters_plot = image_plot_pipeline.prepare_parameters(probe, scan_plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fe2bf0b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparison of beamforming grid sizes:\n",
      "Loss pipeline grid size (sos): 25 x 25\n",
      "Image pipeline grid size (B-mode): 71 x 71\n"
     ]
    }
   ],
   "source": [
    "print(\"Comparison of beamforming grid sizes:\")\n",
    "print(f\"Loss pipeline grid size (sos): {scan['grid_size_x']} x {scan['grid_size_z']}\")\n",
    "print(f\"Image pipeline grid size (B-mode): {scan_plot['grid_size_x']} x {scan_plot['grid_size_z']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c623ec2f",
   "metadata": {},
   "source": [
    "## Set up speed of sound grid\n",
    "Here we define the grid of sound speed voxels that will be optimized.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7bd95a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "sos_grid_x = ops.linspace(x_min, x_max, 40)\n",
    "sos_grid_z = ops.linspace(z_min, z_max, 40)\n",
    "initial_sound_speed = 1460\n",
    "sos_map = initial_sound_speed * ops.ones((sos_grid_z.shape[0], sos_grid_x.shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f04728",
   "metadata": {},
   "source": [
    "## Optimizer\n",
    "Here we define the optimization schedule parameters.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "378ca51a",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "num_iterations = 200\n",
    "step_size = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e44042d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_fn, update_fn, get_params_fn = adam(step_size)\n",
    "opt_state = init_fn(sos_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb81416",
   "metadata": {},
   "source": [
    "## Loss function and optimization loop\n",
    "Here we combine the pixelwise loss with some regularizers along the lateral and axial dimensions to aid the optimization. Furthermore, we introduce some helper functions for the optimization loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7f662c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(\n",
    "    sos_map,\n",
    "    sos_grid_x,\n",
    "    sos_grid_z,\n",
    "    loss_pipeline,\n",
    "    parameters,\n",
    "    data_frame,\n",
    "    flatgrid,\n",
    "):\n",
    "    dx_sos = sos_grid_x[1] - sos_grid_x[0]\n",
    "    dz_sos = sos_grid_z[1] - sos_grid_z[0]\n",
    "    parameters[\"flatgrid\"] = flatgrid\n",
    "    out = loss_pipeline(\n",
    "        data=data_frame,\n",
    "        sos_map=sos_map,\n",
    "        sos_grid_x=sos_grid_x,\n",
    "        sos_grid_z=sos_grid_z,\n",
    "        **parameters,\n",
    "    )\n",
    "    metric = out[\"data\"]\n",
    "    metric_safe = ops.nan_to_num(metric, nan=0.0)\n",
    "    metric_loss = ops.mean(metric_safe)\n",
    "    tvz = ops.mean(ops.square(ops.diff(sos_map, axis=0)))\n",
    "    tvx = ops.mean(ops.square(ops.diff(sos_map, axis=1)))\n",
    "    variation_loss = (tvx + tvz) * 1e2 * dx_sos * dz_sos\n",
    "    total_loss = metric_loss + variation_loss\n",
    "    return total_loss\n",
    "\n",
    "\n",
    "loss_fn_caller = AutoGrad()\n",
    "loss_fn_caller.set_function(loss_fn)\n",
    "\n",
    "\n",
    "def compute_gradients(sos_map, data_frame, flatgrid):\n",
    "    kwargs = dict(\n",
    "        sos_grid_x=sos_grid_x,\n",
    "        sos_grid_z=sos_grid_z,\n",
    "        loss_pipeline=loss_pipeline,\n",
    "        parameters=parameters,\n",
    "        data_frame=data_frame,\n",
    "        flatgrid=flatgrid,\n",
    "    )\n",
    "    grad, loss = loss_fn_caller.gradient_and_value(sos_map, **kwargs)\n",
    "\n",
    "    return grad, loss\n",
    "\n",
    "\n",
    "def apply_gradients(opt_state, grad):\n",
    "    new_sos_grid, m, v, i = update_fn(grad, opt_state)\n",
    "    new_opt_state = (new_sos_grid, m, v, i)\n",
    "    return new_sos_grid, new_opt_state\n",
    "\n",
    "\n",
    "def resample_grid(parameters, xlims, zlims):\n",
    "    seed_generator = keras.random.SeedGenerator(int(time.time() * 1e6) % (2**32 - 1))\n",
    "    n_pix = parameters[\"flatgrid\"].shape[0]\n",
    "    x = keras.random.uniform(\n",
    "        shape=(n_pix,), minval=xlims[0] + 5e-3, maxval=xlims[1], seed=seed_generator\n",
    "    )\n",
    "    y = ops.zeros_like(x)\n",
    "    z = keras.random.uniform(shape=(n_pix,), minval=zlims[0], maxval=zlims[1], seed=seed_generator)\n",
    "\n",
    "    coords = ops.stack([x, y, z], axis=-1)\n",
    "    return coords"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21a18024",
   "metadata": {},
   "source": [
    "Let's set up the plotting, with three subplots for the B-mode, loss map, and speed of sound map. We will update these after every few iterations to visualize the optimization process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1ae1d40f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "fig, (ax_bmode, ax_lossmap, ax_img) = plt.subplots(\n",
    "    1,\n",
    "    3,\n",
    "    figsize=(15, 4),\n",
    "    dpi=100,\n",
    ")\n",
    "\n",
    "extent = [\n",
    "    ops.min(sos_grid_x) * 1000,\n",
    "    ops.max(sos_grid_x) * 1000,\n",
    "    ops.min(sos_grid_z) * 1000,\n",
    "    ops.max(sos_grid_z) * 1000,\n",
    "]\n",
    "bmodeim = ax_bmode.imshow(\n",
    "    np.zeros((grid_size_x, grid_size_z)),\n",
    "    extent=extent,\n",
    "    cmap=\"gray\",\n",
    "    vmin=-60,\n",
    "    vmax=0,\n",
    ")\n",
    "lossim = ax_lossmap.imshow(\n",
    "    np.zeros((grid_size_x, grid_size_z)),\n",
    "    extent=extent,\n",
    "    cmap=\"gray\",\n",
    "    vmin=0,\n",
    "    vmax=1,\n",
    ")\n",
    "im = ax_img.imshow(\n",
    "    np.zeros((grid_size_x, grid_size_z)),\n",
    "    extent=extent,\n",
    "    cmap=\"jet\",\n",
    "    origin=\"lower\",\n",
    ")\n",
    "im.set_clim(1440, 1500)\n",
    "\n",
    "ax_bmode.set_title(\"Beamformed image\")\n",
    "ax_bmode.set_xlabel(\"x [mm]\")\n",
    "ax_bmode.set_ylabel(\"z [mm]\")\n",
    "\n",
    "ax_lossmap.set_title(\"CMPE Loss plot\")\n",
    "ax_lossmap.set_xlabel(\"x [mm]\")\n",
    "ax_lossmap.set_ylabel(\"z [mm]\")\n",
    "ax_lossmap.set_yticks([])\n",
    "\n",
    "ax_img.set_title(\"Speed of Sound (SOS) Estimate\")\n",
    "ax_img.set_xlabel(\"x [mm]\")\n",
    "ax_img.set_ylabel(\"z [mm]\")\n",
    "ax_img.invert_yaxis()\n",
    "ax_img.set_yticks([])\n",
    "\n",
    "fig.colorbar(bmodeim, ax=ax_bmode, fraction=0.05, pad=0.02)\n",
    "fig.colorbar(lossim, ax=ax_lossmap, fraction=0.05, pad=0.02)\n",
    "fig.colorbar(im, ax=ax_img, fraction=0.05, pad=0.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9ec2bdd",
   "metadata": {},
   "source": [
    "Now we can finally iteratively update the sound speed grid to minimize the common midpoint phase error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d3c10284",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m200/200\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m47s\u001b[0m 211ms/step - loss: 0.1738\n",
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[32mSuccessfully saved GIF to -> \u001b[33msos_optim.gif\u001b[0m\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "viz_frames = []\n",
    "progbar = Progbar(num_iterations)\n",
    "for i in range(num_iterations):\n",
    "    flatgrid = resample_grid(parameters, xlims, zlims)\n",
    "    grad, loss = compute_gradients(sos_map, data_frame, flatgrid)\n",
    "    sos_map, opt_state = apply_gradients(opt_state, grad)\n",
    "    progbar.update(i + 1, [(\"loss\", loss)])\n",
    "    if (i + 1) % 5 == 0 or i == num_iterations - 1:\n",
    "        bmode = image_plot_pipeline(\n",
    "            data=data_frame,\n",
    "            sos_map=sos_map,\n",
    "            sos_grid_x=sos_grid_x,\n",
    "            sos_grid_z=sos_grid_z,\n",
    "            **parameters_plot,\n",
    "        )[\"data\"][0].reshape(grid_size_x, grid_size_z)\n",
    "\n",
    "        lossimage = loss_pipeline(\n",
    "            data=data_frame,\n",
    "            sos_map=sos_map,\n",
    "            sos_grid_x=sos_grid_x,\n",
    "            sos_grid_z=sos_grid_z,\n",
    "            **parameters_plot,\n",
    "        )[\"data\"][0].reshape(grid_size_x, grid_size_z)\n",
    "\n",
    "        bmodeim.set_data(bmode)\n",
    "        lossim.set_data(lossimage)\n",
    "        im.set_data(ops.convert_to_numpy(sos_map))\n",
    "        viz_frames.append(matplotlib_figure_to_numpy(fig))\n",
    "\n",
    "plt.close(fig)\n",
    "save_to_gif(viz_frames, \"sos_optim.gif\", shared_color_palette=True, fps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e0ded4c",
   "metadata": {},
   "source": [
    "![Speed of sound optimization progress](sos_optim.gif)"
   ]
  }
 ],
 "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}