Merge pull request #276 from brownbaerchen/resilience_schroedinger

Resilience in Schrödinger equation

Author: pancetta

pySDC/implementations/controller_classes/controller_nonMPI.py

@@ -39,7 +39,7 @@ def __init__(self, num_procs, controller_params, description):
             for _ in range(num_procs - 1):
                 self.MS.append(dill.copy(self.MS[0]))
         # if this fails (e.g. due to un-picklable data in the steps), initialize seperately
-        except dill.PicklingError and TypeError:
+        except (dill.PicklingError, TypeError):
             self.logger.warning('Need to initialize steps separately due to pickling error')
             for _ in range(num_procs - 1):
                 self.MS.append(stepclass.step(description))

pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py

@@ -154,7 +154,7 @@ def solve_system(self, rhs, factor, u0, t):
 
         return me
 
-    def u_exact(self, t):
+    def u_exact(self, t, **kwargs):
         """
         Routine to compute the exact solution at time t, see (1.3) https://arxiv.org/pdf/nlin/0702010.pdf for details
 

pySDC/projects/Resilience/Lorenz.py

@@ -66,7 +66,7 @@ def run_Lorenz(
     # initialize controller parameters
     controller_params = dict()
     controller_params['logger_level'] = 30
-    controller_params['hook_class'] = hook_collection + hook_class if type(hook_class) == list else [hook_class]
+    controller_params['hook_class'] = hook_collection + (hook_class if type(hook_class) == list else [hook_class])
     controller_params['mssdc_jac'] = False
 
     if custom_controller_params is not None:

pySDC/projects/Resilience/ResilienceLorenz.py

@@ -0,0 +1,149 @@
+import numpy as np
+from pathlib import Path
+from mpi4py import MPI
+
+from pySDC.helpers.stats_helper import get_sorted
+
+from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
+from pySDC.implementations.problem_classes.NonlinearSchroedinger_MPIFFT import nonlinearschroedinger_imex
+from pySDC.implementations.transfer_classes.TransferMesh_MPIFFT import fft_to_fft
+from pySDC.projects.Resilience.hook import LogData, hook_collection
+
+
+def run_Schroedinger(
+    custom_description=None,
+    num_procs=1,
+    Tend=1.0,
+    hook_class=LogData,
+    fault_stuff=None,
+    custom_controller_params=None,
+    custom_problem_params=None,
+    use_MPI=False,
+    space_comm=None,
+    **kwargs,
+):
+    """
+    Run a Schroedinger problem with default parameters.
+
+    Args:
+        custom_description (dict): Overwrite presets
+        num_procs (int): Number of steps for MSSDC
+        Tend (float): Time to integrate to
+        hook_class (pySDC.Hook): A hook to store data
+        fault_stuff (dict): A dictionary with information on how to add faults
+        custom_controller_params (dict): Overwrite presets
+        custom_problem_params (dict): Overwrite presets
+        use_MPI (bool): Whether or not to use MPI
+
+    Returns:
+        dict: The stats object
+        controller: The controller
+        Tend: The time that was supposed to be integrated to
+    """
+
+    space_comm = MPI.COMM_WORLD if space_comm is None else space_comm
+    rank = space_comm.Get_rank()
+
+    # initialize level parameters
+    level_params = dict()
+    level_params['restol'] = 1e-08
+    level_params['dt'] = 1e-01 / 2
+    level_params['nsweeps'] = 1
+
+    # initialize sweeper parameters
+    sweeper_params = dict()
+    sweeper_params['quad_type'] = 'RADAU-RIGHT'
+    sweeper_params['num_nodes'] = 3
+    sweeper_params['QI'] = 'IE'
+    sweeper_params['initial_guess'] = 'spread'
+
+    # initialize problem parameters
+    problem_params = dict()
+    problem_params['nvars'] = (128, 128)
+    problem_params['spectral'] = False
+    problem_params['comm'] = space_comm
+
+    if custom_problem_params is not None:
+        problem_params = {**problem_params, **custom_problem_params}
+
+    # initialize step parameters
+    step_params = dict()
+    step_params['maxiter'] = 50
+
+    # initialize controller parameters
+    controller_params = dict()
+    controller_params['logger_level'] = 30 if rank == 0 else 99
+    controller_params['hook_class'] = hook_collection + (hook_class if type(hook_class) == list else [hook_class])
+    controller_params['mssdc_jac'] = False
+
+    # fill description dictionary for easy step instantiation
+    if custom_controller_params is not None:
+        controller_params = {**controller_params, **custom_controller_params}
+
+    description = dict()
+    description['problem_params'] = problem_params
+    description['problem_class'] = nonlinearschroedinger_imex
+    description['sweeper_class'] = imex_1st_order
+    description['sweeper_params'] = sweeper_params
+    description['level_params'] = level_params
+    description['step_params'] = step_params
+
+    if custom_description is not None:
+        for k in custom_description.keys():
+            if type(custom_description[k]) == dict:
+                description[k] = {**description.get(k, {}), **custom_description.get(k, {})}
+            else:
+                description[k] = custom_description[k]
+
+    # set time parameters
+    t0 = 0.0
+
+    # instantiate controller
+    assert use_MPI == False, "MPI version in time not implemented"
+    controller = controller_nonMPI(num_procs=num_procs, controller_params=controller_params, description=description)
+
+    # get initial values on finest level
+    P = controller.MS[0].levels[0].prob
+    uinit = P.u_exact(t0)
+
+    # insert faults
+    if fault_stuff is not None:
+        from pySDC.projects.Resilience.fault_injection import prepare_controller_for_faults
+
+        nvars = [me / 2 for me in problem_params['nvars']]
+        nvars[0] += 1
+
+        rnd_args = {'iteration': 5, 'problem_pos': nvars, 'min_node': 1}
+        args = {'time': 0.3, 'target': 0}
+        prepare_controller_for_faults(controller, fault_stuff, rnd_args, args)
+
+    # call main function to get things done...
+    uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
+
+    return stats, controller, Tend
+
+
+def plot_solution(stats):  # pragma: no cover
+    import matplotlib.pyplot as plt
+
+    u = get_sorted(stats, type='u')
+    fig, axs = plt.subplots(1, 2, figsize=(12, 5))
+    axs[0].imshow(np.abs(u[0][1]))
+    axs[0].set_title(f't={u[0][0]:.2f}')
+    for i in range(len(u)):
+        axs[1].cla()
+        axs[1].imshow(np.abs(u[i][1]))
+        axs[1].set_title(f't={u[i][0]:.2f}')
+        plt.pause(1e-1)
+    fig.tight_layout()
+    plt.show()
+
+
+def main():
+    stats, _, _ = run_Schroedinger(space_comm=MPI.COMM_WORLD)
+    plot_solution(stats)
+
+
+if __name__ == "__main__":
+    main()

pySDC/projects/Resilience/Schroedinger.py

@@ -113,7 +113,7 @@ def multiple_runs(
     if dt_list is not None:
         pass
     elif Tend_fixed:
-        dt_list = 0.1 * 10.0 ** -(np.arange(5) / 2)
+        dt_list = 0.1 * 10.0 ** -(np.arange(3) / 2)
     else:
         dt_list = 0.01 * 10.0 ** -(np.arange(20) / 10.0)
 
@@ -377,7 +377,7 @@ def check_order_with_adaptivity():
     we expect is 1 + 1/k.
     """
     setup_mpl()
-    ks = [4, 3, 2]
+    ks = [3, 2]
     for serial in [True, False]:
         fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))
         plot_all_errors(ax, ks, serial, Tend_fixed=5e-1, var='e_tol', dt_list=[1e-5, 1e-6, 1e-7], avoid_restarts=True)