treams_core/smatrix/
array.rs

1//! Plane-port radiation of periodic multipole arrays.
2//!
3//! Upstream: `treams.SMatrices.from_array`.
4
5mod saved;
6
7use faer::MatRef;
8use nalgebra::DMatrix;
9
10use super::{Blocks, checked_dimension, dimension};
11use crate::{
12    Complex, Error, Result,
13    linalg::{product, product_adjoint_right, product_views, view},
14    numerics::finite,
15};
16
17/// The four multipole-by-plane-wave channel arrays of [`from_array`], in the order
18/// incident up, incident down, emitted up, emitted down.
19///
20/// Each array has one row per multipole and one column per plane-wave mode. Unlike
21/// [`Blocks`], the arrays need not be square.
22pub type Channels = [DMatrix<Complex>; 4];
23
24/// What [`from_array`] saves for its pullback: the response, the channel arrays and
25/// the scattered fields.
26#[derive(Debug)]
27pub struct FromArrayResidual {
28    response: DMatrix<Complex>,
29    channels: Channels,
30    scattered: [DMatrix<Complex>; 2],
31}
32
33/// Compose an effective response with incident/emitted, up/down channel arrays.
34///
35/// Returns four plane-wave scattering blocks, including direct transmission.
36///
37/// Upstream: `treams.SMatrices.from_array`, given the channels of
38/// [`spherical_channels`](crate::channels::spherical_channels) or
39/// [`cylindrical_channels`](crate::channels::cylindrical_channels).
40pub fn from_array(
41    response: DMatrix<Complex>,
42    channels: Channels,
43) -> Result<(Blocks, FromArrayResidual)> {
44    let d = response.nrows();
45    let c = channels[0].ncols();
46    if d == 0
47        || !response.is_square()
48        || c == 0
49        || response.iter().any(|&v| !finite(v))
50        || channels
51            .iter()
52            .any(|a| a.shape() != (d, c) || a.iter().any(|&v| !finite(v)))
53    {
54        return Err(Error::InvalidInput(
55            "require a finite square response and four matching multipole-by-plane channel arrays"
56                .into(),
57        ));
58    }
59    let scattered = std::array::from_fn(|direction| product(&response, &channels[direction]));
60    let mut value = std::array::from_fn(|b| {
61        product_views(
62            view(&channels[2 + b / 2]).transpose(),
63            view(&scattered[b % 2]),
64        )
65    });
66    for block in [0, 3] {
67        for i in 0..c {
68            value[block][(i, i)] += 1.0;
69        }
70    }
71    dimension(&value).map_err(|_| Error::NonFinite("non-finite array scattering matrix".into()))?;
72    Ok((
73        value,
74        FromArrayResidual {
75            response,
76            channels,
77            scattered,
78        },
79    ))
80}
81
82/// Cotangents of the inputs of [`from_array`], in the order of its arguments.
83#[derive(Debug)]
84pub struct FromArrayGradient {
85    /// Cotangent of the effective multipole response.
86    pub response: DMatrix<Complex>,
87    /// Cotangents of the four channel arrays.
88    pub channels: Channels,
89}
90
91impl FromArrayResidual {
92    /// The shape `(n, n)` of each scattering block, with `n` plane-wave modes in each
93    /// direction.
94    #[must_use]
95    pub fn shape(&self) -> (usize, usize) {
96        (self.channels[0].ncols(), self.channels[0].ncols())
97    }
98
99    /// Multipole and plane-wave mode counts of the channel inputs.
100    #[must_use]
101    pub fn input_shape(&self) -> (usize, usize) {
102        self.channels[0].shape()
103    }
104
105    /// Apply response and channel tangents to the recorded scattered fields.
106    pub fn pushforward(&self, response: &DMatrix<Complex>, channels: &Channels) -> Result<Blocks> {
107        if response.shape() != self.response.shape()
108            || response.iter().any(|&z| !finite(z))
109            || channels
110                .iter()
111                .any(|a| a.shape() != self.input_shape() || a.iter().any(|&z| !finite(z)))
112        {
113            return Err(Error::InvalidInput(
114                "invalid array scattering tangent".into(),
115            ));
116        }
117        let scattered: [DMatrix<Complex>; 2] = std::array::from_fn(|direction| {
118            product(response, &self.channels[direction])
119                + product(&self.response, &channels[direction])
120        });
121        Ok(std::array::from_fn(|b| {
122            product_views(
123                view(&channels[2 + b / 2]).transpose(),
124                view(&self.scattered[b % 2]),
125            ) + product_views(
126                view(&self.channels[2 + b / 2]).transpose(),
127                view(&scattered[b % 2]),
128            )
129        }))
130    }
131
132    /// Effective-response and four-channel cotangents.
133    pub fn pullback(&self, cotangent: &Blocks) -> Result<FromArrayGradient> {
134        let c = checked_dimension(
135            cotangent,
136            self.channels[0].ncols(),
137            "invalid array scattering cotangent shape",
138        )?;
139        // Blocks are C_{2+e}^T S_s with S_s = R C_s, so the emission cotangents
140        // conjugate without transposing: bar S_s = conj(C_2) g_s + conj(C_3) g_{2+s}.
141        let conjugate =
142            |a: &DMatrix<Complex>, b: MatRef<'_, Complex>| product_views(view(a).conjugate(), b);
143        let mut response = DMatrix::zeros(self.response.nrows(), self.response.ncols());
144        let mut channels: Channels =
145            std::array::from_fn(|_| DMatrix::zeros(self.response.nrows(), c));
146        for direction in 0..2 {
147            let adjoint = conjugate(&self.channels[2], view(&cotangent[direction]))
148                + conjugate(&self.channels[3], view(&cotangent[2 + direction]));
149            response += product_adjoint_right(&adjoint, &self.channels[direction]);
150            channels[direction] = product_views(view(&self.response).adjoint(), view(&adjoint));
151            channels[2 + direction] = conjugate(
152                &self.scattered[0],
153                view(&cotangent[2 * direction]).transpose(),
154            ) + conjugate(
155                &self.scattered[1],
156                view(&cotangent[2 * direction + 1]).transpose(),
157            );
158        }
159        Ok(FromArrayGradient { response, channels })
160    }
161}
treams_core/smatrix/array.rs

treams_core/smatrix/
array.rs
