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()