treams_core/lattice/
mod.rs

1//! Lattice sums of outgoing spherical and cylindrical waves, and lattice geometry.
2//!
3//! Upstream: `treams.lattice` (`lattice/_esum.pyx`, `_dsum.pyx` and `_misc.pyx` of
4//! tfp-photonics/treams, MIT), whose formulae and conventions the sums follow, and the
5//! Brillouin-zone helpers of `treams.misc`.
6//!
7//! # Definition
8//!
9//! A lattice of dimension `d` holds the points `R = n_1 a_1 + ... + n_d a_d` with integer
10//! `n_i`. For a wavenumber `k`, a real Bloch vector `kpar` (the wavevector components
11//! along the lattice) and a shift `r`, [`sum`] returns
12//!
13//! ```text
14//! D_lm(k, kpar, r) = Σ'_R h_l^(1)(k |r + R|) Y_lm(-r - R) exp(i kpar · R)        spherical
15//! D_m(k, kpar, r)  = Σ'_R H_m^(1)(k |r + R|) exp(i m φ(-r - R)) exp(i kpar · R)   cylindrical
16//! ```
17//!
18//! as in `treams.lattice`. `Y_lm(-r - R)` is the spherical harmonic in the direction of
19//! `-r - R` and `φ` its azimuth. The prime drops the point with `r + R = 0`, if there is
20//! one. [`Family`] selects the sum and its labels: degree `l` and order `m` for spherical
21//! sums (`d <= 3`), order `m` for cylindrical sums (`d <= 2`).
22//!
23//! # Frame
24//!
25//! The lattice lies along fixed Cartesian axes:
26//!
27//! | Sum | Lattice axes |
28//! |---|---|
29//! | 1D spherical | z |
30//! | 1D cylindrical | x |
31//! | 2D spherical and cylindrical | x and y (the plane `z = 0`) |
32//! | 3D spherical | x, y and z |
33//!
34//! [`BlochLattice`] takes the lattice vectors and `kpar` as components along these axes.
35//! The shift is a Cartesian 3-vector; cylindrical shifts lie in the xy plane.
36//!
37//! # Branch and Bloch sign
38//!
39//! The waves are outgoing: `h_l^(1)` and `H_m^(1)`, which decay for `Im k > 0`. The Bloch
40//! factor is `exp(+i kpar · R)`, so the sum at `r + R` is `exp(-i kpar · R)` times the
41//! sum at `r`. The sum depends on the lattice and on `kpar` modulo the reciprocal
42//! lattice, but not on the primitive basis or on the Bloch cell. Exact diffraction
43//! thresholds (`|kpar + G| = k` for a reciprocal lattice vector `G`) are singular.
44//!
45//! # Ewald split
46//!
47//! The Ewald split parameter `eta` divides each sum into a real-space part, whose terms
48//! fall like `exp(-Re((k eta)^2) |R|^2 / 2)`, and a reciprocal-space part, whose terms
49//! fall like `exp(-Re(1 / (k eta)^2) |kpar + G|^2 / 2)`. The exact sum does not depend
50//! on `eta`. `eta = 0` selects the automatic split
51//!
52//! ```text
53//! eta = sqrt(2 pi) / (k L) max(|k L| / 8, 1),    L = V^(1/d)
54//! ```
55//!
56//! with the cell measure `V` (length, area or volume). Derivatives hold the split fixed.
57//! The automatic split has the modulus `sqrt(2 pi) / 8 = 0.31` from `|k L| = 8` on, and
58//! more below. A split is small when it lies below every automatic one,
59//! `Re(1 / (2 eta^2)) > 16 / pi`: there the two parts grow like `exp(Re(1 / (2 eta^2)))`
60//! and cancel.
61//!
62//! Ewald sums require `Im k >= 0`, and spherical sums additionally require `Re k >= 0`.
63//! Within this domain every sum has the same value at every split where its Ewald parts
64//! converge. Inputs outside it give [`InvalidInput`]: the gain-side continuation is not
65//! supported, and negative-real spherical wavenumbers select the wrong Kambe branch.
66//! Finite direct shells accept every finite nonzero complex wavenumber.
67//!
68//! # Crosswalk
69//!
70//! Each `treams.lattice` function is one call of [`sum_part`] (or [`sum`] for `Full`).
71//! `r` is the shift, `i` the shell of a direct sum, and `dim` the lattice dimension of
72//! [`BlochLattice`].
73//!
74//! | treams | Family | `dim` | Shift |
75//! |---|---|---|---|
76//! | `lsumsw1d(l, k, kpar, a, r, eta)` | `Spherical { l, m: 0 }` | 1 | `[0, 0, r]` |
77//! | `lsumsw1d_shift(l, m, ...)` | `Spherical { l, m }` | 1 | `r` |
78//! | `lsumsw2d(l, m, ...)` | `Spherical { l, m }` | 2 | `[r_x, r_y, 0]` |
79//! | `lsumsw2d_shift(l, m, ...)` | `Spherical { l, m }` | 2 | `r` |
80//! | `lsumsw3d(l, m, ...)` | `Spherical { l, m }` | 3 | `r` |
81//! | `lsumcw1d(m, ...)` | `Cylindrical { m }` | 1 | `[r, 0, 0]` |
82//! | `lsumcw1d_shift(m, ...)` | `Cylindrical { m }` | 1 | `[r_x, r_y, 0]` |
83//! | `lsumcw2d(m, ...)` | `Cylindrical { m }` | 2 | `[r_x, r_y, 0]` |
84//!
85//! The prefix selects the [`SumPart`]:
86//!
87//! | treams prefix | [`SumPart`] |
88//! |---|---|
89//! | `lsum*` | `Full` |
90//! | `realsum*` | `Real` |
91//! | `recsum*`, with `zero2d` or `zero3d` added at `r = 0` | `Reciprocal`, which adds the self term at `r = 0` for every degree |
92//! | `dsum*(..., i)` | `Direct(i)`; `eta` is not used |
93//!
94//! The geometry helpers map one to one:
95//!
96//! | treams | treams-rs |
97//! |---|---|
98//! | `lattice.volume`, `lattice.area` | [`volume`] |
99//! | `lattice.reciprocal` | [`reciprocal()`] |
100//! | `lattice.cube(d, n)`, `lattice.cubeedge(d, n)` | [`cube`] with `edge` false or true |
101//! | `lattice.diffr_orders_circle(b, rmax)` | [`diffraction_orders`] |
102//! | `misc.firstbrillouin1d` | [`first_brillouin_1d`] |
103//! | `misc.firstbrillouin2d`, `misc.firstbrillouin3d` | [`first_brillouin`] |
104//!
105//! [`derivatives`], [`derivatives_part`] and [`SumResidual`] have no upstream
106//! counterpart: they give the derivatives and gradients of the sums.
107//!
108//! # Differences from `treams.lattice`
109//!
110//! - **Automatic split in 3D.** The cell length is `L = V^(1/3)`. Upstream `_check_eta`
111//!   divides by `k V^(2/3)`. Both choices give the same sum wherever the parts converge;
112//!   they differ in cost and in which sums converge.
113//! - **Convergence.** Each Ewald part adds integer cube shells until two consecutive
114//!   shells add less than `2e-13 max(|P|, 1)` of the part `P`, and fails after 200, 32
115//!   or 16 shells in one, two or three dimensions (up to four times that where a large
116//!   `|k| L` or a small `|k eta|` needs it). Upstream stops once two shells change a part
117//!   by less than `1e-10`, or after 200, 20 or 10 shells, and reports no failure.
118//! - **Reduced bases.** Both Ewald parts sum over shells of reduced bases, the shift is
119//!   reduced into a cell of the direct lattice and `kpar` into a cell of the reciprocal
120//!   lattice. Upstream sums over the given basis, whose shells can miss short lattice
121//!   vectors of a skewed basis.
122//! - **Outgoing sheets.** For `Im k >= 0` the reciprocal integrals at a complex split
123//!   take the sheet that continues the sum from the automatic split, so the value does
124//!   not depend on the split. Upstream takes the principal branches.
125//! - **Small splits.** Sums at small splits predict the rounding their cancelling parts
126//!   leave and fail where it exceeds `1e-3` of the scale of a component. Upstream returns
127//!   the cancelled sum.
128//! - **Exact zeros.** For odd `l` (spherical) or odd `m` (cylindrical), `kpar = 0` and a
129//!   shift on the lattice line or plane with `2 r` a lattice vector, inversion maps the
130//!   sum onto its negative. The sum is then exactly zero at every split. Upstream
131//!   returns exact zeros only at `r = 0` (and, in 1D, at `r = ±a / 2`); at the other
132//!   such shifts it returns the rounding of its parts.
133//! - **Spectral series.** Off the axis the Ewald parts of 1D spherical sums cancel by up
134//!   to `exp(w^2)` (see the glossary). There each component whose rounding exceeds
135//!   `1e-12 max(|S|, 1)` comes from the spectral series, a sum over diffraction orders
136//!   of cylindrical waves, where that is more accurate.
137//! - **Derivatives.** [`derivatives`] returns the analytic derivatives in `k`, `eta`,
138//!   the shift, `kpar` and the lattice vectors, from jets (see the glossary) and the
139//!   derivative identities of the integrals.
140//!
141//! # Glossary
142//!
143//! | Term | Meaning |
144//! |---|---|
145//! | shell | The integer points `n` with `max_i \|n_i\| = s` of shell `s`; sums add shell by shell |
146//! | split | The Ewald split parameter `eta` |
147//! | small split | A split below every automatic one, `Re(1 / (2 eta^2)) > 16 / pi` |
148//! | diffraction order | `q = kpar + G` for a reciprocal lattice vector `G` |
149//! | `v` | `(q · q / k^2 - 1) / (2 eta^2)` of a diffraction order `q` |
150//! | `w` | `k rho eta`, with the distance `rho` of the shift from the lattice plane or axis |
151//! | `t` | `w^2` |
152//! | `F_n(v, t)` | The reduced integral `(-1)^n ∫_1^∞ s^(n-1) exp(-v s - t / (2 s)) ds` of the reciprocal terms; `F_n(v, 0) = Γ(n, v) / (-v)^n` (upstream `_redincgamma`, `_redintkambe`) |
153//! | `I_n(z, eta)` | The Kambe integral `∫_eta^∞ t^n exp(-z^2 t^2 / 2 + 1 / (2 t^2)) dt` of the real-space terms, at `z = k \|r + R\|` (upstream `intkambe`) |
154//! | sheet | The branch of `ln v` on which a reciprocal term takes `F_n` |
155//! | self term | The correction at `r = 0` that removes the missing `R = 0` term from the reciprocal part; upstream adds it as `zero2d` or `zero3d` |
156//! | jet | A value carried through the arithmetic together with its first derivatives |
157//! | settled | A real-space shell sum that reaches its shell limit while the Bloch phases cancel each of its last two shells; it is kept only if it agrees with the sum at the automatic split |
158//!
159//! # Failure messages
160//!
161//! | Message | Error | Cause |
162//! |---|---|---|
163//! | `Ewald sums require Im(k) >= 0; gain media are unsupported` | [`InvalidInput`] | An amplifying wavenumber |
164//! | `spherical Ewald sums require Re(k) >= 0` | [`InvalidInput`] | A spherical wavenumber on the unsupported Kambe branch |
165//! | `Ewald sum did not converge within the shell limit` | [`NotConverged`] | A part needs more shells than its limit |
166//! | ... `; use a larger split (eta = 0 selects one)` | [`NotConverged`] | The same, at a small split |
167//! | ... `, neither at this split nor at the automatic one` | [`NotConverged`] | A settled sum whose check at the automatic split does not converge |
168//! | `Ewald split too small: its cancelling parts lose about ...; use a larger split (eta = 0 selects one)` | [`NotConverged`] | The parts at a small split cancel by more than about `1e-3` of the scale of the value or a derivative |
169//! | `Ewald sum lost its accuracy to cancelling Kambe integrals; reduce the split parameter` | [`NotConverged`] | A 1D spherical sum off the axis rounds by more than `1e-3 max(\|S\|, 1)`, and its spectral series does not help |
170//! | `non-finite Ewald summand; change the split parameter or reduce the order` | [`NonFinite`] | A term is not finite |
171//! | `non-finite direct lattice term` | [`NonFinite`] | A term of a direct sum is not finite |
172//! | `lattice sum is at a diffraction threshold; supply a limiting complex wavenumber` | [`InvalidInput`] | A diffraction order lies on its threshold, `\|v\| <= 1e-14` |
173//! | `component adjoints require an explicit nonzero Ewald split` | [`InvalidInput`] | Derivatives of the real or reciprocal part at `eta = 0` |
174//! | `direct half-cell shell grouping changes discontinuously; use the Ewald sum for adjoints` | [`InvalidInput`] | Derivatives of a direct 1D shell at a half-cell shift |
175//!
176//! Invalid labels, dimensions and non-finite inputs give [`InvalidInput`] with a message
177//! that names the input. A value-only call can return a sum whose derivatives fail.
178//!
179//! # Tuning constants
180//!
181//! | Name | Value | Meaning | Why |
182//! |---|---|---|---|
183//! | `resolve_split` | `1 / 8`, `sqrt(2 pi)` | Automatic split `sqrt(2 pi) / (k L) max(\|k L\| / 8, 1)` | As upstream `_check_eta`, except for `L` in 3D |
184//! | `SHELL_TOLERANCE` | `2e-13` | Part of `max(\|P\|, 1)` that each of the last two shells may add | About 900 ulps of the scale |
185//! | `shells` | 200, 32, 16 | Shell limit in 1D, 2D, 3D, raised up to four times where the terms need it (`shells_within`) | 200 in 1D as upstream; above upstream's 20 and 10 in 2D and 3D |
186//! | `TERM_DECAY` | `50` | Shell limits reach every term above `exp(-50)` | `exp(-50) = 2e-22`, far below `SHELL_TOLERANCE` |
187//! | `AUTOMATIC_GROWTH` | `16 / pi (1 + 1e-12)` | Largest `Re(1 / (2 eta^2))` of the automatic split; larger values mark a small split | Exact, with a margin for rounding |
188//! | `MAX_LOSS` | `1e-3` | Largest predicted loss of a sum, relative to the scale of a component, before it fails | Calibrated: the bounds exceed the errors by a median factor of 17 to 87 |
189//! | `EARLY_FAILURE_LOSS` | `2e-3` | Loss at which a complete sum at a small split fails before its far shells | Calibrated: without it, a 3D sum at the split 0.03 takes 29 s to fail |
190//! | `TERM_ULPS` | `8` | Rounding of each Ewald term and of the self term, in ulps | The `MAX_LOSS` calibration compares the bounds built from it with the errors |
191//! | `REAL_ULPS` | `64` | Rounding of the real-space part of a 1D spherical sum, in ulps | Its Kambe integrals of even order keep their relative accuracy |
192//! | `SPECTRAL_SWITCH_LOSS` | `1e-12` | Rounding of a 1D spherical component, relative to its scale, from which the spectral series replaces it | Far below `MAX_LOSS`, so the series takes over long before a sum fails |
193//! | `SPECTRAL_FIRST_W` | `2.5` | `w = \|k\| rho \|eta\|` from which a 1D spherical sum tries the spectral series first | The reciprocal part cancels by `exp(w^2) = 5e2` there |
194//! | `SERIES_SHELLS` | `512` | Diffraction orders the spectral series may take on each side | Caps the cost of one series; beyond it the Ewald sum is used |
195//! | `SERIES_DECAY` | `45` | The spectral series runs only where its terms fall by `exp(-45)` within `SERIES_SHELLS` | `exp(-45) = 3e-20`, below `QUIET_SHELL` |
196//! | `QUIET_SHELL` | `1e-18` | Part of the summed moduli below which a shell of the spectral series is quiet | The series stops after two quiet shells past its peak; `1e-18` is below the rounding of the sum |
197//! | `SERIES_T` | `1e-2` | Largest `\|t\|` at which integer orders of `F_n` come from their series in `t` | Below about `\|w\| = 1e-50` the Kambe integrals at `eta = -i / w` leave the range of `f64`; the series loses only a few ulps |
198//! | `SERIES_REACH` | `1` | The same, for the orders below the Kambe base pair | Each of these orders would need a gamma-function series of its own |
199//! | `HALF_INTEGER_SERIES_T` | `1` | Largest `\|t\|` at which half-integer orders come from their series in `t` | Calibrated: within `1e-13` of 30-digit references for `\|t\|` from 0.3 to 1 |
200//! | `SERIES_TERMS` | `64` | Terms of the series in `t` at most | `\|t\| <= 1` stops within 20 |
201//! | `NEAR_THRESHOLD` | `1 / 16` | Bound on the real and imaginary parts of `q · q / k^2 - 1` below which the zeroth order takes compensated sums | Rounded sums lose 50 ulps of `v` there |
202//! | `THRESHOLD_DISTANCE` | `1e-14` | `\|v\|` (reciprocal terms) or `\|k_q^2 / k^2\|` (spectral series) at or below which a diffraction order counts as on its threshold | The sums are singular on the threshold `\|q\| = k` |
203//! | `SHEET_MARGIN` | `0.5` | Angle margin within which a diffraction order needs the argument of its sheet | Far above the rounding of the angle (`1e-2` next to a threshold) |
204//! | `axial_cylindrical_reciprocal` | `4` | `\|v\|` above which the gamma functions of value-only 1D cylindrical sums on the axis recur upward | A downward step there subtracts the leading asymptotic term |
205//! | `kambe_arguments` | `1e-100` | Imaginary part that puts a real Kambe argument on the outgoing side | Any positive part selects the upper half plane |
206//! | `reduce_basis` | `1e-8`, `2^20`, 64 | Relative shortening a reduction step needs, largest coefficient, most passes | Only strict decreases count, so ties such as those of a hexagonal basis keep it as given |
207//! | `first_brillouin` | `1e-13`, `1e-14` | Relative and absolute tolerance on squared lengths of a reduced 2D basis | The absolute one is upstream's |
208//! | `PARALLEL` | 8 | Batched sums run in parallel from this many outputs | Fewer outputs run on the calling thread |
209//!
210//! # References
211//!
212//! - K. Kambe, Z. Naturforsch. A 22 (1967), <https://doi.org/10.1515/zna-1967-0305>
213//! - K. Kambe, Z. Naturforsch. A 23 (1968), <https://doi.org/10.1515/zna-1968-0908>
214//! - C. M. Linton, SIAM Rev. 52, 630 (2010), <https://doi.org/10.1137/09075130X>
215//! - D. Beutel et al., J. Opt. Soc. Am. B 38 (2021), <https://doi.org/10.1364/JOSAB.419645>
216//! - NIST DLMF, chapter 8 (incomplete gamma functions), <https://dlmf.nist.gov/8>
217//!
218//! [`NotConverged`]: crate::Error::NotConverged
219//! [`NonFinite`]: crate::Error::NonFinite
220//! [`InvalidInput`]: crate::Error::InvalidInput
221#![allow(clippy::float_cmp, clippy::indexing_slicing)] // Validated dimensions, exact branch points.
222
223use crate::{
224    Complex, Error, Result,
225    numerics::{Jet, finite},
226};
227
228mod accuracy;
229mod batch;
230mod cell;
231mod direct;
232mod ewald;
233mod geometry;
234mod inputs;
235mod real;
236mod reciprocal;
237mod reduced;
238mod saved;
239mod sheets;
240mod shells;
241mod spectral;
242mod wave;
243
244// Cell helpers of treams.lattice (volume, area, reciprocal, cube, cubeedge,
245// diffr_orders_circle) and treams.misc (firstbrillouin1d/2d/3d).
246pub use geometry::{
247    cube, diffraction_orders, first_brillouin, first_brillouin_1d, reciprocal, volume,
248};
249
250pub use batch::{SumResidual, sum_array};
251pub use cell::BlochLattice;
252pub(crate) use cell::resolve_split;
253pub use wave::Family;
254// Read by the property tests in properties/lattice/.
255#[cfg(test)]
256pub(crate) use ewald::probes;
257#[cfg(test)]
258pub(crate) use reduced::HALF_INTEGER_SERIES_T;
259#[cfg(test)]
260pub(crate) use shells::SHELL_TOLERANCE;
261
262use accuracy::below_automatic;
263use direct::direct_shell;
264use ewald::{SphericalCache, ewald};
265use inputs::{Evaluation, Inputs, unpack};
266use shells::{NOT_CONVERGED, is_shell_limit};
267
268/// The lattice sum of `wave` at the wavenumber `k`, the shift `shift` and the Ewald split
269/// parameter `eta`, as `treams.lattice.lsum*` returns it.
270///
271/// `shift` is the `r` of treams, and the lattice vectors `a` and the Bloch vector `kpar`
272/// come with `lattice`. `eta = 0` selects the automatic split. The module docs give the
273/// definition, the frame, the crosswalk to treams and the differences from it.
274///
275/// Upstream: the `lsum*` functions of `treams.lattice`, from `lsumsw1d` to `lsumcw2d`.
276///
277/// # Errors
278///
279/// [`Error::InvalidInput`] for invalid labels, dimensions or inputs and at a diffraction
280/// threshold, [`Error::NotConverged`] where an Ewald part does not converge or the parts
281/// cancel beyond use, and [`Error::NonFinite`] for a non-finite term. The module docs
282/// list the messages.
283pub fn sum(
284    wave: Family,
285    k: Complex,
286    lattice: &BlochLattice,
287    shift: [f64; 3],
288    eta: Complex,
289) -> Result<Complex> {
290    sum_part(wave, k, lattice, shift, eta, SumPart::Full)
291}
292
293/// Part of an Ewald sum, or one unaccelerated integer cube shell: the `lsum*`,
294/// `realsum*`, `recsum*` and `dsum*` functions of `treams.lattice`.
295#[derive(Clone, Copy, Debug)]
296pub enum SumPart {
297    /// Complete outgoing lattice sum.
298    Full,
299    /// Real-space Ewald contribution.
300    Real,
301    /// Reciprocal-space Ewald contribution, including the self correction.
302    Reciprocal,
303    /// One cube boundary of the given basis (which is not reduced); axial half-cell
304    /// shifts pair equidistant images instead.
305    Direct(i64),
306}
307
308/// Evaluate an Ewald component or one direct-summation shell.
309///
310/// Upstream: the `lsum*`, `realsum*`, `recsum*` and `dsum*` functions of
311/// `treams.lattice`, one per [`SumPart`].
312///
313/// # Errors
314///
315/// As [`sum`].
316pub fn sum_part(
317    wave: Family,
318    k: Complex,
319    lattice: &BlochLattice,
320    shift: [f64; 3],
321    eta: Complex,
322    part: SumPart,
323) -> Result<Complex> {
324    Ok(evaluate::<0>(wave, k, lattice, shift, eta, Evaluation::Part(part))?.value)
325}
326
327/// Local derivatives of one scalar Ewald sum.
328///
329/// Where the shift is a lattice point, the derivatives are those of the sum without that
330/// point, which stays excluded as the inputs move.
331#[derive(Clone, Copy, Debug)]
332pub struct Derivatives {
333    /// Sum value.
334    pub value: Complex,
335    /// Complex wavenumber derivative.
336    pub k: Complex,
337    /// Ewald split derivative; zero for the full and direct sums.
338    pub eta: Complex,
339    /// Cartesian shift derivatives.
340    #[doc(alias = "r")]
341    pub shift: [Complex; 3],
342    /// Bloch derivatives in the lattice frame of [`BlochLattice`] (unused components are zero).
343    pub kpar: [Complex; 3],
344    /// Row lattice-vector derivatives in the lattice frame (unused components are zero).
345    #[doc(alias = "a")]
346    pub vectors: [[Complex; 3]; 3],
347}
348
349/// Differentiate an Ewald sum using analytic integral identities and Cartesian polynomials.
350///
351/// Derivatives of 1D spherical sums carry rounding bounds of their own, from the Kambe
352/// orders one above and one below the value's. Each derivative that loses more than
353/// `1e-12` of its scale comes from the spectral series (see the module docs). Where only
354/// derivatives switch, the value stays the Ewald sum's, whose bound is then within
355/// `1e-12 max(|S|, 1)`.
356///
357/// # Errors
358///
359/// As [`sum`].
360pub fn derivatives(
361    wave: Family,
362    k: Complex,
363    lattice: &BlochLattice,
364    shift: [f64; 3],
365    eta: Complex,
366) -> Result<Derivatives> {
367    derivatives_part(wave, k, lattice, shift, eta, SumPart::Full)
368}
369
370/// Differentiate the selected component, holding a direct shell's index fixed.
371///
372/// # Errors
373///
374/// As [`sum`], and [`Error::InvalidInput`] for the real or reciprocal part at `eta = 0`,
375/// whose split derivative needs an explicit split.
376pub fn derivatives_part(
377    wave: Family,
378    k: Complex,
379    lattice: &BlochLattice,
380    shift: [f64; 3],
381    eta: Complex,
382    part: SumPart,
383) -> Result<Derivatives> {
384    let component = matches!(part, SumPart::Real | SumPart::Reciprocal);
385    if component && eta == Complex::default() {
386        return Err(Error::InvalidInput(
387            "component adjoints require an explicit nonzero Ewald split".into(),
388        ));
389    }
390    let evaluation = Evaluation::Part(part);
391    let dim = lattice.dim;
392    let mut result = match dim {
393        1 => unpack(
394            evaluate::<6>(wave, k, lattice, shift, eta, evaluation)?,
395            dim,
396        ),
397        2 => unpack(
398            evaluate::<10>(wave, k, lattice, shift, eta, evaluation)?,
399            dim,
400        ),
401        _ => unpack(
402            evaluate::<16>(wave, k, lattice, shift, eta, evaluation)?,
403            dim,
404        ),
405    };
406    if component {
407        // The complete sum is independent of the split: the reciprocal part
408        // carries the negative split derivative of the real part.
409        let real =
410            evaluate::<1>(wave, k, lattice, shift, eta, Evaluation::EtaDerivative)?.derivative[0];
411        result.eta = if matches!(part, SumPart::Real) {
412            real
413        } else {
414            -real
415        };
416    }
417    Ok(result)
418}
419
420/// Validates the inputs and evaluates what `evaluation` selects: a direct shell, or the
421/// Ewald sum at the resolved split. `N` is the number of jet slots: 0 for a value, 1 for
422/// the split derivative of the real-space part, and for all derivatives 6, 10 or 16 in
423/// one, two or three dimensions (`k`, three shift components, `dim` Bloch components and
424/// `dim * dim` lattice-vector components).
425fn evaluate<const N: usize>(
426    wave: Family,
427    k: Complex,
428    lattice: &BlochLattice,
429    r: [f64; 3],
430    eta: Complex,
431    evaluation: Evaluation,
432) -> Result<Jet<N>> {
433    evaluate_shared(wave, k, lattice, r, eta, evaluation, None)
434}
435
436/// Evaluates spherical orders sharing one degree's real-space samples. Other callers
437/// pass no cache and retain the ordinary scalar path.
438fn evaluate_shared<const N: usize>(
439    wave: Family,
440    k: Complex,
441    lattice: &BlochLattice,
442    r: [f64; 3],
443    eta: Complex,
444    evaluation: Evaluation,
445    shared: Option<&SphericalCache<N>>,
446) -> Result<Jet<N>> {
447    wave.validate(lattice.dim)?;
448    if let Evaluation::Part(SumPart::Direct(shell)) = evaluation
449        && !(0..=i64::from(i32::MAX)).contains(&shell)
450    {
451        return Err(Error::InvalidInput(
452            "direct shell index must be in [0, i32::MAX]".into(),
453        ));
454    }
455    if !finite(k) || k == Complex::default() || !finite(eta) || r.iter().any(|v| !v.is_finite()) {
456        return Err(Error::InvalidInput(
457            "lattice wavenumber must be finite and nonzero; shift and eta must be finite".into(),
458        ));
459    }
460    if matches!(wave, Family::Cylindrical { .. }) && r[2] != 0.0 {
461        return Err(Error::InvalidInput(
462            "scalar cylindrical lattice shifts lie in the xy plane".into(),
463        ));
464    }
465    let direct = matches!(evaluation, Evaluation::Part(SumPart::Direct(_)));
466    if !direct {
467        if k.im < 0.0 {
468            return Err(Error::InvalidInput(
469                "Ewald sums require Im(k) >= 0; gain media are unsupported".into(),
470            ));
471        }
472        if matches!(wave, Family::Spherical { .. }) && k.re < 0.0 {
473            return Err(Error::InvalidInput(
474                "spherical Ewald sums require Re(k) >= 0".into(),
475            ));
476        }
477    }
478    let variable = !matches!(evaluation, Evaluation::EtaDerivative);
479    let inputs = Inputs::new(wave, k, lattice, r, variable, !direct);
480    if let Evaluation::Part(SumPart::Direct(shell)) = evaluation {
481        return direct_shell(wave, lattice, &inputs, shell);
482    }
483    let eta = resolve_split(k, lattice, eta);
484    // Ewald powers and zero-shift jets require a representable squared split.
485    // In particular, 0 * an overflowing eta squared would corrupt exact zeros.
486    if !finite(eta * eta) {
487        return Err(Error::InvalidInput(
488            "Ewald split squared must be finite; reduce the split parameter".into(),
489        ));
490    }
491    ewald(wave, lattice, inputs, eta, evaluation, shared).map_err(|error| larger_split(error, eta))
492}
493
494/// Consecutive orders of one spherical degree, sharing their
495/// real-space radial integrals while retaining each harmonic's convergence policy.
496pub(crate) fn spherical_degree(
497    l: i32,
498    orders: std::ops::RangeInclusive<i32>,
499    k: Complex,
500    lattice: &BlochLattice,
501    shift: [f64; 3],
502    eta: Complex,
503) -> Result<Vec<Complex>> {
504    let shared = SphericalCache::default();
505    orders
506        .map(|m| {
507            Ok(evaluate_shared::<0>(
508                Family::Spherical { l, m },
509                k,
510                lattice,
511                shift,
512                eta,
513                Evaluation::Part(SumPart::Full),
514                Some(&shared),
515            )?
516            .value)
517        })
518        .collect()
519}
520
521/// Derivatives of the requested orders of one spherical degree. The jet cache shares
522/// the radial derivatives as well as the values, with no change to any order's shell
523/// stopping rule or rounding checks.
524pub(crate) fn spherical_degree_derivatives(
525    l: i32,
526    orders: &[i32],
527    k: Complex,
528    lattice: &BlochLattice,
529    shift: [f64; 3],
530    eta: Complex,
531) -> Result<Vec<Derivatives>> {
532    fn degree<const N: usize>(
533        l: i32,
534        orders: &[i32],
535        k: Complex,
536        lattice: &BlochLattice,
537        shift: [f64; 3],
538        eta: Complex,
539    ) -> Result<Vec<Derivatives>> {
540        let shared = SphericalCache::default();
541        orders
542            .iter()
543            .map(|&m| {
544                evaluate_shared::<N>(
545                    Family::Spherical { l, m },
546                    k,
547                    lattice,
548                    shift,
549                    eta,
550                    Evaluation::Part(SumPart::Full),
551                    Some(&shared),
552                )
553                .map(|value| unpack(value, lattice.dim))
554            })
555            .collect()
556    }
557    match lattice.dim {
558        1 => degree::<6>(l, orders, k, lattice, shift, eta),
559        2 => degree::<10>(l, orders, k, lattice, shift, eta),
560        _ => degree::<16>(l, orders, k, lattice, shift, eta),
561    }
562}
563
564/// `error`, advising a larger split where a sum at a split `eta` below every automatic
565/// one does not converge: its real-space shells converge sooner at a larger one. A sum
566/// whose shells settle but whose automatic split, which verifies it, does not converge
567/// either reports that instead (see [`ewald`]).
568///
569/// [`ewald`]: ewald::ewald
570fn larger_split(error: Error, eta: Complex) -> Error {
571    if is_shell_limit(&error) && below_automatic(eta) {
572        Error::NotConverged(format!(
573            "{NOT_CONVERGED}; use a larger split (eta = 0 selects one)"
574        ))
575    } else {
576        error
577    }
578}
579
580/// Cotangents of the continuous Ewald parameters under the real Hermitian pairing.
581///
582/// The fields run (k, eta, shift, kpar, vectors); `treams_rs.diff.lattice_sum` returns
583/// the same cotangents in the treams argument order (k, kpar, a, r, eta).
584#[derive(Clone, Copy, Debug, Default)]
585pub struct SumGradient {
586    /// Complex wavenumber cotangent.
587    pub k: Complex,
588    /// Complex split cotangent; zero for the full and direct sums.
589    pub eta: Complex,
590    /// Real shift cotangents.
591    #[doc(alias = "r")]
592    pub shift: [f64; 3],
593    /// Real Bloch cotangents in the lattice frame of [`BlochLattice`].
594    pub kpar: [f64; 3],
595    /// Real row lattice-vector cotangents in the lattice frame.
596    #[doc(alias = "a")]
597    pub vectors: [[f64; 3]; 3],
598}
599
600impl SumGradient {
601    pub(crate) fn add(&mut self, other: Self) {
602        self.k += other.k;
603        self.eta += other.eta;
604        for i in 0..3 {
605            self.shift[i] += other.shift[i];
606            self.kpar[i] += other.kpar[i];
607            for j in 0..3 {
608                self.vectors[i][j] += other.vectors[i][j];
609            }
610        }
611    }
612}
613
614impl Derivatives {
615    /// Directional derivative in all continuous inputs of this scalar sum.
616    #[must_use]
617    pub fn pushforward(&self, tangent: &SumTangent) -> Complex {
618        self.k * tangent.k
619            + self.eta * tangent.eta
620            + self
621                .shift
622                .into_iter()
623                .zip(tangent.shift)
624                .map(|(d, t)| d * t)
625                .sum::<Complex>()
626            + self
627                .kpar
628                .into_iter()
629                .zip(tangent.kpar)
630                .map(|(d, t)| d * t)
631                .sum::<Complex>()
632            + self
633                .vectors
634                .into_iter()
635                .flatten()
636                .zip(tangent.vectors.into_iter().flatten())
637                .map(|(d, t)| d * t)
638                .sum::<Complex>()
639    }
640
641    pub(crate) fn pullback(&self, cotangent: Complex) -> SumGradient {
642        SumGradient {
643            k: self.k.conj() * cotangent,
644            eta: self.eta.conj() * cotangent,
645            shift: self.shift.map(|x| (cotangent.conj() * x).re),
646            kpar: self.kpar.map(|x| (cotangent.conj() * x).re),
647            vectors: self
648                .vectors
649                .map(|row| row.map(|x| (cotangent.conj() * x).re)),
650        }
651    }
652}
653
654/// A direction in the continuous inputs of a lattice sum.
655///
656/// Geometry components outside the lattice or shift dimension are zero.
657#[derive(Clone, Copy, Debug, Default)]
658pub struct SumTangent {
659    /// Complex wavenumber direction.
660    pub k: Complex,
661    /// Complex Ewald-split direction.
662    pub eta: Complex,
663    /// Cartesian shift direction.
664    pub shift: [f64; 3],
665    /// Bloch-wavevector direction in the lattice frame.
666    pub kpar: [f64; 3],
667    /// Row lattice-vector directions in the lattice frame.
668    pub vectors: [[f64; 3]; 3],
669}
670
671impl SumTangent {
672    /// Validate and pad the geometry direction of a periodic expansion.
673    pub(crate) fn with_lattice(
674        kpar: &[f64],
675        vectors: &nalgebra::DMatrix<f64>,
676        dim: usize,
677    ) -> Result<Self> {
678        if kpar.len() != dim
679            || vectors.shape() != (dim, dim)
680            || kpar.iter().chain(vectors.iter()).any(|x| !x.is_finite())
681        {
682            return Err(Error::InvalidInput(
683                "lattice tangent must be finite and match the lattice dimension".into(),
684            ));
685        }
686        let mut result = Self::default();
687        result.kpar[..dim].copy_from_slice(kpar);
688        for i in 0..dim {
689            for j in 0..dim {
690                result.vectors[i][j] = vectors[(i, j)];
691            }
692        }
693        Ok(result)
694    }
695
696    pub(crate) fn is_finite(&self) -> bool {
697        finite(self.k)
698            && finite(self.eta)
699            && self
700                .shift
701                .iter()
702                .chain(&self.kpar)
703                .chain(self.vectors.iter().flatten())
704                .all(|x| x.is_finite())
705    }
706}
707
708#[cfg(test)]
709mod tests {
710    //! Sharing real-space samples must preserve the independently converged scalar
711    //! sum and every jet component, including the special split and spectral paths.
712
713    use super::{
714        BlochLattice, Derivatives, Family, derivatives, spherical_degree,
715        spherical_degree_derivatives, sum,
716    };
717    use crate::Complex;
718
719    fn components(derivatives: &Derivatives) -> Vec<Complex> {
720        [derivatives.value, derivatives.k, derivatives.eta]
721            .into_iter()
722            .chain(derivatives.shift)
723            .chain(derivatives.kpar)
724            .chain(derivatives.vectors.into_iter().flatten())
725            .collect()
726    }
727
728    #[test]
729    fn shared_spherical_orders_match_scalar_sums_and_derivatives() {
730        for dim in 1..=3 {
731            let vectors: Vec<Vec<f64>> = (0..dim)
732                .map(|i| (0..dim).map(|j| if i == j { 2.0 } else { 0.2 }).collect())
733                .collect();
734            let lattice = BlochLattice::new(&vectors, &vec![0.13; dim]).unwrap();
735            let k = Complex::new(1.2, 0.08);
736            let l = if dim == 2 { 5 } else { 2 };
737            for shift in [[0.0; 3], [0.2, -0.1, 0.0], [0.3, -0.2, 0.4]] {
738                for eta in [
739                    Complex::default(),
740                    Complex::new(0.7, -0.03),
741                    // The 3D real part crosses the cached-prefix boundary.
742                    Complex::new(0.4, 0.0),
743                    Complex::new(0.25, 0.0),
744                ] {
745                    let shared = spherical_degree(l, -l..=l, k, &lattice, shift, eta);
746                    let scalar = (-l..=l)
747                        .map(|m| sum(Family::Spherical { l, m }, k, &lattice, shift, eta))
748                        .collect::<crate::Result<Vec<_>>>();
749                    assert_eq!(
750                        shared.map_err(|error| error.to_string()),
751                        scalar.map_err(|error| error.to_string()),
752                        "dim={dim}, shift={shift:?}, eta={eta}"
753                    );
754                    // Sparse orders are the pullback's ordinary path when some
755                    // harmonic cotangents vanish.
756                    let orders = [-l, 0, l];
757                    let shared = spherical_degree_derivatives(l, &orders, k, &lattice, shift, eta);
758                    let scalar = orders
759                        .into_iter()
760                        .map(|m| derivatives(Family::Spherical { l, m }, k, &lattice, shift, eta))
761                        .collect::<crate::Result<Vec<_>>>();
762                    assert_eq!(
763                        shared
764                            .map(|values| values.iter().map(components).collect::<Vec<_>>())
765                            .map_err(|error| error.to_string()),
766                        scalar
767                            .map(|values| values.iter().map(components).collect::<Vec<_>>())
768                            .map_err(|error| error.to_string()),
769                        "dim={dim}, shift={shift:?}, eta={eta}"
770                    );
771                }
772            }
773        }
774    }
775
776    #[test]
777    fn shared_spherical_orders_keep_outgoing_branches_and_reject_gain() {
778        let lattice = BlochLattice::new(&[vec![2.0, 0.2], vec![0.2, 2.0]], &[0.13, -0.08]).unwrap();
779        let shift = [0.3, -0.2, 0.0];
780        let l = 4;
781        for k in [Complex::new(1.2, 0.0), Complex::new(1.2, -0.08)] {
782            for eta in [Complex::default(), Complex::new(-0.7, 0.03)] {
783                // Real k takes the outgoing branch cut. Re(k eta) < 0 takes the
784                // principal sheet; gain inputs fail before either cache is used.
785                let shared = spherical_degree(l, -l..=l, k, &lattice, shift, eta);
786                let scalar = (-l..=l)
787                    .map(|m| sum(Family::Spherical { l, m }, k, &lattice, shift, eta))
788                    .collect::<crate::Result<Vec<_>>>();
789                assert_eq!(
790                    shared.map_err(|error| error.to_string()),
791                    scalar.map_err(|error| error.to_string())
792                );
793                let orders = [-l, 0, l];
794                let shared = spherical_degree_derivatives(l, &orders, k, &lattice, shift, eta);
795                let scalar = orders
796                    .into_iter()
797                    .map(|m| derivatives(Family::Spherical { l, m }, k, &lattice, shift, eta))
798                    .collect::<crate::Result<Vec<_>>>();
799                let derivative_result = |result: crate::Result<Vec<Derivatives>>| {
800                    result
801                        .map(|values| values.iter().map(components).collect::<Vec<_>>())
802                        .map_err(|error| error.to_string())
803                };
804                assert_eq!(derivative_result(shared), derivative_result(scalar));
805            }
806        }
807    }
808}