treams_core/smatrix/
illuminate.rs

1//! Illumination of a pair of adjacent stacks: the outgoing fields and the internal
2//! fields between them, with the pullback of both.
3//!
4//! Upstream: `treams.SMatrices.illuminate` with a second S-matrix (`smat=`).
5
6mod saved;
7
8use faer::MatRef;
9use nalgebra::DMatrix;
10use rayon::prelude::*;
11
12use super::{
13    Blocks, StoredBlock, any_nonfinite, checked_dimension,
14    solve::{AdjointSolve, InternalSolve},
15};
16use crate::{
17    Complex, Error, Result,
18    linalg::{product, product_adjoint_right, product_views, view},
19    numerics::finite,
20};
21
22/// What [`illuminate`] saves for its pullback: copies of the blocks of both stacks, the
23/// incoming fields, the internal-field solve and the downgoing internal fields.
24///
25/// The solve keeps a dense LU factorization unless it solved iteratively. The iterative
26/// (Krylov) solve needs at least 512 rows, at most 8 illumination columns, an internal
27/// operator that is a proven contraction, and convergence within 16 steps.
28#[derive(Debug)]
29pub struct IlluminateResidual {
30    lower: [StoredBlock; 4],
31    upper: [StoredBlock; 4],
32    incoming: [DMatrix<Complex>; 2],
33    solve: InternalSolve,
34    down: DMatrix<Complex>,
35}
36
37/// Rows from which the eight recorded input blocks are copied concurrently, when
38/// the Rayon pool has more than one thread.
39const PARALLEL_SNAPSHOT_ROWS: usize = 128;
40
41/// Rows from which a one-thread pool copies the recorded input blocks on its
42/// worker rather than on the calling thread.
43const WORKER_SNAPSHOT_ROWS: usize = 512;
44
45/// Rows from which the finite check of the illumination inputs scans the ten
46/// matrices in parallel.
47const PARALLEL_SCAN_ROWS: usize = 256;
48
49/// Illuminate a pair of stacks. Returns outgoing up/down and internal up/down fields;
50/// each input and output matrix has one column per independent illumination.
51///
52/// The residual keeps copies of the borrowed blocks, in their memory order, for the
53/// pullback.
54///
55/// Upstream: `treams.SMatrices.illuminate`, called as `lower.illuminate(illu, illu2, smat=upper)`.
56pub fn illuminate(
57    lower: [MatRef<'_, Complex>; 4],
58    upper: [MatRef<'_, Complex>; 4],
59    incoming: [DMatrix<Complex>; 2],
60) -> Result<([DMatrix<Complex>; 4], IlluminateResidual)> {
61    // The pages of fresh copies fault and fill at memory speed; copy the eight
62    // blocks concurrently before the solve takes the workers.
63    let store = |blocks: [MatRef<'_, Complex>; 4]| {
64        let [a, b, c, d] = blocks;
65        let ((a, b), (c, d)) = crate::threads::join(
66            || crate::threads::join(|| StoredBlock::copy_view(a), || StoredBlock::copy_view(b)),
67            || crate::threads::join(|| StoredBlock::copy_view(c), || StoredBlock::copy_view(d)),
68        );
69        [a, b, c, d]
70    };
71    // A one-thread pool gains nothing from concurrency. Large recorded snapshots
72    // are allocated on its worker, alongside the numerical work; small snapshots
73    // avoid dispatching to the worker just to copy their inputs.
74    let threshold = if crate::threads::current_num_threads() > 1 {
75        PARALLEL_SNAPSHOT_ROWS
76    } else {
77        WORKER_SNAPSHOT_ROWS
78    };
79    let (lower, upper) = if lower[0].nrows() >= threshold {
80        crate::threads::join(|| store(lower), || store(upper))
81    } else {
82        (
83            lower.map(StoredBlock::copy_view),
84            upper.map(StoredBlock::copy_view),
85        )
86    };
87    let (top, bottom, down, solve) = illumination_fields(
88        lower.each_ref().map(StoredBlock::view),
89        upper.each_ref().map(StoredBlock::view),
90        incoming.each_ref().map(view),
91    )?;
92    let value = [top, bottom, solve.value.clone(), down.clone()];
93    Ok((
94        value,
95        IlluminateResidual {
96            lower,
97            upper,
98            incoming,
99            solve,
100            down,
101        },
102    ))
103}
104
105/// The fields of [`illuminate`] without the data its pullback needs.
106/// Inputs can be row- or column-major; only the operator and thin fields are owned.
107///
108/// Upstream: `treams.SMatrices.illuminate`, called as
109/// `lower.illuminate(illu, illu2, smat=upper)`.
110pub fn illuminate_value(
111    lower: [MatRef<'_, Complex>; 4],
112    upper: [MatRef<'_, Complex>; 4],
113    incoming: [MatRef<'_, Complex>; 2],
114) -> Result<[DMatrix<Complex>; 4]> {
115    let (top, bottom, down, solve) = illumination_fields(lower, upper, incoming)?;
116    Ok([top, bottom, solve.value, down])
117}
118
119/// Outgoing up, outgoing down and internal down fields, and the internal up solve.
120type InternalFields = (
121    DMatrix<Complex>,
122    DMatrix<Complex>,
123    DMatrix<Complex>,
124    InternalSolve,
125);
126
127/// Validate the inputs and compute every field of [`illuminate`].
128fn illumination_fields(
129    lower: [MatRef<'_, Complex>; 4],
130    upper: [MatRef<'_, Complex>; 4],
131    incoming: [MatRef<'_, Complex>; 2],
132) -> Result<InternalFields> {
133    let n = lower[0].nrows();
134    let p = incoming[0].ncols();
135    let nonfinite = |a: &MatRef<'_, Complex>| any_nonfinite(*a);
136    if n == 0
137        || p == 0
138        || lower.iter().chain(&upper).any(|a| a.shape() != (n, n))
139        || incoming.iter().any(|a| a.shape() != (n, p))
140        || if n >= PARALLEL_SCAN_ROWS {
141            crate::threads::install(|| {
142                lower
143                    .par_iter()
144                    .chain(upper.par_iter())
145                    .chain(incoming.par_iter())
146                    .any(nonfinite)
147            })
148        } else {
149            lower.iter().chain(&upper).chain(&incoming).any(nonfinite)
150        }
151    {
152        return Err(Error::InvalidInput(
153            "require matching finite S matrices and mode-by-illumination inputs".into(),
154        ));
155    }
156    let direct = product_views(upper[3], incoming[1]);
157    let rhs = product_views(lower[0], incoming[0]) + product_views(lower[1], view(&direct));
158    let solve = InternalSolve::new(lower[1], upper[2], rhs)?;
159    let down = product_views(upper[2], view(&solve.value)) + direct;
160    let top = product_views(upper[0], view(&solve.value)) + product_views(upper[1], incoming[1]);
161    let bottom = product_views(lower[2], incoming[0]) + product_views(lower[3], view(&down));
162    if top
163        .iter()
164        .chain(bottom.iter())
165        .chain(down.iter())
166        .any(|&z| !finite(z))
167    {
168        return Err(Error::Singular);
169    }
170    Ok((top, bottom, down, solve))
171}
172
173/// Cotangents of the inputs of [`illuminate`], in the order of its arguments.
174#[derive(Debug)]
175pub struct IlluminateGradient {
176    /// Cotangents of the blocks of the lower S-matrix.
177    pub lower: Blocks,
178    /// Cotangents of the blocks of the upper S-matrix.
179    pub upper: Blocks,
180    /// Cotangents of the incoming up and down amplitudes.
181    pub incoming: [DMatrix<Complex>; 2],
182}
183
184impl IlluminateResidual {
185    /// Mode and independent-illumination counts.
186    #[must_use]
187    pub fn shape(&self) -> (usize, usize) {
188        self.solve.value.shape()
189    }
190
191    /// Propagate stack and incident-field tangents through the recorded solve.
192    /// Products remain mode-by-illumination arrays for a thin illumination batch.
193    pub fn pushforward(
194        &self,
195        lower: &Blocks,
196        upper: &Blocks,
197        incoming: &[DMatrix<Complex>; 2],
198    ) -> Result<[DMatrix<Complex>; 4]> {
199        let (n, _) = self.shape();
200        checked_dimension(lower, n, "invalid lower illumination S matrix tangent")?;
201        checked_dimension(upper, n, "invalid upper illumination S matrix tangent")?;
202        if incoming
203            .iter()
204            .any(|a| a.shape() != self.shape() || a.iter().any(|&z| !finite(z)))
205        {
206            return Err(Error::InvalidInput(
207                "invalid illumination input tangent".into(),
208            ));
209        }
210        let up = &self.solve.value;
211        let direct = product(&upper[2], up)
212            + product(&upper[3], &self.incoming[1])
213            + product_views(self.upper[3].view(), view(&incoming[1]));
214        let rhs = product(&lower[0], &self.incoming[0])
215            + product_views(self.lower[0].view(), view(&incoming[0]))
216            + product(&lower[1], &self.down)
217            + product_views(self.lower[1].view(), view(&direct));
218        let tangent_up =
219            self.solve
220                .solve_forward(self.lower[1].view(), self.upper[2].view(), rhs)?;
221        let tangent_down = direct + product_views(self.upper[2].view(), view(&tangent_up));
222        let top = product(&upper[0], up)
223            + product_views(self.upper[0].view(), view(&tangent_up))
224            + product(&upper[1], &self.incoming[1])
225            + product_views(self.upper[1].view(), view(&incoming[1]));
226        let bottom = product(&lower[2], &self.incoming[0])
227            + product_views(self.lower[2].view(), view(&incoming[0]))
228            + product(&lower[3], &self.down)
229            + product_views(self.lower[3].view(), view(&tangent_down));
230        Ok([top, bottom, tangent_up, tangent_down])
231    }
232
233    /// Return lower/upper S matrices and incoming up/down amplitude cotangents.
234    pub fn pullback(&self, cotangent: &[DMatrix<Complex>; 4]) -> Result<IlluminateGradient> {
235        if cotangent
236            .iter()
237            .any(|a| a.shape() != self.shape() || a.iter().any(|&z| !finite(z)))
238        {
239            return Err(Error::InvalidInput(
240                "invalid field coefficient cotangent".into(),
241            ));
242        }
243        let down =
244            &cotangent[3] + product_views(self.lower[3].view().adjoint(), view(&cotangent[1]));
245        let up = &cotangent[2]
246            + product_views(self.upper[0].view().adjoint(), view(&cotangent[0]))
247            + product_views(self.upper[2].view().adjoint(), view(&down));
248        let AdjointSolve {
249            adjoint: rhs,
250            value: internal_up,
251        } = self
252            .solve
253            .solve_adjoint(self.lower[1].view(), self.upper[2].view(), up)?;
254        let direct = down + product_views(self.lower[1].view().adjoint(), view(&rhs));
255        let incoming = [
256            product_views(self.lower[2].view().adjoint(), view(&cotangent[1]))
257                + product_views(self.lower[0].view().adjoint(), view(&rhs)),
258            product_views(self.upper[1].view().adjoint(), view(&cotangent[0]))
259                + product_views(self.upper[3].view().adjoint(), view(&direct)),
260        ];
261        let lower = [
262            product_adjoint_right(&rhs, &self.incoming[0]),
263            product_adjoint_right(&rhs, &self.down),
264            product_adjoint_right(&cotangent[1], &self.incoming[0]),
265            product_adjoint_right(&cotangent[1], &self.down),
266        ];
267        let upper = [
268            product_adjoint_right(&cotangent[0], internal_up),
269            product_adjoint_right(&cotangent[0], &self.incoming[1]),
270            product_adjoint_right(&direct, internal_up),
271            product_adjoint_right(&direct, &self.incoming[1]),
272        ];
273        Ok(IlluminateGradient {
274            lower,
275            upper,
276            incoming,
277        })
278    }
279}
treams_core/smatrix/illuminate.rs

treams_core/smatrix/
illuminate.rs
