treams_core/coeffs/
mie_cyl.rs

1//! Mie coefficients of multilayer chiral cylinders and their analytic pullback.
2//!
3//! Upstream: `treams.coeffs.mie_cyl`. The pullback is a treams-rs extension.
4#![allow(clippy::indexing_slicing)] // Fixed four-channel interface matrices.
5
6mod saved;
7
8use super::{LayerGradient, Material, Matrix2, Matrix4, Matrix42};
9use crate::{
10    Complex, Error, MAX_DEGREE, Result,
11    numerics::finite,
12    special::{Radial, RadialJet, cylindrical_radial},
13};
14
15/// One boundary between two media, as the pullback reads it.
16#[derive(Clone, Debug)]
17struct Interface {
18    /// The inner medium at the boundary.
19    inside: Side,
20    /// The outer medium at the boundary.
21    outside: Side,
22    /// The inverse of the outer boundary matrix.
23    inverse_outer: Matrix4,
24    /// The map from inner to outer coefficients: `inverse_outer` times the inner
25    /// boundary matrix.
26    transfer: Matrix4,
27    /// The product of the transfers inside this boundary, applied to the two regular
28    /// core channels.
29    before: Matrix42,
30}
31
32/// What [`mie_cyl`] saves for its pullback: the four-channel interface states and
33/// radial functions, so the pullback evaluates no Bessel functions.
34#[derive(Clone, Debug)]
35pub struct MieCylResidual {
36    order: i32,
37    kz: f64,
38    k0: f64,
39    radii: Vec<f64>,
40    materials: Vec<Material>,
41    interfaces: Vec<Interface>,
42    inverse: Matrix2,
43    value: Matrix2,
44}
45
46/// Gradients of the inputs of [`mie_cyl`], in the order of its arguments.
47#[derive(Clone, Debug)]
48pub struct MieCylGradient {
49    /// Gradient of the axial wavenumber.
50    pub kz: f64,
51    /// Gradient of the vacuum wavenumber.
52    pub k0: f64,
53    /// Gradients of the radii and the materials.
54    pub layers: LayerGradient,
55}
56
57/// A direction of change of one medium at one boundary: of its helicity wavenumbers,
58/// the axial wavenumber, the radius and its impedance.
59#[derive(Clone, Copy, Debug, Default)]
60struct Direction {
61    ks: [Complex; 2],
62    kz: Complex,
63    radius: Complex,
64    impedance: Complex,
65}
66
67/// One medium at a boundary: its helicity wavenumbers, impedance, radial wavenumbers
68/// `kr` and arguments `x = kr radius`, and the regular and (except in the core)
69/// singular radial functions of each helicity.
70#[derive(Clone, Copy, Debug)]
71struct Side {
72    ks: [Complex; 2],
73    impedance: Complex,
74    kr: [Complex; 2],
75    x: [Complex; 2],
76    radial: [[Option<RadialJet>; 2]; 2],
77}
78
79impl Side {
80    fn new(
81        order: i32,
82        kz: f64,
83        ks: [Complex; 2],
84        radius: f64,
85        impedance: Complex,
86        regular_only: bool,
87    ) -> Result<Self> {
88        let mut side = Self {
89            ks,
90            impedance,
91            kr: [Complex::default(); 2],
92            x: [Complex::default(); 2],
93            radial: [[None; 2]; 2],
94        };
95        for (pol, &k) in ks.iter().enumerate() {
96            let kr = (k * k - kz * kz).sqrt();
97            if k.norm_sqr() == 0.0 || kr.norm_sqr() == 0.0 {
98                return Err(Error::InvalidInput(
99                    "cylindrical cutoff k_rho=0 is not supported".into(),
100                ));
101            }
102            let x = kr * radius;
103            (side.kr[pol], side.x[pol]) = (kr, x);
104            let kinds: &[Radial] = if regular_only {
105                &[Radial::Regular]
106            } else {
107                &[Radial::Regular, Radial::Singular]
108            };
109            for (kind, &radial) in kinds.iter().enumerate() {
110                side.radial[pol][kind] = Some(cylindrical_radial(order, x, radial)?);
111            }
112        }
113        Ok(side)
114    }
115
116    /// The columns `(pol, kind)` of the boundary matrix with their radial function.
117    fn columns(&self) -> impl Iterator<Item = (usize, usize, RadialJet)> + '_ {
118        (0..2).flat_map(move |kind| {
119            (0..2).filter_map(move |pol| self.radial[pol][kind].map(|f| (pol, kind, f)))
120        })
121    }
122
123    /// The boundary matrix: tangential field components of each wave.
124    fn matrix(&self, order: i32, kz: f64) -> Matrix4 {
125        let mut value = Matrix4::zeros();
126        for (pol, kind, f) in self.columns() {
127            let (k, kr, x) = (self.ks[pol], self.kr[pol], self.x[pol]);
128            let phi = -kr * f.value / k;
129            let sign = if pol == 0 { 1.0 } else { -1.0 };
130            let z = -kz * f64::from(order) * f.value / (k * x) + sign * f.first;
131            let col = pol + 2 * kind;
132            value[(0, col)] = phi;
133            value[(1, col)] = z;
134            value[(2, col)] = -sign * phi / self.impedance;
135            value[(3, col)] = -sign * z / self.impedance;
136        }
137        value
138    }
139
140    /// The directional derivative of [`Self::matrix`], linear in the direction.
141    fn derivative(&self, order: i32, kz: f64, radius: f64, direction: Direction) -> Matrix4 {
142        let mut derivative = Matrix4::zeros();
143        for (pol, kind, f) in self.columns() {
144            let (k, kr, x) = (self.ks[pol], self.kr[pol], self.x[pol]);
145            let dk = direction.ks[pol];
146            let dkr = (k * dk - kz * direction.kz) / kr;
147            let dx = dkr * radius + kr * direction.radius;
148            let phi = -kr * f.value / k;
149            let dphi = -(dkr * f.value + kr * f.first * dx) / k - phi * dk / k;
150            let sign = if pol == 0 { 1.0 } else { -1.0 };
151            let z = -kz * f64::from(order) * f.value / (k * x) + sign * f.first;
152            let dz = -f64::from(order) / (k * x)
153                * (direction.kz * f.value + kz * f.first * dx - kz * f.value * (dk / k + dx / x))
154                + sign * f.second * dx;
155            let (impedance, dimpedance) = (self.impedance, direction.impedance);
156            let col = pol + 2 * kind;
157            derivative[(0, col)] = dphi;
158            derivative[(1, col)] = dz;
159            derivative[(2, col)] = -sign * (dphi - phi * dimpedance / impedance) / impedance;
160            derivative[(3, col)] = -sign * (dz - z * dimpedance / impedance) / impedance;
161        }
162        derivative
163    }
164}
165
166/// Mie coefficients of concentric chiral cylinders at the axial wavenumber `kz` and the
167/// azimuthal order `m` (`order`).
168///
169/// `radii` are the boundary radii, from inner to outer, and `materials` has one more
170/// entry, the embedding medium last.
171///
172/// Upstream: `treams.coeffs.mie_cyl`.
173pub fn mie_cyl(
174    kz: f64,
175    order: i32,
176    k0: f64,
177    radii: &[f64],
178    materials: &[Material],
179) -> Result<MieCylResidual> {
180    crate::coeffs::validate_layers(radii, materials)?;
181    if !kz.is_finite()
182        || !k0.is_finite()
183        || k0 <= 0.0
184        || order.unsigned_abs() > MAX_DEGREE.unsigned_abs()
185    {
186        return Err(Error::InvalidInput(
187            // 128 is MAX_DEGREE.
188            "require finite kz, positive k0 and |m| <= 128".into(),
189        ));
190    }
191    let mut q = Matrix42::zeros();
192    q[(0, 0)] = Complex::new(1.0, 0.0);
193    q[(1, 1)] = Complex::new(1.0, 0.0);
194    let mut interfaces = Vec::with_capacity(radii.len());
195    for (i, &radius) in radii.iter().enumerate() {
196        let side = |m: Material, regular_only| {
197            let ks = [k0 * (m.index() - m.kappa), k0 * (m.index() + m.kappa)];
198            Side::new(order, kz, ks, radius, m.impedance(), regular_only)
199        };
200        let inside = side(materials[i], i == 0)?;
201        let outside = side(materials[i + 1], false)?;
202        let inverse_outer = outside
203            .matrix(order, kz)
204            .try_inverse()
205            .ok_or(Error::Singular)?;
206        let transfer = inverse_outer * inside.matrix(order, kz);
207        interfaces.push(Interface {
208            inside,
209            outside,
210            inverse_outer,
211            transfer,
212            before: q,
213        });
214        q = transfer * q;
215    }
216    let inverse = q
217        .fixed_rows::<2>(0)
218        .into_owned()
219        .try_inverse()
220        .ok_or(Error::Singular)?;
221    let value = q.fixed_rows::<2>(2) * inverse;
222    if value.iter().any(|&z| !finite(z)) {
223        return Err(Error::Singular);
224    }
225    Ok(MieCylResidual {
226        order,
227        kz,
228        k0,
229        radii: radii.to_vec(),
230        materials: materials.to_vec(),
231        interfaces,
232        inverse,
233        value,
234    })
235}
236impl MieCylResidual {
237    /// The number of recorded layer boundaries.
238    #[must_use]
239    pub fn boundaries(&self) -> usize {
240        self.radii.len()
241    }
242
243    /// The coefficients in negative, positive helicity order, which the pullback reads.
244    #[must_use]
245    pub const fn value(&self) -> &Matrix2 {
246        &self.value
247    }
248
249    /// The coefficient tangent along changes of `kz`, `k0`, the radii and the
250    /// materials. Material entries store changes of epsilon, mu and kappa.
251    /// Cached interface inverses and radial derivatives are reused for one forward
252    /// pass, with no additional factorizations or Bessel evaluations.
253    pub fn pushforward(
254        &self,
255        kz: f64,
256        k0: f64,
257        radii: &[f64],
258        materials: &[Material],
259    ) -> Result<Matrix2> {
260        super::validate_layer_tangents(self.radii.len(), radii, materials)?;
261        if !kz.is_finite() || !k0.is_finite() {
262            return Err(Error::InvalidInput(
263                "wavenumber tangents must be finite".into(),
264            ));
265        }
266        let optical: Vec<_> = self
267            .materials
268            .iter()
269            .zip(materials)
270            .map(|(&material, &direction)| material.pushforward(direction))
271            .collect();
272        let mut dq = Matrix42::zeros();
273        for (i, interface) in self.interfaces.iter().enumerate() {
274            let derivative = |side: usize, boundary: &Side| {
275                let tangent = optical[i + side];
276                let dn = [tangent.index - tangent.kappa, tangent.index + tangent.kappa];
277                boundary.derivative(
278                    self.order,
279                    self.kz,
280                    self.radii[i],
281                    Direction {
282                        ks: std::array::from_fn(|pol| {
283                            boundary.ks[pol] * (k0 / self.k0) + self.k0 * dn[pol]
284                        }),
285                        kz: kz.into(),
286                        radius: radii[i].into(),
287                        impedance: tangent.impedance,
288                    },
289                )
290            };
291            let inner = derivative(0, &interface.inside);
292            let outer = derivative(1, &interface.outside);
293            let transfer = interface.inverse_outer * (inner - outer * interface.transfer);
294            dq = transfer * interface.before + interface.transfer * dq;
295        }
296        Ok((dq.fixed_rows::<2>(2) - self.value * dq.fixed_rows::<2>(0)) * self.inverse)
297    }
298
299    /// Gradients of `kz`, `k0`, the radii and the materials from `cotangent`, the
300    /// gradient of a real loss with respect to the coefficient matrix.
301    ///
302    /// The pullback runs on the calling thread and adds the boundaries in order, from
303    /// the outermost in.
304    pub fn pullback(&self, cotangent: &Matrix2) -> Result<MieCylGradient> {
305        if cotangent.iter().any(|&z| !finite(z)) {
306            return Err(Error::InvalidInput("cotangent must be finite".into()));
307        }
308        let mut result = MieCylGradient {
309            kz: 0.0,
310            k0: 0.0,
311            layers: LayerGradient::zeros(self.radii.len()),
312        };
313        let lower = cotangent * self.inverse.adjoint();
314        let mut gq = Matrix42::zeros();
315        gq.fixed_rows_mut::<2>(0)
316            .copy_from(&(-self.value.adjoint() * lower));
317        gq.fixed_rows_mut::<2>(2).copy_from(&lower);
318        for (i, interface) in self.interfaces.iter().enumerate().rev() {
319            let gf = gq * interface.before.adjoint();
320            let inner = interface.inverse_outer.adjoint() * gf;
321            let outer = -inner * interface.transfer.adjoint();
322            let sides = [(&interface.inside, inner), (&interface.outside, outer)];
323            for (side, (boundary, g)) in sides.into_iter().enumerate() {
324                let material = self.materials[i + side];
325                let n = material.index();
326                let contract = |direction: Direction| -> Complex {
327                    let df = boundary.derivative(self.order, self.kz, self.radii[i], direction);
328                    g.dotc(&df).conj()
329                };
330                let one = Complex::new(1.0, 0.0);
331                result.layers.radii[i] += contract(Direction {
332                    radius: one,
333                    ..Direction::default()
334                })
335                .re;
336                result.kz += contract(Direction {
337                    kz: one,
338                    ..Direction::default()
339                })
340                .re;
341                result.k0 += contract(Direction {
342                    ks: [n - material.kappa, n + material.kappa],
343                    ..Direction::default()
344                })
345                .re;
346                let layers = &mut result.layers;
347                let targets = [&mut layers.epsilon, &mut layers.mu, &mut layers.kappa];
348                for (tangent, target) in material.tangents().into_iter().zip(targets) {
349                    target[i + side] += contract(Direction {
350                        ks: [
351                            self.k0 * (tangent.index - tangent.kappa),
352                            self.k0 * (tangent.index + tangent.kappa),
353                        ],
354                        impedance: tangent.impedance,
355                        ..Direction::default()
356                    });
357                }
358            }
359            gq = interface.transfer.adjoint() * gq;
360        }
361        Ok(result)
362    }
363}
treams_core/coeffs/mie_cyl.rs

treams_core/coeffs/
mie_cyl.rs
