398 lines
15 KiB
Python
398 lines
15 KiB
Python
|
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)
|