feos-dft 0.9.5

use nalgebra::DVector;
use ndarray::*;
use num_dual::DualNum;
// use rustfft::num_complex::Complex;
use std::f64::consts::{FRAC_PI_3, PI};
use std::ops::Mul;

/// A weight function corresponding to a single weighted density.
#[derive(Clone)]
pub struct WeightFunction<T> {
    /// Factor in front of normalized weight function
    pub prefactor: DVector<T>,
    /// Kernel radius of the convolution
    pub kernel_radius: DVector<T>,
    /// Shape of the weight function (Dirac delta, Heaviside, etc.)
    pub shape: WeightFunctionShape,
}

impl<T: DualNum<f64> + Copy> WeightFunction<T> {
    /// Create a new weight function without prefactor
    pub fn new_unscaled(kernel_radius: DVector<T>, shape: WeightFunctionShape) -> Self {
        Self {
            prefactor: DVector::from_element(kernel_radius.len(), T::one()),
            kernel_radius,
            shape,
        }
    }

    /// Create a new weight function with weight constant = 1
    pub fn new_scaled(kernel_radius: DVector<T>, shape: WeightFunctionShape) -> Self {
        let unscaled = Self::new_unscaled(kernel_radius, shape);
        let weight_constants = unscaled.scalar_weight_constants(T::zero());
        let weight_constants = DVector::from(weight_constants.to_vec());
        Self {
            prefactor: weight_constants.map(|w| w.recip()),
            kernel_radius: unscaled.kernel_radius,
            shape,
        }
    }

    /// Calculates the value of the scalar weight function depending on its shape
    /// `k_abs` describes the absolute value of the Fourier variable
    pub(crate) fn fft_scalar_weight_functions<D: Dimension, T2: DualNum<f64>>(
        &self,
        k_abs: &Array<T2, D>,
        lanczos: &Option<Array<T2, D>>,
    ) -> Array<T, D::Larger>
    where
        T: Mul<T2, Output = T>,
        D::Larger: Dimension<Smaller = D>,
    {
        // Allocate vector for weight functions
        let mut d = vec![self.prefactor.len()];
        k_abs.shape().iter().for_each(|&x| d.push(x));
        let mut w = Array::zeros(d.into_dimension())
            .into_dimensionality()
            .unwrap();

        // Calculate weight function for each component
        for i in 0..self.prefactor.len() {
            let radius = self.kernel_radius[i];
            let p = self.prefactor[i];
            let rik = k_abs.mapv(|k| radius * k);
            let mut w_i = w.index_axis_mut(Axis(0), i);
            w_i.assign(&match self.shape {
                WeightFunctionShape::Theta => rik.mapv(|rik| {
                    (rik.sph_j0() + rik.sph_j2()) * 4.0 * FRAC_PI_3 * radius.powi(3) * p
                }),
                WeightFunctionShape::Delta => {
                    rik.mapv(|rik| rik.sph_j0() * 4.0 * PI * radius.powi(2) * p)
                }
                WeightFunctionShape::KR1 => {
                    rik.mapv(|rik| (rik.sph_j0() + rik.cos()) * 0.5 * radius * p)
                }
                WeightFunctionShape::KR0 => rik.mapv(|rik| (rik * rik.sin() * 0.5 + rik.cos()) * p),
                WeightFunctionShape::DeltaVec => unreachable!(),
            });

            // Apply Lanczos sigma factor
            if let Some(l) = lanczos {
                let w_i_l = &w_i * l;
                w_i.assign(&w_i_l);
            }
        }

        // Return real part of weight function
        w
    }

