Fixes for Storing the Conserved Quantities before and after the poloidal Step with Arakawa method

plotting/plot_diagnostics_arakawa.py

@@ -4,7 +4,7 @@
 import json
 
 
-def plot_diagnostics(foldername, save_plot=True, show_plot=False):
+def plot_diagnostics(foldername, v_loc, save_plot=True, show_plot=False):
     """
     Plot the relative differences of mass, l^2-norm, and energy against time from the
     akw_consv.txt in foldername.
@@ -13,6 +13,9 @@ def plot_diagnostics(foldername, save_plot=True, show_plot=False):
         foldername : str
             name of the directory where the data is saved.
 
+        v_loc : str
+            '0' or 'max'; where in the velocity distribution the slice (idx_v) is
+
         save_plot : bool
             if the plots should be saved
 
@@ -24,10 +27,16 @@ def plot_diagnostics(foldername, save_plot=True, show_plot=False):
     with open(foldername + "initParams.json") as file:
         dt = json.load(file)["dt"]
 
-    with open(foldername + "akw_consv.txt") as file:
-        for line in file:
-            for entry in line.split():
-                data.append(entry)
+    if v_loc == '0':
+        with open(foldername + "akw_consv_v_0.txt") as file:
+            for line in file:
+                for entry in line.split():
+                    data.append(entry)
+    elif v_loc == 'max':
+        with open(foldername + "akw_consv_v_max.txt") as file:
+            for line in file:
+                for entry in line.split():
+                    data.append(entry)
 
     data = np.array(data).reshape(-1, 6)
 
@@ -40,14 +49,34 @@ def plot_diagnostics(foldername, save_plot=True, show_plot=False):
     plt.plot(times, np.abs(np.divide(data[:, 1] - data[:, 4], data[:, 1])),
              label='$l^2$-norm')
     plt.plot(times, np.abs(np.divide(data[:, 2] - data[:, 5], data[:, 2])),
-             label='energy')
+             label='potential energy')
     plt.legend()
     plt.xlabel('time')
     plt.ylabel('error')
-    plt.title('Relative error in mass, $l^2$-norm, and energy')
+    plt.title('Relative error in mass, $l^2$-norm, and potential energy')
+
+    if save_plot:
+        if v_loc == '0':
+            plt.savefig(foldername + 'plots_v_0/akw_consv_rel_err.png')
+        elif v_loc == 'max':
+            plt.savefig(foldername + 'plots_v_max/akw_consv_rel_err.png')
+
+    if show_plot:
+        plt.show()
+
+    plt.close()
+
+    plt.plot(times, data[:, 2], label='potential energy')
+    plt.legend()
+    plt.xlabel('time')
+    plt.ylabel(r'E_{pot}')
+    plt.title('The Potential Energy as a Function of time')
 
     if save_plot:
-        plt.savefig(foldername + 'plots/akw_conservation.png')
+        if v_loc == '0':
+            plt.savefig(foldername + 'plots_v_0/akw_consv_en_pot.png')
+        elif v_loc == 'max':
+            plt.savefig(foldername + 'plots_v_max/akw_consv_en_pot.png')
 
     if show_plot:
         plt.show()
@@ -64,9 +93,12 @@ def main():
     while True:
         foldername = 'simulation_' + str(k) + '/'
         if os.path.exists(foldername):
-            if not os.path.exists(foldername + 'plots/') and os.path.exists(foldername + 'akw_consv.txt'):
-                os.mkdir(foldername + 'plots/')
-                plot_diagnostics(foldername)
+            if not os.path.exists(foldername + 'plots_v_0/') and os.path.exists(foldername + 'akw_consv_v_0.txt'):
+                os.mkdir(foldername + 'plots_v_0/')
+                plot_diagnostics(foldername, '0')
+            if not os.path.exists(foldername + 'plots_v_max/') and os.path.exists(foldername + 'akw_consv_v_max.txt'):
+                os.mkdir(foldername + 'plots_v_max/')
+                plot_diagnostics(foldername, 'max')
             k += 1
         else:
             break

plotting/postprocessing_utils.py

@@ -12,7 +12,7 @@
 from pygyro import splines as spl
 from pygyro.tools.getSlice import get_grid_slice, get_flux_surface_grid_slice
 from pygyro.tools.getPhiSlice import get_phi_slice
-from pygyro.arakawa.utilities import compute_int_f, compute_int_f_squared, get_total_energy
+from pygyro.arakawa.utilities import compute_int_f, compute_int_f_squared, get_potential_energy
 
 from pygyro.model.process_grid import compute_2d_process_grid
 from pygyro.model.grid import Grid
@@ -311,7 +311,7 @@ def plot_conservation(foldername):
 
             int_f.append(compute_int_f(f.ravel(), dtheta, dr, r))
             int_f2.append(compute_int_f_squared(f.ravel(), dtheta, dr, r))
-            int_en.append(get_total_energy(
+            int_en.append(get_potential_energy(
                 f.ravel(), phi.ravel(), dtheta, dr, r))
 
     p = np.array(t_f).argsort()

pygyro/advection/advection.py

@@ -5,6 +5,7 @@
 from math import pi
 from scipy.sparse.linalg import spsolve
 from scipy.sparse import diags
+from mpi4py import MPI
 
 from ..arakawa.discrete_brackets_polar import assemble_bracket_arakawa, assemble_row_columns_akw_bracket_4th_order_extrapolation, update_bracket_4th_order_dirichlet_extrapolation
 from ..splines.splines import BSplines, Spline1D, Spline2D
@@ -17,7 +18,7 @@
     poloidal_advection_step_expl, \
     poloidal_advection_step_impl
 from ..initialisation.initialiser_funcs import f_eq
-from ..arakawa.utilities import compute_int_f, compute_int_f_squared, get_total_energy
+from ..arakawa.utilities import compute_int_f, compute_int_f_squared, get_potential_energy
 
 
 def fieldline(theta, z_diff, iota, r, R0):
@@ -642,6 +643,9 @@ def __init__(self, eta_vals: list, constants,
                  equilibrium_outside: bool = True, verbose: bool = False,
                  explicit: bool = False,
                  save_conservation: bool = False, foldername=''):
+        comm = MPI.COMM_WORLD
+        rank = comm.Get_rank()
+
         self._points = eta_vals[1::-1]
         self._points_theta = self._points[0]
         self._points_r = self._points[1]
@@ -675,17 +679,23 @@ def __init__(self, eta_vals: list, constants,
             with open(foldername + "/initParams.json") as file:
                 self._dt = json.load(file)["dt"]
 
-            self._conservation_savefile = "{0}/akw_consv.txt".format(
+            self._consv_savefile_v_0 = "{0}/akw_consv_v_0.txt".format(
                 foldername)
 
             # Create savefile if it does not already exist from previous simulation
-            if not os.path.exists(self._conservation_savefile):
-                with open(self._conservation_savefile, 'w') as savefile:
+            if not os.path.exists(self._consv_savefile_v_0) and rank == 0:
+                with open(self._consv_savefile_v_0, 'w') as savefile:
                     savefile.write(
-                        "int_f before\t\t\tint_f_sqd before\t\tenergy before\t\t\tint_f after\t\t\t\tint_f_sqd after\t\t\tenergy after\n")
-            else:
-                with open(self._conservation_savefile, 'w') as savefile:
-                    savefile.write("\n")
+                        "int_f before\t\t\tl2_norm_f before\t\ten_pot before\t\t\tint_f after\t\t\t\tint_f_sqd after\t\t\ten_pot after\n")
+
+            self._consv_savefile_v_max = "{0}/akw_consv_v_max.txt".format(
+                foldername)
+
+            # Create savefile if it does not already exist from previous simulation
+            if not os.path.exists(self._consv_savefile_v_max) and rank == 0:
+                with open(self._consv_savefile_v_max, 'w') as savefile:
+                    savefile.write(
+                        "int_f before\t\t\tl2_norm_f before\t\ten_pot before\t\t\tint_f after\t\t\t\tint_f_sqd after\t\t\ten_pot after\n")
 
         self.bc = bc
 
@@ -1046,8 +1056,13 @@ def gridStep(self, f: Grid, phi: Grid, dt: float):
                     if self._save_conservation and dt == self._dt:
                         global_v = global_inds_v[i]
                         global_z = global_inds_z[j]
-                        if global_v == (f.eta_grid[3].size // 2) and global_z == 0:
-                            self._save_consv_to_file('before', f, i, j, phi)
+                        if global_z == 0:
+                            if global_v == (f.eta_grid[3].size // 2):
+                                self._save_consv_to_file(
+                                    'before', '0', f, i, j, phi)
+                            elif global_v == f.eta_grid[3].size - 1:
+                                self._save_consv_to_file(
+                                    'before', 'max', f, i, j, phi)
 
                     # assume phi equals 0 outside
                     values_phi = np.zeros(self.order)
@@ -1065,8 +1080,13 @@ def gridStep(self, f: Grid, phi: Grid, dt: float):
                     if self._save_conservation and dt == self._dt:
                         global_v = global_inds_v[i]
                         global_z = global_inds_z[j]
-                        if global_v == (f.eta_grid[3].size // 2) and global_z == 0:
-                            self._save_consv_to_file('after', f, i, j, phi)
+                        if global_z == 0:
+                            if global_v == (f.eta_grid[3].size // 2):
+                                self._save_consv_to_file(
+                                    'after', '0', f, i, j, phi)
+                            elif global_v == f.eta_grid[3].size - 1:
+                                self._save_consv_to_file(
+                                    'after', 'max', f, i, j, phi)
 
         else:
             # Do step
@@ -1078,8 +1098,13 @@ def gridStep(self, f: Grid, phi: Grid, dt: float):
                     if self._save_conservation and dt == self._dt:
                         global_v = global_inds_v[i]
                         global_z = global_inds_z[j]
-                        if global_v == (f.eta_grid[3].size // 2) and global_z == 0:
-                            self._save_consv_to_file('before', f, i, j, phi)
+                        if global_z == 0:
+                            if global_v == (f.eta_grid[3].size // 2):
+                                self._save_consv_to_file(
+                                    'before', '0', f, i, j, phi)
+                            elif global_v == f.eta_grid[3].size - 1:
+                                self._save_consv_to_file(
+                                    'before', 'max', f, i, j, phi)
 
                     self.step_normal(f.get2DSlice(i, j),
                                      dt, phi.get2DSlice(j))
@@ -1088,10 +1113,15 @@ def gridStep(self, f: Grid, phi: Grid, dt: float):
                     if self._save_conservation and dt == self._dt:
                         global_v = global_inds_v[i]
                         global_z = global_inds_z[j]
-                        if global_v == (f.eta_grid[3].size // 2) and global_z == 0:
-                            self._save_consv_to_file('after', f, i, j, phi)
-
-    def _save_consv_to_file(self, boa: str, f: Grid, idx_v: int, idx_z: int, phi: Grid):
+                        if global_z == 0:
+                            if global_v == (f.eta_grid[3].size // 2):
+                                self._save_consv_to_file(
+                                    'after', '0', f, i, j, phi)
+                            elif global_v == f.eta_grid[3].size - 1:
+                                self._save_consv_to_file(
+                                    'after', 'max', f, i, j, phi)
+
+    def _save_consv_to_file(self, boa: str, v_loc: str, f: Grid, idx_v: int, idx_z: int, phi: Grid):
         """
         Compute the conserved quantities (integral of f and its square, energy) and save it to the savefile.
 
