use feos_core::StateHD;
use nalgebra::DVector;
use num_dual::DualNum;
use std::borrow::Cow;
use std::f64::consts::FRAC_PI_6;
#[cfg(feature = "dft")]
mod dft;
#[cfg(feature = "dft")]
pub use dft::{FMTContribution, FMTFunctional, FMTVersion, HardSphereParameters};
pub enum MonomerShape<'a, D> {
Spherical(usize),
NonSpherical(DVector<D>),
Heterosegmented([DVector<D>; 4], &'a DVector<usize>),
}
pub trait HardSphereProperties {
fn monomer_shape<D: DualNum<f64> + Copy>(&self, temperature: D) -> MonomerShape<'_, D>;
fn hs_diameter<D: DualNum<f64> + Copy>(&self, temperature: D) -> DVector<D>;
fn component_index(&self) -> Cow<'_, [usize]> {
match self.monomer_shape(1.0) {
MonomerShape::Spherical(n) => Cow::Owned((0..n).collect()),
MonomerShape::NonSpherical(m) => Cow::Owned((0..m.len()).collect()),
MonomerShape::Heterosegmented(_, component_index) => {
Cow::Borrowed(component_index.as_slice())
}
}
}
fn geometry_coefficients<D: DualNum<f64> + Copy>(&self, temperature: D) -> [DVector<D>; 4] {
match self.monomer_shape(temperature) {
MonomerShape::Spherical(n) => {
let m = DVector::from_element(n, D::from(1.0));
[m.clone(), m.clone(), m.clone(), m]
}
MonomerShape::NonSpherical(m) => [m.clone(), m.clone(), m.clone(), m],
MonomerShape::Heterosegmented(g, _) => g,
}
}
fn zeta<D: DualNum<f64> + Copy, const N: usize>(
&self,
temperature: D,
partial_density: &DVector<D>,
k: [i32; N],
) -> [D; N] {
let component_index = self.component_index();
let geometry_coefficients = self.geometry_coefficients(temperature);
let diameter = self.hs_diameter(temperature);
let mut zeta = [D::zero(); N];
for i in 0..diameter.len() {
for (z, &k) in zeta.iter_mut().zip(k.iter()) {
*z += partial_density[component_index[i]]
* diameter[i].powi(k)
* (geometry_coefficients[k as usize][i] * FRAC_PI_6);
}
}
zeta
}
}
pub struct HardSphere;
impl HardSphere {
#[inline]
pub fn helmholtz_energy_density_and_properties<
D: DualNum<f64> + Copy,
P: HardSphereProperties,
>(
&self,
parameters: &P,
state: &StateHD<D>,
) -> (D, [D; 4], DVector<D>) {
let p = parameters;
let diameter = p.hs_diameter(state.temperature);
let component_index = p.component_index();
let geometry_coefficients = p.geometry_coefficients(state.temperature);
let mut zeta = [D::zero(); 4];
for i in 0..diameter.len() {
for (z, &k) in zeta.iter_mut().zip([0, 1, 2, 3].iter()) {
*z += state.molefracs[component_index[i]]
* diameter[i].powi(k)
* (geometry_coefficients[k as usize][i] * FRAC_PI_6);
}
}
let zeta_23 = zeta[2] / zeta[3];
let density = state.partial_density.sum();
zeta.iter_mut().for_each(|z| *z *= density);
let frac_1mz3 = -(zeta[3] - 1.0).recip();
let a = (zeta[1] * zeta[2] * frac_1mz3 * 3.0
+ zeta[2].powi(2) * frac_1mz3.powi(2) * zeta_23
+ (zeta[2] * zeta_23.powi(2) - zeta[0]) * (-zeta[3]).ln_1p())
/ std::f64::consts::FRAC_PI_6;
(a, zeta, diameter)
}
#[inline]
pub fn helmholtz_energy_density<D: DualNum<f64> + Copy, P: HardSphereProperties>(
&self,
parameters: &P,
state: &StateHD<D>,
) -> D {
self.helmholtz_energy_density_and_properties(parameters, state)
.0
}
}