Coverage for pySDC/implementations/problem_classes/AllenCahn_Temp_MPIFFT.py: 81%
123 statements
« prev ^ index » next coverage.py v7.6.7, created at 2024-11-16 14:51 +0000
1import numpy as np
2from mpi4py_fft import PFFT
4from pySDC.core.errors import ProblemError
5from pySDC.core.problem import Problem
6from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
8from mpi4py_fft import newDistArray
11class allencahn_temp_imex(Problem):
12 r"""
13 This class implements the :math:`N`-dimensional Allen-Cahn equation with periodic boundary conditions :math:`u \in [0, 1]^2`
15 .. math::
16 \frac{\partial u}{\partial t} = D \Delta u - \frac{2}{\varepsilon^2} u (1 - u) (1 - 2u)
17 - 6 d_w \frac{u - T_M}{T_M}u (1 - u)
19 on a spatial domain :math:`[-\frac{L}{2}, \frac{L}{2}]^2`, with driving force :math:`d_w`, and :math:`N=2,3`. :math:`D` and
20 :math:`T_M` are fixed parameters. Different initial conditions can be used, for example, circles of the form
22 .. math::
23 u({\bf x}, 0) = \tanh\left(\frac{r - \sqrt{(x_i-0.5)^2 + (y_j-0.5)^2}}{\sqrt{2}\varepsilon}\right),
25 for :math:`i, j=0,..,N-1`, where :math:`N` is the number of spatial grid points. For time-stepping, the problem is treated
26 *semi-implicitly*, i.e., the nonlinear system is solved by Fast-Fourier Tranform (FFT) and the linear parts in the right-hand
27 side will be treated explicitly using ``mpi4py-fft`` [1]_ to solve them.
29 Attributes
30 ----------
31 nvars : List of int tuples, optional
32 Number of unknowns in the problem, e.g. ``nvars=[(128, 128), (64, 64)]``.
33 eps : float, optional
34 Scaling parameter :math:`\varepsilon`.
35 radius : float, optional
36 Radius of the circles.
37 spectral : bool, optional
38 Indicates if spectral initial condition is used.
39 TM : float, optional
40 Problem parameter :math:`T_M`.
41 D : float, optional
42 Problem parameter :math:`D`.
43 dw : float, optional
44 Driving force.
45 L : float, optional
46 Denotes the period of the function to be approximated for the Fourier transform.
47 init_type : str, optional
48 Initialises type of initial state.
49 comm : bool, optional
50 Communicator.
52 Attributes
53 ----------
54 fft : fft object
55 Object for FFT.
56 X : np.ogrid
57 Grid coordinates in real space.
58 K2 : np.ndarray
59 Laplace operator in spectral space.
60 dx : float
61 Mesh width in x direction.
62 dy : float
63 Mesh width in y direction.
65 References
66 ----------
67 .. [1] https://mpi4py-fft.readthedocs.io/en/latest/
68 """
70 dtype_u = mesh
71 dtype_f = imex_mesh
73 def __init__(
74 self,
75 nvars=None,
76 eps=0.04,
77 radius=0.25,
78 spectral=None,
79 TM=1.0,
80 D=10.0,
81 dw=0.0,
82 L=1.0,
83 init_type='circle',
84 comm=None,
85 ):
86 """Initialization routine"""
88 if nvars is None:
89 nvars = [(128, 128)]
91 if not (isinstance(nvars, tuple) and len(nvars) > 1):
92 raise ProblemError('Need at least two dimensions')
94 # creating FFT structure
95 ndim = len(nvars)
96 axes = tuple(range(ndim))
97 self.fft = PFFT(comm, list(nvars), axes=axes, dtype=np.float64, collapse=True)
99 # get test data to figure out type and dimensions
100 tmp_u = newDistArray(self.fft, spectral)
102 # add two components to contain field and temperature
103 self.ncomp = 2
104 sizes = tmp_u.shape + (self.ncomp,)
106 # invoke super init, passing the communicator and the local dimensions as init
107 super().__init__(init=(sizes, comm, tmp_u.dtype))
108 self._makeAttributeAndRegister(
109 'nvars',
110 'eps',
111 'radius',
112 'spectral',
113 'TM',
114 'D',
115 'dw',
116 'L',
117 'init_type',
118 'comm',
119 localVars=locals(),
120 readOnly=True,
121 )
123 L = np.array([self.L] * ndim, dtype=float)
125 # get local mesh
126 X = list(np.ogrid[self.fft.local_slice(False)])
127 N = self.fft.global_shape()
128 for i in range(len(N)):
129 X[i] = X[i] * L[i] / N[i]
130 self.X = [np.broadcast_to(x, self.fft.shape(False)) for x in X]
132 # get local wavenumbers and Laplace operator
133 s = self.fft.local_slice()
134 N = self.fft.global_shape()
135 k = [np.fft.fftfreq(n, 1.0 / n).astype(int) for n in N[:-1]]
136 k.append(np.fft.rfftfreq(N[-1], 1.0 / N[-1]).astype(int))
137 K = [ki[si] for ki, si in zip(k, s)]
138 Ks = list(np.meshgrid(*K, indexing='ij', sparse=True))
139 Lp = 2 * np.pi / L
140 for i in range(ndim):
141 Ks[i] = (Ks[i] * Lp[i]).astype(float)
142 K = [np.broadcast_to(k, self.fft.shape(True)) for k in Ks]
143 K = np.array(K).astype(float)
144 self.K2 = np.sum(K * K, 0, dtype=float)
146 # Need this for diagnostics
147 self.dx = self.L / nvars[0]
148 self.dy = self.L / nvars[1]
150 def eval_f(self, u, t):
151 """
152 Routine to evaluate the right-hand side of the problem.
154 Parameters
155 ----------
156 u : dtype_u
157 Current values of the numerical solution.
158 t : float
159 Current time of the numerical solution is computed.
161 Returns
162 -------
163 f : dtype_f
164 The right-hand side of the problem.
165 """
167 f = self.dtype_f(self.init)
169 if self.spectral:
170 f.impl[..., 0] = -self.K2 * u[..., 0]
172 if self.eps > 0:
173 tmp_u = newDistArray(self.fft, False)
174 tmp_T = newDistArray(self.fft, False)
175 tmp_u = self.fft.backward(u[..., 0], tmp_u)
176 tmp_T = self.fft.backward(u[..., 1], tmp_T)
177 tmpf = -2.0 / self.eps**2 * tmp_u * (1.0 - tmp_u) * (1.0 - 2.0 * tmp_u) - 6.0 * self.dw * (
178 tmp_T - self.TM
179 ) / self.TM * tmp_u * (1.0 - tmp_u)
180 f.expl[..., 0] = self.fft.forward(tmpf)
182 f.impl[..., 1] = -self.D * self.K2 * u[..., 1]
183 f.expl[..., 1] = f.impl[..., 0] + f.expl[..., 0]
185 else:
186 u_hat = self.fft.forward(u[..., 0])
187 lap_u_hat = -self.K2 * u_hat
188 f.impl[..., 0] = self.fft.backward(lap_u_hat, f.impl[..., 0])
190 if self.eps > 0:
191 f.expl[..., 0] = -2.0 / self.eps**2 * u[..., 0] * (1.0 - u[..., 0]) * (1.0 - 2.0 * u[..., 0])
192 f.expl[..., 0] -= 6.0 * self.dw * (u[..., 1] - self.TM) / self.TM * u[..., 0] * (1.0 - u[..., 0])
194 u_hat = self.fft.forward(u[..., 1])
195 lap_u_hat = -self.D * self.K2 * u_hat
196 f.impl[..., 1] = self.fft.backward(lap_u_hat, f.impl[..., 1])
197 f.expl[..., 1] = f.impl[..., 0] + f.expl[..., 0]
199 return f
201 def solve_system(self, rhs, factor, u0, t):
202 """
203 Simple FFT solver for the diffusion part.
205 Parameters
206 ----------
207 rhs : dtype_f
208 Right-hand side for the linear system.
209 factor : float
210 Abbrev. for the node-to-node stepsize (or any other factor required).
211 u0 : dtype_u
212 Initial guess for the iterative solver (not used here so far).
213 t : float
214 Current time (e.g. for time-dependent BCs).
216 Returns
217 -------
218 me : dtype_u
219 The solution as mesh.
220 """
222 if self.spectral:
223 me = self.dtype_u(self.init)
224 me[..., 0] = rhs[..., 0] / (1.0 + factor * self.K2)
225 me[..., 1] = rhs[..., 1] / (1.0 + factor * self.D * self.K2)
227 else:
228 me = self.dtype_u(self.init)
229 rhs_hat = self.fft.forward(rhs[..., 0])
230 rhs_hat /= 1.0 + factor * self.K2
231 me[..., 0] = self.fft.backward(rhs_hat)
232 rhs_hat = self.fft.forward(rhs[..., 1])
233 rhs_hat /= 1.0 + factor * self.D * self.K2
234 me[..., 1] = self.fft.backward(rhs_hat)
236 return me
238 def u_exact(self, t):
239 r"""
240 Routine to compute the exact solution at time :math:`t`.
242 Parameters
243 ----------
244 t : float
245 Time of the exact solution.
247 Returns
248 -------
249 me : dtype_u
250 The exact solution.
251 """
253 def circle():
254 tmp_me = newDistArray(self.fft, self.spectral)
255 r2 = (self.X[0] - 0.5) ** 2 + (self.X[1] - 0.5) ** 2
256 if self.spectral:
257 tmp = 0.5 * (1.0 + np.tanh((self.radius - np.sqrt(r2)) / (np.sqrt(2) * self.eps)))
258 tmp_me[:] = self.fft.forward(tmp)
259 else:
260 tmp_me[:] = 0.5 * (1.0 + np.tanh((self.radius - np.sqrt(r2)) / (np.sqrt(2) * self.eps)))
261 return tmp_me
263 def circle_rand():
264 tmp_me = newDistArray(self.fft, self.spectral)
265 ndim = len(tmp_me.shape)
266 L = int(self.L)
267 # get random radii for circles/spheres
268 np.random.seed(1)
269 lbound = 3.0 * self.eps
270 ubound = 0.5 - self.eps
271 rand_radii = (ubound - lbound) * np.random.random_sample(size=tuple([L] * ndim)) + lbound
272 # distribute circles/spheres
273 tmp = newDistArray(self.fft, False)
274 if ndim == 2:
275 for i in range(0, L):
276 for j in range(0, L):
277 # build radius
278 r2 = (self.X[0] + i - L + 0.5) ** 2 + (self.X[1] + j - L + 0.5) ** 2
279 # add this blob, shifted by 1 to avoid issues with adding up negative contributions
280 tmp += np.tanh((rand_radii[i, j] - np.sqrt(r2)) / (np.sqrt(2) * self.eps)) + 1
281 # normalize to [0,1]
282 tmp *= 0.5
283 assert np.all(tmp <= 1.0)
284 if self.spectral:
285 tmp_me[:] = self.fft.forward(tmp)
286 else:
287 tmp_me[:] = tmp[:]
288 return tmp_me
290 def sine():
291 tmp_me = newDistArray(self.fft, self.spectral)
292 if self.spectral:
293 tmp = 0.5 * (2.0 + 0.0 * np.sin(2 * np.pi * self.X[0]) * np.sin(2 * np.pi * self.X[1]))
294 tmp_me[:] = self.fft.forward(tmp)
295 else:
296 tmp_me[:] = 0.5 * (2.0 + 0.0 * np.sin(2 * np.pi * self.X[0]) * np.sin(2 * np.pi * self.X[1]))
297 return tmp_me
299 assert t == 0, 'ERROR: u_exact only valid for t=0'
300 me = self.dtype_u(self.init, val=0.0)
301 if self.init_type == 'circle':
302 me[..., 0] = circle()
303 me[..., 1] = sine()
304 elif self.init_type == 'circle_rand':
305 me[..., 0] = circle_rand()
306 me[..., 1] = sine()
307 else:
308 raise NotImplementedError('type of initial value not implemented, got %s' % self.init_type)
310 return me
313# class allencahn_temp_imex_timeforcing(allencahn_temp_imex):
314# """
315# Example implementing Allen-Cahn equation in 2-3D using mpi4py-fft for solving linear parts, IMEX time-stepping,
316# time-dependent forcing
317# """
318#
319# def eval_f(self, u, t):
320# """
321# Routine to evaluate the RHS
322#
323# Args:
324# u (dtype_u): current values
325# t (float): current time
326#
327# Returns:
328# dtype_f: the RHS
329# """
330#
331# f = self.dtype_f(self.init)
332#
333# if self.spectral:
334#
335# f.impl = -self.K2 * u
336#
337# tmp = newDistArray(self.fft, False)
338# tmp[:] = self.fft.backward(u, tmp)
339#
340# if self.eps > 0:
341# tmpf = -2.0 / self.eps ** 2 * tmp * (1.0 - tmp) * (1.0 - 2.0 * tmp)
342# else:
343# tmpf = self.dtype_f(self.init, val=0.0)
344#
345# # build sum over RHS without driving force
346# Rt_local = float(np.sum(self.fft.backward(f.impl) + tmpf))
347# if self.comm is not None:
348# Rt_global = self.comm.allreduce(sendobj=Rt_local, op=MPI.SUM)
349# else:
350# Rt_global = Rt_local
351#
352# # build sum over driving force term
353# Ht_local = float(np.sum(6.0 * tmp * (1.0 - tmp)))
354# if self.comm is not None:
355# Ht_global = self.comm.allreduce(sendobj=Ht_local, op=MPI.SUM)
356# else:
357# Ht_global = Rt_local
358#
359# # add/substract time-dependent driving force
360# if Ht_global != 0.0:
361# dw = Rt_global / Ht_global
362# else:
363# dw = 0.0
364#
365# tmpf -= 6.0 * dw * tmp * (1.0 - tmp)
366# f.expl[:] = self.fft.forward(tmpf)
367#
368# else:
369#
370# u_hat = self.fft.forward(u)
371# lap_u_hat = -self.K2 * u_hat
372# f.impl[:] = self.fft.backward(lap_u_hat, f.impl)
373#
374# if self.eps > 0:
375# f.expl = -2.0 / self.eps ** 2 * u * (1.0 - u) * (1.0 - 2.0 * u)
376#
377# # build sum over RHS without driving force
378# Rt_local = float(np.sum(f.impl + f.expl))
379# if self.comm is not None:
380# Rt_global = self.comm.allreduce(sendobj=Rt_local, op=MPI.SUM)
381# else:
382# Rt_global = Rt_local
383#
384# # build sum over driving force term
385# Ht_local = float(np.sum(6.0 * u * (1.0 - u)))
386# if self.comm is not None:
387# Ht_global = self.comm.allreduce(sendobj=Ht_local, op=MPI.SUM)
388# else:
389# Ht_global = Rt_local
390#
391# # add/substract time-dependent driving force
392# if Ht_global != 0.0:
393# dw = Rt_global / Ht_global
394# else:
395# dw = 0.0
396#
397# f.expl -= 6.0 * dw * u * (1.0 - u)
398#
399# return f