diff --git a/doc/arithmetic-convention.nblink b/doc/arithmetic-convention.nblink new file mode 100644 index 00000000..39928e92 --- /dev/null +++ b/doc/arithmetic-convention.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/arithmetic-convention.ipynb" +} diff --git a/doc/index.rst b/doc/index.rst index fd7f9ed8..70b8b439 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -108,10 +108,10 @@ This package is published under MIT license. :caption: User Guide user-guide + arithmetic-convention creating-variables creating-expressions creating-constraints - coordinate-alignment sos-constraints piecewise-linear-constraints piecewise-linear-constraints-tutorial diff --git a/examples/arithmetic-convention.ipynb b/examples/arithmetic-convention.ipynb new file mode 100644 index 00000000..97e93167 --- /dev/null +++ b/examples/arithmetic-convention.ipynb @@ -0,0 +1,460 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c68183ce878b22db", + "metadata": {}, + "source": "# Arithmetic Convention\n\nlinopy enforces strict defaults for coordinate alignment so that mismatches never silently produce wrong results.\n\n**Rule \u2014 Exact label matching on shared dimensions**\n\nWhen two operands share a dimension, their coordinate labels on that dimension must match exactly (`join=\"exact\"`). A `ValueError` is raised on mismatch.\n\n**Broadcasting** \u2014 When dimensions are *not* shared, operands broadcast freely over the missing dimensions \u2014 for both expressions and constants. This preserves all standard algebraic laws (commutativity, associativity, distributivity).\n\nInspired by [pyoframe](https://github.com/Bravos-Power/pyoframe)." + }, + { + "cell_type": "code", + "id": "4251ba8271bff255", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.502763Z", + "start_time": "2026-03-09T19:45:36.697630Z" + } + }, + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import xarray as xr\n", + "\n", + "import linopy" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "c9d84bb1c59f2690", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "id": "57506c7b4bf9f4bf", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.555658Z", + "start_time": "2026-03-09T19:45:37.505351Z" + } + }, + "source": [ + "m = linopy.Model()\n", + "\n", + "time = pd.RangeIndex(5, name=\"time\")\n", + "techs = pd.Index([\"solar\", \"wind\", \"gas\"], name=\"tech\")\n", + "scenarios = pd.Index([\"low\", \"high\"], name=\"scenario\")\n", + "\n", + "x = m.add_variables(lower=0, coords=[time], name=\"x\")\n", + "y = m.add_variables(lower=0, coords=[time], name=\"y\")\n", + "gen = m.add_variables(lower=0, coords=[time, techs], name=\"gen\")\n", + "risk = m.add_variables(lower=0, coords=[techs, scenarios], name=\"risk\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "61636799d26f4d99", + "metadata": {}, + "source": [ + "## What works by default" + ] + }, + { + "cell_type": "code", + "id": "1f7af87e662800c", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.567659Z", + "start_time": "2026-03-09T19:45:37.558480Z" + } + }, + "source": [ + "# Same coords \u2014 just works\n", + "x + y" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "985ade4e21e26271", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.586804Z", + "start_time": "2026-03-09T19:45:37.581356Z" + } + }, + "source": [ + "# Constant with matching coords\n", + "factor = xr.DataArray([2, 3, 4, 5, 6], dims=[\"time\"], coords={\"time\": time})\n", + "x * factor" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "8f6a99d864238dbb", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.600600Z", + "start_time": "2026-03-09T19:45:37.592617Z" + } + }, + "source": [ + "# Constant with fewer dims \u2014 broadcasts freely\n", + "cost = xr.DataArray([1.0, 0.5, 3.0], dims=[\"tech\"], coords={\"tech\": techs})\n", + "gen * cost # cost broadcasts over time" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "d417bfa628cb280a", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.618719Z", + "start_time": "2026-03-09T19:45:37.608423Z" + } + }, + "source": [ + "# Expression + Expression with non-shared dims \u2014 broadcasts freely\n", + "gen + risk # (time, tech) + (tech, scenario) \u2192 (time, tech, scenario)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "400b4084ef94eb35", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.626365Z", + "start_time": "2026-03-09T19:45:37.621393Z" + } + }, + "source": [ + "# Scalar \u2014 always fine\n", + "x + 5" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "2e4640266401ba61", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.642965Z", + "start_time": "2026-03-09T19:45:37.630809Z" + } + }, + "source": [ + "# Constraints \u2014 RHS with fewer dims broadcasts naturally\n", + "capacity = xr.DataArray([100, 80, 50], dims=[\"tech\"], coords={\"tech\": techs})\n", + "m.add_constraints(gen <= capacity, name=\"cap\") # capacity broadcasts over time" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "c4e9c6dbcec7c0d9", + "metadata": {}, + "source": "## What raises an error" + }, + { + "cell_type": "code", + "id": "fe1b95f337be4e9f", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.653963Z", + "start_time": "2026-03-09T19:45:37.648263Z" + } + }, + "source": [ + "# Mismatched coordinates on shared dimension\n", + "y_short = m.add_variables(\n", + " lower=0, coords=[pd.RangeIndex(3, name=\"time\")], name=\"y_short\"\n", + ")\n", + "\n", + "try:\n", + " x + y_short # time coords don't match\n", + "except ValueError as e:\n", + " print(\"ValueError:\", e)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "5a0bb6e7d4b175c5", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.662586Z", + "start_time": "2026-03-09T19:45:37.658665Z" + } + }, + "source": "# Constant introduces new dimensions \u2014 broadcasts in arithmetic\nprofile = xr.DataArray(\n np.ones((3, 5)), dims=[\"tech\", \"time\"], coords={\"tech\": techs, \"time\": time}\n)\nx + profile # x[time] broadcasts over tech", + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "e0f899f096773d96", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.681087Z", + "start_time": "2026-03-09T19:45:37.677125Z" + } + }, + "source": "# Multiplication with mismatched coordinates\npartial = xr.DataArray([10, 20, 30], dims=[\"time\"], coords={\"time\": [0, 1, 2]})\ntry:\n x * partial # time coords [0..4] vs [0,1,2]\nexcept ValueError as e:\n print(\"ValueError:\", e)", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ValueError: cannot align objects with join='exact' where index/labels/sizes are not equal along these coordinates (dimensions): 'time' ('time',)\n", + "Use .add()/.sub()/.mul()/.div() with an explicit join= parameter:\n", + " .add(other, join=\"inner\") # intersection of coordinates\n", + " .add(other, join=\"outer\") # union of coordinates (with fill)\n", + " .add(other, join=\"left\") # keep left operand's coordinates\n", + " .add(other, join=\"override\") # positional alignment\n" + ] + } + ], + "execution_count": null + }, + { + "cell_type": "code", + "id": "aa03d3184a0e8b65", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.697975Z", + "start_time": "2026-03-09T19:45:37.694178Z" + } + }, + "source": [ + "# Constraint RHS with mismatched coordinates\n", + "partial_rhs = xr.DataArray([10, 20, 30], dims=[\"time\"], coords={\"time\": [0, 1, 2]})\n", + "\n", + "try:\n", + " x <= partial_rhs\n", + "except ValueError as e:\n", + " print(\"ValueError:\", e)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "64a6f983ce55547e", + "metadata": {}, + "source": [ + "## Escape hatches\n", + "\n", + "When coordinates don't match, linopy provides several ways to state your intent explicitly." + ] + }, + { + "cell_type": "markdown", + "id": "709150bc01fc8c3", + "metadata": {}, + "source": [ + "### 1. `.sel()` \u2014 Subset before operating\n", + "\n", + "The cleanest way to restrict to matching coordinates. No need for an inner join \u2014 explicitly select what you want." + ] + }, + { + "cell_type": "code", + "id": "b4f5bf23a8ee17d5", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.712616Z", + "start_time": "2026-03-09T19:45:37.704269Z" + } + }, + "source": [ + "x.sel(time=[0, 1, 2]) + y_short # select matching coords first" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "f12b0cb6d0e31651", + "metadata": {}, + "source": "### 2. Named methods with `join=`\n\nAll arithmetic operations have named-method equivalents that accept a `join` parameter:\n\n| `join` | Coordinates kept | Fill |\n|--------|-----------------|------|\n| `\"exact\"` | Must match | `ValueError` if different |\n| `\"inner\"` | Intersection | \u2014 |\n| `\"outer\"` | Union | Zero (arithmetic) / NaN (constraints) |\n| `\"left\"` | Left operand's | Zero / NaN for missing right |\n| `\"right\"` | Right operand's | Zero for missing left |\n| `\"override\"` | Left operand's | Positional alignment |" + }, + { + "cell_type": "code", + "id": "78c967671819ef0c", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.746402Z", + "start_time": "2026-03-09T19:45:37.720673Z" + } + }, + "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\")\n", + "\n", + "print(\"inner:\", list(a.add(b, join=\"inner\").coords[\"i\"].values)) # [1, 2]\n", + "print(\"outer:\", list(a.add(b, join=\"outer\").coords[\"i\"].values)) # [0, 1, 2, 3]\n", + "print(\"left: \", list(a.add(b, join=\"left\").coords[\"i\"].values)) # [0, 1, 2]\n", + "print(\"right:\", list(a.add(b, join=\"right\").coords[\"i\"].values)) # [1, 2, 3]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "424610ceccde798a", + "metadata": {}, + "source": "### 3. `linopy.align()` \u2014 Explicit pre-alignment\n\nFor complex multi-operand alignment. Linopy types automatically use correct sentinel fill values (labels/vars=-1, coeffs=NaN) while `fill_value` applies to `const`:" + }, + { + "cell_type": "code", + "id": "23f414e973e33c34", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.765307Z", + "start_time": "2026-03-09T19:45:37.756074Z" + } + }, + "source": "a_aligned, b_aligned = linopy.align(a, b, join=\"outer\")\na_aligned + b_aligned", + "outputs": [ + { + "data": { + "text/plain": [ + "LinearExpression [i: 4]:\n", + "------------------------\n", + "[0]: +1 a[0]\n", + "[1]: +1 a[1] + 1 b[1]\n", + "[2]: +1 a[2] + 1 b[2]\n", + "[3]: +1 b[3]" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "e64caf260c82ea6d", + "metadata": {}, + "source": "## Positional alignment\n\nSometimes two operands have the same shape but different coordinate labels \u2014 e.g., data from different sources, or time series with different start dates. The exact join will raise. There are several ways to handle this:\n\n### Option 1: `.assign_coords()` (recommended)\n\nExplicitly relabel one operand to match the other. This is the clearest \u2014 the reader sees exactly which mapping is intended." + }, + { + "cell_type": "code", + "id": "9a513a6be9e5925e", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.787197Z", + "start_time": "2026-03-09T19:45:37.776535Z" + } + }, + "source": "c = m2.add_variables(coords=[[\"x\", \"y\", \"z\"]], name=\"c\")\nd = m2.add_variables(coords=[[\"p\", \"q\", \"r\"]], name=\"d\")\n\n# Relabel d's coordinates to match c\nc + d.assign_coords(dim_0=c.coords[\"dim_0\"])", + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "262eaf85fa44e152", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-09T19:45:37.803806Z", + "start_time": "2026-03-09T19:45:37.795935Z" + } + }, + "source": "### Option 2: `join=\"override\"`\n\nUses the left operand's coordinates positionally. Shorter, but less explicit about the mapping. Requires same size on the shared dimension." + }, + { + "cell_type": "code", + "id": "8lk83w4yydw", + "source": "c.add(d, join=\"override\")", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "ysdlzpp192", + "source": "**Prefer `.assign_coords()`** \u2014 it makes the intent explicit and keeps coordinate metadata intact. Use `join=\"override\"` as a shorthand when the positional mapping is obvious.", + "metadata": {} + }, + { + "cell_type": "markdown", + "id": "cd0ef5ca04e57be", + "metadata": {}, + "source": [ + "## Working with pandas\n", + "\n", + "Under the strict convention, pandas objects must have **named indices** to avoid dimension name mismatches. A `pd.Series` without a named index becomes `dim_0` and will fail the exact join against a named variable dimension.\n", + "\n", + "```python\n", + "# Bad \u2014 index name is None, becomes \"dim_0\"\n", + "cost = pd.Series([10, 20], index=[\"wind\", \"solar\"])\n", + "\n", + "# Good \u2014 explicit dimension name\n", + "cost = pd.Series([10, 20], index=pd.Index([\"wind\", \"solar\"], name=\"tech\"))\n", + "```\n", + "\n", + "Consider using `force_dim_names=True` on the model to catch unnamed dimension issues at variable creation time." + ] + }, + { + "cell_type": "markdown", + "id": "f0c3e862b0430c11", + "metadata": {}, + "source": "## Summary\n\n| Situation | Behavior | How to handle |\n|---|---|---|\n| Shared dims, matching coords | \u2713 Proceeds | `x + y` |\n| Non-shared dims, expr + expr | \u2713 Broadcasts | `gen[time,tech] + risk[tech,scenario]` |\n| Constant with subset dims | \u2713 Broadcasts | `cost[tech] * gen[time,tech]` |\n| Constant introduces new dims | \u2713 Broadcasts | `x[time] + profile[time,tech]` |\n| Shared dims, mismatching coords | \u2717 Raises | `.sel()` or `.add(y, join=\"outer\")` |\n| Pandas without named index | \u2717 Raises on dim mismatch | Name the index |" + }, + { + "cell_type": "markdown", + "id": "d56kb3o89nb", + "source": "## Algebraic Properties\n\nAll standard algebraic laws hold for linopy arithmetic. This means you can freely refactor expressions without worrying about dimension ordering.\n\nLet `x[A]`, `y[A]`, `z[A]` be linopy variables with matching dims, `g[A,B]` a variable with extra dims, `c[B]` a constant (DataArray), and `s` a scalar.\n\n| Property | Example |\n|---|---|\n| **Commutativity of +** | `x + y == y + x` |\n| **Commutativity of \u00d7** | `x * c == c * x` |\n| **Associativity of +** | `(x + y) + z == x + (y + z)` |\n| **Associativity with constant** | `(x[A] + c[B]) + g[A,B] == x[A] + (c[B] + g[A,B])` |\n| **Scalar distributivity** | `s * (x + y) == s*x + s*y` |\n| **Constant distributivity** | `c[B] * (x[A] + g[A,B]) == c[B]*x[A] + c[B]*g[A,B]` |\n| **Additive identity** | `x + 0 == x` |\n| **Multiplicative identity** | `x * 1 == x` |\n| **Negation** | `x - y == x + (-y)` |\n| **Double negation** | `-(-x) == x` |\n| **Zero** | `x * 0 == 0` |\n\n### Limitation: constant preparation\n\nThese guarantees only hold for operations that involve at least one linopy object. Operations between plain constants (`DataArray + DataArray`, `Series + Series`) happen **outside** linopy and use their library's own alignment rules \u2014 see the \"Preparing constants\" section above. To maintain algebraic consistency end-to-end, convert constants to `xr.DataArray` with explicit coordinates early and consider setting `xr.set_options(arithmetic_join=\"exact\")`.", + "metadata": {} + }, + { + "cell_type": "markdown", + "id": "e7u7uhbm1dl", + "source": "## Broadcasting in constraints\n\nBroadcasting is allowed everywhere, including constraints. This can lead to two situations worth being aware of:\n\n| Constraint type | Example | What happens | Feedback |\n|---|---|---|---|\n| `<=` / `>=` | `x[time] <= rhs[time, scenario]` | Creates one constraint per (time, scenario). Only the tightest bound is active \u2014 the rest are redundant. | No issue \u2014 solver ignores slack constraints. |\n| `==` | `x[time] == rhs[time, scenario]` | Creates one equality per (time, scenario). If `rhs` differs across `scenario`, the variable must simultaneously equal multiple values. | Solver reports **infeasible** \u2014 clear feedback. |\n\nlinopy does **not** raise an error in these cases because:\n- Redundant inequality constraints are harmless (just slightly wasteful).\n- Infeasible equality constraints are caught by the solver with a clear diagnostic.\n- Blocking these would break algebraic equivalences \u2014 e.g., `x <= rhs` must behave the same as `x - rhs <= 0`, which involves arithmetic broadcasting.", + "metadata": {} + }, + { + "cell_type": "markdown", + "id": "bpoepi5bcn8", + "source": "## Preparing constants\n\nlinopy enforces exact matching for all operations involving linopy objects. However, operations between plain constants (DataArrays, pandas, numpy) **before** they enter linopy use their own alignment rules, which can silently produce wrong results:\n\n| Library | Default alignment | Risk |\n|---|---|---|\n| **xarray** | Inner join \u2014 drops mismatched coords | Silent data loss |\n| **pandas** | Outer join \u2014 fills with NaN | Silent NaN propagation |\n| **numpy** | Positional \u2014 no coord checks | Wrong results if shapes match by accident |\n\nTo protect xarray operations, set the global arithmetic join to `\"exact\"`:\n\n```python\nxr.set_options(arithmetic_join=\"exact\")\n```\n\nFor pandas and numpy, there is no equivalent setting \u2014 prepare constants carefully and convert to `xr.DataArray` with explicit coords early.", + "metadata": {} + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/coordinate-alignment.ipynb b/examples/coordinate-alignment.ipynb deleted file mode 100644 index 1547bd9d..00000000 --- a/examples/coordinate-alignment.ipynb +++ /dev/null @@ -1,488 +0,0 @@ -{ - "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/linopy/common.py b/linopy/common.py index 746459b4..21f851df 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -10,7 +10,7 @@ import operator import os from collections.abc import Callable, Generator, Hashable, Iterable, Sequence -from functools import partial, reduce, wraps +from functools import reduce, wraps from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload from warnings import warn @@ -1205,27 +1205,9 @@ def deco(cls: Any) -> Any: return deco -def check_common_keys_values(list_of_dicts: list[dict[str, Any]]) -> bool: - """ - Check if all common keys among a list of dictionaries have the same value. - - Parameters - ---------- - list_of_dicts : list of dict - A list of dictionaries. - - Returns - ------- - bool - True if all common keys have the same value across all dictionaries, False otherwise. - """ - common_keys = set.intersection(*(set(d.keys()) for d in list_of_dicts)) - return all(len({d[k] for d in list_of_dicts if k in d}) == 1 for k in common_keys) - - def align( *objects: LinearExpression | QuadraticExpression | Variable | T_Alignable, - join: JoinOptions = "inner", + join: JoinOptions = "exact", copy: bool = True, indexes: Any = None, exclude: str | Iterable[Hashable] = frozenset(), @@ -1288,36 +1270,38 @@ def align( from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable - finisher: list[partial[Any] | Callable[[Any], Any]] = [] + # Extract underlying Datasets for index computation. das: list[Any] = [] for obj in objects: - if isinstance(obj, LinearExpression | QuadraticExpression): - finisher.append(partial(obj.__class__, model=obj.model)) - das.append(obj.data) - elif isinstance(obj, Variable): - finisher.append( - partial( - obj.__class__, - model=obj.model, - name=obj.data.attrs["name"], - skip_broadcast=True, - ) - ) + if isinstance(obj, LinearExpression | QuadraticExpression | Variable): das.append(obj.data) else: - finisher.append(lambda x: x) das.append(obj) exclude = frozenset(exclude).union(HELPER_DIMS) - aligned = xr_align( - *das, - join=join, - copy=copy, - indexes=indexes, - exclude=exclude, - fill_value=fill_value, + + # Compute target indexes. + target_aligned = xr_align( + *das, join=join, copy=False, indexes=indexes, exclude=exclude ) - return tuple([f(da) for f, da in zip(finisher, aligned)]) + + # Reindex each object to target indexes. + reindex_kwargs: dict[str, Any] = {} + if fill_value is not dtypes.NA: + reindex_kwargs["fill_value"] = fill_value + results: list[Any] = [] + for obj, target in zip(objects, target_aligned): + indexers = { + dim: target.indexes[dim] + for dim in target.dims + if dim not in exclude and dim in target.indexes + } + # Variable.reindex has no fill_value — it always uses sentinels + if isinstance(obj, Variable): + results.append(obj.reindex(indexers)) + else: + results.append(obj.reindex(indexers, **reindex_kwargs)) + return tuple(results) LocT = TypeVar( diff --git a/linopy/expressions.py b/linopy/expressions.py index e1fbe1a9..32bf781b 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from itertools import product, zip_longest -from typing import TYPE_CHECKING, Any, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, Self, TypeVar, cast, overload from warnings import warn import numpy as np @@ -48,7 +48,6 @@ LocIndexer, as_dataarray, assign_multiindex_safe, - check_common_keys_values, check_has_nulls, check_has_nulls_polars, fill_missing_coords, @@ -528,6 +527,7 @@ def _align_constant( other: DataArray, fill_value: float = 0, join: str | None = None, + default_join: str = "exact", ) -> tuple[DataArray, DataArray, bool]: """ Align a constant DataArray with self.const. @@ -539,7 +539,10 @@ def _align_constant( fill_value : float, default: 0 Fill value for missing coordinates. join : str, optional - Alignment method. If None, uses size-aware default behavior. + Alignment method. If None, uses default_join. + default_join : str, default: "exact" + Default join mode when join is None. Use "exact" for add/sub, + "inner" for mul/div. Returns ------- @@ -551,19 +554,32 @@ def _align_constant( 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 + join = default_join + + if join == "override": + return self.const, other.assign_coords(coords=self.coords), False + elif join == "left": 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 - ) + try: + self_const, aligned = xr.align( + self.const, other, join=join, fill_value=fill_value + ) + except ValueError as e: + if "exact" in str(e): + raise ValueError( + f"{e}\n" + "Use .add()/.sub()/.mul()/.div() with an explicit join= parameter:\n" + ' .add(other, join="inner") # intersection of coordinates\n' + ' .add(other, join="outer") # union of coordinates (with fill)\n' + ' .add(other, join="left") # keep left operand\'s coordinates\n' + ' .add(other, join="override") # positional alignment' + ) from None + raise return self_const, aligned, True def _add_constant( @@ -573,11 +589,12 @@ def _add_constant( 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 + da, fill_value=0, join=join, default_join="exact" ) if needs_data_reindex: + fv = {**self._fill_value, "const": 0} return self.__class__( - self.data.reindex_like(self_const, fill_value=self._fill_value).assign( + self.data.reindex_like(self_const, fill_value=fv).assign( const=self_const + da ), self.model, @@ -593,10 +610,11 @@ def _apply_constant_op( ) -> 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 + factor, fill_value=fill_value, join=join, default_join="exact" ) if needs_data_reindex: - data = self.data.reindex_like(self_const, fill_value=self._fill_value) + fv = {**self._fill_value, "const": 0} + data = self.data.reindex_like(self_const, fill_value=fv) return self.__class__( assign_multiindex_safe( data, coeffs=op(data.coeffs, factor), const=op(self_const, factor) @@ -1076,13 +1094,39 @@ def to_constraint( ) 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." + effective_join = join if join is not None else "exact" + if effective_join == "override": + aligned_rhs = rhs.assign_coords(coords=self.const.coords) + expr_const = self.const + expr_data = self.data + elif effective_join == "left": + aligned_rhs = rhs.reindex_like(self.const, fill_value=np.nan) + expr_const = self.const + expr_data = self.data + else: + try: + expr_const_aligned, aligned_rhs = xr.align( + self.const, rhs, join=effective_join, fill_value=np.nan + ) + except ValueError as e: + if "exact" in str(e): + raise ValueError( + f"{e}\n" + "Use .le()/.ge()/.eq() with an explicit join= parameter:\n" + ' .le(rhs, join="inner") # intersection of coordinates\n' + ' .le(rhs, join="left") # keep expression coordinates (NaN fill)\n' + ' .le(rhs, join="override") # positional alignment' + ) from None + raise + expr_const = expr_const_aligned.fillna(0) + expr_data = self.data.reindex_like( + expr_const_aligned, fill_value=self._fill_value ) - rhs = rhs.reindex_like(self.const, fill_value=np.nan) + constraint_rhs = aligned_rhs - expr_const + data = assign_multiindex_safe( + expr_data[["coeffs", "vars"]], sign=sign, rhs=constraint_rhs + ) + return constraints.Constraint(data, model=self.model) elif isinstance(rhs, np.ndarray | pd.Series | pd.DataFrame) and rhs.ndim > len( self.coord_dims ): @@ -1440,9 +1484,45 @@ def _sum( set_index = exprwrap(Dataset.set_index) - reindex = exprwrap(Dataset.reindex, fill_value=_fill_value) + def reindex( + self, + indexers: Mapping[Any, Any] | None = None, + fill_value: float = np.nan, + **indexers_kwargs: Any, + ) -> Self: + """ + Reindex the expression. - reindex_like = exprwrap(Dataset.reindex_like, fill_value=_fill_value) + ``fill_value`` sets the constant for missing coordinates (default NaN). + Variable labels and coefficients always use sentinel values + (vars=-1, coeffs=NaN). + """ + fv = {**self._fill_value, "const": fill_value} + return self.__class__( + self.data.reindex(indexers, fill_value=fv, **indexers_kwargs), self.model + ) + + def reindex_like( + self, + other: Any, + fill_value: float = np.nan, + **kwargs: Any, + ) -> Self: + """ + Reindex like another object. + + ``fill_value`` sets the constant for missing coordinates (default NaN). + Variable labels and coefficients always use sentinel values. + """ + fv = {**self._fill_value, "const": fill_value} + return self.__class__( + self.data.reindex_like( + other if isinstance(other, Dataset) else other.data, + fill_value=fv, + **kwargs, + ), + self.model, + ) rename = exprwrap(Dataset.rename) @@ -2320,16 +2400,6 @@ def merge( model = exprs[0].model - 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 - ] - override = check_common_keys_values(coord_dims) # type: ignore - else: - override = False - data = [e.data if isinstance(e, linopy_types) else e for e in exprs] data = [fill_missing_coords(ds, fill_helper_dims=True) for ds in data] @@ -2345,23 +2415,51 @@ def merge( if join is not None: kwargs["join"] = join - elif override: - kwargs["join"] = "override" else: - kwargs.setdefault("join", "outer") - - if dim == TERM_DIM: - ds = xr.concat([d[["coeffs", "vars"]] for d in data], dim, **kwargs) - subkwargs = {**kwargs, "fill_value": 0} - const = xr.concat([d["const"] for d in data], dim, **subkwargs).sum(TERM_DIM) - ds = assign_multiindex_safe(ds, const=const) - elif dim == FACTOR_DIM: - ds = xr.concat([d[["vars"]] for d in data], dim, **kwargs) - coeffs = xr.concat([d["coeffs"] for d in data], dim, **kwargs).prod(FACTOR_DIM) - const = xr.concat([d["const"] for d in data], dim, **kwargs).prod(FACTOR_DIM) - ds = assign_multiindex_safe(ds, coeffs=coeffs, const=const) - else: - ds = xr.concat(data, dim, **kwargs) + kwargs["join"] = "exact" + + try: + if dim == TERM_DIM: + ds = xr.concat([d[["coeffs", "vars"]] for d in data], dim, **kwargs) + subkwargs = {**kwargs, "fill_value": 0} + const = xr.concat([d["const"] for d in data], dim, **subkwargs).sum( + TERM_DIM + ) + ds = assign_multiindex_safe(ds, const=const) + elif dim == FACTOR_DIM: + ds = xr.concat([d[["vars"]] for d in data], dim, **kwargs) + coeffs = xr.concat([d["coeffs"] for d in data], dim, **kwargs).prod( + FACTOR_DIM + ) + const = xr.concat([d["const"] for d in data], dim, **kwargs).prod( + FACTOR_DIM + ) + ds = assign_multiindex_safe(ds, coeffs=coeffs, const=const) + else: + # Pre-pad helper dims to same size before concat + fill = kwargs.get("fill_value", FILL_VALUE) + for helper_dim in HELPER_DIMS: + sizes = [d.sizes.get(helper_dim, 0) for d in data] + max_size = max(sizes) if sizes else 0 + if max_size > 0 and min(sizes) < max_size: + data = [ + d.reindex({helper_dim: range(max_size)}, fill_value=fill) + if d.sizes.get(helper_dim, 0) < max_size + else d + for d in data + ] + ds = xr.concat(data, dim, **kwargs) + except ValueError as e: + if "exact" in str(e): + raise ValueError( + f"{e}\n" + "Use .add()/.sub()/.mul()/.div() with an explicit join= parameter:\n" + ' .add(other, join="inner") # intersection of coordinates\n' + ' .add(other, join="outer") # union of coordinates (with fill)\n' + ' .add(other, join="left") # keep left operand\'s coordinates\n' + ' .add(other, join="override") # positional alignment' + ) from None + raise for d in set(HELPER_DIMS) & set(ds.coords): ds = ds.reset_index(d, drop=True) diff --git a/linopy/piecewise.py b/linopy/piecewise.py index 5128d1e5..c31204f6 100644 --- a/linopy/piecewise.py +++ b/linopy/piecewise.py @@ -560,6 +560,8 @@ def _add_pwl_incremental( if n_segments >= 2: delta_lo = delta_var.isel({seg_dim: slice(None, -1)}, drop=True) delta_hi = delta_var.isel({seg_dim: slice(1, None)}, drop=True) + # Align coords for positional comparison (lo=[0..n-2], hi=[1..n-1]) + delta_hi = delta_hi.assign_coords({seg_dim: delta_lo.coords[seg_dim].values}) fill_con = model.add_constraints(delta_hi <= delta_lo, name=fill_name) bp0 = breakpoints.isel({dim: 0}) diff --git a/linopy/variables.py b/linopy/variables.py index 823c9a59..c11f52cf 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -398,10 +398,7 @@ def __mul__(self, other: SideLike) -> ExpressionLike: Multiply variables with a coefficient, variable, or expression. """ try: - if isinstance(other, Variable | ScalarVariable): - return self.to_linexpr() * other - - return self.to_linexpr(other) + return self.to_linexpr() * other except TypeError: return NotImplemented @@ -449,13 +446,7 @@ def __div__( """ Divide variables with a coefficient. """ - if isinstance(other, expressions.LinearExpression | Variable): - raise TypeError( - "unsupported operand type(s) for /: " - f"{type(self)} and {type(other)}. " - "Non-linear expressions are not yet supported." - ) - return self.to_linexpr()._divide_by_constant(other) + return self.to_linexpr() / other def __truediv__( self, coefficient: ConstantLike | LinearExpression | Variable @@ -1230,6 +1221,39 @@ def equals(self, other: Variable) -> bool: shift = varwrap(Dataset.shift, fill_value=_fill_value) + def reindex( + self, + indexers: Mapping[Any, Any] | None = None, + **indexers_kwargs: Any, + ) -> Variable: + """ + Reindex the variable, filling with sentinel values. + + Always fills with labels=-1, lower=NaN, upper=NaN to preserve + valid label references. + """ + return self.__class__( + self.data.reindex(indexers, fill_value=self._fill_value, **indexers_kwargs), + self.model, + self.name, + ) + + def reindex_like( + self, + other: Any, + **kwargs: Any, + ) -> Variable: + """Reindex like another object, filling with sentinel values.""" + return self.__class__( + self.data.reindex_like( + other if isinstance(other, Dataset) else other.data, + fill_value=self._fill_value, + **kwargs, + ), + self.model, + self.name, + ) + swap_dims = varwrap(Dataset.swap_dims) set_index = varwrap(Dataset.set_index) diff --git a/test/test_algebraic_properties.py b/test/test_algebraic_properties.py new file mode 100644 index 00000000..09548bf3 --- /dev/null +++ b/test/test_algebraic_properties.py @@ -0,0 +1,211 @@ +""" +Algebraic properties of linopy arithmetic. + +All standard algebraic laws should hold for linopy expressions. +This file serves as both specification and test suite. + +Notation: + x[A], y[A], z[A] — linopy variables with dimension A + g[A,B] — linopy variable with dimensions A and B + c[B] — constant (DataArray) with dimension B + s — scalar (int/float) + +SPECIFICATION +============= + +1. Commutativity + a + b == b + a for any linopy operands a, b + a * c == c * a for variable/expression a, constant c + +2. Associativity + (a + b) + c == a + (b + c) for any linopy operands a, b, c + Including mixed: (x[A] + c[B]) + g[A,B] == x[A] + (c[B] + g[A,B]) + +3. Distributivity + c * (a + b) == c*a + c*b for constant c, linopy operands a, b + s * (a + b) == s*a + s*b for scalar s + +4. Identity + a + 0 == a additive identity + a * 1 == a multiplicative identity + +5. Negation + a - b == a + (-b) subtraction is addition of negation + -(-a) == a double negation + +6. Zero + a * 0 == 0 multiplication by zero +""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model +from linopy.expressions import LinearExpression + + +@pytest.fixture +def m(): + return Model() + + +@pytest.fixture +def time(): + return pd.RangeIndex(3, name="time") + + +@pytest.fixture +def tech(): + return pd.Index(["solar", "wind"], name="tech") + + +@pytest.fixture +def x(m, time): + """Variable with dims [time].""" + return m.add_variables(lower=0, coords=[time], name="x") + + +@pytest.fixture +def y(m, time): + """Variable with dims [time].""" + return m.add_variables(lower=0, coords=[time], name="y") + + +@pytest.fixture +def z(m, time): + """Variable with dims [time].""" + return m.add_variables(lower=0, coords=[time], name="z") + + +@pytest.fixture +def g(m, time, tech): + """Variable with dims [time, tech].""" + return m.add_variables(lower=0, coords=[time, tech], name="g") + + +@pytest.fixture +def c(tech): + """Constant (DataArray) with dims [tech].""" + return xr.DataArray([2.0, 3.0], dims=["tech"], coords={"tech": tech}) + + +def assert_linequal(a: LinearExpression, b: LinearExpression) -> None: + """Assert two linear expressions are algebraically equivalent.""" + assert set(a.dims) == set(b.dims), f"dims differ: {a.dims} vs {b.dims}" + for dim in a.dims: + if dim.startswith("_"): + continue + np.testing.assert_array_equal( + sorted(a.coords[dim].values), sorted(b.coords[dim].values) + ) + assert a.const.sum().item() == pytest.approx(b.const.sum().item()) + + +# ============================================================ +# 1. Commutativity +# ============================================================ + + +class TestCommutativity: + def test_add_expr_expr(self, x, y): + """X + y == y + x""" + assert_linequal(x + y, y + x) + + def test_mul_expr_constant(self, g, c): + """G * c == c * g""" + assert_linequal(g * c, c * g) + + def test_add_expr_constant(self, g, c): + """G + c == c + g""" + assert_linequal(g + c, c + g) + + +# ============================================================ +# 2. Associativity +# ============================================================ + + +class TestAssociativity: + def test_add_same_dims(self, x, y, z): + """(x + y) + z == x + (y + z)""" + assert_linequal((x + y) + z, x + (y + z)) + + def test_add_with_constant(self, x, g, c): + """(x[A] + c[B]) + g[A,B] == x[A] + (c[B] + g[A,B])""" + assert_linequal((x + c) + g, x + (c + g)) + + +# ============================================================ +# 3. Distributivity +# ============================================================ + + +class TestDistributivity: + def test_scalar(self, x, y): + """S * (x + y) == s*x + s*y""" + assert_linequal(3 * (x + y), 3 * x + 3 * y) + + def test_constant_subset_dims(self, g, c): + """c[B] * (g[A,B] + g[A,B]) == c*g + c*g""" + assert_linequal(c * (g + g), c * g + c * g) + + def test_constant_mixed_dims(self, x, g, c): + """c[B] * (x[A] + g[A,B]) == c*x + c*g""" + assert_linequal(c * (x + g), c * x + c * g) + + +# ============================================================ +# 4. Identity +# ============================================================ + + +class TestIdentity: + def test_additive(self, x): + """X + 0 == x""" + result = x + 0 + assert isinstance(result, LinearExpression) + assert (result.const == 0).all() + np.testing.assert_array_equal(result.coeffs.squeeze().values, [1, 1, 1]) + + def test_multiplicative(self, x): + """X * 1 == x""" + result = x * 1 + assert isinstance(result, LinearExpression) + np.testing.assert_array_equal(result.coeffs.squeeze().values, [1, 1, 1]) + + +# ============================================================ +# 5. Negation +# ============================================================ + + +class TestNegation: + def test_subtraction_is_add_negation(self, x, y): + """X - y == x + (-y)""" + assert_linequal(x - y, x + (-y)) + + def test_subtraction_definition(self, x, y): + """X - y == x + (-1) * y""" + assert_linequal(x - y, x + (-1) * y) + + def test_double_negation(self, x): + """-(-x) has same coefficients as x""" + result = -(-x) + np.testing.assert_array_equal( + result.coeffs.squeeze().values, + (1 * x).coeffs.squeeze().values, + ) + + +# ============================================================ +# 6. Zero +# ============================================================ + + +class TestZero: + def test_multiplication_by_zero(self, x): + """X * 0 has zero coefficients""" + result = x * 0 + assert (result.coeffs == 0).all() diff --git a/test/test_common.py b/test/test_common.py index 267fbf76..4b84755a 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -674,22 +674,11 @@ def test_get_dims_with_index_levels() -> None: assert get_dims_with_index_levels(ds5) == [] -def test_align(x: Variable, u: Variable) -> None: # noqa: F811 +def test_align(x: Variable) -> None: # noqa: F811 alpha = xr.DataArray([1, 2], [[1, 2]]) - beta = xr.DataArray( - [1, 2, 3], - [ - ( - "dim_3", - pd.MultiIndex.from_tuples( - [(1, "b"), (2, "b"), (1, "c")], names=["level1", "level2"] - ), - ) - ], - ) # inner join - x_obs, alpha_obs = align(x, alpha) + x_obs, alpha_obs = align(x, alpha, join="inner") assert isinstance(x_obs, Variable) assert x_obs.shape == alpha_obs.shape == (1,) assert_varequal(x_obs, x.loc[[1]]) @@ -701,16 +690,9 @@ def test_align(x: Variable, u: Variable) -> None: # noqa: F811 assert_varequal(x_obs, x) assert_equal(alpha_obs, DataArray([np.nan, 1], [[0, 1]])) - # multiindex - beta_obs, u_obs = align(beta, u) - assert u_obs.shape == beta_obs.shape == (2,) - assert isinstance(u_obs, Variable) - assert_varequal(u_obs, u.loc[[(1, "b"), (2, "b")]]) - assert_equal(beta_obs, beta.loc[[(1, "b"), (2, "b")]]) - # with linear expression expr = 20 * x - x_obs, expr_obs, alpha_obs = align(x, expr, alpha) + x_obs, expr_obs, alpha_obs = align(x, expr, alpha, join="inner") assert x_obs.shape == alpha_obs.shape == (1,) assert expr_obs.shape == (1, 1) # _term dim assert isinstance(expr_obs, LinearExpression) diff --git a/test/test_constraints.py b/test/test_constraints.py index e5da08d4..b20b18cf 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -166,23 +166,14 @@ def test_constraint_rhs_lower_dim(rhs_factory) -> None: assert c.shape == (10, 10) -@pytest.mark.parametrize( - "rhs_factory", - [ - pytest.param(lambda m: np.ones((5, 3)), id="numpy"), - pytest.param( - lambda m: xr.DataArray(np.ones((5, 3)), dims=["dim_0", "extra"]), - id="dataarray", - ), - pytest.param(lambda m: pd.DataFrame(np.ones((5, 3))), id="dataframe"), - ], -) -def test_constraint_rhs_higher_dim_constant_raises(rhs_factory) -> None: +def test_constraint_rhs_higher_dim_constant_broadcasts() -> None: m = Model() x = m.add_variables(coords=[range(5)], name="x") - with pytest.raises(ValueError, match="dimensions"): - m.add_constraints(x >= rhs_factory(m)) + # DataArray RHS with extra dims broadcasts (creates redundant constraints) + rhs = xr.DataArray(np.ones((5, 3)), dims=["dim_0", "extra"]) + c = m.add_constraints(x >= rhs, name="broadcast_con") + assert "extra" in c.dims @pytest.mark.parametrize( diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 2af1a8ea..ed808e78 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -7,8 +7,6 @@ from __future__ import annotations -from typing import Any - import numpy as np import pandas as pd import polars as pl @@ -222,6 +220,7 @@ def test_linear_expression_with_multiplication(x: Variable) -> None: expr = np.array(1) * x assert isinstance(expr, LinearExpression) + # Constants with extra dims broadcast freely expr = xr.DataArray(np.array([[1, 2], [2, 3]])) * x assert isinstance(expr, LinearExpression) @@ -317,9 +316,9 @@ def test_linear_expression_with_constant_multiplication( assert isinstance(obs, LinearExpression) assert (obs.const == 10).all() + # Constants with extra dims broadcast freely obs = expr * pd.Series([1, 2, 3], index=pd.RangeIndex(3, name="new_dim")) assert isinstance(obs, LinearExpression) - assert obs.shape == (2, 3, 1) def test_linear_expression_multi_indexed(u: Variable) -> None: @@ -443,8 +442,12 @@ def test_linear_expression_sum( assert_linequal(expr.sum(["dim_0", TERM_DIM]), expr.sum("dim_0")) - # test special case otherride coords - expr = v.loc[:9] + v.loc[10:] + # disjoint coords now raise with exact default + with pytest.raises(ValueError, match="exact"): + v.loc[:9] + v.loc[10:] + + # positional alignment via assign_coords + expr = v.loc[:9] + v.loc[10:].assign_coords(dim_2=v.loc[:9].coords["dim_2"]) assert expr.nterm == 2 assert len(expr.coords["dim_2"]) == 10 @@ -467,8 +470,12 @@ def test_linear_expression_sum_with_const( assert_linequal(expr.sum(["dim_0", TERM_DIM]), expr.sum("dim_0")) - # test special case otherride coords - expr = v.loc[:9] + v.loc[10:] + # disjoint coords now raise with exact default + with pytest.raises(ValueError, match="exact"): + v.loc[:9] + v.loc[10:] + + # positional alignment via assign_coords + expr = v.loc[:9] + v.loc[10:].assign_coords(dim_2=v.loc[:9].coords["dim_2"]) assert expr.nterm == 2 assert len(expr.coords["dim_2"]) == 10 @@ -577,7 +584,18 @@ def test_linear_expression_multiplication_invalid( expr / x -class TestSubsetCoordinateAlignment: +class TestExactAlignmentDefault: + """ + Test the alignment convention: exact for all operations (+, -, *, /). + + v has dim_2=[0..19] (20 entries). + subset has dim_2=[1, 3] (2 entries, subset of v's coords). + superset has dim_2=[0..24] (25 entries, superset of v's coords). + + Each test shows the operation, verifies the exact default (raises), + then shows the explicit join= that recovers the desired result. + """ + @pytest.fixture def subset(self) -> xr.DataArray: return xr.DataArray([10.0, 30.0], dims=["dim_2"], coords={"dim_2": [1, 3]}) @@ -588,329 +606,336 @@ def superset(self) -> xr.DataArray: np.arange(25, dtype=float), dims=["dim_2"], coords={"dim_2": range(25)} ) + @pytest.fixture + def matching(self) -> xr.DataArray: + return xr.DataArray( + np.arange(20, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(20)}, + ) + @pytest.fixture def expected_fill(self) -> np.ndarray: + """Old expected result: 20-entry array with values at positions 1,3.""" 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: Any, make_rhs: Any - ) -> None: - assert_linequal(make_lhs(v, subset), make_rhs(v, subset)) + # --- Addition / subtraction with subset constant --- 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() + # now raises + with pytest.raises(ValueError, match="exact"): + v + subset + + # explicit join="left" recovers old behavior: 20 entries, fill 0 + result = v.add(subset, join="left") + assert result.sizes["dim_2"] == 20 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) + with pytest.raises(ValueError, match="exact"): + v - subset - def test_subset_sub_var(self, v: Variable, subset: xr.DataArray) -> None: - assert_linequal(subset - v, -v + subset) + result = v.sub(subset, join="left") + assert result.sizes["dim_2"] == 20 + np.testing.assert_array_equal(result.const.values, -expected_fill) 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() + with pytest.raises(ValueError, match="exact"): + (v + 5) + subset + + result = (v + 5).add(subset, join="left") + assert result.sizes["dim_2"] == 20 np.testing.assert_array_equal(result.const.values, expected_fill + 5) - def test_expr_sub_subset( + # --- Addition with superset constant --- + + def test_var_add_superset(self, v: Variable, superset: xr.DataArray) -> None: + with pytest.raises(ValueError, match="exact"): + v + superset + + result = v.add(superset, join="left") + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.const.values).any() + + # --- Addition / multiplication with disjoint coords --- + + def test_disjoint_add(self, v: Variable) -> None: + disjoint = xr.DataArray( + [100.0, 200.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + with pytest.raises(ValueError, match="exact"): + v + disjoint + + result = v.add(disjoint, join="outer") + assert result.sizes["dim_2"] == 22 # union of [0..19] and [50, 60] + + def test_disjoint_mul(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + with pytest.raises(ValueError, match="exact"): + v * disjoint + + # explicit join="left": 20 entries, all zeros + result = v.mul(disjoint, join="left") + assert result.sizes["dim_2"] == 20 + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.zeros(20)) + + def test_disjoint_div(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + with pytest.raises(ValueError, match="exact"): + v / disjoint + + # --- Multiplication / division with subset constant --- + + def test_var_mul_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) + with pytest.raises(ValueError, match="exact"): + v * subset + + # explicit join="inner": 2 entries (intersection) + result = v.mul(subset, join="inner") + assert result.sizes["dim_2"] == 2 + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(10.0) + assert result.coeffs.squeeze().sel(dim_2=3).item() == pytest.approx(30.0) + + # explicit join="left" recovers old behavior: 20 entries, fill 0 + result = v.mul(subset, join="left") + assert result.sizes["dim_2"] == 20 + np.testing.assert_array_equal(result.coeffs.squeeze().values, 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_expr_mul_subset(self, v: Variable, subset: xr.DataArray) -> None: + with pytest.raises(ValueError, match="exact"): + (1 * v) * 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"] + result = (1 * v).mul(subset, join="inner") + assert result.sizes["dim_2"] == 2 + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(10.0) + + def test_var_mul_superset(self, v: Variable, superset: xr.DataArray) -> None: + with pytest.raises(ValueError, match="exact"): + v * superset + + result = v.mul(superset, join="inner") + assert result.sizes["dim_2"] == 20 assert not np.isnan(result.coeffs.values).any() + + def test_var_div_subset(self, v: Variable, subset: xr.DataArray) -> None: + with pytest.raises(ValueError, match="exact"): + v / subset + + # explicit join="inner": 2 entries + result = v.div(subset, join="inner") + assert result.sizes["dim_2"] == 2 + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(0.1) + assert result.coeffs.squeeze().sel(dim_2=3).item() == pytest.approx(1.0 / 30) + + # explicit join="left": 20 entries, fill 1 + result = v.div(subset, join="left") + assert result.sizes["dim_2"] == 20 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) + # --- Constraints with subset RHS --- + 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()) + with pytest.raises(ValueError, match="exact"): + v <= subset - @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"] + # explicit join="left": 20 entries, NaN where RHS missing + con = v.to_linexpr().le(subset, join="left") + assert con.sizes["dim_2"] == 20 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()) 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"] + with pytest.raises(ValueError, match="exact"): + expr <= subset + + con = expr.le(subset, join="left") + assert con.sizes["dim_2"] == 20 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 + @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) + def test_var_comparison_subset( + self, v: Variable, subset: xr.DataArray, sign: str ) -> 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) + with pytest.raises(ValueError, match="exact"): + if sign == "<=": + v <= subset + elif sign == ">=": + v >= subset + else: + v == subset + + def test_constraint_le_join_inner(self, v: Variable, subset: xr.DataArray) -> None: + con = v.to_linexpr().le(subset, join="inner") + assert con.sizes["dim_2"] == 2 + assert con.rhs.sel(dim_2=1).item() == 10.0 + assert con.rhs.sel(dim_2=3).item() == 30.0 - def test_superset_mul_var(self, v: Variable, superset: xr.DataArray) -> None: - assert_linequal(superset * v, v * superset) + # --- Matching coordinates: unchanged behavior --- - @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"] + def test_add_matching_unchanged(self, v: Variable, matching: xr.DataArray) -> None: + result = v + matching + assert result.sizes["dim_2"] == 20 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_mul_matching_unchanged(self, v: Variable, matching: xr.DataArray) -> None: + result = v * matching + assert result.sizes["dim_2"] == 20 - 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_le_matching_unchanged(self, v: Variable, matching: xr.DataArray) -> None: + con = v <= matching + assert con.sizes["dim_2"] == 20 - def test_subset_sub_var_coefficients( - self, v: Variable, subset: xr.DataArray + def test_add_commutativity_matching( + self, v: Variable, matching: xr.DataArray ) -> None: - result = subset - v - np.testing.assert_array_equal(result.coeffs.squeeze().values, -np.ones(20)) + assert_linequal(v + matching, matching + v) - @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_mul_commutativity(self, v: Variable, subset: xr.DataArray) -> None: + with pytest.raises(ValueError, match="exact"): + v * subset + with pytest.raises(ValueError, match="exact"): + subset * v - 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() + # --- Explicit join modes --- - 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)}, + def test_add_join_inner(self, v: Variable, subset: xr.DataArray) -> None: + result = v.add(subset, join="inner") + assert result.sizes["dim_2"] == 2 + assert result.const.sel(dim_2=1).item() == 10.0 + assert result.const.sel(dim_2=3).item() == 30.0 + + def test_add_join_outer(self, v: Variable, subset: xr.DataArray) -> None: + result = v.add(subset, join="outer") + assert result.sizes["dim_2"] == 20 + assert result.const.sel(dim_2=1).item() == 10.0 + assert result.const.sel(dim_2=0).item() == 0.0 + + def test_add_positional_assign_coords(self, v: Variable) -> None: + disjoint = xr.DataArray( + np.ones(20), dims=["dim_2"], coords={"dim_2": range(50, 70)} ) - result = v / superset_nonzero - assert result.sizes["dim_2"] == v.sizes["dim_2"] - assert not np.isnan(result.coeffs.values).any() + result = v + disjoint.assign_coords(dim_2=v.coords["dim_2"]) + assert result.sizes["dim_2"] == 20 + assert list(result.coords["dim_2"].values) == list(range(20)) + + # --- Quadratic expressions --- 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) + with pytest.raises(ValueError, match="exact"): + qexpr + subset - def test_quadexpr_sub_subset( - self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray - ) -> None: - qexpr = v * v - result = qexpr - subset + result = qexpr.add(subset, join="left") 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) + assert result.sizes["dim_2"] == 20 + 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 + with pytest.raises(ValueError, match="exact"): + qexpr * subset + + # explicit join="inner": 2 entries + result = qexpr.mul(subset, join="inner") 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) + assert result.sizes["dim_2"] == 2 - def test_subset_mul_quadexpr( - self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray - ) -> None: - qexpr = v * v - result = subset * qexpr + # explicit join="left": 20 entries + result = qexpr.mul(subset, join="left") assert isinstance(result, QuadraticExpression) - assert result.sizes["dim_2"] == v.sizes["dim_2"] - assert not np.isnan(result.coeffs.values).any() + assert result.sizes["dim_2"] == 20 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) + # --- Multi-dimensional --- 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 + + with pytest.raises(ValueError, match="exact"): + w * subset_2d + + # explicit join="inner": 2x2 + result = w.mul(subset_2d, join="inner") + assert result.sizes["a"] == 2 + assert result.sizes["b"] == 2 + + # explicit join="left": 4x5, zeros at non-subset positions + result = w.mul(subset_2d, join="left") 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: + with pytest.raises(ValueError, match="exact"): + w + subset_2d + + # --- Edge cases --- + + def test_constraint_rhs_mismatched_coords_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"): + # Raises because dim_2 coords [0,1] don't match v's [0..19] (exact join) + with pytest.raises(ValueError, match="exact"): v <= rhs + def test_add_constant_extra_dims_broadcasts(self, v: Variable) -> None: + # Constant with only new dims (no shared dim overlap) broadcasts freely + da = xr.DataArray([1.0, 2.0, 3.0], dims=["extra"]) + result = v + da + assert "extra" in result.dims + result = v - da + assert "extra" in result.dims + result = v * da + assert "extra" in result.dims + 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 # type: ignore[operator] - 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]) @@ -931,7 +956,8 @@ def test_subset_constraint_solve_integration(self) -> None: 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") + # exact default raises — use explicit join="left" (NaN = no constraint) + m.add_constraints(x.to_linexpr().le(subset_ub, join="left"), name="subset_ub") m.add_objective(x.sum(), sense="max") m.solve(solver_name=available_solvers[0]) sol = m.solution["x"] @@ -996,7 +1022,8 @@ def test_linear_expression_isnull(v: Variable) -> None: expr = np.arange(20) * v filter = (expr.coeffs >= 10).any(TERM_DIM) expr = expr.where(filter) - assert expr.isnull().sum() == 10 + # Entries where filter is False are null (coeffs=NaN, const=NaN) + assert expr.isnull().sum() == 10 # first 10 entries (coeff 0..9) are null def test_linear_expression_flat(v: Variable) -> None: @@ -1126,6 +1153,7 @@ def test_linear_expression_fillna(v: Variable) -> None: filled = filtered.fillna(10) assert isinstance(filled, LinearExpression) + # fillna replaces NaN const values (10 entries × 10) + kept values (10 × 10) assert filled.const.sum() == 200 assert filled.coeffs.isnull().sum() == 10 @@ -1789,10 +1817,12 @@ def b(self, m2: Model) -> Variable: 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_join_none_raises_on_mismatch(self, a: Variable, b: Variable) -> None: + # a has i=[0,1,2], b has i=[1,2,3] — exact default raises + with pytest.raises(ValueError, match="exact"): + a.to_linexpr() + b.to_linexpr() + with pytest.raises(ValueError, match="exact"): + a.to_linexpr().add(b.to_linexpr(), join=None) def test_add_expr_join_inner(self, a: Variable, b: Variable) -> None: result = a.to_linexpr().add(b.to_linexpr(), join="inner") @@ -1820,10 +1850,10 @@ def test_add_constant_join_outer(self, a: Variable) -> None: 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: + def test_add_constant_positional(self, a: 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") + const = xr.DataArray([10, 20, 30], dims=["i"], coords={"i": [5, 6, 7]}) + result = expr + const.assign_coords(i=expr.coords["i"]) assert list(result.data.indexes["i"]) == [0, 1, 2] assert (result.const.values == const.values).all() @@ -1880,8 +1910,8 @@ def test_merge_join_parameter(self, a: Variable, b: Variable) -> None: result: LinearExpression = 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") + def test_same_shape_add_assign_coords(self, a: Variable, c: Variable) -> None: + result = a.to_linexpr() + c.to_linexpr().assign_coords(i=a.coords["i"]) assert list(result.data.indexes["i"]) == [0, 1, 2] def test_add_expr_outer_const_values(self, a: Variable, b: Variable) -> None: @@ -1919,17 +1949,17 @@ def test_add_constant_inner_fill_values(self, a: Variable) -> None: 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: + def test_add_constant_positional_different_coords(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") + result = expr + other.assign_coords(i=expr.coords["i"]) 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: + def test_sub_constant_positional(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") + result = expr - other.assign_coords(i=expr.coords["i"]) assert list(result.coords["i"].values) == [0, 1, 2] np.testing.assert_array_equal(result.const.values, [-5, -15, -25]) @@ -1943,10 +1973,10 @@ def test_sub_expr_outer_const_values(self, a: Variable, b: Variable) -> None: 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: + def test_mul_constant_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") + result = expr * other.assign_coords(i=expr.coords["i"]) 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]) @@ -1963,10 +1993,10 @@ def test_mul_constant_outer_fill_values(self, a: Variable) -> None: 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: + def test_div_constant_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") + result = expr / other.assign_coords(i=expr.coords["i"]) assert list(result.coords["i"].values) == [0, 1, 2] np.testing.assert_array_equal(result.const.values, [5.0, 2.0, 1.0]) @@ -1990,16 +2020,16 @@ def test_variable_add_outer_values(self, a: Variable, b: Variable) -> None: assert set(result.coords["i"].values) == {0, 1, 2, 3} assert result.nterm == 2 - def test_variable_mul_override(self, a: Variable) -> None: + def test_variable_mul_positional(self, a: Variable) -> None: other = xr.DataArray([2, 3, 4], dims=["i"], coords={"i": [5, 6, 7]}) - result = a.mul(other, join="override") + result = a * other.assign_coords(i=a.coords["i"]) 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: + def test_variable_div_positional(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") + result = a / other.assign_coords(i=a.coords["i"]) assert isinstance(result, LinearExpression) assert list(result.coords["i"].values) == [0, 1, 2] np.testing.assert_array_almost_equal( @@ -2013,34 +2043,38 @@ def test_merge_outer_join(self, a: Variable, b: Variable) -> None: 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"]: + for join in ["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]) + # assign_coords also works when coords already match + result = expr_a + const.assign_coords(i=expr_a.coords["i"]) + 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: + def test_add_scalar(self, a: Variable) -> None: expr = 1 * a + 5 - result = expr.add(10, join="override") + result = expr + 10 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() + def test_quadratic_add_constant_join_inner(self, a: Variable, c: Variable) -> None: + quad = a.to_linexpr() * c.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] + assert list(result.data.indexes["i"]) == [1, 2] - def test_quadratic_add_expr_join_inner(self, a: Variable) -> None: - quad = a.to_linexpr() * a.to_linexpr() + def test_quadratic_add_expr_join_inner(self, a: Variable, c: Variable) -> None: + quad = a.to_linexpr() * c.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() + def test_quadratic_mul_constant_join_inner(self, a: Variable, c: Variable) -> None: + quad = a.to_linexpr() * c.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] + assert list(result.data.indexes["i"]) == [1, 2] def test_merge_join_left(self, a: Variable, b: Variable) -> None: result: LinearExpression = merge([a.to_linexpr(), b.to_linexpr()], join="left") diff --git a/test/test_optimization.py b/test/test_optimization.py index 492d703a..6bcb1627 100644 --- a/test/test_optimization.py +++ b/test/test_optimization.py @@ -186,8 +186,8 @@ def model_with_non_aligned_variables() -> Model: lower = pd.Series(0, range(8)) y = m.add_variables(lower=lower, coords=[lower.index], name="y") - m.add_constraints(x + y, GREATER_EQUAL, 10.5) - m.objective = 1 * x + 0.5 * y + m.add_constraints(x.add(y, join="outer"), GREATER_EQUAL, 10.5) + m.objective = x.add(0.5 * y, join="outer") return m diff --git a/test/test_typing.py b/test/test_typing.py index 99a27033..2375dc72 100644 --- a/test/test_typing.py +++ b/test/test_typing.py @@ -6,20 +6,41 @@ def test_operations_with_data_arrays_are_typed_correctly() -> None: m = linopy.Model() + s: xr.DataArray = xr.DataArray(5.0) + + v: linopy.Variable = m.add_variables(lower=0.0, name="v") + e: linopy.LinearExpression = v * 1.0 + q = v * v + + _ = s * v + _ = v * s + _ = v + s + + _ = s * e + _ = e * s + _ = e + s + + _ = s * q + _ = q * s + _ = q + s + + +def test_constant_with_extra_dims_broadcasts() -> None: + m = linopy.Model() + a: xr.DataArray = xr.DataArray([1, 2, 3]) v: linopy.Variable = m.add_variables(lower=0.0, name="v") e: linopy.LinearExpression = v * 1.0 q = v * v - _ = a * v - _ = v * a - _ = v + a + # Constants can introduce new dimensions (broadcasting) + result_v = a * v + assert "dim_0" in result_v.dims - _ = a * e - _ = e * a - _ = e + a + result_e = a * e + assert "dim_0" in result_e.dims - _ = a * q - _ = q * a - _ = q + a + # QuadraticExpression also allows constant broadcasting + result_q = a * q + assert isinstance(result_q, linopy.expressions.QuadraticExpression)