def test_when_results_are_xobjs(self):
def fn_ds(a, b):
ds = xr.Dataset(
coords={
'x': [1, 2],
'y': ['foo', 'bar', 'baz'],
},
data_vars={
'apples': (['x', 'y'], np.tile(10 * a + b, (2, 3))),
'lemons': ('y', np.tile(10 * a + b, (3,))),
}
)
return ds
fds = combo_runner_to_ds(
fn_ds,
combos={
'a': [2, 3, 4, 5],
'b': [6, 7, 8, 9, 10],
}, var_names=None)
assert fds.sel(a=4, b=9, x=2, y='bar')['apples'].values == 49
assert fds.sel(a=5, b=7, y='bar')['lemons'].values == 57
assert 'output' not in fds
def test_array_return(self):
ds = xr.Dataset(coords={
'a': [1, 2, 3],
'b': [40, 50],
't': [0.0, 0.1, 0.2, 0.3, 0.4]
})
ds['z1'] = (('a', 'b', 't'),
np.array([[41 + 0.1j * np.arange(5), [np.nan] * 5],
[[np.nan] * 5, 52 + 0.1j * np.arange(5)],
[[np.nan] * 5, [np.nan] * 5]]))
ds['z2'] = (('a', 'b', 't'),
np.array([[41 - 0.1j * np.arange(5), [np.nan] * 5],
[[np.nan] * 5, 52 - 0.1j * np.arange(5)],
[[np.nan] * 5, [np.nan] * 5]]))
settings = {
'var_names': ['z1', 'z2'],
'var_dims': {
('z1', 'z2'): 't'
},
'var_coords': {
't': [0.0, 0.1, 0.2, 0.3, 0.4]
}
}
fill_missing_cases(ds, foo2_zarray1_zarray2, **settings)
dst = combo_runner_to_ds(foo2_zarray1_zarray2,
combos=[('a', [1, 2, 3]), ('b', [40, 50])],
**settings)
assert dst.equals(ds)
def test_constants_to_attrs(self):
ds = combo_runner_to_ds(foo3_scalar,
combos=[('a', [1, 2, 3]),
('c', [100, 200, 300])],
constants={'b': 20},
var_names='x')
assert ds.attrs['b'] == 20
def test_constants_to_attrs(self):
ds = combo_runner_to_ds(foo3_scalar,
combos=[('a', [1, 2, 3]),
('c', [100, 200, 300])],
constants={'b': 20},
var_names='x')
assert ds.attrs['b'] == 20
def test_const_array_to_coord(self):
ds = combo_runner_to_ds(foo_array_input,
combos=[('a', [1, 2, 3])],
constants={'t': [10, 20, 30]},
var_names=['x'],
var_dims=[['t']])
assert 't' in ds.dims
assert 't' not in ds.attrs
def test_const_array_to_coord(self):
ds = combo_runner_to_ds(foo_array_input,
combos=[('a', [1, 2, 3])],
constants={'t': [10, 20, 30]},
var_names=['x'],
var_dims=[['t']])
assert 't' in ds.dims
assert 't' not in ds.attrs
def test_double_array_return_with_no_given_dimensions(self):
ds = combo_runner_to_ds(foo2_array_array,
combos=[('a', [1, 2]), ('b', [30, 40])],
var_names=['array1', 'array2'],
var_dims=[['auto'], ['auto']])
assert (ds['auto'].data.dtype == int
or ds['auto'].data.dtype == np.int64)
assert_allclose(ds['auto'].data, [0, 1, 2, 3, 4])
def test_multiresult(self):
ds = combo_runner_to_ds(foo3_float_bool, _test_combos1,
var_names=['bananas', 'cakes'])
assert ds.bananas.data.dtype == int
assert ds.cakes.data.dtype == bool
assert ds.sel(a=2, b=30, c=400)['bananas'].data == 432
assert ds.sel(a=1, b=10, c=100)['bananas'].data == 111
assert ds.sel(a=2, b=30, c=400)['cakes'].data
assert not ds.sel(a=1, b=10, c=100)['cakes'].data
def test_double_array_return_with_no_given_dimensions(self):
ds = combo_runner_to_ds(foo2_array_array,
combos=[('a', [1, 2]),
('b', [30, 40])],
var_names=['array1', 'array2'],
var_dims=[['auto'], ['auto']])
assert (ds['auto'].data.dtype == int or
ds['auto'].data.dtype == np.int64)
assert_allclose(ds['auto'].data, [0, 1, 2, 3, 4])
def test_multiresult(self):
ds = combo_runner_to_ds(foo3_float_bool,
_test_combos1,
var_names=['bananas', 'cakes'])
assert ds.bananas.data.dtype == int
assert ds.cakes.data.dtype == bool
assert ds.sel(a=2, b=30, c=400)['bananas'].data == 432
assert ds.sel(a=1, b=10, c=100)['bananas'].data == 111
assert ds.sel(a=2, b=30, c=400)['cakes'].data
assert not ds.sel(a=1, b=10, c=100)['cakes'].data
def test_single_string_var_names_with_no_var_dims(self):
combos = ('a', [1, 2, 3])
ds = combo_runner_to_ds(foo3_scalar,
combos,
constants={
'b': 10,
'c': 100
},
var_names='sum')
assert_allclose(ds['sum'].data, np.array([111, 112, 113]))
def test_array_and_single_result(self):
combos = (('a', [1, 2]),
('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array_bool, combos,
var_names=['bananas', 'ripe'],
var_dims=(['sugar'], []),
var_coords={'sugar': [*range(10, 20)]})
assert ds.ripe.data.dtype == bool
assert ds.sel(a=2, b=30, sugar=14)['bananas'].data == 32.4
with pytest.raises((KeyError, ValueError)):
ds['ripe'].sel(sugar=12)
def test_arrayresult(self):
combos = (('a', [1, 2]),
('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array, combos,
var_names='bananas',
var_dims={'bananas': ['sugar']},
var_coords={'sugar': [*range(10)]})
assert ds.bananas.data.dtype == float
assert_allclose(ds.sel(a=2, b=30)['bananas'].data,
[32.0, 32.1, 32.2, 32.3, 32.4,
32.5, 32.6, 32.7, 32.8, 32.9])
def test_array_and_single_result(self):
combos = (('a', [1, 2]), ('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array_bool,
combos,
var_names=['bananas', 'ripe'],
var_dims=(['sugar'], []),
var_coords={'sugar': [*range(10, 20)]})
assert ds.ripe.data.dtype == bool
assert ds.sel(a=2, b=30, sugar=14)['bananas'].data == 32.4
with pytest.raises((KeyError, ValueError)):
ds['ripe'].sel(sugar=12)
def test_arrayresult(self):
combos = (('a', [1, 2]), ('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array,
combos,
var_names='bananas',
var_dims={'bananas': ['sugar']},
var_coords={'sugar': [*range(10)]})
assert ds.bananas.data.dtype == float
assert_allclose(
ds.sel(a=2, b=30)['bananas'].data,
[32.0, 32.1, 32.2, 32.3, 32.4, 32.5, 32.6, 32.7, 32.8, 32.9])
def test_double_array_return_with_same_dimensions(self):
combos = (('a', [1, 2]),
('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array_array, combos,
var_names=['apples', 'oranges'],
var_dims={('apples', 'oranges'): ['seeds']},
var_coords={'seeds': [*range(5)]})
assert ds.oranges.data.dtype == int
assert_allclose(ds.sel(a=2, b=30).apples.data, [30, 32, 34, 36, 38])
assert_allclose(ds.sel(a=2, b=30).oranges.data, [30, 28, 26, 24, 22])
assert 'seeds' in ds.apples.coords
assert 'seeds' in ds.oranges.coords
def test_complex_output(self):
ds = combo_runner_to_ds(foo2_zarray1_zarray2,
combos=[('a', [1, 2]), ('b', [30, 40])],
var_names=['array1', 'array2'],
var_dims=[['auto'], ['auto']])
assert ds['array1'].data.size == 2 * 2 * 5
assert ds['array2'].data.size == 2 * 2 * 5
assert ds['array1'].data.dtype == complex
assert ds['array2'].data.dtype == complex
assert_allclose(ds['array1'].sel(a=2, b=30).data,
32 + np.arange(5) * 0.1j)
assert_allclose(ds['array2'].sel(a=2, b=30).data,
32 - np.arange(5) * 0.1j)
def test_double_array_return_with_same_dimensions(self):
combos = (('a', [1, 2]), ('b', [10, 20, 30]))
ds = combo_runner_to_ds(foo2_array_array,
combos,
var_names=['apples', 'oranges'],
var_dims={('apples', 'oranges'): ['seeds']},
var_coords={'seeds': [*range(5)]})
assert ds.oranges.data.dtype == int
assert_allclose(ds.sel(a=2, b=30).apples.data, [30, 32, 34, 36, 38])
assert_allclose(ds.sel(a=2, b=30).oranges.data, [30, 28, 26, 24, 22])
assert 'seeds' in ds.apples.coords
assert 'seeds' in ds.oranges.coords
def test_complex_output(self):
ds = combo_runner_to_ds(foo2_zarray1_zarray2,
combos=[('a', [1, 2]),
('b', [30, 40])],
var_names=['array1', 'array2'],
var_dims=[['auto'], ['auto']])
assert ds['array1'].data.size == 2 * 2 * 5
assert ds['array2'].data.size == 2 * 2 * 5
assert ds['array1'].data.dtype == complex
assert ds['array2'].data.dtype == complex
assert_allclose(ds['array1'].sel(a=2, b=30).data,
32 + np.arange(5) * 0.1j)
assert_allclose(ds['array2'].sel(a=2, b=30).data,
32 - np.arange(5) * 0.1j)
def test_array_return(self):
ds = xr.Dataset(coords={'a': [1, 2, 3], 'b': [40, 50],
't': [0.0, 0.1, 0.2, 0.3, 0.4]})
ds['z1'] = (('a', 'b', 't'),
np.array([[41 + 0.1j * np.arange(5), [np.nan] * 5],
[[np.nan] * 5, 52 + 0.1j * np.arange(5)],
[[np.nan] * 5, [np.nan] * 5]]))
ds['z2'] = (('a', 'b', 't'),
np.array([[41 - 0.1j * np.arange(5), [np.nan] * 5],
[[np.nan] * 5, 52 - 0.1j * np.arange(5)],
[[np.nan] * 5, [np.nan] * 5]]))
settings = {'var_names': ['z1', 'z2'],
'var_dims': {('z1', 'z2'): 't'},
'var_coords': {'t': [0.0, 0.1, 0.2, 0.3, 0.4]}}
fill_missing_cases(ds, foo2_zarray1_zarray2, **settings)
dst = combo_runner_to_ds(foo2_zarray1_zarray2,
combos=[('a', [1, 2, 3]),
('b', [40, 50])], **settings)
assert dst.equals(ds)
def test_when_results_are_xobjs(self):
def fn_ds(a, b):
ds = xr.Dataset(coords={
'x': [1, 2],
'y': ['foo', 'bar', 'baz'],
},
data_vars={
'apples': (['x',
'y'], np.tile(10 * a + b, (2, 3))),
'lemons': ('y', np.tile(10 * a + b, (3, ))),
})
return ds
fds = combo_runner_to_ds(fn_ds,
combos={
'a': [2, 3, 4, 5],
'b': [6, 7, 8, 9, 10],
},
var_names=None)
assert fds.sel(a=4, b=9, x=2, y='bar')['apples'].values == 49
assert fds.sel(a=5, b=7, y='bar')['lemons'].values == 57
assert 'output' not in fds
def test_single_string_var_names_with_no_var_dims(self):
combos = ('a', [1, 2, 3])
ds = combo_runner_to_ds(foo3_scalar, combos,
constants={'b': 10, 'c': 100},
var_names='sum')
assert_allclose(ds['sum'].data, np.array([111, 112, 113]))
def test_basic(self):
combos = _test_combos1
ds = combo_runner_to_ds(foo3_scalar, combos, var_names=['bananas'])
assert ds.sel(a=2, b=30, c=400)['bananas'].data == 432
def test_basic(self):
combos = _test_combos1
ds = combo_runner_to_ds(foo3_scalar, combos, var_names=['bananas'])
assert ds.sel(a=2, b=30, c=400)['bananas'].data == 432