{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "30a5e819",
   "metadata": {},
   "source": [
    "# REFoCUS: Synthetic Aperture Recovery for Ultrasound Imaging\n",
    "In this notebook, we demonstrate how to use the [Refocus](../../_autosummary/zea.ops.pipeline.rst) operation to decode plane-wave transmit data into synthetic aperture (SA) data, and how it affects the resulting B-mode image quality.\n",
    "\n",
    "REFoCUS works by inverting the transmit encoding model in the frequency domain, effectively recovering a full-matrix capture (multistatic) dataset from a conventional plane-wave or focused transmit sequence. This can improve lateral resolution and contrast compared to standard delay-and-sum beamforming.\n",
    "\n",
    "> **References**\n",
    ">\n",
    "> Bottenus, N. (2018). *Recovery of the complete data set from focused transmit beams.* IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control, 65(1), 30–38. https://doi.org/10.1109/TUFFC.2017.2773495\n",
    ">\n",
    "> Ali, R., Dahl, J., & Bottenus, N. (2019). *Extending Retrospective Encoding for Robust Recovery of the Multistatic Dataset.* IEEE TUFFC, 67(5), 943–956.\n",
    ">\n",
    "> Reference implementation: https://github.com/nbottenus/REFoCUS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0525856",
   "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/refocus_pipeline_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/refocus_pipeline_example.ipynb)\n",
    " \n",
    "[![Hugging Face dataset](https://img.shields.io/badge/Hugging%20Face-Dataset-yellow?logo=huggingface)](https://huggingface.co/datasets/zeahub/picmus)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36097b2e",
   "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": "2892316a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a3e312fd",
   "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": "markdown",
   "id": "8ca0649c",
   "metadata": {},
   "source": [
    "We'll import all necessary libraries and modules, including the `Refocus` and `Pipeline` classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "311c11ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'jax'\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import zea\n",
    "from zea.display import to_8bit\n",
    "from zea.ops import (\n",
    "    Pipeline,\n",
    "    Cast,\n",
    "    Refocus,\n",
    "    ApplyWindow,\n",
    "    Demodulate,\n",
    "    Beamform,\n",
    "    EnvelopeDetect,\n",
    "    Normalize,\n",
    "    LogCompress,\n",
    ")\n",
    "from zea.visualize import set_mpl_style"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d72a4c",
   "metadata": {},
   "source": [
    "We will work with the GPU if available, and initialize using `init_device` to pick the best available device. We also set the matplotlib style for consistent plot styling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "45f293f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "zea.init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c189f821",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "We'll load a real RF dataset from the [PICMUS](https://www.creatis.insa-lyon.fr/Challenge/IEEE_IUS_2016/home) challenge. This is a contrast-speckle phantom acquired with a plane-wave sequence using 75 transmit angles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "params-cell-refocus",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "num_transmits = 75\n",
    "grid_size_x = None\n",
    "grid_size_z = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd1aa79b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6657bd5b4924425f94d420ae08785b6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "database/experiments/contrast_speckle/co(…):   0%|          | 0.00/127M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_frames: 1\n",
      "n_tx:     75\n",
      "n_ax:     3328\n",
      "n_el:     128\n",
      "n_ch:     1\n"
     ]
    }
   ],
   "source": [
    "path = \"hf://zeahub/picmus/database/experiments/contrast_speckle/contrast_speckle_expe_dataset_rf/contrast_speckle_expe_dataset_rf.hdf5\"\n",
    "\n",
    "with zea.File(path, mode=\"r\", revision=\"v0.1.0\") as file:\n",
    "    data = file.data.raw_data[0]\n",
    "    parameters = file.load_parameters()\n",
    "\n",
    "data = data[None, ...]\n",
    "\n",
    "parameters.set_transmits(num_transmits)\n",
    "data = data[:, parameters.selected_transmits]\n",
    "\n",
    "if grid_size_x is not None:\n",
    "    parameters.grid_size_x = grid_size_x\n",
    "if grid_size_z is not None:\n",
    "    parameters.grid_size_z = grid_size_z\n",
    "\n",
    "print(\n",
    "    f\"n_frames: {data.shape[0]}\\n\"\n",
    "    f\"n_tx:     {data.shape[1]}\\n\"\n",
    "    f\"n_ax:     {data.shape[2]}\\n\"\n",
    "    f\"n_el:     {data.shape[3]}\\n\"\n",
    "    f\"n_ch:     {data.shape[4]}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d1b04c4",
   "metadata": {},
   "source": [
    "Since we'll be displaying B-mode images several times, let's define a small helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8f9b213f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_bmode(ax, image, parameters, title=\"\", dynamic_range=(-60, 0)):\n",
    "    \"\"\"Display a B-mode image with physical axis labels.\"\"\"\n",
    "    xlims_mm = [v * 1e3 for v in parameters.xlims]\n",
    "    zlims_mm = [v * 1e3 for v in parameters.zlims]\n",
    "    ax.imshow(\n",
    "        to_8bit(image, dynamic_range=dynamic_range),\n",
    "        cmap=\"gray\",\n",
    "        extent=[xlims_mm[0], xlims_mm[1], zlims_mm[1], zlims_mm[0]],\n",
    "        aspect=\"equal\",\n",
    "    )\n",
    "    ax.set_xlabel(\"x (mm)\")\n",
    "    ax.set_ylabel(\"z (mm)\")\n",
    "    ax.set_title(title)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701d99e1",
   "metadata": {},
   "source": [
    "## Baseline: standard DAS pipeline\n",
    "\n",
    "First, let's run the default `Pipeline` without REFoCUS to establish a baseline B-mode image using standard delay-and-sum (DAS) beamforming."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2083bbe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = Pipeline(\n",
    "    [\n",
    "        Cast(dtype=\"float32\"),\n",
    "        Demodulate(),\n",
    "        Beamform(\n",
    "            beamformer=\"delay_and_sum\",\n",
    "            num_patches=200,\n",
    "            enable_pfield=False,\n",
    "        ),\n",
    "        EnvelopeDetect(),\n",
    "        Normalize(),\n",
    "        LogCompress(),\n",
    "    ]\n",
    ")\n",
    "inputs = pipeline.prepare_parameters(parameters)\n",
    "\n",
    "inputs[pipeline.key] = data\n",
    "outputs = pipeline(**inputs)\n",
    "image_das = outputs[pipeline.output_key]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4, 6))\n",
    "plot_bmode(ax, image_das[0], parameters, title=\"DAS (no REFoCUS)\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"refocus_bmode_das.png\", dpi=150, bbox_inches=\"tight\")\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9fbf545",
   "metadata": {},
   "source": [
    "![DAS baseline B-mode](refocus_bmode_das.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afbc27e1",
   "metadata": {},
   "source": [
    "## REFoCUS pre-processing\n",
    "\n",
    "The `Refocus` operation is prepended to any existing pipeline. It decodes the plane-wave data into a synthetic aperture dataset, which is then passed to the downstream beamformer as if it came from individual element firings.\n",
    "\n",
    "It supports four inversion methods:\n",
    "\n",
    "| Method | Description |\n",
    "|--------|-------------|\n",
    "| `adjoint` | Matched-filter pseudo-inverse. Fast and parameter-free. With `param=None` a ramp filter is applied; set `param=0` to disable it. |\n",
    "| `tikhonov` | Tikhonov-regularized least-squares inverse. `param` controls regularization strength (default `0.01`). |\n",
    "| `tsvd` | Truncated SVD. Singular values below `param × σ_max` are discarded. |\n",
    "| `rsvd` | Regularized SVD, similar to Tikhonov but in the SVD domain. |\n",
    "\n",
    "Below we run all four methods and compare the resulting images side-by-side against the baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1127fb87",
   "metadata": {},
   "outputs": [],
   "source": [
    "refocus_configs = [\n",
    "    {\"method\": \"adjoint\", \"param\": None, \"label\": \"adjoint + ramp filter\"},\n",
    "    {\"method\": \"adjoint\", \"param\": 0, \"label\": \"adjoint\"},\n",
    "    {\"method\": \"tikhonov\", \"param\": 0.01, \"label\": \"tikhonov (λ=0.01)\"},\n",
    "    {\"method\": \"tsvd\", \"param\": 0.1, \"label\": \"tsvd (λ=0.1)\"},\n",
    "    {\"method\": \"rsvd\", \"param\": 0.05, \"label\": \"rsvd (λ=0.05)\"},\n",
    "]\n",
    "\n",
    "n_cols = 2\n",
    "n_rows = (len(refocus_configs) + 2) // n_cols\n",
    "fig, axes = plt.subplots(n_rows, n_cols, figsize=(10, n_rows * 5))\n",
    "axes = axes.flatten()\n",
    "\n",
    "# Baseline (no REFoCUS)\n",
    "plot_bmode(axes[0], image_das[0], parameters, title=\"DAS (baseline)\")\n",
    "\n",
    "for i, cfg in enumerate(refocus_configs):\n",
    "    pipeline = Pipeline(\n",
    "        [\n",
    "            Cast(dtype=\"float32\"),\n",
    "            ApplyWindow(size=64, start=64),\n",
    "            Refocus(method=cfg[\"method\"], param=cfg[\"param\"], jit_compile=False),\n",
    "            Demodulate(),\n",
    "            Beamform(\n",
    "                beamformer=\"delay_and_sum\",\n",
    "                num_patches=200,\n",
    "                enable_pfield=False,\n",
    "            ),\n",
    "            EnvelopeDetect(),\n",
    "            Normalize(),\n",
    "            LogCompress(),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    inputs = pipeline.prepare_parameters(parameters)\n",
    "\n",
    "    outputs = pipeline(**{pipeline.key: data}, **inputs)\n",
    "    image = outputs[pipeline.output_key]\n",
    "\n",
    "    plot_bmode(axes[i + 1], image[0], parameters, title=f\"REFoCUS\\n{cfg['label']}\")\n",
    "\n",
    "# Hide any unused axes\n",
    "for ax in axes[len(refocus_configs) + 1 :]:\n",
    "    ax.set_visible(False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"refocus_example.png\", dpi=150, bbox_inches=\"tight\")\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "349aa994",
   "metadata": {},
   "source": [
    "![refocus_example](refocus_example.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa64825f",
   "metadata": {},
   "source": []
  }
 ],
 "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
}