treams_core/smatrix/
compose.rs

1//! Redheffer composition of two adjacent S-matrices.
2//!
3//! Upstream: `treams.SMatrices.add`.
4
5mod saved;
6
7use std::ops::AddAssign;
8
9use faer::MatRef;
10use nalgebra::DMatrix;
11
12use super::{
13    Blocks, checked_dimension, dimension,
14    solve::{AdjointSolve, InternalSolve},
15};
16use crate::{
17    Complex, Error, Result,
18    linalg::{product, product_adjoint_right, product_views, view, view_mut},
19};
20
21/// What [`add`] saves for its pullback: four of the eight input blocks, the internal
22/// fields and their factorization.
23///
24/// The pullback reads only blocks 1 and 3 of the lower S-matrix and blocks 0 and 2
25/// of the upper one; the other four input blocks are released after the forward pass.
26#[derive(Debug)]
27pub struct AddResidual {
28    /// Lower blocks 1 and 3: reflection of waves incident from the positive side (above)
29    /// and downward transmission.
30    lower: [DMatrix<Complex>; 2],
31    /// Upper blocks 0 and 2: upward transmission and reflection of waves incident from
32    /// the negative side (below).
33    upper: [DMatrix<Complex>; 2],
34    /// Upgoing internal fields `(I - L1 U2)^-1 [L0, L1 U3]` and their factorization.
35    solve: InternalSolve,
36    /// Downgoing internal fields `U2 up + [0, U3]`.
37    down: DMatrix<Complex>,
38}
39
40/// Place `upper` above `lower`, eliminating their internal incident fields.
41///
42/// Upstream: `treams.SMatrices.add`, called as `lower.add(upper)`.
43pub fn add(lower: Blocks, upper: Blocks) -> Result<(Blocks, AddResidual)> {
44    let n = dimension(&lower)?;
45    if dimension(&upper)? != n {
46        return Err(Error::InvalidInput("S matrix dimensions must match".into()));
47    }
48    let [l0, l1, mut l2, l3] = lower;
49    let [u0, mut u1, u2, u3] = upper;
50    let mut rhs = DMatrix::zeros(n, 2 * n);
51    rhs.columns_mut(0, n).copy_from(&l0);
52    rhs.columns_mut(n, n).copy_from(&product(&l1, &u3));
53    // The Krylov solve needs at least ITERATIVE_MIN_ROWS = 512 rows and at most
54    // ITERATIVE_MAX_COLUMNS = 8 right-hand sides, so 2n columns never qualify and
55    // the direct solve gives the same bits as `InternalSolve::new`.
56    let solve = InternalSolve::direct(view(&l1), view(&u2), rhs)?;
57    let up = view(&solve.value);
58    let mut down = product_views(view(&u2), up);
59    down.columns_mut(n, n).add_assign(&u3);
60    let down_view = view(&down);
61    let value = [
62        product_views(view(&u0), up.subcols(0, n)),
63        {
64            accumulate_product(&mut u1, view(&u0), up.subcols(n, n));
65            u1
66        },
67        {
68            accumulate_product(&mut l2, view(&l3), down_view.subcols(0, n));
69            l2
70        },
71        product_views(view(&l3), down_view.subcols(n, n)),
72    ];
73    dimension(&value).map_err(|_| Error::Singular)?;
74    Ok((
75        value,
76        AddResidual {
77            lower: [l1, l3],
78            upper: [u0, u2],
79            solve,
80            down,
81        },
82    ))
83}
84
85/// `target += left * right`.
86fn accumulate_product(
87    target: &mut DMatrix<Complex>,
88    left: MatRef<'_, Complex>,
89    right: MatRef<'_, Complex>,
90) {
91    let (m, n, k) = (target.nrows(), target.ncols(), left.ncols());
92    crate::threads::product(m, n, k, |par| {
93        faer::linalg::matmul::matmul(
94            view_mut(target),
95            faer::Accum::Add,
96            left,
97            right,
98            Complex::new(1.0, 0.0),
99            par,
100        );
101    });
102}
103
104/// Cotangents of the inputs of [`add`], in the order of its arguments.
105#[derive(Debug)]
106pub struct AddGradient {
107    /// Cotangents of the blocks of the lower S-matrix.
108    pub lower: Blocks,
109    /// Cotangents of the blocks of the upper S-matrix.
110    pub upper: Blocks,
111}
112
113impl AddResidual {
114    /// The shape `(n, n)` of each scattering block, with `n` modes in each direction.
115    #[must_use]
116    pub fn shape(&self) -> (usize, usize) {
117        (self.down.nrows(), self.down.nrows())
118    }
119
120    /// Propagate all eight block tangents through the recorded internal-field solve.
121    /// The forward factorization is reused; no dense Jacobian is formed.
122    pub fn pushforward(&self, lower: &Blocks, upper: &Blocks) -> Result<Blocks> {
123        let n = self.down.nrows();
124        checked_dimension(lower, n, "invalid lower S matrix tangent")?;
125        checked_dimension(upper, n, "invalid upper S matrix tangent")?;
126        let up = view(&self.solve.value);
127        let mut direct = product_views(view(&upper[2]), up);
128        direct.columns_mut(n, n).add_assign(&upper[3]);
129        let mut rhs = product_views(view(&lower[1]), view(&self.down))
130            + product_views(view(&self.lower[0]), view(&direct));
131        rhs.columns_mut(0, n).add_assign(&lower[0]);
132        let tangent_up =
133            self.solve
134                .solve_forward(view(&self.lower[0]), view(&self.upper[1]), rhs)?;
135        let tangent_down = direct + product(&self.upper[1], &tangent_up);
136        let mut top = product_views(view(&upper[0]), up) + product(&self.upper[0], &tangent_up);
137        top.columns_mut(n, n).add_assign(&upper[1]);
138        let mut bottom = product(&lower[3], &self.down) + product(&self.lower[1], &tangent_down);
139        bottom.columns_mut(0, n).add_assign(&lower[2]);
140        Ok([
141            top.columns(0, n).into_owned(),
142            top.columns(n, n).into_owned(),
143            bottom.columns(0, n).into_owned(),
144            bottom.columns(n, n).into_owned(),
145        ])
146    }
147
148    /// Input cotangents under `dL = Re(sum(conj(g) * dx))`.
149    pub fn pullback(&self, cotangent: &Blocks) -> Result<AddGradient> {
150        let n = checked_dimension(
151            cotangent,
152            self.down.nrows(),
153            "invalid S matrix cotangent shape",
154        )?;
155        let [reflection, transmission] = &self.lower;
156        let [top_transmission, bottom_reflection] = &self.upper;
157        let mut top = DMatrix::zeros(n, 2 * n);
158        top.columns_mut(0, n).copy_from(&cotangent[0]);
159        top.columns_mut(n, n).copy_from(&cotangent[1]);
160        let mut bottom = DMatrix::zeros(n, 2 * n);
161        bottom.columns_mut(0, n).copy_from(&cotangent[2]);
162        bottom.columns_mut(n, n).copy_from(&cotangent[3]);
163        let down = product_views(view(transmission).adjoint(), view(&bottom));
164        let adjoint = product_views(view(top_transmission).adjoint(), view(&top))
165            + product_views(view(bottom_reflection).adjoint(), view(&down));
166        let AdjointSolve { adjoint, value: up } =
167            self.solve
168                .solve_adjoint(view(reflection), view(bottom_reflection), adjoint)?;
169        let incident = down + product_views(view(reflection).adjoint(), view(&adjoint));
170        let lower = [
171            adjoint.columns(0, n).into_owned(),
172            product_adjoint_right(&adjoint, &self.down),
173            cotangent[2].clone(),
174            product_adjoint_right(&bottom, &self.down),
175        ];
176        let upper = [
177            product_adjoint_right(&top, up),
178            cotangent[1].clone(),
179            product_adjoint_right(&incident, up),
180            incident.columns(n, n).into_owned(),
181        ];
182        Ok(AddGradient { lower, upper })
183    }
184}
treams_core/smatrix/compose.rs

treams_core/smatrix/
compose.rs
