scattering

Functions related to scattering matrix computation for the FMM algorithm.

Copyright (c) Meta Platforms, Inc. and affiliates.

class fmmax.scattering.ScatteringMatrix(s11: Array, s12: Array, s21: Array, s22: Array, start_layer_solve_result: LayerSolveResult, start_layer_thickness: Array, end_layer_solve_result: LayerSolveResult, end_layer_thickness: Array)

Stores the scattering matrix for a stack of layers.

The first layer in a stack is the “start” layer, and the last layer in the stack is the “end” layer.

The scattering matrix relates the forward-going and backward-going waves on the two sides of a layer stack, which are labeled a and b respectively.

Note that forward going fields are defined at the start of a layer while backward-going fields are defined at the end of a layer, as depicted below. This is discussed near equation 4.1 in [1999 Whittaker].

| | | |

layer 0 | layer 1 | … | layer N |

start layer | | | end layer |

| | | | -> a_0 -> a_N b_0 <- b_N <-

Following the convention of [1999 Whittaker], the terms a_N and b_0 are obtained from,

a_N = s11 @ a_0 + s12 @ b_N b_0 = s21 @ a_0 + s22 @ b_N

Besides the actual scattering matrix element, the ScatteringMatrix stores information about the start and end layers, which are needed to extend the scattering matrix to include more layers.

s11

Relates forward-going fields at start to forward-going fields at end.

Type:

jax.Array

s12

Relates backward-going fields at end to forward-going fields at end.

Type:

jax.Array

s21

Relates forward-going fields at start to backward-going fields at start.

Type:

jax.Array

s22

Relates backward-going fields at end to backward-going fields at start.

Type:

jax.Array

start_layer_solve_result

The eigensolve result for the start layer.

Type:

fmmax.fmm.LayerSolveResult

start_layer_thickness

The start layer thickness.

Type:

jax.Array

end_layer_solve_result

The eigensolve result for the end layer.

Type:

fmmax.fmm.LayerSolveResult

end_layer_thickness

The end layer thickness.

Type:

jax.Array

fmmax.scattering.append_layer(s_matrix: ScatteringMatrix, next_layer_solve_result: LayerSolveResult, next_layer_thickness: Array, force_x64_solve: bool = False) → ScatteringMatrix

Returns new scattering matrix for the stack with an appended layer.

Parameters:

• s_matrix – The existing scattering matrix.

• next_layer_solve_result – The eigensolve result for the layer to append.

• next_layer_thickness – The thickness for the layer to append.

• force_x64_solve – If True, matrix solves will be done with 64 bit precision.

Returns:

The new ScatteringMatrix.

fmmax.scattering.prepend_layer(s_matrix: ScatteringMatrix, next_layer_solve_result: LayerSolveResult, next_layer_thickness: Array, force_x64_solve: bool = False) → ScatteringMatrix

Returns new scattering matrix for the stack with a prepended layer.

Parameters:

• s_matrix – The existing scattering matrix.

• next_layer_solve_result – The eigensolve result for the layer to append.

• next_layer_thickness – The thickness for the layer to append.

• force_x64_solve – If True, matrix solves will be done with 64 bit precision.

Returns:

The new ScatteringMatrix.

fmmax.scattering.redheffer_star_product(a: ScatteringMatrix, b: ScatteringMatrix, force_x64_solve: bool = False) → ScatteringMatrix

Compute the Redheffer star product of two scattering matrices.

fmmax.scattering.set_end_layer_thickness(s_matrix: ScatteringMatrix, thickness: Array) → ScatteringMatrix

Returns a new ScatteringMatrix with a modified end layer thickness.

Parameters:

• s_matrix – The initial ScatteringMatrix.

• thickness – The desired thickness of the layer.

Returns:

The new ScatteringMatrix.

fmmax.scattering.set_start_layer_thickness(s_matrix: ScatteringMatrix, thickness: Array) → ScatteringMatrix

Returns a new ScatteringMatrix with a modified start layer thickness.

Parameters:

• s_matrix – The initial ScatteringMatrix.

• thickness – The desired thickness of the layer.

Returns:

The new ScatteringMatrix.

fmmax.scattering.stack_s_matrices_interior(layer_solve_results: Sequence[LayerSolveResult], layer_thicknesses: Sequence[Array], force_x64_solve: bool = False) → Tuple[Tuple[ScatteringMatrix, ScatteringMatrix], ...]

Computes scattering matrices before and after each layer in the stack.

Specifically, for each layer i two ScatteringMatrix are returned. The first relates fields in the substack 0, …, i, while the second relates the fields in the substack i, …, N, where N is the maximum layer index. These two scattering matrices are denoted as the corresponding to the “before” substack and the “after” substack.

Parameters:

• layer_solve_results – The eigensolve results for layers in the stack.

• layer_thicknesses – The thicknesses for layers in the stack.

• force_x64_solve – If True, matrix solves will be done with 64 bit precision.

Returns:

The tuple of (scattering_matrix_before, scattering_matrix_after).

fmmax.scattering.stack_s_matrix(layer_solve_results: Sequence[LayerSolveResult], layer_thicknesses: Sequence[Array], force_x64_solve: bool = False) → ScatteringMatrix

Computes the s-matrix for a stack of layers.

If only a single layer is provided, the scattering matrix is just the identity matrix, and start and end layer data is for the same layer.

Parameters:

• layer_solve_results – The eigensolve results for layers in the stack.

• layer_thicknesses – The thicknesses for layers in the stack.

• force_x64_solve – If True, matrix solves will be done with 64 bit precision.

Returns:

The ScatteringMatrix.

fmmax.scattering.stack_s_matrix_scan(layer_solve_results: LayerSolveResult, layer_thicknesses: Array, force_x64_solve: bool = False) → ScatteringMatrix

Computes the scattering matrix for a stack of layers by a scan operation.

Unlike stack_s_matrix, this function uses a scan operation rather than a python for loop, which can lead to significantly smaller programs and shorter compile times. However, it requires that the layer solve results and thicknesses for be represented in the leading batch dimension of the arguments.

This function is best used when the eigensolve for each layer is also done in a batched manner.

Parameters:

• layer_solve_results – The layer solve results for all layers in the stack.

• layer_thicknesses – The layer thicknesses for all layers in the stack.

• force_x64_solve – If True, matrix solves will be done with 64 bit precision.

Returns:

The scattering matrix for the stack.