{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Multigroup operator\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import jax\n", "jax.config.update(\"jax_enable_x64\", True)\n", "jax.config.update(\"jax_platforms\", \"cpu\")\n", "import jax.numpy as jnp" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "A multigroup operator again requires picking a strategy. However, to actually solve the multigroup equation:\n", "\n", "$$ (\\mathbf{I} - \\mathbf{M}\\mathbf{T}^{-1}\\mathbf{H}\\mathbf{D}\\mathbf{S})[\\phi_{g\\ell\\epsilon i}]=\\mathbf{M}\\mathbf{T}^{-1}\\mathbf{D}\\mathbf{H}[q_{g\\ell\\epsilon i}] $$\n", "\n", "there are actually *two* multigroup operators requires: one including scattering:\n", "\n", "$$ \\text{Multigroup operator}: (\\mathbf{I} - \\mathbf{M}\\mathbf{T}^{-1}\\mathbf{H}\\mathbf{D}\\mathbf{S})$$\n", "\n", "and one that does *not* include scattering:\n", "\n", "$$ \\text{Multigroup chain operator}: \\mathbf{M}\\mathbf{T}^{-1}\\mathbf{D}\\mathbf{H}$$\n", "\n", "They can both be constructed and have their own strategy types.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.multi_group_operator import MultiGroupStrategy, MultiGroupChainStrategy" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "Currently, the main implementations just use a composed strategy that applies a single-group strategy and then maps it over the groups (using a serial scan or a vmap), or in other words:\n", "\n", "$$ \\text{for all } g: \\quad \\mathbf{I} - \\underbrace{\\mathbf{M}_g \\mathbf{T}^{-1}_g \\mathbf{H}_g \\mathbf{D}_g}_{\\text{chain}}\\underbrace{\\mathbf{S}_g}_{\\text{gather-scatter}} $$\n", "\n", "Note that $\\mathbf{S}$ requires the full multi-group unknown vector: it gathers the scattering output of all groups into the specific group $g$ (i.e. maps $G \\times L \\times E \\times B \\to L\\times E \\times B$), while the other operators just act on a single-group unknown vector.\n", "\n", "This gather scatter is simply \n", "\n", "\n", "$$ \\mathbf{S}_g[\\phi_{\\tilde{g}\\ell ei}] = \\sum_{\\tilde{g}}\\sigma_{\\ell g \\tilde{g} e} \\phi_{\\tilde{g\\ell ei}}$$\n", "\n", "while the single-group multi-angle transport operator, moment, mass, and discrete operators were all discussed in previous notebooks." ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "Constructing a single-group chain strategy is usually done by combining the suboperators, where each one has its own strategy too (i.e., a 'tower' of strategies)" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.multi_group_operator import MultiGroupChainScanStrategy, MultiGroupChainGatherScanStrategy\n", "from jax_sn.operator_creation.default_strategies import *\n", "mg_chain_gather_strategy = MultiGroupChainGatherScanStrategy(\n", " single_group_chain_strategy = ComposedSingleGroupChainStrategy(\n", " transport_strategy = VmapAnglesForward(VmapForward()),\n", " transport_inversion_strategy = VmapAnglesInversion(VmapSweep()),\n", " moment_strategy = EinsumFluxMoment(),\n", " mass_strategy = LaxMass(),\n", " discrete_strategy = EinsumFluxMoment()\n", " ),\n", " single_group_gather_scatter_strategy = BroadcastedSingleGroupGather() \n", ")\n", "\n", "mg_chain_strategy = MultiGroupChainScanStrategy(\n", " single_group_chain_strategy = mg_chain_gather_strategy.single_group_chain_strategy,\n", ")\n", "print(mg_chain_gather_strategy)\n" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "This is somewhat cumbersome, and in practice you'd probably not construct it like this. We just show it here to illustrate how strategies are used internally.\n", "\n", "Instead, you could construct one of the defaults:" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operator_creation.default_strategies import DefaultMultiGroupStrategy\n", "\n", "print(\"Vmap strategy:\\n\", DefaultMultiGroupStrategy.default_multi_group_strategy())\n", "print(\"Warp strategy:\\n\", DefaultMultiGroupStrategy.default_multi_group_warp_strategy(batch_size = 128))" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "Recall that the `WarpLevelAngles` is a *fused* multi-angle strategy and therefore has no inner strategy!" ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "Let's now construct a MultiGroupOperator. First we again make a similar tokamak domain:" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "from jax_sbgeom.flux_surfaces.miller_tokamak import MillerTokamak\n", "import jax_sbgeom as jsb\n", "from jax_sn.domain import Domain, CrossSectionsIndexed, CrossSectionData\n", "from jax_sn.solution_domain import BasixLagrangianSimplex\n", "from jax_sn.tools.plotly import plot_domain\n", "\n", "tokamak = MillerTokamak(rminor = 1.0, rmajor = 1.2, elongation = 0.5, triangularity = 0.2)\n", "\n", "def construct_domain_solution_element_quadrature_set(n_s, n_theta):\n", "\n", "\n", " mesh_tokamak = jsb.flux_surfaces.flux_surface_meshing.mesh_poloidal_plane(tokamak, jnp.linspace(0.0,1.0,n_s), 0.0, n_theta)\n", " # Convert it to a 2D mesh!\n", " mesh_tokamak = (jnp.stack([mesh_tokamak[0][:,0], mesh_tokamak[0][:,2]], axis=-1), mesh_tokamak[1])\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", " water_idx, iron_idx, void_idx = 0, 1, 2\n", " cross_section_indices = jnp.concatenate([jnp.full((mesh_tokamak[1].shape[0]//3,), water_idx), jnp.full((mesh_tokamak[1].shape[0]//3,), iron_idx), jnp.full((mesh_tokamak[1].shape[0] - 2*(mesh_tokamak[1].shape[0]//3),), void_idx)])\n", " cross_sections_tokamak = CrossSectionsIndexed(indices = cross_section_indices, cross_section_data = CrossSectionData.from_list([water, iron, void]))\n", "\n", "\n", " element_order = 2\n", " solution_element = BasixLagrangianSimplex(element_order, 2)\n", "\n", " quadrature_order = 3\n", " from jax_sn.quadrature_set import create_tn_quadrature_set\n", " quadrature_set = create_tn_quadrature_set(quadrature_order)\n", " domain_tokamak = Domain.from_mesh_and_cross_sections(mesh_tokamak, solution_element.face_template, cross_sections_tokamak)\n", " return domain_tokamak, quadrature_set, solution_element\n", "\n", "\n", "\n", "n_s = 4 \n", "n_theta = 7\n", "\n", "domain_tokamak, quadrature_set, solution_element = construct_domain_solution_element_quadrature_set(n_s, n_theta)\n", "\n", "plot_domain(domain_tokamak, width=800, height=800, material_names = [\"water\", \"iron\", \"void\"])" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "Then we construct the {py:class}`~jax_sn.operators.operator_data.OperatorData` object, which contains all data required for the multigroup operator. This contains both the {py:class}`~jax_sn.operators.operator_data.DiscretisationData` which includes all spatial/angular/energy discretisation information, and the {py:class}`~jax_sn.operators.operator_data.NonDifferentiableData`, which contains (non-operator-defining) data that cannot be differentiated and is more like setup data (i.e. topological ordering). The strategy builds that tier itself, since what it needs of it (here: a plain sweep ordering) is the strategy's business:" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": {}, "outputs": [], "source": [ "from jax_sn.operators.operator_data import OperatorData, DiscretisationData, SpatialData, AngularData, NonDifferentiableData\n", "from jax_sn.solution_domain import SolutionDomain\n", "from jax_sn.operators.multi_group_operator import MultiGroupOperator\n", "\n", "\n", "operator_data = OperatorData(\n", " DiscretisationData(\n", " SolutionDomain.from_element_and_domain(solution_element, domain_tokamak).spatial_data, \n", " AngularData.from_quadrature_set(quadrature_set, domain_tokamak.cross_sections_indexed.l_max)\n", " ),\n", " mg_chain_gather_strategy.build_default_non_differentiable_data(domain_tokamak, quadrature_set.angles)\n", ")\n", "\n", "mg_operator = MultiGroupOperator(operator_data, mg_chain_gather_strategy)\n", "\n", "print(mg_operator)" ] }, { "cell_type": "markdown", "id": "14", "metadata": {}, "source": [ "We can apply the multigroup operator as any other operator:" ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": {}, "outputs": [], "source": [ "print(\"Shape of application: [n_groups, n_moments, n_elements, n_basis]\",\n", " mg_operator.mv(jnp.ones_like(mg_operator.in_structure())).shape)" ] }, { "cell_type": "markdown", "id": "16", "metadata": {}, "source": [ "The meaning of an application is not as clear here though: in some sense, it means, given a flux $\\phi_{g\\ell ei}$, how much of that flux had to come from an external source (and not from scattering its own distribution).\n", "\n", "The main use of the multigroup operator is solving the multigroup equation, after which its result is just the flux moments." ] }, { "cell_type": "markdown", "id": "17", "metadata": {}, "source": [ "We have now shown every part of the multigroup operator construction. Before showing how to solve it and post-process it, we will show a much easier and more efficient way to construct it from the set of quantities that fully define it. \n", "\n" ] } ], "metadata": { "kernelspec": { "display_name": "blanketopt", "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.13.12" } }, "nbformat": 4, "nbformat_minor": 5 }