{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Quadrature set and moments\n", "\n", "In this notebook, we will construct a quadrature set, plot it, and learn what the representation in moments and discrete angles mean and construct the operators\n", "\n", "$$\\begin{aligned}\\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]\\end{aligned}$$\n", "\n", "in the code." ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "import numpy as onp\n", "jax.config.update(\"jax_platforms\", \"cpu\")\n", "jax.config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "An angle $k$ is just a 3D vector $\\Omega_k$; a quadrature set is a collection of angles $\\{\\Omega_k\\}$ with associated weights $\\{w_k\\}$, which are used to approximate integrals over the unit sphere:\n", "\n", "$$ \\int_{4\\pi} f(\\Omega) d\\Omega \\approx \\sum_k w_k f(\\Omega_k)$$" ] }, { "cell_type": "markdown", "id": "3", "metadata": {}, "source": [ "We can create a quadrature set using either just manually creating angles and weights, or by selecting a particular precomputed quadrature set." ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "from jax_sn.quadrature_set import create_tn_quadrature_set\n", "\n", "tn_quadrature_set = create_tn_quadrature_set(N = 4)\n", "print(tn_quadrature_set)\n" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "We can plot the quadrature set, coloring the weight." ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "from jax_sn.tools.plotly import plot_quadrature_set\n", "\n", "plot_quadrature_set(tn_quadrature_set)" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "## Integration" ] }, { "cell_type": "markdown", "id": "8", "metadata": {}, "source": [ "To integrate a particular function, it is just a simple weighted sum over the quadrature set's angles and weights:\n", "\n", "$$ \\int_{4\\pi} f(\\Omega)\\, d\\Omega \\approx \\sum_k w_k f(\\Omega_k)$$" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "def integrate_over_sphere(f, quadrature_set):\n", " return jnp.sum(quadrature_set.weights * f(quadrature_set.angles))" ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "Let's integrate a family of functions with a known closed form: $f(\\Omega) = (\\Omega \\cdot n)^p$ for a fixed unit vector $n$ and integer $p \\geq 0$. Writing $\\mu = \\Omega \\cdot n = \\cos\\theta$ and $d\\Omega = d\\mu\\, d\\phi$:\n", "\n", "$$\\int_{4\\pi} (\\Omega \\cdot n)^p \\, d\\Omega = \\int_0^{2\\pi}\\!\\!\\int_{-1}^{1} \\mu^p \\, d\\mu\\, d\\phi = 2\\pi \\int_{-1}^1 \\mu^p\\, d\\mu = \\begin{cases} \\dfrac{4\\pi}{p+1} & p \\text{ even} \\\\ 0 & p \\text{ odd} \\end{cases}$$\n", "\n", "independent of $n$." ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "def analytic_integral(p):\n", " return 0.0 if p % 2 else 4 * jnp.pi / (p + 1)\n", "\n", "n = jnp.array([0.0, 1/jnp.sqrt(2), 1/jnp.sqrt(2)])\n", "\n", "for p in range(12):\n", " f = lambda omega, p=p: jnp.dot(omega, n) ** p\n", " numeric = float(integrate_over_sphere(f, tn_quadrature_set))\n", " analytic = analytic_integral(p)\n", " print(f\"p={p:2d} numeric={numeric:10.5f} analytic={analytic:10.5f} |diff|={abs(numeric - analytic):.2e}\")" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "The quadrature is only exact for $p=1$; you have discretization errors for higher orders." ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "## Moments\n", "\n", "As mentioned in the {ref}`theory section `, the isotropicity assumption of the scattering kernel immediately determines that there are two representations of the angular flux relevant for the discrete ordinates method: \n", "\n", "- The angular flux on the discrete angles $\\psi_{k} = \\psi(\\Omega_k)$\n", "- The spherical harmonic moments of the angular flux $\\phi_{lm} = \\int_{4\\pi} Y_{lm}(\\Omega) \\psi(\\Omega) d\\Omega$ used for computing the scattering source.\n", "\n", "The quadrature set and its weight can perform precisely the integration needed to convert between them. Given a set of angles $\\{\\Omega_k\\}$ and weights $\\{w_k\\}$, the discrete-to-moments transformation is\n", "\n", "$$ \\mathbf{M}[\\psi_{k}] = \\int Y_{lm}(\\Omega) \\psi(\\Omega) d\\Omega = \\sum_k w_k Y_{lm}(\\Omega_k) \\psi_k = \\mathcal{M}\\cdot \\mathbf{\\psi} $$\n", "\n", "This $\\mathcal{M} = w_k Y_{lm}(\\Omega_k)$ matrix can be computed from the quadrature set alone (given a maximum moment order $L$).\n", "\n", "Given these moments, we can re-compute the angular flux on any angle by expanding the spherical harmonics:\n", "\n", "$$ \\psi(\\Omega) = \\sum_{l=0}^{L} \\sum_{m=-l}^{l} Y_{lm}(\\Omega) \\phi_{lm} = \\mathcal{Y}\\cdot \\mathbf{\\phi} $$\n", "\n", "For the special case that we are expanding it back to the original quadrature set (as is always the case in our transport simulations -- the scattering source is needed on the quadrature set again!), this is the moments-to-discrete operator:\n", "\n", "$$ \\mathbf{D}[\\phi_{lm}] = \\sum_{l=0}^{L} \\sum_{m=-l}^{l} Y_{lm}(\\Omega_k) \\phi_{lm} = \\mathcal{D}\\cdot \\mathbf{\\phi} $$\n" ] }, { "cell_type": "markdown", "id": "14", "metadata": {}, "source": [ "### Visualizing moments \n", "Let's use a transport solution from an isotropic slab:\n", "\n", "$$\\psi(x,\\Omega) = \\frac{q}{\\sigma}\\left(1-e^{-\\sigma s(x,\\Omega)}\\right), \\qquad s(x,\\Omega) = \\begin{cases}x/\\mu & \\mu>0 \\\\ (L-x)/(-\\mu) & \\mu<0 \\\\ \\infty & \\mu=0\\end{cases}$$\n", "\n", "and compute the moments from the discrete angles and weights." ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": {}, "outputs": [], "source": [ "from jax_sn.analytical_solutions import psi_isotropic_box\n", "\n", "box_size, sigma, q = 1.0, 2.0, 1.0\n", "x0 = 0.3 * box_size # off-centre, so psi differs between the +z and -z facing directions\n", "\n", "psi_k = jnp.array([psi_isotropic_box(jnp.array([[x0]]), box_size, i, sigma, q) for i in tn_quadrature_set.angles])[:,0]\n", "\n", "l_max = 4\n", "moments = jnp.einsum(\"k, kl -> l\", psi_k, tn_quadrature_set.wk_ylm_real(l_max))\n" ] }, { "cell_type": "markdown", "id": "16", "metadata": {}, "source": [ "We will now visualize these moments in three different ways:\n", "\n", "1. The original function on the quadrature set angles\n", "2. The spherical harmonic moments, shaded with how large they are after integrating.\n", "3. The re-interpolated description of the angular flux on a fine grid: \n", "\n", "$$ \\psi(\\Omega) = \\sum_{l=0}^{L} \\sum_{m=-l}^{l} Y_{lm}(\\Omega) \\phi_{lm}$$" ] }, { "cell_type": "code", "execution_count": null, "id": "17", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "from plotly.subplots import make_subplots\n", "from jax_sn.tools.plotly import plot_quadrature_set, plot_moments_spheres, plot_spherical_field, add_axis_triad\n", "\n", "fig_angles = plot_quadrature_set(tn_quadrature_set, values=psi_k, show_unit_sphere=True,\n", " colorbar_title=\"ψ(Ω)\")\n", "fig_moments = plot_moments_spheres(moments, l_max)\n", "fig_field = plot_spherical_field(l_max, moments, colorbar_title=\"Σ φ_lm·Y_lm(Ω)\")\n", "\n", "shared_max = max(float(jnp.max(jnp.abs(psi_k))), float(onp.max(onp.abs(fig_field.data[0].surfacecolor))))\n", "fig_angles.data[-1].marker.update(cmin=0.0, cmax=shared_max, colorscale=\"Viridis\")\n", "fig_field.data[0].update(cmin=0.0, cmax=shared_max, colorscale=\"Viridis\")\n", "\n", "# keep each scene's own colour bar clear of the others\n", "fig_angles.data[-1].marker.colorbar.update(x=1.05, len=0.28, y=0.90)\n", "fig_moments.data[0].colorbar.update(x=1.05, len=0.28, y=0.50)\n", "fig_field.data[0].colorbar.update(x=1.05, len=0.28, y=0.12)\n", "\n", "fig = make_subplots(rows=3, cols=1, specs=[[{\"type\": \"scene\"}], [{\"type\": \"scene\"}], [{\"type\": \"scene\"}]],\n", " row_heights=[0.32, 0.40, 0.28], vertical_spacing=0.06)\n", "\n", "for trace in fig_angles.data:\n", " fig.add_trace(trace, row=1, col=1)\n", "for trace in fig_moments.data:\n", " fig.add_trace(trace, row=2, col=1)\n", "for trace in fig_field.data:\n", " fig.add_trace(trace, row=3, col=1)\n", "\n", "\n", "locked_camera = dict(eye=dict(x=0, y=0, z=3.2), up=dict(x=0, y=1, z=0), projection=dict(type=\"orthographic\"))\n", "hidden = dict(visible=False)\n", "fig.update_scenes(aspectmode=\"data\", xaxis=hidden, yaxis=hidden, zaxis=hidden, camera=locked_camera)\n", "\n", "add_axis_triad(fig, origin=(-1.5, -1.5, -1.3), length=0.7, row=1, col=1)\n", "add_axis_triad(fig, origin=(-(l_max + 2.3) * 2.2, 1.3 * 2.2, 0.0), length=1.7, row=2, col=1)\n", "add_axis_triad(fig, origin=(-1.5, -1.5, -1.3), length=0.7, row=3, col=1)\n", "\n", "# connecting arrows for the M/D operators, centred in the gap between panels\n", "gaps = [(fig.layout.scene.domain.y[0], fig.layout.scene2.domain.y[1]),\n", " (fig.layout.scene2.domain.y[0], fig.layout.scene3.domain.y[1])]\n", "for (y_top, y_bottom), label in zip(gaps, [\"M\", \"Sum\"]):\n", " fig.add_shape(type=\"line\", xref=\"paper\", yref=\"paper\", x0=0.5, x1=0.5, y0=y_top, y1=y_bottom,\n", " line=dict(color=\"#8B0000\", width=3))\n", " fig.add_annotation(x=0.5, y=y_bottom, xref=\"paper\", yref=\"paper\", showarrow=False, yanchor=\"top\",\n", " text=\"▼\", font=dict(size=18, color=\"#8B0000\"))\n", " fig.add_annotation(x=0.56, y=(y_top + y_bottom) / 2, xref=\"paper\", yref=\"paper\", showarrow=False,\n", " text=f\"{label}\", font=dict(size=20, color=\"#8B0000\"))\n", "\n", "fig.update_layout(width=750, height=1650, showlegend=False, margin=dict(l=0, r=100, t=20, b=0))\n", "fig" ] }, { "cell_type": "markdown", "id": "18", "metadata": {}, "source": [ "We can see that the moments contributing are the isotropic moment ($l=0$), the forward-backward moment $(l=1, m=1)$ and two $(l=2)$ moments.\n", "\n", "The final *Sum* operation is exactly the moments-to-discrete operator referred to as $\\mathbf{D}$. It is used for re-interpolating the scattering source on the original quadrature set after multiplying all terms with the scattering cross sections." ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "## In the code" ] }, { "cell_type": "markdown", "id": "20", "metadata": {}, "source": [ "In the code, these operations converting between the two representations are standardized operators:\n", "\n", "{py:class}`~jax_sn.operators.moment_operator.MomentOperator` and {py:class}`~jax_sn.operators.moment_operator.DiscreteOperator`. They require a *strategy* or how to compute the batched matrix multiplications above. Since this strategy is so similar, we have this as a shared {py:class}`~jax_sn.operators.moment_operator.FluxMomentStrategy`." ] }, { "cell_type": "code", "execution_count": null, "id": "21", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.moment_operator import MomentOperator, DiscreteOperator, EinsumFluxMoment, MomentData, DiscreteData\n", "\n", "moment_operator = MomentOperator(MomentData(tn_quadrature_set.wk_ylm_real(l_max)), (), EinsumFluxMoment()) \n", "#() is elements basis shape or additional batching dimensions, not used here since we have a single set of values instead \n", "# of defined on the whole domain.\n", "discrete_operator = DiscreteOperator(DiscreteData(tn_quadrature_set.ylm_real(l_max)), (), EinsumFluxMoment())" ] }, { "cell_type": "markdown", "id": "22", "metadata": {}, "source": [ "In practice, you'd never construct them as such: there are convenience functions and they normally only appear within a composed multigroup or single-group operator.\n", "\n", "They compute the same as above:" ] }, { "cell_type": "code", "execution_count": null, "id": "23", "metadata": {}, "outputs": [], "source": [ "onp.testing.assert_allclose(moment_operator.mv(psi_k), moments)" ] }, { "cell_type": "markdown", "id": "24", "metadata": {}, "source": [ "Finally, they are not inverses of each other! They cannot be: they are not necessarily square. It is more a projection operator, which loses some information." ] }, { "cell_type": "code", "execution_count": null, "id": "25", "metadata": {}, "outputs": [], "source": [ "try:\n", " onp.testing.assert_allclose(discrete_operator.mv(moment_operator.mv(psi_k)), psi_k)\n", "except AssertionError as e:\n", " print(\"The discrete operator is not the inverse of the moment operator, as expected.\")\n", " print(e)" ] } ], "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 }