# Coverage for pySDC/implementations/problem_classes/AllenCahn_Temp_MPIFFT.py: 81%

## 123 statements

, created at 2024-09-19 09:13 +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.

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 ----------

68 """

70 dtype_u = mesh

71 dtype_f = imex_mesh

73 def __init__(

74 self,

75 nvars=None,

76 eps=0.04,

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',

112 'spectral',

113 'TM',

114 'D',

115 'dw',

116 'L',

117 'init_type',

118 'comm',

119 localVars=locals(),

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):

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