Skip to main content

vonkarman_core/
field.rs

1use crate::float::Float;
2use ndarray::{Array2, Array3};
3use serde::{Deserialize, Serialize};
4
5/// Grid metadata for a 3D periodic domain.
6#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
7pub struct GridSpec {
8    pub nx: usize,
9    pub ny: usize,
10    pub nz: usize,
11    pub lx: f64,
12    pub ly: f64,
13    pub lz: f64,
14}
15
16impl GridSpec {
17    /// Create a cubic grid with side length `l` and `n` points per axis.
18    pub fn cubic(n: usize, l: f64) -> Self {
19        Self {
20            nx: n,
21            ny: n,
22            nz: n,
23            lx: l,
24            ly: l,
25            lz: l,
26        }
27    }
28
29    /// Grid spacing (assumes uniform, uses x-direction).
30    pub fn dx(&self) -> f64 {
31        self.lx / self.nx as f64
32    }
33
34    pub fn dy(&self) -> f64 {
35        self.ly / self.ny as f64
36    }
37
38    pub fn dz(&self) -> f64 {
39        self.lz / self.nz as f64
40    }
41
42    /// Total number of physical-space grid points.
43    pub fn total_points(&self) -> usize {
44        self.nx * self.ny * self.nz
45    }
46
47    /// Spectral-space shape after R2C transform: (nx, ny, nz/2+1).
48    pub fn spectral_shape(&self) -> (usize, usize, usize) {
49        (self.nx, self.ny, self.nz / 2 + 1)
50    }
51
52    /// 3/2-padded grid for dealiasing.
53    pub fn padded_3half(&self) -> Self {
54        Self {
55            nx: 3 * self.nx / 2,
56            ny: 3 * self.ny / 2,
57            nz: 3 * self.nz / 2,
58            lx: self.lx,
59            ly: self.ly,
60            lz: self.lz,
61        }
62    }
63
64    /// Volume element (uniform grid).
65    pub fn dv(&self) -> f64 {
66        self.dx() * self.dy() * self.dz()
67    }
68}
69
70/// Grid metadata for an axisymmetric (r, z) domain.
71#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
72pub struct AxiGridSpec {
73    pub nr: usize,
74    pub nz: usize,
75    pub r_max: f64,
76    pub lz: f64,
77}
78
79/// A 3D scalar field on a periodic grid.
80#[derive(Debug, Clone)]
81pub struct ScalarField<F: Float> {
82    pub data: Array3<F>,
83    pub grid: GridSpec,
84}
85
86impl<F: Float> ScalarField<F> {
87    pub fn zeros(grid: GridSpec) -> Self {
88        Self {
89            data: Array3::from_elem((grid.nx, grid.ny, grid.nz), F::ZERO),
90            grid,
91        }
92    }
93}
94
95/// A 3D vector field (3 components) on a periodic grid.
96#[derive(Debug, Clone)]
97pub struct VectorField<F: Float> {
98    pub data: [Array3<F>; 3],
99    pub grid: GridSpec,
100}
101
102impl<F: Float> VectorField<F> {
103    pub fn zeros(grid: GridSpec) -> Self {
104        let shape = (grid.nx, grid.ny, grid.nz);
105        Self {
106            data: [
107                Array3::from_elem(shape, F::ZERO),
108                Array3::from_elem(shape, F::ZERO),
109                Array3::from_elem(shape, F::ZERO),
110            ],
111            grid,
112        }
113    }
114
115    /// Convenience accessors.
116    pub fn x(&self) -> &Array3<F> {
117        &self.data[0]
118    }
119    pub fn y(&self) -> &Array3<F> {
120        &self.data[1]
121    }
122    pub fn z(&self) -> &Array3<F> {
123        &self.data[2]
124    }
125    pub fn x_mut(&mut self) -> &mut Array3<F> {
126        &mut self.data[0]
127    }
128    pub fn y_mut(&mut self) -> &mut Array3<F> {
129        &mut self.data[1]
130    }
131    pub fn z_mut(&mut self) -> &mut Array3<F> {
132        &mut self.data[2]
133    }
134}
135
136/// A 2D scalar field on an axisymmetric (r, z) grid.
137#[derive(Debug, Clone)]
138pub struct AxiScalarField<F: Float> {
139    pub data: Array2<F>,
140    pub grid: AxiGridSpec,
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn grid_spec_basic() {
149        let g = GridSpec::cubic(64, 2.0 * std::f64::consts::PI);
150        assert_eq!(g.nx, 64);
151        assert_eq!(g.ny, 64);
152        assert_eq!(g.nz, 64);
153        let dx = g.dx();
154        assert!((dx - 2.0 * std::f64::consts::PI / 64.0).abs() < 1e-14);
155    }
156
157    #[test]
158    fn grid_spec_spectral_shape() {
159        let g = GridSpec::cubic(64, 2.0 * std::f64::consts::PI);
160        let (snx, sny, snz) = g.spectral_shape();
161        assert_eq!(snx, 64);
162        assert_eq!(sny, 64);
163        assert_eq!(snz, 33); // N/2 + 1
164    }
165
166    #[test]
167    fn grid_spec_padded() {
168        let g = GridSpec::cubic(64, 2.0 * std::f64::consts::PI);
169        let pg = g.padded_3half();
170        assert_eq!(pg.nx, 96); // 3*64/2
171        assert_eq!(pg.ny, 96);
172        assert_eq!(pg.nz, 96);
173    }
174
175    #[test]
176    fn scalar_field_zeros() {
177        let g = GridSpec::cubic(8, 1.0);
178        let s = ScalarField::<f64>::zeros(g);
179        assert_eq!(s.data.shape(), &[8, 8, 8]);
180        assert_eq!(s.data[[0, 0, 0]], 0.0);
181    }
182
183    #[test]
184    fn vector_field_zeros() {
185        let g = GridSpec::cubic(8, 1.0);
186        let v = VectorField::<f64>::zeros(g);
187        assert_eq!(v.x().shape(), &[8, 8, 8]);
188        assert_eq!(v.y().shape(), &[8, 8, 8]);
189        assert_eq!(v.z().shape(), &[8, 8, 8]);
190    }
191
192    #[test]
193    fn total_points() {
194        let g = GridSpec::cubic(32, 1.0);
195        assert_eq!(g.total_points(), 32 * 32 * 32);
196    }
197}