@@ -1100,6 +1130,9 @@ def _save_consv_to_file(self, boa: str, f: Grid, idx_v: int, idx_z: int, phi: Gr
             boa : str
                 'before' or 'after'; if computation is before or after doing the grid_step
 
+            v_loc : str
+                '0' or 'max'; where in the velocity distribution the slice (idx_v) is
+
             f : pygyro.model.grid.Grid
                 Grid object that characterizes the distribution function
 
@@ -1113,26 +1146,36 @@ def _save_consv_to_file(self, boa: str, f: Grid, idx_v: int, idx_z: int, phi: Gr
                 Grid object that characterizes the electric field
         """
         if boa == 'before':
-            int_f_before = compute_int_f(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
-                                         self._points_r, method='trapz')
-            int_f_squared_before = compute_int_f_squared(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
-                                                         self._points_r, method='trapz')
-            energy_before = get_total_energy(f.get2DSlice(idx_v, idx_z), phi.get2DSlice(idx_z), self._dtheta, self._dr,
-                                             self._points_r, method='trapz')
-            with open(self._conservation_savefile, 'a') as savefile:
-                savefile.write(
-                    format(int_f_before, '.15E') + "\t" + format(int_f_squared_before, '.15E') + "\t" + format(energy_before, '.15E') + "\t")
+            int_f = compute_int_f(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
+                                  self._points_r, method='trapz')
+            l2_norm_f = compute_int_f_squared(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
+                                              self._points_r, method='trapz')
+            en_pot = get_potential_energy(f.get2DSlice(idx_v, idx_z), phi.get2DSlice(idx_z), self._dtheta, self._dr,
+                                          self._points_r, method='trapz')
+            if v_loc == '0':
+                with open(self._consv_savefile_v_0, 'a') as savefile:
+                    savefile.write(
+                        format(int_f, '.15E') + "\t" + format(l2_norm_f, '.15E') + "\t" + format(en_pot, '.15E') + "\t")
+            elif v_loc == 'max':
+                with open(self._consv_savefile_v_max, 'a') as savefile:
+                    savefile.write(
+                        format(int_f, '.15E') + "\t" + format(l2_norm_f, '.15E') + "\t" + format(en_pot, '.15E') + "\t")
 
         elif boa == 'after':
-            int_f_after = compute_int_f(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
-                                        self._points_r, method='trapz')
-            int_f_squared_after = compute_int_f_squared(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
-                                                        self._points_r, method='trapz')
-            energy_after = get_total_energy(f.get2DSlice(idx_v, idx_z), phi.get2DSlice(idx_z), self._dtheta, self._dr,
-                                            self._points_r, method='trapz')
-            with open(self._conservation_savefile, 'a') as savefile:
-                savefile.write(
-                    format(int_f_after, '.15E') + "\t" + format(int_f_squared_after, '.15E') + "\t" + format(energy_after, '.15E') + "\n")
+            int_f = compute_int_f(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
+                                  self._points_r, method='trapz')
+            l2_norm_f = compute_int_f_squared(f.get2DSlice(idx_v, idx_z), self._dtheta, self._dr,
+                                              self._points_r, method='trapz')
+            en_pot = get_potential_energy(f.get2DSlice(idx_v, idx_z), phi.get2DSlice(idx_z), self._dtheta, self._dr,
+                                          self._points_r, method='trapz')
+            if v_loc == '0':
+                with open(self._consv_savefile_v_0, 'a') as savefile:
+                    savefile.write(
+                        format(int_f, '.15E') + "\t" + format(l2_norm_f, '.15E') + "\t" + format(en_pot, '.15E') + "\n")
+            elif v_loc == 'max':
+                with open(self._consv_savefile_v_max, 'a') as savefile:
+                    savefile.write(
+                        format(int_f, '.15E') + "\t" + format(l2_norm_f, '.15E') + "\t" + format(en_pot, '.15E') + "\n")
 
         else:
             raise NotImplementedError(

pygyro/arakawa/test_discrete_brackets.py

@@ -3,7 +3,7 @@
 from scipy.sparse.linalg import spsolve
 from scipy.sparse import diags
 
-from .utilities import compute_int_f, compute_int_f_squared, get_total_energy
+from .utilities import compute_int_f, compute_int_f_squared, get_potential_energy
 from .discrete_brackets_polar import assemble_bracket_arakawa
 
 
@@ -250,8 +250,8 @@ def test_conservation(bc, order, int_method, tol=1e-10, iter_tol=1e-10):
                                r_grid, method=int_method)
     int_f_squared_init = compute_int_f_squared(f, d_theta, d_r,
                                                r_grid, method=int_method)
-    total_energy_init = get_total_energy(f, phi, d_theta, d_r,
-                                         r_grid, method=int_method)
+    total_energy_init = get_potential_energy(f, phi, d_theta, d_r,
+                                             r_grid, method=int_method)
 
     for _ in range(N):
         # scaling is only found in the identity
@@ -265,8 +265,8 @@ def test_conservation(bc, order, int_method, tol=1e-10, iter_tol=1e-10):
                               r_grid, method=int_method)
         int_f_squared = compute_int_f_squared(f, d_theta, d_r,
                                               r_grid, method=int_method)
-        total_energy = get_total_energy(f, phi, d_theta, d_r,
-                                        r_grid, method=int_method)
+        total_energy = get_potential_energy(f, phi, d_theta, d_r,
+                                            r_grid, method=int_method)
 
         assert np.abs(int_f - int_f_init)/int_f_init < iter_tol
         assert np.abs(int_f_squared - int_f_squared_init) / \

pygyro/arakawa/utilities.py

@@ -240,10 +240,10 @@ def compute_int_f_squared(f: np.ndarray,
     return res
 
 
-def get_total_energy(f: np.ndarray, phi: np.ndarray,
-                     d_theta: float, d_r: float,
-                     r_grid: np.ndarray,
-                     method: str = 'sum'):
+def get_potential_energy(f: np.ndarray, phi: np.ndarray,
+                         d_theta: float, d_r: float,
+                         r_grid: np.ndarray,
+                         method: str = 'sum'):
     """
     Compute the total energy, i.e. the integral of f times phi over the whole
     domain in polar coordinates. A uniform grid is assumed.

testSetups/iota0.json

@@ -21,8 +21,8 @@
 "m" : 5,
 "n" : 1,
 "iotaVal" : 0.0,
-"npts" : [32, 32, 8, 32],
-"dt": 16,
+"npts" : [16, 16, 20, 20],
+"dt": 40,
 "_variableNames" : ["r", "theta", "z", "v"],
 "poloidalAdvectionMethod" : ["akw"]
 }