feos_dft/
weight_functions.rs

1use ndarray::*;
2use num_dual::DualNum;
3// use rustfft::num_complex::Complex;
4use std::f64::consts::{FRAC_PI_3, PI};
5use std::ops::Mul;
6
7/// A weight function corresponding to a single weighted density.
8#[derive(Clone)]
9pub struct WeightFunction<T> {
10    /// Factor in front of normalized weight function
11    pub prefactor: Array1<T>,
12    /// Kernel radius of the convolution
13    pub kernel_radius: Array1<T>,
14    /// Shape of the weight function (Dirac delta, Heaviside, etc.)
15    pub shape: WeightFunctionShape,
16}
17
18impl<T: DualNum<f64> + Copy> WeightFunction<T> {
19    /// Create a new weight function without prefactor
20    pub fn new_unscaled(kernel_radius: Array1<T>, shape: WeightFunctionShape) -> Self {
21        Self {
22            prefactor: Array::ones(kernel_radius.raw_dim()),
23            kernel_radius,
24            shape,
25        }
26    }
27
28    /// Create a new weight function with weight constant = 1
29    pub fn new_scaled(kernel_radius: Array1<T>, shape: WeightFunctionShape) -> Self {
30        let unscaled = Self::new_unscaled(kernel_radius, shape);
31        let weight_constants = unscaled.scalar_weight_constants(T::zero());
32        Self {
33            prefactor: weight_constants.mapv(|w| w.recip()),
34            kernel_radius: unscaled.kernel_radius,
35            shape,
36        }
37    }
38
39    /// Calculates the value of the scalar weight function depending on its shape
40    /// `k_abs` describes the absolute value of the Fourier variable
41    pub(crate) fn fft_scalar_weight_functions<D: Dimension, T2: DualNum<f64>>(
42        &self,
43        k_abs: &Array<T2, D>,
44        lanczos: &Option<Array<T2, D>>,
45    ) -> Array<T, D::Larger>
46    where
47        T: Mul<T2, Output = T>,
48        D::Larger: Dimension<Smaller = D>,
49    {
50        // Allocate vector for weight functions
51        let mut d = vec![self.prefactor.len()];
52        k_abs.shape().iter().for_each(|&x| d.push(x));
53        let mut w = Array::zeros(d.into_dimension())
54            .into_dimensionality()
55            .unwrap();
56
57        // Calculate weight function for each component
58        for i in 0..self.prefactor.len() {
59            let radius = self.kernel_radius[i];
60            let p = self.prefactor[i];
61            let rik = k_abs.mapv(|k| radius * k);
62            let mut w_i = w.index_axis_mut(Axis(0), i);
63            w_i.assign(&match self.shape {
64                WeightFunctionShape::Theta => rik.mapv(|rik| {
65                    (rik.sph_j0() + rik.sph_j2()) * 4.0 * FRAC_PI_3 * radius.powi(3) * p
66                }),
67                WeightFunctionShape::Delta => {
68                    rik.mapv(|rik| rik.sph_j0() * 4.0 * PI * radius.powi(2) * p)
69                }
70                WeightFunctionShape::KR1 => {
71                    rik.mapv(|rik| (rik.sph_j0() + rik.cos()) * 0.5 * radius * p)
72                }
73                WeightFunctionShape::KR0 => rik.mapv(|rik| (rik * rik.sin() * 0.5 + rik.cos()) * p),
74                WeightFunctionShape::DeltaVec => unreachable!(),
75            });
76
77            // Apply Lanczos sigma factor
78            if let Some(l) = lanczos {
79                w_i.assign(&(&w_i * l));
80            }
81        }
82
83        // Return real part of weight function
84        w
85    }
86
87    /// Calculates the value of the vector weight function depending on its shape
88    /// `k_abs` describes the absolute value of the Fourier variable
89    /// `k` describes the (potentially multi-dimensional) Fourier variable
90    pub(crate) fn fft_vector_weight_functions<D: Dimension, T2: DualNum<f64>>(
91        &self,
92        k_abs: &Array<T2, D>,
93        k: &Array<T2, D::Larger>,
94        lanczos: &Option<Array<T2, D>>,
95    ) -> Array<T, <D::Larger as Dimension>::Larger>
96    where
97        D::Larger: Dimension<Smaller = D>,
98        <D::Larger as Dimension>::Larger: Dimension<Smaller = D::Larger>,
99        T: Mul<T2, Output = T>,
100    {
101        // Allocate vector for weight functions
102        let mut d = vec![k.shape()[0], self.prefactor.len()];
103        k.shape().iter().skip(1).for_each(|&x| d.push(x));
104        let mut w = Array::zeros(d.into_dimension())
105            .into_dimensionality()
106            .unwrap();
107
108        // Iterate all dimensions
109        for (k_x, mut w_x) in k.outer_iter().zip(w.outer_iter_mut()) {
110            // Calculate weight function for each component
111            for i in 0..self.prefactor.len() {
112                let radius = self.kernel_radius[i];
113                let p = self.prefactor[i];
114                let rik = k_abs.mapv(|k| radius * k);
115                let mut w_i = w_x.index_axis_mut(Axis(0), i);
116                w_i.assign(&match self.shape {
117                    WeightFunctionShape::DeltaVec => {
118                        &rik.mapv(|rik| {
119                            (rik.sph_j0() + rik.sph_j2()) * (radius.powi(3) * 4.0 * FRAC_PI_3 * p)
120                        }) * &k_x
121                    }
122                    _ => unreachable!(),
123                });
124
125                // Apply Lanczos sigma factor
126                if let Some(l) = lanczos {
127                    w_i.assign(&(&w_i * l));
128                }
129            }
130        }
131        // Return imaginary part of weight function
132        w
133    }
134
135    /// Scalar weights for the bulk convolver (for the bulk convolver only the prefactor
136    /// of the normalized weight functions are required)
137    pub fn scalar_weight_constants(&self, k: T) -> Array1<T> {
138        let k = arr0(k);
139        self.fft_scalar_weight_functions(&k, &None)
140    }
141
142    /// Vector weights for the bulk convolver (for the bulk convolver only the prefactor
143    /// of the normalized weight functions are required)
144    pub fn vector_weight_constants(&self, k: T) -> Array1<T> {
145        let k_abs = arr0(k);
146        let k = arr1(&[k]);
147        self.fft_vector_weight_functions(&k_abs, &k, &None)
148            .index_axis_move(Axis(0), 0)
149    }
150}
151
152/// Possible weight function shapes.
153#[derive(Clone, Copy, PartialEq, Eq, Debug)]
154pub enum WeightFunctionShape {
155    /// Heaviside step function
156    Theta,
157    /// Dirac delta function
158    Delta,
159    /// Combination of first & second derivative of Dirac delta function (with different prefactor; only in Kierlik-Rosinberg functional)
160    KR0,
161    /// First derivative of Dirac delta function (with different prefactor; only in Kierlik-Rosinberg functional)
162    KR1,
163    /// Vector-shape as combination of Dirac delta and outward normal
164    DeltaVec,
165}
166
167/// Information about weight functions
168pub struct WeightFunctionInfo<T> {
169    /// Index of the component that each individual segment belongs to.
170    pub(crate) component_index: Array1<usize>,
171    /// Flag if local density is required in the functional
172    pub(crate) local_density: bool,
173    /// Container for scalar component-wise weighted densities
174    pub(crate) scalar_component_weighted_densities: Vec<WeightFunction<T>>,
175    /// Container for vector component-wise weighted densities
176    pub(crate) vector_component_weighted_densities: Vec<WeightFunction<T>>,
177    /// Container for scalar FMT weighted densities
178    pub(crate) scalar_fmt_weighted_densities: Vec<WeightFunction<T>>,
179    /// Container for vector FMT weighted densities
180    pub(crate) vector_fmt_weighted_densities: Vec<WeightFunction<T>>,
181}
182
183impl<T> WeightFunctionInfo<T> {
184    /// Calculates the total number of weighted densities for each functional
185    /// from multiple weight functions depending on dimension.
186    pub fn n_weighted_densities(&self, dimensions: usize) -> usize {
187        let segments = self.component_index.len();
188        (if self.local_density { segments } else { 0 })
189            + self.scalar_component_weighted_densities.len() * segments
190            + self.vector_component_weighted_densities.len() * segments * dimensions
191            + self.scalar_fmt_weighted_densities.len()
192            + self.vector_fmt_weighted_densities.len() * dimensions
193    }
194}
195
196impl<T> WeightFunctionInfo<T> {
197    /// Initializing empty `WeightFunctionInfo`.
198    pub fn new(component_index: Array1<usize>, local_density: bool) -> Self {
199        Self {
200            component_index,
201            local_density,
202            scalar_component_weighted_densities: Vec::new(),
203            vector_component_weighted_densities: Vec::new(),
204            scalar_fmt_weighted_densities: Vec::new(),
205            vector_fmt_weighted_densities: Vec::new(),
206        }
207    }
208
209    /// Adds and sorts [WeightFunction] depending on information
210    /// about {FMT, component} & {scalar-valued, vector-valued}.
211    pub fn add(mut self, weight_function: WeightFunction<T>, fmt: bool) -> Self {
212        let segments = self.component_index.len();
213
214        // Check size of `kernel_radius`
215        if segments != weight_function.kernel_radius.len() {
216            panic!(
217                "Number of segments is fixed to {}; `kernel_radius` has {} entries.",
218                segments,
219                weight_function.kernel_radius.len()
220            );
221        }
222
223        // Check size of `prefactor`
224        if segments != weight_function.prefactor.len() {
225            panic!(
226                "Number of segments is fixed to {}; `prefactor` has {} entries.",
227                segments,
228                weight_function.prefactor.len()
229            );
230        }
231
232        // Add new `WeightFunction`
233        match (fmt, weight_function.shape) {
234            // {Component, vector}
235            (false, WeightFunctionShape::DeltaVec) => self
236                .vector_component_weighted_densities
237                .push(weight_function),
238            // {Component, scalar}
239            (false, _) => self
240                .scalar_component_weighted_densities
241                .push(weight_function),
242            // {FMT, vector}
243            (true, WeightFunctionShape::DeltaVec) => {
244                self.vector_fmt_weighted_densities.push(weight_function)
245            }
246            // {FMT, scalar}
247            (true, _) => self.scalar_fmt_weighted_densities.push(weight_function),
248        };
249
250        // Return
251        self
252    }
253
254    /// Adds and sorts multiple [WeightFunction]s.
255    pub fn extend(mut self, weight_functions: Vec<WeightFunction<T>>, fmt: bool) -> Self {
256        // Add each element of vector
257        for wf in weight_functions {
258            self = self.add(wf, fmt);
259        }
260
261        // Return
262        self
263    }
264
265    /// Expose weight functions outside of this crate
266    pub fn as_slice(&self) -> [&Vec<WeightFunction<T>>; 4] {
267        [
268            &self.scalar_component_weighted_densities,
269            &self.vector_component_weighted_densities,
270            &self.scalar_fmt_weighted_densities,
271            &self.vector_fmt_weighted_densities,
272        ]
273    }
274}
275
276impl<T: DualNum<f64> + Copy> WeightFunctionInfo<T> {
277    /// calculates the matrix of weight constants for this set of weighted densities
278    pub fn weight_constants(&self, k: T, dimensions: usize) -> Array2<T> {
279        let segments = self.component_index.len();
280        let n_wd = self.n_weighted_densities(dimensions);
281        let mut weight_constants = Array::zeros([n_wd, segments]);
282        let mut j = 0;
283        if self.local_density {
284            weight_constants
285                .slice_mut(s![j..j + segments, ..])
286                .into_diag()
287                .assign(&Array::ones(segments));
288            j += segments;
289        }
290        for w in &self.scalar_component_weighted_densities {
291            weight_constants
292                .slice_mut(s![j..j + segments, ..])
293                .into_diag()
294                .assign(&w.scalar_weight_constants(k));
295            j += segments;
296        }
297        if dimensions == 1 {
298            for w in &self.vector_component_weighted_densities {
299                weight_constants
300                    .slice_mut(s![j..j + segments, ..])
301                    .into_diag()
302                    .assign(&w.vector_weight_constants(k));
303                j += segments;
304            }
305        }
306        for w in &self.scalar_fmt_weighted_densities {
307            weight_constants
308                .slice_mut(s![j, ..])
309                .assign(&w.scalar_weight_constants(k));
310            j += 1;
311        }
312        if dimensions == 1 {
313            for w in &self.vector_fmt_weighted_densities {
314                weight_constants
315                    .slice_mut(s![j, ..])
316                    .assign(&w.vector_weight_constants(k));
317                j += 1;
318            }
319        }
320        weight_constants
321    }
322}
feos_dft/weight_functions.rs

feos_dft/
weight_functions.rs
