def test_subset_imask_squeeze(self, model_data, imask_subset_config):
"""
Include an additional dimension in 'where', which we then subset on (and squeeze out)
"""
# foreach doesn't have this additional dimension
foreach = ["techs"]
imask = _imask_foreach(model_data, foreach)
assert imask.dims == ("techs",)
# on using 'where', the 'nodes' dimension is added
imask = _imask_where(model_data, "foo", ["node_tech"], imask, "and_")
assert sorted(imask.dims) == sorted(["nodes", "techs"])
imask_subset = _subset_imask("foo", imask_subset_config(foreach), imask)
assert imask_subset.dims == ("techs",)
assert imask_subset.equals(imask.loc[{"nodes": "foo"}].drop_vars("nodes"))
def test_imask_where_initial_imask(
self, model_data, where_array, results, initial_operator
):
foreach = ["nodes", "techs", "carriers"]
imask = _imask_foreach(model_data, foreach)
where_imask = _imask_where(
model_data, "foo", where_array, imask, initial_operator
)
locs = [{"nodes": "foo", "techs": "foo"}, {"nodes": "bar", "techs": "foo"}]
for i in range(2):
if initial_operator == "and_":
assert (
imask.loc[locs[i]] * results[i] == where_imask.loc[locs[i]]
).all()
elif initial_operator == "or_":
assert (
imask.loc[locs[i]] + results[i] == where_imask.loc[locs[i]]
).all()
def test_imask_where_incorrect_where(self, model_data):
with pytest.raises(ValueError) as excinfo:
_imask_where(model_data, "foo", ["node_tech", "inheritance(bar)"])
assert check_error_or_warning(excinfo, "'where' array for set `foo` must")
def test_imask_where_array(self, model_data, where_array, results):
where_imask = _imask_where(model_data, "foo", where_array)
assert (
where_imask
== [[results[0], False, False, False], [results[1], False, False, False]]
).all()
def test_imask_where_not(self, model_data):
where_imask = _imask_where(model_data, "foo", ["with_inf"])
not_where_imask = _imask_where(model_data, "foo", ["not with_inf"])
assert where_imask.equals(~not_where_imask)
def test_imask_where_val_is(self, model_data, val, result):
where_imask = _imask_where(model_data, "foo", [val])
assertion = where_imask == result
if isinstance(result, list):
assertion = assertion.all()
assert assertion
def test_imask_where_param_exists(self, model_data, param):
where_imask = _imask_where(model_data, "foo", [param])
if param in model_data:
assert (where_imask == _param_exists(model_data, param)).all()
else:
assert where_imask is False
def test_imask_where_ineritance(self, model_data):
where_imask = _imask_where(model_data, "foo", ["inheritance(bar)"])
assert (
where_imask == [[True, False, False, False], [True, False, False, False]]
).all()