    /// Calculates the value of the vector weight function depending on its shape
    /// `k_abs` describes the absolute value of the Fourier variable
    /// `k` describes the (potentially multi-dimensional) Fourier variable
    pub(crate) fn fft_vector_weight_functions<D: Dimension, T2: DualNum<f64>>(
        &self,
        k_abs: &Array<T2, D>,
        k: &Array<T2, D::Larger>,
        lanczos: &Option<Array<T2, D>>,
    ) -> Array<T, <D::Larger as Dimension>::Larger>
    where
        D::Larger: Dimension<Smaller = D>,
        <D::Larger as Dimension>::Larger: Dimension<Smaller = D::Larger>,
        T: Mul<T2, Output = T>,
    {
        // Allocate vector for weight functions
        let mut d = vec![k.shape()[0], self.prefactor.len()];
        k.shape().iter().skip(1).for_each(|&x| d.push(x));
        let mut w = Array::zeros(d.into_dimension())
            .into_dimensionality()
            .unwrap();

        // Iterate all dimensions
        for (k_x, mut w_x) in k.outer_iter().zip(w.outer_iter_mut()) {
            // Calculate weight function for each component
            for i in 0..self.prefactor.len() {
                let radius = self.kernel_radius[i];
                let p = self.prefactor[i];
                let rik = k_abs.mapv(|k| radius * k);
                let mut w_i = w_x.index_axis_mut(Axis(0), i);
                w_i.assign(&match self.shape {
                    WeightFunctionShape::DeltaVec => {
                        &rik.mapv(|rik| {
                            (rik.sph_j0() + rik.sph_j2()) * (radius.powi(3) * 4.0 * FRAC_PI_3 * p)
                        }) * &k_x
                    }
                    _ => unreachable!(),
                });

                // Apply Lanczos sigma factor
                if let Some(l) = lanczos {
                    let w_i_l = &w_i * l;
                    w_i.assign(&w_i_l)
                }
            }
        }
        // Return imaginary part of weight function
        w
    }

    /// Scalar weights for the bulk convolver (for the bulk convolver only the prefactor
    /// of the normalized weight functions are required)
    pub fn scalar_weight_constants(&self, k: T) -> Array1<T> {
        let k = arr0(k);
        self.fft_scalar_weight_functions(&k, &None)
    }

    /// Vector weights for the bulk convolver (for the bulk convolver only the prefactor
    /// of the normalized weight functions are required)
    pub fn vector_weight_constants(&self, k: T) -> Array1<T> {
        let k_abs = arr0(k);
        let k = arr1(&[k]);
        self.fft_vector_weight_functions(&k_abs, &k, &None)
            .index_axis_move(Axis(0), 0)
    }
}

/// Possible weight function shapes.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum WeightFunctionShape {
    /// Heaviside step function
    Theta,
    /// Dirac delta function
    Delta,
    /// Combination of first & second derivative of Dirac delta function (with different prefactor; only in Kierlik-Rosinberg functional)
    KR0,
    /// First derivative of Dirac delta function (with different prefactor; only in Kierlik-Rosinberg functional)
    KR1,
    /// Vector-shape as combination of Dirac delta and outward normal
    DeltaVec,
}

/// Information about weight functions
pub struct WeightFunctionInfo<T> {
    /// Index of the component that each individual segment belongs to.
    pub(crate) component_index: DVector<usize>,
    /// Flag if local density is required in the functional
    pub(crate) local_density: bool,
    /// Container for scalar component-wise weighted densities
    pub(crate) scalar_component_weighted_densities: Vec<WeightFunction<T>>,
    /// Container for vector component-wise weighted densities
    pub(crate) vector_component_weighted_densities: Vec<WeightFunction<T>>,
    /// Container for scalar FMT weighted densities
    pub(crate) scalar_fmt_weighted_densities: Vec<WeightFunction<T>>,
    /// Container for vector FMT weighted densities
    pub(crate) vector_fmt_weighted_densities: Vec<WeightFunction<T>>,
}

impl<T> WeightFunctionInfo<T> {
    /// Calculates the total number of weighted densities for each functional
    /// from multiple weight functions depending on dimension.
    pub fn n_weighted_densities(&self, dimensions: usize) -> usize {
        let segments = self.component_index.len();
        (if self.local_density { segments } else { 0 })
            + self.scalar_component_weighted_densities.len() * segments
            + self.vector_component_weighted_densities.len() * segments * dimensions
            + self.scalar_fmt_weighted_densities.len()
            + self.vector_fmt_weighted_densities.len() * dimensions
    }
}

