Improve grid interpolation for restart files

boutdata/data.py

@@ -695,9 +695,9 @@ def __init__(
                         comments = []
                     if inline_comment is not None:
                         parent_section.inline_comments[sectionname] = inline_comment
-                        parent_section._comment_whitespace[
-                            sectionname
-                        ] = comment_whitespace
+                        parent_section._comment_whitespace[sectionname] = (
+                            comment_whitespace
+                        )
                 else:
                     # A key=value pair
 

boutdata/gen_surface.py

@@ -1,6 +1,7 @@
 """Flux surface generator for tokamak grid files
 
 """
+
 from __future__ import print_function
 
 import numpy as np

boutdata/mayavi2.py

@@ -44,7 +44,6 @@ def aligned_points(grid, nz=1, period=1.0, maxshift=0.4):
 
 
 def create_grid(grid, data, period=1):
-
     s = np.shape(data)
 
     nx = grid["nx"]

boutdata/mms.py

@@ -408,7 +408,6 @@ def write(self, nx, ny, output, MXG=2):
             ("sinty", self.sinty),
             ("zShift", self.zShift),
         ]:
-
             # Note: This is slow, and could be improved using something like lambdify
             values = zeros([ngx, ngy])
             for i, x in enumerate(xarr):

boutdata/restart.py

@@ -174,7 +174,6 @@ def is_pow2(x):
 
         # Open the restart file in read mode and create the new file
         with DataFile(f) as old, DataFile(new_f, write=True, create=True) as new:
-
             # Find the dimension
             for var in old.list():
                 # Read the data
@@ -223,7 +222,6 @@ def is_pow2(x):
 
                 # Find 3D variables
                 if old.ndims(var) == 3:
-
                     # Asynchronous call (locks first at .get())
                     jobs.append(
                         pool.apply_async(
@@ -974,6 +972,7 @@ def change_grid(
     output=".",
     interpolator="nearest",
     show=False,
+    floor=True,
 ):
     """
     Convert a set of restart files from one grid to another
@@ -992,9 +991,11 @@ def change_grid(
     output : str, optional
          Directory where output restart files will be written
     interpolator : str, optional
-         Interpolation method to use. Options are 'nearest', 'CloughTocher', 'RBF'
+         Interpolation method to use. Options are 'nearest', 'CloughTocher', 'RBF', 'linear'
     show : bool, optional
          Display the interpolated fields using Matplotlib
+    floor : bool, optional
+         Apply floor to interpolated densities and pressures
 
     """
 
@@ -1047,6 +1048,9 @@ def change_grid(
     # Read information from a restart file
     with DataFile(file_list[0]) as f:
         for var in copy_vars:
+            if var not in f:
+                print(f"{var} not present - skipping")
+                continue
             copy_data[var] = f[var]
 
         # Get a list of variables
@@ -1109,14 +1113,22 @@ def extrapolate_yguards(data2d, myg):
                 from_data.flatten(),
                 neighbors=50,
             )
-
         elif interpolator == "nearest":
             # Nearest neighbour. Tends to be robust
             from scipy.interpolate import NearestNDInterpolator
 
             interp = NearestNDInterpolator(
                 list(zip(from_Rxy.flatten(), from_Zxy.flatten())), from_data.flatten()
             )
+        elif interpolator == "linear":
+            # Linear interpolator
+            from scipy.interpolate import LinearNDInterpolator
+
+            interp = LinearNDInterpolator(
+                list(zip(from_Rxy.flatten(), from_Zxy.flatten())),
+                from_data.flatten(),
+                fill_value=0.0,
+            )
         else:
             raise ValueError("Invalid interpolator")
 
@@ -1132,6 +1144,17 @@ def extrapolate_yguards(data2d, myg):
                 np.amax(to_data),
             )
         )
+
+        if floor:
+            if var[0] == "N" and var[1] != "V":
+                # A density
+                to_data = np.clip(to_data, 1e-5, None)
+                print("\t\t floor -> {}:{}".format(np.amin(to_data), np.amax(to_data)))
+            elif var[0] == "P":
+                # Pressure
+                to_data = np.clip(to_data, 1e-8, None)
+                print("\t\t floor -> {}:{}".format(np.amin(to_data), np.amax(to_data)))
+
         if show:
             import matplotlib.pyplot as plt