first solar system problem up and running

pancetta · pancetta · commit c64e391c9f3b · 2018-04-17T13:15:32.000+02:00
diff --git a/pySDC/implementations/problem_classes/OuterSolarSystem.py b/pySDC/implementations/problem_classes/OuterSolarSystem.py
@@ -0,0 +1,115 @@
+from __future__ import division
+
+import numpy as np
+
+from pySDC.core.Problem import ptype
+from pySDC.implementations.datatype_classes.particles import particles, fields, acceleration
+from pySDC.core.Errors import ParameterError, ProblemError
+
+
+# noinspection PyUnusedLocal
+class outer_solar_system(ptype):
+    """
+    Example implementing the harmonic oscillator
+    """
+
+    def __init__(self, problem_params, dtype_u, dtype_f):
+        """
+        Initialization routine
+
+        Args:
+            problem_params (dict): custom parameters for the example
+            dtype_u: particle data type (will be passed parent class)
+            dtype_f: acceleration data type (will be passed parent class)
+        """
+
+        # these parameters will be used later, so assert their existence
+        essential_keys = []
+        for key in essential_keys:
+            if key not in problem_params:
+                msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys()))
+                raise ParameterError(msg)
+
+        # invoke super init, passing nparts, dtype_u and dtype_f
+        super(outer_solar_system, self).__init__((3, 6), dtype_u, dtype_f, problem_params)
+        self.G = 2.95912208286E-4
+
+    def eval_f(self, u, t):
+        """
+        Routine to compute the RHS
+
+        Args:
+            u (dtype_u): the particles
+            t (float): current time (not used here)
+        Returns:
+            dtype_f: RHS
+        """
+        me = acceleration(self.init, val=0.0)
+
+        for i in range(self.init[-1]):
+            for j in range(i):
+                dx = u.pos.values[:, i] - u.pos.values[:, j]
+                r = np.sqrt(np.dot(dx, dx))
+                df = self.G * dx / (r ** 3)
+                me.values[:, i] -= u.m[j] * df
+                me.values[:, j] += u.m[i] * df
+
+        return me
+
+    def u_exact(self, t):
+        """
+        Routine to compute the exact trajectory at time t
+
+        Args:
+            t (float): current time
+        Returns:
+            dtype_u: exact position and velocity
+        """
+        assert t == 0.0, 'error, u_exact only works for the initial time t0=0'
+        me = particles(self.init)
+
+        me.pos.values[:, 0] = [0.0, 0.0, 0.0]
+        me.pos.values[:, 1] = [-3.5025653,-3.8169847,-1.5507963]
+        me.pos.values[:, 2] = [9.0755314,-3.0458353,-1.6483708]
+        me.pos.values[:, 3] = [8.3101420,-16.2901086,-7.2521278]
+        me.pos.values[:, 4] = [11.4707666,-25.7294829,-10.8169456]
+        me.pos.values[:, 5] = [-15.5387357,-25.2225594,-3.1902382]
+
+        me.vel.values[:, 0] = [0.0, 0.0, 0.0]
+        me.vel.values[:, 1] = [0.00565429,-0.00412490,-0.00190589]
+        me.vel.values[:, 2] = [0.00168318, 0.00483525, 0.00192462]
+        me.vel.values[:, 3] = [0.00354178, 0.00137102, 0.00055029]
+        me.vel.values[:, 4] = [0.00288930, 0.00114527, 0.00039677]
+        me.vel.values[:, 5] = [0.00276725, -0.0017072,-0.00136504]
+
+        me.m[0] = 1.00000597682
+        me.m[1] = 0.000954786104043
+        me.m[2] = 0.000285583733151
+        me.m[3] = 0.0000437273164546
+        me.m[4] = 0.0000517759138449
+        me.m[5] = 1.0 / 130000000.0
+
+        return me
+
+    def eval_hamiltonian(self, u):
+        """
+        Routine to compute the Hamiltonian
+
+        Args:
+            u (dtype_u): the particles
+        Returns:
+            float: hamiltonian
+        """
+
+        ham = 0
+
+        for i in range(self.init[-1]):
+            ham += 0.5 * u.m[i] * np.dot(u.vel.values[:, i], u.vel.values[:, i])
+
+        for i in range(self.init[-1]):
+            for j in range(i):
+                dx = u.pos.values[:, i] - u.pos.values[:, j]
+                r = np.sqrt(np.dot(dx, dx))
+                ham -= self.G * u.m[i] * u.m[j] / r
+
+        return ham
diff --git a/pySDC/implementations/sweeper_classes/verlet.py b/pySDC/implementations/sweeper_classes/verlet.py
@@ -156,6 +156,8 @@ def integrate(self):
             for j in range(1, self.coll.num_nodes + 1):
                 p[-1].pos += L.dt * (L.dt * self.QQ[m, j] * L.f[j]) + L.dt * self.coll.Qmat[m, j] * L.u[0].vel
                 p[-1].vel += L.dt * self.coll.Qmat[m, j] * L.f[j]
+                p[-1].m = L.u[0].m
+                p[-1].q = L.u[0].q
 
         return p
 
