1use crate::float::Float;
2use ndarray::{Array2, Array3};
3use serde::{Deserialize, Serialize};
4
5#[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 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 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 pub fn total_points(&self) -> usize {
44 self.nx * self.ny * self.nz
45 }
46
47 pub fn spectral_shape(&self) -> (usize, usize, usize) {
49 (self.nx, self.ny, self.nz / 2 + 1)
50 }
51
52 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 pub fn dv(&self) -> f64 {
66 self.dx() * self.dy() * self.dz()
67 }
68}
69
70#[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#[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#[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 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#[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); }
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); 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}