treams_core/smatrix/
mod.rs

1//! Planar scattering matrices of plane-wave modes and their analytic pullbacks.
2//!
3//! Covers Redheffer composition, chiral Fresnel and Cartesian interfaces,
4//! propagation, multilayer stacks, internal illumination between stacks, periodic
5//! transfer matrices and Bloch bands, radiation of periodic multipole arrays, port
6//! transmittance and reflectance, and chirality-density forms.
7//!
8//! Upstream: `treams.SMatrices`, `treams.coeffs.fresnel` and `treams.chirality_density`
9//! (see [Upstream counterparts](#upstream-counterparts)).
10//!
11//! # Sides and directions
12//!
13//! - The *negative side* lies below the structure (`z < 0` for a z normal); in [`add`]
14//!   it holds the lower operand. The *positive side* lies above and holds the upper
15//!   operand.
16//! - Direction 0 is *up*, towards the positive side; direction 1 is *down*.
17//! - [`Blocks`] index outgoing, then incident direction: `[S00, S01, S10, S11]`.
18//! - [`fresnel`], [`interface()`] and [`layer_stack`] take their media in the order
19//!   (negative side, positive side).
20//! - [`tr()`] numbers its ports the other way: port 0 is the positive side, as in the
21//!   material order `(above, below)` of upstream `SMatrices`.
22//!
23//! # Upstream counterparts
24//!
25//! | treams-rs | treams |
26//! |---|---|
27//! | [`add`] | `SMatrices.add` |
28//! | [`periodic()`], [`bands`] | `SMatrices.periodic`, `SMatrices.bands_kz` |
29//! | [`from_array`] | `SMatrices.from_array` |
30//! | [`illuminate()`], [`illuminate_value`] | `SMatrices.illuminate` |
31//! | [`tr()`], [`tr_value`] | `SMatrices.tr` |
32//! | [`fresnel`] | `coeffs.fresnel` |
33//! | [`interface()`], [`propagation`] | `SMatrices.interface`, `SMatrices.propagation` |
34//! | [`layer_stack`] | `SMatrices.stack` of interfaces and propagations |
35//! | [`chirality_density`], [`oriented_chirality`] | `chirality_density` |
36//! | [`Blocks`] | the four `SMatrix` blocks of one `SMatrices` |
37#![allow(clippy::indexing_slicing)] // Four blocks, validated equal matrix dimensions.
38
39mod array;
40mod chirality;
41mod compose;
42mod illuminate;
43mod interface;
44mod layers;
45mod periodic;
46mod saved;
47mod solve;
48mod tr;
49
50use faer::MatRef;
51use nalgebra::DMatrix;
52
53pub use array::{Channels, FromArrayGradient, FromArrayResidual, from_array};
54pub use chirality::{
55    ChiralityDensityGradient, ChiralityDensityResidual, OrientedChiralityGradient,
56    OrientedChiralityResidual, chirality_density, oriented_chirality,
57};
58pub use compose::{AddGradient, AddResidual, add};
59pub use illuminate::{IlluminateGradient, IlluminateResidual, illuminate, illuminate_value};
60pub use interface::{
61    FresnelGradient, FresnelResidual, InterfaceGradient, InterfaceResidual, PropagationGradient,
62    PropagationResidual, fresnel, interface, propagation,
63};
64pub use layers::{LayerStackGradient, LayerStackResidual, layer_stack};
65pub use periodic::{BandsGradient, BandsResidual, PeriodicResidual, bands, periodic};
66pub use tr::{TrGradient, TrPorts, TrResidual, tr, tr_value};
67
68use crate::{Complex, Error, Result, linalg::view, numerics::finite};
69
70/// The four scattering blocks of one S-matrix.
71///
72/// These are the four `ScatteringBlock`s of one `treams_rs.SMatrix`, each `S[out][in]`
73/// indexed by propagation direction (0 = up, 1 = down).
74///
75/// The order is `[S00, S01, S10, S11]`: upward transmission, reflection of waves
76/// incident from the positive side (above), reflection of waves incident from the
77/// negative side (below), and downward transmission.
78///
79/// Upstream: the four `SMatrix` blocks of one `treams.SMatrices`.
80pub type Blocks = [DMatrix<Complex>; 4];
81
82/// Owned square block that preserves its input memory order.
83#[derive(Debug)]
84struct StoredBlock {
85    values: DMatrix<Complex>,
86    row_major: bool,
87}
88
89/// The matrix in its memory order: the transpose of row-major storage, which then
90/// reads column by column, and whether it was transposed.
91fn storage_order(matrix: MatRef<'_, Complex>) -> (MatRef<'_, Complex>, bool) {
92    if matrix.row_stride() != 1 && matrix.col_stride() == 1 {
93        (matrix.transpose(), true)
94    } else {
95        (matrix, false)
96    }
97}
98
99/// Whether any entry is not finite, scanning contiguous storage without early exit
100/// so the scan vectorizes.
101fn any_nonfinite(matrix: MatRef<'_, Complex>) -> bool {
102    let (matrix, _) = storage_order(matrix);
103    if let Some(contiguous) = matrix.try_as_col_major() {
104        (0..matrix.ncols()).any(|j| {
105            contiguous
106                .col(j)
107                .as_slice()
108                .iter()
109                .fold(false, |bad, &z| bad | !finite(z))
110        })
111    } else {
112        (0..matrix.ncols()).any(|j| (0..matrix.nrows()).any(|i| !finite(matrix[(i, j)])))
113    }
114}
115
116impl StoredBlock {
117    /// Copy `matrix` column by column in its memory order.
118    fn copy_view(matrix: MatRef<'_, Complex>) -> Self {
119        let (matrix, row_major) = storage_order(matrix);
120        let mut values = Vec::with_capacity(matrix.nrows() * matrix.ncols());
121        if let Some(contiguous) = matrix.try_as_col_major() {
122            for j in 0..matrix.ncols() {
123                values.extend_from_slice(contiguous.col(j).as_slice());
124            }
125        } else {
126            for j in 0..matrix.ncols() {
127                for i in 0..matrix.nrows() {
128                    values.push(matrix[(i, j)]);
129                }
130            }
131        }
132        Self {
133            values: DMatrix::from_vec(matrix.nrows(), matrix.ncols(), values),
134            row_major,
135        }
136    }
137
138    /// The copied matrix, transposed back when it was row-major.
139    fn view(&self) -> MatRef<'_, Complex> {
140        let data = view(&self.values);
141        if self.row_major {
142            data.transpose()
143        } else {
144            data
145        }
146    }
147}
148
149/// The common dimension `n` of four finite, nonempty `n`-by-`n` blocks.
150fn dimension(blocks: &Blocks) -> Result<usize> {
151    let n = blocks[0].nrows();
152    if n == 0
153        || blocks
154            .iter()
155            .any(|m| m.shape() != (n, n) || m.iter().any(|&z| !finite(z)))
156    {
157        return Err(Error::InvalidInput(
158            "S matrices require four finite, equally sized nonempty square blocks".into(),
159        ));
160    }
161    Ok(n)
162}
163
164/// The dimension of cotangent blocks, which must match the forward `expected` one.
165fn checked_dimension(g: &Blocks, expected: usize, message: &str) -> Result<usize> {
166    if dimension(g)? != expected {
167        return Err(Error::InvalidInput(message.into()));
168    }
169    Ok(expected)
170}
171
172/// Scattering blocks of `n` unscattered modes: identity transmission, no reflection.
173pub(crate) fn identity_blocks(n: usize) -> Blocks {
174    std::array::from_fn(|b| {
175        if b == 0 || b == 3 {
176            DMatrix::identity(n, n)
177        } else {
178            DMatrix::zeros(n, n)
179        }
180    })
181}