init
This commit is contained in:
@@ -0,0 +1,397 @@
|
||||
import cftime
|
||||
import logging
|
||||
import numpy as np
|
||||
from os.path import basename, splitext, exists
|
||||
import xarray as xr
|
||||
from vtkmodules.vtkCommonCore import (
|
||||
vtkVariant,
|
||||
)
|
||||
from vtkmodules.vtkCommonDataModel import (
|
||||
vtkDataObject
|
||||
)
|
||||
from vtkmodules.vtkCommonExecutionModel import (
|
||||
vtkAlgorithm,
|
||||
vtkStreamingDemandDrivenPipeline
|
||||
)
|
||||
from vtkmodules.vtkIONetCDF import vtkNetCDFCFReader, vtkXArrayAccessor
|
||||
from vtkmodules.util import numpy_support
|
||||
from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase
|
||||
|
||||
@xr.register_dataset_accessor("vtk")
|
||||
class VtkAccessor:
|
||||
def __init__(self, dsxr):
|
||||
self._dsxr = dsxr
|
||||
|
||||
def create_reader(self):
|
||||
'''
|
||||
Returns a vtkXArrayCFReader that reads data from the XArray
|
||||
(using zero-copy when possible). At the moment, data is copied
|
||||
for coordinates (because they are converted to double in the reader)
|
||||
and for certain data that is subset either in XArray or in VTK.
|
||||
Lazy loading in XArray is respected, that is data is accessed only when
|
||||
it is needed.
|
||||
Time is passed to VTK either as an int64 for datetime64 or timedelta64,
|
||||
or as a double (using cftime.toordinal) for cftime.
|
||||
'''
|
||||
reader = vtkXArrayCFReader()
|
||||
reader.SetXArray(self._dsxr)
|
||||
return reader
|
||||
|
||||
|
||||
class vtkXArrayCFReader(VTKPythonAlgorithmBase):
|
||||
'''Reads data from a file using the XArray readers and then connects
|
||||
the XArray data to the vtkNetCDFCFREader (using zero-copy when
|
||||
possible). At the moment, data is copied for coordinates (because
|
||||
they are converted to double in the reader) and for certain data
|
||||
that is subset either in XArray or in VTK. Lazy loading in XArray
|
||||
is respected, that is data is accessed only when it is needed.
|
||||
Time is passed to VTK either as an int64 for datetime64 or
|
||||
timedelta64, or as a double (using cftime.toordinal) for cftime.
|
||||
'''
|
||||
|
||||
_FORWARD_GET = {
|
||||
"GetAccessor",
|
||||
"GetAllDimensions",
|
||||
|
||||
"GetNumberOfVariableArrays",
|
||||
"GetAllVariableArrayNames",
|
||||
"GetVariableArrayName",
|
||||
"GetVariableArrayStatus",
|
||||
|
||||
"GetTimeDimensionName",
|
||||
"GetLatitudeDimensionName",
|
||||
"GetLongitudeDimensionName",
|
||||
"GetVerticalDimensionName",
|
||||
|
||||
"GetOutput",
|
||||
"GetOutputType",
|
||||
"GetSphericalCoordinates",
|
||||
|
||||
"GetReplaceFillValueWithNan",
|
||||
|
||||
"GetVariableDimensions",
|
||||
"GetVerticalBias",
|
||||
"GetVerticalScale",
|
||||
"PrintSelf",
|
||||
}
|
||||
_FORWARD_SET = {
|
||||
"SetDimensions",
|
||||
|
||||
"SetTimeDimensionName",
|
||||
"SetLatitudeDimensionName",
|
||||
"SetLongitudeDimensionName",
|
||||
"SetVerticalDimensionName",
|
||||
|
||||
|
||||
"SetSphericalCoordinates",
|
||||
"SphericalCoordinatesOn",
|
||||
"SphericalCoordinatesOff",
|
||||
|
||||
"SetReplaceFillValueWithNan",
|
||||
"ReplaceFillValueWithNanOn",
|
||||
"ReplaceFillValueWithNanOff",
|
||||
|
||||
"SetOutputType",
|
||||
"SetOutputTypeToAutomatic",
|
||||
"SetOutputTypeToImage",
|
||||
"SetOutputTypeToRectilinear",
|
||||
"SetOutputTypeToStructured",
|
||||
"SetOutputTypeToUnstructured",
|
||||
|
||||
"SetVariableArrayStatus",
|
||||
"SetVerticalBias",
|
||||
"SetVerticalScale",
|
||||
"UpdateMetaData",
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
VTKPythonAlgorithmBase.__init__(
|
||||
self, nInputPorts=0, nOutputPorts=1, outputType="vtkDataObject"
|
||||
)
|
||||
self._log = logging.getLogger("vtkXArrayCFReader")
|
||||
self._filename = None
|
||||
self._timesteps = None
|
||||
self._timeindex = None
|
||||
self._node = None
|
||||
self._dsxr = None
|
||||
self._reader = vtkNetCDFCFReader()
|
||||
self._ndarray_cftime_toordinal = np.frompyfunc(vtkXArrayCFReader._cftime_toordinal, 1, 1)
|
||||
# reference to contiguous arrays so that they are not dealocated
|
||||
self._arrays = {}
|
||||
|
||||
|
||||
def __getattr__(self, name):
|
||||
in_set = name in self._FORWARD_SET
|
||||
in_get = name in self._FORWARD_GET
|
||||
if in_set or in_get:
|
||||
if in_set:
|
||||
self.Modified()
|
||||
return getattr(self._reader, name)
|
||||
else:
|
||||
raise AttributeError()
|
||||
|
||||
def SetFileName(self, name):
|
||||
"""Specify filename for the file to read."""
|
||||
if self._filename != name:
|
||||
self._filename = name
|
||||
self.Modified()
|
||||
|
||||
def GetFileName(self):
|
||||
return self._filename
|
||||
|
||||
def CanReadFile(self, filepath):
|
||||
ext = splitext(filepath)[1]
|
||||
filename = basename(filepath)
|
||||
correct_name = False
|
||||
if ext == '.nc' or ext == '.grib' or ext == '.h5':
|
||||
correct_name = True
|
||||
else:
|
||||
if ext == '' and filename == '.zgroup':
|
||||
correct_name = True
|
||||
if correct_name and exists(filepath):
|
||||
return 1
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def SetNode(self, node):
|
||||
if self._node != node:
|
||||
self._node = node
|
||||
self.Modified()
|
||||
|
||||
def GetNode(self):
|
||||
return self._node
|
||||
|
||||
def SetXArray(self, dsxr):
|
||||
self._dsxr = dsxr
|
||||
self._update_accessor()
|
||||
self.Modified()
|
||||
|
||||
def GetXArray(self):
|
||||
return self._dsxr
|
||||
|
||||
def RequestDataObject(self, request, inInfo, outInfo):
|
||||
self._log.debug(f"DataObject ======================================================================")
|
||||
if not self._dsxr:
|
||||
if self._node:
|
||||
tree = xr.open_datatree(self._filename)
|
||||
self._dsxr = tree[self._node].to_dataset()
|
||||
else:
|
||||
self._dsxr = xr.open_dataset(self._filename, decode_timedelta=True)
|
||||
self._update_accessor()
|
||||
self._reader.UpdateDataObject()
|
||||
roi = self._reader.GetOutputInformation(0)
|
||||
if roi.Has(vtkDataObject.DATA_OBJECT()):
|
||||
rdata = roi.Get(vtkDataObject.DATA_OBJECT())
|
||||
else:
|
||||
self._log.error("vtkNetCDFCFReader did not create the dataset")
|
||||
rdata = None
|
||||
oi = outInfo.GetInformationObject(0)
|
||||
oi.Set(vtkDataObject.DATA_OBJECT(), rdata)
|
||||
return 1
|
||||
|
||||
def RequestInformation(self, request, inInfo, outInfo):
|
||||
self._log.debug(f"Information ======================================================================")
|
||||
oi = outInfo.GetInformationObject(0)
|
||||
self._reader.UpdateInformation()
|
||||
roi = self._reader.GetOutputInformation(0)
|
||||
if roi.Has(vtkStreamingDemandDrivenPipeline.TIME_STEPS()):
|
||||
self._timesteps = roi.Get(vtkStreamingDemandDrivenPipeline.TIME_STEPS())
|
||||
oi.Set(vtkStreamingDemandDrivenPipeline.TIME_STEPS(), self._timesteps, len(self._timesteps))
|
||||
oi.Set(vtkStreamingDemandDrivenPipeline.TIME_RANGE(), [self._timesteps[0], self._timesteps[-1]], 2)
|
||||
self._timesteps = np.asarray(self._timesteps)
|
||||
if roi.Has(vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT()):
|
||||
ext = roi.Get(vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT())
|
||||
self._log.debug("Whole extent: {}".format(ext))
|
||||
oi.Set(vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), ext, 6)
|
||||
if roi.Has(vtkAlgorithm.CAN_HANDLE_PIECE_REQUEST()):
|
||||
oi.Set(vtkAlgorithm.CAN_HANDLE_PIECE_REQUEST(), 1)
|
||||
if roi.Has(vtkAlgorithm.CAN_PRODUCE_SUB_EXTENT()):
|
||||
oi.Set(vtkAlgorithm.CAN_PRODUCE_SUB_EXTENT(), 1)
|
||||
return 1
|
||||
|
||||
def RequestUpdateExtent(self, request, inInfo, outInfo):
|
||||
self._log.debug(f"UpdateExtent ======================================================================")
|
||||
oi = outInfo.GetInformationObject(0)
|
||||
if oi.Has(vtkStreamingDemandDrivenPipeline.UPDATE_TIME_STEP()):
|
||||
utime = oi.Get(vtkStreamingDemandDrivenPipeline.UPDATE_TIME_STEP())
|
||||
timeindex = (np.abs(self._timesteps - utime)).argmin()
|
||||
if timeindex != self._timeindex:
|
||||
self._log.debug(f"Time index = {timeindex}")
|
||||
self._timeindex = timeindex
|
||||
self.Modified()
|
||||
if oi.Has(vtkStreamingDemandDrivenPipeline.UPDATE_EXTENT()):
|
||||
ext = [0, 0, 0, 0, 0, 0]
|
||||
oi.Get(vtkStreamingDemandDrivenPipeline.UPDATE_EXTENT(), ext)
|
||||
self._log.debug("Update extent: {}".format(ext))
|
||||
roi = self._reader.GetOutputInformation(0)
|
||||
roi.Set(vtkStreamingDemandDrivenPipeline.UPDATE_EXTENT(), ext, 6)
|
||||
self._reader.PropagateUpdateExtent()
|
||||
return 1
|
||||
|
||||
|
||||
def RequestData(self, request, inInfo, outInfo):
|
||||
self._log.debug(f"Data ======================================================================")
|
||||
if self._timeindex:
|
||||
dsxr = self._dsxr.isel({self.GetTimeDimensionName() : self._timeindex})
|
||||
else:
|
||||
# no time, so no aditional selection is needed
|
||||
dsxr = self._dsxr
|
||||
accessor = self._reader.GetAccessor()
|
||||
self._set_data_vars(accessor, dsxr)
|
||||
self._reader.Update()
|
||||
# self._reader's data is already set for this's data so no ShallowCopy is needed
|
||||
return 1
|
||||
|
||||
@staticmethod
|
||||
def _get_nc_type(numpy_array_type):
|
||||
"""Returns a nc_type given a numpy array."""
|
||||
NC_BYTE = 1 # 1 byte integer
|
||||
NC_CHAR = 2 # iso/ascii character
|
||||
NC_SHORT = 3 # 2 byte integer
|
||||
NC_INT = 4 # 4 byte integer
|
||||
NC_LONG = NC_INT
|
||||
NC_FLOAT = 5
|
||||
NC_DOUBLE = 6
|
||||
NC_UBYTE = 7
|
||||
NC_USHORT = 8
|
||||
NC_UINT = 9
|
||||
NC_INT64 = 10 # 8 bypte integer
|
||||
NC_UINT64 = 11
|
||||
NC_STRING = 12
|
||||
_np_nc = {
|
||||
np.uint8: NC_UBYTE,
|
||||
np.uint16: NC_USHORT,
|
||||
np.uint32: NC_UINT,
|
||||
np.uint64: NC_UINT64,
|
||||
np.int8: NC_BYTE,
|
||||
np.int16: NC_SHORT,
|
||||
np.int32: NC_INT,
|
||||
np.int64: NC_INT64,
|
||||
np.float32: NC_FLOAT,
|
||||
np.float64: NC_DOUBLE,
|
||||
np.datetime64: NC_INT64,
|
||||
np.timedelta64: NC_INT64,
|
||||
np.str_: NC_STRING,
|
||||
np.bytes_: NC_CHAR,
|
||||
}
|
||||
for key, nc_type in _np_nc.items():
|
||||
if (
|
||||
numpy_array_type == key
|
||||
or np.issubdtype(numpy_array_type, key)
|
||||
or numpy_array_type == np.dtype(key)
|
||||
):
|
||||
return nc_type
|
||||
raise TypeError(
|
||||
"Could not find a suitable NetCDF type for %s" % (str(numpy_array_type))
|
||||
)
|
||||
|
||||
def _update_accessor(self):
|
||||
accessor, timename = self._get_accessor()
|
||||
self._reader.SetAccessor(accessor)
|
||||
if timename:
|
||||
self._reader.SetTimeDimensionName(timename)
|
||||
|
||||
def _get_accessor(self):
|
||||
acclog = logging.getLogger("_get_accessor_")
|
||||
acclog.setLevel(logging.WARNING)
|
||||
accessor = vtkXArrayAccessor()
|
||||
time_name = None
|
||||
time_names = []
|
||||
# Set Dim and DimLen
|
||||
dimNameToIndex = {k: i for i, k in enumerate(self._dsxr.sizes.keys())}
|
||||
accessor.SetDim(list(self._dsxr.sizes.keys()))
|
||||
accessor.SetDimLen(list(self._dsxr.sizes.values()))
|
||||
|
||||
# Set Var
|
||||
varList = list(self._dsxr.data_vars.keys()) + list(self._dsxr.coords.keys())
|
||||
varNameToIndex = {k: i for i, k in enumerate(varList)}
|
||||
is_coord = [0] * len(self._dsxr.data_vars)
|
||||
is_coord = is_coord + [1] * len(self._dsxr.coords)
|
||||
coords_bounds = self._get_coords_bounds()
|
||||
accessor.SetVar(varList, is_coord)
|
||||
for i, v in enumerate(varList):
|
||||
# data_vars are set after array selection and time selection to
|
||||
# take advantage of xarray lazy loading
|
||||
# https://docs.xarray.dev/en/latest/internals/internal-design.html
|
||||
if is_coord[i] or v in coords_bounds:
|
||||
# if there is subsetting in xarray, self._dsxr[v].values is
|
||||
# not contiguous. If the array is not contiguous, a contiguous
|
||||
# copy is created otherwise the contiguous array is simply returned
|
||||
v_data = np.ascontiguousarray(self._dsxr[v].values)
|
||||
if (
|
||||
v_data.dtype.type == np.datetime64
|
||||
or v_data.dtype.type == np.timedelta64
|
||||
):
|
||||
un = np.datetime_data(v_data.dtype)
|
||||
# unit = ns and 1 base unit
|
||||
if un[0] == "ns" and un[1] == 1:
|
||||
time_names.append(v)
|
||||
if v_data.dtype.char == "O":
|
||||
# object array, assume cftime
|
||||
# copy cftime array to a doubles array
|
||||
self._arrays[v] = self._ndarray_cftime_toordinal(v_data).astype(np.float64)
|
||||
time_names.append(v)
|
||||
v_data = self._arrays[v]
|
||||
else:
|
||||
self._arrays[v] = v_data
|
||||
acclog.debug(f"{v=} {v_data.shape=} {v_data.dtype} {self._dsxr[v].dims=}")
|
||||
acclog.debug(f"address:{hex(v_data.ctypes.data)}")
|
||||
accessor.SetVarValue(i, v_data)
|
||||
accessor.SetVarType(i, vtkXArrayCFReader._get_nc_type(v_data.dtype))
|
||||
else:
|
||||
accessor.SetVarType(i, vtkXArrayCFReader._get_nc_type(self._dsxr[v].variable.dtype))
|
||||
accessor.SetVarDims(i, [dimNameToIndex[name] for name in self._dsxr[v].dims])
|
||||
accessor.SetVarCoords(
|
||||
i, [varNameToIndex[name] for name in self._dsxr[v].coords]
|
||||
)
|
||||
|
||||
acclog.debug("Attributes:")
|
||||
for item in self._dsxr[v].attrs.items():
|
||||
acclog.debug(
|
||||
"name: {} value: {} type: {}".format(
|
||||
item[0], item[1], type(item[1])
|
||||
)
|
||||
)
|
||||
if np.issubdtype(type(item[1]), np.integer):
|
||||
accessor.SetAtt(i, item[0], vtkVariant(int(item[1])))
|
||||
elif np.issubdtype(type(item[1]), np.floating):
|
||||
accessor.SetAtt(i, item[0], vtkVariant(float(item[1])))
|
||||
elif isinstance(item[1], np.ndarray):
|
||||
accessor.SetAtt(
|
||||
i, item[0], vtkVariant(numpy_support.numpy_to_vtk(item[1]))
|
||||
)
|
||||
else:
|
||||
accessor.SetAtt(i, item[0], vtkVariant(item[1]))
|
||||
if len(time_names) >= 1:
|
||||
for name in time_names:
|
||||
if accessor.IsCOARDSCoordinate(name):
|
||||
time_name = name
|
||||
break
|
||||
return accessor, time_name
|
||||
|
||||
def _set_data_vars(self, accessor, dsxr):
|
||||
# data_vars are listed first in the list of data_vars,coords so we don't
|
||||
# need to add coords to the list, and still get the corect indexes
|
||||
varList = list(dsxr.data_vars.keys())
|
||||
for i, v in enumerate(varList):
|
||||
if self._reader.GetVariableArrayStatus(v):
|
||||
v_data = np.ascontiguousarray(dsxr[v].values)
|
||||
accessor.SetVarValue(i, v_data)
|
||||
self._arrays[v] = v_data
|
||||
|
||||
def _get_coords_bounds(self):
|
||||
'''
|
||||
Special data_vars associated coords
|
||||
'''
|
||||
b=set()
|
||||
for coord in list(self._dsxr.coords):
|
||||
bounds_attr = 'bounds'
|
||||
if bounds_attr in self._dsxr[coord].attrs:
|
||||
b.add(self._dsxr[coord].attrs[bounds_attr])
|
||||
return b
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _cftime_toordinal(o):
|
||||
return o.toordinal(fractional=True)
|
||||
Reference in New Issue
Block a user