pyCalculations/derivatives, VlsvWriter updates

pyCalculations/calculations.py

@@ -56,4 +56,4 @@
 import fit
 from fieldtracer import static_field_tracer
 from fieldtracer import dynamic_field_tracer
-
+from derivatives import fg_divPoynting, fg_PoyntingFlux, fg_vol_curl, ballooning_crit,vfield3_curvature

pyCalculations/derivatives.py

@@ -0,0 +1,146 @@
+# 
+# This file is part of Analysator.
+# Copyright 2013-2016 Finnish Meteorological Institute
+# Copyright 2017-2018 University of Helsinki
+# 
+# For details of usage, see the COPYING file and read the "Rules of the Road"
+# at http://www.physics.helsinki.fi/vlasiator/
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+# 
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+# 
+
+import numpy as np
+import scipy as sp
+import pytools as pt
+import warnings
+from scipy import interpolate
+
+def fg_PoyntingFlux(reader):
+   b = reader.read_fg_variable_as_volumetric('fg_b')
+   print('b.shape=',b.shape)
+   e = reader.read_fg_variable_as_volumetric('fg_e')
+   print('e.shape=',e.shape)
+
+   mu_0 = 1.25663706144e-6
+   S = np.cross(e,b)/mu_0
+   return S
+
+
+def fg_divPoynting(reader):
+   S = fg_PoyntingFlux(reader)
+   dx = reader.get_fsgrid_cell_size()
+   divS = ((np.roll(S[:,:,:,0],-1, 0) - np.roll(S[:,:,:,0], 1, 0))/(2*dx[0]) +
+           (np.roll(S[:,:,:,1],-1, 1) - np.roll(S[:,:,:,1], 1, 1))/(2*dx[1]) +
+           (np.roll(S[:,:,:,2],-1, 2) - np.roll(S[:,:,:,2], 1, 2))/(2*dx[2])
+          )
+   print('divS.shape', divS.shape)
+
+   return divS
+
+def fg_vol_jacobian(reader, b):
+   # Return the jacobian of an fsgrid volumetric variable
+   # result is a 3x3 array for each cell
+   # Last index of output gives the vector component,
+   # second-to-last the direction of derivative
+   dx = reader.get_fsgrid_cell_size()
+   dFx_dx, dFx_dy, dFx_dz = np.gradient(b[:,:,:,0], dx[0])
+   dFy_dx, dFy_dy, dFy_dz = np.gradient(b[:,:,:,1], dx[1])
+   dFz_dx, dFz_dy, dFz_dz = np.gradient(b[:,:,:,2], dx[2])
+
+   dFx = np.stack([dFx_dx, dFx_dy, dFx_dz], axis=-1)
+   dFy = np.stack([dFy_dx, dFy_dy, dFy_dz], axis=-1)
+   dFz = np.stack([dFz_dx, dFz_dy, dFz_dz], axis=-1)
+
+   return np.stack([dFx,dFy,dFz], axis=-1)
+
+def fg_vol_curl(reader, array):
+   dx = reader.get_fsgrid_cell_size()
+   dummy,  dFx_dy, dFx_dz = np.gradient(array[:,:,:,0], dx)
+   dFy_dx, dummy,  dFy_dz = np.gradient(array[:,:,:,1], dx)
+   dFz_dx, dFz_dy, dummy  = np.gradient(array[:,:,:,2], dx)
+
+   rotx = dFz_dy - dFy_dz
+   roty = dFx_dz - dFz_dx
+   rotz = dFy_dx - dFx_dy
+
+   return np.stack([rotx, roty, rotz], axis=-1)
+
+def fg_vol_div(reader, array):
+   dx = reader.get_fsgrid_cell_size()
+   dFx_dx = np.gradient(array[:,:,:,0], dx[0], axis=0)
+   dFy_dy = np.gradient(array[:,:,:,1], dx[1], axis=1)
+   dFz_dz = np.gradient(array[:,:,:,2], dx[2], axis=2)
+
+   return dFx_dx+dFy_dy+dFz_dz
+
+def vfield3_dot(a, b):
+    """Calculates dot product of vectors a and b in 3D vector field"""
+
+    return (a*b).sum(-1)
+    #return (
+    #    a[:, :, :, 0] * b[:, :, :, 0]
+    #    + a[:, :, :, 1] * b[:, :, :, 1]
+    #    + a[:, :, :, 2] * b[:, :, :, 2]
+    #)
+
+def vfield3_matder(reader, a, b):
+    """Calculates material derivative of 3D vector fields a and b"""
+    dr = reader.get_fsgrid_cell_size()
+    bx = b[:, :, :, 0]
+    by = b[:, :, :, 1]
+    bz = b[:, :, :, 2]
+
+    grad_bx = np.gradient(bx, dr[0])
+    grad_by = np.gradient(by, dr[1])
+    grad_bz = np.gradient(bz, dr[2])
+
+    resx = vfield3_dot(a, grad_bx)
+    resy = vfield3_dot(a, grad_by)
+    resz = vfield3_dot(a, grad_bz)
+
+    return np.stack((resx, resy, resz), axis=-1)
+
+def vfield3_normalise(a):
+
+    amag = np.linalg.norm(a, axis=-1)
+
+    res=a/amag
+    return res
+    #resx = a[:, :, :, 0] / amag
+    #resy = a[:, :, :, 1] / amag
+    #resz = a[:, :, :, 2] / amag
+
+    #return np.stack((resx, resy, resz), axis=-1)
+
+def vfield3_curvature(reader, a):
+   dr = reader.get_fsgrid_cell_size()
+   a = vfield3_normalise(a)
+   return vfield3_matder(a, a, dr)
+
+def ballooning_crit(reader, B, P, beta):
+    dr = reader.get_fsgrid_cell_size()
+    n = vfield3_curvature(B, dr)
+
+    nnorm = vfield3_normalise(n)
+
+    gradP = np.gradient(P,dr)
+
+    kappaP = vfield3_dot(nnorm, gradP) / P
+
+    kappaC = vfield3_dot(nnorm, n)
+
+    balloon = (2 + beta) / 4.0 * kappaP / (kappaC + 1e-27)
+
+    return (balloon, kappaC, gradP)

