treams_core/lattice/
direct.rs

1//! Direct sums over cube shells, without Ewald acceleration.
2//!
3//! Upstream: `dsumsw*` and `dsumcw*` of `lattice/_dsum.pyx`, with the terms `_fsw*` and
4//! `_fcw*`.
5
6use super::{
7    cell::BlochLattice,
8    geometry::visit_cube,
9    inputs::Inputs,
10    wave::{Family, solid_jet},
11};
12use crate::{
13    Complex, Error, Result,
14    numerics::Jet,
15    special::{self, Bessel, Radial, harmonic_normalization},
16};
17
18/// One image of a direct sum: the outgoing wave `h_l^(1)(k |r|) Y_lm(r / |r|)` or
19/// `H_m^(1)(k |r|) e^(i m phi)` at the displacement `r = -shift - R` (upstream `_fsw*`
20/// and `_fcw*`). Values take `special::bessel`; jets take the radial jets, whose first
21/// derivative the chain rule carries. The spherical harmonic is
22/// `harmonic_normalization(l, m)` times the solid harmonic of the unit vector.
23fn direct_term<const N: usize>(wave: Family, k: Jet<N>, r: [Jet<N>; 3]) -> Result<Jet<N>> {
24    let radius = r.into_iter().map(|v| v * v).sum::<Jet<N>>().sqrt();
25    let argument = k * radius;
26    let (degree, spherical) = match wave {
27        Family::Spherical { l, .. } => (l, true),
28        Family::Cylindrical { m } => (m, false),
29    };
30    let radial = if N == 0 {
31        Jet::constant(special::bessel(
32            f64::from(degree),
33            argument.value,
34            Bessel::H1,
35            spherical,
36            0,
37        )?)
38    } else {
39        let radial = if spherical {
40            special::spherical_radial(
41                u32::try_from(degree).map_err(|_| Error::InvalidInput("invalid degree".into()))?,
42                argument.value,
43                Radial::Singular,
44            )?
45        } else {
46            special::cylindrical_radial(degree, argument.value, Radial::Singular)?
47        };
48        argument.chain(radial.value, radial.first)
49    };
50    let angular = match wave {
51        Family::Spherical { l, m } => {
52            harmonic_normalization(l, m) * solid_jet(l, m, r.map(|v| v / radius))
53        }
54        Family::Cylindrical { m } => ((r[0] + Complex::i() * r[1]) / radius).powi(m),
55    };
56    let result = radial * angular;
57    if !result.finite() {
58        return Err(Error::NonFinite("non-finite direct lattice term".into()));
59    }
60    Ok(result)
61}
62
63/// One unaccelerated cube shell of images, in the given lattice basis (upstream
64/// `dsum*(..., i)` with `i = shell`).
65///
66/// A 1D sum whose shift lies on the axis at half a period pairs the images
67/// `direction * shell` and `-direction * (shell + 1)`, which lie at equal distances, so
68/// each shell adds a symmetric pair. That grouping jumps when the shift moves off the
69/// half cell, so jets fail there.
70pub(super) fn direct_shell<const N: usize>(
71    wave: Family,
72    lattice: &BlochLattice,
73    inputs: &Inputs<N>,
74    shell: i64,
75) -> Result<Jet<N>> {
76    let (dim, axis, r) = (inputs.dim, inputs.axes[0], &inputs.r);
77    let half_cell = dim == 1
78        && r.iter()
79            .enumerate()
80            .all(|(j, v)| j == axis || v.value == Complex::default())
81        && (r[axis].value.re / lattice.direct[0][0]).abs() == 0.5;
82    if N > 0 && half_cell {
83        return Err(Error::InvalidInput("direct half-cell shell grouping changes discontinuously; use the Ewald sum for adjoints".into()));
84    }
85    let mut result = Jet::default();
86    let mut term = |n| {
87        let point = inputs.point(&inputs.direct, n);
88        let shift = inputs.image(&point);
89        if shift.iter().all(|v| v.value == Complex::default()) {
90            return Ok(());
91        }
92        result += (Complex::i() * inputs.phase(&point)).exp() * direct_term(wave, inputs.k, shift)?;
93        Ok(())
94    };
95    if half_cell {
96        let direction = if r[axis].value.re / lattice.direct[0][0] > 0.0 {
97            1
98        } else {
99            -1
100        };
101        term([direction * shell, 0, 0])?;
102        term([-direction * (shell + 1), 0, 0])?;
103    } else {
104        visit_cube(dim, shell, true, term)?;
105    }
106    Ok(result)
107}
treams_core/lattice/direct.rs

treams_core/lattice/
direct.rs
