use super::pore3d::{calculate_distance2, evaluate_lj_potential};
use crate::profile::{CUTOFF_RADIUS, MAX_POTENTIAL};
use crate::Geometry;
use feos_core::ReferenceSystem;
use gauss_quad::GaussLegendre;
use ndarray::{Array1, Array2, Zip};
use quantity::Length;
use std::f64::consts::PI;
#[expect(clippy::too_many_arguments)]
pub fn calculate_fea_potential(
grid: &Array1<f64>,
mi: f64,
coordinates: &Length<Array2<f64>>,
sigma_sf: Array1<f64>,
epsilon_k_sf: Array1<f64>,
pore_center: &[f64; 3],
system_size: &[Length; 3],
n_grid: &[usize; 2],
temperature: f64,
geometry: Geometry,
cutoff_radius: Option<f64>,
) -> Array1<f64> {
let mut potential: Array1<f64> = Array1::zeros(grid.len());
let cutoff_radius2 = cutoff_radius.unwrap_or(CUTOFF_RADIUS).powi(2);
let coordinates = Array2::from_shape_fn(coordinates.raw_dim(), |(i, j)| {
(coordinates.get((i, j))).to_reduced()
});
let system_size = [
system_size[0].to_reduced(),
system_size[1].to_reduced(),
system_size[2].to_reduced(),
];
let (nodes1, weights1) = match geometry {
Geometry::Cartesian => {
let nodes = Array1::linspace(
0.5 * system_size[1] / n_grid[0] as f64,
system_size[1] - 0.5 * system_size[1] / n_grid[0] as f64,
n_grid[0],
);
let weights = Array1::from_elem(n_grid[0], system_size[1] / n_grid[0] as f64);
(nodes, weights)
}
Geometry::Spherical | Geometry::Cylindrical => {
let (unscaled_nodes, unscaled_weights) =
GaussLegendre::new(n_grid[0]).unwrap().into_iter().unzip();
let nodes = PI + Array1::from_vec(unscaled_nodes) * PI;
let weights = Array1::from_vec(unscaled_weights) * PI;
(nodes, weights)
}
};
let (nodes2, weights2) = match geometry {
Geometry::Cylindrical | Geometry::Cartesian => {
let nodes = Array1::linspace(
0.5 * system_size[2] / n_grid[1] as f64,
system_size[2] - 0.5 * system_size[2] / n_grid[1] as f64,
n_grid[1],
);
let weights = Array1::from_elem(n_grid[1], system_size[2] / n_grid[1] as f64);
(nodes, weights)
}
Geometry::Spherical => {
let (unscaled_nodes, unscaled_weights) =
GaussLegendre::new(n_grid[1]).unwrap().into_iter().unzip();
let nodes = PI / 2.0 + Array1::from_vec(unscaled_nodes) * PI / 2.0;
let weights = Array1::from_vec(unscaled_weights) * PI / 2.0
* Array1::from_shape_fn(n_grid[1], |i| nodes[i].sin());
(nodes, weights)
}
};
let weights = Array2::from_shape_fn((n_grid[0], n_grid[1]), |(i, j)| weights1[i] * weights2[j]);
let weights_sum = weights.sum();
Zip::indexed(&mut potential).par_for_each(|i0, f| {
let mut potential_2d: Array2<f64> = Array2::zeros((n_grid[0], n_grid[1]));
for (i1, &n1) in nodes1.iter().enumerate() {
for (i2, &n2) in nodes2.iter().enumerate() {
let point = match geometry {
Geometry::Cartesian => [grid[i0], n1, n2],
Geometry::Cylindrical => [
pore_center[0] + grid[i0] * n1.cos(),
pore_center[1] + grid[i0] * n1.sin(),
n2,
],
Geometry::Spherical => [
pore_center[0] + grid[i0] * n2.sin() * n1.cos(),
pore_center[1] + grid[i0] * n2.sin() * n1.sin(),
pore_center[2] + grid[i0] * n2.cos(),
],
};
let distance2 = calculate_distance2(point, &coordinates, system_size);
let potential_sum: f64 = (0..sigma_sf.len())
.map(|alpha| {
mi * evaluate_lj_potential(
distance2[alpha],
sigma_sf[alpha],
epsilon_k_sf[alpha],
cutoff_radius2,
) / temperature
})
.sum();
potential_2d[[i1, i2]] = (-potential_sum.min(MAX_POTENTIAL)).exp();
}
}
*f = (potential_2d * &weights).sum();
});
-temperature * potential.map(|p| (p / weights_sum).ln())
}