treams_core/
basis.rs

1//! Multipole bases of modes at expansion positions, and the expansion gradients shared by
2//! the spherical and cylindrical families.
3//!
4//! Upstream: `treams.SphericalWaveBasis` and `treams.CylindricalWaveBasis`. The gradients
5//! are a treams-rs extension.
6//!
7//! An expansion entry depends on the displacement `destination - source` between two
8//! positions. Its position cotangent therefore adds to the destination position and
9//! subtracts from the source position ([`ExpansionGradient`]).
10#![allow(clippy::indexing_slicing)] // Basis and matrix dimensions validated on entry.
11
12use crate::{Complex, Error, Result, numerics::finite};
13use nalgebra::DMatrix;
14
15/// The mode label of one multipole family: [`crate::sw::Mode`] or [`crate::cw::Mode`].
16pub trait ModeLabel: Copy + std::fmt::Debug {
17    /// Reject labels outside the supported range.
18    fn validate(self) -> Result<()>;
19    /// Polarization index: 1 or 0, as in treams.
20    fn pol(self) -> u8;
21}
22
23/// Modes of one multipole family at Cartesian expansion positions; [`crate::sw::Basis`]
24/// and [`crate::cw::Basis`] name the two families.
25///
26/// Upstream: `treams.SphericalWaveBasis` and `treams.CylindricalWaveBasis`.
27#[derive(Clone, Debug)]
28pub struct Basis<M> {
29    /// Position index (upstream `pidx`) and mode label of each matrix axis entry.
30    pub modes: Vec<(usize, M)>,
31    /// Expansion centres, one per position index (upstream `positions`).
32    pub positions: Vec<[f64; 3]>,
33}
34
35impl<M: ModeLabel> Basis<M> {
36    pub(crate) fn validate(&self) -> Result<()> {
37        if self.modes.is_empty() || self.positions.iter().flatten().any(|v| !v.is_finite()) {
38            return Err(Error::InvalidInput(
39                "basis must be nonempty with finite positions".into(),
40            ));
41        }
42        for &(pidx, mode) in &self.modes {
43            mode.validate()?;
44            if pidx >= self.positions.len() {
45                return Err(Error::InvalidInput(
46                    "basis position index outside positions".into(),
47                ));
48            }
49        }
50        Ok(())
51    }
52
53    /// Position index and polarization index of basis entry `i`.
54    pub(crate) fn position_pol(&self, i: usize) -> (usize, usize) {
55        let (pidx, mode) = self.modes[i];
56        (pidx, usize::from(mode.pol()))
57    }
58}
59
60/// Check the medium wavenumbers `ks` (negative, positive helicity) of a basis
61/// operation: both finite, nonzero unless `allow_zero`, and equal for parity
62/// polarizations (`helicity` false), because parity waves exist only in an achiral
63/// medium. Operations that never mix polarizations pass `helicity = true`.
64pub(crate) fn validate_wavenumbers(
65    ks: [Complex; 2],
66    helicity: bool,
67    allow_zero: bool,
68) -> Result<()> {
69    if ks
70        .iter()
71        .any(|&k| !finite(k) || (!allow_zero && k == Complex::default()))
72    {
73        return Err(Error::InvalidInput(
74            if allow_zero {
75                "medium wavenumbers must be finite"
76            } else {
77                "medium wavenumbers must be finite and nonzero"
78            }
79            .into(),
80        ));
81    }
82    if !helicity && ks[0] != ks[1] {
83        return Err(Error::InvalidInput(
84            "parity polarization requires equal wavenumbers (an achiral medium)".into(),
85        ));
86    }
87    Ok(())
88}
89
90/// A spherical or cylindrical basis, for operations that accept either family: field
91/// evaluation, plane-wave expansion and the EBCM surface waves.
92#[derive(Clone, Debug)]
93pub enum MultipoleBasis {
94    /// Spherical vector waves.
95    Spherical(crate::sw::Basis),
96    /// Cylindrical vector waves; axial wavenumbers are fixed mode labels.
97    Cylindrical(crate::cw::Basis),
98}
99
100impl From<crate::sw::Basis> for MultipoleBasis {
101    fn from(value: crate::sw::Basis) -> Self {
102        Self::Spherical(value)
103    }
104}
105
106impl From<crate::cw::Basis> for MultipoleBasis {
107    fn from(value: crate::cw::Basis) -> Self {
108        Self::Cylindrical(value)
109    }
110}
111
112impl MultipoleBasis {
113    pub(crate) fn validate(&self) -> Result<()> {
114        match self {
115            Self::Spherical(b) => b.validate(),
116            Self::Cylindrical(b) => b.validate(),
117        }
118    }
119
120    pub(crate) fn positions(&self) -> &[[f64; 3]] {
121        match self {
122            Self::Spherical(b) => &b.positions,
123            Self::Cylindrical(b) => &b.positions,
124        }
125    }
126
127    pub(crate) fn len(&self) -> usize {
128        match self {
129            Self::Spherical(b) => b.modes.len(),
130            Self::Cylindrical(b) => b.modes.len(),
131        }
132    }
133
134    pub(crate) fn position_pol(&self, i: usize) -> (usize, usize) {
135        match self {
136            Self::Spherical(b) => b.position_pol(i),
137            Self::Cylindrical(b) => b.position_pol(i),
138        }
139    }
140}
141
142/// Periodic coupling cotangents including the lattice geometry.
143#[derive(Debug)]
144pub struct LatticeExpansionGradient {
145    /// Position and medium-wavenumber cotangents.
146    pub expansion: ExpansionGradient,
147    /// Bloch wavevector cotangent in lattice coordinates.
148    pub kpar: Vec<f64>,
149    /// Row lattice-vector cotangent.
150    pub vectors: DMatrix<f64>,
151}
152
153impl LatticeExpansionGradient {
154    /// Zero cotangents of `destination` and `source` positions, both wavenumbers and a
155    /// `dim`-dimensional lattice.
156    pub(crate) fn zeros(destination: usize, source: usize, dim: usize) -> Self {
157        Self {
158            expansion: ExpansionGradient::zeros(destination, source),
159            kpar: vec![0.0; dim],
160            vectors: DMatrix::zeros(dim, dim),
161        }
162    }
163
164    /// Add the Bloch-vector and lattice-vector cotangents of one lattice sum.
165    pub(crate) fn add_lattice(&mut self, g: &crate::lattice::SumGradient) {
166        for (j, value) in self.kpar.iter_mut().enumerate() {
167            *value += g.kpar[j];
168            for i in 0..self.vectors.nrows() {
169                self.vectors[(i, j)] += g.vectors[i][j];
170            }
171        }
172    }
173}
174
175/// Expansion cotangents, separating destination and source position dependence.
176#[derive(Clone, Debug)]
177pub struct ExpansionGradient {
178    /// Destination Cartesian position derivatives.
179    pub destination: Vec<[f64; 3]>,
180    /// Source Cartesian position derivatives.
181    pub source: Vec<[f64; 3]>,
182    /// Complex wave number cotangents in negative, positive helicity order.
183    pub ks: [Complex; 2],
184}
185
186impl ExpansionGradient {
187    /// Zero cotangents of `destination` and `source` positions and both wavenumbers.
188    pub(crate) fn zeros(destination: usize, source: usize) -> Self {
189        Self {
190            destination: vec![[0.0; 3]; destination],
191            source: vec![[0.0; 3]; source],
192            ks: [Complex::default(); 2],
193        }
194    }
195
196    /// Add the cotangents of another pullback over the same positions.
197    pub(crate) fn accumulate(&mut self, other: &Self) {
198        for (a, b) in self
199            .destination
200            .iter_mut()
201            .chain(&mut self.source)
202            .flatten()
203            .zip(other.destination.iter().chain(&other.source).flatten())
204        {
205            *a += b;
206        }
207        for (a, b) in self.ks.iter_mut().zip(other.ks) {
208            *a += b;
209        }
210    }
211}
212
213/// Validate the continuous directions shared by finite and lattice expansions.
214pub(crate) fn validate_expansion_tangent(
215    destination: &[[f64; 3]],
216    source: &[[f64; 3]],
217    ks: [Complex; 2],
218    counts: (usize, usize),
219) -> Result<()> {
220    if (destination.len(), source.len()) != counts
221        || destination
222            .iter()
223            .chain(source)
224            .flatten()
225            .any(|x| !x.is_finite())
226        || ks.into_iter().any(|x| !finite(x))
227    {
228        return Err(Error::InvalidInput(
229            "expansion tangents must be finite and match inputs".into(),
230        ));
231    }
232    Ok(())
233}
234
235/// Contract an entry's local derivatives with its position and medium directions.
236pub(crate) fn pair_tangent(
237    destination: &[[f64; 3]],
238    source: &[[f64; 3]],
239    ks: [Complex; 2],
240    [p, q, pol]: [usize; 3],
241    position: [Complex; 3],
242    k: Complex,
243) -> Complex {
244    k * ks[pol]
245        + position
246            .into_iter()
247            .enumerate()
248            .map(|(axis, derivative)| derivative * (destination[p][axis] - source[q][axis]))
249            .sum::<Complex>()
250}
251
252/// Pull the cotangent `cot` of an entry depending on the displacement
253/// `destination[p] - source[q]` (derivative `position`) and the medium wavenumber
254/// of polarization `pol` (derivative `k`) back into `result`.
255pub(crate) fn add_pair_gradient(
256    result: &mut ExpansionGradient,
257    [p, q, pol]: [usize; 3],
258    cot: Complex,
259    position: [Complex; 3],
260    k: Complex,
261) {
262    result.ks[pol] += cot * k.conj();
263    for (axis, derivative) in position.into_iter().enumerate() {
264        let gradient = (cot.conj() * derivative).re;
265        result.destination[p][axis] += gradient;
266        result.source[q][axis] -= gradient;
267    }
268}
treams_core/basis.rs

treams_core/
basis.rs