impl<T> WeightFunctionInfo<T> {
    /// Initializing empty `WeightFunctionInfo`.
    pub fn new(component_index: DVector<usize>, local_density: bool) -> Self {
        Self {
            component_index,
            local_density,
            scalar_component_weighted_densities: Vec::new(),
            vector_component_weighted_densities: Vec::new(),
            scalar_fmt_weighted_densities: Vec::new(),
            vector_fmt_weighted_densities: Vec::new(),
        }
    }

    /// Adds and sorts [WeightFunction] depending on information
    /// about {FMT, component} & {scalar-valued, vector-valued}.
    pub fn add(mut self, weight_function: WeightFunction<T>, fmt: bool) -> Self {
        let segments = self.component_index.len();

        // Check size of `kernel_radius`
        if segments != weight_function.kernel_radius.len() {
            panic!(
                "Number of segments is fixed to {}; `kernel_radius` has {} entries.",
                segments,
                weight_function.kernel_radius.len()
            );
        }

        // Check size of `prefactor`
        if segments != weight_function.prefactor.len() {
            panic!(
                "Number of segments is fixed to {}; `prefactor` has {} entries.",
                segments,
                weight_function.prefactor.len()
            );
        }

        // Add new `WeightFunction`
        match (fmt, weight_function.shape) {
            // {Component, vector}
            (false, WeightFunctionShape::DeltaVec) => self
                .vector_component_weighted_densities
                .push(weight_function),
            // {Component, scalar}
            (false, _) => self
                .scalar_component_weighted_densities
                .push(weight_function),
            // {FMT, vector}
            (true, WeightFunctionShape::DeltaVec) => {
                self.vector_fmt_weighted_densities.push(weight_function)
            }
            // {FMT, scalar}
            (true, _) => self.scalar_fmt_weighted_densities.push(weight_function),
        };

        // Return
        self
    }

    /// Adds and sorts multiple [WeightFunction]s.
    pub fn extend(mut self, weight_functions: Vec<WeightFunction<T>>, fmt: bool) -> Self {
        // Add each element of vector
        for wf in weight_functions {
            self = self.add(wf, fmt);
        }

        // Return
        self
    }

    /// Expose weight functions outside of this crate
    pub fn as_slice(&self) -> [&Vec<WeightFunction<T>>; 4] {
        [
            &self.scalar_component_weighted_densities,
            &self.vector_component_weighted_densities,
            &self.scalar_fmt_weighted_densities,
            &self.vector_fmt_weighted_densities,
        ]
    }
}

impl<T: DualNum<f64> + Copy> WeightFunctionInfo<T> {
    /// calculates the matrix of weight constants for this set of weighted densities
    pub fn weight_constants(&self, k: T, dimensions: usize) -> Array2<T> {
        let segments = self.component_index.len();
        let n_wd = self.n_weighted_densities(dimensions);
        let mut weight_constants = Array::zeros([n_wd, segments]);
        let mut j = 0;
        if self.local_density {
            weight_constants
                .slice_mut(s![j..j + segments, ..])
                .into_diag()
                .assign(&Array::ones(segments));
            j += segments;
        }
        for w in &self.scalar_component_weighted_densities {
            weight_constants
                .slice_mut(s![j..j + segments, ..])
                .into_diag()
                .assign(&w.scalar_weight_constants(k));
            j += segments;
        }
        if dimensions == 1 {
            for w in &self.vector_component_weighted_densities {
                weight_constants
                    .slice_mut(s![j..j + segments, ..])
                    .into_diag()
                    .assign(&w.vector_weight_constants(k));
                j += segments;
            }
        }
        for w in &self.scalar_fmt_weighted_densities {
            weight_constants
                .slice_mut(s![j, ..])
                .assign(&w.scalar_weight_constants(k));
            j += 1;
        }
        if dimensions == 1 {
            for w in &self.vector_fmt_weighted_densities {
                weight_constants
                    .slice_mut(s![j, ..])
                    .assign(&w.vector_weight_constants(k));
                j += 1;
            }
        }
        weight_constants
    }
}