utils.py 4 KB
Newer Older
BorjaEst's avatar
BorjaEst committed
1
2
3
4
5
6
7
8
"""This module offers some utils for code management. For example
utils to change easily from directory or create netCDF files.

Although this module is not expected to be used out of the o3skim
package, some functions might be convenient when using the python 
shell.
"""

9
10
from contextlib import contextmanager
import os
11
12
13
import yaml
import netCDF4
import xarray as xr
BorjaEst's avatar
BorjaEst committed
14
15
16
import logging

logger = logging.getLogger('o3skim.utils')
17
18
19
20


@contextmanager
def cd(newdir):
BorjaEst's avatar
BorjaEst committed
21
22
23
24
25
26
27
    """Changes the directory inside a 'with' context. When the code
    reaches the end of the 'with' block or the code fails, the 
    previous folder is restored.

    :param newdir: Path folder where to change the working directory.
    :type newdir: str
    """
28
29
30
    prevdir = os.getcwd()
    os.chdir(os.path.expanduser(newdir))
    try:
BorjaEst's avatar
BorjaEst committed
31
        logger.debug("Temp dir change to: '%s'", newdir)
32
33
34
        yield
    finally:
        os.chdir(prevdir)
BorjaEst's avatar
BorjaEst committed
35
        logger.debug("Restore directory: '%s'", prevdir)
36

37

38
def return_on_failure(message, default=None):
BorjaEst's avatar
BorjaEst committed
39
40
41
42
43
    """Decorator to do not break but log. Note that the error stack
    is printed as well to do not lose relevant information.

    :param message: Additional message to log when the function fails.
    :type message: str
44
45
46

    :param default: Value to return if fails.
    :type default: any or None, optional
BorjaEst's avatar
BorjaEst committed
47
    """
48
49
50
    def decorate(function):
        def applicator(*args, **kwargs):
            try:
51
                return function(*args, **kwargs)
52
            except:
53
                # Log error with stack using root (not utils)
54
                logging.error(message, exc_info=True)
55
                return default
56
57
58
59
        return applicator
    return decorate


60
def load(yaml_file):
BorjaEst's avatar
BorjaEst committed
61
62
63
64
65
66
67
68
69
70
71
    """Loads the .yaml file and returns a python dictionary with all
    the yaml keys and values. Note a yaml file can have key:values 
    sotored inside values, therefore the returned dictionary might 
    have dictionaries stored inside values. 

    :param yaml_file: Path pointing to the yaml configuration file.
    :type yaml_file: str

    :return: Dictionary with the yaml structure key:value.
    :rtype: dict
    """
72
    with open(yaml_file, "r") as ymlfile:
73
        config = yaml.safe_load(ymlfile)
BorjaEst's avatar
BorjaEst committed
74
75
        logging.debug("Configuration data: %s", config)
        return config
76
77
78


def create_empty_netCDF(fname):
BorjaEst's avatar
BorjaEst committed
79
80
81
82
83
    """Creates a new empty netCDF file.

    :param fname: Name and path where to create the file.
    :type fname: str
    """
84
85
86
87
88
    root_grp = netCDF4.Dataset(fname, 'w', format='NETCDF4')
    root_grp.description = 'Example simulation data'
    root_grp.close()


89
def to_netcdf(dirname, name, dataset, groupby=None):
BorjaEst's avatar
BorjaEst committed
90
91
92
93
94
95
96
97
98
99
    """Creates or appends data to named netCDF files.

    :param path: Location where to find or create the netCDF files.
    :type path: str

    :param name: Name/Prefix for file/s where to store the data.
    :type name: str

    :param dataset: Dataset to write to the netCDF file.
    :type dataset: :class:`xarray.Dataset`
100
101
102

    :param groupby: How to group files (None, year, decade).
    :type groupby: str, optional
BorjaEst's avatar
BorjaEst committed
103
    """
104
105
106
107
108
109
    def split_by_year(dataset):
        """Splits a dataset by year"""
        years, dsx = zip(*dataset.groupby("time.year"))
        fnames = [dirname + "/" + name + "_%s.nc" % y for y in years]
        return fnames, dsx

BorjaEst's avatar
BorjaEst committed
110
111
112
113
114
115
116
    def split_by_decade(dataset):
        """Splits a dataset by decade"""
        decades = dataset.indexes["time"].year//10*10
        decades, dsx = zip(*dataset.groupby(xr.DataArray(decades)))
        fnames = [dirname + "/" + name + "_%s-%s.nc" % (d, d+10) for d in decades]
        return fnames, dsx

117
118
119
    def no_split(dataset):
        """Does not split a dataset"""
        dsx = (dataset,)
BorjaEst's avatar
BorjaEst committed
120
        fnames = [dirname + "/" + name + ".nc"]
121
122
123
        return fnames, dsx

    split_by = {
BorjaEst's avatar
BorjaEst committed
124
        "year": split_by_year,
BorjaEst's avatar
BorjaEst committed
125
        "decade": split_by_decade
126
127
    }
    fnames, dsx = split_by.get(groupby, no_split)(dataset)
BorjaEst's avatar
BorjaEst committed
128

BorjaEst's avatar
BorjaEst committed
129
    logging.info("Save dataset into: %s", fnames)
130
131
    [create_empty_netCDF(fn) for fn in fnames if not os.path.isfile(fn)]
    xr.save_mfdataset(dsx, fnames, mode='a')