Added Lorenz attractor problem

Thomas Baumann · Thomas Baumann · commit 70fa7135656d · 2023-01-19T16:37:22.000+01:00
diff --git a/pySDC/implementations/problem_classes/Lorenz.py b/pySDC/implementations/problem_classes/Lorenz.py
@@ -0,0 +1,166 @@
+import numpy as np
+from pySDC.core.Problem import ptype
+from pySDC.implementations.datatype_classes.mesh import mesh
+
+
+class LorenzAttractor(ptype):
+    """
+    Simple script to run a Lorenz attractor problem.
+
+    The Lorenz attractor is a system of three ordinary differential equations that exhibits some chaotic behaviour.
+    It is well known for the "Butterfly Effect", because the solution looks like a butterfly (solve to Tend = 100) or
+    so to see this with these initial conditions) and because of the chaotic nature.
+
+    Since the problem is non-linear, we need to use a Newton solver.
+
+    Problem and initial conditions does not originate from, but was taken from doi.org/10.2140/camcos.2015.10.1
+    """
+
+    def __init__(self, problem_params):
+        """
+        Initialization routine
+
+        Args:
+            problem_params (dict): custom parameters for the example
+        """
+
+        # take care of essential parameters in defaults such that they always exist
+        defaults = {
+            'sigma': 10.0,
+            'rho': 28.0,
+            'beta': 8.0 / 3.0,
+            'nvars': 3,
+            'newton_tol': 1e-9,
+            'newton_maxiter': 99,
+        }
+
+        for key in defaults.keys():
+            problem_params[key] = problem_params.get(key, defaults[key])
+
+        # invoke super init, passing number of dofs, dtype_u and dtype_f
+        super().__init__(
+            init=(problem_params['nvars'], None, np.dtype('float64')),
+            dtype_u=mesh,
+            dtype_f=mesh,
+            params=problem_params,
+        )
+
+    def eval_f(self, u, t):
+        """
+        Routine to evaluate the RHS
+
+        Args:
+            u (dtype_u): current values
+            t (float): current time
+
+        Returns:
+            dtype_f: the RHS
+        """
+        # abbreviations
+        sigma = self.params.sigma
+        rho = self.params.rho
+        beta = self.params.beta
+
+        f = self.dtype_f(self.init)
+
+        f[0] = sigma * (u[1] - u[0])
+        f[1] = rho * u[0] - u[1] - u[0] * u[2]
+        f[2] = u[0] * u[1] - beta * u[2]
+        return f
+
+    def solve_system(self, rhs, dt, u0, t):
+        """
+        Simple Newton solver for the nonlinear system.
+        Notice that I did not go through the trouble of inverting the Jacobian beforehand. If you have some time on your
+        hands feel free to do that! In the current implementation it is inverted using `numpy.linalg.solve`, which is a
+        bit more expensive than simple matrix-vector multiplication.
+
+        Args:
+            rhs (dtype_f): right-hand side for the linear system
+            dt (float): abbrev. for the local stepsize (or any other factor required)
+            u0 (dtype_u): initial guess for the iterative solver
+            t (float): current time (e.g. for time-dependent BCs)
+
+        Returns:
+            dtype_u: solution as mesh
+        """
+
+        # abbreviations
+        sigma = self.params.sigma
+        rho = self.params.rho
+        beta = self.params.beta
+
+        # start Newton iterations
+        u = self.dtype_u(u0)
+        res = np.inf
+        for n in range(0, self.params.newton_maxiter):
+
+            # assemble G such that G(u) = 0 at the solution to the step
+            G = np.array(
+                [
+                    u[0] - dt * sigma * (u[1] - u[0]) - rhs[0],
+                    u[1] - dt * (rho * u[0] - u[1] - u[0] * u[2]) - rhs[1],
+                    u[2] - dt * (u[0] * u[1] - beta * u[2]) - rhs[2],
+                ]
+            )
+
+            # compute the residual and determine if we are done
+            res = np.linalg.norm(G, np.inf)
+            if res <= self.params.newton_tol or np.isnan(res):
+                break
+
+            # assemble Jacobian J of G
+            J = np.array(
+                [
+                    [1.0 + dt * sigma, -dt * sigma, 0],
+                    [-dt * (rho - u[2]), 1 + dt, dt * u[0]],
+                    [-dt * u[1], -dt * u[0], 1.0 + dt * beta],
+                ]
+            )
+
+            # solve the linear system for the Newton correction J delta = G
+            delta = np.linalg.solve(J, G)
+
+            # update solution
+            u = u - delta
+
+        return u
+
+    def u_exact(self, t, u_init=None, t_init=None):
+        """
+        Routine to return initial conditions or to approximate exact solution using scipy.
+
+        Args:
+            t (float): current time
+            u_init (pySDC.implementations.problem_classes.Lorenz.dtype_u): initial conditions for getting the exact solution
+            t_init (float): the starting time
+
+        Returns:
+            dtype_u: exact solution
+        """
+        me = self.dtype_u(self.init)
+
+        if t > 0:
+            from scipy.integrate import solve_ivp
+
+            def rhs(t, u):
+                """
+                Evaluate the right hand side, but switch the arguments for scipy.
+
+                Args:
+                    t (float): Time
+                    u (numpy.ndarray): Solution at time t
+
+                Returns:
+                    (numpy.ndarray): Right hand side evaluation
+                """
+                return self.eval_f(u, t)
+
+            tol = 100 * np.finfo(float).eps
+            u_init = self.u_exact(t=0) if u_init is None else u_init
+            t_init = 0 if t_init is None else t_init
+
+            me[:] = solve_ivp(rhs, (t_init, t), u_init, rtol=tol, atol=tol).y[:, -1]
+        else:
+            me[:] = 1.0
+        return me
diff --git a/pySDC/projects/Resilience/Lorentz.py b/pySDC/projects/Resilience/Lorentz.py
@@ -0,0 +1,189 @@
+# script to run a Lorenz attractor problem
+import numpy as np
+import matplotlib.pyplot as plt
+
+from pySDC.helpers.stats_helper import get_sorted, get_list_of_types
+from pySDC.implementations.problem_classes.Lorenz import LorenzAttractor
+from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
+from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
+from pySDC.implementations.convergence_controller_classes.adaptivity import Adaptivity
+from pySDC.core.Errors import ProblemError
+from pySDC.projects.Resilience.hook import log_data, hook_collection
+
+
+def run_Lorenz(
+    custom_description=None,
+    num_procs=1,
+    Tend=1.0,
+    hook_class=log_data,
+    fault_stuff=None,
+    custom_controller_params=None,
+    custom_problem_params=None,
+    use_MPI=False,
+    **kwargs,
+):
+    """
+    Run a Lorenz attractor 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
+    """
+
+    # initialize level parameters
+    level_params = dict()
+    level_params['dt'] = 1e-2
+
+    # initialize sweeper parameters
+    sweeper_params = dict()
+    sweeper_params['quad_type'] = 'RADAU-RIGHT'
+    sweeper_params['num_nodes'] = 3
+    sweeper_params['QI'] = 'IE'
+
+    problem_params = {
+        'newton_tol': 1e-9,
+        'newton_maxiter': 99,
+    }
+
+    if custom_problem_params is not None:
+        problem_params = {**problem_params, **custom_problem_params}
+
+    # initialize step parameters
+    step_params = dict()
+    step_params['maxiter'] = 4
+
+    # initialize controller parameters
+    controller_params = dict()
+    controller_params['logger_level'] = 30
+    controller_params['hook_class'] = hook_collection + [hook_class]
+    controller_params['mssdc_jac'] = False
+
+    if custom_controller_params is not None:
+        controller_params = {**controller_params, **custom_controller_params}
+
+    # fill description dictionary for easy step instantiation
+    description = dict()
+    description['problem_class'] = LorenzAttractor
+    description['problem_params'] = problem_params
+    description['sweeper_class'] = generic_implicit
+    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 k == 'sweeper_class':
+                description[k] = custom_description[k]
+                continue
+            description[k] = {**description.get(k, {}), **custom_description.get(k, {})}
+
+    # set time parameters
+    t0 = 0.0
+
+    # instantiate controller
+    if use_MPI:
+        from mpi4py import MPI
+        from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
+
+        comm = kwargs.get('comm', MPI.COMM_WORLD)
+        controller = controller_MPI(controller_params=controller_params, description=description, comm=comm)
+
+        # get initial values on finest level
+        P = controller.S.levels[0].prob
+        uinit = P.u_exact(t0)
+    else:
+        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:
+        raise NotImplementedError
+
+    # call main function to get things done...
+    try:
+        uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
+    except ProblemError:
+        stats = controller.hooks.return_stats()
+
+    return stats, controller, Tend
+
+
+def plot_solution(stats):
+    """
+    Plot the solution in 3D.
+
+    Args:
+        stats (dict): The stats object of the run
+
+    Returns:
+        None
+    """
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    u = get_sorted(stats, type='u')
+    ax.plot([me[1][0] for me in u], [me[1][1] for me in u], [me[1][2] for me in u])
+    ax.set_xlabel('x')
+    ax.set_ylabel('y')
+    ax.set_zlabel('z')
+    plt.show()
+
+
+def check_solution(stats, controller, thresh=5e-4):
+    """
+    Check if the global error solution wrt. a scipy reference solution is tolerable.
+
+    Args:
+        stats (dict): The stats object of the run
+        controller (pySDC.Controller.controller): The controller
+        thresh (float): Threshold for accepting the accuracy
+
+    Returns:
+        None
+    """
+    u = get_sorted(stats, type='u')
+    u_exact = controller.MS[0].levels[0].prob.u_exact(t=u[-1][0])
+    error = np.linalg.norm(u[-1][1] - u_exact)
+
+    assert error < thresh, f"Error too large, got e={error:.2e}"
+
+
+def main(plotting=True):
+    """
+    Make a test run and see if the accuracy checks out.
+
+    Args:
+        plotting (bool): Plot the solution or not
+
+    Returns:
+        None
+    """
+    custom_description = {}
+    custom_description['convergence_controllers'] = {Adaptivity: {'e_tol': 1e-5}}
+    custom_controller_params = {'logger_level': 30}
+    stats, controller, _ = run_Lorenz(
+        custom_description=custom_description, custom_controller_params=custom_controller_params, Tend=10
+    )
+    check_solution(stats, controller, 5e-4)
+    if plotting:
+        plot_solution(stats)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pySDC/tests/test_projects/test_resilience/test_Lorenz.py b/pySDC/tests/test_projects/test_resilience/test_Lorenz.py
@@ -0,0 +1,8 @@
+import pytest
+
+
+@pytest.mark.base
+def test_main():
+    from pySDC.projects.Resilience.Lorentz import main
+
+    main(False)