pyPlots/plot_colormap.py

@@ -816,7 +816,7 @@ def exprMA_cust(exprmaps, requestvariables=False):
     if lin is None:
         # Special SymLogNorm case
         if symlog is not None:
-            if LooseVersion(matplotlib.__version__) < LooseVersion("3.3.0"):
+            if LooseVersion(matplotlib.__version__) < LooseVersion("3.2.0"):
                 norm = SymLogNorm(linthresh=linthresh, linscale = 1.0, vmin=vminuse, vmax=vmaxuse, clip=True)
                 print("WARNING: colormap SymLogNorm uses base-e but ticks are calculated with base-10.")
                 #TODO: copy over matplotlib 3.3.0 implementation of SymLogNorm into pytools/analysator

pyPlots/plot_helpers.py

@@ -661,6 +661,12 @@ def expr_flowcompression(pass_maps, requestvariables=False):
         return ['V']
     Vmap = TransposeVectorArray(pass_maps['V']) # Bulk flow
     return numdiv(Vmap).T
+
+def expr_divPoynting(pass_maps, requestvariables=False):
+    if requestvariables==True:
+        return ['poynting']
+    dSmap = TransposeVectorArray(pass_maps['poynting']) # Bulk flow
+    return numdiv(dSmap).T
 
 # def expr_gradB_aniso(pass_maps):
 #     Bmap = TransposeVectorArray(pass_maps['B']) # Magnetic field

pyVlsv/vlsvreader.py

@@ -1073,6 +1073,44 @@ def calcLocalSize(globalCells, ntasks, my_n):
 
        return np.squeeze(orderedData)
 
