Less redundant approach to the DFT wrapper struct (#27)

feos-org · prehner · Apr 12, 2022 · Mar 9, 2022 · Mar 10, 2022 · Mar 10, 2022
commit 42a8fbec2d4e45545f79e850e2ced2d0d010cf52
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `DFTSolver` now uses `Verbosity` instead of a `bool` to control its output. [#19](https://github.com/feos-org/feos-dft/pull/19)
 - `SurfaceTensionDiagram` now uses the new `StateVec` struct to access properties of the bulk phases. [#19](https://github.com/feos-org/feos-dft/pull/19)
 - `Pore1D::initialize` and `Pore3D::initialize` now accept initial values for the density profiles as optional arguments. [#24](https://github.com/feos-org/feos-dft/pull/24)
+- Internally restructured the `DFT` structure to avoid redundant data. [#24](https://github.com/feos-org/feos-dft/pull/24)
+- Removed the `m` function in `FluidParameters`, it is instead inferred from `HelmholtzEnergyFunctional` which is now a supertrait of `FluidParameters`. [#24](https://github.com/feos-org/feos-dft/pull/24)
 
 ### Packaging
 - Updated `pyo3` and `numpy` dependencies to 0.16.

diff --git a/src/adsorption/external_potential.rs b/src/adsorption/external_potential.rs
@@ -1,4 +1,5 @@
 use crate::adsorption::fea_potential::calculate_fea_potential;
+use crate::functional::HelmholtzEnergyFunctional;
 use crate::geometry::Geometry;
 use feos_core::EosUnit;
 use libc::c_double;
@@ -55,10 +56,9 @@ pub enum ExternalPotential<U> {
 }
 
 /// Parameters of the fluid required to evaluate the external potential.
-pub trait FluidParameters {
+pub trait FluidParameters: HelmholtzEnergyFunctional {
     fn epsilon_k_ff(&self) -> Array1<f64>;
     fn sigma_ff(&self) -> &Array1<f64>;
-    fn m(&self) -> Array1<f64>;
 }
 
 impl<U: EosUnit> ExternalPotential<U> {

diff --git a/src/adsorption/pore.rs b/src/adsorption/pore.rs
@@ -1,6 +1,6 @@
 use crate::adsorption::{ExternalPotential, FluidParameters};
 use crate::convolver::ConvolverFFT;
-use crate::functional::{HelmholtzEnergyFunctional, DFT};
+use crate::functional::{HelmholtzEnergyFunctional, MoleculeShape, DFT};
 use crate::functional_contribution::FunctionalContribution;
 use crate::geometry::{Axis, Geometry, Grid};
 use crate::profile::{DFTProfile, CUTOFF_RADIUS, MAX_POTENTIAL};
@@ -178,13 +178,12 @@ impl<U: EosUnit> PoreSpecification<U, Ix1> for Pore1D<U> {
         density: Option<&QuantityArray2<U>>,
         external_potential: Option<&Array2<f64>>,
     ) -> EosResult<PoreProfile1D<U, F>> {
-        let dft = &bulk.eos;
+        let dft: &F = &bulk.eos;
         let n_grid = self.n_grid.unwrap_or(DEFAULT_GRID_POINTS);
 
         let axis = match self.geometry {
             Geometry::Cartesian => {
-                let potential_offset =
-                    POTENTIAL_OFFSET * bulk.eos.functional.sigma_ff().max().unwrap();
+                let potential_offset = POTENTIAL_OFFSET * bulk.eos.sigma_ff().max().unwrap();
                 Axis::new_cartesian(n_grid, 0.5 * self.pore_size, Some(potential_offset))?
             }
             Geometry::Cylindrical => Axis::new_polar(n_grid, self.pore_size)?,
@@ -198,7 +197,7 @@ impl<U: EosUnit> PoreSpecification<U, Ix1> for Pore1D<U> {
                     self.pore_size,
                     bulk.temperature,
                     &self.potential,
-                    &bulk.eos.functional,
+                    dft,
                     &axis,
                     self.potential_cutoff,
                 )
@@ -209,7 +208,7 @@ impl<U: EosUnit> PoreSpecification<U, Ix1> for Pore1D<U> {
         // initialize convolver
         let grid = Grid::new_1d(axis);
         let t = bulk.temperature.to_reduced(U::reference_temperature())?;
-        let weight_functions = dft.functional.weight_functions(t);
+        let weight_functions = dft.weight_functions(t);
         let convolver = ConvolverFFT::plan(&grid, &weight_functions, Some(1));
 
         Ok(PoreProfile {
@@ -231,7 +230,7 @@ impl<U: EosUnit> PoreSpecification<U, Ix3> for Pore3D<U> {
         density: Option<&QuantityArray4<U>>,
         external_potential: Option<&Array4<f64>>,
     ) -> EosResult<PoreProfile3D<U, F>> {
-        let dft = &bulk.eos;
+        let dft: &F = &bulk.eos;
 
         // generate grid
         let x = Axis::new_cartesian(self.n_grid[0], self.system_size[0], None)?;
@@ -252,7 +251,7 @@ impl<U: EosUnit> PoreSpecification<U, Ix3> for Pore3D<U> {
         let external_potential = external_potential.map_or_else(
             || {
                 external_potential_3d(
-                    &bulk.eos.functional,
+                    dft,
                     [&x, &y, &z],
                     self.system_size,
                     coordinates,
@@ -268,7 +267,7 @@ impl<U: EosUnit> PoreSpecification<U, Ix3> for Pore3D<U> {
 
         // initialize convolver
         let grid = Grid::Periodical3(x, y, z);
-        let weight_functions = dft.functional.weight_functions(t);
+        let weight_functions = dft.weight_functions(t);
         let convolver = ConvolverFFT::plan(&grid, &weight_functions, Some(1));
 
         Ok(PoreProfile {
@@ -453,7 +452,7 @@ impl Helium {
     fn new() -> DFT<Self> {
         let epsilon = arr1(&[EPSILON_HE]);
         let sigma = arr1(&[SIGMA_HE]);
-        DFT::new_homosegmented(Self { epsilon, sigma }, &Array1::ones(1))
+        (Self { epsilon, sigma }).into()
     }
 }
 
@@ -469,6 +468,10 @@ impl HelmholtzEnergyFunctional for Helium {
     fn compute_max_density(&self, _: &Array1<f64>) -> f64 {
         1.0
     }
+
+    fn molecule_shape(&self) -> MoleculeShape {
+        MoleculeShape::Spherical(1)
+    }
 }
 
 impl FluidParameters for Helium {
@@ -479,8 +482,4 @@ impl FluidParameters for Helium {
     fn sigma_ff(&self) -> &Array1<f64> {
         &self.sigma
     }
-
-    fn m(&self) -> Array1<f64> {
-        arr1(&[1.0])
-    }
 }
diff --git a/src/functional.rs b/src/functional.rs
@@ -12,50 +12,40 @@ use petgraph::graph::{Graph, UnGraph};
 use petgraph::visit::EdgeRef;
 use petgraph::Directed;
 use quantity::{QuantityArray, QuantityArray1, QuantityScalar};
+use std::borrow::Cow;
 use std::fmt;
-use std::ops::{AddAssign, MulAssign};
+use std::ops::{AddAssign, Deref, MulAssign};
 use std::rc::Rc;
 
 /// Wrapper struct for the [HelmholtzEnergyFunctional] trait.
+///
+/// Needed (for now) to generically implement the `EquationOfState`
+/// trait for Helmholtz energy functionals.
 #[derive(Clone)]
-pub struct DFT<T> {
-    /// Helmholtz energy functional
-    pub functional: T,
-    /// map segment -> component
-    pub component_index: Array1<usize>,
-    /// chain lengths of individual components
-    pub m: Array1<f64>,
-    /// ideal chain contribution
-    pub ideal_chain_contribution: IdealChainContribution,
-}
+pub struct DFT<F>(F);
 
-impl<T> DFT<T> {
-    /// Create a new DFT struct for a homosegmented Helmholtz energy functional.
-    pub fn new_homosegmented(functional: T, m: &Array1<f64>) -> Self {
-        let component_index = Array1::from_shape_fn(m.len(), |i| i);
-        Self::new(functional, &component_index, m)
+impl<F> From<F> for DFT<F> {
+    fn from(functional: F) -> Self {
+        Self(functional)
     }
+}
 
-    /// Create a new DFT struct for a heterosegmented Helmholtz energy functional.
-    pub fn new_heterosegmented(functional: T, component_index: &Array1<usize>) -> Self {
-        let m = Array1::ones(component_index.len());
-        Self::new(functional, component_index, &m)
+impl<F> DFT<F> {
+    pub fn into<F2: From<F>>(self) -> DFT<F2> {
+        DFT(self.0.into())
     }
+}
 
-    /// Create a new DFT struct for a general Helmholtz energy functional.
-    pub fn new(functional: T, component_index: &Array1<usize>, m: &Array1<f64>) -> Self {
-        Self {
-            functional,
-            component_index: component_index.clone(),
-            m: m.clone(),
-            ideal_chain_contribution: IdealChainContribution::new(component_index, m),
-        }
+impl<F> Deref for DFT<F> {
+    type Target = F;
+    fn deref(&self) -> &F {
+        &self.0
     }
 }
 
 impl<T: MolarWeight<U>, U: EosUnit> MolarWeight<U> for DFT<T> {
     fn molar_weight(&self) -> QuantityArray1<U> {
-        self.functional.molar_weight()
+        (self as &T).molar_weight()
     }
 }
 
@@ -74,15 +64,15 @@ impl fmt::Display for DefaultIdealGasContribution {
 
 impl<T: HelmholtzEnergyFunctional> EquationOfState for DFT<T> {
     fn components(&self) -> usize {
-        self.component_index[self.component_index.len() - 1] + 1
+        self.component_index()[self.component_index().len() - 1] + 1
     }
 
     fn subset(&self, component_list: &[usize]) -> Self {
-        self.functional.subset(component_list)
+        (self as &T).subset(component_list)
     }
 
     fn compute_max_density(&self, moles: &Array1<f64>) -> f64 {
-        self.functional.compute_max_density(moles)
+        (self as &T).compute_max_density(moles)
     }
 
     fn residual(&self) -> &[Box<dyn HelmholtzEnergy>] {
@@ -93,12 +83,11 @@ impl<T: HelmholtzEnergyFunctional> EquationOfState for DFT<T> {
     where
         dyn HelmholtzEnergy: HelmholtzEnergyDual<D>,
     {
-        self.functional
-            .contributions()
+        self.contributions()
             .iter()
             .map(|c| (c as &dyn HelmholtzEnergy).helmholtz_energy(state))
             .sum::<D>()
-            + self.ideal_chain_contribution.helmholtz_energy(state)
+            + self.ideal_chain_contribution().helmholtz_energy(state)
     }
 
     fn evaluate_residual_contributions<D: DualNum<f64>>(
@@ -109,7 +98,6 @@ impl<T: HelmholtzEnergyFunctional> EquationOfState for DFT<T> {
         dyn HelmholtzEnergy: HelmholtzEnergyDual<D>,
     {
         let mut res: Vec<(String, D)> = self
-            .functional
             .contributions()
             .iter()
             .map(|c| {
@@ -120,23 +108,36 @@ impl<T: HelmholtzEnergyFunctional> EquationOfState for DFT<T> {
             })
             .collect();
         res.push((
-            self.ideal_chain_contribution.to_string(),
-            self.ideal_chain_contribution.helmholtz_energy(state),
+            self.ideal_chain_contribution().to_string(),
+            self.ideal_chain_contribution().helmholtz_energy(state),
         ));
         res
     }
 
     fn ideal_gas(&self) -> &dyn IdealGasContribution {
-        self.functional.ideal_gas()
+        (self as &T).ideal_gas()
     }
 }
 
+/// Different representations for molecules within DFT.
+pub enum MoleculeShape<'a> {
+    /// For spherical molecules, the number of components.
+    Spherical(usize),
+    /// For non-spherical molecules in a homosegmented approach, the chain length parameter $m$.
+    NonSpherical(&'a Array1<f64>),
+    /// For non-spherical molecules in a heterosegmented approach, the component index for every segment.
+    Heterosegmented(&'a Array1<usize>),
+}
+
 /// A general Helmholtz energy functional.
 pub trait HelmholtzEnergyFunctional: Sized {
     /// Return a slice of [FunctionalContribution]s.
     fn contributions(&self) -> &[Box<dyn FunctionalContribution>];
 
-    /// Return a [DFT] for the specified subset of components.
+    /// Return the shape of the molecules and the necessary specifications.
+    fn molecule_shape(&self) -> MoleculeShape;
+
+    /// Return a functional for the specified subset of components.
     fn subset(&self, component_list: &[usize]) -> DFT<Self>;
 
     /// Return the maximum density in Angstrom^-3.
@@ -170,6 +171,28 @@ pub trait HelmholtzEnergyFunctional: Sized {
             .map(|c| c.weight_functions(temperature))
             .collect()
     }
+
+    fn m(&self) -> Cow<Array1<f64>> {
+        match self.molecule_shape() {
+            MoleculeShape::Spherical(n) => Cow::Owned(Array1::ones(n)),
+            MoleculeShape::NonSpherical(m) => Cow::Borrowed(m),
+            MoleculeShape::Heterosegmented(component_index) => {
+                Cow::Owned(Array1::ones(component_index.len()))
+            }
+        }
+    }
+
+    fn component_index(&self) -> Cow<Array1<usize>> {
+        match self.molecule_shape() {
+            MoleculeShape::Spherical(n) => Cow::Owned(Array1::from_shape_fn(n, |i| i)),
+            MoleculeShape::NonSpherical(m) => Cow::Owned(Array1::from_shape_fn(m.len(), |i| i)),
+            MoleculeShape::Heterosegmented(component_index) => Cow::Borrowed(component_index),
+        }
+    }
+
+    fn ideal_chain_contribution(&self) -> IdealChainContribution {
+        IdealChainContribution::new(&self.component_index(), &self.m())
+    }
 }
 
 impl<T: HelmholtzEnergyFunctional> DFT<T> {
@@ -191,11 +214,15 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         let (mut f, dfdrho) = self.functional_derivative(t, &rho, convolver)?;
 
         // calculate the grand potential density
-        for ((rho, dfdrho), &m) in rho.outer_iter().zip(dfdrho.outer_iter()).zip(self.m.iter()) {
+        for ((rho, dfdrho), &m) in rho
+            .outer_iter()
+            .zip(dfdrho.outer_iter())
+            .zip(self.m().iter())
+        {
             f -= &((&dfdrho + m) * &rho);
         }
 
-        let bond_lengths = self.functional.bond_lengths(t);
+        let bond_lengths = self.bond_lengths(t);
         for segment in bond_lengths.node_indices() {
             let n = bond_lengths.neighbors(segment).count();
             f += &(&rho.index_axis(Axis(0), segment.index()) * (0.5 * n as f64));
@@ -214,7 +241,7 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         D::Larger: Dimension<Smaller = D>,
     {
         let n = self.components();
-        let ig = self.functional.ideal_gas();
+        let ig = self.ideal_gas();
         let lambda = ig.de_broglie_wavelength(temperature, n);
         let mut phi = Array::zeros(density.raw_dim().remove_axis(Axis(0)));
         for (i, rhoi) in density.outer_iter().enumerate() {
@@ -233,7 +260,7 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         D::Larger: Dimension<Smaller = D>,
     {
         let n = self.components();
-        let ig = self.functional.ideal_gas();
+        let ig = self.ideal_gas();
         let lambda = ig.de_broglie_wavelength(temperature, n);
         let mut phi = Array::zeros(density.raw_dim().remove_axis(Axis(0)));
         for (i, rhoi) in density.outer_iter().enumerate() {
@@ -256,9 +283,9 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
     {
         let density_dual = density.mapv(N::from);
         let weighted_densities = convolver.weighted_densities(&density_dual);
-        let functional_contributions = self.functional.contributions();
+        let functional_contributions = self.contributions();
         let mut helmholtz_energy_density: Array<N, D> = self
-            .ideal_chain_contribution
+            .ideal_chain_contribution()
             .calculate_helmholtz_energy_density(&density.mapv(N::from))?;
         for (c, wd) in functional_contributions.iter().zip(weighted_densities) {
             let nwd = wd.shape()[0];
@@ -319,11 +346,11 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         let density_dual = density.mapv(Dual64::from);
         let temperature_dual = Dual64::from(temperature).derive();
         let weighted_densities = convolver.weighted_densities(&density_dual);
-        let functional_contributions = self.functional.contributions();
+        let functional_contributions = self.contributions();
         let mut helmholtz_energy_density: Vec<Array<Dual64, D>> =
             Vec::with_capacity(functional_contributions.len() + 1);
         helmholtz_energy_density.push(
-            self.ideal_chain_contribution
+            self.ideal_chain_contribution()
                 .calculate_helmholtz_energy_density(&density.mapv(Dual64::from))?,
         );
 
@@ -389,7 +416,7 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         D::Larger: Dimension<Smaller = D>,
     {
         let weighted_densities = convolver.weighted_densities(density);
-        let contributions = self.functional.contributions();
+        let contributions = self.contributions();
         let mut partial_derivatives = Vec::with_capacity(contributions.len());
         let mut helmholtz_energy_density = Array::zeros(density.raw_dim().remove_axis(Axis(0)));
         for (c, wd) in contributions.iter().zip(weighted_densities) {
@@ -424,7 +451,7 @@ impl<T: HelmholtzEnergyFunctional> DFT<T> {
         D::Larger: Dimension<Smaller = D>,
     {
         // calculate weight functions
-        let bond_lengths = self.functional.bond_lengths(temperature).into_edge_type();
+        let bond_lengths = self.bond_lengths(temperature).into_edge_type();
         let mut bond_weight_functions = bond_lengths.map(
             |_, _| (),
             |_, &l| WeightFunction::new_scaled(arr1(&[l]), WeightFunctionShape::Delta),