use crate::weight_functions::WeightFunctionInfo;
use feos_core::{FeosResult, StateHD};
use ndarray::RemoveAxis;
use ndarray::prelude::*;
use num_dual::*;
use num_traits::Zero;
pub trait FunctionalContribution: Sync + Send {
fn name(&self) -> &'static str;
fn weight_functions<N: DualNum<f64> + Copy>(&self, temperature: N) -> WeightFunctionInfo<N>;
fn weight_functions_pdgt<N: DualNum<f64> + Copy>(
&self,
temperature: N,
) -> WeightFunctionInfo<N> {
self.weight_functions(temperature)
}
fn helmholtz_energy_density<N: DualNum<f64> + Copy>(
&self,
temperature: N,
weighted_densities: ArrayView2<N>,
) -> FeosResult<Array1<N>>;
fn bulk_helmholtz_energy_density<N: DualNum<f64> + Copy>(&self, state: &StateHD<N>) -> N {
let weight_functions = self.weight_functions(state.temperature);
let density: Array1<_> = weight_functions
.component_index
.iter()
.map(|&c| state.partial_density[c])
.collect();
let weight_constants = weight_functions.weight_constants(Zero::zero(), 0);
let weighted_densities = weight_constants.dot(&density).insert_axis(Axis(1));
self.helmholtz_energy_density(state.temperature, weighted_densities.view())
.unwrap()[0]
}
fn first_partial_derivatives<N: DualNum<f64> + Copy>(
&self,
temperature: N,
weighted_densities: Array2<N>,
mut helmholtz_energy_density: ArrayViewMut1<N>,
mut first_partial_derivative: ArrayViewMut2<N>,
) -> FeosResult<()> {
let mut wd = weighted_densities.mapv(Dual::from_re);
let t = Dual::from_re(temperature);
let mut phi = Array::zeros(weighted_densities.raw_dim().remove_axis(Axis(0)));
for i in 0..wd.shape()[0] {
wd.index_axis_mut(Axis(0), i)
.map_inplace(|x| x.eps = N::one());
phi = self.helmholtz_energy_density(t, wd.view())?;
first_partial_derivative
.index_axis_mut(Axis(0), i)
.assign(&phi.mapv(|p| p.eps));
wd.index_axis_mut(Axis(0), i)
.map_inplace(|x| x.eps = N::zero());
}
helmholtz_energy_density.assign(&phi.mapv(|p| p.re));
Ok(())
}
fn second_partial_derivatives(
&self,
temperature: f64,
weighted_densities: ArrayView2<f64>,
mut helmholtz_energy_density: ArrayViewMut1<f64>,
mut first_partial_derivative: ArrayViewMut2<f64>,
mut second_partial_derivative: ArrayViewMut3<f64>,
) -> FeosResult<()> {
let mut wd = weighted_densities.mapv(HyperDual64::from);
let t = HyperDual64::from(temperature);
let mut phi = Array::zeros(weighted_densities.raw_dim().remove_axis(Axis(0)));
for i in 0..wd.shape()[0] {
wd.index_axis_mut(Axis(0), i).map_inplace(|x| x.eps1 = 1.0);
for j in 0..=i {
wd.index_axis_mut(Axis(0), j).map_inplace(|x| x.eps2 = 1.0);
phi = self.helmholtz_energy_density(t, wd.view())?;
let p = phi.mapv(|p| p.eps1eps2);
second_partial_derivative
.index_axis_mut(Axis(0), i)
.index_axis_mut(Axis(0), j)
.assign(&p);
if i != j {
second_partial_derivative
.index_axis_mut(Axis(0), j)
.index_axis_mut(Axis(0), i)
.assign(&p);
}
wd.index_axis_mut(Axis(0), j).map_inplace(|x| x.eps2 = 0.0);
}
first_partial_derivative
.index_axis_mut(Axis(0), i)
.assign(&phi.mapv(|p| p.eps1));
wd.index_axis_mut(Axis(0), i).map_inplace(|x| x.eps1 = 0.0);
}
helmholtz_energy_density.assign(&phi.mapv(|p| p.re));
Ok(())
}
}