{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Beamforming to a 3D grid with `zea.Pipeline`\n",
    "\n",
    "In this notebook, we demonstrate beamforming 3D data acquired with a matrix probe using a `zea.Pipeline`."
   ]
  },
  {
   "cell_type": "markdown",
   "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/3d_beamforming_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/3d_beamforming_example.ipynb)\n",
    " \n",
    "[![Hugging Face dataset](https://img.shields.io/badge/Hugging%20Face-Dataset-yellow?logo=huggingface)](https://huggingface.co/datasets/zeahub/CIRS_3d_focused)"
   ]
  },
  {
   "cell_type": "markdown",
   "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,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\"\n",
    "os.environ[\"ZEA_DISABLE_CACHE\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: Using backend 'jax'\n"
     ]
    }
   ],
   "source": [
    "from zea import init_device\n",
    "from zea.data import load_file\n",
    "from zea.ops import (\n",
    "    Pipeline,\n",
    "    Demodulate,\n",
    "    Map,\n",
    "    EnvelopeDetect,\n",
    "    ReshapeGrid,\n",
    "    Normalize,\n",
    "    LogCompress,\n",
    "    TOFCorrection,\n",
    "    DelayAndSum,\n",
    ")\n",
    "from zea.visualize import set_mpl_style\n",
    "\n",
    "init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Set rate to downscale grid resolution for more efficient beamforming\n",
    "downscale_rate = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we download an RF data tensor acquired from a CIRS040 phantom using an 8MHz 32x32 element Matrix probe. We then load the first frame, which is of shape `(1, 56, 1280, 1024, 1)`, corresponding to 1 frame, 56 transmit events, with 1280 axial samples across 1024 channels, and 1 final dimension to indicate that the data is real-valued."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RF data shape = (1, 56, 1280, 1024, 1)\n"
     ]
    }
   ],
   "source": [
    "path = \"hf://zeahub/phantoms/2025_12_16_cirs_focused_3d.hdf5\"\n",
    "\n",
    "rf_data, scan, probe = load_file(\n",
    "    path=path,\n",
    "    indices=[0],\n",
    "    data_type=\"raw_data\",\n",
    ")\n",
    "\n",
    "# index the first frame\n",
    "print(f\"RF data shape = {rf_data.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we specify our desired beamforming parameters by modifying attributes of the `scan` object. This defines our 3D beamforming grid, which is of shape `(203, 94, 103, 3)`, corresponding to `203` axial, `94` lateral, and `103` elevational voxels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3D grid shape = (254, 94, 103, 3)\n"
     ]
    }
   ],
   "source": [
    "scan.n_ch = 2  # IQ data, should be stored in file but isn't currently\n",
    "scan.zlims = (0, 25e-3)  # reduce z-limits a bit for better visualization\n",
    "scan.grid_size_x = scan.grid_size_x // downscale_rate\n",
    "scan.grid_size_y = scan.grid_size_y // downscale_rate\n",
    "scan.grid_size_z = scan.grid_size_z // downscale_rate\n",
    "print(f\"3D grid shape = {scan.grid.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create a standard delay-and-sum beamforming pipeline. We use the `Map` operation to break the time-of-flight correction and summing into a number of chunks which are processed one at a time to avoid running out of GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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 transmit origins provided, using zeros\n"
     ]
    }
   ],
   "source": [
    "pipeline = Pipeline(\n",
    "    [\n",
    "        Demodulate(),\n",
    "        Map(\n",
    "            [TOFCorrection(), DelayAndSum()],\n",
    "            argnames=\"flatgrid\",\n",
    "            chunks=1024,  # Increase the number of chunks if you run out of memory\n",
    "        ),\n",
    "        ReshapeGrid(),\n",
    "        EnvelopeDetect(),\n",
    "        Normalize(),\n",
    "        LogCompress(),\n",
    "    ],\n",
    "    with_batch_dim=True,\n",
    ")\n",
    "parameters = pipeline.prepare_parameters(probe, scan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can beamform and visualize the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[33mDEBUG\u001b[0m [zea.Pipeline] The following input keys are not used by the pipeline: {'zlims', 'xlims', 'center_frequency', 'n_el'}. Make sure this is intended. This warning will only be shown once.\n"
     ]
    }
   ],
   "source": [
    "out = pipeline(data=rf_data, **parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[32mSuccessfully saved GIF to -> \u001b[33m./cirs_volume_rotation_2.gif\u001b[0m\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from zea.internal.notebooks import animate_volume_mip\n",
    "\n",
    "animate_volume_mip(\n",
    "    out[\"data\"],\n",
    "    f\"./cirs_volume_rotation_{downscale_rate}.gif\",\n",
    "    n_frames=60,  # 60 frames for smooth rotation\n",
    "    interval=200,  # 200ms per frame (5 fps)\n",
    "    cmap=\"gray\",\n",
    "    axis=0,  # Rotate around vertical axis\n",
    "    zoom=0.8,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Rotating volume](./cirs_volume_rotation_2.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": 2
}