utils.py 1.34 KB
Newer Older
1
2
3
"""This module offers some utils for code management"""
from contextlib import contextmanager
import os
4
5
6
import yaml
import netCDF4
import xarray as xr
BorjaEst's avatar
BorjaEst committed
7
8
9
import logging

logger = logging.getLogger('o3skim.utils')
10
11
12
13
14
15
16
17


@contextmanager
def cd(newdir):
    """Changes the directory inside a 'with' context"""
    prevdir = os.getcwd()
    os.chdir(os.path.expanduser(newdir))
    try:
BorjaEst's avatar
BorjaEst committed
18
        logger.debug("Temp dir change to: '%s'", newdir)
19
20
21
        yield
    finally:
        os.chdir(prevdir)
BorjaEst's avatar
BorjaEst committed
22
        logger.debug("Restore directory: '%s'", prevdir)
23
24
25
26

def load(yaml_file):
    """Loads the .yaml file with the sources configurations"""
    with open(yaml_file, "r") as ymlfile:
BorjaEst's avatar
BorjaEst committed
27
28
29
        config = yaml.load(ymlfile)
        logging.debug("Configuration data: %s", config)
        return config
30
31
32
33
34
35
36
37
38
39
40
41
42


def create_empty_netCDF(fname):
    """Creates a new empty netCDF file"""
    root_grp = netCDF4.Dataset(fname, 'w', format='NETCDF4')
    root_grp.description = 'Example simulation data'
    root_grp.close()


def to_netcdf(path, name, dataset):
    """Creates or appends data to named netcdf files"""
    years, dsx = zip(*dataset.groupby("time.year"))
    fnames = [path + "/" + name + "_%s.nc" % y for y in years]
BorjaEst's avatar
BorjaEst committed
43
    logging.info("Save dataset into: %s", fnames)
44
45
    [create_empty_netCDF(fn) for fn in fnames if not os.path.isfile(fn)]
    xr.save_mfdataset(dsx, fnames, mode='a')