More plots for the RBC paper (#594)

* Added order script

* Fixes

* Fixed plotting

* Fixed recording

* Added order plot for 3D RBC

* Added spectrum plot

* Refactor

* Refactor

---------

Co-authored-by: Thomas Baumann <[email protected]>

Author: brownbaerchen

pySDC/projects/GPU/analysis_scripts/RBC3D_order.py

@@ -0,0 +1,151 @@
+import os
+import pickle
+import numpy as np
+from pySDC.helpers.fieldsIO import FieldsIO
+from pySDC.projects.GPU.configs.base_config import get_config
+from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D
+from mpi4py import MPI
+import matplotlib.pyplot as plt
+from pySDC.helpers.plot_helper import figsize_by_journal
+from pySDC.projects.GPU.analysis_scripts.RBC3D_plotting_utils import get_plotting_style, savefig
+
+step_sizes = {
+    'RBC3DG4R4Ra1e5': [8e-2, 4e-2, 2e-2, 1e-2, 5e-3],
+    'RBC3DG4R4SDC23Ra1e5': [5e-3 * 2**i for i in range(8)],
+    'RBC3DG4R4SDC34Ra1e5': [5e-3 * 2**i for i in range(8)],
+    'RBC3DG4R4SDC44Ra1e5': [5e-3 * 2**i for i in range(8)],
+    'RBC3DG4R4RKRa1e5': [5e-3 * 2**i for i in range(8)],
+    'RBC3DG4R4EulerRa1e5': [5e-3 * 2**i for i in range(8)],
+}
+n_freefall_times = {}
+
+
+def no_logging_hook(*args, **kwargs):
+    return None
+
+
+def get_path(args):
+    config = get_config(args)
+    fname = config.get_file_name()
+    return f'{fname[:fname.index('dt')]}order.pickle'
+
+
+def compute_errors(args, dts, Tend):
+    errors = {'u': [], 'v': [], 'w': [], 'T': [], 'p': [], 'dt': []}
+    prob = RayleighBenard3D(nx=4, ny=4, nz=4, comm=MPI.COMM_SELF)
+
+    dts = np.sort(dts)[::-1]
+    ref = run(args, dts[-1], Tend)
+    for dt in dts[:-1]:
+        u = run(args, dt, Tend)
+        e = u - ref
+        for comp in ['u', 'v', 'w', 'T', 'p']:
+            i = prob.index(comp)
+            e_comp = np.max(np.abs(e[i])) / np.max(np.abs(ref[i]))
+            e_comp = MPI.COMM_WORLD.allreduce(e_comp, op=MPI.MAX)
+            errors[comp].append(float(e_comp))
+        errors['dt'].append(dt)
+
+    path = get_path(args)
+    if MPI.COMM_WORLD.rank == 0:
+        with open(path, 'wb') as file:
+            pickle.dump(errors, file)
+            print(f'Saved errors to {path}', flush=True)
+
+
+def plot_error_all_components(args):  # pragma: no cover
+    fig, ax = plt.subplots()
+    with open(get_path(args), 'rb') as file:
+        errors = pickle.load(file)
+
+    for comp in ['u', 'v', 'w', 'T', 'p']:
+        e = np.array(errors[comp])
+        dt = np.array(errors['dt'])
+        order = np.log(e[1:] / e[:-1]) / np.log(dt[1:] / dt[:-1])
+        ax.loglog(errors['dt'], errors[comp], label=f'{comp} order {np.mean(order):.1f}')
+
+    ax.loglog(errors['dt'], np.array(errors['dt']) ** 4, label='Order 4', ls='--')
+    ax.loglog(errors['dt'], np.array(errors['dt']) ** 2, label='Order 2', ls='--')
+    ax.legend()
+    ax.set_xlabel(r'$\Delta t$')
+    ax.set_ylabel(r'$e$')
+
+
+def compare_order(Ra):  # pragma: no cover
+    fig, ax = plt.subplots(figsize=figsize_by_journal('Nature_CS', 1, 0.6))
+    if Ra == 1e5:
+        names = ['RK', 'Euler', 'SDC23', 'SDC34', 'SDC44'][::-1]
+        configs = [f'RBC3DG4R4{me}Ra1e5' for me in names]
+        paths = [f'./data/RBC3DG4R4{me}Ra1e5-res-1-order.pickle' for me in names]
+
+    else:
+        raise NotImplementedError
+
+    for path, config in zip(paths, configs, strict=True):
+        with open(path, 'rb') as file:
+            errors = pickle.load(file)
+
+        e = np.array(errors['T'])
+        dt = np.array(errors['dt'])
+        order = np.log(e[1:] / e[:-1]) / np.log(dt[1:] / dt[:-1])
+        print(f'{config}: order: mean={np.mean(order):.1f} / median={np.median(order):.1f}')
+        ax.loglog(dt, e, **get_plotting_style(config))
+
+    for _dt in dt:
+        for i in [1, 3, 4]:
+            ax.text(_dt, _dt**i, i, fontweight='bold', fontsize=14, ha='center', va='center')
+            ax.loglog(dt, dt**i, ls=':', color='black')
+
+    ax.legend(frameon=False)
+    ax.set_xlabel(r'$\Delta t$')
+    ax.set_ylabel(r'$e$')
+    savefig(fig, 'RBC3D_order_Ra1e5')
+
+
+def run(args, dt, Tend):
+    from pySDC.projects.GPU.run_experiment import run_experiment
+    from pySDC.core.errors import ConvergenceError
+
+    args['mode'] = 'run'
+    args['dt'] = dt
+
+    config = get_config(args)
+    config.Tend = n_freefall_times.get(type(config).__name__, 3)
+
+    desc = config.get_description(res=args['res'])
+    prob = desc['problem_class'](**desc['problem_params'])
+
+    ic_config_name = type(config).__name__
+    for name in ['RK', 'Euler', 'O3', 'O4', 'SDC23', 'SDC34', 'SDC44']:
+        ic_config_name = ic_config_name.replace(name, 'SDC34')
+
+    ic_config = get_config({**args, 'config': ic_config_name})
+    config.ic_config['config'] = type(ic_config)
+    config.ic_config['res'] = ic_config.res
+    config.ic_config['dt'] = ic_config.dt
+
+    config.get_LogToFile = no_logging_hook
+    config.Tend = Tend
+
+    u_hat = run_experiment(args, config)
+    u = prob.itransform(u_hat)
+
+    return u
+
+
+if __name__ == '__main__':
+    from pySDC.projects.GPU.run_experiment import parse_args
+
+    args = parse_args()
+
+    if args['mode'] == 'run':
+        config = get_config(args)
+        dts = step_sizes[type(config).__name__]
+        compute_errors(args, dts, max(dts))
+
+    if args['config'] is not None:
+        plot_error_all_components(args)
+
+    compare_order(1e5)
+    if MPI.COMM_WORLD.rank == 0:
+        plt.show()

pySDC/projects/GPU/analysis_scripts/RBC3D_plotting_utils.py

@@ -0,0 +1,45 @@
+from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl
+import warnings
+
+setup_mpl()
+
+
+def get_plotting_style(config):  # pragma: no cover
+
+    args = {'color': None, 'ls': None, 'marker': None, 'markersize': 6, 'label': None}
+
+    if config == 'RBC3DG4R4SDC23Ra1e5':
+        args['color'] = 'tab:blue'
+        args['ls'] = '-'
+        args['marker'] = 'o'
+        args['label'] = 'SDC23'
+    elif config == 'RBC3DG4R4SDC34Ra1e5':
+        args['color'] = 'tab:orange'
+        args['ls'] = '-'
+        args['marker'] = '<'
+        args['label'] = 'SDC34'
+    elif config in ['RBC3DG4R4SDC44Ra1e5', 'RBC3DG4R4Ra1e5']:
+        args['color'] = 'tab:green'
+        args['ls'] = '-'
+        args['marker'] = 'x'
+        args['label'] = 'SDC44'
+    elif config == 'RBC3DG4R4EulerRa1e5':
+        args['color'] = 'tab:purple'
+        args['ls'] = '--'
+        args['marker'] = '.'
+        args['label'] = 'Euler'
+    elif config == 'RBC3DG4R4RKRa1e5':
+        args['color'] = 'tab:red'
+        args['ls'] = '--'
+        args['marker'] = '>'
+        args['label'] = 'RK444'
+    else:
+        warnings.warn(f'No plotting style for {config=!r}')
+
+    return args
+
+
+def savefig(fig, name, format='pdf', base_path='./plots', **kwargs):  # pragma: no cover
+    path = f'{base_path}/{name}.{format}'
+    fig.savefig(path, bbox_inches='tight', **kwargs)
+    print(f'Saved figure {path!r}')

pySDC/projects/GPU/analysis_scripts/RBC3D_spectrum.py

@@ -0,0 +1,33 @@
+from pySDC.projects.GPU.analysis_scripts.process_RBC3D_data import get_pySDC_data
+from pySDC.projects.GPU.analysis_scripts.RBC3D_plotting_utils import figsize_by_journal, get_plotting_style, savefig
+import matplotlib.pyplot as plt
+
+
+def plot_spectrum(res, dt, config_name, ax):  # pragma: no cover
+    data = get_pySDC_data(res=res, dt=dt, config_name=config_name)
+
+    spectrum = data['avg_spectrum']
+    k = data['k']
+    ax.loglog(k[spectrum > 1e-16], spectrum[spectrum > 1e-16], **get_plotting_style(config_name), markevery=5)
+    ax.set_xlabel('$k$')
+    ax.set_ylabel(r'$\|\hat{u}_x\|$')
+
+
+def plot_spectra_Ra1e5():  # pragma: no cover
+    fig, ax = plt.subplots(figsize=figsize_by_journal('Nature_CS', 1, 0.6))
+
+    configs = [f'RBC3DG4R4{name}Ra1e5' for name in ['SDC34', 'SDC23', '', 'Euler', 'RK']]
+    dts = [0.06, 0.06, 0.06, 0.02, 0.04]
+    res = 32
+
+    for config, dt in zip(configs, dts, strict=True):
+        plot_spectrum(res, dt, config, ax)
+
+    ax.legend(frameon=False)
+    savefig(fig, 'RBC3D_spectrum_Ra1e5')
+
+
+if __name__ == '__main__':
+    plot_spectra_Ra1e5()
+
+    plt.show()

pySDC/projects/GPU/analysis_scripts/plot_Nu.py

@@ -3,18 +3,11 @@
 import numpy as np
 from scipy import integrate
 from pySDC.helpers.plot_helper import figsize_by_journal, setup_mpl
+from pySDC.projects.GPU.analysis_scripts.process_RBC3D_data import get_pySDC_data
 
 setup_mpl()
 
 
-def get_pySDC_data(res=-1, dt=-1, config_name='RBC3DG4', base_path='data/RBC_time_averaged'):
-    path = f'{base_path}/{config_name}-res{res}-dt{dt:.0e}.pickle'
-    with open(path, 'rb') as file:
-        data = pickle.load(file)
-
-    return data
-
-
 def interpolate_NuV_to_reference_times(data, reference_data, order=12):
     from qmat.lagrange import getSparseInterpolationMatrix
 

pySDC/projects/GPU/analysis_scripts/process_RBC3D_data.py

@@ -204,5 +204,13 @@ def process_RBC3D_data(base_path='./data/RBC_time_averaged', plot=True, args=Non
     return path
 
 
+def get_pySDC_data(res=-1, dt=-1, config_name='RBC3DG4', base_path='data/RBC_time_averaged'):
+    path = f'{base_path}/{config_name}-res{res}-dt{dt:.0e}.pickle'
+    with open(path, 'rb') as file:
+        data = pickle.load(file)
+
+    return data
+
+
 if __name__ == '__main__':
     process_RBC3D_data()

pySDC/projects/GPU/configs/RBC3D_configs.py

@@ -221,6 +221,15 @@ def get_description(self, *args, **kwargs):
         return desc
 
 
+class RBC3DM4K4(RBC3Dverification):
+
+    def get_description(self, *args, **kwargs):
+        desc = super().get_description(*args, **kwargs)
+        desc['level_params']['nsweeps'] = 4
+        desc['sweeper_params']['num_nodes'] = 4
+        return desc
+
+
 class RBC3DverificationRK(RBC3Dverification):
 
     def get_description(self, *args, res=-1, dt=-1, **kwargs):
@@ -268,6 +277,13 @@ class RBC3DG4R4SDC34Ra1e5(RBC3DM3K4):
     converged = 50
 
 
+class RBC3DG4R4SDC44Ra1e5(RBC3DM4K4):
+    Tend = 200
+    dt = 6e-2
+    res = 32
+    converged = 50
+
+
 class RBC3DG4R4RKRa1e5(RBC3DverificationRK):
     Tend = 200
     dt = 8e-2

pySDC/projects/GPU/tests/test_RBC_3D_analysis.py

@@ -26,10 +26,10 @@ def get_config(args):
     return config
 
 
-def generate_simulation_file(path):
+def generate_simulation_file(path, args=None):
     from pySDC.projects.GPU.run_experiment import run_experiment
 
-    args = get_args(path)
+    args = {**get_args(path), **args} if args is not None else get_args(path)
     config = get_config(args)
 
     run_experiment(args, config)
@@ -80,7 +80,7 @@ def test_processing(tmp_processed_data):
 
 
 def test_get_pySDC_data(tmp_processed_data, tmp_path):
-    from pySDC.projects.GPU.analysis_scripts.plot_Nu import get_pySDC_data
+    from pySDC.projects.GPU.analysis_scripts.process_RBC3D_data import get_pySDC_data
 
     args = get_args(tmp_path)
     data = get_pySDC_data(res=args['res'], dt=args['dt'], config_name=args['config'], base_path=tmp_path)
@@ -112,3 +112,24 @@ def _get_Nu(_t):
     tI, NuI = interpolate_NuV_to_reference_times(data, ref_data, order=4)
     assert not np.allclose(NuI, ref_data['Nu']['V'])
     assert np.allclose(tI, ref_data['t'])
+
+
+def test_error_computation(tmp_sim_data, tmp_path):
+    from pySDC.projects.GPU.analysis_scripts.RBC3D_order import compute_errors, get_path
+    from pySDC.projects.GPU.configs.RBC3D_configs import RBC3DG4R4SDC34Ra1e5
+    import numpy as np
+    import pickle
+
+    args = get_args(tmp_path)
+
+    generate_simulation_file(tmp_path, args={'config': args['config'].replace('SDC23', 'SDC34')})
+    RBC3DG4R4SDC34Ra1e5.res = args['res']
+    RBC3DG4R4SDC34Ra1e5.dt = args['dt']
+
+    dts = [1e-2, 5e-4]
+    compute_errors(args, dts, np.max(dts))
+
+    with open(get_path(args), 'rb') as file:
+        errors = pickle.load(file)
+    for comp in ['T', 'p']:
+        assert max(errors[comp]) < 1e-12