+   def read_fg_variable_as_volumetric(self, name, centering=None, operator="pass"):
+      fgdata = self.read_fsgrid_variable(name, operator)
+
+      fssize=list(self.get_fsgrid_mesh_size())
+      if 1 in fssize:
+         #expand to have a singleton dimension for a reduced dim - lets roll happen with ease
+         singletons = [i for i, sz in enumerate(fssize) if sz == 1]
+         for dim in singletons:
+            fgdata=np.expand_dims(fgdata, dim)
+      #print('read in fgdata with shape', fgdata.shape, name)
+      celldata = np.zeros_like(fgdata)
+      known_centerings = {"fg_b":"face", "fg_e":"edge"}
+      if centering is None:
+         try:
+            centering = known_centerings[name.lower()]
+         except KeyError:
+            print("A variable ("+name+") with unknown centering! Aborting.")
+            return False
+
+      #vector variable
+      if fgdata.shape[-1] == 3:
+         if centering=="face":
+            celldata[:,:,:,0] = (fgdata[:,:,:,0] + np.roll(fgdata[:,:,:,0],-1, 0))/2.0
+            celldata[:,:,:,1] = (fgdata[:,:,:,1] + np.roll(fgdata[:,:,:,1],-1, 1))/2.0
+            celldata[:,:,:,2] = (fgdata[:,:,:,2] + np.roll(fgdata[:,:,:,2],-1, 2))/2.0
+            # Use Leo's reconstuction for fg_b instead
+         elif centering=="edge":
+            celldata[:,:,:,0] = (fgdata[:,:,:,0] + np.roll(fgdata[:,:,:,0],-1, 1) + np.roll(fgdata[:,:,:,0],-1, 2) + np.roll(fgdata[:,:,:,0],-1, (1,2)))/4.0
+            celldata[:,:,:,1] = (fgdata[:,:,:,1] + np.roll(fgdata[:,:,:,1],-1, 0) + np.roll(fgdata[:,:,:,1],-1, 2) + np.roll(fgdata[:,:,:,1],-1, (0,2)))/4.0
+            celldata[:,:,:,2] = (fgdata[:,:,:,2] + np.roll(fgdata[:,:,:,2],-1, 0) + np.roll(fgdata[:,:,:,2],-1, 1) + np.roll(fgdata[:,:,:,2],-1, (0,1)))/4.0
+         else:
+            print("Unknown centering ('" +centering+ "')! Aborting.")
+            return False
+      else:
+         print("A scalar variable! I don't know what to do with this! Aborting.")
+         return False
+      return celldata
+
    def read_variable(self, name, cellids=-1,operator="pass"):
       ''' Read variables from the open vlsv file. 
       Arguments:
@@ -1202,6 +1240,120 @@ def get_amr_level(self,cellid):
          AMR_count += 1
       return AMR_count - 1 
 
+   def get_cell_dx(self, cellid):
+      '''Returns the dx of a given cell defined by its cellid
+
+      :param cellid:        The cell's cellid
+      :returns:             The cell's size [dx, dy, dz]
+      '''
+      return np.array([self.__dx,self.__dy,self.__dz])/2**self.get_amr_level(cellid)
+
+   def get_cell_bbox(self, cellid):
+      '''Returns the bounding box of a given cell defined by its cellid
+
+      :param cellid:        The cell's cellid
+      :returns:             The cell's bbox [xmin,ymin,zmin],[xmax,ymax,zmax]
+      '''
+
+      hdx = self.get_cell_dx(cellid)*0.5
+      mid = self.get_cell_coordinates(cellid)
+      #print('halfdx:', hdx, 'mid:', mid, 'low:', mid-hdx, 'hi:', mid+hdx)
+      return mid-hdx, mid+hdx
+
+   def get_cell_fsgrid_slicemap(self, cellid):
+      '''Returns a slice tuple of fsgrid indices that are contained in the SpatialGrid
+      cell.
+      '''
+      low, up = self.get_cell_bbox(cellid)
+      lowi, upi = self.get_fsgrid_slice_indices(low, up)
+      return lowi, upi
+
+   def get_bbox_fsgrid_slicemap(self, low, up):
+      '''Returns a slice tuple of fsgrid indices that are contained in the (low, up) bounding box.
+      '''
+      lowi, upi = self.get_fsgrid_slice_indices(low, up)
+      return lowi, upi
+
+   def get_cell_fsgrid_subarray(self, cellid, array):
+      '''Returns a subarray of the fsgrid array, corresponding to the fsgrid
+      covered by the SpatialGrid cellid. [untested]
+      '''
+      lowi, upi = self.get_cell_fsgrid_slicemap(cellid)
+      #print('subarray:',lowi, upi)
+      if array.ndim == 4:
+         return array[lowi[0]:upi[0]+1, lowi[1]:upi[1]+1, lowi[2]:upi[2]+1, :]
+      else:
+         return array[lowi[0]:upi[0]+1, lowi[1]:upi[1]+1, lowi[2]:upi[2]+1]
+
+   def get_bbox_fsgrid_subarray(self, low, up, array):
+      '''Returns a subarray of the fsgrid array, corresponding to the (low, up) bounding box.
+      '''
+      lowi, upi = self.get_bbox_fsgrid_slicemap(low,up)
+      #print('subarray:',lowi, upi)
+      if array.ndim == 4:
+         return array[lowi[0]:upi[0]+1, lowi[1]:upi[1]+1, lowi[2]:upi[2]+1, :]
+      else:
+         return array[lowi[0]:upi[0]+1, lowi[1]:upi[1]+1, lowi[2]:upi[2]+1]
+
+
+   def downsample_fsgrid_subarray(self, cellid, array):
+      '''Returns a mean value of fsgrid values underlying the SpatialGrid cellid.
+      '''
+      fsarr = self.get_cell_fsgrid_subarray(cellid, array)
+      n = fsarr.size
+      if fsarr.ndim == 4:
+         n = n/3
+      ncells = 8**(self.get_max_refinement_level()-self.get_amr_level(cellid))
+      if(n != ncells):
+         print("Warning: weird fs subarray size", n, 'for amrlevel', self.get_amr_level(cellid), 'expect', ncells)
+      return np.mean(fsarr,axis=(0,1,2))
+
+   def upsample_fsgrid_subarray(self, cellid, var, array):
+      '''Set the elements of the fsgrid array to the value of corresponding SpatialGrid
+      cellid. Mutator for array.
+      '''
+      lowi, upi = self.get_cell_fsgrid_slicemap(cellid)
+      #print(lowi,upi)
+      value = self.read_variable(var, cellids=[cellid])
+      if array.ndim == 4:
+         #print(value)
+         array[lowi[0]:upi[0]+1,lowi[1]:upi[1]+1,lowi[2]:upi[2]+1,:] = value
+         #print(array[lowi[0]:upi[0]+1,lowi[1]:upi[1]+1,lowi[2]:upi[2]+1,:])
+      else:
+         array[lowi[0]:upi[0]+1,lowi[1]:upi[1]+1,lowi[2]:upi[2]+1] = value
+      return
+
+   def read_variable_as_fg(self, var):
+      cellids = self.read_variable('CellID')
+      sz = self.get_fsgrid_mesh_size()
+      testvar = self.read_variable(var, [cellids[0]])
+      varsize = testvar.size
+      if(varsize > 1):
+         fgarr = np.zeros([sz[0], sz[1], sz[2], varsize], dtype=testvar.dtype)
+      else:
+         fgarr = np.zeros(sz, dtype=testvar.dtype)
+      for c in cellids:
+         self.upsample_fsgrid_subarray(c, var, fgarr)
+         #print
+      return fgarr
+
+
+   def get_cell_fsgrid(self, cellid):
+      '''Returns a slice tuple of fsgrid indices that are contained in the SpatialGrid
+      cell.
+      '''
+      low, up = self.get_cell_bbox(cellid)
+      lowi, upi = self.get_fsgrid_slice_indices(low, up)
+      return lowi, upi
+
+   def get_fsgrid_coordinates(self, ri):
+      '''Returns real-space center coordinates of the fsgrid 3-index.
+      '''
+      lowerlimit = self.get_fsgrid_mesh_extent()[0:3]
+      dxs = self.get_fsgrid_cell_size()
+
+      return lowerlimit+dxs*(np.array(ri)+0.5)
+
    def get_unique_cellids(self, coords):
       ''' Returns a list of cellids containing all the coordinates in coords,
           with no duplicate cellids. Relative order of elements is conserved.
