{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# FEEC data structures: StencilVector and StencilMatrix\n",
    "\n",
    "This tutorial is the third FEEC tutorial in the developer guide. It focuses on the core linear-algebra data structures used for FEEC coefficients and operators in Struphy.\n",
    "\n",
    "If you are new to Struphy FEEC, start with:\n",
    "1. [FEEC basics](dev_tutorial_feec_basics.ipynb)\n",
    "2. [FEEC boundary conditions](dev_tutorial_feec_bcs.ipynb)\n",
    "\n",
    "Here we drill down into:\n",
    "1. `StencilVector`: storage of scalar FE coefficients\n",
    "2. `StencilMatrix`: storage of stencil-like sparse operators\n",
    "\n",
    "Particle data structures are intentionally not covered in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## 1. Setup with the current Derham API\n",
    "\n",
    "We use the current API based on `TensorProductGrid` and `DerhamOptions`, as introduced in the [FEEC basics tutorial](dev_tutorial_feec_basics.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from feectools.linalg.stencil import StencilMatrix, StencilVector\n",
    "\n",
    "from struphy.feec.psydac_derham import Derham\n",
    "from struphy.io.options import DerhamOptions\n",
    "from struphy.topology.grids import TensorProductGrid\n",
    "\n",
    "grid = TensorProductGrid(num_elements=(8, 6, 4))\n",
    "derham = Derham(grid, DerhamOptions())\n",
    "\n",
    "print('Discrete spaces available in coeff_spaces:')\n",
    "for key in ['H1', 'Hcurl', 'Hdiv', 'L2']:\n",
    "    print(f'  {key:5s} -> {type(derham.coeff_spaces[key]).__name__}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## 2. StencilVector in practice\n",
    "\n",
    "A `StencilVector` stores scalar FE coefficients in a distributed layout with ghost regions.\n",
    "\n",
    "In FEEC, 0-form (`H1`) and 3-form (`L2`) coefficients are `StencilVector` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option A: build a zero vector from a coefficient space\n",
    "x = derham.coeff_spaces['H1'].zeros()\n",
    "print('x type from zeros():', type(x))\n",
    "\n",
    "# Option B: project a function with the commuting projector P0\n",
    "f = lambda e1, e2, e3: np.cos(2 * np.pi * e1) + 0.25 * np.sin(2 * np.pi * e2)\n",
    "xh = derham.P0(f)\n",
    "print('xh type from P0:', type(xh))\n",
    "\n",
    "assert isinstance(xh, StencilVector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Starts (global index owned by this rank):', xh.starts)\n",
    "print('Ends   (global index owned by this rank):', xh.ends)\n",
    "print('Pads   (ghost widths per direction):', xh.pads)\n",
    "print('shape  (logical tensor shape):', xh.shape)\n",
    "print('_data shape (includes ghost layers):', xh._data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### 2.1 Global indexing vs local storage\n",
    "\n",
    "Using `xh[i, j, k]` addresses coefficients by global FE indices.\n",
    "\n",
    "Using `xh._data[...]` addresses the local storage array directly, where ghost layers are explicit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill owned coefficients with a recognizable pattern through global indexing\n",
    "s = xh.starts\n",
    "e = xh.ends\n",
    "\n",
    "for i in range(s[0], e[0] + 1):\n",
    "    for j in range(s[1], e[1] + 1):\n",
    "        for k in range(s[2], e[2] + 1):\n",
    "            xh[i, j, k] = i + 10 * j + 100 * k\n",
    "\n",
    "print('Sample values with global indexing:')\n",
    "print('xh[s0, s1, s2] =', xh[s[0], s[1], s[2]])\n",
    "print('xh[e0, e1, e2] =', xh[e[0], e[1], e[2]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare compact coefficient view with local storage\n",
    "flat = xh.toarray()\n",
    "print('xh.toarray() shape:', flat.shape)\n",
    "print('First 8 compact coefficients:', flat[:8])\n",
    "\n",
    "print('Direct local array shape (with ghosts):', xh._data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "### 2.2 Relation to vector-valued FE spaces\n",
    "\n",
    "For `Hcurl` and `Hdiv`, coefficient vectors are block-structured (one block per component).\n",
    "Each block is itself stencil-based.\n",
    "\n",
    "For projector calls like `P1`, vector-valued arguments should be passed as a **3-tuple of scalar callables** `(f1, f2, f3)`, one callable per component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "g1 = lambda e1, e2, e3: e1 + e2\n",
    "g2 = lambda e1, e2, e3: e2 + e3\n",
    "g3 = lambda e1, e2, e3: e3 + e1\n",
    "u_hcurl = derham.P1((g1, g2, g3))\n",
    "\n",
    "print('Type of Hcurl coefficients:', type(u_hcurl))\n",
    "print('Number of blocks/components:', len(u_hcurl.blocks))\n",
    "print('Type of first block:', type(u_hcurl.blocks[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## 3. StencilMatrix in practice\n",
    "\n",
    "A `StencilMatrix` represents a linear operator between stencil vector spaces.\n",
    "It stores row ownership plus a compact set of local diagonals (stencil offsets), not a dense matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "V0 = derham.coeff_spaces['H1']\n",
    "A = StencilMatrix(V0, V0)\n",
    "\n",
    "print('A type:', type(A))\n",
    "print('A codomain starts:', A.codomain.starts)\n",
    "print('A codomain ends  :', A.codomain.ends)\n",
    "print('A pads:', A.pads)\n",
    "print('A._data shape:', A._data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### 3.1 Understanding matrix indexing\n",
    "\n",
    "In 3D, indexing is `A[i, j, k, s1, s2, s3]`:\n",
    "- `(i, j, k)` are global row indices owned by the current rank\n",
    "- `(s1, s2, s3)` are stencil offsets (diagonal at `(0, 0, 0)`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = A.codomain.starts\n",
    "e = A.codomain.ends\n",
    "\n",
    "# Build a simple 1D second-derivative-like stencil in eta1 direction\n",
    "for i in range(s[0], e[0] + 1):\n",
    "    for j in range(s[1], e[1] + 1):\n",
    "        for k in range(s[2], e[2] + 1):\n",
    "            A[i, j, k, 0, 0, 0] = 2.0\n",
    "            A[i, j, k, -1, 0, 0] = -1.0\n",
    "            A[i, j, k, 1, 0, 0] = -1.0\n",
    "\n",
    "print('Stencil entries at one sample row:')\n",
    "print('diag      =', A[s[0], s[1], s[2], 0, 0, 0])\n",
    "print('left diag =', A[s[0], s[1], s[2], -1, 0, 0])\n",
    "print('rightdiag =', A[s[0], s[1], s[2], 1, 0, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply matrix to a vector and inspect the result\n",
    "y = A.dot(xh)\n",
    "print('Type of y = A.dot(xh):', type(y))\n",
    "print('y.toarray() shape:', y.toarray().shape)\n",
    "print('First 8 entries of y.toarray():', y.toarray()[:8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## 4. Practical developer notes\n",
    "\n",
    "1. `toarray()` is ideal for debugging small problems, but avoid it in production-scale runs because it materializes owned coefficients as a dense 1D array.\n",
    "2. Prefer operations at the vector/matrix level (`A.dot(x)`, projector calls, operator composition) over manual loops in performance-critical code.\n",
    "3. In parallel runs, ghost regions and ownership become essential for correctness. This notebook keeps the focus on conceptual usage, while MPI details can be explored in specialized FEEC tests/examples.\n",
    "\n",
    "For context on FEEC spaces and projectors, revisit the [basics tutorial](dev_tutorial_feec_basics.ipynb). For boundary-condition effects on projected fields, see the [boundary-conditions tutorial](dev_tutorial_feec_bcs.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (3.12.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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}