1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use super::pore3d::{calculate_distance2, evaluate_lj_potential};
use crate::profile::{CUTOFF_RADIUS, MAX_POTENTIAL};
use crate::Geometry;
use feos_core::EosUnit;
use gauss_quad::GaussLegendre;
use ndarray::{Array1, Array2, Zip};
use quantity::si::{SIArray2, SINumber, SIUnit};
use std::f64::consts::PI;
use std::usize;

// Calculate free-energy average potential for given solid structure.
pub fn calculate_fea_potential(
    grid: &Array1<f64>,
    mi: f64,
    coordinates: &SIArray2,
    sigma_sf: Array1<f64>,
    epsilon_k_sf: Array1<f64>,
    pore_center: &[f64; 3],
    system_size: &[SINumber; 3],
    n_grid: &[usize; 2],
    temperature: f64,
    geometry: Geometry,
    cutoff_radius: Option<f64>,
) -> Array1<f64> {
    // allocate external potential
    let mut potential: Array1<f64> = Array1::zeros(grid.len());

    // calculate squared cutoff radius
    let cutoff_radius2 = cutoff_radius.unwrap_or(CUTOFF_RADIUS).powi(2);

    // dimensionless solid coordinates
    let coordinates = Array2::from_shape_fn(coordinates.raw_dim(), |(i, j)| {
        (coordinates.get((i, j)))
            .to_reduced(SIUnit::reference_length())
            .unwrap()
    });

    let system_size = [
        system_size[0]
            .to_reduced(SIUnit::reference_length())
            .unwrap(),
        system_size[1]
            .to_reduced(SIUnit::reference_length())
            .unwrap(),
        system_size[2]
            .to_reduced(SIUnit::reference_length())
            .unwrap(),
    ];

    // Create secondary axis:
    // Cartesian coordinates => y
    // Cylindrical coordinates => phi
    // Spherical coordinates => phi
    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 nodes = PI + Array1::from_vec(GaussLegendre::nodes_and_weights(n_grid[0]).0) * PI;
            let weights = Array1::from_vec(GaussLegendre::nodes_and_weights(n_grid[0]).1) * PI;
            (nodes, weights)
        }
    };

    // Create tertiary axis
    // Cartesian coordinates => z
    // Cylindrical coordinates => z
    // Spherical coordinates => theta
    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 nodes = PI / 2.0
                + Array1::from_vec(GaussLegendre::nodes_and_weights(n_grid[1]).0) * PI / 2.0;
            let weights = Array1::from_vec(GaussLegendre::nodes_and_weights(n_grid[1]).1) * PI
                / 2.0
                * Array1::from_shape_fn(n_grid[1], |i| nodes[i].sin());
            (nodes, weights)
        }
    };

    // calculate weights
    let weights = Array2::from_shape_fn((n_grid[0], n_grid[1]), |(i, j)| weights1[i] * weights2[j]);

    // calculate sum of weights
    let weights_sum = weights.sum();

    // calculate FEA potential
    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())
}