@@ -1799,6 +1951,52 @@ def get_fsgrid_mesh_extent(self):
       '''
       return np.array([self.__xmin, self.__ymin, self.__zmin, self.__xmax, self.__ymax, self.__zmax])
 
+   def get_fsgrid_cell_size(self):
+      ''' Read fsgrid cell size
+
+      :returns: Maximum and minimum coordinates of the mesh, [dx, dy, dz]
+      '''
+      size = self.get_fsgrid_mesh_size()
+      ext = self.get_fsgrid_mesh_extent()
+      ext = ext[3:6]-ext[0:3]
+      return ext/size
+
+   def get_fsgrid_indices(self, coords):
+      ''' Convert spatial coordinates coords to an index array [xi, yi, zi] for fsgrid
+
+      :returns 3-tuple of integers [xi, yi, zi] corresponding to fsgrid cell containing coords (low-inclusive)
+      Example:
+      ii = f.get_fsgrid_mesh_extent(coords)
+      fsvar_at_coords = fsvar_array.item(ii)
+      '''
+      lower = self.get_fsgrid_mesh_extent()[0:3]
+      dx = self.get_fsgrid_cell_size()
+      r0 = coords-lower
+      ri = np.floor(r0/dx).astype(int)
+      if (ri < 0).any() or (ri>sz-1).any():
+         print("get_fsgrid_indices: Resulting index out of bounds, returning None")
+         return None
+      return tuple(ri)
+
+   def get_fsgrid_slice_indices(self, lower, upper, eps=1e-3):
+      ''' Get indices for a subarray of an fsgrid variable, in the cuboid from lower to upper.
+      This is meant for mapping a set of fsgrid cells to a given SpatialGrid cell.
+      Shifts the corners (lower, upper) by dx_fsgrid*eps inward, if direct low-inclusive behaviour
+      is required, set kword eps = 0.
+
+
+      :returns two 3-tuples of integers.
+      Example:
+      ii = f.get_fsgrid_mesh_extent(coords)
+      fsvar_at_coords = fsvar_array.item(ii)
+      '''
+      dx = self.get_fsgrid_cell_size()
+      eps = dx*eps
+      loweri = self.get_fsgrid_indices(lower+eps)
+      upperi = self.get_fsgrid_indices(upper-eps)
+      return loweri, upperi
+
+
    def get_velocity_mesh_size(self, pop="proton"):
       ''' Read velocity mesh size