Skip to content

Commit

Permalink
fix: unsupported dimension names (PyPSA#359)
Browse files Browse the repository at this point in the history
Add function to ensure that added variables do not use unvalid dimension names.
  • Loading branch information
tom-welfonder committed Oct 10, 2024
1 parent e11c575 commit b88c434
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 0 deletions.
32 changes: 32 additions & 0 deletions linopy/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -357,6 +357,37 @@ def check_force_dim_names(self, ds: DataArray | Dataset) -> None:
else:
return

def check_valid_dim_names(self, ds: DataArray | Dataset) -> None:
"""
Ensure that the added data does not lead to a naming conflict.
Parameters
----------
model : linopy.Model
ds : xr.DataArray/Variable/LinearExpression
Data that should be added to the model.
Raises
------
ValueError
If broadcasted data leads to unsupported dimension names.
Returns
-------
None.
"""
unsupported_dim_names = ["labels", "coeffs", "vars", "sign", "rhs"]
contains_unsupported_dim_names = any(
dim in unsupported_dim_names for dim in list(ds.dims)
)
if contains_unsupported_dim_names:
raise ValueError(
"Added data contains unsupported dimension names. "
"Dimensions cannot be named 'labels', 'coeffs', 'vars', 'sign' or 'rhs'."
)
else:
return

def add_variables(
self,
lower: Any = -inf,
Expand Down Expand Up @@ -464,6 +495,7 @@ def add_variables(
)
(data,) = xr.broadcast(data)
self.check_force_dim_names(data)
self.check_valid_dim_names(data)

if mask is not None:
mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool)
Expand Down
18 changes: 18 additions & 0 deletions test/test_variable_assignment.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,15 @@ def test_variable_assignment_without_coords_and_dims_names():
assert x.dims == ("i", "j")


def test_variable_assignment_without_coords_and_invalid_dims_names():
# setting bounds without explicit coords
m = Model()
lower = np.zeros((10, 10))
upper = np.ones((10, 10))
with pytest.raises(ValueError):
m.add_variables(lower, upper, name="x", dims=["sign", "j"])


def test_variable_assignment_without_coords_in_bounds():
# setting bounds without explicit coords
m = Model()
Expand All @@ -141,6 +150,15 @@ def test_variable_assignment_without_coords_in_bounds():
assert x.dims == ("i", "j")


def test_variable_assignment_without_coords_in_bounds_invalid_dims_names():
# setting bounds without explicit coords
m = Model()
lower = xr.DataArray(np.zeros((10, 10)), dims=["i", "sign"])
upper = xr.DataArray(np.ones((10, 10)), dims=["i", "sign"])
with pytest.raises(ValueError):
m.add_variables(lower, upper, name="x")


def test_variable_assignment_without_coords_pandas_types():
# setting bounds without explicit coords
m = Model()
Expand Down

0 comments on commit b88c434

Please sign in to comment.