treams_core/smatrix/
periodic.rs

1//! Periodic transfer matrices of an S-matrix and the Bloch bands of its repetition.
2//!
3//! Upstream: `treams.SMatrices.periodic` and `treams.SMatrices.bands_kz`.
4
5mod saved;
6
7use nalgebra::DMatrix;
8
9use super::{Blocks, checked_dimension, dimension};
10use crate::{
11    Complex, Error, Result,
12    linalg::{
13        self, SolveGradient, SolveResidual, product, product_adjoint_right, product_views, view,
14    },
15    numerics::{finite, ratio},
16};
17
18/// What [`periodic()`] saves for its pullback: the upper rows of the transfer matrix, the
19/// reflection block `S10` and the solve with `S11`.
20#[derive(Debug)]
21pub struct PeriodicResidual {
22    top: DMatrix<Complex>,
23    reflection: DMatrix<Complex>,
24    solve: SolveResidual,
25}
26
27/// Convert scattering blocks to the transfer matrix used for periodic bands.
28///
29/// Upstream: `treams.SMatrices.periodic`.
30pub fn periodic(blocks: Blocks) -> Result<(DMatrix<Complex>, PeriodicResidual)> {
31    let n = dimension(&blocks)?;
32    let [up, reflection_up, reflection_down, down] = blocks;
33    let mut top = DMatrix::zeros(n, 2 * n);
34    top.columns_mut(0, n).copy_from(&up);
35    top.columns_mut(n, n).copy_from(&reflection_up);
36    let mut rhs = -product(&reflection_down, &top);
37    for i in 0..n {
38        rhs[(i, n + i)] += 1.0;
39    }
40    let solve = linalg::solve_owned(down, rhs)?;
41    let mut value = DMatrix::zeros(2 * n, 2 * n);
42    value.rows_mut(0, n).copy_from(&top);
43    value.rows_mut(n, n).copy_from(solve.value());
44    Ok((
45        value,
46        PeriodicResidual {
47            top,
48            reflection: reflection_down,
49            solve,
50        },
51    ))
52}
53
54impl PeriodicResidual {
55    /// The shape of the transfer matrix: both propagation directions of every
56    /// plane-wave mode, `2 n` by `2 n`.
57    #[must_use]
58    pub fn shape(&self) -> (usize, usize) {
59        (2 * self.top.nrows(), 2 * self.top.nrows())
60    }
61
62    /// Transfer-matrix tangent from the four scattering-block tangents, reusing LU.
63    pub fn pushforward(&self, blocks: &Blocks) -> Result<DMatrix<Complex>> {
64        let n = self.top.nrows();
65        checked_dimension(blocks, n, "invalid periodic transfer tangent")?;
66        let mut top = DMatrix::zeros(n, 2 * n);
67        top.columns_mut(0, n).copy_from(&blocks[0]);
68        top.columns_mut(n, n).copy_from(&blocks[1]);
69        let rhs = -product(&blocks[2], &self.top) - product(&self.reflection, &top);
70        let bottom = self.solve.pushforward(&blocks[3], rhs)?;
71        let mut value = DMatrix::zeros(2 * n, 2 * n);
72        value.rows_mut(0, n).copy_from(&top);
73        value.rows_mut(n, n).copy_from(&bottom);
74        Ok(value)
75    }
76
77    /// Pull back a transfer-matrix cotangent into all four scattering blocks.
78    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<Blocks> {
79        let n = self.top.nrows();
80        if cotangent.shape() != self.shape() || cotangent.iter().any(|&z| !finite(z)) {
81            return Err(Error::InvalidInput(
82                "invalid periodic transfer cotangent".into(),
83            ));
84        }
85        let SolveGradient {
86            operator: down,
87            rhs,
88        } = self.solve.pullback(cotangent.rows(n, n).into_owned())?;
89        let reflection = -product_adjoint_right(&rhs, &self.top);
90        let top =
91            cotangent.rows(0, n) - product_views(view(&self.reflection).adjoint(), view(&rhs));
92        Ok([
93            top.columns(0, n).into_owned(),
94            top.columns(n, n).into_owned(),
95            reflection,
96            down,
97        ])
98    }
99}
100
101/// What [`bands`] saves for its pullback: the residuals of the transfer matrix and of
102/// its eigensystem, the period and the Bloch wavenumbers.
103#[derive(Debug)]
104pub struct BandsResidual {
105    periodic: PeriodicResidual,
106    eigen: linalg::EigResidual,
107    period: f64,
108    wavenumbers: Vec<Complex>,
109}
110
111/// Bloch modes of a periodically repeated S matrix with positive repeat distance.
112///
113/// Upstream: `treams.SMatrices.bands_kz(az)` with `az = period`.
114pub fn bands(blocks: Blocks, period: f64) -> Result<BandsResidual> {
115    if !period.is_finite() || period <= 0.0 {
116        return Err(Error::InvalidInput(
117            "band period must be finite and positive".into(),
118        ));
119    }
120    let (transfer, periodic) = periodic(blocks)?;
121    let eigen = linalg::eig(&transfer)?;
122    let wavenumbers: Vec<_> = eigen
123        .values()
124        .iter()
125        .map(|v| -Complex::i() * v.ln() / period)
126        .collect();
127    if wavenumbers.iter().any(|&z| !finite(z)) {
128        return Err(Error::NonFinite(
129            "zero or non-finite Bloch multiplier".into(),
130        ));
131    }
132    Ok(BandsResidual {
133        periodic,
134        eigen,
135        period,
136        wavenumbers,
137    })
138}
139
140/// Cotangents of the inputs of [`bands`], in the order of its arguments.
141#[derive(Debug)]
142pub struct BandsGradient {
143    /// Cotangents of the four scattering blocks.
144    pub blocks: Blocks,
145    /// Cotangent of the repeat distance.
146    pub period: f64,
147}
148
149impl BandsResidual {
150    /// Normal Bloch wavenumbers from the principal complex logarithm, which the
151    /// pullback reads.
152    #[must_use]
153    pub fn wavenumbers(&self) -> &[Complex] {
154        &self.wavenumbers
155    }
156
157    /// Unit right Bloch eigenvectors, with the largest component real positive.
158    #[must_use]
159    pub const fn vectors(&self) -> &DMatrix<Complex> {
160        self.eigen.vectors()
161    }
162
163    /// Bloch wavenumber and eigenvector tangents with a fixed logarithm branch.
164    /// The recorded eigensystem defines the normalization and phase convention.
165    pub fn pushforward(
166        &self,
167        blocks: &Blocks,
168        period: f64,
169    ) -> Result<(Vec<Complex>, DMatrix<Complex>)> {
170        if !period.is_finite() {
171            return Err(Error::InvalidInput("invalid band period tangent".into()));
172        }
173        let transfer = self.periodic.pushforward(blocks)?;
174        let logarithm_derivatives: Vec<_> = self
175            .eigen
176            .values()
177            .iter()
178            .map(|&v| ratio(-Complex::i() / self.period, v))
179            .collect();
180        let (values, vectors) = self.eigen.pushforward(&transfer)?;
181        let wavenumbers = values
182            .iter()
183            .zip(&logarithm_derivatives)
184            .zip(&self.wavenumbers)
185            .map(|((&value, &derivative), &wavenumber)| {
186                derivative * value - wavenumber * (period / self.period)
187            })
188            .collect();
189        Ok((wavenumbers, vectors))
190    }
191
192    /// Return S-matrix and repeat-distance cotangents; logarithm branch is fixed.
193    pub fn pullback(
194        &self,
195        wavenumbers: &[Complex],
196        vectors: DMatrix<Complex>,
197    ) -> Result<BandsGradient> {
198        if wavenumbers.len() != self.wavenumbers.len() || wavenumbers.iter().any(|&z| !finite(z)) {
199            return Err(Error::InvalidInput(
200                "invalid Bloch wavenumber cotangent".into(),
201            ));
202        }
203        let period = -wavenumbers
204            .iter()
205            .zip(&self.wavenumbers)
206            .map(|(g, k)| (g.conj() * k).re)
207            .sum::<f64>()
208            / self.period;
209        let values: Vec<_> = wavenumbers
210            .iter()
211            .zip(self.eigen.values())
212            .map(|(g, l)| g * ratio(-Complex::i() / self.period, *l).conj())
213            .collect();
214        let matrix = self.eigen.pullback(&values, vectors)?;
215        Ok(BandsGradient {
216            blocks: self.periodic.pullback(&matrix)?,
217            period,
218        })
219    }
220}
treams_core/smatrix/periodic.rs

treams_core/smatrix/
periodic.rs
