{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Multigroup solution\n", "\n", "In this notebook, we will show how to solve a multigroup neutron transport problem, using both the total-group solver and the per-group solver." ] }, { "cell_type": "markdown", "id": "1", "metadata": {}, "source": [ "We'll re-use the same materials as before and create the domain from the following regions:\n", "\n", "```text\n", "0 - 0.2 m: Water\n", "0.2 - 0.3 m: Water with source\n", "0.3 - 0.8 m: Water\n", "0.8 - 1.0 m: Iron\n", "1.0 - 1.1 m: Void\n", "1.1 - 1.2 m: Water\n", "1.2 - 1.3 m: Void\n", "```\n", "\n", "We just use a simple discretization with a set number of elements per 10 cm and build from there (we hide the \n", "details of creating the mesh and material indices)" ] }, { "cell_type": "code", "execution_count": null, "id": "2", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "jax.config.update(\"jax_platforms\", \"cpu\")\n", "jax.config.update(\"jax_enable_x64\", True)\n", "import jax_sn\n", "import equinox as eqx\n", "from typing import List\n", "from jax_sn.domain.cross_sections import 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", "\n", "\n", "def base_region_stats(n_elements_per_10_cm : int):\n", " base_region_boundaries = jnp.array([0.0, 0.2, 0.3, 0.8, 1.0, 1.1, 1.2, 1.3])\n", " base_region_sizes = [a - b for a,b in zip(base_region_boundaries[1:], base_region_boundaries[:-1])]\n", " base_elements_per_region = [int(round( size * 10 * n_elements_per_10_cm)) for size in base_region_sizes]\n", " return base_region_boundaries, base_region_sizes, base_elements_per_region\n", "\n", "def create_vertices(region_boundaries, elements_per_region):\n", " region_sizes = [a - b for a,b in zip(region_boundaries[1:], region_boundaries[:-1])] \n", "\n", " vertices = jnp.concatenate([\n", " jnp.linspace(start, start + size, n_elements, endpoint=False) \n", " for start, size, n_elements in zip(region_boundaries[:-1], region_sizes, elements_per_region)\n", " ] + [jnp.array([region_boundaries[-1]])])[:, None]\n", " return vertices\n", "\n", "def create_mesh_and_source(n_elements_per_10_cm : int):\n", " region_boundaries, region_sizes, elements_per_region = base_region_stats(n_elements_per_10_cm)\n", " vertices = create_vertices(region_boundaries, elements_per_region)\n", " connectivity = jnp.array([[i, i+1] for i in range(sum(elements_per_region))])\n", "\n", " water_idx, iron_idx, void_idx = 0, 1, 2 \n", " material_indices = [water_idx, water_idx, water_idx, iron_idx, void_idx, water_idx, void_idx]\n", "\n", " material_indices_elements =jnp.concatenate([\n", " jnp.full(n_elements, material) for n_elements, material in zip(elements_per_region, material_indices)\n", " ])\n", " source_regions = [0, 1, 0, 0, 0, 0, 0]\n", "\n", " source_base = jnp.concatenate([\n", " jnp.full(n_elements, source) for n_elements, source in zip(elements_per_region, source_regions)\n", " ])\n", " \n", " source_values = jnp.array([9.7702e4, 4.5451e5]) \n", " source = source_base[None, :] * source_values[:, None] \n", " cross_sections = CrossSectionsIndexed(\n", " indices = material_indices_elements,\n", " cross_section_data = CrossSectionData.from_list([water, iron, void]) \n", " )\n", " return vertices, connectivity, cross_sections, source" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operator_creation import MultiGroupBuilder, DefaultMultiGroupStrategy\n", "from jax_sn.solution_domain import BasixLagrangianSimplex\n", "from jax_sn.quadrature_set import create_tn_quadrature_set\n", "n_elements_per_10_cm = 2\n", "degree = 1\n", "tn_order = 3\n", "\n", "\n", "vertices, connectivity, cross_sections, source = create_mesh_and_source(n_elements_per_10_cm )\n", "strategy = DefaultMultiGroupStrategy.default_multi_group_strategy()\n", "element = BasixLagrangianSimplex(degree = degree, dimension = 1)\n", "quadrature_set = create_tn_quadrature_set(tn_order)\n", "\n", "mg_builder = MultiGroupBuilder(vertices, connectivity, element, cross_sections, quadrature_set, strategy)" ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "from jax_sn.tools.plotly import plot_domain\n", "plot_domain(mg_builder.domain, material_names = ['Water', 'Iron', 'Void'], width =800)" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "Before solving we create an isotropic source (we repeat along the basis to have a flat source within each element):" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "from jax_sn.solvers.right_hand_side import TotalSource\n", "\n", "def create_source(mg_builder, elementwise_source):\n", " isotropic_src = jnp.repeat(elementwise_source[..., None], mg_builder.operator_data.n_basis, axis=-1)\n", " total_source = TotalSource.from_operator_and_isotropic_source(mg_builder.multi_group_operator, isotropic_src)\n", " return total_source\n", "\n", "source_p1 = create_source(mg_builder, source)" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "## Solving\n", "\n", "Solving is easy in this formulation! We first create the correct right-hand-side:\n", "\n", "$$ \\mathbf{b} = \\mathbf{M} \\mathbf{T}^{-1} \\mathbf{D} \\mathbf{H} [q_{g\\ell ei}] $$\n", "\n", "This is precisely the chain strategy in multigroup form. We could create a separate multigroup version of it, but here we just scan the single-group chain strategy we already made.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.multi_group_operator.composed_multi_group_strategies import MultiGroupChainScanStrategy\n", "mg_chain = MultiGroupChainScanStrategy(strategy.single_group_chain_strategy)" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "We will create a top-level function to solve, including the right-hand-side operator, since this allows for a \n", "consistent derivative later:" ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "import lineax as lx\n", "import equinox as eqx\n", "\n", "@eqx.filter_jit\n", "def solve_anisotropic(mgop , total_source : TotalSource, solver : lx.AbstractLinearSolver, mg_chain : MultiGroupChainScanStrategy): \n", " rhs = mg_chain.forward(mgop.data, total_source)\n", " return lx.linear_solve(operator = mgop, vector = rhs, solver = solver)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "solution = solve_anisotropic(mg_builder.multi_group_operator,source_p1, lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)\n", "print(\"Solution converged in\", solution.stats['num_steps'], \"steps\")" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "Let's plot it:" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "from jax_sn.tools.plotly import plot_1d_solution\n", "\n", "group_names = [\"E₀\", \"E₁\"]\n", "# the domain's regions, shaded behind the flux so the material layout is readable off the plot itself\n", "regions = [(0.0, 0.2, \"Water\"), (0.2, 0.3, \"Source\"), (0.3, 0.8, \"Water\"), (0.8, 1.0, \"Iron\"),\n", " (1.0, 1.1, \"Void\"), (1.1, 1.2, \"Water\"), (1.2, 1.3, \"Void\")]\n", "region_style = {\"Water\" : dict(fillcolor = \"#8C9BAB\", opacity = 0.13),\n", " \"Source\": dict(fillcolor = \"#00A67E\", opacity = 0.22),\n", " \"Iron\" : dict(fillcolor = \"#4A5560\", opacity = 0.30),\n", " \"Void\" : dict(fillcolor = \"#FFFFFF\", opacity = 0.00)}\n", "\n", "def interpolate_and_plot(solution_domain, solution, fig = None, dash = None, label = \"\"):\n", " '''Scalar flux (the l=0 moment) per energy group. `plot_1d_solution` samples every element\n", " separately, so inter-element jumps show as real gaps rather than being smoothed over.\n", " Pass an existing `fig` (and a `dash`/`label`) to overlay a second discretisation.'''\n", " new_fig = plot_1d_solution(solution_domain, solution.value[:, 0], xaxis_title = \"x [m]\", yaxis_title = \"ϕ\",\n", " solution_names = [name + label for name in group_names], width = 800)\n", " new_fig.update_traces(line = dict(dash = dash, width = 2.5))\n", " if fig is not None:\n", " return fig.add_traces(new_fig.data)\n", "\n", " for x0, x1, name in regions:\n", " new_fig.add_vrect(x0 = x0, x1 = x1, line_width = 0, layer = \"below\", annotation_text = name,\n", " annotation_position = \"outside top\", annotation_font_size = 10,\n", " **region_style[name])\n", " new_fig.update_xaxes(range = [0.0, 1.3])\n", " new_fig.update_yaxes(type = \"log\", exponentformat = \"power\", showgrid = True)\n", " new_fig.update_layout(template = \"plotly_white\", width = 800, height = 480, hovermode = \"x unified\",\n", " legend = dict(x = 0.99, y = 0.97, xanchor = \"right\", yanchor = \"top\",\n", " bgcolor = \"rgba(255,255,255,0.75)\", bordercolor = \"#C8C8C8\",\n", " borderwidth = 1),\n", " margin = dict(l = 60, r = 20, t = 45, b = 50))\n", " return new_fig" ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": {}, "outputs": [], "source": [ "interpolate_and_plot(mg_builder.solution_domain, solution)" ] }, { "cell_type": "markdown", "id": "15", "metadata": {}, "source": [ "Clearly, the linear elements are not so smooth in this regime!\n", "Let's do third order elements, re-using the same strategies." ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [ "mg_p3 = mg_builder.updated(element = BasixLagrangianSimplex(degree = 3, dimension=1))\n", "solution_p3 = solve_anisotropic(mg_p3.multi_group_operator, create_source(mg_p3, source), lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)\n", "\n", "fig = interpolate_and_plot(mg_builder.solution_domain, solution, dash = \"dot\", label = \" p=1\")\n", "interpolate_and_plot(mg_p3.solution_domain, solution_p3, fig = fig, label = \" p=3\")" ] }, { "cell_type": "markdown", "id": "17", "metadata": {}, "source": [ "Naturally the third order elements are significantly smoother and have less discontinuity." ] }, { "cell_type": "markdown", "id": "18", "metadata": {}, "source": [ "### Exploiting downscatter structure" ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "The above solved the multigroup equation by applying the multigroup operator. These cross sections do however not include any upscattering, and as such, the resulting multigroup structure is lower triangular:" ] }, { "cell_type": "code", "execution_count": null, "id": "20", "metadata": {}, "outputs": [], "source": [ "cross_sections.cross_section_data.downscattering_only" ] }, { "cell_type": "markdown", "id": "21", "metadata": {}, "source": [ "This can be exploited by using a more specialized lower-triangular block solver." ] }, { "cell_type": "code", "execution_count": null, "id": "22", "metadata": {}, "outputs": [], "source": [ "from jax_sn.solvers import MultiGroupLowerTriangularBlockSolver, SolverDiagnostics" ] }, { "cell_type": "markdown", "id": "23", "metadata": {}, "source": [ "For each particular single group block, the equation is:\n", "\n", "$$( \\mathbf{I} - \\mathbf{M}_g \\mathbf{T}_g^{-1} \\mathbf{H}_g \\mathbf{D}_g \\mathbf{S}_{gg})[\\phi_{g\\ell ei}] = \\mathbf{M}_g \\mathbf{T}_g^{-1} \\mathbf{H}_g \\mathbf{D}_g[q_{g\\ell ei} + \\mathbf{S}_{g}[\\phi_{g\\ell ei}]] $$\n", "\n", "where $\\mathbf{S}_g$ is the familiar gather-scatter operator (which gathers contributions from all, in this case only higher, energy groups.) and $\\mathbf{S}_{gg}$ is the same operator, but only for the current group (or equivalently, the inscatter operator)" ] }, { "cell_type": "markdown", "id": "24", "metadata": {}, "source": [ "Therefore, to use a lower-triangular block solver, we need strategies for the single-group versions. These can be the same as for the MultiGroupOperator, or just arbitrary single-group strategies. The solver *is not differentiated*, so the single-group strategies don't even need to be differentiable, as long as they are consistent with the MultiGroupOperator!\n", "\n", "It needs three:\n", "\n", "- `SingleGroupStrategy`: $\\mathbf{I}- \\mathbf{M}_g \\mathbf{T}_g^{-1} \\mathbf{H}_g \\mathbf{D}_g \\mathbf{S}_{gg}$\n", "- `SingleGroupChain`: $\\mathbf{M}_g \\mathbf{T}_g^{-1} \\mathbf{H}_g \\mathbf{D}_g$\n", "- `SingleGroupGatherScatter`: $\\mathbf{S}_g$" ] }, { "cell_type": "markdown", "id": "25", "metadata": {}, "source": [ "In practice, we can just create it from a composed multigroup operator easily, which just then requires the inscatter strategy (this has no multigroup analogue)." ] }, { "cell_type": "code", "execution_count": null, "id": "26", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.scatter_operator import IndexedScatter\n", "\n", "single_group_solver = lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 50)\n", "\n", "block_solver = MultiGroupLowerTriangularBlockSolver.from_composed(mg_builder.strategy, single_group_solver, \n", " scatter_strategy = IndexedScatter(), \n", " diagnostics = SolverDiagnostics(True, True))\n", "\n", "print(block_solver)" ] }, { "cell_type": "markdown", "id": "27", "metadata": {}, "source": [ "The single-group strategies are then created from the multigroup strategy!" ] }, { "cell_type": "code", "execution_count": null, "id": "28", "metadata": {}, "outputs": [], "source": [ "solution_block = solve_anisotropic(mg_builder.multi_group_operator, create_source(mg_builder, source), block_solver, mg_chain)" ] }, { "cell_type": "markdown", "id": "29", "metadata": {}, "source": [ "Naturally, this are the same (up to machine precision; per-group tolerances are different then total tolerance) as the original result" ] }, { "cell_type": "code", "execution_count": null, "id": "30", "metadata": {}, "outputs": [], "source": [ "import numpy as onp\n", "onp.testing.assert_allclose(solution_block.value, solution.value, atol = 1e-12)" ] }, { "cell_type": "markdown", "id": "31", "metadata": {}, "source": [ "## Taking derivatives" ] }, { "cell_type": "markdown", "id": "32", "metadata": {}, "source": [ "Let's now take some derivatives. We parametrize the region boundaries and create a function that solves and sums the square of the result[^1]:\n", "\n", "[^1]: We square to make the function non-linear; the compiler will optimize a linear function (since the cotangent values then don't actually depend on the primals) to not solve the adjoint problem after the forward solve but concurrently (which isn't pedagogically useful since it's not the general situation)" ] }, { "cell_type": "markdown", "id": "33", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "id": "34", "metadata": {}, "outputs": [], "source": [ "from typing import Callable\n", "from jax_sn.operators.multi_group_operator import MultiGroupChainStrategy\n", "n_elements_per_10_cm = 2\n", "base_region_boundaries, _ ,elements_per_region = base_region_stats(n_elements_per_10_cm)\n", "\n", "\n", "@eqx.filter_jit\n", "def solve_and_sum(theta : jax.Array, f_vertices : Callable, builder : MultiGroupBuilder,\n", " elementwise_source : jax.Array, \n", " solver : lx.AbstractLinearSolver, \n", " mg_chain : MultiGroupChainStrategy):\n", " new_vertices = f_vertices(theta)\n", " new_builder = builder.updated(vertices = new_vertices)\n", " new_source = create_source(new_builder, elementwise_source)\n", " new_solution = solve_anisotropic(new_builder.multi_group_operator, new_source, solver, mg_chain)\n", " return jnp.sum(new_solution.value**2)\n", "\n", "f_vertices = lambda theta : create_vertices(theta, elements_per_region)\n", "grad_solve_and_sum = eqx.filter_jit(eqx.filter_grad(solve_and_sum))" ] }, { "cell_type": "markdown", "id": "35", "metadata": {}, "source": [ "These can accept both the total solver and the block solver:" ] }, { "cell_type": "code", "execution_count": null, "id": "36", "metadata": {}, "outputs": [], "source": [ "args_total_solver = (base_region_boundaries, f_vertices, mg_builder, source, lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)\n", "\n", "print(\"Result\", solve_and_sum(*args_total_solver))\n", "\n", "print(\"Gradient\", grad_solve_and_sum(*args_total_solver))" ] }, { "cell_type": "code", "execution_count": null, "id": "37", "metadata": {}, "outputs": [], "source": [ "args_block_solver = (base_region_boundaries, f_vertices, mg_builder, source, block_solver, mg_chain)\n", "print(\"Result\", solve_and_sum(*args_block_solver))\n", "print(\"Gradient\", grad_solve_and_sum(*args_block_solver))" ] }, { "cell_type": "markdown", "id": "38", "metadata": {}, "source": [ "Again, up to numerical precision, both results are identical! We can see that the adjoint solve of the block solver goes the other direction than the forward solve (solving the upper triangular transpose)." ] }, { "cell_type": "markdown", "id": "39", "metadata": {}, "source": [ "This concludes the brief tour of the entire code. The {ref}`examples` directory contains some more involved examples." ] } ], "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 }