{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-0001-4a10-9f00-000000000001",
   "metadata": {},
   "source": [
    "# Streaming POD vector space tutorial\n",
    "\n",
    "In this tutorial you will learn:\n",
    "- How to construct a POD vector space from snapshots that are loaded *on demand* through a `snapshot_loader`, so that the full snapshot matrix is never held in memory at once\n",
    "- The roles of `block_size`, `max_basis_dimension`, and the basis `truncater`\n",
    "- How to use the streaming vector space with **no scaling and no orthogonalizer** (the defaults)\n",
    "- How to use the streaming vector space **with a scaler and an orthogonalizer**, and why the two are paired"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-0002-4a10-9f00-000000000002",
   "metadata": {},
   "source": [
    "## The math behind streaming POD\n",
    "\n",
    "Let the snapshot matrix be $X \\in \\mathbb{R}^{N \\times M}$, where each column is a state snapshot. For the tensor data used here, $N = n_{\\text{var}} \\, n_x$ and $M$ is the number of snapshots. Proper orthogonal decomposition (POD) seeks a rank-$k$ orthonormal basis $\\Phi \\in \\mathbb{R}^{N \\times k}$ that best approximates the snapshots,\n",
    "\n",
    "$$\n",
    "\\min_{\\Phi^\\top \\Phi = I_k} \\; \\lVert X - \\Phi \\Phi^\\top X \\rVert_F ,\n",
    "$$\n",
    "\n",
    "whose solution is the leading $k$ left singular vectors of $X = U \\Sigma V^\\top$.\n",
    "\n",
    "For large datasets, $X$ may not fit in memory. **Streaming POD** never materializes all of $X$: it reads one column block $X_{[s:e]}$ at a time through a `snapshot_loader` callable, so the memory footprint scales with the block size, not with the total number of snapshots $M$.\n",
    "\n",
    "`romtools` uses a randomized two-pass algorithm. Given a maximum candidate-basis dimension $r_{\\max}$, draw a random matrix $\\Omega \\in \\mathbb{R}^{M \\times r_{\\max}}$:\n",
    "\n",
    "- **Pass 1** accumulate the range sketch one block at a time,\n",
    "\n",
    "$$\n",
    "Y = X \\Omega = \\sum_{\\text{blocks } [s:e]} X_{[s:e]} \\, \\Omega_{[s:e]} ,\n",
    "$$\n",
    "\n",
    "  then orthonormalize $Y$ to obtain $Q$ (with $Q^\\top Q = I$) that approximately spans the range of $X$.\n",
    "- **Pass 2** form the small core matrix, again one block at a time,\n",
    "\n",
    "$$\n",
    "B = Q^\\top X = \\sum_{\\text{blocks } [s:e]} Q^\\top X_{[s:e]} ,\n",
    "$$\n",
    "\n",
    "  compute its (inexpensive) SVD $B = \\tilde{U}\\, \\Sigma\\, V^\\top$ and form the candidate POD basis $\\Phi_{\\mathrm{candidate}} = Q \\tilde{U}$. The supplied truncater then selects the retained columns.\n",
    "\n",
    "Because $Q$ and $\\tilde{U}$ both have orthonormal columns, the candidate basis is orthonormal by construction. For fixed-rank truncation, choosing `max_basis_dimension` larger than the retained dimension provides headroom that improves the randomized range approximation. An energy-based truncater instead chooses the smallest available dimension that reaches its requested fraction of the complete snapshot energy; it raises an error if the maximum candidate dimension is too small. The snapshots are read **twice** (one pass each); a data-derived scaler (see below) adds **one** more pass to compute its scalings.\n",
    "\n",
    "This tutorial runs in serial. For very large problems the same construction is row-distributed across MPI ranks by passing a communicator `comm` and an SVD functor `svdFnc=vector_space.utils.SvdMethodOfSnapshots(comm)`; each rank's loader then returns only its local rows for the requested snapshot range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-0003-4a10-9f00-000000000003",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, let's import the relevant modules:\n",
    "import inspect\n",
    "import romtools\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from romtools import vector_space\n",
    "\n",
    "streaming_pod_parameters = inspect.signature(\n",
    "    vector_space.VectorSpaceFromStreamingPOD\n",
    ").parameters\n",
    "if 'max_basis_dimension' not in streaming_pod_parameters:\n",
    "    raise RuntimeError(\n",
    "        'This notebook imported an outdated romtools package from '\n",
    "        f'{romtools.__file__}. Reinstall this checkout and restart the kernel.'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-0004-4a10-9f00-000000000004",
   "metadata": {},
   "source": [
    "## Loading data through a `snapshot_loader`\n",
    "\n",
    "The streaming vector space does not take a snapshot array directly. Instead it takes a **`snapshot_loader`**: a callable `loader(start, end)` that returns the block of snapshots for the half-open range `[start, end)`. The returned array must be three-dimensional with shape `(n_var, n_dofs, end - start)` (the last axis is the snapshot axis), and the leading axes must be the same for every call.\n",
    "\n",
    "In a real large-scale problem the loader would read each block from disk so that the full matrix is never resident in memory. Any callable conforming to `romtools.vector_space.utils.SnapshotLoader` can provide those blocks. Here, for illustration, we simply slice the in-memory 1D Euler snapshots used throughout these tutorials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-0005-4a10-9f00-000000000005",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load pre-computed snapshots of the 1D Euler equations (from pressio-demo-apps).\n",
    "snapshots = np.load('snapshots.npz')['snapshots']\n",
    "\n",
    "# The snapshots are in tensor form:\n",
    "n_vars, nx, nt = snapshots.shape\n",
    "print('snapshot tensor shape (n_vars, nx, nt):', snapshots.shape)\n",
    "\n",
    "# A snapshot_loader returns the 3D block for the half-open range [start, end).\n",
    "# In practice this would stream each block from disk; here we slice the\n",
    "# in-memory array purely for illustration.\n",
    "def my_snapshot_loader(start, end):\n",
    "    return snapshots[..., start:end]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-0006-4a10-9f00-000000000006",
   "metadata": {},
   "source": [
    "## Case 1: streaming POD with no scaling and no orthogonalizer\n",
    "\n",
    "In the simplest case we pass the loader, `block_size`, `max_basis_dimension`, and a truncater. Here `BasisSizeTruncater(basis_dimension)` retains 20 modes from a 25-mode randomized candidate basis. The five-mode headroom improves the randomized range approximation. The scaler and orthogonalizer default to no-ops.\n",
    "\n",
    "When no shifter is supplied, the shift (affine offset) vector is zero. Case 2 demonstrates a shifter initialized directly from streamed snapshot blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-0007-4a10-9f00-000000000007",
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size         = 8    # maximum number of snapshots held in memory at once\n",
    "basis_dimension    = 20   # number of POD modes to retain\n",
    "max_basis_dimension = 25  # size of the randomized candidate basis\n",
    "truncater = vector_space.utils.BasisSizeTruncater(basis_dimension)\n",
    "\n",
    "my_streaming_vector_space = vector_space.VectorSpaceFromStreamingPOD(\n",
    "    snapshot_loader=my_snapshot_loader,\n",
    "    block_size=block_size,\n",
    "    n_snapshots=nt,\n",
    "    max_basis_dimension=max_basis_dimension,\n",
    "    truncater=truncater)\n",
    "\n",
    "# We can view the retained basis and the full candidate spectrum:\n",
    "basis = my_streaming_vector_space.get_basis()\n",
    "print('The dimension of the vector space is', my_streaming_vector_space.extents())\n",
    "print('Candidate singular values:', my_streaming_vector_space.get_singular_values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-0008-4a10-9f00-000000000008",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can look at the density component of the first basis vector:\n",
    "plt.plot(basis[0, :, 0])\n",
    "plt.xlabel(r'Index')\n",
    "plt.ylabel(r'$\\rho$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-0009-4a10-9f00-000000000009",
   "metadata": {},
   "outputs": [],
   "source": [
    "# With no orthogonalizer, the streaming basis is already orthonormal in the\n",
    "# Euclidean inner product (Phi^T Phi = I) by construction:\n",
    "is_identity = np.einsum('ijk,ijl->kl', basis, basis)\n",
    "assert np.allclose(is_identity, np.eye(basis.shape[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-000a-4a10-9f00-00000000000a",
   "metadata": {},
   "source": [
    "## Case 2: streaming POD with shifting, scaling, and orthogonalization\n",
    "\n",
    "An average **streaming shifter** first computes the temporal mean in one loader pass and subtracts it from every subsequent snapshot block. The 1D Euler state also has variables of very different magnitudes (density, momentum, and energy), so a **scaler** non-dimensionalizes the shifted data before the SVD. Here we use a per-variable `VariableScaler(\"variance\")`.\n",
    "\n",
    "Scaling changes the basis after the SVD (`post_scale`), which in general destroys the Euclidean orthonormality of the modes. We therefore pair the scaler with an **orthogonalizer** to restore $\\Phi^\\top \\Phi = I$. `EuclideanL2Orthogonalizer` re-orthonormalizes in the standard $L^2$ inner product; if you need a physics-based inner product (e.g. weighting by cell volumes or quadrature weights $w$), use `EuclideanVectorWeightedL2Orthogonalizer(w)`, which enforces $\\Phi^\\top \\mathrm{diag}(w)\\, \\Phi = I$.\n",
    "\n",
    "Initialization follows the same order as standard POD: shift first, then scale. The average shifter and `VariableScaler` each perform one initialization pass before the two POD passes, for four passes total. A first-snapshot streaming shifter instead needs only `loader(0, 1)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-000b-4a10-9f00-00000000000b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an average shifter, a per-variable scaler, and an orthogonalizer:\n",
    "my_shifter        = vector_space.utils.create_streaming_average_shifter()\n",
    "my_scaler         = vector_space.utils.VariableScaler('variance')\n",
    "my_orthogonalizer = vector_space.utils.EuclideanL2Orthogonalizer()\n",
    "\n",
    "my_scaled_vector_space = vector_space.VectorSpaceFromStreamingPOD(\n",
    "    snapshot_loader=my_snapshot_loader,\n",
    "    block_size=block_size,\n",
    "    n_snapshots=nt,\n",
    "    max_basis_dimension=max_basis_dimension,\n",
    "    truncater=vector_space.utils.BasisSizeTruncater(basis_dimension),\n",
    "    shifter=my_shifter,\n",
    "    scaler=my_scaler,\n",
    "    orthogonalizer=my_orthogonalizer)\n",
    "\n",
    "scaled_basis = my_scaled_vector_space.get_basis()\n",
    "print('The dimension of the scaled vector space is', my_scaled_vector_space.extents())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-000c-4a10-9f00-00000000000c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Even though scaling was applied, the orthogonalizer guarantees the basis is\n",
    "# orthonormal in the Euclidean inner product:\n",
    "is_identity = np.einsum('ijk,ijl->kl', scaled_basis, scaled_basis)\n",
    "assert np.allclose(is_identity, np.eye(scaled_basis.shape[-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1c2d3-000d-4a10-9f00-00000000000d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare the density component of the first mode with and without scaling:\n",
    "plt.plot(basis[0, :, 0], label='no scaling')\n",
    "plt.plot(scaled_basis[0, :, 0], '--', label='variance scaling')\n",
    "plt.xlabel(r'Index')\n",
    "plt.ylabel(r'$\\rho$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1c2d3-000e-4a10-9f00-00000000000e",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- `VectorSpaceFromStreamingPOD` builds a POD basis while holding only one block of snapshots in memory at a time, loading data on demand through a `snapshot_loader`.\n",
    "- `block_size` trades memory footprint against the number of loader calls; `max_basis_dimension` sets the randomized candidate rank, and a truncater selects the retained modes. Leaving headroom above a fixed retained rank improves randomized accuracy.\n",
    "- With no scaler and no orthogonalizer, the basis is orthonormal by construction.\n",
    "- Streaming average and first-snapshot shifters initialize their affine offsets directly from loader blocks; shifting occurs before scaling, as in standard POD.\n",
    "- A scaler (e.g. `VariableScaler`) should be paired with an orthogonalizer to restore orthonormality after scaling.\n",
    "\n",
    "See the API documentation for details:\n",
    "- [`VectorSpaceFromStreamingPOD`](https://pressio.github.io/rom-tools-and-workflows/romtools/vector_space.html)\n",
    "- [scalers](https://pressio.github.io/rom-tools-and-workflows/romtools/vector_space/utils/scaler.html) and [orthogonalizers](https://pressio.github.io/rom-tools-and-workflows/romtools/vector_space/utils/orthogonalizer.html)"
   ]
  }
 ],
 "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}