{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65d76646",
   "metadata": {},
   "source": [
    "# How fine should the beamforming grid be?\n",
    "\n",
    "A pixel pitch of $\\lambda/2$ is a common choice, but the pitch an image needs follows from its k-space support, which depends on the aperture, the steering, the bandwidth, and whether the image is RF or IQ.\n",
    "\n",
    "In this notebook we derive the limits `zea` checks against, compare them to a beamformed point target, and see when $\\lambda/2$ is too coarse or needlessly fine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b724dd",
   "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/kspace_aliasing_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/kspace_aliasing_example.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa98e3b8",
   "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": "00ee7afb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:44.085226Z",
     "iopub.status.busy": "2026-11-05T15:34:44.085032Z",
     "iopub.status.idle": "2026-11-05T15:34:45.261627Z",
     "shell.execute_reply": "2026-11-05T15:34:45.260075Z"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install zea"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fdbb8299",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:45.265008Z",
     "iopub.status.busy": "2026-11-05T15:34:45.264753Z",
     "iopub.status.idle": "2026-11-05T15:34:45.269209Z",
     "shell.execute_reply": "2026-11-05T15:34:45.268202Z"
    }
   },
   "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": "9770222a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:45.271259Z",
     "iopub.status.busy": "2026-11-05T15:34:45.271075Z",
     "iopub.status.idle": "2026-11-05T15:34:48.147946Z",
     "shell.execute_reply": "2026-11-05T15:34:48.146465Z"
    }
   },
   "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",
    "from keras import ops\n",
    "from matplotlib.patches import Rectangle\n",
    "from scipy.signal import hilbert\n",
    "from scipy.signal.windows import tukey\n",
    "\n",
    "import zea\n",
    "from zea import Parameters\n",
    "from zea.beamform.beamformer import fnum_window_fn_tukey, fnumber_mask\n",
    "from zea.beamform.delays import compute_t0_delays_planewave\n",
    "from zea.beamform.pixelgrid import aliasing_limits\n",
    "from zea.func import log_compress\n",
    "from zea.probes import create_probe_geometry\n",
    "from zea.visualize import set_mpl_style\n",
    "\n",
    "zea.init_device(verbose=False)\n",
    "set_mpl_style()\n",
    "\n",
    "BLUE, CYAN, GREEN, RED = \"#4fc3f7\", \"#00e5ff\", \"#66bb6a\", \"#ef5350\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ed37de",
   "metadata": {},
   "source": [
    "## 1. The k-space support of a beamformed image\n",
    "\n",
    "A transmit along the unit vector $\\mathbf{t}$ and a receive along the unit vector $\\mathbf{r}$ at frequency $f$ contribute the spatial frequency\n",
    "\n",
    "$$\\mathbf{k} = \\frac{f}{c}\\left(\\mathbf{t} + \\mathbf{r}\\right) = \\frac{\\mathbf{t} + \\mathbf{r}}{\\lambda}$$\n",
    "\n",
    "to the image ([Walker and Trahey](https://doi.org/10.1109/58.677599)). Here $\\mathbf{k}$ is a vector with components $k_x$ and $k_z$, in cycles per metre. When $\\mathbf{t}$ and $\\mathbf{r}$ both point along $z$, this gives $k_z = 2/\\lambda$, twice the spatial frequency of the wave itself, because a pixel $\\lambda/2$ deeper adds $\\lambda$ to the pulse-echo path. The same holds along any direction they share, and $|\\mathbf{k}|$ gets smaller the further apart they point.\n",
    "\n",
    "The support follows from sweeping each term:\n",
    "- the receive angle over $\\tan\\theta_r \\le 1/(2F)$, with $F$ the f-number\n",
    "- the frequency over the band\n",
    "- the transmit over the steering angles\n",
    "\n",
    "Here is an example for a 5 MHz pulse with 3 MHz of bandwidth, $F = 0.5$, and five plane waves up to $\\pm 25°$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fc52d833",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:48.150820Z",
     "iopub.status.busy": "2026-11-05T15:34:48.150602Z",
     "iopub.status.idle": "2026-11-05T15:34:48.616108Z",
     "shell.execute_reply": "2026-11-05T15:34:48.614561Z"
    }
   },
   "outputs": [],
   "source": [
    "sound_speed = 1540.0\n",
    "center_frequency = 5e6\n",
    "bandwidth = 3e6\n",
    "wavelength = sound_speed / center_frequency\n",
    "f_number = 0.5\n",
    "steering = np.deg2rad([-25, -12.5, 0, 12.5, 25])\n",
    "\n",
    "\n",
    "def k_support(frequencies, polar_angles, f_number):\n",
    "    \"\"\"All k = f (t + r) / c in cycles per wavelength, with shape (n_points, 2).\"\"\"\n",
    "    max_angle = np.arctan(1 / (2 * f_number))\n",
    "    r = np.linspace(-max_angle, max_angle, 301)\n",
    "    f, t, r = np.meshgrid(frequencies, polar_angles, r)\n",
    "    k_x = f * (np.sin(t) + np.sin(r)) / sound_speed\n",
    "    k_z = f * (np.cos(t) + np.cos(r)) / sound_speed\n",
    "    return np.stack([k_x.ravel(), k_z.ravel()], axis=-1) * wavelength\n",
    "\n",
    "\n",
    "frequencies = center_frequency + np.linspace(-0.5, 0.5, 101) * bandwidth\n",
    "steps = [\n",
    "    (\"1. Receive angles\", [center_frequency], [0.0]),\n",
    "    (\"2. Add the band\", frequencies, [0.0]),\n",
    "    (\"3. Add steering\", frequencies, steering),\n",
    "]\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(11, 3.6), sharex=True, sharey=True)\n",
    "for ax, (title, f, t) in zip(axes, steps):\n",
    "    k = k_support(f, t, f_number)\n",
    "    ax.scatter(*k.T, s=4, color=BLUE, linewidths=0, rasterized=True)\n",
    "    ax.set_title(title, loc=\"left\")\n",
    "    ax.set_xlabel(\"$k_x$ [cycles per wavelength]\")\n",
    "    ax.set_aspect(\"equal\")\n",
    "axes[0].set_xlim(-1.8, 1.8)\n",
    "axes[0].set_ylim(2.9, 0.9)\n",
    "axes[0].set_ylabel(\"$k_z$ [cycles per wavelength]\")\n",
    "fig.suptitle(\"K-space support of a beamformed image\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"kspace_support_build_up.png\", dpi=150)\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a879489",
   "metadata": {},
   "source": [
    "![The k-space support built up in three steps](kspace_support_build_up.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00848d73",
   "metadata": {},
   "source": [
    "## 2. Comparison with a beamformed point target\n",
    "\n",
    "We simulate a point scatterer with these settings and beamform it on a $\\lambda/16$ grid, so that its 2D Fourier transform shows the support. This grid is well within all the limits below, so we call the result the reference image: what a coarser grid should reproduce. The RF is band-limited to give the band known edges, and we use a smooth receive aperture instead of the default hard f-number mask, see section 5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a7ec8744",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:48.618880Z",
     "iopub.status.busy": "2026-11-05T15:34:48.618765Z",
     "iopub.status.idle": "2026-11-05T15:34:58.188085Z",
     "shell.execute_reply": "2026-11-05T15:34:58.186201Z"
    }
   },
   "outputs": [],
   "source": [
    "n_el = 192\n",
    "element_pitch = wavelength / 2\n",
    "probe_geometry = create_probe_geometry(n_el, element_pitch)\n",
    "\n",
    "fine_pitch = wavelength / 16\n",
    "n_grid = 513\n",
    "half_extent = (n_grid - 1) / 2 * fine_pitch\n",
    "depth = 10e-3\n",
    "sampling_frequency = 16 * center_frequency\n",
    "n_ax = 4096\n",
    "\n",
    "\n",
    "def transmit_geometry(polar_angles):\n",
    "    n_tx = len(polar_angles)\n",
    "    return dict(\n",
    "        focus_distances=np.full(n_tx, np.inf),\n",
    "        probe_geometry=probe_geometry,\n",
    "        tx_apodizations=np.ones((n_tx, n_el)),\n",
    "        polar_angles=np.asarray(polar_angles, dtype=float),\n",
    "        transmit_origins=np.zeros((n_tx, 3)),\n",
    "    )\n",
    "\n",
    "\n",
    "def make_parameters(polar_angles, **kwargs):\n",
    "    n_tx = len(polar_angles)\n",
    "    settings = dict(\n",
    "        xlims=(-half_extent, half_extent),\n",
    "        zlims=(depth - half_extent, depth + half_extent),\n",
    "        grid_size_x=n_grid,\n",
    "        grid_size_z=n_grid,\n",
    "    )\n",
    "    return Parameters(\n",
    "        n_tx=n_tx,\n",
    "        n_el=n_el,\n",
    "        n_ax=n_ax,\n",
    "        center_frequency=center_frequency,\n",
    "        sampling_frequency=sampling_frequency,\n",
    "        sound_speed=sound_speed,\n",
    "        element_width=element_pitch,\n",
    "        # Required by Simulate, but unused without lens correction.\n",
    "        lens_thickness=1e-3,\n",
    "        lens_sound_speed=1000.0,\n",
    "        t0_delays=compute_t0_delays_planewave(\n",
    "            probe_geometry=probe_geometry,\n",
    "            polar_angles=np.asarray(polar_angles),\n",
    "            sound_speed=sound_speed,\n",
    "        ),\n",
    "        initial_times=np.zeros(n_tx),\n",
    "        **transmit_geometry(polar_angles),\n",
    "        **{**settings, **kwargs},\n",
    "    )\n",
    "\n",
    "\n",
    "def band_limit(rf):\n",
    "    \"\"\"Keep only center_frequency +- bandwidth / 2.\"\"\"\n",
    "    f = np.fft.rfftfreq(rf.shape[1], 1 / sampling_frequency)\n",
    "    in_band = np.abs(f - center_frequency) <= bandwidth / 2\n",
    "    gain = np.zeros_like(f)\n",
    "    gain[in_band] = tukey(in_band.sum(), 0.5)\n",
    "    spectrum = np.fft.rfft(rf, axis=1) * gain[None, :, None, None]\n",
    "    return np.fft.irfft(spectrum, n=rf.shape[1], axis=1)\n",
    "\n",
    "\n",
    "def point_target_image(polar_angles, f_number, iq=True):\n",
    "    \"\"\"Beamform one scatterer at (0, depth), returning a real (RF) or complex (IQ) image.\"\"\"\n",
    "    # f_number=0 disables the hard mask, the smooth one below replaces it.\n",
    "    parameters = make_parameters(polar_angles, f_number=0.0)\n",
    "    parameters.flat_receive_apodization = ops.convert_to_numpy(\n",
    "        fnumber_mask(parameters.flatgrid, probe_geometry, f_number, fnum_window_fn_tukey)\n",
    "    )[..., 0]\n",
    "\n",
    "    simulator = zea.Pipeline([zea.ops.Simulate()], with_batch_dim=False)\n",
    "    inputs = simulator.prepare_parameters(\n",
    "        parameters,\n",
    "        scatterer_positions=np.array([[0.0, 0.0, depth]], dtype=\"float32\"),\n",
    "        scatterer_magnitudes=np.ones(1, dtype=\"float32\"),\n",
    "    )\n",
    "    rf = simulator(return_numpy=True, **inputs)[simulator.output_key]\n",
    "    rf = band_limit(rf).astype(\"float32\")\n",
    "\n",
    "    layers = [zea.ops.Beamform(enable_receive_apodization=True)]\n",
    "    if iq:\n",
    "        layers.insert(0, zea.ops.Demodulate())\n",
    "    pipeline = zea.Pipeline(layers, with_batch_dim=False)\n",
    "    inputs = pipeline.prepare_parameters(parameters)\n",
    "    inputs[pipeline.key] = rf\n",
    "    image = pipeline(return_numpy=True, **inputs)[pipeline.output_key]\n",
    "    image = image[..., 0] + 1j * image[..., 1] if iq else image[..., 0]\n",
    "    # Fade out the edges, so that the crop itself does not smear the spectrum.\n",
    "    return image * np.outer(tukey(n_grid, 0.3), tukey(n_grid, 0.3))\n",
    "\n",
    "\n",
    "def limits_for(polar_angles, f_number, iq=True):\n",
    "    return aliasing_limits(\n",
    "        sound_speed,\n",
    "        center_frequency,\n",
    "        f_number,\n",
    "        **transmit_geometry(polar_angles),\n",
    "        demodulated=iq,\n",
    "        bandwidth=bandwidth,\n",
    "        element_width=element_pitch,\n",
    "    )\n",
    "\n",
    "\n",
    "compounded = point_target_image(steering, f_number)\n",
    "limits = limits_for(steering, f_number)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f594c4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:58.191905Z",
     "iopub.status.busy": "2026-11-05T15:34:58.191663Z",
     "iopub.status.idle": "2026-11-05T15:34:58.433977Z",
     "shell.execute_reply": "2026-11-05T15:34:58.432823Z"
    }
   },
   "outputs": [],
   "source": [
    "k_axis = np.fft.fftshift(np.fft.fftfreq(n_grid, fine_pitch)) * wavelength\n",
    "k_extent = [k_axis[0], k_axis[-1], k_axis[-1], k_axis[0]]\n",
    "zoom = int(2e-3 / fine_pitch)\n",
    "zoom_extent = np.array([-1, 1, 1, -1]) * zoom * fine_pitch * 1e3\n",
    "\n",
    "\n",
    "def spectrum_db(image):\n",
    "    magnitude = np.abs(np.fft.fftshift(np.fft.fft2(image)))\n",
    "    return log_compress(np.float32(magnitude / magnitude.max()))\n",
    "\n",
    "\n",
    "def show_spectrum(ax, image, ylim=(3, 0)):\n",
    "    ax.imshow(spectrum_db(image), extent=k_extent, cmap=\"magma\", vmin=-50, vmax=0, aspect=\"auto\")\n",
    "    ax.set_box_aspect(1)\n",
    "    ax.set_xlim(-3, 3)\n",
    "    ax.set_ylim(ylim)\n",
    "    ax.set_xlabel(\"$k_x$ [cycles per wavelength]\")\n",
    "    ax.set_ylabel(\"$k_z$ [cycles per wavelength]\")\n",
    "\n",
    "\n",
    "def envelope(image):\n",
    "    return np.abs(hilbert(image, axis=0)) if np.isrealobj(image) else np.abs(image)\n",
    "\n",
    "\n",
    "def show_image(ax, image, reference=None):\n",
    "    image_envelope = envelope(image)\n",
    "    reference = image_envelope.max() if reference is None else reference\n",
    "    center = n_grid // 2\n",
    "    crop = image_envelope[center - zoom : center + zoom, center - zoom : center + zoom]\n",
    "    ax.imshow(\n",
    "        log_compress(np.float32(crop / reference)),\n",
    "        extent=zoom_extent,\n",
    "        cmap=\"gray\",\n",
    "        vmin=-50,\n",
    "        vmax=0,\n",
    "    )\n",
    "    ax.set_xlabel(\"$x$ [mm]\")\n",
    "    ax.set_ylabel(\"$z - z_0$ [mm]\")\n",
    "\n",
    "\n",
    "def draw_box(ax, limits):\n",
    "    x_lo, x_hi = np.array(limits[\"kx_range\"]) * wavelength\n",
    "    z_lo, z_hi = np.array(limits[\"kz_range\"]) * wavelength\n",
    "    ax.add_patch(Rectangle((x_lo, z_lo), x_hi - x_lo, z_hi - z_lo, fill=False, ec=CYAN, lw=1.6))\n",
    "\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))\n",
    "show_image(axes[0], compounded)\n",
    "axes[0].set_title(\"Image of a point\", loc=\"left\")\n",
    "show_spectrum(axes[1], compounded)\n",
    "draw_box(axes[1], limits)\n",
    "axes[1].set_title(\"Its 2D Fourier transform\", loc=\"left\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"kspace_measured_support.png\", dpi=150)\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1058f7bf",
   "metadata": {},
   "source": [
    "![A beamformed point and its measured k-space](kspace_measured_support.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c652912d",
   "metadata": {},
   "source": [
    "The measured spectrum fills the predicted region, with its outer edges faded by the smooth aperture. The cyan box is the bounding box from `aliasing_limits`. Only its width along each axis matters for the pitch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bbb6629",
   "metadata": {},
   "source": [
    "## 3. Pitch limits\n",
    "\n",
    "A grid with pitch $d$ repeats the spectrum every $1/d$, so each side of the box has to fit in one period. With $k_\\mathrm{lo}$ and $k_\\mathrm{hi}$ the edges of the box along each axis, $\\lambda_\\mathrm{min} = c / (f_c + B/2)$ the shortest wavelength in the band, and $B$ the bandwidth:\n",
    "\n",
    "| Axis | Limit | Approximately |\n",
    "|---|---|---|\n",
    "| Lateral | $d_x \\le 1 / (k_{x,\\mathrm{hi}} - k_{x,\\mathrm{lo}})$ | $\\lambda_\\mathrm{min} F / 2$ for equal transmit and receive apertures |\n",
    "| Axial, RF | $d_z \\le 1 / (2 k_{z,\\mathrm{hi}})$ | $\\lambda_\\mathrm{min} / 4$ |\n",
    "| Axial, IQ | $d_z \\le 1 / (k_{z,\\mathrm{hi}} - k_{z,\\mathrm{lo}})$ | $c / (2B)$ on axis |\n",
    "\n",
    "An RF image is real, so axially its mirrored band counts too. Laterally that band sits at the opposite $k_z$, so RF and IQ have the same limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "180bca0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:58.436438Z",
     "iopub.status.busy": "2026-11-05T15:34:58.436340Z",
     "iopub.status.idle": "2026-11-05T15:34:58.493834Z",
     "shell.execute_reply": "2026-11-05T15:34:58.492764Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Largest pitch, in wavelengths   Lateral  Axial RF  Axial IQ\n",
      "One plane wave, F=2                1.59      0.19      0.82\n",
      "One plane wave, F=1                0.86      0.19      0.78\n",
      "Compounded ±25°, F=1               0.44      0.19      0.75\n",
      "Compounded ±25°, F=0.5             0.34      0.19      0.68\n"
     ]
    }
   ],
   "source": [
    "acquisitions = {\n",
    "    \"One plane wave, F=2\": ([0.0], 2.0),\n",
    "    \"One plane wave, F=1\": ([0.0], 1.0),\n",
    "    \"Compounded ±25°, F=1\": (steering, 1.0),\n",
    "    \"Compounded ±25°, F=0.5\": (steering, 0.5),\n",
    "}\n",
    "\n",
    "print(f\"{'Largest pitch, in wavelengths':<30}{'Lateral':>9}{'Axial RF':>10}{'Axial IQ':>10}\")\n",
    "for name, (angles, f) in acquisitions.items():\n",
    "    iq, rf = limits_for(angles, f, iq=True), limits_for(angles, f, iq=False)\n",
    "    print(\n",
    "        f\"{name:<30}{iq['dx_max'] / wavelength:>9.2f}\"\n",
    "        f\"{rf['dz_max'] / wavelength:>10.2f}{iq['dz_max'] / wavelength:>10.2f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19a98b99",
   "metadata": {},
   "source": [
    "So $\\lambda/2$ can be too coarse, as for RF and the compounded acquisitions, or needlessly fine, as for the lateral pitch of a single plane wave."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c1ca0b0",
   "metadata": {},
   "source": [
    "## 4. Aliasing on the grid\n",
    "\n",
    "Delay-and-sum computes each pixel independently, with no anti-aliasing filter, so a coarser grid samples the same image at fewer points, and frequencies beyond the limit fold back instead of being smoothed away. We can therefore mimic a coarser grid by keeping every $M$-th row or column of pixels in the reference image.\n",
    "\n",
    "In the examples below, we recover the image with ideal band-limited interpolation and report the error energy relative to the reference image.\n",
    "\n",
    "First the lateral pitch of the compounded acquisition, which has a limit of $0.34\\lambda$. We compare a pitch below the limit, $\\lambda/4$, with one above it, $\\lambda/2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d0002db3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:58.496012Z",
     "iopub.status.busy": "2026-11-05T15:34:58.495769Z",
     "iopub.status.idle": "2026-11-05T15:34:58.908488Z",
     "shell.execute_reply": "2026-11-05T15:34:58.907307Z"
    }
   },
   "outputs": [],
   "source": [
    "def subsample_and_recover(image, axis, step, center=0.0):\n",
    "    \"\"\"Keep every step-th row (axis=0) or column (axis=1), recover by band-limiting.\n",
    "\n",
    "    Returns the sampled image and the recovered image. `center` is the middle of the kept\n",
    "    period in cycles per wavelength.\n",
    "    \"\"\"\n",
    "    sampled = np.zeros_like(image)\n",
    "    keep = [slice(None), slice(None)]\n",
    "    keep[axis] = slice(None, None, step)\n",
    "    sampled[tuple(keep)] = step * image[tuple(keep)]\n",
    "\n",
    "    k = np.fft.fftfreq(n_grid, fine_pitch) * wavelength\n",
    "    half_period = wavelength / (2 * step * fine_pitch)\n",
    "    gain = np.abs(k - center) <= half_period\n",
    "    gain = gain[:, None] if axis == 0 else gain[None, :]\n",
    "    recovered = np.fft.ifft2(np.fft.fft2(sampled) * gain)\n",
    "    return sampled, recovered.real if np.isrealobj(image) else recovered\n",
    "\n",
    "\n",
    "def error_db(recovered, image):\n",
    "    return 20 * np.log10(np.linalg.norm(recovered - image) / np.linalg.norm(image))\n",
    "\n",
    "\n",
    "def compare(rows, filename, ylim=(3, 0)):\n",
    "    \"\"\"One row per (title, image, axis, step, center, color): spectrum, then recovered image.\"\"\"\n",
    "    fig, axes = plt.subplots(len(rows), 2, figsize=(9, 4.2 * len(rows)))\n",
    "    for ax_row, (title, image, axis, step, center, color) in zip(axes, rows):\n",
    "        sampled, recovered = subsample_and_recover(image, axis, step, center)\n",
    "        show_spectrum(ax_row[0], sampled, ylim=ylim)\n",
    "        ax_row[0].set_title(title, loc=\"left\", color=color)\n",
    "        show_image(ax_row[1], recovered, reference=envelope(image).max())\n",
    "        error = error_db(recovered, image)\n",
    "        ax_row[1].set_title(f\"Recovered image, error {error:.0f} dB\", loc=\"left\", color=color)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(filename, dpi=150)\n",
    "    plt.close()\n",
    "\n",
    "\n",
    "compare(\n",
    "    [\n",
    "        (\"Lateral pitch λ/4\", compounded, 1, 4, 0.0, GREEN),\n",
    "        (\"Lateral pitch λ/2\", compounded, 1, 8, 0.0, RED),\n",
    "    ],\n",
    "    \"kspace_lateral_aliasing.png\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a34a5f",
   "metadata": {},
   "source": [
    "![Lateral sampling within and beyond the limit](kspace_lateral_aliasing.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fcbca03",
   "metadata": {},
   "source": [
    "Next, one plane wave at $F = 1$ with an axial pitch of $\\lambda/2$, as IQ and as RF:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0576541c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:34:58.911256Z",
     "iopub.status.busy": "2026-11-05T15:34:58.911157Z",
     "iopub.status.idle": "2026-11-05T15:35:01.131457Z",
     "shell.execute_reply": "2026-11-05T15:35:01.130202Z"
    }
   },
   "outputs": [],
   "source": [
    "single_angle = [0.0]\n",
    "single_iq = point_target_image(single_angle, 1.0, iq=True)\n",
    "single_rf = point_target_image(single_angle, 1.0, iq=False)\n",
    "kz_lo, kz_hi = limits_for(single_angle, 1.0, iq=True)[\"kz_range\"]\n",
    "band_center = (kz_lo + kz_hi) / 2 * wavelength\n",
    "\n",
    "compare(\n",
    "    [\n",
    "        (\"IQ, axial pitch λ/2\", single_iq, 0, 8, band_center, GREEN),\n",
    "        (\"RF, axial pitch λ/2\", single_rf, 0, 8, 0.0, RED),\n",
    "    ],\n",
    "    \"kspace_axial_aliasing.png\",\n",
    "    ylim=(3, -3),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5be95948",
   "metadata": {},
   "source": [
    "![IQ and RF at the same axial pitch](kspace_axial_aliasing.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8611fe",
   "metadata": {},
   "source": [
    "The two spectra look alike, but each RF lobe is the band and its mirror stacked on top of each other, which is why it looks symmetric. Only the IQ image can be recovered."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14524030",
   "metadata": {},
   "source": [
    "## 5. In `zea`\n",
    "\n",
    "`Parameters` checks a cartesian grid against these limits whenever it computes the grid, and warns if the pitch is too large. You can also call the check directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "05d0ee31",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-11-05T15:35:01.134439Z",
     "iopub.status.busy": "2026-11-05T15:35:01.134336Z",
     "iopub.status.idle": "2026-11-05T15:35:01.302479Z",
     "shell.execute_reply": "2026-11-05T15:35:01.301346Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m Lateral grid pitch 154.0 um exceeds 104.9 um (lambda_min = 236.9 um, x of the transmit directions in [-0.423, 0.423], sin(theta_r) = +-0.707, element directivity sin(theta) <= 0.928). Consider increasing grid_size_x to 95 or more (for an automatically sized grid, raise pixels_per_wavelength), or narrowing the apertures (f_number / transmit steering) so the data stays inside the grid's lateral band.\n"
     ]
    }
   ],
   "source": [
    "coarse = make_parameters(\n",
    "    steering,\n",
    "    f_number=f_number,\n",
    "    grid_size_x=int(2 * half_extent / (wavelength / 2)) + 1,\n",
    ")\n",
    "coarse.check_for_aliasing(bandwidth=bandwidth);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec4a19d",
   "metadata": {},
   "source": [
    "The parameters cannot tell two things, so those are up to you:\n",
    "\n",
    "- **RF or IQ.** The check assumes IQ, the output of the default pipeline. If you beamform RF, call `check_for_aliasing(demodulated=False)`.\n",
    "- **Bandwidth.** Set `bandwidth` if you know it. Otherwise the check assumes `probe_bandwidth_percent` of the center frequency, 100% by default, which gives an axial IQ limit close to $\\lambda/2$.\n",
    "\n",
    "The limits come from a simple model, so leave some margin:\n",
    "\n",
    "- **Hard apertures leak.** The default f-number mask and the discrete element positions put faint energy outside the box.\n",
    "- **The edges are soft.** Some energy spills just past the box, especially at high f-numbers where the box is narrow, so the error does not drop to zero exactly at the limit."
   ]
  }
 ],
 "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.13"
  },
  "zea": {
   "tags": [
    "beamforming",
    "simulation"
   ],
   "level": "intermediate"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}