gsi_utils.py

# gsi_utils.py
# a collection of functions, classes, etc.
# used for the GSI global analysis


def isTrue(str_in):
    """ isTrue(str_in)
    - function to translate shell variables to python logical variables

    input: str_in - string (should be like 'YES', 'TRUE', etc.)
    returns: status (logical True or False)

    """
    str_in = str_in.upper()
    if str_in in ['YES', '.TRUE.']:
        status = True
    else:
        status = False
    return status


def link_file(from_file, to_file):
    """ link_file(from_file, to_file)
    - function to check if a path exists, and if not, make a symlink
    input: from_file - string path
            to_file   - string path
    """
    import os
    if not os.path.exists(to_file):
        if not os.path.islink(to_file):
            os.symlink(from_file, to_file)
        else:
            print(to_file + " exists, unlinking.")
            os.unlink(to_file)
            os.symlink(from_file, to_file)
        print("ln -s " + from_file + " " + to_file)


def copy_file(from_file, to_file):
    import shutil
    shutil.copy(from_file, to_file)
    print("cp " + from_file + " " + to_file)


def make_dir(directory):
    import os
    os.makedirs(directory)
    print("mkdir -p " + directory)


def write_nml(nml_dict, nml_file):
    """ write_nml(nml_dict, nml_file)
    - function to write out namelist dictionary nml_dict to file nml_file
    input: nml_dict - dictionary of dictionaries
                        first dictionary is &nml, second is nmlvar='value'
                        NOTE: this shoudl be an OrderedDict or else it might fail
           nml_file - string path to write namelist file to
    """
    nfile = open(nml_file, 'w')

    for nml, nmlvars in nml_dict.items():
        nfile.write('&' + nml + '\n')
        for var, val in nmlvars.items():
            nfile.write('  ' + str(var) + ' = ' + str(val) + '\n')
        nfile.write('/\n\n')
    nfile.close()


def get_ncdims(ncfile):
    """ get_ncdims(ncfile)
    - function to return dictionary of netCDF file dimensions and their lengths
    input: ncfile - string to path to netCDF file
    output: ncdims - dictionary where key is the name of a dimension and the
                        value is the length of that dimension

                        ex:  ncdims['pfull'] = 127
    """
    try:
        import netCDF4 as nc
    except ImportError as err:
        raise ImportError(f"Unable to import netCDF4 module")
    ncf = nc.Dataset(ncfile)
    ncdims = {}
    for d in ncf.dimensions.keys():
        ncdims[d] = int(len(ncf.dimensions[d]))
    ncf.close()

    return ncdims


def get_nemsdims(nemsfile, nemsexe):
    """ get_nemsdims(nemsfile,nemsexe)
    - function to return dictionary of NEMSIO file dimensions for use
    input:  nemsfile - string to path nemsio file
            nemsexe  - string to path nemsio_get executable
    output: nemsdims - dictionary where key is the name of a dimension and the
                       value is the length of that dimension
                       ex: nemsdims['pfull'] = 127
    """
    import subprocess
    ncdims = {
        'dimx': 'grid_xt',
                'dimy': 'grid_yt',
                'dimz': 'pfull',
    }
    nemsdims = {}
    for dim in ['dimx', 'dimy', 'dimz']:
        out = subprocess.Popen([nemsexe, nemsfile, dim], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        stdout, stderr = out.communicate()
        nemsdims[ncdims[dim]] = int(stdout.split(' ')[-1].rstrip())
    return nemsdims


def get_timeinfo(ncfile):
    """ get_timeinfo(ncfile)
     - function to return datetime objects of initialized time and valid time
     input: ncfile - string to path to netCDF file
     returns: inittime, validtime - datetime objects
              nfhour - integer forecast hour
    """
    try:
        import netCDF4 as nc
    except ImportError as err:
        raise ImportError(f"Unable to import netCDF4 module\n{err}")
    import datetime as dt
    import re
    ncf = nc.Dataset(ncfile)
    time_units = ncf['time'].units
    date_str = time_units.split('since ')[1]
    date_str = re.sub("[^0-9]", "", date_str)
    initstr = date_str[0:10]
    inittime = dt.datetime.strptime(initstr, "%Y%m%d%H")
    nfhour = int(ncf['time'][0])
    validtime = inittime + dt.timedelta(hours=nfhour)
    ncf.close()

    return inittime, validtime, nfhour