diff --git a/pyCalculations/calculations.py b/pyCalculations/calculations.py
index 483f8bc7..cb114147 100644
--- a/pyCalculations/calculations.py
+++ b/pyCalculations/calculations.py
@@ -56,4 +56,4 @@
import fit
from fieldtracer import static_field_tracer
from fieldtracer import dynamic_field_tracer
-
+from differentialops import fg_divPoynting, fg_PoyntingFlux, fg_vol_curl, ballooning_crit,vfield3_curvature
diff --git a/pyCalculations/differentialops.py b/pyCalculations/differentialops.py
new file mode 100644
index 00000000..93f8dc1a
--- /dev/null
+++ b/pyCalculations/differentialops.py
@@ -0,0 +1,144 @@
+#
+# 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)
+
+ return (kappaC, gradP)
\ No newline at end of file
diff --git a/pyPlots/plot_colormap.py b/pyPlots/plot_colormap.py
index fd78700a..8f92e936 100644
--- a/pyPlots/plot_colormap.py
+++ b/pyPlots/plot_colormap.py
@@ -249,7 +249,11 @@ def exprMA_cust(exprmaps, requestvariables=False):
if filename:
f=pt.vlsvfile.VlsvReader(filename)
elif (filedir and step is not None):
- filename = glob.glob(filedir+'bulk*'+str(step).rjust(7,'0')+'.vlsv')[0]
+ try:
+ filename = glob.glob(filedir+'bulk*'+str(step).rjust(7,'0')+'.vlsv')[0]
+ except IndexError as err:
+ print("IndexError for filedir =", filedir, ", step = ", step)
+ return
#filename = filedir+'bulk.'+str(step).rjust(7,'0')+'.vlsv'
f=pt.vlsvfile.VlsvReader(filename)
elif vlsvobj:
@@ -816,7 +820,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
diff --git a/pyPlots/plot_helpers.py b/pyPlots/plot_helpers.py
index 2a22b4c4..39edf51c 100644
--- a/pyPlots/plot_helpers.py
+++ b/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
diff --git a/pyVlsv/vlsvreader.py b/pyVlsv/vlsvreader.py
index fab118b6..6e11a791 100644
--- a/pyVlsv/vlsvreader.py
+++ b/pyVlsv/vlsvreader.py
@@ -25,7 +25,7 @@
import xml.etree.ElementTree as ET
import ast
import numpy as np
-import os
+import os,sys
import numbers
import vlsvvariables
from reduction import datareducers,multipopdatareducers,data_operators,v5reducers,multipopv5reducers
@@ -54,7 +54,11 @@ def __init__(self, file_name):
file_name = os.path.abspath(file_name)
self.file_name = file_name
- self.__fptr = open(self.file_name,"rb")
+ try:
+ self.__fptr = open(self.file_name,"rb")
+ except FileNotFoundError as e:
+ print("File not found: ", self.file_name)
+ raise e
self.__xml_root = ET.fromstring("")
self.__fileindex_for_cellid={}
self.__max_spatial_amr_level = -1
@@ -671,7 +675,7 @@ def read(self, name="", tag="", mesh="", operator="pass", cellids=-1):
reducer_multipop = multipopdatareducers
# If this is a variable that can be summed over the populations (Ex. rho, PTensorDiagonal, ...)
- if self.check_variable(self.active_populations[0]+'/'+name):
+ if len(self.active_populations) > 0 and self.check_variable(self.active_populations[0]+'/'+name):
tmp_vars = []
for pname in self.active_populations:
vlsvvariables.activepopulation = pname
@@ -975,6 +979,22 @@ def read_interpolated_variable(self, name, coordinates, operator="pass",periodic
# Done.
return ret_array
+ def read_fsgrid_variable_cellid(self, name, cellids=-1, operator="pass"):
+ ''' Reads fsgrid variables from the open vlsv file.
+ Arguments:
+ :param name: Name of the variable
+ :param cellids: SpatialGrid cellids for which to fetch data. Default: return full fsgrid data
+ :param operator: Datareduction operator. "pass" does no operation on data
+ :returns: *ordered* list of numpy arrays with the data
+
+ ... seealso:: :func:`read_fsgrid_variable`
+ '''
+ var = self.read_fsgrid_variable(name, operator=operator)
+ if cellids == -1:
+ return var
+ else:
+ return [self.downsample_fsgrid_subarray(cid, var) for cid in cellids]
+
def read_fsgrid_variable(self, name, operator="pass"):
''' Reads fsgrid variables from the open vlsv file.
Arguments:
@@ -1073,6 +1093,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 +1260,154 @@ 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.
+ '''
+ 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 fsgrid_array_to_vg(self, array):
+ cellIds=self.read_variable("CellID")
+
+ #if not hasattr(self, 'fsCellIdTargets'):
+ # self.fsCellIdTargets = np.zeros(self.get_fsgrid_mesh_size())
+ # for cid in cellIds:
+ # lowi, upi = self.get_cell_fsgrid_slicemap(cid)
+ # self.fsCellIdTargets[lowi[0]:upi[0]+1, lowi[1]:upi[1]+1, lowi[2]:upi[2]+1] = cid
+
+ vgarr = [self.downsample_fsgrid_subarray(cid, array) for cid in cellIds]
+ return vgarr
+
+ def vg_uniform_grid_process(self, variable, expr, exprtuple):
+ cellIds=self.read_variable("CellID")
+ array = self.read_variable_as_fg(variable)
+ array = expr(*exprtuple)
+ return self.fsgrid_array_to_vg(array)
+
+ def get_cellid_at_fsgrid_index(self, i,j,k):
+ coords = self.get_fsgrid_coordinates([i,j,k])
+ return self.get_cellid(coords)
+
+ 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):
+ vg_cellids = self.read_variable('CellID')
+ sz = self.get_fsgrid_mesh_size()
+ sz_amr = self.get_spatial_mesh_size()
+ vg_var = self.read_variable(var)
+ varsize = vg_var[0].size
+ if(varsize > 1):
+ fg_var = np.zeros([sz[0], sz[1], sz[2], varsize], dtype=vg_var.dtype)
+ else:
+ fg_var = np.zeros(sz, dtype=vg_var.dtype)
+ vg_cellids_on_fg = np.zeros(sz, dtype=np.int64) + 1000000000 # big number to catch errors in the latter code, 0 is not good for that
+ current_amr_level = self.get_amr_level(np.min(vg_cellids))
+ max_amr_level = int(np.log2(sz[0] / sz_amr[0]))
+ current_max_cellid = 0
+ for level in range(current_amr_level+1):
+ current_max_cellid += 2**(3*(level))*(self.__xcells*self.__ycells*self.__zcells)
+ for id in np.argsort(vg_cellids):
+ if vg_cellids[id] > current_max_cellid:
+ current_amr_level += 1
+ current_max_cellid += 2**(3*(current_amr_level))*(self.__xcells*self.__ycells*self.__zcells)
+ this_cell_indices = np.array(self.get_cell_indices(vg_cellids[id], current_amr_level), dtype=np.int64)
+ refined_ids_start = this_cell_indices * 2**(max_amr_level-current_amr_level)
+ refined_ids_end = refined_ids_start + 2**(max_amr_level-current_amr_level)
+ vg_cellids_on_fg[refined_ids_start[0]:refined_ids_end[0],refined_ids_start[1]:refined_ids_end[1],refined_ids_start[2]:refined_ids_end[2]] = id
+ fg_var = vg_var[vg_cellids_on_fg]
+ return fg_var
+
+ 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 +2005,53 @@ 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)
+ sz = self.get_fsgrid_mesh_size()
+ 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
diff --git a/pyVlsv/vlsvwriter.py b/pyVlsv/vlsvwriter.py
index 222d7103..7f2c2fb5 100644
--- a/pyVlsv/vlsvwriter.py
+++ b/pyVlsv/vlsvwriter.py
@@ -41,8 +41,11 @@ def __init__(self, vlsvReader, file_name ):
:param file_name: Name of the vlsv file where to input data
'''
self.file_name = os.path.abspath(file_name)
- self.__fptr = open(self.file_name,"wb")
-
+ try:
+ self.__fptr = open(self.file_name,"wb")
+ except FileNotFoundError as e:
+ print("No such path: ", self.file_name)
+ raise e
self.__xml_root = ET.fromstring("")
self.__fileindex_for_cellid={}
@@ -69,6 +72,8 @@ def __initialize( self, vlsvReader ):
tags['MESH_NODE_CRDS_Z'] = ''
tags['MESH'] = ''
tags['MESH_DOMAIN_SIZES'] = ''
+ tags['MESH_GHOST_DOMAINS'] = ''
+ tags['MESH_GHOST_LOCALIDS'] = ''
tags['CellID'] = ''
tags['MESH_BBOX'] = ''
tags['COORDS'] = ''
@@ -79,7 +84,7 @@ def __initialize( self, vlsvReader ):
if 'name' in child.attrib: name = child.attrib['name']
else: name = ''
if 'mesh' in child.attrib: mesh = child.attrib['mesh']
- else: mesh = ''
+ else: mesh = None
tag = child.tag
extra_attribs = {}
for i in child.attrib.items():
@@ -90,10 +95,15 @@ def __initialize( self, vlsvReader ):
self.write( data=data, name=name, tag=tag, mesh=mesh, extra_attribs=extra_attribs )
- def copy_variables( self, vlsvReader ):
- ''' Copies all variables from vlsv reader to the file
-
+ def copy_variables( self, vlsvReader, varlist=None ):
+ ''' Copies variables from vlsv reader to the file.
+ varlist = None: list of variables to copy;
'''
+
+ # Delegate to the variable list handler
+ if (varlist is not None):
+ self.copy_variables_list(vlsvReader, varlist)
+
# Get the xml sheet:
xml_root = vlsvReader._VlsvReader__xml_root
@@ -123,6 +133,44 @@ def copy_variables( self, vlsvReader ):
self.write( data=data, name=name, tag=tag, mesh=mesh, extra_attribs=extra_attribs )
return
+ def copy_variables_list( self, vlsvReader, vars ):
+ ''' Copies variables in the list vars from vlsv reader to the file
+
+ '''
+ # Get the xml sheet:
+ xml_root = vlsvReader._VlsvReader__xml_root
+
+ # Get list of tags to write:
+ tags = {}
+ tags['VARIABLE'] = ''
+
+ # Copy the xml root and write variables
+ for child in xml_root:
+ if child.tag in tags:
+ if 'name' in child.attrib:
+ name = child.attrib['name']
+ if not name in vars:
+ continue
+ else:
+ continue
+ if 'mesh' in child.attrib:
+ mesh = child.attrib['mesh']
+ else:
+ if tag in ['VARIABLE']:
+ print('MESH required')
+ return
+ mesh = None
+ tag = child.tag
+ # Copy extra attributes:
+ extra_attribs = {}
+ for i in child.attrib.items():
+ if i[0] != 'name' and i[0] != 'mesh':
+ extra_attribs[i[0]] = i[1]
+ data = vlsvReader.read( name=name, tag=tag, mesh=mesh )
+ # Write the data:
+ self.write( data=data, name=name, tag=tag, mesh=mesh, extra_attribs=extra_attribs )
+ return
+
def write_velocity_space( self, vlsvReader, cellid, blocks_and_values ):
''' Writes given velocity space into vlsv file
@@ -162,9 +210,10 @@ def write(self, data, name, tag, mesh, extra_attribs={}):
datatype = ''
# Add the data into the xml data:
- child = ET.SubElement(parent=self.__xml_root, tag=tag)
+ child = ET.SubElement(self.__xml_root, tag)
child.attrib["name"] = name
- child.attrib["mesh"] = mesh
+ if mesh is not None:
+ child.attrib["mesh"] = mesh
child.attrib["arraysize"] = len(np.atleast_1d(data))
if extra_attribs != '':
for i in extra_attribs.items():
@@ -177,7 +226,14 @@ def write(self, data, name, tag, mesh, extra_attribs={}):
return False
else:
child.attrib["vectorsize"] = 1
- datatype = str(type(data[0]))
+ if(len(data) == 0):
+ if tag=="MESH_GHOST_DOMAINS" or tag=="MESH_GHOST_LOCALIDS":
+ datatype="int32"
+ else:
+ print("Trying to extract datatype from an empty array. I will fail as usual, since this is not the special case that is guarded against!")
+ datatype = str(type(data[0]))
+ else:
+ datatype = str(type(data[0]))
# Parse the data types:
if 'uint' in datatype:
@@ -187,7 +243,7 @@ def write(self, data, name, tag, mesh, extra_attribs={}):
elif 'float' in datatype:
child.attrib["datatype"] = "float"
else:
- print("BAD DATATYPE")
+ print("BAD DATATYPE: " + datatype)
return False
if '64' in datatype:
@@ -207,6 +263,34 @@ def write(self, data, name, tag, mesh, extra_attribs={}):
# write the xml footer:
self.__write_xml_footer()
+ def write_variable(self, data, name, mesh, variableLaTex, unit, unitLaTeX, unitConversion, extra_attribs={}):
+ ''' Writes an array into the vlsv file as a variable; requires input of metadata required by VlsvReader
+
+ :param name: Name of the data array
+ :param mesh: Mesh for the data array
+ :param variableLaTeX: LaTeX string representation of variable
+ :param unit: plaintext string represenation of the unit
+ :param unitLaTeX: LaTeX string represenation of the unit
+ :param unitConversion: string represenation of the unit conversion to get to SI
+ :param extra_attribs: Dictionary with whatever xml attributes that should be defined in the array that aren't name, tag, or mesh
+
+ :returns: True if the data was written successfully
+
+ '''
+
+ return self.write(data, name, 'VARIABLE', mesh, extra_attribs={'variableLaTeX':variableLaTex, 'unit':unit, 'unitLaTeX':unitLaTeX, 'unitConversion':unitConversion})
+
+
+ def write_fgarray_to_SpatialGrid(self, reader, data, name, extra_attribs={}):
+ # get a reader for the target file
+ #print(data.shape[0:3], reader.get_fsgrid_mesh_size(), (data.shape[0:3] == reader.get_fsgrid_mesh_size()))
+ if not (data.shape[0:3] == reader.get_fsgrid_mesh_size()).all():
+ print("Data shape does not match target fsgrid mesh")
+ return
+ cids = reader.read_variable("CellID")
+ vgdata = [reader.downsample_fsgrid_subarray(cid, data) for cid in cids]
+ self.write(vgdata, name, "VARIABLE", "SpatialGrid",extra_attribs)
+
def __write_xml_footer( self ):
# Write the xml footer:
max_xml_size = 1000000
@@ -221,6 +305,7 @@ def __write_xml_footer( self ):
for i in child.attrib.items():
child.attrib[i[0]] = str(child.attrib[i[0]])
tree = ET.ElementTree( self.__xml_root)
+ xml_footer_indent( self.__xml_root)
tree.write(fptr)
# Write the offset (first 8 bytes = endianness):
offset_endianness = 8
@@ -234,3 +319,17 @@ def close( self ):
self.__write_xml_footer()
self.__fptr.close()
+def xml_footer_indent(elem, level=0):
+ i = "\n" + level*" "
+ if len(elem):
+ if not elem.text or not elem.text.strip():
+ elem.text = i + " "
+ if not elem.tail or not elem.tail.strip():
+ elem.tail = i
+ for elem in elem:
+ xml_footer_indent(elem, level+1)
+ if not elem.tail or not elem.tail.strip():
+ elem.tail = i
+ else:
+ if level and (not elem.tail or not elem.tail.strip()):
+ elem.tail = i
\ No newline at end of file