treams_core/lattice/
inputs.rs

1//! The inputs of one sum as forward-mode jets, and their jet slots.
2//!
3//! treams-rs extension: treams has no derivatives of its lattice sums.
4
5use std::f64::consts::PI;
6
7use super::{
8    Derivatives, SumPart,
9    cell::{BlochLattice, Reduction},
10    wave::Family,
11};
12use crate::{Complex, numerics::Jet};
13
14/// Jet slot of the wavenumber.
15pub(super) const K: usize = 0;
16/// First of the three Cartesian shift slots.
17const POSITION: usize = 1;
18/// First of the `dim` Bloch slots, which the `dim * dim` lattice-vector slots follow.
19pub(super) const KPAR: usize = 4;
20
21/// Jet slot of component `j` of lattice vector `i`.
22const fn vector_slot(dim: usize, i: usize, j: usize) -> usize {
23    KPAR + dim + dim * i + j
24}
25
26/// The lattice vector `(i, j)` of a jet slot at or after the vector slots.
27const fn vector_component(dim: usize, slot: usize) -> (usize, usize) {
28    ((slot - KPAR - dim) / dim, (slot - KPAR - dim) % dim)
29}
30
31pub(super) fn unpack<const N: usize>(result: Jet<N>, dim: usize) -> Derivatives {
32    Derivatives {
33        value: result.value,
34        k: result.derivative[K],
35        eta: Complex::default(),
36        shift: std::array::from_fn(|i| result.derivative[POSITION + i]),
37        kpar: std::array::from_fn(|i| {
38            if i < dim {
39                result.derivative[KPAR + i]
40            } else {
41                Complex::default()
42            }
43        }),
44        vectors: std::array::from_fn(|i| {
45            std::array::from_fn(|j| {
46                if i < dim && j < dim {
47                    result.derivative[vector_slot(dim, i, j)]
48                } else {
49                    Complex::default()
50                }
51            })
52        }),
53    }
54}
55
56/// What one evaluation of `evaluate` computes.
57#[derive(Clone, Copy, Debug)]
58pub(super) enum Evaluation {
59    /// A part of the sum, differentiated in every jet slot.
60    Part(SumPart),
61    /// The split derivative of the real-space part in slot 0, all other inputs fixed.
62    EtaDerivative,
63}
64
65impl Evaluation {
66    /// Whether the real-space part contributes.
67    pub(super) fn real(self) -> bool {
68        matches!(
69            self,
70            Self::Part(SumPart::Full | SumPart::Real) | Self::EtaDerivative
71        )
72    }
73
74    /// Whether the reciprocal part and the self correction contribute.
75    pub(super) fn reciprocal(self) -> bool {
76        matches!(self, Self::Part(SumPart::Full | SumPart::Reciprocal))
77    }
78}
79
80/// The inputs of one sum as jets in the slots above (or constants for the split
81/// derivative), in lattice coordinates except for the Cartesian shift.
82pub(super) struct Inputs<const N: usize> {
83    pub(super) dim: usize,
84    /// Cartesian axis of each lattice coordinate.
85    pub(super) axes: [usize; 3],
86    pub(super) k: Jet<N>,
87    pub(super) r: [Jet<N>; 3],
88    pub(super) kpar: [Jet<N>; 3],
89    pub(super) direct: [[Jet<N>; 3]; 3],
90    /// Reciprocal rows; `a_i · b_j = 2π δ_ij` gives their vector derivatives.
91    pub(super) reciprocal: [[Jet<N>; 3]; 3],
92    /// Cell measure `|det a|`.
93    pub(super) measure: Jet<N>,
94}
95
96impl<const N: usize> Inputs<N> {
97    /// Without `with_reciprocal` the reciprocal rows and the measure, which only the Ewald
98    /// parts read, stay zero: the direct shells skip their jets. Always inlined, so the
99    /// caller builds the inputs in its own frame instead of copying a returned value.
100    #[allow(clippy::inline_always)]
101    #[inline(always)]
102    pub(super) fn new(
103        wave: Family,
104        k: Complex,
105        lattice: &BlochLattice,
106        r: [f64; 3],
107        variable: bool,
108        with_reciprocal: bool,
109    ) -> Self {
110        let dim = lattice.dim;
111        let jet = |value: Complex, slot| {
112            if variable {
113                Jet::<N>::variable(value, slot)
114            } else {
115                Jet::constant(value)
116            }
117        };
118        // A lattice function whose derivative in vector component (p, q) is `slope(p, q)`.
119        let lattice_jet = |value: f64, slope: &dyn Fn(usize, usize) -> f64| Jet::<N> {
120            value: Complex::new(value, 0.0),
121            derivative: std::array::from_fn(|slot| {
122                if slot < KPAR + dim {
123                    return Complex::default();
124                }
125                let (p, q) = vector_component(dim, slot);
126                Complex::new(if p < dim && q < dim { slope(p, q) } else { 0.0 }, 0.0)
127            }),
128        };
129        let direct = &lattice.direct;
130        let reciprocal = &lattice.reciprocal;
131        // Filled in place, so no `from_fn` temporaries are copied into the result.
132        let mut inputs = Self {
133            dim,
134            axes: wave.axes(dim),
135            k: jet(k, K),
136            r: [Jet::default(); 3],
137            kpar: [Jet::default(); 3],
138            direct: [[Jet::default(); 3]; 3],
139            reciprocal: [[Jet::default(); 3]; 3],
140            measure: Jet::default(),
141        };
142        for (i, row) in inputs.direct.iter_mut().enumerate() {
143            inputs.r[i] = jet(r[i].into(), POSITION + i);
144            if i < dim {
145                inputs.kpar[i] = jet(lattice.kpar[i].into(), KPAR + i);
146            }
147            for (j, entry) in row.iter_mut().enumerate() {
148                *entry = if i < dim && j < dim {
149                    jet(direct[i][j].into(), vector_slot(dim, i, j))
150                } else {
151                    Jet::constant(direct[i][j])
152                };
153            }
154        }
155        if with_reciprocal {
156            for (i, row) in inputs.reciprocal.iter_mut().enumerate() {
157                for (j, entry) in row.iter_mut().enumerate() {
158                    *entry = lattice_jet(reciprocal[i][j], &|p, q| {
159                        -reciprocal[i][q] * reciprocal[p][j] / (2.0 * PI)
160                    });
161                }
162            }
163            inputs.measure = lattice_jet(lattice.measure, &|p, q| {
164                lattice.measure * reciprocal[p][q] / (2.0 * PI)
165            });
166        }
167        inputs
168    }
169
170    /// Integer combination `Σ_i n_i rows_i` of lattice rows.
171    #[allow(clippy::cast_precision_loss)] // Shell indices are bounded by i32::MAX, exactly representable.
172    pub(super) fn point(&self, rows: &[[Jet<N>; 3]; 3], n: [i64; 3]) -> [Jet<N>; 3] {
173        self.combination(rows, n.map(|n| n as f64))
174    }
175
176    /// Combination `Σ_i c_i rows_i` of lattice rows with integer-valued `c`.
177    pub(super) fn combination(&self, rows: &[[Jet<N>; 3]; 3], c: [f64; 3]) -> [Jet<N>; 3] {
178        std::array::from_fn(|j| (0..self.dim).map(|i| c[i] * rows[i][j]).sum())
179    }
180
181    /// The reduced rows `U rows`, or `rows` when they are reduced already.
182    pub(super) fn reduce(
183        &self,
184        rows: &[[Jet<N>; 3]; 3],
185        reduction: Option<&Reduction>,
186    ) -> [[Jet<N>; 3]; 3] {
187        reduction.map_or(*rows, |reduction| {
188            std::array::from_fn(|i| self.combination(rows, reduction.rows[i]))
189        })
190    }
191
192    /// Cartesian displacement `-r - R` of the image at lattice point `R`.
193    pub(super) fn image(&self, point: &[Jet<N>; 3]) -> [Jet<N>; 3] {
194        let mut shift = self.r.map(|v| -v);
195        for j in 0..self.dim {
196            shift[self.axes[j]] -= point[j];
197        }
198        shift
199    }
200
201    /// Bloch phase `q · R` of lattice point `R`.
202    pub(super) fn phase(&self, point: &[Jet<N>; 3]) -> Jet<N> {
203        (0..self.dim).map(|j| self.kpar[j] * point[j]).sum()
204    }
205}
treams_core/lattice/inputs.rs

treams_core/lattice/
inputs.rs
