-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmechanisms.py
194 lines (159 loc) · 7.01 KB
/
mechanisms.py
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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
"""
Mechanisms to be implemented in a `Model` instance.
"""
import cupy as cp
import cupyx.scipy.fft as cufft
import scipy.fft
scipy.fft.set_global_backend(cufft)
fft_lib = scipy.fft
class Advection:
"""Advection implemented without dealiasing.
"""
solution_mode = 'approximate'
def __init__(self, model):
self.model = model
def __call__(self):
self.model.uk = -self.model.iky * self.model.psik
self.model.vk = self.model.ikx * self.model.psik
self.model.u = fft_lib.irfft2(self.model.uk)
self.model.v = fft_lib.irfft2(self.model.vk)
self.model.z_x = fft_lib.irfft2(self.model.ikx * self.model.zk)
self.model.z_y = fft_lib.irfft2(self.model.iky * self.model.zk)
self.model.nlk = (fft_lib.rfft2(self.model.u * self.model.z_x)
+ fft_lib.rfft2(self.model.v * self.model.z_y))
self.model.rhs = -self.model.nlk
class AdvectionWithTruncation:
"""Advection implemented with 2/3 rule for approximate dealiasing.
"""
solution_mode = 'approximate'
def __init__(self, model):
self.model = model
self.mask = (model.kx < model.n_x / 3.) & (
cp.abs(model.ky) < model.n_x / 3.)
def __call__(self):
self.model.uk = -self.model.iky * self.model.psik
self.model.vk = self.model.ikx * self.model.psik
self.model.u = fft_lib.irfft2(self.model.uk * self.mask)
self.model.v = fft_lib.irfft2(self.model.vk * self.mask)
self.model.z_x = fft_lib.irfft2(
self.model.ikx * self.model.zk * self.mask)
self.model.z_y = fft_lib.irfft2(
self.model.iky * self.model.zk * self.mask)
self.model.nlk = (fft_lib.rfft2(self.model.u * self.model.z_x)
+ fft_lib.rfft2(self.model.v * self.model.z_y))
self.model.rhs = -self.model.nlk
class DealiasedAdvection:
"""Advection dealiased using the 3/2 rule.
"""
solution_mode = 'approximate'
def __init__(self, model):
self.model = model
self.pad = 3. / 2.
self.m_x = int(self.pad * self.model.n_x)
self.m_k = int(self.pad * self.model.n_kx)
self.padder = cp.ones(self.m_x, dtype=bool)
self.padder[int(self.model.n_x / 2):
int(self.model.n_x * (self.pad - 0.5)):] = False
def __call__(self):
self.nlk = cp.zeros(self.model.zk.shape,
dtype=self.model.complex_dtype)
self.model.uk = -self.model.iky * self.model.psik
self.model.vk = self.model.ikx * self.model.psik
# Create padded arrays
self.uk_padded = cp.zeros((self.m_x, self.m_k),
dtype=self.model.complex_dtype)
self.vk_padded = cp.zeros((self.m_x, self.m_k),
dtype=self.model.complex_dtype)
self.z_xk_padded = cp.zeros((self.m_x, self.m_k),
dtype=self.model.complex_dtype)
self.z_yk_padded = cp.zeros((self.m_x, self.m_k),
dtype=self.model.complex_dtype)
# Enter known coefficients, leaving padded entries equal to zero
self.uk_padded[self.padder, :self.model.n_kx] = self.model.uk[:, :]
self.vk_padded[self.padder, :self.model.n_kx] = self.model.vk[:, :]
self.z_xk_padded[self.padder, :self.model.n_kx] = (
self.model.ikx * self.model.zk)[:, :]
self.z_yk_padded[self.padder, :self.model.n_kx] = (
self.model.iky * self.model.zk)[:, :]
# Inverse transform padded arrays
self.u_padded = fft_lib.irfft2(self.uk_padded)
self.v_padded = fft_lib.irfft2(self.vk_padded)
self.z_x_padded = fft_lib.irfft2(self.z_xk_padded)
self.z_y_padded = fft_lib.irfft2(self.z_yk_padded)
# Calculate Jacobian term
self.nlk[:, :] = fft_lib.rfft2(
(self.u_padded * self.z_x_padded + self.v_padded * self.z_y_padded
))[self.padder, :self.model.n_kx] * self.pad ** 2
self.model.rhs = -self.nlk
class Diffusion:
"""Diffusion with `order` to be specified. Order refers to the power of the
Laplacian. `order=1.` gives standard Newtonian viscosity; `order>1.` gives
hyperviscosity; `order=0.` gives linear drag; `order=-1.` gives
large-scale/hyper friction, etc.
"""
solution_mode = 'exact'
def __init__(self, model, order=1., coefficient=None):
self.model = model
self.order = order
self.coefficient = coefficient
if self.order >= 0.:
self.multiplier = cp.exp(
-self.coefficient * self.model.timestepper.dt
* self.model.wv2 ** self.order)
else:
self.multiplier = cp.exp(
-self.coefficient * self.model.timestepper.dt
* self.model.wv2i ** (-self.order))
def __call__(self):
self.model.zk *= self.multiplier
class Beta:
"""Beta plane.
"""
solution_mode = 'exact'
def __init__(self, model, beta):
self.model = model
self.beta = beta
self.solution_multiplier = cp.exp(self.model.timestepper.dt * self.beta
* self.model.ikx * self.model.wv2i)
def __call__(self):
self.model.zk *= self.solution_multiplier
class StochasticRingForcing:
"""White-in-time stochastic forcing concentrated on a band of wavenumbers.
Energy is input at a mean rate which is uniform across forced wavenumbers.
"""
solution_mode = 'discrete'
def __init__(self, model, k_f, dk_f, energy_input_rate, seed=1):
self.model = model
self.k_f = k_f
self.dk_f = dk_f
self.energy_input_rate = energy_input_rate
self.rng = cp.random.default_rng(seed=seed)
self.band_filter = ((self.model.wv >= self.k_f - self.dk_f)
& (self.model.wv <= self.k_f + self.dk_f))
self.fk_vars = (self.energy_input_rate * self.band_filter
* self.model.n_x ** 4 * self.model.wv * 0.5
/ cp.sum(self.band_filter * self.model.wv2i ** 0.5))
def __call__(self):
self.fk = cp.reshape(self.rng.standard_normal(
size=self.model.wv.size, dtype=self.model.real_dtype)
+ 1j * self.rng.standard_normal(
size=self.model.wv.size,
dtype=self.model.real_dtype),
self.model.wv.shape) * self.fk_vars ** 0.5
self.model.zk += self.fk * self.model.timestepper.dt ** 0.5
class SpectralFilter:
"""Exponential spectral filter for applying highly scale-selective but
non-physical dissipation at small scales.
"""
solution_mode = 'exact'
def __init__(self, model):
filterfac = 23.6
cphi = 0.65 * cp.pi
wvx = cp.sqrt(
(model.kx * model.dx) ** 2. + (model.ky * model.dx) ** 2.)
exp_filter = cp.exp(-filterfac * (wvx - cphi) ** 4.)
exp_filter[wvx <= cphi] = 1.
self.exp_filter = exp_filter
self.model = model
def __call__(self):
self.model.zk *= self.exp_filter