{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c12c99d2",
   "metadata": {},
   "source": [
    "# L3A -- Raviart--Thomas elements for the mixed Laplacian\n",
    "\n",
    "Companion notebook to Lecture 3 (cells are cited on the slides by their\n",
    "`Cell N` labels). Tags: `[LECTURE]` cells are run live; `[SELF-STUDY Sk]`\n",
    "and `[ADD-BACK n]` cells are for after the lecture or for optional live\n",
    "insertion.\n",
    "\n",
    "Setting: the mixed Laplacian of Lectures 1--2 on the unit square,\n",
    "$u=\\sin(\\pi x)\\sin(\\pi y)$, $\\kappa=1$ unless stated. The boundary\n",
    "condition $u=0$ is **natural** in the mixed form: none of the mixed solves\n",
    "below contains a `DirichletBC` (Lecture 3, slide 6).\n",
    "\n",
    "> **Warning (Firedrake / FEniCS indexing).** Firedrake indexes the\n",
    "> Raviart--Thomas family shifted by one relative to the course notation\n",
    "> (Boffi--Brezzi--Fortin): `FunctionSpace(mesh, \"RT\", 1)` is $RT_0$.\n",
    "> The $BDM$ indices agree: `(\"BDM\", 1)` is $BDM_1$. Cross-check\n",
    "> element definitions on [DefElement](https://defelement.org).\n",
    "\n",
    "Provenance: every expected number quoted in the code comments is pinned by\n",
    "the pure-NumPy reference implementation `verify_rt0_2d.py` distributed with\n",
    "the course materials; the printouts of this notebook must match those values\n",
    "to the digits shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c8b4612",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1 [SETUP] -- Firedrake import guard (FEM on Colab).\n",
    "# On a local Firedrake installation this cell only imports and reports.\n",
    "try:\n",
    "    import firedrake\n",
    "except ImportError:\n",
    "    try:\n",
    "        import google.colab  # noqa: F401\n",
    "        !wget \"https://fem-on-colab.github.io/releases/firedrake-install-release-real.sh\" -O \"/tmp/firedrake-install.sh\" && bash \"/tmp/firedrake-install.sh\"\n",
    "        import firedrake\n",
    "    except ImportError:\n",
    "        raise RuntimeError(\n",
    "            \"Firedrake is not available. Install it locally or run this \"\n",
    "            \"notebook on Google Colab (the cell above then installs it).\")\n",
    "try:\n",
    "    from importlib.metadata import version\n",
    "    print(\"Firedrake version:\", version(\"firedrake\"))\n",
    "except Exception:\n",
    "    print(\"Firedrake imported (version string not available).\")\n",
    "print(\"Firedrake location:\", firedrake.__file__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f39cff01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2 [LECTURE] -- imports, data, mixed solver, helpers.\n",
    "import warnings\n",
    "# NumPy >= 2.5 emits a DeprecationWarning from petsc4py/Firedrake internals\n",
    "# (assignment to ndarray.shape) during dense matrix extraction; it is\n",
    "# harmless here and silenced by message (technique from the L2 notebook).\n",
    "warnings.filterwarnings(\n",
    "    \"ignore\", message=\".*ndarray.shape.*\", category=DeprecationWarning)\n",
    "\n",
    "import numpy as np\n",
    "from firedrake import *\n",
    "\n",
    "# tsfc warns that the estimated quadrature degree (12, from the sine\n",
    "# products in the error integrands) exceeds the polynomial degrees. The\n",
    "# estimate is correct and the cost is negligible. We silence the logger\n",
    "# rather than pin per-form quadrature degrees: the conservation identity of\n",
    "# cell 7 holds at machine precision only because the residual form and the\n",
    "# solve use the SAME quadrature rule, so the automatic estimate stays in\n",
    "# charge everywhere.\n",
    "# Caution: \"from firedrake import *\" exports firedrake.logging, shadowing\n",
    "# the standard-library module of the same name; the alias avoids relying on\n",
    "# import order.\n",
    "import logging as pylogging\n",
    "pylogging.getLogger(\"tsfc\").setLevel(pylogging.ERROR)\n",
    "\n",
    "def exact(mesh):\n",
    "    x, y = SpatialCoordinate(mesh)\n",
    "    u  = sin(pi*x)*sin(pi*y)\n",
    "    sig = as_vector([pi*cos(pi*x)*sin(pi*y), pi*sin(pi*x)*cos(pi*y)])\n",
    "    f  = 2*pi**2*sin(pi*x)*sin(pi*y)\n",
    "    return u, sig, f\n",
    "\n",
    "def solve_mixed(N, family=\"RT\", degree=1, kappa=None):\n",
    "    \"\"\"Mixed Laplacian on the unit square; natural condition u = 0 on the\n",
    "    whole boundary (there is no DirichletBC in this function: the Dirichlet\n",
    "    datum is natural in the mixed form -- Lecture 3, slide 6).\n",
    "    family/degree: (\"RT\", 1) is RT_0 in course notation; (\"BDM\", 1) is BDM_1.\n",
    "    \"\"\"\n",
    "    mesh = UnitSquareMesh(N, N)\n",
    "    V = FunctionSpace(mesh, family, degree)\n",
    "    Q = FunctionSpace(mesh, \"DG\", 0)\n",
    "    W = V * Q\n",
    "    u_ex, sig_ex, f = exact(mesh)\n",
    "    kinv = 1.0 if kappa is None else 1.0/kappa(mesh)\n",
    "    sigma, u = TrialFunctions(W)\n",
    "    tau, v = TestFunctions(W)\n",
    "    a = (kinv*inner(sigma, tau) + div(tau)*u + div(sigma)*v)*dx\n",
    "    L = -f*v*dx\n",
    "    w = Function(W)\n",
    "    solve(a == L, w, solver_parameters={\n",
    "        \"ksp_type\": \"preonly\", \"pc_type\": \"lu\",\n",
    "        \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "    return w, mesh, (u_ex, sig_ex, f)\n",
    "\n",
    "def l2err(expr):\n",
    "    return sqrt(assemble(inner(expr, expr)*dx))\n",
    "\n",
    "def dense(form):\n",
    "    \"\"\"Dense NumPy array of an assembled bilinear form (technique of the\n",
    "    L2 notebook, cell 2). Use on small meshes only.\"\"\"\n",
    "    A = assemble(form).petscmat\n",
    "    return A.convert(\"dense\").getDenseArray().copy()\n",
    "\n",
    "# smoke test\n",
    "w, mesh, (u_ex, sig_ex, f) = solve_mixed(8)\n",
    "sig_h, u_h = w.subfunctions\n",
    "print(\"smoke test, N=8:  ||sigma - sigma_h||_0 =\", f\"{l2err(sig_ex - sig_h):.4e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ecbe8d2",
   "metadata": {},
   "source": [
    "## Demo 1a -- convergence of $RT_0$--$P_0$ (slide 23)\n",
    "\n",
    "Rates for the flux, its divergence, the scalar, and the distance to the elementwise means $P_hu$. The last column is the superconvergence of Theorem 26."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0eafc39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4 [LECTURE] -- Demo 1a: convergence of RT0-P0.\n",
    "# Expected rates (verify_rt0_2d.py): 1.00 / 1.00 / 1.00 / 1.99, and 2.00\n",
    "# for the postprocessed scalar of notebook L3B.\n",
    "rows, Ns = [], [4, 8, 16, 32]\n",
    "for N in Ns:\n",
    "    w, mesh, (u_ex, sig_ex, f) = solve_mixed(N)\n",
    "    sig_h, u_h = w.subfunctions\n",
    "    Q = u_h.function_space()\n",
    "    Phu = Function(Q).project(u_ex)          # elementwise means of u\n",
    "    rows.append([l2err(sig_ex - sig_h),\n",
    "                 l2err(div(sig_h) + f),      # = ||f - P_h f|| (Prop. 24)\n",
    "                 l2err(u_ex - u_h),\n",
    "                 l2err(Phu - u_h)])\n",
    "rows = np.array(rows)\n",
    "hdr = [\"sigma\", \"div sigma\", \"u\", \"P_h u - u_h\"]\n",
    "print(\"N     \" + \"\".join(f\"{h:>14s}\" for h in hdr))\n",
    "for N, r in zip(Ns, rows):\n",
    "    print(f\"{N:<5d}\" + \"\".join(f\"{e:14.4e}\" for e in r))\n",
    "rates = np.log2(rows[:-1]/rows[1:])\n",
    "for r in rates:\n",
    "    print(\"rate \" + \"\".join(f\"{x:14.2f}\" for x in r))\n",
    "print(\"\\nThe divergence error equals ||f - P_h f|| exactly (Theorem 25):\")\n",
    "print(\"compare with the identical column for BDM_1 in cell 11.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8af1383b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5 [LECTURE] -- Demo 1b: the constants, measured.\n",
    "# beta_h from the pencil  B G_V^{-1} B^T q = beta_h^2 M_Q q  (Lecture 2,\n",
    "# Exercise 1(a)); alpha_h as the Rayleigh quotient of a on ker B.\n",
    "# Expected (verify_rt0_2d.py): beta_h = 0.97597 -> 0.97560 from above,\n",
    "# limit sqrt(2 pi^2/(1+2 pi^2)) = 0.9755932; alpha_h = 1 exactly.\n",
    "beta = float(np.sqrt(2*np.pi**2/(1 + 2*np.pi**2)))\n",
    "print(f\"continuous beta = sqrt(2 pi^2/(1+2 pi^2)) = {beta:.7f}\\n\")\n",
    "for N in [4, 8, 12, 16]:\n",
    "    mesh = UnitSquareMesh(N, N)\n",
    "    V = FunctionSpace(mesh, \"RT\", 1)     # RT_0 (indexing warning, header)\n",
    "    Q = FunctionSpace(mesh, \"DG\", 0)\n",
    "    sigma, tau = TrialFunction(V), TestFunction(V)\n",
    "    u, v = TrialFunction(Q), TestFunction(Q)\n",
    "    Bm = dense(div(TrialFunction(V))*TestFunction(Q)*dx)\n",
    "    Gv = dense((inner(sigma, tau) + div(sigma)*div(tau))*dx)\n",
    "    Mq = dense(u*v*dx)\n",
    "    d = 1.0/np.sqrt(np.diag(Mq))         # M_Q is diagonal for DG0\n",
    "    S = Bm @ np.linalg.solve(Gv, Bm.T)\n",
    "    lam = np.linalg.eigvalsh(d[:, None]*S*d[None, :])\n",
    "    # alpha_h on ker B (dense SVD; small N only)\n",
    "    from scipy.linalg import eigh\n",
    "    U_, sv, Vt = np.linalg.svd(Bm)\n",
    "    Z = Vt[Bm.shape[0]:].T\n",
    "    Am = dense(inner(sigma, tau)*dx)\n",
    "    r = eigh(Z.T @ Am @ Z, Z.T @ Gv @ Z, eigvals_only=True)\n",
    "    print(f\"N = {N:<3d}  beta_h = {np.sqrt(lam[0]):.7f}   \"\n",
    "          f\"alpha_h on K_h in [{r.min():.6f}, {r.max():.6f}]   \"\n",
    "          f\"dim K_h = {Z.shape[1]} (= N^2+2N = {N*N+2*N})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4309a2db",
   "metadata": {},
   "source": [
    "## Demo 1c -- local conservation (slides 20 and 25)\n",
    "\n",
    "The balance defect $\\int_{\\partial K}\\boldsymbol q_h\\cdot\\boldsymbol n\\,ds-\\int_K f\\,dx$ per element, for the mixed solution and for the $P_1$ Lagrange solution of the same problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67c4d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7 [LECTURE] -- Demo 1c: elementwise balance, RT0-P0 vs P1 Lagrange.\n",
    "# The balance functional per element K:  r_K = int_dK q_h.n ds - int_K f dx,\n",
    "# with q_h the (Darcy) flux of each method. For a DG0 test function v the\n",
    "# facet form below assembles exactly (r_K)_K.\n",
    "def cell_balance(q, f, mesh):\n",
    "    Q0 = FunctionSpace(mesh, \"DG\", 0)\n",
    "    v = TestFunction(Q0)\n",
    "    n = FacetNormal(mesh)\n",
    "    form = (v('+')*dot(q('+'), n('+')) + v('-')*dot(q('-'), n('-')))*dS \\\n",
    "         + v*dot(q, n)*ds - v*f*dx\n",
    "    return assemble(form).dat.data_ro.copy()\n",
    "\n",
    "N = 32\n",
    "# mixed\n",
    "w, mesh, (u_ex, sig_ex, f) = solve_mixed(N)\n",
    "sig_h, u_h = w.subfunctions\n",
    "r_mixed = cell_balance(-sig_h, f, mesh)\n",
    "# primal P1 (essential boundary condition; this code does use DirichletBC)\n",
    "Vp = FunctionSpace(mesh, \"CG\", 1)\n",
    "up, vp = TrialFunction(Vp), TestFunction(Vp)\n",
    "u_p = Function(Vp)\n",
    "solve(inner(grad(up), grad(vp))*dx == f*vp*dx, u_p,\n",
    "      bcs=DirichletBC(Vp, 0, \"on_boundary\"))\n",
    "r_p1 = cell_balance(-grad(u_p), f, mesh)\n",
    "# the element source scale\n",
    "Q0 = FunctionSpace(mesh, \"DG\", 0)\n",
    "intf = assemble(TestFunction(Q0)*f*dx).dat.data_ro\n",
    "print(f\"max_K |int_K f|                    : {np.abs(intf).max():.3e}\")\n",
    "print(f\"RT0-P0    max_K |balance defect|   : {np.abs(r_mixed).max():.3e}\")\n",
    "print(f\"P1 CG     max_K |balance defect|   : {np.abs(r_p1).max():.3e}\")\n",
    "print(\"\\nFor P1 the flux -grad u_h is elementwise constant, so its net flux\")\n",
    "print(\"through every closed element boundary is zero: the defect equals\")\n",
    "print(\"|int_K f| on every element (Lecture 3, slide 20). A weaker balance\")\n",
    "print(\"on vertex patches survives; no single-valued flux exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "633b278b",
   "metadata": {},
   "source": [
    "## Demo 1d -- a discontinuous coefficient (slide 25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f23c745",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 9 [LECTURE] -- Demo 1d: a layered coefficient.\n",
    "# kappa = 1 above the line y = 1/2 and 10^-3 below; f = 1. Conservation is\n",
    "# unaffected: the constraint rows do not involve kappa.\n",
    "import matplotlib.pyplot as plt\n",
    "from firedrake.pyplot import tripcolor, quiver\n",
    "\n",
    "def kappa_layered(mesh):\n",
    "    x, y = SpatialCoordinate(mesh)\n",
    "    return conditional(gt(y, 0.5), 1.0, 1.0e-3)\n",
    "\n",
    "N = 32\n",
    "mesh = UnitSquareMesh(N, N)\n",
    "V = FunctionSpace(mesh, \"RT\", 1)\n",
    "Q = FunctionSpace(mesh, \"DG\", 0)\n",
    "W = V * Q\n",
    "x, y = SpatialCoordinate(mesh)\n",
    "kinv = 1.0/kappa_layered(mesh)\n",
    "sigma, u = TrialFunctions(W)\n",
    "tau, v = TestFunctions(W)\n",
    "a = (kinv*inner(sigma, tau) + div(tau)*u + div(sigma)*v)*dx\n",
    "w = Function(W)\n",
    "solve(a == -1.0*v*dx, w, solver_parameters={\n",
    "    \"ksp_type\": \"preonly\", \"pc_type\": \"lu\",\n",
    "    \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "sig_h, u_h = w.subfunctions\n",
    "r = cell_balance(-sig_h, Constant(1.0), mesh)\n",
    "print(f\"layered kappa: max_K |balance defect| = {np.abs(r).max():.3e}\")\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(10, 4))\n",
    "c0 = tripcolor(u_h, axes=axs[0])\n",
    "axs[0].axhline(0.5, color=\"w\", ls=\"--\", lw=1)\n",
    "axs[0].set_title(\"$u_h$ (layered $\\\\kappa$)\")\n",
    "fig.colorbar(c0, ax=axs[0])\n",
    "q_h = Function(VectorFunctionSpace(mesh, \"DG\", 0)).project(-sig_h)\n",
    "c1 = quiver(q_h, axes=axs[1])\n",
    "axs[1].axhline(0.5, color=\"k\", ls=\"--\", lw=1)\n",
    "axs[1].set_title(\"$q_h = -\\\\sigma_h$: normal flux continuous\")\n",
    "for a_ in axs:\n",
    "    a_.set_aspect(\"equal\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80b3579f",
   "metadata": {},
   "source": [
    "## Self-study cells\n",
    "\n",
    "S1 is add-back 2 (the $BDM_1$ profile, backup B3); S2 examines conformity on the facets; S3 is add-back 4 (patch balances, backup B4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "504e586c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 11 [SELF-STUDY S1 / ADD-BACK 2] -- BDM_1-P_0 on the same problems.\n",
    "# Expected rates (verify_rt0_2d.py): 2.00 / 1.00 / 1.00 / 1.99 -- the\n",
    "# measured one-dimensional profile of Lecture 1 (frame 11 there), now\n",
    "# obtained in two dimensions. The divergence column repeats cell 4 to all\n",
    "# digits: both equal ||f - P_h f||.\n",
    "rows, Ns = [], [4, 8, 16, 32]\n",
    "for N in Ns:\n",
    "    w, mesh, (u_ex, sig_ex, f) = solve_mixed(N, family=\"BDM\", degree=1)\n",
    "    sig_h, u_h = w.subfunctions\n",
    "    Phu = Function(u_h.function_space()).project(u_ex)\n",
    "    rows.append([l2err(sig_ex - sig_h), l2err(div(sig_h) + f),\n",
    "                 l2err(u_ex - u_h), l2err(Phu - u_h)])\n",
    "rows = np.array(rows)\n",
    "for N, r in zip(Ns, rows):\n",
    "    print(f\"{N:<5d}\" + \"\".join(f\"{e:14.4e}\" for e in r))\n",
    "for r in np.log2(rows[:-1]/rows[1:]):\n",
    "    print(\"rate \" + \"\".join(f\"{x:14.2f}\" for x in r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca89560",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 12 [SELF-STUDY S2] -- what conformity in H(div) means, numerically.\n",
    "# The normal component of sigma_h is single-valued across interior edges;\n",
    "# the tangential component is not. Both statements are one facet integral.\n",
    "w, mesh, _ = solve_mixed(16)\n",
    "sig_h, _ = w.subfunctions\n",
    "n = FacetNormal(mesh)\n",
    "# jump(sig_h, n) = dot(sig_h('+'), n('+')) + dot(sig_h('-'), n('-')):\n",
    "# exactly the normal jump of Lemma 15.\n",
    "jn = sqrt(assemble(jump(sig_h, n)**2*dS))\n",
    "# tangential component, restricted to one side of each facet:\n",
    "tp = as_vector([-n('+')[1], n('+')[0]])\n",
    "jt = sqrt(assemble((dot(sig_h('+'), tp) - dot(sig_h('-'), tp))**2*dS))\n",
    "print(f\"||[sigma_h . n]||_(interior edges) = {jn:.3e}   (zero: conforming)\")\n",
    "print(f\"||[sigma_h . t]||_(interior edges) = {jt:.3e}   (nonzero: allowed)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79c82243",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 13 [SELF-STUDY S3 / ADD-BACK 4] -- balances on unions of elements.\n",
    "# Summing the elementwise identity over any set of cells telescopes (shared\n",
    "# edges cancel): int_(d omega) q_h.n = int_omega f, exactly, for every union\n",
    "# omega. Verified here on random unions, reusing the residual vector.\n",
    "rng = np.random.default_rng(1)\n",
    "w, mesh, (u_ex, sig_ex, f) = solve_mixed(32)\n",
    "sig_h, _ = w.subfunctions\n",
    "r = cell_balance(-sig_h, f, mesh)          # per-element defects\n",
    "for trial in range(5):\n",
    "    sel = rng.random(len(r)) < 0.3         # a random union of ~30% of cells\n",
    "    print(f\"random union {trial}: {sel.sum():4d} cells, \"\n",
    "          f\"|patch defect| = {abs(r[sel].sum()):.3e}\")\n",
    "print(\"\\nThe patch defect is the sum of element defects: machine zero.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c451160d",
   "metadata": {},
   "source": [
    "Pointers: Exercise 3 uses cells 7, 9, 13. The reference implementation `verify_rt0_2d.py` repeats every computation of this notebook in pure NumPy, including the assembly; reading it is the fastest way to see the orientation handling of slide 10 in full."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Firedrake",
   "language": "python",
   "name": "firedrake"
  },
  "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
}