{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff73dffb",
   "metadata": {},
   "source": [
    "# Beamforming with a cartesian grid and a polar grid\n",
    "\n",
    "In this notebook, we will demonstrate how you can do beamforming with `zea` using a cartesian grid and a polar grid. We will use the `ScanConvert` to convert the polar data to cartesian data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ef993a3",
   "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/polar_grid_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/polar_grid_example.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a88e145",
   "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": "bda59bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "866adf5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\"\n",
    "os.environ[\"ZEA_DISABLE_CACHE\"] = \"1\"\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d04ebab3",
   "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",
    "import numpy as np\n",
    "\n",
    "import zea\n",
    "from zea import ops\n",
    "from zea.beamform.delays import compute_t0_delays_focused\n",
    "from zea.beamform.phantoms import fish\n",
    "from zea.probes import Probe\n",
    "from zea.scan import Scan\n",
    "from zea.visualize import pad_or_crop_extent, set_mpl_style\n",
    "from zea.internal.core import DEFAULT_DYNAMIC_RANGE\n",
    "\n",
    "zea.init_device(verbose=False)\n",
    "set_mpl_style()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23da309d",
   "metadata": {},
   "source": [
    "## Define `zea.Probe` and `zea.Scan`\n",
    "\n",
    "Let's initialize a linear ultrasound probe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b44250f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_el = 128\n",
    "aperture = 30e-3\n",
    "probe_geometry = np.stack(\n",
    "    [\n",
    "        np.linspace(-aperture / 2, aperture / 2, n_el),\n",
    "        np.zeros(n_el),\n",
    "        np.zeros(n_el),\n",
    "    ],\n",
    "    axis=1,\n",
    ")\n",
    "\n",
    "probe = Probe(\n",
    "    probe_geometry=probe_geometry,\n",
    "    center_frequency=2.5e6,\n",
    "    sampling_frequency=10e6,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b00c985",
   "metadata": {},
   "source": [
    "We will use a focused scan for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea9ea615",
   "metadata": {},
   "outputs": [],
   "source": [
    "sound_speed = 1540.0\n",
    "n_tx = 8\n",
    "\n",
    "tx_apodizations = np.ones((n_tx, probe.n_el)) * np.hanning(probe.n_el)[None]\n",
    "angles = np.linspace(30, -30, n_tx) * np.pi / 180\n",
    "focus_distances = np.ones(n_tx) * 15e-3\n",
    "transmit_origins = np.zeros((n_tx, 3))\n",
    "t0_delays = compute_t0_delays_focused(\n",
    "    transmit_origins=transmit_origins,\n",
    "    focus_distances=focus_distances,\n",
    "    probe_geometry=probe.probe_geometry,\n",
    "    polar_angles=angles,\n",
    "    sound_speed=sound_speed,\n",
    ")\n",
    "\n",
    "scan = Scan(\n",
    "    n_el=n_el,\n",
    "    center_frequency=probe.center_frequency,\n",
    "    sampling_frequency=probe.sampling_frequency,\n",
    "    probe_geometry=probe.probe_geometry,\n",
    "    t0_delays=t0_delays,\n",
    "    tx_apodizations=tx_apodizations,\n",
    "    focus_distances=focus_distances,\n",
    "    transmit_origins=transmit_origins,\n",
    "    polar_angles=angles,\n",
    "    initial_times=np.ones(n_tx) * 1e-6,\n",
    "    n_ax=1024,\n",
    "    lens_sound_speed=1000,\n",
    "    lens_thickness=1e-3,\n",
    "    sound_speed=sound_speed,\n",
    "    xlims=(-20e-3, 20e-3),\n",
    "    zlims=(0, 35e-3),\n",
    "    n_tx=n_tx,\n",
    "    n_ch=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b15da7f",
   "metadata": {},
   "source": [
    "Finally we initialize a scatterer phantom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "009c6fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the fish phantom\n",
    "scat_positions = fish()\n",
    "n_scat = len(scat_positions)\n",
    "simulation_parameters = dict(\n",
    "    scatterer_positions=scat_positions.astype(np.float32),\n",
    "    scatterer_magnitudes=np.ones(n_scat, dtype=np.float32),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33b7b4cc",
   "metadata": {},
   "source": [
    "## Initialize the pipeline\n",
    "\n",
    "We initialize the default beamforming pipeline and prepend the simulator as the first operation.\n",
    "Finally, we normalize the beamformed data to [0, 255] range for visualization purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0dde4c4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mDEPRECATED\u001b[0m Beamformer name 'das' is deprecated. Please use 'delay_and_sum' instead.\n"
     ]
    }
   ],
   "source": [
    "pipeline = ops.Pipeline.from_default(beamformer=\"das\", with_batch_dim=False)\n",
    "pipeline.prepend(ops.Simulate())\n",
    "pipeline.append(ops.Normalize(input_range=DEFAULT_DYNAMIC_RANGE, output_range=(0, 255)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "457d383e",
   "metadata": {},
   "source": [
    "## Beamforming with a cartesian grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a952f37a",
   "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 azimuth angles provided, using zeros\n",
      "\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: {'n_el', 'zlims', 'xlims'}. Make sure this is intended. This warning will only be shown once.\n"
     ]
    }
   ],
   "source": [
    "# Prepare parameters for the pipeline\n",
    "scan.grid_type = \"cartesian\"  # cartesian grid is the default\n",
    "parameters = pipeline.prepare_parameters(probe, scan)\n",
    "\n",
    "# Run the pipeline\n",
    "image_cart = pipeline(**parameters, **simulation_parameters)[\"data\"]\n",
    "\n",
    "# Define the extent for the cartesian grid\n",
    "extent_cart = scan.extent.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a90ba7ae",
   "metadata": {},
   "source": [
    "## Beamforming with a polar grid\n",
    "\n",
    "We add the scan convert operation to the pipeline because we will now use a polar grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "165584fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m GPU support for order > 1 is not available. Disabling jit for ScanConvert.\n"
     ]
    },
    {
     "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: {'n_el', 'zlims', 'xlims'}. Make sure this is intended. This warning will only be shown once.\n"
     ]
    }
   ],
   "source": [
    "# Append ScanConvert to the pipeline\n",
    "pipeline_sc = pipeline.copy()\n",
    "pipeline_sc.append(ops.ScanConvert(order=3))\n",
    "pipeline_sc.append(zea.ops.keras_ops.Clip(x_min=0, x_max=255))\n",
    "\n",
    "# Prepare parameters for the pipeline\n",
    "scan.grid_type = \"polar\"  # update grid type to polar\n",
    "parameters = pipeline_sc.prepare_parameters(probe, scan)\n",
    "\n",
    "image_polar = pipeline_sc(**parameters, **simulation_parameters)[\"data\"]\n",
    "\n",
    "# Define the extent for the polar grid\n",
    "extent_polar = scan.extent.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2c2dd36",
   "metadata": {},
   "source": [
    "## Visualize the beamformed data\n",
    "\n",
    "Let's put the images side by side to compare the beamformed data on a\n",
    "cartesian grid and a polar grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "11fea932",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure the polar image has the same extent as the cartesian image\n",
    "image_polar_corrected = pad_or_crop_extent(image_polar, extent_polar, extent_cart)\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "axs[0].imshow(image_cart, cmap=\"gray\", extent=extent_cart, vmin=0, vmax=255)\n",
    "axs[0].set_xlabel(\"X (mm)\")\n",
    "axs[0].set_ylabel(\"Z (mm)\")\n",
    "axs[0].set_title(\"Cartesian\")\n",
    "axs[0].locator_params(nbins=4)\n",
    "\n",
    "axs[1].imshow(image_polar_corrected, cmap=\"gray\", extent=extent_cart, vmin=0, vmax=255)\n",
    "axs[1].set_title(\"Polar\")\n",
    "axs[1].set_xlabel(\"X (mm)\")\n",
    "axs[1].set_ylabel(\"Z (mm)\")\n",
    "axs[1].locator_params(nbins=4)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"polar_grid_fish.png\")\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f37eb64c",
   "metadata": {},
   "source": [
    "![polar_grid_fish](polar_grid_fish.png)"
   ]
  }
 ],
 "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
}