utils.py 1.04 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
7
8
9
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:
        yield
    finally:
        os.chdir(prevdir)
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


def load(yaml_file):
    """Loads the .yaml file with the sources configurations"""
    with open(yaml_file, "r") as ymlfile:
        return yaml.load(ymlfile)


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]
    [create_empty_netCDF(fn) for fn in fnames if not os.path.isfile(fn)]
    xr.save_mfdataset(dsx, fnames, mode='a')