{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# RAM and disk caching for huge images\n",
    "\n",
    "A long SAR capture can be tens of gigabytes: neither the raw sweeps nor the\n",
    "range compressed data fit in RAM, let alone in VRAM. `torchbp.data` provides\n",
    "lazy data sources that stand in for the range compressed `data` tensor:\n",
    "\n",
    "- `MemmapData` wraps any sliceable on-disk array (a numpy memmap, an h5py\n",
    "  dataset, a zarr array) and loads sweep ranges on demand.\n",
    "- `CallbackData` calls a user function to load a sweep range. For example a\n",
    "  safetensors `get_slice`, or a reader that seeks into a custom capture format.\n",
    "- `CachedData` caches another source's transformed output in RAM or in a\n",
    "  memory-mapped disk file, so an expensive per-chunk pipeline runs only once\n",
    "  per sweep even when a consumer reads the data many times.\n",
    "\n",
    "Every backprojection function accepts a lazy source in place of the tensor.\n",
    "`torchbp.ops.ffbp` and the `gpga`/`gpga_tde` autofocus consume it in bounded\n",
    "chunks, so their peak data residency stays far below the full extent. The\n",
    "other functions (direct backprojection, `afbp`, `cfbp`) accept a lazy source\n",
    "but materialize the full extent at entry.\n",
    "\n",
    "This notebook demonstrates the machinery on a miniature simulated capture that\n",
    "plays the role of a file too large to load. The sizes are small because the\n",
    "notebook runs on the documentation build server, but nothing in the code\n",
    "depends on the data being small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import os\n",
    "import tempfile\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import torchbp\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import hamming\n",
    "from torchbp.util import detrend\n",
    "\n",
    "torch.manual_seed(125)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Simulated raw data on disk\n",
    "\n",
    "We simulate a small FMCW acquisition: a platform flying along the $y$ axis\n",
    "over a grid of point targets, with a smooth low-frequency motion error that the\n",
    "autofocus will recover at the end of the notebook. The simulated raw dechirped\n",
    "sweeps are written to a file on disk and the in-memory copy is deleted.\n",
    "From here on, the file is the only place the raw data exists, exactly as it\n",
    "would be for a real capture that does not fit in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 6e9        # RF center frequency (Hz)\n",
    "bw = 300e6      # RF bandwidth (Hz)\n",
    "tsweep = 200e-6 # Sweep length (s)\n",
    "fs = 3e6        # Sampling frequency (Hz)\n",
    "nsweeps = 512   # Number of measurements (aperture positions)\n",
    "nsamples = int(fs * tsweep)  # Time-domain samples per sweep\n",
    "oversample = 2  # Range FFT oversampling (reduces interpolation error)\n",
    "nfft = nsamples * oversample\n",
    "\n",
    "wl = 3e8 / fc\n",
    "r_res = 3e8 / (2 * bw * oversample)  # Range bin size in the compressed data\n",
    "\n",
    "# Imaging grid. Azimuth angle \"theta\" is sin of radians.\n",
    "nr = 2 * int(1 + (70 - 30) / (3e8 / (2 * bw)))\n",
    "grid_polar = {\"r\": (30, 70), \"theta\": (-0.7, 0.7), \"nr\": nr, \"ntheta\": 384}\n",
    "\n",
    "# 3x3 grid of point targets.\n",
    "targets = [[x, y, 0] for x in np.linspace(35, 65, 3) for y in np.linspace(-15, 15, 3)]\n",
    "target_pos = torch.tensor(targets, dtype=torch.float32, device=device)\n",
    "target_rcs = torch.ones(target_pos.shape[0], dtype=torch.float32, device=device)\n",
    "\n",
    "# Nominal straight trajectory and the true one with a band-limited error.\n",
    "pos_nominal = torch.zeros([nsweeps, 3], dtype=torch.float32, device=device)\n",
    "pos_nominal[:, 1] = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps) * 0.25 * wl\n",
    "pos_nominal[:, 2] = 20\n",
    "\n",
    "def bandlimited_error(scale, cutoff=nsweeps // 32):\n",
    "    r = scale * torch.rand(nsweeps, device=device) * wl\n",
    "    fdata = torch.fft.fft(r)\n",
    "    fdata[cutoff // 2:-cutoff // 2] = 0\n",
    "    e = torch.fft.ifft(fdata).real\n",
    "    return detrend(e - torch.mean(e))\n",
    "\n",
    "pos = pos_nominal.clone()\n",
    "pos[:, 0] += bandlimited_error(10)\n",
    "pos[:, 1] += bandlimited_error(2.5)\n",
    "pos[:, 2] += bandlimited_error(5)\n",
    "\n",
    "raw = torchbp.util.generate_fmcw_data(\n",
    "    target_pos, target_rcs, pos, fc, bw, tsweep, fs, rvp=False)\n",
    "\n",
    "# Write the raw sweeps to disk and drop the in-memory copy.\n",
    "tmpdir = tempfile.mkdtemp()\n",
    "raw_path = os.path.join(tmpdir, \"raw_sweeps.npy\")\n",
    "np.save(raw_path, raw.cpu().numpy())\n",
    "del raw\n",
    "print(f\"Raw capture on disk: {os.path.getsize(raw_path) / 1e6:.1f} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Lazy source with a range compression transform\n",
    "\n",
    "`MemmapData` wraps the on-disk array. The `transform` argument is the\n",
    "per-chunk processing pipeline: it is called as `transform(chunk, start, stop)`\n",
    "after the chunk has been moved to the compute device, where `start`/`stop` are\n",
    "the absolute sweep indices of the chunk, so per-sweep factors. Here the\n",
    "azimuth window over the full aperture can be indexed correctly no matter\n",
    "which chunk is being processed. On a GPU the pipeline runs there, so range\n",
    "compression is accelerated for free.\n",
    "\n",
    "The transform may change the dtype and the number of samples per sweep (here:\n",
    "windowing plus an oversampled inverse FFT), so the declared `shape` and\n",
    "`dtype` describe the transform's output, not the raw file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "wr = torch.tensor(hamming(nsamples)[None, :], dtype=torch.float32, device=device)\n",
    "wa = torch.tensor(hamming(nsweeps)[:, None], dtype=torch.float32, device=device)\n",
    "\n",
    "# Counters to observe when (and how much) data is actually loaded.\n",
    "stats = {\"chunks\": 0, \"sweeps\": 0, \"max_chunk\": 0}\n",
    "\n",
    "def range_compress(chunk, start, stop):\n",
    "    stats[\"chunks\"] += 1\n",
    "    stats[\"sweeps\"] += stop - start\n",
    "    stats[\"max_chunk\"] = max(stats[\"max_chunk\"], stop - start)\n",
    "    x = chunk * wa[start:stop] * wr\n",
    "    return torch.fft.ifft(x, dim=-1, n=nfft)\n",
    "\n",
    "raw_mm = np.load(raw_path, mmap_mode=\"r\")\n",
    "lazy = torchbp.MemmapData(raw_mm, device=device, dtype=torch.complex64,\n",
    "                          transform=range_compress, shape=(nsweeps, nfft))\n",
    "print(lazy.shape, lazy.dtype, lazy.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The lazy source mimics the small part of the Tensor interface the library\n",
    "uses. Slicing is free, it returns a view restricted to a sweep range without\n",
    "touching the disk, and `load()` materializes a view by reading and\n",
    "transforming exactly that range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "view = lazy[100:200]\n",
    "print(\"After slicing:  \", stats)\n",
    "chunk = view.load()\n",
    "print(\"After view.load:\", stats, \"-> tensor\", tuple(chunk.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "As a reference for the rest of the notebook, run the same range compression\n",
    "eagerly on the whole array at once. This is what a script that loads\n",
    "everything up front would compute, and what the lazy path must reproduce."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "data_eager = range_compress(\n",
    "    torch.as_tensor(np.array(raw_mm)).to(device), 0, nsweeps)\n",
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "\n",
    "err = torch.linalg.norm(lazy.load() - data_eager) / torch.linalg.norm(data_eager)\n",
    "print(f\"Lazy vs eager relative error: {err.item():.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Streaming fast factorized backprojection\n",
    "\n",
    "`ffbp` slices the lazy source down its subaperture recursion tree and\n",
    "materializes data only at the leaves. With `stages=4` and 512 sweeps each leaf\n",
    "is 64 sweeps, so at no point is more than 1/8 of the data resident.\n",
    "The data is read once, in leaf-sized chunks, and the\n",
    "image matches the one formed from the eager tensor.\n",
    "\n",
    "The platform motion error is still uncompensated here, so the point targets\n",
    "are smeared. The autofocus at the end of the notebook fixes this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "img_lazy = torchbp.ops.ffbp(lazy, grid_polar, fc, r_res, pos_nominal, stages=4)\n",
    "print(f\"{stats['chunks']} chunks of at most {stats['max_chunk']} sweeps \"\n",
    "      f\"({stats['sweeps']} sweeps in total)\")\n",
    "\n",
    "img_eager = torchbp.ops.ffbp(data_eager, grid_polar, fc, r_res, pos_nominal, stages=4)\n",
    "err = torch.linalg.norm(img_lazy - img_eager) / torch.linalg.norm(img_eager)\n",
    "print(f\"Image relative error lazy vs eager: {err.item():.2e}\")\n",
    "\n",
    "img_db = 20 * torch.log10(torch.abs(img_lazy)).cpu()\n",
    "m = torch.max(img_db)\n",
    "extent = [*grid_polar[\"r\"], *grid_polar[\"theta\"]]\n",
    "plt.figure()\n",
    "plt.title(\"Uncorrected image from the lazy source\")\n",
    "plt.imshow(img_db.numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "           extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Caching: run the pipeline once, read many times\n",
    "\n",
    "A single `ffbp` call reads each sweep exactly once, so transforming on the fly\n",
    "costs nothing extra. Iterative consumers are different: `gpga`/`gpga_tde`\n",
    "re-read the data on every autofocus iteration, which would re-run the range\n",
    "compression each time. `CachedData` sits between the source and the consumer\n",
    "and stores the transformed sweeps, in one of three places:\n",
    "\n",
    "- `storage=\"ram\"` on the compute device (the default `storage_device`): once\n",
    "  filled, this is equivalent to an eagerly transformed tensor. Loads are\n",
    "  zero-copy views of the cache, so the lazy plumbing adds no per-load\n",
    "  overhead.\n",
    "- `storage=\"ram\", storage_device=\"cpu\"` with a CUDA compute device: for data\n",
    "  that fits in system RAM but not in VRAM. The cache stages in pinned CPU\n",
    "  memory and each load copies its chunk to the GPU at full transfer bandwidth.\n",
    "- `storage=\"disk\"`: a memory-mapped scratch file, for data larger than RAM.\n",
    "  RAM use stays bounded. The operating system's page cache keeps the hot\n",
    "  parts fast.\n",
    "\n",
    "The cache is keyed per sweep. The first load of a sweep runs the pipeline,\n",
    "later loads (also overlapping or out-of-order ones) read the cache."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "cached = torchbp.CachedData(lazy, storage=\"ram\")\n",
    "\n",
    "first = cached.load()   # runs the pipeline\n",
    "after_first = dict(stats)\n",
    "second = cached.load()  # reads the cache\n",
    "print(\"After first load: \", after_first)\n",
    "print(\"After second load:\", stats)\n",
    "print(\"Zero-copy views of the cache:\", first.data_ptr() == second.data_ptr())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "`storage=\"disk\"` behaves identically but backs the cache with a file. By\n",
    "default the file is a temporary one that is deleted when the cache is garbage\n",
    "collected (placed in `/var/tmp` if the system temp dir is a RAM-backed tmpfs,\n",
    "so a \"disk\" cache never silently consumes the RAM it is supposed to save), an\n",
    "explicit `path` is kept on disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "cache_path = os.path.join(tmpdir, \"compressed.torchbp-cache\")\n",
    "disk_cached = torchbp.CachedData(lazy, path=cache_path, storage=\"disk\")\n",
    "err = torch.linalg.norm(disk_cached.load() - data_eager) / torch.linalg.norm(data_eager)\n",
    "print(f\"Disk cache on disk: {os.path.getsize(cache_path) / 1e6:.1f} MB, \"\n",
    "      f\"relative error {err.item():.2e}\")\n",
    "del disk_cached"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## `fill()`: compress once, then free the raw data\n",
    "\n",
    "An eager script has a useful memory property: once range compression is done,\n",
    "the raw data can be freed. A lazy pipeline pins it instead, the cache holds\n",
    "the source, the source holds the raw array. `CachedData.fill()` restores the\n",
    "eager timeline: it populates the whole cache in bounded blocks and then drops\n",
    "the source, so the cache becomes the only backing store and the raw data is\n",
    "free to delete. Here we prove it by removing the raw file from disk. The\n",
    "cache keeps working without it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "cached = torchbp.CachedData(lazy, storage=\"ram\").fill()\n",
    "print(f\"fill() transformed each sweep exactly once: {stats['sweeps']} sweeps\")\n",
    "\n",
    "# The cache no longer references the source: delete the raw data entirely.\n",
    "del lazy, raw_mm\n",
    "gc.collect()\n",
    "os.remove(raw_path)\n",
    "\n",
    "err = torch.linalg.norm(cached.load() - data_eager) / torch.linalg.norm(data_eager)\n",
    "print(f\"Raw file deleted. Cache still serves the data, relative error {err.item():.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Choosing the storage tier automatically\n",
    "\n",
    "`torchbp.data.available_ram()` reports the available system RAM in bytes,\n",
    "counting reclaimable caches as available (Linux `MemAvailable`, Windows\n",
    "`GlobalMemoryStatusEx`). A processing script can use it, together with\n",
    "`torch.cuda.mem_get_info()`, to pick the tier: compute device when the\n",
    "compressed data fits comfortably, pinned CPU RAM when VRAM is short, disk when\n",
    "RAM is short too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def choose_storage(nbytes):\n",
    "    if nbytes > torchbp.data.available_ram() // 2:\n",
    "        return {\"storage\": \"disk\"}\n",
    "    if device == \"cuda\" and nbytes > torch.cuda.mem_get_info()[0] // 2:\n",
    "        return {\"storage\": \"ram\", \"storage_device\": \"cpu\"}\n",
    "    return {\"storage\": \"ram\"}\n",
    "\n",
    "nbytes = nsweeps * nfft * data_eager.element_size()\n",
    "print(f\"Available RAM: {torchbp.data.available_ram() / 2**30:.1f} GiB\")\n",
    "print(f\"This capture ({nbytes / 1e6:.1f} MB) ->\", choose_storage(nbytes))\n",
    "print(\"A 40 GB capture ->\", choose_storage(40 * 10**9))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Autofocus on the cached data\n",
    "\n",
    "`gpga_tde` recovers the platform motion error from the data. It re-reads the\n",
    "data every iteration. The image is re-formed with updated positions, and the\n",
    "estimation stage reads the data in bounded chunks, so it is exactly the\n",
    "consumer the cache exists for. With `algorithm=\"ffbp\"` the image formation\n",
    "inside the autofocus streams the data leaf by leaf as well.\n",
    "\n",
    "The transform counters stay at zero throughout: the pipeline never runs again,\n",
    "every read is served from the cache."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats.update(chunks=0, sweeps=0, max_chunk=0)\n",
    "img_focused, pos_solved = torchbp.autofocus.gpga_tde(\n",
    "    None, cached, pos_nominal, fc, r_res, grid_polar,\n",
    "    azimuth_divisions=4, range_divisions=4,\n",
    "    algorithm=\"ffbp\", image_opts={\"stages\": 4}, estimate_z=True)\n",
    "print(\"Pipeline runs during autofocus:\", stats[\"chunks\"])\n",
    "\n",
    "rms = lambda a, b: torch.sqrt(torch.mean(torch.square(a - b))).item()\n",
    "print(f\"RMS position error before autofocus: {1e3 * rms(pos_nominal, pos):.1f} mm\")\n",
    "print(f\"RMS position error after autofocus:  {1e3 * rms(pos_solved, pos):.1f} mm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_db = 20 * torch.log10(torch.abs(img_focused)).cpu()\n",
    "m = torch.max(img_db)\n",
    "plt.figure()\n",
    "plt.title(\"Autofocused image\")\n",
    "plt.imshow(img_db.numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "           extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "## Notes and limitations\n",
    "\n",
    "- Direct backprojection, `afbp` and `cfbp` accept a lazy source but load the\n",
    "  full extent at entry, only `ffbp` and the `gpga`/`gpga_tde` autofocus are\n",
    "  memory efficient with one.\n",
    "- For BP, manual image tiling and data partitioning can be used to generate a\n",
    "  huge image, but at this time there is not support for it in torchbp.\n",
    "- Gradients do not flow to the data through a lazy source\n",
    "  (`requires_grad` is always False). Gradients with respect to positions are\n",
    "  unaffected.\n",
    "- `CachedData` supports only 2D `[nsweeps, samples]` sources.\n",
    "- For a real out-of-core run, combine the pieces as in this notebook:\n",
    "  a `MemmapData`/`CallbackData` over the capture file with the range\n",
    "  compression as `transform`, wrapped in a `CachedData` with the tier chosen\n",
    "  from `available_ram()` and filled with `fill()` before autofocus. See\n",
    "  `examples/sar_process_safetensor_gpga.py` for a full script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}