Commit db1495cb authored by BorjaEst's avatar BorjaEst
Browse files

Move sources.py utils functions to utils.py

parent 35ad9caf
"""This module creates the sources objects""" """This module creates the sources objects"""
import glob import glob
import yaml
import netCDF4
import xarray as xr import xarray as xr
import os.path import os.path
from . import utils
def load(yaml_file):
"""Loads the .yaml file with the sources configurations"""
with open(yaml_file, "r") as ymlfile:
return yaml.load(ymlfile)
class Source: class Source:
...@@ -32,16 +25,16 @@ class Model: ...@@ -32,16 +25,16 @@ class Model:
"""Standarised model with standarised variables""" """Standarised model with standarised variables"""
def __init__(self, variables): def __init__(self, variables):
self.get_tco3_zm(**variables) self.__get_tco3_zm(**variables)
self.get_vrm_zm(**variables) self.__get_vrm_zm(**variables)
def skim(self, path): def skim(self, path):
if hasattr(self, '_tco3_zm'): if hasattr(self, '_tco3_zm'):
to_netcdf(path, "tco3_zm", self._tco3_zm) utils.to_netcdf(path, "tco3_zm", self._tco3_zm)
if hasattr(self, '_vrm_zm'): if hasattr(self, '_vrm_zm'):
to_netcdf(path, "vrm_zm", self._vrm_zm) utils.to_netcdf(path, "vrm_zm", self._vrm_zm)
def get_tco3_zm(self, tco3_zm=None, **kwarg): def __get_tco3_zm(self, tco3_zm=None, **kwarg):
"""Gets and standarises the tco3_zm data""" """Gets and standarises the tco3_zm data"""
if tco3_zm: if tco3_zm:
fnames = glob.glob(tco3_zm['dir'] + "/*.nc") fnames = glob.glob(tco3_zm['dir'] + "/*.nc")
...@@ -53,7 +46,7 @@ class Model: ...@@ -53,7 +46,7 @@ class Model:
tco3_zm['coordinades']['lon']: 'lon' tco3_zm['coordinades']['lon']: 'lon'
})['tco3_zm'].to_dataset() })['tco3_zm'].to_dataset()
def get_vrm_zm(self, vrm_zm=None, **kwarg): def __get_vrm_zm(self, vrm_zm=None, **kwarg):
"""Gets and standarises the vrm_zm data""" """Gets and standarises the vrm_zm data"""
if vrm_zm: if vrm_zm:
fnames = glob.glob(vrm_zm['dir'] + "/*.nc") fnames = glob.glob(vrm_zm['dir'] + "/*.nc")
...@@ -66,16 +59,3 @@ class Model: ...@@ -66,16 +59,3 @@ class Model:
})['tco3_zm'].to_dataset() })['tco3_zm'].to_dataset()
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')
"""This module offers some utils for code management""" """This module offers some utils for code management"""
from contextlib import contextmanager from contextlib import contextmanager
import os import os
import yaml
import netCDF4
import xarray as xr
@contextmanager @contextmanager
...@@ -12,3 +15,24 @@ def cd(newdir): ...@@ -12,3 +15,24 @@ def cd(newdir):
yield yield
finally: finally:
os.chdir(prevdir) os.chdir(prevdir)
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')
...@@ -5,7 +5,7 @@ import numpy as np ...@@ -5,7 +5,7 @@ import numpy as np
import netCDF4 import netCDF4
import os.path import os.path
import datetime import datetime
from o3skim import sources from o3skim import utils
base = datetime.datetime(2000, 1, 1) base = datetime.datetime(2000, 1, 1)
...@@ -40,4 +40,4 @@ def dataset(name, coordinades): ...@@ -40,4 +40,4 @@ def dataset(name, coordinades):
def netcdf(path, name, coordinades, **kwarg): def netcdf(path, name, coordinades, **kwarg):
"""Creates or appends data to a mock netcdf file""" """Creates or appends data to a mock netcdf file"""
ds = dataset(name, coordinades) ds = dataset(name, coordinades)
sources.to_netcdf(path, name, ds) utils.to_netcdf(path, name, ds)
...@@ -16,7 +16,7 @@ class TestO3SKIM_sources(unittest.TestCase): ...@@ -16,7 +16,7 @@ class TestO3SKIM_sources(unittest.TestCase):
def setUp(self): def setUp(self):
"""Loads and creates the test folders and files from test_sources.yaml""" """Loads and creates the test folders and files from test_sources.yaml"""
self.config = sources.load("tests/test_sources.yaml") self.config = utils.load("tests/test_sources.yaml")
self.assertTrue(type(self.config) is dict) self.assertTrue(type(self.config) is dict)
self.create_mock_datasets() self.create_mock_datasets()
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment