From 7862e7c11eb6edbcbc446f8b44019ecc90558771 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Sat, 16 Nov 2024 11:58:50 +0100 Subject: [PATCH 1/3] Added test for generic mpi4py-fft Laplacian class --- .../test_problems/test_generic_MPIFFT.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 pySDC/tests/test_problems/test_generic_MPIFFT.py 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 0000000000..74a8accd62 --- /dev/null +++ b/pySDC/tests/test_problems/test_generic_MPIFFT.py @@ -0,0 +1,40 @@ +import pytest + + +@pytest.mark.mpi4py +@pytest.mark.parametrize('nx', [16, 32]) +@pytest.mark.parametrize('ny', [16, 32]) +@pytest.mark.parametrize('nz', [0]) +@pytest.mark.parametrize('f', [1, 3]) +@pytest.mark.parametrize('direction', [0, 1, 10]) +def test_derivative(nx, ny, nz, f, direction): + from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT + + nvars = (nx, ny) + if nz > 0: + nvars.append(nz) + prob = IMEX_Laplacian_MPIFFT(nvars=nvars) + + 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 + + du = prob.eval_f(u, 0).impl + assert xp.allclose(du, du_expect), 'Got unexpected derivative' + + _u = prob.solve_system(du, factor=1e8, u0=du, t=0) * -1e8 + assert xp.allclose(_u, u, atol=1e-7), 'Got unexpected inverse derivative' + + +if __name__ == '__main__': + test_derivative(32, 32, 0, 1, 1) From 1196767111e847c3d5ee783b4ac2abd87cc4c0bc Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Sat, 16 Nov 2024 12:49:50 +0100 Subject: [PATCH 2/3] Expanded tests --- .../generic_MPIFFT_Laplacian.py | 29 +++++++------ .../test_problems/test_generic_MPIFFT.py | 43 +++++++++++++------ 2 files changed, 45 insertions(+), 27 deletions(-) diff --git a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py index e1fc8f6812..5487f7e641 100644 --- a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py +++ b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py @@ -1,13 +1,11 @@ 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""" @@ -99,14 +97,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 +125,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/tests/test_problems/test_generic_MPIFFT.py b/pySDC/tests/test_problems/test_generic_MPIFFT.py index 74a8accd62..cd82cf4115 100644 --- a/pySDC/tests/test_problems/test_generic_MPIFFT.py +++ b/pySDC/tests/test_problems/test_generic_MPIFFT.py @@ -2,39 +2,56 @@ @pytest.mark.mpi4py -@pytest.mark.parametrize('nx', [16, 32]) -@pytest.mark.parametrize('ny', [16, 32]) -@pytest.mark.parametrize('nz', [0]) +@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, direction): +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.append(nz) - prob = IMEX_Laplacian_MPIFFT(nvars=nvars) + nvars += (nz,) + prob = IMEX_Laplacian_MPIFFT(nvars=nvars, spectral=spectral) xp = prob.xp if direction == 0: - u = xp.sin(f * prob.X[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]) + _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]) + _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 - du = prob.eval_f(u, 0).impl + 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' - _u = prob.solve_system(du, factor=1e8, u0=du, t=0) * -1e8 - assert xp.allclose(_u, u, atol=1e-7), 'Got unexpected inverse 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(32, 32, 0, 1, 1) + test_derivative(6, 6, 6, 3, False, 1) From c275e02121aced4c6cef4159078b0769eade7003 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Sat, 16 Nov 2024 13:58:03 +0100 Subject: [PATCH 3/3] 3D initial conditions in Gray-Scott --- .../problem_classes/GrayScott_MPIFFT.py | 13 ++-- .../generic_MPIFFT_Laplacian.py | 1 + pySDC/projects/GPU/configs/GS_configs.py | 63 ++++++++++++++----- 3 files changed, 54 insertions(+), 23 deletions(-) diff --git a/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py b/pySDC/implementations/problem_classes/GrayScott_MPIFFT.py index ba14ea762d..796a359b14 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 5487f7e641..8f4ee38efb 100644 --- a/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py +++ b/pySDC/implementations/problem_classes/generic_MPIFFT_Laplacian.py @@ -11,6 +11,7 @@ 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 ---------- diff --git a/pySDC/projects/GPU/configs/GS_configs.py b/pySDC/projects/GPU/configs/GS_configs.py index 6ee1374d5d..f57cdc7ea2 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