solar system visualization

Author: pancetta

pySDC/projects/Hamiltonian/README.rst

@@ -21,5 +21,6 @@ For this, only six planets are modeled, namely the Sun (which in its mass contai
 The acceleration due to gravitational forces are computed using a simple N^2 solver.
 All this is implemented within the ``OuterSolarSystem`` problem class.
 Coarsening can be done using only the sun for computing the acceleration.
+Note that the tests/autogenerated results only run the visualization to save time.
 
 .. include:: doc_solar_system.rst

pySDC/projects/Hamiltonian/data/outer_solar_system.dat

@@ -3,8 +3,6 @@
 import os
 import dill
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib import animation
 from mpl_toolkits.mplot3d import Axes3D
 
 import pySDC.helpers.plot_helper as plt_helper
@@ -120,7 +118,7 @@ def run_simulation(prob=None):
     out = '   Std and var for number of iterations: %4.2f -- %4.2f' % (float(np.std(niters)), float(np.var(niters)))
     print(out)
 
-    # assert np.mean(niters) <= 4.0, 'Mean number of iterations is too high, got %s' % np.mean(niters)
+    assert np.mean(niters) <= 4.0, 'Mean number of iterations is too high, got %s' % np.mean(niters)
 
     fname = 'data/' + prob + '.dat'
     f = open(fname, 'wb')
@@ -144,6 +142,9 @@ def show_results(prob=None, cwd=''):
     stats = dill.load(f)
     f.close()
 
+    plt_helper.mpl.style.use('classic')
+    plt_helper.setup_mpl()
+
     # extract error in hamiltonian and prepare for plotting
     extract_stats = filter_stats(stats, type='err_hamiltonian')
     result = defaultdict(list)
@@ -152,8 +153,6 @@ def show_results(prob=None, cwd=''):
     for k, v in result.items():
         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)
 
     # Rearrange data for easy plotting
@@ -163,8 +162,7 @@ def show_results(prob=None, cwd=''):
         ham = [item[1] for item in v]
         err_ham = ham[-1]
         plt_helper.plt.semilogy(time, ham, '-', lw=1, label='Iter ' + str(k))
-    print(err_ham)
-    # assert err_ham < 3.0E-15, 'Error in the Hamiltonian is too large for %s, got %s' % (prob, err_ham)
+    assert err_ham < 7.5E-15, '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')
@@ -177,35 +175,43 @@ def show_results(prob=None, cwd=''):
     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'
 
-    # extract error in hamiltonian and prepare for plotting
+    # extract positions and prepare for plotting
     extract_stats = filter_stats(stats, type='position')
     result = sort_stats(extract_stats, sortby='time')
 
-    plt_helper.mpl.style.use('classic')
-    plt_helper.setup_mpl()
     fig = plt_helper.plt.figure()
     ax = fig.add_subplot(111, projection='3d')
 
     # Rearrange data for easy plotting
+    nparts = len(result[1][1][0])
+    ndim = len(result[1][1])
     nsteps = len(result)
-    pos = np.zeros((6, 3, nsteps))
+    pos = np.zeros((nparts, ndim, nsteps))
 
     for idx, item in enumerate(result):
-        for n in range(6):
-            pos[n, 0, idx] = item[1][0][n]
-            pos[n, 1, idx] = item[1][1][n]
-            pos[n, 2, idx] = item[1][2][n]
-
-    for n in range(6):
-        ax.plot(pos[n, 0, :], pos[n, 1, :], zs=pos[n, 2, :])
+        for n in range(nparts):
+            for m in range(ndim):
+                pos[n, m, idx] = item[1][m][n]
+
+    for n in range(nparts):
+        if ndim == 2:
+            ax.plot(pos[n, 0, :], pos[n, 1, :])
+        elif ndim == 3:
+            ax.plot(pos[n, 0, :], pos[n, 1, :], pos[n, 2, :])
+        else:
+            raise NotImplemented('Wrong number of dimensions for plotting, got %s' % ndim)
 
     fname = 'data/' + prob + '_positions'
-    plt_helper.plt.savefig(fname)
+    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)
+    run_simulation(prob)
     show_results(prob)
 
 

pySDC/projects/Hamiltonian/solar_system.py

@@ -1,4 +1,5 @@
-from pySDC.projects.Hamiltonian.solar_system import main
+from pySDC.projects.Hamiltonian.solar_system import show_results
 
 def test_main():
-    main()
+    prob = 'outer_solar_system'
+    show_results(prob)