Commit cc82b554 authored by BorjaEst's avatar BorjaEst
Browse files

Fix mockup to do not use Dataset subclass

parent 07bdfce1
"""Mockup sub-package for tests"""
import xarray as xr
import numpy as np
# import o3skim.utils as utils
import xarray as xr
from tests.mockup import coordinate
def tco3(year_range):
for year in year_range:
xr.Dataset(
data_vars=dict(
tco3=(["lon", "lat", "time"],
np.random.rand(21, 11, 12))
),
coords=dict(
lon=coordinate.lon(num=21),
lat=coordinate.lat(num=11),
plev=coordinate.plev(num=4),
time=coordinate.time(num=12, start=year)
),
attrs=dict(description="Test tco3 dataset")
).to_netcdf("tco3_{}.nc".format(year))
def vmro3(year_range):
for year in year_range:
xr.Dataset(
data_vars=dict(
vmro3=(["lon", "lat", "plev", "time"],
np.random.rand(21, 11, 4, 12))
),
coords=dict(
lon=coordinate.lon(num=21),
lat=coordinate.lat(num=11),
plev=coordinate.plev(num=4),
time=coordinate.time(num=12, start=year)
),
attrs=dict(description="Test vmro3 dataset")
).to_netcdf("vmro3_{}.nc".format(year))
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 combined(year_range):
for year in year_range:
xr.Dataset(
data_vars=dict(
tco3=(["lon", "lat", "time"],
np.random.rand(21, 11, 12)),
vmro3=(["lon", "lat", "plev", "time"],
np.random.rand(21, 11, 4, 12))
),
coords=dict(
lon=coordinate.lon(num=21),
lat=coordinate.lat(num=11),
plev=coordinate.plev(num=4),
time=coordinate.time(num=12, start=year)
),
attrs=dict(description="Test ozone dataset")
).to_netcdf("merged_{}.nc".format(year))
def _coord_data(coordinates):
coord_names, shape = zip(*[(k, len(x)) for k, x in coordinates])
return coord_names, np.ones(shape)
def noise(name='noise'):
"""Some noise that blocks '*' wild card"""
xr.Dataset(
data_vars=dict(noise=(["time"], np.random.rand(12))),
coords=dict(time=coordinate.time(num=12)),
attrs=dict(description="Test noise dataset")
).to_netcdf(name + ".nc")
"""Module with coordinate mockups"""
import datetime
import pandas as pd
import numpy as np
def __timeline(start, end=None, step=10):
if not end:
end = start
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(num=12, start=2000, end=None):
end = start + 1 if not end else end
start = pd.Timestamp("{}-01-01".format(start))
end = pd.Timestamp("{}-01-01".format(end))
t = np.linspace(start.value, end.value, num + 1)[0:-1]
return ('time', pd.to_datetime(t))
def time(start=2000, end=None, step=30):
end = start if not end else None
return ('time', [t for t in __timeline(start, end, step)])
def lat(num=11):
return ('lat', np.linspace(-90, 90, num))
def lat(delta=10):
return ('lat', [x for x in range(-90, 90, delta)])
def lon(num=21):
return ('lon', np.linspace(-180, 180, num))
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)])
def plev(num=4):
return ('plev', np.logspace(1, 4, num))
"""This modules creates mockup data for testing"""
import os
from o3skim import utils
from tests import mockup
from tests.mockup import coordinate
def _tco3(year, name='tco3'):
return mockup.Dataset(
var_coords={
name: [
coordinate.time(year),
coordinate.lat(),
coordinate.lon()
]
},
coords=[
coordinate.time(year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
)
def _vmro3(year, name='vmro3'):
return mockup.Dataset(
var_coords={
name: [
coordinate.time(year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
},
coords=[
coordinate.time(year),
coordinate.plev(),
coordinate.lat(),
coordinate.lon()
]
)
def tco3(year_range, name='tco3'):
for year in year_range:
_tco3(year, name).to_netcdf(name + "_" + str(year) + ".nc")
def vmro3(year_range, name='vmro3'):
for year in year_range:
_vmro3(year, name).to_netcdf(name + "_" + str(year) + ".nc")
_datasets = {
'tco3': _tco3,
'vmro3': _vmro3
}
def combined(year_range, vars=['tco3', 'vmro3']):
for year in year_range:
dataset = mockup.Dataset({}, [])
for var in vars:
load = _datasets[var](year)
dataset = dataset.merge(load)
dataset.to_netcdf("merged_" + str(year) + ".nc")
def noise(name='noise'):
"""Some noise that blocks '*' wild card"""
mockup.Dataset(
var_coords={'noise': [coordinate.time(2000, 2010)]},
coords=[coordinate.time(2000, 2010)]
).to_netcdf(name + ".nc")
......@@ -5,7 +5,7 @@ import pytest
import o3skim
import xarray
import tests.mockup_data as mockup_data
import tests.mockup as mockup_data
# Test configurations ----------------------------------------------
......@@ -29,11 +29,12 @@ def data_dir(tmpdir_factory):
with o3skim.utils.cd(source_dir):
if source == 'SourceMerged':
mockup_data.combined(year_line)
mockup_data.noise(name='merged_noise.nc')
if source == 'SourceSplit':
mockup_data.tco3(year_line)
mockup_data.noise(name='tco3_noise.nc')
mockup_data.vmro3(year_line)
mockup_data.noise(name='vmro3_noise.nc')
with o3skim.utils.cd(data_dir):
yield data_dir
......
......@@ -5,7 +5,7 @@ import pytest
import o3skim
import xarray
import tests.mockup_data as mockup_data
import tests.mockup as mockup_data
# Test configurations ----------------------------------------------
configuration_file = "tests/sources_example.yaml"
......@@ -36,7 +36,6 @@ def data_dir(tmpdir_factory):
mockup_data.noise(name='tco3_noise.nc')
mockup_data.vmro3(year_line)
mockup_data.noise(name='vmro3_noise.nc')
with o3skim.utils.cd(data_dir):
yield data_dir
......
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