Added "Nearest" interp method for AMR grid, both in regular and irreg…

pyCalculations/calculations.py

@@ -59,4 +59,4 @@
 from fieldtracer import static_field_tracer, static_field_tracer_3d
 from fieldtracer import dynamic_field_tracer
 from non_maxwellianity import epsilon_M
-from interpolator_amr import AMRInterpolator
+from interpolator_amr import AMRInterpolator, supported_amr_interpolators

pyCalculations/interpolator_amr.py

@@ -213,28 +213,30 @@ def __call__(self, pt, cellids = None):
          return fp
 
 
+supported_amr_interpolators = {'linear','rbf','delaunay'}
+
 class AMRInterpolator(object):
    ''' Wrapper class for interpolators, esp. at refinement interfaces.
    Supported methods:
-   Trilinear
+   linear
       - (nearly) C0 continuous, regular-grid trilinear interpolant extended to collapsed hexahedral cells.
       - Non-parametric
       - Exact handling of multiply-degenerate hexahedra is missing, with the enabling hack causing errors in trilinear coordinate on the order of 1m
 
 
-   Radial Basis Functions, RBF
+   Radial Basis Functions, `rbf`
       - Accurate, slow-ish, but hard to make properly continuous and number of neighbors on which to base the interpolant is not trivial to find.
       - Not continuous with regular-grid trilinear interpolants, needs to be used in the entire interpolation domain.
       - kword options: "neighbors" for number of neighbors (64)
       - basis function uses SciPy default. A free parameter.
 
-   Delaunay (not recommended)
+   delaunay (not recommended)
       - Choice of triangulation is not unique with regular grids, including refinement interfaces.
       - kword options as in qhull; "qhull_options" : "QJ" (breaks degeneracies)
 
    '''
 
-   def __init__(self, reader, method = "Trilinear", cellids=np.array([1,2,3,4,5],dtype=np.int64)):
+   def __init__(self, reader, method = "linear", cellids=np.array([1,2,3,4,5],dtype=np.int64)):
       self.__reader = reader
       self.__cellids = np.array(list(set(cellids)),dtype=np.int64)
       self.duals = {}
@@ -249,24 +251,29 @@ def get_points(self):
       return self.__reader.get_cell_coordinates(self.__cellids)
 
    def get_interpolator(self, name, operator, coords, 
-                        method="Trilinear", 
+                        method="linear", 
                         methodargs={
                            "RBF":{"neighbors":64}, # Harrison-Stetson number of neighbors
                            "Delaunay":{"qhull_options":"QJ"}
                            }):
-      methodargs["Trilinear"] = {"reader":self.__reader, "var" : name, "op":operator}
+      # Check for old aliases
+      if method.lower() == "trilinear":
+         warnings.warn("'trilinear' interpolator method renamed to 'linear' for consistency")
+         method = "linear"
+
+      methodargs["linear"] = {"reader":self.__reader, "var" : name, "op":operator}
 
       pts = self.__reader.get_cell_coordinates(self.__cellids)
       vals = self.__reader.read_variable(name, self.__cellids, operator=operator)
-      if method == "Delaunay":
+      if method == "delaunay":
          self.__Delaunay = Delaunay(self.reader.get_cell_coordinates(self.__cellids),**methodargs[method])
          return LinearNDInterpolator(self.__Delaunay, vals)
-      elif method == "RBF":
+      elif method == "rbf":
          try:
             return RBFInterpolator(pts, vals, **methodargs[method])
          except Exception as e:
             warnings.warn("RBFInterpolator could not be imported. SciPy >= 1.7 is required for this class. Falling back to Hexahedral trilinear interpolator. Error given was " + str(e))
-            return HexahedralTrilinearInterpolator(pts, vals, **methodargs["Trilinear"])
-      elif method == "Trilinear":
+            return HexahedralTrilinearInterpolator(pts, vals, **methodargs["linear"])
+      elif method == "linear":
          return HexahedralTrilinearInterpolator(pts, vals, **methodargs[method])
 

pyVlsv/vlsvreader.py

@@ -42,9 +42,12 @@
 from variable import get_data
 import warnings
 import time
-from interpolator_amr import AMRInterpolator
+from interpolator_amr import AMRInterpolator, supported_amr_interpolators
 from operator import itemgetter
 
+
+interp_method_aliases = {"trilinear":"linear"}
+
 class PicklableFile(object):
    def __init__(self, fileobj):
       self.fileobj = fileobj
@@ -1169,7 +1172,7 @@ def read_metadata(self, name="", tag="", mesh=""):
       return -1
 
 
-   def read_interpolated_fsgrid_variable(self, name, coordinates, operator="pass",periodic=[True,True,True]):
+   def read_interpolated_fsgrid_variable(self, name, coordinates, operator="pass",periodic=[True,True,True], method="linear"):
       ''' Read a linearly interpolated FSgrid variable value from the open vlsv file. Feel free to vectorize!
       Note that this does not account for varying centerings of fsgrid data.
       Arguments:
@@ -1182,6 +1185,9 @@ def read_interpolated_fsgrid_variable(self, name, coordinates, operator="pass",p
       .. seealso:: :func:`read` :func:`read_variable_info`
       '''
 