diff --git a/pySDC/projects/Hamiltonian/solar_system.py b/pySDC/projects/Hamiltonian/solar_system.py
@@ -0,0 +1,156 @@
+from __future__ import division
+from collections import defaultdict
+import os
+import dill
+
+import pySDC.helpers.plot_helper as plt_helper
+
+from pySDC.implementations.collocation_classes.gauss_lobatto import CollGaussLobatto
+from pySDC.implementations.controller_classes.allinclusive_classic_nonMPI import allinclusive_classic_nonMPI
+from pySDC.implementations.sweeper_classes.verlet import verlet
+from pySDC.implementations.datatype_classes.particles import particles, acceleration
+from pySDC.implementations.problem_classes.OuterSolarSystem import outer_solar_system
+from pySDC.implementations.transfer_classes.TransferParticles_NoCoarse import particles_to_particles
+
+from pySDC.helpers.stats_helper import filter_stats
+from pySDC.projects.Hamiltonian.hamiltonian_output import hamiltonian_output
+
+
+def setup_outer_solar_system():
+    """
+    Helper routine for setting up everything for the outer solar system problem
+
+    Returns:
+        description (dict): description of the controller
+        controller_params (dict): controller parameters
+    """
+
+    # initialize level parameters
+    level_params = dict()
+    level_params['restol'] = 1E-10
+    level_params['dt'] = 100.0
+
+    # initialize sweeper parameters
+    sweeper_params = dict()
+    sweeper_params['collocation_class'] = CollGaussLobatto
+    sweeper_params['num_nodes'] = [5]
+    sweeper_params['spread'] = False
+
+    # initialize problem parameters for the Penning trap
+    problem_params = dict()
+
+    # initialize step parameters
+    step_params = dict()
+    step_params['maxiter'] = 50
+
+    # initialize controller parameters
+    controller_params = dict()
+    controller_params['hook_class'] = hamiltonian_output  # specialized hook class for more statistics and output
+    controller_params['logger_level'] = 30
+    controller_params['predict'] = False
+
+    # Fill description dictionary for easy hierarchy creation
+    description = dict()
+    description['problem_class'] = outer_solar_system
+    description['problem_params'] = problem_params
+    description['dtype_u'] = particles
+    description['dtype_f'] = acceleration
+    description['sweeper_class'] = verlet
+    description['sweeper_params'] = sweeper_params
+    description['level_params'] = level_params
+    description['step_params'] = step_params
+    # description['space_transfer_class'] = particles_to_particles
+
+    return description, controller_params
+
+
+def run_simulation(prob=None):
+    """
+    Routine to run the simulation of a second order problem
+
+    Args:
+        prob (str): name of the problem
+
+    """
+
+    if prob == 'outer_solar_system':
+        description, controller_params = setup_outer_solar_system()
+        # set time parameters
+        t0 = 0.0
+        Tend = 100000.0
+        num_procs = 1
+    else:
+        raise NotImplemented('Problem type not implemented, got %s' % prob)
+
+    # instantiate the controller
+    controller = allinclusive_classic_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(t=t0)
+
+    # call main function to get things done...
+    uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
+
+    fname = 'data/' + prob + '.dat'
+    f = open(fname, 'wb')
+    dill.dump(stats, f)
+    f.close()
+
+    assert os.path.isfile(fname), 'Run for %s did not create stats file' % prob
+
+
+def show_results(prob=None, cwd=''):
+    """
+    Helper function to plot the error of the Hamiltonian
+
+    Args:
+        prob (str): name of the problem
+        cwd (str): current working directory
+    """
+    f = open(cwd + 'data/' + prob + '.dat', 'rb')
+    stats = dill.load(f)
+    f.close()
+
+    extract_stats = filter_stats(stats, type='err_hamiltonian')
+    result = defaultdict(list)
+    for k, v in extract_stats.items():
+        result[k.iter].append((k.time, v))
+    for k, v in result.items():
+        # assert k <= 6, 'Number of iterations is too high for %s, got %s' % (prob, k)
+        result[k] = sorted(result[k], key=lambda x: x[0])
+
+    plt_helper.mpl.style.use('classic')
+    plt_helper.setup_mpl()
+    plt_helper.newfig(textwidth=238.96, scale=0.89)
+
+    err_ham = 1
+    for k, v in result.items():
+        time = [item[0] for item in v]
+        ham = [item[1] for item in v]
+        err_ham = ham[-1]
+        plt_helper.plt.semilogy(time, ham, '-', lw=1, label='Iter ' + str(k))
+
+    # assert err_ham < 2.3E-08, 'Error in the Hamiltonian is too large for %s, got %s' % (prob, err_ham)
+
+    plt_helper.plt.xlabel('Time')
+    plt_helper.plt.ylabel('Error in Hamiltonian')
+    plt_helper.plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+
+    fname = 'data/' + prob + '_hamiltonian'
+    plt_helper.savefig(fname)
+
+    assert os.path.isfile(fname + '.pdf'), 'ERROR: plotting did not create PDF file'
+    assert os.path.isfile(fname + '.pgf'), 'ERROR: plotting did not create PGF file'
+    assert os.path.isfile(fname + '.png'), 'ERROR: plotting did not create PNG file'
+
+
+def main():
+    prob = 'outer_solar_system'
+    run_simulation(prob)
+    show_results(prob)
+
+
+if __name__ == "__main__":
+    main()
