From 4218591245bd7e16e7ff98ff82e39b4370b861ec Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Sat, 16 Nov 2024 15:24:45 +0100 Subject: [PATCH] 3d Gray-Scott (#506) * Added test for generic mpi4py-fft Laplacian class * Expanded tests * 3D initial conditions in Gray-Scott --- .../problem_classes/GrayScott_MPIFFT.py | 13 ++-- .../generic_MPIFFT_Laplacian.py | 30 ++++----- pySDC/projects/GPU/configs/GS_configs.py | 63 ++++++++++++++----- .../test_problems/test_generic_MPIFFT.py | 57 +++++++++++++++++ 4 files changed, 126 insertions(+), 37 deletions(-) create mode 100644 pySDC/tests/test_problems/test_generic_MPIFFT.py diff --git a/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py b/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py index ba14ea762..796a359b1 100644 --- a/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py +++ b/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py @@ -200,7 +200,6 @@ def u_exact(self, t, seed=10700000): Exact solution. """ assert t == 0.0, 'Exact solution only valid as initial condition' - assert self.ndim == 2, 'The initial conditions are 2D for now..' xp = self.xp @@ -222,12 +221,13 @@ def u_exact(self, t, seed=10700000): _v[...] = xp.sqrt(F) * (A + xp.sqrt(A**2 - 4)) / 2 for _ in range(-self.num_blobs): - x0, y0 = rng.random(size=2) * self.L[0] - self.L[0] / 2 - lx, ly = rng.random(size=2) * self.L[0] / self.nvars[0] * 30 + x0 = rng.random(size=self.ndim) * self.L[0] - self.L[0] / 2 + l = rng.random(size=self.ndim) * self.L[0] / self.nvars[0] * 30 - mask_x = xp.logical_and(self.X[0] > x0, self.X[0] < x0 + lx) - mask_y = xp.logical_and(self.X[1] > y0, self.X[1] < y0 + ly) - mask = xp.logical_and(mask_x, mask_y) + masks = [xp.logical_and(self.X[i] > x0[i], self.X[i] < x0[i] + l[i]) for i in range(self.ndim)] + mask = masks[0] + for m in masks[1:]: + mask = xp.logical_and(mask, m) _u[mask] = rng.random() _v[mask] = rng.random() @@ -236,6 +236,7 @@ def u_exact(self, t, seed=10700000): """ Blobs as in https://www.chebfun.org/examples/pde/GrayScott.html """ + assert self.ndim == 2, 'The initial conditions are 2D for now..' inc = self.L[0] / (self.num_blobs + 1) diff --git a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py index e1fc8f681..8f4ee38ef 100644 --- a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py +++ b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py @@ -1,18 +1,17 @@ import numpy as np from mpi4py import MPI -from mpi4py_fft import PFFT +from mpi4py_fft import PFFT, newDistArray from pySDC.core.errors import ProblemError from pySDC.core.problem import Problem, WorkCounter from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh -from mpi4py_fft import newDistArray - class IMEX_Laplacian_MPIFFT(Problem): r""" Generic base class for IMEX problems using a spectral method to solve the Laplacian implicitly and a possible rest explicitly. The FFTs are done with``mpi4py-fft`` [1]_. + Works in two and three dimensions. Parameters ---------- @@ -99,14 +98,24 @@ def __init__( 'nvars', 'spectral', 'L', 'alpha', 'comm', 'x0', 'useGPU', localVars=locals(), readOnly=True ) - # get local mesh + self.getLocalGrid() + self.getLaplacian() + + # Need this for diagnostics + self.dx = self.L[0] / nvars[0] + self.dy = self.L[1] / nvars[1] + + # work counters + self.work_counters['rhs'] = WorkCounter() + + def getLocalGrid(self): X = list(self.xp.ogrid[self.fft.local_slice(False)]) N = self.fft.global_shape() for i in range(len(N)): - X[i] = x0 + (X[i] * L[i] / N[i]) + X[i] = self.x0 + (X[i] * self.L[i] / N[i]) self.X = [self.xp.broadcast_to(x, self.fft.shape(False)) for x in X] - # get local wavenumbers and Laplace operator + def getLaplacian(self): s = self.fft.local_slice() N = self.fft.global_shape() k = [self.xp.fft.fftfreq(n, 1.0 / n).astype(int) for n in N] @@ -117,14 +126,7 @@ def __init__( Ks[i] = (Ks[i] * Lp[i]).astype(float) K = [self.xp.broadcast_to(k, self.fft.shape(True)) for k in Ks] K = self.xp.array(K).astype(float) - self.K2 = self.xp.sum(K * K, 0, dtype=float) # Laplacian in spectral space - - # Need this for diagnostics - self.dx = self.L[0] / nvars[0] - self.dy = self.L[1] / nvars[1] - - # work counters - self.work_counters['rhs'] = WorkCounter() + self.K2 = self.xp.sum(K * K, 0, dtype=float) def eval_f(self, u, t): """ diff --git a/pySDC/projects/GPU/configs/GS_configs.py b/pySDC/projects/GPU/configs/GS_configs.py index 6ee1374d5..f57cdc7ea 100644 --- a/pySDC/projects/GPU/configs/GS_configs.py +++ b/pySDC/projects/GPU/configs/GS_configs.py @@ -27,6 +27,7 @@ class GrayScott(Config): num_frames = 200 sweeper_type = 'IMEX' res_per_blob = 2**7 + ndim = 3 def get_LogToFile(self, ranks=None): import numpy as np @@ -49,16 +50,14 @@ def process_solution(L): 't': L.time + L.dt, 'u': uend[0].get().view(np.ndarray), 'v': uend[1].get().view(np.ndarray), - 'X': L.prob.X[0].get().view(np.ndarray), - 'Y': L.prob.X[1].get().view(np.ndarray), + 'X': L.prob.X.get().view(np.ndarray), } else: return { 't': L.time + L.dt, 'u': uend[0], 'v': uend[1], - 'X': L.prob.X[0], - 'Y': L.prob.X[1], + 'X': L.prob.X, } def logging_condition(L): @@ -75,7 +74,7 @@ def logging_condition(L): LogToFile.logging_condition = logging_condition return LogToFile - def plot(self, P, idx, n_procs_list): # pragma: no cover + def plot(self, P, idx, n_procs_list, projection='xy'): # pragma: no cover import numpy as np from matplotlib import ticker as tkr @@ -99,19 +98,49 @@ def plot(self, P, idx, n_procs_list): # pragma: no cover vmax['u'] = max([vmax['u'], buffer[f'u-{rank}']['u'].real.max()]) for rank in range(n_procs_list[2]): - im = ax.pcolormesh( - buffer[f'u-{rank}']['X'], - buffer[f'u-{rank}']['Y'], - buffer[f'u-{rank}']['v'].real, - vmin=vmin['v'], - vmax=vmax['v'], - cmap='binary', - ) + if len(buffer[f'u-{rank}']['X']) == 2: + ax.set_xlabel('$x$') + ax.set_ylabel('$y$') + im = ax.pcolormesh( + buffer[f'u-{rank}']['X'][0], + buffer[f'u-{rank}']['X'][1], + buffer[f'u-{rank}']['v'].real, + vmin=vmin['v'], + vmax=vmax['v'], + cmap='binary', + ) + else: + v3d = buffer[f'u-{rank}']['v'].real + + if projection == 'xy': + slices = [slice(None), slice(None), v3d.shape[2] // 2] + x = buffer[f'u-{rank}']['X'][0][*slices] + y = buffer[f'u-{rank}']['X'][1][*slices] + ax.set_xlabel('$x$') + ax.set_ylabel('$y$') + elif projection == 'xz': + slices = [slice(None), v3d.shape[1] // 2, slice(None)] + x = buffer[f'u-{rank}']['X'][0][*slices] + y = buffer[f'u-{rank}']['X'][2][*slices] + ax.set_xlabel('$x$') + ax.set_ylabel('$z$') + elif projection == 'yz': + slices = [v3d.shape[0] // 2, slice(None), slice(None)] + x = buffer[f'u-{rank}']['X'][1][*slices] + y = buffer[f'u-{rank}']['X'][2][*slices] + ax.set_xlabel('$y$') + ax.set_ylabel('$z$') + + im = ax.pcolormesh( + x, + y, + v3d[*slices], + vmin=vmin['v'], + vmax=vmax['v'], + cmap='binary', + ) fig.colorbar(im, cax, format=tkr.FormatStrFormatter('%.1f')) ax.set_title(f't={buffer[f"u-{rank}"]["t"]:.2f}') - ax.set_xlabel('$x$') - ax.set_ylabel('$y$') - ax.set_aspect(1.0) ax.set_aspect(1.0) return fig @@ -130,7 +159,7 @@ def get_description(self, *args, res=-1, **kwargs): desc['sweeper_params']['QI'] = 'MIN-SR-S' desc['sweeper_params']['QE'] = 'PIC' - desc['problem_params']['nvars'] = (2**8 if res == -1 else res,) * 2 + desc['problem_params']['nvars'] = (2**8 if res == -1 else res,) * self.ndim desc['problem_params']['Du'] = 0.00002 desc['problem_params']['Dv'] = 0.00001 desc['problem_params']['A'] = 0.04 diff --git a/pySDC/tests/test_problems/test_generic_MPIFFT.py b/pySDC/tests/test_problems/test_generic_MPIFFT.py new file mode 100644 index 000000000..cd82cf411 --- /dev/null +++ b/pySDC/tests/test_problems/test_generic_MPIFFT.py @@ -0,0 +1,57 @@ +import pytest + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [8, 16]) +@pytest.mark.parametrize('ny', [8, 16]) +@pytest.mark.parametrize('nz', [0, 8]) +@pytest.mark.parametrize('f', [1, 3]) +@pytest.mark.parametrize('spectral', [True, False]) +@pytest.mark.parametrize('direction', [0, 1, 10]) +def test_derivative(nx, ny, nz, f, spectral, direction): + from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT + + nvars = (nx, ny) + if nz > 0: + nvars += (nz,) + prob = IMEX_Laplacian_MPIFFT(nvars=nvars, spectral=spectral) + + xp = prob.xp + + if direction == 0: + _u = xp.sin(f * prob.X[0]) + du_expect = -(f**2) * xp.sin(f * prob.X[0]) + elif direction == 1: + _u = xp.sin(f * prob.X[1]) + du_expect = -(f**2) * xp.sin(f * prob.X[1]) + elif direction == 10: + _u = xp.sin(f * prob.X[1]) + xp.cos(f * prob.X[0]) + du_expect = -(f**2) * xp.sin(f * prob.X[1]) - f**2 * xp.cos(f * prob.X[0]) + else: + raise + + if spectral: + u = prob.fft.forward(_u) + else: + u = _u + + _du = prob.eval_f(u, 0).impl + + if spectral: + du = prob.fft.backward(_du) + else: + du = _du + assert xp.allclose(du, du_expect), 'Got unexpected derivative' + + u2 = prob.solve_system(_du, factor=1e8, u0=du, t=0) * -1e8 + + if spectral: + _u2 = prob.fft.backward(u2) + else: + _u2 = u2 + + assert xp.allclose(_u2, _u, atol=1e-7), 'Got unexpected inverse derivative' + + +if __name__ == '__main__': + test_derivative(6, 6, 6, 3, False, 1)