+      if method != "Linear":
+         raise NotImplementedError("interpolation method "+method+" not implemented for read_interpolated_fsgrid_variable, only linear supported so far.")
+
       warnings.warn("read_interpolated_fsgrid_variable: face- vs. edge- centered variables not accounted for!")
 
       if name[0:3] != 'fg_':
@@ -1274,20 +1280,20 @@ def interpolateSingle(r):
          ret.append(interpolateSingle(r))
       return np.asarray(ret)
 
-   def read_interpolated_ionosphere_variable(self, name, coordinates, operator="pass"):
+   def read_interpolated_ionosphere_variable(self, name, coordinates, operator="pass", method="linear"):
       ''' Read a linearly interpolated ionosphere variable value from the open vlsv file.
       Arguments:
       :param name: Name of the (ionosphere) variable
       :param coords: Coordinates (x,y,z) from which to read data 
       :param operator: Datareduction operator. "pass" does no operation on data
+      :param method: Interpolation method. Not implemented; barycentric interp would fall under linear.
       :returns: numpy array with the data
 
       .. seealso:: :func:`read` :func:`read_variable_info`
       '''
 
       # At this stage, this function has not yet been implemented -- logging.info a warning and exit
-      logging.info('Interpolation of ionosphere variables has not yet been implemented; exiting.')
-      return -1
+      raise NotImplementedError('Interpolation of ionosphere variables has not yet been implemented; exiting.')
 
    # These are the 8 cells that span the upper corner vertex on a regular grid
    def get_vg_regular_interp_neighbors(self, cellids):
@@ -1320,13 +1326,15 @@ def get_vg_regular_interp_neighbors(self, cellids):
 
       return cellid_neighbors
 
-   def read_interpolated_variable(self, name, coords, operator="pass",periodic=[True, True, True], method="Trilinear"):
+   def read_interpolated_variable(self, name, coords, operator="pass",periodic=[True, True, True], method="linear"):
       ''' Read a linearly interpolated variable value from the open vlsv file.
       Arguments:
       :param name: Name of the variable
       :param coords: Coordinates from which to read data 
       :param periodic: Periodicity of the system. Default is periodic in all dimension
       :param operator: Datareduction operator. "pass" does no operation on data
+      :param method: Interpolation method, default "linear", options: ["nearest", "linear"]
+
       :returns: numpy array with the data
 
       .. seealso:: :func:`read` :func:`read_variable_info`
@@ -1337,11 +1345,17 @@ def read_interpolated_variable(self, name, coords, operator="pass",periodic=[Tru
       if (len(periodic)!=3):
             raise ValueError("Periodic must be a list of 3 booleans.")
 
+      if method.lower() in interp_method_aliases.keys():
+         warnings.warn("Updated alias " +method+" -> "+interp_method_aliases[method.lower()])
+         method = interp_method_aliases[method.lower()]
+
       # First test whether the requested variable is on the FSgrid or ionosphre, and redirect to the dedicated function if needed
       if name[0:3] == 'fg_':
-         return self.read_interpolated_fsgrid_variable(name, coords, operator, periodic)
+         return self.read_interpolated_fsgrid_variable(name, coords, operator, periodic, method)
       if name[0:3] == 'ig_':
-         return self.read_interpolated_ionosphere_variable(name, coords, operator, periodic)
+         return self.read_interpolated_ionosphere_variable(name, coords, operator, periodic, method)
+
+      # case vg
 
       coordinates = get_data(coords)
       coordinates = np.array(coordinates)
@@ -1365,6 +1379,18 @@ def read_interpolated_variable(self, name, coords, operator="pass",periodic=[Tru
          raise IndexError("Coordinates are required to be three-dimensional (coords.shape[1]==3 or convertible to such))")
       closest_cell_ids = self.get_cellid(coordinates)
 
+      if method.lower() == "nearest":
+         final_values = self.read_variable(name, cellids=closest_cell_ids, operator=operator)
+         if stack:
+            return final_values.squeeze()
+         else:
+            if value_length == 1:
+               return final_values.squeeze()[()] # The only special case to return a scalar instead of an array
+            else:
+               return final_values.squeeze()
+      elif method.lower() != 'linear':
+         raise NotImplementedError(method + ' is not a valid interpolation method')
+
       batch_closest_cell_coordinates=self.get_cell_coordinates(closest_cell_ids)
 
       offsets = np.zeros(coordinates.shape,dtype=np.int32)
@@ -1419,7 +1445,7 @@ def read_interpolated_variable(self, name, coords, operator="pass",periodic=[Tru
       refs0 = np.reshape(self.get_amr_level(cellid_neighbors),(-1,8))
       if np.any(np.any(refs0 != refs0[:,0][:,np.newaxis],axis =1)):
          irregs = np.any(refs0 != refs0[:,0][:,np.newaxis],axis =1)[unique_cell_indices]
-         final_values[irregs,:] = np.reshape(self.read_interpolated_variable_irregular(name, coordinates[irregs], operator, method=method),(-1,value_length))
+         final_values[irregs,:] = np.reshape(self.read_interpolated_variable_irregular(name, coordinates[irregs], operator, method=method.lower()),(-1,value_length))
          # warnings.warn("Interpolation across refinement levels. Results are now better, but some discontinuitues might appear. If that bothers, try the read_interpolated_variable_irregular variant directly.",UserWarning)
 
       if stack:
@@ -1444,20 +1470,20 @@ def get_duals(self,cids):
 
 
    def read_interpolated_variable_irregular(self, name, coords, operator="pass",periodic=[True, True, True],
-                                            method="Trilinear",
+                                            method="linear",
                                             methodargs={
-                                             "RBF":{"neighbors":64},
-                                             "Delaunay":{"qhull_options":"QJ"}
+                                             "rbf":{"neighbors":64},
+                                             "delaunay":{"qhull_options":"QJ"}
                                              }):
       ''' Read a linearly interpolated variable value from the open vlsv file.
       Arguments:
       :param name:         Name of the variable
       :param coords:       Coordinates from which to read data 
       :param periodic:     Periodicity of the system. Default is periodic in all dimension
       :param operator:     Datareduction operator. "pass" does no operation on data
