Commit a09497b6 authored by BorjaEst's avatar BorjaEst
Browse files

Add mockup data generation using tests __init__

parent a959cc73
......@@ -71,3 +71,4 @@ coverage.xml
.pytest_cache/
cover/
flake8.log
tests/data
\ No newline at end of file
"""Unit test package for torch_cd_example."""
import os
from o3skim import utils
data_dir = "./tests/data"
# Generation of mockup data
os.makedirs(data_dir, exist_ok=True)
with utils.cd(data_dir):
from tests import mockup_data
"""Mockup sub-package for tests"""
import xarray as xr
import numpy as np
# import o3skim.utils as utils
class Dataset(xr.Dataset):
def __init__(self, var_coords, coords):
super().__init__(
data_vars={k: _coord_data(cx) for k, cx in var_coords.items()},
coords=dict(coords)
)
def _coord_data(coordinates):
coord_names, shape = zip(*[(k, len(x)) for k, x in coordinates])
return coord_names, np.ones(shape)
"""Module with coordinate mockups"""
import datetime
def __timeline(start, end, step=10):
start = datetime.datetime(start, 1, 1)
end = datetime.datetime(end, 12, 31)
step = datetime.timedelta(days=step)
while start < end:
yield start
start += step
def time(start=2000, end=2010, step=30):
return ('time', [t for t in __timeline(start, end, step)])
def lat(delta=10):
return ('lat', [x for x in range(-90, 90, delta)])
def lon(delta=20):
return ('lon', [x for x in range(-180, 180, delta)])
def plev(delta=100):
return ('plev', [x for x in range(1, 1000, delta)])
"""This modules creates mockup data for testing"""
import xarray as xr
import numpy as np
import datetime
import os
from o3skim import utils
base = datetime.datetime(2000, 1, 1)
indexes = {
'time': [base + datetime.timedelta(days=10*i) for i in range(99)],
'plev': [x for x in range(1, 1000, 100)],
'lat': [x for x in range(-90, 90, 10)],
'lon': [x for x in range(-180, 180, 20)]
}
from o3skim import utils
from tests import mockup
from tests.mockup import coordinate
def data_vars(coordinates):
"""Creates a mock n-array with coordinate values"""
alias = tuple(coordinates.values())
dim = tuple(len(indexes[c]) for c in coordinates)
return alias, np.ones(dim),
model1_dir = "./variable_per_file"
os.makedirs(model1_dir, exist_ok=True)
with utils.cd(model1_dir):
def map_coord(coordinates):
"""Creates a mock array by coordinate"""
return {alias: indexes[c] for c, alias in coordinates.items()}
path = "toz"
os.makedirs(path, exist_ok=True)
with utils.cd(path):
for year in range(2000, 2002):
mockup.Dataset(
var_coords={
'toz': [
coordinate.time(year,year),
coordinate.lat(),
coordinate.lon()
]
},
coords=[
coordinate.time(year, year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
).to_netcdf("toz_" + str(year) + ".nc")
path = "vmro3"
os.makedirs(path, exist_ok=True)
with utils.cd(path):
for year in range(2000, 2002):
mockup.Dataset(
var_coords={
'vmro3': [
coordinate.time(year,year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
},
coords=[
coordinate.time(year, year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
).to_netcdf("vmro3_" + str(year) + ".nc")
def dataset(name, coordinates):
"""Creates a dataset according to the global module indexes"""
return xr.Dataset(
{name: data_vars(coordinates)},
coords=map_coord(coordinates)
)
model2_dir = "./variables_merged"
os.makedirs(model2_dir, exist_ok=True)
with utils.cd(model2_dir):
def netcdf(dirname, name, coordinates, **kwarg):
"""Creates or appends data to a mock netcdf file"""
ds = dataset(name, coordinates)
utils.to_netcdf(dirname, name, ds, groupby="year")
for year in range(2000, 2002):
mockup.Dataset(
var_coords={
'toz': [
coordinate.time(year,year),
coordinate.lat(),
coordinate.lon()
],
'vmro3': [
coordinate.time(year,year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
},
coords=[
coordinate.time(year, year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
).to_netcdf("toz_" + str(year) + ".nc")
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