Quadrature set and moments¶

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

\[\begin{split}\begin{aligned}\mathbf{M}[a_{\dots k \dots}] &= \sum_k w_k Y_{\ell k} a_{\dots k \dots}, \\[0.5em] \mathbf{D}[b_{\dots \ell \dots}] &= \sum_\ell Y_{\ell k} b_{\dots \ell \dots}, \\[0.5em]\end{aligned}\end{split}\]

in the code.

import jax
import jax.numpy as jnp
import numpy as onp
jax.config.update("jax_platforms", "cpu")
jax.config.update("jax_enable_x64", True)

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:

\[ \int_{4\pi} f(\Omega) d\Omega \approx \sum_k w_k f(\Omega_k)\]

We can create a quadrature set using either just manually creating angles and weights, or by selecting a particular precomputed quadrature set.

from jax_sn.quadrature_set import create_tn_quadrature_set

tn_quadrature_set = create_tn_quadrature_set(N  = 4)
print(tn_quadrature_set)

QuadratureSet(angles=f64[128,3], weights=f64[128])

We can plot the quadrature set, coloring the weight.

from jax_sn.tools.plotly import plot_quadrature_set

plot_quadrature_set(tn_quadrature_set)

Integration¶

To integrate a particular function, it is just a simple weighted sum over the quadrature set’s angles and weights:

\[ \int_{4\pi} f(\Omega)\, d\Omega \approx \sum_k w_k f(\Omega_k)\]

def integrate_over_sphere(f, quadrature_set):
    return jnp.sum(quadrature_set.weights * f(quadrature_set.angles))

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\):

\[\begin{split}\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}\end{split}\]

independent of \(n\).

def analytic_integral(p):
    return 0.0 if p % 2 else 4 * jnp.pi / (p + 1)

n = jnp.array([0.0, 1/jnp.sqrt(2), 1/jnp.sqrt(2)])

for p in range(12):
    f = lambda omega, p=p: jnp.dot(omega, n) ** p
    numeric  = float(integrate_over_sphere(f, tn_quadrature_set))
    analytic = analytic_integral(p)
    print(f"p={p:2d}   numeric={numeric:10.5f}   analytic={analytic:10.5f}   |diff|={abs(numeric - analytic):.2e}")

p= 0   numeric=  12.56637   analytic=  12.56637   |diff|=3.55e-15
p= 1   numeric=   0.00000   analytic=   0.00000   |diff|=0.00e+00
p= 2   numeric=   4.18879   analytic=   4.18879   |diff|=1.78e-15
p= 3   numeric=   0.00000   analytic=   0.00000   |diff|=2.43e-17
p= 4   numeric=   2.50798   analytic=   2.51327   |diff|=5.29e-03
p= 5   numeric=   0.00000   analytic=   0.00000   |diff|=0.00e+00
p= 6   numeric=   1.79049   analytic=   1.79520   |diff|=4.70e-03
p= 7   numeric=  -0.00000   analytic=   0.00000   |diff|=1.73e-18
p= 8   numeric=   1.39674   analytic=   1.39626   |diff|=4.79e-04
p= 9   numeric=   0.00000   analytic=   0.00000   |diff|=0.00e+00
p=10   numeric=   1.14931   analytic=   1.14240   |diff|=6.92e-03

p=11   numeric=   0.00000   analytic=   0.00000   |diff|=4.34e-19

The quadrature is only exact for \(p=1\); you have discretization errors for higher orders.

Moments¶

As mentioned in the 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:

The angular flux on the discrete angles \(\psi_{k} = \psi(\Omega_k)\)
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.

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

\[ \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} \]

This \(\mathcal{M} = w_k Y_{lm}(\Omega_k)\) matrix can be computed from the quadrature set alone (given a maximum moment order \(L\)).

Given these moments, we can re-compute the angular flux on any angle by expanding the spherical harmonics:

\[ \psi(\Omega) = \sum_{l=0}^{L} \sum_{m=-l}^{l} Y_{lm}(\Omega) \phi_{lm} = \mathcal{Y}\cdot \mathbf{\phi} \]

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:

\[ \mathbf{D}[\phi_{lm}] = \sum_{l=0}^{L} \sum_{m=-l}^{l} Y_{lm}(\Omega_k) \phi_{lm} = \mathcal{D}\cdot \mathbf{\phi} \]

Visualizing moments¶

Let’s use a transport solution from an isotropic slab:

\[\begin{split}\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}\end{split}\]

and compute the moments from the discrete angles and weights.

from jax_sn.analytical_solutions import psi_isotropic_box

box_size, sigma, q = 1.0, 2.0, 1.0
x0 = 0.3 * box_size  # off-centre, so psi differs between the +z and -z facing directions

psi_k = jnp.array([psi_isotropic_box(jnp.array([[x0]]), box_size, i, sigma, q) for i in tn_quadrature_set.angles])[:,0]

l_max = 4
moments = jnp.einsum("k, kl -> l", psi_k, tn_quadrature_set.wk_ylm_real(l_max))

We will now visualize these moments in three different ways:

The original function on the quadrature set angles
The spherical harmonic moments, shaded with how large they are after integrating.
The re-interpolated description of the angular flux on a fine grid:

\[ \psi(\Omega) = \sum_{l=0}^{L} \sum_{m=-l}^{l} Y_{lm}(\Omega) \phi_{lm}\]

Show code cell source

Hide code cell source

from plotly.subplots import make_subplots
from jax_sn.tools.plotly import plot_quadrature_set, plot_moments_spheres, plot_spherical_field, add_axis_triad

