ioda.py

Go to the documentation of this file.

import ioda._ioda_python as _iodapy
import datetime as dt
import numpy as np
import os
 
class ObsSpace:
 
    def __repr__(self):
        return f"ObsSpace({self.name},{self.iodafile})"
 
    def __str__(self):
        return f"IODA ObsSpace Object - {self.name}"
 
    def __init__(self, path, dim_dict=None, mode='r', name="NoName", iodalayout=0):
        self.namename = name
        if mode not in ['r', 'w' ,'rw']:
            raise TypeError(f"{mode} not one of 'r','w','rw'")
        self.writewrite = True if 'w' in mode else False
        self.readread = True if 'r' in mode else False
        if  os.path.isfile(path) and self.readread:
            self.iodafileiodafile = path
        elif os.path.exists(os.path.dirname(path)) and self.writewrite:
            self.iodafileiodafile = path
        else:
            raise OSError(f"{path} does not specify a valid path to a file")
        self.iodalayoutiodalayout = iodalayout
        # read obsgroup from IODA
        if self.readread:
            self.file, self.obsgroupobsgroup = self._read_obsgroup_read_obsgroup(self.iodafileiodafile)
            self._def_classvars_def_classvars()
        else:
            if dim_dict is not None:
                self.file, self.obsgroupobsgroup = self._create_obsgroup_create_obsgroup(self.iodafileiodafile, dim_dict)
            else:
                raise TypeError("dim_dict is not defined")
 
    def _def_classvars(self):
        # define some variables
        self.dimensionsdimensions = self.obsgroupobsgroup.vars.list()
        self.groupsgroups = self.file.listGroups(recurse=True)
        self.attrsattrs = self.obsgroupobsgroup.atts.list()
        self.nlocsnlocs = len(self.obsgroupobsgroup.vars.open('nlocs').readVector.int())
        self.variablesvariables = self.file.listVars(recurse=True)
        self.nvarsnvars = len(self.variablesvariables)
 
    def write_attr(self, attrName, attrVal):
        # get type of variable
        try:
            attrType = self.NumpyToIodaDtypeNumpyToIodaDtype(attrVal)
            if attrType == _iodapy.Types.float:
                if np.array(attrVal).size == 1:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     [1]).writeDatum.float(attrVal)
                else:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     len(attrVal)).writeVector.float(attrVal)
            elif attrType == _iodapy.Types.double:
                if np.array(attrVal).size == 1:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     [1]).writeDatum.double(attrVal)
                else:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     len(attrVal)).writeVector.double(attrVal)
            elif attrType == _iodapy.Types.int64:
                if np.array(attrVal).size == 1:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     [1]).writeDatum.int64(attrVal)
                else:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     len(attrVal)).writeVector.int64(attrVal)
            elif attrType == _iodapy.Types.int32:
                if np.array(attrVal).size == 1:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     [1]).writeDatum.int32(attrVal)
                else:
                    self.obsgroupobsgroup.atts.create(attrName, attrType,
                                     len(attrVal)).writeVector.int32(attrVal)
            # add other elif here TODO
        except AttributeError:  # if string
            if (type(attrVal) == str):
                attrType = _iodapy.Types.str
                self.obsgroupobsgroup.atts.create(
                    attrName, attrType, [1]).writeDatum.str(attrVal)
 
    def _read_obsgroup(self, path):
        # open the obs group for a specified file and return file and obsgroup objects
        if self.writewrite:
            iodamode = _iodapy.Engines.BackendOpenModes.Read_Write
        else:
            iodamode = _iodapy.Engines.BackendOpenModes.Read_Only
        g = _iodapy.Engines.HH.openFile(
                        name=path,
                        mode=iodamode,
                        )
        dlp = _iodapy.DLP.DataLayoutPolicy.generate(
                  _iodapy.DLP.DataLayoutPolicy.Policies(self.iodalayoutiodalayout))
        og = _iodapy.ObsGroup(g, dlp)
        return g, og
 
    def _create_obsgroup(self, path, dim_dict):
        # open a new file for writing and return file and obsgroup objects
        g = _iodapy.Engines.HH.createFile(
                        name=path,
                        mode=_iodapy.Engines.BackendCreateModes.Truncate_If_Exists,
                        )
        # create list of dims in obs group
        _dim_list = []
        for key, value in dim_dict.items():
            if key == 'nlocs':
                _nlocs_var = _iodapy.NewDimensionScale.int32(
                    'nlocs', value, _iodapy.Unlimited, value)
                _dim_list.append(_nlocs_var)
            else:
                _dim_list.append(_iodapy.NewDimensionScale.int32(
                key, value, value, value))
        # create obs group
        og = _iodapy.ObsGroup.generate(g, _dim_list)
        # default compression option
        self._p1_p1 = _iodapy.VariableCreationParameters()
        self._p1_p1.compressWithGZIP()
        return g, og
 
    def create_var(self, varname, groupname=None, dtype=np.dtype('float32'), dim_list=['nlocs']):
        dtype_tmp = np.array([],dtype=dtype)
        typeVar = self.NumpyToIodaDtypeNumpyToIodaDtype(dtype_tmp)
        _varstr = varname
        if groupname is not None:
            _varstr = f"{groupname}/{varname}"
        # get list of dimension variables
        dims = [self.obsgroupobsgroup.vars.open(dim) for dim in dim_list]
        newVar = self.file.vars.create(_varstr, typeVar,
                                       scales=dims, params=self._p1_p1)
 
    def Variable(self, varname, groupname=None):
        return self._Variable_Variable(self, varname, groupname)
 
    def NumpyToIodaDtype(self, NumpyArr):
        ############################################################
        # This method converts the numpy data type to the
        # corresponding ioda datatype
 
        NumpyDtype = NumpyArr.dtype
 
        if (NumpyDtype == np.dtype('float64')):
            IodaDtype = _iodapy.Types.double    # convert double to float
        elif (NumpyDtype == np.dtype('float32')):
            IodaDtype = _iodapy.Types.float
        elif (NumpyDtype == np.dtype('int64')):
            IodaDtype = _iodapy.Types.int64    # convert long to int
        elif (NumpyDtype == np.dtype('int32')):
            IodaDtype = _iodapy.Types.int32
        elif (NumpyDtype == np.dtype('int16')):
            IodaDtype = _iodapy.Types.int16
        elif (NumpyDtype == np.dtype('int8')):
            IodaDtype = _iodapy.Types.int16
        elif (NumpyDtype == np.dtype('S1')):
            IodaDtype = _iodapy.Types.str
        elif (NumpyDtype == np.dtype('object')):
            IodaDtype = _iodapy.Types.str
        else:
            try:
                a = str(NumpyArr[0])
                IodaDtype = _iodapy.Types.str
            except TypeError:
                print("ERROR: Unrecognized numpy data type: ", NumpyDtype)
                exit(-2)
        return IodaDtype
 
    class _Variable:
        def __repr__(self):
            return f"IODA variable ({self._varstr})"
 
        def __str__(self):
            return f"IODA variable object - {self._varstr}"
 
        def __init__(self, obsspace, varname, groupname=None):
            self.obsspaceobsspace = obsspace
            self._varstr_varstr = varname
            if groupname is not None:
                self._varstr_varstr = f"{groupname}/{varname}"
            self._iodavar_iodavar = self.obsspaceobsspace.obsgroup.vars.open(self._varstr_varstr)
            self.attrsattrs = self._iodavar_iodavar.atts.list()
 
        def write_data(self, npArray):
            datatype = self.obsspaceobsspace.NumpyToIodaDtype(npArray)
            if datatype == _iodapy.Types.float:
                self._iodavar_iodavar.writeNPArray.float(npArray)
            elif datatype == _iodapy.Types.double:
                self._iodavar_iodavar.writeNPArray.double(npArray)
            elif datatype == _iodapy.Types.int64:
                self._iodavar_iodavar.writeNPArray.int64(npArray)
            elif datatype == _iodapy.Types.int32:
                self._iodavar_iodavar.writeNPArray.int32(npArray)
            elif datatype == _iodapy.Types.str:
                self._iodavar_iodavar.writeVector.str(npArray)
            # add other elif here TODO
 
        def write_attr(self, attrName, attrVal):
            # get type of variable
            try:
                attrType = self.obsspaceobsspace.NumpyToIodaDtype(attrVal)
                if attrType == _iodapy.Types.float:
                    if len(attrVal) == 1:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         [1]).writeDatum.float(attrVal)
                    else:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         len(attrVal)).writeVector.float(attrVal)
                elif attrType == _iodapy.Types.double:
                    if len(attrVal) == 1:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         [1]).writeDatum.double(attrVal)
                    else:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         len(attrVal)).writeVector.double(attrVal)
                elif attrType == _iodapy.Types.int64:
                    if len(attrVal) == 1:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         [1]).writeDatum.int64(attrVal)
                    else:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         len(attrVal)).writeVector.int64(attrVal)
                elif attrType == _iodapy.Types.int32:
                    if len(attrVal) == 1:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         [1]).writeDatum.int32(attrVal)
                    else:
                        self._iodavar_iodavar.atts.create(attrName, attrType,
                                         len(attrVal)).writeVector.int32(attrVal)
                # add other elif here TODO
            except AttributeError:  # if string
                if (type(attrVal) == str):
                    attrType = _iodapy.Types.str
                    self._iodavar_iodavar.atts.create(
                        attrName, attrType, [1]).writeDatum.str(attrVal)
 
        def read_data(self):
            """
              returns numpy array/python list of variable data
            """
            # figure out what datatype the variable is
            if self._iodavar_iodavar.isA2(_iodapy.Types.float):
                # float
                data = self._iodavar_iodavar.readNPArray.float()
                data[np.abs(data) > 9e36] = np.nan # undefined values
            elif self._iodavar_iodavar.isA2(_iodapy.Types.int):
                # integer
                data = self._iodavar_iodavar.readNPArray.int()
            elif self._iodavar_iodavar.isA2(_iodapy.Types.str):
                # string
                data = self._iodavar_iodavar.readVector.str() # ioda cannot read str to datetime...
                # convert to datetimes if applicable
                if "datetime" in self._varstr_varstr:
                    data = self._str_to_datetime_str_to_datetime(data)
            else:
                raise TypeError(f"Variable {self._varstr} type not supported")
            # reshape if it is a 2D array with second dim size 1
            try:
                if data.shape[1] == 1:
                    data = data.flatten()
            except (IndexError, AttributeError):
                pass
            return data
 
        def _str_to_datetime(self, datain):
            # comes as a list of strings, need to
            # make into datetime objects
            datetimes = [dt.datetime.strptime(x, "%Y-%m-%dT%H:%M:%SZ") for x in datain]
            return np.array(datetimes)