feos_dft/
geometry.rs

1use feos_core::ReferenceSystem;
2use ndarray::{Array1, Array2};
3use quantity::{Angle, Length, Quantity};
4use std::f64::consts::{FRAC_PI_3, PI};
5
6/// Grids with up to three dimensions.
7#[derive(Clone)]
8pub enum Grid {
9    Cartesian1(Axis),
10    Cartesian2(Axis, Axis),
11    Periodical2(Axis, Axis, Angle),
12    Cartesian3(Axis, Axis, Axis),
13    Periodical3(Axis, Axis, Axis, [Angle; 3]),
14    Spherical(Axis),
15    Polar(Axis),
16    Cylindrical { r: Axis, z: Axis },
17}
18
19impl Grid {
20    pub fn new_1d(axis: Axis) -> Self {
21        match axis.geometry {
22            Geometry::Cartesian => Self::Cartesian1(axis),
23            Geometry::Cylindrical => Self::Polar(axis),
24            Geometry::Spherical => Self::Spherical(axis),
25        }
26    }
27
28    pub fn axes(&self) -> Vec<&Axis> {
29        match self {
30            Self::Cartesian1(x) => vec![x],
31            Self::Cartesian2(x, y) | Self::Periodical2(x, y, _) => vec![x, y],
32            Self::Cartesian3(x, y, z) | Self::Periodical3(x, y, z, _) => vec![x, y, z],
33            Self::Spherical(r) | Self::Polar(r) => vec![r],
34            Self::Cylindrical { r, z } => vec![r, z],
35        }
36    }
37
38    pub fn axes_mut(&mut self) -> Vec<&mut Axis> {
39        match self {
40            Self::Cartesian1(x) => vec![x],
41            Self::Cartesian2(x, y) | Self::Periodical2(x, y, _) => vec![x, y],
42            Self::Cartesian3(x, y, z) | Self::Periodical3(x, y, z, _) => vec![x, y, z],
43            Self::Spherical(r) | Self::Polar(r) => vec![r],
44            Self::Cylindrical { r, z } => vec![r, z],
45        }
46    }
47
48    pub fn grids(&self) -> Vec<&Array1<f64>> {
49        self.axes().iter().map(|ax| &ax.grid).collect()
50    }
51
52    pub(crate) fn integration_weights(&self) -> (Vec<&Array1<f64>>, f64) {
53        (
54            self.axes()
55                .iter()
56                .map(|ax| &ax.integration_weights)
57                .collect(),
58            self.functional_determinant(),
59        )
60    }
61
62    pub(crate) fn functional_determinant(&self) -> f64 {
63        match &self {
64            Self::Periodical2(_, _, alpha) => alpha.sin(),
65            Self::Periodical3(_, _, _, [alpha, beta, gamma]) => {
66                let xi = (alpha.cos() - gamma.cos() * beta.cos()) / gamma.sin();
67                gamma.sin() * (1.0 - beta.cos().powi(2) - xi * xi).sqrt()
68            }
69            _ => 1.0,
70        }
71    }
72}
73
74/// Geometries of individual axes.
75#[derive(Copy, Clone, PartialEq)]
76pub enum Geometry {
77    Cartesian,
78    Cylindrical,
79    Spherical,
80}
81
82impl Geometry {
83    /// Return the number of spatial dimensions for this geometry.
84    pub fn dimension(&self) -> i32 {
85        match self {
86            Self::Cartesian => 1,
87            Self::Cylindrical => 2,
88            Self::Spherical => 3,
89        }
90    }
91}
92
93/// An individual discretized axis.
94#[derive(Clone)]
95pub struct Axis {
96    pub geometry: Geometry,
97    pub grid: Array1<f64>,
98    pub edges: Array1<f64>,
99    integration_weights: Array1<f64>,
100    potential_offset: f64,
101}
102
103impl Axis {
104    /// Create a new (equidistant) cartesian axis.
105    ///
106    /// The potential_offset is required to make sure that particles
107    /// can not interact through walls.
108    pub fn new_cartesian(points: usize, length: Length, potential_offset: Option<f64>) -> Self {
109        let potential_offset = potential_offset.unwrap_or(0.0);
110        let l = length.to_reduced() + potential_offset;
111        let cell_size = l / points as f64;
112        let grid = Array1::linspace(0.5 * cell_size, l - 0.5 * cell_size, points);
113        let edges = Array1::linspace(0.0, l, points + 1);
114        let integration_weights = Array1::from_elem(points, cell_size);
115        Self {
116            geometry: Geometry::Cartesian,
117            grid,
118            edges,
119            integration_weights,
120            potential_offset,
121        }
122    }
123
124    /// Create a new (equidistant) spherical axis.
125    pub fn new_spherical(points: usize, length: Length) -> Self {
126        let l = length.to_reduced();
127        let cell_size = l / points as f64;
128        let grid = Array1::linspace(0.5 * cell_size, l - 0.5 * cell_size, points);
129        let edges = Array1::linspace(0.0, l, points + 1);
130        let integration_weights = Array1::from_shape_fn(points, |k| {
131            4.0 * FRAC_PI_3 * cell_size.powi(3) * (3 * k * k + 3 * k + 1) as f64
132        });
133        Self {
134            geometry: Geometry::Spherical,
135            grid,
136            edges,
137            integration_weights,
138            potential_offset: 0.0,
139        }
140    }
141
142    /// Create a new logarithmically scaled cylindrical axis.
143    pub fn new_polar(points: usize, length: Length) -> Self {
144        let l = length.to_reduced();
145
146        let mut alpha = 0.002_f64;
147        for _ in 0..20 {
148            alpha = -(1.0 - (-alpha).exp()).ln() / (points - 1) as f64;
149        }
150        let x0 = 0.5 * ((-alpha * points as f64).exp() + (-alpha * (points - 1) as f64).exp());
151        let grid = (0..points)
152            .map(|i| l * x0 * (alpha * i as f64).exp())
153            .collect();
154        let edges = (0..=points)
155            .map(|i| {
156                if i == 0 {
157                    0.0
158                } else {
159                    l * (-alpha * (points - i) as f64).exp()
160                }
161            })
162            .collect();
163
164        let k0 = (2.0 * alpha).exp() * (2.0 * alpha.exp() + (2.0 * alpha).exp() - 1.0)
165            / ((1.0 + alpha.exp()).powi(2) * ((2.0 * alpha).exp() - 1.0));
166        let integration_weights = (0..points)
167            .map(|i| {
168                (match i {
169                    0 => k0 * (2.0 * alpha).exp(),
170                    1 => ((2.0 * alpha).exp() - k0) * (2.0 * alpha).exp(),
171                    _ => (2.0 * alpha * i as f64).exp() * ((2.0 * alpha).exp() - 1.0),
172                }) * ((-2.0 * alpha * points as f64).exp() * PI * l * l)
173            })
174            .collect();
175
176        Self {
177            geometry: Geometry::Cylindrical,
178            grid,
179            edges,
180            integration_weights,
181            potential_offset: 0.0,
182        }
183    }
184
185    /// Returns the total length of the axis.
186    ///
187    /// This includes the `potential_offset` and used e.g.
188    /// to determine the correct frequency vector in FFT.
189    pub fn length(&self) -> f64 {
190        self.edges[self.grid.len()] - self.edges[0]
191    }
192
193    /// Returns the volume of the axis.
194    ///
195    /// Depending on the geometry, the result is in m, m² or m³.
196    /// The `potential_offset` is not included in the volume, as
197    /// it is mainly used to calculate excess properties.
198    pub fn volume(&self) -> f64 {
199        let length = self.edges[self.grid.len()] - self.potential_offset - self.edges[0];
200        (match self.geometry {
201            Geometry::Cartesian => 1.0,
202            Geometry::Cylindrical => 4.0 * PI,
203            Geometry::Spherical => 4.0 * FRAC_PI_3,
204        }) * length.powi(self.geometry.dimension())
205    }
206
207    /// Interpolate a function on the given axis.
208    pub fn interpolate<U>(
209        &self,
210        x: f64,
211        y: &Quantity<Array2<f64>, U>,
212        i: usize,
213    ) -> Quantity<f64, U> {
214        let n = self.grid.len();
215        y.get((
216            i,
217            if x >= self.edges[n] {
218                n - 1
219            } else {
220                match self.geometry {
221                    Geometry::Cartesian | Geometry::Spherical => (x / self.edges[1]) as usize,
222                    Geometry::Cylindrical => {
223                        if x < self.edges[1] {
224                            0
225                        } else {
226                            (n as f64
227                                - (n - 1) as f64 * (x / self.edges[n]).ln()
228                                    / (self.edges[1] / self.edges[n]).ln())
229                                as usize
230                        }
231                    }
232                }
233            },
234        ))
235    }
236}
feos_dft/geometry.rs

feos_dft/
geometry.rs