fig_angles  = plot_quadrature_set(tn_quadrature_set, values=psi_k, show_unit_sphere=True,
                                   colorbar_title="ψ(Ω)")
fig_moments = plot_moments_spheres(moments, l_max)
fig_field   = plot_spherical_field(l_max, moments, colorbar_title="Σ φ_lm·Y_lm(Ω)")

shared_max = max(float(jnp.max(jnp.abs(psi_k))), float(onp.max(onp.abs(fig_field.data[0].surfacecolor))))
fig_angles.data[-1].marker.update(cmin=0.0, cmax=shared_max, colorscale="Viridis")
fig_field.data[0].update(cmin=0.0, cmax=shared_max, colorscale="Viridis")

# keep each scene's own colour bar clear of the others
fig_angles.data[-1].marker.colorbar.update(x=1.05, len=0.28, y=0.90)
fig_moments.data[0].colorbar.update(x=1.05, len=0.28, y=0.50)
fig_field.data[0].colorbar.update(x=1.05, len=0.28, y=0.12)

fig = make_subplots(rows=3, cols=1, specs=[[{"type": "scene"}], [{"type": "scene"}], [{"type": "scene"}]],
                     row_heights=[0.32, 0.40, 0.28], vertical_spacing=0.06)

for trace in fig_angles.data:
    fig.add_trace(trace, row=1, col=1)
for trace in fig_moments.data:
    fig.add_trace(trace, row=2, col=1)
for trace in fig_field.data:
    fig.add_trace(trace, row=3, col=1)


locked_camera = dict(eye=dict(x=0, y=0, z=3.2), up=dict(x=0, y=1, z=0), projection=dict(type="orthographic"))
hidden = dict(visible=False)
fig.update_scenes(aspectmode="data", xaxis=hidden, yaxis=hidden, zaxis=hidden, camera=locked_camera)

add_axis_triad(fig, origin=(-1.5, -1.5, -1.3), length=0.7, row=1, col=1)
add_axis_triad(fig, origin=(-(l_max + 2.3) * 2.2, 1.3 * 2.2, 0.0), length=1.7, row=2, col=1)
add_axis_triad(fig, origin=(-1.5, -1.5, -1.3), length=0.7, row=3, col=1)

# connecting arrows for the M/D operators, centred in the gap between panels
gaps = [(fig.layout.scene.domain.y[0], fig.layout.scene2.domain.y[1]),
        (fig.layout.scene2.domain.y[0], fig.layout.scene3.domain.y[1])]
for (y_top, y_bottom), label in zip(gaps, ["M", "Sum"]):
    fig.add_shape(type="line", xref="paper", yref="paper", x0=0.5, x1=0.5, y0=y_top, y1=y_bottom,
                  line=dict(color="#8B0000", width=3))
    fig.add_annotation(x=0.5, y=y_bottom, xref="paper", yref="paper", showarrow=False, yanchor="top",
                        text="▼", font=dict(size=18, color="#8B0000"))
    fig.add_annotation(x=0.56, y=(y_top + y_bottom) / 2, xref="paper", yref="paper", showarrow=False,
                        text=f"<b>{label}</b>", font=dict(size=20, color="#8B0000"))

fig.update_layout(width=750, height=1650, showlegend=False, margin=dict(l=0, r=100, t=20, b=0))
fig

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.

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.

In the code¶

In the code, these operations converting between the two representations are standardized operators:

MomentOperator and 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 FluxMomentStrategy.

from jax_sn.operators.moment_operator import MomentOperator, DiscreteOperator, EinsumFluxMoment, MomentData, DiscreteData

moment_operator = MomentOperator(MomentData(tn_quadrature_set.wk_ylm_real(l_max)), (), EinsumFluxMoment()) 
#() is elements basis shape or additional batching dimensions, not used here since we have a single set of values instead 
# of defined on the whole domain.
discrete_operator = DiscreteOperator(DiscreteData(tn_quadrature_set.ylm_real(l_max)), (), EinsumFluxMoment())

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.

They compute the same as above:

onp.testing.assert_allclose(moment_operator.mv(psi_k), moments)

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.

try:
    onp.testing.assert_allclose(discrete_operator.mv(moment_operator.mv(psi_k)), psi_k)
except AssertionError as e:
    print("The discrete operator is not the inverse of the moment operator, as expected.")
    print(e)

The discrete operator is not the inverse of the moment operator, as expected.

Not equal to tolerance rtol=1e-07, atol=0

Mismatched elements: 128 / 128 (100%)
First 5 mismatches are at indices:
 [0]: 0.5366489649469993 (ACTUAL), 0.49999963830122846 (DESIRED)
 [1]: 0.536648964946999 (ACTUAL), 0.49999963830122846 (DESIRED)
 [2]: 0.5366489649469993 (ACTUAL), 0.49999963830122846 (DESIRED)
 [3]: 0.536648964946999 (ACTUAL), 0.49999963830122846 (DESIRED)
 [4]: 0.5139213247425937 (ACTUAL), 0.4988324528030406 (DESIRED)
Max absolute difference among violations: 0.03929608
Max relative difference among violations: 0.12660701
 ACTUAL: array([0.536649, 0.536649, 0.536649, 0.536649, 0.513921, 0.513921,
       0.513921, 0.513921, 0.471039, 0.471039, 0.471039, 0.471039,
       0.365837, 0.365837, 0.365837, 0.365837, 0.522193, 0.522193,...
 DESIRED: array([0.5     , 0.5     , 0.5     , 0.5     , 0.498832, 0.498832,
       0.498832, 0.498832, 0.470887, 0.470887, 0.470887, 0.470887,
       0.352179, 0.352179, 0.352179, 0.352179, 0.498684, 0.498684,...