diff --git a/doc/api.rst b/doc/api.rst index f0afc322..a7612927 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -248,6 +248,7 @@ Structure expressions.LinearExpression.coeffs expressions.LinearExpression.const expressions.LinearExpression.nterm + expressions.LinearExpression.has_terms Conversion ---------- @@ -288,6 +289,7 @@ Structure expressions.QuadraticExpression.coeffs expressions.QuadraticExpression.const expressions.QuadraticExpression.nterm + expressions.QuadraticExpression.has_terms Conversion ---------- diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ca582462..58c86a72 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -5,6 +5,7 @@ Upcoming Version ---------------- * Add documentation about `LinearExpression.where` with `drop=True`. Add `BaseExpression.variable_names` property. +* Add ``BaseExpression.has_terms`` property: boolean array, true at slots with at least one live term (`#741 `_). **Features** diff --git a/linopy/expressions.py b/linopy/expressions.py index 673eaba9..7d40fad1 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1338,6 +1338,37 @@ def nterm(self) -> int: """ return len(self.data._term) + @property + def has_terms(self) -> DataArray: + """ + Get a boolean array which is true at slots with at least one live term. + + A term is live when it references a variable (``vars != -1``). Slots + without any live term arise from outer joins in + :func:`merge `, from reindexing past the + original coordinates, or from masking. In contrast to + :meth:`isnull`, the constant is ignored: a slot carrying only a + constant has no terms. + + Returns + ------- + xr.DataArray + + Examples + -------- + Mask out constraint rows whose left-hand side has no terms: + + >>> import linopy + >>> import pandas as pd + >>> m = linopy.Model() + >>> x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + >>> lhs = (1 * x).reindex(i=pd.RangeIndex(5, name="i")) + >>> lhs.has_terms.values + array([ True, True, True, False, False]) + """ + helper_dims = set(self.vars.dims).intersection(HELPER_DIMS) + return (self.vars != -1).any(helper_dims).rename("has_terms") + @property def variable_names(self) -> set[str]: """ diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 82aba70e..2580f033 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -1156,6 +1156,56 @@ def test_linear_expression_isnull(v: Variable) -> None: assert expr.isnull().sum() == 10 +class TestHasTerms: + """has_terms: true at slots with at least one live term, regardless of the constant.""" + + def test_basic_and_masking(self, v: Variable) -> None: + expr = np.arange(20) * v + assert expr.has_terms.all() + + filter = (expr.coeffs >= 10).any(TERM_DIM) + masked = expr.where(filter) + assert_equal(masked.has_terms, filter.rename("has_terms")) + + def test_ignores_const(self, v: Variable) -> None: + # has_terms differs from isnull() at slots whose constant was revived by + # fillna: no longer null, but still without terms + expr = np.arange(20) * v + filter = (expr.coeffs >= 10).any(TERM_DIM) + masked = expr.where(filter) + assert_equal(masked.isnull(), ~masked.has_terms) + + filled = masked.fillna(0) + assert not filled.isnull().any() + assert_equal(filled.has_terms, filter.rename("has_terms")) + + def test_merge_reindex(self, x: Variable, y: Variable) -> None: + # the nodal-balance pattern: outer merge, then reindex to a superset of + # coordinates; slots beyond the original coordinates carry no terms + lhs = merge([1 * x, 1 * y], join="outer").reindex( + dim_0=pd.RangeIndex(4, name="dim_0") + ) + assert lhs.has_terms.values.tolist() == [True, True, False, False] + + def test_constant_only(self, m: Model) -> None: + expr = LinearExpression(xr.DataArray([1, 2], dims=["dim_0"]), m) + assert expr.nterm == 0 + assert not expr.has_terms.any() + + def test_quadratic(self, v: Variable) -> None: + # linear terms inside a quadratic expression carry one factor == -1; + # they must still count as live terms + quad = v * v + 2 * v + assert quad.has_terms.all() + assert TERM_DIM not in quad.has_terms.dims + + filter = xr.DataArray( + np.arange(20) >= 10, dims="dim_2", coords={"dim_2": range(20)} + ) + masked = quad.where(filter) + assert_equal(masked.has_terms, filter.rename("has_terms")) + + def test_linear_expression_flat(v: Variable) -> None: coeff = np.arange(1, 21) # use non-zero coefficients expr = coeff * v