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)