use feos_core::{FeosResult, ResidualDyn, StateHD, Subset};
use feos_dft::adsorption::FluidParameters;
use feos_dft::solvation::PairPotential;
use feos_dft::{
FunctionalContribution, HelmholtzEnergyFunctional, HelmholtzEnergyFunctionalDyn, MoleculeShape,
WeightFunction, WeightFunctionInfo, WeightFunctionShape,
};
use nalgebra::DVector;
use ndarray::*;
use num_dual::DualNum;
use std::f64::consts::PI;
use super::{HardSphereProperties, MonomerShape};
const PI36M1: f64 = 1.0 / (36.0 * PI);
const N3_CUTOFF: f64 = 1e-5;
#[derive(Clone, Copy, PartialEq)]
pub enum FMTVersion {
WhiteBear,
KierlikRosinberg,
AntiSymWhiteBear,
}
pub struct FMTContribution<'p, P> {
pub properties: &'p P,
version: FMTVersion,
}
impl<'p, P> FMTContribution<'p, P> {
pub fn new(properties: &'p P, version: FMTVersion) -> Self {
Self {
properties,
version,
}
}
}
impl<'p, P: HardSphereProperties + Send + Sync> FunctionalContribution for FMTContribution<'p, P> {
fn name(&self) -> &'static str {
match self.version {
FMTVersion::WhiteBear => "FMT functional (WB)",
FMTVersion::KierlikRosinberg => "FMT functional (KR)",
FMTVersion::AntiSymWhiteBear => "FMT functional (AntiSymWB)",
}
}
fn weight_functions<N: DualNum<f64> + Copy>(&self, temperature: N) -> WeightFunctionInfo<N> {
let r = self.properties.hs_diameter(temperature) * N::from(0.5);
let [c0, c1, c2, c3] = self.properties.geometry_coefficients(temperature);
match (self.version, r.len()) {
(FMTVersion::WhiteBear | FMTVersion::AntiSymWhiteBear, 1) => WeightFunctionInfo::new(
self.properties.component_index().into_owned().into(),
false,
)
.extend(
vec![
WeightFunctionShape::Delta,
WeightFunctionShape::Theta,
WeightFunctionShape::DeltaVec,
]
.into_iter()
.zip([c2, c3.clone(), c3])
.map(|(s, c)| WeightFunction {
prefactor: c,
kernel_radius: r.clone(),
shape: s,
})
.collect(),
false,
),
(FMTVersion::WhiteBear | FMTVersion::AntiSymWhiteBear, _) => WeightFunctionInfo::new(
self.properties.component_index().into_owned().into(),
false,
)
.add(
WeightFunction {
prefactor: c0.zip_map(&r, |c, r| r.powi(-2) * c / (4.0 * PI)),
kernel_radius: r.clone(),
shape: WeightFunctionShape::Delta,
},
true,
)
.add(
WeightFunction {
prefactor: c1.zip_map(&r, |c, r| r.recip() * c / (4.0 * PI)),
kernel_radius: r.clone(),
shape: WeightFunctionShape::Delta,
},
true,
)
.add(
WeightFunction {
prefactor: c2,
kernel_radius: r.clone(),
shape: WeightFunctionShape::Delta,
},
true,
)
.add(
WeightFunction {
prefactor: c3.clone(),
kernel_radius: r.clone(),
shape: WeightFunctionShape::Theta,
},
true,
)
.add(
WeightFunction {
prefactor: c3.zip_map(&r, |c, r| r.recip() * c / (4.0 * PI)),
kernel_radius: r.clone(),
shape: WeightFunctionShape::DeltaVec,
},
true,
)
.add(
WeightFunction {
prefactor: c3,
kernel_radius: r,
shape: WeightFunctionShape::DeltaVec,
},
true,
),
(FMTVersion::KierlikRosinberg, _) => WeightFunctionInfo::new(
self.properties.component_index().into_owned().into(),
false,
)
.extend(
vec![
WeightFunctionShape::KR0,
WeightFunctionShape::KR1,
WeightFunctionShape::Delta,
WeightFunctionShape::Theta,
]
.into_iter()
.zip(self.properties.geometry_coefficients(temperature))
.map(|(s, c)| WeightFunction {
prefactor: c,
kernel_radius: r.clone(),
shape: s,
})
.collect(),
true,
),
}
}
fn helmholtz_energy_density<N: DualNum<f64> + Copy>(
&self,
temperature: N,
weighted_densities: ArrayView2<N>,
) -> FeosResult<Array1<N>> {
let pure_component_weighted_densities = matches!(
self.version,
FMTVersion::WhiteBear | FMTVersion::AntiSymWhiteBear
) && self.properties.component_index().len() == 1;
let (n2, n3) = if pure_component_weighted_densities {
(
weighted_densities.index_axis(Axis(0), 0),
weighted_densities.index_axis(Axis(0), 1),
)
} else {
(
weighted_densities.index_axis(Axis(0), 2),
weighted_densities.index_axis(Axis(0), 3),
)
};
let (n0, n1) = if pure_component_weighted_densities {
let r = self.properties.hs_diameter(temperature)[0] * 0.5;
(
n2.mapv(|n2| n2 / (r * r * 4.0 * PI)),
n2.mapv(|n2| n2 / (r * 4.0 * PI)),
)
} else {
(
weighted_densities.index_axis(Axis(0), 0).to_owned(),
weighted_densities.index_axis(Axis(0), 1).to_owned(),
)
};
let (n1n2, n2n2) = match self.version {
FMTVersion::WhiteBear | FMTVersion::AntiSymWhiteBear => {
let (n1v, n2v) = if pure_component_weighted_densities {
let r = self.properties.hs_diameter(temperature)[0] * 0.5;
let n2v = weighted_densities.slice_axis(Axis(0), Slice::new(2, None, 1));
(n2v.mapv(|n2v| n2v / (r * 4.0 * PI)), n2v)
} else {
let dim = ((weighted_densities.shape()[0] - 4) / 2) as isize;
(
weighted_densities
.slice_axis(Axis(0), Slice::new(4, Some(4 + dim), 1))
.to_owned(),
weighted_densities
.slice_axis(Axis(0), Slice::new(4 + dim, Some(4 + 2 * dim), 1)),
)
};
match self.version {
FMTVersion::WhiteBear => (
&n1 * &n2 - (&n1v * &n2v).sum_axis(Axis(0)),
&n2 * &n2 - (&n2v * &n2v).sum_axis(Axis(0)) * 3.0,
),
FMTVersion::AntiSymWhiteBear => {
let mut xi2 = (&n2v * &n2v).sum_axis(Axis(0)) / n2.map(|n| n.powi(2));
xi2.iter_mut().for_each(|x| {
if x.re() > 1.0 {
*x = N::one()
}
});
(
&n1 * &n2 - (&n1v * &n2v).sum_axis(Axis(0)),
&n2 * &n2 * xi2.mapv(|x| (-x + 1.0).powi(3)),
)
}
_ => unreachable!(),
}
}
FMTVersion::KierlikRosinberg => (&n1 * &n2, &n2 * &n2),
};
let ln31 = n3.mapv(|n3| (-n3).ln_1p());
let n3rec = n3.mapv(|n3| n3.recip());
let n3m1 = n3.mapv(|n3| -n3 + 1.0);
let n3m1rec = n3m1.mapv(|n3m1| n3m1.recip());
let mut f3 = (&n3m1 * &n3m1 * &ln31 + n3) * &n3rec * n3rec * &n3m1rec * &n3m1rec;
f3.iter_mut().zip(n3).for_each(|(f3, &n3)| {
if n3.re() < N3_CUTOFF {
*f3 = (((n3 * 35.0 / 6.0 + 4.8) * n3 + 3.75) * n3 + 8.0 / 3.0) * n3 + 1.5;
}
});
Ok(-(&n0 * &ln31) + n1n2 * &n3m1rec + n2n2 * n2 * PI36M1 * f3)
}
}
pub struct HardSphereParameters {
sigma: DVector<f64>,
}
impl HardSphereProperties for HardSphereParameters {
fn monomer_shape<N>(&self, _: N) -> MonomerShape<'_, N> {
MonomerShape::Spherical(self.sigma.len())
}
fn hs_diameter<N: DualNum<f64>>(&self, _: N) -> DVector<N> {
self.sigma.map(N::from)
}
}
pub struct FMTFunctional {
properties: HardSphereParameters,
version: FMTVersion,
}
impl FMTFunctional {
pub fn new(sigma: DVector<f64>, version: FMTVersion) -> Self {
let properties = HardSphereParameters { sigma };
Self {
properties,
version,
}
}
}
impl Subset for FMTFunctional {
fn subset(&self, component_list: &[usize]) -> Self {
let sigma = DVector::from_vec(
component_list
.iter()
.map(|&c| self.properties.sigma[c])
.collect(),
);
Self::new(sigma, self.version)
}
}
impl ResidualDyn for FMTFunctional {
fn components(&self) -> usize {
self.properties.sigma.len()
}
fn compute_max_density<D: DualNum<f64> + Copy>(&self, molefracs: &DVector<D>) -> D {
molefracs.dot(&self.properties.sigma.map(D::from)).recip() * 1.2
}
fn reduced_helmholtz_energy_density_contributions<D: DualNum<f64> + Copy>(
&self,
state: &StateHD<D>,
) -> Vec<(&'static str, D)> {
self.evaluate_bulk(state)
}
}
impl HelmholtzEnergyFunctionalDyn for FMTFunctional {
type Contribution<'a>
= FMTContribution<'a, HardSphereParameters>
where
Self: 'a;
fn contributions<'a>(
&'a self,
) -> impl Iterator<Item = FMTContribution<'a, HardSphereParameters>> {
std::iter::once(FMTContribution::new(&self.properties, self.version))
}
fn molecule_shape(&self) -> MoleculeShape<'_> {
MoleculeShape::Spherical(self.properties.sigma.len())
}
}
impl PairPotential for FMTFunctional {
fn pair_potential(&self, i: usize, r: &Array1<f64>, _: f64) -> Array2<f64> {
let s = &self.properties.sigma;
Array::from_shape_fn((s.len(), r.len()), |(j, k)| {
if r[k] > 0.5 * (s[i] + s[j]) {
0.0
} else {
f64::INFINITY
}
})
}
}
impl FluidParameters for FMTFunctional {
fn epsilon_k_ff(&self) -> DVector<f64> {
DVector::zeros(self.properties.sigma.len())
}
fn sigma_ff(&self) -> DVector<f64> {
self.properties.sigma.clone()
}
}