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