diff --git a/CLAUDE.md b/CLAUDE.md index 67155ae3..1f696a0b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -110,27 +110,6 @@ When modifying the codebase, maintain consistency with these patterns and ensure * Always create a feature branch for new features or bug fixes. * Use the github cli (gh) to interact with the Github repository. -### GitHub Claude Code Integration - -This repository includes Claude Code GitHub Actions for automated assistance: - -1. **Automated PR Reviews** (`claude-code-review.yml`): - - Automatically reviews PRs only when first created (opened) - - Subsequent reviews require manual `@claude` mention - - Focuses on Python best practices, xarray patterns, and optimization correctness - - Can run tests and linting as part of the review - - **Skip initial review by**: Adding `[skip-review]` or `[WIP]` to PR title, or using draft PRs - -2. **Manual Claude Assistance** (`claude.yml`): - - Trigger by mentioning `@claude` in any: - - Issue comments - - Pull request comments - - Pull request reviews - - New issue body or title - - Claude can help with bug fixes, feature implementation, code explanations, etc. - -**Note**: Both workflows require the `ANTHROPIC_API_KEY` secret to be configured in the repository settings. - ## Development Guidelines @@ -140,3 +119,4 @@ This repository includes Claude Code GitHub Actions for automated assistance: 4. Use type hints and mypy for type checking. 5. Always write tests into the `test` directory, following the naming convention `test_*.py`. 6. Always write temporary and non git-tracked code in the `dev-scripts` directory. +7. In test scripts use linopy assertions from the testing.py module where useful (assert_linequal, assert_varequal, etc.) diff --git a/doc/index.rst b/doc/index.rst index bff9fa65..c575fc60 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -111,6 +111,7 @@ This package is published under MIT license. creating-variables creating-expressions creating-constraints + coordinate-alignment sos-constraints manipulating-models testing-framework diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 13df0267..b75c8728 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -6,6 +6,12 @@ Upcoming Version * Fix docs (pick highs solver) * Add the `sphinx-copybutton` to the documentation +* Harmonize coordinate alignment for operations with subset/superset objects: + - Multiplication and division fill missing coords with 0 (variable doesn't participate) + - Addition and subtraction of constants fill missing coords with 0 (identity element) and pin result to LHS coords + - Comparison operators (``==``, ``<=``, ``>=``) fill missing RHS coords with NaN (no constraint created) + - Fixes crash on ``subset + var`` / ``subset + expr`` reverse addition + - Fixes superset DataArrays expanding result coords beyond the variable's coordinate space * Add ``auto_mask`` parameter to ``Model`` class that automatically masks variables and constraints where bounds, coefficients, or RHS values contain NaN. This eliminates the need to manually create mask arrays when working with sparse or incomplete data. * Speed up LP file writing by 2-2.7x on large models through Polars streaming engine, join-based constraint assembly, and reduced per-constraint overhead * Fix multiplication of constant-only ``LinearExpression`` with other expressions diff --git a/examples/coordinate-alignment.ipynb b/examples/coordinate-alignment.ipynb new file mode 100644 index 00000000..1547bd9d --- /dev/null +++ b/examples/coordinate-alignment.ipynb @@ -0,0 +1,488 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Coordinate Alignment\n", + "\n", + "Since linopy builds on xarray, coordinate alignment matters when combining variables or expressions that live on different coordinates. By default, linopy aligns operands automatically and fills missing entries with sensible defaults. This guide shows how alignment works and how to control it with the ``join`` parameter." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import xarray as xr\n", + "\n", + "import linopy" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Default Alignment Behavior\n", + "\n", + "When two operands share a dimension but have different coordinates, linopy keeps the **larger** (superset) coordinate range and fills missing positions with zeros (for addition) or zero coefficients (for multiplication)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "\n", + "time = pd.RangeIndex(5, name=\"time\")\n", + "x = m.add_variables(lower=0, coords=[time], name=\"x\")\n", + "\n", + "subset_time = pd.RangeIndex(3, name=\"time\")\n", + "y = m.add_variables(lower=0, coords=[subset_time], name=\"y\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Adding ``x`` (5 time steps) and ``y`` (3 time steps) gives an expression over all 5 time steps. Where ``y`` has no entry (time 3, 4), the coefficient is zero — i.e. ``y`` simply drops out of the sum at those positions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "x + y" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The same applies when multiplying by a constant that covers only a subset of coordinates. Missing positions get a coefficient of zero:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "factor = xr.DataArray([2, 3, 4], dims=[\"time\"], coords={\"time\": [0, 1, 2]})\n", + "x * factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Adding a constant subset also fills missing coordinates with zero:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "x + factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Constraints with Subset RHS\n", + "\n", + "For constraints, missing right-hand-side values are filled with ``NaN``, which tells linopy to **skip** the constraint at those positions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rhs = xr.DataArray([10, 20, 30], dims=[\"time\"], coords={\"time\": [0, 1, 2]})\n", + "con = x <= rhs\n", + "con" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The constraint only applies at time 0, 1, 2. At time 3 and 4 the RHS is ``NaN``, so no constraint is created." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Same-Shape Operands: Positional Alignment\n\nWhen two operands have the **same shape** on a shared dimension, linopy uses **positional alignment** by default — coordinate labels are ignored and the left operand's labels are kept. This is a performance optimization but can be surprising:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "offset_const = xr.DataArray(\n", + " [10, 20, 30, 40, 50], dims=[\"time\"], coords={\"time\": [5, 6, 7, 8, 9]}\n", + ")\n", + "x + offset_const" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "Even though ``offset_const`` has coordinates ``[5, 6, 7, 8, 9]`` and ``x`` has ``[0, 1, 2, 3, 4]``, the result uses ``x``'s labels. The values are aligned by **position**, not by label. The same applies when adding two variables or expressions of identical shape:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "z = m.add_variables(lower=0, coords=[pd.RangeIndex(5, 10, name=\"time\")], name=\"z\")\n", + "x + z" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "``x`` (time 0–4) and ``z`` (time 5–9) share no coordinate labels, yet the result has 5 entries under ``x``'s coordinates — because they have the same shape, positions are matched directly.\n\nTo force **label-based** alignment, pass an explicit ``join``:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "x.add(z, join=\"outer\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "With ``join=\"outer\"``, the result spans all 10 time steps (union of 0–4 and 5–9), filling missing positions with zeros. This is the correct label-based alignment. The same-shape positional shortcut is equivalent to ``join=\"override\"`` — see below." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## The ``join`` Parameter\n", + "\n", + "For explicit control over alignment, use the ``.add()``, ``.sub()``, ``.mul()``, and ``.div()`` methods with a ``join`` parameter. The supported values follow xarray conventions:\n", + "\n", + "- ``\"inner\"`` — intersection of coordinates\n", + "- ``\"outer\"`` — union of coordinates (with fill)\n", + "- ``\"left\"`` — keep left operand's coordinates\n", + "- ``\"right\"`` — keep right operand's coordinates\n", + "- ``\"override\"`` — positional alignment, ignore coordinate labels\n", + "- ``\"exact\"`` — coordinates must match exactly (raises on mismatch)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m2 = linopy.Model()\n", + "\n", + "i_a = pd.Index([0, 1, 2], name=\"i\")\n", + "i_b = pd.Index([1, 2, 3], name=\"i\")\n", + "\n", + "a = m2.add_variables(coords=[i_a], name=\"a\")\n", + "b = m2.add_variables(coords=[i_b], name=\"b\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Inner join** — only shared coordinates (i=1, 2):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.add(b, join=\"inner\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Outer join** — union of coordinates (i=0, 1, 2, 3):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.add(b, join=\"outer\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Left join** — keep left operand's coordinates (i=0, 1, 2):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.add(b, join=\"left\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Right join** — keep right operand's coordinates (i=1, 2, 3):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.add(b, join=\"right\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "**Override** — positional alignment, ignore coordinate labels. The result uses the left operand's coordinates. Here ``a`` has i=[0, 1, 2] and ``b`` has i=[1, 2, 3], so positions are matched as 0↔1, 1↔2, 2↔3:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.add(b, join=\"override\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Multiplication with ``join``\n", + "\n", + "The same ``join`` parameter works on ``.mul()`` and ``.div()``. When multiplying by a constant that covers a subset, ``join=\"inner\"`` restricts the result to shared coordinates only, while ``join=\"left\"`` fills missing values with zero:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "const = xr.DataArray([2, 3, 4], dims=[\"i\"], coords={\"i\": [1, 2, 3]})\n", + "\n", + "a.mul(const, join=\"inner\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.mul(const, join=\"left\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Alignment in Constraints\n", + "\n", + "The ``.le()``, ``.ge()``, and ``.eq()`` methods create constraints with explicit coordinate alignment. They accept the same ``join`` parameter:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rhs = xr.DataArray([10, 20], dims=[\"i\"], coords={\"i\": [0, 1]})\n", + "\n", + "a.le(rhs, join=\"inner\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With ``join=\"inner\"``, the constraint only exists at the intersection (i=0, 1). Compare with ``join=\"left\"``:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a.le(rhs, join=\"left\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With ``join=\"left\"``, the result covers all of ``a``'s coordinates (i=0, 1, 2). At i=2, where the RHS has no value, the RHS becomes ``NaN`` and the constraint is masked out.\n", + "\n", + "The same methods work on expressions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expr = 2 * a + 1\n", + "expr.eq(rhs, join=\"inner\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Practical Example\n\nConsider a generation dispatch model where solar availability follows a daily profile and a minimum demand constraint only applies during peak hours." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m3 = linopy.Model()\n", + "\n", + "hours = pd.RangeIndex(24, name=\"hour\")\n", + "techs = pd.Index([\"solar\", \"wind\", \"gas\"], name=\"tech\")\n", + "\n", + "gen = m3.add_variables(lower=0, coords=[hours, techs], name=\"gen\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Capacity limits apply to all hours and techs — standard broadcasting handles this:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "capacity = xr.DataArray([100, 80, 50], dims=[\"tech\"], coords={\"tech\": techs})\n", + "m3.add_constraints(gen <= capacity, name=\"capacity_limit\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "For solar, we build a full 24-hour availability profile — zero at night, sine-shaped during daylight (hours 6–18). Since this covers all hours, standard alignment works directly and solar is properly constrained to zero at night:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solar_avail = np.zeros(24)\n", + "solar_avail[6:19] = 100 * np.sin(np.linspace(0, np.pi, 13))\n", + "solar_availability = xr.DataArray(solar_avail, dims=[\"hour\"], coords={\"hour\": hours})\n", + "\n", + "solar_gen = gen.sel(tech=\"solar\")\n", + "m3.add_constraints(solar_gen <= solar_availability, name=\"solar_avail\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "Now suppose a minimum demand of 120 MW must be met, but only during peak hours (8–20). The demand array covers a subset of hours, so we use ``join=\"inner\"`` to restrict the constraint to just those hours:" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "peak_hours = pd.RangeIndex(8, 21, name=\"hour\")\n", + "peak_demand = xr.DataArray(\n", + " np.full(len(peak_hours), 120.0), dims=[\"hour\"], coords={\"hour\": peak_hours}\n", + ")\n", + "\n", + "total_gen = gen.sum(\"tech\")\n", + "m3.add_constraints(total_gen.ge(peak_demand, join=\"inner\"), name=\"peak_demand\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "The demand constraint only applies during peak hours (8–20). Outside that range, no minimum generation is required." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "| ``join`` | Coordinates | Fill behavior |\n", + "|----------|------------|---------------|\n", + "| ``None`` (default) | Auto-detect (keeps superset) | Zeros for arithmetic, NaN for constraint RHS |\n", + "| ``\"inner\"`` | Intersection only | No fill needed |\n", + "| ``\"outer\"`` | Union | Fill with operation identity (0 for add, 0 for mul) |\n", + "| ``\"left\"`` | Left operand's | Fill right with identity |\n", + "| ``\"right\"`` | Right operand's | Fill left with identity |\n", + "| ``\"override\"`` | Left operand's (positional) | Positional alignment, ignore labels |\n", + "| ``\"exact\"`` | Must match exactly | Raises error if different |" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/creating-constraints.ipynb b/examples/creating-constraints.ipynb index b46db1bc..55251233 100644 --- a/examples/creating-constraints.ipynb +++ b/examples/creating-constraints.ipynb @@ -231,6 +231,12 @@ "source": [ "m.constraints[\"my-constraint\"]" ] + }, + { + "cell_type": "markdown", + "id": "r0wxi7v1m7l", + "source": "## Coordinate Alignment in Constraints\n\nAs an alternative to the ``<=``, ``>=``, ``==`` operators, linopy provides ``.le()``, ``.ge()``, and ``.eq()`` methods on variables and expressions. These methods accept a ``join`` parameter (``\"inner\"``, ``\"outer\"``, ``\"left\"``, ``\"right\"``) for explicit control over how coordinates are aligned when creating constraints. See the :doc:`coordinate-alignment` guide for details.", + "metadata": {} } ], "metadata": { diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index aafd8a09..1d808b07 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -193,6 +193,12 @@ "x + b" ] }, + { + "cell_type": "markdown", + "id": "a8xsfdqrcrn", + "source": ".. tip::\n\n\tFor explicit control over how coordinates are aligned during arithmetic, use the `.add()`, `.sub()`, `.mul()`, and `.div()` methods with a ``join`` parameter (``\"inner\"``, ``\"outer\"``, ``\"left\"``, ``\"right\"``). See the :doc:`coordinate-alignment` guide for details.", + "metadata": {} + }, { "attachments": {}, "cell_type": "markdown", diff --git a/linopy/expressions.py b/linopy/expressions.py index 649989f7..ee56b477 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -9,6 +9,7 @@ import functools import logging +import operator from abc import ABC, abstractmethod from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence from dataclasses import dataclass, field @@ -93,17 +94,6 @@ from linopy.model import Model from linopy.variables import ScalarVariable, Variable -SUPPORTED_CONSTANT_TYPES = ( - np.number, - int, - float, - DataArray, - pd.Series, - pd.DataFrame, - np.ndarray, - pl.Series, -) - FILL_VALUE = {"vars": -1, "coeffs": np.nan, "const": np.nan} @@ -533,31 +523,107 @@ def _multiply_by_linear_expression( res = res + self.reset_const() * other.const return res + def _align_constant( + self: GenericExpression, + other: DataArray, + fill_value: float = 0, + join: str | None = None, + ) -> tuple[DataArray, DataArray, bool]: + """ + Align a constant DataArray with self.const. + + Parameters + ---------- + other : DataArray + The constant to align. + fill_value : float, default: 0 + Fill value for missing coordinates. + join : str, optional + Alignment method. If None, uses size-aware default behavior. + + Returns + ------- + self_const : DataArray + The expression's const, potentially reindexed. + aligned : DataArray + The aligned constant. + needs_data_reindex : bool + Whether the expression's data needs reindexing. + """ + if join is None: + if other.sizes == self.const.sizes: + return self.const, other.assign_coords(coords=self.coords), False + return ( + self.const, + other.reindex_like(self.const, fill_value=fill_value), + False, + ) + elif join == "override": + return self.const, other.assign_coords(coords=self.coords), False + else: + self_const, aligned = xr.align( + self.const, other, join=join, fill_value=fill_value + ) + return self_const, aligned, True + + def _add_constant( + self: GenericExpression, other: ConstantLike, join: str | None = None + ) -> GenericExpression: + if np.isscalar(other) and join is None: + return self.assign(const=self.const + other) + da = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + self_const, da, needs_data_reindex = self._align_constant( + da, fill_value=0, join=join + ) + if needs_data_reindex: + return self.__class__( + self.data.reindex_like(self_const, fill_value=self._fill_value).assign( + const=self_const + da + ), + self.model, + ) + return self.assign(const=self_const + da) + + def _apply_constant_op( + self: GenericExpression, + other: ConstantLike, + op: Callable[[DataArray, DataArray], DataArray], + fill_value: float, + join: str | None = None, + ) -> GenericExpression: + factor = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + self_const, factor, needs_data_reindex = self._align_constant( + factor, fill_value=fill_value, join=join + ) + if needs_data_reindex: + data = self.data.reindex_like(self_const, fill_value=self._fill_value) + return self.__class__( + assign_multiindex_safe( + data, coeffs=op(data.coeffs, factor), const=op(self_const, factor) + ), + self.model, + ) + return self.assign(coeffs=op(self.coeffs, factor), const=op(self_const, factor)) + def _multiply_by_constant( - self: GenericExpression, other: ConstantLike + self: GenericExpression, other: ConstantLike, join: str | None = None + ) -> GenericExpression: + return self._apply_constant_op(other, operator.mul, fill_value=0, join=join) + + def _divide_by_constant( + self: GenericExpression, other: ConstantLike, join: str | None = None ) -> GenericExpression: - multiplier = as_dataarray(other, coords=self.coords, dims=self.coord_dims) - coeffs = self.coeffs * multiplier - assert all(coeffs.sizes[d] == s for d, s in self.coeffs.sizes.items()) - const = self.const * multiplier - return self.assign(coeffs=coeffs, const=const) + return self._apply_constant_op(other, operator.truediv, fill_value=1, join=join) def __div__(self: GenericExpression, other: SideLike) -> GenericExpression: try: - if isinstance( - other, - variables.Variable - | variables.ScalarVariable - | LinearExpression - | ScalarLinearExpression - | QuadraticExpression, - ): + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): raise TypeError( "unsupported operand type(s) for /: " f"{type(self)} and {type(other)}" "Non-linear expressions are not yet supported." ) - return self._multiply_by_constant(other=1 / other) + return self._divide_by_constant(other) except TypeError: return NotImplemented @@ -584,36 +650,160 @@ def __lt__(self, other: Any) -> NotImplementedType: ) def add( - self: GenericExpression, other: SideLike + self: GenericExpression, + other: SideLike, + join: str | None = None, ) -> GenericExpression | QuadraticExpression: """ Add an expression to others. - """ - return self.__add__(other) + + Parameters + ---------- + other : expression-like + The expression to add. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + if join is None: + return self.__add__(other) + if isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other, join=join) + other = as_expression(other, model=self.model, dims=self.coord_dims) + if isinstance(other, LinearExpression) and isinstance( + self, QuadraticExpression + ): + other = other.to_quadexpr() + return merge([self, other], cls=self.__class__, join=join) def sub( - self: GenericExpression, other: SideLike + self: GenericExpression, + other: SideLike, + join: str | None = None, ) -> GenericExpression | QuadraticExpression: """ Subtract others from expression. + + Parameters + ---------- + other : expression-like + The expression to subtract. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. """ - return self.__sub__(other) + return self.add(-other, join=join) def mul( - self: GenericExpression, other: SideLike + self: GenericExpression, + other: SideLike, + join: str | None = None, ) -> GenericExpression | QuadraticExpression: """ Multiply the expr by a factor. - """ - return self.__mul__(other) + + Parameters + ---------- + other : expression-like + The factor to multiply by. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + if join is None: + return self.__mul__(other) + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): + raise TypeError( + "join parameter is not supported for expression-expression multiplication" + ) + return self._multiply_by_constant(other, join=join) def div( - self: GenericExpression, other: VariableLike | ConstantLike + self: GenericExpression, + other: VariableLike | ConstantLike, + join: str | None = None, ) -> GenericExpression | QuadraticExpression: """ Divide the expr by a factor. + + Parameters + ---------- + other : constant-like + The divisor. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + if join is None: + return self.__div__(other) + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): + raise TypeError( + "unsupported operand type(s) for /: " + f"{type(self)} and {type(other)}. " + "Non-linear expressions are not yet supported." + ) + return self._divide_by_constant(other, join=join) + + def le( + self: GenericExpression, + rhs: SideLike, + join: str | None = None, + ) -> Constraint: """ - return self.__div__(other) + Less than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(LESS_EQUAL, rhs, join=join) + + def ge( + self: GenericExpression, + rhs: SideLike, + join: str | None = None, + ) -> Constraint: + """ + Greater than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(GREATER_EQUAL, rhs, join=join) + + def eq( + self: GenericExpression, + rhs: SideLike, + join: str | None = None, + ) -> Constraint: + """ + Equality constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_constraint(EQUAL, rhs, join=join) def pow(self, other: int) -> QuadraticExpression: """ @@ -854,7 +1044,9 @@ def cumsum( dim_dict = {dim_name: self.data.sizes[dim_name] for dim_name in dim} return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) - def to_constraint(self, sign: SignLike, rhs: SideLike) -> Constraint: + def to_constraint( + self, sign: SignLike, rhs: SideLike, join: str | None = None + ) -> Constraint: """ Convert a linear expression to a constraint. @@ -863,7 +1055,14 @@ def to_constraint(self, sign: SignLike, rhs: SideLike) -> Constraint: sign : str, array-like Sign(s) of the constraints. rhs : constant, Variable, LinearExpression - Right-hand side of the constraint. + Right-hand side of the constraint. If a DataArray, it is + reindexed to match expression coordinates (fill_value=np.nan). + Extra dimensions in the RHS not present in the expression + raise a ValueError. NaN entries in the RHS mean "no constraint". + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. Returns ------- @@ -876,7 +1075,16 @@ def to_constraint(self, sign: SignLike, rhs: SideLike) -> Constraint: f"Both sides of the constraint are constant. At least one side must contain variables. {self} {rhs}" ) - all_to_lhs = (self - rhs).data + if isinstance(rhs, DataArray): + extra_dims = set(rhs.dims) - set(self.coord_dims) + if extra_dims: + raise ValueError( + f"RHS DataArray has dimensions {extra_dims} not present " + f"in the expression. Cannot create constraint." + ) + rhs = rhs.reindex_like(self.const, fill_value=np.nan) + + all_to_lhs = self.sub(rhs, join=join).data data = assign_multiindex_safe( all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=-all_to_lhs.const ) @@ -1312,11 +1520,11 @@ def __add__( return other.__add__(self) try: - if np.isscalar(other): - return self.assign(const=self.const + other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) - return merge([self, other], cls=self.__class__) + if isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other) + else: + other = as_expression(other, model=self.model, dims=self.coord_dims) + return merge([self, other], cls=self.__class__) except TypeError: return NotImplemented @@ -1824,13 +2032,7 @@ def __mul__(self, other: SideLike) -> QuadraticExpression: """ Multiply the expr by a factor. """ - if isinstance( - other, - BaseExpression - | ScalarLinearExpression - | variables.Variable - | variables.ScalarVariable, - ): + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): raise TypeError( "unsupported operand type(s) for *: " f"{type(self)} and {type(other)}. " @@ -1852,15 +2054,15 @@ def __add__(self, other: SideLike) -> QuadraticExpression: dimension names of self will be filled in other """ try: - if np.isscalar(other): - return self.assign(const=self.const + other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) + if isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other) + else: + other = as_expression(other, model=self.model, dims=self.coord_dims) - if isinstance(other, LinearExpression): - other = other.to_quadexpr() + if isinstance(other, LinearExpression): + other = other.to_quadexpr() - return merge([self, other], cls=self.__class__) + return merge([self, other], cls=self.__class__) except TypeError: return NotImplemented @@ -1878,13 +2080,7 @@ def __sub__(self, other: SideLike) -> QuadraticExpression: dimension names of self will be filled in other """ try: - if np.isscalar(other): - return self.assign(const=self.const - other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) - if type(other) is LinearExpression: - other = other.to_quadexpr() - return merge([self, -other], cls=self.__class__) + return self.__add__(-other) except TypeError: return NotImplemented @@ -1906,13 +2102,7 @@ def __matmul__( """ Matrix multiplication with other, similar to xarray dot. """ - if isinstance( - other, - BaseExpression - | ScalarLinearExpression - | variables.Variable - | variables.ScalarVariable, - ): + if isinstance(other, SUPPORTED_EXPRESSION_TYPES): raise TypeError( "Higher order non-linear expressions are not yet supported." ) @@ -2065,6 +2255,7 @@ def merge( ], dim: str = TERM_DIM, cls: type[GenericExpression] = None, # type: ignore + join: str | None = None, **kwargs: Any, ) -> GenericExpression: """ @@ -2084,6 +2275,10 @@ def merge( Dimension along which the expressions should be concatenated. cls : type Explicitly set the type of the resulting expression (So that the type checker will know the return type) + join : str, optional + How to align coordinates. One of "outer", "inner", "left", "right", + "exact", "override". When None (default), auto-detects based on + expression shapes. **kwargs Additional keyword arguments passed to xarray.concat. Defaults to {coords: "minimal", compat: "override"} or, in the special case described @@ -2118,7 +2313,9 @@ def merge( model = exprs[0].model - if cls in linopy_types and dim in HELPER_DIMS: + if join is not None: + override = join == "override" + elif cls in linopy_types and dim in HELPER_DIMS: coord_dims = [ {k: v for k, v in e.sizes.items() if k not in HELPER_DIMS} for e in exprs ] @@ -2139,7 +2336,9 @@ def merge( elif cls == variables.Variable: kwargs["fill_value"] = variables.FILL_VALUE - if override: + if join is not None: + kwargs["join"] = join + elif override: kwargs["join"] = "override" else: kwargs.setdefault("join", "outer") @@ -2317,3 +2516,23 @@ def to_linexpr(self) -> LinearExpression: vars = xr.DataArray(list(self.vars), dims=TERM_DIM) ds = xr.Dataset({"coeffs": coeffs, "vars": vars}) return LinearExpression(ds, self.model) + + +SUPPORTED_CONSTANT_TYPES = ( + np.number, + np.bool_, + int, + float, + DataArray, + pd.Series, + pd.DataFrame, + np.ndarray, + pl.Series, +) + +SUPPORTED_EXPRESSION_TYPES = ( + BaseExpression, + ScalarLinearExpression, + variables.Variable, + variables.ScalarVariable, +) diff --git a/linopy/model.py b/linopy/model.py index a2fa8e4e..a8ef8782 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -741,6 +741,16 @@ def add_constraints( # TODO: add a warning here, routines should be safe against this data = data.drop_vars(drop_dims) + rhs_nan = data.rhs.isnull() + if rhs_nan.any(): + data["rhs"] = data.rhs.fillna(0) + rhs_mask = ~rhs_nan + mask = ( + rhs_mask + if mask is None + else (as_dataarray(mask).astype(bool) & rhs_mask) + ) + data["labels"] = -1 (data,) = xr.broadcast(data, exclude=[TERM_DIM]) diff --git a/linopy/monkey_patch_xarray.py b/linopy/monkey_patch_xarray.py index dc60608c..1e526c92 100644 --- a/linopy/monkey_patch_xarray.py +++ b/linopy/monkey_patch_xarray.py @@ -1,37 +1,45 @@ from __future__ import annotations from collections.abc import Callable -from functools import partialmethod, update_wrapper -from types import NotImplementedType +from functools import update_wrapper from typing import Any from xarray import DataArray from linopy import expressions, variables - -def monkey_patch(cls: type[DataArray], pass_unpatched_method: bool = False) -> Callable: - def deco(func: Callable) -> Callable: - func_name = func.__name__ - wrapped = getattr(cls, func_name) - update_wrapper(func, wrapped) - if pass_unpatched_method: - func = partialmethod(func, unpatched_method=wrapped) # type: ignore - setattr(cls, func_name, func) - return func - - return deco - - -@monkey_patch(DataArray, pass_unpatched_method=True) -def __mul__( - da: DataArray, other: Any, unpatched_method: Callable -) -> DataArray | NotImplementedType: - if isinstance( - other, - variables.Variable - | expressions.LinearExpression - | expressions.QuadraticExpression, - ): - return NotImplemented - return unpatched_method(da, other) +_LINOPY_TYPES = ( + variables.Variable, + variables.ScalarVariable, + expressions.LinearExpression, + expressions.ScalarLinearExpression, + expressions.QuadraticExpression, +) + + +def _make_patched_op(op_name: str) -> None: + """Patch a DataArray operator to return NotImplemented for linopy types, enabling reflected operators.""" + original = getattr(DataArray, op_name) + + def patched( + da: DataArray, other: Any, unpatched_method: Callable = original + ) -> Any: + if isinstance(other, _LINOPY_TYPES): + return NotImplemented + return unpatched_method(da, other) + + update_wrapper(patched, original) + setattr(DataArray, op_name, patched) + + +for _op in ( + "__mul__", + "__add__", + "__sub__", + "__truediv__", + "__le__", + "__ge__", + "__eq__", +): + _make_patched_op(_op) +del _op diff --git a/linopy/variables.py b/linopy/variables.py index d90a4775..3ba563da 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -315,6 +315,7 @@ def to_linexpr( Linear expression with the variables and coefficients. """ coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) + coefficient = coefficient.reindex_like(self.labels, fill_value=0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( TERM_DIM, -1 ) @@ -454,7 +455,7 @@ def __div__( f"{type(self)} and {type(other)}. " "Non-linear expressions are not yet supported." ) - return self.to_linexpr(1 / other) + return self.to_linexpr()._divide_by_constant(other) def __truediv__( self, coefficient: float | int | LinearExpression | Variable @@ -544,29 +545,110 @@ def __lt__(self, other: Any) -> NotImplementedType: def __contains__(self, value: str) -> bool: return self.data.__contains__(value) - def add(self, other: Variable) -> LinearExpression: + def add(self, other: SideLike, join: str | None = None) -> LinearExpression: """ Add variables to linear expressions or other variables. + + Parameters + ---------- + other : expression-like + The expression to add. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. """ - return self.__add__(other) + return self.to_linexpr().add(other, join=join) - def sub(self, other: Variable) -> LinearExpression: + def sub(self, other: SideLike, join: str | None = None) -> LinearExpression: """ Subtract linear expressions or other variables from the variables. + + Parameters + ---------- + other : expression-like + The expression to subtract. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. """ - return self.__sub__(other) + return self.to_linexpr().sub(other, join=join) - def mul(self, other: int) -> LinearExpression: + def mul(self, other: ConstantLike, join: str | None = None) -> LinearExpression: """ Multiply variables with a coefficient. + + Parameters + ---------- + other : constant-like + The coefficient to multiply by. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. """ - return self.__mul__(other) + return self.to_linexpr().mul(other, join=join) - def div(self, other: int) -> LinearExpression: + def div(self, other: ConstantLike, join: str | None = None) -> LinearExpression: """ Divide variables with a coefficient. + + Parameters + ---------- + other : constant-like + The divisor. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. """ - return self.__div__(other) + return self.to_linexpr().div(other, join=join) + + def le(self, rhs: SideLike, join: str | None = None) -> Constraint: + """ + Less than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_linexpr().le(rhs, join=join) + + def ge(self, rhs: SideLike, join: str | None = None) -> Constraint: + """ + Greater than or equal constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_linexpr().ge(rhs, join=join) + + def eq(self, rhs: SideLike, join: str | None = None) -> Constraint: + """ + Equality constraint. + + Parameters + ---------- + rhs : expression-like + Right-hand side of the constraint. + join : str, optional + How to align coordinates. One of "outer", "inner", "left", + "right", "exact", "override". When None (default), uses the + current default behavior. + """ + return self.to_linexpr().eq(rhs, join=join) def pow(self, other: int) -> QuadraticExpression: """ diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 0da9ec7f..d956ef1f 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -575,6 +575,371 @@ def test_linear_expression_multiplication_invalid( expr / x +class TestSubsetCoordinateAlignment: + @pytest.fixture + def subset(self) -> xr.DataArray: + return xr.DataArray([10.0, 30.0], dims=["dim_2"], coords={"dim_2": [1, 3]}) + + @pytest.fixture + def superset(self) -> xr.DataArray: + return xr.DataArray( + np.arange(25, dtype=float), dims=["dim_2"], coords={"dim_2": range(25)} + ) + + @pytest.fixture + def expected_fill(self) -> np.ndarray: + arr = np.zeros(20) + arr[1] = 10.0 + arr[3] = 30.0 + return arr + + def test_var_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v * subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_expr_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = 1 * v + result = expr * subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + @pytest.mark.parametrize( + "make_lhs,make_rhs", + [ + (lambda v, s: s * v, lambda v, s: v * s), + (lambda v, s: s * (1 * v), lambda v, s: (1 * v) * s), + (lambda v, s: s + v, lambda v, s: v + s), + (lambda v, s: s + (v + 5), lambda v, s: (v + 5) + s), + ], + ids=["subset*var", "subset*expr", "subset+var", "subset+expr"], + ) + def test_commutativity( + self, v: Variable, subset: xr.DataArray, make_lhs: object, make_rhs: object + ) -> None: + assert_linequal(make_lhs(v, subset), make_rhs(v, subset)) + + def test_var_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v + subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill) + + def test_var_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v - subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, -expected_fill) + + def test_subset_sub_var(self, v: Variable, subset: xr.DataArray) -> None: + assert_linequal(subset - v, -v + subset) + + def test_expr_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = v + 5 + result = expr + subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill + 5) + + def test_expr_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = v + 5 + result = expr - subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, 5 - expected_fill) + + def test_subset_sub_expr(self, v: Variable, subset: xr.DataArray) -> None: + expr = v + 5 + assert_linequal(subset - expr, -(expr - subset)) + + def test_var_div_subset(self, v: Variable, subset: xr.DataArray) -> None: + result = v / subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(0.1) + assert result.coeffs.squeeze().sel(dim_2=0).item() == pytest.approx(1.0) + + def test_var_le_subset(self, v: Variable, subset: xr.DataArray) -> None: + con = v <= subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == 10.0 + assert con.rhs.sel(dim_2=3).item() == 30.0 + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) + def test_var_comparison_subset( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = v <= subset + elif sign == ">=": + con = v >= subset + else: + con = v == subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == 10.0 + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + def test_expr_le_subset(self, v: Variable, subset: xr.DataArray) -> None: + expr = v + 5 + con = expr <= subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == pytest.approx(5.0) + assert con.rhs.sel(dim_2=3).item() == pytest.approx(25.0) + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + def test_add_commutativity_full_coords(self, v: Variable) -> None: + full = xr.DataArray( + np.arange(20, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(20)}, + ) + assert_linequal(v + full, full + v) + + def test_superset_addition_pins_to_lhs( + self, v: Variable, superset: xr.DataArray + ) -> None: + result = v + superset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + + def test_superset_add_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset + v, v + superset) + + def test_superset_sub_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset - v, -v + superset) + + def test_superset_mul_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset * v, v * superset) + + @pytest.mark.parametrize("sign", ["<=", ">="]) + def test_superset_comparison_var( + self, v: Variable, superset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = superset <= v + else: + con = superset >= v + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(con.lhs.coeffs.values).any() + assert not np.isnan(con.rhs.values).any() + + def test_disjoint_addition_pins_to_lhs(self, v: Variable) -> None: + disjoint = xr.DataArray( + [100.0, 200.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v + disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, np.zeros(20)) + + def test_expr_div_subset(self, v: Variable, subset: xr.DataArray) -> None: + expr = 1 * v + result = expr / subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(0.1) + assert result.coeffs.squeeze().sel(dim_2=0).item() == pytest.approx(1.0) + + def test_subset_add_var_coefficients( + self, v: Variable, subset: xr.DataArray + ) -> None: + result = subset + v + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.ones(20)) + + def test_subset_sub_var_coefficients( + self, v: Variable, subset: xr.DataArray + ) -> None: + result = subset - v + np.testing.assert_array_equal(result.coeffs.squeeze().values, -np.ones(20)) + + @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) + def test_subset_comparison_var( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = subset <= v + elif sign == ">=": + con = subset >= v + else: + con = subset == v + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert np.isnan(con.rhs.sel(dim_2=0).item()) + assert con.rhs.sel(dim_2=1).item() == pytest.approx(10.0) + + def test_superset_mul_pins_to_lhs( + self, v: Variable, superset: xr.DataArray + ) -> None: + result = v * superset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + + def test_superset_div_pins_to_lhs(self, v: Variable) -> None: + superset_nonzero = xr.DataArray( + np.arange(1, 26, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(25)}, + ) + result = v / superset_nonzero + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + + def test_quadexpr_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr + subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill) + + def test_quadexpr_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr - subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, -expected_fill) + + def test_quadexpr_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr * subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_subset_mul_quadexpr( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = subset * qexpr + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_subset_add_quadexpr(self, v: Variable, subset: xr.DataArray) -> None: + qexpr = v * v + assert_quadequal(subset + qexpr, qexpr + subset) + + def test_multidim_subset_mul(self, m: Model) -> None: + coords_a = pd.RangeIndex(4, name="a") + coords_b = pd.RangeIndex(5, name="b") + w = m.add_variables(coords=[coords_a, coords_b], name="w") + + subset_2d = xr.DataArray( + [[2.0, 3.0], [4.0, 5.0]], + dims=["a", "b"], + coords={"a": [1, 3], "b": [0, 4]}, + ) + result = w * subset_2d + assert result.sizes["a"] == 4 + assert result.sizes["b"] == 5 + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(a=1, b=0).item() == pytest.approx(2.0) + assert result.coeffs.squeeze().sel(a=3, b=4).item() == pytest.approx(5.0) + assert result.coeffs.squeeze().sel(a=0, b=0).item() == pytest.approx(0.0) + assert result.coeffs.squeeze().sel(a=1, b=2).item() == pytest.approx(0.0) + + def test_multidim_subset_add(self, m: Model) -> None: + coords_a = pd.RangeIndex(4, name="a") + coords_b = pd.RangeIndex(5, name="b") + w = m.add_variables(coords=[coords_a, coords_b], name="w") + + subset_2d = xr.DataArray( + [[2.0, 3.0], [4.0, 5.0]], + dims=["a", "b"], + coords={"a": [1, 3], "b": [0, 4]}, + ) + result = w + subset_2d + assert result.sizes["a"] == 4 + assert result.sizes["b"] == 5 + assert not np.isnan(result.const.values).any() + assert result.const.sel(a=1, b=0).item() == pytest.approx(2.0) + assert result.const.sel(a=3, b=4).item() == pytest.approx(5.0) + assert result.const.sel(a=0, b=0).item() == pytest.approx(0.0) + + def test_constraint_rhs_extra_dims_raises(self, v: Variable) -> None: + rhs = xr.DataArray( + [[1.0, 2.0]], dims=["extra", "dim_2"], coords={"dim_2": [0, 1]} + ) + with pytest.raises(ValueError, match="not present in the expression"): + v <= rhs + + def test_da_truediv_var_raises(self, v: Variable) -> None: + da = xr.DataArray(np.ones(20), dims=["dim_2"], coords={"dim_2": range(20)}) + with pytest.raises(TypeError): + da / v + + def test_disjoint_mul_produces_zeros(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v * disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.zeros(20)) + + def test_disjoint_div_preserves_coeffs(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v / disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.ones(20)) + + def test_da_eq_da_still_works(self) -> None: + da1 = xr.DataArray([1, 2, 3]) + da2 = xr.DataArray([1, 2, 3]) + result = da1 == da2 + assert result.values.all() + + def test_da_eq_scalar_still_works(self) -> None: + da = xr.DataArray([1, 2, 3]) + result = da == 2 + np.testing.assert_array_equal(result.values, [False, True, False]) + + def test_subset_constraint_solve_integration(self) -> None: + from linopy import available_solvers + + if not available_solvers: + pytest.skip("No solver available") + m = Model() + coords = pd.RangeIndex(5, name="i") + x = m.add_variables(lower=0, upper=100, coords=[coords], name="x") + subset_ub = xr.DataArray([10.0, 20.0], dims=["i"], coords={"i": [1, 3]}) + m.add_constraints(x <= subset_ub, name="subset_ub") + m.add_objective(x.sum(), sense="max") + m.solve(solver_name=available_solvers[0]) + sol = m.solution["x"] + assert sol.sel(i=1).item() == pytest.approx(10.0) + assert sol.sel(i=3).item() == pytest.approx(20.0) + assert sol.sel(i=0).item() == pytest.approx(100.0) + assert sol.sel(i=2).item() == pytest.approx(100.0) + assert sol.sel(i=4).item() == pytest.approx(100.0) + + def test_expression_inherited_properties(x: Variable, y: Variable) -> None: expr = 10 * x + y assert isinstance(expr.attrs, dict) @@ -1399,3 +1764,286 @@ def test_constant_only_expression_mul_linexpr_with_vars_and_const( assert not result_rev.is_constant assert (result_rev.coeffs == expected_coeffs).all() assert (result_rev.const == expected_const).all() + + +class TestJoinParameter: + @pytest.fixture + def m2(self) -> Model: + m = Model() + m.add_variables(coords=[pd.Index([0, 1, 2], name="i")], name="a") + m.add_variables(coords=[pd.Index([1, 2, 3], name="i")], name="b") + m.add_variables(coords=[pd.Index([0, 1, 2], name="i")], name="c") + return m + + @pytest.fixture + def a(self, m2: Model) -> Variable: + return m2.variables["a"] + + @pytest.fixture + def b(self, m2: Model) -> Variable: + return m2.variables["b"] + + @pytest.fixture + def c(self, m2: Model) -> Variable: + return m2.variables["c"] + + def test_add_join_none_preserves_default(self, a: Variable, b: Variable) -> None: + result_default = a.to_linexpr() + b.to_linexpr() + result_none = a.to_linexpr().add(b.to_linexpr(), join=None) + assert_linequal(result_default, result_none) + + def test_add_expr_join_inner(self, a: Variable, b: Variable) -> None: + result = a.to_linexpr().add(b.to_linexpr(), join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_add_expr_join_outer(self, a: Variable, b: Variable) -> None: + result = a.to_linexpr().add(b.to_linexpr(), join="outer") + assert list(result.data.indexes["i"]) == [0, 1, 2, 3] + + def test_add_expr_join_left(self, a: Variable, b: Variable) -> None: + result = a.to_linexpr().add(b.to_linexpr(), join="left") + assert list(result.data.indexes["i"]) == [0, 1, 2] + + def test_add_expr_join_right(self, a: Variable, b: Variable) -> None: + result = a.to_linexpr().add(b.to_linexpr(), join="right") + assert list(result.data.indexes["i"]) == [1, 2, 3] + + def test_add_constant_join_inner(self, a: Variable) -> None: + const = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().add(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_add_constant_join_outer(self, a: Variable) -> None: + const = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().add(const, join="outer") + assert list(result.data.indexes["i"]) == [0, 1, 2, 3] + + def test_add_constant_join_override(self, a: Variable, c: Variable) -> None: + expr = a.to_linexpr() + const = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [0, 1, 2]}) + result = expr.add(const, join="override") + assert list(result.data.indexes["i"]) == [0, 1, 2] + assert (result.const.values == const.values).all() + + def test_sub_expr_join_inner(self, a: Variable, b: Variable) -> None: + result = a.to_linexpr().sub(b.to_linexpr(), join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_mul_constant_join_inner(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().mul(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_mul_constant_join_outer(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().mul(const, join="outer") + assert list(result.data.indexes["i"]) == [0, 1, 2, 3] + assert result.coeffs.sel(i=0).item() == 0 + assert result.coeffs.sel(i=1).item() == 2 + assert result.coeffs.sel(i=2).item() == 3 + + def test_div_constant_join_inner(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().div(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_div_constant_join_outer(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.to_linexpr().div(const, join="outer") + assert list(result.data.indexes["i"]) == [0, 1, 2, 3] + + def test_variable_add_join(self, a: Variable, b: Variable) -> None: + result = a.add(b, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_variable_sub_join(self, a: Variable, b: Variable) -> None: + result = a.sub(b, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_variable_mul_join(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.mul(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_variable_div_join(self, a: Variable) -> None: + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = a.div(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_mul_expr_with_join_raises(self, a: Variable, b: Variable) -> None: + with pytest.raises(TypeError, match="join parameter is not supported"): + a.to_linexpr().mul(b.to_linexpr(), join="inner") + + def test_merge_join_parameter(self, a: Variable, b: Variable) -> None: + result = merge([a.to_linexpr(), b.to_linexpr()], join="inner") + assert list(result.data.indexes["i"]) == [1, 2] + + def test_same_shape_add_join_override(self, a: Variable, c: Variable) -> None: + result = a.to_linexpr().add(c.to_linexpr(), join="override") + assert list(result.data.indexes["i"]) == [0, 1, 2] + + def test_add_expr_outer_const_values(self, a: Variable, b: Variable) -> None: + expr_a = 1 * a + 5 + expr_b = 2 * b + 10 + result = expr_a.add(expr_b, join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.const.sel(i=0).item() == 5 + assert result.const.sel(i=1).item() == 15 + assert result.const.sel(i=2).item() == 15 + assert result.const.sel(i=3).item() == 10 + + def test_add_expr_inner_const_values(self, a: Variable, b: Variable) -> None: + expr_a = 1 * a + 5 + expr_b = 2 * b + 10 + result = expr_a.add(expr_b, join="inner") + assert list(result.coords["i"].values) == [1, 2] + assert result.const.sel(i=1).item() == 15 + assert result.const.sel(i=2).item() == 15 + + def test_add_constant_outer_fill_values(self, a: Variable) -> None: + expr = 1 * a + 5 + const = xr.DataArray([10, 20], dims=["i"], coords={"i": [1, 3]}) + result = expr.add(const, join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.const.sel(i=0).item() == 5 + assert result.const.sel(i=1).item() == 15 + assert result.const.sel(i=2).item() == 5 + assert result.const.sel(i=3).item() == 20 + + def test_add_constant_inner_fill_values(self, a: Variable) -> None: + expr = 1 * a + 5 + const = xr.DataArray([10, 20], dims=["i"], coords={"i": [1, 3]}) + result = expr.add(const, join="inner") + assert list(result.coords["i"].values) == [1] + assert result.const.sel(i=1).item() == 15 + + def test_add_constant_override_positional(self, a: Variable) -> None: + expr = 1 * a + 5 + other = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [5, 6, 7]}) + result = expr.add(other, join="override") + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.const.values, [15, 25, 35]) + + def test_sub_constant_override(self, a: Variable) -> None: + expr = 1 * a + 5 + other = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [5, 6, 7]}) + result = expr.sub(other, join="override") + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.const.values, [-5, -15, -25]) + + def test_sub_expr_outer_const_values(self, a: Variable, b: Variable) -> None: + expr_a = 1 * a + 5 + expr_b = 2 * b + 10 + result = expr_a.sub(expr_b, join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.const.sel(i=0).item() == 5 + assert result.const.sel(i=1).item() == -5 + assert result.const.sel(i=2).item() == -5 + assert result.const.sel(i=3).item() == -10 + + def test_mul_constant_override_positional(self, a: Variable) -> None: + expr = 1 * a + 5 + other = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [5, 6, 7]}) + result = expr.mul(other, join="override") + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.const.values, [10, 15, 20]) + np.testing.assert_array_equal(result.coeffs.squeeze().values, [2, 3, 4]) + + def test_mul_constant_outer_fill_values(self, a: Variable) -> None: + expr = 1 * a + 5 + other = xr.DataArray([2, 3], dims=["i"], coords={"i": [1, 3]}) + result = expr.mul(other, join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.const.sel(i=0).item() == 0 + assert result.const.sel(i=1).item() == 10 + assert result.const.sel(i=2).item() == 0 + assert result.const.sel(i=3).item() == 0 + assert result.coeffs.squeeze().sel(i=1).item() == 2 + assert result.coeffs.squeeze().sel(i=0).item() == 0 + + def test_div_constant_override_positional(self, a: Variable) -> None: + expr = 1 * a + 10 + other = xr.DataArray([2.0, 5.0, 10.0], dims=["i"], coords={"i": [5, 6, 7]}) + result = expr.div(other, join="override") + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.const.values, [5.0, 2.0, 1.0]) + + def test_div_constant_outer_fill_values(self, a: Variable) -> None: + expr = 1 * a + 10 + other = xr.DataArray([2.0, 5.0], dims=["i"], coords={"i": [1, 3]}) + result = expr.div(other, join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.const.sel(i=1).item() == pytest.approx(5.0) + assert result.coeffs.squeeze().sel(i=1).item() == pytest.approx(0.5) + assert result.const.sel(i=0).item() == pytest.approx(10.0) + assert result.coeffs.squeeze().sel(i=0).item() == pytest.approx(1.0) + + def test_div_expr_with_join_raises(self, a: Variable, b: Variable) -> None: + with pytest.raises(TypeError): + a.to_linexpr().div(b.to_linexpr(), join="outer") + + def test_variable_add_outer_values(self, a: Variable, b: Variable) -> None: + result = a.add(b, join="outer") + assert isinstance(result, LinearExpression) + assert set(result.coords["i"].values) == {0, 1, 2, 3} + assert result.nterm == 2 + + def test_variable_mul_override(self, a: Variable) -> None: + other = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [5, 6, 7]}) + result = a.mul(other, join="override") + assert isinstance(result, LinearExpression) + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.coeffs.squeeze().values, [2, 3, 4]) + + def test_variable_div_override(self, a: Variable) -> None: + other = xr.DataArray([2.0, 5.0, 10.0], dims=["i"], coords={"i": [5, 6, 7]}) + result = a.div(other, join="override") + assert isinstance(result, LinearExpression) + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_almost_equal( + result.coeffs.squeeze().values, [0.5, 0.2, 0.1] + ) + + def test_merge_outer_join(self, a: Variable, b: Variable) -> None: + result = merge([a.to_linexpr(), b.to_linexpr()], join="outer") + assert set(result.coords["i"].values) == {0, 1, 2, 3} + + def test_add_same_coords_all_joins(self, a: Variable, c: Variable) -> None: + expr_a = 1 * a + 5 + const = xr.DataArray([1, 2, 3], dims=["i"], coords={"i": [0, 1, 2]}) + for join in ["override", "outer", "inner"]: + result = expr_a.add(const, join=join) + assert list(result.coords["i"].values) == [0, 1, 2] + np.testing.assert_array_equal(result.const.values, [6, 7, 8]) + + def test_add_scalar_with_explicit_join(self, a: Variable) -> None: + expr = 1 * a + 5 + result = expr.add(10, join="override") + np.testing.assert_array_equal(result.const.values, [15, 15, 15]) + assert list(result.coords["i"].values) == [0, 1, 2] + + def test_quadratic_add_constant_join_inner(self, a: Variable, b: Variable) -> None: + quad = a.to_linexpr() * b.to_linexpr() + const = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [1, 2, 3]}) + result = quad.add(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2, 3] + + def test_quadratic_add_expr_join_inner(self, a: Variable) -> None: + quad = a.to_linexpr() * a.to_linexpr() + const = xr.DataArray([10, 20], dims=["i"], coords={"i": [0, 1]}) + result = quad.add(const, join="inner") + assert list(result.data.indexes["i"]) == [0, 1] + + def test_quadratic_mul_constant_join_inner(self, a: Variable, b: Variable) -> None: + quad = a.to_linexpr() * b.to_linexpr() + const = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [1, 2, 3]}) + result = quad.mul(const, join="inner") + assert list(result.data.indexes["i"]) == [1, 2, 3] + + def test_merge_join_left(self, a: Variable, b: Variable) -> None: + result = merge([a.to_linexpr(), b.to_linexpr()], join="left") + assert list(result.data.indexes["i"]) == [0, 1, 2] + + def test_merge_join_right(self, a: Variable, b: Variable) -> None: + result = merge([a.to_linexpr(), b.to_linexpr()], join="right") + assert list(result.data.indexes["i"]) == [1, 2, 3]