feos_dft/
weight_functions.rs

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

feos_dft/
weight_functions.rs
