{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Spatial Discretisation\n",
    "\n",
    "In this notebook, the spatial discretisation and the process for constructing it in the code will be shown.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "We recap the total operator equation:\n",
    "\n",
    "$$\n",
    "(\\mathbf{I} - \\mathbf{M}\\mathbf{T}^{-1}\\mathbf{H}\\mathbf{D}\\mathbf{S})\n",
    "[\\phi_{g\\ell\\epsilon i}]\n",
    "=\n",
    "\\mathbf{M}\\mathbf{T}^{-1}\\mathbf{H}[q_{gk\\epsilon i}]\n",
    "$$\n",
    "\n",
    "with the following operator definitions:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{M}[a_{\\dots k \\dots}]\n",
    "&= \\sum_k w_k Y_{\\ell k} a_{\\dots k \\dots}, \\\\[0.5em]\n",
    "\\mathbf{D}[b_{\\dots \\ell \\dots}]\n",
    "&= \\sum_\\ell Y_{\\ell k} b_{\\dots \\ell \\dots}, \\\\[0.5em]\n",
    "\\mathbf{H}[c_{\\dots \\epsilon i}]\n",
    "&= \\sum_i c_{\\dots \\epsilon i}\n",
    "\\int_{\\epsilon} H_{\\epsilon j} H_{\\epsilon i}\\, dV, \\\\[0.5em]\n",
    "\\mathbf{S}[d_{\\tilde g \\ell \\epsilon \\dots}]\n",
    "&= \\sum_{\\tilde g}\n",
    "\\sigma^s_{\\ell g \\tilde g \\epsilon}\n",
    "d_{\\tilde g \\ell \\epsilon \\dots}, \\\\[0.5em]\n",
    "\\mathbf{T}[\\psi_{gk\\epsilon i}]\n",
    "&=\n",
    "\\mathbf{G}_{gk\\epsilon}[\\psi_{gk\\epsilon i}]\n",
    "-\n",
    "\\sum_{\\beta\\in\\beta^-_{\\epsilon k}}\n",
    "\\mathbf{F}_{k\\epsilon,\\epsilon(\\beta)}\n",
    "[\\psi_{gk\\epsilon(\\beta)i}] \\\\\n",
    "&=\n",
    "\\sum_i \\psi_{gk\\epsilon i}\n",
    "\\left[\n",
    "\\int_{\\epsilon}\n",
    "H_{\\epsilon i}\n",
    "\\left(\n",
    "\\sigma_g H_{\\epsilon j}\n",
    "-\n",
    "\\Omega_k\\cdot\\nabla H_{\\epsilon j}\n",
    "\\right)\n",
    "\\,dV\n",
    "+\n",
    "\\sum_{\\beta\\in\\beta^+_{\\epsilon k}}\n",
    "|\\Omega_k\\cdot\\mathbf{n}_\\beta|\n",
    "\\int_{\\beta}\n",
    "H_{\\epsilon i}H_{\\epsilon j}\\,dA\n",
    "\\right] \\\\\n",
    "&\\quad\n",
    "-\n",
    "\\sum_{\\beta\\in\\beta^-_{\\epsilon k}}\n",
    "\\sum_i\n",
    "\\psi_{gk\\epsilon(\\beta)i}\n",
    "|\\Omega_k\\cdot\\mathbf{n}_\\beta|\n",
    "\\int_{\\beta}\n",
    "H_{\\epsilon(\\beta)i}\n",
    "H_{\\epsilon j}\\,dA.\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "This assumed a particular spatial discretisation with *elements* $\\epsilon$ with a number of faces $\\beta_\\epsilon$ and basis functions $H_{\\epsilon i}$ defined on each element (local to each element).\n",
    "\n",
    "In particular, the sweep algorithm requires the faces to be unambiguously upwind or downwind ($\\Omega_k \\cdot \\mathbf{n}_\\beta$ positive or negative) and the element to be *convex*. Here, we show the currently standard [Basix](https://github.com/FEniCS/basix) Lagrangian simplex implementation, although different kinds of elements could be implemented as long as they satisfy the above assumptions.\n",
    "\n",
    "Note that of these operators, only the $\\mathbf{T}$ and $\\mathbf{H}$ operators depend on the spatial discretisation: they are the only ones involving basis functions $H_{\\epsilon i}$ and integrals over the elements $\\epsilon$ and faces $\\beta$. The other operators are purely algebraic and do not depend on the spatial discretisation, so we'll focus our attention on these.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "###  Simplex Elements\n",
    "\n",
    "Lagrangian elements mean that there are a set of *nodes* $\\mathbf{x}_{\\epsilon i}$ associated with each element $\\epsilon$, and the basis functions $H_{\\epsilon i}$ are defined such that they are equal to 1 at their associated node and 0 at all other nodes. The basis functions are polynomials of a given order $p$. Let's plot them for both 1D and 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "jax.config.update(\"jax_platforms\",   \"cpu\")\n",
    "from jax_sn.solution_domain import BasixLagrangianSimplex\n",
    "from jax_sn.tools.plotly import plot_element_basis_functions\n",
    "plot_element_basis_functions(element_1 := BasixLagrangianSimplex(degree = 4, dimension = 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_element_basis_functions(element_2 := BasixLagrangianSimplex(degree = 3, dimension = 2), mode='all', n_subdivisions=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "These basis functions don't actually show up anywhere, except in the integrals. For simplex elements, the integrals are even easier: they have a linear Jacobian (mapping a real element to the reference element), so the integrals can be pre-computed and just multiplied by the Jacobian, as explained {ref}`here <simplex_elements>`:\n",
    "\n",
    "$$ \\int H_{\\epsilon i} H_{\\epsilon j}\\, dV = J_\\epsilon \\int H_{i} H_{j}\\, dV. $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Do the elements have a linear Jacobian?\")\n",
    "print(f\"Element 1: {element_1.has_linear_jacobian}\")\n",
    "print(f\"Element 2: {element_2.has_linear_jacobian}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Geometric Domain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "These integrals therefore require some precomputed data. Since this depends on the particular type of element, each element has a way to calculate it from a {py:class}`~jax_sn.domain.GeometricDomain`. Of course, they have to be compatible: you cannot compute it for a triangle using a tetrahedral domain. A {py:class}`~jax_sn.domain.GeometricDomain` only stores adjacency information, cross sections, and nodes / connectivity —  it doesn't know about which type of element it is.\n",
    "\n",
    "A detail is that a domain requires a face template: which nodes form which face. You can input it normally, but this has to be compatible with the face template of the finite element later. Commonly, we just create them both at the same time to avoid mismatching. \n",
    "\n",
    "The face template doesn't depend on the degree of the finite element!\n",
    "\n",
    "We create a simple 2D domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "vertices = jnp.array([\n",
    "    [0.0, 0.0],  # 0 \\\n",
    "    [1.0, 0.0],  # 1  | box corners\n",
    "    [1.0, 1.0],  # 2  |\n",
    "    [0.0, 1.0],  # 3 /\n",
    "    [2.0, 0.0],  # 4 \\ top triangle\n",
    "    [0.0, 2.0],  # 5 \\ right triangle\n",
    "])\n",
    "\n",
    "connectivity = jnp.array([\n",
    "    [0, 1, 2],  # box triangle A\n",
    "    [0, 2, 3],  # box triangle B\n",
    "    [1, 2 ,4],  # top triangle\n",
    "    [2, 3, 5],  # right triangle\n",
    "])\n",
    "\n",
    "solution_element = BasixLagrangianSimplex(degree=2, dimension=2)\n",
    "\n",
    "print(\"Face template:\", solution_element.face_template)\n",
    "\n",
    "from jax_sn.domain import GeometricDomain\n",
    "geometric_domain = GeometricDomain.from_mesh((vertices, connectivity), solution_element.face_template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "and plot it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.tools.plotly import plot_domain\n",
    "plot_domain(geometric_domain, color_by_material= False, show_element_numbers = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "The geometric domain has the following data structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(geometric_domain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "- `face_normals` are just the normals of each face: 4 elements, 3 faces per element, 3 dimensional. (minus one means not connected)\n",
    "- `etoe` is which element is connected to which element through its faces\n",
    "- `etof` is which face number each element is connected to (of the other element)\n",
    "- `face_perm_idx` is the permutation of the face nodes to match the other element's face nodes.\n",
    "\n",
    "These are calculated automatically by the above constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(geometric_domain.etoe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "Here we see element 0 is connected to element 2 through face 0 and element 1 through face 1, with no elements connected through face 2."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Precomputing integral quantities\n",
    "\n",
    "Using this geometric domain, we can precompute all geometric Jacobians (including edge Jacobians, inverse Jacobians, and determinants):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution_element.precompute_geometric_data(geometric_domain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "For the computation of the integrals, besides the Jacobian, we also need the integrals over the reference element, e.g. \n",
    "\n",
    "$$ \\int H_i H_j\\, dV $$\n",
    "over the reference element (as well as the face / derivative versions). \n",
    "\n",
    "These are stored in a {py:class}`~jax_sn.solution_domain.base_fem.FEMMatrices` and can be computed for a particular element easily:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.solution_domain import FEMMatrices\n",
    "\n",
    "fem_matrices = FEMMatrices.from_element(solution_element)\n",
    "print(fem_matrices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "All matrices have a size of 6 because a 2nd degree triangle has 6 basis functions.\n",
    "\n",
    "- `mass_matrix` is the $\\int H_i H_j\\, dV$ integral over the reference element\n",
    "- `integration_coeff` is the $\\int H_i \\, dV$ integral over the reference element\n",
    "- `derivative_matrix` are the $\\int H_i \\frac{\\partial H_j}{\\partial q_k} \\, dV$ matrices (2 reference coordinate dimension)\n",
    "- `outflow_matrices` are the $\\int_{\\beta} H_i H_j\\, dA$ matrices for each face $\\beta$ of the reference element (3 faces for a triangle)\n",
    "- `inflow_matrices` are the $\\int_{\\beta} H^p_i H_j\\, dA$ matrices for each face in the connected element $p$ (3 faces), each permutation of the face nodes (2), and each face of the reference element (3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "### Domain and cross sections"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "To actually solve the equations, we need the cross sections in there as well ($\\sigma_g \\epsilon $ and $\\sigma^s_{\\ell g \\tilde g \\epsilon}$). Therefore, to create the container for all precomputed data, the {py:class}`~jax_sn.solution_domain.SolutionDomain`, we need to construct the domain including cross section information.\n",
    "\n",
    "Let's first create some cross sections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.domain import Domain, CrossSectionData, CrossSectionsIndexed\n",
    "\n",
    "water = CrossSectionData(\n",
    "    total = jnp.array([15.2090, 24.4140]),\n",
    "    scattering = jnp.array([\n",
    "        [\n",
    "            [5.5069, 0.0],\n",
    "            [6.6227, 11.4080]\n",
    "        ],\n",
    "        [\n",
    "            [1.158567, 0.0],\n",
    "            [1.1234767, 2.4733]\n",
    "        ]\n",
    "    ])\n",
    ")\n",
    "iron = CrossSectionData(\n",
    "    total = jnp.array([30.3110, 26.8760]),\n",
    "    scattering = jnp.array([\n",
    "        [\n",
    "            [18.2500, 0.0],\n",
    "            [6.45294, 21.3570]\n",
    "        ],\n",
    "        [\n",
    "            [4.36, 0.0],\n",
    "            [-0.003508, 2.3457]\n",
    "        ]        \n",
    "    ])\n",
    ")\n",
    "void = CrossSectionData(\n",
    "    total = jnp.array([1e-10, 1e-10]),\n",
    "    scattering = jnp.zeros((2, 2, 2))\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "Now we point each element to a specific material using indexing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "water_idx = 0 \n",
    "iron_idx  = 1\n",
    "void_idx  = 2\n",
    "material_idx = jnp.array([water_idx, water_idx, iron_idx, void_idx])\n",
    "# Create a batched version of the data:\n",
    "total_xs               = CrossSectionData.from_list([water, iron, void])\n",
    "cross_sections_indexed = CrossSectionsIndexed(material_idx, total_xs)\n",
    "\n",
    "print(cross_sections_indexed)\n",
    "\n",
    "domain = Domain.from_geometric_domain_and_cross_sections(geometric_domain, cross_sections_indexed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "and plot the domain again, colorbing by material:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_domain(domain, material_names = [\"Water\", \"Iron\", \"Void\"], show_element_numbers = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "### Solution domain and spatial data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "We can now actually create the solution domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.solution_domain import SolutionDomain\n",
    "solution_domain = SolutionDomain.from_element_and_domain(solution_element, domain)\n",
    "print(solution_domain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "It's simply a container for all the precomputed data, the geometric domain, and the element type. \n",
    "\n",
    "One thing to note is that if you modify the vertices (e.g. using a `eqx.tree_at`), this might make the precomputed data inconsistent! \n",
    "\n",
    "See the ... for how to cleanly update the vertices instead.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {},
   "source": [
    "## Mass Operator\n",
    "\n",
    "The mass operator $\\mathbf{H}$\n",
    "\n",
    "$$\\mathbf{H}[c_{\\dots \\epsilon i}]\n",
    "= \\sum_i c_{\\dots \\epsilon i}\n",
    "\\int_{\\epsilon} H_{\\epsilon j} H_{\\epsilon i}\\, dV= \\sum_i c_{\\dots \\epsilon i}\\mathcal{J}_\\epsilon\n",
    "\\int_{\\epsilon} H_{j} H_{ i}\\, dV\\\\[0.5em]$$\n",
    "\n",
    "is quite simple; it's just a batched matrix product and requires the FEM matrices and the precomputed data. \n",
    "\n",
    "It needs additional shape information to have proper checking of input/output arguments (e.g. group and moment indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sn.operators.mass_operator import MassOperator , LaxMass, MassData\n",
    "\n",
    "mass_data = MassData(solution_domain.fem_matrices, solution_domain.precomputed_geometric_data)\n",
    "mass_operator = MassOperator(mass_data, additional_shape = (), strategy = LaxMass())\n",
    "\n",
    "print(mass_operator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36",
   "metadata": {},
   "source": [
    "In practice you'd never construct it this way; it's a part of a composed operator and there are convenience functions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37",
   "metadata": {},
   "source": [
    "### Connecting spatial discretisation to other operators\n",
    "\n",
    "Every other operator also 'sees' the spatial discretisation, but in a simpler way: they are all *local* operators that just operator on each element and value separately.\n",
    "\n",
    "In other words, they are just applied to each element and value directly, without considering any integrals or basis functions. They don't see the spatial discretisation except for the additional dimensions they operate on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jaxsbgeom",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}