treams_core/special/integrals/
mod.rs

1//! Incomplete gamma and Kambe integrals used by the Ewald lattice sums.
2//! Recurrences follow DLMF 8.8 and Kambe (1968), as in treams.
3//!
4//! Upstream: `treams.special.incgamma` and `intkambe` (`_integrals.pyx`).
5//!
6//! [`incgamma`] and [`intkambe`] evaluate one integral after checking its inputs.
7//! [`incgamma_array`] and [`intkambe_array`] evaluate them over arrays and return an
8//! [`IncgammaResidual`] or [`IntkambeResidual`], which keeps the arguments for the
9//! gradient.
10//! The private submodules hold the numerics:
11//!
12//! - `gamma`: the upper incomplete gamma function and its scaled ladder;
13//! - `kambe`: the Kambe integrals;
14//! - `double`: double-double arithmetic for the small-split Kambe series.
15//!
16//! The Kambe integrals have several entry points, one per kind of caller:
17//!
18//! | Entry point | Returns | Caller |
19//! |---|---|---|
20//! | [`intkambe`] | `I_n(z, eta)` after checking the inputs | `treams_rs.special.intkambe`, [`intkambe_array`] |
21//! | `kambe` | `I_n(z, eta)` without checks | [`intkambe`], the [`IntkambeResidual`] pullback (order `n + 2`), real-space terms of the Ewald sums (`lattice::real`) |
22//! | `kambe_with_bound` | `I_n` and a bound on its rounding error | integer-order reduced integrals of the 1D spherical sums (`lattice::reduced`) |
23//! | `kambe_series`, `kambe_series_with` | the series over scaled incomplete gamma functions | `kambe`, `kambe_with_bound` and `SmallSplitKambe` only |
24//! | `EvenKambe` | the even orders of one argument from one recurrence | half-integer reduced integrals of the 2D spherical and 1D cylindrical sums (`lattice::reduced`) |
25//! | `SmallSplitKambe` | the orders of real-space terms at a split below every automatic one, with bounds | real-space terms at small splits (`lattice::real`) |
26#![allow(clippy::float_cmp, clippy::while_float, clippy::indexing_slicing)] // Half-integer degrees step exactly by one; order tables cover each requested order.
27
28mod double;
29mod gamma;
30mod kambe;
31
32pub(crate) use gamma::{ScaledGammaLadder, upper_gamma};
33pub(crate) use kambe::{EvenKambe, SmallSplitKambe, kambe, kambe_argument, kambe_with_bound};
34
35use super::MAX_LABEL;
36use crate::{
37    Complex, Error, MAX_DEGREE, Result,
38    numerics::{
39        broadcast::{self, element},
40        finite,
41        parallel::Parallel,
42    },
43};
44use gamma::gamma_derivative;
45
46/// Upper incomplete gamma function `Gamma(n, z) = integral_z^infinity t^(n-1) e^-t dt`.
47///
48/// The degree `n` is an integer or half-integer. The branch cut lies on the negative
49/// real axis, and the sign of a zero `Im z` selects its side.
50///
51/// Upstream: `treams.special.incgamma`. Differences: `|n|` above 128 and non-finite `z`
52/// give an error; treams runs one degree recurrence, while this function chooses among
53/// four methods the one that keeps relative accuracy at `z`.
54pub fn incgamma(n: f64, z: Complex) -> Result<Complex> {
55    if !n.is_finite() || (2.0 * n).fract() != 0.0 || n.abs() > f64::from(MAX_DEGREE) || !finite(z) {
56        return Err(Error::InvalidInput(
57            // 128 is MAX_DEGREE.
58            "gamma degree must be an integer or half-integer in [-128, 128], and z finite".into(),
59        ));
60    }
61    Ok(upper_gamma(n, z))
62}
63
64/// The Kambe integral `I_n(z, eta) = integral_eta^infinity t^n exp(-z^2 t^2/2 + 1/(2t^2)) dt`.
65///
66/// Upstream: `treams.special.intkambe`. Differences: `|n|` above 260 and non-finite
67/// arguments give an error. The odd orders -3 and -1 come from a series over scaled
68/// incomplete gamma functions, and lower orders from that series or the downward
69/// recurrence of treams, whichever bounds the rounding error more tightly.
70pub fn intkambe(n: i32, z: Complex, eta: Complex) -> Result<Complex> {
71    if !(-MAX_LABEL..=MAX_LABEL).contains(&n) || !finite(z) || !finite(eta) {
72        return Err(Error::InvalidInput(
73            // 260 is MAX_LABEL.
74            "Kambe order must be in [-260, 260] and arguments finite".into(),
75        ));
76    }
77    Ok(kambe(n, z, eta))
78}
79
80/// Elementwise integral evaluations and pullbacks run in parallel from this many elements.
81const PARALLEL: Parallel = Parallel::AtLeast(1024);
82const BROADCAST: &str = "arrays must have equal lengths or scalar inputs";
83const COTANGENT: &str = "integral cotangent must be finite and match output";
84
85/// [`incgamma`] over arrays, each input one value or one value per element, with the
86/// residual of their argument gradient.
87///
88/// Upstream: `treams.special.incgamma` over arrays.
89pub fn incgamma_array(
90    degrees: Vec<f64>,
91    arguments: Vec<Complex>,
92) -> Result<(Vec<Complex>, IncgammaResidual)> {
93    let residual = IncgammaResidual::new(degrees, arguments)?;
94    let value = broadcast::map(residual.size, PARALLEL, |i| {
95        incgamma(
96            element(&residual.degrees, i),
97            element(&residual.arguments, i),
98        )
99    })?;
100    Ok((value, residual))
101}
102
103/// What [`incgamma_array`] saves for its pullback: the degrees and the arguments. Only
104/// the arguments get gradients.
105#[derive(Debug)]
106pub struct IncgammaResidual {
107    degrees: Vec<f64>,
108    arguments: Vec<Complex>,
109    size: usize,
110}
111impl IncgammaResidual {
112    /// Keep the broadcast inputs needed by derivatives without evaluating values.
113    pub fn new(degrees: Vec<f64>, arguments: Vec<Complex>) -> Result<Self> {
114        let size = broadcast::size(&[degrees.len(), arguments.len()], BROADCAST)?;
115        Ok(Self {
116            degrees,
117            arguments,
118            size,
119        })
120    }
121
122    /// Directional derivative with respect to the complex argument.
123    pub fn pushforward(&self, tangents: [&[Complex]; 1]) -> Result<Vec<Complex>> {
124        broadcast::pushforward(
125            tangents,
126            self.size,
127            "gamma tangent must be finite and broadcast to output",
128            PARALLEL,
129            |i, [tangent]| {
130                if tangent == Complex::default() {
131                    return Ok(tangent);
132                }
133                let value = tangent
134                    * gamma_derivative(element(&self.degrees, i), element(&self.arguments, i));
135                if !finite(value) {
136                    return Err(Error::SpecialFunction(
137                        "gamma derivative is singular or overflowed".into(),
138                    ));
139                }
140                Ok(value)
141            },
142        )
143    }
144
145    /// The gradient with respect to the arguments, given the gradient `cotangent` with
146    /// respect to the values: `cotangent` times the conjugate of `dGamma/dz`.
147    pub fn pullback(&self, cotangent: &[Complex]) -> Result<Vec<Complex>> {
148        let [gradient] = broadcast::pullback(
149            cotangent,
150            self.size,
151            COTANGENT,
152            [self.arguments.len()],
153            PARALLEL,
154            |i, &g| {
155                if g == Complex::default() {
156                    return Ok([g]);
157                }
158                let gradient =
159                    g * gamma_derivative(element(&self.degrees, i), element(&self.arguments, i))
160                        .conj();
161                if !finite(gradient) {
162                    return Err(Error::SpecialFunction(
163                        "gamma derivative is singular or overflowed".into(),
164                    ));
165                }
166                Ok([gradient])
167            },
168        )?;
169        Ok(gradient)
170    }
171}
172
173/// [`intkambe`] over arrays, each input one value or one value per element, with the
174/// residual of their `z` and `eta` gradients; the integer order has none.
175///
176/// Upstream: `treams.special.intkambe` over arrays.
177pub fn intkambe_array(
178    orders: Vec<i32>,
179    arguments: [Vec<Complex>; 2],
180) -> Result<(Vec<Complex>, IntkambeResidual)> {
181    let residual = IntkambeResidual::new(orders, arguments)?;
182    let value = broadcast::map(residual.size, PARALLEL, |i| {
183        intkambe(
184            element(&residual.orders, i),
185            element(&residual.arguments[0], i),
186            element(&residual.arguments[1], i),
187        )
188    })?;
189    Ok((value, residual))
190}
191
192/// What [`intkambe_array`] saves for its pullback: the orders and the arguments. The
193/// pullback uses the adjacent-order recurrence and the endpoint integrand.
194#[derive(Debug)]
195pub struct IntkambeResidual {
196    orders: Vec<i32>,
197    arguments: [Vec<Complex>; 2],
198    size: usize,
199}
200impl IntkambeResidual {
201    /// Keep the broadcast inputs needed by derivatives without evaluating values.
202    pub fn new(orders: Vec<i32>, arguments: [Vec<Complex>; 2]) -> Result<Self> {
203        let size = broadcast::size(
204            &[orders.len(), arguments[0].len(), arguments[1].len()],
205            BROADCAST,
206        )?;
207        Ok(Self {
208            orders,
209            arguments,
210            size,
211        })
212    }
213
214    /// Directional derivative in `z` and `eta`. Inactive arguments do not evaluate
215    /// their derivative, so a direction in `eta` remains valid when only `dz` is singular.
216    pub fn pushforward(&self, tangents: [&[Complex]; 2]) -> Result<Vec<Complex>> {
217        broadcast::pushforward(
218            tangents,
219            self.size,
220            "Kambe tangents must be finite and broadcast to output",
221            PARALLEL,
222            |i, [tz, teta]| {
223                if tz == Complex::default() && teta == Complex::default() {
224                    return Ok(Complex::default());
225                }
226                let n = element(&self.orders, i);
227                let [z, eta] = self.arguments.each_ref().map(|v| element(v, i));
228                let dz = if tz == Complex::default() {
229                    Complex::default()
230                } else {
231                    tz * kambe_z_derivative(n, z, eta)
232                };
233                let deta = if teta == Complex::default() {
234                    Complex::default()
235                } else {
236                    teta * kambe_eta_derivative(n, z, eta)
237                };
238                let value = dz + deta;
239                if !finite(value) {
240                    return Err(Error::SpecialFunction(
241                        "Kambe derivative is singular or overflowed".into(),
242                    ));
243                }
244                Ok(value)
245            },
246        )
247    }
248
249    /// The gradients with respect to `z` and `eta`, given the gradient `cotangent` with
250    /// respect to the values. At `z = 0` the `z` gradient of orders `n <= -3` is its
251    /// limit, zero. Orders `n >= -2` have no `z` gradient at `z = 0` and give an error.
252    pub fn pullback(&self, cotangent: &[Complex]) -> Result<[Vec<Complex>; 2]> {
253        let lengths = self.arguments.each_ref().map(Vec::len);
254        broadcast::pullback(
255            cotangent,
256            self.size,
257            COTANGENT,
258            lengths,
259            PARALLEL,
260            |i, &g| {
261                if g == Complex::default() {
262                    return Ok([g; 2]);
263                }
264                let n = element(&self.orders, i);
265                let [z, eta] = self.arguments.each_ref().map(|v| element(v, i));
266                let dz = kambe_z_derivative(n, z, eta);
267                let deta = kambe_eta_derivative(n, z, eta);
268                let result = [dz, deta].map(|d| g * d.conj());
269                if result.iter().any(|&v| !finite(v)) {
270                    return Err(Error::SpecialFunction(
271                        "Kambe derivative is singular or overflowed".into(),
272                    ));
273                }
274                Ok(result)
275            },
276        )
277    }
278}
279
280fn kambe_z_derivative(n: i32, z: Complex, eta: Complex) -> Complex {
281    if z == Complex::default() && n <= -3 {
282        Complex::default()
283    } else {
284        -z * kambe(n + 2, z, eta)
285    }
286}
287
288fn kambe_eta_derivative(n: i32, z: Complex, eta: Complex) -> Complex {
289    -eta.powi(n) * (0.5 * (1.0 / (eta * eta) - z * z * eta * eta)).exp()
290}
291
292/// Reference cases of a `n re(z) im(z) [...]: re im` table with `#` comments, for the
293/// gamma and Kambe tests.
294#[cfg(test)]
295fn reference(text: &str) -> Vec<(Vec<f64>, Complex)> {
296    crate::test_support::table::<f64, f64>(text)
297        .into_iter()
298        .map(|(key, value)| (key, Complex::new(value[0], value[1])))
299        .collect()
300}
treams_core/special/integrals/mod.rs

treams_core/special/integrals/
mod.rs
