Sweeping a single angle¶
This notebook will show a transport operator inversion and the sweep of a single angle in both 1D and 2D.
import os
os.environ["JAX_PLATFORM_NAME"] = "cpu"
import jax
import jax.numpy as jnp
import numpy as onp
We are first going to create a simple slab domain in 1D:
n_vertices = 7
vertices = jnp.linspace(0, 1, n_vertices)[:, None]
connectivity = jnp.array([[i, i + 1] for i in range(n_vertices - 1)])
mesh = (vertices, connectivity)
from jax_sn.domain import CrossSectionData, CrossSectionsIndexed, Domain
from jax_sn.solution_domain import BasixLagrangianSimplex
cross_section = CrossSectionData(
total = jnp.array([[5.0]]),
scattering= jnp.array([[[[0.0]]]]))
cross_sections_indexed = CrossSectionsIndexed(indices = jnp.full((n_vertices - 1,), 0), cross_section_data = cross_section)
domain = Domain.from_mesh_and_cross_sections(mesh, BasixLagrangianSimplex(1,1).face_template, cross_sections_indexed) # face template is [[0],[1]] first face is first node, second is second.
Let’s pick a particular angle to invert the transport operator
angle = jnp.array([-1.0, 0.0, 0.0])
Technically, a transport operator may require different data: a level-based sweep is different then a topologically ordered one. Here just create the basic variant:
from jax_sn.operators.transport_operator import BasicSweepData
sweep_data = BasicSweepData.from_domain_angle(domain, angle)
print("Data structure ", sweep_data)
print("Levels of each element:", sweep_data.levels)
print("Topological sort: ", sweep_data.levels_argsort)
Data structure BasicSweepData(levels=i32[6], levels_argsort=i32[6])
Levels of each element: [5 4 3 2 1 0]
Topological sort: [5 4 3 2 1 0]
Unsurprisingly, with a negative angle, the level of the left-most element is 18 (5) while the topological sort determines that the first element to be solved is 18 (i.e. the right-most-element).
We can also visualize this (later we will do the same for 2D which is actually interesting):
from jax_sn.tools.plotly import plot_domain
plot_domain(domain, show_graph=True, graph_angle = jnp.array([1.0,0.0,0.0]), show_levels=True)
As discussed in the preceding spatial discretization notebook, the transport operator requires a spatial discretization, so we create a SolutionDomain object containing all spatial discretization data.
from jax_sn.solution_domain import BasixLagrangianSimplex, SolutionDomain
element_order = 1
element = BasixLagrangianSimplex(element_order, 1)
solution_domain = SolutionDomain.from_element_and_domain(element, domain)
print(solution_domain)
SolutionDomain(
spatial_data=SpatialData(
domain=Domain(
vertices=f32[7,1],
connectivity=i32[6,2],
face_normals=f32[6,2,3],
etoe=weak_i32[6,2],
etof=weak_i32[6,2],
face_perm_idx=i32[6,2],
face_template=i32[2,1],
cross_sections_indexed=CrossSectionsIndexed(
indices=weak_i32[6],
cross_section_data=CrossSectionData(
total=f32[1,1], scattering=f32[1,1,1,1]
)
)
),
fem_matrices=FEMMatrices(
mass_matrix=f32[2,2],
integration_coeff=f32[2],
derivative_matrix=f32[1,2,2],
outflow_matrices=f32[2,2,2],
inflow_matrices=f32[2,1,2,2,2]
),
precomputed_geometric_data=PrecomputedGeometricData(
vol_jacobians=f32[6,1,1],
vol_jacobians_inv=f32[6,1,1],
vol_jacobians_det=f32[6],
face_jacobians_det=f32[6,2]
),
element=<jax_sn.solution_domain.basix_fem.BasixLagrangianSimplex object at 0x7f3803c13500>
)
)
Now let’s actually create the transport operator. This requires the spatial data in the solution domain, the energy group, the angle, and the sweep data for that angle (together the TransportData). Besides that, it needs a strategy for both forward and inversion applications (forward is required to differentiate through the inversion by the adjoint method). The inversion strategy builds the sweep data it needs (applying any knobs of its own along the way); in this case, that is just the basic version.
from jax_sn.operators.transport_operator import TransportOperatorSingleAngle, VmapForward, VmapSweep, TransportData
inversion_strategy = VmapSweep()
sweep_data = inversion_strategy.build_default_sweep_data(domain, angle)
print("Sweep data this strategy builds:", type(sweep_data).__name__)
transport_data = TransportData(spatial_data = solution_domain.spatial_data,
energy_group = 0,
angle = angle,
sweep_data = sweep_data)
transport_operator = TransportOperatorSingleAngle(data = transport_data,
strategy = VmapForward(),
inversion_strategy = VmapSweep())
Sweep data this strategy builds: BasicSweepData
This is somewhat cumbersome, so there are convenience functions to directly do this from just a solution domain and strategies:
Inverting the transport operator¶
The transport operator can just be inverted as
transport_operator.inv(jnp.ones_like(transport_operator.out_structure()))
Array([[2.3843722 , 2.3663392 ],
[2.3638349 , 2.3221047 ],
[2.316309 , 2.2197418 ],
[2.2063296 , 1.9828639 ],
[1.9518268 , 1.434704 ],
[1.3628811 , 0.16620493]], dtype=float32)
However, to compare with the analytical solution, we should first apply the mass operator to the source, since the Galerkin equation is:
such that
For the domain we constructed, the cross section is constant, and thus is the analytical solution of an entirely flat source of value \(q\)
def analytical_solution(x, angle, sigma_t, q ):
return jnp.where(angle[0] >0, q/sigma_t * (1.0 - jnp.exp(-sigma_t * x)), q/sigma_t * (1.0 - jnp.exp(-sigma_t * (1.0 - x))))
We can now plot the transport operator inversion result and compare it to the analytical solution:
rhs = jnp.ones_like(solution_domain.elements_basis_shape)
from jax_sn.operators.mass_operator import MassOperator, MassData, LaxMass
massop = MassOperator(MassData(solution_domain.spatial_data.fem_matrices,
solution_domain.spatial_data.precomputed_geometric_data),
strategy = LaxMass(), additional_shape = ())
psi = transport_operator.inv(massop.mv(rhs))
from jax_sn.tools.plotly import plot_1d_solution
plot_1d_solution(
solution_domain,
psi,
reference=lambda x: analytical_solution(x, angle, cross_section.total[0, 0], 1.0),
yaxis_title="ψ",
)
The discontinuous elements can be seen clearly here! It also matches the general shape of the analytical solution. Let’s use a higher order finite element now:
solution_domain_p3 = SolutionDomain.from_element_and_domain(BasixLagrangianSimplex(3,1), domain)
transport_operator_p3 = TransportOperatorSingleAngle(TransportData(solution_domain_p3.spatial_data, 0,angle,
VmapSweep().build_default_sweep_data(solution_domain_p3.domain, angle)),
VmapForward(), VmapSweep())
mass_operator_p3 = MassOperator(MassData(solution_domain_p3.spatial_data.fem_matrices,
solution_domain_p3.spatial_data.precomputed_geometric_data),
strategy = LaxMass(), additional_shape = ())
rhs_p3 = jnp.ones_like(solution_domain_p3.elements_basis_shape)
psi_p3 = transport_operator_p3.inv(mass_operator_p3.mv(rhs_p3))
plot_1d_solution(
solution_domain_p3,
psi_p3,
reference=lambda x: analytical_solution(x, angle, cross_section.total[0, 0], 1.0),
yaxis_title="ψ",
)
This matches a lot closer!
Constructing these operators is rather tedious, so see … for a convenient way to do this.
2D cross section¶
We will do the same problem for a circular cross section now of a Miller-like tokamak using the jax_sbgeom package.
from jax_sbgeom.flux_surfaces.miller_tokamak import MillerTokamak
import jax_sbgeom as jsb
tokamak = MillerTokamak(rminor = 1.0, rmajor = 1.2, elongation = 0.5, triangularity = 0.2)
We make a very low-res version to show the concept clearly:
mesh_tokamak = jsb.flux_surfaces.flux_surface_meshing.mesh_poloidal_plane(tokamak, jnp.linspace(0.0,1.0,4), 0.0, 7)
# Convert it to a 2D mesh!
mesh_tokamak = (jnp.stack([mesh_tokamak[0][:,0], mesh_tokamak[0][:,2]], axis=-1), mesh_tokamak[1])
cross_sections_tokamak = CrossSectionsIndexed(indices = jnp.full((mesh_tokamak[1].shape[0],), 0), cross_section_data = cross_section)
domain_tokamak = Domain.from_mesh_and_cross_sections(mesh_tokamak, BasixLagrangianSimplex(1, 2).face_template, cross_sections_tokamak)
plot_domain(domain_tokamak, width=800, height=800)
Now let’s pick an angle and show the corresponding graph:
angle = jnp.array([1/jnp.sqrt(2), 1/jnp.sqrt(2), 0.0])
plot_domain(domain_tokamak, color_by_material = False, show_graph = True, graph_angle = angle, show_levels=True, mesh_opacity = 0.3, width=800, height=800)
which is no longer trivial. Let us solve and show the transport operator solution for an angle with both 1 and 2 degree finite elements:
def _solve_and_show_2d(deg : int ):
solution_domain_tokamak = SolutionDomain.from_element_and_domain(BasixLagrangianSimplex(deg,2), domain_tokamak)
transport_operator = TransportOperatorSingleAngle(
TransportData(solution_domain_tokamak.spatial_data, 0, angle, VmapSweep().build_default_sweep_data(solution_domain_tokamak.domain, angle)), VmapForward(), VmapSweep())
mass_operator = MassOperator(MassData(solution_domain_tokamak.spatial_data.fem_matrices,
solution_domain_tokamak.spatial_data.precomputed_geometric_data),
strategy = LaxMass(), additional_shape = ())
psi_2d = transport_operator.inv(mass_operator.mv(jnp.ones_like(solution_domain_tokamak.elements_basis_shape)))
from jax_sn.tools.plotly import plot_2d_solution
return plot_2d_solution(
solution_domain_tokamak,
psi_2d,
width=800, height=800, elevate = True, show_domain_mesh = True, show_element_boundaries = True, angle =angle, n_subdivisions = 1 if deg == 1 else 6
)
_solve_and_show_2d(1)
_solve_and_show_2d(2)
Qualitatively, the result is similar to the 1D: increasing to a higher value in the direction of the angle.