treams_core/tmatrix/
metric.rs

1//! Circular dichroism, duality breaking and electromagnetic chirality of helicity T-matrices.
2//!
3//! Upstream: `treams.TMatrix.cd`, `treams.TMatrix.db` and `treams.TMatrix.chi`.
4#![allow(clippy::indexing_slicing)] // Validated square matrix, polarizations and helicity groups.
5
6mod saved;
7
8use nalgebra::DMatrix;
9
10use crate::linalg::SvdvalsResidual;
11use crate::{Complex, Error, Result, numerics::finite};
12
13/// Normalized measures of a global helicity T-matrix.
14#[derive(Clone, Copy, Debug)]
15pub enum Metric {
16    /// Absorption circular dichroism in a real embedding medium. Upstream:
17    /// `treams.TMatrix.cd`.
18    CircularDichroism,
19    /// Fraction of scattering that changes helicity. Upstream: `treams.TMatrix.db`.
20    DualityBreaking,
21    /// Electromagnetic chirality from the singular values of the four helicity blocks.
22    /// Upstream: `treams.TMatrix.chi`.
23    Chirality,
24}
25
26/// What [`metric`] saves for its pullback: the gradient itself, since the value is a
27/// real scalar; no singular vectors are kept.
28#[derive(Debug)]
29pub struct MetricResidual {
30    /// The matrix gradient, or why it is undefined: chirality may be zero where its
31    /// derivative is undefined.
32    gradient: Result<DMatrix<Complex>>,
33    ks_gradient: [f64; 2],
34    dimension: usize,
35}
36
37/// Gradients of the inputs of [`metric`], in the order of its arguments.
38#[derive(Clone, Debug)]
39pub struct MetricGradient {
40    /// Gradient of the helicity T-matrix.
41    pub matrix: DMatrix<Complex>,
42    /// Gradients of the real embedding wavenumbers; zero except for circular
43    /// dichroism.
44    pub ks: [f64; 2],
45}
46
47impl MetricResidual {
48    /// The shape of the recorded T-matrix.
49    #[must_use]
50    pub const fn shape(&self) -> (usize, usize) {
51        (self.dimension, self.dimension)
52    }
53
54    /// The real metric tangent, from the saved scalar gradient and changes of the
55    /// T-matrix and embedding wavenumbers, using the real complex inner product.
56    pub fn pushforward(&self, matrix: &DMatrix<Complex>, ks: [f64; 2]) -> Result<f64> {
57        if matrix.shape() != (self.dimension, self.dimension)
58            || matrix.iter().any(|&z| !finite(z))
59            || ks.iter().any(|k| !k.is_finite())
60        {
61            return Err(Error::InvalidInput(
62                "metric tangents must be finite and match the recorded inputs".into(),
63            ));
64        }
65        if matrix.iter().all(|&z| z == Complex::default()) && ks.iter().all(|&k| k == 0.0) {
66            return Ok(0.0);
67        }
68        Ok(self
69            .gradient
70            .as_ref()
71            .map_err(Clone::clone)?
72            .dotc(matrix)
73            .re
74            + self
75                .ks_gradient
76                .iter()
77                .zip(ks)
78                .map(|(g, d)| g * d)
79                .sum::<f64>())
80    }
81
82    /// Gradients of the matrix and the embedding wavenumbers from `cotangent`, the
83    /// gradient of a real loss with respect to the metric.
84    ///
85    /// The pullback scales the saved gradient on the calling thread, so the Rayon pool
86    /// does not change it.
87    pub fn pullback(&self, cotangent: f64) -> Result<MetricGradient> {
88        if !cotangent.is_finite() {
89            return Err(Error::InvalidInput(
90                "metric cotangent must be finite".into(),
91            ));
92        }
93        if cotangent == 0.0 {
94            return Ok(MetricGradient {
95                matrix: DMatrix::zeros(self.dimension, self.dimension),
96                ks: [0.0; 2],
97            });
98        }
99        Ok(MetricGradient {
100            matrix: self
101                .gradient
102                .as_ref()
103                .map_err(Clone::clone)?
104                .map(|z| z * cotangent),
105            ks: self.ks_gradient.map(|k| k * cotangent),
106        })
107    }
108}
109
110/// One [`Metric`] of a global helicity T-matrix.
111///
112/// `pol` is the polarization index of each row and `ks` holds the negative and positive
113/// helicity wavenumbers of the embedding, which enter circular dichroism only.
114///
115/// Upstream: `treams.TMatrix.cd`, `treams.TMatrix.db` and `treams.TMatrix.chi`.
116/// Differences: circular dichroism requires positive real wavenumbers where treams
117/// requires a real material; the other two scale the matrix to a largest entry of 1
118/// first; the chirality gradient is an error at zero helicity contrast.
119pub fn metric(
120    matrix: &DMatrix<Complex>,
121    pol: &[u8],
122    ks: [f64; 2],
123    kind: Metric,
124) -> Result<(f64, MetricResidual)> {
125    let dimension = matrix.nrows();
126    if dimension == 0
127        || !matrix.is_square()
128        || pol.len() != dimension
129        || pol.iter().any(|&p| p > 1)
130        || matrix.iter().any(|&z| !finite(z))
131        || ks.iter().any(|k| !k.is_finite() || *k <= 0.0)
132    {
133        return Err(Error::InvalidInput("require a finite square helicity matrix, matching polarizations and positive real wavenumbers".into()));
134    }
135    if matches!(kind, Metric::CircularDichroism) {
136        return circular_dichroism(matrix, pol, ks);
137    }
138    // Normalized metrics are invariant to a common matrix scale. Work at unit
139    // scale so norm squares do not overflow or underflow for weak scattering.
140    let scale = matrix.iter().map(|z| z.norm()).fold(0.0, f64::max);
141    if scale == 0.0 {
142        return Err(Error::InvalidInput(
143            "normalized metric requires nonzero scattering".into(),
144        ));
145    }
146    let matrix = matrix.map(|z| z / scale);
147    let (value, gradient) = if matches!(kind, Metric::DualityBreaking) {
148        let (value, gradient) = duality_breaking(&matrix, pol);
149        (value, Ok(gradient))
150    } else {
151        chirality(&matrix, pol)?
152    };
153    Ok((
154        value,
155        MetricResidual {
156            gradient: gradient.map(|g| g.map(|z| z / scale)),
157            ks_gradient: [0.0; 2],
158            dimension,
159        },
160    ))
161}
162
163/// Absorption circular dichroism with its matrix and wavenumber gradients.
164fn circular_dichroism(
165    matrix: &DMatrix<Complex>,
166    pol: &[u8],
167    ks: [f64; 2],
168) -> Result<(f64, MetricResidual)> {
169    let dimension = matrix.nrows();
170    let weights: Vec<_> = pol.iter().map(|&p| ks[usize::from(p)].powi(-2)).collect();
171    let mut absorption = [0.0; 2];
172    for j in 0..dimension {
173        absorption[usize::from(pol[j])] -= matrix[(j, j)].re * weights[j]
174            + (0..dimension)
175                .map(|i| matrix[(i, j)].norm_sqr() * weights[i])
176                .sum::<f64>();
177    }
178    let total = absorption.iter().sum::<f64>();
179    if total == 0.0 || !total.is_finite() {
180        return Err(Error::InvalidInput(
181            "circular dichroism requires nonzero total absorption".into(),
182        ));
183    }
184    let value = (absorption[1] - absorption[0]) / total;
185    let g = [
186        -2.0 * absorption[1] / total / total,
187        2.0 * absorption[0] / total / total,
188    ];
189    let mut gradient = DMatrix::from_fn(dimension, dimension, |i, j| {
190        -2.0 * g[usize::from(pol[j])] * weights[i] * matrix[(i, j)]
191    });
192    let mut gks = [0.0; 2];
193    for i in 0..dimension {
194        let p = usize::from(pol[i]);
195        gradient[(i, i)] -= g[p] * weights[i];
196        gks[p] += 2.0 * weights[i] / ks[p]
197            * (g[p] * matrix[(i, i)].re
198                + (0..dimension)
199                    .map(|j| g[usize::from(pol[j])] * matrix[(i, j)].norm_sqr())
200                    .sum::<f64>());
201    }
202    Ok((
203        value,
204        MetricResidual {
205            gradient: Ok(gradient),
206            ks_gradient: gks,
207            dimension,
208        },
209    ))
210}
211
212/// The fraction of scattering that flips helicity and its gradient, for a matrix at
213/// unit scale.
214fn duality_breaking(matrix: &DMatrix<Complex>, pol: &[u8]) -> (f64, DMatrix<Complex>) {
215    let dimension = matrix.nrows();
216    let norm = matrix.norm_squared();
217    let flipped = (0..dimension)
218        .flat_map(|j| (0..dimension).map(move |i| (i, j)))
219        .filter(|&(i, j)| pol[i] != pol[j])
220        .map(|(i, j)| matrix[(i, j)].norm_sqr())
221        .sum::<f64>();
222    let value = flipped / norm;
223    let gradient = DMatrix::from_fn(dimension, dimension, |i, j| {
224        matrix[(i, j)] * (2.0 * (f64::from(pol[i] != pol[j]) - value) / norm)
225    });
226    (value, gradient)
227}
228
229/// One helicity block of the chirality metric and its place in the singular-value
230/// differences.
231#[derive(Clone, Copy)]
232struct HelicityBlock {
233    /// Helicity of the rows.
234    row: usize,
235    /// Helicity of the columns.
236    col: usize,
237    /// The difference of spectra the block enters: 0 is `spectrum(1, 1) -
238    /// spectrum(0, 0)`, 1 is `spectrum(1, 0) - spectrum(0, 1)`.
239    difference: usize,
240    /// The sign of the block's spectrum in that difference.
241    sign: f64,
242}
243
244/// The four helicity blocks, ordered so that block `3 - p` minus block `p` is
245/// difference `p`, as treams pairs them in `TMatrix.chi`.
246const CHIRALITY_BLOCKS: [HelicityBlock; 4] = [
247    HelicityBlock {
248        row: 0,
249        col: 0,
250        difference: 0,
251        sign: -1.0,
252    },
253    HelicityBlock {
254        row: 0,
255        col: 1,
256        difference: 1,
257        sign: -1.0,
258    },
259    HelicityBlock {
260        row: 1,
261        col: 0,
262        difference: 1,
263        sign: 1.0,
264    },
265    HelicityBlock {
266        row: 1,
267        col: 1,
268        difference: 0,
269        sign: 1.0,
270    },
271];
272
273/// Electromagnetic chirality from the singular values of the four helicity blocks,
274/// for a matrix at unit scale. Its gradient is undefined at zero contrast.
275fn chirality(matrix: &DMatrix<Complex>, pol: &[u8]) -> Result<(f64, Result<DMatrix<Complex>>)> {
276    let indices: [Vec<_>; 2] = std::array::from_fn(|p| {
277        pol.iter()
278            .enumerate()
279            .filter_map(|(i, &q)| (usize::from(q) == p).then_some(i))
280            .collect()
281    });
282    let count = indices[0].len();
283    if count == 0 || count != indices[1].len() {
284        return Err(Error::InvalidInput(
285            "chirality requires equally sized nonempty helicity groups".into(),
286        ));
287    }
288    let block = |row: usize, col: usize| {
289        DMatrix::from_fn(count, count, |i, j| {
290            matrix[(indices[row][i], indices[col][j])]
291        })
292    };
293    let spectra = CHIRALITY_BLOCKS
294        .iter()
295        .map(|&HelicityBlock { row, col, .. }| crate::linalg::svdvals(&block(row, col)))
296        .collect::<Result<Vec<_>>>()?;
297    let differences: [Vec<_>; 2] = std::array::from_fn(|p| {
298        spectra[3 - p]
299            .values()
300            .iter()
301            .zip(spectra[p].values())
302            .map(|(a, b)| a - b)
303            .collect()
304    });
305    let contrast = differences
306        .iter()
307        .flatten()
308        .map(|x| x * x)
309        .sum::<f64>()
310        .sqrt();
311    let value = contrast / matrix.norm_squared().sqrt();
312    // The value stays available where the gradient is undefined.
313    let gradient = if contrast == 0.0 {
314        Err(Error::Derivative(
315            crate::DerivativeError::ZeroChiralityContrast,
316        ))
317    } else {
318        chirality_gradient(matrix, &indices, spectra, &differences, contrast)
319    };
320    Ok((value, gradient))
321}
322
323/// The chirality gradient through the singular values of the helicity blocks.
324fn chirality_gradient(
325    matrix: &DMatrix<Complex>,
326    indices: &[Vec<usize>; 2],
327    spectra: Vec<SvdvalsResidual>,
328    differences: &[Vec<f64>; 2],
329    contrast: f64,
330) -> Result<DMatrix<Complex>> {
331    let norm = matrix.norm_squared();
332    let value = contrast / norm.sqrt();
333    let mut gradient = matrix.map(|z| -value * z / norm);
334    for (
335        residual,
336        &HelicityBlock {
337            row,
338            col,
339            difference,
340            sign,
341        },
342    ) in spectra.into_iter().zip(&CHIRALITY_BLOCKS)
343    {
344        let weights: Vec<_> = differences[difference]
345            .iter()
346            .map(|x| sign * x / contrast / norm.sqrt())
347            .collect();
348        let block = residual.pullback(&weights)?;
349        for (j, &target_col) in indices[col].iter().enumerate() {
350            for (i, &target_row) in indices[row].iter().enumerate() {
351                gradient[(target_row, target_col)] += block[(i, j)];
352            }
353        }
354    }
355    Ok(gradient)
356}
treams_core/tmatrix/metric.rs

treams_core/tmatrix/
metric.rs