-      :param method:       Method for interpolation, default "RBF" ("Delaunay" is available)
-      :param methodargs:   Dict of dicts to pass kwargs to interpolators. Default values for "RBF", "Delaunay";
-                           see scipy.interpolate.RBFInterpolator for RBF and scipy.interpolate.LinearNDInterpolator for Delaunay
+      :param method:       Method for interpolation, default "linear" ("nearest", "rbf, "delaunay")
+      :param methodargs:   Dict of dicts to pass kwargs to interpolators. Default values for "rbf", "delaunay";
+                           see scipy.interpolate.RBFInterpolator for rbf and scipy.interpolate.LinearNDInterpolator for delaunay
       :returns: numpy array with the data
 
       .. seealso:: :func:`read` :func:`read_variable_info`
@@ -1471,20 +1497,49 @@ def read_interpolated_variable_irregular(self, name, coords, operator="pass",per
       if (len(periodic)!=3):
             raise ValueError("Periodic must be a list of 3 booleans.")
 
+      if method.lower() in interp_method_aliases.keys():
+         warnings.warn("Updated alias " +method+" -> "+interp_method_aliases[method])
+         method = interp_method_aliases[method]
+
+
       # First test whether the requested variable is on the FSgrid or ionosphre, and redirect to the dedicated function if needed
       if name[0:3] == 'fg_':
-         return self.read_interpolated_fsgrid_variable(name, coords, operator, periodic)
+         return self.read_interpolated_fsgrid_variable(name, coords, operator, periodic, method)
       if name[0:3] == 'ig_':
-         return self.read_interpolated_ionosphere_variable(name, coords, operator, periodic)
+         return self.read_interpolated_ionosphere_variable(name, coords, operator, periodic, method)
+
+      # Default case: AMR grid
 
       coordinates = get_data(coords)
       coordinates = np.array(coordinates)
+
 
       ncoords = coordinates.shape[0]
       if(coordinates.shape[1] != 3):
          raise IndexError("Coordinates are required to be three-dimensional (coords.shape[1]==3 or convertible to such))")
       cellids = self.get_cellid(coordinates)
-      # containing_cells = self.get_unique_cellids(coordinates)
+
+      if method == "nearest":
+         # Check one value for the length
+         test_variable = self.read_variable(name,cellids=[1],operator=operator)
+         if isinstance(test_variable,np.ma.core.MaskedConstant):
+            value_length=1
+         elif isinstance(test_variable, Iterable):
+            value_length=len(test_variable)
+         else:
+            value_length=1
+         final_values = self.read_variable(name, cellids=cellids, operator=operator)
+         if stack:
+            return final_values.squeeze()
+         else:
+            if value_length == 1:
+               return final_values.squeeze()[()] # The only special case to return a scalar instead of an array
+            else:
+               return final_values.squeeze()      # Other methods
+      else:
+         if method.lower() not in supported_amr_interpolators:
+            raise NotImplementedError(method + ' is not a valid interpolation method for AMR grids')
+
       containing_cells = np.unique(cellids)
       self.build_duals(containing_cells)
       duals = self.get_duals(containing_cells)
@@ -1494,7 +1549,7 @@ def read_interpolated_variable_irregular(self, name, coords, operator="pass",per
 
       cells_set.discard(0)
       intp_wrapper = AMRInterpolator(self,cellids=np.array(list(cells_set)))
-      intp = intp_wrapper.get_interpolator(name,operator, coords, method=method, methodargs=methodargs)
+      intp = intp_wrapper.get_interpolator(name,operator, coords, method=method.lower(), methodargs=methodargs)
 
       final_values = intp(coords, cellids=cellids)[:,np.newaxis]