diff --git a/.github/dependabot.yml b/.github/dependabot.yml index f8f779b59..e7750e981 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -12,3 +12,30 @@ updates: github-actions: patterns: - '*' + +# Pinned ``[benchmarks]`` extra in pyproject.toml. One PR per dep bump +# → CodSpeed CI runs and attributes any perf delta to that specific +# bump. Keeps the cross-version ``sweep`` baseline (lockfile-pinned) +# stable while still surfacing upstream perf changes per-PR with +# eyes-open review. Loose ``[project.dependencies]`` (numpy, scipy, ...) +# have no version specifier so Dependabot leaves them alone — only the +# ``==`` pins in ``[benchmarks]`` produce PRs. +- package-ecosystem: pip + directory: / + schedule: + interval: monthly + open-pull-requests-limit: 5 + groups: + # Measurement scaffolding + CLI/notebook tooling. Perf-irrelevant — + # they don't move CodSpeed signal, so batching into one PR cuts + # review noise. Perf-relevant deps (numpy, xarray, highspy, …) stay + # un-grouped so each gets its own attributed CodSpeed delta. + benchmark-tooling: + patterns: + - pytest + - pytest-benchmark + - pytest-memray + - pytest-codspeed + - nbconvert + - typer + - plotly diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 6a439e802..a97cd4ef3 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,10 +1,14 @@ Closes # (if applicable). + + ## Changes proposed in this Pull Request ## Checklist +- [ ] AI-generated content is marked (see [`AGENTS.md`](../blob/master/AGENTS.md)). - [ ] Code changes are sufficiently documented; i.e. new functions contain docstrings and further explanations may be given in `doc`. - [ ] Unit tests for new features were added (if applicable). - [ ] A note for the release notes `doc/release_notes.rst` of the upcoming release is included. diff --git a/.github/workflows/benchmark-smoke.yml b/.github/workflows/benchmark-smoke.yml new file mode 100644 index 000000000..96e75ae1b --- /dev/null +++ b/.github/workflows/benchmark-smoke.yml @@ -0,0 +1,39 @@ +name: Benchmark smoke + +# Builds every spec and fires every phase once (--benchmark-disable): +# a "did a refactor break a spec?" check, not timing. + +on: + push: + branches: [ master ] + pull_request: + branches: [ '*' ] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + smoke: + name: Benchmark smoke (quick) + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 # setuptools_scm + + - name: Set up Python 3.13 + uses: actions/setup-python@v6 + with: + python-version: "3.13" + + - name: Install package and benchmark dependencies + run: | + python -m pip install uv + uv pip install --system -e ".[dev,benchmarks]" + + - name: Run benchmark smoke + # Every spec builds at one size and every phase fires once, no timings. + run: | + pytest benchmarks/ --benchmark-disable -q diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index 49a03d15b..95cfc2a3b 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -31,7 +31,7 @@ jobs: id-token: write steps: - name: Checkout repository - uses: actions/checkout@v6 + uses: actions/checkout@v7 with: fetch-depth: 1 diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 7c066b733..44a120caf 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -55,7 +55,7 @@ jobs: # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages steps: - name: Checkout repository - uses: actions/checkout@v6 + uses: actions/checkout@v7 # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml new file mode 100644 index 000000000..a70170d18 --- /dev/null +++ b/.github/workflows/codspeed.yml @@ -0,0 +1,76 @@ +name: CodSpeed + +# Both instruments in one workflow (CodSpeed aggregates per-commit results from a +# single workflow). Memory: heap tracking on a free GitHub runner, every PR + +# master. Walltime: bare-metal macro runner, master + the maintainer-gated +# `trigger:benchmark` label only. + +on: + push: + branches: [ master ] + pull_request: + branches: [ master ] + types: [ opened, synchronize, reopened, labeled ] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + memory: + name: Memory (heap) + runs-on: ubuntu-latest + continue-on-error: true # informational, never blocks a merge + permissions: + contents: read # actions/checkout + id-token: write # OIDC auth with CodSpeed — no token secret + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 # setuptools_scm + - name: Set up Python 3.13 + uses: actions/setup-python@v6 + with: + python-version: "3.13" + - name: Install pinned benchmark environment + run: | + python -m pip install uv + uv pip install --system -e ".[dev,benchmarks]" + - name: Run benchmarks under CodSpeed (memory) + uses: CodSpeedHQ/action@v4 + with: + mode: memory + run: | + pytest benchmarks/ --codspeed + + walltime: + name: Walltime (macro runner) + # Master push / dispatch always; PRs only when explicitly labelled — macro + # minutes are metered and bare-metal shouldn't run arbitrary PR code. + if: >- + ${{ github.event_name != 'pull_request' || + contains(github.event.pull_request.labels.*.name, 'trigger:benchmark') }} + runs-on: codspeed-macro + continue-on-error: true + permissions: + contents: read + id-token: write + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 # setuptools_scm + - name: Set up Python 3.13 + uses: actions/setup-python@v6 + with: + python-version: "3.13" + - name: Install pinned benchmark environment + run: | + python -m pip install uv + uv pip install --system -e ".[dev,benchmarks]" + - name: Run benchmarks under CodSpeed (walltime) + uses: CodSpeedHQ/action@v4 + with: + mode: walltime + run: | + pytest benchmarks/ --codspeed diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 54d9a2112..573293a63 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -11,7 +11,7 @@ jobs: name: Build and verify package runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: hynek/build-and-inspect-python-package@v2 release: @@ -20,8 +20,8 @@ jobs: needs: [build] runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 - - uses: softprops/action-gh-release@v2 + - uses: actions/checkout@v7 + - uses: softprops/action-gh-release@v3 with: generate_release_notes: true @@ -36,7 +36,7 @@ jobs: permissions: id-token: write steps: - - uses: actions/download-artifact@v7 + - uses: actions/download-artifact@v8 with: name: Packages path: dist diff --git a/.github/workflows/test-models.yml b/.github/workflows/test-models.yml index d5c14d4af..82d26568e 100644 --- a/.github/workflows/test-models.yml +++ b/.github/workflows/test-models.yml @@ -5,8 +5,7 @@ on: branches: - master pull_request: - branches: - - master + branches: ["*"] schedule: - cron: "0 5 * * *" @@ -24,14 +23,14 @@ jobs: matrix: version: - master - # - latest # Activate when v0.14.0 is released + # - latest defaults: run: shell: bash -l {0} steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: repository: PyPSA/pypsa-eur ref: master @@ -43,11 +42,23 @@ jobs: latest_tag=$(git describe --tags `git rev-list --tags --max-count=1`) git checkout $latest_tag + - name: Setup Pixi + uses: prefix-dev/setup-pixi@v0.9.6 + with: + pixi-version: v0.68.1 + cache: true + # Do not cache in branches + cache-write: ${{ github.event_name == 'push' && github.ref_name == 'master' }} + + - name: Setup cache keys + run: | + echo "WEEK=$(date +'%Y%U')" >> $GITHUB_ENV # data and cutouts + # Only run check if package is not pinned - - name: Check if inhouse package is pinned + - name: Check if linopy package is pinned run: | - grep_line=$(grep -- '- pypsa' envs/environment.yaml) - if [[ $grep_line == *"<"* || $grep_line == *"=="* ]]; then + grep_line=$(grep -- '${{ github.event.repository.name }} = ' pixi.toml) + if [[ $grep_line == *"<* || $grep_line == *"==* ]]; then echo "pinned=true" >> $GITHUB_ENV else echo "pinned=false" >> $GITHUB_ENV @@ -67,41 +78,27 @@ jobs: cutouts key: data-cutouts-${{ env.week }} - - uses: conda-incubator/setup-miniconda@v3 - if: env.pinned == 'false' - with: - activate-environment: pypsa-eur - - - name: Cache Conda env - if: env.pinned == 'false' - uses: actions/cache@v5 - with: - path: ${{ env.CONDA }}/envs - key: conda-pypsa-eur-${{ env.week }}-${{ hashFiles('envs/linux-64.lock.yaml') }} - id: cache-env - - - name: Update environment - if: env.pinned == 'false' && steps.cache-env.outputs.cache-hit != 'true' - run: conda env update -n pypsa-eur -f envs/linux-64.lock.yaml - - name: Install package from ref if: env.pinned == 'false' run: | - python -m pip install git+https://github.com/${{ github.repository }}@${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + pixi remove pypsa + pixi remove linopy + pixi add --pypi --git https://github.com/${{ github.repository }}.git ${{ github.event.repository.name }} --rev ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + pixi add --pypi pypsa - name: Run snakemake test workflows if: env.pinned == 'false' run: | - make test + pixi run integration-tests - name: Run unit tests if: env.pinned == 'false' run: | - make unit-test + pixi run unit-tests - name: Upload artifacts if: env.pinned == 'false' - uses: actions/upload-artifact@v6 + uses: actions/upload-artifact@v7 with: name: results-pypsa-eur-${{ matrix.version }} path: | @@ -109,3 +106,7 @@ jobs: .snakemake/log results retention-days: 3 + + - name: Show remaining disk space + if: always() + run: df -h diff --git a/.github/workflows/test-notebooks.yml b/.github/workflows/test-notebooks.yml new file mode 100644 index 000000000..9c9fdcd06 --- /dev/null +++ b/.github/workflows/test-notebooks.yml @@ -0,0 +1,58 @@ +name: Test Notebooks + +on: + push: + branches: [ master ] + pull_request: + branches: [ '*' ] + schedule: + - cron: "0 5 * * TUE" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + notebooks: + name: Test documentation notebooks + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + + - name: Set up Python 3.12 + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install package and dependencies + run: | + python -m pip install uv + uv pip install --system -e ".[docs]" + + - name: Execute notebooks + run: | + EXIT_CODE=0 + for notebook in examples/*.ipynb; do + name=$(basename "$notebook") + + # Skip notebooks that require credentials or special setup + case "$name" in + solve-on-oetc.ipynb|solve-on-remote.ipynb) + echo "Skipping $name (requires credentials or special setup)" + continue + ;; + esac + + echo "::group::Running $name" + if jupyter nbconvert --to notebook --execute --ExecutePreprocessor.timeout=600 "$notebook"; then + echo "✓ $name passed" + else + echo "::error::✗ $name failed" + EXIT_CODE=1 + fi + echo "::endgroup::" + done + exit $EXIT_CODE diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2253d2cfb..7ae92b4c1 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -18,7 +18,7 @@ jobs: name: Build and verify package runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 # Needed for setuptools_scm - uses: hynek/build-and-inspect-python-package@v2 @@ -42,7 +42,7 @@ jobs: - windows-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 # Needed for setuptools_scm @@ -64,7 +64,7 @@ jobs: - name: Set up windows package manager if: matrix.os == 'windows-latest' - uses: crazy-max/ghaction-chocolatey@v3 + uses: crazy-max/ghaction-chocolatey@v4 with: args: -h @@ -74,7 +74,7 @@ jobs: choco install glpk - name: Download package - uses: actions/download-artifact@v7 + uses: actions/download-artifact@v8 with: name: Packages path: dist @@ -82,7 +82,7 @@ jobs: - name: Install package and dependencies run: | python -m pip install uv - uv pip install --system "$(ls dist/*.whl)[dev,solvers]" + uv pip install --system "$(ls dist/*.whl)[dev,solvers,oetc]" - name: Test with pytest env: @@ -92,7 +92,7 @@ jobs: - name: Upload code coverage report if: matrix.os == 'ubuntu-latest' - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} @@ -102,7 +102,7 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 # Needed for setuptools_scm @@ -112,7 +112,7 @@ jobs: python-version: 3.12 - name: Download package - uses: actions/download-artifact@v7 + uses: actions/download-artifact@v8 with: name: Packages path: dist diff --git a/.gitignore b/.gitignore index 7b962a6b1..7e6d63e2f 100644 --- a/.gitignore +++ b/.gitignore @@ -32,6 +32,10 @@ ENV/ env.bak/ venv.bak/ +# Benchmarks (new pytest-benchmark suite) +.benchmarks/ + +# Benchmarks (old Snakemake suite in benchmark/) benchmark/*.pdf benchmark/benchmarks benchmark/.snakemake @@ -41,6 +45,10 @@ benchmark/scripts/__pycache__ benchmark/scripts/benchmarks-pypsa-eur/__pycache__ benchmark/scripts/leftovers/ +# Benchmarks (internal suite): regenerable .ipynb viewing artifacts +benchmarks/walkthrough.ipynb +benchmarks/.ipynb_checkpoints/ + # IDE .idea/ @@ -49,4 +57,4 @@ benchmark/scripts/leftovers/ # direnv .envrc -AGENTS.md +coverage.xml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 635784cdd..cc0f6720f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,13 +3,13 @@ ci: repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: end-of-file-fixer - id: trailing-whitespace - id: check-merge-conflict - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.12.2 + rev: v0.15.9 hooks: - id: ruff args: [--fix] @@ -21,19 +21,22 @@ repos: # - id: docformatter # args: [--in-place, --make-summary-multi-line, --pre-summary-newline] - repo: https://github.com/keewis/blackdoc - rev: v0.4.1 + rev: v0.4.6 hooks: - id: blackdoc exclude: ^dev-scripts/ additional_dependencies: ['black==24.8.0'] - repo: https://github.com/codespell-project/codespell - rev: v2.4.1 + rev: v2.4.2 hooks: - id: codespell + args: ['--ignore-words-list=coo'] types_or: [python, rst, markdown] files: ^(linopy|doc)/ -- repo: https://github.com/aflc/pre-commit-jupyter - rev: v1.2.1 +- repo: https://github.com/kynan/nbstripout + rev: 0.8.1 hooks: - - id: jupyter-notebook-cleanup + - id: nbstripout + args: + - --extra-keys=cell.metadata.ExecuteTime cell.metadata.execution exclude: examples/solve-on-remote.ipynb diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..635788a1e --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,73 @@ +# AGENTS.md + +Contribution rules and conventions for linopy, for humans and AI agents alike. + +## AI-assisted contributions + +You're welcome to help author code, pull requests and issues. But the +conversation around a change — PR and issue descriptions, comments and review +discussions — is how maintainers and contributors understand each other, so it +must stay human and honest. Three rules: + +1. **AI-generated content must be marked**, preferably with a GitHub note alert + at the top of the generated section. Put verbose additional context (logs, generated analysis, long + reproductions) in a collapsed `
` block so it doesn't drown the + discussion. + + Don't silently mix hand-written and generated prose. + +2. **The human writes their own intent.** If a pull request or issue is more + than a self-documenting bug report, the author must write — or at least + rewrite in their own words — their personal intent and motivation by hand, + with everything else placed below it. Don't submit an AI-drafted description + as your own; the maintainers want the human's voice, concise and to the + point. + + ```markdown + Thats why this is important for me (handwritten) + + > [!NOTE] + > The following content was generated by AI. + + What was implemented and how + +
+ Full benchmark output + + ...verbose content here... + +
+ ``` + +3. **Conversations are not held through an agent.** Posting bare information via + an agent — e.g. a log, a benchmark result, a reproduction — is fine as long + as it's marked (see rule 1). But discussion itself — replies, answers to + questions, review back-and-forth — must be written by the human. Don't let an + agent argue, agree or decide on your behalf. + +## Development workflow + +- Manage the environment with [`uv`](https://docs.astral.sh/uv/) and work inside + the project virtualenv. +- Run the test suite with `pytest`, lint and format with `ruff` + (`ruff check --fix .`), and type-check with `mypy`. +- For the authoritative, detailed commands (extras, coverage, GPU and benchmark + runs) see [`doc/contributing.rst`](doc/contributing.rst). + +## Project conventions + +- Branch off `master` for every change and open pull requests via the GitHub CLI + (`gh`). +- Write tests for new features and bug fixes under `test/` as `test_*.py`, using + linopy's own assertions (`assert_linequal`, `assert_varequal`, …) where + useful. Run the tests after making changes and make sure they pass. +- Keep temporary, non-tracked scripts in `dev-scripts/`. + +## Architecture in one paragraph + +linopy is a linear optimization library built on xarray: variables, constraints +and expressions are dimension-labelled N-dimensional structures. Arithmetic on +them broadcasts and aligns by dimension and builds expressions lazily, which are +turned into the solver format only at solve time. Solvers sit behind an abstract +interface with both file-based and direct-API backends. Keep new features +consistent with this xarray-based, lazily-evaluated design. diff --git a/CLAUDE.md b/CLAUDE.md index 67155ae3d..276bc76e2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,142 +1,8 @@ # CLAUDE.md -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. +This project's guidance for AI agents — contribution rules, development workflow +and project conventions — lives in [`AGENTS.md`](AGENTS.md), imported below so +it loads automatically. Read it before making changes, opening pull requests or +issues, or writing comments. -## Common Development Commands - -### Running Tests -```bash -# Run all tests (excluding GPU tests by default) -pytest - -# Run tests with coverage -pytest --cov=./ --cov-report=xml linopy --doctest-modules test - -# Run a specific test file -pytest test/test_model.py - -# Run a specific test function -pytest test/test_model.py::test_model_creation - -# Run GPU tests (requires GPU hardware and cuPDLPx installation) -pytest --run-gpu - -# Run only GPU tests -pytest -m gpu --run-gpu -``` - -**GPU Testing**: Tests that require GPU hardware (e.g., cuPDLPx solver) are automatically skipped by default since CI machines typically don't have GPUs. To run GPU tests locally, use the `--run-gpu` flag. The tests are automatically marked with `@pytest.mark.gpu` based on solver capabilities. - -### Linting and Type Checking -```bash -# Run linter (ruff) -ruff check . -ruff check --fix . # Auto-fix issues - -# Run formatter -ruff format . - -# Run type checker -mypy . - -# Run all pre-commit hooks -pre-commit run --all-files -``` - -### Development Setup -```bash -# Create virtual environment and install development dependencies -python -m venv venv -source venv/bin/activate # On Windows: venv\Scripts\activate -pip install uv -uv pip install -e .[dev,solvers] -``` - -## High-Level Architecture - -linopy is a linear optimization library built on top of xarray, providing N-dimensional labeled arrays for variables and constraints. The architecture follows these key principles: - -### Core Components - -1. **Model** (`model.py`): Central container for optimization problems - - Manages variables, constraints, and objective - - Handles solver integration through abstract interfaces - - Supports chunked operations for memory efficiency - - Provides matrix representations for solver APIs - -2. **Variables** (`variables.py`): Multi-dimensional decision variables - - Built on xarray.Dataset with labels, lower, and upper bounds - - Arithmetic operations automatically create LinearExpressions - - Support for continuous and binary variables - - Container class (Variables) manages collections with dict-like access - -3. **Constraints** (`constraints.py`): Linear inequality/equality constraints - - Store coefficients, variable references, signs, and RHS values - - Support ≤, ≥, and = constraints - - Container class (Constraints) provides organized access - -4. **Expressions** (`expressions.py`): Linear combinations of variables - - LinearExpression: coeffs × vars + const - - QuadraticExpression: for non-linear optimization - - Support full arithmetic operations with automatic broadcasting - - Special `_term` dimension for handling multiple terms - -5. **Solvers** (`solvers.py`): Abstract interface with multiple implementations - - File-based solvers: Write LP/MPS files, call solver, parse results - - Direct API solvers: Use Python bindings (e.g., gurobipy) - - Automatic solver detection based on installed packages - -### Data Flow Pattern - -1. User creates Model and adds Variables with coordinates (dimensions) -2. Variables combined into LinearExpressions through arithmetic -3. Expressions used to create Constraints and Objective -4. Model.solve() converts to solver format and retrieves solution -5. Solution stored back in xarray format with original dimensions - -### Key Design Patterns - -- **xarray Integration**: All data structures use xarray for dimension handling -- **Lazy Evaluation**: Expressions built symbolically before solving -- **Broadcasting**: Operations automatically align dimensions -- **Solver Abstraction**: Clean separation between model and solver specifics -- **Memory Efficiency**: Support for dask arrays and chunked operations - -When modifying the codebase, maintain consistency with these patterns and ensure new features integrate naturally with the xarray-based architecture. - -## Working with the Github Repository - -* The main branch is `master`. -* 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 - -1. Always write tests for new features or bug fixes. -2. Always run the tests after making changes and ensure they pass. -3. Always use ruff for linting and formatting, run `ruff check --fix .` to auto-fix issues. -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. +@AGENTS.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d1e26cc36..2c41201d7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1 +1,3 @@ Linopy's contributor guidelines can be found in the official [documentation](https://linopy.readthedocs.io/en/latest/contributing.html). + +If you use AI tools when contributing, please read [`AGENTS.md`](AGENTS.md) for how AI-generated content must be marked and what we expect you to write by hand. diff --git a/README.md b/README.md index 644b556c6..870bbfb4e 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ So far **linopy** is available on the PyPI repository ```bash -pip install linopy +uv pip install linopy ``` or on conda-forge @@ -143,13 +143,14 @@ Fri 0 4 * [Cbc](https://projects.coin-or.org/Cbc) * [GLPK](https://www.gnu.org/software/glpk/) -* [HiGHS](https://www.maths.ed.ac.uk/hall/HiGHS/) +* [HiGHS](https://highs.dev/) * [Gurobi](https://www.gurobi.com/) * [Xpress](https://www.fico.com/en/products/fico-xpress-solver) * [Cplex](https://www.ibm.com/de-de/analytics/cplex-optimizer) * [MOSEK](https://www.mosek.com/) * [COPT](https://www.shanshu.ai/copt) * [cuPDLPx](https://github.com/MIT-Lu-Lab/cuPDLPx) +* [Knitro](https://www.artelys.com/solvers/knitro/) Note that these do have to be installed by the user separately. @@ -158,10 +159,8 @@ Note that these do have to be installed by the user separately. To set up a local development environment for linopy and to run the same tests that are run in the CI, you can run: ```sh -python -m venv venv -source venv/bin/activate -pip install uv -uv pip install -e .[dev,solvers] +uv sync --extra dev --extra solvers +source .venv/bin/activate pytest ``` diff --git a/benchmark/benchmark_auto_mask.py b/benchmark/benchmark_auto_mask.py new file mode 100644 index 000000000..d478e9501 --- /dev/null +++ b/benchmark/benchmark_auto_mask.py @@ -0,0 +1,639 @@ +#!/usr/bin/env python3 +""" +Benchmark comparing manual masking vs auto_mask for models with NaN coefficients. + +This creates a realistic scenario: a multi-period dispatch model where: +- Not all generators are available in all time periods (NaN in capacity bounds) +- Not all transmission lines exist between all regions (NaN in flow limits) +""" + +import sys +from pathlib import Path + +# Ensure we use the local linopy installation +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + +import time # noqa: E402 +from typing import Any # noqa: E402 + +import numpy as np # noqa: E402 +import pandas as pd # noqa: E402 + +from linopy import GREATER_EQUAL, Model # noqa: E402 + + +def create_nan_data( + n_generators: int = 500, + n_periods: int = 100, + n_regions: int = 20, + nan_fraction_gen: float = 0.3, # 30% of generator-period combinations unavailable + nan_fraction_lines: float = 0.7, # 70% of region pairs have no direct line + seed: int = 42, +) -> dict[str, Any]: + """Create realistic input data with NaN patterns.""" + rng = np.random.default_rng(seed) + + generators = pd.Index(range(n_generators), name="generator") + periods = pd.Index(range(n_periods), name="period") + regions = pd.Index(range(n_regions), name="region") + + # Generator capacities - some generators unavailable in some periods (maintenance, etc.) + gen_capacity = pd.DataFrame( + rng.uniform(50, 500, size=(n_generators, n_periods)), + index=generators, + columns=periods, + ) + # Set random entries to NaN (generator unavailable) + nan_mask_gen = rng.random((n_generators, n_periods)) < nan_fraction_gen + gen_capacity.values[nan_mask_gen] = np.nan + + # Generator costs + gen_cost = pd.Series(rng.uniform(10, 100, n_generators), index=generators) + + # Generator to region mapping + gen_region = pd.Series(rng.integers(0, n_regions, n_generators), index=generators) + + # Demand per region per period + demand = pd.DataFrame( + rng.uniform(100, 1000, size=(n_regions, n_periods)), + index=regions, + columns=periods, + ) + + # Transmission line capacities - sparse network (not all regions connected) + # Use distinct dimension names to avoid xarray duplicate dimension issues + regions_from = pd.Index(range(n_regions), name="region_from") + regions_to = pd.Index(range(n_regions), name="region_to") + + line_capacity = pd.DataFrame( + np.nan, + index=regions_from, + columns=regions_to, + dtype=float, # Start with all NaN + ) + # Only some region pairs have lines + for i in range(n_regions): + for j in range(n_regions): + if i != j and rng.random() > nan_fraction_lines: + line_capacity.loc[i, j] = rng.uniform(100, 500) + + return { + "generators": generators, + "periods": periods, + "regions": regions, + "regions_from": regions_from, + "regions_to": regions_to, + "gen_capacity": gen_capacity, + "gen_cost": gen_cost, + "gen_region": gen_region, + "demand": demand, + "line_capacity": line_capacity, + } + + +def build_model_manual_mask(data: dict[str, Any]) -> Model: + """Build model using manual masking (traditional approach).""" + m = Model() + + generators = data["generators"] + periods = data["periods"] + regions = data["regions"] + regions_from = data["regions_from"] + regions_to = data["regions_to"] + gen_capacity = data["gen_capacity"] + gen_cost = data["gen_cost"] + gen_region = data["gen_region"] + demand = data["demand"] + line_capacity = data["line_capacity"] + + # Generator dispatch variables - manually mask where capacity is NaN + gen_mask = gen_capacity.notnull() + dispatch = m.add_variables( + lower=0, + upper=gen_capacity, + coords=[generators, periods], + name="dispatch", + mask=gen_mask, + ) + + # Flow variables between regions - manually mask where no line exists + flow_mask = line_capacity.notnull() + flow = m.add_variables( + lower=-line_capacity.abs(), + upper=line_capacity.abs(), + coords=[regions_from, regions_to], + name="flow", + mask=flow_mask, + ) + + # Energy balance constraint per region per period + for r in regions: + gens_in_region = generators[gen_region == r] + gen_sum = dispatch.loc[gens_in_region, :].sum("generator") + + # Net flow into region + flow_in = flow.loc[:, r].sum("region_from") + flow_out = flow.loc[r, :].sum("region_to") + + m.add_constraints( + gen_sum + flow_in - flow_out, + GREATER_EQUAL, + demand.loc[r], + name=f"balance_r{r}", + ) + + # Objective: minimize generation cost + obj = (dispatch * gen_cost).sum() + m.add_objective(obj) + + return m + + +def build_model_auto_mask(data: dict[str, Any]) -> Model: + """Build model using auto_mask=True (new approach).""" + m = Model(auto_mask=True) + + generators = data["generators"] + periods = data["periods"] + regions = data["regions"] + regions_from = data["regions_from"] + regions_to = data["regions_to"] + gen_capacity = data["gen_capacity"] + gen_cost = data["gen_cost"] + gen_region = data["gen_region"] + demand = data["demand"] + line_capacity = data["line_capacity"] + + # Generator dispatch variables - auto-masked where capacity is NaN + dispatch = m.add_variables( + lower=0, + upper=gen_capacity, # NaN values will be auto-masked + coords=[generators, periods], + name="dispatch", + ) + + # Flow variables between regions - auto-masked where no line exists + flow = m.add_variables( + lower=-line_capacity.abs(), + upper=line_capacity.abs(), # NaN values will be auto-masked + coords=[regions_from, regions_to], + name="flow", + ) + + # Energy balance constraint per region per period + for r in regions: + gens_in_region = generators[gen_region == r] + gen_sum = dispatch.loc[gens_in_region, :].sum("generator") + + # Net flow into region + flow_in = flow.loc[:, r].sum("region_from") + flow_out = flow.loc[r, :].sum("region_to") + + m.add_constraints( + gen_sum + flow_in - flow_out, + GREATER_EQUAL, + demand.loc[r], + name=f"balance_r{r}", + ) + + # Objective: minimize generation cost + obj = (dispatch * gen_cost).sum() + m.add_objective(obj) + + return m + + +def build_model_no_mask(data: dict[str, Any]) -> Model: + """Build model WITHOUT any masking (NaN values left in place).""" + m = Model() + + generators = data["generators"] + periods = data["periods"] + regions = data["regions"] + regions_from = data["regions_from"] + regions_to = data["regions_to"] + gen_capacity = data["gen_capacity"] + gen_cost = data["gen_cost"] + gen_region = data["gen_region"] + demand = data["demand"] + line_capacity = data["line_capacity"] + + # Generator dispatch variables - NO masking, NaN bounds left in place + dispatch = m.add_variables( + lower=0, + upper=gen_capacity, # Contains NaN values + coords=[generators, periods], + name="dispatch", + ) + + # Flow variables between regions - NO masking + flow = m.add_variables( + lower=-line_capacity.abs(), + upper=line_capacity.abs(), # Contains NaN values + coords=[regions_from, regions_to], + name="flow", + ) + + # Energy balance constraint per region per period + for r in regions: + gens_in_region = generators[gen_region == r] + gen_sum = dispatch.loc[gens_in_region, :].sum("generator") + + # Net flow into region + flow_in = flow.loc[:, r].sum("region_from") + flow_out = flow.loc[r, :].sum("region_to") + + m.add_constraints( + gen_sum + flow_in - flow_out, + GREATER_EQUAL, + demand.loc[r], + name=f"balance_r{r}", + ) + + # Objective: minimize generation cost + obj = (dispatch * gen_cost).sum() + m.add_objective(obj) + + return m + + +def benchmark( + n_generators: int = 500, + n_periods: int = 100, + n_regions: int = 20, + n_runs: int = 3, + solve: bool = True, +) -> dict[str, Any]: + """Run benchmark comparing no masking, manual masking, and auto masking.""" + print("=" * 70) + print("BENCHMARK: No Masking vs Manual Masking vs Auto-Masking") + print("=" * 70) + print("\nModel size:") + print(f" - Generators: {n_generators}") + print(f" - Time periods: {n_periods}") + print(f" - Regions: {n_regions}") + print(f" - Potential dispatch vars: {n_generators * n_periods:,}") + print(f" - Potential flow vars: {n_regions * n_regions:,}") + print(f"\nRunning {n_runs} iterations each...\n") + + # Generate data once + data = create_nan_data( + n_generators=n_generators, + n_periods=n_periods, + n_regions=n_regions, + ) + + # Count NaN entries + gen_nan_count = data["gen_capacity"].isna().sum().sum() + gen_total = data["gen_capacity"].size + line_nan_count = data["line_capacity"].isna().sum().sum() + line_total = data["line_capacity"].size + + print("NaN statistics:") + print( + f" - Generator capacity: {gen_nan_count:,}/{gen_total:,} " + f"({100 * gen_nan_count / gen_total:.1f}% NaN)" + ) + print( + f" - Line capacity: {line_nan_count:,}/{line_total:,} " + f"({100 * line_nan_count / line_total:.1f}% NaN)" + ) + print() + + # Benchmark NO masking (baseline) + no_mask_times = [] + for i in range(n_runs): + start = time.perf_counter() + m_no_mask = build_model_no_mask(data) + elapsed = time.perf_counter() - start + no_mask_times.append(elapsed) + if i == 0: + # Can't use nvars directly as it will fail with NaN values + # Instead count total variable labels (including those with NaN bounds) + no_mask_nvars = sum( + m_no_mask.variables[k].labels.size for k in m_no_mask.variables + ) + no_mask_ncons = m_no_mask.ncons + + # Benchmark manual masking + manual_times = [] + for i in range(n_runs): + start = time.perf_counter() + m_manual = build_model_manual_mask(data) + elapsed = time.perf_counter() - start + manual_times.append(elapsed) + if i == 0: + manual_nvars = m_manual.nvars + manual_ncons = m_manual.ncons + + # Benchmark auto masking + auto_times = [] + for i in range(n_runs): + start = time.perf_counter() + m_auto = build_model_auto_mask(data) + elapsed = time.perf_counter() - start + auto_times.append(elapsed) + if i == 0: + auto_nvars = m_auto.nvars + auto_ncons = m_auto.ncons + + # Results + print("-" * 70) + print("RESULTS: Model Building Time") + print("-" * 70) + + print("\nNo masking (baseline):") + print(f" - Mean time: {np.mean(no_mask_times):.3f}s") + print(f" - Variables: {no_mask_nvars:,} (includes NaN-bounded vars)") + print(f" - Constraints: {no_mask_ncons:,}") + + print("\nManual masking:") + print(f" - Mean time: {np.mean(manual_times):.3f}s") + print(f" - Variables: {manual_nvars:,}") + print(f" - Constraints: {manual_ncons:,}") + manual_overhead = np.mean(manual_times) - np.mean(no_mask_times) + print(f" - Overhead vs no-mask: {manual_overhead * 1000:+.1f}ms") + + print("\nAuto masking:") + print(f" - Mean time: {np.mean(auto_times):.3f}s") + print(f" - Variables: {auto_nvars:,}") + print(f" - Constraints: {auto_ncons:,}") + auto_overhead = np.mean(auto_times) - np.mean(no_mask_times) + print(f" - Overhead vs no-mask: {auto_overhead * 1000:+.1f}ms") + + # Comparison + print("\nComparison (Auto vs Manual):") + speedup = np.mean(manual_times) / np.mean(auto_times) + diff = np.mean(auto_times) - np.mean(manual_times) + if speedup > 1: + print(f" - Auto-mask is {speedup:.2f}x FASTER than manual") + else: + print(f" - Auto-mask is {1 / speedup:.2f}x SLOWER than manual") + print(f" - Time difference: {diff * 1000:+.1f}ms") + + # Verify models are equivalent + print("\nVerification:") + print(f" - Manual == Auto variables: {manual_nvars == auto_nvars}") + print(f" - Manual == Auto constraints: {manual_ncons == auto_ncons}") + print(f" - Variables masked out: {no_mask_nvars - manual_nvars:,}") + + results = { + "n_generators": n_generators, + "n_periods": n_periods, + "potential_vars": n_generators * n_periods, + "no_mask_time": np.mean(no_mask_times), + "manual_time": np.mean(manual_times), + "auto_time": np.mean(auto_times), + "nvars": manual_nvars, + "masked_out": no_mask_nvars - manual_nvars, + } + + # LP file write benchmark + print("\n" + "-" * 70) + print("RESULTS: LP File Write Time & Size") + print("-" * 70) + + import os + import tempfile + + # Write LP file for manual masked model + with tempfile.NamedTemporaryFile(suffix=".lp", delete=False) as f: + manual_lp_path = f.name + start = time.perf_counter() + m_manual.to_file(manual_lp_path) + manual_write_time = time.perf_counter() - start + manual_lp_size = os.path.getsize(manual_lp_path) / (1024 * 1024) # MB + os.unlink(manual_lp_path) + + # Write LP file for auto masked model + with tempfile.NamedTemporaryFile(suffix=".lp", delete=False) as f: + auto_lp_path = f.name + start = time.perf_counter() + m_auto.to_file(auto_lp_path) + auto_write_time = time.perf_counter() - start + auto_lp_size = os.path.getsize(auto_lp_path) / (1024 * 1024) # MB + os.unlink(auto_lp_path) + + print("\nManual masking:") + print(f" - Write time: {manual_write_time:.3f}s") + print(f" - File size: {manual_lp_size:.2f} MB") + + print("\nAuto masking:") + print(f" - Write time: {auto_write_time:.3f}s") + print(f" - File size: {auto_lp_size:.2f} MB") + + print(f"\nFiles identical: {abs(manual_lp_size - auto_lp_size) < 0.01}") + + results["manual_write_time"] = manual_write_time + results["auto_write_time"] = auto_write_time + results["lp_size_mb"] = manual_lp_size + + # Quick solve comparison + if solve: + print("\n" + "-" * 70) + print("RESULTS: Solve Time (single run)") + print("-" * 70) + + start = time.perf_counter() + m_manual.solve("highs", io_api="direct") + manual_solve = time.perf_counter() - start + + start = time.perf_counter() + m_auto.solve("highs", io_api="direct") + auto_solve = time.perf_counter() - start + + print(f"\nManual masking solve: {manual_solve:.3f}s") + print(f"Auto masking solve: {auto_solve:.3f}s") + + if m_manual.objective.value is not None and m_auto.objective.value is not None: + print( + f"Objective values match: " + f"{np.isclose(m_manual.objective.value, m_auto.objective.value)}" + ) + print(f" - Manual: {m_manual.objective.value:.2f}") + print(f" - Auto: {m_auto.objective.value:.2f}") + + return results + + +def benchmark_code_simplicity() -> None: + """Show the code simplicity benefit of auto_mask.""" + print("\n" + "=" * 70) + print("CODE COMPARISON: Manual vs Auto-Mask") + print("=" * 70) + + manual_code = """ +# Manual masking - must create mask explicitly +gen_mask = gen_capacity.notnull() +dispatch = m.add_variables( + lower=0, + upper=gen_capacity, + coords=[generators, periods], + name="dispatch", + mask=gen_mask, # Extra step required +) +""" + + auto_code = """ +# Auto masking - just pass the data with NaN +m = Model(auto_mask=True) +dispatch = m.add_variables( + lower=0, + upper=gen_capacity, # NaN auto-masked + coords=[generators, periods], + name="dispatch", +) +""" + + print("\nManual masking approach:") + print(manual_code) + print("Auto-mask approach:") + print(auto_code) + print("Benefits of auto_mask:") + print(" - Less boilerplate code") + print(" - No need to manually track which arrays have NaN") + print(" - Reduces risk of forgetting to mask") + print(" - Cleaner, more declarative style") + + +def benchmark_constraint_masking(n_runs: int = 3) -> None: + """Benchmark auto-masking of constraints with NaN in RHS.""" + print("\n" + "=" * 70) + print("BENCHMARK: Constraint Auto-Masking (NaN in RHS)") + print("=" * 70) + + n_vars = 1000 + n_constraints = 5000 + nan_fraction = 0.3 + + rng = np.random.default_rng(42) + idx = pd.Index(range(n_vars), name="i") + con_idx = pd.Index(range(n_constraints), name="c") + + # Create RHS with NaN values + rhs = pd.Series(rng.uniform(1, 100, n_constraints), index=con_idx) + nan_mask = rng.random(n_constraints) < nan_fraction + rhs.values[nan_mask] = np.nan + + print("\nModel size:") + print(f" - Variables: {n_vars}") + print(f" - Potential constraints: {n_constraints}") + print(f" - NaN in RHS: {nan_mask.sum()} ({100 * nan_fraction:.0f}%)") + print(f"\nRunning {n_runs} iterations each...\n") + + # Manual masking + manual_times = [] + for i in range(n_runs): + start = time.perf_counter() + m = Model() + x = m.add_variables(lower=0, coords=[idx], name="x") + coeffs = pd.DataFrame( + rng.uniform(0.1, 1, (n_constraints, n_vars)), index=con_idx, columns=idx + ) + con_mask = rhs.notnull() # Manual mask creation + m.add_constraints((coeffs * x).sum("i"), GREATER_EQUAL, rhs, mask=con_mask) + m.add_objective(x.sum()) + elapsed = time.perf_counter() - start + manual_times.append(elapsed) + if i == 0: + manual_ncons = m.ncons + + # Auto masking + auto_times = [] + for i in range(n_runs): + start = time.perf_counter() + m = Model(auto_mask=True) + x = m.add_variables(lower=0, coords=[idx], name="x") + coeffs = pd.DataFrame( + rng.uniform(0.1, 1, (n_constraints, n_vars)), index=con_idx, columns=idx + ) + m.add_constraints((coeffs * x).sum("i"), GREATER_EQUAL, rhs) # No mask needed + m.add_objective(x.sum()) + elapsed = time.perf_counter() - start + auto_times.append(elapsed) + if i == 0: + auto_ncons = m.ncons + + print("-" * 70) + print("RESULTS: Constraint Building Time") + print("-" * 70) + print("\nManual masking:") + print(f" - Mean time: {np.mean(manual_times):.3f}s") + print(f" - Active constraints: {manual_ncons:,}") + + print("\nAuto masking:") + print(f" - Mean time: {np.mean(auto_times):.3f}s") + print(f" - Active constraints: {auto_ncons:,}") + + overhead = np.mean(auto_times) - np.mean(manual_times) + print(f"\nOverhead: {overhead * 1000:.1f}ms") + print(f"Same constraint count: {manual_ncons == auto_ncons}") + + +def print_summary_table(results: list[dict[str, Any]]) -> None: + """Print a summary table of all benchmark results.""" + print("\n" + "=" * 110) + print("SUMMARY TABLE: Model Building & LP Write Times") + print("=" * 110) + print( + f"{'Model':<12} {'Pot.Vars':>10} {'Act.Vars':>10} {'Masked':>8} " + f"{'No-Mask':>9} {'Manual':>9} {'Auto':>9} {'Diff':>8} " + f"{'LP Write':>9} {'LP Size':>9}" + ) + print("-" * 110) + for r in results: + name = f"{r['n_generators']}x{r['n_periods']}" + lp_write = r.get("manual_write_time", 0) * 1000 + lp_size = r.get("lp_size_mb", 0) + print( + f"{name:<12} {r['potential_vars']:>10,} {r['nvars']:>10,} " + f"{r['masked_out']:>8,} {r['no_mask_time'] * 1000:>8.0f}ms " + f"{r['manual_time'] * 1000:>8.0f}ms {r['auto_time'] * 1000:>8.0f}ms " + f"{(r['auto_time'] - r['manual_time']) * 1000:>+7.0f}ms " + f"{lp_write:>8.0f}ms {lp_size:>8.1f}MB" + ) + print("-" * 110) + print("Pot.Vars = Potential variables, Act.Vars = Active (non-masked) variables") + print("Masked = Variables masked out due to NaN bounds") + print("Diff = Auto-mask time minus Manual mask time (negative = faster)") + print("LP Write = Time to write LP file, LP Size = LP file size in MB") + + +if __name__ == "__main__": + all_results = [] + + # Run benchmarks with different sizes + print("\n### SMALL MODEL ###") + all_results.append( + benchmark(n_generators=100, n_periods=50, n_regions=10, n_runs=5, solve=False) + ) + + print("\n\n### MEDIUM MODEL ###") + all_results.append( + benchmark(n_generators=500, n_periods=100, n_regions=20, n_runs=3, solve=False) + ) + + print("\n\n### LARGE MODEL ###") + all_results.append( + benchmark(n_generators=1000, n_periods=200, n_regions=30, n_runs=3, solve=False) + ) + + print("\n\n### VERY LARGE MODEL ###") + all_results.append( + benchmark(n_generators=2000, n_periods=500, n_regions=40, n_runs=3, solve=False) + ) + + print("\n\n### EXTRA LARGE MODEL ###") + all_results.append( + benchmark(n_generators=5000, n_periods=500, n_regions=50, n_runs=2, solve=False) + ) + + # Print summary table + print_summary_table(all_results) + + # Run constraint benchmark + benchmark_constraint_masking() + + # Show code comparison + benchmark_code_simplicity() diff --git a/benchmark/notebooks/plot-benchmarks.py.ipynb b/benchmark/notebooks/plot-benchmarks.py.ipynb index f1099a0b1..f61d12e57 100644 --- a/benchmark/notebooks/plot-benchmarks.py.ipynb +++ b/benchmark/notebooks/plot-benchmarks.py.ipynb @@ -3,7 +3,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9a85db47", + "id": "0", "metadata": {}, "outputs": [], "source": [ @@ -19,7 +19,7 @@ { "cell_type": "code", "execution_count": null, - "id": "709bdf49", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -31,7 +31,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f36897fb", + "id": "2", "metadata": {}, "outputs": [], "source": [ @@ -65,7 +65,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c5c93666", + "id": "3", "metadata": {}, "outputs": [], "source": [ diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..aa293a154 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,81 @@ +# Internal Performance Benchmarks + +End-to-end performance tracking for `linopy` — build → matrix generation → +LP / netCDF (de)serialization → solver handoff → a fixed PyPSA model. Solver +algorithm runtime is out of scope. + +The suite is a set of `pytest-benchmark` tests driven by a model registry. +**CodSpeed** measures them in CI (walltime on dedicated runners, memory on every +PR); locally you just run `pytest`. + +> `benchmark/` (singular) is the legacy external-framework suite. +> `benchmarks/` (plural) is this internal suite. + +## Layout + +- `registry.py`, `phases.py`, `conftest.py` — the harness (specs, measured + verbs, pytest wiring). +- `models/`, `patterns/` — the subjects; each file self-registers one `BenchSpec`. +- `drivers/` — one `test_.py` per measured phase. + +## Models vs patterns + +Two kinds of benchmark spec, same harness and same phases, distinguished by +their sweep axis: + +- **Models** (`models/`, `REGISTRY`) — whole `linopy.Model`s swept over + `size` (axis `n`): "how does cost scale with the problem?" +- **Patterns** (`patterns/`, `PATTERNS`) — fragments of realistic modelling + code (a balance constraint, a KVL contraction) swept over `severity` + (0–100, axis `severity`): "how does cost respond as one data shape goes + from benign to pathological?" + +Both kinds build a complete `linopy.Model`, so both run the **same phases** and +share the phase drivers (`drivers/test_build.py`, `drivers/test_matrices.py`, …) +— they're just more `(spec, value)` rows, tagged by `axis`. There is no separate +pattern driver. Running a pattern through `build` *and* `to_lp` shows whether a +dense-`_term` blow-up propagates to export or collapses. + +Patterns target the operations where the dense-`_term` representation forces +materialisation — `groupby().sum()` padding, sparse `@` densification — so a +`severity` sweep draws the cost cliff. Adding either kind is one file: drop it +in `models/` or `patterns/`, call `register(...)` / `register_pattern(...)`. + +## Install + +```bash +uv sync --extra dev --extra benchmarks +source .venv/bin/activate +``` + +`pypsa` is optional — `pypsa_scigrid` and `drivers/test_pypsa_carbon_management.py` +skip gracefully without it: `uv pip install pypsa`. + +The `[benchmarks]` extra in `pyproject.toml` pins every direct dep that affects +measurement (`numpy`, `scipy`, `xarray`, `pandas`, `polars`, `dask`, …) so +run-to-run deltas reflect linopy changes, not dependency bumps. + +## Running + +```bash +pytest benchmarks/ # the suite +pytest benchmarks/ --benchmark-disable -q # smoke: every spec builds once +pytest benchmarks/ --pipeline # + the opt-in end-to-end pipeline test +``` + +Each spec declares one `sizes` (models) / `severities` (patterns) tuple — a +small representative set, kept tight because CodSpeed measures it on every PR. +Need a scaling curve? That's a local pytest-benchmem job, not this suite. + +## CI + +- **Smoke** (`benchmark-smoke.yml`) — every PR: every spec builds and every + phase fires once under `--benchmark-disable`. A "did a refactor break a + spec?" check, not timing. +- **CodSpeed** (`codspeed.yml`) — two jobs: **memory** (heap-allocation + tracking, every PR, free GitHub runner) and **walltime** (bare-metal macro + runner, on `master` or a PR labelled `trigger:benchmark`). Informational, + non-gating. + +Activating CodSpeed upstream needs a maintainer to connect the repo to the +CodSpeed app (OIDC auth, no token secret); the workflows are already wired. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..48c26ef03 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1,15 @@ +""" +Linopy benchmark suite — run with ``pytest benchmarks/``. + +The model registry it drives is reusable on its own:: + + from benchmarks import REGISTRY + model = REGISTRY["basic"].build(100) +""" + +# Importing the models / patterns packages triggers each module's +# ``register(...)`` / ``register_pattern(...)`` call at import time. +from benchmarks import models, patterns # noqa: F401 +from benchmarks.registry import PATTERNS, REGISTRY + +__all__ = ["PATTERNS", "REGISTRY"] diff --git a/benchmarks/conftest.py b/benchmarks/conftest.py new file mode 100644 index 000000000..b9ef60148 --- /dev/null +++ b/benchmarks/conftest.py @@ -0,0 +1,104 @@ +"""Benchmark configuration and shared test helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from benchmarks.registry import iter_params, spec_param_id + +if TYPE_CHECKING: + import linopy + from benchmarks.registry import BenchSpec + +# Test modules the CodSpeed instruments measure (edit to change coverage). +# build + the two export paths: to_lp (LP text) and to_solver (direct handoff, +# which also exercises matrix-gen). matrices is dropped — a subset of to_solver; +# netcdf excluded — disk I/O, noisy. All still run under the smoke job. +CODSPEED_MODULES = ( + "test_build", + "test_to_lp", + "test_to_solver", +) + + +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--pipeline", + action="store_true", + default=False, + help=( + "Include the opt-in end-to-end pipeline benchmark (build → matrices " + "→ lp in one measured region). Off by default — it re-runs the " + "per-phase work and includes the build." + ), + ) + + +def pytest_collection_modifyitems( + config: pytest.Config, items: list[pytest.Item] +) -> None: + """ + ``test_pipeline`` (end-to-end) is opt-in — deselected unless ``--pipeline``. + ``--codspeed`` narrows the run to ``CODSPEED_MODULES`` (drops netcdf/matrices). + """ + if not config.getoption("--pipeline"): + dropped = [i for i in items if i.path.stem == "test_pipeline"] + if dropped: + config.hook.pytest_deselected(items=dropped) + items[:] = [i for i in items if i.path.stem != "test_pipeline"] + + if getattr(config.option, "codspeed", False): + deselected = [i for i in items if i.path.stem not in CODSPEED_MODULES] + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = [i for i in items if i.path.stem in CODSPEED_MODULES] + + +def cases(phase: str) -> pytest.MarkDecorator: + """Parametrize a phase driver over every ``(spec, n)`` that phase runs.""" + params = iter_params(phase) + return pytest.mark.parametrize( + ("spec", "n"), + params, + ids=[spec_param_id(s.name, s.axis, v) for s, v in params], + ) + + +def require(spec: BenchSpec) -> None: + """``importorskip`` a spec's optional dependencies before it runs.""" + for mod in spec.requires: + pytest.importorskip(mod) + + +def build_model(spec: BenchSpec, n: int) -> linopy.Model: + """Build ``spec`` at ``n`` — the untimed setup, after the requires-skip.""" + require(spec) + return spec.build(n) + + +@pytest.fixture(autouse=True) +def _benchmem_dims(request: pytest.FixtureRequest, benchmark: object) -> None: + """ + Mirror each case's ``spec``/``phase``/``axis`` into pytest-benchmark + ``extra_info`` as analysis dims, so a ``--benchmark-json`` run plots cleanly + under pytest-benchmem — which reads dims from ``params``/``extra_info`` and + can see neither the (unserialisable) spec param nor the phase, which lives in + the test-function name. The numeric ``n`` is already a clean param. No-op + under CodSpeed, whose fixture carries no ``extra_info``. + """ + callspec = getattr(request.node, "callspec", None) + info = getattr(benchmark, "extra_info", None) + func = getattr(request, "function", None) + if ( + callspec is None + or info is None + or func is None + or "spec" not in callspec.params + ): + return + spec = callspec.params["spec"] + info.update( + spec=spec.name, phase=func.__name__.removeprefix("test_"), axis=spec.axis + ) diff --git a/benchmarks/drivers/__init__.py b/benchmarks/drivers/__init__.py new file mode 100644 index 000000000..4e9099f9f --- /dev/null +++ b/benchmarks/drivers/__init__.py @@ -0,0 +1,8 @@ +""" +Phase drivers — one ``test_.py`` per measured phase. + +Each driver is parametrised over every ``(spec, value)`` the phase runs (via +``benchmarks.conftest.cases``), does untimed setup, then wraps the phase verb +from :mod:`benchmarks.phases` in ``benchmark(...)``. Shared across models and +patterns alike — a pattern is just more rows tagged ``axis="severity"``. +""" diff --git a/benchmarks/drivers/test_build.py b/benchmarks/drivers/test_build.py new file mode 100644 index 000000000..8d6e536a4 --- /dev/null +++ b/benchmarks/drivers/test_build.py @@ -0,0 +1,18 @@ +"""Benchmarks for model construction speed.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from benchmarks.conftest import cases, require +from benchmarks.registry import BUILD + +if TYPE_CHECKING: + from benchmarks.registry import BenchSpec + + +@cases(BUILD) +def test_build(benchmark: Callable[..., object], spec: BenchSpec, n: int) -> None: + require(spec) + benchmark(lambda: spec.build(n)) diff --git a/benchmarks/drivers/test_matrices.py b/benchmarks/drivers/test_matrices.py new file mode 100644 index 000000000..a7e61b05f --- /dev/null +++ b/benchmarks/drivers/test_matrices.py @@ -0,0 +1,19 @@ +"""Benchmarks for matrix generation (model -> sparse matrices).""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from benchmarks.conftest import build_model, cases +from benchmarks.phases import touch_matrices +from benchmarks.registry import MATRICES + +if TYPE_CHECKING: + from benchmarks.registry import BenchSpec + + +@cases(MATRICES) +def test_matrices(benchmark: Callable[..., object], spec: BenchSpec, n: int) -> None: + m = build_model(spec, n) + benchmark(lambda: touch_matrices(m)) diff --git a/benchmarks/drivers/test_netcdf.py b/benchmarks/drivers/test_netcdf.py new file mode 100644 index 000000000..3764e31d2 --- /dev/null +++ b/benchmarks/drivers/test_netcdf.py @@ -0,0 +1,39 @@ +""" +Benchmarks for the netCDF persistence round-trip. + +We track ``to_netcdf`` and ``read_netcdf`` separately because the cost split +matters in practice: distributed workflows tend to do many reads of a single +written artifact. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from benchmarks.conftest import build_model, cases +from benchmarks.phases import read_netcdf, write_netcdf +from benchmarks.registry import FROM_NETCDF, TO_NETCDF + +if TYPE_CHECKING: + from pathlib import Path + + from benchmarks.registry import BenchSpec + + +@cases(TO_NETCDF) +def test_to_netcdf( + benchmark: Callable[..., object], spec: BenchSpec, n: int, tmp_path: Path +) -> None: + m = build_model(spec, n) + benchmark(lambda: write_netcdf(m, tmp_path / "model.nc")) + + +@cases(FROM_NETCDF) +def test_from_netcdf( + benchmark: Callable[..., object], spec: BenchSpec, n: int, tmp_path: Path +) -> None: + m = build_model(spec, n) + path = tmp_path / "model.nc" + write_netcdf(m, path) # setup — untimed + benchmark(lambda: read_netcdf(path)) diff --git a/benchmarks/drivers/test_pipeline.py b/benchmarks/drivers/test_pipeline.py new file mode 100644 index 000000000..1033ad2eb --- /dev/null +++ b/benchmarks/drivers/test_pipeline.py @@ -0,0 +1,38 @@ +""" +End-to-end pipeline benchmark: build → matrices → LP write in one region. + +Opt-in (deselected unless ``--pipeline``): it re-runs the per-phase work and, +unlike the individual phase benchmarks, *includes the model build* — so it +captures the end-to-end cost/peak a real build-then-export session hits, which +can't be recovered by summing the marginal per-phase numbers. Parametrized over +the ``to_lp`` specs (it ends in an LP write). +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from benchmarks.conftest import cases, require +from benchmarks.phases import touch_matrices, write_lp +from benchmarks.registry import TO_LP + +if TYPE_CHECKING: + from pathlib import Path + + from benchmarks.registry import BenchSpec + + +@cases(TO_LP) +def test_pipeline( + benchmark: Callable[..., object], spec: BenchSpec, n: int, tmp_path: Path +) -> None: + require(spec) + path = tmp_path / "model.lp" + + def pipeline() -> None: + m = spec.build(n) + touch_matrices(m) + write_lp(m, path) + + benchmark(pipeline) diff --git a/benchmarks/drivers/test_pypsa_carbon_management.py b/benchmarks/drivers/test_pypsa_carbon_management.py new file mode 100644 index 000000000..209416ba8 --- /dev/null +++ b/benchmarks/drivers/test_pypsa_carbon_management.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import pytest + +import linopy as lp + +# pypsa is an optional benchmark dep. Skip the whole module if it's missing +# so the rest of the suite stays collectable without it. +pypsa = pytest.importorskip("pypsa") + + +@pytest.fixture(scope="module") +def network() -> Any: + try: + return pypsa.examples.carbon_management() + except Exception as exc: # network / example-data drift, not a linopy signal + pytest.skip(f"pypsa example data unavailable: {exc}") + + +def test_create_model_frozen(benchmark: Callable[..., object], network: Any) -> None: + benchmark(network.optimize.create_model, freeze_constraints=True) + + +def test_create_model_mutable(benchmark: Callable[..., object], network: Any) -> None: + benchmark(network.optimize.create_model, freeze_constraints=False) + + +@pytest.fixture(scope="module") +def model_frozen(network: Any) -> Any: + return network.optimize.create_model(freeze_constraints=True) + + +@pytest.fixture(scope="module") +def model_mutable(network: Any) -> Any: + return network.optimize.create_model(freeze_constraints=False) + + +def test_to_highspy_frozen(benchmark: Callable[..., object], model_frozen: Any) -> None: + benchmark(lp.io.to_highspy, model_frozen) + + +def test_to_highspy_mutable( + benchmark: Callable[..., object], model_mutable: Any +) -> None: + benchmark(lp.io.to_highspy, model_mutable) + + +def test_to_highspy_mutable_no_names( + benchmark: Callable[..., object], model_mutable: Any +) -> None: + benchmark(lp.io.to_highspy, model_mutable, set_names=False) + + +def test_to_highspy_frozen_no_names( + benchmark: Callable[..., object], model_frozen: Any +) -> None: + benchmark(lp.io.to_highspy, model_frozen, set_names=False) diff --git a/benchmarks/drivers/test_to_lp.py b/benchmarks/drivers/test_to_lp.py new file mode 100644 index 000000000..2303d7cb0 --- /dev/null +++ b/benchmarks/drivers/test_to_lp.py @@ -0,0 +1,24 @@ +"""Benchmarks for LP file writing speed.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from benchmarks.conftest import build_model, cases +from benchmarks.phases import write_lp +from benchmarks.registry import TO_LP + +if TYPE_CHECKING: + from pathlib import Path + + from benchmarks.registry import BenchSpec + + +@cases(TO_LP) +def test_to_lp( + benchmark: Callable[..., object], spec: BenchSpec, n: int, tmp_path: Path +) -> None: + m = build_model(spec, n) + path = tmp_path / "model.lp" + benchmark(lambda: write_lp(m, path)) diff --git a/benchmarks/drivers/test_to_solver.py b/benchmarks/drivers/test_to_solver.py new file mode 100644 index 000000000..defb14d26 --- /dev/null +++ b/benchmarks/drivers/test_to_solver.py @@ -0,0 +1,50 @@ +""" +Benchmarks for solver handoff (model -> native solver instance). + +Times each ``linopy.io.to_`` wrapper. These wrappers delegate to the +same direct-API build path as the new stateful Solver API +(``Solver.from_name(name, model, io_api="direct")``), so the numbers serve +double duty: regression tracking for the wrappers, *and* for the underlying +``Solver._build_direct`` paths. They've also been available for many releases +— using them keeps the suite runnable on older linopy versions. + +The actual ``Solver.solve()`` runtime (i.e. solver-side algorithm time) is +intentionally not benchmarked. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +import pytest + +from benchmarks.conftest import build_model +from benchmarks.phases import SOLVER_HANDOFFS +from benchmarks.registry import iter_params, spec_param_id +from linopy.solvers import available_solvers + +if TYPE_CHECKING: + from benchmarks.registry import BenchSpec + +# One case per (available solver wrapper) × (spec, value) it applies to. +_PARAMS = [ + (name, wrapper, spec, n) + for name, tag, wrapper in SOLVER_HANDOFFS + for spec, n in iter_params(tag) +] +_IDS = [f"{name}-{spec_param_id(s.name, s.axis, v)}" for name, _w, s, v in _PARAMS] + + +@pytest.mark.parametrize(("name", "wrapper", "spec", "n"), _PARAMS, ids=_IDS) +def test_to_solver( + benchmark: Callable[..., object], + name: str, + wrapper: Callable[..., object], + spec: BenchSpec, + n: int, +) -> None: + if name not in available_solvers: + pytest.skip(f"{name} not installed") + m = build_model(spec, n) + benchmark(lambda: wrapper(m)) diff --git a/benchmarks/models/__init__.py b/benchmarks/models/__init__.py new file mode 100644 index 000000000..66c9a7c76 --- /dev/null +++ b/benchmarks/models/__init__.py @@ -0,0 +1,25 @@ +""" +Model builders for benchmarks. + +Importing this package triggers every submodule's ``register(...)`` call, +populating :data:`benchmarks.registry.REGISTRY`. Each submodule exposes a +``build_(size) -> linopy.Model`` callable and a module-level ``SPEC`` +:class:`~benchmarks.registry.BenchSpec`. The documented access path is +``REGISTRY[""]``; submodule re-exports are intentionally not exposed +here so that adding a new model is one new file plus one import below. +""" + +# Side-effect imports — each module calls ``register(...)`` at import time. +from benchmarks.models import ( # noqa: F401 + basic, + expression_arithmetic, + knapsack, + masked, + milp, + piecewise, + pypsa_scigrid, + qp, + sos, + sparse_network, + storage, +) diff --git a/benchmarks/models/basic.py b/benchmarks/models/basic.py new file mode 100644 index 000000000..554ad05e7 --- /dev/null +++ b/benchmarks/models/basic.py @@ -0,0 +1,28 @@ +"""Basic benchmark model: 2*N^2 variables and constraints (continuous LP).""" + +from __future__ import annotations + +import linopy +from benchmarks.registry import BenchSpec, register + +SIZES = (10, 250) + + +def build_basic(n: int) -> linopy.Model: + """Build a basic N*N model with 2*N^2 vars and 2*N^2 constraints.""" + m = linopy.Model() + x = m.add_variables(coords=[range(n), range(n)], dims=["i", "j"], name="x") + y = m.add_variables(coords=[range(n), range(n)], dims=["i", "j"], name="y") + m.add_constraints(x + y <= 10, name="upper") + m.add_constraints(x - y >= -5, name="lower") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + +SPEC = register( + BenchSpec( + name="basic", + build=build_basic, + sweep=SIZES, + ) +) diff --git a/benchmarks/models/expression_arithmetic.py b/benchmarks/models/expression_arithmetic.py new file mode 100644 index 000000000..0d5af581d --- /dev/null +++ b/benchmarks/models/expression_arithmetic.py @@ -0,0 +1,40 @@ +"""Expression arithmetic benchmark: stress-tests +, *, sum, broadcasting.""" + +from __future__ import annotations + +import numpy as np + +import linopy +from benchmarks.registry import BenchSpec, register + +SIZES = (10, 250) + + +def build_expression_arithmetic(n: int) -> linopy.Model: + """Build a model that exercises expression arithmetic heavily.""" + m = linopy.Model() + + # Variables on different dimensions to trigger broadcasting + x = m.add_variables(coords=[range(n), range(n)], dims=["i", "j"], name="x") + y = m.add_variables(coords=[range(n)], dims=["i"], name="y") + z = m.add_variables(coords=[range(n)], dims=["j"], name="z") + + # Expression arithmetic: broadcasting y (dim i) and z (dim j) against x (dim i,j) + coeffs = np.linspace(-1, 1, n * n).reshape(n, n) + expr1 = x * coeffs + y - z + expr2 = 2 * x - 3 * y + z + combined = expr1 + expr2 + + m.add_constraints(combined <= 100, name="combined") + m.add_constraints(expr1.sum("j") >= -10, name="row_sum") + m.add_objective(combined.sum()) + return m + + +SPEC = register( + BenchSpec( + name="expression_arithmetic", + build=build_expression_arithmetic, + sweep=SIZES, + ) +) diff --git a/benchmarks/models/knapsack.py b/benchmarks/models/knapsack.py new file mode 100644 index 000000000..fe01ad8b7 --- /dev/null +++ b/benchmarks/models/knapsack.py @@ -0,0 +1,34 @@ +"""Knapsack benchmark model: N binary variables, 1 constraint (MILP, binary).""" + +from __future__ import annotations + +import numpy as np + +import linopy +from benchmarks.registry import DEFAULT_PHASES, BenchSpec, register + +SIZES = (100, 10_000) + + +def build_knapsack(n: int) -> linopy.Model: + """Build a knapsack model with N items.""" + rng = np.random.default_rng(42) + weights = rng.integers(1, 100, size=n) + values = rng.integers(1, 100, size=n) + capacity = int(weights.sum() * 0.5) + + m = linopy.Model() + x = m.add_variables(coords=[range(n)], dims=["item"], binary=True, name="x") + m.add_constraints((x * weights).sum() <= capacity, name="capacity") + m.add_objective(-(x * values).sum()) + return m + + +SPEC = register( + BenchSpec( + name="knapsack", + build=build_knapsack, + sweep=SIZES, + phases=DEFAULT_PHASES, # HiGHS handles binary; matrices handles MILP + ) +) diff --git a/benchmarks/models/masked.py b/benchmarks/models/masked.py new file mode 100644 index 000000000..eb9255fbe --- /dev/null +++ b/benchmarks/models/masked.py @@ -0,0 +1,86 @@ +""" +Masked-variables benchmark: transportation with sparse allowed routes. + +A standard transportation LP, but only a sparse subset of (origin, dest) pairs +are valid routes. The ``mask=`` keyword on ``add_variables`` skips the rest, +keeping the variable count sub-quadratic. + +Decision variables: + x[origin, dest] >= 0 continuous, only created for allowed routes + +Constraints: + sum_dest x[o, .] <= supply[o] + sum_orig x[., d] == demand[d] + +Objective: + minimize sum cost[o, d] * x[o, d] + +The mask is dense at small sizes and sparser at large sizes, mimicking +real-world transport networks where each origin only serves a fixed +fan-out regardless of total node count. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import linopy +from benchmarks.registry import ( + DEFAULT_PHASES, + BenchSpec, + register, +) + +SIZES = (10, 100) + + +def build_masked(n: int) -> linopy.Model: + rng = np.random.default_rng(42) + origins = np.arange(n) + dests = np.arange(n) + + # Each origin serves at most ~min(20, n) destinations. + fan_out = min(20, n) + mask_np = np.zeros((n, n), dtype=bool) + for o in range(n): + # Deterministic fan-out so size determines connectivity. + targets = rng.choice(n, size=fan_out, replace=False) + mask_np[o, targets] = True + + mask = xr.DataArray(mask_np, coords=[("origin", origins), ("dest", dests)]) + cost = xr.DataArray( + rng.uniform(1, 10, size=(n, n)), + coords=[("origin", origins), ("dest", dests)], + ) + + # Supply scaled so the problem stays feasible at any size: + # each origin can ship up to ``demand_per_dest * fan_out`` units. + demand_per_dest = 5.0 + supply_per_origin = demand_per_dest * n # plenty of slack + supply = xr.DataArray(np.full(n, supply_per_origin), coords=[("origin", origins)]) + demand = xr.DataArray(np.full(n, demand_per_dest), coords=[("dest", dests)]) + + m = linopy.Model() + x = m.add_variables( + lower=0, + coords=[("origin", origins), ("dest", dests)], + mask=mask, + name="x", + ) + + m.add_constraints(x.sum("dest") <= supply, name="supply", mask=mask.any("dest")) + m.add_constraints(x.sum("origin") == demand, name="demand", mask=mask.any("origin")) + + m.add_objective((cost * x).sum()) + return m + + +SPEC = register( + BenchSpec( + name="masked", + build=build_masked, + sweep=SIZES, + phases=DEFAULT_PHASES, + ) +) diff --git a/benchmarks/models/milp.py b/benchmarks/models/milp.py new file mode 100644 index 000000000..f6058cc8a --- /dev/null +++ b/benchmarks/models/milp.py @@ -0,0 +1,75 @@ +""" +MILP benchmark: capacitated facility location with general integers. + +Decision variables: + y_f in {0,1,...,K} integer "modules" to open at facility f + x_{f,c} >= 0 continuous flow from facility f to customer c + +Constraints: + sum_c x_{f,c} <= cap * y_f (capacity per facility) + sum_f x_{f,c} == d_c (demand at each customer) + +Objective: + minimize sum_{f,c} t_{f,c} * x_{f,c} + sum_f f_f * y_f + +The general-integer ``y`` exercises the matrix accessor's MIP integer-section +path and the LP-writer's general-integer block — neither the binary knapsack +nor the continuous LPs hit those paths. +""" + +from __future__ import annotations + +import numpy as np + +import linopy +from benchmarks.registry import ( + DEFAULT_PHASES, + BenchSpec, + register, +) + +SIZES = (10, 50) + + +def build_milp(n: int) -> linopy.Model: + rng = np.random.default_rng(42) + facilities = np.arange(n) + customers = np.arange(n) + + cap = 100.0 # capacity per module + Y_MAX = 5 # max modules per facility + transport = rng.uniform(1, 20, size=(n, n)) # per-unit shipping cost + fixed = rng.uniform(50, 200, size=n) # cost per facility module + demand = rng.uniform(20, 80, size=n) # demand at each customer + + m = linopy.Model() + y = m.add_variables( + lower=0, + upper=Y_MAX, + coords=[facilities], + dims=["facility"], + integer=True, + name="y", + ) + x = m.add_variables( + lower=0, + coords=[facilities, customers], + dims=["facility", "customer"], + name="x", + ) + + m.add_constraints(x.sum("customer") - cap * y <= 0, name="capacity") + m.add_constraints(x.sum("facility") == demand, name="demand") + + m.add_objective((transport * x).sum() + (fixed * y).sum()) + return m + + +SPEC = register( + BenchSpec( + name="milp", + build=build_milp, + sweep=SIZES, + phases=DEFAULT_PHASES, + ) +) diff --git a/benchmarks/models/piecewise.py b/benchmarks/models/piecewise.py new file mode 100644 index 000000000..895e854a1 --- /dev/null +++ b/benchmarks/models/piecewise.py @@ -0,0 +1,89 @@ +""" +Piecewise-linear benchmark: generation with piecewise fuel-cost curves. + +Each generator has a piecewise fuel cost curve pinned via +``add_piecewise_formulation``. The default ``method="auto"`` picks an +SOS2 or incremental expansion, generating auxiliary variables and +constraints — that overhead is what we want to measure. + +Decision variables: + power[gen] in [0, 100] (continuous) + fuel[gen] in [0, inf) (continuous, pinned to piecewise curve) + +Constraints: + sum_gen power[gen] >= demand + piecewise: fuel[gen] = f(power[gen]) for each gen + +Objective: + minimize sum_gen fuel[gen] +""" + +from __future__ import annotations + +import warnings + +import linopy +from benchmarks.registry import ( + DEFAULT_PHASES, + BenchSpec, + register, +) + +SIZES = (10, 1_000) + +_API_AVAILABLE = hasattr(linopy.Model, "add_piecewise_formulation") and hasattr( + linopy, "EvolvingAPIWarning" +) + + +def build_piecewise(n_gens: int) -> linopy.Model: + # Shared breakpoints, broadcast across generators. + x_pts = [0.0, 30.0, 60.0, 100.0] + y_pts = [0.0, 36.0, 84.0, 170.0] # convex-ish fuel curve + + m = linopy.Model() + power = m.add_variables( + lower=0, + upper=100, + coords=[range(n_gens)], + dims=["gen"], + name="power", + ) + fuel = m.add_variables( + lower=0, + coords=[range(n_gens)], + dims=["gen"], + name="fuel", + ) + + demand = 0.5 * n_gens * x_pts[-1] + m.add_constraints(power.sum() >= demand, name="demand") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=linopy.EvolvingAPIWarning) + m.add_piecewise_formulation( + (power, x_pts), + (fuel, y_pts), + ) + + m.add_objective(fuel.sum()) + return m + + +# ``add_piecewise_formulation`` is a recent (still-evolving) API. Skip +# registration silently on older linopy so the rest of the suite stays usable. +SPEC: BenchSpec | None +if _API_AVAILABLE: + SPEC = register( + BenchSpec( + name="piecewise", + build=build_piecewise, + sweep=SIZES, + # Monotonic breakpoints + ``method="auto"`` → incremental + # reformulation (pure MILP with binaries), which every supported + # solver handles. + phases=DEFAULT_PHASES, + ) + ) +else: + SPEC = None diff --git a/benchmarks/models/pypsa_scigrid.py b/benchmarks/models/pypsa_scigrid.py new file mode 100644 index 000000000..bb6e86535 --- /dev/null +++ b/benchmarks/models/pypsa_scigrid.py @@ -0,0 +1,36 @@ +"""PyPSA SciGrid-DE benchmark model (requires pypsa).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from benchmarks.registry import BenchSpec, register + +if TYPE_CHECKING: + import linopy + +SIZES = (10, 50) # small networks — PyPSA import already dominates the cost + + +def build_pypsa_scigrid(snapshots: int = 100) -> linopy.Model: + """Build PyPSA SciGrid model. Requires pypsa to be installed.""" + import pypsa + import pytest + + try: + n = pypsa.examples.scigrid_de() + except Exception as exc: # network / example-data drift, not a linopy signal + pytest.skip(f"pypsa example data unavailable: {exc}") + n.set_snapshots(n.snapshots[:snapshots]) + n.optimize.create_model() # the linopy build under benchmark — unguarded + return n.model + + +SPEC = register( + BenchSpec( + name="pypsa_scigrid", + build=build_pypsa_scigrid, + sweep=SIZES, + requires=("pypsa",), + ) +) diff --git a/benchmarks/models/qp.py b/benchmarks/models/qp.py new file mode 100644 index 000000000..50e39e7b9 --- /dev/null +++ b/benchmarks/models/qp.py @@ -0,0 +1,61 @@ +""" +QP benchmark: continuous quadratic objective on a portfolio-style model. + +Decision variables: + x_i >= 0 (weight on asset i, continuous) + +Constraints: + sum_i x_i == 1 + x_i <= 0.3 (no asset > 30% of portfolio) + +Objective: + minimize sum_i q_i * x_i^2 - sum_i r_i * x_i + +A pure diagonal quadratic — enough to exercise the QP build / write / matrix +paths without paying for cross-terms. Cross-term coupling needs single-term +factors on both sides (see ``LinearExpression._multiply_by_linear_expression``), +which is awkward to set up cleanly via the public API. +""" + +from __future__ import annotations + +import numpy as np + +import linopy +from benchmarks.registry import ( + DEFAULT_PHASES, + BenchSpec, + register, +) + +SIZES = (10, 1_000) + + +def build_qp(n_assets: int) -> linopy.Model: + rng = np.random.default_rng(42) + q = rng.uniform(0.5, 2.0, size=n_assets) + r = rng.uniform(0.05, 0.15, size=n_assets) + + m = linopy.Model() + x = m.add_variables( + lower=0, + upper=0.3, + coords=[range(n_assets)], + dims=["asset"], + name="x", + ) + + m.add_constraints(x.sum() == 1, name="budget") + + m.add_objective((q * x**2).sum() - (r * x).sum()) + return m + + +SPEC = register( + BenchSpec( + name="qp", + build=build_qp, + sweep=SIZES, + phases=DEFAULT_PHASES, + ) +) diff --git a/benchmarks/models/sos.py b/benchmarks/models/sos.py new file mode 100644 index 000000000..3c1e2db86 --- /dev/null +++ b/benchmarks/models/sos.py @@ -0,0 +1,96 @@ +""" +SOS1 benchmark: multi-mode generation with at-most-one-mode-per-generator. + +Each generator has ``n_modes`` operating modes (different cap/cost tradeoff). +SOS1 over the ``mode`` dimension enforces that each generator picks at most +one mode. + +Decision variables: + y[gen, mode] >= 0 continuous output per (generator, mode) + +Constraints: + y[gen, mode] <= cap[mode] + sum_{gen,mode} y >= demand_total + SOS1 over "mode" for each gen + +This benchmark exercises ``Model.add_sos_constraints`` (commits be6d3a3 / +8aa8d0c) and the LP-writer's SOS section. In linopy, native SOS support is +declared by Gurobi / Cplex / Xpress only (see ``SolverFeature.SOS_CONSTRAINTS``). +HiGHS and Mosek would need ``apply_sos_reformulation()`` first. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import linopy +from benchmarks.registry import ( + BUILD, + FROM_NETCDF, + MATRICES, + TO_GUROBIPY, + TO_LP, + TO_NETCDF, + TO_XPRESS, + BenchSpec, + register, +) + +SIZES = (10, 1_000) + +_N_MODES = 5 +_API_AVAILABLE = hasattr(linopy.Model, "add_sos_constraints") + + +def build_sos(n_gens: int) -> linopy.Model: + modes = np.arange(_N_MODES) + cap = xr.DataArray(np.linspace(20.0, 100.0, _N_MODES), coords=[("mode", modes)]) + cost = xr.DataArray(np.linspace(1.0, 8.0, _N_MODES), coords=[("mode", modes)]) + + m = linopy.Model() + y = m.add_variables( + lower=0, + upper=float(cap.max()), + coords=[range(n_gens), modes], + dims=["gen", "mode"], + name="y", + ) + + m.add_constraints(y <= cap, name="mode_cap") + demand_total = 0.4 * n_gens * float(cap.max()) + m.add_constraints(y.sum() >= demand_total, name="demand") + + m.add_sos_constraints(y, sos_type=1, sos_dim="mode") + + m.add_objective((cost * y).sum()) + return m + + +# ``add_sos_constraints`` is a recent API. On older linopy we silently skip +# registering this model — the rest of the suite stays usable. +SPEC: BenchSpec | None +if _API_AVAILABLE: + SPEC = register( + BenchSpec( + name="sos", + build=build_sos, + sweep=SIZES, + # HiGHS / Mosek lack native SOS in linopy — would need + # ``reformulate_sos=True``, which mutates the model and defeats + # the benchmark. Only solvers with native SOS appear here. + phases=frozenset( + { + BUILD, + MATRICES, + TO_LP, + TO_NETCDF, + FROM_NETCDF, + TO_GUROBIPY, + TO_XPRESS, + } + ), + ) + ) +else: + SPEC = None diff --git a/benchmarks/models/sparse_network.py b/benchmarks/models/sparse_network.py new file mode 100644 index 000000000..13d6c3ad9 --- /dev/null +++ b/benchmarks/models/sparse_network.py @@ -0,0 +1,60 @@ +"""Sparse network benchmark: variables on mismatched coordinate subsets.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +import linopy +from benchmarks.registry import BenchSpec, register + +SIZES = (10, 250) + + +def build_sparse_network(n_buses: int) -> linopy.Model: + """Build a ring network model with mismatched bus/line coordinate subsets.""" + rng = np.random.default_rng(42) + n_lines = n_buses # ring topology + n_time = min(n_buses, 24) + + buses = pd.RangeIndex(n_buses, name="bus") + lines = pd.RangeIndex(n_lines, name="line") + time = pd.RangeIndex(n_time, name="time") + + # Ring topology: line i connects bus i -> bus (i+1) % n + bus_from = np.arange(n_lines) + bus_to = (bus_from + 1) % n_buses + + m = linopy.Model() + + # Bus-level variables (bus × time) + gen = m.add_variables(lower=0, coords=[buses, time], name="gen") + + # Line-level variables (line × time) + flow = m.add_variables(lower=-100, upper=100, coords=[lines, time], name="flow") + + # Incidence matrix (bus × line): +1 for incoming, -1 for outgoing + incidence = np.zeros((n_buses, n_lines)) + incidence[bus_to, np.arange(n_lines)] = 1 # incoming + incidence[bus_from, np.arange(n_lines)] = -1 # outgoing + incidence_da = xr.DataArray(incidence, coords=[buses, lines]) + + # Vectorized flow balance: gen - demand + incidence @ flow == 0 + demand = xr.DataArray( + rng.uniform(10, 100, size=(n_buses, n_time)), coords=[buses, time] + ) + net_flow = (flow * incidence_da).sum("line") + m.add_constraints(gen + net_flow == demand, name="balance") + + m.add_objective(gen.sum()) + return m + + +SPEC = register( + BenchSpec( + name="sparse_network", + build=build_sparse_network, + sweep=SIZES, + ) +) diff --git a/benchmarks/models/storage.py b/benchmarks/models/storage.py new file mode 100644 index 000000000..5e8417280 --- /dev/null +++ b/benchmarks/models/storage.py @@ -0,0 +1,53 @@ +""" +Storage state-of-charge model — intertemporal coupling via ``.shift()``. + +A fleet of storage units, each with a bidiagonal SoC recursion +``soc[t] - decay*soc[t-1] - eff*charge[t] + discharge[t]/eff == 0`` built with +``soc.shift(time=1)`` (``t=0`` falls off as the boundary). This is the one op +family no other model exercises — the ``.shift()``/``.isel()`` intertemporal +coupling that PyPSA's SoC and flixopt's ``charge_state.isel`` recursion lean on. + +It is a *model*, not a pattern: each balance row has a fixed ~4 terms regardless +of horizon or unit count, so it scales with ``size`` (units × timesteps) and has +no benign→worst data-shape dial. ``size`` is the number of storage units. +""" + +from __future__ import annotations + +import pandas as pd + +import linopy +from benchmarks.registry import BenchSpec, register + +SIZES = (10, 250) +N_TIME = 168 +DECAY = 0.99 +ETA = 0.95 + + +def build_storage(n_storage: int) -> linopy.Model: + storages = pd.RangeIndex(n_storage, name="storage") + time = pd.RangeIndex(N_TIME, name="time") + + m = linopy.Model() + soc = m.add_variables(lower=0, upper=100, coords=[storages, time], name="soc") + charge = m.add_variables(lower=0, upper=50, coords=[storages, time], name="charge") + discharge = m.add_variables( + lower=0, upper=50, coords=[storages, time], name="discharge" + ) + + prev = soc.shift(time=1) # soc[t-1]; t=0 shifted out (initial-SoC boundary) + m.add_constraints( + soc - DECAY * prev - ETA * charge + discharge / ETA == 0, name="soc_balance" + ) + m.add_objective((charge + discharge).sum()) + return m + + +SPEC = register( + BenchSpec( + name="storage", + build=build_storage, + sweep=SIZES, + ) +) diff --git a/benchmarks/patterns/__init__.py b/benchmarks/patterns/__init__.py new file mode 100644 index 000000000..090976741 --- /dev/null +++ b/benchmarks/patterns/__init__.py @@ -0,0 +1,22 @@ +""" +Benchmark *patterns* — realistic modelling idioms swept over a severity dial. + +A pattern is a fragment of real modelling code (a balance constraint, a KVL +contraction), not a whole model and not an isolated method call. Each is +measured the same way a model is — time and peak memory, through the shared +phases — but parametrised by ``severity`` (0–100, how pathological the data +shape is) instead of ``size``. See :class:`benchmarks.registry.BenchSpec`. + +Importing this package registers every idiom into +:data:`benchmarks.registry.PATTERNS` (mirrors :mod:`benchmarks.models`); adding +a pattern is one new file plus one import below. +""" + +# Side-effect imports — each module calls ``register_pattern(...)`` at import. +from benchmarks.patterns import ( # noqa: F401 + cumsum, + kvl_cycles, + merge_balance, + nodal_balance, + rolling, +) diff --git a/benchmarks/patterns/cumsum.py b/benchmarks/patterns/cumsum.py new file mode 100644 index 000000000..212e96e7c --- /dev/null +++ b/benchmarks/patterns/cumsum.py @@ -0,0 +1,44 @@ +""" +Cumulative-sum fold — ``.cumsum(dim)`` stacks a growing window into ``_term``. + +A running total over time — cumulative energy, a rolling budget: +``(1 * x).cumsum("time")``. linopy currently routes ``cumsum`` through +``rolling(window=full_dim)`` (``expressions.py``), so its ``_term`` grows +triangularly to the dim size. It is benchmarked as its own op — not folded into +``rolling`` — because it is a distinct public op and a natural de-densification +target (a prefix sum need not materialise the triangle), so this is the +instrument that would show such a kernel change land. ``severity`` dials the +size of the cumulated dimension. +""" + +from __future__ import annotations + +import pandas as pd + +import linopy +from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern + +N_ROW = 64 # broadcast/volume dim — the triangular fold is on t, not row +DIM_MAX = 200 + + +def build_cumsum(severity: int) -> linopy.Model: + rows = pd.RangeIndex(N_ROW, name="row") + n = max(2, round(severity / 100 * DIM_MAX)) + + m = linopy.Model() + x = m.add_variables(coords=[rows, pd.RangeIndex(n, name="t")], name="x") + running = (1 * x).cumsum("t") # (row, t); _term grows triangularly to n + m.add_constraints(running == 0, name="cumulative") + m.add_objective((1 * x).sum()) + return m + + +SPEC = register_pattern( + BenchSpec( + name="cumsum", + build=build_cumsum, + sweep=SEVERITIES, + axis="severity", + ) +) diff --git a/benchmarks/patterns/kvl_cycles.py b/benchmarks/patterns/kvl_cycles.py new file mode 100644 index 000000000..5657eedd8 --- /dev/null +++ b/benchmarks/patterns/kvl_cycles.py @@ -0,0 +1,73 @@ +""" +KVL-cycles pattern — sparse ``@`` densifies the result to a full ``_term`` (#748). + +The idiom: contract a per-branch flow against a (branch × cycle) cycle matrix — +Kirchhoff's voltage law, ``flow @ C``. ``__matmul__`` is ``(flow * C).sum(...)``, +which stacks *every* branch into ``_term`` regardless of whether ``C`` is zero +there. ``severity`` dials ``C``'s sparsity: at 0 it is dense (every branch in +every cycle — nothing to gain), at 100 only ~3 branches per cycle carry a +nonzero (the real grid shape), yet the current kernel still produces +``_term == n_branch``. So the *cost is flat* across severity on today's kernel +— the win from a sparse-aware ``@`` is what grows with it. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +import linopy +from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern + +N_BRANCH = 300 +N_CYCLE = 100 +N_TIME = 168 # snapshot horizon — sets the always-paid flat level (the +# densification width is branch; severity dials C's sparsity, which today's +# kernel ignores, so memory stays flat across severity) +MIN_PER_CYCLE = 3 + + +def _cycle_matrix(severity: int, branches: pd.Index, cycles: pd.Index) -> xr.DataArray: + """ + Branch×cycle incidence whose density falls as ``severity`` rises. + + - ``severity == 0`` → dense: every branch participates in every cycle. + - ``severity == 100`` → ~``MIN_PER_CYCLE`` branches per cycle (real KVL). + + Entries are ±1. The number of nonzeros per cycle interpolates linearly + between ``N_BRANCH`` (dense) and ``MIN_PER_CYCLE`` (sparse). + """ + rng = np.random.default_rng(0) + n_branch = len(branches) + per_cycle = round(n_branch - severity / 100 * (n_branch - MIN_PER_CYCLE)) + per_cycle = max(MIN_PER_CYCLE, per_cycle) + c_mat = np.zeros((n_branch, len(cycles))) + for col in range(len(cycles)): + idx = rng.choice(n_branch, size=per_cycle, replace=False) + c_mat[idx, col] = rng.choice([-1.0, 1.0], size=per_cycle) + return xr.DataArray(c_mat, coords=[branches, cycles]) + + +def build_kvl_cycles(severity: int) -> linopy.Model: + branches = pd.RangeIndex(N_BRANCH, name="branch") + cycles = pd.RangeIndex(N_CYCLE, name="cycle") + time = pd.RangeIndex(N_TIME, name="time") + + m = linopy.Model() + flow = m.add_variables(lower=-100, upper=100, coords=[time, branches], name="flow") + cycle_matrix = _cycle_matrix(severity, branches, cycles) + kvl = (flow * cycle_matrix).sum("branch") + m.add_constraints(kvl == 0.0, name="kvl") + m.add_objective(flow.sum()) + return m + + +SPEC = register_pattern( + BenchSpec( + name="kvl_cycles", + build=build_kvl_cycles, + sweep=SEVERITIES, + axis="severity", + ) +) diff --git a/benchmarks/patterns/merge_balance.py b/benchmarks/patterns/merge_balance.py new file mode 100644 index 000000000..84bb1d918 --- /dev/null +++ b/benchmarks/patterns/merge_balance.py @@ -0,0 +1,57 @@ +""" +Ragged merge — concat of mixed-width blocks pads all to the global max (#749). + +The documented build peak: a balance assembled by merging sub-expressions of +*different* ``_term`` widths along a shared dim. PyPSA's nodal balance does +``merge(gen + storage + lines + links, join="outer")`` (the single largest +allocation in a SciGRID build); flixopt's bus balance is the sibling +``sum([flow_rate for flow in flows])``. Merging along a non-``_term`` dim makes +linopy align the ``_term`` axes by padding every block to the widest one — so +one fat block leaves the narrow blocks mostly fill. ``severity`` dials the +widest block's term count. +""" + +from __future__ import annotations + +import pandas as pd + +import linopy +from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern + +N_BLOCKS = 30 +N_ROW = 128 # broadcast/volume dim — the ragged padding is on _term, not row +NARROW = 3 +WIDE = 200 + + +def _block( + m: linopy.Model, rows: pd.Index, name: str, width: int +) -> linopy.LinearExpression: + """A ``(row,)`` expression with ``width`` terms (a ``(row, k)`` var folded over ``k``).""" + k = pd.RangeIndex(width, name=f"k_{name}") + x = m.add_variables(coords=[rows, k], name=name) + return (1 * x).sum(f"k_{name}") + + +def build_merge_balance(severity: int) -> linopy.Model: + rows = pd.RangeIndex(N_ROW, name="row") + widest = max(NARROW, round(NARROW + severity / 100 * (WIDE - NARROW))) + + m = linopy.Model() + blocks = [_block(m, rows, f"narrow{i}", NARROW) for i in range(N_BLOCKS - 1)] + blocks.append(_block(m, rows, "wide", widest)) + + lhs = linopy.merge(blocks, dim="block", join="outer") + m.add_constraints(lhs == 0, name="balance") + m.add_objective(blocks[0]) + return m + + +SPEC = register_pattern( + BenchSpec( + name="merge_balance", + build=build_merge_balance, + sweep=SEVERITIES, + axis="severity", + ) +) diff --git a/benchmarks/patterns/nodal_balance.py b/benchmarks/patterns/nodal_balance.py new file mode 100644 index 000000000..458df39a4 --- /dev/null +++ b/benchmarks/patterns/nodal_balance.py @@ -0,0 +1,72 @@ +""" +Nodal-balance pattern — grouped-sum padding under bus-connectivity skew (#745). + +The idiom: sum each bus's generators (``groupby(bus).sum()``) and balance the +result against demand. ``LinearExpression.groupby(...).sum()`` pads every group +to the largest group's term count, so as generators concentrate on one hub the +result's ``_term`` axis blows up — most of it fill. ``severity`` dials that +skew; the build's peak memory is expected to climb steeply with it on the +current (dense) kernel. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +import linopy +from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern + +N_GEN = 2000 +N_BUS = 50 +N_TIME = 8 # broadcast/volume dim — the groupby pathology is on gen, not time + + +def _bus_of_gen(severity: int) -> np.ndarray: + """ + Assign each generator to a bus, skewed toward one hub by ``severity``. + + - ``severity == 0`` → round-robin: every bus holds ~``N_GEN / N_BUS``. + - ``severity == 100`` → bus 0 holds almost all generators. + + The first ``N_BUS`` generators anchor one bus each, so no bus is ever empty + — the constraint *shape* (``N_BUS`` rows) is fixed across the sweep and only + the per-group term count (the padding) varies. + """ + rng = np.random.default_rng(0) + bus = np.arange(N_GEN) % N_BUS # uniform baseline + anchor = np.zeros(N_GEN, dtype=bool) + anchor[:N_BUS] = True # pin one generator per bus + move = (~anchor) & (rng.random(N_GEN) < severity / 100) + bus[move] = 0 # reassign a severity-fraction of the rest onto the hub + return bus + + +def build_nodal_balance(severity: int) -> linopy.Model: + gens = pd.RangeIndex(N_GEN, name="gen") + time = pd.RangeIndex(N_TIME, name="time") + buses = pd.RangeIndex(N_BUS, name="bus") + rng = np.random.default_rng(1) + + m = linopy.Model() + gen = m.add_variables(lower=0, coords=[gens, time], name="gen") + + bus_of_gen = pd.Series(_bus_of_gen(severity), index=gens, name="bus") + supply = (1 * gen).groupby(bus_of_gen).sum() + demand = xr.DataArray( + rng.uniform(10.0, 100.0, size=(N_BUS, N_TIME)), coords=[buses, time] + ) + m.add_constraints(supply == demand, name="balance") + m.add_objective(gen.sum()) + return m + + +SPEC = register_pattern( + BenchSpec( + name="nodal_balance", + build=build_nodal_balance, + sweep=SEVERITIES, + axis="severity", + ) +) diff --git a/benchmarks/patterns/rolling.py b/benchmarks/patterns/rolling.py new file mode 100644 index 000000000..30065179d --- /dev/null +++ b/benchmarks/patterns/rolling.py @@ -0,0 +1,46 @@ +""" +Rolling-window coupling — ``rolling(K).sum()`` stacks K terms into ``_term``. + +The *windowed* form of intertemporal coupling (unlike the 1-step storage SoC, +this one has a real density dial): minimum up/down time and windowed energy / +ramp limits sum a variable over a sliding window of K timesteps +(PyPSA ``status.rolling(K).sum()`` for min-up-time, ``constraints.py:450``). +``rolling(K).sum()`` builds a result with **K terms per row** — so the window +width is a clean severity dial. ``severity`` dials K from a single step to the +full horizon. +""" + +from __future__ import annotations + +import pandas as pd + +import linopy +from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern + +N_UNIT = 8 # broadcast dim — the window densification is on time, not unit +N_TIME = 1000 +MIN_WINDOW = 1 + + +def build_rolling(severity: int) -> linopy.Model: + units = pd.RangeIndex(N_UNIT, name="unit") + time = pd.RangeIndex(N_TIME, name="time") + window = max(MIN_WINDOW, round(MIN_WINDOW + severity / 100 * (N_TIME - MIN_WINDOW))) + + m = linopy.Model() + status = m.add_variables(lower=0, upper=1, coords=[units, time], name="status") + # min-up-time style: every K-step window carries at most K active steps. + windowed = status.rolling(time=window).sum() + m.add_constraints(windowed <= window, name="window_limit") + m.add_objective(status.sum()) + return m + + +SPEC = register_pattern( + BenchSpec( + name="rolling", + build=build_rolling, + sweep=SEVERITIES, + axis="severity", + ) +) diff --git a/benchmarks/phases.py b/benchmarks/phases.py new file mode 100644 index 000000000..983fb9c19 --- /dev/null +++ b/benchmarks/phases.py @@ -0,0 +1,74 @@ +""" +The measured operations — what each benchmark phase *does to a model*. + +The ``test_.py`` drivers wrap these verbs in ``benchmark(...)``; setup +(building the model, scratch files) stays in the driver, only the verb itself +lives here. +""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable +from pathlib import Path + +import linopy +import linopy.io as lio +from benchmarks.registry import TO_GUROBIPY, TO_HIGHSPY, TO_MOSEK, TO_XPRESS +from linopy import read_netcdf + +# linopy <0.4.1's ``to_file`` doesn't accept ``progress``. Checked once at import +# so the suite stays runnable against older linopy (e.g. cross-version sweeps), +# and the benchmark loop stays branchless. +_TO_FILE_HAS_PROGRESS = "progress" in inspect.signature(linopy.Model.to_file).parameters + +# Re-export so a driver can ``from benchmarks.phases import read_netcdf``. +__all__ = [ + "SOLVER_HANDOFFS", + "read_netcdf", + "touch_matrices", + "write_lp", + "write_netcdf", +] + + +def touch_matrices(m: linopy.Model) -> None: + """Force every matrix block to materialise — the thing we measure.""" + mats = m.matrices + for attr in ("A", "b", "c", "lb", "ub", "sense", "vlabels", "clabels"): + getattr(mats, attr) + if m.is_quadratic: + mats.Q + + +def write_lp(m: linopy.Model, path: Path) -> None: + """ + Write the model as an LP file. + + Where supported, ``progress=False`` is pinned so the progress bar's overhead + doesn't leak into the measurement; linopy <0.4.1 doesn't accept the kwarg. + """ + if _TO_FILE_HAS_PROGRESS: + m.to_file(path, progress=False) + else: + m.to_file(path) + + +def write_netcdf(m: linopy.Model, path: Path) -> None: + m.to_netcdf(path) + + +# (solver_name, registry phase tag, wrapper) — consumed by test_to_solver.py. +# Each wrapper is fetched via ``getattr`` so the tuple silently drops any wrapper +# missing from the installed linopy (e.g. ``to_xpress`` is absent before linopy +# 0.7.1) — keeping the suite runnable on older releases for cross-version sweeps. +SOLVER_HANDOFFS: tuple[tuple[str, str, Callable[[linopy.Model], object]], ...] = tuple( + (name, tag, wrapper) + for name, tag, wrapper in ( + ("highs", TO_HIGHSPY, getattr(lio, "to_highspy", None)), + ("gurobi", TO_GUROBIPY, getattr(lio, "to_gurobipy", None)), + ("mosek", TO_MOSEK, getattr(lio, "to_mosek", None)), + ("xpress", TO_XPRESS, getattr(lio, "to_xpress", None)), + ) + if wrapper is not None +) diff --git a/benchmarks/registry.py b/benchmarks/registry.py new file mode 100644 index 000000000..5f3f98ef1 --- /dev/null +++ b/benchmarks/registry.py @@ -0,0 +1,137 @@ +""" +Registry of benchmark models and patterns. + +A :class:`BenchSpec` declares how to build a model and which values (sizes for a +model, ``axis="n"``; severities for a pattern, ``axis="severity"``) and phases +it runs; ``register`` / ``register_pattern`` add it to :data:`REGISTRY` / +:data:`PATTERNS`:: + + from benchmarks import REGISTRY + model = REGISTRY["basic"].build(100) +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from dataclasses import dataclass + +import linopy + +# --- Phase tags ------------------------------------------------------------- + +BUILD = "build" +MATRICES = "matrices" +TO_LP = "to_lp" +TO_NETCDF = "to_netcdf" +FROM_NETCDF = "from_netcdf" +TO_HIGHSPY = "to_highspy" +TO_GUROBIPY = "to_gurobipy" +TO_MOSEK = "to_mosek" +TO_XPRESS = "to_xpress" + +ALL_PHASES = frozenset( + { + BUILD, + MATRICES, + TO_LP, + TO_NETCDF, + FROM_NETCDF, + TO_HIGHSPY, + TO_GUROBIPY, + TO_MOSEK, + TO_XPRESS, + } +) + +# The default phase set; a spec overrides with a narrower one when the default +# solvers can't ingest it natively (e.g. native SOS for HiGHS). +DEFAULT_PHASES = ALL_PHASES + +# The severity sweep every pattern runs (axis "severity"). +SEVERITIES: tuple[int, ...] = (0, 50, 100) + + +@dataclass(frozen=True, repr=False) +class BenchSpec: + """ + One benchmark spec. A model is swept over ``sweep`` sizes (``axis="n"``); a + pattern over a 0–100 severity dial (``axis="severity"``). Both build a + :class:`linopy.Model` from one integer and run the same ``phases`` — the + model-vs-pattern distinction lives in :func:`register` vs + :func:`register_pattern` (and the ``models/`` vs ``patterns/`` dirs). + """ + + name: str + build: Callable[[int], linopy.Model] + sweep: tuple[int, ...] + axis: str = "n" + phases: frozenset[str] = DEFAULT_PHASES + requires: tuple[str, ...] = () + + def applies_to(self, phase: str) -> bool: + return phase in self.phases + + def __repr__(self) -> str: + return f"BenchSpec({self.name!r}, axis={self.axis!r}, sweep={self.sweep})" + + +REGISTRY: dict[str, BenchSpec] = {} +PATTERNS: dict[str, BenchSpec] = {} + + +def _validate(spec: BenchSpec, registry: dict[str, BenchSpec], kind: str) -> None: + if spec.name in registry: + raise ValueError(f"{kind} {spec.name!r} already registered") + unknown = spec.phases - ALL_PHASES + if unknown: + raise ValueError(f"{kind} {spec.name!r}: unknown phases {sorted(unknown)}") + + +def register(spec: BenchSpec) -> BenchSpec: + """Add a model ``spec`` to :data:`REGISTRY`. Returns it for chaining.""" + _validate(spec, REGISTRY, "model") + REGISTRY[spec.name] = spec + return spec + + +def register_pattern(spec: BenchSpec) -> BenchSpec: + """Add a pattern ``spec`` (``axis="severity"``) to :data:`PATTERNS`.""" + _validate(spec, PATTERNS, "pattern") + if spec.axis != "severity" or not all(0 <= s <= 100 for s in spec.sweep): + raise ValueError( + f"pattern {spec.name!r}: needs axis='severity' and sweep in [0, 100], " + f"got axis={spec.axis!r} sweep={spec.sweep}" + ) + PATTERNS[spec.name] = spec + return spec + + +def all_specs() -> list[BenchSpec]: + """Every spec in the suite — models then patterns.""" + return [*REGISTRY.values(), *PATTERNS.values()] + + +def iter_params( + phase: str, specs: Iterable[BenchSpec] | None = None +) -> list[tuple[BenchSpec, int]]: + """ + Flatten ``(spec, value)`` pairs for one phase — the pytest parametrize + source. ``specs`` defaults to every model and pattern in the suite. + """ + specs = all_specs() if specs is None else specs + return [ + (spec, value) + for spec in specs + if spec.applies_to(phase) + for value in spec.sweep + ] + + +def spec_param_id(name: str, axis: str, value: object) -> str: + """ + The ``-=`` fragment that fills a test id's ``[...]``. + + Single source of truth for the parametrize-id shape — the pytest param + ids and the solver-handoff ids all build on it. + """ + return f"{name}-{axis}={value}" diff --git a/codecov.yml b/codecov.yml index 69cb76019..74a549c12 100644 --- a/codecov.yml +++ b/codecov.yml @@ -1 +1,4 @@ comment: false + +ignore: + - "benchmarks/**" diff --git a/doc/api.rst b/doc/api.rst index 6011aa810..6fb3434f2 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -4,159 +4,581 @@ API reference ############# -This page provides an auto-generated summary of linopy's API. +Reference for linopy's public API. Most workflows start at +:class:`~linopy.model.Model` — :class:`~linopy.variables.Variable`, +:class:`~linopy.constraints.Constraint`, and +:class:`~linopy.objective.Objective` are all built through +:meth:`Model.add_variables `, +:meth:`Model.add_constraints `, +:meth:`Model.add_objective `, +and accessed through the matching +:attr:`Model.variables `, +:attr:`Model.constraints `, and +:attr:`Model.objective ` accessors. +The supporting classes below cover those types in detail. + +.. contents:: + :local: + :depth: 2 + + +Model +===== + +Central container for an optimization problem. Most of linopy's +surface lives here. +.. autosummary:: + :toctree: generated/ + model.Model -Creating a model -================ +Building a model +---------------- + +.. autosummary:: + :toctree: generated/ + + model.Model.add_variables + model.Model.add_constraints + model.Model.add_objective + model.Model.add_sos_constraints + model.Model.add_piecewise_formulation + +Inspecting a model +------------------ + +.. autosummary:: + :toctree: generated/ + + model.Model.variables + model.Model.constraints + model.Model.objective + model.Model.sense + model.Model.type + model.Model.is_linear + model.Model.is_quadratic + +Modifying a model +----------------- + +.. autosummary:: + :toctree: generated/ + + model.Model.remove_variables + model.Model.remove_constraints + model.Model.remove_objective + model.Model.remove_sos_constraints + model.Model.copy + model.Model.reformulate_sos_constraints + +Solving +------- + +.. autosummary:: + :toctree: generated/ + + model.Model.solve + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + model.Model.solution + model.Model.dual + model.Model.status + model.Model.termination_condition + +Diagnostics +----------- .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + model.Model.compute_infeasibilities + model.Model.format_infeasibilities - model.Model - model.Model.add_variables - model.Model.add_constraints - model.Model.add_objective - model.Model.linexpr - model.Model.remove_constraints +IO +-- + +.. autosummary:: + :toctree: generated/ + model.Model.to_file + model.Model.to_netcdf + io.read_netcdf -Classes under the hook -====================== Variable --------- +======== + +Subclass of ``xarray.DataArray`` carrying labels for a multi-dimensional +decision variable. + +.. autosummary:: + :toctree: generated/ + + variables.Variable + +Attributes +---------- + +.. autosummary:: + :toctree: generated/ -``Variable`` is a subclass of ``xarray.DataArray`` and contains all labels referring to a multi-dimensional variable. + variables.Variable.lower + variables.Variable.upper + variables.Variable.type + variables.Variable.solution + +Modification +------------ + +``Variable.update`` is the canonical mutation API. The legacy ``lower`` / +``upper`` setters still forward to ``update`` but emit a +``DeprecationWarning`` and will be removed in a future release. + +.. autosummary:: + :toctree: generated/ + + variables.Variable.update + variables.Variable.fix + variables.Variable.unfix + variables.Variable.relax + variables.Variable.unrelax + +Operations +---------- + +.. autosummary:: + :toctree: generated/ + + variables.Variable.sum + variables.Variable.where + +Conversion +---------- .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + variables.Variable.to_linexpr + variables.Variable.to_polars - variables.Variable - variables.Variable.lower - variables.Variable.upper - variables.Variable.sum - variables.Variable.where - variables.Variable.sanitize - variables.Variables - variables.ScalarVariable Variables +========= + +Container for the collection of variables on a model. Accessed via +``model.variables``. + +.. autosummary:: + :toctree: generated/ + + variables.Variables + +Attributes +---------- + +.. autosummary:: + :toctree: generated/ + + variables.Variables.lower + variables.Variables.upper + variables.Variables.solution + +Modification +------------ + +.. autosummary:: + :toctree: generated/ + + variables.Variables.fix + variables.Variables.unfix + variables.Variables.relax + variables.Variables.unrelax + +Inventory +--------- + +.. autosummary:: + :toctree: generated/ + + variables.Variables.continuous + variables.Variables.binaries + variables.Variables.integers + variables.Variables.semi_continuous + variables.Variables.sos + + +LinearExpression +================ + +Linear combination of variables. Arithmetic on ``Variable`` / +``LinearExpression`` returns a ``LinearExpression``. + +.. autosummary:: + :toctree: generated/ + + expressions.LinearExpression + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + expressions.LinearExpression.solution + +Operations +---------- + +.. autosummary:: + :toctree: generated/ + + expressions.LinearExpression.sum + expressions.LinearExpression.where + expressions.LinearExpression.groupby + expressions.LinearExpression.rolling + +Structure --------- -``Variables`` is a container for multiple N-D labeled variables. It is automatically added to a ``Model`` instance when initialized. +.. autosummary:: + :toctree: generated/ + + expressions.LinearExpression.vars + expressions.LinearExpression.coeffs + expressions.LinearExpression.const + expressions.LinearExpression.nterm + expressions.LinearExpression.has_terms + +Conversion +---------- + +.. autosummary:: + :toctree: generated/ + + expressions.LinearExpression.to_polars + +Construction +------------ .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + expressions.LinearExpression.from_tuples + expressions.merge + - variables.Variables - variables.Variables.add - variables.Variables.remove - variables.Variables.continuous - variables.Variables.integers - variables.Variables.binaries - variables.Variables.integers - variables.Variables.flat +QuadraticExpression +=================== + +Quadratic combination of variables, returned when squared +``Variable`` / ``LinearExpression`` arithmetic is performed. + +.. autosummary:: + :toctree: generated/ + expressions.QuadraticExpression -LinearExpressions +Structure +--------- + +.. autosummary:: + :toctree: generated/ + + expressions.QuadraticExpression.vars + expressions.QuadraticExpression.coeffs + expressions.QuadraticExpression.const + expressions.QuadraticExpression.nterm + expressions.QuadraticExpression.has_terms + +Conversion +---------- + +.. autosummary:: + :toctree: generated/ + + expressions.QuadraticExpression.to_matrix + expressions.QuadraticExpression.to_polars + +Post-solve access ----------------- .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + expressions.QuadraticExpression.solution - expressions.LinearExpression - expressions.LinearExpression.sum - expressions.LinearExpression.where - expressions.LinearExpression.groupby - expressions.LinearExpression.rolling - expressions.LinearExpression.from_tuples - expressions.merge - expressions.ScalarLinearExpression Constraint +========== + +Subclass of ``xarray.DataArray`` carrying labels for a multi-dimensional +constraint. + +.. autosummary:: + :toctree: generated/ + + constraints.Constraint + +Structure +--------- + +.. autosummary:: + :toctree: generated/ + + constraints.Constraint.lhs + constraints.Constraint.sign + constraints.Constraint.rhs + constraints.Constraint.coeffs + constraints.Constraint.vars + +Modification +------------ + +``Constraint.update`` is the canonical mutation API. The legacy ``lhs`` / +``sign`` / ``rhs`` / ``coeffs`` / ``vars`` setters still forward to +``update`` but emit a ``DeprecationWarning`` and will be removed in a +future release. + +.. autosummary:: + :toctree: generated/ + + constraints.Constraint.update + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + constraints.Constraint.dual + +Conversion ---------- -``Constraint`` is a subclass of ``xarray.DataArray`` and contains all labels referring to a multi-dimensional constraint. +.. autosummary:: + :toctree: generated/ + + constraints.Constraint.to_polars + + +CSRConstraint +============= + +Memory-efficient, immutable constraint representation backed by a scipy +CSR sparse matrix. Opt in via ``Model(freeze_constraints=True)`` or +``Model.add_constraints(..., freeze=True)``. See the +:doc:`creating-constraints` guide for usage. .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + constraints.CSRConstraint - constraints.Constraint - constraints.Constraint.coeffs - constraints.Constraint.vars - constraints.Constraint.lhs - constraints.Constraint.sign - constraints.Constraint.rhs - constraints.Constraint.flat +Structure +--------- + +.. autosummary:: + :toctree: generated/ + + constraints.CSRConstraint.coeffs + constraints.CSRConstraint.vars + constraints.CSRConstraint.sign + constraints.CSRConstraint.rhs + constraints.CSRConstraint.ncons + constraints.CSRConstraint.nterm + +Post-solve access +----------------- + +.. autosummary:: + :toctree: generated/ + + constraints.CSRConstraint.dual + +Conversion +---------- + +.. autosummary:: + :toctree: generated/ + + constraints.CSRConstraint.to_polars Constraints ------------ +=========== + +Container for the collection of constraints on a model. Accessed via +``model.constraints``. .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + constraints.Constraints - constraints.Constraints - constraints.Constraints.add - constraints.Constraints.remove - constraints.Constraints.coefficientrange - constraints.Constraints.inequalities - constraints.Constraints.equalities - constraints.Constraints.sanitize_missings - constraints.Constraints.flat - constraints.Constraints.to_matrix +Inventory +--------- + +.. autosummary:: + :toctree: generated/ + constraints.Constraints.inequalities + constraints.Constraints.equalities -IO functions -============ +Aggregate access +---------------- .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + constraints.Constraints.coeffs + constraints.Constraints.vars + constraints.Constraints.sign + constraints.Constraints.rhs + constraints.Constraints.dual + +Conversion +---------- + +.. autosummary:: + :toctree: generated/ + + constraints.Constraints.to_matrix + + +Objective +========= + +Wraps the objective expression on a model. Accessed via +``model.objective``. + +.. autosummary:: + :toctree: generated/ + + objective.Objective + objective.Objective.expression + objective.Objective.sense + objective.Objective.value + objective.Objective.is_linear + objective.Objective.is_quadratic + + +Piecewise +========= + +Construction helpers +-------------------- + +.. autosummary:: + :toctree: generated/ + + piecewise.breakpoints + piecewise.segments + piecewise.Slopes - model.Model.get_problem_file - model.Model.get_solution_file - model.Model.to_file - model.Model.to_netcdf - io.read_netcdf +PiecewiseFormulation +-------------------- -Solver utilities -================= +Returned by :func:`Model.add_piecewise_formulation`. .. autosummary:: - :toctree: generated/ + :toctree: generated/ - solvers.available_solvers - solvers.quadratic_solvers - solvers.Solver + piecewise.PiecewiseFormulation + piecewise.PiecewiseFormulation.method + piecewise.PiecewiseFormulation.convexity + piecewise.PiecewiseFormulation.variables + piecewise.PiecewiseFormulation.constraints + +Low-level helper +---------------- + +.. autosummary:: + :toctree: generated/ + + piecewise.tangent_lines + +Type aliases +------------ + +.. autosummary:: + :toctree: generated/ + + constants.PWL_METHOD + constants.PWL_METHODS + constants.PWL_CONVEXITY + constants.PWL_CONVEXITIES Solvers -======= +======== .. autosummary:: - :toctree: generated/ + :toctree: generated/ + + solvers.available_solvers + solvers.CBC + solvers.COPT + solvers.Cplex + solvers.GLPK + solvers.Gurobi + solvers.Highs + solvers.Knitro + solvers.MindOpt + solvers.Mosek + solvers.SCIP + solvers.Xpress + solvers.cuPDLPx + + +Remote solving +============== - solvers.CBC - solvers.Cplex - solvers.GLPK - solvers.Gurobi - solvers.Highs - solvers.Mosek - solvers.SCIP - solvers.Xpress +.. autosummary:: + :toctree: generated/ + remote.RemoteHandler -Solving + +Solver status and result types +============================== + +Types returned by or compared against :attr:`Model.status`, +:attr:`Model.termination_condition`, and :attr:`Model.solution`. + +.. autosummary:: + :toctree: generated/ + + constants.SolverStatus + constants.TerminationCondition + constants.Status + constants.Solution + constants.Result + + +Utilities +========= + +.. autosummary:: + :toctree: generated/ + + align + options + + +Warnings ======== +These warning classes can be silenced or filtered via +:func:`warnings.filterwarnings`. + .. autosummary:: - :toctree: generated/ + :toctree: generated/ - model.Model.solve - constants.SolverStatus - constants.TerminationCondition - constants.Status - constants.Solution - constants.Result + EvolvingAPIWarning + PerformanceWarning diff --git a/doc/benchmark.rst b/doc/benchmark.rst index da2a98e97..f56d5dffa 100644 --- a/doc/benchmark.rst +++ b/doc/benchmark.rst @@ -1,7 +1,7 @@ .. _benchmark: -Benchmarks -========== +Performance comparison +====================== Linopy's performance scales well with the problem size. Its overall speed is comparable with the famous `JuMP `_ package written in `Julia `_. It even outperforms `JuMP` in total memory efficiency when it comes to large models. Compared to `Pyomo `_, the common optimization package in python, one can expect @@ -13,9 +13,9 @@ for large problems. The following figure shows the memory usage and speed for so .. math:: - & \min \;\; \sum_{i,j} 2 x_{i,j} \; y_{i,j} \\ + & \min \;\; \sum_{i,j} 2 x_{i,j} + y_{i,j} \\ s.t. & \\ - & x_{i,j} - y_{i,j} \; \ge \; i \qquad \forall \; i,j \in \{1,...,N\} \\ + & x_{i,j} - y_{i,j} \; \ge \; i-1 \qquad \forall \; i,j \in \{1,...,N\} \\ & x_{i,j} + y_{i,j} \; \ge \; 0 \qquad \forall \; i,j \in \{1,...,N\} diff --git a/doc/conf.py b/doc/conf.py index d33175e11..c6b3b90c1 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -13,21 +13,18 @@ # import os # import sys # sys.path.insert(0, os.path.abspath('.')) -import pkg_resources # part of setuptools +import linopy # -- Project information ----------------------------------------------------- project = "linopy" -copyright = "2021, Fabian Hofmann" -author = "Fabian Hofmann" +copyright = "2021-2026, Fabian Hofmann, Felix Bumann" +author = "Fabian Hofmann, Felix Bumann" # The full version, including alpha/beta/rc tags -version = pkg_resources.get_distribution("linopy").version +version = linopy.__version__ release = "master" if "dev" in version else version -# For some reason is this needed, otherwise autosummary does fail on RTD but not locally -import linopy # noqa - # -- General configuration --------------------------------------------------- # Add any Sphinx extension module names here, as strings. They can be @@ -75,6 +72,16 @@ autosummary_generate = True autodoc_typehints = "none" +# Intersphinx — resolve :class:`xarray.DataArray`, :func:`numpy.ndarray`, etc. +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "numpy": ("https://numpy.org/doc/stable", None), + "pandas": ("https://pandas.pydata.org/docs", None), + "xarray": ("https://docs.xarray.dev/en/stable", None), + "scipy": ("https://docs.scipy.org/doc/scipy", None), + "dask": ("https://docs.dask.org/en/stable", None), +} + # Napoleon configurations napoleon_google_docstring = False @@ -94,16 +101,13 @@ """ -nbsphinx_allow_errors = True +nbsphinx_allow_errors = False nbsphinx_execute = "auto" nbsphinx_execute_arguments = [ "--InlineBackend.figure_formats={'svg', 'pdf'}", "--InlineBackend.rc={'figure.dpi': 96}", ] -# Exclude notebooks that require credentials or special setup -nbsphinx_execute_never = ["**/solve-on-oetc*"] - # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for diff --git a/doc/contributing.rst b/doc/contributing.rst index 02162694d..6c2ee2de3 100644 --- a/doc/contributing.rst +++ b/doc/contributing.rst @@ -10,17 +10,25 @@ contributing code. You are invited to submit pull requests / issues to our `Github repository `_. +AI-assisted contributions +========================= + +We're happy for you to use AI tools when contributing to linopy, but the +conversation around a change — pull request and issue descriptions, comments and +review discussions — must stay human and honest. The rules (mark AI-generated +content, write your own intent by hand) apply to humans and AI agents alike and +are kept in a single place: see `AGENTS.md +`_ in the repository +root. If you use an AI agent here, point it at that file. + Development Setup ================= -For linting, formatting and checking your code contributions -against our guidelines (e.g. we use `Black `_ as code style -and use `pre-commit `_: +For linting and formatting, we use `ruff `_ +and run it via `pre-commit `_: -1. Installation ``conda install -c conda-forge pre-commit`` or ``pip install pre-commit`` -2. Usage: - * To automatically activate ``pre-commit`` on every ``git commit``: Run ``pre-commit install`` - * To manually run it: ``pre-commit run --all`` +* Install the git hook (once): ``pre-commit install`` +* Run manually: ``pre-commit run --all-files`` Running Tests ============= @@ -35,7 +43,7 @@ To run the test suite: .. code-block:: bash # Install development dependencies - pip install -e .[dev,solvers] + uv sync --extra dev --extra solvers # Run all tests pytest @@ -68,6 +76,30 @@ GPU tests are automatically detected based on solver capabilities - no manual ma See the :doc:`gpu-acceleration` guide for more information about GPU solver setup and usage. +Performance Benchmarks +====================== + +When working on performance-sensitive code, use the internal benchmark suite in ``benchmarks/`` to check for regressions. + +.. code-block:: bash + + # Install benchmark dependencies + uv sync --extra benchmarks + + # Quick timing benchmarks + pytest benchmarks/ --quick + + # Compare timing between branches + pytest benchmarks/test_build.py --benchmark-save=master + pytest benchmarks/test_build.py --benchmark-save=my-feature --benchmark-compare=0001_master + + # Compare peak memory between branches + python benchmarks/memory.py save master --quick + python benchmarks/memory.py save my-feature --quick + python benchmarks/memory.py compare master my-feature + +See ``benchmarks/README.md`` for full details on models, phases, and usage. + Contributing examples ===================== @@ -98,23 +130,25 @@ Then for every notebook: e.g. `Edit -> Clear all output` in JupyterLab. 3. Provide a link to the documentation: - Include a file ``foo.nblink`` located in ``doc/examples/foo.nblink`` + Include a file ``foo.nblink`` located in ``doc/foo.nblink`` with this content - .. code-block: - { - 'path' : '../../examples/foo.ipynb' - } + .. code-block:: json + + { + "path": "../examples/foo.ipynb" + } + + Adjust the path for your file's name. + This ``nblink`` allows us to link your notebook into the documentation. - Adjust the path for your file's name. - This ``nblink`` allows us to link your notebook into the documentation. 4. Link your file in the documentation: Either - * Include your ``examples/foo.nblink`` directly into one of - the documentations toctrees; or - * Tell us where in the documentation you want your example to show up + * Include your ``foo.nblink`` directly into one of + the documentation's toctrees; or + * Tell us where in the documentation you want your example to show up 5. Commit your changes. If the precommit hook you installed above kicks in, confirm diff --git a/doc/coordinate-alignment.nblink b/doc/coordinate-alignment.nblink new file mode 100644 index 000000000..ef588b91f --- /dev/null +++ b/doc/coordinate-alignment.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/coordinate-alignment.ipynb" +} diff --git a/doc/gpu-acceleration.rst b/doc/gpu-acceleration.rst index a9048973d..2498993a6 100644 --- a/doc/gpu-acceleration.rst +++ b/doc/gpu-acceleration.rst @@ -24,7 +24,7 @@ To install it, you have to have the `CUDA Toolkit =0.1.2 + uv pip install "cupdlpx>=0.1.2" **Features:** diff --git a/doc/index.rst b/doc/index.rst index bff9fa65e..398466074 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -23,7 +23,7 @@ Main features `xarray `__ which allows for many flexible data-handling features: -- Define (arrays of) contnuous or binary variables with +- Define (arrays of) continuous or binary variables with **coordinates**, e.g. time, consumers, etc. - Apply **arithmetic operations** on the variables like adding, subtracting, multiplying with all the **broadcasting** potentials of @@ -42,7 +42,7 @@ flexible data-handling features: - Support of various solvers - `Cbc `__ - `GLPK `__ - - `HiGHS `__ + - `HiGHS `__ - `MindOpt `__ - `Gurobi `__ - `Xpress `__ @@ -81,7 +81,7 @@ A BibTeX entry for LaTeX users is License ------- -Copyright 2021-2023 Fabian Hofmann +Copyright 2021-2026 Fabian Hofmann, Felix Bumann This package is published under MIT license. @@ -111,22 +111,47 @@ This package is published under MIT license. creating-variables creating-expressions creating-constraints - sos-constraints + coordinate-alignment manipulating-models - testing-framework + +.. toctree:: + :hidden: + :maxdepth: 2 + :caption: Examples + transport-tutorial - infeasible-model + migrating-from-pyomo + +.. toctree:: + :hidden: + :maxdepth: 2 + :caption: Advanced Features + + sos-constraints + piecewise-linear-constraints + testing-framework + +.. toctree:: + :hidden: + :maxdepth: 2 + :caption: Solving + solve-on-remote solve-on-oetc gpu-acceleration - migrating-from-pyomo - gurobi-double-logging +.. toctree:: + :hidden: + :maxdepth: 2 + :caption: Troubleshooting + + infeasible-model + gurobi-double-logging .. toctree:: :hidden: :maxdepth: 2 - :caption: Benchmarking + :caption: Comparisons benchmark syntax @@ -134,7 +159,7 @@ This package is published under MIT license. .. toctree:: :hidden: :maxdepth: 2 - :caption: References + :caption: Reference api release_notes diff --git a/doc/piecewise-inequality-bounds-tutorial.nblink b/doc/piecewise-inequality-bounds-tutorial.nblink new file mode 100644 index 000000000..698826c74 --- /dev/null +++ b/doc/piecewise-inequality-bounds-tutorial.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/piecewise-inequality-bounds.ipynb" +} diff --git a/doc/piecewise-linear-constraints-tutorial.nblink b/doc/piecewise-linear-constraints-tutorial.nblink new file mode 100644 index 000000000..ea48e11f5 --- /dev/null +++ b/doc/piecewise-linear-constraints-tutorial.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/piecewise-linear-constraints.ipynb" +} diff --git a/doc/piecewise-linear-constraints.rst b/doc/piecewise-linear-constraints.rst new file mode 100644 index 000000000..c57d26943 --- /dev/null +++ b/doc/piecewise-linear-constraints.rst @@ -0,0 +1,645 @@ +.. _piecewise-linear-constraints: + +Piecewise Linear Constraints +============================ + +Piecewise linear (PWL) constraints approximate nonlinear functions as +connected linear pieces, allowing you to model cost curves, efficiency +curves, or production functions within a linear programming framework. + +Throughout this page: a **breakpoint** is a knot where the slope can +change; a **piece** is the linear part between adjacent breakpoints; a +**segment** is a disjoint operating region in the disjunctive +formulation. Per-tuple breakpoint arrays are paired by index — the +``i``-th entries across all tuples together define one knot. + +.. contents:: + :local: + :depth: 2 + + +Quick Start +----------- + +**Equality — lock variables onto the piecewise curve:** + +.. code-block:: python + + import linopy + + m = linopy.Model() + power = m.add_variables(name="power", lower=0, upper=100) + fuel = m.add_variables(name="fuel") + + # fuel = f(power) on the piecewise curve defined by these breakpoints + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, [0, 36, 84, 170]), + ) + +**Inequality — bound one expression by the curve:** + +.. code-block:: python + + # fuel >= f(power) on the same heat-rate curve as above. + m.add_piecewise_formulation( + (fuel, [0, 36, 84, 170], ">="), + (power, [0, 30, 60, 100]), + ) + +Over-fuelling is physically admissible but wasteful, so minimising +fuel pulls the operating point onto the curve. ``method="auto"`` +picks the cheapest correct formulation: pure LP (chord constraints) +here, since convex + ``">="`` is LP-applicable; SOS2/incremental +otherwise. + +Each ``(expression, breakpoints[, sign])`` tuple pairs a variable with +its breakpoint values. The optional sign (default ``"=="``) is ``"<="`` +or ``">="`` to mark that expression as bounded by the curve. With every +sign ``"=="``, all tuples land on the same point of the piecewise curve +— see *Per-tuple sign* below for the geometry of the inequality cases. + + +API +--- + +``add_piecewise_formulation`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + m.add_piecewise_formulation( + (expr1, breakpoints1), # sign defaults to "==" (equality role) + (expr2, breakpoints2, "<="), # or with an explicit sign + ..., + method="auto", # "auto", "sos2", "incremental", or "lp" + active=None, # binary variable to gate the constraint + name=None, # base name for generated variables/constraints + ) + +Adds constraints — and, depending on the resolved method, auxiliary +variables — for either an all-equality joint (every tuple at the same +point on the curve, the default) or a one-sided bound on a single +tuple. The pure-LP path adds chord and domain constraints only; SOS2, +incremental, and disjunctive also add interpolation weights and/or +binaries (see *Formulation Methods* below). + + +Breakpoint Construction +----------------------- + +Each tuple's breakpoints come from :func:`~linopy.breakpoints` (a +single connected curve) or :func:`~linopy.segments` (disjoint +operating bands). :class:`~linopy.Slopes` can stand in for +:func:`~linopy.breakpoints` when per-piece slopes are the natural +input — it resolves to a breakpoints array. + +``breakpoints()`` — a connected curve +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Values along a single **connected** piecewise curve — the default case +for efficiency curves, heat rates, and cost curves. + +The simplest form passes a Python list directly in the tuple: + +.. code-block:: python + + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, [0, 36, 84, 170]), + ) + +Equivalent, but explicit about the DataArray construction: + +.. code-block:: python + + m.add_piecewise_formulation( + (power, linopy.breakpoints([0, 30, 60, 100])), + (fuel, linopy.breakpoints([0, 36, 84, 170])), + ) + +**Per-entity curves.** Different generators can have different +curves. Pass a dict to ``breakpoints()`` with entity names as keys: + +.. code-block:: python + + m.add_piecewise_formulation( + ( + power, + linopy.breakpoints( + {"gas": [0, 30, 60, 100], "coal": [0, 50, 100, 150]}, dim="gen" + ), + ), + ( + fuel, + linopy.breakpoints( + {"gas": [0, 40, 90, 180], "coal": [0, 55, 130, 225]}, dim="gen" + ), + ), + ) + +Ragged lengths are NaN-padded automatically. Breakpoints are auto- +broadcast over remaining dimensions (e.g. ``time``). + +**Specifying by slopes.** :class:`linopy.Slopes` resolves to a +breakpoint array from per-piece slopes plus an initial ``y0``, +instead of from absolute y-values — useful when slopes are the +natural input (e.g. marginal costs). The x grid is borrowed from +the sibling tuple, so the y breakpoints don't have to be computed +by hand: + +.. code-block:: python + + m.add_piecewise_formulation( + (power, [0, 50, 100, 150]), + (cost, linopy.Slopes([1.1, 1.5, 1.9], y0=0)), + ) + # cost breakpoints resolve to: [0, 55, 130, 225] + +For standalone resolution outside ``add_piecewise_formulation``, call +:meth:`linopy.Slopes.to_breakpoints` with an explicit x grid:: + + bp = linopy.Slopes([1.1, 1.5, 1.9], y0=0).to_breakpoints([0, 50, 100, 150]) + +``segments()`` — disjoint operating bands +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For equipment with disconnected operating bands. Each segment is one +band's ``(range, curve)``; a binary picks exactly one per operating +point, with continuous interpolation within the chosen band. + +.. code-block:: python + + # Stepped pump with two speed bands. + m.add_piecewise_formulation( + (flow, linopy.segments([(5, 25), (40, 100)])), + (power, linopy.segments([(1, 7), (15, 50)])), + ) + +Bounded tuples (``"<="`` / ``">="``) are supported on disjunctive +curves too. + +For a single on/off gate on one continuous curve, prefer ``active=...`` +(see :ref:`piecewise-active`) — using a degenerate ``(0, 0)`` segment +to encode "off" mixes the disjunctive concept with on/off logic. + +N-variable linking +~~~~~~~~~~~~~~~~~~ + +Independent of the building block used, any number of variables can be +linked through shared breakpoints (joint equality): + +.. code-block:: python + + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, [0, 40, 85, 160]), + (heat, [0, 25, 55, 95]), + ) + +All variables are symmetric here; every feasible point is the same +``λ``-weighted combination of breakpoints across all three. Sign +restrictions apply (see *Per-tuple sign* below) — for ``N ≥ 3`` tuples +all signs must be ``"=="``. + + +Per-tuple sign — equality vs inequality +---------------------------------------- + +Roles and restrictions +~~~~~~~~~~~~~~~~~~~~~~ + +Each tuple's optional third element is a sign: + +- ``"=="`` (default) — **equality role**: the tuple enters as an + equality. +- ``"<="`` / ``">="`` — **bounded**: the expression undershoots / + overshoots the curve. + +.. note:: + + **Current restrictions.** + + - At most one tuple may carry a non-equality sign — a single bounded side. + - With **3 or more** tuples, all signs must be ``"=="``. + + Multi-bounded and N≥3-inequality use cases aren't supported yet. + If you have a concrete use case, please open an issue at + https://github.com/PyPSA/linopy/issues so we can scope it properly. + +Geometry +~~~~~~~~ + +What the formulation actually constrains depends on the tuple count and +signs: + +- **All-equality (default).** Shared interpolation weights put the + joint :math:`(e_1, \ldots, e_N)` exactly on the curve. +- **One bounded + one equality (2 tuples).** The joint :math:`(x, y)` + lies in the **hypograph** (``"<="`` on a concave :math:`f`) or + **epigraph** (``">="`` on a convex :math:`f`): + + .. math:: + + \{ (x, y) \ :\ x_{\min} \le x \le x_{\max},\ y \le f(x) \} + \qquad \text{(hypograph)} + + The equality axis is just confined to its breakpoint domain + :math:`[x_{\min}, x_{\max}]` — a single coordinate can't locate a + curve point. Same projection in LP (enforced directly) and + SOS2/incremental (enforced via the weight link). +- **Mismatched sign + curvature** (convex + ``"<="``, or concave + + ``">="``) describes a *non-convex* region — ``method="auto"`` falls + back to SOS2/incremental and ``method="lp"`` raises. + +.. code-block:: python + + # All-equality: joint (x, y) on the curve. + m.add_piecewise_formulation((y, y_pts), (x, x_pts)) + + # Bounded: joint (x, y) in the hypograph — y ≤ f(x), x ∈ [x_min, x_max]. + m.add_piecewise_formulation((y, y_pts, "<="), (x, x_pts)) + + # Bounded: joint (x, y) in the epigraph — y ≥ f(x), x ∈ [x_min, x_max]. + m.add_piecewise_formulation((y, y_pts, ">="), (x, x_pts)) + + # 3-variable all-equality (CHP): joint (power, fuel, heat) on the curve. + m.add_piecewise_formulation((power, p_pts), (fuel, f_pts), (heat, h_pts)) + +Choice of bounded tuple and sign +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pick the sign matching the physically admissible direction for that +expression: + +- ``"<="`` for a quantity with a controllable *dissipation path* — heat + rejection via cooling tower (*thermal curtailment*), electrical + curtailment, emissions after post-treatment — so undershooting the + curve is realisable. +- ``">="`` for an *input* whose over-supply is admissible but wasteful — + fuel, raw materials — so overshooting the curve is realisable + (objective pressure then pulls the operating point onto the curve). + +The wrong direction (``"<="`` on fuel, ``">="`` on a non-curtailable +output) yields a valid but **loose** formulation that admits operating +points the plant cannot physically realise; an objective rewarding the +wrong direction may then find a non-physical optimum — safe only when +no such objective pressure exists. + +When is a one-sided bound wanted? +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For *continuous* curves, the main reason to reach for ``"<="`` / +``">="`` is to unlock the **LP chord formulation** — no SOS2, no +binaries, just pure LP. On a convex/concave curve with a matching +sign, the chord inequalities are as tight as SOS2, so you get the same +optimum with a cheaper model. Inequality formulations also tighten +the LP relaxation of SOS2/incremental, which can reduce branch-and- +bound work even when LP itself is not applicable. + +For *disjunctive* curves (``segments(...)``), the per-tuple sign is a +first-class tool in its own right: disconnected operating regions with +a bounded output, always exact regardless of segment curvature (see +the disjunctive section below). + +If the curvature doesn't match the sign (convex + ``"<="``, or concave + +``">="``), LP is not applicable — ``method="auto"`` falls back to +SOS2/incremental with the signed link, which gives a valid but much +more expensive model. In that case prefer ``"=="`` unless you +genuinely need the one-sided semantics. See the +:doc:`piecewise-inequality-bounds-tutorial` notebook for a full +walkthrough. + +Formulation Methods +------------------- + +Pass ``method="auto"`` (the default) and linopy picks the cheapest correct +formulation based on ``sign``, curvature and breakpoint layout: + +- **2-variable inequality on a convex/concave curve** → ``lp`` (chord lines, + no auxiliary variables) +- **All breakpoints monotonic** → ``incremental`` +- **Otherwise** → ``sos2`` +- **Disjunctive (segments)** → SOS2 applied per segment with binary + segment selection (the disjunctive formulation in the table below). + +The resolved choice is exposed on the returned ``PiecewiseFormulation`` +via ``.method`` (and ``.convexity`` when well-defined). An +``INFO``-level log line explains the resolution whenever +``method="auto"`` is in play. + +At-a-glance comparison: + +.. list-table:: + :header-rows: 1 + :widths: 26 18 18 18 20 + + * - Property + - ``lp`` + - ``incremental`` + - ``sos2`` + - Disjunctive + * - Curve layout + - Connected + - Connected + - Connected + - Disconnected + * - Supported per-tuple sign + - one ``<=`` or ``>=`` (required) + - all ``==`` or one ``<=``/``>=`` + - all ``==`` or one ``<=``/``>=`` + - all ``==`` or one ``<=``/``>=`` + * - Number of tuples + - Exactly 2 + - ≥ 2 (3+ requires all ``==``) + - ≥ 2 (3+ requires all ``==``) + - ≥ 2 (3+ requires all ``==``) + * - Breakpoint order + - Strictly monotonic + - Strictly monotonic + - Any + - Any (per segment) + * - Curvature requirement + - Concave (``<=``) or convex (``>=``) + - None + - None + - None + * - Auxiliary variables + - **None** + - Continuous + binary + - Continuous + SOS2 + - Continuous + binary + SOS2 + * - ``active=`` supported + - No + - Yes + - Yes + - Yes + * - Solver requirement + - **Any LP solver** + - MIP-capable + - SOS2-capable (or MIP via :ref:`Big-M reformulation `) + - SOS2 + MIP (or MIP via :ref:`Big-M reformulation `) + +.. note:: + + Disjunctive formulations report ``method="sos2"`` (the underlying + per-segment encoding uses SOS2), but the table treats them as a + separate column because the per-segment binaries change the + auxiliary-variable structure and solver requirements. + +LP (chord-line) formulation +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For **2-variable inequality** on a **convex** or **concave** curve. +Adds one chord inequality per piece plus a domain bound — no auxiliary +variables and no MIP relaxation: + +.. math:: + + &y \ \text{sign}\ m_k \cdot x + c_k + \quad \forall\ \text{pieces } k + + &x_{\min} \le x \le x_{\max} + +where :math:`m_k = (y_{k+1} - y_k)/(x_{k+1} - x_k)` and +:math:`c_k = y_k - m_k\, x_k`. The domain bound uses :math:`x_{\min}` +and :math:`x_{\max}` rather than the first/last breakpoint so that +descending x grids work too — strictly-monotonic breakpoints are +accepted in either order. For concave :math:`f` with ``sign="<="``, +the intersection of all chord inequalities equals the hypograph of +:math:`f` on its domain. + +The LP dispatch requires curvature and sign to match: ``sign="<="`` +needs concave (or linear); ``sign=">="`` needs convex (or linear). A +mismatch is *not* just a loose bound — it describes the wrong region +(see the :doc:`piecewise-inequality-bounds-tutorial`). +``method="auto"`` detects this and falls back; ``method="lp"`` raises. + +.. code-block:: python + + # y <= f(x) on a concave f — auto picks LP + m.add_piecewise_formulation((y, yp, "<="), (x, xp)) + + # Or explicitly: + m.add_piecewise_formulation((y, yp, "<="), (x, xp), method="lp") + +**Not supported with** ``method="lp"``: all-equality, more than 2 +tuples, and ``active``. ``method="auto"`` falls back to +SOS2/incremental in all three cases. + +Chord expressions as a building block +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The underlying chord expressions are also exposed as a standalone +helper, :func:`~linopy.tangent_lines`, which returns the per-piece +chord as a :class:`~linopy.expressions.LinearExpression` with no +variables created. Use it directly when you want to compose the chord +bound with other constraints by hand, without the domain bound that +``method="lp"`` adds automatically: + +.. code-block:: python + + chord = linopy.tangent_lines(x, x_pts, y_pts) + m.add_constraints(y <= chord + slack) + +Incremental (Delta) formulation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The default MIP encoding when ``method="auto"`` is in play and breakpoints +are **strictly monotonic** — produces a tight MIP directly, without going +through an SOS2 layer. Uses fill-fraction variables :math:`\delta_i` with +binary indicators :math:`z_i`: + +.. math:: + + &\delta_i \in [0, 1], \quad z_i \in \{0, 1\} + + &\delta_{i+1} \le \delta_i, \quad z_{i+1} \le \delta_i, \quad \delta_i \le z_i + + &e_j = B_{j,0} + \sum_{i=1}^{n} \delta_i \, (B_{j,i} - B_{j,i-1}) + +With a bounded tuple, the link to that tuple's expression flips to the +requested sign while the equality-signed tuples keep the equality above. + +.. code-block:: python + + m.add_piecewise_formulation((power, xp), (fuel, yp), method="incremental") + +**Limitation:** breakpoint sequences must be strictly monotonic. + +SOS2 (Convex combination) +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Fallback when breakpoints aren't strictly monotonic (the only case +``method="auto"`` does not pick incremental for a connected curve). +Introduces interpolation weights :math:`\lambda_i` with an SOS2 +adjacency constraint: + +.. math:: + + &\sum_{i=0}^{n} \lambda_i = 1, \qquad + \text{SOS2}(\lambda_0, \ldots, \lambda_n) + + &e_j = \sum_{i=0}^{n} \lambda_i \, B_{j,i} + \quad \text{for each expression } j + +The SOS2 constraint ensures at most two adjacent :math:`\lambda_i` are +non-zero, so every expression is interpolated within the same piece. +With a bounded tuple, the bounded link flips to the requested sign as +above. + +.. note:: + + Solvers with native SOS2 support handle the adjacency constraint via + branch-and-bound. Solvers without it see the Big-M reformulation + linopy applies (controlled by ``reformulate_sos=`` on ``solve``) — + see :ref:`sos-reformulation` for the reformulated MIP form, which is + the model those solvers actually receive. When breakpoints are + monotonic, prefer ``method="incremental"`` (or just ``"auto"``): it + builds a similar MIP encoding directly and does not depend on + solver SOS2 support or the reformulation step. + +.. code-block:: python + + m.add_piecewise_formulation((power, xp), (fuel, yp), method="sos2") + +Disjunctive (Disaggregated convex combination) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For **disconnected segments** (gaps between operating regions). Binary +indicators :math:`z_k` select exactly one segment; SOS2 applies within it: + +.. math:: + + &z_k \in \{0, 1\}, \quad \sum_{k} z_k = 1 + + &\sum_{i} \lambda_{k,i} = z_k, \qquad + e_j = \sum_{k} \sum_{i} \lambda_{k,i} \, B_{j,k,i} + +No big-M constants are needed, giving a tight LP relaxation. + +**Disjunctive + bounded tuple.** A per-tuple ``"<="`` / ``">="`` works +here too, applied to the bounded tuple exactly as for the continuous +methods. Because the disjunctive machinery already carries a +per-segment binary, there is **no curvature requirement** on the +segments — inequality is always exact on the hypograph (or epigraph) of +the active segment, whatever its slope pattern. This makes disjunctive +plus a bounded tuple a first-class tool for "bounded output on +disconnected operating regions" that ``method="lp"`` cannot handle. + + +Advanced Features +----------------- + +.. _piecewise-active: + +Active parameter (unit commitment) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``active`` parameter gates the piecewise function with a binary variable. +When ``active=0``, all auxiliary variables (and thus the linked expressions) +are forced to zero: + +.. code-block:: python + + commit = m.add_variables(name="commit", binary=True, coords=[time]) + m.add_piecewise_formulation( + (power, [30, 60, 100]), + (fuel, [40, 90, 170]), + active=commit, + ) + +- ``commit=1``: power operates in [30, 100], fuel = f(power) +- ``commit=0``: power = 0, fuel = 0 + +Not supported with ``method="lp"`` (gating needs a binary). Use +``method="auto"``, or *Chord expressions as a building block* for +manual gating. + +.. warning:: + + With a bounded tuple and ``active=0``, the output is only forced to + ``0`` on the signed side — the complementary bound still comes from + the output variable's own lower/upper bound. In the common case of + non-negative outputs (fuel, cost, heat), set ``lower=0`` on that + variable: combined with the ``y ≤ 0`` constraint from deactivation, + this forces ``y = 0`` automatically. + +Partial gates +^^^^^^^^^^^^^ + +``active`` must cover the formulation's full coordinate; a gate defined +over only a subset (or with masked entries) is rejected unless +``active_fill`` is set. ``active_fill`` gates the missing entries as +always-active (``1``) or always-off (``0``) — handy when one formulation +mixes committable and non-committable units sharing a single ``status``: + +.. code-block:: python + + m.add_piecewise_formulation( + (power, [30, 60, 100]), (fuel, [40, 90, 170]), active=status, active_fill=1 + ) + +``active_fill`` is transitional: under v1 semantics, pad ``active`` +explicitly with ``active.reindex(coords).fillna(value)`` instead. + +Auto-broadcasting +~~~~~~~~~~~~~~~~~ + +Breakpoints are automatically broadcast to match expression dimensions — you +don't need ``expand_dims``: + +.. code-block:: python + + time = pd.Index([1, 2, 3], name="time") + x = m.add_variables(name="x", lower=0, upper=100, coords=[time]) + y = m.add_variables(name="y", coords=[time]) + + # 1D breakpoints auto-expand to match x's time dimension + m.add_piecewise_formulation((x, [0, 50, 100]), (y, [0, 70, 150])) + +NaN masking +~~~~~~~~~~~ + +Trailing NaN values in breakpoints mask the corresponding lambda / delta +variables (and, for LP, the corresponding chord constraints). This is useful +for per-entity breakpoints with ragged lengths: + +.. code-block:: python + + # gen1 has 3 breakpoints, gen2 has 2 (NaN-padded) + bp = linopy.breakpoints({"gen1": [0, 50, 100], "gen2": [0, 80]}, dim="gen") + +Interior NaN values (gaps in the middle) are not supported and raise an error. + +Inspecting generated objects +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The returned :class:`PiecewiseFormulation` exposes ``.variables`` and +``.constraints`` as live views into the model — use them to introspect +exactly what was generated, rather than relying on documented name +conventions: + +.. code-block:: python + + f = m.add_piecewise_formulation((y, y_pts, "<="), (x, x_pts)) + print(f) # method, convexity, vars/cons summary + +The comparison table above describes the *kind* of auxiliary objects each +method creates (continuous + SOS2, binary + SOS2, none, …); exact name +suffixes are an implementation detail and may evolve. + + +Tutorials +--------- + +.. toctree:: + :maxdepth: 1 + + piecewise-linear-constraints-tutorial + piecewise-inequality-bounds-tutorial + +See Also +-------- + +- :doc:`sos-constraints` — low-level SOS1/SOS2 constraint API diff --git a/doc/prerequisites.rst b/doc/prerequisites.rst index 23b178971..7a9e83bd9 100644 --- a/doc/prerequisites.rst +++ b/doc/prerequisites.rst @@ -12,11 +12,11 @@ Before you start, make sure you have the following: Install Linopy -------------- -You can install Linopy using pip or conda. Here are the commands for each method: +You can install Linopy using uv or conda. Here are the commands for each method: .. code-block:: bash - pip install linopy + uv pip install linopy or @@ -35,26 +35,28 @@ CPU-based solvers - `Cbc `__ - open source, free, fast - `GLPK `__ - open source, free, not very fast -- `HiGHS `__ - open source, free, fast +- `HiGHS `__ - open source, free, fast +- `SCIP `__ - open source (Apache-2.0), fast MIP solver - `Gurobi `__ - closed source, commercial, very fast - `Xpress `__ - closed source, commercial, very fast (GPU acceleration available in v9.8+) - `Cplex `__ - closed source, commercial, very fast -- `MOSEK `__ -- `MindOpt `__ - +- `MOSEK `__ - closed source, commercial, strong on conic/QP +- `MindOpt `__ - closed source, commercial - `COPT `__ - closed source, commercial, very fast -For a subset of the solvers, Linopy provides a wrapper. +The ``linopy[solvers]`` extra installs the Python clients for the +supported solvers (HiGHS, SCIP, Gurobi, CPLEX, MOSEK, MindOpt, COPT, +Xpress, Knitro). For the commercial ones a separate license is still +required: .. code:: bash - pip install linopy[solvers] + uv pip install "linopy[solvers]" -We recommend to install the HiGHS solver if possible, which is free and open source but not yet available on all platforms. - -.. code:: bash - - pip install highspy +We recommend the HiGHS solver, which is free, open source, and fast +across a wide range of problem sizes. It is included in both the +``solvers`` and ``dev`` extras. GPU-accelerated solvers @@ -68,7 +70,7 @@ For large-scale optimization problems, GPU-accelerated solvers can provide signi .. code:: bash - pip install cupdlpx + uv pip install cupdlpx For most of the other solvers, please click on the links to get further installation information. diff --git a/doc/release_notes.rst b/doc/release_notes.rst index b727c22d1..2981c9dd5 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,8 +4,232 @@ Release Notes Upcoming Version ---------------- -* Fix docs (pick highs solver) -* Add the `sphinx-copybutton` to the documentation +**Features** + + +*In-place solver updates (persistent re-solve)* + +* A built solver can now be re-solved against a mutated ``Model`` without a full rebuild. Construct with ``Solver.from_name(..., track_updates=True)`` and re-call ``solver.solve(model)`` after edits — the diff against the previous build is applied in place when the backend supports it, falling back to a rebuild otherwise. Supported on HiGHS, Gurobi, Xpress, and Mosek (``io_api="direct"``). +* Pass ``disallow_rebuild=True`` to ``solve(model, ...)`` to guarantee an in-place update or raise ``RebuildRequiredError``. Inspect ``solver._last_rebuild_reason`` (a ``RebuildReason``, or ``None`` after an in-place update) to understand why a rebuild was triggered. +* New ``linopy.persistent`` module exposes ``ModelSnapshot``, ``ModelDiff``, and ``RebuildReason`` for users who want to introspect or build the diff themselves. ``ModelDiff.from_snapshot`` / ``from_models`` return the ``RebuildReason`` directly when the change cannot be applied in place. + +*Improved IO* + +* ``Model.to_netcdf`` now records the writing linopy version in the ``_linopy_version`` dataset attribute. Files written by older versions (without the attribute) continue to read unchanged. + +*Other* + +* Default internal integer labels to ``int32`` (configurable via ``linopy.options["label_dtype"]``, set to ``np.int64`` for the old behavior), cutting memory ~25% and speeding up model build 10-35%. Raises ``ValueError`` if labels exceed the int32 maximum. +* ``add_variables(binary=True, ...)`` now accepts ``lower``/``upper`` bounds, as long as they are 0 or 1. Previously binary bounds could only be set via the ``.lower``/``.upper`` setters after creation. (https://github.com/PyPSA/linopy/issues/776) +* ``add_piecewise_formulation`` gained an ``active_fill`` parameter that gates a partial ``active`` (defined over a subset of the indexed dimension, or masked) as always-active (``1``) or always-off (``0``); without it, a partial ``active`` — which was previously zeroed silently — now raises. Useful when one formulation mixes gated and ungated entities (e.g. committable and non-committable units sharing a ``status``). ``active_fill`` is transitional and will be removed once v1 semantics make ``active.reindex(coords).fillna(value)`` sufficient. (https://github.com/PyPSA/linopy/issues/796) + +**Deprecations** + +* Mutation via assignment to ``Variable.lower`` / ``Variable.upper`` / ``Constraint.coeffs`` / ``Constraint.vars`` / ``Constraint.lhs`` / ``Constraint.sign`` / ``Constraint.rhs`` is deprecated and emits a ``DeprecationWarning``. Use ``Variable.update(...)`` / ``Constraint.update(...)`` instead — the canonical mutation API with one validation path and one place that flips the persistent-solver dirty flag. Read access to these properties is unchanged. The setters will be removed in a future release. +* Passing a raw ``DataArray`` of integer labels to ``Constraint.vars = ...`` setter is deprecated and emits a ``FutureWarning``. Pass a ``Variable`` to ``Constraint.update()`` instead — it is the supported input. The ``DataArray`` path will be removed in a future release. + +**Bug fixes** + +* LP file export now honors bounds tightened below ``[0, 1]`` on a binary variable via the ``.lower``/``.upper`` setters after creation (e.g. ``upper = 0``). Previously such bounds were written only by ``io_api="direct"`` and dropped by ``io_api="lp"``. (https://github.com/PyPSA/linopy/issues/776) +* Freezing an empty constraint group (e.g. an empty ``isel`` slice) no longer raises ``ValueError: cannot reshape array of size 0``. ``Model(freeze_constraints=True)`` and ``Constraint.freeze()`` now round-trip zero-row constraints losslessly. +* ``Variable.where`` no longer raises ``ValueError: exact match required for all data variable names`` once a solution is attached (after ``Model.solve``) or the variable is fixed. The fill value now covers auxiliary data variables (``solution``, stashed bounds) instead of only ``labels``/``lower``/``upper``. + +Version 0.8.0 +------------- + +**Features** + +*Constraints — CSR-backed storage* + +* Add ``CSRConstraint``: a memory-efficient immutable constraint representation backed by scipy CSR sparse matrices. Up to 90% memory savings for constraints with many terms and 30–120× faster matrix generation for direct solver APIs. +* Opt in globally via ``Model(freeze_constraints=True)`` or per-call via ``model.add_constraints(..., freeze=True)``. +* Lossless conversion both ways with ``Constraint.freeze()`` / ``CSRConstraint.mutable()``. + +*Inspect the solver after solving* + +* After ``model.solve()``, the solver object stays available on ``model.solver``. You can inspect it, reuse it, or release the underlying solver (and its license) by calling ``model.solver.close()`` or assigning ``model.solver = None``. It is also released automatically when the model is garbage-collected. +* New ``SolverReport`` on the result (``result.report``) reports runtime, MIP gap, dual (best) bound, and iteration counts. It is shown in ``repr(result)`` and currently populated by CBC, HiGHS, Gurobi, Knitro, and cuPDLPx. + +*Dualize LP models* + +* New ``Model.dualize()`` constructs the LP dual as a standalone model, lifting finite variable bounds into explicit constraints so they are reflected in the dual. Dual variables are named after the primal constraints. Works for linear problems with linear objective and constraints. (https://github.com/PyPSA/linopy/pull/626) + +*A new way to call a solver (advanced)* + +Most users should keep calling ``model.solve(...)``. If you want more control, you can now build the solver yourself and run it in two steps: + +.. code-block:: python + + solver = Solver.from_name("gurobi", model, io_api="direct", options=...) + result = solver.solve() + model.assign_result(result) # write the solution back + +``Solver`` is now a dataclass, so writing a new solver backend is simpler — subclasses just override the hooks they need (``_build_direct``, ``_run_direct``, ``_run_file``). + +*Querying solver capabilities* + +* Ask a solver class what it can do via ``Gurobi.supports(SolverFeature.MIP)`` (or any other ``SolverFeature``). ``SolverFeature`` is importable from ``linopy``. +* ``linopy.solver_capabilities`` still works (re-exports ``SolverFeature`` and ``solver_supports``), but the new ``SolverClass.supports(...)`` API is preferred. + +*Knowing which solvers you can actually use* + +* ``linopy.available_solvers`` no longer tries to acquire licenses at import time, so importing linopy is faster and doesn't grab a license from solvers like Gurobi or Mosek. **Note:** membership now means "the package is installed", not "I have a working license" (see Breaking Changes). Call ``available_solvers.refresh()`` to re-scan. Same for ``quadratic_solvers``. +* New ``linopy.licensed_solvers``: the subset of installed solvers that currently pass a license check. Handy in tests and for picking a solver at runtime. +* New helpers for explicit license checks: ``linopy.solvers.check_solver_licenses("gurobi", "mosek")``, ``Gurobi.license_status()``, ``Gurobi.is_available()``. They return a ``LicenseStatus`` dataclass (``name``, ``ok``, ``message``). + +*Constraints — indicator constraints* + +* Add indicator constraints for solvers that support them. They are part of the unified constraints container: ``model.add_indicator_constraints`` returns a ``Constraint`` and the constraint is stored in ``model.constraints`` (filterable via ``model.constraints.indicator`` / ``model.constraints.regular``), so it round-trips through netCDF and ``model.copy()``. + +*Compact multi-key grouping* + +* ``LinearExpressionGroupby.sum`` gains a pandas-style ``observed`` parameter for grouping by a list of coordinate names: ``expr.groupby(["period", "season"]).sum(observed=True)`` keeps the result stacked over only the observed key combinations (a ``MultiIndex`` ``group`` dimension) instead of unstacking into one dimension per key, which materialises the dense cartesian grid. The default ``observed=False`` mirrors xarray. When the grid would be mostly fill values, a ``UserWarning`` points to ``observed=True``. + +*Other additions* + +* Add ``BaseExpression.has_terms`` property: boolean array, true at slots with at least one live term (`#741 `_). +* Add ``BaseExpression.variable_names`` property, and documentation for ``LinearExpression.where`` with ``drop=True``. + +**Performance** + +* ~10× faster direct solver communication (``io_api="direct"``), thanks to the new CSR-based matrix construction. Conversion helpers like ``to_highspy`` benefit too. +* Writing the solution back to the model after solving is faster: it no longer rebuilds the constraint matrix, and now uses positional (rather than label-based) indexing — roughly 2× faster overall. +* Xpress now supports ``io_api="direct"``: the linopy model is loaded via the native ``loadproblem`` array API instead of being serialised through an LP/MPS file, with SOS constraints attached in-place. Adds ``model.to_xpress()`` matching the existing ``to_gurobipy`` / ``to_highspy`` / ``to_mosek`` helpers. + +**Deprecations** + +* ``Solver.solve_problem``, ``Solver.solve_problem_from_model``, and ``Solver.solve_problem_from_file`` still work but emit a ``DeprecationWarning``. Use ``Solver.from_name(...).solve()`` (or simply ``model.solve(...)``) instead. They will be removed in a future release. +* **Implicit MultiIndex-level projection is deprecated.** Passing an input indexed by a *level* of a stacked-``MultiIndex`` dimension (e.g. per-``period`` bounds onto a ``(period, timestep)`` ``snapshot`` index) emits an ``EvolvingAPIWarning`` — in arithmetic and in ``add_variables`` / ``add_constraints`` — and will raise under the upcoming v1 convention. Project the input onto the dimension explicitly (select with the dimension's level values) to keep current behavior. Affects PyPSA multi-investment models. See the level-projection entry under Bug Fixes for the new alignment behavior. + +**Bug Fixes** + +* **⚠ Behavior change:** ``add_variables`` / ``add_constraints``: extends 0.7.0's coords-as-truth rule to ``lower``, ``upper`` and ``mask`` for every bound type and dim order. Pandas ``Series`` / ``DataFrame`` bounds or masks missing a dimension are now broadcast to ``coords`` instead of being silently dropped (`#709 `__) — a model that previously *ignored* such a partial bound now *applies* it, silently, with no error — review partial pandas bounds/masks when upgrading. The variable's dimension order always follows ``coords`` (`#706 `__); bare-tuple coord entries (``coords=[(0, 1, 2)]``) now behave like lists. Mismatched values or extra dims raise ``ValueError`` with a labelled message; sparse-coord masks (formerly a v0.6.3 ``FutureWarning``, #580) raise ``ValueError``, and masks with dims not in the data raise ``ValueError`` instead of ``AssertionError``. +* **⚠ Behavior change:** Fix Mosek interface to inspect both the basic and IPM solutions and pick the one with the better status, so that an optimal crossover solution is not discarded when IPM terminates with a (near-)Farkas certificate. Mosek may now return a different (better-status) solution than 0.7.0 for the same model. +* Pandas inputs whose index names *levels* of a stacked-``MultiIndex`` ``coords`` dimension are now projected onto that dimension: a level subset broadcasts across the others, the full set aligns element-wise. In ``add_variables`` / ``add_constraints`` the input must provide a value for every level combination of the MultiIndex or a ``ValueError`` is raised (the error lists the missing combinations); aligning the full level set with full coverage stays silent. Strict validation also rejects a ``MultiIndex`` input with *unnamed* levels whose combinations don't match ``coords`` (previously a silent bypass, as such inputs can't be projected by level name). Implicit level projections are deprecated (see Deprecations). +* ``LinearExpression.groupby`` now accepts a **non-dimension** coordinate as the key -- by name (``expr.groupby("period").sum()``, where ``period`` labels another dimension) or as the coordinate ``DataArray`` -- which previously raised ``ValueError: ... already exists``. Grouping by a dimension or a ``MultiIndex`` level already worked (`#750 `__). +* ``LinearExpression.groupby`` with a **list of coordinate names** (``expr.groupby(["period", "season"]).sum()``) now takes the fast reindex path instead of silently falling back to the slow xarray implementation, returning one dimension per key as before (`#753 `__). See ``observed`` under Features to keep the result compact instead. +* SOS constraints on masked variables no longer cause solver-specific failures (Gurobi ``IndexError``, Xpress ``?404 Invalid column number``, LP parse errors, silent set corruption). ``Model.solve()`` and ``Model.to_file()`` now raise a clear ``NotImplementedError`` referring users to `#688 `__; pass ``reformulate_sos=True`` as a workaround. +* ``Model.solve(..., reformulate_sos=True)`` now actually reformulates SOS constraints even when the solver supports them natively. Previously it was silently ignored with a warning. +* ``add_piecewise_formulation`` now produces a reproducible dimension order in the broadcast breakpoint array. The previous set-based expansion gave a hash-randomized order that varied between processes. +* ``Variable.fix(value)`` now places ``value`` correctly fix binary variables and correctly work on variables with named dimensions; previously array values could be misaligned. + +**Breaking Changes** + +* ``result.solution.primal`` and ``result.solution.dual`` are now ``numpy`` arrays indexed by linopy's integer labels (with ``NaN`` for slots without a value), instead of pandas Series keyed by variable/constraint name. If you accessed them by name, use ``model.variables[name].solution`` (or ``model.constraints[name].dual``) instead. +* ``Model.solver_model`` and ``Model.solver_name`` are now read-only properties that delegate to ``model.solver``. You can't reassign them (only ``= None`` is allowed, which closes the solver), and ``solver_name`` is ``None`` before the first solve. +* ``available_solvers`` now lists all *installed* solvers, even ones without a working license. If you used it to decide "can I actually solve with X?", switch to ``linopy.licensed_solvers`` or ``SolverClass.license_status()``. +* Drop Python 3.10 support. Minimum supported version is now Python 3.11. +* ``add_variables`` / ``add_constraints``: the v0.6.3 ``mask`` deprecations (#580) are now hard ``ValueError``\ s; an unnamed ``pd.MultiIndex`` in sequence-form ``coords`` raises ``TypeError`` unless paired with ``dims=[i]``. See Bug Fixes above. +* Sequence-form ``coords`` entries can no longer be ``xarray.DataArray`` objects — they raise ``TypeError``. Pass the underlying index instead: ``variable.indexes[dim]`` (a ``pd.Index``). + +**Internal** + +* New module ``linopy.alignment`` owns conversion, broadcasting, and alignment of user input against coordinates (moved out of ``linopy.common``): ``as_dataarray`` (convert only), ``broadcast_to_coords`` (convert and broadcast against ``coords``; ``strict=True`` by default raises on any mismatch, naming ``label`` in the error), ``validate_alignment``, and ``align``. +* Each ``Solver`` subclass now overrides at most three hooks: ``_build_direct`` (build the native model), ``_run_direct`` (run it), and ``_run_file`` (run the solver on an LP/MPS file). File-only solvers (CBC, GLPK, CPLEX, SCIP, Knitro, COPT, MindOpt) only override ``_run_file``. +* New ``ConstraintLabelIndex`` cached on ``Model.constraints`` (mirrors the existing ``Variables.label_index``); ``ConstraintBase`` gains ``active_labels()`` and a ``range`` property; ``CSRConstraint`` exposes ``coords``. +* ``linopy.common`` gains ``values_to_lookup_array``; the legacy pandas-based helpers ``series_to_lookup_array`` and ``lookup_vals`` are removed. +* ``model.to_gurobipy()`` / ``model.to_highspy()`` / ``to_cupdlpx(model)`` (and similar) all return the underlying solver model as before; internally they now go through ``Solver.from_model(model, io_api="direct")``. No user-visible change. +* Adopt Python 3.11 type-syntax: the status enums (``ModelStatus``, ``SolverStatus``, ``TerminationCondition``) are now ``StrEnum``, and classmethods plus the expression base class use ``Self`` instead of string forward-references and a self-typed ``TypeVar``. No user-visible change — ``Model.solve()`` still returns ``(status, termination_condition)`` as plain strings. +* ``Variable.fix()`` now fixes a variable by collapsing its bounds (``lower = upper = value``) instead of adding a ``__fix__`` equality constraint; ``unfix()`` restores the original bounds (`#769 `_). A fix outside the current bounds now warns and overrides instead of raising, and its shadow price appears as the variable's reduced cost rather than a constraint dual. + +Version 0.7.0 +------------- + +**Features** + +*Piecewise linear constraints (new)* + +* ``Model.add_piecewise_formulation((power, x_pts), (fuel, y_pts))`` adds piecewise constraints with SOS2, incremental, disjunctive, or pure-LP formulations and automatic method dispatch. Supports N-variable linking (e.g. CHP) and per-entity breakpoints; emits :class:`linopy.EvolvingAPIWarning` while the API stabilises. +* One-sided bounds: append ``"<="`` / ``">="`` to a tuple, e.g. ``(fuel, y_pts, "<=")``. On matching convex/concave curves this dispatches to a pure-LP chord formulation. +* Unit-commitment gating via ``active``: when zero, deactivates the piecewise relation. +* ``PiecewiseFormulation`` exposes ``.method`` / ``.convexity`` (persisted across netCDF round-trip). +* Construction helpers: ``linopy.breakpoints()``, ``linopy.segments()``, ``linopy.Slopes`` for per-piece slopes, and ``tangent_lines()``. + +*Variables* + +* ``fix()`` / ``unfix()`` / ``fixed`` for fixing variables to values via equality constraints (rounds integers/binaries). +* ``relax()`` / ``unrelax()`` / ``relaxed`` for LP relaxation; supports partial relaxation (e.g. ``m.variables.integers.relax()``). +* Semi-continuous variables on solvers that support them. + +*Model* + +* ``Model.copy()`` for a deep copy of a model, optionally including the solution; supports the ``copy`` protocol. +* SOS1 / SOS2 reformulations for solvers without native SOS, applied automatically by ``Model.solve()`` when needed. +* ``format_labels()`` / ``format_infeasibilities()`` return strings instead of printing; deprecates the ``print_*`` siblings. + +*Expressions* + +* Coordinate alignment between subset/superset operands: missing coords fill with 0 in arithmetic and NaN in comparisons. Fixes ``subset + var`` reverse-addition and result coords expanding past the variable's space. + +*Solvers* + +* OETC: ``Model.solve()`` forwards solver options to the handler; ``OetcSettings.from_env()`` reads ``OETC_*``. +* SCIP supports quadratic problems on Windows. + +**Performance** + +* Faster solution unpacking in ``Model.solve()``. + +**Bug Fixes** + +* ``Model.solve()`` raises a clear ``ValueError`` when no objective is set. +* ``add_variables`` no longer ignores ``coords`` when ``lower`` / ``upper`` are DataArrays, and handles MultiIndex coords correctly with scalar bounds. +* ``Model.to_netcdf`` no longer fails on the scipy netCDF backend when variables or constraints have MultiIndex coords; level names are now serialised as a JSON string (the legacy list form remains readable). +* CPLEX no longer errors on quality attributes that aren't always available. + +**Breaking Changes** + +* ``google-cloud-storage`` and ``requests`` are now optional. Install ``linopy[oetc]`` to keep the previous behaviour. + + +Version 0.6.7 +------------- + +* Fix Xpress IIS label mapping for masked constraints and add a regression test for matching infeasible coordinates. +* Fix ``Model.compute_infeasibilities`` returning a flattened, deduplicated union of all IIS when Xpress found more than one. The Xpress path now computes a single IIS (via ``firstIIS``), matching the Gurobi path. +* Use ``xarray.Dataset.copy`` instead of constructor for compatibility with the latest xarray version. +* Blacklist highspy 1.14.0 which produces wrong results due to broken presolve and crashes on Windows (`HiGHS#2964 `_). + +Version 0.6.6 +------------- + +* Free the knitro context and compute necessary quantities within linopy. Knitro context is not exposed anymore. + + +Version 0.6.5 +------------- + +* Expose the knitro context to allow for more flexible use of the knitro python API. + + +Version 0.6.4 +-------------- + +* Add support for the `knitro` solver via the knitro python API + +Version 0.6.3 +-------------- + +**Fix Regression** + +* Reinsert broadcasting logic of mask object to be fully compatible with performance improvements in version 0.6.2 using `np.where` instead of `xr.where`. + + +Version 0.6.2 +-------------- + +**Features** + +* 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. + +**Performance** + +* 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 + +**Bug Fixes** + +* Fix multiplication of constant-only ``LinearExpression`` with other expressions +* Fix docs and Gurobi license handling Version 0.6.1 -------------- @@ -657,7 +881,7 @@ Version 0.0.5 * The `Variable` class now has a `lower` and `upper` accessor, which allows to inspect and modify the lower and upper bounds of a assigned variable. * The `Constraint` class now has a `lhs`, `vars`, `coeffs`, `rhs` and `sign` accessor, which allows to inspect and modify the left-hand-side, the signs and right-hand-side of a assigned constraint. * Constraints can now be build combining linear expressions with right-hand-side via a `>=`, `<=` or a `==` operator. This creates an `AnonymousConstraint` which can be passed to `Model.add_constraints`. -* Add support of the HiGHS open source solver https://www.maths.ed.ac.uk/hall/HiGHS/ (https://github.com/PyPSA/linopy/pull/8, https://github.com/PyPSA/linopy/pull/17). +* Add support of the HiGHS open source solver https://highs.dev/ (https://github.com/PyPSA/linopy/pull/8, https://github.com/PyPSA/linopy/pull/17). **Breaking changes** diff --git a/doc/sos-constraints.rst b/doc/sos-constraints.rst index 37dd72d26..caa4b5e56 100644 --- a/doc/sos-constraints.rst +++ b/doc/sos-constraints.rst @@ -75,7 +75,7 @@ Method Signature .. code-block:: python - Model.add_sos_constraints(variable, sos_type, sos_dim) + Model.add_sos_constraints(variable, sos_type, sos_dim, big_m=None) **Parameters:** @@ -85,6 +85,8 @@ Method Signature Type of SOS constraint (1 or 2) - ``sos_dim`` : str Name of the dimension along which the SOS constraint applies +- ``big_m`` : float | None + Custom Big-M value for reformulation (see :ref:`sos-reformulation`) **Requirements:** @@ -254,12 +256,94 @@ SOS constraints are supported by most modern mixed-integer programming solvers t - MOSEK - MindOpt +For unsupported solvers, use automatic reformulation (see below). + +.. _sos-reformulation: + +SOS Reformulation +----------------- + +For solvers without native SOS support, linopy can reformulate SOS constraints +as binary + linear constraints using the Big-M method. + +.. code-block:: python + + # Automatic reformulation during solve + m.solve(solver_name="highs", reformulate_sos=True) + + # Or reformulate manually + m.reformulate_sos_constraints() + m.solve(solver_name="highs") + +**Requirements:** + +- Variables must have **non-negative lower bounds** (lower >= 0) +- Big-M values are derived from variable upper bounds +- For infinite upper bounds, specify custom values via the ``big_m`` parameter + +.. code-block:: python + + # Finite bounds (default) + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + # Infinite upper bounds: specify Big-M + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + +The reformulation uses the tighter of ``big_m`` and variable upper bound. + +Mathematical Formulation +~~~~~~~~~~~~~~~~~~~~~~~~ + +**SOS1 Reformulation:** + +Original constraint: :math:`\text{SOS1}(\{x_1, x_2, \ldots, x_n\})` means at most one +:math:`x_i` can be non-zero. + +Given :math:`x = (x_1, \ldots, x_n) \in \mathbb{R}^n_+`, introduce binary +:math:`y = (y_1, \ldots, y_n) \in \{0,1\}^n`: + +.. math:: + + x_i &\leq M_i \cdot y_i \quad \forall i \in \{1, \ldots, n\} \\ + x_i &\geq 0 \quad \forall i \in \{1, \ldots, n\} \\ + \sum_{i=1}^{n} y_i &\leq 1 \\ + y_i &\in \{0, 1\} \quad \forall i \in \{1, \ldots, n\} + +where :math:`M_i \geq \sup\{x_i\}` (upper bound on :math:`x_i`). + +**SOS2 Reformulation:** + +Original constraint: :math:`\text{SOS2}(\{x_1, x_2, \ldots, x_n\})` means at most two +:math:`x_i` can be non-zero, and if two are non-zero, they must have consecutive indices. + +Given :math:`x = (x_1, \ldots, x_n) \in \mathbb{R}^n_+`, introduce binary +:math:`y = (y_1, \ldots, y_{n-1}) \in \{0,1\}^{n-1}`: + +.. math:: + + x_1 &\leq M_1 \cdot y_1 \\ + x_i &\leq M_i \cdot (y_{i-1} + y_i) \quad \forall i \in \{2, \ldots, n-1\} \\ + x_n &\leq M_n \cdot y_{n-1} \\ + x_i &\geq 0 \quad \forall i \in \{1, \ldots, n\} \\ + \sum_{i=1}^{n-1} y_i &\leq 1 \\ + y_i &\in \{0, 1\} \quad \forall i \in \{1, \ldots, n-1\} + +where :math:`M_i \geq \sup\{x_i\}`. Interpretation: :math:`y_i = 1` activates interval +:math:`[i, i+1]`, allowing :math:`x_i` and :math:`x_{i+1}` to be non-zero. + Common Patterns --------------- Piecewise Linear Cost Function ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. note:: + + For a higher-level API that handles all the SOS2 bookkeeping automatically, + see :doc:`piecewise-linear-constraints`. + .. code-block:: python def add_piecewise_cost(model, variable, breakpoints, costs): diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 494ac2407..8b7ee5bd5 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -3,10 +3,58 @@ Overview ======== -Welcome to the User Guide for Linopy. This guide is designed to help you understand and effectively use Linopy's features to solve your optimization problems, complementing the ``Getting Started`` section. +In :doc:`Getting Started ` you installed linopy, built +a first scalar model, and saw N-D variables on coordinates. The User +Guide reopens each of those pieces in depth and adds the rest of the +modelling surface. -In the following sections, we will take a closer look at how to create and manipulate models, variables, and constraints, and how to solve these models to find optimal solutions. Each section includes detailed explanations and code examples to help you understand the concepts and apply them to your own projects. +Each page is a runnable Jupyter notebook — read it top to bottom, or +use it as a reference once you know what you're looking for. -If you are completely new to Linopy, consider to first have a look at the `Getting Started` section. -Let's get started! +Core building blocks +-------------------- + +The four notebooks below cover the model object you'll interact with +most. Read them in order the first time; come back to them whenever +you're unsure what a particular operator or argument does. + +- :doc:`creating-variables` — declaring decision variables, with bounds + and coordinates. Continuous, integer, binary, and semi-continuous. +- :doc:`creating-expressions` — combining variables into linear (and + quadratic) expressions; arithmetic, broadcasting, ``sum``, + ``groupby``, ``rolling``, ``where``. +- :doc:`creating-constraints` — turning expressions into ``≤`` / ``≥`` + / ``==`` constraints, and the ``CSRConstraint`` memory-efficient + alternative. +- :doc:`coordinate-alignment` — how linopy lines up operands that live + on different coordinates, and how to control it with ``join``. + +After these four you can build any LP/MIP/QP linopy supports. + + +Working with an existing model +------------------------------ + +Once you've built a model, you'll often want to inspect it, change a +bound, swap a constraint, or copy it for what-if analysis. + +- :doc:`manipulating-models` — modifying or removing variables and + constraints in place; ``Model.copy()``; ``fix`` / ``relax`` for + variables. + + +Where to go next +---------------- + +- **Examples** — end-to-end problem walkthroughs: + :doc:`transport-tutorial`, :doc:`migrating-from-pyomo`. +- **Advanced features** — :doc:`sos-constraints`, + :doc:`piecewise-linear-constraints`, and the + :doc:`testing-framework` for asserting structural properties of a + model. +- **Solving** — :doc:`solve-on-remote` (SSH), + :doc:`solve-on-oetc` (OET Cloud), :doc:`gpu-acceleration` (cuPDLPx). +- **Troubleshooting** — :doc:`infeasible-model` (diagnosing infeasible + problems), :doc:`gurobi-double-logging` (and other solver quirks). +- **Reference** — the full :doc:`api` listing. diff --git a/examples/coordinate-alignment.ipynb b/examples/coordinate-alignment.ipynb new file mode 100644 index 000000000..54de243a5 --- /dev/null +++ b/examples/coordinate-alignment.ipynb @@ -0,0 +1,512 @@ +{ + "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", + "\n", + "When 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", + "\n", + "To 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", + "\n", + "Consider 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/create-a-model-with-coordinates.ipynb b/examples/create-a-model-with-coordinates.ipynb index f2b12eed0..56c374811 100644 --- a/examples/create-a-model-with-coordinates.ipynb +++ b/examples/create-a-model-with-coordinates.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "4db583af", + "id": "0", "metadata": {}, "source": [ "# Use Coordinates\n", @@ -16,7 +16,7 @@ }, { "cell_type": "markdown", - "id": "comparable-talent", + "id": "1", "metadata": {}, "source": [ "Minimize:\n", @@ -36,7 +36,7 @@ }, { "cell_type": "markdown", - "id": "proprietary-receipt", + "id": "2", "metadata": {}, "source": [ "In order to formulate the new problem with linopy, we start again by initializing a model." @@ -45,7 +45,7 @@ { "cell_type": "code", "execution_count": null, - "id": "close-maximum", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -56,7 +56,7 @@ }, { "cell_type": "markdown", - "id": "positive-appearance", + "id": "4", "metadata": {}, "source": [ "Again, we define `x` and `y` using the `add_variables` function, but now we are adding a `coords` argument. This automatically creates optimization variables for all coordinates, in this case time-steps." @@ -65,7 +65,7 @@ { "cell_type": "code", "execution_count": null, - "id": "included-religious", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -83,7 +83,7 @@ }, { "cell_type": "markdown", - "id": "terminal-ethernet", + "id": "6", "metadata": {}, "source": [ "Following the previous example, we write the constraints out using the syntax from above, while multiplying the rhs with `t`. Note that the coordinates from the lhs and the rhs have to match. \n", @@ -95,7 +95,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c24d120a", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -106,7 +106,7 @@ }, { "cell_type": "markdown", - "id": "f09803f4", + "id": "8", "metadata": {}, "source": [ "It always helps to write out the constraints before adding them to the model. Since they look good, let's assign them." @@ -115,7 +115,7 @@ { "cell_type": "code", "execution_count": null, - "id": "comprehensive-blend", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -126,7 +126,7 @@ }, { "cell_type": "markdown", - "id": "induced-professor", + "id": "10", "metadata": {}, "source": [ "Now, when it comes to the objective, we use the `sum` function of `linopy.LinearExpression`. This stacks all terms all terms of the `time` dimension and writes them into one big expression. " @@ -135,7 +135,7 @@ { "cell_type": "code", "execution_count": null, - "id": "alternate-story", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -146,16 +146,16 @@ { "cell_type": "code", "execution_count": null, - "id": "outer-presence", + "id": "12", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")" + "m.solve(solver_name=\"highs\", output_flag=False)" ] }, { "cell_type": "markdown", - "id": "495cd082", + "id": "13", "metadata": {}, "source": [ "In order to inspect the solution. You can go via the variables, i.e. `y.solution` or via the `solution` aggregator of the model, which combines the solution of all variables. This can sometimes be helpful." @@ -164,7 +164,7 @@ { "cell_type": "code", "execution_count": null, - "id": "monthly-census", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -173,11 +173,21 @@ }, { "cell_type": "markdown", - "id": "owned-europe", + "id": "15", "metadata": {}, "source": [ "Alright! Now you learned how to set up linopy variables and expressions with coordinates. In the User Guide, which follows, we are going to see, how the representation of variables with coordinates allows us to formulate more advanced operations." ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## Where to next\n", + "\n", + "You've now seen the full path from declaring variables on coordinates to solving the model. The [User Guide overview](user-guide.rst) reopens each piece in depth and points you at every topic from here." + ] } ], "metadata": { diff --git a/examples/create-a-model.ipynb b/examples/create-a-model.ipynb index a158e0cf1..943029d6b 100644 --- a/examples/create-a-model.ipynb +++ b/examples/create-a-model.ipynb @@ -3,7 +3,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "4db583af", + "id": "0", "metadata": {}, "source": [ "# Solve a Basic Model\n", @@ -14,7 +14,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "together-ocean", + "id": "1", "metadata": {}, "source": [ "Minimize:\n", @@ -31,7 +31,7 @@ { "cell_type": "code", "execution_count": null, - "id": "dramatic-cannon", + "id": "2", "metadata": {}, "outputs": [], "source": [] @@ -39,7 +39,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "43949d36", + "id": "3", "metadata": {}, "source": [ "### Initializing a `Model`\n", @@ -50,7 +50,7 @@ { "cell_type": "code", "execution_count": null, - "id": "technical-conducting", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -62,7 +62,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "e5b16d53", + "id": "5", "metadata": {}, "source": [ "This creates a new Model object, which you can then use to define your optimization problem." @@ -71,7 +71,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "rolled-delicious", + "id": "6", "metadata": {}, "source": [ "\n", @@ -84,7 +84,7 @@ { "cell_type": "code", "execution_count": null, - "id": "protecting-power", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -95,7 +95,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "featured-maria", + "id": "8", "metadata": {}, "source": [ "`x` and `y` are linopy variables of the class `linopy.Variable`. Each of them contain all relevant information that define it. The `name` parameter is optional but can be useful for referencing the variables later." @@ -104,7 +104,7 @@ { "cell_type": "code", "execution_count": null, - "id": "virtual-anxiety", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -114,7 +114,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "sonic-rebate", + "id": "10", "metadata": {}, "source": [ "Since both `x` and `y` are scalar variables (meaning they don't have any dimensions), their underlying data contain only one optimization variable each. \n", @@ -128,7 +128,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fbb46cad", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -138,7 +138,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "f4666bee", + "id": "12", "metadata": {}, "source": [ "Note, we can also mix the constant and the variable expression, like this" @@ -147,7 +147,7 @@ { "cell_type": "code", "execution_count": null, - "id": "60f41b76", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -157,7 +157,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "02abd938", + "id": "14", "metadata": {}, "source": [ "... and linopy will automatically take over the separation of variables expression on the lhs, and constant values on the rhs.\n", @@ -168,7 +168,7 @@ { "cell_type": "code", "execution_count": null, - "id": "hollywood-production", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -179,7 +179,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "global-maple", + "id": "16", "metadata": {}, "source": [ "## Adding the Objective \n", @@ -190,7 +190,7 @@ { "cell_type": "code", "execution_count": null, - "id": "overall-exhibition", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -200,7 +200,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "6f9692aa", + "id": "18", "metadata": {}, "source": [ "## Solving the Model\n", @@ -211,17 +211,17 @@ { "cell_type": "code", "execution_count": null, - "id": "pressing-copying", + "id": "19", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")" + "m.solve(solver_name=\"highs\", output_flag=False)" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "preceding-limit", + "id": "20", "metadata": {}, "source": [ "The solution of the linear problem assigned to the variables under `solution` in form of a `xarray.Dataset`. " @@ -230,7 +230,7 @@ { "cell_type": "code", "execution_count": null, - "id": "electric-duration", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -240,7 +240,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e6d31751", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -250,7 +250,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "e296f641", + "id": "23", "metadata": {}, "source": [ "Well done! You solved your first linopy model!" diff --git a/examples/creating-constraints.ipynb b/examples/creating-constraints.ipynb index b46db1bcb..d504deb34 100644 --- a/examples/creating-constraints.ipynb +++ b/examples/creating-constraints.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "e8249281", + "id": "0", "metadata": {}, "source": [ "# Creating Constraints\n", @@ -18,7 +18,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e0c196e4", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -29,7 +29,7 @@ }, { "cell_type": "markdown", - "id": "043c0b06", + "id": "2", "metadata": {}, "source": [ "Given a variable `x` which has to by lower than 10/3, the constraint would be formulated as \n", @@ -51,7 +51,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6b496b92", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -60,7 +60,7 @@ }, { "cell_type": "markdown", - "id": "73541c03", + "id": "4", "metadata": {}, "source": [ "When applying one of the operators `<=`, `>=`, `==` to the expression, an unassigned constraint is built:" @@ -69,7 +69,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4c8aba7e", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -79,7 +79,7 @@ }, { "cell_type": "markdown", - "id": "0d75781d", + "id": "6", "metadata": {}, "source": [ "Unasssigned means, it is not yet added to the model. We can inspect the elements of the anonymous constraint: " @@ -88,7 +88,7 @@ { "cell_type": "code", "execution_count": null, - "id": "01f182b5", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -98,7 +98,7 @@ { "cell_type": "code", "execution_count": null, - "id": "783287b3", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -107,7 +107,7 @@ }, { "cell_type": "markdown", - "id": "aac468c3", + "id": "9", "metadata": {}, "source": [ "We can now add the constraint to the model by passing the unassigned `Constraint` to the `.add_constraint` function. " @@ -116,7 +116,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0adf929b", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -126,7 +126,7 @@ }, { "cell_type": "markdown", - "id": "e78c2635", + "id": "11", "metadata": {}, "source": [ "The same output would be generated if passing lhs, sign and rhs as separate arguments to the function:" @@ -135,7 +135,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c084adec", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -144,7 +144,7 @@ }, { "cell_type": "markdown", - "id": "2b4db4d5", + "id": "13", "metadata": {}, "source": [ "Note that the return value of the operation is a `Constraint` which contains the reference labels to the constraints in the optimization model. Also is redirects to its lhs, sign and rhs, for example we can call" @@ -153,7 +153,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ea6e990c", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -162,7 +162,7 @@ }, { "cell_type": "markdown", - "id": "e6ae2a19", + "id": "15", "metadata": {}, "source": [ "to inspect the lhs of a defined constraint." @@ -170,7 +170,7 @@ }, { "cell_type": "markdown", - "id": "efb74da3", + "id": "16", "metadata": {}, "source": [ "When moving the constant value to the left hand side in the initialization, it will be pulled to the right hand side as soon as the constraint is defined" @@ -179,7 +179,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e582051e", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -189,7 +189,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e2c2dbb3", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -198,7 +198,7 @@ }, { "cell_type": "markdown", - "id": "15909055", + "id": "19", "metadata": {}, "source": [ "Like this, the all defined constraints have a clear separation between variable on the left, and constants on the right. " @@ -206,7 +206,7 @@ }, { "cell_type": "markdown", - "id": "b9d31509", + "id": "20", "metadata": {}, "source": [ "All constraints are added to the `.constraints` container from where all assigned constraints can be accessed." @@ -215,7 +215,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d205e695", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -225,12 +225,163 @@ { "cell_type": "code", "execution_count": null, - "id": "cc5baaf4", + "id": "22", "metadata": {}, "outputs": [], "source": [ "m.constraints[\"my-constraint\"]" ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": {}, + "source": [ + "## Coordinate Alignment in Constraints\n", + "\n", + "As 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." + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "## CSR Backend (Advanced)\n", + "\n", + "By default, linopy stores each constraint as an `xarray.Dataset` (`Constraint`). This is flexible and allows full label-based indexing, but can use significant memory when constraints have many terms.\n", + "\n", + "For large models, linopy provides an alternative **CSR backend** via the `CSRConstraint` class. It stores the constraint coefficients as a [scipy CSR sparse matrix](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_array.html) with flat numpy arrays for the right-hand side and signs. This can reduce memory usage by up to **90%** and speeds up matrix generation for direct solver APIs by **30–120x**.\n", + "\n", + "`CSRConstraint` is **immutable** — once frozen, the constraint data cannot be modified in place. You can always convert back to the mutable xarray-backed `Constraint` if needed." + ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "### Freezing individual constraints\n", + "\n", + "Pass `freeze=True` to `add_constraints` to store a single constraint in CSR format:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "from linopy import Model\n", + "\n", + "m2 = Model()\n", + "y = m2.add_variables(coords=[np.arange(100)], name=\"y\")\n", + "\n", + "m2.add_constraints(y <= 10, name=\"upper\", freeze=True)\n", + "\n", + "print(type(m2.constraints[\"upper\"]))\n", + "m2.constraints[\"upper\"]" + ] + }, + { + "cell_type": "markdown", + "id": "27", + "metadata": {}, + "source": [ + "### Freezing all constraints globally\n", + "\n", + "Set `freeze_constraints=True` on the `Model` to automatically freeze every constraint added via `add_constraints`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "m3 = Model(freeze_constraints=True)\n", + "z = m3.add_variables(coords=[np.arange(50)], name=\"z\")\n", + "m3.add_constraints(z >= 0, name=\"lower\")\n", + "m3.add_constraints(z <= 100, name=\"upper\")\n", + "\n", + "print(type(m3.constraints[\"lower\"]))\n", + "print(type(m3.constraints[\"upper\"]))" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "### Converting between representations\n", + "\n", + "Use `.freeze()` and `.mutable()` to convert between the two representations. The conversion is lossless:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "frozen = m3.constraints[\"lower\"]\n", + "print(f\"Frozen type: {type(frozen).__name__}\")\n", + "\n", + "thawed = frozen.mutable()\n", + "print(f\"Mutable type: {type(thawed).__name__}\")\n", + "\n", + "refrozen = thawed.freeze()\n", + "print(f\"Re-frozen type: {type(refrozen).__name__}\")" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": {}, + "source": [ + "### API differences from `Constraint`\n", + "\n", + "`CSRConstraint` deliberately exposes a narrower API than the xarray-backed `Constraint`:\n", + "\n", + "- **No in-place mutation.** `Constraint.update(...)` is only available on `Constraint`. (The legacy setters — `con.coeffs = ...`, `con.vars = ...`, `con.sign = ...`, `con.rhs = ...`, `con.lhs = ...` — still forward to `update` on `Constraint` but emit a `DeprecationWarning` and will be removed in a future release.)\n", + "- **No label-based indexing.** `con.loc[...]` is only available on `Constraint`.\n", + "- **Accessing `.coeffs` / `.vars` triggers reconstruction.** On a `CSRConstraint` these properties rebuild the full xarray `Dataset` on demand and emit a `PerformanceWarning`. For solver-oriented workflows prefer `con.to_matrix()` or work with the CSR data directly.\n", + "\n", + "If you need any of the above, call `.mutable()` first to get a `Constraint`:\n", + "\n", + "```python\n", + "con = m.constraints[\"my_constraint\"].mutable()\n", + "con.loc[{\"time\": 0}] # label-based indexing now available\n", + "con.update(rhs=5) # mutation now available\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "32", + "metadata": {}, + "source": [ + "### When to use the CSR backend\n", + "\n", + "The CSR backend is most beneficial when:\n", + "\n", + "- Your model has **many constraints with many terms**.\n", + "- **Memory** is a bottleneck.\n", + "- You use a **direct solver API** (e.g. HiGHS, Gurobi Python bindings) rather than file-based I/O.\n", + "\n", + "For small models the overhead is negligible and the default xarray-backed `Constraint` is perfectly fine.\n", + "\n", + "Additionally, if you don't need variable and constraint names in the solver (e.g. for batch solves), you can disable name export for extra speed:\n", + "\n", + "```python\n", + "m = Model(freeze_constraints=True, set_names_in_solver_io=False)\n", + "```" + ] } ], "metadata": { diff --git a/examples/creating-expressions.ipynb b/examples/creating-expressions.ipynb index aafd8a09d..cb41a2c66 100644 --- a/examples/creating-expressions.ipynb +++ b/examples/creating-expressions.ipynb @@ -3,7 +3,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "4db583af", + "id": "0", "metadata": {}, "source": [ "# Creating Expressions\n", @@ -25,7 +25,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "f3712718", + "id": "1", "metadata": {}, "source": [ ".. hint::\n", @@ -35,7 +35,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "95b25d13", + "id": "2", "metadata": {}, "source": [ "Let's start by creating a model." @@ -44,7 +44,7 @@ { "cell_type": "code", "execution_count": null, - "id": "close-maximum", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -65,7 +65,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "582a0cad", + "id": "4", "metadata": {}, "source": [ "## Arithmetic Operations\n", @@ -78,7 +78,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0aec195a", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -89,7 +89,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "8c4b9ea6", + "id": "6", "metadata": {}, "source": [ ".. note::\n", @@ -99,7 +99,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "a1de2b9a", + "id": "7", "metadata": {}, "source": [ "Similarly, you can subtract `y` from `x` or multiply `x` and `y` as follows:" @@ -108,7 +108,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c56761cd", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -119,7 +119,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b59fa397", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -130,7 +130,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "0ca00a73", + "id": "10", "metadata": {}, "source": [ "In all cases, the returned shape is the same. Note that, the output type of the multiplication is a `QuadraticExpression` and not a `LinearExpression`.\n" @@ -139,7 +139,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "3ff0d1cd", + "id": "11", "metadata": {}, "source": [ "The `z` expression, which carries along `x` and `y`, has different attributes such as `coord_dims`, `dims`, `size`." @@ -148,7 +148,7 @@ { "cell_type": "code", "execution_count": null, - "id": "35c7331f", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -158,14 +158,17 @@ { "attachments": {}, "cell_type": "markdown", - "id": "f7578221", + "id": "13", "metadata": {}, - "source": ".. important::\n\n\tWhen combining variables or expression with dimensions of the same name and size, the first object will determine the coordinates of the resulting expression. For example:" + "source": [ + ".. important::\n", + " When combining variables or expressions with dimensions of the same name and size, the first object determines the coordinates of the resulting expression. For example:" + ] }, { "cell_type": "code", "execution_count": null, - "id": "8c511f35", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -177,7 +180,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "59f43468", + "id": "15", "metadata": {}, "source": [ "`b` has the same shape as `x`, but they have different coordinates. When we combine `x` and `b` the coordinates on dimension `time` will be taken from the first object and the coordinates of the subsequent object will be ignored:" @@ -186,17 +189,26 @@ { "cell_type": "code", "execution_count": null, - "id": "26edd6ab", + "id": "16", "metadata": {}, "outputs": [], "source": [ "x + b" ] }, + { + "cell_type": "markdown", + "id": "17", + "metadata": {}, + "source": [ + ".. tip::\n", + " For 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." + ] + }, { "attachments": {}, "cell_type": "markdown", - "id": "de6d3073", + "id": "18", "metadata": {}, "source": [ "## Using `.loc` to select a subset\n", @@ -209,7 +221,7 @@ { "cell_type": "code", "execution_count": null, - "id": "93119cfc", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -219,7 +231,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ae6a1b29", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -229,7 +241,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "14b02f7d", + "id": "21", "metadata": {}, "source": [ "which is the same as" @@ -238,7 +250,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7281f08c", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -249,7 +261,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "c3c97abf", + "id": "23", "metadata": {}, "source": [ "In combination with the overwrite of the coordinates, this is useful when you need to combine different selections, like" @@ -258,7 +270,7 @@ { "cell_type": "code", "execution_count": null, - "id": "27063ea9", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -268,7 +280,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "cdb53f49", + "id": "25", "metadata": {}, "source": [ "## Using `.where` to select active variables or expressions\n", @@ -281,7 +293,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ab8f59fd", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -292,7 +304,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "3ee10a58", + "id": "27", "metadata": {}, "source": [ "We can use this to make a conditional summation:" @@ -301,17 +313,113 @@ { "cell_type": "code", "execution_count": null, - "id": "21fa9664", + "id": "28", "metadata": {}, "outputs": [], "source": [ "(x + y).where(mask) + xr.DataArray(5, coords=[time]).where(~mask, 0)" ] }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "Sometimes `.where` may lead to a situation where some of the variables are completely masked" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "mask_a = xr.DataArray(False, coords=[time])\n", + "mask_b = xr.DataArray(time > 2, coords=[time])\n", + "\n", + "z = (x.where(mask_a) + y).where(mask_b)\n", + "z" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": {}, + "source": [ + "In this example you can see that many of the elements of the LinearExpression are None. If you want to remove all the None terms, you can use `.where(.., drop=True)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": {}, + "outputs": [], + "source": [ + "z = z.where(mask_b, drop=True)\n", + "z" + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": {}, + "source": [ + "That looks nicer!
" + ] + }, + { + "cell_type": "markdown", + "id": "34", + "metadata": {}, + "source": [ + "You may notice that the variable `x` is not used at all. The expression still contains two terms (one of them is unused) but it only has one variable `y`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35", + "metadata": {}, + "outputs": [], + "source": [ + "z.nterm" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "z.variable_names" + ] + }, + { + "cell_type": "markdown", + "id": "37", + "metadata": {}, + "source": [ + "You can get rid of the unused term with `.simplify()`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "38", + "metadata": {}, + "outputs": [], + "source": [ + "z = z.simplify()\n", + "z.nterm" + ] + }, { "attachments": {}, "cell_type": "markdown", - "id": "652973ea", + "id": "39", "metadata": {}, "source": [ "## Using `.shift` to shift the Variable along one dimension\n", @@ -324,7 +432,7 @@ { "cell_type": "code", "execution_count": null, - "id": "organized-hampshire", + "id": "40", "metadata": {}, "outputs": [], "source": [ @@ -334,7 +442,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "eaf3f38c", + "id": "41", "metadata": {}, "source": [ "## Using `.groupby` to group by a key and apply operations on the groups\n", @@ -347,7 +455,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5170d187", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -358,7 +466,7 @@ { "attachments": {}, "cell_type": "markdown", - "id": "7ded9a54", + "id": "43", "metadata": {}, "source": [ "## Using `.rolling` to perform a rolling operation\n", @@ -371,7 +479,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d703fb70", + "id": "44", "metadata": {}, "outputs": [], "source": [ diff --git a/examples/creating-variables.ipynb b/examples/creating-variables.ipynb index 8e8793481..1c8b53d06 100644 --- a/examples/creating-variables.ipynb +++ b/examples/creating-variables.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "e8249281", + "id": "0", "metadata": {}, "source": [ "# Creating Variables\n", @@ -18,7 +18,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e0c196e4", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -33,7 +33,13 @@ }, { "cell_type": "markdown", - "id": "6c6420a7", + "id": "2", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "id": "3", "metadata": {}, "source": [ "First of all it is crucial to know, that the return value of the `.add_variables` function is a `linopy.Variable` which itself contains all important information and provides helpful functions. It can have an arbitrary number of labeled dimensions. For each combination of coordinates, exactly one representative scalar variable is defined and, in the end, passed to the solver. \n", @@ -63,7 +69,7 @@ }, { "cell_type": "markdown", - "id": "a2283b9a", + "id": "4", "metadata": {}, "source": [ "Let's start by creating a simple variable:\n", @@ -74,7 +80,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ee589323", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -84,7 +90,7 @@ }, { "cell_type": "markdown", - "id": "708920a3", + "id": "6", "metadata": {}, "source": [ "which is a variable without any coordinates and with just one optimization variable. The variable name is set by `name = 'x'`. " @@ -92,7 +98,7 @@ }, { "cell_type": "markdown", - "id": "b276e45d", + "id": "7", "metadata": {}, "source": [ "Like this the variable appears with its name when defining expression with it:" @@ -101,7 +107,7 @@ { "cell_type": "code", "execution_count": null, - "id": "68d7e7da", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -110,7 +116,7 @@ }, { "cell_type": "markdown", - "id": "528a7c40", + "id": "9", "metadata": {}, "source": [ "We can alter the lower and upper bounds of the variable by assigning scalar values to them." @@ -119,7 +125,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a1c080e0", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -128,7 +134,7 @@ }, { "cell_type": "markdown", - "id": "885ac764", + "id": "11", "metadata": {}, "source": [ "### Variable Types\n", @@ -139,7 +145,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8a5d4543", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -148,7 +154,7 @@ }, { "cell_type": "markdown", - "id": "3af58bc4", + "id": "13", "metadata": {}, "source": [ ".. note::\n", @@ -161,7 +167,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ff64db81", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -170,7 +176,7 @@ }, { "cell_type": "markdown", - "id": "7b432107", + "id": "15", "metadata": {}, "source": [ "### Working with dimensions\n", @@ -181,7 +187,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b4dfc46d", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -192,7 +198,7 @@ }, { "cell_type": "markdown", - "id": "1ff347f9", + "id": "17", "metadata": {}, "source": [ "The returned `Variable` now has the same shape as the `lower` bound that we passed to the initialization. Since we did not specify any dimension name, it defaults to `dim_0`. In order to give the dimension a proper name we can use the `dims` argument. " @@ -201,7 +207,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f93e5c08", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -211,7 +217,7 @@ }, { "cell_type": "markdown", - "id": "d4b20bb5", + "id": "19", "metadata": {}, "source": [ "You can arbitrarily broadcast dimensions when passing DataArray's with different set of dimensions. Let's do it and give `lower` another dimension than `upper`:" @@ -220,7 +226,7 @@ { "cell_type": "code", "execution_count": null, - "id": "71584630", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -231,7 +237,7 @@ }, { "cell_type": "markdown", - "id": "3e4e48c0", + "id": "21", "metadata": {}, "source": [ "Now instead of a single dimension, we end up with two dimensions `my-dim` and `my-dim-2` in the variable. This kind of **broadcasting** is a deeply incorporated in the functionality of linopy. " @@ -239,7 +245,7 @@ }, { "cell_type": "markdown", - "id": "41893f11", + "id": "22", "metadata": {}, "source": [ "We recall that, in order to improve the inspection, it is encouraged to define a `name` when creating a variable. So in your model you would rather write something like:" @@ -248,7 +254,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e8857233", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -259,7 +265,7 @@ }, { "cell_type": "markdown", - "id": "437cc7a1", + "id": "24", "metadata": {}, "source": [ "#### Initializing variables with numpy arrays\n", @@ -270,10 +276,8 @@ { "cell_type": "code", "execution_count": null, - "id": "0fe33c34", - "metadata": { - "scrolled": true - }, + "id": "25", + "metadata": {}, "outputs": [], "source": [ "lower = np.array([1, 2])\n", @@ -283,7 +287,7 @@ }, { "cell_type": "markdown", - "id": "2ab6d301", + "id": "26", "metadata": {}, "source": [ "This is equivalent to the following" @@ -292,7 +296,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b8313ce0", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -304,7 +308,7 @@ }, { "cell_type": "markdown", - "id": "5052b9b5", + "id": "28", "metadata": {}, "source": [ "Note that \n", @@ -323,7 +327,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5f0994da", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -333,7 +337,7 @@ }, { "cell_type": "markdown", - "id": "dff8126d", + "id": "30", "metadata": {}, "source": [ "The dimension is now called `dim_0`, any new assignment of variable without dimension names, will also use that dimension name. When combining the variables to expressions it is important that you make sure that dimension names represent what they should. \n", @@ -345,7 +349,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d133a7a4", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -358,7 +362,7 @@ }, { "cell_type": "markdown", - "id": "9203ff16", + "id": "32", "metadata": {}, "source": [ "#### Initializing variables with Pandas objects\n", @@ -369,7 +373,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2cf719be", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -380,7 +384,7 @@ }, { "cell_type": "markdown", - "id": "3a4cf2d4", + "id": "34", "metadata": {}, "source": [ "or naming the indexes and columns of the pandas objects directly, e.g." @@ -389,7 +393,7 @@ { "cell_type": "code", "execution_count": null, - "id": "61896a6f", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -400,7 +404,7 @@ }, { "cell_type": "markdown", - "id": "2b462ff9", + "id": "36", "metadata": {}, "source": [ ".. note::\n", @@ -410,10 +414,8 @@ { "cell_type": "code", "execution_count": null, - "id": "6ffd5a4e", - "metadata": { - "scrolled": true - }, + "id": "37", + "metadata": {}, "outputs": [], "source": [ "lower = pd.Series([1, 1]).rename_axis(\"my-dim\")\n", @@ -423,7 +425,7 @@ }, { "cell_type": "markdown", - "id": "31bbdbab", + "id": "38", "metadata": {}, "source": [ "Now instead of 2 variables, 4 variables were defined. \n", @@ -434,7 +436,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fa2adc81", + "id": "39", "metadata": {}, "outputs": [], "source": [ @@ -445,7 +447,7 @@ }, { "cell_type": "markdown", - "id": "8b1734df", + "id": "40", "metadata": {}, "source": [ "Again, one is always safer when explicitly naming the dimensions:" @@ -454,7 +456,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21f7db15", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -465,14 +467,17 @@ }, { "cell_type": "markdown", - "id": "77e264e2", + "id": "42", "metadata": {}, - "source": ".. important::\n\n **New in version 0.3.6**\n\n As pandas objects always have indexes, the `coords` argument is not required and is ignored is passed. Before, it was used to overwrite the indexes of the pandas objects. A warning is raised if `coords` is passed and if these are not aligned with the pandas object." + "source": [ + ".. note::\n", + " As pandas objects already carry indexes, the ``coords`` argument is ignored when supplied alongside them. A warning is raised if a ``coords`` value is passed that does not align with the pandas object." + ] }, { "cell_type": "code", "execution_count": null, - "id": "d0fc67cf", + "id": "43", "metadata": {}, "outputs": [], "source": [ @@ -482,7 +487,7 @@ }, { "cell_type": "markdown", - "id": "49de1cc3", + "id": "44", "metadata": {}, "source": [ "### Masking Arrays\n", @@ -497,7 +502,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9d802903", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -513,7 +518,7 @@ { "cell_type": "code", "execution_count": null, - "id": "447d8a8a", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -525,7 +530,7 @@ }, { "cell_type": "markdown", - "id": "df1c551c", + "id": "47", "metadata": {}, "source": [ "Now the diagonal values, for example at the variable at [a,a], are `None`. " @@ -533,7 +538,7 @@ }, { "cell_type": "markdown", - "id": "23a040d4", + "id": "48", "metadata": {}, "source": [ "### Accessing assigned variables\n", @@ -544,7 +549,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2946a80c", + "id": "49", "metadata": {}, "outputs": [], "source": [ @@ -553,7 +558,7 @@ }, { "cell_type": "markdown", - "id": "45cf0755", + "id": "50", "metadata": {}, "source": [ "You can always access the variables from the `.variables` container either by get-item, i.e." @@ -562,7 +567,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d974727d", + "id": "51", "metadata": {}, "outputs": [], "source": [ @@ -571,7 +576,7 @@ }, { "cell_type": "markdown", - "id": "03182836", + "id": "52", "metadata": {}, "source": [ "or by attribute accessing" @@ -580,7 +585,7 @@ { "cell_type": "code", "execution_count": null, - "id": "308b879a", + "id": "53", "metadata": {}, "outputs": [], "source": [ diff --git a/examples/infeasible-model.ipynb b/examples/infeasible-model.ipynb index 7dae15181..8766ac785 100644 --- a/examples/infeasible-model.ipynb +++ b/examples/infeasible-model.ipynb @@ -32,7 +32,10 @@ "\n", "m.add_constraints(x <= 5)\n", "m.add_constraints(y <= 5)\n", - "m.add_constraints(x + y >= 12)" + "m.add_constraints(x + y >= 12)\n", + "\n", + "# A trivial objective is required; the model is solved purely to check feasibility.\n", + "m.add_objective(0 * x)" ] }, { @@ -122,8 +125,7 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.3" - }, - "orig_nbformat": 4 + } }, "nbformat": 4, "nbformat_minor": 2 diff --git a/examples/manipulating-models.ipynb b/examples/manipulating-models.ipynb index 1a35cd19d..bd86399ea 100644 --- a/examples/manipulating-models.ipynb +++ b/examples/manipulating-models.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "37a85c22", + "id": "0", "metadata": {}, "source": [ "# Modifying Models\n", @@ -15,7 +15,7 @@ { "cell_type": "code", "execution_count": null, - "id": "16a41836", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -28,7 +28,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8f4d182f", + "id": "2", "metadata": {}, "outputs": [], "source": [ @@ -48,16 +48,16 @@ "con2 = m.add_constraints(5 * x + 2 * y >= 3 * factor, name=\"con2\")\n", "\n", "m.add_objective(x + 2 * y)\n", - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "\n", - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] }, { "cell_type": "markdown", - "id": "d3f4b966", + "id": "3", "metadata": {}, "source": [ "The figure above shows the optimal values of `x(t)` and `y(t)`. \n", @@ -70,20 +70,23 @@ { "cell_type": "code", "execution_count": null, - "id": "f7db57f8", + "id": "4", "metadata": {}, "outputs": [], "source": [ - "x.lower = 1" + "x.update(lower=1)" ] }, { "cell_type": "markdown", - "id": "66b8be86", + "id": "5", "metadata": {}, "source": [ ".. note::\n", - " The same could have been achieved by calling `m.variables.x.lower = 1`\n", + " Assignment via the ``x.lower = 1`` setter still works but is\n", + " deprecated and will be removed in a future release. Use\n", + " ``Variable.update`` instead — it is the canonical mutation API\n", + " with a single validation path.\n", "\n", "Let's solve it again!" ] @@ -91,11 +94,11 @@ { "cell_type": "code", "execution_count": null, - "id": "c37add87", + "id": "6", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] @@ -103,7 +106,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b5be8d00", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -112,7 +115,7 @@ }, { "cell_type": "markdown", - "id": "d35e3309", + "id": "8", "metadata": {}, "source": [ "We see that the new lower bound of x is binding across all time steps.\n", @@ -123,28 +126,28 @@ { "cell_type": "code", "execution_count": null, - "id": "451aba93", + "id": "9", "metadata": {}, "outputs": [], "source": [ - "x.lower = xr.DataArray(range(10, 0, -1), coords=(time,))" + "x.update(lower=xr.DataArray(range(10, 0, -1), coords=(time,)))" ] }, { "cell_type": "code", "execution_count": null, - "id": "e25f26a1", + "id": "10", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] }, { "cell_type": "markdown", - "id": "4d991939", + "id": "11", "metadata": {}, "source": [ "You can manipulate the upper bound of a variable in the same way." @@ -152,33 +155,39 @@ }, { "cell_type": "markdown", - "id": "de29c28e", + "id": "12", "metadata": {}, "source": [ "## Varying Constraints\n", "\n", - "A similar functionality is implemented for constraints. Here we can modify the left-hand-side, the sign and the right-hand-side.\n", + "A similar functionality is implemented for constraints. We use\n", + "``Constraint.update`` to change the left-hand-side, the sign,\n", + "and the right-hand-side.\n", "\n", - "Assume we want to relax the right-hand-side of the first constraint `con1` to `8 * factor`. This would translate to:" + "Assume we want to relax the right-hand-side of the first constraint\n", + "``con1`` to ``8 * factor``. This translates to:" ] }, { "cell_type": "code", "execution_count": null, - "id": "18d1bf4b", + "id": "13", "metadata": {}, "outputs": [], "source": [ - "con1.rhs = 8 * factor" + "con1.update(rhs=8 * factor)" ] }, { "cell_type": "markdown", - "id": "5499b3b4", + "id": "14", "metadata": {}, "source": [ ".. note::\n", - " The same could have been achieved by calling `m.constraints.con1.rhs = 8 * factor`\n", + " Assignment via the ``con1.rhs = 8 * factor`` setter still works\n", + " but is deprecated and will be removed in a future release. Use\n", + " ``Constraint.update`` instead — it is the canonical mutation API\n", + " with a single validation path.\n", "\n", "Let's solve it again!" ] @@ -186,18 +195,18 @@ { "cell_type": "code", "execution_count": null, - "id": "e4d34142", + "id": "15", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] }, { "cell_type": "markdown", - "id": "bc683e13", + "id": "16", "metadata": {}, "source": [ "In contrast to previous figure, we now see that the optimal value of `y` does not reach values above 10 in the end. \n", @@ -208,28 +217,34 @@ { "cell_type": "code", "execution_count": null, - "id": "f8e81d20", + "id": "17", "metadata": {}, "outputs": [], "source": [ - "con1.lhs = 3 * x + 8 * y" + "con1.update(lhs=3 * x + 8 * y)" ] }, { "cell_type": "markdown", - "id": "cc377d95", + "id": "18", "metadata": {}, "source": [ "**Note:**\n", - "The same could have been achieved by calling \n", - "```python \n", - "m.constraints['con1'].lhs = 3 * x + 8 * y\n", + "Assignment via the ``con1.lhs = 3 * x + 8 * y`` setter still works\n", + "but is deprecated and will be removed in a future release. Use\n", + "``Constraint.update`` instead — it is the canonical mutation API\n", + "with a single validation path.\n", + "\n", + "``Constraint.update`` also accepts a full constraint expression in one call:\n", + "\n", + "```python\n", + "con1.update(3 * x + 8 * y <= 8 * factor) # replaces lhs / sign / rhs at once\n", "```" ] }, { "cell_type": "markdown", - "id": "633d463b", + "id": "19", "metadata": {}, "source": [ "which leads to" @@ -238,18 +253,18 @@ { "cell_type": "code", "execution_count": null, - "id": "9b73250d", + "id": "20", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] }, { "cell_type": "markdown", - "id": "e509d5d7", + "id": "21", "metadata": {}, "source": [ "## Varying the objective \n", @@ -262,7 +277,7 @@ { "cell_type": "code", "execution_count": null, - "id": "44689b5b", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -272,18 +287,18 @@ { "cell_type": "code", "execution_count": null, - "id": "2144af8e", + "id": "23", "metadata": {}, "outputs": [], "source": [ - "m.solve(solver_name=\"highs\")\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", "sol = m.solution.to_dataframe()\n", "sol.plot(grid=True, ylabel=\"Optimal Value\")" ] }, { "cell_type": "markdown", - "id": "f1faa095", + "id": "24", "metadata": {}, "source": [ "As a consequence, `y` stays at zero for all time steps." @@ -292,12 +307,87 @@ { "cell_type": "code", "execution_count": null, - "id": "85cbd60b", + "id": "25", "metadata": {}, "outputs": [], "source": [ "m.objective" ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "## Fixing Variables and Extracting MILP Duals\n", + "\n", + "A common workflow in mixed-integer programming is to solve the MILP, then fix the integer/binary variables to their optimal values and re-solve as an LP to obtain dual values (shadow prices).\n", + "\n", + "Let's extend our model with a binary variable `z` that activates an additional capacity constraint on `x`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "z = m.add_variables(binary=True, coords=[time], name=\"z\")\n", + "\n", + "# x can only exceed 5 when z is active: x <= 5 + 100 * z\n", + "m.add_constraints(x <= 5 + 100 * z, name=\"capacity\")\n", + "\n", + "# Penalize activation of z in the objective\n", + "m.objective = x + 3 * y + 10 * z\n", + "\n", + "m.solve(solver_name=\"highs\", output_flag=False)" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "Now fix the binary variable `z` to its optimal values and relax its integrality. This converts the model into an LP, which allows us to extract dual values." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "m.variables.binaries.fix()\n", + "m.variables.binaries.relax()\n", + "m.solve(solver_name=\"highs\", output_flag=False)\n", + "\n", + "# Dual values are now available on the constraints\n", + "m.constraints[\"con1\"].dual" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": {}, + "source": [ + "Calling `unfix()` on all variables restores their original bounds and `unrelax()` restores the integrality of `z`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "m.variables.unfix()\n", + "m.variables.unrelax()\n", + "\n", + "# z is binary again\n", + "m.variables[\"z\"].attrs[\"binary\"]" + ] } ], "metadata": { diff --git a/examples/migrating-from-pyomo.ipynb b/examples/migrating-from-pyomo.ipynb index 3d34ce600..c3535a40f 100644 --- a/examples/migrating-from-pyomo.ipynb +++ b/examples/migrating-from-pyomo.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "ded90143", + "id": "0", "metadata": {}, "source": [ "## Migrating from Pyomo\n", @@ -13,7 +13,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19f3b954", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -30,7 +30,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2bbfd13b", + "id": "2", "metadata": {}, "outputs": [], "source": [ @@ -39,7 +39,7 @@ }, { "cell_type": "markdown", - "id": "a1631a76", + "id": "3", "metadata": {}, "source": [ ".. important::\n", @@ -53,7 +53,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4ed6eafb", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -70,7 +70,7 @@ }, { "cell_type": "markdown", - "id": "4faecead", + "id": "5", "metadata": {}, "source": [ "Note that the function's first argument has to be the model itself, even though it might not be used in the function." @@ -78,7 +78,7 @@ }, { "cell_type": "markdown", - "id": "d7368607", + "id": "6", "metadata": {}, "source": [ "This functionality is also supported by the `.add_constraints` function. When passing a function as a first argument, `.add_constraints` expects `coords` to by non-empty. The function itself has to return a `AnonymousScalarConstraint`, as done by " @@ -87,7 +87,7 @@ { "cell_type": "code", "execution_count": null, - "id": "eeebb710", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -97,7 +97,7 @@ { "cell_type": "code", "execution_count": null, - "id": "087203ad", + "id": "8", "metadata": {}, "outputs": [], "source": [ diff --git a/examples/piecewise-inequality-bounds.ipynb b/examples/piecewise-inequality-bounds.ipynb new file mode 100644 index 000000000..59ec23b9a --- /dev/null +++ b/examples/piecewise-inequality-bounds.ipynb @@ -0,0 +1,266 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Creating Piecewise Inequality Bounds\n", + "\n", + "When you only need a one-sided bound from a piecewise curve — `y ≤ f(x)` for a concave upper envelope, `y ≥ f(x)` for a convex lower envelope — `add_piecewise_formulation` accepts an optional sign as the third tuple element:\n", + "\n", + "```python\n", + "m.add_piecewise_formulation(\n", + " (fuel, fuel_pts, \">=\"), # fuel ≥ f(power) — over-fuelling admissible\n", + " (power, power_pts), # equality role\n", + ")\n", + "```\n", + "\n", + "The pay-off is a pure-LP encoding when the curve's curvature matches the sign — no SOS2, no binaries. This notebook covers the geometry of the feasible region, the curvature × sign combinations that unlock the LP path, and what happens when they don't match.\n", + "\n", + "For the formulation math see the [reference page](piecewise-linear-constraints.rst); for the all-equality variant and other features see [Creating Piecewise Linear Constraints](piecewise-linear-constraints-tutorial.nblink).\n", + "\n", + "## Tuple roles\n", + "\n", + "| Tuple form | Role | What it constrains |\n", + "|---|---|---|\n", + "| `(expr, breaks)` | `==` (equality) | With 2+ equality tuples sharing weights, the joint point lies on the curve. With 1 equality + 1 bounded, the equality tuple's marginal feasible set is just its breakpoint domain `[x_min, x_max]` — one coordinate alone can't locate a curve point. |\n", + "| `(expr, breaks, \"<=\")` | bounded above | `expr ≤ f(other tuples)`. |\n", + "| `(expr, breaks, \">=\")` | bounded below | `expr ≥ f(other tuples)`. |\n", + "\n", + "Currently at most one tuple may carry a non-equality sign, and 3+ tuples must all be equality. Multi-bounded and N≥3 inequality cases aren't supported yet — if you have a concrete use case, please open an issue at https://github.com/PyPSA/linopy/issues so we can scope it properly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import linopy\n", + "\n", + "# Silence the evolving-API warning for cleaner tutorial output.\n", + "warnings.filterwarnings(\"ignore\", category=linopy.EvolvingAPIWarning)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup — a convex heat-rate curve\n", + "\n", + "A convex, monotonically increasing curve maps power output to the fuel required (the classic heat-rate curve). Bounding `fuel` by this curve with `>=` says the unit must consume *at least* the design fuel for a given power output — over-fuelling is physically admissible but wasteful, so an objective that minimises fuel pulls the operating point onto the curve. Convex + `>=` is exactly the combination that lets the LP method apply.\n", + "\n", + "The breakpoint arrays:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "power_pts = np.array([0.0, 30.0, 60.0, 100.0])\n", + "fuel_pts = np.array([0.0, 36.0, 84.0, 170.0]) # slopes 1.2, 1.6, 2.15 (convex)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Three methods, identical feasible region\n", + "\n", + "With `fuel` bounded `>=` and our convex curve, the three methods give the **same** feasible region for `power ∈ [0, 100]`:\n", + "\n", + "- **`method=\"lp\"`** — tangent lines + domain bounds. No auxiliary variables.\n", + "- **`method=\"sos2\"`** — lambdas + SOS2 + a split link: equality for the equality-signed tuple, signed for the bounded one. Solver picks the piece.\n", + "- **`method=\"incremental\"`** — delta fractions + binaries + split link. Same mathematics, MIP encoding instead of SOS2.\n", + "\n", + "`method=\"auto\"` dispatches to `\"lp\"` whenever applicable — it's always preferable because it's pure LP.\n", + "\n", + "Let's verify they produce the same solution at `power=60`, where `f(60)=84`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def solve(method, power_val):\n", + " m = linopy.Model()\n", + " power = m.add_variables(lower=0, upper=100, name=\"power\")\n", + " fuel = m.add_variables(lower=0, upper=200, name=\"fuel\")\n", + " m.add_piecewise_formulation(\n", + " (fuel, fuel_pts, \">=\"), # fuel ≥ f(power) — over-fuelling admissible\n", + " (power, power_pts), # equality role (domain-bounded to [0, 100])\n", + " method=method,\n", + " )\n", + " m.add_constraints(power == power_val)\n", + " m.add_objective(fuel) # minimise fuel against the lower bound\n", + " m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + " return float(m.solution[\"fuel\"]), list(m.variables), list(m.constraints)\n", + "\n", + "\n", + "for method in [\"lp\", \"sos2\", \"incremental\"]:\n", + " fuel_val, vars_, cons_ = solve(method, 60)\n", + " print(f\"{method:12}: fuel={fuel_val:.2f} vars={vars_} cons={cons_}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All three give `fuel=84` at `power=60` (which is `f(60)` exactly) — the math is equivalent. The LP method is strictly cheaper: no auxiliary variables, just three chord constraints and two domain bounds.\n", + "\n", + "The SOS2 and incremental methods create lambdas (or deltas + binaries) and split the link into an equality constraint for the equality-signed tuple plus a signed link for the bounded tuple — but the feasible region is the same." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visualising the feasible region\n", + "\n", + "The feasible region for `(power, fuel)` with `fuel` bounded `>=` is the **epigraph** of `f` restricted to the power domain:\n", + "\n", + "$$\\{ (\\mathrm{power}, \\mathrm{fuel}) : 0 \\le \\mathrm{power} \\le 100,\\ \\mathrm{fuel} \\ge f(\\mathrm{power}) \\}$$\n", + "\n", + "Below we colour feasible points green, infeasible ones red:\n", + "\n", + "- `(60, 100)` — above the curve, `100 ≥ f(60)=84` ✓\n", + "- `(60, 84)` — on the curve ✓\n", + "- `(60, 70)` — below `f(60)`, infeasible ✗\n", + "- `(120, 100)` — power beyond domain, infeasible ✗" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def in_epigraph(px, fy):\n", + " if px < power_pts[0] or px > power_pts[-1]:\n", + " return False\n", + " return fy >= np.interp(px, power_pts, fuel_pts)\n", + "\n", + "\n", + "xx, yy = np.meshgrid(np.linspace(-10, 130, 200), np.linspace(-10, 200, 200))\n", + "region = np.vectorize(in_epigraph)(xx, yy)\n", + "\n", + "test_points = [(60, 100), (60, 84), (60, 70), (120, 100)]\n", + "\n", + "fig, ax = plt.subplots(figsize=(6, 5))\n", + "ax.contourf(xx, yy, region, levels=[0.5, 1], colors=[\"lightsteelblue\"], alpha=0.5)\n", + "ax.plot(power_pts, fuel_pts, \"o-\", color=\"C0\", lw=2, label=\"f(power)\")\n", + "for px, fy in test_points:\n", + " feas = in_epigraph(px, fy)\n", + " ax.scatter(\n", + " [px], [fy], color=\"green\" if feas else \"red\", zorder=5, s=80, edgecolors=\"black\"\n", + " )\n", + " ax.annotate(f\"({px}, {fy})\", (px, fy), textcoords=\"offset points\", xytext=(8, 5))\n", + "ax.set(\n", + " xlabel=\"power\",\n", + " ylabel=\"fuel\",\n", + " title=\"sign='>=' feasible region — epigraph of f(power) on [0, 100]\",\n", + ")\n", + "ax.grid(alpha=0.3)\n", + "ax.legend()\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## When is LP the right choice?\n", + "\n", + "`tangent_lines` imposes the **intersection** of chord inequalities. Whether that intersection matches the true hypograph/epigraph of `f` depends on the curvature × sign combination:\n", + "\n", + "| curvature | bounded `<=` | bounded `>=` |\n", + "|-----------|--------------|--------------|\n", + "| **concave** | **hypograph (exact ✓)** | **wrong region** — requires `y ≥ max_k chord_k(x) > f(x)` |\n", + "| **convex** | **wrong region** — requires `y ≤ min_k chord_k(x) < f(x)` | **epigraph (exact ✓)** |\n", + "| linear | exact | exact |\n", + "| mixed (non-convex) | convex hull of `f` (wrong for exact hypograph) | concave hull of `f` (wrong for exact epigraph) |\n", + "\n", + "In the ✗ cases, tangent lines do **not** give a loose relaxation — they give a **strictly wrong feasible region** that rejects points satisfying the true constraint. Example: for a concave `f` with `y ≥ f(x)`, the chord of any piece extrapolated over another piece's x-range lies *above* `f`, so `y ≥ max_k chord_k(x)` forbids `y = f(x)` itself.\n", + "\n", + "`method=\"auto\"` dispatches to LP only in the two **exact** cases (concave + `<=` or convex + `>=`). For the other combinations it falls back to SOS2 or incremental, which encode the hypograph/epigraph exactly via discrete piece selection.\n", + "\n", + "`method=\"lp\"` explicitly forces LP and raises on a mismatched curvature rather than silently producing a wrong feasible region.\n", + "\n", + "For **non-convex** curves with either sign, the only exact option is a piecewise formulation. That's what the bounded-tuple path does internally: it falls back to SOS2/incremental with the sign on the bounded link. No relaxation, no wrong bounds." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 1. Non-convex curve: auto falls back (LP relaxation would be loose)\n", + "power_nc = [0, 30, 60, 100]\n", + "fuel_nc = [0, 50, 30, 80] # slopes change sign → mixed convexity\n", + "\n", + "m1 = linopy.Model()\n", + "power1 = m1.add_variables(lower=0, upper=100, name=\"power\")\n", + "fuel1 = m1.add_variables(lower=0, upper=200, name=\"fuel\")\n", + "f1 = m1.add_piecewise_formulation((fuel1, fuel_nc, \">=\"), (power1, power_nc))\n", + "print(f\"non-convex + '>=' → {f1.method}\")\n", + "\n", + "# 2. Convex curve + sign='<=': LP would be loose → auto falls back to MIP\n", + "m2 = linopy.Model()\n", + "power2 = m2.add_variables(lower=0, upper=100, name=\"power\")\n", + "fuel2 = m2.add_variables(lower=0, upper=200, name=\"fuel\")\n", + "f2 = m2.add_piecewise_formulation(\n", + " (fuel2, list(fuel_pts), \"<=\"), (power2, list(power_pts))\n", + ")\n", + "print(f\"convex + '<=' → {f2.method}\")\n", + "\n", + "# 3. Explicit method=\"lp\" with mismatched curvature raises\n", + "try:\n", + " m3 = linopy.Model()\n", + " power3 = m3.add_variables(lower=0, upper=100, name=\"power\")\n", + " fuel3 = m3.add_variables(lower=0, upper=200, name=\"fuel\")\n", + " m3.add_piecewise_formulation(\n", + " (fuel3, list(fuel_pts), \"<=\"), (power3, list(power_pts)), method=\"lp\"\n", + " )\n", + "except ValueError as e:\n", + " print(f\"lp(convex, '<=') → raises: {e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "- **One bounded tuple + a 2-variable formulation** gives a hypograph (`<=`) or epigraph (`>=`) feasible region.\n", + "- **Curvature × sign matching** — concave + `<=` or convex + `>=` — lets `method=\"auto\"` skip MIP entirely. Mismatched combinations fall back to SOS2/incremental with a signed link.\n", + "- **`method=\"lp\"` is strict** — it raises on a mismatched curvature rather than silently encoding the wrong region.\n", + "- At most one tuple may carry a non-`==` sign, and 3+ tuples must all be `==`. Multi-bounded / N≥3 inequalities — open an issue at https://github.com/PyPSA/linopy/issues.\n", + "\n", + "**See also**: [reference page](piecewise-linear-constraints.rst) for the formulation math, [Creating Piecewise Linear Constraints](piecewise-linear-constraints-tutorial.nblink) for all-equality, unit commitment, CHP, fleets, slopes." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/piecewise-linear-constraints.ipynb b/examples/piecewise-linear-constraints.ipynb new file mode 100644 index 000000000..ef26a55b1 --- /dev/null +++ b/examples/piecewise-linear-constraints.ipynb @@ -0,0 +1,446 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Creating Piecewise Linear Constraints\n", + "\n", + "`add_piecewise_formulation` links variables through a shared piecewise-linear curve. Pair each variable with its breakpoint values; the solver puts every variable on the *same* point of the curve at every feasible solution.\n", + "\n", + "```python\n", + "m.add_piecewise_formulation(\n", + " (power, [0, 30, 60, 100]),\n", + " (fuel, [0, 36, 84, 170]),\n", + ")\n", + "```\n", + "\n", + "This tutorial walks through the main features of `add_piecewise_formulation`. For the formulation math see the [reference page](piecewise-linear-constraints.rst); for the inequality variant in depth see [Creating Piecewise Inequality Bounds](piecewise-inequality-bounds-tutorial.nblink).\n", + "\n", + "**Roadmap**\n", + "\n", + "1. Getting started — the basic 2-variable equality.\n", + "2. Picking a method — `\"auto\"`, `\"sos2\"`, `\"incremental\"`, `\"lp\"`.\n", + "3. Disjunctive segments — disconnected operating regions with `segments()`.\n", + "4. Inequality bounds — `<=` / `>=` per-tuple sign.\n", + "5. Unit commitment — gating with `active=...`.\n", + "6. N-variable linking — CHP plants and beyond.\n", + "7. Per-entity breakpoints — fleets with different curves.\n", + "8. Specifying with slopes — `linopy.Slopes`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "import xarray as xr\n", + "\n", + "import linopy\n", + "\n", + "# Silence the evolving-API warning for cleaner tutorial output.\n", + "warnings.filterwarnings(\"ignore\", category=linopy.EvolvingAPIWarning)\n", + "\n", + "time = pd.Index([1, 2, 3], name=\"time\")\n", + "\n", + "\n", + "def plot_curve(\n", + " bp_x, bp_y, operating_x, operating_y, *, ax=None, xlabel=\"power\", ylabel=\"fuel\"\n", + "):\n", + " \"\"\"PWL curve with solver's operating points overlaid.\"\"\"\n", + " ax = ax if ax is not None else plt.subplots(figsize=(4.5, 3.5))[1]\n", + " ax.plot(bp_x, bp_y, \"o-\", color=\"C0\", label=\"breakpoints\")\n", + " ax.plot(operating_x, operating_y, \"D\", color=\"C1\", ms=10, label=\"solved\", alpha=0.8)\n", + " ax.set(xlabel=xlabel, ylabel=ylabel)\n", + " ax.legend()\n", + " return ax" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Getting started\n", + "\n", + "A gas turbine with a convex heat rate. Each `(variable, breakpoints)` tuple pairs a variable with its breakpoint values. All tuples share interpolation weights, so at any feasible point every variable corresponds to the *same* point on the curve." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "demand = xr.DataArray([50, 80, 30], coords=[time])\n", + "\n", + "m = linopy.Model()\n", + "power = m.add_variables(name=\"power\", lower=0, upper=100, coords=[time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + "\n", + "x_pts = [0, 30, 60, 100]\n", + "y_pts = [0, 36, 84, 170]\n", + "pwf = m.add_piecewise_formulation((power, x_pts), (fuel, y_pts))\n", + "m.add_constraints(power == demand, name=\"demand\")\n", + "m.add_objective(fuel.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "\n", + "print(pwf) # inspect the auto-resolved method\n", + "m.solution[[\"power\", \"fuel\"]].to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_curve(x_pts, y_pts, m.solution[\"power\"].values, m.solution[\"fuel\"].values);" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Picking a method\n", + "\n", + "`method=\"auto\"` (default) picks the cheapest correct formulation based on sign, curvature and breakpoint layout. The explicit options are `\"sos2\"`, `\"incremental\"`, and `\"lp\"`; the choice is about **cost** (auxiliary variables, solver capability), not correctness — on cases where they all apply they give the same optimum.\n", + "\n", + "For now: a quick sanity check that all applicable methods yield the same fuel dispatch on the convex curve from §1.\n", + "\n", + "A full comparison — when each method dispatches, what sign/curvature/breakpoint patterns each requires — lives in §3 (disjunctive), §4 (inequalities), and the [reference page's \"Formulation Methods\" section](piecewise-linear-constraints.rst)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def solve_method(method):\n", + " m = linopy.Model()\n", + " power = m.add_variables(name=\"power\", lower=0, upper=100, coords=[time])\n", + " fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + " m.add_piecewise_formulation((power, x_pts), (fuel, y_pts), method=method)\n", + " m.add_constraints(power == demand, name=\"demand\")\n", + " m.add_objective(fuel.sum())\n", + " m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + " return m.solution[\"fuel\"].to_pandas()\n", + "\n", + "\n", + "pd.DataFrame({m: solve_method(m) for m in [\"auto\", \"sos2\", \"incremental\"]})" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Disjunctive segments — discrete operating bands\n", + "\n", + "Some equipment has **disjoint operating ranges** rather than a continuous one. A stepped pump has two speed bands with a forbidden zone between them — the pump physically can't operate in that gap. `segments()` models this directly: one segment per band, a binary picks exactly one per operating point.\n", + "\n", + "Below: two pumps in parallel, each with a low band (5–25 m³/h) and a high band (40–100 m³/h). Demands that land in the single-pump gap or above its maximum force the optimiser to combine bands across the two pumps.\n", + "\n", + "(For an on/off gate on a single continuous curve, use `active=...` instead; see §5.)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pumps = pd.Index([\"p1\", \"p2\"], name=\"pump\")\n", + "\n", + "m = linopy.Model()\n", + "flow = m.add_variables(name=\"flow\", lower=0, upper=100, coords=[pumps, time])\n", + "power = m.add_variables(name=\"power\", lower=0, coords=[pumps, time])\n", + "\n", + "# Each pump has two operating bands; the gap between them is a forbidden zone.\n", + "m.add_piecewise_formulation(\n", + " (flow, linopy.segments([(5, 25), (40, 100)])),\n", + " (power, linopy.segments([(1, 7), (15, 50)])),\n", + ")\n", + "m.add_constraints(flow.sum(\"pump\") == xr.DataArray([30, 75, 150], coords=[time]))\n", + "m.add_objective(power.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "\n", + "# Flat columns: flow_p1, flow_p2, power_p1, power_p2 per timestep.\n", + "sol = m.solution[[\"flow\", \"power\"]].to_dataframe().unstack(\"pump\")\n", + "sol.columns = [f\"{var}_{p}\" for var, p in sol.columns]\n", + "sol" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Every timestep is single-pump-infeasible:\n", + "\n", + "- *t=1*, demand=30: in the single-pump gap (25, 40). Both pumps run in low band, splitting the load.\n", + "- *t=2*, demand=75: too much for low+low (max 50), too little for high+high (min 80). The low pump tops out at 25 m³/h; the high pump covers the remaining 50.\n", + "- *t=3*, demand=150: above a single pump's maximum (100). Both pumps run in high band." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Inequality bounds — per-tuple sign\n", + "\n", + "Append a third tuple element (`\"<=\"` or `\">=\"`) to mark a single expression as **bounded** by the curve instead of entering as an equality. The 2-variable hypograph (`y ≤ f(x)`) and epigraph (`y ≥ f(x)`) are the canonical cases.\n", + "\n", + "On a concave curve with `<=` (or convex with `>=`), `method=\"auto\"` dispatches to a pure LP chord formulation — no binaries, no SOS2.\n", + "\n", + "Below: the same heat-rate curve as §1, now read as `fuel ≥ f(power)` — over-fuelling is admissible but wasteful (the curve is the design minimum), so an objective that minimises fuel pulls the operating point onto the curve. See [Creating Piecewise Inequality Bounds](piecewise-inequality-bounds-tutorial.nblink) for mismatched curvature, auto-dispatch fallbacks, and more geometry." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "power = m.add_variables(name=\"power\", lower=0, upper=100, coords=[time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + "\n", + "# Same convex heat-rate curve as §1, now bounded with \">=\"\n", + "pwf = m.add_piecewise_formulation(\n", + " (fuel, [0, 36, 84, 170], \">=\"), # fuel ≥ f(power) — over-fuelling allowed\n", + " (power, [0, 30, 60, 100]), # equality role\n", + ")\n", + "m.add_constraints(power == demand)\n", + "m.add_objective(fuel.sum()) # minimise fuel against the lower bound\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "\n", + "print(f\"resolved method={pwf.method}, curvature={pwf.convexity}\")\n", + "m.solution[[\"power\", \"fuel\"]].to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_curve(\n", + " [0, 30, 60, 100],\n", + " [0, 36, 84, 170],\n", + " m.solution[\"power\"].values,\n", + " m.solution[\"fuel\"].values,\n", + ");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Unit commitment — `active`\n", + "\n", + "A binary variable gates the whole formulation. `active=0` forces the PWL variables (and thus all linked outputs) to zero. Combined with the natural `lower=0` on cost/fuel/heat, this gives a clean on/off coupling:\n", + "\n", + "- `active=1`: the unit operates in its full range, outputs tied to the curve.\n", + "- `active=0`: `power = 0`, `fuel = 0`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "p_min, p_max = 30, 100\n", + "\n", + "power = m.add_variables(name=\"power\", lower=0, upper=p_max, coords=[time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + "backup = m.add_variables(name=\"backup\", lower=0, coords=[time])\n", + "commit = m.add_variables(name=\"commit\", binary=True, coords=[time])\n", + "\n", + "x_pts = [p_min, 60, p_max]\n", + "y_pts = [40, 90, 170]\n", + "m.add_piecewise_formulation(\n", + " (power, x_pts),\n", + " (fuel, y_pts),\n", + " active=commit,\n", + ")\n", + "# demand below p_min at t=1 — commit must be 0 and backup covers it\n", + "m.add_constraints(power + backup == xr.DataArray([15, 80, 40], coords=[time]))\n", + "m.add_objective(fuel.sum() + 50 * commit.sum() + 200 * backup.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "m.solution[[\"commit\", \"power\", \"fuel\", \"backup\"]].to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_curve(x_pts, y_pts, m.solution[\"power\"].values, m.solution[\"fuel\"].values);" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. N-variable linking — CHP plant\n", + "\n", + "More than two variables can share the same interpolation — useful for combined heat-and-power plants where power, fuel and heat are all functions of a single operating point." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "power = m.add_variables(name=\"power\", lower=0, upper=100, coords=[time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + "heat = m.add_variables(name=\"heat\", lower=0, coords=[time])\n", + "\n", + "x_pts = [0, 30, 60, 100]\n", + "y_pts = [0, 40, 85, 160]\n", + "z_pts = [0, 25, 55, 95]\n", + "m.add_piecewise_formulation(\n", + " (power, x_pts),\n", + " (fuel, y_pts),\n", + " (heat, z_pts),\n", + ")\n", + "m.add_constraints(fuel == xr.DataArray([20, 100, 160], coords=[time]))\n", + "m.add_objective(power.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "m.solution[[\"power\", \"fuel\", \"heat\"]].to_pandas().round(2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 2, figsize=(8, 3))\n", + "plot_curve(\n", + " x_pts, y_pts, m.solution[\"power\"].values, m.solution[\"fuel\"].values, ax=axes[0]\n", + ")\n", + "plot_curve(\n", + " x_pts,\n", + " z_pts,\n", + " m.solution[\"power\"].values,\n", + " m.solution[\"heat\"].values,\n", + " ylabel=\"heat\",\n", + " ax=axes[1],\n", + ");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Per-entity breakpoints — a fleet of generators\n", + "\n", + "Pass a dict to `breakpoints()` with entity names as keys for different curves per entity. Ragged lengths are NaN-padded automatically, and breakpoints broadcast over any remaining dimensions (here, `time`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gens = pd.Index([\"gas\", \"coal\"], name=\"gen\")\n", + "x_gen = linopy.breakpoints(\n", + " {\"gas\": [0, 30, 60, 100], \"coal\": [0, 50, 100, 150]}, dim=\"gen\"\n", + ")\n", + "y_gen = linopy.breakpoints(\n", + " {\"gas\": [0, 40, 90, 180], \"coal\": [0, 55, 130, 225]}, dim=\"gen\"\n", + ")\n", + "\n", + "m = linopy.Model()\n", + "power = m.add_variables(name=\"power\", lower=0, upper=150, coords=[gens, time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[gens, time])\n", + "m.add_piecewise_formulation((power, x_gen), (fuel, y_gen))\n", + "m.add_constraints(power.sum(\"gen\") == xr.DataArray([80, 120, 50], coords=[time]))\n", + "m.add_objective(fuel.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "m.solution[[\"power\", \"fuel\"]].to_dataframe()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Specifying with slopes — `Slopes`\n", + "\n", + "When marginal costs (slopes) are more natural than absolute y-values, wrap them in `linopy.Slopes`. The x grid is borrowed from the sibling tuple — no need to repeat it. Same curve as section 1:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "power = m.add_variables(name=\"power\", lower=0, upper=100, coords=[time])\n", + "fuel = m.add_variables(name=\"fuel\", lower=0, coords=[time])\n", + "\n", + "m.add_piecewise_formulation(\n", + " (power, [0, 30, 60, 100]),\n", + " (fuel, linopy.Slopes([1.2, 1.6, 2.15], y0=0)),\n", + ")\n", + "m.add_constraints(power == demand, name=\"demand\")\n", + "m.add_objective(fuel.sum())\n", + "m.solve(solver_name=\"highs\", reformulate_sos=\"auto\", output_flag=False)\n", + "\n", + "m.solution[[\"power\", \"fuel\"]].to_pandas()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## When to use what\n", + "\n", + "| Pattern | API |\n", + "|---|---|\n", + "| `y` is a function of `x` | `(x, x_pts), (y, y_pts)` — all-equality |\n", + "| `y` bounded by `f(x)` on a convex/concave curve | `(y, y_pts, \"<=\"` or `\">=\"), (x, x_pts)` — LP if curvature matches |\n", + "| Disconnected operating regions | `linopy.segments(...)` per tuple |\n", + "| Unit on/off coupling | `active=binary_var` |\n", + "| Multiple synchronized outputs (e.g. CHP) | 3+ tuples, all `\"==\"` |\n", + "| Different curves per entity | `linopy.breakpoints({...}, dim=...)` |\n", + "| Slopes more natural than absolute y-values | `linopy.Slopes(...)` |\n", + "\n", + "For the formulation math, see the [reference page](piecewise-linear-constraints.rst). For inequality bounds in depth, see [Creating Piecewise Inequality Bounds](piecewise-inequality-bounds-tutorial.nblink)." + ] + } + ], + "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.13.2" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/solve-on-oetc.ipynb b/examples/solve-on-oetc.ipynb index 7459bdb96..28e1c04d8 100644 --- a/examples/solve-on-oetc.ipynb +++ b/examples/solve-on-oetc.ipynb @@ -30,6 +30,13 @@ "All of these steps are handled automatically by linopy's `OetcHandler`." ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **Note:** This notebook requires Google Cloud credentials and access to the OETC platform. It is not executed during the documentation build, so no cell outputs are shown. To run it yourself, install the `linopy[oetc]` extra and configure your credentials." + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -79,7 +86,12 @@ "source": [ "## Configure OETC Settings\n", "\n", - "Next, we need to configure the OETC settings including credentials and compute requirements:" + "There are two ways to configure OETC settings:\n", + "\n", + "1. **Manual construction** — build `OetcCredentials` and `OetcSettings` explicitly\n", + "2. **`OetcSettings.from_env()`** — resolve credentials and options from environment variables\n", + "\n", + "### Option 1: Manual Construction" ] }, { @@ -123,6 +135,48 @@ "print(f\"Disk space: {settings.disk_space_gb} GB\")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Option 2: Create Settings from Environment Variables\n", + "\n", + "`OetcSettings.from_env()` reads configuration from environment variables,\n", + "with optional keyword overrides. This is the recommended approach for\n", + "CI/CD pipelines and production deployments.\n", + "\n", + "| Environment Variable | Required | Description |\n", + "|---|---|---|\n", + "| `OETC_EMAIL` | Yes | Account email |\n", + "| `OETC_PASSWORD` | Yes | Account password |\n", + "| `OETC_NAME` | Yes | Job name |\n", + "| `OETC_AUTH_URL` | Yes | Authentication server URL |\n", + "| `OETC_ORCHESTRATOR_URL` | Yes | Orchestrator server URL |\n", + "| `OETC_CPU_CORES` | No | CPU cores (default: 2) |\n", + "| `OETC_DISK_SPACE_GB` | No | Disk space in GB (default: 10) |\n", + "| `OETC_DELETE_WORKER_ON_ERROR` | No | Delete worker on error (default: false) |\n", + "\n", + "Keyword arguments take precedence over environment variables." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create settings from environment variables\n", + "# All required env vars must be set: OETC_EMAIL, OETC_PASSWORD,\n", + "# OETC_NAME, OETC_AUTH_URL, OETC_ORCHESTRATOR_URL\n", + "settings = OetcSettings.from_env()\n", + "\n", + "# Or override specific values via keyword arguments\n", + "settings = OetcSettings.from_env(\n", + " cpu_cores=8,\n", + " disk_space_gb=50,\n", + ")" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -221,7 +275,14 @@ "\n", "### Solver Options\n", "\n", - "You can pass solver-specific options through the `solver_options` parameter:" + "Solver name and options can be configured at two levels:\n", + "\n", + "1. **Settings level** — defaults stored in `OetcSettings.solver` and `OetcSettings.solver_options`\n", + "2. **Call level** — passed via `m.solve(solver_name=..., **solver_options)`\n", + "\n", + "Call-level options **override** settings-level options. The two dicts are\n", + "merged (call-time takes precedence), and the original settings are never\n", + "mutated." ] }, { @@ -230,28 +291,32 @@ "metadata": {}, "outputs": [], "source": [ - "# Example with advanced solver options\n", + "# Settings-level defaults\n", "advanced_settings = OetcSettings(\n", " credentials=credentials,\n", " name=\"advanced-linopy-job\",\n", " authentication_server_url=\"https://auth.oetcloud.com\",\n", " orchestrator_server_url=\"https://orchestrator.oetcloud.com\",\n", - " solver=\"gurobi\", # Using Gurobi solver\n", + " solver=\"gurobi\",\n", " solver_options={\n", - " \"TimeLimit\": 600, # 10 minutes\n", - " \"MIPGap\": 0.01, # 1% optimality gap\n", - " \"Threads\": 4, # Use 4 threads\n", - " \"OutputFlag\": 1, # Enable solver output\n", + " \"TimeLimit\": 600,\n", + " \"MIPGap\": 0.01,\n", " },\n", - " cpu_cores=8, # More CPU cores for larger problems\n", - " disk_space_gb=50, # More disk space\n", + " cpu_cores=8,\n", + " disk_space_gb=50,\n", ")\n", "\n", - "print(\"Advanced OETC settings:\")\n", - "print(f\"Solver: {advanced_settings.solver}\")\n", - "print(f\"Solver options: {advanced_settings.solver_options}\")\n", - "print(f\"CPU cores: {advanced_settings.cpu_cores}\")\n", - "print(f\"Disk space: {advanced_settings.disk_space_gb} GB\")" + "advanced_handler = OetcHandler(advanced_settings)\n", + "\n", + "# Call-level overrides: solver_name and solver_options are forwarded\n", + "# to OETC and merged with the settings defaults.\n", + "# Here MIPGap from settings (0.01) is kept, TimeLimit is overridden to 300.\n", + "status, condition = m.solve(\n", + " remote=advanced_handler,\n", + " solver_name=\"gurobi\",\n", + " TimeLimit=300,\n", + " Threads=4,\n", + ")" ] }, { @@ -356,6 +421,9 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" + }, + "nbsphinx": { + "execute": "never" } }, "nbformat": 4, diff --git a/examples/solve-on-remote.ipynb b/examples/solve-on-remote.ipynb index 16e01b41c..73e6346bf 100644 --- a/examples/solve-on-remote.ipynb +++ b/examples/solve-on-remote.ipynb @@ -4,12 +4,47 @@ "cell_type": "markdown", "id": "4db583af", "metadata": {}, - "source": ["# Remote Solving with SSH", "\n", - "This example demonstrates how linopy can solve optimization models on remote machines using SSH connections. This is one of two remote solving options available in linopy", + "source": [ + "# Remote Solving with SSH\n", + "\n", + "This example demonstrates how linopy can solve optimization models on remote machines using SSH connections. This is one of two remote solving options available in linopy:\n", + "\n", + "1. **SSH Remote Solving** (this example) - Connect to your own servers via SSH\n", + "2. **OETC Cloud Solving** - Use cloud-based optimization services (see [OETC notebook](solve-on-oetc.ipynb))\n", + "\n", + "## SSH Remote Solving\n", + "\n", + "SSH remote solving is ideal when you have:\n", + "\n", + "* Access to dedicated servers with optimization solvers installed\n", + "* Full control over the computing environment\n", + "* Existing infrastructure for optimization workloads\n", + "\n", + "## What you need for SSH remote solving\n", + "\n", + "* The `remote` extra installed on your local machine (`uv pip install \"linopy[remote]\"`), which pulls in `paramiko`\n", + "* A remote server with a working installation of linopy (e.g., in a conda environment)\n", + "* SSH access to that machine\n", + "\n", + "## How SSH Remote Solving Works\n", + "\n", + "The workflow consists of the following steps, most of which linopy handles automatically:\n", "\n", - "1. **SSH Remote Solving** (this example) - Connect to your own servers via SSH\\n2. **OETC Cloud Solving** - Use cloud-based optimization services (see `solve-on-oetc.ipynb`)", - "\n\n", - "## SSH Remote Solving\\n\\nSSH remote solving is ideal when you have:\\n* Access to dedicated servers with optimization solvers installed\\n* Full control over the computing environment\\n* Existing infrastructure for optimization workloads\\n\\n## What you need for SSH remote solving:\\n* A running installation of paramiko on your local machine (`pip install paramiko`)\\n* A remote server with a working installation of linopy (e.g., in a conda environment)\\n* SSH access to that machine\\n\\n## How SSH Remote Solving Works\\n\\nThe workflow consists of the following steps, most of which linopy handles automatically:\\n\\n1. Define a model on the local machine\\n2. Save the model on the remote machine via SSH\\n3. Load, solve and write out the model on the remote machine\\n4. Copy the solved model back to the local machine\\n5. Load the solved model on the local machine\\n\\nThe model initialization happens locally, while the actual solving happens remotely.\""] + "1. Define a model on the local machine\n", + "2. Save the model on the remote machine via SSH\n", + "3. Load, solve and write out the model on the remote machine\n", + "4. Copy the solved model back to the local machine\n", + "5. Load the solved model on the local machine\n", + "\n", + "The model initialization happens locally, while the actual solving happens remotely.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **Note:** This notebook requires SSH access to a remote server with a solver installed. It is not executed during the documentation build, so no cell outputs are shown. To run it yourself, configure SSH access and install a solver on the remote machine." + ] }, { "cell_type": "markdown", @@ -311,7 +346,7 @@ "\n", ".xr-section-summary-in + label:before {\n", " display: inline-block;\n", - " content: '►';\n", + " content: '\u25ba';\n", " font-size: 11px;\n", " width: 15px;\n", " text-align: center;\n", @@ -322,7 +357,7 @@ "}\n", "\n", ".xr-section-summary-in:checked + label:before {\n", - " content: '▼';\n", + " content: '\u25bc';\n", "}\n", "\n", ".xr-section-summary-in:checked + label > span {\n", @@ -610,6 +645,9 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.3" + }, + "nbsphinx": { + "execute": "never" } }, "nbformat": 4, diff --git a/examples/testing-framework.ipynb b/examples/testing-framework.ipynb index e8181330e..5517557dd 100644 --- a/examples/testing-framework.ipynb +++ b/examples/testing-framework.ipynb @@ -132,8 +132,7 @@ "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.3" - }, - "orig_nbformat": 4 + } }, "nbformat": 4, "nbformat_minor": 2 diff --git a/examples/transport-tutorial.ipynb b/examples/transport-tutorial.ipynb index b42e67e8e..8fdfbe9b9 100644 --- a/examples/transport-tutorial.ipynb +++ b/examples/transport-tutorial.ipynb @@ -76,9 +76,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "collapsed": true - }, + "metadata": {}, "outputs": [], "source": [ "# Import of linopy and related modules\n", @@ -114,9 +112,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "collapsed": true - }, + "metadata": {}, "outputs": [], "source": [ "## Define sets ##\n", @@ -222,9 +218,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "collapsed": true - }, + "metadata": {}, "outputs": [], "source": [ "# Parameter c(i,j) transport cost in thousands of dollars per case ;\n", @@ -313,9 +307,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "collapsed": true - }, + "metadata": {}, "outputs": [], "source": [ "## Define contraints ##\n", @@ -360,9 +352,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "collapsed": true - }, + "metadata": {}, "outputs": [], "source": [ "## Define Objective and solve ##\n", @@ -417,7 +407,7 @@ "outputs": [], "source": [ "# Solve the model\n", - "m.solve(solver_name=\"highs\")" + "m.solve(solver_name=\"highs\", output_flag=False)" ] }, { diff --git a/linopy/__init__.py b/linopy/__init__.py index 3efc297aa..b813f71d5 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -12,33 +12,68 @@ # Note: For intercepting multiplications between xarray dataarrays, Variables and Expressions # we need to extend their __mul__ functions with a quick special case import linopy.monkey_patch_xarray # noqa: F401 -from linopy.common import align +from linopy.alignment import align from linopy.config import options -from linopy.constants import EQUAL, GREATER_EQUAL, LESS_EQUAL -from linopy.constraints import Constraint, Constraints +from linopy.constants import ( + EQUAL, + GREATER_EQUAL, + LESS_EQUAL, + EvolvingAPIWarning, + PerformanceWarning, +) +from linopy.constraints import ( + Constraint, + ConstraintBase, + Constraints, + CSRConstraint, +) from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf -from linopy.model import Model, Variable, Variables, available_solvers +from linopy.model import Model, Variable, Variables from linopy.objective import Objective -from linopy.remote import OetcHandler, RemoteHandler +from linopy.piecewise import ( + PiecewiseFormulation, + Slopes, + breakpoints, + segments, + tangent_lines, +) +from linopy.remote import RemoteHandler +from linopy.solvers import SolverFeature, available_solvers, licensed_solvers + +try: + from linopy.remote import OetcCredentials, OetcHandler, OetcSettings # noqa: F401 +except ImportError: + pass __all__ = ( - "Constraint", + "CSRConstraint", + "ConstraintBase", "Constraints", + "Constraint", "EQUAL", + "PerformanceWarning", + "EvolvingAPIWarning", "GREATER_EQUAL", "LESS_EQUAL", "LinearExpression", "Model", "Objective", "OetcHandler", + "PiecewiseFormulation", "QuadraticExpression", "RemoteHandler", + "Slopes", + "SolverFeature", "Variable", "Variables", - "available_solvers", "align", + "available_solvers", + "licensed_solvers", + "breakpoints", "merge", "options", "read_netcdf", + "segments", + "tangent_lines", ) diff --git a/linopy/alignment.py b/linopy/alignment.py new file mode 100644 index 000000000..13126b852 --- /dev/null +++ b/linopy/alignment.py @@ -0,0 +1,938 @@ +#!/usr/bin/env python3 +""" +Conversion, broadcasting, and alignment of user input against coordinates. + +This module owns the seam between what users pass (scalars, numpy arrays, +pandas / polars objects, DataArrays) and what linopy stores (labelled +DataArrays conforming to a model's coordinates): + +- :func:`as_dataarray` — convert only (type dispatch + positional labeling). +- :func:`broadcast_to_coords` — convert and broadcast against ``coords``; + ``strict=True`` (default) raises on any mismatch, ``strict=False`` passes + mismatches through for downstream xarray alignment. +- :func:`validate_alignment` — the validation primitive behind the strict mode. +- :func:`align` — the symmetric counterpart, wrapping :func:`xarray.align` + for any number of linopy objects. + +Terminology for stacked MultiIndex dimensions: a dim has *levels* (its +component index names, e.g. ``period`` / ``timestep``) and *level +combinations* (its elements — one tuple per position, e.g. ``(2030, 't1')``). +""" + +from __future__ import annotations + +from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, overload +from warnings import warn + +import numpy as np +import pandas as pd +import polars as pl +from numpy import arange +from xarray import Coordinates, DataArray, Dataset, broadcast +from xarray import align as xr_align +from xarray.core import dtypes +from xarray.core.types import JoinOptions, T_Alignable +from xarray.namedarray.utils import is_dict_like + +try: + from xarray.core.coordinates import CoordinateValidationError +except ImportError: + # Added in xarray 2025.6.0; it subclasses ValueError on newer versions. + CoordinateValidationError = ValueError # type: ignore[assignment, misc] + +from linopy.constants import ( + HELPER_DIMS, + EvolvingAPIWarning, +) +from linopy.types import CoordsLike, DimsLike + +if TYPE_CHECKING: + from linopy.expressions import LinearExpression, QuadraticExpression + from linopy.variables import Variable + + +def _coords_to_dict( + coords: Sequence[Sequence | pd.Index] | Mapping, + dims: DimsLike | None = None, +) -> dict[Hashable, Any]: + """ + Normalize coords to a dict mapping dim names to coordinate values. + + Container forms: + + - ``xarray.Coordinates`` → kept dim entries only (MultiIndex level + coords dropped). + - ``Mapping`` → returned as a shallow ``dict`` copy. + - sequence-of-entries → each entry handled per the rules below. + + Sequence-entry rules (``i`` is the position in ``coords``, ``dims[i]`` + is the matching entry in ``dims`` when one exists). An entry is + *unlabeled* if it's an unnamed ``pd.Index`` or a bare ``list`` / + ``range`` / ``ndarray``. A ``tuple`` is **not** unlabeled: following + xarray, it is read as ``(dim_name, values[, attrs])`` — the first + element names the dimension. + + +---------------------------------+-----------------------+-----------+ + | Entry | Naming source | Outcome | + +=================================+=======================+===========+ + | ``pd.Index`` with ``.name`` | ``.name`` | accepted | + +---------------------------------+-----------------------+-----------+ + | unlabeled entry | ``dims[i]`` | accepted | + +---------------------------------+-----------------------+-----------+ + | unlabeled entry | — (no ``dims[i]``) | skipped | + | | | — xarray | + | | | assigns | + | | | ``dim_0`` | + | | | etc. | + +---------------------------------+-----------------------+-----------+ + | ``(name, values)`` tuple | ``name`` (1st elem) | accepted | + | | | (xarray | + | | | form) | + +---------------------------------+-----------------------+-----------+ + | tuple of length < 2 | — | TypeError | + +---------------------------------+-----------------------+-----------+ + | ``pd.MultiIndex`` with ``.name``| ``.name`` | accepted | + +---------------------------------+-----------------------+-----------+ + | ``pd.MultiIndex`` w/o ``.name`` | ``dims[i]`` | accepted | + | | | (named on | + | | | a copy) | + +---------------------------------+-----------------------+-----------+ + | ``pd.MultiIndex`` w/o ``.name`` | — (no ``dims[i]``) | TypeError | + +---------------------------------+-----------------------+-----------+ + | anything else (e.g. DataArray) | — | TypeError | + +---------------------------------+-----------------------+-----------+ + """ + if isinstance(coords, Coordinates): + # Coordinates iterates over every coord variable, including + # MultiIndex level coords. Keep only the entries that are dims. + return {d: coords[d] for d in coords.dims if d in coords} + if isinstance(coords, Mapping): + return dict(coords) + dim_names: list[Any] | None = None + if dims is not None: + dim_names = list(dims) if isinstance(dims, list | tuple) else [dims] + result: dict[Hashable, Any] = {} + for i, c in enumerate(coords): + if isinstance(c, pd.MultiIndex): + name = c.name or ( + dim_names[i] if dim_names and i < len(dim_names) else None + ) + if name is None: + raise TypeError( + "MultiIndex coords entries must have .name set so " + "xarray can use it as the dimension name. Set it via " + "`idx.name = 'my_dim'`, or pass `dims=[...]` to name " + "entries by position." + ) + if c.name is None: + c = c.copy() + c.name = name + result[name] = c + elif isinstance(c, pd.Index): + name = ( + c.name + if c.name + else (dim_names[i] if dim_names and i < len(dim_names) else None) + ) + if name is not None: + result[name] = c if c.name == name else c.rename(name) + elif isinstance(c, tuple): + if ( + len(c) < 2 + or not isinstance(c[0], Hashable) + or isinstance(c[0], list | tuple | np.ndarray) + ): + raise TypeError( + f"tuple coords entries follow xarray's (dim_name, values) " + f"convention; got {c!r}. Pass a list for a bare sequence " + f"of coordinate values." + ) + name, values = c[0], c[1] + try: + result[name] = pd.Index(values, name=name) + except TypeError as err: + raise TypeError( + f"tuple coords entries follow xarray's (dim_name, values) " + f"convention with array-like values; got {c!r}. Pass a " + f"list for a bare sequence of coordinate values." + ) from err + elif isinstance(c, list | range | np.ndarray): + if dim_names and i < len(dim_names): + result[dim_names[i]] = pd.Index(c, name=dim_names[i]) + else: + raise TypeError( + f"coords entries must be pd.Index, an unlabeled sequence " + f"(list / range / numpy.ndarray), or a (dim_name, values) " + f"tuple; got {type(c).__name__}. For an xarray DataArray " + f"coord, pass `variable.indexes[]` (a pd.Index) instead." + ) + return result + + +def _as_index(coord_values: Any) -> pd.Index: + return ( + coord_values if isinstance(coord_values, pd.Index) else pd.Index(coord_values) + ) + + +def _as_multiindex(coord_values: Any) -> pd.MultiIndex | None: + """Return the backing ``pd.MultiIndex`` of a coords entry, or ``None``.""" + if isinstance(coord_values, pd.MultiIndex): + return coord_values + if isinstance(coord_values, DataArray): + idx = coord_values.to_index() + if isinstance(idx, pd.MultiIndex): + return idx + return None + + +def get_from_iterable(lst: DimsLike | None, index: int) -> Any | None: + """ + Returns the element at the specified index of the list, or None if the index + is out of bounds. + """ + if lst is None: + return None + if isinstance(lst, Sequence | Iterable): + lst = list(lst) + else: + lst = [lst] + return lst[index] if 0 <= index < len(lst) else None + + +def pandas_to_dataarray( + arr: pd.DataFrame | pd.Series, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + **kwargs: Any, +) -> DataArray: + """ + Convert a pandas DataFrame or Series to a DataArray. + + As pandas objects already have a concept of coordinates, the + coordinates (index, columns) will be used as coordinates for the DataArray. + Solely the dimension names can be specified. + + Parameters + ---------- + arr (Union[pd.DataFrame, pd.Series]): + The input pandas DataFrame or Series. + coords (Union[dict, list, None]): + The coordinates for the DataArray. If None, default coordinates will be used. + dims (Union[list, None]): + The dimensions for the DataArray. If None, the column names of the DataFrame or the index names of the Series will be used. + **kwargs: + Additional keyword arguments to be passed to the DataArray constructor. + + Returns + ------- + DataArray: + The converted DataArray. + """ + dims = [ + axis.name or get_from_iterable(dims, i) or f"dim_{i}" + for i, axis in enumerate(arr.axes) + ] + return DataArray(arr, coords=None, dims=dims, **kwargs) + + +def numpy_to_dataarray( + arr: np.ndarray, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + **kwargs: Any, +) -> DataArray: + """ + Convert a numpy array to a DataArray. + + Parameters + ---------- + arr (np.ndarray): + The input numpy array. + coords (Union[dict, list, None]): + The coordinates for the DataArray. If None, default coordinates will be used. + dims (Union[list, None]): + The dimensions for the DataArray. If None, the dimensions will be automatically generated. + **kwargs: + Additional keyword arguments to be passed to the DataArray constructor. + + Returns + ------- + DataArray: + The converted DataArray. + """ + # fallback case for zero dim arrays + if arr.ndim == 0: + if dims is None and is_dict_like(coords): + dims = list(coords.keys()) + return DataArray(arr.item(), coords=coords, dims=dims, **kwargs) + + if isinstance(dims, Iterable | Sequence): + dims = list(dims) + elif dims is not None: + dims = [dims] + + if dims is not None and len(dims): + dims = [get_from_iterable(dims, i) or f"dim_{i}" for i in range(arr.ndim)] + + if dims is not None and len(dims) and coords is not None: + if isinstance(coords, list): + coords = dict(zip(dims, coords[: arr.ndim])) + elif is_dict_like(coords): + coords = {k: v for k, v in coords.items() if k in dims} + + return DataArray(arr, coords=coords, dims=dims, **kwargs) + + +def _named_pandas_to_dataarray(arr: pd.Series | pd.DataFrame) -> DataArray | None: + """ + Convert a pandas Series or DataFrame with fully named axes to a DataArray. + + Returns ``None`` if any axis (or MultiIndex level) is unnamed or + non-string, so the caller can fall back to ``as_dataarray``. + """ + names = list(arr.index.names) + if isinstance(arr, pd.DataFrame): + names += list(arr.columns.names) + if any(not isinstance(n, str) for n in names): + return None + + if isinstance(arr, pd.DataFrame): + if isinstance(arr.index, pd.MultiIndex) or isinstance( + arr.columns, pd.MultiIndex + ): + arr = arr.stack(list(range(arr.columns.nlevels)), future_stack=True) + return arr.to_xarray() + return DataArray(arr) + + return arr.to_xarray() + + +@overload +def fill_missing_coords(ds: DataArray, fill_helper_dims: bool = False) -> DataArray: ... + + +@overload +def fill_missing_coords(ds: Dataset, fill_helper_dims: bool = False) -> Dataset: ... + + +def fill_missing_coords( + ds: DataArray | Dataset, fill_helper_dims: bool = False +) -> Dataset | DataArray: + """ + Fill coordinates of a xarray Dataset or DataArray with integer coordinates. + + This function fills in the integer coordinates for all dimensions of a + Dataset or DataArray that have no coordinates assigned yet. + + Parameters + ---------- + ds : xarray.DataArray or xarray.Dataset + fill_helper_dims : bool, optional + Whether to fill in integer coordinates for helper dimensions, by default False. + + """ + ds = ds.copy() + if not isinstance(ds, Dataset | DataArray): + raise TypeError(f"Expected xarray.DataArray or xarray.Dataset, got {type(ds)}.") + + skip_dims = [] if fill_helper_dims else HELPER_DIMS + + # Fill in missing integer coordinates + for dim in ds.dims: + if dim not in ds.coords and dim not in skip_dims: + ds.coords[dim] = arange(ds.sizes[dim]) + + return ds + + +def as_dataarray( + arr: Any, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + **kwargs: Any, +) -> DataArray: + """ + Convert ``arr`` to a DataArray. + + Picks the right constructor for each supported input type (pandas, + polars, numpy, scalar, DataArray) and labels positional axes with + ``dims`` / ``coords``. The result is not reshaped against ``coords``: + dims are neither expanded, reordered, nor projected onto MultiIndex + dims. Use :func:`broadcast_to_coords` when + ``coords`` should govern the result's shape. + + Parameters + ---------- + arr + The input to convert. + coords + Coordinate values used to label positional axes. + dims + Dimension names used to label positional axes. + **kwargs + Forwarded to the underlying DataArray construction. + + Returns + ------- + DataArray + The converted input, dims and entries as ``arr`` provides them. + """ + if isinstance(arr, pd.Series | pd.DataFrame): + arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) + elif isinstance(arr, np.ndarray): + arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs) + elif isinstance(arr, pl.Series): + arr = numpy_to_dataarray(arr.to_numpy(), coords=coords, dims=dims, **kwargs) + elif isinstance(arr, np.number | int | float | str | bool | list): + if isinstance(arr, np.number): + arr = float(arr) + if dims is None: + if isinstance(coords, Coordinates): + dims = coords.dims + elif is_dict_like(coords) and np.ndim(arr) == 0: + dims = list(coords.keys()) + arr = DataArray(arr, coords=coords, dims=dims, **kwargs) + + elif not isinstance(arr, DataArray): + supported_types = [ + np.number, + str, + bool, + list, + pd.Series, + pd.DataFrame, + np.ndarray, + DataArray, + pl.Series, + ] + supported_types_str = ", ".join([t.__name__ for t in supported_types]) + raise TypeError( + f"Unsupported type of arr: {type(arr)}. Supported types are: {supported_types_str}" + ) + + arr = fill_missing_coords(arr) + return arr + + +class _LevelProjection(NamedTuple): + """ + Record of one MultiIndex-level projection performed by ``_broadcast_to_coords``. + + Terminology: a stacked MultiIndex dim has *levels* (its component index + names, e.g. ``period`` / ``timestep``) and *level combinations* (its + elements — one tuple per position, e.g. ``(2030, 't1')``). + """ + + dim: Hashable + levels: list[Hashable] + is_partial: bool # input carried only a subset of the MI's levels + has_gap: bool # some level combinations of the MI dim got no value (NaN) + missing: list[Any] # the level combinations that got no value + + +def _project_onto_multiindex_levels( + arr: DataArray, + expected: dict[Hashable, Any], +) -> tuple[DataArray, list[_LevelProjection]]: + """ + Map ``arr`` dims that name levels of a stacked-MultiIndex coords dim onto it. + + For every level combination of the MultiIndex dim, select the ``arr`` + value at that combination's level values. A subset of levels broadcasts + across the remaining ones; the full set aligns element-wise. ``arr`` is + returned unchanged when it carries no level dims. + + Raises ``ValueError`` only on structural errors: a level name owned by + two MI dims, or a level value missing from ``arr``. Partial projections + and coverage gaps are recorded in the returned ``_LevelProjection`` list; + the caller decides how to treat them. + """ + level_owner: dict[Hashable, Hashable] = {} + owner_mi: dict[Hashable, pd.MultiIndex] = {} + for dim, coord_values in expected.items(): + mi = _as_multiindex(coord_values) + if mi is None: + continue + owner_mi[dim] = mi + for level in mi.names: + if level is None: + continue + if level in level_owner: + raise ValueError( + f"Level {level!r} is shared by MultiIndex dimensions " + f"{level_owner[level]!r} and {dim!r}; cannot resolve which " + f"to align to." + ) + level_owner[level] = dim + + groups: dict[Hashable, list[Hashable]] = {} + for d in arr.dims: + if d in expected: + continue + owner = level_owner.get(d) + if owner is not None: + groups.setdefault(owner, []).append(d) + + projections: list[_LevelProjection] = [] + for dim, levels in groups.items(): + mi = owner_mi[dim] + selectors = { + level: DataArray(np.asarray(mi.get_level_values(level)), dims=[dim]) + for level in levels + } + try: + arr = arr.sel(selectors) + except KeyError as err: + raise ValueError( + f"Cannot align level(s) {levels} onto MultiIndex dimension " + f"{dim!r}: value {err} is missing." + ) from err + arr = arr.assign_coords(Coordinates.from_pandas_multiindex(mi, dim)) + # A level combination is "missing" when the projection gave it no + # value at any position of the other dims. + null_mask = arr.isnull() + other_dims = [d for d in arr.dims if d != dim] + if other_dims: + null_mask = null_mask.any(other_dims) + has_gap = bool(null_mask.any()) + missing = list(arr.indexes[dim][null_mask.values]) if has_gap else [] + projections.append( + _LevelProjection( + dim=dim, + levels=levels, + is_partial=len(levels) < sum(name is not None for name in mi.names), + has_gap=has_gap, + missing=missing, + ) + ) + + return arr, projections + + +def _warn_implicit_projections(projections: list[_LevelProjection]) -> None: + """ + Deprecation warnings for implicit MultiIndex-level projections. + + The same check in every mode (scenario B of the #732 / #737 discussion): + implicit projection is deprecated and raises under the v1 convention. The + strict path raises on coverage gaps before reaching here, so only partial + levels warn there; the non-strict path warns for both. + + TODO(#738): migrate to ``warn_legacy()`` / ``LinopySemanticsWarning`` + once the v1 semantics infrastructure (#717) lands. + """ + for p in projections: + if p.is_partial or p.has_gap: + kind = ( + f"broadcasting level subset {p.levels}" + if p.is_partial + else f"filling uncovered level combinations with NaN " + f"(from level(s) {p.levels})" + ) + warn( + f"multiindex-projection: implicitly {kind} onto MultiIndex " + f"dimension {p.dim!r}. This is deprecated and will raise under " + f"the v1 convention; project the input onto the dimension " + f"explicitly (select with the dimension's level values) to " + f"keep current behavior.", + EvolvingAPIWarning, + stacklevel=3, + ) + + +def _broadcast_to_coords( + arr: Any, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + **kwargs: Any, +) -> tuple[DataArray, list[_LevelProjection]]: + """ + Convert ``arr`` and broadcast it against ``coords`` (shared mechanics). + + Returns the broadcast DataArray together with the MultiIndex-level + projections performed along the way, so the public entry points can + apply their own policy (warn or raise) to partial projections and + coverage gaps. + """ + if coords is None: + return as_dataarray(arr, coords, dims, **kwargs), [] + + expected = _coords_to_dict(coords, dims=dims) + if not expected: + return as_dataarray(arr, coords, dims, **kwargs), [] + + if isinstance(arr, pd.Series | pd.DataFrame): + converted = _named_pandas_to_dataarray(arr) + if converted is not None: + arr = converted + + if not isinstance(arr, DataArray): + # numpy/polars/unnamed-pandas inputs are positional — their only + # meaningful information is the values; any axis labels are + # auto-generated. Default dims to coords' keys so the conversion + # labels axes correctly (instead of dim_0/dim_1), then re-assign + # coords from expected so positional inputs align to coords by + # position. A shape mismatch surfaces here as a clear xarray + # "conflicting sizes" error rather than a confusing + # "coordinates do not match" further down. + if dims is None: + dims = list(expected) + arr = as_dataarray(arr, coords, dims=dims, **kwargs) + # Skip MultiIndex dims — re-assigning a PandasMultiIndex coord emits + # a FutureWarning and isn't needed (the conversion already used it). + arr = arr.assign_coords( + { + d: expected[d] + for d in arr.dims + if d in expected and not isinstance(arr.indexes.get(d), pd.MultiIndex) + } + ) + + arr, projections = _project_onto_multiindex_levels(arr, expected) + + for dim, coord_values in expected.items(): + if dim not in arr.dims: + continue + if isinstance(arr.indexes.get(dim), pd.MultiIndex): + continue + expected_idx = _as_index(coord_values) + actual_idx = arr.coords[dim].to_index() + if actual_idx.equals(expected_idx): + continue + # Same values, different order → reindex to match expected order. + # Different value sets are left alone for downstream xarray alignment. + if len(actual_idx) == len(expected_idx) and set(actual_idx) == set( + expected_idx + ): + arr = arr.reindex({dim: expected_idx}) + + # expand_dims prepends new dimensions and their coordinate variables; + # the subsequent transpose restores coords order. Both are no-ops when + # the array already matches. Reconstruct so the DataArray's coords + # iteration order also follows coords (a Dataset built from this picks + # up its dim order from coord insertion). + expand = {k: v for k, v in expected.items() if k not in arr.dims} + if expand: + # expand_dims drops the level coords of a MultiIndex-backed dim, + # leaving a degenerate flat index that fails to align downstream. + # Broadcast against a proper Coordinates template instead. + plain = {} + for dim, coord_values in expand.items(): + mi = _as_multiindex(coord_values) + # Fall back to expand_dims when arr already carries one of the + # MultiIndex's level names as its own coord: broadcasting against + # the level coords would raise on the conflicting index. + if mi is None or set(mi.names) & (set(arr.coords) | set(arr.dims)): + plain[dim] = coord_values + continue + template = DataArray( + np.zeros(len(mi)), + coords=Coordinates.from_pandas_multiindex(mi, dim), + dims=[dim], + ) + arr, _ = broadcast(arr, template) + if plain: + arr = arr.expand_dims(plain) + + target_dims = tuple(d for d in expected if d in arr.dims) + tuple( + d for d in arr.dims if d not in expected + ) + arr = arr.transpose(*target_dims) + + coord_order = [c for c in target_dims if c in arr.coords] + [ + c for c in arr.coords if c not in target_dims + ] + if list(arr.coords) != coord_order: + arr = DataArray( + arr.variable, + coords={c: arr.coords[c] for c in coord_order}, + name=arr.name, + ) + + return arr, projections + + +@overload +def broadcast_to_coords( + arr: Any, + coords: CoordsLike | None = ..., + dims: DimsLike | None = ..., + *, + strict: Literal[True] = ..., + label: str, + **kwargs: Any, +) -> DataArray: ... + + +@overload +def broadcast_to_coords( + arr: Any, + coords: CoordsLike | None = ..., + dims: DimsLike | None = ..., + *, + strict: Literal[False], + label: None = ..., + **kwargs: Any, +) -> DataArray: ... + + +def broadcast_to_coords( + arr: Any, + coords: CoordsLike | None = None, + dims: DimsLike | None = None, + *, + strict: bool = True, + label: str | None = None, + **kwargs: Any, +) -> DataArray: + """ + Convert ``arr`` to a DataArray and broadcast it against ``coords``. + + When ``coords`` carries named dimensions, the result is aligned with + them: positional inputs are labeled by position, shared dims with equal + values in a different order are reindexed, dims missing from ``arr`` + are expanded, dims naming levels of a stacked-MultiIndex coords dim are + projected onto it, and the result is transposed to ``coords`` order. + + ``strict`` decides what happens to anything broadcasting alone cannot + resolve — extra dims, disagreeing coord values, and MultiIndex coverage + gaps: + + - ``strict=True`` (default): raise, naming ``label`` in the error. + - ``strict=False``: pass through unchanged so downstream xarray + alignment can handle them. + + A stacked-MultiIndex dim of ``coords`` has *levels* (its component + index names, e.g. ``period`` / ``timestep``) and *level combinations* + (its elements — one tuple per position, e.g. ``(2030, 't1')``). Inputs + indexed by levels instead of the dim itself are implicitly projected + onto the dim's level combinations. These projections are deprecated in + both modes and emit an :class:`~linopy.EvolvingAPIWarning`; the v1 + convention will require them to be explicit. Two cases: + + - input misses a whole level → broadcasts across it; warns in both modes. + - input gives some level combinations no value (a *coverage gap*) → + warns under ``strict=False``, raises under ``strict=True`` (the error + lists the missing combinations). + + Parameters + ---------- + arr + The input to convert and broadcast. + coords + Coordinate values the result is broadcast against. ``None`` falls + back to plain conversion. + dims + Dimension names used to label positional axes. + strict + Check that the result stays within ``coords`` (raise on violation) + instead of passing violations through. + label + Name of the input in error messages (e.g. ``"lower bound"``). + Required when ``strict=True``, not accepted otherwise. + **kwargs + Forwarded to the underlying DataArray construction. + + Returns + ------- + DataArray + Broadcast against ``coords``. + """ + if not strict: + da, projections = _broadcast_to_coords(arr, coords, dims, **kwargs) + _warn_implicit_projections(projections) + return da + + if label is None: + raise TypeError( + "broadcast_to_coords(strict=True) requires `label` to name the " + "input in error messages, e.g. label='lower bound'." + ) + subject = label + if coords is not None: + _coords_to_dict(coords, dims=dims) + try: + da, projections = _broadcast_to_coords(arr, coords, dims=dims, **kwargs) + except TypeError as err: + raise TypeError(f"{subject} could not be aligned to coords: {err}") from err + except (ValueError, CoordinateValidationError) as err: + raise ValueError(f"{subject} could not be aligned to coords: {err}") from err + for p in projections: + if p.has_gap: + preview = ", ".join(str(c) for c in p.missing[:5]) + if len(p.missing) > 5: + preview += f", … ({len(p.missing)} in total)" + raise ValueError( + f"{subject} could not be aligned to coords: no value for " + f"{len(p.missing)} level combination(s) of MultiIndex dimension " + f"{p.dim!r}: {preview}. The input is indexed by level(s) " + f"{p.levels} and must cover every combination." + ) + _warn_implicit_projections(projections) + validate_alignment(da, coords, dims=dims, label=label) + return da + + +def validate_alignment( + arr: DataArray, + coords: CoordsLike | None, + dims: DimsLike | None = None, + *, + label: str | None = None, +) -> None: + """ + Raise ``ValueError`` if ``arr`` is incompatible with ``coords``. + + ``arr`` is compatible with ``coords`` when both of the following hold: + + - every dim in ``arr.dims`` is also a dim in ``coords`` (no extras); + - for every dim shared between ``arr`` and ``coords``, the coord + values are equal. + + ``dims`` mirrors the ``dims`` argument of ``as_dataarray``: it names + unnamed entries in a sequence-form ``coords`` by position, so + ``coords=[[1, 2, 3]], dims=["x"]`` is enforced the same way as + ``coords={"x": [1, 2, 3]}``. + + ``label`` names the argument in error messages (e.g. ``"lower bound"``). + + No-op when ``coords`` is ``None`` or carries no named dimensions. + """ + if coords is None: + return + expected = _coords_to_dict(coords, dims=dims) + if not expected: + return + subject = label or "Value" + expected_dims = set(expected) + extra = set(arr.dims) - expected_dims + if extra: + raise ValueError( + f"{subject} has dimension(s) {sorted(extra, key=str)} not declared in coords " + f"({sorted(expected_dims, key=str)}). Add them to coords or remove them from " + f"{subject.lower()}." + ) + for dim, coord_values in expected.items(): + if dim not in arr.dims: + continue + expected_mi = _as_multiindex(coord_values) + actual_mi = _as_multiindex(arr.indexes.get(dim)) + if expected_mi is not None or actual_mi is not None: + if ( + expected_mi is None + or actual_mi is None + or not actual_mi.equals(expected_mi) + ): + raise ValueError( + f"{subject}: MultiIndex for dimension {dim!r} does not " + f"match coords." + ) + continue + expected_idx = _as_index(coord_values) + actual_idx = arr.coords[dim].to_index() + if not actual_idx.equals(expected_idx): + raise ValueError( + f"{subject}: coordinate values for dimension {dim!r} do not match " + f"coords — expected {expected_idx.tolist()}, got " + f"{actual_idx.tolist()}." + ) + + +def align( + *objects: LinearExpression | QuadraticExpression | Variable | T_Alignable, + join: JoinOptions = "inner", + copy: bool = True, + indexes: Any = None, + exclude: str | Iterable[Hashable] = frozenset(), + fill_value: Any = dtypes.NA, +) -> tuple[LinearExpression | QuadraticExpression | Variable | T_Alignable, ...]: + """ + Given any number of Variables, Expressions, Dataset and/or DataArray objects, + returns new objects with aligned indexes and dimension sizes. + + Array from the aligned objects are suitable as input to mathematical + operators, because along each dimension they have the same index and size. + + Missing values (if ``join != 'inner'``) are filled with ``fill_value``. + The default fill value is NaN. + + This functions essentially wraps the xarray function + :py:func:`xarray.align`. + + Parameters + ---------- + *objects : Variable, LinearExpression, Dataset or DataArray + Objects to align. + join : {"outer", "inner", "left", "right", "exact", "override"}, optional + Method for joining the indexes of the passed objects along each + dimension: + + - "outer": use the union of object indexes + - "inner": use the intersection of object indexes + - "left": use indexes from the first object with each dimension + - "right": use indexes from the last object with each dimension + - "exact": instead of aligning, raise `ValueError` when indexes to be + aligned are not equal + - "override": if indexes are of same size, rewrite indexes to be + those of the first object with that dimension. Indexes for the same + dimension must have the same size in all objects. + + copy : bool, default: True + If ``copy=True``, data in the return values is always copied. If + ``copy=False`` and reindexing is unnecessary, or can be performed with + only slice operations, then the output may share memory with the input. + In either case, new xarray objects are always returned. + indexes : dict-like, optional + Any indexes explicitly provided with the `indexes` argument should be + used in preference to the aligned indexes. + exclude : str, iterable of hashable or None, optional + Dimensions that must be excluded from alignment + fill_value : scalar or dict-like, optional + Value to use for newly missing values. If a dict-like, maps + variable names to fill values. Use a data array's name to + refer to its values. + + Returns + ------- + aligned : tuple of DataArray or Dataset + Tuple of objects with the same type as `*objects` with aligned + coordinates. + + + """ + from linopy.expressions import LinearExpression, QuadraticExpression + from linopy.variables import Variable + + finisher: list[partial[Any] | Callable[[Any], Any]] = [] + 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, + ) + ) + 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, + ) + return tuple([f(da) for f, da in zip(finisher, aligned)]) diff --git a/linopy/common.py b/linopy/common.py index 7dd97b654..2c45c999c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -8,9 +8,8 @@ from __future__ import annotations import operator -import os from collections.abc import Callable, Generator, Hashable, Iterable, Sequence -from functools import partial, reduce, wraps +from functools import cached_property, reduce, wraps from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload from warnings import warn @@ -18,29 +17,27 @@ import numpy as np import pandas as pd import polars as pl -from numpy import arange, signedinteger +from numpy import nan, signedinteger +from polars.datatypes import DataTypeClass from xarray import DataArray, Dataset, apply_ufunc, broadcast from xarray import align as xr_align -from xarray.core import dtypes, indexing -from xarray.core.types import JoinOptions, T_Alignable +from xarray.core import indexing from xarray.namedarray.utils import is_dict_like from linopy.config import options from linopy.constants import ( - HELPER_DIMS, SIGNS, SIGNS_alternative, SIGNS_pretty, sign_replace_dict, ) from linopy.types import ( - CoordsLike, - DimsLike, + CONSTANT_TYPES, SideLike, ) if TYPE_CHECKING: - from linopy.constraints import Constraint + from linopy.constraints import ConstraintBase from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable @@ -114,178 +111,6 @@ def format_string_as_variable_name(name: Hashable) -> str: return str(name).replace(" ", "_").replace("-", "_") -def get_from_iterable(lst: DimsLike | None, index: int) -> Any | None: - """ - Returns the element at the specified index of the list, or None if the index - is out of bounds. - """ - if lst is None: - return None - if isinstance(lst, Sequence | Iterable): - lst = list(lst) - else: - lst = [lst] - return lst[index] if 0 <= index < len(lst) else None - - -def pandas_to_dataarray( - arr: pd.DataFrame | pd.Series, - coords: CoordsLike | None = None, - dims: DimsLike | None = None, - **kwargs: Any, -) -> DataArray: - """ - Convert a pandas DataFrame or Series to a DataArray. - - As pandas objects already have a concept of coordinates, the - coordinates (index, columns) will be used as coordinates for the DataArray. - Solely the dimension names can be specified. - - Parameters - ---------- - arr (Union[pd.DataFrame, pd.Series]): - The input pandas DataFrame or Series. - coords (Union[dict, list, None]): - The coordinates for the DataArray. If None, default coordinates will be used. - dims (Union[list, None]): - The dimensions for the DataArray. If None, the column names of the DataFrame or the index names of the Series will be used. - **kwargs: - Additional keyword arguments to be passed to the DataArray constructor. - - Returns - ------- - DataArray: - The converted DataArray. - """ - dims = [ - axis.name or get_from_iterable(dims, i) or f"dim_{i}" - for i, axis in enumerate(arr.axes) - ] - if coords is not None: - pandas_coords = dict(zip(dims, arr.axes)) - if isinstance(coords, Sequence): - coords = dict(zip(dims, coords)) - shared_dims = set(pandas_coords.keys()) & set(coords.keys()) - non_aligned = [] - for dim in shared_dims: - coord = coords[dim] - if not isinstance(coord, pd.Index): - coord = pd.Index(coord) - if not pandas_coords[dim].equals(coord): - non_aligned.append(dim) - if any(non_aligned): - warn( - f"coords for dimension(s) {non_aligned} is not aligned with the pandas object. " - "Previously, the indexes of the pandas were ignored and overwritten in " - "these cases. Now, the pandas object's coordinates are taken considered" - " for alignment." - ) - - return DataArray(arr, coords=None, dims=dims, **kwargs) - - -def numpy_to_dataarray( - arr: np.ndarray, - coords: CoordsLike | None = None, - dims: DimsLike | None = None, - **kwargs: Any, -) -> DataArray: - """ - Convert a numpy array to a DataArray. - - Parameters - ---------- - arr (np.ndarray): - The input numpy array. - coords (Union[dict, list, None]): - The coordinates for the DataArray. If None, default coordinates will be used. - dims (Union[list, None]): - The dimensions for the DataArray. If None, the dimensions will be automatically generated. - **kwargs: - Additional keyword arguments to be passed to the DataArray constructor. - - Returns - ------- - DataArray: - The converted DataArray. - """ - # fallback case for zero dim arrays - if arr.ndim == 0: - return DataArray(arr.item(), coords=coords, dims=dims, **kwargs) - - ndim = max(arr.ndim, 0 if coords is None else len(coords)) - if isinstance(dims, Iterable | Sequence): - dims = list(dims) - elif dims is not None: - dims = [dims] - - if dims is not None and len(dims): - # fill up dims with default names to match the number of dimensions - dims = [get_from_iterable(dims, i) or f"dim_{i}" for i in range(ndim)] - - if isinstance(coords, list) and dims is not None and len(dims): - coords = dict(zip(dims, coords)) - - return DataArray(arr, coords=coords, dims=dims, **kwargs) - - -def as_dataarray( - arr: Any, - coords: CoordsLike | None = None, - dims: DimsLike | None = None, - **kwargs: Any, -) -> DataArray: - """ - Convert an object to a DataArray. - - Parameters - ---------- - arr: - The input object. - coords (Union[dict, list, None]): - The coordinates for the DataArray. If None, default coordinates will be used. - dims (Union[list, None]): - The dimensions for the DataArray. If None, the dimensions will be automatically generated. - **kwargs: - Additional keyword arguments to be passed to the DataArray constructor. - - Returns - ------- - DataArray: - The converted DataArray. - """ - if isinstance(arr, pd.Series | pd.DataFrame): - arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs) - elif isinstance(arr, np.ndarray): - arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs) - elif isinstance(arr, pl.Series): - arr = numpy_to_dataarray(arr.to_numpy(), coords=coords, dims=dims, **kwargs) - elif isinstance(arr, np.number): - arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs) - elif isinstance(arr, int | float | str | bool | list): - arr = DataArray(arr, coords=coords, dims=dims, **kwargs) - - elif not isinstance(arr, DataArray): - supported_types = [ - np.number, - str, - bool, - list, - pd.Series, - pd.DataFrame, - np.ndarray, - DataArray, - pl.Series, - ] - supported_types_str = ", ".join([t.__name__ for t in supported_types]) - raise TypeError( - f"Unsupported type of arr: {type(arr)}. Supported types are: {supported_types_str}" - ) - - arr = fill_missing_coords(arr) - return arr - - # TODO: rename to to_pandas_dataframe def to_dataframe( ds: Dataset, @@ -320,7 +145,7 @@ def check_has_nulls(df: pd.DataFrame, name: str) -> None: raise ValueError(f"Fields {name} contains nan's in field(s) {fields}") -def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]: +def infer_schema_polars(ds: Dataset) -> dict[str, DataTypeClass]: """ Infer the polars data schema from a xarray dataset. @@ -332,21 +157,20 @@ def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]: ------- dict: A dictionary mapping column names to their corresponding Polars data types. """ - schema = {} - np_major_version = int(np.__version__.split(".")[0]) - use_int32 = os.name == "nt" and np_major_version < 2 + schema: dict[str, DataTypeClass] = {} for name, array in ds.items(): + name = str(name) if np.issubdtype(array.dtype, np.integer): - schema[name] = pl.Int32 if use_int32 else pl.Int64 + schema[name] = pl.Int32 if array.dtype.itemsize <= 4 else pl.Int64 elif np.issubdtype(array.dtype, np.floating): - schema[name] = pl.Float64 # type: ignore + schema[name] = pl.Float64 elif np.issubdtype(array.dtype, np.bool_): - schema[name] = pl.Boolean # type: ignore + schema[name] = pl.Boolean elif np.issubdtype(array.dtype, np.object_): - schema[name] = pl.Object # type: ignore + schema[name] = pl.Object else: - schema[name] = pl.Utf8 # type: ignore - return schema # type: ignore + schema[name] = pl.Utf8 + return schema def to_polars(ds: Dataset, **kwargs: Any) -> pl.DataFrame: @@ -422,7 +246,7 @@ def filter_nulls_polars(df: pl.DataFrame) -> pl.DataFrame: if "labels" in df.columns: cond.append(pl.col("labels").ne(-1)) - cond = reduce(operator.and_, cond) # type: ignore + cond = reduce(operator.and_, cond) # type: ignore[arg-type] return df.filter(cond) @@ -449,6 +273,25 @@ def group_terms_polars(df: pl.DataFrame) -> pl.DataFrame: return df +def maybe_group_terms_polars(df: pl.DataFrame) -> pl.DataFrame: + """ + Group terms only if there are duplicate (labels, vars) pairs. + + This avoids the expensive group_by operation when terms already + reference distinct variables (e.g. ``x - y`` has ``_term=2`` but + no duplicates). When skipping, columns are reordered to match the + output of ``group_terms_polars``. + """ + varcols = [c for c in df.columns if c.startswith("vars")] + keys = [c for c in ["labels"] + varcols if c in df.columns] + key_count = df.select(pl.struct(keys).n_unique()).item() + if key_count < df.height: + return group_terms_polars(df) + # Match column order of group_terms (group-by keys, coeffs, rest) + rest = [c for c in df.columns if c not in keys and c != "coeffs"] + return df.select(keys + ["coeffs"] + rest) + + def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: """ Join multiple xarray Dataarray's to a Dataset and warn if coordinates are not equal. @@ -462,7 +305,7 @@ def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: ) arrs = xr_align(*dataarrays, join="outer") if integer_dtype: - arrs = tuple([ds.fillna(-1).astype(int) for ds in arrs]) + arrs = tuple([ds.fillna(-1).astype(options["label_dtype"]) for ds in arrs]) return Dataset({ds.name: ds for ds in arrs}) @@ -490,45 +333,7 @@ def assign_multiindex_safe(ds: Dataset, **fields: Any) -> Dataset: return Dataset({**ds[remainders], **fields}, attrs=ds.attrs) -@overload -def fill_missing_coords(ds: DataArray, fill_helper_dims: bool = False) -> DataArray: ... - - -@overload -def fill_missing_coords(ds: Dataset, fill_helper_dims: bool = False) -> Dataset: ... - - -def fill_missing_coords( - ds: DataArray | Dataset, fill_helper_dims: bool = False -) -> Dataset | DataArray: - """ - Fill coordinates of a xarray Dataset or DataArray with integer coordinates. - - This function fills in the integer coordinates for all dimensions of a - Dataset or DataArray that have no coordinates assigned yet. - - Parameters - ---------- - ds : xarray.DataArray or xarray.Dataset - fill_helper_dims : bool, optional - Whether to fill in integer coordinates for helper dimensions, by default False. - - """ - ds = ds.copy() - if not isinstance(ds, Dataset | DataArray): - raise TypeError(f"Expected xarray.DataArray or xarray.Dataset, got {type(ds)}.") - - skip_dims = [] if fill_helper_dims else HELPER_DIMS - - # Fill in missing integer coordinates - for dim in ds.dims: - if dim not in ds.coords and dim not in skip_dims: - ds.coords[dim] = arange(ds.sizes[dim]) - - return ds - - -T = TypeVar("T", Dataset, "Variable", "LinearExpression", "Constraint") +T = TypeVar("T", Dataset, "Variable", "LinearExpression", "ConstraintBase") @overload @@ -557,10 +362,10 @@ def iterate_slices( @overload def iterate_slices( - ds: Constraint, + ds: ConstraintBase, slice_size: int | None = 10_000, slice_dims: list | None = None, -) -> Generator[Constraint, None, None]: ... +) -> Generator[ConstraintBase, None, None]: ... def iterate_slices( @@ -629,7 +434,7 @@ def iterate_slices( start = i * chunk_size end = min(start + chunk_size, size_of_leading_dim) slice_dict = {leading_dim: slice(start, end)} - yield ds.isel(slice_dict) + yield ds.isel(slice_dict) # type: ignore[attr-defined] def _remap(array: np.ndarray, mapping: np.ndarray) -> np.ndarray: @@ -913,6 +718,98 @@ def find_single(value: int) -> tuple[str, dict] | tuple[None, None]: raise ValueError("Array's with more than two dimensions is not supported") +class VariableLabelIndex: + """ + Index for O(1) mapping between variable labels and dense positions. + + Both arrays are computed lazily and cached: + - ``vlabels``: active variable labels in encounter order, shape (n_active_vars,) + - ``label_to_pos``: derived from vlabels; size _xCounter, maps label -> position (-1 if masked) + + Invalidated by clearing the instance ``__dict__`` when variables are added or removed. + """ + + def __init__(self, variables: Any) -> None: + self._variables = variables + + @cached_property + def vlabels(self) -> np.ndarray: + """Active variable labels in encounter order, shape (n_active_vars,).""" + label_lists = [] + for _, var in self._variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + label_lists.append(labels[mask]) + return ( + np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp) + ) + + @cached_property + def label_to_pos(self) -> np.ndarray: + """ + Mapping from variable label to dense position, shape (_xCounter,). + + Position i in the active variable array corresponds to label vlabels[i]. + Masked or unused labels map to -1. + """ + vlabels = self.vlabels + n = self._variables.model._xCounter + label_to_pos = np.full(n, -1, dtype=np.intp) + label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp) + return label_to_pos + + @property + def n_active_vars(self) -> int: + """Number of active (non-masked) variables.""" + return len(self.vlabels) + + def invalidate(self) -> None: + """Clear cached arrays so they are recomputed on next access.""" + self.__dict__.pop("vlabels", None) + self.__dict__.pop("label_to_pos", None) + + +class ConstraintLabelIndex: + """ + Index for O(1) mapping between constraint labels and dense positions. + + Mirrors VariableLabelIndex on the constraint side, but without building + the full constraint matrix — only labels and the row mask are computed. + """ + + def __init__(self, constraints: Any) -> None: + self._constraints = constraints + + @cached_property + def clabels(self) -> np.ndarray: + """Active constraint labels in build order, shape (n_active_cons,).""" + label_lists = [ + c.active_labels() + for c in self._constraints.data.values() + if not c.is_indicator + ] + return ( + np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp) + ) + + @cached_property + def label_to_pos(self) -> np.ndarray: + """Mapping from constraint label to dense position, shape (_cCounter,).""" + clabels = self.clabels + n = self._constraints.model._cCounter + label_to_pos = np.full(n, -1, dtype=np.intp) + label_to_pos[clabels] = np.arange(len(clabels), dtype=np.intp) + return label_to_pos + + @property + def n_active_cons(self) -> int: + return len(self.clabels) + + def invalidate(self) -> None: + self.__dict__.pop("clabels", None) + self.__dict__.pop("label_to_pos", None) + + def get_label_position( obj: Any, values: int | np.ndarray, @@ -960,7 +857,7 @@ def get_label_position( raise ValueError("Array's with more than two dimensions is not supported") -def print_coord(coord: dict[str, Any] | Iterable[Any]) -> str: +def format_coord(coord: dict[str, Any] | Iterable[Any]) -> str: """ Format coordinates into a string representation. @@ -973,11 +870,11 @@ def print_coord(coord: dict[str, Any] | Iterable[Any]) -> str: with nested coordinates grouped in parentheses. Examples: - >>> print_coord({"x": 1, "y": 2}) + >>> format_coord({"x": 1, "y": 2}) '[1, 2]' - >>> print_coord([1, 2, 3]) + >>> format_coord([1, 2, 3]) '[1, 2, 3]' - >>> print_coord([(1, 2), (3, 4)]) + >>> format_coord([(1, 2), (3, 4)]) '[(1, 2), (3, 4)]' """ # Handle empty input @@ -998,7 +895,7 @@ def print_coord(coord: dict[str, Any] | Iterable[Any]) -> str: return f"[{', '.join(formatted)}]" -def print_single_variable(model: Any, label: int) -> str: +def format_single_variable(model: Any, label: int) -> str: if label == -1: return "None" @@ -1017,10 +914,10 @@ def print_single_variable(model: Any, label: int) -> str: else: bounds = f" ∈ [{lower:.4g}, {upper:.4g}]" - return f"{name}{print_coord(coord)}{bounds}" + return f"{name}{format_coord(coord)}{bounds}" -def print_single_expression( +def format_single_expression( c: np.ndarray, v: np.ndarray, const: float, @@ -1032,7 +929,7 @@ def print_single_expression( c, v = np.atleast_1d(c), np.atleast_1d(v) # catch case that to many terms would be printed - def print_line( + def format_line( expr: list[tuple[float, tuple[str, Any] | list[tuple[str, Any]]]], const: float ) -> str: res = [] @@ -1046,11 +943,11 @@ def print_line( var_string = "" for name, coords in var: if name is not None: - coord_string = print_coord(coords) + coord_string = format_coord(coords) var_string += f" {name}{coord_string}" else: name, coords = var - coord_string = print_coord(coords) + coord_string = format_coord(coords) var_string = f" {name}{coord_string}" res.append(f"{coeff_string}{var_string}") @@ -1077,7 +974,7 @@ def print_line( truncate = max_terms // 2 positions = model.variables.get_label_position(v[..., :truncate]) expr = list(zip(c[:truncate], positions)) - res = print_line(expr, const) + res = format_line(expr, const) res += " ... " expr = list( zip( @@ -1085,15 +982,15 @@ def print_line( model.variables.get_label_position(v[-truncate:]), ) ) - residual = print_line(expr, const) + residual = format_line(expr, const) if residual != " None": res += residual return res expr = list(zip(c, model.variables.get_label_position(v))) - return print_line(expr, const) + return format_line(expr, const) -def print_single_constraint(model: Any, label: int) -> str: +def format_single_constraint(model: Any, label: int) -> str: constraints = model.constraints name, coord = constraints.get_label_position(label) @@ -1102,10 +999,10 @@ def print_single_constraint(model: Any, label: int) -> str: sign = model.constraints[name].sign.sel(coord).item() rhs = model.constraints[name].rhs.sel(coord).item() - expr = print_single_expression(coeffs, vars, 0, model) + expr = format_single_expression(coeffs, vars, 0, model) sign = SIGNS_pretty[sign] - return f"{name}{print_coord(coord)}: {expr} {sign} {rhs:.12g}" + return f"{name}{format_coord(coord)}: {expr} {sign} {rhs:.12g}" def has_optimized_model(func: Callable[..., Any]) -> Callable[..., Any]: @@ -1177,110 +1074,13 @@ def check_common_keys_values(list_of_dicts: list[dict[str, Any]]) -> bool: 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", - copy: bool = True, - indexes: Any = None, - exclude: str | Iterable[Hashable] = frozenset(), - fill_value: Any = dtypes.NA, -) -> tuple[LinearExpression | QuadraticExpression | Variable | T_Alignable, ...]: - """ - Given any number of Variables, Expressions, Dataset and/or DataArray objects, - returns new objects with aligned indexes and dimension sizes. - - Array from the aligned objects are suitable as input to mathematical - operators, because along each dimension they have the same index and size. - - Missing values (if ``join != 'inner'``) are filled with ``fill_value``. - The default fill value is NaN. - - This functions essentially wraps the xarray function - :py:func:`xarray.align`. - - Parameters - ---------- - *objects : Variable, LinearExpression, Dataset or DataArray - Objects to align. - join : {"outer", "inner", "left", "right", "exact", "override"}, optional - Method for joining the indexes of the passed objects along each - dimension: - - - "outer": use the union of object indexes - - "inner": use the intersection of object indexes - - "left": use indexes from the first object with each dimension - - "right": use indexes from the last object with each dimension - - "exact": instead of aligning, raise `ValueError` when indexes to be - aligned are not equal - - "override": if indexes are of same size, rewrite indexes to be - those of the first object with that dimension. Indexes for the same - dimension must have the same size in all objects. - - copy : bool, default: True - If ``copy=True``, data in the return values is always copied. If - ``copy=False`` and reindexing is unnecessary, or can be performed with - only slice operations, then the output may share memory with the input. - In either case, new xarray objects are always returned. - indexes : dict-like, optional - Any indexes explicitly provided with the `indexes` argument should be - used in preference to the aligned indexes. - exclude : str, iterable of hashable or None, optional - Dimensions that must be excluded from alignment - fill_value : scalar or dict-like, optional - Value to use for newly missing values. If a dict-like, maps - variable names to fill values. Use a data array's name to - refer to its values. - - Returns - ------- - aligned : tuple of DataArray or Dataset - Tuple of objects with the same type as `*objects` with aligned - coordinates. - - - """ - from linopy.expressions import LinearExpression, QuadraticExpression - from linopy.variables import Variable - - finisher: list[partial[Any] | Callable[[Any], Any]] = [] - 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, - ) - ) - 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, - ) - return tuple([f(da) for f, da in zip(finisher, aligned)]) - - LocT = TypeVar( "LocT", "Dataset", "Variable", "LinearExpression", "QuadraticExpression", - "Constraint", + "ConstraintBase", ) @@ -1298,7 +1098,7 @@ def __getitem__( # expand the indexer so we can handle Ellipsis labels = indexing.expanded_indexer(key, self.object.ndim) key = dict(zip(self.object.dims, labels)) - return self.object.sel(key) + return self.object.sel(key) # type: ignore[attr-defined] class EmptyDeprecationWrapper: @@ -1332,6 +1132,60 @@ def __call__(self) -> bool: return self.value +def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: + """ + Serialize a list of pd.Index (including MultiIndex) to a DataArray dict. + + Suitable for embedding coordinate metadata as plain data variables in a + Dataset that has its own unrelated dimensions (e.g. CSR netcdf format). + Reconstruct with :func:`coords_from_dataset`. + """ + data_vars: dict[str, DataArray] = {} + for c in coords: + if isinstance(c, pd.MultiIndex): + for level_name, level_values in zip(c.names, c.levels): + data_vars[f"_coord_{c.name}_level_{level_name}"] = DataArray( + np.array(level_values), + dims=[f"_coorddim_{c.name}_level_{level_name}"], + ) + data_vars[f"_coord_{c.name}_codes"] = DataArray( + np.array(c.codes).T, + dims=[f"_coorddim_{c.name}", f"_coorddim_{c.name}_nlevels"], + ) + else: + data_vars[f"_coord_{c.name}"] = DataArray( + np.array(c), dims=[f"_coorddim_{c.name}"] + ) + return data_vars + + +def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: + """ + Deserialize a list of pd.Index (including MultiIndex) from a Dataset. + + Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`. + """ + coords = [] + for d in coord_dims: + if f"_coord_{d}_codes" in ds: + codes_2d = ds[f"_coord_{d}_codes"].values.T + level_names = [ + str(k)[len(f"_coord_{d}_level_") :] + for k in ds + if str(k).startswith(f"_coord_{d}_level_") + ] + arrays = [ + ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]] + for i, ln in enumerate(level_names) + ] + mi = pd.MultiIndex.from_arrays(arrays, names=level_names) + mi.name = d + coords.append(mi) + else: + coords.append(pd.Index(ds[f"_coord_{d}"].values, name=d)) + return coords + + def is_constant(x: SideLike) -> bool: """ Check if the given object is a constant type or an expression type without @@ -1351,7 +1205,6 @@ def is_constant(x: SideLike) -> bool: True if the object is constant-like, False otherwise. """ from linopy.expressions import ( - SUPPORTED_CONSTANT_TYPES, LinearExpression, QuadraticExpression, ) @@ -1361,9 +1214,42 @@ def is_constant(x: SideLike) -> bool: return False if isinstance(x, LinearExpression | QuadraticExpression): return x.is_constant - if isinstance(x, SUPPORTED_CONSTANT_TYPES): + if isinstance(x, CONSTANT_TYPES): return True raise TypeError( "Expected a constant, variable, or expression on the constraint side, " f"got {type(x)}." ) + + +def values_to_lookup_array( + values: np.ndarray, labels: np.ndarray, size: int | None = None +) -> np.ndarray: + """ + Build a dense NaN-padded lookup array from values and integer labels. + + Non-negative labels are placed at their corresponding positions; negative + labels are skipped. Gaps are filled with NaN. + + Parameters + ---------- + values : np.ndarray + Values to place into the lookup array. + labels : np.ndarray + Integer labels giving the target position for each value. + size : int, optional + Length of the returned array. Defaults to ``max(labels) + 1`` if any + non-negative label is present, otherwise 0. + + Returns + ------- + np.ndarray + Dense float lookup array. + """ + labels = np.asarray(labels, dtype=int) + mask = labels >= 0 + if size is None: + size = int(labels[mask].max()) + 1 if mask.any() else 0 + arr = np.full(size, nan, dtype=float) + arr[labels[mask]] = values[mask] + return arr diff --git a/linopy/config.py b/linopy/config.py index c098709d1..6a28d43fe 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -9,28 +9,38 @@ from typing import Any +import numpy as np + +_VALID_LABEL_DTYPES = {np.int32, np.int64} + class OptionSettings: - def __init__(self, **kwargs: int) -> None: + """Runtime configuration knobs (e.g. display widths). Use as a context manager or set values directly via ``options(key=value)``.""" + + def __init__(self, **kwargs: Any) -> None: self._defaults = kwargs self._current_values = kwargs.copy() - def __call__(self, **kwargs: int) -> None: + def __call__(self, **kwargs: Any) -> None: self.set_value(**kwargs) - def __getitem__(self, key: str) -> int: + def __getitem__(self, key: str) -> Any: return self.get_value(key) - def __setitem__(self, key: str, value: int) -> None: + def __setitem__(self, key: str, value: Any) -> None: return self.set_value(**{key: value}) - def set_value(self, **kwargs: int) -> None: + def set_value(self, **kwargs: Any) -> None: for k, v in kwargs.items(): if k not in self._defaults: raise KeyError(f"{k} is not a valid setting.") + if k == "label_dtype" and v not in _VALID_LABEL_DTYPES: + raise ValueError( + f"label_dtype must be one of {_VALID_LABEL_DTYPES}, got {v}" + ) self._current_values[k] = v - def get_value(self, name: str) -> int: + def get_value(self, name: str) -> Any: if name in self._defaults: return self._current_values[name] else: @@ -57,4 +67,8 @@ def __repr__(self) -> str: return f"OptionSettings:\n {settings}" -options = OptionSettings(display_max_rows=14, display_max_terms=6) +options = OptionSettings( + display_max_rows=14, + display_max_terms=6, + label_dtype=np.int32, +) diff --git a/linopy/constants.py b/linopy/constants.py index 021a9a10d..41753c3c4 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -5,11 +5,10 @@ import logging from dataclasses import dataclass, field -from enum import Enum -from typing import Any, Union +from enum import StrEnum +from typing import Any, Literal, Self, TypeAlias, get_args import numpy as np -import pandas as pd logger = logging.getLogger(__name__) @@ -18,6 +17,11 @@ GREATER_EQUAL = ">=" LESS_EQUAL = "<=" + +class PerformanceWarning(UserWarning): + """Warning raised when an operation triggers expensive Dataset reconstruction.""" + + long_EQUAL = "==" short_GREATER_EQUAL = ">" short_LESS_EQUAL = "<" @@ -33,8 +37,44 @@ short_LESS_EQUAL: LESS_EQUAL, } +STASHED_LOWER = "_stashed_lower" +STASHED_UPPER = "_stashed_upper" +STASHED_ATTRS: list[str] = [STASHED_LOWER, STASHED_UPPER] + TERM_DIM = "_term" STACKED_TERM_DIM = "_stacked_term" + +PWL_LAMBDA_SUFFIX = "_lambda" +PWL_CONVEX_SUFFIX = "_convex" +PWL_LINK_SUFFIX = "_link" +PWL_DELTA_SUFFIX = "_delta" +PWL_FILL_ORDER_SUFFIX = "_fill_order" +PWL_SEGMENT_BINARY_SUFFIX = "_segment_binary" +PWL_SELECT_SUFFIX = "_select" +PWL_ORDER_BINARY_SUFFIX = "_order_binary" +PWL_DELTA_BOUND_SUFFIX = "_delta_bound" +PWL_BINARY_ORDER_SUFFIX = "_binary_order" +PWL_ACTIVE_BOUND_SUFFIX = "_active_bound" +PWL_OUTPUT_LINK_SUFFIX = "_output_link" +PWL_CHORD_SUFFIX = "_chord" +PWL_DOMAIN_LO_SUFFIX = "_domain_lo" +PWL_DOMAIN_HI_SUFFIX = "_domain_hi" + +PWL_METHOD: TypeAlias = Literal["sos2", "lp", "incremental", "auto"] +"""Allowed values for the ``method`` argument of :meth:`Model.add_piecewise_formulation`.""" + +PWL_METHODS: frozenset[str] = frozenset(get_args(PWL_METHOD)) +"""Set of valid :data:`~linopy.constants.PWL_METHOD` values.""" + +PWL_CONVEXITY: TypeAlias = Literal["convex", "concave", "linear", "mixed"] +"""Possible values for :attr:`~linopy.piecewise.PiecewiseFormulation.convexity`.""" + +PWL_CONVEXITIES: frozenset[str] = frozenset(get_args(PWL_CONVEXITY)) +"""Set of valid :data:`~linopy.constants.PWL_CONVEXITY` values.""" +BREAKPOINT_DIM = "_breakpoint" +SEGMENT_DIM = "_segment" +LP_PIECE_DIM = f"{BREAKPOINT_DIM}_piece" +PWL_LINK_DIM = "_pwl_var" GROUPED_TERM_DIM = "_grouped_term" GROUP_DIM = "_group" FACTOR_DIM = "_factor" @@ -49,8 +89,43 @@ CV_DIM, ] +# SOS constraint attribute keys +SOS_TYPE_ATTR = "sos_type" +SOS_DIM_ATTR = "sos_dim" +SOS_BIG_M_ATTR = "big_m_upper" + +# Indicator constraint attribute keys +INDICATOR_BINARY_VAR_ATTR = "indicator_binary_var" +INDICATOR_BINARY_VAL_ATTR = "indicator_binary_val" + + +class EvolvingAPIWarning(FutureWarning): + """ + Signals a newly-added API whose details may evolve in minor releases. + + Subclasses :class:`FutureWarning` so it is visible by default. Each + emit prefixes its message with the affected feature (e.g. + ``"piecewise: ..."``) so message-regex filters can target a single + feature without hiding warnings from other features. + + Silence globally with:: + + import warnings + import linopy + + warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning) + + Or only one feature:: -class ModelStatus(Enum): + warnings.filterwarnings( + "ignore", + category=linopy.EvolvingAPIWarning, + message=r"^piecewise:", + ) + """ + + +class ModelStatus(StrEnum): """ Model status. @@ -66,7 +141,7 @@ class ModelStatus(Enum): initialized = "initialized" -class SolverStatus(Enum): +class SolverStatus(StrEnum): """ Solver status. """ @@ -78,7 +153,7 @@ class SolverStatus(Enum): unknown = "unknown" @classmethod - def process(cls, status: str) -> "SolverStatus": + def process(cls, status: str) -> Self: try: return cls(status) except ValueError: @@ -87,14 +162,14 @@ def process(cls, status: str) -> "SolverStatus": @classmethod def from_termination_condition( cls, termination_condition: "TerminationCondition" - ) -> "SolverStatus": + ) -> Self: for status in STATUS_TO_TERMINATION_CONDITION_MAP: if termination_condition in STATUS_TO_TERMINATION_CONDITION_MAP[status]: return status return cls("unknown") -class TerminationCondition(Enum): +class TerminationCondition(StrEnum): """ Termination condition of the solver. """ @@ -126,9 +201,7 @@ class TerminationCondition(Enum): licensing_problems = "licensing_problems" @classmethod - def process( - cls, termination_condition: Union["TerminationCondition", str] - ) -> "TerminationCondition": + def process(cls, termination_condition: Self | str) -> Self: if isinstance(termination_condition, TerminationCondition): termination_condition = termination_condition.value try: @@ -176,7 +249,7 @@ class Status: legacy_status: tuple[str, str] | str = "" @classmethod - def process(cls, status: str, termination_condition: str) -> "Status": + def process(cls, status: str, termination_condition: str) -> Self: return cls( status=SolverStatus.process(status), termination_condition=TerminationCondition.process(termination_condition), @@ -185,9 +258,11 @@ def process(cls, status: str, termination_condition: str) -> "Status": @classmethod def from_termination_condition( - cls, termination_condition: Union["TerminationCondition", str] - ) -> "Status": - termination_condition = TerminationCondition.process(termination_condition) + cls, termination_condition: TerminationCondition | str | None + ) -> Self: + termination_condition = TerminationCondition.process( + termination_condition if termination_condition is not None else "unknown" + ) solver_status = SolverStatus.from_termination_condition(termination_condition) return cls(solver_status, termination_condition) @@ -196,21 +271,35 @@ def is_ok(self) -> bool: return self.status == SolverStatus.ok -def _pd_series_float() -> pd.Series: - return pd.Series(dtype=float) - - @dataclass class Solution: """ Solution returned by the solver. + + ``primal`` and ``dual`` are dense float arrays indexed by linopy label: + ``primal[label]`` is the value for variable ``label``, with ``NaN`` where + no value is available (masked labels, vars dropped by the solver, etc.). + Each solver is responsible for emitting arrays in this label-indexed form. """ - primal: pd.Series = field(default_factory=_pd_series_float) - dual: pd.Series = field(default_factory=_pd_series_float) + primal: np.ndarray = field(default_factory=lambda: np.array([], dtype=float)) + dual: np.ndarray = field(default_factory=lambda: np.array([], dtype=float)) objective: float = field(default=np.nan) +@dataclass +class SolverReport: + """ + Solver-reported performance metrics. + """ + + runtime: float | None = None + mip_gap: float | None = None + dual_bound: float | None = None + barrier_iterations: int | None = None + simplex_iterations: int | None = None + + @dataclass class Result: """ @@ -220,6 +309,8 @@ class Result: status: Status solution: Solution | None = None solver_model: Any = None + solver_name: str = "" + report: SolverReport | None = None def __repr__(self) -> str: solver_model_string = ( @@ -232,10 +323,21 @@ def __repr__(self) -> str: ) else: solution_string = "Solution: None\n" + solver_name_string = f"Solver: {self.solver_name}\n" if self.solver_name else "" + report_string = "" + if self.report is not None: + if self.report.runtime is not None: + report_string += f"Runtime: {self.report.runtime:.2f}s\n" + if self.report.mip_gap is not None: + report_string += f"MIP gap: {self.report.mip_gap:.2e}\n" + if self.report.dual_bound is not None: + report_string += f"Dual bound: {self.report.dual_bound:.2e}\n" return ( f"Status: {self.status.status.value}\n" f"Termination condition: {self.status.termination_condition.value}\n" + solution_string + + solver_name_string + + report_string + f"Solver model: {solver_model_string}\n" f"Solver message: {self.status.legacy_status}" ) diff --git a/linopy/constraints.py b/linopy/constraints.py index 291beb1d1..45b2fa34e 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -7,7 +7,9 @@ from __future__ import annotations import functools -from collections.abc import Callable, Hashable, ItemsView, Iterator, Sequence +import warnings +from abc import ABC, abstractmethod +from collections.abc import Callable, Generator, Hashable, ItemsView, Iterator, Sequence from dataclasses import dataclass from itertools import product from typing import ( @@ -15,6 +17,7 @@ Any, overload, ) +from warnings import warn import numpy as np import pandas as pd @@ -29,27 +32,29 @@ from linopy import expressions, variables from linopy.common import ( + ConstraintLabelIndex, LabelPositionIndex, LocIndexer, + VariableLabelIndex, align_lines_by_delimiter, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, + coords_from_dataset, + coords_to_dataset_vars, filter_nulls_polars, + format_coord, + format_single_constraint, + format_single_expression, format_string_as_variable_name, generate_indices_for_printout, get_dims_with_index_levels, get_label_position, - group_terms_polars, has_optimized_model, - infer_schema_polars, iterate_slices, + maybe_group_terms_polars, maybe_replace_signs, - print_coord, - print_single_constraint, - print_single_expression, replace_by_map, - require_constant, save_join, to_dataframe, to_polars, @@ -61,10 +66,12 @@ HELPER_DIMS, LESS_EQUAL, TERM_DIM, + PerformanceWarning, SIGNS_pretty, ) from linopy.types import ( ConstantLike, + ConstraintLike, CoordsLike, ExpressionLike, SignLike, @@ -74,6 +81,7 @@ if TYPE_CHECKING: from linopy.model import Model + FILL_VALUE = {"labels": -1, "rhs": np.nan, "coeffs": 0, "vars": -1, "sign": "="} @@ -97,199 +105,258 @@ def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: return _conwrap -def _con_unwrap(con: Constraint | Dataset) -> Dataset: - return con.data if isinstance(con, Constraint) else con +def _con_unwrap(con: ConstraintBase | Dataset) -> Dataset: + return con.data if isinstance(con, ConstraintBase) else con -class Constraint: +class ConstraintBase(ABC): """ - Projection to a single constraint in a model. + Abstract base class for Constraint and CSRConstraint. - The Constraint class is a subclass of xr.DataArray hence most xarray - functions can be applied to it. + Provides all read-only properties and methods shared by both the frozen + Constraint (CSR-backed) and the standard Constraint (Dataset-backed). """ - __slots__ = ("_data", "_model", "_assigned") - _fill_value = FILL_VALUE - def __init__( - self, - data: Dataset, - model: Model, - name: str = "", - skip_broadcast: bool = False, - ) -> None: - """ - Initialize the Constraint. + @property + @abstractmethod + def data(self) -> Dataset: + """Get the underlying xarray Dataset representation.""" - Parameters - ---------- - labels : xarray.DataArray - labels of the constraint. - model : linopy.Model - Underlying model. - name : str - Name of the constraint. - """ + @property + @abstractmethod + def model(self) -> Model: + """Get the model reference.""" - from linopy.model import Model + @property + @abstractmethod + def name(self) -> str: + """Get the constraint name.""" - if not isinstance(data, Dataset): - raise ValueError(f"data must be a Dataset, got {type(data)}") + @property + @abstractmethod + def is_assigned(self) -> bool: + """Whether the constraint has been assigned labels by the model.""" - if not isinstance(model, Model): - raise ValueError(f"model must be a Model, got {type(model)}") + @property + @abstractmethod + def labels(self) -> DataArray: + """Get the labels DataArray.""" - # check that `labels`, `lower` and `upper`, `sign` and `mask` are in data - for attr in ("coeffs", "vars", "sign", "rhs"): - if attr not in data: - raise ValueError(f"missing '{attr}' in data") + @property + @abstractmethod + def range(self) -> tuple[int, int]: + """Return the label range of the constraint.""" - data = data.assign_attrs(name=name) + @property + @abstractmethod + def coeffs(self) -> DataArray: + """Get the LHS coefficients DataArray.""" - if not skip_broadcast: - (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + @property + @abstractmethod + def vars(self) -> DataArray: + """Get the LHS variable labels DataArray.""" - self._assigned = "labels" in data - self._data = data - self._model = model + @property + @abstractmethod + def sign(self) -> DataArray: + """Get the constraint sign DataArray.""" + + @property + @abstractmethod + def rhs(self) -> DataArray: + """Get the RHS DataArray.""" + + @property + @abstractmethod + def dual(self) -> DataArray: + """Get the dual values DataArray.""" + + @dual.setter + @abstractmethod + def dual(self, value: DataArray) -> None: + """Set the dual values DataArray.""" + + @property + @abstractmethod + def is_indicator(self) -> bool: + """Whether the constraint is an indicator constraint.""" + + @property + @abstractmethod + def binary_var(self) -> DataArray | None: + """Get the indicator binary variable labels, or None.""" + + @property + @abstractmethod + def binary_val(self) -> int | np.ndarray | None: + """Get the indicator triggering value(s), or None.""" + + @property + def data_attrs(self) -> list[str]: + """Data variables that define this constraint's persistent state.""" + base = list(Constraints.dataset_attrs) + return base + ["binary_var", "binary_val"] if self.is_indicator else base + + @abstractmethod + def has_variable(self, variable: variables.Variable) -> bool: + """Check if the constraint references any of the given variable labels.""" + + @abstractmethod + def sanitize_zeros(self) -> ConstraintBase: + """Remove terms with zero or near-zero coefficients.""" + + @abstractmethod + def sanitize_missings(self) -> ConstraintBase: + """Mask out rows where all variables are missing (-1).""" + + @abstractmethod + def sanitize_infinities(self) -> ConstraintBase: + """Mask out rows with invalid infinite RHS values.""" + + @abstractmethod + def to_polars(self) -> pl.DataFrame: + """Convert constraint to a polars DataFrame.""" + + @abstractmethod + def freeze(self) -> CSRConstraint: + """Return an immutable Constraint (CSR-backed).""" + + @abstractmethod + def mutable(self) -> Constraint: + """Return a mutable Constraint.""" + + @abstractmethod + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """ + Return (csr, con_labels, b, sense) in one pass. + + Avoids computing the CSR matrix twice when both the matrix and + the RHS/sense vectors are needed. + """ + + @abstractmethod + def active_labels(self) -> np.ndarray: + """Active constraint labels in build order, without building the CSR.""" def __getitem__( self, selector: str | int | slice | list | tuple | dict ) -> Constraint: """ Get selection from the constraint. - This is a wrapper around the xarray __getitem__ method. It returns a - new object with the selected data. + Returns a Constraint with the selected data. """ - data = Dataset({k: self.data[k][selector] for k in self.data}, attrs=self.attrs) - return self.__class__(data, self.model, self.name) + data = Dataset( + {k: self.data[k][selector] for k in self.data}, attrs=self.data.attrs + ) + return Constraint(data, self.model, self.name) @property def attrs(self) -> dict[str, Any]: - """ - Get the attributes of the constraint. - """ + """Get the attributes of the constraint.""" return self.data.attrs @property def coords(self) -> DatasetCoordinates: - """ - Get the coordinates of the constraint. - """ + """Get the coordinates of the constraint.""" return self.data.coords @property def indexes(self) -> Indexes: - """ - Get the indexes of the constraint. - """ + """Get the indexes of the constraint.""" return self.data.indexes @property def dims(self) -> Frozen[Hashable, int]: - """ - Get the dimensions of the constraint. - """ + """Get the dimensions of the constraint.""" return self.data.dims @property def sizes(self) -> Frozen[Hashable, int]: - """ - Get the sizes of the constraint. - """ + """Get the sizes of the constraint.""" return self.data.sizes @property def nterm(self) -> int: - """ - Get the number of terms in the constraint. - """ - return self.lhs.nterm + """Get the number of terms in the constraint.""" + return self.data.sizes.get(TERM_DIM, 1) @property def ndim(self) -> int: - """ - Get the number of dimensions of the constraint. - """ + """Get the number of dimensions of the constraint.""" return self.rhs.ndim @property def shape(self) -> tuple[int, ...]: - """ - Get the shape of the constraint. - """ + """Get the shape of the constraint.""" return self.rhs.shape @property def size(self) -> int: - """ - Get the size of the constraint. - """ + """Get the size of the constraint.""" return self.rhs.size @property - def loc(self) -> LocIndexer: - return LocIndexer(self) - - @property - def data(self) -> Dataset: - """ - Get the underlying DataArray. - """ - return self._data + @abstractmethod + def ncons(self) -> int: + """Get the number of active constraints (non-masked, with at least one valid variable).""" @property - def labels(self) -> DataArray: - """ - Get the labels of the constraint. - """ - return self.data.get("labels", DataArray([])) + def coord_dims(self) -> tuple[Hashable, ...]: + return tuple(k for k in self.dims if k not in HELPER_DIMS) @property - def model(self) -> Model: - """ - Get the model of the constraint. - """ - return self._model + def coord_sizes(self) -> dict[Hashable, int]: + return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} @property - def name(self) -> str: - """ - Return the name of the constraint. - """ - return self.attrs["name"] + def coord_names(self) -> list[str]: + """Get the names of the coordinates.""" + return get_dims_with_index_levels(self.data, self.coord_dims) @property - def coord_dims(self) -> tuple[Hashable, ...]: - return tuple(k for k in self.dims if k not in HELPER_DIMS) + def type(self) -> str: + """Get the type string of the constraint.""" + return "Constraint" if self.is_assigned else "Constraint (unassigned)" @property - def coord_sizes(self) -> dict[Hashable, int]: - return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} + def term_dim(self) -> str: + """Return the term dimension of the constraint.""" + return TERM_DIM @property - def coord_names(self) -> list[str]: + def mask(self) -> DataArray | None: """ - Get the names of the coordinates. + Get the mask of the constraint. + + The mask indicates on which coordinates the constraint is enabled + (True) and disabled (False). """ - return get_dims_with_index_levels(self.data, self.coord_dims) + if self.is_assigned: + result: DataArray = self.labels != FILL_VALUE["labels"] # type: ignore[assignment] + return result.astype(bool) + return None @property - def is_assigned(self) -> bool: - return self._assigned + @abstractmethod + def lhs(self) -> expressions.LinearExpression: + """Get the left-hand-side linear expression of the constraint.""" + + def __contains__(self, value: Any) -> bool: + return self.data.__contains__(value) def __repr__(self) -> str: - """ - Print the constraint arrays. - """ + """Print the constraint arrays.""" max_lines = options["display_max_rows"] dims = list(self.coord_sizes.keys()) ndim = len(dims) dim_names = self.coord_names dim_sizes = list(self.coord_sizes.values()) - size = np.prod(dim_sizes) # that the number of theoretical printouts + size = np.prod(dim_sizes) masked_entries = (~self.mask).sum().values if self.mask is not None else 0 lines = [] @@ -305,7 +372,7 @@ def __repr__(self) -> str: for i, ind in enumerate(indices) ] if self.mask is None or self.mask.values[indices]: - expr = print_single_expression( + expr = format_single_expression( self.coeffs.values[indices], self.vars.values[indices], 0, @@ -313,9 +380,9 @@ def __repr__(self) -> str: ) sign = SIGNS_pretty[self.sign.values[indices]] rhs = self.rhs.values[indices] - line = print_coord(coord) + f": {expr} {sign} {rhs}" + line = format_coord(coord) + f": {expr} {sign} {rhs}" else: - line = print_coord(coord) + ": None" + line = format_coord(coord) + ": None" lines.append(line) lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) @@ -324,7 +391,7 @@ def __repr__(self) -> str: underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") elif size == 1: - expr = print_single_expression( + expr = format_single_expression( self.coeffs.values, self.vars.values, 0, self.model ) lines.append( @@ -352,143 +419,1088 @@ def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None ) print(self) - def __contains__(self, value: Any) -> bool: - return self.data.__contains__(value) + @property + def flat(self) -> pd.DataFrame: + """ + Convert the constraint to a pandas DataFrame. + + The resulting DataFrame represents a long table format of the all + non-masked constraints with non-zero coefficients. It contains the + columns `labels`, `coeffs`, `vars`, `rhs`, `sign`. + + .. deprecated:: + Use ``to_polars()`` instead. + """ + warnings.warn( + "Constraint.flat is deprecated, use to_polars() instead.", + DeprecationWarning, + stacklevel=2, + ) + ds = self.data + + def mask_func(data: dict) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + if "labels" in data: + mask &= data["labels"] != -1 + return mask + + df = to_dataframe(ds, mask_func=mask_func) + + # Group repeated variables in the same constraint + agg_custom = {k: "first" for k in list(df.columns)} + agg_standards = dict(coeffs="sum", rhs="first", sign="first") + agg = {**agg_custom, **agg_standards} + df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) + check_has_nulls(df, name=f"{self.type} {self.name}") + return df + + @abstractmethod + def to_matrix( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray]: + """ + Construct a CSR matrix for this constraint. + + Only active (non-masked) rows are included. Column indices are dense + positions in the active variable array, as given by ``label_index``. + + Returns + ------- + csr : scipy.sparse.csr_array + Shape (n_active_cons, n_active_vars). + con_labels : np.ndarray + Active constraint labels in row order. + """ + + def to_netcdf_ds(self) -> Dataset: + """Return a Dataset representation suitable for netcdf serialization.""" + return self.data + + iterate_slices = iterate_slices + + +def _equal_nnz_slices( + indptr: np.ndarray, slice_size: int +) -> Generator[slice, None, None]: + """Yield row slices such that each slice contains at most slice_size non-zeros.""" + n_rows = len(indptr) - 1 + start = 0 + while start < n_rows: + offset = np.searchsorted( + indptr[start + 1 :], indptr[start] + slice_size, side="left" + ) + end = min(start + 1 + offset, n_rows) + yield slice(start, end) + start = end + + +class CSRConstraint(ConstraintBase): + """ + Frozen constraint backed by a CSR sparse matrix. + + Parameters + ---------- + csr : scipy.sparse.csr_array + Shape (n_flat, model._xCounter). Each row is a flat position in the + constraint grid (including masked/empty rows). + rhs : np.ndarray + Shape (n_flat,). Right-hand-side values. + sign : str or np.ndarray + Constraint sign. Either a single str ('=', '<=', '>=') for uniform + signs, or a per-row np.ndarray of sign strings for mixed signs. + coords : list of pd.Index + One index per coordinate dimension defining the constraint grid. + model : Model + The linopy model this constraint belongs to. + name : str + Name of the constraint. + cindex : int or None + Starting label assigned by the model. None if not yet assigned. + dual : np.ndarray or None + Shape (n_flat,). Dual values after solving, or None. + """ + + __slots__ = ( + "_csr", + "_con_labels", + "_rhs", + "_sign", + "_coords", + "_model", + "_name", + "_cindex", + "_dual", + "_binvar_labels", + "_binval", + ) + + def __init__( + self, + csr: scipy.sparse.csr_array, + con_labels: np.ndarray, + rhs: np.ndarray, + sign: str | np.ndarray, + coords: list[pd.Index], + model: Model, + name: str = "", + cindex: int | None = None, + dual: np.ndarray | None = None, + binvar_labels: np.ndarray | None = None, + binval: int | np.ndarray | None = None, + ) -> None: + self._csr = csr + self._con_labels = con_labels + self._rhs = rhs + self._sign = sign + self._coords = coords + self._model = model + self._name = name + self._cindex = cindex + self._dual = dual + self._binvar_labels = binvar_labels + self._binval = binval + + @property + def model(self) -> Model: + return self._model + + @property + def name(self) -> str: + return self._name + + @property + def is_assigned(self) -> bool: + return self._cindex is not None + + @property + def shape(self) -> tuple[int, ...]: + return tuple(len(c) for c in self._coords) + + @property + def full_size(self) -> int: + return int(np.prod(shape)) if (shape := self.shape) else 1 + + @property + def range(self) -> tuple[int, int]: + """Return the (start, end) label range of the constraint.""" + if self._cindex is None: + raise AttributeError("Constraint has not been assigned labels yet.") + return (self._cindex, self._cindex + self.full_size) + + @property + def ncons(self) -> int: + return self._csr.shape[0] + + @property + def attrs(self) -> dict[str, Any]: + d: dict[str, Any] = {"name": self._name} + if self._cindex is not None: + d["label_range"] = (self._cindex, self._cindex + self.full_size) + return d + + @property + def coords(self) -> DatasetCoordinates: + return Dataset(coords={c.name: c for c in self._coords}).coords + + @property + def dims(self) -> Frozen[Hashable, int]: + d: dict[Hashable, int] = {c.name: len(c) for c in self._coords} + d[TERM_DIM] = self.nterm + return Frozen(d) + + @property + def active_positions(self) -> np.ndarray: + """Flat positions of active (non-masked) rows in the original coord shape.""" + if self._cindex is None: + return np.arange(self._csr.shape[0]) + return self._con_labels - self._cindex + + @property + def sizes(self) -> Frozen[Hashable, int]: + return self.dims + + @property + def indexes(self) -> Indexes: + return Indexes({c.name: c for c in self._coords}) + + @property + def nterm(self) -> int: + return int(np.diff(self._csr.indptr).max()) if self._csr.nnz > 0 else 1 + + @property + def coord_names(self) -> list[str]: + return [str(c.name) for c in self._coords] + + def _active_to_dataarray( + self, active_values: np.ndarray, fill: float | int | str = -1 + ) -> DataArray: + full = np.full(self.full_size, fill, dtype=active_values.dtype) + full[self.active_positions] = active_values + return DataArray(full.reshape(self.shape), coords=self._coords) + + @property + def labels(self) -> DataArray: + """Get labels DataArray, shape (*coord_dims).""" + if self._cindex is None: + return DataArray([]) + return self._active_to_dataarray(self._con_labels, fill=-1) + + @property + def coeffs(self) -> DataArray: + """Get coefficients DataArray, shape (*coord_dims, _term).""" + warnings.warn( + "Accessing .coeffs on a Constraint triggers full Dataset reconstruction. " + "Use .to_matrix() for efficient access.", + PerformanceWarning, + stacklevel=2, + ) + return self.data.coeffs + + @property + def vars(self) -> DataArray: + """Get variable labels DataArray, shape (*coord_dims, _term).""" + warnings.warn( + "Accessing .vars on a Constraint triggers full Dataset reconstruction. " + "Use .to_matrix() for efficient access.", + PerformanceWarning, + stacklevel=2, + ) + return self.data.vars + + @property + def sign(self) -> DataArray: + """Get sign DataArray.""" + if isinstance(self._sign, str): + return DataArray(np.full(self.shape, self._sign), coords=self._coords) + return self._active_to_dataarray(self._sign, fill="") + + @property + def rhs(self) -> DataArray: + """Get RHS DataArray, shape (*coord_dims).""" + return self._active_to_dataarray(self._rhs, fill=np.nan) + + @rhs.setter + def rhs(self, value: ConstantLike) -> None: + raise AttributeError( + "CSRConstraint.rhs is read-only; call .mutable() to modify." + ) + + @property + def lhs(self) -> expressions.LinearExpression: + """Get LHS as LinearExpression (triggers Dataset reconstruction).""" + ds = self._to_dataset(self.nterm) + return expressions.LinearExpression(ds[["coeffs", "vars"]], self._model) + + @lhs.setter + def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: + raise AttributeError( + "CSRConstraint.lhs is read-only; call .mutable() to modify term structure." + ) + + @property + @has_optimized_model + def dual(self) -> DataArray: + """Get dual values DataArray, shape (*coord_dims).""" + if self._dual is None: + raise AttributeError( + "Underlying is optimized but does not have dual values stored." + ) + return self._active_to_dataarray(self._dual, fill=np.nan) + + @dual.setter + def dual(self, value: DataArray) -> None: + """Set dual values from a DataArray aligned with the full coord shape.""" + vals = np.asarray(value).ravel() + self._dual = vals[self.active_positions] + + @property + def is_indicator(self) -> bool: + return self._binvar_labels is not None + + @property + def binary_var(self) -> DataArray | None: + if self._binvar_labels is None: + return None + return self._active_to_dataarray(self._binvar_labels, fill=-1) + + @property + def binary_val(self) -> int | np.ndarray | None: + return self._binval + + def _to_dataset(self, nterm: int) -> Dataset: + """ + Reconstruct labels/coeffs/vars Dataset from the CSR matrix. + + Parameters + ---------- + nterm : int + Number of terms per row (width of the dense term block). + + Returns + ------- + Dataset with variables ``labels``, ``coeffs``, ``vars``. + """ + csr = self._csr + counts = np.diff(csr.indptr) + shape = self.shape + full_size = self.full_size + + # Map active row i -> flat position in full shape via con_labels + active_positions = self.active_positions + coeffs_2d = np.zeros((full_size, nterm), dtype=csr.dtype) + vars_2d = np.full((full_size, nterm), -1, dtype=options["label_dtype"]) + if csr.nnz > 0: + row_indices = np.repeat(active_positions, counts) + term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + # csr.indices are column positions into vlabels; map back to variable labels + vlabels = self._model.variables.label_index.vlabels + vars_2d[row_indices, term_cols] = vlabels[csr.indices] + coeffs_2d[row_indices, term_cols] = csr.data + + dim_names = self.coord_names + xr_coords = {c.name: c for c in self._coords} + dims_with_term = dim_names + [TERM_DIM] + coeffs_da = DataArray( + coeffs_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dims_with_term, + ) + vars_da = DataArray( + vars_2d.reshape(shape + (nterm,)), + coords=xr_coords, + dims=dims_with_term, + ) + ds = Dataset({"coeffs": coeffs_da, "vars": vars_da}) + if self._cindex is not None: + labels_flat = np.full(full_size, -1, dtype=options["label_dtype"]) + labels_flat[active_positions] = self._con_labels + ds = assign_multiindex_safe( + ds, + labels=DataArray(labels_flat.reshape(shape), coords=self._coords), + ) + return ds + + @property + def data(self) -> Dataset: + """Reconstruct the xarray Dataset from the CSR representation.""" + ds = self._to_dataset(self.nterm) + extra: dict[str, Any] = {"sign": self.sign, "rhs": self.rhs} + if self._dual is not None: + extra["dual"] = self._active_to_dataarray(self._dual, fill=np.nan) + if self._binvar_labels is not None: + extra["binary_var"] = self._active_to_dataarray( + self._binvar_labels, fill=-1 + ) + binval = self._binval + if isinstance(binval, np.ndarray) and binval.ndim > 0: + extra["binary_val"] = self._active_to_dataarray(binval, fill=-1) + else: + extra["binary_val"] = binval + return assign_multiindex_safe(ds, **extra).assign_attrs(self.attrs) + + def __repr__(self) -> str: + """Print the constraint without reconstructing the full Dataset.""" + max_lines = options["display_max_rows"] + coords = self._coords + shape = self.shape + dim_names = self.coord_names + size = self.full_size + # Map active rows (CSR is active-only) back to full flat positions + csr = self._csr + nterm = self.nterm + active_positions = self.active_positions + masked_entries = size - len(active_positions) + + # pos_to_row[flat_idx] -> CSR row index, or -1 if masked + # active_positions is sorted (labels assigned in order) + pos_to_row = np.full(size, -1, dtype=np.intp) + pos_to_row[active_positions] = np.arange(len(active_positions), dtype=np.intp) + + header_string = f"{self.type} `{self._name}`" if self._name else f"{self.type}" + lines = [] + + def row_expr(row: int) -> str: + start, end = int(csr.indptr[row]), int(csr.indptr[row + 1]) + vars_row = np.full(nterm, -1, dtype=np.int64) + coeffs_row = np.zeros(nterm, dtype=csr.dtype) + vars_row[: end - start] = csr.indices[start:end] + coeffs_row[: end - start] = csr.data[start:end] + sign = self._sign if isinstance(self._sign, str) else self._sign[row] + return f"{format_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}" + + if size > 1: + for indices in generate_indices_for_printout(shape, max_lines): + if indices is None: + lines.append("\t\t...") + else: + coord = [coords[i][int(ind)] for i, ind in enumerate(indices)] + flat_idx = int(np.ravel_multi_index(indices, shape)) + row = pos_to_row[flat_idx] + body = row_expr(row) if row >= 0 else "None" + lines.append(format_coord(coord) + f": {body}") + lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) + + shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, shape)) + mask_str = f" - {masked_entries} masked entries" if masked_entries else "" + underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) + lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") + elif size == 1: + body = row_expr(0) if len(active_positions) > 0 else "None" + lines.append(f"{header_string}\n{'-' * len(header_string)}\n{body}") + else: + lines.append(f"{header_string}\n{'-' * len(header_string)}\n") + + return "\n".join(lines) + + def to_matrix( + self, label_index: VariableLabelIndex | None = None + ) -> tuple[scipy.sparse.csr_array, np.ndarray]: + """Return the stored CSR matrix and con_labels.""" + return self._csr, self._con_labels + + def to_netcdf_ds(self) -> Dataset: + """Return a Dataset with raw CSR components for netcdf serialization.""" + csr = self._csr + data_vars: dict[str, DataArray] = { + "indptr": DataArray(csr.indptr, dims=["_indptr"]), + "indices": DataArray(csr.indices, dims=["_nnz"]), + "data": DataArray(csr.data, dims=["_nnz"]), + "rhs": DataArray(self._rhs, dims=["_flat"]), + "_con_labels": DataArray(self._con_labels, dims=["_flat"]), + } + if isinstance(self._sign, np.ndarray): + data_vars["_sign"] = DataArray(self._sign, dims=["_flat"]) + data_vars.update(coords_to_dataset_vars(self._coords)) + if self._dual is not None: + data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) + dim_names = [c.name for c in self._coords] + attrs: dict[str, Any] = { + "_linopy_format": "csr", + "cindex": self._cindex if self._cindex is not None else -1, + "shape": list(csr.shape), + "coord_dims": dim_names, + "name": self._name, + } + if isinstance(self._sign, str): + attrs["sign"] = self._sign + if self._binvar_labels is not None: + attrs["is_indicator"] = True + data_vars["_binvar_labels"] = DataArray(self._binvar_labels, dims=["_flat"]) + binval = self._binval + if isinstance(binval, np.ndarray) and binval.ndim > 0: + data_vars["_binval"] = DataArray(binval, dims=["_flat"]) + else: + attrs["binval"] = int(binval) # type: ignore[arg-type] + return Dataset(data_vars, attrs=attrs) + + @classmethod + def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: + """Reconstruct a Constraint from a netcdf Dataset (CSR format).""" + attrs = ds.attrs + shape = tuple(attrs["shape"]) + csr = scipy.sparse.csr_array( + (ds["data"].values, ds["indices"].values, ds["indptr"].values), + shape=shape, + ) + rhs = ds["rhs"].values + sign: str | np.ndarray = ds["_sign"].values if "_sign" in ds else attrs["sign"] + _cindex_raw = int(attrs["cindex"]) + cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None + coord_dims = attrs["coord_dims"] + if isinstance(coord_dims, str): + coord_dims = [coord_dims] + coords = coords_from_dataset(ds, coord_dims) + dual = ds["dual"].values if "dual" in ds else None + if "_con_labels" in ds: + con_labels = ds["_con_labels"].values + elif cindex is not None: + con_labels = np.arange(cindex, cindex + len(rhs), dtype=np.intp) + else: + con_labels = np.arange(len(rhs), dtype=np.intp) + binvar_labels: np.ndarray | None = None + binval: int | np.ndarray | None = None + if "_binvar_labels" in ds: + binvar_labels = ds["_binvar_labels"].values + binval = ds["_binval"].values if "_binval" in ds else attrs["binval"] + return cls( + csr, + con_labels, + rhs, + sign, + coords, + model, + name, + cindex=cindex, + dual=dual, + binvar_labels=binvar_labels, + binval=binval, + ) + + def has_variable(self, variable: variables.Variable) -> bool: + vlabels = self._model.variables.label_index.vlabels + return bool( + np.isin(vlabels[self._csr.indices], variable.labels.values.ravel()).any() + ) + + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """Return (csr, con_labels, b, sense) — all pre-stored, no recomputation.""" + if isinstance(self._sign, str): + sense = np.full(len(self._rhs), self._sign[0]) + else: + sense = np.array([s[0] for s in self._sign]) + return self._csr, self._con_labels, self._rhs, sense + + def active_labels(self) -> np.ndarray: + return self._con_labels + + def sanitize_zeros(self) -> CSRConstraint: + """ + Remove terms with zero or near-zero coefficients. + + Copy-on-write: rebinds ``_csr`` instead of mutating its arrays, so + external holders of the previous arrays (e.g. a ModelSnapshot + sharing them) keep a valid baseline. + """ + zeros = np.abs(self._csr.data) <= 1e-10 + if zeros.any(): + csr = self._csr.copy() + csr.data[zeros] = 0 + csr.eliminate_zeros() + self._csr = csr + return self + + def sanitize_missings(self) -> CSRConstraint: + """No-op: missing rows are already excluded during freezing.""" + return self + + def sanitize_infinities(self) -> CSRConstraint: + """Mask out rows with invalid infinite RHS values (mutates in-place).""" + if isinstance(self._sign, str): + if self._sign == LESS_EQUAL: + invalid = self._rhs == np.inf + elif self._sign == GREATER_EQUAL: + invalid = self._rhs == -np.inf + else: + return self + else: + invalid = ((self._sign == LESS_EQUAL) & (self._rhs == np.inf)) | ( + (self._sign == GREATER_EQUAL) & (self._rhs == -np.inf) + ) + if not invalid.any(): + return self + keep = ~invalid + self._csr = self._csr[keep] + self._con_labels = self._con_labels[keep] + self._rhs = self._rhs[keep] + if not isinstance(self._sign, str): + self._sign = self._sign[keep] + return self + + def freeze(self) -> CSRConstraint: + """Return self (already immutable).""" + return self + + def mutable(self) -> Constraint: + """Convert to a Constraint.""" + return Constraint(self.data, self._model, self._name) + + def to_polars(self) -> pl.DataFrame: + """Convert frozen constraint to polars DataFrame directly from CSR.""" + csr = self._csr + sign_dtype = pl.Enum(["=", "<=", ">="]) + if csr.nnz == 0: + return pl.DataFrame( + schema={ + "labels": pl.Int64, + "coeffs": pl.Float64, + "vars": pl.Int64, + "sign": sign_dtype, + "rhs": pl.Float64, + } + ) + + rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) + vlabels = self._model.variables.label_index.vlabels + + data: dict[str, Any] = { + "labels": self._con_labels[rows], + "coeffs": csr.data, + "vars": vlabels[csr.indices], + "rhs": self._rhs[rows], + } + sign_expr: pl.Expr | pl.Series = ( + pl.lit(self._sign, dtype=sign_dtype) + if isinstance(self._sign, str) + else pl.Series("sign", self._sign[rows], dtype=sign_dtype) + ) + df = pl.DataFrame(data).with_columns(sign=sign_expr) + return df[["labels", "coeffs", "vars", "sign", "rhs"]] + + def iterate_slices( + self, + slice_size: int | None = 2_000_000, + slice_dims: list | None = None, + ) -> Generator[CSRConstraint, None, None]: + """ + Yield row-batched sub-Constraints without Dataset reconstruction. + + Batches are raw CSR slices suitable only for ``to_polars()``. They are + yielded with ``coords=[]`` because batches cover contiguous active rows, + not a contiguous slice of the coordinate grid, so the original coords + would be misleading. Do not call ``.data``, ``.mutable()``, or any + coord-dependent property on batch slices. + """ + nnz = self._csr.nnz + if slice_size is None or nnz <= slice_size: + yield self + return + + for rows in _equal_nnz_slices(self._csr.indptr, slice_size): + sign = self._sign if isinstance(self._sign, str) else self._sign[rows] + yield CSRConstraint( + csr=self._csr[rows], + con_labels=self._con_labels[rows], + rhs=self._rhs[rows], + sign=sign, + coords=[], + model=self._model, + name=self._name, + ) + + @classmethod + def from_mutable( + cls, + con: Constraint, + cindex: int | None = None, + ) -> CSRConstraint: + """ + Create a CSRConstraint from a Constraint. + + Parameters + ---------- + con : Constraint + cindex : int or None + Starting label index, if assigned. + """ + label_index = con.model.variables.label_index + csr, con_labels = con.to_matrix(label_index) + csr.eliminate_zeros() + coords = [con.indexes[d] for d in con.coord_dims] + # Build active_mask aligned with con_labels (rows in csr) + # Use same filter as to_matrix: label != -1 AND at least one var != -1 + labels_flat = con.labels.values.ravel() + vars_flat = con.vars.values.reshape(len(labels_flat), con.nterm) + active_mask = (labels_flat != -1) & (vars_flat != -1).any(axis=1) + rhs = con.rhs.values.ravel()[active_mask] + sign_vals = con.sign.values.ravel() + active_signs = sign_vals[active_mask] + unique_signs = np.unique(active_signs) + if len(unique_signs) == 0: + sign: str | np.ndarray = "=" + elif len(unique_signs) == 1: + sign = str(unique_signs[0]) + else: + sign = active_signs + dual = ( + con.data["dual"].values.ravel()[active_mask] if "dual" in con.data else None + ) + binvar_labels: np.ndarray | None = None + binval: int | np.ndarray | None = None + binary_var = con.binary_var + if binary_var is not None: + binvar_labels = binary_var.values.ravel()[active_mask] + bv = con.binary_val + if isinstance(bv, np.ndarray) and bv.ndim > 0: + binval = bv.ravel()[active_mask] + else: + binval = bv + return cls( + csr, + con_labels, + rhs, + sign, + coords, + con.model, + con.name, + cindex=cindex, + dual=dual, + binvar_labels=binvar_labels, + binval=binval, + ) + + +class Constraint(ConstraintBase): + """ + Constraint backed by an xarray Dataset. + + Supports setters, xarray operations via conwrap, and from_rule construction. + """ + + __slots__ = ("_data", "_model", "_assigned", "_coef_dirty") + + def __init__( + self, + data: Dataset, + model: Model, + name: str = "", + skip_broadcast: bool = False, + ) -> None: + from linopy.model import Model + + if not isinstance(data, Dataset): + raise ValueError(f"data must be a Dataset, got {type(data)}") + + if not isinstance(model, Model): + raise ValueError(f"model must be a Model, got {type(model)}") + + for attr in ("coeffs", "vars", "sign", "rhs"): + if attr not in data: + raise ValueError(f"missing '{attr}' in data") + + data = data.assign_attrs(name=name) + + if not skip_broadcast: + (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + + self._assigned = "labels" in data + self._data = data + self._model = model + self._coef_dirty = False + + @property + def data(self) -> Dataset: + return self._data + + @property + def model(self) -> Model: + return self._model + + @property + def name(self) -> str: + return self.attrs["name"] @property - def type(self) -> str: - """ - Get the type of the constraint. - """ - return "Constraint" if self.is_assigned else "Constraint (unassigned)" + def is_assigned(self) -> bool: + return self._assigned + + @property + def ncons(self) -> int: + """Get the number of active constraints (non-masked, with at least one valid variable).""" + labels = self.labels.values + vars_arr = self.vars.values + if labels.ndim == 0: + return int(labels != FILL_VALUE["labels"] and (vars_arr != -1).any()) + return int( + ((labels != FILL_VALUE["labels"]) & (vars_arr != -1).any(axis=-1)).sum() + ) @property def range(self) -> tuple[int, int]: - """ - Return the range of the constraint. - """ + """Return the range of the constraint.""" return self.data.attrs["label_range"] @property - def term_dim(self) -> str: - """ - Return the term dimension of the constraint. - """ - return TERM_DIM + def loc(self) -> LocIndexer: + return LocIndexer(self) @property - def mask(self) -> DataArray | None: - """ - Get the mask of the constraint. - - The mask indicates on which coordinates the constraint is enabled - (True) and disabled (False). - - Returns - ------- - xr.DataArray - """ - if self.is_assigned: - return (self.data.labels != FILL_VALUE["labels"]).astype(bool) - return None + def labels(self) -> DataArray: + return self.data.get("labels", DataArray([])) @property def coeffs(self) -> DataArray: - """ - Get the left-hand-side coefficients of the constraint. - - The function raises an error in case no model is set as a - reference. - """ return self.data.coeffs @coeffs.setter def coeffs(self, value: ConstantLike) -> None: - value = DataArray(value).broadcast_like(self.vars, exclude=[self.term_dim]) - self._data = assign_multiindex_safe(self.data, coeffs=value) + """Syntactic sugar for :meth:`Constraint.update`. Do not add logic here; mutate via ``update`` so the contract stays single-sourced.""" + warn( + "Constraint.coeffs setter is deprecated and will be removed in a " + "future release; use Constraint.update(coeffs=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(coeffs=value) @property def vars(self) -> DataArray: - """ - Get the left-hand-side variables of the constraint. - - The function raises an error in case no model is set as a - reference. - """ return self.data.vars @vars.setter def vars(self, value: variables.Variable | DataArray) -> None: - if isinstance(value, variables.Variable): - value = value.labels - if not isinstance(value, DataArray): - raise TypeError("Expected value to be of type DataArray or Variable") - value = value.broadcast_like(self.coeffs, exclude=[self.term_dim]) - self._data = assign_multiindex_safe(self.data, vars=value) + """Syntactic sugar for :meth:`Constraint.update`. Do not add logic here; mutate via ``update`` so the contract stays single-sourced.""" + warn( + "Constraint.vars setter is deprecated and will be removed in a " + "future release; use Constraint.update(variables=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(variables=value) @property - def lhs(self) -> expressions.LinearExpression: - """ - Get the left-hand-side linear expression of the constraint. + def sign(self) -> DataArray: + return self.data.sign - The function raises an error in case no model is set as a - reference. - """ + @sign.setter + def sign(self, value: SignLike) -> None: + """Syntactic sugar for :meth:`Constraint.update`. Do not add logic here; mutate via ``update`` so the contract stays single-sourced.""" + warn( + "Constraint.sign setter is deprecated and will be removed in a " + "future release; use Constraint.update(sign=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(sign=value) + + @property + def rhs(self) -> DataArray: + return self.data.rhs + + @rhs.setter + def rhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: + """Syntactic sugar for :meth:`Constraint.update`. Do not add logic here; mutate via ``update`` so the contract stays single-sourced.""" + warn( + "Constraint.rhs setter is deprecated and will be removed in a " + "future release; use Constraint.update(rhs=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(rhs=value) + + @property + def is_indicator(self) -> bool: + return "binary_var" in self._data + + @property + def binary_var(self) -> DataArray | None: + return self._data.get("binary_var") + + @property + def binary_val(self) -> int | np.ndarray | None: + if "binary_val" in self._data: + return self._data["binary_val"].values + return None + + @property + def lhs(self) -> expressions.LinearExpression: data = self.data[["coeffs", "vars"]].rename({self.term_dim: TERM_DIM}) return expressions.LinearExpression(data, self.model) @lhs.setter def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: - value = expressions.as_expression( - value, self.model, coords=self.coords, dims=self.coord_dims + """Syntactic sugar for :meth:`Constraint.update`. Do not add logic here; mutate via ``update`` so the contract stays single-sourced.""" + warn( + "Constraint.lhs setter is deprecated and will be removed in a " + "future release; use Constraint.update(lhs=...) instead.", + DeprecationWarning, + stacklevel=2, ) + self.update(lhs=value) + + def _assign_lhs( + self, expr: expressions.LinearExpression, rhs: DataArray | None = None + ) -> None: + """ + Internal: replace coeffs/vars from ``expr``, adjusting rhs for + the expression's constant part. Sets ``_coef_dirty``. + """ + base_rhs = self.rhs if rhs is None else rhs self._data = self.data.drop_vars(["coeffs", "vars"]).assign( - coeffs=value.coeffs, vars=value.vars, rhs=self.rhs - value.const + coeffs=expr.coeffs, + vars=expr.vars, + rhs=base_rhs - expr.const, ) + self._coef_dirty = True - @property - def sign(self) -> DataArray: + def _update_data(self, **fields: Any) -> None: """ - Get the signs of the constraint. + Internal: write ``fields`` into ``self._data`` and update dirty bookkeeping. - The function raises an error in case no model is set as a - reference. + Writes that touch the lhs structure (``coeffs``, ``vars``) flip + ``_coef_dirty``. Other fields (``rhs``, ``sign``, …) leave it alone. """ - return self.data.sign - - @sign.setter - @require_constant - def sign(self, value: SignLike) -> None: - value = maybe_replace_signs(DataArray(value)).broadcast_like(self.sign) - self._data = assign_multiindex_safe(self.data, sign=value) + self._data = assign_multiindex_safe(self.data, **fields) + if "coeffs" in fields or "vars" in fields: + self._coef_dirty = True - @property - def rhs(self) -> DataArray: + def update( + self, + constraint: ConstraintLike | None = None, + *, + lhs: ExpressionLike | VariableLike | ConstantLike | None = None, + rhs: ExpressionLike | VariableLike | ConstantLike | None = None, + sign: SignLike | None = None, + coeffs: ConstantLike | None = None, + variables: variables.Variable | DataArray | None = None, + ) -> Constraint: """ - Get the right hand side constants of the constraint. + Update the constraint in place. - The function raises an error in case no model is set as a - reference. - """ - return self.data.rhs + The only mutation API; setters forward here. Two call shapes: - @rhs.setter - def rhs(self, value: ExpressionLike) -> None: - value = expressions.as_expression( - value, self.model, coords=self.coords, dims=self.coord_dims - ) - self.lhs = self.lhs - value.reset_const() - self._data = assign_multiindex_safe(self.data, rhs=value.const) + * ``c.update(x + 5 <= 3)`` — pass a complete constraint + expression (mirroring ``add_constraints``). Replaces lhs, + sign, and rhs at once. + * ``c.update(lhs=, rhs=, sign=, coeffs=, variables=)`` — pass + only what you want to change. + + Use the keyword form for targeted changes — it skips the + unchanged attributes entirely. The positional form always + rewrites lhs / sign / rhs (and flips ``_coef_dirty``), so it + is the wrong shape for hot loops that only touch one part: + + .. code-block:: python + + # Hot loop, rhs is the only thing changing per iteration: + for k in scenarios: + c.update(rhs=rhs_k) # ← targeted, cheap + + # Same loop written positionally rebuilds lhs every + # iteration even though it never changes: + for k in scenarios: + c.update(big_lhs_expr <= rhs_k) # ← avoid + + Parameters + ---------- + constraint : ConstraintLike, optional + A complete constraint expression (e.g. ``x + 5 <= 3``). + Mutually exclusive with the keyword arguments below. + lhs : ExpressionLike / VariableLike / ConstantLike, optional + Replace the LHS expression. Any constant part is moved to + ``rhs`` so ``c.lhs`` stays pure-variable. Cannot be combined + with ``coeffs`` / ``variables``. Sets the internal + ``_coef_dirty`` flag. + rhs : ExpressionLike / VariableLike / ConstantLike, optional + New right-hand side. + + * Constant rhs (scalar, array, DataArray) → assigned directly + to ``c.rhs``; ``c.lhs`` is untouched. + * Variable / Expression rhs → rearranged onto the lhs to + preserve the invariant that ``c.rhs`` is constant-only, + matching ``add_constraints``. **This rewrites ``c.lhs``.** + + Example — the two calls below produce the same final state:: + + # Form A: explicit, only changes rhs + c.update(rhs=5) + + # Form B: rhs carries a variable, so lhs is rewritten too. + # Starting from `2*x <= 3`, this gives `2*x - y <= 5`: + c.update(rhs=y + 5) + + If you want the rewrite to be loud, use the positional form + (``c.update(2*x - y <= 5)``) which makes both sides explicit. + sign : SignLike, optional + New sign. One of ``"<=" / "==" / ">="`` (or their ``< > =`` + aliases). + coeffs : ConstantLike, optional + Replace coefficient values (same sparsity / term structure). + Lower-level than ``lhs=``; sets ``_coef_dirty``. + variables : Variable, optional + Replace variable label array (same sparsity / term + structure). Lower-level than ``lhs=``; sets ``_coef_dirty``. + + A raw ``DataArray`` of integer labels is still accepted + for back-compat but emits a ``FutureWarning`` — pass a + ``Variable`` instead. The DataArray path will be removed + in a future release. + + Returns + ------- + Constraint + ``self`` for chaining. + """ + if constraint is not None: + if any(x is not None for x in (lhs, rhs, sign, coeffs, variables)): + raise TypeError( + "Constraint.update: positional `constraint` argument " + "cannot be combined with keyword arguments." + ) + con: ConstraintBase + if isinstance(constraint, AnonymousScalarConstraint): + con = constraint.to_constraint() + elif isinstance(constraint, ConstraintBase): + con = constraint + else: + raise TypeError( + "Constraint.update: positional argument must be a " + "ConstraintLike (e.g. `x + 5 <= 3`); got " + f"{type(constraint).__name__}." + ) + lhs, sign, rhs = con.lhs, con.sign, con.rhs + + if all(v is None for v in (lhs, rhs, sign, coeffs, variables)): + return self + + if lhs is not None and (coeffs is not None or variables is not None): + raise TypeError( + "Constraint.update: pass either `lhs=` (replace the whole " + "expression) or `coeffs=` / `variables=` (partial array " + "replacement), not both." + ) + + # 1. lhs replacement first so subsequent rhs= rearrangement sees the new lhs. + if lhs is not None: + expr = expressions.as_expression( + lhs, self.model, coords=self.coords, dims=self.coord_dims + ) + if isinstance(expr, expressions.QuadraticExpression): + raise TypeError( + "Constraint.update: lhs must be linear; got a quadratic expression." + ) + self._assign_lhs(expr) + + # 2. rhs (rearranges non-constant part onto lhs). + if rhs is not None: + expr = expressions.as_expression( + rhs, self.model, coords=self.coords, dims=self.coord_dims + ) + residual = expr.reset_const() + if residual.nterm != 0: + self._assign_lhs(self.lhs - residual, rhs=expr.const) + else: + self._update_data(rhs=expr.const) + + # 3. coeffs / variables partial updates (only valid without lhs=). + if coeffs is not None: + new_coeffs = DataArray(coeffs).broadcast_like( + self.vars, exclude=[self.term_dim] + ) + self._update_data(coeffs=new_coeffs) + if variables is not None: + from linopy.variables import Variable as _Variable + + if isinstance(variables, _Variable): + v = variables.labels + elif isinstance(variables, DataArray): + warnings.warn( + "Passing a DataArray to Constraint.update(variables=...) " + "is deprecated and will be removed in a future release; " + "pass a Variable instead.", + FutureWarning, + stacklevel=2, + ) + v = variables + else: + raise TypeError( + "Constraint.update(variables=...) expects a Variable; " + f"got {type(variables).__name__}." + ) + new_vars = v.broadcast_like(self.coeffs, exclude=[self.term_dim]) + self._update_data(vars=new_vars) + + # 4. sign last so it composes cleanly with the rest. + if sign is not None: + new_sign = maybe_replace_signs(DataArray(sign)).broadcast_like(self.sign) + self._update_data(sign=new_sign) + + return self @property @has_optimized_model def dual(self) -> DataArray: - """ - Get the dual values of the constraint. - - The function raises an error in case no model is set as a - reference or the model status is not okay. - """ if "dual" not in self.data: raise AttributeError( "Underlying is optimized but does not have dual values stored." @@ -497,12 +1509,132 @@ def dual(self) -> DataArray: @dual.setter def dual(self, value: ConstantLike) -> None: - """ - Get the dual values of the constraint. - """ value = DataArray(value).broadcast_like(self.labels) self._data = assign_multiindex_safe(self.data, dual=value) + def has_variable(self, variable: variables.Variable) -> bool: + return bool(self.data["vars"].isin(variable.labels.values.ravel()).any()) + + def _matrix_export_data( + self, label_index: VariableLabelIndex + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + label_to_pos = label_index.label_to_pos + labels_flat = self.labels.values.ravel() + vars_vals = self.vars.values + n_rows = len(labels_flat) + vars_2d = ( + vars_vals.reshape(n_rows, -1) + if n_rows > 0 + else vars_vals.reshape(0, max(1, vars_vals.size)) + ) + + row_mask = (labels_flat != -1) & (vars_2d != -1).any(axis=1) + con_labels = labels_flat[row_mask] + vars_final = vars_2d[row_mask] + valid_final = vars_final != -1 + + coeffs_final = self.coeffs.values.ravel().reshape(vars_2d.shape)[row_mask] + cols = label_to_pos[vars_final[valid_final]] + data = coeffs_final[valid_final] + + counts = valid_final.sum(axis=1) + indptr = np.empty(len(con_labels) + 1, dtype=np.int32) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + return con_labels, row_mask, cols, data, indptr + + def to_matrix( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray]: + """ + Construct a CSR matrix for this constraint. + + Only active (non-masked) rows are included. Column indices are dense + positions in the active variable array, as given by ``label_index``. + + Returns + ------- + csr : scipy.sparse.csr_array + Shape (n_active_cons, n_active_vars). + con_labels : np.ndarray + Active constraint labels in row order. + """ + con_labels, _, cols, data, indptr = self._matrix_export_data(label_index) + csr = scipy.sparse.csr_array( + (data, cols, indptr), shape=(len(con_labels), label_index.n_active_vars) + ) + csr.sum_duplicates() + return csr, con_labels + + def active_labels(self) -> np.ndarray: + labels_flat = self.labels.values.ravel() + vars_vals = self.vars.values + n_rows = len(labels_flat) + vars_2d = ( + vars_vals.reshape(n_rows, -1) + if n_rows > 0 + else vars_vals.reshape(0, max(1, vars_vals.size)) + ) + row_mask = (labels_flat != -1) & (vars_2d != -1).any(axis=1) + return labels_flat[row_mask] + + def to_matrix_with_rhs( + self, label_index: VariableLabelIndex + ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: + """Return (csr, con_labels, b, sense) in one pass.""" + con_labels, row_mask, cols, data, indptr = self._matrix_export_data(label_index) + csr = scipy.sparse.csr_array( + (data, cols, indptr), shape=(len(con_labels), label_index.n_active_vars) + ) + csr.sum_duplicates() + + b = self.rhs.values.ravel()[row_mask] + sign_flat = self.sign.values.ravel()[row_mask] + unique_signs = np.unique(sign_flat) + if len(unique_signs) == 1: + sense = np.full(len(con_labels), str(unique_signs[0])[0], dtype="U1") + else: + sense = sign_flat.astype("U1") + return csr, con_labels, b, sense + + def sanitize_zeros(self) -> Constraint: + """Remove terms with zero or near-zero coefficients.""" + not_zero = abs(self.coeffs) > 1e-10 + self._update_data( + vars=self.vars.where(not_zero, -1), + coeffs=self.coeffs.where(not_zero), + ) + return self + + def sanitize_missings(self) -> Constraint: + """Mask out rows where all variables are missing (-1).""" + contains_non_missing = (self.vars != -1).any(self.term_dim) + labels = self.labels.where(contains_non_missing, -1) + self._data = assign_multiindex_safe(self.data, labels=labels) + return self + + def sanitize_infinities(self) -> Constraint: + """Mask out rows with invalid infinite RHS values.""" + valid_infinity_values = ((self.sign == LESS_EQUAL) & (self.rhs == np.inf)) | ( + (self.sign == GREATER_EQUAL) & (self.rhs == -np.inf) + ) + labels = self.labels.where(~valid_infinity_values, -1) + self._data = assign_multiindex_safe(self.data, labels=labels) + return self + + def freeze(self) -> CSRConstraint: + """Convert to an immutable Constraint.""" + cindex = ( + int(self.data.attrs["label_range"][0]) + if "label_range" in self.data.attrs + else None + ) + return CSRConstraint.from_mutable(self, cindex=cindex) + + def mutable(self) -> Constraint: + """Return self (already mutable).""" + return self + @classmethod def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constraint: """ @@ -511,7 +1643,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai This functionality mirrors the assignment of constraints as done by Pyomo. - Parameters ---------- model : linopy.Model @@ -529,7 +1660,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai The order and size of coords has to be same as the argument list followed by `model` in function `rule`. - Returns ------- linopy.Constraint @@ -553,7 +1683,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai coords = DataArray(coords=coords).coords shape = list(map(len, coords.values())) - # test output type output = rule(model, *[c.values[0] for c in coords.values()]) if not isinstance(output, AnonymousScalarConstraint) and output is not None: msg = f"`rule` has to return AnonymousScalarConstraint not {type(output)}." @@ -573,37 +1702,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai data = lhs.data.assign(sign=sign, rhs=rhs) return cls(data, model=model) - @property - def flat(self) -> pd.DataFrame: - """ - Convert the constraint to a pandas DataFrame. - - The resulting DataFrame represents a long table format of the all - non-masked constraints with non-zero coefficients. It contains the - columns `labels`, `coeffs`, `vars`, `rhs`, `sign`. - - Returns - ------- - df : pandas.DataFrame - """ - ds = self.data - - def mask_func(data: pd.DataFrame) -> pd.Series: - mask = (data["vars"] != -1) & (data["coeffs"] != 0) - if "labels" in data: - mask &= data["labels"] != -1 - return mask - - df = to_dataframe(ds, mask_func=mask_func) - - # Group repeated variables in the same constraint - agg_custom = {k: "first" for k in list(df.columns)} - agg_standards = dict(coeffs="sum", rhs="first", sign="first") - agg = {**agg_custom, **agg_standards} - df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) - check_has_nulls(df, name=f"{self.type} {self.name}") - return df - def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. @@ -622,72 +1720,61 @@ def to_polars(self) -> pl.DataFrame: long = to_polars(ds[keys]) long = filter_nulls_polars(long) - long = group_terms_polars(long) + if ds.sizes.get("_term", 1) > 1: + long = maybe_group_terms_polars(long) check_has_nulls_polars(long, name=f"{self.type} {self.name}") - short_ds = ds[[k for k in ds if "_term" not in ds[k].dims]] - schema = infer_schema_polars(short_ds) - schema["sign"] = pl.Enum(["=", "<=", ">="]) - short = to_polars(short_ds, schema=schema) - short = filter_nulls_polars(short) - check_has_nulls_polars(short, name=f"{self.type} {self.name}") - - df = pl.concat([short, long], how="diagonal_relaxed").sort(["labels", "rhs"]) - # delete subsequent non-null rhs (happens is all vars per label are -1) - is_non_null = df["rhs"].is_not_null() - prev_non_is_null = is_non_null.shift(1).fill_null(False) - df = df.filter(is_non_null & ~prev_non_is_null | ~is_non_null) + labels_flat = ds["labels"].values.reshape(-1) + mask = labels_flat != -1 + labels_masked = labels_flat[mask] + rhs_flat = np.broadcast_to(ds["rhs"].values, ds["labels"].shape).reshape(-1) + + sign_values = ds["sign"].values + sign_flat = np.broadcast_to(sign_values, ds["labels"].shape).reshape(-1) + all_same_sign = len(sign_flat) > 0 and ( + sign_flat[0] == sign_flat[-1] and (sign_flat[0] == sign_flat).all() + ) + + short_data: dict = { + "labels": labels_masked, + "rhs": rhs_flat[mask], + } + if all_same_sign: + short = pl.DataFrame(short_data).with_columns( + pl.lit(sign_flat[0]).cast(pl.Enum(["=", "<=", ">="])).alias("sign") + ) + else: + short_data["sign"] = pl.Series( + "sign", sign_flat[mask], dtype=pl.Enum(["=", "<=", ">="]) + ) + short = pl.DataFrame(short_data) + + df = long.join(short, on="labels", how="inner") return df[["labels", "coeffs", "vars", "sign", "rhs"]] - # Wrapped function which would convert variable to dataarray + # Wrapped xarray methods — only available on Constraint assign = conwrap(Dataset.assign) - assign_multiindex_safe = conwrap(assign_multiindex_safe) - assign_attrs = conwrap(Dataset.assign_attrs) - assign_coords = conwrap(Dataset.assign_coords) - - # bfill = conwrap(Dataset.bfill) - broadcast_like = conwrap(Dataset.broadcast_like) - chunk = conwrap(Dataset.chunk) - drop_sel = conwrap(Dataset.drop_sel) - drop_isel = conwrap(Dataset.drop_isel) - expand_dims = conwrap(Dataset.expand_dims) - - # ffill = conwrap(Dataset.ffill) - sel = conwrap(Dataset.sel) - isel = conwrap(Dataset.isel) - shift = conwrap(Dataset.shift) - swap_dims = conwrap(Dataset.swap_dims) - set_index = conwrap(Dataset.set_index) - - reindex = conwrap(Dataset.reindex, fill_value=_fill_value) - - reindex_like = conwrap(Dataset.reindex_like, fill_value=_fill_value) - + reindex = conwrap(Dataset.reindex, fill_value=FILL_VALUE) + reindex_like = conwrap(Dataset.reindex_like, fill_value=FILL_VALUE) rename = conwrap(Dataset.rename) - rename_dims = conwrap(Dataset.rename_dims) - roll = conwrap(Dataset.roll) - stack = conwrap(Dataset.stack) - unstack = conwrap(Dataset.unstack) - iterate_slices = iterate_slices - @dataclass(repr=False) class Constraints: @@ -695,9 +1782,10 @@ class Constraints: A constraint container used for storing multiple constraint arrays. """ - data: dict[str, Constraint] + data: dict[str, ConstraintBase] model: Model _label_position_index: LabelPositionIndex | None = None + _constraint_label_index: ConstraintLabelIndex | None = None dataset_attrs = ["labels", "coeffs", "vars", "sign", "rhs"] dataset_names = [ @@ -716,25 +1804,34 @@ def _formatted_names(self) -> dict[str, str]: """ return {format_string_as_variable_name(n): n for n in self} - def __repr__(self) -> str: - """ - Return a string representation of the linopy model. - """ - r = "linopy.model.Constraints" - line = "-" * len(r) - r += f"\n{line}\n" - + def _format_items(self, exclude: set[str] | None = None) -> str: + """Format constraint items, optionally excluding names in a group.""" + r = "" + count = 0 for name, ds in self.items(): + if exclude and name in exclude: + continue + count += 1 coords = ( " (" + ", ".join([str(c) for c in ds.coords.keys()]) + ")" if ds.coords else "" ) r += f" * {name}{coords}\n" - if not len(list(self)): + if count == 0: r += "\n" return r + def __repr__(self) -> str: + """ + Return a string representation of the constraints container. + """ + r = "linopy.model.Constraints" + line = "-" * len(r) + r += f"\n{line}\n" + r += self._format_items() + return r + @overload def __getitem__(self, names: str) -> Constraint: ... @@ -743,10 +1840,10 @@ def __getitem__(self, names: list[str]) -> Constraints: ... def __getitem__(self, names: str | list[str]) -> Constraint | Constraints: if isinstance(names, str): - return self.data[names] + return self.data[names] # type: ignore[return-value] return Constraints({name: self.data[name] for name in names}, self.model) - def __getattr__(self, name: str) -> Constraint: + def __getattr__(self, name: str) -> ConstraintBase: # If name is an attribute of self (including methods and properties), return that if name in self.data: return self.data[name] @@ -776,7 +1873,7 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[str]: return self.data.__iter__() - def items(self) -> ItemsView[str, Constraint]: + def items(self) -> ItemsView[str, ConstraintBase]: return self.data.items() def _ipython_key_completions_(self) -> list[str]: @@ -789,12 +1886,15 @@ def _ipython_key_completions_(self) -> list[str]: """ return list(self) - def add(self, constraint: Constraint) -> None: + def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBase: """ - Add a constraint to the constraints constrainer. + Add a constraint to the constraints container. """ + if freeze and isinstance(constraint, Constraint): + constraint = constraint.freeze() self.data[constraint.name] = constraint self._invalidate_label_position_index() + return constraint def remove(self, name: str) -> None: """ @@ -807,6 +1907,15 @@ def _invalidate_label_position_index(self) -> None: """Invalidate the label position index cache.""" if self._label_position_index is not None: self._label_position_index.invalidate() + if self._constraint_label_index is not None: + self._constraint_label_index.invalidate() + + @property + def label_index(self) -> ConstraintLabelIndex: + """Index for O(1) label->position mapping and compact clabels array.""" + if self._constraint_label_index is None: + self._constraint_label_index = ConstraintLabelIndex(self) + return self._constraint_label_index @property def labels(self) -> Dataset: @@ -881,33 +1990,7 @@ def ncons(self) -> int: This excludes constraints with missing labels or where all variables are masked (vars == -1). """ - total = 0 - for con in self.data.values(): - labels = con.labels.values - vars_arr = con.vars.values - - # Handle scalar constraint (single constraint, labels is 0-d) - if labels.ndim == 0: - # Scalar: valid if label != -1 and any var != -1 - if labels != -1 and (vars_arr != -1).any(): - total += 1 - continue - - # Array constraint: labels has constraint dimensions, vars has - # constraint dimensions + _term dimension - valid_labels = labels != -1 - - # Check if any variable in each constraint is valid (not -1) - # vars has shape (..., n_terms) where ... matches labels shape - has_valid_var = (vars_arr != -1).any(axis=-1) - - active = valid_labels & has_valid_var - - if con.mask is not None: - active = active & con.mask.values - - total += int(active.sum()) - return total + return sum(con.ncons for con in self.data.values() if not con.is_indicator) @property def inequalities(self) -> Constraints: @@ -923,38 +2006,45 @@ def equalities(self) -> Constraints: """ return self[[n for n, s in self.items() if (s.sign == EQUAL).all()]] + @property + def indicator(self) -> Constraints: + """ + Get the subset of constraints which are indicator constraints. + """ + return self[[n for n, c in self.items() if c.is_indicator]] + + @property + def regular(self) -> Constraints: + """ + Get the subset of constraints which are not indicator constraints. + """ + return self[[n for n, c in self.items() if not c.is_indicator]] + def sanitize_zeros(self) -> None: """ Filter out terms with zero and close-to-zero coefficient. """ - for name in self: - not_zero = abs(self[name].coeffs) > 1e-10 - con = self[name] - con.vars = self[name].vars.where(not_zero, -1) - con.coeffs = self[name].coeffs.where(not_zero) + for con in self.data.values(): + con.sanitize_zeros() def sanitize_missings(self) -> None: """ Set constraints labels to -1 where all variables in the lhs are missing. """ - for name in self: - con = self[name] - contains_non_missing = (con.vars != -1).any(con.term_dim) - labels = self[name].labels.where(contains_non_missing, -1) - con._data = assign_multiindex_safe(con.data, labels=labels) + for con in self.data.values(): + con.sanitize_missings() def sanitize_infinities(self) -> None: """ - Replace infinite values in the constraints with a large value. + Remove constraints whose RHS is an invalid infinity. + + Constraints with ``rhs == inf`` and sign ``<=``, or ``rhs == -inf`` + and sign ``>=``, are trivially satisfied and are masked out (label set + to -1). """ - for name in self: - con = self[name] - valid_infinity_values = ((con.sign == LESS_EQUAL) & (con.rhs == np.inf)) | ( - (con.sign == GREATER_EQUAL) & (con.rhs == -np.inf) - ) - labels = con.labels.where(~valid_infinity_values, -1) - con._data = assign_multiindex_safe(con.data, labels=labels) + for con in self.data.values(): + con.sanitize_infinities() def get_name_by_label(self, label: int | float) -> str: """ @@ -1000,29 +2090,47 @@ def get_label_position( self._label_position_index = LabelPositionIndex(self) return get_label_position(self, values, self._label_position_index) - def print_labels( + def format_labels( self, values: Sequence[int], display_max_terms: int | None = None - ) -> None: + ) -> str: """ - Print a selection of labels of the constraints. + Get a string representation of a selection of constraint labels. Parameters ---------- values : list, array-like One dimensional array of constraint labels. + display_max_terms : int, optional + Maximum number of terms to display per constraint. If ``None``, + uses the global ``linopy.options.display_max_terms`` setting. + + Returns + ------- + str + String representation of the selected constraints. """ with options as opts: if display_max_terms is not None: opts.set_value(display_max_terms=display_max_terms) - res = [print_single_constraint(self.model, v) for v in values] + res = [format_single_constraint(self.model, v) for v in values] - output = "\n".join(res) - try: - print(output) - except UnicodeEncodeError: - # Replace Unicode math symbols with ASCII equivalents for Windows console - output = output.replace("≤", "<=").replace("≥", ">=").replace("≠", "!=") - print(output) + return "\n".join(res) + + def print_labels( + self, values: Sequence[int], display_max_terms: int | None = None + ) -> None: + """ + Print a selection of labels of the constraints. + + .. deprecated:: + Use :meth:`format_labels` instead. + """ + warn( + "`Constraints.print_labels` is deprecated. Use `Constraints.format_labels` instead.", + DeprecationWarning, + stacklevel=2, + ) + print(self.format_labels(values, display_max_terms=display_max_terms)) def set_blocks(self, block_map: np.ndarray) -> None: """ @@ -1040,6 +2148,8 @@ def set_blocks(self, block_map: np.ndarray) -> None: N = block_map.max() for name, constraint in self.items(): + if not isinstance(constraint, Constraint): + self.data[name] = constraint = constraint.mutable() res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype) entries = replace_by_map(constraint.vars, block_map) @@ -1071,43 +2181,69 @@ def flat(self) -> pd.DataFrame: return pd.DataFrame(columns=["coeffs", "vars", "labels", "key"]) df = pd.concat(dfs, ignore_index=True) unique_labels = df.labels.unique() - map_labels = pd.Series(np.arange(len(unique_labels)), index=unique_labels) + map_labels = pd.Series( + np.arange(len(unique_labels), dtype=options["label_dtype"]), + index=unique_labels, + ) df["key"] = df.labels.map(map_labels) return df - def to_matrix(self, filter_missings: bool = True) -> scipy.sparse.csc_matrix: + def to_matrix(self) -> tuple[scipy.sparse.csr_array, np.ndarray]: """ - Construct a constraint matrix in sparse format. + Construct a constraint matrix in sparse format by stacking per-constraint CSR matrices. - Missing values, i.e. -1 in labels and vars, are ignored filtered - out. - """ - # TODO: rename "filter_missings" to "~labels_as_coordinates" - cons = self.flat + Each per-constraint CSR is already dense: rows are active constraints + only, column indices are dense variable positions (not raw labels). + Shape is ``(n_active_cons, n_active_vars)``. + Returns + ------- + matrix : scipy.sparse.csr_array + Shape ``(n_active_cons, n_active_vars)``. + con_labels : np.ndarray + Shape ``(n_active_cons,)``, maps each matrix row to the + original constraint label. + """ if not len(self): raise ValueError("No constraints available to convert to matrix.") - if filter_missings: - vars = self.model.variables.flat - shape = (cons.key.max() + 1, vars.key.max() + 1) - cons["vars"] = cons.vars.map(vars.set_index("labels").key) - return scipy.sparse.csc_matrix( - (cons.coeffs, (cons.key, cons.vars)), shape=shape - ) - else: - shape = self.model.shape - return scipy.sparse.csc_matrix( - (cons.coeffs, (cons.labels, cons.vars)), shape=shape - ) + label_index = self.model.variables.label_index + csrs = [] + con_labels_list = [] + for c in self.data.values(): + if c.is_indicator: + continue + csr, con_labels = c.to_matrix(label_index) + csrs.append(csr) + con_labels_list.append(con_labels) + if not csrs: + raise ValueError("No constraints available to convert to matrix.") + return scipy.sparse.vstack(csrs, format="csr"), np.concatenate(con_labels_list) def reset_dual(self) -> None: """ Reset the stored solution of variables. """ for k, c in self.items(): - if "dual" in c: - c._data = c.data.drop_vars("dual") + if isinstance(c, CSRConstraint): + if c._dual is not None: + self.data[k] = CSRConstraint( + c._csr, + c._con_labels, + c._rhs, + c._sign, + c._coords, + c._model, + c._name, + cindex=c._cindex, + dual=None, + ) + elif isinstance(c, Constraint): + if "dual" in c.data: + c._data = c.data.drop_vars("dual") + else: + msg = f"reset_dual encountered an unknown constraint type: {type(c)}" + raise NotImplementedError(msg) class AnonymousScalarConstraint: @@ -1141,7 +2277,7 @@ def __repr__(self) -> str: """ Get the representation of the AnonymousScalarConstraint. """ - expr_string = print_single_expression( + expr_string = format_single_expression( np.array(self.lhs.coeffs), np.array(self.lhs.vars), 0, self.lhs.model ) return f"AnonymousScalarConstraint: {expr_string} {self.sign} {self.rhs}" diff --git a/linopy/dualization.py b/linopy/dualization.py new file mode 100644 index 000000000..edcfffdf0 --- /dev/null +++ b/linopy/dualization.py @@ -0,0 +1,673 @@ +""" +Linopy dualization module. + +This module contains implementations for constructing the dual of a linear optimization problem. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import xarray as xr + +from linopy.expressions import LinearExpression + +if TYPE_CHECKING: + from linopy.model import Model + +logger = logging.getLogger(__name__) + + +def _skip( + da: xr.DataArray, component_type: Literal["variable", "constraint"], name: str +) -> bool: + """Return True if the label array is empty or entirely masked (all -1).""" + if da.size == 0: + logger.debug(f"Skipping empty {component_type} '{name}'.") + return True + + if (da == -1).all(): + logger.debug(f"{component_type} '{name}' is fully masked, skipping.") + return True + return False + + +def _lift_bounds_to_constraints(m: Model) -> None: + """ + Convert finite variable bounds to explicit ``>=`` / ``<=`` constraints. + + Each finite lower bound becomes a ``'{var_name}-bound-lower'`` constraint and + each finite upper bound becomes a ``'{var_name}-bound-upper'`` constraint. + The variable bounds are then relaxed to ``[-inf, inf]`` to avoid + double-counting when the dual is formed. + + Bounds may vary elementwise within one variable block. Infinite bounds are + not converted into constraints. + + Parameters + ---------- + m : Model + Model to mutate in-place. + """ + for var_name, var in m.variables.items(): + label_mask = var.labels != -1 + lb = var.lower + ub = var.upper + + finite_lb = xr.DataArray(np.isfinite(lb.values), coords=lb.coords, dims=lb.dims) + finite_ub = xr.DataArray(np.isfinite(ub.values), coords=ub.coords, dims=ub.dims) + + bound_specs = [ + ("lower", var >= lb, label_mask & finite_lb, var.lower, -np.inf), + ("upper", var <= ub, label_mask & finite_ub, var.upper, np.inf), + ] + for suffix, con, bound_mask, bound, relaxed in bound_specs: + con_name = f"{var_name}-bound-{suffix}" + if con_name in m.constraints or not bool(bound_mask.any()): + continue + m.add_constraints(con, name=con_name, mask=bound_mask) + bound.values[bound_mask.values] = relaxed + + +def _dual_bounds_from_constraint_signs( + signs: xr.DataArray, + labels: xr.DataArray, + primal_is_min: bool, +) -> tuple[xr.DataArray, xr.DataArray, xr.DataArray]: + """ + Return elementwise dual-variable bounds for constraint signs. + + ``signs`` is broadcast to ``labels``. Valid signs are ``=``, ``<=``, and + ``>=``. ``valid_sign`` is True where the broadcast sign is valid. + """ + signs = signs.broadcast_like(labels) + + is_eq = signs == "=" + is_le = signs == "<=" + is_ge = signs == ">=" + valid_sign = is_eq | is_le | is_ge + + lower = xr.zeros_like(labels, dtype=float) + upper = xr.zeros_like(labels, dtype=float) + + lower = lower.where(~is_eq, -np.inf) + upper = upper.where(~is_eq, np.inf) + + if primal_is_min: + lower = lower.where(~is_le, -np.inf) + upper = upper.where(~is_le, 0.0) + + lower = lower.where(~is_ge, 0.0) + upper = upper.where(~is_ge, np.inf) + else: + lower = lower.where(~is_le, 0.0) + upper = upper.where(~is_le, np.inf) + + lower = lower.where(~is_ge, -np.inf) + upper = upper.where(~is_ge, 0.0) + + return lower, upper, valid_sign + + +def _add_dual_variables(m: Model, m_dual: Model) -> dict: + """ + Add one dual variable to m_dual for each included constraint in m. + + Dual variable bounds encode the sign convention for each constraint type + and primal objective sense: + + ============ =========== ================ ================ + Constraint Primal sense lower upper + ============ =========== ================ ================ + = min / max -inf +inf (free) + <= min -inf 0 + <= max 0 +inf + >= min 0 +inf + >= max -inf 0 + ============ =========== ================ ================ + + Fully masked or empty constraints are skipped. Constraint arrays may contain + mixed signs; dual variable bounds are assigned elementwise. + + Parameters + ---------- + m : Model + Primal model. + m_dual : Model + Dual model to populate. + + Returns + ------- + dict + ``{constraint_name: dual_variable}`` for every dualized constraint. + """ + primal_is_min = m.objective.sense == "min" + + dual_vars = {} + for name, con in m.constraints.items(): + if _skip(con.labels, "constraint", name): + continue + + lower, upper, valid_sign = _dual_bounds_from_constraint_signs( + con.sign, + con.labels, + primal_is_min, + ) + + unmasked = con.labels != -1 + invalid_unmasked = unmasked & ~valid_sign + if bool(invalid_unmasked.any()): + logger.warning( + f"Constraint '{name}' has unrecognized signs; invalid entries " + "will be skipped." + ) + + mask = unmasked & valid_sign + if not bool(mask.any()): + logger.warning( + f"Constraint '{name}' has no entries with valid signs, skipping." + ) + continue + + logger.debug( + f"Adding dual variable for constraint '{name}' " + f"with shape {con.shape} and dims {con.labels.dims}." + ) + coords = ( + [con.indexes[dim] for dim in con.labels.dims] if con.coord_dims else None + ) + dual_vars[name] = m_dual.add_variables( + lower=lower, + upper=upper, + coords=coords, + name=name, + mask=mask, + ) + + return dual_vars + + +def _build_flat_con_to_dual_label_lookup(m: Model, dual_vars: dict) -> np.ndarray: + """ + Build a lookup from flat primal constraint labels to flat dual variable labels. + + The returned array maps each flat constraint label to the corresponding + flat dual variable label from ``dual_vars``. Entries are -1 for constraints + not in ``dual_vars`` or with masked labels. + + Parameters + ---------- + m : Model + Primal model. + dual_vars : dict + ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``. + + Returns + ------- + np.ndarray of int64 + Lookup array of length ``max_flat_con_label + 1``; empty if no valid + constraint labels exist. + """ + max_flat_con = -1 + for con_name in dual_vars: + flat = m.constraints[con_name].labels.values.ravel() + valid = flat[flat != -1] + if len(valid): + max_flat_con = max(max_flat_con, int(valid.max())) + + if max_flat_con < 0: + return np.array([], dtype=np.int64) + + lookup = np.full(max_flat_con + 1, -1, dtype=np.int64) + for con_name, dv in dual_vars.items(): + con_flat = m.constraints[con_name].labels.values.ravel().astype(np.int64) + dv_flat = dv.labels.values.ravel().astype(np.int64) + valid = (con_flat != -1) & (dv_flat != -1) + if valid.any(): + lookup[con_flat[valid]] = dv_flat[valid] + + return lookup + + +def _extract_dual_feas_entries( + A: Any, + vlabels: np.ndarray, + clabels: np.ndarray, + flat_con_to_dual: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Return ``(flat_var_label, flat_dual_label, coeff)`` for each included + nonzero entry of ``A``. + + Converts sparse matrix ``A`` to coordinate format, maps rows/columns to + flat labels via ``clabels`` / ``vlabels``, and drops entries that are masked + or belong to a constraint without a dual variable. + + Parameters + ---------- + A : scipy sparse matrix + Primal constraint matrix in any scipy sparse format. + vlabels : np.ndarray of int64 + Flat variable label per column of ``A``; -1 for masked columns. + clabels : np.ndarray of int64 + Flat constraint label per row of ``A``; -1 for masked rows. + flat_con_to_dual : np.ndarray of int64 + Lookup array as returned by ``_build_flat_con_to_dual_label_lookup()``. + + Returns + ------- + flat_v : np.ndarray of int64 + Flat primal variable label for each retained nonzero entry. + flat_d : np.ndarray of int64 + Corresponding flat dual variable label. + coeffs : np.ndarray of float64 + Corresponding coefficient from ``A``. + """ + A_coo = A.tocoo() + flat_v = vlabels[A_coo.col].astype(np.int64) + flat_c = clabels[A_coo.row].astype(np.int64) + coeffs = A_coo.data + + has_labels = (flat_v != -1) & (flat_c != -1) + flat_v, flat_c, coeffs = flat_v[has_labels], flat_c[has_labels], coeffs[has_labels] + + flat_d = _gather_with_default(flat_c, flat_con_to_dual, -1, np.int64) + + has_dual = flat_d != -1 + return flat_v[has_dual], flat_d[has_dual], coeffs[has_dual] + + +def _build_obj_coeff_lookup(vlabels: np.ndarray, c_vec: np.ndarray) -> np.ndarray: + """ + Build a lookup from flat variable labels to objective coefficients. + + The returned array maps each valid flat variable label to its objective + coefficient. Labels that do not occur in ``vlabels`` but fall within the + lookup range have value 0.0. + + Parameters + ---------- + vlabels : np.ndarray of int64 + Flat variable label per column of ``A``; -1 for masked columns. + c_vec : np.ndarray of float64 + Objective coefficient per column, in the same ordering as ``vlabels``. + + Returns + ------- + np.ndarray of float64 + Lookup array of length ``max_flat_var_label + 1``; empty if no valid + variable labels exist. + """ + valid_mask = vlabels != -1 + valid_vlabels = vlabels[valid_mask] + if not len(valid_vlabels): + return np.array([], dtype=np.float64) + + lookup = np.zeros(int(valid_vlabels.max()) + 1, dtype=np.float64) + lookup[valid_vlabels.astype(np.int64)] = c_vec[valid_mask] + return lookup + + +def _build_label_to_flat_index_lookup(labels: np.ndarray) -> np.ndarray: + """ + Build a lookup from flat labels to flat indices. + + ``labels`` is expected to contain nonnegative integer labels, with -1 used + as the masked-label sentinel. + + ``lookup[label]`` gives the flat index of ``label`` in ``labels``. + Labels that do not occur in ``labels`` map to -1. If no valid labels exist, + returns an empty int64 array. + """ + valid = labels != -1 + if not valid.any(): + return np.array([], dtype=np.int64) + + lookup = np.full(int(labels[valid].max()) + 1, -1, dtype=np.int64) + lookup[labels[valid].astype(np.int64)] = np.where(valid)[0].astype(np.int64) + return lookup + + +def _gather_with_default( + labels: np.ndarray, lookup: np.ndarray, default: float, dtype: type +) -> np.ndarray: + """ + Gather ``lookup[labels]`` elementwise, using ``default`` where ``labels`` fall + outside ``lookup``'s range (including the -1 masked sentinel). + """ + out: np.ndarray = np.full(len(labels), default, dtype=dtype) + in_range = (labels >= 0) & (labels < len(lookup)) + if in_range.any(): + out[in_range] = lookup[labels[in_range].astype(np.int64)] + return out + + +def _lookup_flat_indices(labels: np.ndarray, lookup: np.ndarray) -> np.ndarray: + """ + Look up flat indices for flat labels. + + ``labels`` is expected to contain nonnegative integer labels, with -1 used + as the masked-label sentinel. ``lookup`` is expected to be an array as + returned by ``_build_label_to_flat_index_lookup()``. + + Labels outside the lookup range, including masked labels, map to -1. + """ + return _gather_with_default(labels, lookup, -1, np.int64) + + +def _term_slots_for_sorted_flat_indices(sorted_flat_indices: np.ndarray) -> np.ndarray: + """ + Return the term slot within each run of equal sorted flat indices. + + ``sorted_flat_indices`` must be non-empty and sorted in nondecreasing order. + Each run of equal values is treated as one group. + + Example: + ------- + ``[2, 2, 2, 5, 5, 9]`` becomes ``[0, 1, 2, 0, 1, 0]``. + """ + group_start = np.empty(len(sorted_flat_indices), dtype=bool) + group_start[0] = True + group_start[1:] = sorted_flat_indices[1:] != sorted_flat_indices[:-1] + + group_ids = np.cumsum(group_start) - 1 + group_start_pos = np.where(group_start)[0] + + return ( + np.arange(len(sorted_flat_indices), dtype=np.int64) - group_start_pos[group_ids] + ) + + +def _build_dual_feas_lhs( + var: Any, + flat_v: np.ndarray, + flat_d: np.ndarray, + nnz_data: np.ndarray, + m_dual: Model, +) -> LinearExpression: + """ + Build the dual-feasibility LHS for one primal variable. + + For each coordinate of ``var``, constructs the linear expression + ``sum_j A[j, var] * lambda[j]`` over the corresponding dual variables. + + Parameters + ---------- + var : linopy.Variable + Primal variable whose dual-feasibility LHS is being constructed. + flat_v : np.ndarray of int64 + Flat primal variable labels for all included nonzero entries, as returned + by ``_extract_dual_feas_entries()``. + flat_d : np.ndarray of int64 + Corresponding flat dual variable labels. + nnz_data : np.ndarray of float64 + Corresponding nonzero coefficients from the primal constraint matrix. + m_dual : Model + Dual model owning the dual variables. + + Returns + ------- + LinearExpression + Dual-feasibility LHS over ``var``'s coordinate space. Returns an empty + expression, equivalent to constant zero, when ``var`` has no included + constraint-matrix entries. + """ + var_flat = var.labels.values.ravel().astype(np.int64) + n_elements = len(var_flat) + + var_to_idx = _build_label_to_flat_index_lookup(var_flat) + if not len(var_to_idx): + return LinearExpression(None, m_dual) + + lin_idx = _lookup_flat_indices(flat_v, var_to_idx) + + keep = lin_idx != -1 + if not keep.any(): + return LinearExpression(None, m_dual) + + lin_idx = lin_idx[keep] + dual_labels = flat_d[keep] + coeffs = nnz_data[keep] + + order = np.argsort(lin_idx, kind="stable") + lin_idx = lin_idx[order] + dual_labels = dual_labels[order] + coeffs = coeffs[order] + + col_idx = _term_slots_for_sorted_flat_indices(lin_idx) + max_terms = int(col_idx.max()) + 1 + + dual_labels_2d = np.full((n_elements, max_terms), -1, dtype=np.int64) + dual_coeffs_2d = np.zeros((n_elements, max_terms), dtype=np.float64) + dual_labels_2d[lin_idx, col_idx] = dual_labels + dual_coeffs_2d[lin_idx, col_idx] = coeffs + + target_shape = var.labels.shape + (max_terms,) + dims = list(var.labels.dims) + ["_term"] + ds = xr.Dataset( + { + "vars": xr.DataArray(dual_labels_2d.reshape(target_shape), dims=dims), + "coeffs": xr.DataArray(dual_coeffs_2d.reshape(target_shape), dims=dims), + }, + coords={dim: var.labels.coords[dim] for dim in var.labels.dims}, + ) + return LinearExpression(ds, m_dual) + + +def _add_dual_feasibility_constraints( + m: Model, + m_dual: Model, + dual_vars: dict, +) -> None: + """ + Add the stationarity constraint ``sum_i(A_ji * lambda_i) = c_j`` for each + primal variable. + + Sign conventions are already encoded in the dual variable bounds produced by + ``_add_dual_variables()``, so raw A coefficients are used without adjustment. + Variables with no constraint connections (e.g. unconstrained free variables) + are skipped; the dual is infeasible for such problems if ``c_j != 0``. + + Parameters + ---------- + m : Model + Primal model. + m_dual : Model + Dual model to populate. + dual_vars : dict + ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``. + """ + A = m.matrices.A + if A is None: + raise ValueError("Constraint matrix is None, model has no constraints.") + + vlabels = np.asarray(m.matrices.vlabels, dtype=np.int64) + clabels = np.asarray(m.matrices.clabels, dtype=np.int64) + + flat_con_to_dual = _build_flat_con_to_dual_label_lookup(m, dual_vars) + if not len(flat_con_to_dual): + logger.warning( + "No valid constraint labels found, skipping dual feasibility constraints." + ) + return + + flat_v, flat_d, nnz_data = _extract_dual_feas_entries( + A, vlabels, clabels, flat_con_to_dual + ) + c_lookup = _build_obj_coeff_lookup(vlabels, m.matrices.c) + + logger.debug("Building dual feasibility constraints for each primal variable.") + for var_name, var in m.variables.items(): + if _skip(var.labels, "variable", var_name): + continue + + mask = var.labels != -1 + var_flat = var.labels.values.ravel().astype(np.int64) + + c_arr = _gather_with_default(var_flat, c_lookup, 0.0, np.float64) + c_vals = xr.DataArray( + c_arr.reshape(var.labels.shape), + coords=var.labels.coords, + dims=var.labels.dims, + ) + + lhs = _build_dual_feas_lhs(var, flat_v, flat_d, nnz_data, m_dual) + + if lhs.is_constant: + if np.any(c_arr[mask.values.ravel()] != 0): + logger.warning( + f"Variable '{var_name}' has no constraint connections but has " + "nonzero objective coefficients; the corresponding dual-feasibility " + "condition is infeasible." + ) + else: + logger.debug( + f"Variable '{var_name}' has no constraint connections and zero " + "objective coefficients; skipping redundant dual feasibility constraint." + ) + continue + + m_dual.add_constraints(lhs == c_vals, name=var_name, mask=mask) + + +def _add_dual_objective( + m: Model, + m_dual: Model, + dual_vars: dict, +) -> None: + """ + Construct and add ``sum(rhs_i * lambda_i)`` as the dual objective of m_dual. + + The uniform ``+`` sign is correct because sign conventions are already + encoded in the dual variable bounds: a ``<=`` constraint has ``lambda <= 0``, + so ``rhs * lambda`` contributes negatively without an explicit sign flip. + This aligns with linopy's native dual convention, so + ``m_dual.variables[con_name].solution`` can be compared directly with + ``m.constraints[con_name].dual`` after solving. + + The objective sense is flipped: ``min`` primal → ``max`` dual, and vice versa. + + Parameters + ---------- + m : Model + Primal model. + m_dual : Model + Dual model to populate. + dual_vars : dict + ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``. + """ + dual_obj: LinearExpression = LinearExpression(None, m_dual) + sense = "max" if m.objective.sense == "min" else "min" + + for name, con in m.constraints.items(): + if name not in dual_vars: + continue + + mask = con.labels != -1 + rhs_masked = con.rhs.where(mask, 0) + dual_obj += (rhs_masked * dual_vars[name]).sum() + + logger.debug(f"Constructed dual objective with {len(dual_obj.coeffs)} terms.") + logger.debug("Adding dual objective to model.") + m_dual.add_objective(dual_obj, sense=sense, overwrite=True) + + +def dualize( + m: Model, +) -> Model: + """ + Construct the dual of a linopy LP model. + + Transforms the primal model into its dual equivalent m_dual following + standard LP duality theory. The dual sense is flipped relative to the + primal (min -> max, max -> min), and dual variable bounds depend on + both constraint type and primal objective sense. + + For a minimization primal: + + Primal (min): + min c^T x + s.t. A_eq x = b_eq : λ free + A_leq x <= b_leq : μ <= 0 + A_geq x >= b_geq : ν >= 0 + + Dual (max): + max b_eq^T λ + b_leq^T μ + b_geq^T ν + s.t. A_eq^T λ + A_leq^T μ + A_geq^T ν = c + λ free, μ <= 0, ν >= 0 + + For a maximization primal the dual variable bounds are flipped: + μ >= 0 for <= constraints, ν <= 0 for >= constraints. + + The corresponding bound conventions are: + + ============ =========== ================ ================ + Constraint Primal sense lower upper + ============ =========== ================ ================ + = min / max -inf +inf (free) + <= min -inf 0 + <= max 0 +inf + >= min 0 +inf + >= max -inf 0 + ============ =========== ================ ================ + + Variable bounds are converted to explicit constraints before dualization + via _lift_bounds_to_constraints(), so that they appear in the constraint matrix + A and are correctly reflected in the dual. + + The dual variables in m_dual are named identically to their corresponding + primal constraints and are accessible via m_dual.variables[con_name]. + + Strong duality guarantees that at optimality: + primal objective = dual objective + + Note: This constructs a standalone dual model. Pathological or unsupported + primal formulations may lead to infeasible or unbounded dual models or may + cause this function to raise an error. Only linear primal models with linear + objectives and linear constraints are supported; finite variable bounds are + lifted into explicit constraints before dualization. + + Parameters + ---------- + m : Model + Primal linopy model to dualize. Must have a linear objective and either + linear constraints or finite variable bounds. + + Returns + ------- + Model + Dual model whose variables are named after the primal constraints and + constraints are named after the primal variables. + + Example + ------- + .. code-block:: python + + m_dual = m.dualize() + m_dual.solve(solver_name="gurobi", Method=2, Crossover=1) + gap = abs(m.objective.value - m_dual.objective.value) + """ + from linopy.model import Model + + m1 = m.copy() + m_dual = Model() + + if not m1.variables: + logger.warning("Primal model has no variables. Returning empty dual model.") + return m_dual + + _lift_bounds_to_constraints(m1) + + if not m1.constraints: + logger.warning( + "Primal model has no constraints after lifting variable bounds. " + "Returning empty dual model." + ) + return m_dual + + dual_vars = _add_dual_variables(m1, m_dual) + _add_dual_feasibility_constraints(m1, m_dual, dual_vars) + _add_dual_objective(m1, m_dual, dual_vars) + return m_dual diff --git a/linopy/expressions.py b/linopy/expressions.py index 10e243de8..c59f05d7d 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -9,11 +9,12 @@ 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 from itertools import product, zip_longest -from typing import TYPE_CHECKING, Any, TypeVar, overload +from typing import TYPE_CHECKING, Any, Self, TypeAlias, TypeVar, cast, overload from warnings import warn import numpy as np @@ -29,6 +30,7 @@ from xarray import Coordinates, DataArray, Dataset, IndexVariable from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates from xarray.core.indexes import Indexes +from xarray.core.types import JoinOptions from xarray.core.utils import Frozen try: @@ -37,21 +39,22 @@ from xarray.computation.rolling import DatasetRolling except ImportError: import xarray.core.rolling - from xarray.core.rolling import DatasetRolling # type: ignore + from xarray.core.rolling import DatasetRolling # type: ignore[no-redef] from types import EllipsisType, NotImplementedType from linopy import constraints, variables +from linopy.alignment import as_dataarray, broadcast_to_coords, fill_missing_coords from linopy.common import ( EmptyDeprecationWrapper, LocIndexer, - as_dataarray, assign_multiindex_safe, check_common_keys_values, check_has_nulls, check_has_nulls_polars, - fill_missing_coords, filter_nulls_polars, + format_coord, + format_single_expression, forward_as_properties, generate_indices_for_printout, get_dims_with_index_levels, @@ -60,8 +63,7 @@ has_optimized_model, is_constant, iterate_slices, - print_coord, - print_single_expression, + maybe_group_terms_polars, to_dataframe, to_polars, ) @@ -79,6 +81,7 @@ TERM_DIM, ) from linopy.types import ( + CONSTANT_TYPES, ConstantLike, DimsLike, ExpressionLike, @@ -88,21 +91,13 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + ) 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} @@ -141,6 +136,67 @@ def _expr_unwrap( logger = logging.getLogger(__name__) +def _resolve_group(group: Any, data: Dataset) -> Any: + """ + Normalize a groupby key. + + Unwrap a single-element key list to the scalar key, and resolve a string + naming a coordinate to that coordinate -- so ``groupby("name")`` behaves + like ``groupby(data["name"])``, mirroring xarray. Other inputs (Series, + DataFrame, DataArray, multi-key lists) are returned unchanged. + """ + if isinstance(group, (list, tuple)) and len(group) == 1: + group = group[0] + if isinstance(group, str) and group in data.coords: + group = data[group] + return group + + +def _multikey_value_frame(group: Any, data: Dataset) -> pd.DataFrame | None: + """ + Gather a multi-key list of coordinate names into a value frame. + + Return a DataFrame of the named coordinates when all keys are 1-D + coordinates sharing a single dimension -- so the list rides the fast + reindex path -- otherwise None. + """ + is_name_list = ( + isinstance(group, (list, tuple)) + and len(group) > 1 + and all(isinstance(g, str) and g in data.coords for g in group) + ) + if not is_name_list: + return None + coord_dims = {data[g].dims for g in group} + if len(coord_dims) != 1 or len(next(iter(coord_dims))) != 1: + return None + names = list(group) + return data[names].to_dataframe()[names] + + +def _unstack_multikey(ds: Dataset, dim: str) -> Dataset: + """ + Unstack a stacked multi-key group dimension into one dimension per key. + + Warn before materialising the grid when most cells would be fill values, + pointing to ``observed=True`` for a compact result. + """ + mi = ds.indexes[dim].remove_unused_levels() + observed = len(mi) + grid = int(np.prod([len(level) for level in mi.levels])) + if grid > 2 * observed and grid - observed > 10_000: + warn( + f"Grouping a LinearExpression by {list(mi.names)} produces a dense " + f"{grid:,}-cell grid, but only {observed:,} of those combinations " + f"occur -- the {grid - observed:,} absent ones are materialised as " + f"fill values. Pass `observed=True` to keep the result compact over " + f"only the observed combinations.", + UserWarning, + stacklevel=3, + ) + return ds.unstack(dim, fill_value=LinearExpression._fill_value) + + @dataclass @forward_as_properties(groupby=["dims", "groups"]) class LinearExpressionGroupby: @@ -163,17 +219,25 @@ def groupby(self) -> xarray.core.groupby.DatasetGroupBy: xarray.core.groupby.DataArrayGroupBy The groupby object. """ - if isinstance(self.group, pd.DataFrame): + data = self.data + group = _resolve_group(self.group, data) + + if isinstance(group, pd.DataFrame): raise ValueError( "Grouping by a DataFrame only supported for `sum` operation with `use_fallback=False`." ) - if isinstance(self.group, pd.Series): - group_name = self.group.name or "group" - group = DataArray(self.group, name=group_name) - else: - group = self.group # type: ignore + if isinstance(group, pd.Series): + group = DataArray(group, name=group.name or "group") + + # detach an attached free coordinate (never an indexed/level coord) + if ( + isinstance(group, DataArray) + and group.name in set(data.coords) - set(data.dims) + and group.name not in data.xindexes + ): + data = data.drop_vars([group.name]) - return self.data.groupby(group=group, **self.kwargs) + return data.groupby(group=group, **self.kwargs) def map( self, @@ -205,7 +269,9 @@ def map( self.groupby.map(func, shortcut=shortcut, args=args, **kwargs), self.model ) - def sum(self, use_fallback: bool = False, **kwargs: Any) -> LinearExpression: + def sum( + self, use_fallback: bool = False, observed: bool = False, **kwargs: Any + ) -> LinearExpression: """ Sum the groupby object. @@ -223,6 +289,13 @@ def sum(self, use_fallback: bool = False, **kwargs: Any) -> LinearExpression: Whether to use the fallback implementation, which is a sort of default xarray implementation. If set to False, the operation will be much faster but keyword arguments are ignored. Defaults to False. + observed : bool + Only applies when grouping by a list of coordinate names. If True, + keep the result stacked over the observed key combinations (a + ``MultiIndex`` ``group`` dimension) instead of unstacking into one + dimension per key, which materialises the dense cartesian grid. + Defaults to False, mirroring xarray. Not supported together with + `use_fallback`. **kwargs Arbitrary keyword arguments. @@ -231,9 +304,22 @@ def sum(self, use_fallback: bool = False, **kwargs: Any) -> LinearExpression: LinearExpression The sum of the groupby object. """ + if observed and use_fallback: + raise ValueError( + "`observed=True` is not supported with `use_fallback=True`." + ) + + group = _resolve_group(self.group, self.data) + + # a list of coord names rides the fast path as a value frame + multikey_frame = ( + None if use_fallback else _multikey_value_frame(group, self.data) + ) + if multikey_frame is not None: + group = multikey_frame + non_fallback_types = (pd.Series, pd.DataFrame, xr.DataArray) - if isinstance(self.group, non_fallback_types) and not use_fallback: - group: pd.Series | pd.DataFrame | xr.DataArray = self.group + if isinstance(group, non_fallback_types) and not use_fallback: if isinstance(group, pd.DataFrame): # dataframes do not have a name, so we need to set it final_group_name = "group" @@ -259,10 +345,12 @@ def sum(self, use_fallback: bool = False, **kwargs: Any) -> LinearExpression: arrays = [group, group.groupby(group).cumcount()] idx = pd.MultiIndex.from_arrays(arrays, names=[GROUP_DIM, GROUPED_TERM_DIM]) new_coords = Coordinates.from_pandas_multiindex(idx, group_dim) - coords = self.data.indexes[group_dim] - names_to_drop = [coords.name] - if isinstance(coords, pd.MultiIndex): - names_to_drop += list(coords.names) + # collapsing group_dim invalidates every coordinate aligned to it + names_to_drop = [ + name + for name, coord in self.data.coords.items() + if group_dim in coord.dims + ] ds = self.data.drop_vars(names_to_drop).assign_coords(new_coords) ds = ds.unstack(group_dim, fill_value=LinearExpression._fill_value) ds = LinearExpression._sum(ds, dim=GROUPED_TERM_DIM) @@ -272,9 +360,11 @@ def sum(self, use_fallback: bool = False, **kwargs: Any) -> LinearExpression: index.names = [str(col) for col in orig_group.columns] index.name = GROUP_DIM new_coords = Coordinates.from_pandas_multiindex(index, GROUP_DIM) - ds = xr.Dataset(ds.assign_coords(new_coords)) + ds = ds.assign_coords(new_coords) ds = ds.rename({GROUP_DIM: final_group_name}) + if multikey_frame is not None and not observed: + ds = _unstack_multikey(ds, final_group_name) return LinearExpression(ds, self.model) def func(ds: Dataset) -> Dataset: @@ -342,25 +432,28 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: if data is None: da = xr.DataArray([], dims=[TERM_DIM]) data = Dataset({"coeffs": da, "vars": da, "const": 0.0}) - elif isinstance(data, SUPPORTED_CONSTANT_TYPES): + elif isinstance(data, CONSTANT_TYPES): const = as_dataarray(data) da = xr.DataArray([], dims=[TERM_DIM]) data = Dataset({"coeffs": da, "vars": da, "const": const}) elif not isinstance(data, Dataset): supported_types = ", ".join( - map(lambda s: s.__qualname__, (*SUPPORTED_CONSTANT_TYPES, Dataset)) + map(lambda s: s.__qualname__, (*CONSTANT_TYPES, Dataset)) ) raise ValueError( f"data must be an instance of {supported_types}, got {type(data)}" ) + data = cast(Dataset, data) if not set(data).issuperset({"coeffs", "vars"}): raise ValueError( "data must contain the fields 'coeffs' and 'vars' or 'const'" ) if np.issubdtype(data.vars, np.floating): - data = assign_multiindex_safe(data, vars=data.vars.fillna(-1).astype(int)) + data = assign_multiindex_safe( + data, vars=data.vars.fillna(-1).astype(options["label_dtype"]) + ) if not np.issubdtype(data.coeffs, np.floating): data["coeffs"].values = data.coeffs.values.astype(float) @@ -375,12 +468,12 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: data = assign_multiindex_safe(data, const=data.const.astype(float)) (data,) = xr.broadcast(data, exclude=HELPER_DIMS) + data = cast(Dataset, data) (coeffs_vars,) = xr.broadcast(data[["coeffs", "vars"]], exclude=[FACTOR_DIM]) coeffs_vars_dict = {str(k): v for k, v in coeffs_vars.items()} data = assign_multiindex_safe(data, **coeffs_vars_dict) - # transpose with new Dataset to really ensure correct order - data = Dataset(data.transpose(..., TERM_DIM)) + data = data.transpose(..., TERM_DIM) # ensure helper dimensions are not set as coordinates if drop_dims := set(HELPER_DIMS).intersection(data.coords): @@ -391,7 +484,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: raise ValueError("model must be an instance of linopy.Model") self._model = model - self._data = data + self._data = cast(Dataset, data) def __repr__(self) -> str: """ @@ -417,16 +510,16 @@ def __repr__(self) -> str: self.data.indexes[dims[i]][ind] for i, ind in enumerate(indices) ] if self.mask is None or self.mask.values[indices]: - expr = print_single_expression( + expr = format_single_expression( self.coeffs.values[indices], self.vars.values[indices], self.const.values[indices], self.model, ) - line = print_coord(coord) + f": {expr}" + line = format_coord(coord) + f": {expr}" else: - line = print_coord(coord) + ": None" + line = format_coord(coord) + ": None" lines.append(line) shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) @@ -434,7 +527,7 @@ def __repr__(self) -> str: underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") elif size == 1: - expr = print_single_expression( + expr = format_single_expression( self.coeffs.values, self.vars.values, self.const.item(), self.model ) lines.append(f"{header_string}\n{'-' * len(header_string)}\n{expr}") @@ -466,40 +559,30 @@ def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None print(self) @abstractmethod - def __add__( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: ... + def __add__(self, other: SideLike) -> Self | QuadraticExpression: ... @abstractmethod - def __radd__(self: GenericExpression, other: SideLike) -> GenericExpression: ... + def __radd__(self, other: SideLike) -> Self: ... @abstractmethod - def __sub__( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: ... + def __sub__(self, other: SideLike) -> Self | QuadraticExpression: ... @abstractmethod - def __rsub__(self: GenericExpression, other: SideLike) -> GenericExpression: ... + def __rsub__(self, other: SideLike) -> Self: ... @abstractmethod - def __mul__( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: ... + def __mul__(self, other: SideLike) -> Self | QuadraticExpression: ... @abstractmethod - def __rmul__( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: ... + def __rmul__(self, other: SideLike) -> Self | QuadraticExpression: ... @abstractmethod - def __matmul__( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: ... + def __matmul__(self, other: SideLike) -> Self | QuadraticExpression: ... @abstractmethod def __pow__(self, other: int) -> QuadraticExpression: ... - def __neg__(self: GenericExpression) -> GenericExpression: + def __neg__(self) -> Self: """ Get the negative of the expression. """ @@ -507,54 +590,158 @@ def __neg__(self: GenericExpression) -> GenericExpression: def _multiply_by_linear_expression( self, other: LinearExpression | ScalarLinearExpression - ) -> QuadraticExpression: + ) -> LinearExpression | QuadraticExpression: if isinstance(other, ScalarLinearExpression): other = other.to_linexpr() if other.nterm > 1: raise TypeError("Multiplication of multiple terms is not supported.") + + if other.is_constant: + return cast(LinearExpression, self._multiply_by_constant(other.const)) + if self.is_constant: + return cast(LinearExpression, other._multiply_by_constant(self.const)) + # multiplication: (v1 + c1) * (v2 + c2) = v1 * v2 + c1 * v2 + c2 * v1 + c1 * c2 # with v being the variables and c the constants # merge on factor dimension only returns v1 * v2 + c1 * c2 ds = other.data[["coeffs", "vars"]].sel(_term=0).broadcast_like(self.data) ds = assign_multiindex_safe(ds, const=other.const) - res = merge([self, ds], dim=FACTOR_DIM, cls=QuadraticExpression) # type: ignore + res = merge([self, ds], dim=FACTOR_DIM, cls=QuadraticExpression) # deal with cross terms c1 * v2 + c2 * v1 if self.has_constant: res = res + self.const * other.reset_const() if other.has_constant: res = res + self.reset_const() * other.const - return res + return cast(QuadraticExpression, res) + + def _align_constant( + self, + other: DataArray, + fill_value: float = 0, + join: JoinOptions | 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, other: ConstantLike, join: JoinOptions | None = None + ) -> Self: + # NaN values in self.const or other are filled with 0 (additive identity) + # so that missing data does not silently propagate through arithmetic. + if np.isscalar(other) and join is None: + return self.assign(const=self.const.fillna(0) + other) + da = broadcast_to_coords( + other, coords=self.coords, dims=self.coord_dims, strict=False + ) + self_const, da, needs_data_reindex = self._align_constant( + da, fill_value=0, join=join + ) + da = da.fillna(0) + self_const = self_const.fillna(0) + 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, + other: ConstantLike, + op: Callable[[DataArray, DataArray], DataArray], + fill_value: float, + join: JoinOptions | None = None, + ) -> Self: + """ + Apply a constant operation (mul, div, etc.) to this expression with a scalar or array. + + NaN values are filled with neutral elements before the operation: + - factor (other) is filled with fill_value (0 for mul, 1 for div) + - coeffs and const are filled with 0 (additive identity) + """ + factor = broadcast_to_coords( + other, coords=self.coords, dims=self.coord_dims, strict=False + ) + self_const, factor, needs_data_reindex = self._align_constant( + factor, fill_value=fill_value, join=join + ) + factor = factor.fillna(fill_value) + self_const = self_const.fillna(0) + if needs_data_reindex: + data = self.data.reindex_like(self_const, fill_value=self._fill_value) + coeffs = data.coeffs.fillna(0) + return self.__class__( + assign_multiindex_safe( + data, coeffs=op(coeffs, factor), const=op(self_const, factor) + ), + self.model, + ) + coeffs = self.coeffs.fillna(0) + return self.assign(coeffs=op(coeffs, factor), const=op(self_const, factor)) def _multiply_by_constant( - self: GenericExpression, other: ConstantLike - ) -> 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) - - def __div__(self: GenericExpression, other: SideLike) -> GenericExpression: + self, other: ConstantLike, join: JoinOptions | None = None + ) -> Self: + return self._apply_constant_op(other, operator.mul, fill_value=0, join=join) + + def _divide_by_constant( + self, other: ConstantLike, join: JoinOptions | None = None + ) -> Self: + return self._apply_constant_op(other, operator.truediv, fill_value=1, join=join) + + def __div__(self, other: SideLike) -> Self: 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 - def __truediv__(self: GenericExpression, other: SideLike) -> GenericExpression: + def __truediv__(self, other: SideLike) -> Self: return self.__div__(other) def __le__(self, rhs: SideLike) -> Constraint: @@ -563,7 +750,7 @@ def __le__(self, rhs: SideLike) -> Constraint: def __ge__(self, rhs: SideLike) -> Constraint: return self.to_constraint(GREATER_EQUAL, rhs) - def __eq__(self, rhs: SideLike) -> Constraint: # type: ignore + def __eq__(self, rhs: SideLike) -> Constraint: # type: ignore[override] return self.to_constraint(EQUAL, rhs) def __gt__(self, other: Any) -> NotImplementedType: @@ -577,36 +764,160 @@ def __lt__(self, other: Any) -> NotImplementedType: ) def add( - self: GenericExpression, other: SideLike - ) -> GenericExpression | QuadraticExpression: + self, + other: SideLike, + join: JoinOptions | None = None, + ) -> Self | 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, 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 - ) -> GenericExpression | QuadraticExpression: + self, + other: SideLike, + join: JoinOptions | None = None, + ) -> Self | 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 - ) -> GenericExpression | QuadraticExpression: + self, + other: SideLike, + join: JoinOptions | None = None, + ) -> Self | 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 - ) -> GenericExpression | QuadraticExpression: + self, + other: VariableLike | ConstantLike, + join: JoinOptions | None = None, + ) -> Self | 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, + rhs: SideLike, + join: JoinOptions | 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, + rhs: SideLike, + join: JoinOptions | 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, + rhs: SideLike, + join: JoinOptions | 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: """ @@ -614,17 +925,13 @@ def pow(self, other: int) -> QuadraticExpression: """ return self.__pow__(other) - def dot( - self: GenericExpression, other: ndarray - ) -> GenericExpression | QuadraticExpression: + def dot(self, other: ndarray) -> Self | QuadraticExpression: """ Matrix multiplication with other, similar to xarray dot. """ return self.__matmul__(other) - def __getitem__( - self: GenericExpression, selector: int | tuple[slice, list[int]] | slice - ) -> GenericExpression: + def __getitem__(self, selector: int | tuple[slice, list[int]] | slice) -> Self: """ Get selection from the expression. This is a wrapper around the xarray __getitem__ method. It returns a @@ -703,6 +1010,7 @@ def coord_names(self) -> list[str]: @property def vars(self) -> DataArray: + """Variable labels referenced by each term of the expression.""" return self.data.vars @vars.setter @@ -711,6 +1019,7 @@ def vars(self, value: DataArray) -> None: @property def coeffs(self) -> DataArray: + """Coefficient applied to each term of the expression.""" return self.data.coeffs @coeffs.setter @@ -719,6 +1028,7 @@ def coeffs(self, value: DataArray) -> None: @property def const(self) -> DataArray: + """Constant offset added to the expression.""" return self.data.const @const.setter @@ -740,10 +1050,11 @@ def _map_solution(self) -> DataArray: Replace variable labels by solution values. """ m = self.model - sol = pd.Series(m.matrices.sol, m.matrices.vlabels) + M = m.matrices + sol = pd.Series(M.sol, M.vlabels) sol[-1] = np.nan idx = np.ravel(self.vars) - values = sol[idx].to_numpy().reshape(self.vars.shape) + values = np.asarray(sol[idx]).reshape(self.vars.shape) return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords) @property @@ -759,11 +1070,11 @@ def solution(self) -> DataArray: return sol.rename("solution") def sum( - self: GenericExpression, + self, dim: DimsLike | None = None, drop_zeros: bool = False, **kwargs: Any, - ) -> GenericExpression: + ) -> Self: """ Sum the expression over all or a subset of dimensions. @@ -847,7 +1158,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: JoinOptions | None = None + ) -> Constraint: """ Convert a linear expression to a constraint. @@ -856,7 +1169,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 ------- @@ -869,13 +1189,42 @@ 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, CONSTANT_TYPES): + rhs = broadcast_to_coords( + rhs, coords=self.coords, dims=self.coord_dims, strict=False + ) + + extra_dims = set(rhs.dims) - set(self.coord_dims) + if extra_dims: + logger.warning( + f"Constant RHS contains dimensions {extra_dims} not present " + f"in the expression, which might lead to inefficiencies. " + f"Consider collapsing the dimensions by taking min/max." + ) + rhs = rhs.reindex_like(self.const, fill_value=np.nan) + + # Remember where RHS is NaN (meaning "no constraint") before the + # subtraction, which may fill NaN with 0 as part of normal + # expression arithmetic. + if isinstance(rhs, DataArray): + rhs_nan_mask = rhs.isnull() + else: + rhs_nan_mask = None + + all_to_lhs = self.sub(rhs, join=join).data + computed_rhs = -all_to_lhs.const + + # Restore NaN at positions where the original constant RHS had no + # value so that downstream code still treats them as unconstrained. + if rhs_nan_mask is not None and rhs_nan_mask.any(): + computed_rhs = xr.where(rhs_nan_mask, np.nan, computed_rhs) + data = assign_multiindex_safe( - all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=-all_to_lhs.const + all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=computed_rhs ) return constraints.Constraint(data, model=self.model) - def reset_const(self: GenericExpression) -> GenericExpression: + def reset_const(self) -> Self: """ Reset the constant of the linear expression to zero. """ @@ -893,7 +1242,7 @@ def isnull(self) -> DataArray: return (self.vars == -1).all(helper_dims) & self.const.isnull() def where( - self: GenericExpression, + self, cond: DataArray, other: LinearExpression | int @@ -901,7 +1250,7 @@ def where( | dict[str, float | int | DataArray] | None = None, **kwargs: Any, - ) -> GenericExpression: + ) -> Self: """ Filter variables based on a condition. @@ -945,14 +1294,14 @@ def where( return self.__class__(self.data.where(cond, other=_other, **kwargs), self.model) def fillna( - self: GenericExpression, + self, value: int | float | DataArray | Dataset | LinearExpression | dict[str, float | int | DataArray], - ) -> GenericExpression: + ) -> Self: """ Fill missing values with a given value. @@ -976,7 +1325,7 @@ def fillna( value = {"const": value} return self.__class__(self.data.fillna(value), self.model) - def diff(self: GenericExpression, dim: str, n: int = 1) -> GenericExpression: + def diff(self, dim: str, n: int = 1) -> Self: """ Calculate the n-th order discrete difference along given axis. @@ -1071,6 +1420,57 @@ def nterm(self) -> int: """ return len(self.data._term) + @property + def has_terms(self) -> DataArray: + """ + Get a boolean array which is true at slots with at least one live term. + + A term is live when it references a variable (``vars != -1``). Slots + without any live term arise from outer joins in + :func:`merge `, from reindexing past the + original coordinates, or from masking. In contrast to + :meth:`isnull`, the constant is ignored: a slot carrying only a + constant has no terms. + + Returns + ------- + xr.DataArray + + Examples + -------- + Mask out constraint rows whose left-hand side has no terms: + + >>> import linopy + >>> import pandas as pd + >>> m = linopy.Model() + >>> x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + >>> lhs = (1 * x).reindex(i=pd.RangeIndex(5, name="i")) + >>> lhs.has_terms.values + array([ True, True, True, False, False]) + """ + helper_dims = set(self.vars.dims).intersection(HELPER_DIMS) + return (self.vars != -1).any(helper_dims).rename("has_terms") + + @property + def variable_names(self) -> set[str]: + """ + Get the names of the unique variables present in the expression. + """ + if self.nterm == 0: + return set() + + # Collect all unique labels from the expression (excluding -1) + all_labels = self.vars.values.ravel() + unique_labels = np.unique(all_labels[all_labels != -1]) + + if len(unique_labels) == 0: + return set() + + # Batch lookup variable names for all labels + positions = self.model.variables.get_label_position(unique_labels) + + return {p[0] for p in positions if p[0] is not None} + @property def shape(self) -> tuple[int, ...]: """ @@ -1097,7 +1497,7 @@ def empty(self) -> EmptyDeprecationWrapper: """ return EmptyDeprecationWrapper(not self.size) - def densify_terms(self: GenericExpression) -> GenericExpression: + def densify_terms(self) -> Self: """ Move all non-zero term entries to the front and cut off all-zero entries in the term-axis. @@ -1128,7 +1528,7 @@ def densify_terms(self: GenericExpression) -> GenericExpression: return self.__class__(data.sel({TERM_DIM: slice(0, nterm)}), self.model) - def sanitize(self: GenericExpression) -> GenericExpression: + def sanitize(self) -> Self: """ Sanitize LinearExpression by ensuring int dtype for variables. @@ -1137,7 +1537,7 @@ def sanitize(self: GenericExpression) -> GenericExpression: linopy.LinearExpression """ if not np.issubdtype(self.vars.dtype, np.integer): - return self.assign(vars=self.vars.fillna(-1).astype(int)) + return self.assign(vars=self.vars.fillna(-1).astype(options["label_dtype"])) return self @@ -1305,11 +1705,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, 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 @@ -1426,7 +1826,7 @@ def flat(self) -> pd.DataFrame: """ ds = self.data - def mask_func(data: pd.DataFrame) -> pd.Series: + def mask_func(data: dict) -> pd.Series: mask = (data["vars"] != -1) & (data["coeffs"] != 0) return mask @@ -1463,7 +1863,7 @@ def to_polars(self) -> pl.DataFrame: df = to_polars(self.data) df = filter_nulls_polars(df) - df = group_terms_polars(df) + df = maybe_group_terms_polars(df) check_has_nulls_polars(df, name=self.type) return df @@ -1541,12 +1941,12 @@ def _simplify_row(vars_row: np.ndarray, coeffs_row: np.ndarray) -> np.ndarray: # Combined has dimensions (.., CV_DIM, TERM_DIM) # Drop terms where all vars are -1 (i.e., empty terms across all coordinates) - vars = combined.isel({CV_DIM: 0}).astype(int) + vars = combined.isel({CV_DIM: 0}).astype(options["label_dtype"]) non_empty_terms = (vars != -1).any(dim=[d for d in vars.dims if d != TERM_DIM]) combined = combined.isel({TERM_DIM: non_empty_terms}) # Extract vars and coeffs from the combined result - vars = combined.isel({CV_DIM: 0}).astype(int) + vars = combined.isel({CV_DIM: 0}).astype(options["label_dtype"]) coeffs = combined.isel({CV_DIM: 1}) # Create new dataset with simplified data @@ -1585,7 +1985,7 @@ def from_rule( cls, model: Model, rule: Callable, - coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, + coords: Sequence[Sequence | pd.Index] | Mapping | None = None, ) -> LinearExpression: """ Create a linear expression from a rule and a set of coordinates. @@ -1703,7 +2103,7 @@ def process_one( ) -> LinearExpression: nonlocal model - if isinstance(t, SUPPORTED_CONSTANT_TYPES): + if isinstance(t, CONSTANT_TYPES): if model is None: raise ValueError("Model must be provided when using constants.") expr = LinearExpression(t, model) @@ -1722,7 +2122,7 @@ def process_one( ) expr = v.to_linexpr(c) elif isinstance(v, variables.Variable): - if not isinstance(c, SUPPORTED_CONSTANT_TYPES): + if not isinstance(c, CONSTANT_TYPES): raise TypeError( "Expected constant as coefficient of variable (first element of tuple)." ) @@ -1799,7 +2199,7 @@ def __init__(self, data: Dataset | None, model: Model) -> None: raise ValueError(f"Size of dimension {FACTOR_DIM} must be 2.") # transpose data to have _term as last dimension and _factor as second last - data = xr.Dataset(data.transpose(..., FACTOR_DIM, TERM_DIM)) + data = data.transpose(..., FACTOR_DIM, TERM_DIM) self._data = data @property @@ -1817,13 +2217,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)}. " @@ -1845,15 +2239,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, 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 @@ -1871,13 +2265,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 @@ -1899,13 +2287,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." ) @@ -1926,7 +2308,9 @@ def solution(self) -> DataArray: sol = (self.coeffs * vals.prod(FACTOR_DIM)).sum(TERM_DIM) + self.const return sol.rename("solution") - def to_constraint(self, sign: SignLike, rhs: SideLike) -> NotImplementedType: + def to_constraint( + self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None + ) -> NotImplementedType: raise NotImplementedError( "Quadratic expressions cannot be used in constraints." ) @@ -1941,7 +2325,7 @@ def flat(self) -> DataFrame: ).to_dataset(FACTOR_DIM) ds = self.data.drop_vars("vars").assign(vars) - def mask_func(data: pd.DataFrame) -> pd.Series: + def mask_func(data: dict) -> pd.Series: mask = ((data["vars1"] != -1) | (data["vars2"] != -1)) & ( data["coeffs"] != 0 ) @@ -2026,7 +2410,7 @@ def as_expression( model : linopy.Model, optional Assigned model, by default None **kwargs : - Keyword arguments passed to `linopy.as_dataarray`. + Keyword arguments passed to `linopy.alignment.broadcast_to_coords`. Returns ------- @@ -2043,23 +2427,45 @@ def as_expression( return obj.to_linexpr() else: try: - obj = as_dataarray(obj, **kwargs) + obj = broadcast_to_coords(obj, strict=False, **kwargs) except ValueError as e: raise ValueError("Cannot convert to LinearExpression") from e return LinearExpression(obj, model) +Mergeable: TypeAlias = BaseExpression | variables.Variable | Dataset + + +@overload def merge( - exprs: Sequence[ - LinearExpression | QuadraticExpression | variables.Variable | Dataset - ], - *add_exprs: tuple[ - LinearExpression | QuadraticExpression | variables.Variable | Dataset - ], + exprs: Sequence[Mergeable] | Mergeable, + *add_exprs: Mergeable, + dim: str = ..., + cls: type[GenericExpression], + join: JoinOptions | None = None, + **kwargs: Any, +) -> GenericExpression: ... + + +@overload +def merge( + exprs: Sequence[Mergeable] | Mergeable, + *add_exprs: Mergeable, + dim: str = ..., + cls: None = ..., + join: JoinOptions | None = None, + **kwargs: Any, +) -> BaseExpression: ... + + +def merge( + exprs: Sequence[Mergeable] | Mergeable, + *add_exprs: Mergeable, dim: str = TERM_DIM, - cls: type[GenericExpression] = None, # type: ignore + cls: type[BaseExpression] | None = None, + join: JoinOptions | None = None, **kwargs: Any, -) -> GenericExpression: +) -> BaseExpression: """ Merge multiple expression together. @@ -2077,6 +2483,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 @@ -2086,32 +2496,38 @@ def merge( ------- res : linopy.LinearExpression or linopy.QuadraticExpression """ - if not isinstance(exprs, list) and len(add_exprs): + if not isinstance(exprs, Sequence): warn( "Passing a tuple to the merge function is deprecated. Please pass a list of objects to be merged", DeprecationWarning, ) - exprs = [exprs] + list(add_exprs) # type: ignore + exprs = [exprs] + list(add_exprs) - has_quad_expression = any(type(e) is QuadraticExpression for e in exprs) - has_linear_expression = any(type(e) is LinearExpression for e in exprs) + has_quad_expression = any(isinstance(e, QuadraticExpression) for e in exprs) + has_linear_expression = any(isinstance(e, LinearExpression) for e in exprs) if cls is None: cls = QuadraticExpression if has_quad_expression else LinearExpression - if cls is QuadraticExpression and dim == TERM_DIM and has_linear_expression: + if ( + issubclass(cls, QuadraticExpression) + and dim == TERM_DIM + and has_linear_expression + ): raise ValueError( "Cannot merge linear and quadratic expressions along term dimension." "Convert to QuadraticExpression first." ) - if has_quad_expression and cls is not QuadraticExpression: + if has_quad_expression and not issubclass(cls, QuadraticExpression): raise ValueError("Cannot merge linear expressions to QuadraticExpression") - linopy_types = (variables.Variable, LinearExpression, QuadraticExpression) + linopy_types = (variables.Variable, BaseExpression) model = exprs[0].model - if cls in linopy_types and dim in HELPER_DIMS: + if join is not None: + override = join == "override" + elif issubclass(cls, 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 ] @@ -2127,12 +2543,14 @@ def merge( "coords": "minimal", "compat": "override", } - if cls == LinearExpression: + if issubclass(cls, LinearExpression): kwargs["fill_value"] = FILL_VALUE - elif cls == variables.Variable: + elif issubclass(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") @@ -2178,7 +2596,7 @@ def __init__( self._model = model def __repr__(self) -> str: - expr_string = print_single_expression( + expr_string = format_single_expression( np.array(self.coeffs), np.array(self.vars), 0, self.model ) return f"ScalarLinearExpression: {expr_string}" @@ -2287,7 +2705,9 @@ def __ge__(self, other: int | float) -> AnonymousScalarConstraint: return constraints.AnonymousScalarConstraint(self, GREATER_EQUAL, other) - def __eq__(self, other: int | float) -> AnonymousScalarConstraint: # type: ignore + def __eq__( # type: ignore[override] + self, other: int | float + ) -> AnonymousScalarConstraint: if not isinstance(other, int | float | np.number): raise TypeError( f"unsupported operand type(s) for ==: {type(self)} and {type(other)}" @@ -2310,3 +2730,11 @@ 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_EXPRESSION_TYPES = ( + BaseExpression, + ScalarLinearExpression, + variables.Variable, + variables.ScalarVariable, +) diff --git a/linopy/io.py b/linopy/io.py index 56fe033dd..bccbdecf7 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -5,11 +5,14 @@ from __future__ import annotations +import copy as _copy +import json import logging import shutil import time import warnings -from collections.abc import Callable +from collections.abc import Callable, Iterable +from importlib.metadata import version from io import BufferedWriter from pathlib import Path from tempfile import TemporaryDirectory @@ -19,13 +22,11 @@ import pandas as pd import polars as pl import xarray as xr -from numpy import ones_like, zeros_like -from scipy.sparse import tril, triu from tqdm import tqdm from linopy import solvers from linopy.common import to_polars -from linopy.constants import CONCAT_DIM +from linopy.constants import CONCAT_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR from linopy.objective import Objective if TYPE_CHECKING: @@ -33,10 +34,13 @@ from highspy.highs import Highs from linopy.model import Model + from linopy.variables import Variable logger = logging.getLogger(__name__) +NETCDF_VERSION_ATTR = "_linopy_version" + ufunc_kwargs = dict(vectorize=True) concat_kwargs = dict(dim=CONCAT_DIM, coords="minimal") @@ -54,6 +58,21 @@ def clean_name(name: str) -> str: coord_sanitizer = str.maketrans("[,]", "(,)", " ") +def _format_and_write( + df: pl.DataFrame, columns: list[pl.Expr], f: BufferedWriter +) -> None: + """ + Format columns via concat_str and write to file. + + Uses Polars streaming engine for better memory efficiency. + """ + df.lazy().select(pl.concat_str(columns, ignore_nulls=True)).collect( + engine="streaming" + ).write_csv( + f, separator=" ", null_value="", quote_style="never", include_header=False + ) + + def signed_number(expr: pl.Expr) -> tuple[pl.Expr, pl.Expr]: """ Return polars expressions for a signed number string, handling -0.0 correctly. @@ -74,39 +93,50 @@ def signed_number(expr: pl.Expr) -> tuple[pl.Expr, pl.Expr]: ) -def print_coord(coord: str) -> str: - from linopy.common import print_coord +def format_coord(coord: str) -> str: + from linopy.common import format_coord - coord = print_coord(coord).translate(coord_sanitizer) + coord = format_coord(coord).translate(coord_sanitizer) return coord def get_printers_scalar( m: Model, explicit_coordinate_names: bool = False ) -> tuple[Callable, Callable]: - """Get printer functions for scalar values (non-polars).""" + """ + Get batch printer functions for numpy label arrays (non-polars). + + Returns two callables that take an int64 numpy array of labels and return + a list of name strings. + """ if explicit_coordinate_names: - def print_variable(var: Any) -> str: + def _fmt_var(var: Any) -> str: name, coord = m.variables.get_label_position(var) name = clean_name(name) - return f"{name}{print_coord(coord)}#{var}" + return f"{name}{format_coord(coord)}#{var}" - def print_constraint(cons: Any) -> str: + def _fmt_con(cons: Any) -> str: name, coord = m.constraints.get_label_position(cons) name = clean_name(name) # type: ignore - return f"{name}{print_coord(coord)}#{cons}" # type: ignore + return f"{name}{format_coord(coord)}#{cons}" # type: ignore + + def print_variables(labels: np.ndarray) -> list[str]: + return np.vectorize(_fmt_var)(labels).tolist() + + def print_constraints(labels: np.ndarray) -> list[str]: + return np.vectorize(_fmt_con)(labels).tolist() - return print_variable, print_constraint + return print_variables, print_constraints else: - def print_variable(var: Any) -> str: - return f"x{var}" + def print_variables(labels: np.ndarray) -> list[str]: + return ("x" + pl.Series(labels).cast(pl.String)).to_list() - def print_constraint(cons: Any) -> str: - return f"c{cons}" + def print_constraints(labels: np.ndarray) -> list[str]: + return ("c" + pl.Series(labels).cast(pl.String)).to_list() - return print_variable, print_constraint + return print_variables, print_constraints def get_printers( @@ -118,12 +148,12 @@ def get_printers( def print_variable(var: Any) -> str: name, coord = m.variables.get_label_position(var) name = clean_name(name) - return f"{name}{print_coord(coord)}#{var}" + return f"{name}{format_coord(coord)}#{var}" def print_constraint(cons: Any) -> str: name, coord = m.constraints.get_label_position(cons) name = clean_name(name) # type: ignore - return f"{name}{print_coord(coord)}#{cons}" # type: ignore + return f"{name}{format_coord(coord)}#{cons}" # type: ignore def print_variable_series(series: pl.Series) -> tuple[pl.Expr, pl.Series]: return pl.lit(" "), series.map_elements( @@ -155,10 +185,7 @@ def objective_write_linear_terms( *signed_number(pl.col("coeffs")), *print_variable(pl.col("vars")), ] - df = df.select(pl.concat_str(cols, ignore_nulls=True)) - df.write_csv( - f, separator=" ", null_value="", quote_style="never", include_header=False - ) + _format_and_write(df, cols, f) def objective_write_quadratic_terms( @@ -171,10 +198,7 @@ def objective_write_quadratic_terms( *print_variable(pl.col("vars2")), ] f.write(b"+ [\n") - df = df.select(pl.concat_str(cols, ignore_nulls=True)) - df.write_csv( - f, separator=" ", null_value="", quote_style="never", include_header=False - ) + _format_and_write(df, cols, f) f.write(b"] / 2\n") @@ -215,6 +239,17 @@ def objective_to_file( objective_write_quadratic_terms(f, quads, print_variable) +def _binary_has_nondefault_bounds(var: Variable) -> bool: + """ + Whether a binary variable carries bounds other than the implied (0, 1). + + Scans the raw bound values (a single vectorised pass each), so masked + slots are tolerated: a false positive only routes the variable through + the bounds loop, where masked labels are dropped before writing. + """ + return bool((var.lower.values != 0).any() or (var.upper.values != 1).any()) + + def bounds_to_file( m: Model, f: BufferedWriter, @@ -225,7 +260,16 @@ def bounds_to_file( """ Write out variables of a model to a lp file. """ - names = list(m.variables.continuous) + list(m.variables.integers) + names = ( + list(m.variables.continuous) + + list(m.variables.integers) + + list(m.variables.semi_continuous) + + [ + n + for n in m.variables.binaries + if _binary_has_nondefault_bounds(m.variables[n]) + ] + ) if not len(list(names)): return @@ -254,11 +298,7 @@ def bounds_to_file( *signed_number(pl.col("upper")), ] - kwargs: Any = dict( - separator=" ", null_value="", quote_style="never", include_header=False - ) - formatted = df.select(pl.concat_str(columns, ignore_nulls=True)) - formatted.write_csv(f, **kwargs) + _format_and_write(df, columns, f) def binaries_to_file( @@ -296,11 +336,45 @@ def binaries_to_file( *print_variable(pl.col("labels")), ] - kwargs: Any = dict( - separator=" ", null_value="", quote_style="never", include_header=False - ) - formatted = df.select(pl.concat_str(columns, ignore_nulls=True)) - formatted.write_csv(f, **kwargs) + _format_and_write(df, columns, f) + + +def semi_continuous_to_file( + m: Model, + f: BufferedWriter, + progress: bool = False, + slice_size: int = 2_000_000, + explicit_coordinate_names: bool = False, +) -> None: + """ + Write out semi-continuous variables of a model to a lp file. + """ + names = m.variables.semi_continuous + if not len(list(names)): + return + + print_variable, _ = get_printers( + m, explicit_coordinate_names=explicit_coordinate_names + ) + + f.write(b"\n\nsemi-continuous\n\n") + if progress: + names = tqdm( + list(names), + desc="Writing semi-continuous variables.", + colour=TQDM_COLOR, + ) + + for name in names: + var = m.variables[name] + for var_slice in var.iterate_slices(slice_size): + df = var_slice.to_polars() + + columns = [ + *print_variable(pl.col("labels")), + ] + + _format_and_write(df, columns, f) def integers_to_file( @@ -339,11 +413,7 @@ def integers_to_file( *print_variable(pl.col("labels")), ] - kwargs: Any = dict( - separator=" ", null_value="", quote_style="never", include_header=False - ) - formatted = df.select(pl.concat_str(columns, ignore_nulls=True)) - formatted.write_csv(f, **kwargs) + _format_and_write(df, columns, f) def sos_to_file( @@ -374,16 +444,23 @@ def sos_to_file( for name in names: var = m.variables[name] - sos_type = var.attrs["sos_type"] - sos_dim = var.attrs["sos_dim"] + sos_type = int(var.attrs[SOS_TYPE_ATTR]) # type: ignore[call-overload] + sos_dim = str(var.attrs[SOS_DIM_ATTR]) other_dims = [dim for dim in var.labels.dims if dim != sos_dim] for var_slice in var.iterate_slices(slice_size, other_dims): ds = var_slice.labels.to_dataset() - ds["sos_labels"] = ds["labels"].isel({sos_dim: 0}) + # Per-set id = max member label: unique per set (labels are globally + # unique); a fully-masked set reduces to -1 and is dropped below. + ds["sos_labels"] = ds["labels"].max(sos_dim) ds["weights"] = ds.coords[sos_dim] df = to_polars(ds) + # Drop masked members + df = df.filter((pl.col("labels") != -1) & (pl.col("sos_labels") != -1)) + if df.is_empty(): + continue + df = df.group_by("sos_labels").agg( pl.concat_str( *print_variable(pl.col("labels")), pl.lit(":"), pl.col("weights") @@ -399,11 +476,62 @@ def sos_to_file( pl.col("var_weights"), ] - kwargs: Any = dict( - separator=" ", null_value="", quote_style="never", include_header=False + _format_and_write(df, columns, f) + + +def indicator_constraints_to_file( + m: Model, + f: BufferedWriter, + explicit_coordinate_names: bool = False, +) -> None: + """ + Write indicator constraints to the s.t. section of an LP file. + + Indicator constraints appear in the Subject To section with the format: + ``ic0: x0 = 1 -> +1.0 x1 <= 5.0`` + """ + if not len(m.constraints.indicator): + return + + if not len(m.constraints.regular): + f.write(b"\n\ns.t.\n\n") + + print_variable_scalar, _ = get_printers_scalar( + m, explicit_coordinate_names=explicit_coordinate_names + ) + + for con in m.constraints.indicator.data.values(): + ic_data = con.data + labels_flat = ic_data.labels.values.flatten() + binary_var_flat = ic_data.binary_var.values.flatten() + binary_val_flat = np.broadcast_to( + ic_data.binary_val.values, labels_flat.shape + ).flatten() + coeffs_flat = ic_data.coeffs.values.reshape(len(labels_flat), -1) + vars_flat = ic_data.vars.values.reshape(len(labels_flat), -1) + sign_flat = np.broadcast_to(ic_data.sign.values, labels_flat.shape).flatten() + rhs_flat = np.broadcast_to(ic_data.rhs.values, labels_flat.shape).flatten() + + for i in range(len(labels_flat)): + if labels_flat[i] == -1: + continue + + bvar_name = print_variable_scalar(binary_var_flat[i : i + 1])[0] + valid = vars_flat[i] != -1 + var_names = print_variable_scalar(vars_flat[i][valid]) + + terms = [] + for coeff, var_name in zip(coeffs_flat[i][valid], var_names): + coeff = float(coeff) + prefix = "+" if coeff >= 0 else "" + terms.append(f"{prefix}{coeff} {var_name}") + + lhs_str = " ".join(terms) + line = ( + f"ic{labels_flat[i]}: {bvar_name} = {int(binary_val_flat[i])} -> " + f"{lhs_str} {sign_flat[i]} {float(rhs_flat[i])}\n" ) - formatted = df.select(pl.concat_str(columns, ignore_nulls=True)) - formatted.write_csv(f, **kwargs) + f.write(line.encode()) def constraints_to_file( @@ -414,7 +542,8 @@ def constraints_to_file( slice_size: int = 2_000_000, explicit_coordinate_names: bool = False, ) -> None: - if not len(m.constraints): + regular = m.constraints.regular + if not len(regular): return print_variable, print_constraint = get_printers( @@ -422,10 +551,10 @@ def constraints_to_file( ) f.write(b"\n\ns.t.\n\n") - names = m.constraints + names = list(regular) if progress: names = tqdm( - list(names), + names, desc="Writing constraints.", colour=TQDM_COLOR, ) @@ -433,65 +562,39 @@ def constraints_to_file( # to make this even faster, we can use polars expression # https://docs.pola.rs/user-guide/expressions/plugins/#output-data-types for name in names: - con = m.constraints[name] + con = regular[name] for con_slice in con.iterate_slices(slice_size): df = con_slice.to_polars() if df.height == 0: continue - # Ensure each constraint has both coefficient and RHS terms - analysis = df.group_by("labels").agg( - [ - pl.col("coeffs").is_not_null().sum().alias("coeff_rows"), - pl.col("sign").is_not_null().sum().alias("rhs_rows"), - ] - ) - - valid = analysis.filter( - (pl.col("coeff_rows") > 0) & (pl.col("rhs_rows") > 0) - ) - - if valid.height == 0: - continue - - # Keep only constraints that have both parts - df = df.join(valid.select("labels"), on="labels", how="inner") - # Sort by labels and mark first/last occurrences df = df.sort("labels").with_columns( [ - pl.when(pl.col("labels").is_first_distinct()) - .then(pl.col("labels")) - .otherwise(pl.lit(None)) - .alias("labels_first"), + pl.col("labels").is_first_distinct().alias("is_first_in_group"), (pl.col("labels") != pl.col("labels").shift(-1)) .fill_null(True) .alias("is_last_in_group"), ] ) - row_labels = print_constraint(pl.col("labels_first")) + row_labels = print_constraint(pl.col("labels")) col_labels = print_variable(pl.col("vars")) columns = [ - pl.when(pl.col("labels_first").is_not_null()).then(row_labels[0]), - pl.when(pl.col("labels_first").is_not_null()).then(row_labels[1]), - pl.when(pl.col("labels_first").is_not_null()) - .then(pl.lit(":\n")) - .alias(":"), + pl.when(pl.col("is_first_in_group")).then(row_labels[0]), + pl.when(pl.col("is_first_in_group")).then(row_labels[1]), + pl.when(pl.col("is_first_in_group")).then(pl.lit(":\n")).alias(":"), *signed_number(pl.col("coeffs")), - pl.when(pl.col("vars").is_not_null()).then(col_labels[0]), - pl.when(pl.col("vars").is_not_null()).then(col_labels[1]), + col_labels[0], + col_labels[1], + pl.when(pl.col("is_last_in_group")).then(pl.lit("\n")), pl.when(pl.col("is_last_in_group")).then(pl.col("sign")), pl.when(pl.col("is_last_in_group")).then(pl.lit(" ")), pl.when(pl.col("is_last_in_group")).then(pl.col("rhs").cast(pl.String)), ] - kwargs: Any = dict( - separator=" ", null_value="", quote_style="never", include_header=False - ) - formatted = df.select(pl.concat_str(columns, ignore_nulls=True)) - formatted.write_csv(f, **kwargs) + _format_and_write(df, columns, f) # in the future, we could use lazy dataframes when they support appending # tp existent files @@ -520,6 +623,11 @@ def to_lp_file( slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, ) + indicator_constraints_to_file( + m, + f=f, + explicit_coordinate_names=explicit_coordinate_names, + ) bounds_to_file( m, f=f, @@ -542,6 +650,13 @@ def to_lp_file( slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, ) + semi_continuous_to_file( + m, + f=f, + progress=progress, + slice_size=slice_size, + explicit_coordinate_names=explicit_coordinate_names, + ) sos_to_file( m, f=f, @@ -597,7 +712,9 @@ def to_file( # Use very fast highspy implementation # Might be replaced by custom writer, however needs C/Rust bindings for performance - h = m.to_highspy(explicit_coordinate_names=explicit_coordinate_names) + h = solvers.Highs._build_solver_model( + m, explicit_coordinate_names=explicit_coordinate_names + ) h.writeModel(str(fn)) else: raise ValueError( @@ -608,331 +725,73 @@ def to_file( def to_mosek( - m: Model, task: Any | None = None, explicit_coordinate_names: bool = False + m: Model, + task: Any | None = None, + explicit_coordinate_names: bool = False, + set_names: bool = True, ) -> Any: - """ - Export model to MOSEK. - - Export the model directly to MOSEK without writing files. - - Parameters - ---------- - m : linopy.Model - task : empty MOSEK task - - Returns - ------- - task : MOSEK Task object - """ - if m.variables.sos: - raise NotImplementedError("SOS constraints are not supported by MOSEK.") - + """Build the MOSEK task for `m`.""" import mosek - print_variable, print_constraint = get_printers_scalar( - m, explicit_coordinate_names=explicit_coordinate_names - ) - if task is None: task = mosek.Task() - - task.appendvars(m.nvars) - task.appendcons(m.ncons) - - M = m.matrices - # for j, n in enumerate(("x" + M.vlabels.astype(str).astype(object))): - # task.putvarname(j, n) - - labels = np.vectorize(print_variable)(M.vlabels).astype(object) - task.generatevarnames( - np.arange(0, len(labels)), "%0", [len(labels)], None, [0], list(labels) + return solvers.Mosek._build_solver_model( + m, + task, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, ) - ## Variables - - # MOSEK uses bound keys (free, bounded below or above, ranged and fixed) - # plus bound values (lower and upper), and it is considered an error to - # input an infinite value for a finite bound. - # bkx and bkc define the boundkeys based on upper and lower bound, and blx, - # bux, blc and buc define the finite bounds. The numerical value of a bound - # indicated to be infinite by the bound key is ignored by MOSEK. - bkx = [ - ( - ( - (mosek.boundkey.ra if lb < ub else mosek.boundkey.fx) - if ub < np.inf - else mosek.boundkey.lo - ) - if (lb > -np.inf) - else (mosek.boundkey.up if (ub < np.inf) else mosek.boundkey.fr) - ) - for (lb, ub) in zip(M.lb, M.ub) - ] - blx = [b if b > -np.inf else 0.0 for b in M.lb] - bux = [b if b < np.inf else 0.0 for b in M.ub] - task.putvarboundslice(0, m.nvars, bkx, blx, bux) - - if len(m.binaries.labels) + len(m.integers.labels) > 0: - idx = [i for (i, v) in enumerate(M.vtypes) if v in ["B", "I"]] - task.putvartypelist(idx, [mosek.variabletype.type_int] * len(idx)) - if len(m.binaries.labels) > 0: - bidx = [i for (i, v) in enumerate(M.vtypes) if v == "B"] - task.putvarboundlistconst(bidx, mosek.boundkey.ra, 0.0, 1.0) - - ## Constraints - - if len(m.constraints) > 0: - names = np.vectorize(print_constraint)(M.clabels).astype(object) - for i, n in enumerate(names): - task.putconname(i, n) - bkc = [ - ( - (mosek.boundkey.up if b < np.inf else mosek.boundkey.fr) - if s == "<" - else ( - (mosek.boundkey.lo if b > -np.inf else mosek.boundkey.up) - if s == ">" - else mosek.boundkey.fx - ) - ) - for s, b in zip(M.sense, M.b) - ] - blc = [b if b > -np.inf else 0.0 for b in M.b] - buc = [b if b < np.inf else 0.0 for b in M.b] - # blc = M.b - # buc = M.b - if M.A is not None: - A = M.A.tocsr() - task.putarowslice( - 0, m.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data - ) - task.putconboundslice(0, m.ncons, bkc, blc, buc) - - ## Objective - if M.Q is not None: - Q = (0.5 * tril(M.Q + M.Q.transpose())).tocoo() - task.putqobj(Q.row, Q.col, Q.data) - task.putclist(list(np.arange(m.nvars)), M.c) - - if m.objective.sense == "max": - task.putobjsense(mosek.objsense.maximize) - else: - task.putobjsense(mosek.objsense.minimize) - return task - def to_gurobipy( - m: Model, env: Any | None = None, explicit_coordinate_names: bool = False + m: Model, + env: Any | None = None, + explicit_coordinate_names: bool = False, + set_names: bool = True, ) -> Any: - """ - Export the model to gurobipy. - - This function does not write the model to intermediate files but directly - passes it to gurobipy. Note that for large models this is not - computationally efficient. - - Parameters - ---------- - m : linopy.Model - env : gurobipy.Env - - Returns - ------- - model : gurobipy.Model - """ - import gurobipy - - print_variable, print_constraint = get_printers_scalar( - m, explicit_coordinate_names=explicit_coordinate_names + """Build the gurobipy.Model for `m`.""" + solver = solvers.Gurobi.from_model( + m, + io_api="direct", + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, + env=env, ) + return solver.solver_model - m.constraints.sanitize_missings() - model = gurobipy.Model(env=env) - - M = m.matrices - - names = np.vectorize(print_variable)(M.vlabels).astype(object) - kwargs = {} - if len(m.binaries.labels) + len(m.integers.labels): - kwargs["vtype"] = M.vtypes - x = model.addMVar(M.vlabels.shape, M.lb, M.ub, name=list(names), **kwargs) - - if m.is_quadratic: - model.setObjective(0.5 * x.T @ M.Q @ x + M.c @ x) # type: ignore - else: - model.setObjective(M.c @ x) - - if m.objective.sense == "max": - model.ModelSense = -1 - - if len(m.constraints): - names = np.vectorize(print_constraint)(M.clabels).astype(object) - c = model.addMConstr(M.A, x, M.sense, M.b) # type: ignore - c.setAttr("ConstrName", list(names)) # type: ignore - - if m.variables.sos: - for var_name in m.variables.sos: - var = m.variables.sos[var_name] - sos_type: int = var.attrs["sos_type"] # type: ignore[assignment] - sos_dim: str = var.attrs["sos_dim"] # type: ignore[assignment] - - def add_sos(s: xr.DataArray, sos_type: int, sos_dim: str) -> None: - s = s.squeeze() - indices = s.values.flatten().tolist() - weights = s.coords[sos_dim].values.tolist() - model.addSOS(sos_type, x[indices].tolist(), weights) - - others = [dim for dim in var.labels.dims if dim != sos_dim] - if not others: - add_sos(var.labels, sos_type, sos_dim) - else: - stacked = var.labels.stack(_sos_group=others) - for _, s in stacked.groupby("_sos_group"): - add_sos(s.unstack("_sos_group"), sos_type, sos_dim) - - model.update() - return model - - -def to_highspy(m: Model, explicit_coordinate_names: bool = False) -> Highs: - """ - Export the model to highspy. - - This function does not write the model to intermediate files but directly - passes it to highspy. - - Note, this function does not track variable and constraint labels. - - Parameters - ---------- - m : linopy.Model - Returns - ------- - model : highspy.Highs - """ - if m.variables.sos: - raise NotImplementedError( - "SOS constraints are not supported by the HiGHS direct API. " - "Use io_api='lp' instead." - ) - - import highspy - - print_variable, print_constraint = get_printers_scalar( - m, explicit_coordinate_names=explicit_coordinate_names +def to_highspy( + m: Model, + explicit_coordinate_names: bool = False, + set_names: bool = True, +) -> Highs: + """Build the highspy.Highs instance for `m`.""" + solver = solvers.Highs.from_model( + m, + io_api="direct", + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, ) + return solver.solver_model - M = m.matrices - h = highspy.Highs() - h.addVars(len(M.vlabels), M.lb, M.ub) - if len(m.binaries) + len(m.integers): - vtypes = M.vtypes - labels = np.arange(len(vtypes))[(vtypes == "B") | (vtypes == "I")] - n = len(labels) - h.changeColsIntegrality(n, labels, ones_like(labels)) - if len(m.binaries): - labels = np.arange(len(vtypes))[vtypes == "B"] - n = len(labels) - h.changeColsBounds(n, labels, zeros_like(labels), ones_like(labels)) - - # linear objective - h.changeColsCost(len(M.c), np.arange(len(M.c), dtype=np.int32), M.c) - - # linear constraints - A = M.A - if A is not None: - A = A.tocsr() - num_cons = A.shape[0] - lower = np.where(M.sense != "<", M.b, -np.inf) - upper = np.where(M.sense != ">", M.b, np.inf) - h.addRows(num_cons, lower, upper, A.nnz, A.indptr, A.indices, A.data) - - lp = h.getLp() - lp.col_names_ = np.vectorize(print_variable)(M.vlabels).astype(object) - if len(M.clabels): - lp.row_names_ = np.vectorize(print_constraint)(M.clabels).astype(object) - h.passModel(lp) - - # quadrative objective - Q = M.Q - if Q is not None: - Q = triu(Q) - Q = Q.tocsr() - num_vars = Q.shape[0] - h.passHessian(num_vars, Q.nnz, 1, Q.indptr, Q.indices, Q.data) - - # change objective sense - if m.objective.sense == "max": - h.changeObjectiveSense(highspy.ObjSense.kMaximize) - - return h - - -def to_cupdlpx(m: Model, explicit_coordinate_names: bool = False) -> cupdlpxModel: - """ - Export the model to cupdlpx. - - This function does not write the model to intermediate files but directly - passes it to cupdlpx. - - cuPDLPx does not support named variables and constraints, so the - `explicit_coordinate_names` parameter is ignored. - - Parameters - ---------- - m : linopy.Model - explicit_coordinate_names : bool, optional - Ignored. cuPDLPx does not support named variables/constraints. - - Returns - ------- - model : cupdlpx.Model - """ - import cupdlpx - - if explicit_coordinate_names: - warnings.warn( - "cuPDLPx does not support named variables/constraints. " - "The explicit_coordinate_names parameter is ignored.", - UserWarning, - stacklevel=2, - ) - - # build model using canonical form matrices and vectors - # see https://github.com/MIT-Lu-Lab/cuPDLPx/tree/main/python#modeling - M = m.matrices - if M.A is None: - msg = "Model has no constraints, cannot export to cuPDLPx." - raise ValueError(msg) - A = M.A.tocsr() # cuPDLPx only supports CSR sparse matrix format - # linopy stores constraints as Ax ?= b and keeps track of inequality - # sense in M.sense. Convert to separate lower and upper bound vectors. - l = np.where( - np.logical_or(np.equal(M.sense, ">"), np.equal(M.sense, "=")), - M.b, - -np.inf, - ) - u = np.where( - np.logical_or(np.equal(M.sense, "<"), np.equal(M.sense, "=")), - M.b, - np.inf, - ) - cu_model = cupdlpx.Model( - objective_vector=M.c, - constraint_matrix=A, - constraint_lower_bound=l, - constraint_upper_bound=u, - variable_lower_bound=M.lb, - variable_upper_bound=M.ub, +def to_xpress( + m: Model, + explicit_coordinate_names: bool = False, + set_names: bool = True, +) -> Any: + """Build the xpress.problem instance for `m`.""" + return solvers.Xpress._build_solver_model( + m, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, ) - # change objective sense - if m.objective.sense == "max": - cu_model.ModelSense = cupdlpx.PDLP.MAXIMIZE - return cu_model +def to_cupdlpx(m: Model) -> cupdlpxModel: + """Build the cupdlpx.Model for `m`.""" + solver = solvers.cuPDLPx.from_model(m, io_api="direct") + return solver.solver_model def to_block_files(m: Model, fn: Path) -> None: @@ -1072,7 +931,29 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: Arguments passed to ``xarray.Dataset.to_netcdf``. **kwargs : TYPE Keyword arguments passed to ``xarray.Dataset.to_netcdf``. + + Notes + ----- + The SOS reformulation lifecycle token lives only on the in-memory + Model and is not persisted. If the model has an active SOS + reformulation at serialization time, the netcdf contains the + reformulated MILP form (aux binaries and cardinality constraints) + and a :class:`UserWarning` is emitted to flag that the deserialized + copy will not be able to undo the reformulation. + + ``Model.solve(remote=...)`` invokes ``to_netcdf`` internally on the + reformulated model and suppresses this warning. """ + if m._sos_reformulation_state is not None: + warnings.warn( + "Serializing a model with an active SOS reformulation. The " + "netcdf will contain the reformulated MILP form; the " + "reformulation lifecycle token is not persisted, so a " + "reader cannot undo it. Call `model.undo_sos_reformulation()` " + "first if you want the original SOS form on disk.", + UserWarning, + stacklevel=2, + ) def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: to_rename = set([*ds.dims, *ds.coords, *ds]) @@ -1085,7 +966,8 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: prefix_len = len(prefix) + 1 # leave original index level name names = [n[prefix_len:] for n in ds[dim].to_index().names] ds = ds.reset_index(dim) - ds.attrs[f"{dim}_multiindex"] = list(names) + # scipy netCDF3 backend cannot write unicode-array attrs. + ds.attrs[f"{dim}_multiindex"] = json.dumps(list(names)) return ds @@ -1093,7 +975,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: with_prefix(var.data, f"variables-{name}") for name, var in m.variables.items() ] cons = [ - with_prefix(con.data, f"constraints-{name}") + with_prefix(con.to_netcdf_ds(), f"constraints-{name}") for name, con in m.constraints.items() ] objective = m.objective.data @@ -1106,6 +988,21 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: scalars = {k: getattr(m, k) for k in m.scalar_attrs} ds = xr.merge(vars + cons + obj + params, combine_attrs="drop_conflicts") ds = ds.assign_attrs(scalars) + ds.attrs[NETCDF_VERSION_ATTR] = version("linopy") + if m._relaxed_registry: + ds.attrs["_relaxed_registry"] = json.dumps(m._relaxed_registry) + if m._piecewise_formulations: + ds.attrs["_piecewise_formulations"] = json.dumps( + { + name: { + "method": pwl.method, + "variable_names": pwl.variable_names, + "constraint_names": pwl.constraint_names, + "convexity": pwl.convexity, + } + for name, pwl in m._piecewise_formulations.items() + } + ) ds.attrs = non_bool_dict(ds.attrs) for k in ds: @@ -1128,15 +1025,23 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: Returns ------- m : linopy.Model + + Notes + ----- + The SOS reformulation lifecycle token is not persisted by + :func:`to_netcdf`. If the saved model was in reformulated form, + the deserialized Model is too, but + :meth:`Model.undo_sos_reformulation` is a no-op on it. """ - from linopy.model import ( + from linopy.constraints import ( Constraint, + ConstraintBase, Constraints, - LinearExpression, - Model, - Variable, - Variables, + CSRConstraint, ) + from linopy.expressions import LinearExpression + from linopy.model import Model + from linopy.variables import Variable, Variables if isinstance(path, str): path = Path(path) @@ -1150,11 +1055,20 @@ def has_prefix(k: str, prefix: str) -> bool: def remove_prefix(k: str, prefix: str) -> str: return k[len(prefix) + 1 :] + def parse_multiindex_attr(value: str | Iterable[str]) -> list[str]: + # str = JSON (new); iterable = legacy list from older linopy. + if isinstance(value, str): + return [str(n) for n in json.loads(value)] + return [str(n) for n in value] + def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ds = ds[[k for k in ds if has_prefix(str(k), prefix)]] multiindexes = [] for dim in ds.dims: - for name in ds.attrs.get(f"{dim}_multiindex", []): + attr = ds.attrs.get(f"{dim}_multiindex") + if attr is None: + continue + for name in parse_multiindex_attr(attr): multiindexes.append(prefix + "-" + name) ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes)) to_rename = set([*ds.dims, *ds.coords, *ds]) @@ -1167,8 +1081,8 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: for dim in ds.dims: if f"{dim}_multiindex" in ds.attrs: - names = ds.attrs.pop(f"{dim}_multiindex") - ds = ds.set_index({dim: names}) + names = parse_multiindex_attr(ds.attrs.pop(f"{dim}_multiindex")) + ds = ds.set_index({dim: names}) # type: ignore[dict-item] return ds @@ -1183,10 +1097,14 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: cons = [str(k) for k in ds if str(k).startswith("constraints")] con_names = list({str(k).rsplit("-", 1)[0] for k in cons}) - constraints = {} + constraints: dict[str, ConstraintBase] = {} for k in sorted(con_names): name = remove_prefix(k, "constraints") - constraints[name] = Constraint(get_prefix(ds, k), m, name) + con_ds = get_prefix(ds, k) + if con_ds.attrs.get("_linopy_format") == "csr": + constraints[name] = CSRConstraint.from_netcdf_ds(con_ds, m, name) + else: + constraints[name] = Constraint(con_ds, m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") @@ -1198,6 +1116,148 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: m.parameters = get_prefix(ds, "parameters") for k in m.scalar_attrs: - setattr(m, k, ds.attrs.get(k)) + if k in ds.attrs: + setattr(m, k, ds.attrs[k]) + + if "_relaxed_registry" in ds.attrs: + m._relaxed_registry = json.loads(ds.attrs["_relaxed_registry"]) + + if "_piecewise_formulations" in ds.attrs: + from linopy.piecewise import PiecewiseFormulation + + for name, d in json.loads(ds.attrs["_piecewise_formulations"]).items(): + m._piecewise_formulations[name] = PiecewiseFormulation( + name=name, + method=d["method"], + variable_names=d["variable_names"], + constraint_names=d["constraint_names"], + model=m, + convexity=d["convexity"], + ) return m + + +def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: + """ + Return a copy of this model. + + With ``deep=True`` (default), variables, constraints, objective, + parameters, blocks, and scalar attributes are copied to a fully + independent model. With ``deep=False``, returns a shallow copy. + + :meth:`Model.copy` defaults to deep copy for workflow safety. + In contrast, ``copy.copy(model)`` is shallow via ``__copy__``, and + ``copy.deepcopy(model)`` is deep via ``__deepcopy__``. + + Solver runtime metadata (for example, ``solver_name`` and + ``solver_model``) is intentionally not copied. Solver backend state + is recreated on ``solve()``. + + Parameters + ---------- + m : Model + The model to copy. + include_solution : bool, optional + Whether to include solution and dual values in the copy. + If False (default), solve artifacts are excluded: solution/dual data, + objective value, and solve status are reset to initialized state. + If True, these values are copied when present. For unsolved models, + this has no additional effect. + deep : bool, optional + Whether to return a deep copy (default) or shallow copy. If False, + the returned model uses independent wrapper objects that share + underlying data buffers with the source model. + + Returns + ------- + Model + A deep or shallow copy of the model. + """ + from linopy.constraints import Constraint, ConstraintBase, Constraints + from linopy.expressions import LinearExpression + from linopy.model import Model, Objective + from linopy.variables import Variable, Variables + + SOLVE_STATE_ATTRS = {"status", "termination_condition"} + + new_model = Model( + chunk=m._chunk, + force_dim_names=m._force_dim_names, + auto_mask=m._auto_mask, + freeze_constraints=m.freeze_constraints, + set_names_in_solver_io=m.set_names_in_solver_io, + solver_dir=str(m._solver_dir), + ) + + new_model._variables = Variables( + { + name: Variable( + var.data.copy(deep=deep) + if include_solution + else var.data[m.variables.dataset_attrs].copy(deep=deep), + new_model, + name, + ) + for name, var in m.variables.items() + }, + new_model, + ) + + def _copy_con_data(con: ConstraintBase) -> xr.Dataset: + d = con.mutable().data + if include_solution: + return d.copy(deep=deep) + return d[con.data_attrs].copy(deep=deep) + + new_model._constraints = Constraints( + { + name: Constraint(_copy_con_data(con), new_model, name) + for name, con in m.constraints.items() + }, + new_model, + ) + + obj_expr = LinearExpression(m.objective.expression.data.copy(deep=deep), new_model) + new_model._objective = Objective(obj_expr, new_model, m.objective.sense) + new_model._objective._value = ( + float(m.objective.value) + if (include_solution and m.objective.value is not None) + else None + ) + + new_model._parameters = m._parameters.copy(deep=deep) + new_model._blocks = m._blocks.copy(deep=deep) if m._blocks is not None else None + + for attr in m.scalar_attrs: + if include_solution or attr not in SOLVE_STATE_ATTRS: + setattr(new_model, attr, getattr(m, attr)) + + if m._sos_reformulation_state is not None: + new_model._sos_reformulation_state = _copy.deepcopy(m._sos_reformulation_state) + + return new_model + + +def shallowcopy(m: Model) -> Model: + """ + Support Python's ``copy.copy`` protocol for ``Model``. + + Returns a shallow copy with independent wrapper objects that share + underlying array buffers with ``m``. Solve artifacts are excluded, + matching :meth:`Model.copy` defaults. + """ + return copy(m, include_solution=False, deep=False) + + +def deepcopy(m: Model, memo: dict[int, Any]) -> Model: + """ + Support Python's ``copy.deepcopy`` protocol for ``Model``. + + Returns a deep, structurally independent copy and records it in ``memo`` + as required by Python's copy protocol. Solve artifacts are excluded, + matching :meth:`Model.copy` defaults. + """ + new_model = copy(m, include_solution=False, deep=True) + memo[id(m)] = new_model + return new_model diff --git a/linopy/matrices.py b/linopy/matrices.py index a55bb0bd3..0feba9c90 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -8,172 +8,200 @@ from __future__ import annotations from functools import cached_property -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np -import pandas as pd +import scipy.sparse from numpy import ndarray -from pandas.core.indexes.base import Index -from pandas.core.series import Series -from scipy.sparse._csc import csc_matrix from linopy import expressions +from linopy.constraints import CSRConstraint if TYPE_CHECKING: from linopy.model import Model -def create_vector( - indices: Series | Index, - values: Series | ndarray, - fill_value: str | float | int = np.nan, - shape: int | None = None, -) -> ndarray: - """Create a vector of a size equal to the maximum index plus one.""" - if shape is None: - max_value = indices.max() - if not isinstance(max_value, np.integer | int): - raise ValueError("Indices must be integers.") - shape = max_value + 1 - vector = np.full(shape, fill_value) - vector[indices] = values - return vector +def _stack(csrs: list) -> scipy.sparse.csr_array | None: + """Vertically stack CSR blocks, or None when there are none.""" + if not csrs: + return None + return cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr")) + + +def _concat(arrays: list, dtype: type | None = None) -> ndarray: + """Concatenate arrays, or an empty array when there are none.""" + return np.concatenate(arrays) if arrays else np.array([], dtype=dtype) + + +def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray: + """Broadcast an indicator triggering value to one entry per active row.""" + if np.ndim(binval) == 0: + return np.full(n, int(binval), dtype=np.intp) + return np.asarray(binval, dtype=np.intp).ravel() class MatrixAccessor: """ Helper class to quickly access model related vectors and matrices. + + All arrays are compact — only active (non-masked) entries are included. + Position i in variable-side arrays corresponds to vlabels[i]. + Position i in constraint-side arrays corresponds to clabels[i]. """ def __init__(self, model: Model) -> None: self._parent = model + self._build_vars() + self._build_cons() - def clean_cached_properties(self) -> None: - """Clear the cache for all cached properties of an object""" - - for cached_prop in ["flat_vars", "flat_cons", "sol", "dual"]: - # check existence of cached_prop without creating it - if cached_prop in self.__dict__: - delattr(self, cached_prop) - - @cached_property - def flat_vars(self) -> pd.DataFrame: + def _build_vars(self) -> None: m = self._parent - return m.variables.flat + label_index = m.variables.label_index + self.vlabels: ndarray = label_index.vlabels - @cached_property - def flat_cons(self) -> pd.DataFrame: - m = self._parent - return m.constraints.flat + lb_list = [] + ub_list = [] + vtypes_list = [] - @property - def vlabels(self) -> ndarray: - """Vector of labels of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.labels, -1) + for name, var in m.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 - @property - def vtypes(self) -> ndarray: - """Vector of types of all non-missing variables.""" - m = self._parent - df: pd.DataFrame = self.flat_vars - specs = [] - for name in m.variables: if name in m.binaries: - val = "B" + vtype = "B" elif name in m.integers: - val = "I" + vtype = "I" + elif name in m.semi_continuous: + vtype = "S" else: - val = "C" - specs.append(pd.Series(val, index=m.variables[name].flat.labels)) + vtype = "C" + + lb_list.append(var.lower.values.ravel()[mask]) + ub_list.append(var.upper.values.ravel()[mask]) + vtypes_list.append(np.full(mask.sum(), vtype)) + + if lb_list: + self.lb: ndarray = np.concatenate(lb_list) + self.ub: ndarray = np.concatenate(ub_list) + self.vtypes: ndarray = np.concatenate(vtypes_list) + else: + self.lb = np.array([]) + self.ub = np.array([]) + self.vtypes = np.array([], dtype=object) + + def _build_cons(self) -> None: + m = self._parent + label_index = m.variables.label_index + label_to_pos = label_index.label_to_pos + + reg_csrs, reg_b, reg_sense = [], [], [] + ind_csrs, ind_b, ind_sense, ind_binvar, ind_binval = [], [], [], [], [] + for c in m.constraints.data.values(): + if c.is_indicator: + cc = c if isinstance(c, CSRConstraint) else c.freeze() + csr, _, b, sense = cc.to_matrix_with_rhs(label_index) + ind_csrs.append(csr) + ind_b.append(b) + ind_sense.append(sense) + ind_binvar.append(label_to_pos[cc._binvar_labels]) + binval = cast("int | np.ndarray", cc._binval) + ind_binval.append(_binval_per_row(binval, len(b))) + else: + csr, _, b, sense = c.to_matrix_with_rhs(label_index) + reg_csrs.append(csr) + reg_b.append(b) + reg_sense.append(sense) + + self.clabels: ndarray = m.constraints.label_index.clabels + self.A: scipy.sparse.csr_array | None = _stack(reg_csrs) + self.b: ndarray = _concat(reg_b) + self.sense: ndarray = _concat(reg_sense, dtype=object) + self.indicator_A: scipy.sparse.csr_array | None = _stack(ind_csrs) + self.indicator_b: ndarray = _concat(ind_b) + self.indicator_sense: ndarray = _concat(ind_sense, dtype=object) + self.indicator_binvar: ndarray = _concat(ind_binvar, dtype=np.intp) + self.indicator_binval: ndarray = _concat(ind_binval, dtype=np.intp) - ds = pd.concat(specs) - ds = df.set_index("key").labels.map(ds) - return create_vector(ds.index, ds.to_numpy(), fill_value="") + @cached_property + def c(self) -> ndarray: + """Objective coefficients aligned with vlabels.""" + m = self._parent + result = np.zeros(len(self.vlabels)) + + label_index = m.variables.label_index + label_to_pos = label_index.label_to_pos + expr = m.objective.expression + if isinstance(expr, expressions.QuadraticExpression): + # vars has shape (_factor=2, _term); linear terms have one factor == -1 + vars_2d = expr.data.vars.values # shape (2, n_term) + coeffs_all = expr.data.coeffs.values.ravel() + vars1, vars2 = vars_2d[0], vars_2d[1] + linear = (vars1 == -1) | (vars2 == -1) + var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) + coeffs = coeffs_all[linear] + else: + var_labels = expr.data.vars.values.ravel() + coeffs = expr.data.coeffs.values.ravel() + + mask = var_labels != -1 + np.add.at(result, label_to_pos[var_labels[mask]], coeffs[mask]) + return result - @property - def lb(self) -> ndarray: - """Vector of lower bounds of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.lower) + @cached_property + def Q(self) -> scipy.sparse.csc_matrix | None: + """Quadratic objective matrix, shape (n_active_vars, n_active_vars).""" + m = self._parent + expr = m.objective.expression + if not isinstance(expr, expressions.QuadraticExpression): + return None + return expr.to_matrix()[self.vlabels][:, self.vlabels] @cached_property def sol(self) -> ndarray: - """Vector of solution values of all non-missing variables.""" + """Solution values aligned with vlabels.""" if not self._parent.status == "ok": raise ValueError("Model is not optimized.") - if "solution" not in self.flat_vars: - del self.flat_vars # clear cache - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.solution, fill_value=np.nan) + m = self._parent + result = np.full(len(self.vlabels), np.nan) + label_index = m.variables.label_index + label_to_pos = label_index.label_to_pos + for _, var in m.variables.items(): + labels = var.labels.values.ravel() + mask = labels != -1 + positions = label_to_pos[labels[mask]] + result[positions] = var.solution.values.ravel()[mask] + return result @cached_property def dual(self) -> ndarray: - """Vector of dual values of all non-missing constraints.""" + """Dual values aligned with clabels.""" if not self._parent.status == "ok": raise ValueError("Model is not optimized.") - if "dual" not in self.flat_cons: - del self.flat_cons # clear cache - df: pd.DataFrame = self.flat_cons - if "dual" not in df: + m = self._parent + label_index = m.variables.label_index + dual_list = [] + has_dual = False + for c in m.constraints.data.values(): + if c.is_indicator: + continue + if isinstance(c, CSRConstraint): + # _dual is active-only + if c._dual is not None: + dual_list.append(c._dual) + has_dual = True + else: + dual_list.append(np.full(len(c._con_labels), np.nan)) + else: + csr, _ = c.to_matrix(label_index) + nonempty = np.diff(csr.indptr).astype(bool) + active_rows = np.flatnonzero(nonempty) + if "dual" in c.data: + dual_list.append(c.dual.values.ravel()[active_rows]) + has_dual = True + else: + dual_list.append(np.full(len(active_rows), np.nan)) + if not has_dual: raise AttributeError( "Underlying is optimized but does not have dual values stored." ) - return create_vector(df.key, df.dual, fill_value=np.nan) - - @property - def ub(self) -> ndarray: - """Vector of upper bounds of all non-missing variables.""" - df: pd.DataFrame = self.flat_vars - return create_vector(df.key, df.upper) - - @property - def clabels(self) -> ndarray: - """Vector of labels of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - if df.empty: - return np.array([], dtype=int) - return create_vector(df.key, df.labels, fill_value=-1) - - @property - def A(self) -> csc_matrix | None: - """Constraint matrix of all non-missing constraints and variables.""" - m = self._parent - if not len(m.constraints): - return None - A: csc_matrix = m.constraints.to_matrix(filter_missings=False) - return A[self.clabels][:, self.vlabels] - - @property - def sense(self) -> ndarray: - """Vector of senses of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - return create_vector(df.key, df.sign.astype(np.dtype(" ndarray: - """Vector of right-hand-sides of all non-missing constraints.""" - df: pd.DataFrame = self.flat_cons - return create_vector(df.key, df.rhs) - - @property - def c(self) -> ndarray: - """Vector of objective coefficients of all non-missing variables.""" - m = self._parent - ds = m.objective.flat - if isinstance(m.objective.expression, expressions.QuadraticExpression): - ds = ds[(ds.vars1 == -1) | (ds.vars2 == -1)] - ds["vars"] = ds.vars1.where(ds.vars1 != -1, ds.vars2) - - vars: pd.Series = ds.vars.map(self.flat_vars.set_index("labels").key) - shape: int = self.flat_vars.key.max() + 1 - return create_vector(vars, ds.coeffs, fill_value=0.0, shape=shape) - - @property - def Q(self) -> csc_matrix | None: - """Matrix objective coefficients of quadratic terms of all non-missing variables.""" - m = self._parent - expr = m.objective.expression - if not isinstance(expr, expressions.QuadraticExpression): - return None - return expr.to_matrix()[self.vlabels][:, self.vlabels] + return np.concatenate(dual_list) if dual_list else np.array([]) diff --git a/linopy/model.py b/linopy/model.py index 657b2d45a..b1477b276 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -9,46 +9,62 @@ import logging import os import re +import warnings from collections.abc import Callable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir -from typing import Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload +from warnings import warn import numpy as np import pandas as pd import xarray as xr from deprecation import deprecated -from numpy import inf, nan, ndarray +from numpy import inf from pandas.core.frame import DataFrame from pandas.core.series import Series from xarray import DataArray, Dataset from xarray.core.types import T_Chunks from linopy import solvers +from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( - as_dataarray, assign_multiindex_safe, best_int, maybe_replace_signs, replace_by_map, - set_int_index, to_path, ) +from linopy.config import options from linopy.constants import ( GREATER_EQUAL, HELPER_DIMS, LESS_EQUAL, + SOS_BIG_M_ATTR, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, TERM_DIM, ModelStatus, + Result, TerminationCondition, ) -from linopy.constraints import AnonymousScalarConstraint, Constraint, Constraints +from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + ConstraintBase, + Constraints, + CSRConstraint, +) +from linopy.dualization import dualize from linopy.expressions import ( LinearExpression, QuadraticExpression, ScalarLinearExpression, ) from linopy.io import ( + copy, + deepcopy, + shallowcopy, to_block_files, to_cupdlpx, to_file, @@ -56,15 +72,31 @@ to_highspy, to_mosek, to_netcdf, + to_xpress, ) from linopy.matrices import MatrixAccessor from linopy.objective import Objective -from linopy.remote import OetcHandler, RemoteHandler -from linopy.solver_capabilities import SolverFeature, solver_supports +from linopy.piecewise import ( + add_piecewise_formulation, +) +from linopy.remote import RemoteHandler + +try: + from linopy.remote import OetcHandler +except ImportError: + OetcHandler = None # type: ignore +from linopy.solver_capabilities import solver_supports from linopy.solvers import ( IO_APIS, + SolverFeature, available_solvers, ) +from linopy.sos_reformulation import ( + SOSReformulationResult, + reformulate_sos_constraints, + sos_reformulation_context, + undo_sos_reformulation, +) from linopy.types import ( ConstantLike, ConstraintLike, @@ -75,6 +107,9 @@ ) from linopy.variables import ScalarVariable, Variable, Variables +if TYPE_CHECKING: + from linopy.piecewise import PiecewiseFormulation + logger = logging.getLogger(__name__) @@ -94,8 +129,7 @@ class Model: the optimization process. """ - solver_model: Any - solver_name: str + _solver: solvers.Solver | None _variables: Variables _constraints: Constraints _objective: Objective @@ -108,12 +142,13 @@ class Model: _cCounter: int _varnameCounter: int _connameCounter: int + _pwlCounter: int _blocks: DataArray | None _chunk: T_Chunks _force_dim_names: bool + _freeze_constraints: bool + _set_names_in_solver_io: bool _solver_dir: Path - matrices: MatrixAccessor - __slots__ = ( # containers "_variables", @@ -130,14 +165,20 @@ class Model: "_cCounter", "_varnameCounter", "_connameCounter", + "_pwlCounter", "_blocks", # TODO: check if these should not be mutable "_chunk", "_force_dim_names", + "_auto_mask", + "_freeze_constraints", + "_set_names_in_solver_io", "_solver_dir", - "solver_model", - "solver_name", - "matrices", + "_relaxed_registry", + "_piecewise_formulations", + "_solver", + "_sos_reformulation_state", + "__weakref__", ) def __init__( @@ -145,6 +186,9 @@ def __init__( solver_dir: str | None = None, chunk: T_Chunks = None, force_dim_names: bool = False, + auto_mask: bool = False, + freeze_constraints: bool = False, + set_names_in_solver_io: bool = True, ) -> None: """ Initialize the linopy model. @@ -164,6 +208,16 @@ def __init__( "dim_1" and so on. These helps to avoid unintended broadcasting over dimension. Especially the use of pandas DataFrames and Series may become safer. + auto_mask : bool + Whether to automatically mask variables and constraints where + bounds, coefficients, or RHS values contain NaN. The default is + False. + freeze_constraints : bool + Whether constraints added to the model should be frozen to the + CSR-backed representation by default. The default is False. + set_names_in_solver_io : bool + Whether direct solver exports should include variable and + constraint names by default. The default is True. Returns ------- @@ -180,15 +234,56 @@ def __init__( self._cCounter: int = 0 self._varnameCounter: int = 0 self._connameCounter: int = 0 + self._pwlCounter: int = 0 self._blocks: DataArray | None = None self._chunk: T_Chunks = chunk self._force_dim_names: bool = bool(force_dim_names) + self._auto_mask: bool = bool(auto_mask) + self._freeze_constraints: bool = bool(freeze_constraints) + self._set_names_in_solver_io: bool = bool(set_names_in_solver_io) + self._piecewise_formulations: dict[str, PiecewiseFormulation] = {} + self._relaxed_registry: dict[str, str] = {} self._solver_dir: Path = Path( gettempdir() if solver_dir is None else solver_dir ) + self._solver: solvers.Solver | None = None + self._sos_reformulation_state: SOSReformulationResult | None = None + + @property + def solver(self) -> solvers.Solver | None: + return self._solver + + @solver.setter + def solver(self, value: solvers.Solver | None) -> None: + if self._solver is not None and self._solver is not value: + self._solver.close() + self._solver = value + + @property + def solver_model(self) -> Any: + return self.solver.solver_model if self.solver is not None else None + + @solver_model.setter + def solver_model(self, value: Any) -> None: + if value is not None: + raise AttributeError("solver state is managed via model.solver") + self.solver = None + + @property + def solver_name(self) -> str | None: + return self.solver.solver_name.value if self.solver is not None else None + + @solver_name.setter + def solver_name(self, value: str | None) -> None: + if value is not None: + raise AttributeError("solver state is managed via model.solver") + self.solver = None - self.matrices: MatrixAccessor = MatrixAccessor(self) + @property + def matrices(self) -> MatrixAccessor: + """Matrix representation of the model, computed fresh on each access.""" + return MatrixAccessor(self) @property def variables(self) -> Variables: @@ -204,6 +299,16 @@ def constraints(self) -> Constraints: """ return self._constraints + @property + def indicator_constraints(self) -> Constraints: + """ + Indicator constraints assigned to the model. + + Returns the subset of ``model.constraints`` for which + ``is_indicator`` is True. + """ + return self.constraints.indicator + @property def objective(self) -> Objective: """ @@ -214,12 +319,20 @@ def objective(self) -> Objective: @objective.setter def objective( self, obj: Objective | LinearExpression | QuadraticExpression - ) -> Objective: + ) -> None: + """ + Set the objective function. + + Parameters + ---------- + obj : Objective, LinearExpression, or QuadraticExpression + The objective to assign to the model. If not an Objective instance, + it will be wrapped in an Objective. + """ if not isinstance(obj, Objective): obj = Objective(obj, self) self._objective = obj - return self._objective @property def sense(self) -> str: @@ -230,6 +343,9 @@ def sense(self) -> str: @sense.setter def sense(self, value: str) -> None: + """ + Set the sense of the objective function. + """ self.objective.sense = value @property @@ -244,7 +360,12 @@ def parameters(self) -> Dataset: @parameters.setter def parameters(self, value: Dataset | Mapping) -> None: - self._parameters = Dataset(value) + """ + Set the parameters of the model. + """ + self._parameters = ( + value.copy() if isinstance(value, Dataset) else Dataset(value) + ) @property def solution(self) -> Dataset: @@ -269,6 +390,9 @@ def status(self) -> str: @status.setter def status(self, value: str) -> None: + """ + Set the status of the model. + """ self._status = ModelStatus[value].value @property @@ -280,11 +404,13 @@ def termination_condition(self) -> str: @termination_condition.setter def termination_condition(self, value: str) -> None: - # TODO: remove if-clause, only kept for backward compatibility - if value: - self._termination_condition = TerminationCondition[value].value - else: + """ + Set the termination condition of the model. + """ + if value == "": self._termination_condition = value + else: + self._termination_condition = TerminationCondition[value].value @property def chunk(self) -> T_Chunks: @@ -295,6 +421,9 @@ def chunk(self) -> T_Chunks: @chunk.setter def chunk(self, value: T_Chunks) -> None: + """ + Set the chunk sizes of the model. + """ self._chunk = value @property @@ -312,8 +441,44 @@ def force_dim_names(self) -> bool: @force_dim_names.setter def force_dim_names(self, value: bool) -> None: + """ + Set whether to force custom dimension names for variables and constraints. + """ self._force_dim_names = bool(value) + @property + def auto_mask(self) -> bool: + """ + If True, automatically mask variables and constraints where bounds, + coefficients, or RHS values contain NaN. + """ + return self._auto_mask + + @auto_mask.setter + def auto_mask(self, value: bool) -> None: + """ + Set whether to automatically mask variables and constraints with NaN values. + """ + self._auto_mask = bool(value) + + @property + def freeze_constraints(self) -> bool: + """Whether constraints are frozen to CSR by default when added.""" + return self._freeze_constraints + + @freeze_constraints.setter + def freeze_constraints(self, value: bool) -> None: + self._freeze_constraints = bool(value) + + @property + def set_names_in_solver_io(self) -> bool: + """Whether direct solver exports include names by default.""" + return self._set_names_in_solver_io + + @set_names_in_solver_io.setter + def set_names_in_solver_io(self, value: bool) -> None: + self._set_names_in_solver_io = bool(value) + @property def solver_dir(self) -> Path: """ @@ -323,6 +488,9 @@ def solver_dir(self) -> Path: @solver_dir.setter def solver_dir(self, value: str | Path) -> None: + """ + Set the solver directory of the model. + """ if not isinstance(value, str | Path): raise TypeError("'solver_dir' must path-like.") self._solver_dir = Path(value) @@ -340,22 +508,31 @@ def scalar_attrs(self) -> list[str]: "_cCounter", "_varnameCounter", "_connameCounter", + "_pwlCounter", "force_dim_names", + "auto_mask", + "freeze_constraints", + "set_names_in_solver_io", ] def __repr__(self) -> str: """ Return a string representation of the linopy model. """ - var_string = self.variables.__repr__().split("\n", 2)[2] - con_string = self.constraints.__repr__().split("\n", 2)[2] + from linopy.piecewise import _get_piecewise_groups + from linopy.piecewise import _repr_summary as pwl_repr_summary + + var_names, con_names = _get_piecewise_groups(self) + var_string = self.variables._format_items(exclude=var_names) + con_string = self.constraints._format_items(exclude=con_names) model_string = f"Linopy {self.type} model" return ( f"{model_string}\n{'=' * len(model_string)}\n\n" f"Variables:\n----------\n{var_string}\n" - f"Constraints:\n------------\n{con_string}\n" - f"Status:\n-------\n{self.status}" + f"Constraints:\n------------\n{con_string}" + f"{pwl_repr_summary(self)}" + f"\nStatus:\n-------\n{self.status}" ) def __getitem__(self, key: str) -> Variable: @@ -424,11 +601,12 @@ def add_variables( self, lower: Any = -inf, upper: Any = inf, - coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, + coords: Sequence[Sequence | pd.Index] | Mapping | None = None, name: str | None = None, - mask: DataArray | ndarray | Series | None = None, + mask: MaskLike | None = None, binary: bool = False, integer: bool = False, + semi_continuous: bool = False, **kwargs: Any, ) -> Variable: """ @@ -443,17 +621,32 @@ def add_variables( Parameters ---------- lower : float/array_like, optional - Lower bound of the variable(s). Ignored if `binary` is True. - The default is -inf. + Lower bound of the variable(s). For binary variables it + defaults to 0 and, if given, must be 0 or 1. The default is -inf. upper : TYPE, optional - Upper bound of the variable(s). Ignored if `binary` is True. - The default is inf. - coords : list/xarray.Coordinates, optional - The coords of the variable array. - These are directly passed to the DataArray creation of - `lower` and `upper`. For every single combination of - coordinates a optimization variable is added to the model. - The default is None. + Upper bound of the variable(s). For binary variables it + defaults to 1 and, if given, must be 0 or 1. The default is inf. + coords : list/dict/xarray.Coordinates, optional + The coords of the variable array. When provided with **named + dimensions** (a ``Mapping``, ``xarray.Coordinates``, a + sequence of named ``pd.Index`` objects, or an unnamed + sequence paired with ``dims=`` in ``**kwargs``), ``coords`` + is the source of truth for the variable's dimensions, + order, and values. ``lower``, ``upper`` and ``mask`` are + aligned to this contract: + + - dims of every bound must be a subset of ``coords.dims``; + extra dims raise ``ValueError``; + - dim order in the variable always follows ``coords``; + - shared-dim coordinate values must equal ``coords``; same + values in a different order are auto-reindexed, different + value sets raise ``ValueError``; + - dims listed in ``coords`` but missing from a bound are + broadcast to ``coords`` shape. + + One optimization variable is added per combination of + coordinates. The default is ``None``, in which case the + shape is inferred from the bounds. name : str, optional Reference name of the added variables. The default None results in a name like "var1", "var2" etc. @@ -467,6 +660,11 @@ def add_variables( integer : bool Whether the new variable is a integer variable which are used for Mixed-Integer problems. + semi_continuous : bool + Whether the new variable is a semi-continuous variable. A + semi-continuous variable can take the value 0 or any value + between its lower and upper bounds. Requires a positive lower + bound. **kwargs : Additional keyword arguments are passed to the DataArray creation. @@ -501,6 +699,67 @@ def add_variables( [7]: x[7] ∈ [0, inf] [8]: x[8] ∈ [0, inf] [9]: x[9] ∈ [0, inf] + + Strict coords-as-truth: a bound with an extra dim raises. + + >>> import xarray as xr + >>> m = Model() + >>> bad = xr.DataArray( + ... [[1.0, 2.0, 3.0]] * 2, + ... dims=["extra", "x"], + ... coords={"x": [0, 1, 2]}, + ... ) + >>> m.add_variables(lower=bad, coords=[pd.Index([0, 1, 2], name="x")], name="v") + Traceback (most recent call last): + ... + ValueError: lower bound has dimension(s) ['extra'] not declared in coords ... + + Strict coords-as-truth: a bound whose shared-dim values don't + match raises. + + >>> m = Model() + >>> wrong = xr.DataArray( + ... [1.0, 2.0, 3.0], dims=["x"], coords={"x": [10, 20, 30]} + ... ) + >>> m.add_variables( + ... lower=wrong, coords=[pd.Index([0, 1, 2], name="x")], name="v" + ... ) + Traceback (most recent call last): + ... + ValueError: lower bound: coordinate values for dimension 'x' do not match coords ... + + Strict coords-as-truth, helpful side: a bound whose coord values + match ``coords`` only in a different order is auto-reindexed. + + >>> m = Model() + >>> reordered = xr.DataArray( + ... [3.0, 1.0, 2.0], dims=["x"], coords={"x": ["c", "a", "b"]} + ... ) + >>> v = m.add_variables( + ... lower=reordered, + ... coords=[pd.Index(["a", "b", "c"], name="x")], + ... name="r", + ... ) + >>> v.lower.values.tolist() + [1.0, 2.0, 3.0] + + Unnamed-coords sequence + ``dims=`` opts into the same strict + enforcement as a named index — extra dims still raise. + + >>> m = Model() + >>> m.add_variables(lower=bad, coords=[[0, 1, 2]], dims=["x"], name="w") + Traceback (most recent call last): + ... + ValueError: lower bound has dimension(s) ['extra'] not declared in coords ... + + The same strict contract applies to ``mask`` (including with + ``coords=[[...]], dims=[...]``). + + >>> m = Model() + >>> m.add_variables(mask=bad, coords=[[0, 1, 2]], dims=["x"], name="wm") + Traceback (most recent call last): + ... + ValueError: mask has dimension(s) ['extra'] not declared in coords ... """ if name is None: name = f"var{self._varnameCounter}" @@ -509,19 +768,33 @@ def add_variables( if name in self.variables: raise ValueError(f"Variable '{name}' already assigned to model") - if binary and integer: - raise ValueError("Variable cannot be both binary and integer.") + if sum([binary, integer, semi_continuous]) > 1: + raise ValueError( + "Variable can only be one of binary, integer, or semi-continuous." + ) if binary: - if (lower != -inf) or (upper != inf): - raise ValueError("Binary variables cannot have lower or upper bounds.") - else: - lower, upper = 0, 1 + if np.isscalar(lower) and lower == -inf: + lower = 0 + elif not (np.isin(lower, (0, 1)) | pd.isna(lower)).all(): + raise ValueError("Binary variable lower bounds must be 0 or 1.") + if np.isscalar(upper) and upper == inf: + upper = 1 + elif not (np.isin(upper, (0, 1)) | pd.isna(upper)).all(): + raise ValueError("Binary variable upper bounds must be 0 or 1.") + + if semi_continuous: + if not np.isscalar(lower) or float(lower) <= 0: # type: ignore[arg-type] + raise ValueError( + "Semi-continuous variables require a positive scalar lower bound." + ) + lower_da = broadcast_to_coords(lower, coords, label="lower bound", **kwargs) + upper_da = broadcast_to_coords(upper, coords, label="upper bound", **kwargs) data = Dataset( { - "lower": as_dataarray(lower, coords, **kwargs), - "upper": as_dataarray(upper, coords, **kwargs), + "lower": lower_da, + "upper": upper_da, "labels": -1, } ) @@ -530,18 +803,48 @@ def add_variables( self._check_valid_dim_names(data) if mask is not None: - mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) + mask = broadcast_to_coords( + mask, + coords if coords is not None else data.coords, + label="mask", + **kwargs, + ).astype(bool) + + # Auto-mask based on NaN in bounds (use numpy for speed) + if self.auto_mask: + auto_mask_values = ~np.isnan(data.lower.values) & ~np.isnan( + data.upper.values + ) + auto_mask_arr = DataArray( + auto_mask_values, coords=data.coords, dims=data.dims + ) + if mask is not None: + mask = mask & auto_mask_arr + else: + mask = auto_mask_arr start = self._xCounter end = start + data.labels.size - data.labels.values = np.arange(start, end).reshape(data.labels.shape) + label_dtype = options["label_dtype"] + if end > np.iinfo(label_dtype).max: + raise ValueError( + f"Number of labels ({end}) exceeds the maximum value for " + f"{label_dtype.__name__} ({np.iinfo(label_dtype).max})." + ) + data.labels.values = np.arange( + start, end, dtype=options["label_dtype"] + ).reshape(data.labels.shape) self._xCounter += data.labels.size if mask is not None: - data.labels.values = data.labels.where(mask, -1).values + data.labels.values = np.where(mask.values, data.labels.values, -1) data = data.assign_attrs( - label_range=(start, end), name=name, binary=binary, integer=integer + label_range=(start, end), + name=name, + binary=binary, + integer=integer, + semi_continuous=semi_continuous, ) if self.chunk: @@ -556,6 +859,7 @@ def add_sos_constraints( variable: Variable, sos_type: Literal[1, 2], sos_dim: str, + big_m: float | None = None, ) -> None: """ Add an sos1 or sos2 constraint for one dimension of a variable @@ -569,15 +873,26 @@ def add_sos_constraints( Type of SOS sos_dim : str Which dimension of variable to add SOS constraint to + big_m : float | None, optional + Big-M value for SOS reformulation. Only used when reformulating + SOS constraints for solvers that don't support them natively. + + - None (default): Use variable upper bounds as Big-M + - float: Custom Big-M value + + The reformulation uses the tighter of big_m and variable upper bound: + M = min(big_m, var.upper). + + Tighter Big-M values improve LP relaxation quality and solve time. """ if sos_type not in (1, 2): raise ValueError(f"sos_type must be 1 or 2, got {sos_type}") if sos_dim not in variable.dims: raise ValueError(f"sos_dim must name a variable dimension, got {sos_dim}") - if "sos_type" in variable.attrs or "sos_dim" in variable.attrs: - existing_sos_type = variable.attrs.get("sos_type") - existing_sos_dim = variable.attrs.get("sos_dim") + if SOS_TYPE_ATTR in variable.attrs or SOS_DIM_ATTR in variable.attrs: + existing_sos_type = variable.attrs.get(SOS_TYPE_ATTR) + existing_sos_dim = variable.attrs.get(SOS_DIM_ATTR) raise ValueError( f"variable already has an sos{existing_sos_type} constraint on {existing_sos_dim}" ) @@ -589,7 +904,125 @@ def add_sos_constraints( f"but got {variable.coords[sos_dim].dtype}" ) - variable.attrs.update(sos_type=sos_type, sos_dim=sos_dim) + attrs_update: dict[str, Any] = {SOS_TYPE_ATTR: sos_type, SOS_DIM_ATTR: sos_dim} + if big_m is not None: + if big_m <= 0: + raise ValueError(f"big_m must be positive, got {big_m}") + attrs_update[SOS_BIG_M_ATTR] = float(big_m) + + variable.attrs.update(attrs_update) + + add_piecewise_formulation = add_piecewise_formulation + + def _resolve_constraint_name(self, name: str | None, prefix: str = "con") -> str: + """Validate a constraint name or generate one from ``prefix``.""" + if name in list(self.constraints): + raise ValueError(f"Constraint '{name}' already assigned to model") + if name is None: + name = f"{prefix}{self._connameCounter}" + self._connameCounter += 1 + return name + + def _constraint_data_from_lhs( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None, + rhs: ConstantLike | VariableLike | ExpressionLike | None, + coords: Sequence[Sequence | pd.Index] | Mapping | None = None, + ) -> Dataset: + """Build the constraint Dataset from an ``lhs`` and optional ``sign``/``rhs``.""" + msg_required = ( + f"`sign` and `rhs` are required when `lhs` is a {type(lhs).__name__}." + ) + msg_must_be_none = ( + f"`sign` and `rhs` must be None when `lhs` is a {type(lhs).__name__}." + ) + if isinstance(lhs, LinearExpression): + if sign is None or rhs is None: + raise ValueError(msg_required) + return lhs.to_constraint(sign, rhs).data + elif isinstance(lhs, list | tuple): + if sign is None or rhs is None: + raise ValueError(msg_required) + return self.linexpr(*lhs).to_constraint(sign, rhs).data + elif callable(lhs): + assert coords is not None, "`coords` must be given when lhs is a function" + if sign is not None or rhs is not None: + raise ValueError(msg_must_be_none) + return Constraint.from_rule(self, lhs, coords).data + elif isinstance(lhs, AnonymousScalarConstraint): + if sign is not None or rhs is not None: + raise ValueError(msg_must_be_none) + return lhs.to_constraint().data + elif isinstance(lhs, ConstraintBase): + if sign is not None or rhs is not None: + raise ValueError(msg_must_be_none) + return lhs.data + elif isinstance(lhs, Variable | ScalarVariable | ScalarLinearExpression): + if sign is None or rhs is None: + raise ValueError(msg_required) + return lhs.to_linexpr().to_constraint(sign, rhs).data + else: + raise TypeError( + f"`lhs` must be a LinearExpression, Variable, Constraint, tuple, or " + f"callable, got {type(lhs).__name__}." + ) + + def _allocate_constraint_labels( + self, data: Dataset, name: str, mask: DataArray | None = None + ) -> Dataset: + """Assign label ranges from the constraint counter and apply an optional mask.""" + start = self._cCounter + end = start + data.labels.size + label_dtype = options["label_dtype"] + if end > np.iinfo(label_dtype).max: + raise ValueError( + f"Number of labels ({end}) exceeds the maximum value for " + f"{label_dtype.__name__} ({np.iinfo(label_dtype).max})." + ) + data.labels.values = np.arange(start, end, dtype=label_dtype).reshape( + data.labels.shape + ) + self._cCounter += data.labels.size + if mask is not None: + data.labels.values = np.where(mask.values, data.labels.values, -1) + return data.assign_attrs(label_range=(start, end), name=name) + + @overload + def add_constraints( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None = ..., + rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., + name: str | None = ..., + coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., + mask: MaskLike | None = ..., + freeze: Literal[False] = ..., + ) -> Constraint: ... + + @overload + def add_constraints( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None = ..., + rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., + name: str | None = ..., + coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., + mask: MaskLike | None = ..., + freeze: Literal[True] = ..., + ) -> CSRConstraint: ... def add_constraints( self, @@ -601,9 +1034,10 @@ def add_constraints( sign: SignLike | None = None, rhs: ConstantLike | VariableLike | ExpressionLike | None = None, name: str | None = None, - coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, + coords: Sequence[Sequence | pd.Index] | Mapping | None = None, mask: MaskLike | None = None, - ) -> Constraint: + freeze: bool | None = None, + ) -> ConstraintBase: """ Assign a new, possibly multi-dimensional array of constraints to the model. @@ -615,7 +1049,7 @@ def add_constraints( Parameters ---------- - lhs : linopy.LinearExpression/linopy.Constraint/callable + lhs : linopy.LinearExpression/linopy.ConstraintBase/callable Left hand side of the constraint(s) or optionally full constraint. In case a linear expression is passed, `sign` and `rhs` must not be None. @@ -637,56 +1071,30 @@ def add_constraints( Boolean mask with False values for constraints which are skipped. The shape of the mask has to match the shape the added constraints. Default is None. - + freeze : bool, optional + If True, convert the constraint to an immutable CSR-backed CSRConstraint + for better memory efficiency. If None, uses the model default + ``Model.freeze_constraints`` setting (default False). Returns ------- - labels : linopy.model.Constraint - Array containing the labels of the added constraints. + constraint : linopy.ConstraintBase + The added constraint (Constraint by default, or CSRConstraint if freeze=True). """ - msg_sign_rhs_none = f"Arguments `sign` and `rhs` cannot be None when passing along with a {type(lhs)}." - msg_sign_rhs_not_none = f"Arguments `sign` and `rhs` cannot be None when passing along with a {type(lhs)}." - - if name in list(self.constraints): - raise ValueError(f"Constraint '{name}' already assigned to model") - elif name is None: - name = f"con{self._connameCounter}" - self._connameCounter += 1 + name = self._resolve_constraint_name(name) if sign is not None: sign = maybe_replace_signs(as_dataarray(sign)) - if isinstance(lhs, LinearExpression): - if sign is None or rhs is None: - raise ValueError(msg_sign_rhs_not_none) - data = lhs.to_constraint(sign, rhs).data - elif isinstance(lhs, list | tuple): - if sign is None or rhs is None: - raise ValueError(msg_sign_rhs_none) - data = self.linexpr(*lhs).to_constraint(sign, rhs).data - # directly convert first argument to a constraint - elif callable(lhs): - assert coords is not None, "`coords` must be given when lhs is a function" - rule = lhs - if sign is not None or rhs is not None: - raise ValueError(msg_sign_rhs_none) - data = Constraint.from_rule(self, rule, coords).data - elif isinstance(lhs, AnonymousScalarConstraint): - if sign is not None or rhs is not None: - raise ValueError(msg_sign_rhs_none) - data = lhs.to_constraint().data - elif isinstance(lhs, Constraint): - if sign is not None or rhs is not None: - raise ValueError(msg_sign_rhs_none) - data = lhs.data - elif isinstance(lhs, Variable | ScalarVariable | ScalarLinearExpression): - if sign is None or rhs is None: - raise ValueError(msg_sign_rhs_not_none) - data = lhs.to_linexpr().to_constraint(sign, rhs).data - else: - raise ValueError( - f"Invalid type of `lhs` ({type(lhs)}) or invalid combination of `lhs`, `sign` and `rhs`." - ) + # Capture original RHS for auto-masking before constraint creation + # (NaN values in RHS are lost during constraint creation) + # Use numpy for speed instead of xarray's notnull() + original_rhs_mask = None + if self.auto_mask and rhs is not None: + rhs_da = as_dataarray(rhs) + original_rhs_mask = (rhs_da.coords, rhs_da.dims, ~np.isnan(rhs_da.values)) + + data = self._constraint_data_from_lhs(lhs, sign, rhs, coords) invalid_infinity_values = ( (data.sign == LESS_EQUAL) & (data.rhs == -np.inf) @@ -699,34 +1107,128 @@ 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 = assign_multiindex_safe(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]) if mask is not None: - mask = as_dataarray(mask).astype(bool) - # TODO: simplify - assert set(mask.dims).issubset(data.dims), ( - "Dimensions of mask not a subset of resulting labels dimensions." + mask = broadcast_to_coords(mask, data.coords, label="mask").astype(bool) + + # Auto-mask based on null expressions or NaN RHS (use numpy for speed) + if self.auto_mask: + # Check if expression is null: all vars == -1 + # Use max() instead of all() - if max == -1, all are -1 (since valid vars >= 0) + # This is ~30% faster for large term dimensions + vars_all_invalid = data.vars.values.max(axis=-1) == -1 + auto_mask_values = ~vars_all_invalid + if original_rhs_mask is not None: + coords, dims, rhs_notnull = original_rhs_mask + if rhs_notnull.shape != auto_mask_values.shape: + rhs_da = DataArray(rhs_notnull, coords=coords, dims=dims) + rhs_notnull = rhs_da.broadcast_like(data.labels).values + auto_mask_values = auto_mask_values & rhs_notnull + auto_mask_arr = DataArray( + auto_mask_values, coords=data.labels.coords, dims=data.labels.dims ) + if mask is not None: + mask = mask & auto_mask_arr + else: + mask = auto_mask_arr self.check_force_dim_names(data) - start = self._cCounter - end = start + data.labels.size - data.labels.values = np.arange(start, end).reshape(data.labels.shape) - self._cCounter += data.labels.size - - if mask is not None: - data.labels.values = data.labels.where(mask, -1).values - - data = data.assign_attrs(label_range=(start, end), name=name) + data = self._allocate_constraint_labels(data, name, mask) if self.chunk: data = data.chunk(self.chunk) constraint = Constraint(data, name=name, model=self, skip_broadcast=True) - self.constraints.add(constraint) - return constraint + if freeze is None: + freeze = self.freeze_constraints + return self.constraints.add(constraint, freeze=freeze and not self.chunk) + + def add_indicator_constraints( + self, + binary_var: Variable, + binary_val: int, + lhs: ConstraintLike | ExpressionLike | VariableLike, + sign: SignLike | None = None, + rhs: ConstantLike | None = None, + name: str | None = None, + ) -> ConstraintBase: + """ + Add indicator constraints to the model. + + An indicator constraint has the form: + (binary_var == binary_val) => (linear_constraint) + + The linear constraint is only enforced when binary_var equals + binary_val. These constraints are handled natively by solvers + like Gurobi and CPLEX via general constraints. + + Parameters + ---------- + binary_var : linopy.Variable + Binary variable serving as the indicator. Must have binary=True. + binary_val : int + Triggering value, must be 0 or 1. + lhs : linopy.Constraint, linopy.LinearExpression, or linopy.Variable + The conditionally enforced constraint. If a LinearExpression or + Variable is passed, ``sign`` and ``rhs`` must also be provided. + sign : str, optional + Constraint sign ('<=', '>=', '='). Required when ``lhs`` is an + expression. + rhs : numeric, optional + Right-hand side. Required when ``lhs`` is an expression. + name : str, optional + Name for the indicator constraint group. + + Returns + ------- + linopy.constraints.ConstraintBase + The added indicator constraint. + """ + if not binary_var.attrs.get("binary", False): + raise ValueError( + "Indicator variable must be binary. " + f"Variable '{binary_var.name}' is not binary." + ) + + if binary_val not in (0, 1): + raise ValueError(f"binary_val must be 0 or 1, got {binary_val}.") + + name = self._resolve_constraint_name(name, prefix="indcon") + if sign is not None: + sign = maybe_replace_signs(as_dataarray(sign)) + + data = self._constraint_data_from_lhs(lhs, sign, rhs) + + data["binary_var"] = binary_var.labels + data["binary_val"] = binary_val + + data["labels"] = -1 + (data,) = xr.broadcast(data, exclude=[TERM_DIM]) + + data = self._allocate_constraint_labels(data, name) + + con = Constraint(data, name=name, model=self, skip_broadcast=True) + freeze = self.freeze_constraints + return self.constraints.add(con, freeze=freeze and not self.chunk) + + def remove_indicator_constraints(self, name: str) -> None: + """ + Remove indicator constraint by name. + """ + self.constraints.remove(name) def add_objective( self, @@ -778,18 +1280,28 @@ def remove_variables(self, name: str) -> None: ------- None. """ - labels = self.variables[name].labels - self.variables.remove(name) + variable = self.variables[name] + + self._relaxed_registry.pop(name, None) + + to_remove = [ + k for k, con in self.constraints.items() if con.has_variable(variable) + ] - for k in list(self.constraints): - vars = self.constraints[k].data["vars"] - vars = vars.where(~vars.isin(labels), -1) - self.constraints[k]._data = assign_multiindex_safe( - self.constraints[k].data, vars=vars + if to_remove: + warnings.warn( + f"Removing variable '{name}' also removes constraints {to_remove} " + "because they reference this variable.", + UserWarning, + stacklevel=2, ) + for k in to_remove: + self.constraints.remove(k) + + self.variables.remove(name) self.objective = self.objective.sel( - {TERM_DIM: ~self.objective.vars.isin(labels)} + {TERM_DIM: ~self.objective.vars.isin(variable.labels)} ) def remove_constraints(self, name: str | list[str]) -> None: @@ -830,18 +1342,96 @@ def remove_sos_constraints(self, variable: Variable) -> None: ------- None. """ - if "sos_type" not in variable.attrs or "sos_dim" not in variable.attrs: + if SOS_TYPE_ATTR not in variable.attrs or SOS_DIM_ATTR not in variable.attrs: raise ValueError(f"Variable '{variable.name}' has no SOS constraints") - sos_type = variable.attrs["sos_type"] - sos_dim = variable.attrs["sos_dim"] + sos_type = variable.attrs[SOS_TYPE_ATTR] + sos_dim = variable.attrs[SOS_DIM_ATTR] + + del variable.attrs[SOS_TYPE_ATTR], variable.attrs[SOS_DIM_ATTR] - del variable.attrs["sos_type"], variable.attrs["sos_dim"] + variable.attrs.pop(SOS_BIG_M_ATTR, None) logger.debug( f"Removed sos{sos_type} constraint on {sos_dim} from {variable.name}" ) + reformulate_sos_constraints = reformulate_sos_constraints + + def apply_sos_reformulation(self) -> None: + """ + Reformulate SOS constraints into binary + linear form, in place. + + The reformulation token is stored on the model so it can be reverted + with :meth:`undo_sos_reformulation`. This is the stateful counterpart + to :func:`linopy.sos_reformulation.reformulate_sos_constraints`, where + the caller owns the token. + + Raises + ------ + RuntimeError + If a reformulation has already been applied and not undone. + """ + if self._sos_reformulation_state is not None: + raise RuntimeError( + "SOS reformulation has already been applied to this model. " + "Call `undo_sos_reformulation()` before applying again." + ) + self._sos_reformulation_state = reformulate_sos_constraints(self) + + def undo_sos_reformulation(self) -> None: + """ + Revert a previously applied SOS reformulation. + + Raises + ------ + RuntimeError + If no reformulation is currently applied. + """ + if self._sos_reformulation_state is None: + raise RuntimeError( + "No SOS reformulation is currently applied to this model." + ) + state = self._sos_reformulation_state + self._sos_reformulation_state = None + undo_sos_reformulation(self, state) + + def _resolve_sos_reformulation( + self, + solver_name: str | None, + reformulate_sos: bool | Literal["auto"], + ) -> bool: + """ + Decide whether ``apply_sos_reformulation`` should run. + + Validates ``reformulate_sos`` and returns ``True`` iff the SOS + constraints on this model should be reformulated for the chosen + solver. ``solver_name`` is only consulted when + ``reformulate_sos == "auto"`` (to look up SOS support); for + ``True`` / ``False`` the decision is independent of the solver. + """ + if reformulate_sos not in (True, False, "auto"): + raise ValueError( + f"Invalid value for reformulate_sos: {reformulate_sos!r}. " + "Must be True, False, or 'auto'." + ) + if not self.variables.sos: + return False + + if reformulate_sos is False: + return False + elif reformulate_sos is True: + return True + elif solver_name is None: + raise ValueError( + "`reformulate_sos='auto'` on a model with SOS constraints " + "requires an explicit `solver_name` so we can check " + "whether the chosen solver supports SOS. Pass " + "`solver_name=...` or use `reformulate_sos=True`/`False` " + "to skip the lookup." + ) + return not solver_supports(solver_name, SolverFeature.SOS_CONSTRAINTS) + def remove_objective(self) -> None: """ Remove the objective's linear expression from the model. @@ -873,19 +1463,31 @@ def integers(self) -> Variables: """ return self.variables.integers + @property + def semi_continuous(self) -> Variables: + """ + Get all semi-continuous variables. + """ + return self.variables.semi_continuous + @property def is_linear(self) -> bool: + """Whether the objective is linear.""" return self.objective.is_linear @property def is_quadratic(self) -> bool: + """Whether the objective is quadratic.""" return self.objective.is_quadratic @property def type(self) -> str: - if (len(self.binaries) or len(self.integers)) and len(self.continuous): + """Short string identifying the problem type.""" + if ( + len(self.binaries) or len(self.integers) or len(self.semi_continuous) + ) and len(self.continuous): variable_type = "MI" - elif len(self.binaries) or len(self.integers): + elif len(self.binaries) or len(self.integers) or len(self.semi_continuous): variable_type = "I" else: variable_type = "" @@ -958,7 +1560,7 @@ def calculate_block_maps(self) -> None: @overload def linexpr( - self, *args: Sequence[Sequence | pd.Index | DataArray] | Mapping + self, *args: Sequence[Sequence | pd.Index] | Mapping ) -> LinearExpression: ... @overload @@ -971,7 +1573,7 @@ def linexpr( *args: tuple[ConstantLike, str | Variable | ScalarVariable] | ConstantLike | Callable - | Sequence[Sequence | pd.Index | DataArray] + | Sequence[Sequence | pd.Index] | Mapping, ) -> LinearExpression: """ @@ -1113,6 +1715,7 @@ def solve( solver_name: str | None = None, io_api: str | None = None, explicit_coordinate_names: bool = False, + set_names: bool | None = None, problem_fn: str | Path | None = None, solution_fn: str | Path | None = None, log_fn: str | Path | None = None, @@ -1123,9 +1726,10 @@ def solve( sanitize_zeros: bool = True, sanitize_infinities: bool = True, slice_size: int = 2_000_000, - remote: RemoteHandler | OetcHandler = None, # type: ignore + remote: RemoteHandler | OetcHandler | None = None, progress: bool | None = None, mock_solve: bool = False, + reformulate_sos: bool | Literal["auto"] = False, **solver_options: Any, ) -> tuple[str, str]: """ @@ -1151,6 +1755,11 @@ def solve( this option allows to keep the variable and constraint names in the lp file. This may lead to slower run times. The default is set to False. + set_names : bool, optional + Whether to set variable and constraint names when using the direct + solver API (io_api='direct'). Setting to False can significantly + speed up model export. If None, uses the model default + ``Model.set_names_in_solver_io`` setting (default True). problem_fn : path_like, optional Path of the lp file or output file/directory which is written out during the process. The default None results in a temporary file. @@ -1195,6 +1804,13 @@ def solve( than 10000 variables and constraints. mock_solve : bool, optional Whether to run a mock solve. This will skip the actual solving. Variables will be set to have dummy values + reformulate_sos : bool | Literal["auto"], optional + Whether to reformulate SOS constraints as binary + linear constraints. + If True, always reformulates, even when the solver supports SOS natively. + If "auto", reformulates only when the solver lacks SOS support. + If False, raises if the solver doesn't support SOS. + Reformulation uses the Big-M method and requires all SOS variables + to have finite bounds. Default is False. **solver_options : kwargs Options passed to the solver. @@ -1209,9 +1825,6 @@ def solve( sanitize_zeros=sanitize_zeros, sanitize_infinities=sanitize_infinities ) - # clear cached matrix properties potentially present from previous solve commands - self.matrices.clean_cached_properties() - # check io_api if io_api is not None and io_api not in IO_APIS: raise ValueError( @@ -1219,8 +1832,23 @@ def solve( ) if remote is not None: + # The remote branch short-circuits before reaching Solver.solve(), + # which is where the empty-objective check normally fires. Replicate + # it here. This duplication becomes obsolete once OETC is folded + # into the Solver pipeline (see PyPSA/linopy#683). + if self.objective.expression.empty: + raise ValueError( + "No objective has been set on the model. Use " + "`m.add_objective(...)` first (e.g. `m.add_objective(0 * x)` " + "for a pure feasibility problem)." + ) if isinstance(remote, OetcHandler): - solved = remote.solve_on_oetc(self) + solved = remote.solve_on_oetc( + self, + solver_name=solver_name, + reformulate_sos=reformulate_sos, + **solver_options, + ) else: solved = remote.solve_on_remote( self, @@ -1233,10 +1861,12 @@ def solve( warmstart_fn=warmstart_fn, keep_files=keep_files, sanitize_zeros=sanitize_zeros, + reformulate_sos=reformulate_sos, **solver_options, ) - self.objective.set_value(solved.objective.value) + if solved.objective.value is not None: + self.objective.set_value(float(solved.objective.value)) self.status = solved.status self.termination_condition = solved.termination_condition for k, v in self.variables.items(): @@ -1267,11 +1897,13 @@ def solve( ) logger.info(f"Solver options:\n{options_string}") + solver_class = getattr(solvers, solvers.SolverName(solver_name).name) + if problem_fn is None: problem_fn = self.get_problem_file(io_api=io_api) if solution_fn is None: if ( - solver_supports(solver_name, SolverFeature.SOLUTION_FILE_NOT_NEEDED) + solver_class.supports(SolverFeature.SOLUTION_FILE_NOT_NEEDED) and not keep_files ): # these (solver, keep_files=False) combos do not need a solution file @@ -1279,110 +1911,117 @@ def solve( else: solution_fn = self.get_solution_file() - if sanitize_zeros: - self.constraints.sanitize_zeros() - - if sanitize_infinities: - self.constraints.sanitize_infinities() - - if self.is_quadratic and not solver_supports( - solver_name, SolverFeature.QUADRATIC_OBJECTIVE - ): - raise ValueError( - f"Solver {solver_name} does not support quadratic problems." - ) + with sos_reformulation_context(self, solver_name, reformulate_sos): + if sanitize_zeros: + self.constraints.sanitize_zeros() + if sanitize_infinities: + self.constraints.sanitize_infinities() - # SOS constraints are not supported by all solvers - if self.variables.sos and not solver_supports( - solver_name, SolverFeature.SOS_CONSTRAINTS - ): - raise ValueError(f"Solver {solver_name} does not support SOS constraints.") - - try: - solver_class = getattr(solvers, f"{solvers.SolverName(solver_name).name}") - # initialize the solver as object of solver subclass - solver = solver_class( - **solver_options, - ) - if io_api == "direct": - # no problem file written and direct model is set for solver - result = solver.solve_problem_from_model( + try: + self.solver = None # closes any previous solver + if io_api == "direct": + if set_names is None: + set_names = self.set_names_in_solver_io + build_kwargs: dict[str, Any] = { + "explicit_coordinate_names": explicit_coordinate_names, + "set_names": set_names, + "log_fn": to_path(log_fn), + } + if env is not None: + build_kwargs["env"] = env + else: + build_kwargs = { + "explicit_coordinate_names": explicit_coordinate_names, + "slice_size": slice_size, + "progress": progress, + "problem_fn": to_path(problem_fn), + } + self.solver = solver = solvers.Solver.from_name( + solver_name, model=self, - solution_fn=to_path(solution_fn), - log_fn=to_path(log_fn), - warmstart_fn=to_path(warmstart_fn), - basis_fn=to_path(basis_fn), - env=env, - explicit_coordinate_names=explicit_coordinate_names, - ) - else: - if ( - not solver_supports(solver_name, SolverFeature.LP_FILE_NAMES) - and explicit_coordinate_names - ): - logger.warning( - f"{solver_name} does not support writing names to lp files, disabling it." - ) - explicit_coordinate_names = False - problem_fn = self.to_file( - to_path(problem_fn), io_api=io_api, - explicit_coordinate_names=explicit_coordinate_names, - slice_size=slice_size, - progress=progress, + options=solver_options, + **build_kwargs, ) - result = solver.solve_problem_from_file( - problem_fn=to_path(problem_fn), + if io_api != "direct": + problem_fn = solver._problem_fn + result = solver.solve( solution_fn=to_path(solution_fn), log_fn=to_path(log_fn), warmstart_fn=to_path(warmstart_fn), basis_fn=to_path(basis_fn), env=env, ) + finally: + for fn in (problem_fn, solution_fn): + if fn is not None and (os.path.exists(fn) and not keep_files): + os.remove(fn) - finally: - for fn in (problem_fn, solution_fn): - if fn is not None and (os.path.exists(fn) and not keep_files): - os.remove(fn) + return self.assign_result(result) + + def assign_result( + self, + result: Result, + solver: solvers.Solver | None = None, + ) -> tuple[str, str]: + """ + Write a solver Result back onto the model. + + Copies primal / dual values onto variables / constraints, sets + :attr:`status`, :attr:`termination_condition`, and + :attr:`objective.value`. When ``solver`` is provided, also stores it on + ``self.solver`` so post-solve introspection (``model.solver_model``, + ``compute_infeasibilities()``) works. + + Parameters + ---------- + result : Result + The :class:`linopy.constants.Result` returned by + :meth:`linopy.solvers.Solver.solve`. + solver : Solver, optional + The solver instance that produced the result. Pass it on the + low-level ``Solver.from_name(...).solve()`` path to attach it as + ``self.solver`` for post-solve introspection. ``Model.solve()`` + attaches the solver itself and does not pass this argument. + """ + if solver is not None: + self.solver = solver result.info() - self.objective._value = result.solution.objective - self.status = result.status.status.value - self.termination_condition = result.status.termination_condition.value - self.solver_model = result.solver_model - self.solver_name = solver_name + if result.solution is not None: + self.objective._value = result.solution.objective + + status_value = result.status.status.value + termination_condition = result.status.termination_condition.value + self.status = status_value + self.termination_condition = termination_condition if not result.status.is_ok: - return result.status.status.value, result.status.termination_condition.value + return status_value, termination_condition - # map solution and dual to original shape which includes missing values - sol = result.solution.primal.copy() - sol = set_int_index(sol) - sol.loc[-1] = nan + if result.solution is None or len(result.solution.primal) == 0: + return status_value, termination_condition - for name, var in self.variables.items(): - idx = np.ravel(var.labels) - try: - vals = sol[idx].values.reshape(var.labels.shape) - except KeyError: - vals = sol.reindex(idx).values.reshape(var.labels.shape) - var.solution = xr.DataArray(vals, var.coords) - - if not result.solution.dual.empty: - dual = result.solution.dual.copy() - dual = set_int_index(dual) - dual.loc[-1] = nan - - for name, con in self.constraints.items(): - idx = np.ravel(con.labels) - try: - vals = dual[idx].values.reshape(con.labels.shape) - except KeyError: - vals = dual.reindex(idx).values.reshape(con.labels.shape) - con.dual = xr.DataArray(vals, con.labels.coords) - - return result.status.status.value, result.status.termination_condition.value + primal = result.solution.primal + for _, var in self.variables.items(): + start, end = var.range + var.solution = xr.DataArray( + primal[start:end].reshape(var.shape), var.coords, dims=var.dims + ) + + if len(result.solution.dual): + dual = result.solution.dual + for _, con in self.constraints.items(): + if con.is_indicator: + continue + start, end = con.range + coords = {dim: con.coords[dim] for dim in con.coord_dims} + con.dual = xr.DataArray( + dual[start:end].reshape(con.shape), coords, dims=con.coord_dims + ) + + return status_value, termination_condition def _mock_solve( self, @@ -1391,10 +2030,8 @@ def _mock_solve( ) -> tuple[str, str]: solver_name = "mock" - # clear cached matrix properties potentially present from previous solve commands - self.matrices.clean_cached_properties() - logger.info(f" Solve problem using {solver_name.title()} solver") + self.solver = None # reset result self.reset_solution() @@ -1407,8 +2044,6 @@ def _mock_solve( self.objective._value = 0.0 self.status = "ok" self.termination_condition = TerminationCondition.optimal.value - self.solver_model = None - self.solver_name = solver_name for name, var in self.variables.items(): var.solution = xr.DataArray(0.0, var.coords) @@ -1431,7 +2066,7 @@ def compute_infeasibilities(self) -> list[int]: labels : list[int] Labels of the infeasible constraints. """ - solver_model = getattr(self, "solver_model", None) + solver_model = self.solver_model # Check for Gurobi if "gurobi" in available_solvers: @@ -1460,8 +2095,10 @@ def compute_infeasibilities(self) -> list[int]: # If we get here, either the solver doesn't support IIS or no solver model is available if solver_model is None: # Check if this is a supported solver without a stored model - solver_name = getattr(self, "solver_name", "unknown") - if solver_supports(solver_name, SolverFeature.IIS_COMPUTATION): + solver_name = self.solver_name or "unknown" + if self.solver is not None and self.solver.supports( + SolverFeature.IIS_COMPUTATION + ): raise ValueError( "No solver model available. The model must be solved first with " "a solver that supports IIS computation and the result must be infeasible." @@ -1499,38 +2136,39 @@ def _compute_infeasibilities_gurobi(self, solver_model: Any) -> list[int]: return labels def _compute_infeasibilities_xpress(self, solver_model: Any) -> list[int]: - """Compute infeasibilities for Xpress solver.""" - # Compute all IIS + """ + Compute infeasibilities for Xpress solver. + + This function correctly maps solver constraint positions to linopy + constraint labels, handling masked constraints where some labels may + be skipped (e.g., labels [0, 2, 4] with gaps instead of sequential + [0, 1, 2]). + """ + # Compute a single IIS (matches Gurobi behavior; multiple IIS would + # otherwise get flattened into an ambiguous union). Mode 2 prioritises + # a fast IIS search over minimality. try: # Try new API first - solver_model.IISAll() + solver_model.firstIIS(2) except AttributeError: # Fallback to old API - solver_model.iisall() + solver_model.iisfirst(2) - # Get the number of IIS found - num_iis = solver_model.attributes.numiis - if num_iis == 0: + if solver_model.attributes.numiis == 0: return [] - labels = set() - - # Create constraint mapping for efficient lookups - constraint_to_index = { - constraint: idx - for idx, constraint in enumerate(solver_model.getConstraint()) - } - - # Retrieve each IIS - for iis_num in range(1, num_iis + 1): - iis_constraints = self._extract_iis_constraints(solver_model, iis_num) + clabels = self.constraints.label_index.clabels + constraint_position_map = {} + for position, constraint_obj in enumerate(solver_model.getConstraint()): + if 0 <= position < len(clabels): + constraint_label = clabels[position] + if constraint_label >= 0: + constraint_position_map[constraint_obj] = constraint_label - # Convert constraint objects to indices - for constraint_obj in iis_constraints: - if constraint_obj in constraint_to_index: - labels.add(constraint_to_index[constraint_obj]) - # Note: Silently skip constraints not found in mapping - # This can happen if the model structure changed after solving + labels = set() + for constraint_obj in self._extract_iis_constraints(solver_model, 1): + if constraint_obj in constraint_position_map: + labels.add(constraint_position_map[constraint_obj]) - return sorted(list(labels)) + return sorted(labels) def _extract_iis_constraints(self, solver_model: Any, iis_num: int) -> list[Any]: """ @@ -1600,9 +2238,9 @@ def _extract_iis_constraints(self, solver_model: Any, iis_num: int) -> list[Any] return miisrow - def print_infeasibilities(self, display_max_terms: int | None = None) -> None: + def format_infeasibilities(self, display_max_terms: int | None = None) -> str: """ - Print a list of infeasible constraints. + Return a string representation of infeasible constraints. This function requires that the model was solved using `gurobi` or `xpress` and the termination condition was infeasible. @@ -1610,20 +2248,35 @@ def print_infeasibilities(self, display_max_terms: int | None = None) -> None: Parameters ---------- display_max_terms : int, optional - The maximum number of infeasible terms to display. If `None`, - all infeasible terms will be displayed. + The maximum number of infeasible terms to display. If ``None``, + uses the global ``linopy.options.display_max_terms`` setting. Returns ------- - None - This function does not return anything. It simply prints the - infeasible constraints. + str + String representation of the infeasible constraints. """ labels = self.compute_infeasibilities() - self.constraints.print_labels(labels, display_max_terms=display_max_terms) + return self.constraints.format_labels( + labels, display_max_terms=display_max_terms + ) + + def print_infeasibilities(self, display_max_terms: int | None = None) -> None: + """ + Print a list of infeasible constraints. + + .. deprecated:: + Use :meth:`format_infeasibilities` instead. + """ + warn( + "`Model.print_infeasibilities` is deprecated. Use `Model.format_infeasibilities` instead.", + DeprecationWarning, + stacklevel=2, + ) + print(self.format_infeasibilities(display_max_terms=display_max_terms)) @deprecated( - details="Use `compute_infeasibilities`/`print_infeasibilities` instead." + details="Use `compute_infeasibilities`/`format_infeasibilities` instead." ) def compute_set_of_infeasible_constraints(self) -> Dataset: """ @@ -1652,6 +2305,12 @@ def reset_solution(self) -> None: self.variables.reset_solution() self.constraints.reset_dual() + copy = copy + + __copy__ = shallowcopy + + __deepcopy__ = deepcopy + to_netcdf = to_netcdf to_file = to_file @@ -1664,4 +2323,8 @@ def reset_solution(self) -> None: to_cupdlpx = to_cupdlpx + to_xpress = to_xpress + to_block_files = to_block_files + + dualize = dualize diff --git a/linopy/monkey_patch_xarray.py b/linopy/monkey_patch_xarray.py index dc60608c6..1e526c927 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/objective.py b/linopy/objective.py index b14492707..a51b22076 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -232,10 +232,12 @@ def set_value(self, value: float) -> None: @property def is_linear(self) -> bool: + """Whether the objective expression is linear.""" return type(self.expression) is expressions.LinearExpression @property def is_quadratic(self) -> bool: + """Whether the objective expression is quadratic.""" return type(self.expression) is expressions.QuadraticExpression def to_matrix(self, *args: Any, **kwargs: Any) -> csc_matrix: diff --git a/linopy/persistent/__init__.py b/linopy/persistent/__init__.py new file mode 100644 index 000000000..1058fce49 --- /dev/null +++ b/linopy/persistent/__init__.py @@ -0,0 +1,39 @@ +"""Persistent-solver snapshot and diff primitives.""" + +from __future__ import annotations + +from linopy.persistent.diff import ( + ConSlice, + ModelDiff, + RebuildReason, + VarSlice, +) +from linopy.persistent.errors import ( + RebuildRequiredError, + UnsupportedUpdate, + UpdatesDisabledError, +) +from linopy.persistent.snapshot import ( + ContainerConBuffers, + ContainerVarBuffers, + ModelSnapshot, + StructuralKey, + VarKind, + clear_coef_dirty, +) + +__all__ = [ + "ConSlice", + "ContainerConBuffers", + "ContainerVarBuffers", + "ModelDiff", + "ModelSnapshot", + "RebuildReason", + "RebuildRequiredError", + "StructuralKey", + "UnsupportedUpdate", + "UpdatesDisabledError", + "VarKind", + "VarSlice", + "clear_coef_dirty", +] diff --git a/linopy/persistent/diff.py b/linopy/persistent/diff.py new file mode 100644 index 000000000..38173d12a --- /dev/null +++ b/linopy/persistent/diff.py @@ -0,0 +1,570 @@ +from __future__ import annotations + +import enum +from collections.abc import Iterable +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING + +import numpy as np + +from linopy.constants import short_GREATER_EQUAL, short_LESS_EQUAL +from linopy.constraints import Constraint +from linopy.persistent.snapshot import ( + ContainerConBuffers, + ContainerVarBuffers, + ModelSnapshot, + StructuralKey, + _coord_snapshot, + _extract_con_buffers, + _extract_var_buffers, + _objective_linear_vector, +) + +if TYPE_CHECKING: + from numpy.typing import DTypeLike + + from linopy.common import ConstraintLabelIndex, VariableLabelIndex + from linopy.constraints import ConstraintBase + from linopy.model import Model + from linopy.variables import Variable + + +class RebuildReason(enum.Enum): + STRUCTURAL_LABELS = "vlabels/clabels mismatch" + STRUCTURAL_CONTAINERS = "container set changed" + COORD_REINDEX = "coordinates changed" + SPARSITY = "coefficient sparsity changed" + QUAD_OBJ = "quadratic objective changed" + BACKEND_REJECTED = "backend raised UnsupportedUpdate" + + +@dataclass(frozen=True) +class VarSlice: + bounds: slice + type: slice + + +@dataclass(frozen=True) +class ConSlice: + coef: slice + rhs: slice + sign: slice + + +def _cat(parts: list[np.ndarray], dtype: DTypeLike) -> np.ndarray: + if not parts: + return np.empty(0, dtype=dtype) + return np.concatenate(parts).astype(dtype, copy=False) + + +def _same(a: np.ndarray, b: np.ndarray) -> bool: + return a is b or np.array_equal(a, b) + + +def _coords_equal( + a: dict[str, np.ndarray], b: dict[str, np.ndarray], ignored: frozenset[str] +) -> bool: + keys = a.keys() - ignored + if keys != b.keys() - ignored: + return False + return all(np.array_equal(a[k], b[k]) for k in keys) + + +def _structural_reason(base: StructuralKey, model: Model) -> RebuildReason | None: + if base.var_container_names != tuple( + model.variables + ) or base.con_container_names != tuple(model.constraints): + return RebuildReason.STRUCTURAL_CONTAINERS + if not np.array_equal(base.vlabels, model.variables.label_index.vlabels): + return RebuildReason.STRUCTURAL_LABELS + if not np.array_equal(base.clabels, model.constraints.label_index.clabels): + return RebuildReason.STRUCTURAL_LABELS + return None + + +@dataclass(frozen=True) +class _CoefDelta: + """Coefficient changes of one container, expanded to COO lazily.""" + + buf: ContainerConBuffers + changed_rows: np.ndarray + row_positions: np.ndarray + nnz: int + + +@dataclass +class ModelDiff: + """ + Flat-native delta between two structurally identical model states. + + Instances are produced by :meth:`from_snapshot` / :meth:`from_models`; + any condition that cannot be expressed as an in-place delta is returned + as a :class:`RebuildReason` instead of a diff. + + Coefficient changes are stored per container as ``coef_deltas`` + (changed rows referencing the container's CSR buffers) and expanded to + COO triplets — ``con_coef_rows`` / ``con_coef_cols`` / ``con_coef_vals`` + — on first access. + """ + + var_bounds_indices: np.ndarray + var_bounds_lower: np.ndarray + var_bounds_upper: np.ndarray + var_type_positions: np.ndarray + var_type_kinds: np.ndarray + + coef_deltas: list[_CoefDelta] + n_coef_updates: int + + con_rhs_indices: np.ndarray + con_rhs_values: np.ndarray + con_rhs_signs: np.ndarray + + con_sign_indices: np.ndarray + con_sign_values: np.ndarray + + obj_c_indices: np.ndarray | None + obj_c_values: np.ndarray | None + obj_sense: str | None + + var_slices: dict[str, VarSlice] + con_slices: dict[str, ConSlice] + + #: Snapshot of the diffed (target) model state, assembled from the + #: buffers the diff walk already extracted — adopting it after a + #: successful apply replaces a full re-capture. Note: holding a diff + #: therefore pins all container buffers for its lifetime. + snapshot: ModelSnapshot + + @property + def is_empty(self) -> bool: + return ( + self.var_bounds_indices.size == 0 + and self.var_type_positions.size == 0 + and self.n_coef_updates == 0 + and self.con_rhs_indices.size == 0 + and self.con_sign_indices.size == 0 + and self.obj_c_indices is None + and self.obj_sense is None + ) + + @property + def changed_variables(self) -> set[str]: + return set(self.var_slices) + + @property + def changed_constraints(self) -> set[str]: + return set(self.con_slices) + + @cached_property + def _coef_coo(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + rows = np.empty(self.n_coef_updates, dtype=np.int32) + cols = np.empty(self.n_coef_updates, dtype=np.int32) + vals = np.empty(self.n_coef_updates, dtype=np.float64) + cursor = 0 + for delta in self.coef_deltas: + indptr = delta.buf.indptr + starts = indptr[delta.changed_rows] + counts = indptr[delta.changed_rows + 1] - starts + run_offsets = np.repeat(np.cumsum(counts) - counts, counts) + flat = np.repeat(starts, counts) + np.arange(delta.nnz) - run_offsets + sl = slice(cursor, cursor + delta.nnz) + rows[sl] = np.repeat(delta.row_positions, counts) + cols[sl] = delta.buf.indices[flat] + vals[sl] = delta.buf.data[flat] + cursor += delta.nnz + return rows, cols, vals + + @property + def con_coef_rows(self) -> np.ndarray: + return self._coef_coo[0] + + @property + def con_coef_cols(self) -> np.ndarray: + return self._coef_coo[1] + + @property + def con_coef_vals(self) -> np.ndarray: + return self._coef_coo[2] + + def con_rhs_as_bounds(self) -> tuple[np.ndarray, np.ndarray]: + """Return (lower, upper) row-bounds form of the RHS updates.""" + vals = self.con_rhs_values + signs = self.con_rhs_signs + lower = np.where(signs == short_LESS_EQUAL, -np.inf, vals) + upper = np.where(signs == short_GREATER_EQUAL, np.inf, vals) + return lower, upper + + def summary(self) -> dict[str, int | bool | str | None]: + return { + "var_bounds": int(self.var_bounds_indices.size), + "var_type": int(self.var_type_positions.size), + "con_rhs": int(self.con_rhs_indices.size), + "con_sign": int(self.con_sign_indices.size), + "con_coef_updates": self.n_coef_updates, + "obj_linear_changed": self.obj_c_indices is not None, + "obj_sense_changed_to": self.obj_sense, + } + + def inspect_variable(self, name: str) -> dict[str, object]: + sl = self.var_slices.get(name) + if sl is None: + return {} + entry: dict[str, object] = {} + if sl.bounds.stop > sl.bounds.start: + entry["bounds_indices"] = self.var_bounds_indices[sl.bounds] + entry["lower"] = self.var_bounds_lower[sl.bounds] + entry["upper"] = self.var_bounds_upper[sl.bounds] + if sl.type.stop > sl.type.start: + entry["type_positions"] = self.var_type_positions[sl.type] + entry["type_kinds"] = self.var_type_kinds[sl.type] + return entry + + def inspect_constraint(self, name: str) -> dict[str, object]: + sl = self.con_slices.get(name) + if sl is None: + return {} + entry: dict[str, object] = {} + if sl.coef.stop > sl.coef.start: + entry["coef_rows"] = self.con_coef_rows[sl.coef] + entry["coef_cols"] = self.con_coef_cols[sl.coef] + entry["coef_vals"] = self.con_coef_vals[sl.coef] + if sl.rhs.stop > sl.rhs.start: + entry["rhs_indices"] = self.con_rhs_indices[sl.rhs] + entry["rhs_values"] = self.con_rhs_values[sl.rhs] + entry["rhs_signs"] = self.con_rhs_signs[sl.rhs] + if sl.sign.stop > sl.sign.start: + entry["sign_indices"] = self.con_sign_indices[sl.sign] + entry["sign_values"] = self.con_sign_values[sl.sign] + return entry + + def __repr__(self) -> str: + if self.is_empty: + return "ModelDiff(empty)" + parts = [ + f"{k}={v}" for k, v in self.summary().items() if v not in (0, False, None) + ] + return "ModelDiff(" + ", ".join(parts) + ")" + + @classmethod + def from_snapshot( + cls, + snapshot: ModelSnapshot, + model: Model, + same_model: bool = False, + ignore_dims: Iterable[str] = (), + ) -> ModelDiff | RebuildReason: + """ + Diff ``model`` against a captured ``snapshot``. + + Returns a :class:`ModelDiff` when the change is expressible in + place, or the :class:`RebuildReason` that prevents it. + + Coordinate values are compared on every dim *not* in + ``ignore_dims``; a mismatch triggers + ``RebuildReason.COORD_REINDEX``. Pass ``ignore_dims={"snapshot"}`` + for rolling-horizon use cases where the snapshot coord + legitimately shifts between solves. + + ``same_model`` is a perf hint, **default False**. When True, the + diff trusts ``Constraint._coef_dirty`` to short-circuit the CSR + walk for unchanged containers. That's only safe if every + coefficient mutation went through ``Constraint.update`` (or the + setters that forward there) — direct ``c.coeffs.values[...]`` + writes bypass the flag and would silently miss changes. Pass + ``same_model=True`` only when you own the mutation contract. + """ + reason = _structural_reason(snapshot.structural_key, model) + if reason is not None: + return reason + + builder = _DiffBuilder( + model.variables.label_index, + model.constraints.label_index, + frozenset(ignore_dims), + structural_key=snapshot.structural_key, + ) + + for name, var in model.variables.items(): + reason = builder.diff_var( + name, var, snapshot.var_buffers[name], snapshot.var_coords[name] + ) + if reason is not None: + return reason + + for name, con in model.constraints.items(): + skip = same_model and isinstance(con, Constraint) and not con._coef_dirty + reason = builder.diff_con( + name, + con, + snapshot.con_buffers[name], + snapshot.con_coords[name], + skip_coef_compare=skip, + ) + if reason is not None: + return reason + + reason = builder.diff_objective( + model, snapshot.obj_c, snapshot.obj_quad_present, snapshot.obj_sense + ) + if reason is not None: + return reason + + return builder.finalize() + + @classmethod + def from_models( + cls, + model_a: Model, + model_b: Model, + ignore_dims: Iterable[str] = (), + ) -> ModelDiff | RebuildReason: + """ + Diff two linopy models directly, without capturing a snapshot. + + ``model_a`` is the baseline, ``model_b`` is the target. The + coefficient comparison runs unconditionally — no ``_coef_dirty`` + shortcut applies between independently-built models. Returns a + :class:`ModelDiff` or the :class:`RebuildReason` preventing an + in-place update. + + Captures a snapshot of ``model_a`` and defers to + :meth:`from_snapshot` with ``same_model=False``. + """ + return cls.from_snapshot( + ModelSnapshot.capture(model_a), + model_b, + same_model=False, + ignore_dims=ignore_dims, + ) + + +class _DiffBuilder: + """Accumulates per-container deltas and finalizes them into a ModelDiff.""" + + def __init__( + self, + var_label_index: VariableLabelIndex, + con_label_index: ConstraintLabelIndex, + ignored: frozenset[str], + structural_key: StructuralKey, + ) -> None: + self.var_label_index = var_label_index + self.var_l2p = var_label_index.label_to_pos + self.con_l2p = con_label_index.label_to_pos + self.ignored = ignored + self.structural_key = structural_key + + # Target-state material for the snapshot assembled in finalize(). + self.var_buffers: dict[str, ContainerVarBuffers] = {} + self.con_buffers: dict[str, ContainerConBuffers] = {} + self.var_coords: dict[str, dict[str, np.ndarray]] = {} + self.con_coords: dict[str, dict[str, np.ndarray]] = {} + self._snap_obj_c: np.ndarray | None = None + self._snap_obj_sense: str | None = None + + self.var_bounds_idx: list[np.ndarray] = [] + self.var_bounds_lo: list[np.ndarray] = [] + self.var_bounds_up: list[np.ndarray] = [] + self.var_type_pos: list[np.ndarray] = [] + self.var_type_kinds: list[np.ndarray] = [] + + self.coef_deltas: list[_CoefDelta] = [] + self.con_rhs_idx: list[np.ndarray] = [] + self.con_rhs_vals: list[np.ndarray] = [] + self.con_rhs_signs: list[np.ndarray] = [] + self.con_sign_idx: list[np.ndarray] = [] + self.con_sign_vals: list[np.ndarray] = [] + + self.var_slices: dict[str, VarSlice] = {} + self.con_slices: dict[str, ConSlice] = {} + + self.obj_c_indices: np.ndarray | None = None + self.obj_c_values: np.ndarray | None = None + self.obj_sense: str | None = None + + self._vb_cur = 0 + self._vt_cur = 0 + self._cc_cur = 0 + self._cr_cur = 0 + self._cs_cur = 0 + + def diff_var( + self, + name: str, + var: Variable, + base_buf: ContainerVarBuffers, + base_coords: dict[str, np.ndarray], + ) -> RebuildReason | None: + new_buf = _extract_var_buffers(var) + new_coords = _coord_snapshot(var) + self.var_buffers[name] = new_buf + self.var_coords[name] = new_coords + if not _coords_equal(base_coords, new_coords, self.ignored): + return RebuildReason.COORD_REINDEX + + bound_mask = (new_buf.lower != base_buf.lower) | ( + new_buf.upper != base_buf.upper + ) + bounds_changed = bool(bound_mask.any()) + type_changed = new_buf.type != base_buf.type + if not (bounds_changed or type_changed): + return None + + b_start, t_start = self._vb_cur, self._vt_cur + if bounds_changed: + local_idx = np.flatnonzero(bound_mask) + positions = self.var_l2p[new_buf.active_labels[local_idx]] + self.var_bounds_idx.append(positions.astype(np.int32, copy=False)) + self.var_bounds_lo.append( + new_buf.lower[local_idx].astype(np.float64, copy=False) + ) + self.var_bounds_up.append( + new_buf.upper[local_idx].astype(np.float64, copy=False) + ) + self._vb_cur += local_idx.size + if type_changed: + positions = self.var_l2p[new_buf.active_labels].astype(np.int32, copy=False) + self.var_type_pos.append(positions) + self.var_type_kinds.append( + np.full(positions.size, new_buf.type, dtype=object) + ) + self._vt_cur += positions.size + self.var_slices[name] = VarSlice( + bounds=slice(b_start, self._vb_cur), + type=slice(t_start, self._vt_cur), + ) + return None + + def diff_con( + self, + name: str, + con: ConstraintBase, + base_buf: ContainerConBuffers, + base_coords: dict[str, np.ndarray], + skip_coef_compare: bool, + ) -> RebuildReason | None: + new_buf = _extract_con_buffers(con, self.var_label_index) + new_coords = _coord_snapshot(con) + self.con_buffers[name] = new_buf + self.con_coords[name] = new_coords + if not _coords_equal(base_coords, new_coords, self.ignored): + return RebuildReason.COORD_REINDEX + if not _same(new_buf.indptr, base_buf.indptr): + return RebuildReason.SPARSITY + if not _same(new_buf.indices, base_buf.indices): + return RebuildReason.SPARSITY + + n_rows = new_buf.active_labels.size + if n_rows == 0: + return None + + changed_rows = None + if not (skip_coef_compare or new_buf.data is base_buf.data): + data_diff = new_buf.data != base_buf.data + if data_diff.any(): + nnz_per_row = np.diff(new_buf.indptr) + row_of_nnz = np.repeat(np.arange(n_rows), nnz_per_row) + changed_rows = np.unique(row_of_nnz[data_diff]) + + rhs_idx = None + if new_buf.rhs is not base_buf.rhs: + rhs_idx = np.flatnonzero(new_buf.rhs != base_buf.rhs) + if rhs_idx.size == 0: + rhs_idx = None + sign_idx = None + if new_buf.sign is not base_buf.sign: + sign_idx = np.flatnonzero(new_buf.sign != base_buf.sign) + if sign_idx.size == 0: + sign_idx = None + + if changed_rows is None and rhs_idx is None and sign_idx is None: + return None + + c_start, r_start, s_start = self._cc_cur, self._cr_cur, self._cs_cur + if changed_rows is not None: + row_positions = self.con_l2p[new_buf.active_labels[changed_rows]].astype( + np.int32, copy=False + ) + indptr = new_buf.indptr + nnz = int((indptr[changed_rows + 1] - indptr[changed_rows]).sum()) + self.coef_deltas.append( + _CoefDelta(new_buf, changed_rows, row_positions, nnz) + ) + self._cc_cur += nnz + if rhs_idx is not None: + positions = self.con_l2p[new_buf.active_labels[rhs_idx]] + self.con_rhs_idx.append(positions.astype(np.int32, copy=False)) + self.con_rhs_vals.append( + new_buf.rhs[rhs_idx].astype(np.float64, copy=False) + ) + self.con_rhs_signs.append(new_buf.sign[rhs_idx]) + self._cr_cur += rhs_idx.size + if sign_idx is not None: + positions = self.con_l2p[new_buf.active_labels[sign_idx]] + self.con_sign_idx.append(positions.astype(np.int32, copy=False)) + self.con_sign_vals.append(new_buf.sign[sign_idx]) + self._cs_cur += sign_idx.size + self.con_slices[name] = ConSlice( + coef=slice(c_start, self._cc_cur), + rhs=slice(r_start, self._cr_cur), + sign=slice(s_start, self._cs_cur), + ) + return None + + def diff_objective( + self, + model: Model, + base_obj_c: np.ndarray, + base_obj_quad: bool, + base_obj_sense: str, + ) -> RebuildReason | None: + if model.objective.is_quadratic or base_obj_quad: + return RebuildReason.QUAD_OBJ + + obj_c = _objective_linear_vector(model) + self._snap_obj_c = obj_c + self._snap_obj_sense = model.objective.sense + obj_diff_mask = obj_c != base_obj_c + if obj_diff_mask.any(): + self.obj_c_indices = np.flatnonzero(obj_diff_mask).astype( + np.int32, copy=False + ) + self.obj_c_values = obj_c[self.obj_c_indices].astype(np.float64, copy=False) + if model.objective.sense != base_obj_sense: + self.obj_sense = model.objective.sense + return None + + def finalize(self) -> ModelDiff: + assert self._snap_obj_c is not None and self._snap_obj_sense is not None + snapshot = ModelSnapshot( + structural_key=self.structural_key, + var_buffers=self.var_buffers, + con_buffers=self.con_buffers, + var_coords=self.var_coords, + con_coords=self.con_coords, + obj_c=self._snap_obj_c, + obj_quad_present=False, + obj_sense=self._snap_obj_sense, + ) + return ModelDiff( + snapshot=snapshot, + var_bounds_indices=_cat(self.var_bounds_idx, np.int32), + var_bounds_lower=_cat(self.var_bounds_lo, np.float64), + var_bounds_upper=_cat(self.var_bounds_up, np.float64), + var_type_positions=_cat(self.var_type_pos, np.int32), + var_type_kinds=_cat(self.var_type_kinds, object), + coef_deltas=self.coef_deltas, + n_coef_updates=self._cc_cur, + con_rhs_indices=_cat(self.con_rhs_idx, np.int32), + con_rhs_values=_cat(self.con_rhs_vals, np.float64), + con_rhs_signs=_cat(self.con_rhs_signs, "U1"), + con_sign_indices=_cat(self.con_sign_idx, np.int32), + con_sign_values=_cat(self.con_sign_vals, "U1"), + obj_c_indices=self.obj_c_indices, + obj_c_values=self.obj_c_values, + obj_sense=self.obj_sense, + var_slices=self.var_slices, + con_slices=self.con_slices, + ) diff --git a/linopy/persistent/errors.py b/linopy/persistent/errors.py new file mode 100644 index 000000000..c61592079 --- /dev/null +++ b/linopy/persistent/errors.py @@ -0,0 +1,25 @@ +from __future__ import annotations + + +class UnsupportedUpdate(Exception): + pass + + +class RebuildRequiredError(RuntimeError): + """ + Raised when an in-place update is required but a rebuild is needed. + + Carries the :class:`RebuildReason` that forced the rebuild attempt. + """ + + def __init__(self, reason: object, message: str | None = None) -> None: + self.reason = reason + super().__init__(message or f"rebuild required: {reason}") + + +class UpdatesDisabledError(RuntimeError): + """ + Raised when an in-place update is requested on a solver built with + ``track_updates=False``. Reconstruct the solver with ``track_updates=True`` + to enable diff-based updates. + """ diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py new file mode 100644 index 000000000..fd758ea35 --- /dev/null +++ b/linopy/persistent/snapshot.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import numpy as np + +from linopy import expressions +from linopy.constraints import Constraint + +if TYPE_CHECKING: + from linopy.constraints import ConstraintBase + from linopy.model import Model + from linopy.variables import Variable, VariableLabelIndex + + +class VarKind(enum.Enum): + CONTINUOUS = "continuous" + BINARY = "binary" + INTEGER = "integer" + SEMI_CONTINUOUS = "semi_continuous" + + +def _variable_type(var: Variable) -> VarKind: + attrs = var.attrs + if attrs.get("binary"): + return VarKind.BINARY + if attrs.get("integer"): + return VarKind.INTEGER + if attrs.get("semi_continuous"): + return VarKind.SEMI_CONTINUOUS + return VarKind.CONTINUOUS + + +def _objective_linear_vector(model: Model) -> np.ndarray: + vlabels = model.variables.label_index.vlabels + label_to_pos = model.variables.label_index.label_to_pos + result = np.zeros(len(vlabels), dtype=np.float64) + expr = model.objective.expression + if isinstance(expr, expressions.QuadraticExpression): + vars_2d = expr.data.vars.values + coeffs_all = expr.data.coeffs.values.ravel() + vars1, vars2 = vars_2d[0], vars_2d[1] + linear = (vars1 == -1) | (vars2 == -1) + var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) + coeffs = coeffs_all[linear] + else: + var_labels = expr.data.vars.values.ravel() + coeffs = expr.data.coeffs.values.ravel() + mask = var_labels != -1 + np.add.at(result, label_to_pos[var_labels[mask]], coeffs[mask]) + return result + + +def _extract_var_buffers(var: Variable) -> ContainerVarBuffers: + # Boolean masking copies, so the buffers never alias the live model + # arrays — the snapshot stays a valid baseline even after in-place + # ``.values[...]`` mutations. + labels_flat = var.labels.values.ravel() + mask = labels_flat != -1 + return ContainerVarBuffers( + lower=var.lower.values.ravel()[mask].astype(np.float64, copy=False), + upper=var.upper.values.ravel()[mask].astype(np.float64, copy=False), + type=_variable_type(var), + active_labels=labels_flat[mask].astype(np.int64, copy=False), + ) + + +def _extract_con_buffers( + con: ConstraintBase, var_label_index: VariableLabelIndex +) -> ContainerConBuffers: + """ + Extract flat constraint buffers without copying. + + Mutable ``Constraint`` objects build fresh arrays in + ``to_matrix_with_rhs``, so the buffers are exclusively owned. + ``CSRConstraint`` returns its stored arrays — the buffers share memory + with the constraint, every mutation path rebinds whole arrays + (copy-on-write), and the diff uses object identity to skip comparisons + on untouched containers. + """ + csr, con_labels, b, sense = con.to_matrix_with_rhs(var_label_index) + return ContainerConBuffers( + indptr=csr.indptr, + indices=csr.indices, + data=np.asarray(csr.data, dtype=np.float64), + rhs=np.asarray(b, dtype=np.float64), + sign=np.asarray(sense, dtype="U1"), + active_labels=np.asarray(con_labels, dtype=np.int64), + ) + + +@dataclass(frozen=True) +class StructuralKey: + var_container_names: tuple[str, ...] + con_container_names: tuple[str, ...] + vlabels: np.ndarray + clabels: np.ndarray + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, StructuralKey) + and self.var_container_names == other.var_container_names + and self.con_container_names == other.con_container_names + and np.array_equal(self.vlabels, other.vlabels) + and np.array_equal(self.clabels, other.clabels) + ) + + __hash__ = None # type: ignore[assignment] + + +@dataclass(frozen=True) +class ContainerVarBuffers: + lower: np.ndarray + upper: np.ndarray + type: VarKind + active_labels: np.ndarray + + +@dataclass(frozen=True) +class ContainerConBuffers: + indptr: np.ndarray + indices: np.ndarray + data: np.ndarray + rhs: np.ndarray + sign: np.ndarray + active_labels: np.ndarray + + +def _coord_snapshot(obj: Variable | ConstraintBase) -> dict[str, np.ndarray]: + return {str(name): np.asarray(idx) for name, idx in obj.indexes.items()} + + +def clear_coef_dirty(model: Model) -> None: + """ + Reset ``Constraint._coef_dirty`` on every constraint of ``model``. + + Must be called exactly when a snapshot reflecting the model's current + state is adopted by a tracking solver — clearing without adopting makes + a later ``same_model=True`` diff silently skip changed coefficients. + """ + for con in model.constraints.data.values(): + if isinstance(con, Constraint): + con._coef_dirty = False + + +@dataclass +class ModelSnapshot: + structural_key: StructuralKey + var_buffers: dict[str, ContainerVarBuffers] = field(default_factory=dict) + con_buffers: dict[str, ContainerConBuffers] = field(default_factory=dict) + var_coords: dict[str, dict[str, np.ndarray]] = field(default_factory=dict) + con_coords: dict[str, dict[str, np.ndarray]] = field(default_factory=dict) + obj_c: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=np.float64)) + obj_quad_present: bool = False + obj_sense: str = "min" + + @classmethod + def capture(cls, model: Model) -> ModelSnapshot: + var_label_index = model.variables.label_index + con_label_index = model.constraints.label_index + + structural_key = StructuralKey( + var_container_names=tuple(model.variables), + con_container_names=tuple(model.constraints), + vlabels=var_label_index.vlabels, + clabels=con_label_index.clabels, + ) + + var_buffers = { + name: _extract_var_buffers(var) for name, var in model.variables.items() + } + con_buffers = { + name: _extract_con_buffers(con, var_label_index) + for name, con in model.constraints.items() + } + var_coords = { + name: _coord_snapshot(var) for name, var in model.variables.items() + } + con_coords = { + name: _coord_snapshot(con) for name, con in model.constraints.items() + } + + return cls( + structural_key=structural_key, + var_buffers=var_buffers, + con_buffers=con_buffers, + var_coords=var_coords, + con_coords=con_coords, + obj_c=_objective_linear_vector(model), + obj_quad_present=model.objective.is_quadratic, + obj_sense=model.objective.sense, + ) diff --git a/linopy/piecewise.py b/linopy/piecewise.py new file mode 100644 index 000000000..4099af95d --- /dev/null +++ b/linopy/piecewise.py @@ -0,0 +1,1944 @@ +""" +Piecewise linear constraint formulations. + +Provides SOS2, incremental, pure LP, and disjunctive piecewise linear +constraint methods for use with linopy.Model. +""" + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Sequence +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Literal, TypeAlias, TypeGuard + +import numpy as np +import pandas as pd +import xarray as xr +from xarray import DataArray + +from linopy.constants import ( + BREAKPOINT_DIM, + EQUAL, + GREATER_EQUAL, + HELPER_DIMS, + LESS_EQUAL, + LP_PIECE_DIM, + PWL_ACTIVE_BOUND_SUFFIX, + PWL_BINARY_ORDER_SUFFIX, + PWL_CHORD_SUFFIX, + PWL_CONVEX_SUFFIX, + PWL_CONVEXITY, + PWL_DELTA_BOUND_SUFFIX, + PWL_DELTA_SUFFIX, + PWL_DOMAIN_HI_SUFFIX, + PWL_DOMAIN_LO_SUFFIX, + PWL_FILL_ORDER_SUFFIX, + PWL_LAMBDA_SUFFIX, + PWL_LINK_DIM, + PWL_LINK_SUFFIX, + PWL_METHOD, + PWL_METHODS, + PWL_ORDER_BINARY_SUFFIX, + PWL_OUTPUT_LINK_SUFFIX, + PWL_SEGMENT_BINARY_SUFFIX, + PWL_SELECT_SUFFIX, + SEGMENT_DIM, + SIGNS, + EvolvingAPIWarning, + sign_replace_dict, +) + +if TYPE_CHECKING: + from linopy.constraints import Constraint, Constraints + from linopy.expressions import LinearExpression + from linopy.model import Model + from linopy.types import LinExprLike + from linopy.variables import Variables + +logger = logging.getLogger(__name__) + +# Each user-facing piecewise entry point fires its EvolvingAPIWarning at +# most once per process. Without dedup, a single model build emits the +# verbose warning hundreds of times and drowns out other output. +_EvolvingApiKey: TypeAlias = Literal[ + "tangent_lines", "add_piecewise_formulation", "Slopes" +] +_emitted_evolving_warnings: set[_EvolvingApiKey] = set() + + +def _warn_evolving_api(key: _EvolvingApiKey, message: str, stacklevel: int = 3) -> None: + """ + Emit an :class:`EvolvingAPIWarning` at most once per session per ``key``. + + ``stacklevel`` defaults to 3 (helper → entry-point function → user + code). Pass a larger value when called from one frame deeper than + a function — e.g. from a dataclass ``__post_init__``, which is + itself invoked by an auto-generated ``__init__``. + """ + if key in _emitted_evolving_warnings: + return + _emitted_evolving_warnings.add(key) + warnings.warn(message, category=EvolvingAPIWarning, stacklevel=stacklevel) + + +# Accepted input types for breakpoint-like data +BreaksLike: TypeAlias = ( + Sequence[float] + | np.ndarray + | DataArray + | pd.Series + | pd.DataFrame + | dict[str, Sequence[float]] +) + +# Accepted input types for segment-like data (2D: segments × breakpoints) +SegmentsLike: TypeAlias = ( + Sequence[Sequence[float]] + | np.ndarray + | DataArray + | pd.DataFrame + | dict[str, Sequence[Sequence[float]]] +) + + +# --------------------------------------------------------------------------- +# Deferred slopes spec +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True, repr=False, eq=False) +class Slopes: + """ + Per-piece slopes + initial y-value, deferred until an x grid is known. + + Used as the second element of a tuple in + :func:`add_piecewise_formulation`. When any :class:`Slopes` tuple is + present, **exactly one** other tuple must carry explicit breakpoints — + that tuple's values are the x grid against which all :class:`Slopes` + are integrated:: + + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), # the x grid + (fuel, Slopes([1.2, 1.4, 1.7], y0=0)), # integrated against power + ) + + With two or more non-:class:`Slopes` tuples there is no canonical x + axis, and the call raises :class:`ValueError`. Resolve the + :class:`Slopes` explicitly via :meth:`to_breakpoints` in that case, + or for any standalone use:: + + bp = Slopes([1.2, 1.4, 1.7], y0=0).to_breakpoints([0, 30, 60, 100]) + + Parameters + ---------- + values : BreaksLike + Per-piece slopes. 1D for shared breakpoints; 2D (DataFrame / + dict / DataArray with entity dim) for per-entity slopes. + y0 : float, dict, pd.Series, or DataArray, default 0.0 + y-value at the first breakpoint. Scalar broadcasts to all + entities; dict/Series/DataArray provides per-entity values. + align : {"pieces", "leading"}, default "pieces" + Alignment of ``values`` relative to the x grid. + + - ``"pieces"``: ``len(values) == len(x_points) - 1``; + ``values[i]`` is the slope between ``x[i]`` and ``x[i+1]``. + - ``"leading"``: ``len(values) == len(x_points)``; ``values[0]`` + must be NaN and is dropped, ``values[i]`` for ``i>=1`` is the + slope between ``x[i-1]`` and ``x[i]``. Useful when a marginal + value is tabulated alongside each breakpoint with the first + row's marginal undefined. + dim : str, optional + Entity dimension name. Required when ``values`` is a + ``pd.DataFrame`` or ``dict``. + + Warns + ----- + EvolvingAPIWarning + :class:`Slopes` is part of the newly-added piecewise API. Its + constructor signature and dispatch semantics may be refined. + Silence with ``warnings.filterwarnings("ignore", + category=linopy.EvolvingAPIWarning)``. + """ + + values: BreaksLike + y0: Real | dict[str, Real] | pd.Series | DataArray = 0.0 + align: Literal["pieces", "leading"] = "pieces" + dim: str | None = None + + def __post_init__(self) -> None: + # ``stacklevel=4``: warn → _warn_evolving_api → __post_init__ → + # dataclass-generated ``__init__`` → user code. + _warn_evolving_api( + "Slopes", + "piecewise: Slopes is a new API; the constructor signature and " + "the dispatch rules for inheriting an x grid from sibling tuples " + "may be refined in minor releases.", + stacklevel=4, + ) + + def to_breakpoints(self, x_points: BreaksLike) -> DataArray: + """ + Resolve to a breakpoint :class:`xarray.DataArray`, given an x grid. + + Rarely called directly — typically you pass the :class:`Slopes` + instance to :func:`add_piecewise_formulation` and the x grid is + inherited from a sibling tuple. Use this method for inspection + or when building breakpoints outside the formulation pipeline. + """ + return _breakpoints_from_slopes( + self.values, x_points, self.y0, self.dim, self.align + ) + + def __repr__(self) -> str: + bits = [_summarise_breakslike(self.values), f"y0={self.y0!r}"] + if self.align != "pieces": + bits.append(f"align={self.align!r}") + if self.dim is not None: + bits.append(f"dim={self.dim!r}") + return f"Slopes({', '.join(bits)})" + + def __eq__(self, other: object) -> bool: + """ + Value-equality across the field types accepted by the constructor. + + Two ``Slopes`` are equal iff every field matches: + + * ``align`` and ``dim`` compare with ``==`` (str / None). + * ``y0`` and ``values`` dispatch on type via :func:`_values_equal`: + numeric scalars compare by value across types (``int 0 == + float 0.0 == np.float64(0)``); ``list`` and ``tuple`` are + promoted to ndarray so NaN content compares element-wise + regardless of which NaN object was used; ndarrays use + ``np.array_equal(equal_nan=True)`` (with a fallback for + non-numeric dtypes); ``pd.Series`` / ``pd.DataFrame`` / + ``DataArray`` use ``.equals``; ``dict`` recurses on matching + keys. + + Non-``Slopes`` operands return ``NotImplemented`` per Python + convention. + + Caveats + ------- + * ``Series.equals`` / ``DataFrame.equals`` / ``DataArray.equals`` + are *order-sensitive*: two frames with the same content but + reordered rows / columns / coords compare unequal. + * Cross-container coercion is limited to ``list``/``tuple`` → + ndarray. A ``dict`` and a ``DataFrame`` describing the same + per-entity slopes still compare unequal. + + ``__hash__`` is set to ``None`` (unhashable) since the inner + ``values`` may be a mutable container. + """ + if not isinstance(other, Slopes): + return NotImplemented + return ( + self.align == other.align + and self.dim == other.dim + and _values_equal(self.y0, other.y0) + and _values_equal(self.values, other.values) + ) + + __hash__ = None # type: ignore[assignment] + + +def _is_numeric_scalar(x: object) -> TypeGuard[Real]: + return isinstance(x, Real) and not isinstance(x, bool) + + +def _values_equal(a: object, b: object) -> bool: + """ + Type-dispatched equality for ``Slopes`` field values (NaN-safe). + + Numeric scalars compare by value across types (``int 0 == float 0.0 == + np.float64(0)``); ``bool`` is excluded. Lists / tuples are promoted + to ndarray so in-place ``float('nan')`` content compares NaN-safe. + Non-numeric ndarray dtypes fall back to ``np.array_equal`` without + ``equal_nan``. ``DataFrame`` / ``Series`` / ``DataArray`` use + ``.equals``; ``dict`` recurses on matching keys. + """ + if _is_numeric_scalar(a) and _is_numeric_scalar(b): + af, bf = float(a), float(b) + return af == bf or (af != af and bf != bf) + + if isinstance(a, list | tuple): + a = np.asarray(a) + if isinstance(b, list | tuple): + b = np.asarray(b) + + if isinstance(a, np.ndarray): + if not isinstance(b, np.ndarray) or a.shape != b.shape: + return False + try: + return bool(np.array_equal(a, b, equal_nan=True)) + except TypeError: + return bool(np.array_equal(a, b)) + + if isinstance(a, pd.DataFrame): + return isinstance(b, pd.DataFrame) and bool(a.equals(b)) + if isinstance(a, pd.Series): + return isinstance(b, pd.Series) and bool(a.equals(b)) + if isinstance(a, DataArray): + return isinstance(b, DataArray) and bool(a.equals(b)) + + if isinstance(a, dict): + return ( + isinstance(b, dict) + and a.keys() == b.keys() + and all(_values_equal(a[k], b[k]) for k in a) + ) + + return type(a) is type(b) and bool(a == b) + + +def _summarise_breakslike(v: BreaksLike) -> str: + """Compact one-line summary of a BreaksLike value for use in reprs.""" + if isinstance(v, DataArray): + sizes = ", ".join(f"{d}: {s}" for d, s in v.sizes.items()) + return f"" + if isinstance(v, pd.DataFrame): + return f"" + if isinstance(v, pd.Series): + return f"" + if isinstance(v, dict): + return f"" + + arr = np.asarray(v) + if arr.ndim > 1: + return f"" + seq: list = arr.tolist() + if len(seq) <= 8: + return "[" + ", ".join(_short_num(x) for x in seq) + "]" + head = ", ".join(_short_num(x) for x in seq[:3]) + tail = ", ".join(_short_num(x) for x in seq[-2:]) + return f"[{head}, ..., {tail}] ({len(seq)} items)" + + +def _short_num(x: object) -> str: + """Compact number formatting for repr — ``g`` for floats, ``repr`` else.""" + if isinstance(x, float): + return f"{x:g}" + return repr(x) + + +# Tuple element type covering both eager (DataArray etc.) and deferred (Slopes) bps. +BreaksOrSlopes: TypeAlias = BreaksLike | Slopes + + +# --------------------------------------------------------------------------- +# Result type +# --------------------------------------------------------------------------- + + +@dataclass(slots=True, repr=False) +class PiecewiseFormulation: + """ + Result of ``add_piecewise_formulation``. + + Groups all auxiliary variables and constraints created by a single + piecewise formulation. Stores only names internally; ``variables`` + and ``constraints`` properties return live views from the model. + + Attributes + ---------- + name : str + Formulation name (used as prefix for auxiliary variables and + constraints). + method : PWL_METHOD + Resolved method actually used. Never ``"auto"``; if the caller + passed ``method="auto"``, this holds the method that was chosen. + convexity : PWL_CONVEXITY or None + Shape of the piecewise curve along the breakpoint axis when it is + well-defined (exactly two expressions, non-disjunctive, strictly + monotonic ``x`` breakpoints). ``None`` otherwise. + """ + + name: str + method: PWL_METHOD + """Resolved formulation method (see :data:`~linopy.constants.PWL_METHOD`).""" + variable_names: list[str] + constraint_names: list[str] + model: Model + convexity: PWL_CONVEXITY | None = None + """Shape of the piecewise curve when well-defined (see :data:`~linopy.constants.PWL_CONVEXITY`), else ``None``.""" + + @property + def variables(self) -> Variables: + """View of the auxiliary variables in this formulation.""" + return self.model.variables[self.variable_names] + + @property + def constraints(self) -> Constraints: + """View of the auxiliary constraints in this formulation.""" + return self.model.constraints[self.constraint_names] + + def _user_dims_with_sizes(self) -> dict[str, int]: + """ + User-facing dims across the formulation's variables, with sizes. + + Skips internal ``_``-prefixed dims (e.g. ``_pwl_var``). Insertion + order is preserved, so callers can use the keys as a stable + ordered list. + """ + dims: dict[str, int] = {} + for var in self.variables.data.values(): + for d in var.coords: + ds = str(d) + if not ds.startswith("_") and ds not in dims: + dims[ds] = var.data.sizes[d] + return dims + + def _user_dims(self) -> list[str]: + """User-facing dim names across this formulation's auxiliary variables.""" + return list(self._user_dims_with_sizes()) + + def __repr__(self) -> str: + user_dims = self._user_dims_with_sizes() + dims_str = ", ".join(f"{d}: {s}" for d, s in user_dims.items()) + header = f"PiecewiseFormulation `{self.name}`" + if dims_str: + header += f" [{dims_str}]" + suffix: str = self.method + if self.convexity is not None: + suffix += f", {self.convexity}" + r = f"{header} — {suffix}\n" + r += " Variables:\n" + for vname, var in self.variables.items(): + dims = ", ".join(str(d) for d in var.coords) if var.coords else "" + r += f" * {vname} ({dims})\n" if dims else f" * {vname}\n" + r += " Constraints:\n" + for cname, con in self.constraints.items(): + dims = ", ".join(str(d) for d in con.coords) if con.coords else "" + r += f" * {cname} ({dims})\n" if dims else f" * {cname}\n" + return r + + +def _get_piecewise_groups(model: Model) -> tuple[set[str], set[str]]: + """ + Names of auxiliary variables/constraints that belong to a piecewise + formulation. Returned as separate sets because variables and + constraints live in independent namespaces in the model. + """ + var_names: set[str] = set() + con_names: set[str] = set() + for pwl in model._piecewise_formulations.values(): + var_names.update(pwl.variable_names) + con_names.update(pwl.constraint_names) + return var_names, con_names + + +def _repr_summary(model: Model) -> str: + """ + Render the model-level summary of all piecewise formulations. + + Returns the empty string when the model has no formulations so the + caller can unconditionally concatenate. + """ + if not model._piecewise_formulations: + return "" + r = "\nPiecewise Formulations:\n----------------------\n" + for pwl in model._piecewise_formulations.values(): + n_vars = len(pwl.variable_names) + n_cons = len(pwl.constraint_names) + user_dims = pwl._user_dims() + dims_str = f" ({', '.join(user_dims)})" if user_dims else "" + r += f" * {pwl.name}{dims_str} — {pwl.method}, {n_vars} vars, {n_cons} cons\n" + return r + + +# --------------------------------------------------------------------------- +# DataArray construction helpers +# --------------------------------------------------------------------------- + + +def _strip_nan(vals: Sequence[float] | np.ndarray) -> list[float]: + """Remove NaN values from a sequence.""" + arr = np.asarray(vals, dtype=float) + return arr[~np.isnan(arr)].tolist() + + +def _rename_to_pieces(da: DataArray, piece_index: np.ndarray) -> DataArray: + """Rename breakpoint dim to piece dim and reassign coordinates.""" + da = da.rename({BREAKPOINT_DIM: LP_PIECE_DIM}) + da[LP_PIECE_DIM] = piece_index + return da + + +def _sequence_to_array(values: Sequence[float] | np.ndarray | pd.Series) -> DataArray: + arr = np.asarray(values, dtype=float) + if arr.ndim != 1: + raise ValueError( + f"Expected a 1D sequence of numeric values, got shape {arr.shape}" + ) + return DataArray( + arr, dims=[BREAKPOINT_DIM], coords={BREAKPOINT_DIM: np.arange(len(arr))} + ) + + +def _dict_to_array(d: dict[str, Sequence[float]], dim: str) -> DataArray: + """Convert a dict of ragged sequences to a NaN-padded 2D DataArray.""" + max_len = max(len(v) for v in d.values()) + keys = list(d.keys()) + data = np.full((len(keys), max_len), np.nan) + for i, k in enumerate(keys): + vals = d[k] + data[i, : len(vals)] = vals + return DataArray( + data, + dims=[dim, BREAKPOINT_DIM], + coords={dim: keys, BREAKPOINT_DIM: np.arange(max_len)}, + ) + + +def _dataframe_to_array(df: pd.DataFrame, dim: str) -> DataArray: + # rows = entities (index), columns = breakpoints + data = np.asarray(df.values, dtype=float) + return DataArray( + data, + dims=[dim, BREAKPOINT_DIM], + coords={dim: list(df.index), BREAKPOINT_DIM: np.arange(df.shape[1])}, + ) + + +def _coerce_breaks(values: BreaksLike, dim: str | None = None) -> DataArray: + """Convert any BreaksLike input to a DataArray with BREAKPOINT_DIM.""" + if isinstance(values, DataArray): + if BREAKPOINT_DIM not in values.dims: + raise ValueError( + f"DataArray must have a '{BREAKPOINT_DIM}' dimension, " + f"got dims {list(values.dims)}" + ) + return values + if isinstance(values, pd.DataFrame): + if dim is None: + raise ValueError("'dim' is required when input is a DataFrame") + return _dataframe_to_array(values, dim) + if isinstance(values, pd.Series): + return _sequence_to_array(values) + if isinstance(values, dict): + if dim is None: + raise ValueError("'dim' is required when input is a dict") + return _dict_to_array(values, dim) + # Sequence (list, tuple, etc.) + return _sequence_to_array(values) + + +def _segments_list_to_array(values: Sequence[Sequence[float]]) -> DataArray: + max_len = max(len(seg) for seg in values) + data = np.full((len(values), max_len), np.nan) + for i, seg in enumerate(values): + data[i, : len(seg)] = seg + return DataArray( + data, + dims=[SEGMENT_DIM, BREAKPOINT_DIM], + coords={ + SEGMENT_DIM: np.arange(len(values)), + BREAKPOINT_DIM: np.arange(max_len), + }, + ) + + +def _dict_segments_to_array( + d: dict[str, Sequence[Sequence[float]]], dim: str +) -> DataArray: + parts = [] + for key, seg_list in d.items(): + arr = _segments_list_to_array(seg_list) + parts.append(arr.expand_dims({dim: [key]})) + combined = xr.concat(parts, dim=dim, coords="minimal") + max_bp = max(max(len(seg) for seg in sl) for sl in d.values()) + max_seg = max(len(sl) for sl in d.values()) + if combined.sizes[BREAKPOINT_DIM] < max_bp or combined.sizes[SEGMENT_DIM] < max_seg: + combined = combined.reindex( + {BREAKPOINT_DIM: np.arange(max_bp), SEGMENT_DIM: np.arange(max_seg)}, + fill_value=np.nan, + ) + return combined + + +def _breakpoints_from_slopes( + slopes: BreaksLike, + x_points: BreaksLike, + y0: Real | dict[str, Real] | pd.Series | DataArray, + dim: str | None, + slopes_align: Literal["pieces", "leading"] = "pieces", +) -> DataArray: + """Convert slopes + x_points + y0 into a breakpoint DataArray.""" + slopes_arr = _coerce_breaks(slopes, dim) + xp_arr = _coerce_breaks(x_points, dim) + + if slopes_align == "leading": + if slopes_arr.sizes[BREAKPOINT_DIM] == 0: + raise ValueError("slopes_align='leading' requires at least one slope entry") + first_slope = slopes_arr.isel({BREAKPOINT_DIM: 0}) + if not bool(first_slope.isnull().all()): + raise ValueError( + "slopes_align='leading' requires the first slope of each " + "entity to be NaN" + ) + slopes_arr = slopes_arr.isel({BREAKPOINT_DIM: slice(1, None)}) + + # 1D case: single set of breakpoints + if slopes_arr.ndim == 1: + if not isinstance(y0, Real): + raise TypeError("When 'slopes' is 1D, 'y0' must be a scalar float") + pts = _slopes_to_points(list(xp_arr.values), list(slopes_arr.values), float(y0)) + return _sequence_to_array(pts) + + # Multi-dim case: per-entity slopes + entity_dims = [d for d in slopes_arr.dims if d != BREAKPOINT_DIM] + if len(entity_dims) != 1: + raise ValueError( + f"Expected exactly one entity dimension in slopes, got {entity_dims}" + ) + entity_dim = str(entity_dims[0]) + entity_keys = slopes_arr.coords[entity_dim].values + + # Resolve y0 per entity + if isinstance(y0, Real): + y0_map: dict[str, float] = {str(k): float(y0) for k in entity_keys} + elif isinstance(y0, dict): + y0_map = {str(k): float(y0[k]) for k in entity_keys} + elif isinstance(y0, pd.Series): + y0_map = {str(k): float(y0[k]) for k in entity_keys} + elif isinstance(y0, DataArray): + y0_map = {str(k): float(y0.sel({entity_dim: k}).item()) for k in entity_keys} + else: + raise TypeError( + f"'y0' must be a float, Series, DataArray, or dict, got {type(y0)}" + ) + + computed: dict[str, Sequence[float]] = {} + for key in entity_keys: + sk = str(key) + sl = _strip_nan(slopes_arr.sel({entity_dim: key}).values) + if entity_dim in xp_arr.dims: + xp = _strip_nan(xp_arr.sel({entity_dim: key}).values) + else: + xp = _strip_nan(xp_arr.values) + computed[sk] = _slopes_to_points(xp, sl, y0_map[sk]) + + return _dict_to_array(computed, entity_dim) + + +# --------------------------------------------------------------------------- +# Public factory functions +# --------------------------------------------------------------------------- + + +def _slopes_to_points( + x_points: list[float], slopes: list[float], y0: float +) -> list[float]: + """ + Convert per-piece slopes + initial y-value to y-coordinates at each breakpoint. + + Internal primitive used by ``Slopes.to_breakpoints``. Public callers + should use :class:`Slopes` (DataArray output) instead. + """ + if len(slopes) != len(x_points) - 1: + raise ValueError( + f"len(slopes) must be len(x_points) - 1, " + f"got {len(slopes)} slopes and {len(x_points)} x_points" + ) + y_points: list[float] = [y0] + for i, s in enumerate(slopes): + y_points.append(y_points[-1] + s * (x_points[i + 1] - x_points[i])) + return y_points + + +def breakpoints( + values: BreaksLike, + *, + dim: str | None = None, +) -> DataArray: + """ + Create a breakpoint DataArray for piecewise linear constraints. + + Parameters + ---------- + values : BreaksLike + Breakpoint values. Accepted types: ``Sequence[float]``, + ``pd.Series``, ``pd.DataFrame``, or ``xr.DataArray``. + A 1D input (list, Series) creates 1D breakpoints. + A 2D input (DataFrame, multi-dim DataArray) creates per-entity + breakpoints (``dim`` is required for DataFrame). + dim : str, optional + Entity dimension name. Required when ``values`` is a + ``pd.DataFrame`` or ``dict``. + + Returns + ------- + DataArray + + See Also + -------- + Slopes : per-piece slopes + ``y0`` (deferred or standalone via + :meth:`Slopes.to_breakpoints`). + """ + return _coerce_breaks(values, dim) + + +def _coerce_segments(values: SegmentsLike, dim: str | None = None) -> DataArray: + """Convert any SegmentsLike input to a DataArray with SEGMENT_DIM and BREAKPOINT_DIM.""" + if isinstance(values, DataArray): + if SEGMENT_DIM not in values.dims or BREAKPOINT_DIM not in values.dims: + raise ValueError( + f"DataArray must have both '{SEGMENT_DIM}' and '{BREAKPOINT_DIM}' " + f"dimensions, got dims {list(values.dims)}" + ) + return values + if isinstance(values, pd.DataFrame): + data = np.asarray(values.values, dtype=float) + return DataArray( + data, + dims=[SEGMENT_DIM, BREAKPOINT_DIM], + coords={ + SEGMENT_DIM: np.arange(data.shape[0]), + BREAKPOINT_DIM: np.arange(data.shape[1]), + }, + ) + if isinstance(values, dict): + if dim is None: + raise ValueError("'dim' is required when 'values' is a dict") + return _dict_segments_to_array(values, dim) + # Sequence[Sequence[float]] + return _segments_list_to_array(list(values)) + + +def segments( + values: SegmentsLike, + *, + dim: str | None = None, +) -> DataArray: + """ + Create a segmented breakpoint DataArray for disjunctive piecewise constraints. + + Parameters + ---------- + values : SegmentsLike + Segment breakpoints. Accepted types: ``Sequence[Sequence[float]]``, + ``pd.DataFrame`` (rows=segments, columns=breakpoints), + ``xr.DataArray`` (must have ``SEGMENT_DIM`` and ``BREAKPOINT_DIM``), + or ``dict[str, Sequence[Sequence[float]]]`` (requires ``dim``). + dim : str, optional + Entity dimension name. Required when ``values`` is a dict. + + Returns + ------- + DataArray + """ + return _coerce_segments(values, dim) + + +def _tangent_lines_impl( + x: LinExprLike, + x_points: BreaksLike, + y_points: BreaksLike, +) -> LinearExpression: + """ + Chord-expression math — the body of ``tangent_lines`` without the + :class:`EvolvingAPIWarning`. Called internally by ``_add_lp`` so a + single ``add_piecewise_formulation((y, y_pts, "<="), (x, x_pts))`` + emits exactly one warning, not two. + """ + from linopy.expressions import LinearExpression as LinExpr + from linopy.variables import Variable + + x_points = _coerce_breaks(x_points) + y_points = _coerce_breaks(y_points) + + dx = x_points.diff(BREAKPOINT_DIM) + dy = y_points.diff(BREAKPOINT_DIM) + piece_index = np.arange(dx.sizes[BREAKPOINT_DIM]) + + slopes = _rename_to_pieces(dy / dx, piece_index) + x_base = _rename_to_pieces( + x_points.isel({BREAKPOINT_DIM: slice(None, -1)}), piece_index + ) + y_base = _rename_to_pieces( + y_points.isel({BREAKPOINT_DIM: slice(None, -1)}), piece_index + ) + + intercepts = y_base - slopes * x_base + + if not isinstance(x, Variable | LinExpr): + raise TypeError(f"x must be a Variable or LinearExpression, got {type(x)}") + + return slopes * _to_linexpr(x) + intercepts + + +def tangent_lines( + x: LinExprLike, + x_points: BreaksLike, + y_points: BreaksLike, +) -> LinearExpression: + r""" + Compute tangent-line (chord) expressions for a piecewise linear function. + + Low-level helper returning a :class:`~linopy.expressions.LinearExpression` + with an extra piece dimension. Each element along the piece dimension + is the chord of one piece: :math:`m_k \cdot x + c_k`. No auxiliary + variables are created. + + For most users: prefer :func:`add_piecewise_formulation` with a + bounded tuple ``(y, y_pts, "<=")`` / ``(y, y_pts, ">=")`` — it builds + on this helper and adds the ``x ∈ [x_min, x_max]`` domain bound plus + a curvature-vs-sign check that catches the "wrong region" case. Use + ``tangent_lines`` directly only when you need to compose the chord + expressions manually (e.g. with other linear terms, or without the + domain bound). + + .. code-block:: python + + t = tangent_lines(power, x_pts, y_pts) + m.add_constraints(fuel <= t) # upper bound (concave f) + m.add_constraints(fuel >= t) # lower bound (convex f) + + Parameters + ---------- + x : Variable or LinearExpression + The input expression. + x_points : BreaksLike + Breakpoint x-coordinates (must be strictly monotonic; both + ascending and descending are accepted). + y_points : BreaksLike + Breakpoint y-coordinates. + + Returns + ------- + LinearExpression + Expression with an additional ``_breakpoint_piece`` dimension + (one entry per piece). + + Warns + ----- + EvolvingAPIWarning + ``tangent_lines`` is part of the newly-added piecewise API; the + returned expression shape and piece-dim name may be refined. + Silence with ``warnings.filterwarnings("ignore", + category=linopy.EvolvingAPIWarning)``. + """ + _warn_evolving_api( + "tangent_lines", + "piecewise: tangent_lines is a new API; the returned expression " + "shape and the piece-dim name may be refined in minor releases. " + "Please share your use cases or concerns at " + "https://github.com/PyPSA/linopy/issues — your feedback shapes " + "what stabilises. This warning fires once per session; silence " + "entirely with " + '`warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)`.', + ) + return _tangent_lines_impl(x, x_points, y_points) + + +# --------------------------------------------------------------------------- +# Internal validation and utility functions +# --------------------------------------------------------------------------- + + +def _resolve_active( + active: LinearExpression, reference: DataArray, active_fill: int | None +) -> LinearExpression: + """ + Resolve a possibly-partial ``active`` gate against the formulation. + + A gate defined over only a subset of the indexed dimension (or with + masked entries) would otherwise be gated as if ``active=0`` and forced + to zero. With ``active_fill is None`` such a gate is rejected; otherwise + the gaps are filled with ``active_fill`` (``1`` = always active, ``0`` = + always off). Dimensions absent from ``active`` broadcast and are left + untouched. + """ + skip = {BREAKPOINT_DIM, SEGMENT_DIM} | set(HELPER_DIMS) + indexers = { + d: reference.indexes[d] + for d in active.coord_dims + if d in reference.indexes and d not in skip + } + aligned = active.reindex(indexers) if indexers else active + + if active_fill is not None: + return aligned.where(aligned.has_terms, active_fill) + + term_dims = [d for d in aligned.vars.dims if d not in aligned.coord_dims] + dangling = ((aligned.vars < 0) & aligned.coeffs.notnull()).any(term_dims) + covered = aligned.has_terms | (aligned.const.notnull() & ~dangling) + if not bool(covered.all()): + raise ValueError( + "`active` is not defined over the full coordinate of the " + "piecewise formulation: it is missing labels (a subset of the " + "coordinate) or has masked entries, which would be gated to " + "zero. Pass `active_fill=1` to treat those entries as always " + "active (or `0` as always off), or pass a fully-defined `active`." + ) + return active + + +def _validate_breakpoint_shapes(bp_a: DataArray, bp_b: DataArray) -> bool: + """ + Validate that two breakpoint arrays have compatible shapes. + + Returns whether the formulation is disjunctive (has segment dimension). + """ + for bp in (bp_a, bp_b): + if BREAKPOINT_DIM not in bp.dims: + raise ValueError( + f"Breakpoints are missing the '{BREAKPOINT_DIM}' dimension, " + f"got dims {list(bp.dims)}. " + "Use the breakpoints() or segments() factory." + ) + + if bp_a.sizes[BREAKPOINT_DIM] != bp_b.sizes[BREAKPOINT_DIM]: + raise ValueError( + f"Breakpoints must have same size along '{BREAKPOINT_DIM}', " + f"got {bp_a.sizes[BREAKPOINT_DIM]} and " + f"{bp_b.sizes[BREAKPOINT_DIM]}" + ) + + a_has_seg = SEGMENT_DIM in bp_a.dims + b_has_seg = SEGMENT_DIM in bp_b.dims + if a_has_seg != b_has_seg: + raise ValueError( + "If one breakpoint array has a segment dimension, " + f"both must. Got dims: {list(bp_a.dims)} and {list(bp_b.dims)}." + ) + if a_has_seg and bp_a.sizes[SEGMENT_DIM] != bp_b.sizes[SEGMENT_DIM]: + raise ValueError(f"Breakpoints must have same size along '{SEGMENT_DIM}'") + + return a_has_seg + + +def _validate_numeric_breakpoint_coords(bp: DataArray) -> None: + coord = bp.coords[BREAKPOINT_DIM] + if not pd.api.types.is_numeric_dtype(coord): + raise ValueError( + f"Breakpoint dimension '{BREAKPOINT_DIM}' must have numeric coordinates " + f"for SOS2 weights, but got {coord.dtype}" + ) + values = np.asarray(coord.values) + if len(values) > 1 and not bool(np.all(np.diff(values) > 0)): + raise ValueError( + f"Breakpoint dimension '{BREAKPOINT_DIM}' coordinates must be " + "strictly increasing for SOS2 weights." + ) + + +def _check_strict_monotonicity(bp: DataArray) -> bool: + """Check if breakpoints are strictly monotonic along BREAKPOINT_DIM (ignoring NaN).""" + diffs = bp.diff(BREAKPOINT_DIM) + pos = (diffs > 0) | diffs.isnull() + neg = (diffs < 0) | diffs.isnull() + all_pos_per_slice = pos.all(BREAKPOINT_DIM) + all_neg_per_slice = neg.all(BREAKPOINT_DIM) + has_non_nan = (~diffs.isnull()).any(BREAKPOINT_DIM) + monotonic = (all_pos_per_slice | all_neg_per_slice) & has_non_nan + return bool(monotonic.all()) + + +def _detect_convexity(x_points: DataArray, y_points: DataArray) -> PWL_CONVEXITY: + """ + Classify the shape of a single piecewise curve ``y = f(x)``. + + Invariant to whether breakpoints are listed ascending or descending in + x — same graph, same label. Multi-entity inputs are aggregated across + entities; to classify per entity, iterate at the call site (see + :data:`PWL_CONVEXITIES` for the possible labels). Callers must + enforce strict x-monotonicity per slice upstream. + """ + dx = x_points.diff(BREAKPOINT_DIM) + slopes = y_points.diff(BREAKPOINT_DIM) / dx + # Flip sign when x descends so the classification matches the + # ascending-x traversal. All dx in a strictly-monotonic slice share + # a sign, so the sum resolves direction per entity. + sd = slopes.diff(BREAKPOINT_DIM) * np.sign(dx.sum(BREAKPOINT_DIM)) + + if int((~sd.isnull()).sum()) == 0: + return "linear" + tol = 1e-10 + nonneg = bool(((sd >= -tol) | sd.isnull()).all()) + nonpos = bool(((sd <= tol) | sd.isnull()).all()) + if nonneg and nonpos: + return "linear" + if nonneg: + return "convex" + if nonpos: + return "concave" + return "mixed" + + +def _has_trailing_nan_only(bp: DataArray) -> bool: + """Check that NaN values only appear as trailing entries along BREAKPOINT_DIM.""" + valid = ~bp.isnull() + cummin = np.minimum.accumulate(valid.values, axis=valid.dims.index(BREAKPOINT_DIM)) + cummin_da = DataArray(cummin, coords=valid.coords, dims=valid.dims) + return not bool((valid & ~cummin_da).any()) + + +def _paired_valid_points(*points: DataArray) -> DataArray: + invalid = points[0].isnull() + for point in points[1:]: + invalid = invalid | point.isnull() + return points[0].where(~invalid) + + +def _validate_shared_coords(points: Sequence[DataArray]) -> None: + skip = {BREAKPOINT_DIM, SEGMENT_DIM} | set(HELPER_DIMS) + for i, left in enumerate(points): + for right in points[i + 1 :]: + for dim in (set(left.dims) & set(right.dims)) - skip: + left_index = pd.Index(left.coords[dim].values) + right_index = pd.Index(right.coords[dim].values) + if not left_index.equals(right_index): + raise ValueError( + f"Breakpoint coordinates for dimension '{dim}' must match." + ) + + +def _validate_expr_coords( + points: Sequence[DataArray], exprs: Sequence[LinearExpression] +) -> None: + skip = {BREAKPOINT_DIM, SEGMENT_DIM} | set(HELPER_DIMS) + for point in points: + for expr in exprs: + for dim in (set(point.dims) & set(expr.coord_dims)) - skip: + point_index = pd.Index(point.coords[dim].values) + expr_index = pd.Index(expr.coords[dim].values) + if not point_index.equals(expr_index): + raise ValueError( + f"Breakpoint coordinates for dimension '{dim}' must match " + "the expression coordinates." + ) + + +def _to_linexpr(expr: LinExprLike) -> LinearExpression: + from linopy.expressions import LinearExpression + + if isinstance(expr, LinearExpression): + return expr + return expr.to_linexpr() + + +def _var_coords_from( + points: DataArray, exclude: set[str] | None = None +) -> list[pd.Index]: + """Extract pd.Index coords from points, excluding specified dimensions.""" + excluded = exclude or set() + return [ + pd.Index(points.coords[d].values, name=d) + for d in points.dims + if d not in excluded + ] + + +def _broadcast_points( + points: DataArray, + *exprs: LinExprLike, + disjunctive: bool = False, +) -> DataArray: + """Broadcast points to cover all dimensions from exprs.""" + skip: set[str] = {BREAKPOINT_DIM} | set(HELPER_DIMS) + if disjunctive: + skip.add(SEGMENT_DIM) + + lin_exprs = [_to_linexpr(e) for e in exprs] + + point_dims = {str(d) for d in points.dims} + + # Iterate exprs/dims in order; a set would give a hash-dependent, + # run-varying expanded dimension order. + expand_map: dict[str, list] = {} + for le in lin_exprs: + for dim in le.coord_dims: + d = str(dim) + if d in skip or d in point_dims or d in expand_map: + continue + if d in le.coords: + expand_map[d] = list(le.coords[d].values) + + if expand_map: + points = points.expand_dims(expand_map) + return points + + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- + + +def add_piecewise_formulation( + model: Model, + *pairs: tuple[LinExprLike, BreaksOrSlopes] + | tuple[LinExprLike, BreaksOrSlopes, Literal["==", "<=", ">="]], + method: PWL_METHOD = "auto", + active: LinExprLike | None = None, + active_fill: int | None = None, + name: str | None = None, +) -> PiecewiseFormulation: + r""" + Add piecewise linear constraints. + + Each positional argument is a ``(expression, breakpoints)`` tuple, or + ``(expression, breakpoints, sign)`` to mark that expression as bounded + by the piecewise curve rather than pinned to it. All expressions are + linked through shared interpolation weights so that every operating + point lies on the same piece of the piecewise curve. + + Example — 2 variables (joint equality, the default):: + + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, [0, 36, 84, 170]), + ) + + Example — 3 variables, CHP plant (joint equality):: + + m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, [0, 40, 85, 160]), + (heat, [0, 25, 55, 95]), + ) + + **Per-tuple sign — inequality bounds:** + + Add ``"<="`` or ``">="`` as a third tuple element to mark a single + expression as bounded by the curve instead of pinned to it. The + remaining tuples are still forced to equality (input on the curve). + Reads directly as the relation it encodes: + + .. code-block:: python + + # fuel <= f(power) — concave curve, bounded above + m.add_piecewise_formulation( + (fuel, y_pts, "<="), + (power, x_pts), + ) + + # cost >= g(load) — convex curve, bounded below + m.add_piecewise_formulation( + (cost, y_pts, ">="), + (load, x_pts), + ) + + For 2-variable inequality on convex/concave curves, ``method="auto"`` + automatically selects a pure-LP tangent-line formulation (no auxiliary + variables). Non-convex curves fall back to SOS2/incremental with the + sign applied to the bounded tuple's link constraint. + + **Restrictions on per-tuple sign:** + + - At most one tuple may carry a non-equality sign. All other tuples + default to ``"=="``. + - With **3 or more** tuples, all signs must be ``"=="``. + + Multi-bounded and N≥3-inequality use cases aren't supported yet. If + you have a concrete use case, please open an issue at + https://github.com/PyPSA/linopy/issues so we can scope it properly. + + Parameters + ---------- + *pairs : tuple of (expression, breakpoints) or (expression, breakpoints, sign) + Each pair links an expression (Variable or LinearExpression) to + its breakpoint values. An optional third element ``"<="`` or + ``">="`` marks that expression as bounded by the curve; if + omitted, the expression is pinned (``"=="``). At least two pairs + are required; at most one may carry a non-equality sign; with + 3+ pairs all signs must be ``"=="``. + method : {"auto", "sos2", "incremental", "lp"}, default "auto" + Formulation method. + ``"lp"`` uses tangent lines (pure LP, no variables) and requires + exactly one tuple with ``"<="`` or ``">="`` plus a matching-curvature + curve with exactly two tuples. + ``"auto"`` picks ``"lp"`` when applicable, otherwise + ``"incremental"`` (monotonic breakpoints) or ``"sos2"``. + active : Variable or LinearExpression, optional + Binary variable that gates the piecewise function. When + ``active=0``, all auxiliary variables are forced to zero. + Not supported with ``method="lp"``. + + ``active`` must cover the formulation's full coordinate. A + *partial* gate — one defined over only a subset of the coordinate's + labels, or carrying masked entries — is rejected unless + ``active_fill`` is set (see below). + + With all-equality tuples (the default), the output is then pinned + to ``0``. With a bounded tuple (``"<="`` / ``">="``), deactivation + only pushes the signed bound to ``0`` (the output is ≤ 0 or ≥ 0 + respectively) — the complementary bound still comes from the + output variable's own lower/upper. In the common case where + the output is naturally non-negative (fuel, cost, heat, …), + just set ``lower=0`` on that variable: combined with the + ``y ≤ 0`` constraint from deactivation, this forces ``y = 0`` + automatically. For outputs that genuinely need both signs you + must add the complementary bound yourself (e.g., a big-M + coupling ``y`` with ``active``). + active_fill : int, optional + Fill value for the gap entries of a partial ``active`` — those where + ``active`` has no label (a subset of the coordinate) or is masked: + ``1`` treats them as always active (ungated), ``0`` as always off. + When ``None`` (the default) a partial ``active`` is rejected instead. + Useful when one formulation mixes gated and ungated entities (e.g. + committable and non-committable units sharing a ``status``). + Transitional convenience: under v1 semantics, pad ``active`` + explicitly with ``active.reindex(coords).fillna(value)`` instead — + this parameter is slated for removal then. + name : str, optional + Base name for generated variables/constraints. + + Returns + ------- + PiecewiseFormulation + + Warns + ----- + EvolvingAPIWarning + ``add_piecewise_formulation`` is a newly-added API; details such + as the per-tuple sign convention and ``active`` + non-equality + sign semantics may be refined based on user feedback. Silence + with ``warnings.filterwarnings("ignore", + category=linopy.EvolvingAPIWarning)``. + """ + _warn_evolving_api( + "add_piecewise_formulation", + "piecewise: add_piecewise_formulation is a new API; some details " + "(e.g. the per-tuple sign convention, active+sign semantics) " + "may be refined in minor releases. Please share your use cases " + "or concerns at https://github.com/PyPSA/linopy/issues — your " + "feedback shapes what stabilises. This warning fires once per " + "session; silence entirely with " + '`warnings.filterwarnings("ignore", category=linopy.EvolvingAPIWarning)`.', + ) + + if method not in PWL_METHODS: + raise ValueError(f"method must be one of {sorted(PWL_METHODS)}, got '{method}'") + + if len(pairs) < 2: + raise TypeError( + "add_piecewise_formulation() requires at least 2 " + "(expression, breakpoints[, sign]) pairs." + ) + + # Parse and normalise per-tuple signs. Each pair is either + # (expr, bp) — sign defaults to "==" — or (expr, bp, sign). + parsed: list[tuple[LinExprLike, BreaksOrSlopes, str]] = [] + for i, pair in enumerate(pairs): + if not isinstance(pair, tuple) or len(pair) not in (2, 3): + raise TypeError( + f"Argument {i + 1} must be a (expression, breakpoints) " + f"or (expression, breakpoints, sign) tuple, got {pair!r}." + ) + if len(pair) == 2: + expr, bp = pair + tuple_sign: str = EQUAL + else: + expr, bp, raw_sign = pair + tuple_sign = sign_replace_dict.get(raw_sign, raw_sign) + if tuple_sign not in SIGNS: + raise ValueError( + f"Argument {i + 1}: sign must be one of " + f"{sorted(SIGNS)}, got {raw_sign!r}." + ) + parsed.append((expr, bp, tuple_sign)) + + slopes_set = {i for i, p in enumerate(parsed) if isinstance(p[1], Slopes)} + if slopes_set: + non_slopes_idx = [i for i in range(len(parsed)) if i not in slopes_set] + if not non_slopes_idx: + raise ValueError( + "All tuples are Slopes; at least one tuple must carry an " + "explicit x grid. Pass the x grid via a regular tuple " + "or call Slopes(...).to_breakpoints(x_pts) explicitly." + ) + if len(non_slopes_idx) > 1: + raise ValueError( + f"Slopes tuples present at positions {sorted(slopes_set)}, " + f"but {len(non_slopes_idx)} non-Slopes tuples carry their " + f"own breakpoint values (positions {non_slopes_idx}). " + "There is no canonical x grid for the Slopes to integrate " + "against — borrowing from any one of them would silently " + "depend on tuple order. Either reduce to a single non-Slopes " + "tuple, or resolve the Slopes explicitly by calling " + "Slopes(...).to_breakpoints(x_pts) before passing it in." + ) + x_grid = parsed[non_slopes_idx[0]][1] + parsed = [ + (expr, bp.to_breakpoints(x_grid), sign) + if isinstance(bp, Slopes) + else (expr, bp, sign) + for expr, bp, sign in parsed + ] + + # At most one non-equality sign; with 3+ tuples, none. + bounded_positions = [i for i, p in enumerate(parsed) if p[2] != EQUAL] + if len(bounded_positions) > 1: + raise ValueError( + "At most one tuple may carry a non-equality sign; got " + f"{len(bounded_positions)} (positions {bounded_positions})." + ) + if len(parsed) >= 3 and bounded_positions: + raise ValueError( + "Non-equality signs are not supported with 3+ tuples. " + "Use sign='==' on all tuples (the default), or reduce to 2 tuples. " + "If you have a concrete use case, please open an issue at " + "https://github.com/PyPSA/linopy/issues." + ) + + signed_idx: int | None + if bounded_positions: + bidx = bounded_positions[0] + signed_idx = bidx + sign: str = parsed[bidx][2] + else: + signed_idx = None + sign = EQUAL + + if method == "lp" and sign == EQUAL: + raise ValueError( + "method='lp' requires exactly one tuple with sign='<=' or '>='." + ) + + coerced_bps: list[DataArray] = [] + for _, bp, _s in parsed: + if not isinstance(bp, DataArray): + bp = _coerce_breaks(bp) + scalar_coords = [c for c in bp.coords if c not in bp.dims] + if scalar_coords: + bp = bp.drop_vars(scalar_coords) + coerced_bps.append(bp) + + disjunctive = SEGMENT_DIM in coerced_bps[0].dims + for i in range(1, len(coerced_bps)): + _validate_breakpoint_shapes(coerced_bps[0], coerced_bps[i]) + + raw_exprs = [expr for expr, _, _ in parsed] + lin_exprs = [_to_linexpr(expr) for expr in raw_exprs] + bp_list = [ + _broadcast_points(bp, *raw_exprs, disjunctive=disjunctive) for bp in coerced_bps + ] + _validate_shared_coords(bp_list) + _validate_expr_coords(bp_list, lin_exprs) + + combined_null = bp_list[0].isnull() + for bp in bp_list[1:]: + combined_null = combined_null | bp.isnull() + bp_mask = ~combined_null if bool(combined_null.any()) else None + + if name is None: + name = f"pwl{model._pwlCounter}" + model._pwlCounter += 1 + + from linopy.variables import Variable + + link_coords: list[str] = [] + for i, expr in enumerate(raw_exprs): + if isinstance(expr, Variable) and expr.name: + link_coords.append(expr.name) + else: + # Internal-prefixed fallback so a user variable named e.g. "1" + # can't collide with the synthetic coord for an unnamed expr. + link_coords.append(f"_pwl_{i}") + + if active is None: + if active_fill is not None: + raise ValueError("`active_fill` has no effect without `active`.") + active_expr = None + else: + active_expr = _resolve_active(_to_linexpr(active), bp_list[0], active_fill) + + if signed_idx is None: + inputs = _PwlInputs( + pinned_exprs=lin_exprs, + pinned_bps=bp_list, + pinned_coords=link_coords, + bounded_expr=None, + bounded_bp=None, + bounded_coord=None, + bounded_sign=EQUAL, + bp_mask=bp_mask, + ) + else: + inputs = _PwlInputs( + pinned_exprs=[e for j, e in enumerate(lin_exprs) if j != signed_idx], + pinned_bps=[b for j, b in enumerate(bp_list) if j != signed_idx], + pinned_coords=[c for j, c in enumerate(link_coords) if j != signed_idx], + bounded_expr=lin_exprs[signed_idx], + bounded_bp=bp_list[signed_idx], + bounded_coord=link_coords[signed_idx], + bounded_sign=sign, + bp_mask=bp_mask, + ) + + vars_before = set(model.variables) + cons_before = set(model.constraints) + + if disjunctive: + if method == "incremental": + raise ValueError( + "Incremental method is not supported for disjunctive constraints" + ) + if method == "lp": + raise ValueError( + "method='lp' is not supported for disjunctive (segment) breakpoints" + ) + _add_disjunctive(model, name, inputs, active_expr) + resolved_method: PWL_METHOD = "sos2" + else: + resolved_method = _add_continuous(model, name, inputs, method, active_expr) + + new_vars = [n for n in model.variables if n not in vars_before] + new_cons = [n for n in model.constraints if n not in cons_before] + + if method == "auto": + logger.info( + "piecewise formulation '%s': auto selected method='%s' " + "(sign='%s', %d pair%s)", + name, + resolved_method, + sign, + inputs.n_tuples, + "" if inputs.n_tuples == 1 else "s", + ) + + convexity: PWL_CONVEXITY | None = None + if inputs.n_tuples == 2 and not disjunctive: + if inputs.is_equality: + x_pts = inputs.pinned_bps[1] + y_pts: DataArray = inputs.pinned_bps[0] + else: + assert inputs.bounded_bp is not None + x_pts = inputs.pinned_bps[0] + y_pts = inputs.bounded_bp + if _check_strict_monotonicity(x_pts): + convexity = _detect_convexity(x_pts, y_pts) + + result = PiecewiseFormulation( + name=name, + method=resolved_method, + variable_names=new_vars, + constraint_names=new_cons, + model=model, + convexity=convexity, + ) + model._piecewise_formulations[name] = result + return result + + +def _stack_along_link( + items: Sequence[DataArray | xr.Dataset], + link_coords: list[str], + link_dim: str, +) -> DataArray: + """Expand and concatenate DataArrays/Datasets along a new link dimension.""" + expanded = [ + item.expand_dims({link_dim: [c]}) for item, c in zip(items, link_coords) + ] + return xr.concat(expanded, dim=link_dim, coords="minimal") # type: ignore + + +@dataclass +class _PwlInputs: + """ + Categorised piecewise inputs (post-coercion, post-broadcast). + + ``pinned_*`` are the equality tuples in the user's original order. + ``bounded_*`` is the single non-equality tuple, or ``None``. + ``bounded_sign`` is ``EQUAL`` iff ``bounded_expr is None``. + """ + + pinned_exprs: list[LinearExpression] + pinned_bps: list[DataArray] + pinned_coords: list[str] + bounded_expr: LinearExpression | None + bounded_bp: DataArray | None + bounded_coord: str | None + bounded_sign: str + bp_mask: DataArray | None + link_dim: str = PWL_LINK_DIM + + @property + def is_equality(self) -> bool: + return self.bounded_expr is None + + @property + def n_tuples(self) -> int: + return len(self.pinned_exprs) + (0 if self.is_equality else 1) + + def all_bps(self) -> list[DataArray]: + if self.bounded_bp is None: + return list(self.pinned_bps) + return [self.bounded_bp, *self.pinned_bps] + + def all_coords(self) -> list[str]: + if self.bounded_coord is None: + return list(self.pinned_coords) + return [self.bounded_coord, *self.pinned_coords] + + def all_exprs(self) -> list[LinearExpression]: + if self.bounded_expr is None: + return list(self.pinned_exprs) + return [self.bounded_expr, *self.pinned_exprs] + + +def _lp_eligibility( + inputs: _PwlInputs, + active: LinearExpression | None, +) -> tuple[bool, str]: + """ + Check whether LP tangent-lines dispatch is applicable. + + Returns ``(True, "")`` if LP is applicable, else ``(False, reason)``. + """ + if inputs.n_tuples != 2: + return False, f"{inputs.n_tuples} expressions (LP supports only 2)" + if inputs.is_equality: + return False, "all tuples are equality (LP needs one bounded tuple)" + if active is not None: + return False, "active=... is not supported by LP" + assert inputs.bounded_bp is not None # narrowed by is_equality check + x_pts = inputs.pinned_bps[0] + y_pts = inputs.bounded_bp + paired_x = _paired_valid_points(x_pts, y_pts) + if not _check_strict_monotonicity(paired_x): + return False, "paired x breakpoints are not strictly monotonic" + if not _has_trailing_nan_only(paired_x): + return False, "paired breakpoints contain non-trailing NaN" + convexity = _detect_convexity(x_pts, y_pts) + sign = inputs.bounded_sign + if sign == LESS_EQUAL and convexity not in ("concave", "linear"): + return False, f"sign='<=' needs concave/linear curvature, got '{convexity}'" + if sign == GREATER_EQUAL and convexity not in ("convex", "linear"): + return False, f"sign='>=' needs convex/linear curvature, got '{convexity}'" + return True, "" + + +@dataclass +class _PwlLinks: + """ + Stacked link expressions consumed by SOS2/incremental/disjunctive builders. + """ + + stacked_bp: DataArray + link_dim: str + bp_mask: DataArray | None + sign: str + eq_expr: LinearExpression | None + eq_bp: DataArray | None + signed_expr: LinearExpression | None + signed_bp: DataArray | None + + +def _build_links(model: Model, inputs: _PwlInputs) -> _PwlLinks: + """Stack ``inputs`` into the link representation.""" + from linopy.expressions import LinearExpression + + stacked_bp = _stack_along_link( + inputs.all_bps(), inputs.all_coords(), inputs.link_dim + ) + + if inputs.is_equality: + eq_data = _stack_along_link( + [e.data for e in inputs.pinned_exprs], + inputs.pinned_coords, + inputs.link_dim, + ) + return _PwlLinks( + stacked_bp=stacked_bp, + link_dim=inputs.link_dim, + bp_mask=inputs.bp_mask, + sign=EQUAL, + eq_expr=LinearExpression(eq_data, model), + eq_bp=stacked_bp, + signed_expr=None, + signed_bp=None, + ) + + if inputs.pinned_exprs: + eq_data = _stack_along_link( + [e.data for e in inputs.pinned_exprs], + inputs.pinned_coords, + inputs.link_dim, + ) + eq_expr: LinearExpression | None = LinearExpression(eq_data, model) + eq_bp: DataArray | None = _stack_along_link( + inputs.pinned_bps, inputs.pinned_coords, inputs.link_dim + ) + else: + eq_expr = None + eq_bp = None + + return _PwlLinks( + stacked_bp=stacked_bp, + link_dim=inputs.link_dim, + bp_mask=inputs.bp_mask, + sign=inputs.bounded_sign, + eq_expr=eq_expr, + eq_bp=eq_bp, + signed_expr=inputs.bounded_expr, + signed_bp=inputs.bounded_bp, + ) + + +def _try_lp( + model: Model, + name: str, + inputs: _PwlInputs, + method: str, + active: LinearExpression | None, +) -> bool: + """Dispatch the LP formulation if requested or eligible.""" + if method not in ("lp", "auto"): + return False + if method == "auto" and inputs.is_equality: + return False + + ok, reason = _lp_eligibility(inputs, active) + if not ok: + if method == "lp": + raise ValueError( + f"method='lp' is not applicable: {reason}. Use method='auto'." + ) + logger.info( + "piecewise formulation '%s': LP not applicable (%s); " + "will use SOS2/incremental instead", + name, + reason, + ) + return False + + assert inputs.bounded_expr is not None + assert inputs.bounded_bp is not None + _add_lp( + model, + name, + inputs.pinned_exprs[0], + inputs.bounded_expr, + inputs.pinned_bps[0], + inputs.bounded_bp, + inputs.bounded_sign, + ) + return True + + +def _resolve_sos2_vs_incremental( + method: str, stacked_bp: DataArray +) -> Literal["incremental", "sos2"]: + """ + Validate and (for ``method="auto"``) pick between SOS2 and + incremental based on monotonicity and NaN layout. + """ + trailing_nan_only = _has_trailing_nan_only(stacked_bp) + is_monotonic = _check_strict_monotonicity(stacked_bp) + + if method == "auto": + if not trailing_nan_only: + raise ValueError( + "SOS2 method does not support non-trailing NaN breakpoints." + ) + return "incremental" if is_monotonic else "sos2" + + if method == "incremental": + if not is_monotonic: + raise ValueError( + "Incremental method requires strictly monotonic breakpoints." + ) + if not trailing_nan_only: + raise ValueError( + "Incremental method does not support non-trailing NaN breakpoints." + ) + return "incremental" + + assert method == "sos2" + _validate_numeric_breakpoint_coords(stacked_bp) + if not trailing_nan_only: + raise ValueError("SOS2 method does not support non-trailing NaN breakpoints.") + return "sos2" + + +def _add_continuous( + model: Model, + name: str, + inputs: _PwlInputs, + method: str, + active: LinearExpression | None = None, +) -> PWL_METHOD: + """Returns the resolved method name (``"lp"``, ``"sos2"``, ``"incremental"``).""" + if _try_lp(model, name, inputs, method, active): + return "lp" + + links = _build_links(model, inputs) + resolved = _resolve_sos2_vs_incremental(method, links.stacked_bp) + + if resolved == "sos2": + rhs = active if active is not None else 1 + _add_sos2(model, name, links, rhs) + else: + _add_incremental(model, name, links, active) + return resolved + + +def _add_sos2( + model: Model, + name: str, + links: _PwlLinks, + rhs: LinearExpression | int, +) -> None: + """ + SOS2 formulation. ``links.eq_expr`` is the equality side; + ``links.signed_expr`` (if any) is the output-side link. + """ + dim = BREAKPOINT_DIM + stacked_bp = links.stacked_bp + extra = _var_coords_from(stacked_bp, exclude={dim, links.link_dim}) + lambda_coords = extra + [pd.Index(stacked_bp.coords[dim].values, name=dim)] + + lambda_var = model.add_variables( + lower=0, + upper=1, + coords=lambda_coords, + name=f"{name}{PWL_LAMBDA_SUFFIX}", + mask=links.bp_mask, + ) + model.add_sos_constraints(lambda_var, sos_type=2, sos_dim=dim) + model.add_constraints( + lambda_var.sum(dim=dim) == rhs, name=f"{name}{PWL_CONVEX_SUFFIX}" + ) + + if links.eq_expr is not None and links.eq_bp is not None: + input_weighted = (lambda_var * links.eq_bp).sum(dim=dim) + model.add_constraints( + links.eq_expr == input_weighted, name=f"{name}{PWL_LINK_SUFFIX}" + ) + + if links.signed_expr is not None and links.signed_bp is not None: + output_weighted = (lambda_var * links.signed_bp).sum(dim=dim) + _add_signed_link( + model, + links.signed_expr, + output_weighted, + links.sign, + f"{name}{PWL_OUTPUT_LINK_SUFFIX}", + ) + + +def _add_incremental( + model: Model, + name: str, + links: _PwlLinks, + active: LinearExpression | None, +) -> None: + """ + Incremental formulation. ``links.eq_expr`` is the equality side; + ``links.signed_expr`` (if any) is the output-side link. + """ + dim = BREAKPOINT_DIM + stacked_bp = links.stacked_bp + extra = _var_coords_from(stacked_bp, exclude={dim, links.link_dim}) + + n_pieces = stacked_bp.sizes[dim] - 1 + piece_dim = LP_PIECE_DIM + piece_index = pd.Index(range(n_pieces), name=piece_dim) + delta_coords = extra + [piece_index] + + if links.bp_mask is not None: + mask_lo = links.bp_mask.isel({dim: slice(None, -1)}).rename({dim: piece_dim}) + mask_hi = links.bp_mask.isel({dim: slice(1, None)}).rename({dim: piece_dim}) + mask_lo[piece_dim] = piece_index + mask_hi[piece_dim] = piece_index + delta_mask: DataArray | None = mask_lo & mask_hi + else: + delta_mask = None + + delta_var = model.add_variables( + lower=0, + upper=1, + coords=delta_coords, + name=f"{name}{PWL_DELTA_SUFFIX}", + mask=delta_mask, + ) + + if active is not None: + model.add_constraints( + delta_var <= active, name=f"{name}{PWL_ACTIVE_BOUND_SUFFIX}" + ) + + binary_var = model.add_variables( + binary=True, + coords=delta_coords, + name=f"{name}{PWL_ORDER_BINARY_SUFFIX}", + mask=delta_mask, + ) + model.add_constraints( + delta_var <= binary_var, name=f"{name}{PWL_DELTA_BOUND_SUFFIX}" + ) + + if n_pieces >= 2: + delta_lo = delta_var.isel({piece_dim: slice(None, -1)}, drop=True) + delta_hi = delta_var.isel({piece_dim: slice(1, None)}, drop=True) + model.add_constraints( + delta_hi <= delta_lo, name=f"{name}{PWL_FILL_ORDER_SUFFIX}" + ) + binary_hi = binary_var.isel({piece_dim: slice(1, None)}, drop=True) + model.add_constraints( + binary_hi <= delta_lo, name=f"{name}{PWL_BINARY_ORDER_SUFFIX}" + ) + + def _incremental_weighted(bp: DataArray) -> LinearExpression: + steps = bp.diff(dim).rename({dim: piece_dim}) + steps[piece_dim] = piece_index + bp0 = bp.isel({dim: 0}) + bp0_term: DataArray | LinearExpression = bp0 + if active is not None: + bp0_term = bp0 * active + return (delta_var * steps).sum(dim=piece_dim) + bp0_term + + if links.eq_expr is not None and links.eq_bp is not None: + model.add_constraints( + links.eq_expr == _incremental_weighted(links.eq_bp), + name=f"{name}{PWL_LINK_SUFFIX}", + ) + + if links.signed_expr is not None and links.signed_bp is not None: + _add_signed_link( + model, + links.signed_expr, + _incremental_weighted(links.signed_bp), + links.sign, + f"{name}{PWL_OUTPUT_LINK_SUFFIX}", + ) + + +def _add_disjunctive( + model: Model, + name: str, + inputs: _PwlInputs, + active: LinearExpression | None = None, +) -> None: + """Disjunctive SOS2 formulation.""" + link_dim = inputs.link_dim + links = _build_links(model, inputs) + stacked_bp = links.stacked_bp + bp_mask = inputs.bp_mask + + _validate_numeric_breakpoint_coords(stacked_bp) + if not _has_trailing_nan_only(stacked_bp): + raise ValueError( + "Disjunctive SOS2 does not support non-trailing NaN breakpoints. " + "NaN values must only appear at the end of the breakpoint sequence." + ) + + dim = BREAKPOINT_DIM + extra = _var_coords_from(stacked_bp, exclude={dim, SEGMENT_DIM, link_dim}) + lambda_coords = extra + [ + pd.Index(stacked_bp.coords[SEGMENT_DIM].values, name=SEGMENT_DIM), + pd.Index(stacked_bp.coords[dim].values, name=dim), + ] + binary_coords = extra + [ + pd.Index(stacked_bp.coords[SEGMENT_DIM].values, name=SEGMENT_DIM), + ] + binary_mask = bp_mask.any(dim=dim) if bp_mask is not None else None + + binary_var = model.add_variables( + binary=True, + coords=binary_coords, + name=f"{name}{PWL_SEGMENT_BINARY_SUFFIX}", + mask=binary_mask, + ) + rhs = active if active is not None else 1 + model.add_constraints( + binary_var.sum(dim=SEGMENT_DIM) == rhs, + name=f"{name}{PWL_SELECT_SUFFIX}", + ) + + lambda_var = model.add_variables( + lower=0, + upper=1, + coords=lambda_coords, + name=f"{name}{PWL_LAMBDA_SUFFIX}", + mask=bp_mask, + ) + model.add_sos_constraints(lambda_var, sos_type=2, sos_dim=dim) + model.add_constraints( + lambda_var.sum(dim=dim) == binary_var, + name=f"{name}{PWL_CONVEX_SUFFIX}", + ) + + if links.eq_expr is not None and links.eq_bp is not None: + input_weighted = (lambda_var * links.eq_bp).sum(dim=[SEGMENT_DIM, dim]) + model.add_constraints( + links.eq_expr == input_weighted, name=f"{name}{PWL_LINK_SUFFIX}" + ) + + if links.signed_expr is not None and links.signed_bp is not None: + output_weighted = (lambda_var * links.signed_bp).sum(dim=[SEGMENT_DIM, dim]) + _add_signed_link( + model, + links.signed_expr, + output_weighted, + links.sign, + f"{name}{PWL_OUTPUT_LINK_SUFFIX}", + ) + + +def _add_signed_link( + model: Model, + lhs: LinearExpression, + rhs: LinearExpression, + sign: str, + name: str, + mask: DataArray | None = None, +) -> Constraint: + """Add a link constraint with the requested sign.""" + if sign == EQUAL: + return model.add_constraints(lhs == rhs, name=name, mask=mask) + elif sign == LESS_EQUAL: + return model.add_constraints(lhs <= rhs, name=name, mask=mask) + else: # ">=" + return model.add_constraints(lhs >= rhs, name=name, mask=mask) + + +def _add_lp( + model: Model, + name: str, + x_expr: LinearExpression, + y_expr: LinearExpression, + x_points: DataArray, + y_points: DataArray, + sign: str, +) -> None: + """ + LP tangent-line formulation (no auxiliary variables). + + Adds one chord constraint per piece plus domain bounds on x. + Trailing-NaN pieces (per-entity short curves) are masked out so + they do not contribute spurious ``y ≤ 0`` constraints. + """ + # Per-piece validity: both endpoints must be non-NaN. + bp_valid = ~(x_points.isnull() | y_points.isnull()) + piece_count = x_points.sizes[BREAKPOINT_DIM] - 1 + piece_index = np.arange(piece_count) + full_mask = _rename_to_pieces( + bp_valid.isel({BREAKPOINT_DIM: slice(None, -1)}) + & bp_valid.isel({BREAKPOINT_DIM: slice(1, None)}).values, + piece_index, + ) + piece_mask: DataArray | None = None if bool(full_mask.all()) else full_mask + + # Use the internal impl so we don't fire a second EvolvingAPIWarning — + # ``add_piecewise_formulation`` already warned on entry. + tangents = _tangent_lines_impl(x_expr, x_points, y_points) + _add_signed_link( + model, + y_expr, + tangents, + sign, + f"{name}{PWL_CHORD_SUFFIX}", + mask=piece_mask, + ) + + # Domain bounds: x ∈ [x_min, x_max] over paired-valid breakpoints. + paired_x_points = x_points.where(bp_valid) + x_min = paired_x_points.min(dim=BREAKPOINT_DIM) + x_max = paired_x_points.max(dim=BREAKPOINT_DIM) + model.add_constraints(x_expr >= x_min, name=f"{name}{PWL_DOMAIN_LO_SUFFIX}") + model.add_constraints(x_expr <= x_max, name=f"{name}{PWL_DOMAIN_HI_SUFFIX}") diff --git a/linopy/remote/__init__.py b/linopy/remote/__init__.py index 0ae1df267..d3d5e1620 100644 --- a/linopy/remote/__init__.py +++ b/linopy/remote/__init__.py @@ -8,9 +8,13 @@ - OetcHandler: Cloud-based execution via OET Cloud service """ -from linopy.remote.oetc import OetcCredentials, OetcHandler, OetcSettings from linopy.remote.ssh import RemoteHandler +try: + from linopy.remote.oetc import OetcCredentials, OetcHandler, OetcSettings +except ImportError: + pass + __all__ = [ "RemoteHandler", "OetcHandler", diff --git a/linopy/remote/oetc.py b/linopy/remote/oetc.py index 5bea9c7c9..74f8a9a56 100644 --- a/linopy/remote/oetc.py +++ b/linopy/remote/oetc.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import base64 import gzip import json @@ -7,19 +9,32 @@ import time from dataclasses import dataclass, field from datetime import datetime, timedelta -from enum import Enum +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + from linopy.model import Model -import requests -from google.cloud import storage -from google.oauth2 import service_account -from requests import RequestException +try: + import requests + from google.cloud import storage + from google.oauth2 import service_account + from requests import RequestException + + _oetc_deps_available = True +except ImportError: + _oetc_deps_available = False import linopy +from linopy.sos_reformulation import ( + sos_reformulation_context, + suppress_serialization_warning, +) logger = logging.getLogger(__name__) -class ComputeProvider(str, Enum): +class ComputeProvider(StrEnum): GCP = "GCP" @@ -37,11 +52,97 @@ class OetcSettings: orchestrator_server_url: str compute_provider: ComputeProvider = ComputeProvider.GCP solver: str = "highs" - solver_options: dict = field(default_factory=dict) + solver_options: dict[str, Any] = field(default_factory=dict) cpu_cores: int = 2 disk_space_gb: int = 10 delete_worker_on_error: bool = False + @classmethod + def from_env( + cls, + *, + email: str | None = None, + password: str | None = None, + name: str | None = None, + authentication_server_url: str | None = None, + orchestrator_server_url: str | None = None, + cpu_cores: int | None = None, + disk_space_gb: int | None = None, + delete_worker_on_error: bool | None = None, + ) -> OetcSettings: + required_fields = { + "email": ("OETC_EMAIL", email), + "password": ("OETC_PASSWORD", password), + "name": ("OETC_NAME", name), + "authentication_server_url": ("OETC_AUTH_URL", authentication_server_url), + "orchestrator_server_url": ( + "OETC_ORCHESTRATOR_URL", + orchestrator_server_url, + ), + } + + resolved: dict[str, Any] = {} + missing: list[str] = [] + + for field_name, (env_var, kwarg) in required_fields.items(): + if kwarg is not None: + resolved[field_name] = kwarg + else: + env_val = os.environ.get(env_var, "").strip() + if env_val: + resolved[field_name] = env_val + else: + missing.append(env_var) + + if missing: + raise ValueError( + f"Missing required OETC configuration: {', '.join(missing)}" + ) + + kwargs: dict[str, Any] = { + "credentials": OetcCredentials( + email=resolved["email"], password=resolved["password"] + ), + "name": resolved["name"], + "authentication_server_url": resolved["authentication_server_url"], + "orchestrator_server_url": resolved["orchestrator_server_url"], + } + + if cpu_cores is not None: + kwargs["cpu_cores"] = cpu_cores + elif (cpu_env := os.environ.get("OETC_CPU_CORES")) is not None: + try: + kwargs["cpu_cores"] = int(cpu_env) + except ValueError as e: + raise ValueError( + f"OETC_CPU_CORES is not a valid integer: {cpu_env}" + ) from e + + if disk_space_gb is not None: + kwargs["disk_space_gb"] = disk_space_gb + elif (disk_env := os.environ.get("OETC_DISK_SPACE_GB")) is not None: + try: + kwargs["disk_space_gb"] = int(disk_env) + except ValueError as e: + raise ValueError( + f"OETC_DISK_SPACE_GB is not a valid integer: {disk_env}" + ) from e + + if delete_worker_on_error is not None: + kwargs["delete_worker_on_error"] = delete_worker_on_error + elif (del_env := os.environ.get("OETC_DELETE_WORKER_ON_ERROR")) is not None: + low = del_env.lower() + if low in ("true", "1", "yes"): + kwargs["delete_worker_on_error"] = True + elif low in ("false", "0", "no"): + kwargs["delete_worker_on_error"] = False + else: + raise ValueError( + f"OETC_DELETE_WORKER_ON_ERROR has invalid value: {del_env}" + ) + + return cls(**kwargs) + @dataclass class GcpCredentials: @@ -85,6 +186,11 @@ class JobResult: class OetcHandler: def __init__(self, settings: OetcSettings) -> None: + if not _oetc_deps_available: + raise ImportError( + "The 'google-cloud-storage' and 'requests' packages are required " + "for OetcHandler. Install them with: pip install linopy[oetc]" + ) self.settings = settings self.jwt = self.__sign_in() self.cloud_provider_credentials = self.__get_cloud_provider_credentials() @@ -216,12 +322,16 @@ def __get_gcp_credentials(self) -> GcpCredentials: except Exception as e: raise Exception(f"Error fetching GCP credentials: {e}") - def _submit_job_to_compute_service(self, input_file_name: str) -> str: + def _submit_job_to_compute_service( + self, input_file_name: str, solver: str, solver_options: dict[str, Any] + ) -> str: """ Submit a job to the compute service. Args: input_file_name: Name of the input file uploaded to GCP + solver: Solver name to use + solver_options: Solver options dict Returns: CreateComputeJobResult: The job creation result with UUID @@ -233,8 +343,8 @@ def _submit_job_to_compute_service(self, input_file_name: str) -> str: logger.info("OETC - Submitting compute job...") payload = { "name": self.settings.name, - "solver": self.settings.solver, - "solver_options": self.settings.solver_options, + "solver": solver, + "solver_options": solver_options, "provider": self.settings.compute_provider.value, "cpu_cores": self.settings.cpu_cores, "disk_space_gb": self.settings.disk_space_gb, @@ -524,13 +634,24 @@ def _download_file_from_gcp(self, file_name: str) -> str: except Exception as e: raise Exception(f"Failed to download file from GCP: {e}") - def solve_on_oetc(self, model): # type: ignore + def solve_on_oetc( + self, + model: Model, + solver_name: str | None = None, + *, + reformulate_sos: bool | Literal["auto"] = False, + **solver_options: Any, + ) -> Model: """ Solve a linopy model on the OET Cloud compute app. Parameters ---------- model : linopy.model.Model + solver_name : str, optional + Override the solver from settings. + **solver_options + Override/extend solver_options from settings. Returns ------- @@ -542,17 +663,23 @@ def solve_on_oetc(self, model): # type: ignore Exception: If solving fails at any stage """ try: - # Save model to temporary file and upload - with tempfile.NamedTemporaryFile(prefix="linopy-", suffix=".nc") as fn: - fn.file.close() - model.to_netcdf(fn.name) - input_file_name = self._upload_file_to_gcp(fn.name) - - # Submit job and wait for completion - job_uuid = self._submit_job_to_compute_service(input_file_name) + effective_solver = solver_name or self.settings.solver + merged_solver_options = {**self.settings.solver_options, **solver_options} + + with sos_reformulation_context( + model, effective_solver, reformulate_sos + ) as applied: + with tempfile.NamedTemporaryFile(prefix="linopy-", suffix=".nc") as fn: + fn.file.close() + with suppress_serialization_warning(active=applied): + model.to_netcdf(fn.name) + input_file_name = self._upload_file_to_gcp(fn.name) + + job_uuid = self._submit_job_to_compute_service( + input_file_name, effective_solver, merged_solver_options + ) job_result = self.wait_and_get_job_data(job_uuid) - # Download and load the solution if not job_result.output_files: raise Exception("No output files found in completed job") @@ -562,18 +689,14 @@ def solve_on_oetc(self, model): # type: ignore solution_file_path = self._download_file_from_gcp(output_file_name) - # Load the solved model solved_model = linopy.read_netcdf(solution_file_path) - # Clean up downloaded file os.remove(solution_file_path) logger.info( f"OETC - Model solved successfully. Status: {solved_model.status}" ) - if hasattr(solved_model, "objective") and hasattr( - solved_model.objective, "value" - ): + if solved_model.objective.value is not None: logger.info( f"OETC - Objective value: {solved_model.objective.value:.2e}" ) @@ -581,7 +704,7 @@ def solve_on_oetc(self, model): # type: ignore return solved_model except Exception as e: - raise Exception(f"Error solving model on OETC: {e}") + raise Exception(f"Error solving model on OETC: {e}") from e def _gzip_compress(self, source_path: str) -> str: """ diff --git a/linopy/remote/ssh.py b/linopy/remote/ssh.py index 426ed6463..7c0a06447 100644 --- a/linopy/remote/ssh.py +++ b/linopy/remote/ssh.py @@ -6,12 +6,17 @@ """ import logging +import os import tempfile from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, Literal, Union from linopy.io import read_netcdf +from linopy.sos_reformulation import ( + sos_reformulation_context, + suppress_serialization_warning, +) if TYPE_CHECKING: from linopy.model import Model @@ -169,9 +174,10 @@ def write_model_on_remote(self, model: "Model") -> None: Write a model on the remote machine under `self.model_unsolved_file`. """ logger.info(f"Saving unsolved model at {self.model_unsolved_file} on remote") - with tempfile.NamedTemporaryFile(prefix="linopy", suffix=".nc") as fn: - model.to_netcdf(fn.name) - self.sftp_client.put(fn.name, self.model_unsolved_file) + with tempfile.TemporaryDirectory(prefix="linopy") as tmpdir: + local_path = os.path.join(tmpdir, "model.nc") + model.to_netcdf(local_path) + self.sftp_client.put(local_path, self.model_unsolved_file) def execute(self, cmd: str) -> None: """ @@ -200,43 +206,53 @@ def execute(self, cmd: str) -> None: if exit_status: raise OSError("Execution on remote raised an error, see above.") - def solve_on_remote(self, model: "Model", **kwargs: Any) -> "Model": + def solve_on_remote( + self, + model: "Model", + *, + reformulate_sos: bool | Literal["auto"] = False, + **kwargs: Any, + ) -> "Model": """ Solve a linopy model on the remote machine. - This function - - 1. saves the model to a file on the local machine. - 2. copies that file to the remote machine. - 3. loads, solves and writes out the model, all on the remote machine. - 4. copies the solved model to the local machine. - 5. loads and returns the solved model. + Reformulates SOS constraints locally before serialization when + requested, so the worker just solves a plain MILP and the SOS + lifecycle stays on the caller's model. Parameters ---------- model : linopy.model.Model + reformulate_sos : bool | "auto", optional + Forwarded to ``Model._resolve_sos_reformulation`` to decide + whether to apply SOS reformulation locally before transfer. **kwargs : - Keyword arguments passed to `linopy.model.Model.solve`. + Keyword arguments passed to `linopy.model.Model.solve` on the + remote worker. Returns ------- linopy.model.Model Solved model. """ - self.write_python_file_on_remote(**kwargs) - self.write_model_on_remote(model) + solver_name = kwargs.get("solver_name") + with sos_reformulation_context(model, solver_name, reformulate_sos) as applied: + self.write_python_file_on_remote(**kwargs) + with suppress_serialization_warning(active=applied): + self.write_model_on_remote(model) - command = f"{self.python_executable} {self.python_file}" + command = f"{self.python_executable} {self.python_file}" - logger.info("Solving model on remote.") - self.execute(command) + logger.info("Solving model on remote.") + self.execute(command) - logger.info("Retrieve solved model from remote.") - with tempfile.NamedTemporaryFile(prefix="linopy", suffix=".nc") as fn: - self.sftp_client.get(self.model_solved_file, fn.name) - solved = read_netcdf(fn.name) + logger.info("Retrieve solved model from remote.") + with tempfile.TemporaryDirectory(prefix="linopy") as tmpdir: + local_path = os.path.join(tmpdir, "model.nc") + self.sftp_client.get(self.model_solved_file, local_path) + solved = read_netcdf(local_path) - self.sftp_client.remove(self.python_file) - self.sftp_client.remove(self.model_solved_file) + self.sftp_client.remove(self.python_file) + self.sftp_client.remove(self.model_solved_file) - return solved + return solved diff --git a/linopy/solver_capabilities.py b/linopy/solver_capabilities.py index 0ffea9232..0e7480825 100644 --- a/linopy/solver_capabilities.py +++ b/linopy/solver_capabilities.py @@ -1,297 +1,100 @@ """ -Linopy module for solver capability tracking. +Back-compat shim for legacy solver-capability imports. -This module provides a centralized registry of solver capabilities, -replacing scattered hardcoded checks throughout the codebase. +Capability data is declared on each `Solver` subclass in `linopy.solvers`. +Prefer `Solver.features` / `Solver.supports()` over the helpers in this module. """ from __future__ import annotations -import platform +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass -from enum import Enum, auto -from importlib.metadata import PackageNotFoundError -from importlib.metadata import version as package_version +from enum import Enum from typing import TYPE_CHECKING -from packaging.specifiers import SpecifierSet - if TYPE_CHECKING: - from collections.abc import Sequence - - -def _xpress_supports_gpu() -> bool: - """Check if installed xpress version supports GPU acceleration (>=9.8.0).""" - try: - return package_version("xpress") in SpecifierSet(">=9.8.0") - except PackageNotFoundError: - return False - - -class SolverFeature(Enum): - """Enumeration of all solver capabilities tracked by linopy.""" + from linopy.solvers import Solver, SolverFeature - # Model feature support - INTEGER_VARIABLES = auto() # Support for integer variables +__all__ = ( + "SOLVER_REGISTRY", + "SolverFeature", + "SolverInfo", + "get_available_solvers_with_feature", + "get_solvers_with_feature", + "solver_supports", +) - # Objective function support - QUADRATIC_OBJECTIVE = auto() - # I/O capabilities - DIRECT_API = auto() # Solve directly from Model without writing files - LP_FILE_NAMES = auto() # Support for named variables/constraints in LP files - READ_MODEL_FROM_FILE = auto() # Ability to read models from file - SOLUTION_FILE_NOT_NEEDED = auto() # Solver doesn't need a solution file +def __getattr__(name: str) -> object: + if name == "SolverFeature": + from linopy import solvers as _solvers_mod - # Advanced features - GPU_ACCELERATION = auto() # GPU-accelerated solving - IIS_COMPUTATION = auto() # Irreducible Infeasible Set computation - - # Special constraint types - SOS_CONSTRAINTS = auto() # Special Ordered Sets (SOS1/SOS2) constraints - - # Solver-specific - SOLVER_ATTRIBUTE_ACCESS = auto() # Direct access to solver variable attributes + return _solvers_mod.SolverFeature + raise AttributeError(name) @dataclass(frozen=True) class SolverInfo: - """Information about a solver's capabilities.""" + """Legacy view of a solver's capabilities. Prefer Solver.features / Solver.supports().""" name: str - features: frozenset[SolverFeature] + features: frozenset[Enum] display_name: str = "" def __post_init__(self) -> None: if not self.display_name: object.__setattr__(self, "display_name", self.name.upper()) - def supports(self, feature: SolverFeature) -> bool: - """Check if this solver supports a given feature.""" + def supports(self, feature: Enum) -> bool: return feature in self.features -# Define all solver capabilities -SOLVER_REGISTRY: dict[str, SolverInfo] = { - "gurobi": SolverInfo( - name="gurobi", - display_name="Gurobi", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.DIRECT_API, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - SolverFeature.IIS_COMPUTATION, - SolverFeature.SOS_CONSTRAINTS, - SolverFeature.SOLVER_ATTRIBUTE_ACCESS, - } - ), - ), - "highs": SolverInfo( - name="highs", - display_name="HiGHS", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.DIRECT_API, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - ), - ), - "glpk": SolverInfo( - name="glpk", - display_name="GLPK", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.READ_MODEL_FROM_FILE, - } - ), # No LP_FILE_NAMES support - ), - "cbc": SolverInfo( - name="cbc", - display_name="CBC", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.READ_MODEL_FROM_FILE, - } - ), # No LP_FILE_NAMES support - ), - "cplex": SolverInfo( - name="cplex", - display_name="CPLEX", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOS_CONSTRAINTS, - } - ), - ), - "xpress": SolverInfo( - name="xpress", - display_name="FICO Xpress", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - SolverFeature.GPU_ACCELERATION, - SolverFeature.IIS_COMPUTATION, - } - if _xpress_supports_gpu() - else { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - SolverFeature.IIS_COMPUTATION, - } - ), - ), - "scip": SolverInfo( - name="scip", - display_name="SCIP", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - if platform.system() == "Windows" - else { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - # SCIP has a bug with quadratic models on Windows, see: - # https://github.com/PyPSA/linopy/actions/runs/7615240686/job/20739454099?pr=78 - ), - ), - "mosek": SolverInfo( - name="mosek", - display_name="MOSEK", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.DIRECT_API, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - ), - ), - "copt": SolverInfo( - name="copt", - display_name="COPT", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - ), - ), - "mindopt": SolverInfo( - name="mindopt", - display_name="MindOpt", - features=frozenset( - { - SolverFeature.INTEGER_VARIABLES, - SolverFeature.QUADRATIC_OBJECTIVE, - SolverFeature.LP_FILE_NAMES, - SolverFeature.READ_MODEL_FROM_FILE, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - ), - ), - "cupdlpx": SolverInfo( - name="cupdlpx", - display_name="cuPDLPx", - features=frozenset( - { - SolverFeature.DIRECT_API, - SolverFeature.GPU_ACCELERATION, - SolverFeature.SOLUTION_FILE_NOT_NEEDED, - } - ), - ), -} - +def _solver_class(name: str) -> type[Solver] | None: + from linopy import solvers as _solvers_mod -def solver_supports(solver_name: str, feature: SolverFeature) -> bool: - """ - Check if a solver supports a given feature. + try: + return getattr(_solvers_mod, _solvers_mod.SolverName(name).name, None) + except ValueError: + return None - Parameters - ---------- - solver_name : str - Name of the solver (e.g., "gurobi", "highs") - feature : SolverFeature - The feature to check for - Returns - ------- - bool - True if the solver supports the feature, False otherwise. - Returns False for unknown solvers. - """ - if solver_name not in SOLVER_REGISTRY: - return False - return SOLVER_REGISTRY[solver_name].supports(feature) +def solver_supports(solver_name: str, feature: SolverFeature) -> bool: + cls = _solver_class(solver_name) + return cls is not None and cls.supports(feature) def get_solvers_with_feature(feature: SolverFeature) -> list[str]: - """ - Get all solvers that support a given feature. - - Parameters - ---------- - feature : SolverFeature - The feature to filter by + from linopy.solvers import SolverName - Returns - ------- - list[str] - List of solver names supporting the feature - """ - return [name for name, info in SOLVER_REGISTRY.items() if info.supports(feature)] + return [n.value for n in SolverName if solver_supports(n.value, feature)] def get_available_solvers_with_feature( feature: SolverFeature, available_solvers: Sequence[str] ) -> list[str]: - """ - Get installed solvers that support a given feature. + return [s for s in get_solvers_with_feature(feature) if s in available_solvers] - Parameters - ---------- - feature : SolverFeature - The feature to filter by - available_solvers : Sequence[str] - List of currently available/installed solvers - Returns - ------- - list[str] - List of installed solver names supporting the feature - """ - return [s for s in get_solvers_with_feature(feature) if s in available_solvers] +class _LazyRegistry(Mapping[str, SolverInfo]): + def __getitem__(self, key: str) -> SolverInfo: + cls = _solver_class(key) + if cls is None: + raise KeyError(key) + return SolverInfo( + name=key, + features=cls.supported_features(), + display_name=cls.display_name, + ) + + def __iter__(self) -> Iterator[str]: + from linopy.solvers import SolverName + + return (n.value for n in SolverName) + + def __len__(self) -> int: + from linopy.solvers import SolverName + + return len(SolverName) + + +SOLVER_REGISTRY: Mapping[str, SolverInfo] = _LazyRegistry() diff --git a/linopy/solvers.py b/linopy/solvers.py index fe516b47e..987c7b70c 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -7,57 +7,184 @@ import contextlib import enum +import functools import io import logging import os import re +import shutil import subprocess as sub import sys import threading import warnings -from abc import ABC, abstractmethod +from abc import ABC from collections import namedtuple -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Iterable, Iterator, Sequence +from dataclasses import dataclass, field +from enum import Enum, auto +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar import numpy as np import pandas as pd +from packaging.specifiers import SpecifierSet from packaging.version import parse as parse_version +from scipy.sparse import tril, triu import linopy.io +from linopy.common import count_initial_letters, values_to_lookup_array from linopy.constants import ( + EQUAL, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, Result, Solution, + SolverReport, SolverStatus, Status, TerminationCondition, + short_GREATER_EQUAL, + short_LESS_EQUAL, ) -from linopy.solver_capabilities import ( - SolverFeature, - get_solvers_with_feature, +from linopy.persistent import ( + ModelDiff, + ModelSnapshot, + RebuildReason, + RebuildRequiredError, + UnsupportedUpdate, + UpdatesDisabledError, + VarKind, + clear_coef_dirty, ) + +def _int_list(arr: np.ndarray, dtype: type = np.int64) -> list[int]: + return arr.astype(dtype, copy=False).tolist() + + +def _float_list(arr: np.ndarray) -> list[float]: + return arr.astype(float, copy=False).tolist() + + +def _parse_int_label(name: str) -> int: + """Strip leading non-digits and parse the integer label.""" + s = str(name) + cutoff = count_initial_letters(s) + try: + return int(s[cutoff:]) + except ValueError: + return int(re.sub(r".*#", "", s)) + + +def _names_to_labels(names: Any) -> np.ndarray: + """Vectorised conversion of solver-provided names to integer labels.""" + if len(names) == 0: + return np.array([], dtype=np.int64) + index = pd.Index(names) + if pd.api.types.is_integer_dtype(index.dtype): + return index.to_numpy(dtype=np.int64) + string_index = index.astype(str) + cutoff = count_initial_letters(str(string_index[0])) + try: + return string_index.str[cutoff:].astype(np.int64).to_numpy(dtype=np.int64) + except (TypeError, ValueError): + try: + return ( + string_index.str.replace(r".*#", "", regex=True) + .astype(np.int64) + .to_numpy(dtype=np.int64) + ) + except (TypeError, ValueError): + return np.fromiter( + (_parse_int_label(n) for n in names), dtype=np.int64, count=len(names) + ) + + +def _solution_from_names(values: np.ndarray, names: Any, size: int) -> np.ndarray: + """ + Build a label-indexed dense solution array of length ``size`` from + solver-side names. Used by paths where the solver may iterate in arbitrary + order or drop unused entities (file-based LP solvers, the ``from_file`` + paths of Highs/Gurobi). + """ + if not size: + return np.array([], dtype=float) + return values_to_lookup_array( + np.asarray(values, dtype=float), _names_to_labels(names), size=size + ) + + +def _solution_from_labels( + values: np.ndarray, labels: np.ndarray | None, size: int +) -> np.ndarray: + """Scatter solver-side values into a label-indexed dense array of length ``size``.""" + if not size: + return np.array([], dtype=float) + assert labels is not None + return values_to_lookup_array(np.asarray(values, dtype=float), labels, size=size) + + +def _iter_sos_sets(model: Model) -> Iterator[tuple[int, np.ndarray, np.ndarray]]: + """Yield ``(sos_type, positions, weights)`` per active SOS set in ``model``.""" + label_to_pos = model.variables.label_index.label_to_pos + for var_name in model.variables.sos: + var = model.variables.sos[var_name] + sos_type = int(var.attrs[SOS_TYPE_ATTR]) # type: ignore[call-overload] + sos_dim = str(var.attrs[SOS_DIM_ATTR]) + + labels = var.labels.transpose(sos_dim, ...) + weights = labels.coords[sos_dim].values + arr = labels.values.reshape(labels.shape[0], -1) + + for i in range(arr.shape[1]): + col = arr[:, i] + mask = col != -1 + if mask.any(): + yield sos_type, label_to_pos[col[mask]], weights[mask] + + +class SolverFeature(Enum): + """Enumeration of all solver capabilities tracked by linopy.""" + + INTEGER_VARIABLES = auto() + QUADRATIC_OBJECTIVE = auto() + DIRECT_API = auto() + LP_FILE_NAMES = auto() + READ_MODEL_FROM_FILE = auto() + SOLUTION_FILE_NOT_NEEDED = auto() + GPU_ACCELERATION = auto() + GPU_ONLY = auto() + IIS_COMPUTATION = auto() + SOS_CONSTRAINTS = auto() + INDICATOR_CONSTRAINTS = auto() + SEMI_CONTINUOUS_VARIABLES = auto() + SOLVER_ATTRIBUTE_ACCESS = auto() + MIP_DUAL_BOUND_REPORT = auto() + + +def _installed_version_in(pkg: str, spec: str) -> bool: + """Check whether the installed version of `pkg` satisfies `spec`.""" + try: + return package_version(pkg) in SpecifierSet(spec) + except PackageNotFoundError: + return False + + if TYPE_CHECKING: + import cupdlpx import gurobipy + import highspy + import mosek from linopy.model import Model EnvType = TypeVar("EnvType") -# Generated from solver_capabilities registry for backward compatibility -QUADRATIC_SOLVERS = get_solvers_with_feature(SolverFeature.QUADRATIC_OBJECTIVE) -NO_SOLUTION_FILE_SOLVERS = get_solvers_with_feature( - SolverFeature.SOLUTION_FILE_NOT_NEEDED -) - FILE_IO_APIS = ["lp", "lp-polars", "mps"] IO_APIS = FILE_IO_APIS + ["direct"] -available_solvers = [] - -which = "where" if os.name == "nt" else "which" - def _run_highs_with_keyboard_interrupt(h: Any) -> None: """ @@ -125,47 +252,12 @@ def _target() -> None: h.HandleUserInterrupt = old_handle_user_interrupt -# the first available solver will be the default solver -with contextlib.suppress(ModuleNotFoundError): - import gurobipy - - available_solvers.append("gurobi") -with contextlib.suppress(ModuleNotFoundError): - _new_highspy_mps_layout = None - import highspy - - available_solvers.append("highs") - from importlib.metadata import version - - if parse_version(version("highspy")) < parse_version("1.7.1"): - # Fallback if parse_version is not available or version string is invalid - _new_highspy_mps_layout = False - else: - _new_highspy_mps_layout = True - -if sub.run([which, "glpsol"], stdout=sub.DEVNULL, stderr=sub.STDOUT).returncode == 0: - available_solvers.append("glpk") - - -if sub.run([which, "cbc"], stdout=sub.DEVNULL, stderr=sub.STDOUT).returncode == 0: - available_solvers.append("cbc") - -with contextlib.suppress(ModuleNotFoundError): - import pyscipopt as scip - - available_solvers.append("scip") - -with contextlib.suppress(ModuleNotFoundError): - import cplex - - available_solvers.append("cplex") - +# xpress.Namespaces was added in xpress 9.6. Importing xpress is pure-Python +# and does not acquire a license, so this shim stays eager so downstream code +# can ``from linopy.solvers import xpress_Namespaces``. with contextlib.suppress(ModuleNotFoundError, ImportError): - import xpress - - available_solvers.append("xpress") + import xpress # noqa: F401 - # xpress.Namespaces was added in xpress 9.6 try: from xpress import Namespaces as xpress_Namespaces except ImportError: @@ -176,43 +268,63 @@ class xpress_Namespaces: # type: ignore[no-redef] SET = 3 -with contextlib.suppress(ModuleNotFoundError): - import mosek +class _LazyModule: + """ + Module proxy that imports the underlying package on first attribute access. + + Lets us keep ``gurobipy.Env`` / ``mindoptpy.read`` references throughout the + file while deferring the actual ``import`` (and its license-server side + effects, for mindoptpy/coptpy) until a Solver subclass really needs them. + """ + + __slots__ = ("_name", "_module") + + def __init__(self, name: str) -> None: + self._name = name + self._module: Any = None - with contextlib.suppress(mosek.Error): - t = mosek.Task() - t.optimize() + def __getattr__(self, attr: str) -> Any: + if attr.startswith("__") and attr.endswith("__"): + raise AttributeError(attr) + if self._module is None: + import importlib - available_solvers.append("mosek") + self._module = importlib.import_module(self._name) + return getattr(self._module, attr) -with contextlib.suppress(ModuleNotFoundError): - import mindoptpy - with contextlib.suppress(mindoptpy.MindoptError): - mindoptpy.Env() +gurobipy = _LazyModule("gurobipy") # type: ignore[assignment] +highspy = _LazyModule("highspy") # type: ignore[assignment] +scip = _LazyModule("pyscipopt") +cplex = _LazyModule("cplex") +knitro = _LazyModule("knitro") +mosek = _LazyModule("mosek") +mindoptpy = _LazyModule("mindoptpy") +coptpy = _LazyModule("coptpy") +cupdlpx = _LazyModule("cupdlpx") - available_solvers.append("mindopt") -with contextlib.suppress(ModuleNotFoundError): - import coptpy +def _has_module(name: str) -> bool: + """True if ``name`` is importable, without executing its ``__init__``.""" + import importlib.util try: - coptpy.Envr() - available_solvers.append("copt") - except coptpy.CoptError: - pass + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False -with contextlib.suppress(ModuleNotFoundError): - import cupdlpx +@functools.cache +def _new_highspy_mps_layout() -> bool: + """True for highspy >= 1.7.1 (new MPS coefficient layout).""" + if not _has_module("highspy"): + return False try: - cupdlpx.Model(np.array([0.0]), np.array([[0.0]]), None, None) - available_solvers.append("cupdlpx") - except ImportError: - pass + return parse_version(package_version("highspy")) >= parse_version("1.7.1") + except PackageNotFoundError: + return False -quadratic_solvers = [s for s in QUADRATIC_SOLVERS if s in available_solvers] logger = logging.getLogger(__name__) @@ -239,6 +351,7 @@ class SolverName(enum.Enum): Gurobi = "gurobi" SCIP = "scip" Xpress = "xpress" + Knitro = "knitro" Mosek = "mosek" COPT = "copt" MindOpt = "mindopt" @@ -280,7 +393,7 @@ def maybe_adjust_objective_sign( return solution if np.isnan(solution.objective): return solution - if io_api == "mps" and not _new_highspy_mps_layout: + if io_api == "mps" and not _new_highspy_mps_layout(): logger.info( "Adjusting objective sign due to switched coefficients in MPS file." ) @@ -288,83 +401,586 @@ def maybe_adjust_objective_sign( return solution +@dataclass(frozen=True) +class LicenseStatus: + """Result of :meth:`Solver.license_status` — license/runtime probe outcome.""" + + solver: str + ok: bool + message: str | None = None + + def __bool__(self) -> bool: + return self.ok + + +@dataclass class Solver(ABC, Generic[EnvType]): """ Abstract base class for solving a given linear problem. - All relevant functions are passed on to the specific solver subclasses. - Subclasses must implement the `solve_problem_from_model()` and - `solve_problem_from_file()` methods. + Subclasses provide ``_build_direct`` / ``_run_direct`` (when supporting the + direct API) and ``_run_file`` (when supporting LP/MPS files). Construction + goes via :meth:`Solver.from_name` or :meth:`Solver.from_model`. + + ``track_updates`` toggles persistent-update support: + + * ``False`` (default) — one-shot mode. No :class:`ModelSnapshot` is + captured at build time; any later ``solve(model=...)`` or + ``update(model)`` raises :class:`UpdatesDisabledError`. Use for + throw-away solver instances and high-level ``Model.solve(...)``. + * ``True`` — long-lived mode. A snapshot is captured at build time and + re-captured after each successful in-place update, enabling + diff-based ``solve(model=...)`` / ``update(model)`` across iterations. """ - def __init__( + model: Model | None = None + io_api: str | None = None + options: dict[str, Any] = field(default_factory=dict) + track_updates: bool = False + + # Runtime state — never set via constructor. + status: Status | None = field(init=False, default=None, repr=False) + solution: Solution | None = field(init=False, default=None, repr=False) + report: SolverReport | None = field(init=False, default=None, repr=False) + solver_model: Any = field(init=False, default=None, repr=False) + sense: str | None = field(init=False, default=None, repr=False) + env: Any = field(init=False, default=None, repr=False) + _env_stack: contextlib.ExitStack | None = field( + init=False, default=None, repr=False + ) + _vlabels: np.ndarray | None = field(init=False, default=None, repr=False) + _clabels: np.ndarray | None = field(init=False, default=None, repr=False) + _n_vars: int = field(init=False, default=0, repr=False) + _n_cons: int = field(init=False, default=0, repr=False) + _problem_fn: Path | None = field(init=False, default=None, repr=False) + + snapshot: ModelSnapshot | None = field(init=False, default=None, repr=False) + _rebuilds: int = field(init=False, default=0, repr=False) + _in_place_updates: int = field(init=False, default=0, repr=False) + _last_rebuild_reason: RebuildReason | None = field( + init=False, default=None, repr=False + ) + + display_name: ClassVar[str] = "" + features: ClassVar[frozenset[SolverFeature]] = frozenset() + accepted_io_apis: ClassVar[frozenset[str]] = frozenset() + supports_persistent_update: ClassVar[bool] = False + supports_sign_update: ClassVar[bool] = False + + def __post_init__(self) -> None: + if type(self) is Solver: + raise TypeError( + "Solver is abstract; instantiate a concrete subclass instead." + ) + if not type(self).is_available(): + msg = ( + f"Solver package for '{self.solver_name.value}' is not installed. " + "Please install first to initialize solver instance." + ) + raise ImportError(msg) + self._lock: threading.Lock = threading.Lock() + + def apply_update( self, - **solver_options: Any, + diff: ModelDiff, + var_label_index: Any, + con_label_index: Any, ) -> None: - self.solver_options = solver_options + """ + Apply an in-place :class:`ModelDiff` to the built native model. - # Check for the solver to be initialized whether the package is installed or not. - if self.solver_name.value not in available_solvers: - msg = f"Solver package for '{self.solver_name.value}' is not installed. Please install first to initialize solver instance." - raise ImportError(msg) + Template method: validates the diff up front (a rejected update + leaves the native model untouched), then walks the sections in a + fixed order, dispatching to the per-backend ``_apply_*`` hooks. + """ + if not self.supports_persistent_update: + raise UnsupportedUpdate(type(self).__name__) + self._validate_update(diff) + ctx = self._apply_begin(var_label_index, con_label_index) + if diff.var_bounds_indices.size: + self._apply_var_bounds( + ctx, + diff.var_bounds_indices, + diff.var_bounds_lower, + diff.var_bounds_upper, + ) + if diff.var_type_positions.size: + self._apply_var_types(ctx, diff.var_type_positions, diff.var_type_kinds) + self._reclamp_binary_bounds( + ctx, diff.var_type_positions, diff.var_type_kinds + ) + if diff.con_rhs_indices.size: + self._apply_con_rhs(ctx, diff) + if diff.con_sign_indices.size: + self._apply_con_signs(ctx, diff.con_sign_indices, diff.con_sign_values) + if diff.n_coef_updates: + self._apply_con_coefs( + ctx, diff.con_coef_rows, diff.con_coef_cols, diff.con_coef_vals + ) + if diff.obj_c_indices is not None: + assert diff.obj_c_values is not None + self._apply_obj_linear(ctx, diff.obj_c_indices, diff.obj_c_values) + if diff.obj_sense is not None: + self._apply_obj_sense(ctx, diff.obj_sense) + self.sense = diff.obj_sense + self._apply_end(ctx) + + def _validate_update(self, diff: ModelDiff) -> None: + """Reject unsupported diff content before any native mutation.""" + if diff.con_sign_indices.size and not self.supports_sign_update: + raise UnsupportedUpdate( + f"{self.display_name} does not support in-place constraint sign change" + ) - def safe_get_solution( - self, status: Status, func: Callable[[], Solution] - ) -> Solution: + def _apply_begin(self, var_label_index: Any, con_label_index: Any) -> Any: + """Backend prep + validation; the return value is passed to every hook.""" + return self.solver_model + + def _apply_end(self, ctx: Any) -> None: + return None + + def _apply_var_bounds( + self, ctx: Any, indices: np.ndarray, lower: np.ndarray, upper: np.ndarray + ) -> None: + raise NotImplementedError + + def _apply_var_types( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray + ) -> None: + raise NotImplementedError + + def _reclamp_binary_bounds( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray + ) -> None: """ - Get solution from function call, if status is unknown still try to run it. + Re-clamp variables switched to BINARY to [0, 1]. + + Compensates for backends whose native type system only has a generic + integer kind; backends where the binary type implies the bounds + (Gurobi) override with a no-op. """ - if status.is_ok: - return func() - elif status.status == SolverStatus.unknown: - try: - logger.warning("Solution status unknown. Trying to parse solution.") - sol = func() - status.status = SolverStatus.ok - logger.warning("Solution parsed successfully.") - return sol - except Exception as e: - logger.error(f"Failed to parse solution: {e}") - return Solution() + binary_mask = kinds == VarKind.BINARY + if binary_mask.any(): + bin_positions = positions[binary_mask] + n = bin_positions.size + self._apply_var_bounds(ctx, bin_positions, np.zeros(n), np.ones(n)) - @abstractmethod - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: EnvType | None = None, - explicit_coordinate_names: bool = False, - ) -> Result: + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + raise NotImplementedError + + def _apply_con_signs( + self, ctx: Any, indices: np.ndarray, signs: np.ndarray + ) -> None: + raise NotImplementedError + + def _apply_con_coefs( + self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray + ) -> None: + raise NotImplementedError + + def _apply_obj_linear( + self, ctx: Any, indices: np.ndarray, values: np.ndarray + ) -> None: + raise NotImplementedError + + def _apply_obj_sense(self, ctx: Any, sense: str) -> None: + raise NotImplementedError + + @property + def solver_options(self) -> dict[str, Any]: + return self.options + + @classmethod + @functools.cache + def is_available(cls) -> bool: """ - Abstract method to solve a linear problem from a model. + Return True if this solver's package/binary is importable. - Needs to be implemented in the specific solver subclass. Even if the solver - does not support solving from a model, this method should be implemented and - raise a NotImplementedError. + Must not acquire a license. Subclasses override with the cheapest + possible probe. Base returns False so a forgotten override fails + safe (the solver simply does not show up in ``available_solvers``). """ - pass + return False - @abstractmethod - def solve_problem_from_file( + @classmethod + def license_status(cls) -> LicenseStatus: + """ + Probe license/runtime availability. May acquire a license slot. + + Not cached — license state is mutable (server reachability, expiry). + """ + name = SolverName[cls.__name__].value + if not cls.is_available(): + return LicenseStatus(name, ok=False, message="package not installed") + try: + cls._license_probe() + except Exception as e: + return LicenseStatus(name, ok=False, message=f"{type(e).__name__}: {e}") + return LicenseStatus(name, ok=True) + + @classmethod + def _license_probe(cls) -> None: + """Subclass hook. Default no-op. Raises on failure.""" + return None + + @classmethod + def runtime_features(cls) -> frozenset[SolverFeature]: + """ + Features whose availability depends on the installed solver version + or runtime environment. Override in subclasses; the default is empty. + """ + return frozenset() + + @classmethod + def supported_features(cls) -> frozenset[SolverFeature]: + """All features supported by this solver, static plus runtime.""" + return cls.features | cls.runtime_features() + + @classmethod + def supports(cls, feature: SolverFeature) -> bool: + """Check if this solver supports a given feature.""" + return feature in cls.features or feature in cls.runtime_features() + + @staticmethod + def from_name( + name: str, + model: Model | None = None, + io_api: str | None = None, + options: dict[str, Any] | None = None, + track_updates: bool = False, + **build_kwargs: Any, + ) -> Solver: + """ + Construct the solver subclass registered as ``name``. + + With ``model`` supplied, the solver is built immediately. Without it, + an unbuilt instance is returned and the first ``solve(model, ...)`` + call performs the build. See :class:`Solver` for ``track_updates``. + """ + cls = _solver_class_for(name) + if cls is None: + raise ValueError(f"unknown solver: {name}") + if model is None: + return cls( + model=None, + io_api=io_api, + options=options or {}, + track_updates=track_updates, + ) + return cls.from_model( + model, + io_api=io_api, + options=options or {}, + track_updates=track_updates, + **build_kwargs, + ) + + @classmethod + def from_model( + cls, + model: Model, + io_api: str | None = None, + options: dict[str, Any] | None = None, + track_updates: bool = False, + **build_kwargs: Any, + ) -> Solver: + """Instantiate and build the solver against ``model``.""" + instance = cls( + model=model, + io_api=io_api, + options=options or {}, + track_updates=track_updates, + ) + instance._build(**build_kwargs) + return instance + + def _build(self, **build_kwargs: Any) -> None: + """ + Dispatch to direct or file build based on ``io_api``. + + The Solver never mutates ``self.model``. Constraint sanitization + (``model.constraints.sanitize_zeros()`` / + ``.sanitize_infinities()``) and SOS reformulation + (``model.apply_sos_reformulation()``) are Model-level operations + the caller applies first; this builder consumes whatever shape it + is handed. + """ + if self.model is None: + raise RuntimeError("Solver has no model attached; cannot build.") + self._validate_model() + if self.io_api == "direct": + self._build_direct(**build_kwargs) + if self.track_updates: + self.snapshot = ModelSnapshot.capture(self.model) + clear_coef_dirty(self.model) + else: + self._build_file(**build_kwargs) + + def _validate_model(self) -> None: + """Pre-build checks on whether this solver can handle ``self.model``.""" + model = self.model + assert model is not None + solver_name = self.solver_name.value + cls = type(self) + + if model.is_quadratic and not cls.supports(SolverFeature.QUADRATIC_OBJECTIVE): + raise ValueError( + f"Solver {solver_name} does not support quadratic problems." + ) + + if model.variables.semi_continuous and not cls.supports( + SolverFeature.SEMI_CONTINUOUS_VARIABLES + ): + raise ValueError( + f"Solver {solver_name} does not support semi-continuous variables. " + "Use a solver that supports them (gurobi, cplex, highs)." + ) + + if model.variables.sos and not cls.supports(SolverFeature.SOS_CONSTRAINTS): + raise ValueError( + f"Solver {solver_name} does not support SOS constraints. " + "Reformulate first via `Model.solve(reformulate_sos=True)` or " + "`model.apply_sos_reformulation()`, or use a solver that supports SOS." + ) + + if model.indicator_constraints and not cls.supports( + SolverFeature.INDICATOR_CONSTRAINTS + ): + raise ValueError( + f"Solver {solver_name} does not support indicator constraints. " + "Use a solver that supports them." + ) + + def _build_direct(self, **build_kwargs: Any) -> None: + """Build the native solver model from ``self.model``. Override per-solver.""" + raise NotImplementedError( + f"Solver {self.solver_name.value} does not support direct API model export." + ) + + def _build_file(self, **build_kwargs: Any) -> None: + """Write the LP/MPS file for ``self.model`` and cache its path.""" + model = self.model + assert model is not None + io_api = self.io_api + if io_api is not None and io_api not in FILE_IO_APIS: + raise ValueError( + f"Keyword argument `io_api` has to be one of {IO_APIS} or None" + ) + explicit_coordinate_names = build_kwargs.pop("explicit_coordinate_names", False) + slice_size = build_kwargs.pop("slice_size", 2_000_000) + progress = build_kwargs.pop("progress", None) + problem_fn = build_kwargs.pop("problem_fn", None) + if problem_fn is None: + problem_fn = model.get_problem_file(io_api=io_api) + if not self.supports(SolverFeature.LP_FILE_NAMES) and explicit_coordinate_names: + logger.warning( + f"{self.solver_name.value} does not support writing names to " + "lp files, disabling it." + ) + explicit_coordinate_names = False + problem_fn = model.to_file( + Path(problem_fn) if not isinstance(problem_fn, Path) else problem_fn, + io_api=io_api, + explicit_coordinate_names=explicit_coordinate_names, + slice_size=slice_size, + progress=progress, + ) + self._problem_fn = problem_fn + if self.io_api is None: + self.io_api = read_io_api_from_problem_file(problem_fn) + self._cache_model_sizes(model) + + def solve( self, - problem_fn: Path, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: EnvType | None = None, + model: Model | None = None, + assign: bool = False, + ignore_dims: Iterable[str] = (), + disallow_rebuild: bool = False, + **run_kwargs: Any, ) -> Result: """ - Abstract method to solve a linear problem from a problem file. + Run the prepared solver and return a :class:`Result`. + + With ``model`` supplied, diff against the previous build and either + apply in place or rebuild before running. Requires ``io_api='direct'``. + With ``assign=True`` the Result is written back to the target Model + via :meth:`Model.assign_result`. + + Coordinate alignment is checked on every dim by default. Pass + ``ignore_dims`` to exclude dims whose coord values legitimately shift + between solves. + + Pass ``disallow_rebuild=True`` to guarantee that an existing solver + model is updated in place — any condition that would force a rebuild + (structural change, sparsity change, backend rejection, …) raises + :class:`RebuildRequiredError` instead. The initial build on the first + ``solve(model, ...)`` is still allowed. + + Thread safety: the solver lock is held for the entire call, + including the native run. This is deliberate — diff/apply and the + run must be atomic (otherwise a concurrent apply would change the + problem between apply and run), and native solver handles are not + thread-safe. Concurrent solves therefore serialize per Solver + instance; use separate instances for parallelism. Pure diff + computation (``update(model, apply=False)``) does not take the lock. + """ + if model is not None and self.io_api != "direct": + raise ValueError("solve(model=...) requires io_api='direct'") + + with self._lock: + if model is not None: + if self.solver_model is None: + self.model = model + self._build() + else: + if not self.track_updates and model is self.model: + raise UpdatesDisabledError( + "Solver was constructed with track_updates=False; " + "in-place mutations of the build-time Model cannot " + "be detected without a snapshot. Pass a freshly " + "built Model instance, or reconstruct the solver " + "with Solver.from_name(..., track_updates=True)." + ) + self._apply_locked( + model, + ignore_dims=ignore_dims, + disallow_rebuild=disallow_rebuild, + ) + target = model + else: + target = self.model # type: ignore[assignment] + + if self.model is not None and self.model.objective.expression.empty: + raise ValueError( + "No objective has been set on the model. Use `m.add_objective(...)` " + "first (e.g. `m.add_objective(0 * x)` for a pure feasibility problem)." + ) + if self.io_api == "direct" or self.solver_model is not None: + result = self._run_direct(**run_kwargs) + elif self._problem_fn is not None: + result = self._run_file(**run_kwargs) + else: + raise RuntimeError( + "Solver has not been built; call Solver.from_name(...) or _build() first." + ) - Needs to be implemented in the specific solver subclass. Even if the solver - does not support solving from a file, this method should be implemented and - raise a NotImplementedError. + if assign and target is not None: + target.assign_result(result, solver=self) + return result + + def update( + self, + model: Model, + apply: bool = True, + ignore_dims: Iterable[str] = (), + ) -> ModelDiff | RebuildReason: + """ + Diff ``model`` against the solver state and optionally apply it. + + With ``apply=False`` the diff is computed without taking the solver + lock, so it can overlap a concurrently running solve. The preview + always runs a full comparison (no ``_coef_dirty`` shortcut — a + concurrent apply may clear the flag against a newer snapshot), so it + can report raw in-place ``.values[...]`` mutations that the apply + path, which trusts the flag for the build-time model, would miss. """ - pass + if self.io_api != "direct": + raise ValueError("update requires io_api='direct'") + if self.solver_model is None: + raise RuntimeError("Solver has not been built") + if not self.track_updates and model is self.model: + raise UpdatesDisabledError( + "Solver was constructed with track_updates=False; " + "in-place mutations of the build-time Model cannot be " + "detected without a snapshot. Pass a freshly built Model " + "instance, or reconstruct the solver with " + "Solver.from_name(..., track_updates=True)." + ) + if not apply: + return self._compute_diff(model, ignore_dims, same_model=False) + with self._lock: + return self._apply_locked(model, ignore_dims=ignore_dims) + + def _compute_diff( + self, model: Model, ignore_dims: Iterable[str], same_model: bool + ) -> ModelDiff | RebuildReason: + """ + Diff ``model`` against the solver baseline (the captured snapshot, or + the build-time Model when no snapshot is tracked). + + ``same_model=True`` lets ``from_snapshot`` trust the ``_coef_dirty`` + flag and skip the coefficient re-walk; ``same_model=False`` forces a + full comparison. Snapshot and baseline refs are read once, so the walk + stays consistent even while a concurrent apply swaps ``self.snapshot``; + the ``from_models`` fallback is only consistent if no thread + concurrently mutates either Model. + """ + snapshot = self.snapshot + if snapshot is not None: + return ModelDiff.from_snapshot( + snapshot, model, same_model=same_model, ignore_dims=ignore_dims + ) + baseline = self.model + assert baseline is not None + return ModelDiff.from_models(baseline, model, ignore_dims=ignore_dims) + + def _apply_locked( + self, + model: Model, + ignore_dims: Iterable[str] = (), + disallow_rebuild: bool = False, + ) -> ModelDiff | RebuildReason: + if not self.supports_persistent_update: + if disallow_rebuild: + raise RebuildRequiredError(RebuildReason.BACKEND_REJECTED) + self._rebuild(model, RebuildReason.BACKEND_REJECTED) + return RebuildReason.BACKEND_REJECTED + diff = self._compute_diff(model, ignore_dims, same_model=model is self.model) + if isinstance(diff, RebuildReason): + if disallow_rebuild: + raise RebuildRequiredError(diff) + self._rebuild(model, diff) + return diff + try: + self.apply_update( + diff, + model.variables.label_index, + model.constraints.label_index, + ) + except Exception as exc: + if disallow_rebuild: + raise RebuildRequiredError( + RebuildReason.BACKEND_REJECTED, str(exc) + ) from exc + self._last_rebuild_reason = RebuildReason.BACKEND_REJECTED + self._rebuild(model, RebuildReason.BACKEND_REJECTED) + return diff + self.model = model + if self.track_updates: + self.snapshot = diff.snapshot + clear_coef_dirty(model) + self._in_place_updates += 1 + self._last_rebuild_reason = None + return diff + + def _rebuild(self, model: Model, reason: RebuildReason) -> None: + self.close() + self.model = model + self._build() + self._rebuilds += 1 + self._last_rebuild_reason = reason + + def _run_direct(self, **run_kwargs: Any) -> Result: + """Run the pre-built native solver model. Override per-solver.""" + raise NotImplementedError( + f"Direct API not implemented for {self.solver_name.value}" + ) + + def _run_file(self, **run_kwargs: Any) -> Result: + """Invoke the solver binary on ``self._problem_fn``. Override per-solver.""" + raise NotImplementedError( + f"File-based API not implemented for {self.solver_name.value}" + ) def solve_problem( self, @@ -377,17 +993,19 @@ def solve_problem( env: EnvType | None = None, explicit_coordinate_names: bool = False, ) -> Result: - """ - Solve a linear problem either from a model or a problem file. - - Wraps around `self.solve_problem_from_model()` and - `self.solve_problem_from_file()` and calls the appropriate method - based on the input arguments (`model` or `problem_fn`). - """ + """Deprecated. Use ``Solver.from_name(...).solve(...)`` or ``Model.solve(...)``.""" + warnings.warn( + "Solver.solve_problem is deprecated and will be removed in a future " + "release. Use Solver.from_name(name, model, ...).solve(...) or " + "Model.solve(...) instead.", + DeprecationWarning, + stacklevel=2, + ) if problem_fn is not None and model is not None: - msg = "Both problem file and model are given. Please specify only one." - raise ValueError(msg) - elif model is not None: + raise ValueError( + "Both problem file and model are given. Please specify only one." + ) + if model is not None: return self.solve_problem_from_model( model=model, solution_fn=solution_fn, @@ -397,7 +1015,7 @@ def solve_problem( env=env, explicit_coordinate_names=explicit_coordinate_names, ) - elif problem_fn is not None: + if problem_fn is not None: return self.solve_problem_from_file( problem_fn=problem_fn, solution_fn=solution_fn, @@ -406,30 +1024,7 @@ def solve_problem( basis_fn=basis_fn, env=env, ) - else: - msg = "No problem file or model specified." - raise ValueError(msg) - - @property - def solver_name(self) -> SolverName: - return SolverName[self.__class__.__name__] - - -class CBC(Solver[None]): - """ - Solver subclass for the CBC solver. - - Attributes - ---------- - **solver_options - options for the given solver - """ - - def __init__( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + raise ValueError("No problem file or model specified.") def solve_problem_from_model( self, @@ -438,11 +1033,38 @@ def solve_problem_from_model( log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - env: None = None, + env: EnvType | None = None, explicit_coordinate_names: bool = False, + set_names: bool = True, ) -> Result: - msg = "Direct API not implemented for CBC" - raise NotImplementedError(msg) + """Deprecated shim that builds via ``_build_direct`` and runs via ``_run_direct``.""" + warnings.warn( + "Solver.solve_problem_from_model is deprecated and will be removed in a " + "future release. Use Solver.from_name(name, model, io_api='direct', ...)" + ".solve(...) instead.", + DeprecationWarning, + stacklevel=2, + ) + if not self.supports(SolverFeature.DIRECT_API): + raise NotImplementedError( + f"Direct API not implemented for {self.solver_name.value}" + ) + self.model = model + build_kwargs: dict[str, Any] = { + "explicit_coordinate_names": explicit_coordinate_names, + "set_names": set_names, + "log_fn": log_fn, + } + if env is not None: + build_kwargs["env"] = env + self._build_direct(**build_kwargs) + return self._run_direct( + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + env=env, + ) def solve_problem_from_file( self, @@ -451,34 +1073,160 @@ def solve_problem_from_file( log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - env: None = None, + env: EnvType | None = None, ) -> Result: - """ - Solve a linear problem from a problem file using the CBC solver. - - The function reads the linear problem file and passes it to the solver. - If the solution is successful it returns variable solutions - and constraint dual values. + """Deprecated shim that caches ``problem_fn`` and runs via ``_run_file``.""" + warnings.warn( + "Solver.solve_problem_from_file is deprecated and will be removed in a " + "future release. Use Solver.from_name(name, model, problem_fn=..., ...)" + ".solve(...) instead.", + DeprecationWarning, + stacklevel=2, + ) + problem_fn = ( + Path(problem_fn) if not isinstance(problem_fn, Path) else problem_fn + ) + self._problem_fn = problem_fn + self.io_api = read_io_api_from_problem_file(problem_fn) + return self._run_file( + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + env=env, + ) - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path - Path to the solution file. This is necessary for solving with CBC. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver + def _cache_model_labels(self, model: Model) -> None: + """Cache vlabels/clabels and total label counts for label-indexed solutions.""" + self._vlabels = model.variables.label_index.vlabels + self._clabels = model.constraints.label_index.clabels + self._n_vars = model._xCounter + self._n_cons = model._cCounter + + def _cache_model_sizes(self, model: Model) -> None: + """Cache total label counts only (file-based solvers parse names).""" + self._n_vars = model._xCounter + self._n_cons = model._cCounter + + def update_solver_model(self, model: Model, **kwargs: Any) -> None: + raise NotImplementedError + + def close(self) -> None: + if self._env_stack is not None: + self._env_stack.close() + self.env = None + self.solver_model = None + self._env_stack = None + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.close() + + def __getstate__(self) -> dict[str, Any]: + drop = {"solver_model", "env", "_env_stack", "snapshot", "_lock"} + return {k: v for k, v in self.__dict__.items() if k not in drop} + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self.solver_model = None + self.env = None + self._env_stack = None + self.snapshot = None + self._lock = threading.Lock() + + def __repr__(self) -> str: + status = self.status.status.value if self.status is not None else "unsolved" + parts = [f"name={self.solver_name.value!r}", f"status={status!r}"] + if self.io_api is not None: + parts.append(f"io_api={self.io_api!r}") + if self.solver_model is not None: + parts.append("solver_model=loaded") + if self.env is not None: + parts.append("env=active") + if self.solution is not None: + parts.append(f"objective={self.solution.objective:.4g}") + if self.report is not None and self.report.runtime is not None: + parts.append(f"runtime={self.report.runtime:.3g}s") + return f"{type(self).__name__}({', '.join(parts)})" + + def _make_result( + self, + status: Status, + solution: Solution | None, + solver_model: Any = None, + report: SolverReport | None = None, + ) -> Result: + self.status = status + self.solution = solution + self.report = report + if solver_model is not None: + self.solver_model = solver_model + return Result( + status=status, + solution=solution, + solver_model=solver_model, + solver_name=self.solver_name.value, + report=report, + ) - Returns - ------- - Result + def safe_get_solution( + self, status: Status, func: Callable[[], Solution] + ) -> Solution: + """ + Get solution from function call, if status is unknown still try to run it. """ + if status.is_ok: + return func() + elif status.status == SolverStatus.unknown: + try: + logger.warning("Solution status unknown. Trying to parse solution.") + sol = func() + status.status = SolverStatus.ok + logger.warning("Solution parsed successfully.") + return sol + except Exception as e: + logger.error(f"Failed to parse solution: {e}") + return Solution() + + @property + def solver_name(self) -> SolverName: + return SolverName[self.__class__.__name__] + + +class CBC(Solver[None]): + """ + Solver subclass for the CBC solver. + + Attributes + ---------- + **solver_options + options for the given solver + """ + + display_name: ClassVar[str] = "CBC" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.READ_MODEL_FROM_FILE, + } + ) + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return shutil.which("cbc") is not None + + def _run_file( + self, + solution_fn: Path | None = None, + log_fn: Path | None = None, + warmstart_fn: Path | None = None, + basis_fn: Path | None = None, + env: None = None, + **kw: Any, + ) -> Result: + problem_fn = self._problem_fn + assert problem_fn is not None sense = read_sense_from_problem_file(problem_fn) io_api = read_io_api_from_problem_file(problem_fn) @@ -486,7 +1234,7 @@ def solve_problem_from_file( msg = "No solution file specified. For solving with CBC this is necessary." raise ValueError(msg) - if io_api == "mps" and sense == "max" and _new_highspy_mps_layout: + if io_api == "mps" and sense == "max" and _new_highspy_mps_layout(): msg = ( "CBC does not support maximization in MPS format highspy versions " " >=1.7.1" @@ -577,9 +1325,18 @@ def get_solver_solution() -> Solution: ) variables_b = df.index.isin(variables) - sol = df[variables_b][2] - dual = df[~variables_b][3] - + sol_df = df[variables_b] + dual_df = df[~variables_b] + sol = _solution_from_names( + sol_df[2].to_numpy(dtype=float), + sol_df.index.tolist(), + self._n_vars, + ) + dual = _solution_from_names( + dual_df[3].to_numpy(dtype=float), + dual_df.index.tolist(), + self._n_cons, + ) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) @@ -598,7 +1355,13 @@ def get_solver_solution() -> Solution: runtime = float(m.group(1)) CbcModel = namedtuple("CbcModel", ["mip_gap", "runtime"]) - return Result(status, solution, CbcModel(mip_gap, runtime)) + self.io_api = io_api + return self._make_result( + status, + solution, + solver_model=CbcModel(mip_gap, runtime), + report=SolverReport(runtime=runtime, mip_gap=mip_gap), + ) class GLPK(Solver[None]): @@ -611,64 +1374,30 @@ class GLPK(Solver[None]): options for the given solver """ - def __init( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "GLPK" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.READ_MODEL_FROM_FILE, + } + ) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - msg = "Direct API not implemented for GLPK" - raise NotImplementedError(msg) + @classmethod + @functools.cache + def is_available(cls) -> bool: + return shutil.which("glpsol") is not None - def solve_problem_from_file( + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the glpk solver. - - This function reads the linear problem file and passes it to the - glpk solver. If the solution is successful it returns variable solutions - and constraint dual values. - - For more information on the glpk solver options, see - - https://kam.mff.cuni.cz/~elias/glpk.pdf - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path - Path to the solution file. This is necessary for solving with GLPK. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - - Returns - ------- - Result - """ + problem_fn = self._problem_fn + assert problem_fn is not None CONDITION_MAP = { "integer optimal": "optimal", "integer undefined": "infeasible_or_unbounded", @@ -680,7 +1409,7 @@ def solve_problem_from_file( msg = "No solution file specified. For solving with GLPK this is necessary." raise ValueError(msg) - if io_api == "mps" and sense == "max" and _new_highspy_mps_layout: + if io_api == "mps" and sense == "max" and _new_highspy_mps_layout(): msg = ( "GLPK does not support maximization in MPS format highspy versions " " >=1.7.1" @@ -728,7 +1457,8 @@ def solve_problem_from_file( if not os.path.exists(solution_fn): status = Status(SolverStatus.warning, TerminationCondition.unknown) - return Result(status, Solution()) + self.io_api = io_api + return self._make_result(status, Solution()) f = open(solution_fn) @@ -752,31 +1482,39 @@ def get_solver_solution() -> Solution: dual_io = io.StringIO("".join(read_until_break(f))[:-2]) dual_ = pd.read_fwf(dual_io)[1:].set_index("Row name") if "Marginal" in dual_: - dual = pd.to_numeric(dual_["Marginal"], "coerce").fillna(0) + dual = _solution_from_names( + pd.to_numeric(dual_["Marginal"], "coerce") + .fillna(0) + .to_numpy(dtype=float), + dual_.index.tolist(), + self._n_cons, + ) else: logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + dual = np.array([], dtype=float) sol_io = io.StringIO("".join(read_until_break(f))[:-2]) - sol = ( - pd.read_fwf(sol_io)[1:] - .set_index("Column name")["Activity"] - .astype(float) + sol_df = pd.read_fwf(sol_io)[1:].set_index("Column name") + sol = _solution_from_names( + sol_df["Activity"].astype(float).to_numpy(), + sol_df.index.tolist(), + self._n_vars, ) f.close() return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution) + self.io_api = io_api + return self._make_result(status, solution) class Highs(Solver[None]): """ - Solver subclass for the Highs solver. Highs must be installed - for usage. Find the documentation at https://www.maths.ed.ac.uk/hall/HiGHS/. + Solver subclass for the HiGHS solver. HiGHS must be installed + for usage. Find the documentation at https://highs.dev/. - The full list of solver options is documented at https://www.maths.ed.ac.uk/hall/HiGHS/HighsOptions.set. + The full list of solver options is documented at https://ergo-code.github.io/HiGHS/stable/options/definitions/. Some exemplary options are: @@ -791,50 +1529,83 @@ class Highs(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, + display_name: ClassVar[str] = "HiGHS" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.DIRECT_API, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + SolverFeature.SEMI_CONTINUOUS_VARIABLES, + SolverFeature.MIP_DUAL_BOUND_REPORT, + } + ) + supports_persistent_update: ClassVar[bool] = True + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("highspy") + + @classmethod + @functools.cache + def _vtype_map(cls) -> dict[VarKind, Any]: + return { + VarKind.CONTINUOUS: highspy.HighsVarType.kContinuous, + VarKind.BINARY: highspy.HighsVarType.kInteger, + VarKind.INTEGER: highspy.HighsVarType.kInteger, + VarKind.SEMI_CONTINUOUS: highspy.HighsVarType.kSemiContinuous, + } + + def _apply_var_bounds( + self, ctx: Any, indices: np.ndarray, lower: np.ndarray, upper: np.ndarray ) -> None: - super().__init__(**solver_options) + ctx.changeColsBounds(indices.size, indices, lower, upper) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - """ - Solve a linear problem directly from a linopy model using the Highs solver. - Reads a linear problem file and passes it to the highs solver. - If the solution is feasible the function returns the - objective, solution and dual constraint variables. + def _apply_var_types( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray + ) -> None: + type_map = self._vtype_map() + integrality = np.fromiter( + (int(type_map[k]) for k in kinds), + dtype=np.uint8, + count=positions.size, + ) + ctx.changeColsIntegrality(positions.size, positions, integrality) - Parameters - ---------- - model : linopy.model - Linopy model for the problem. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - explicit_coordinate_names : bool, optional - Transfer variable and constraint names to the solver (default: False) + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + lower, upper = diff.con_rhs_as_bounds() + for pos, lo, up in zip(diff.con_rhs_indices, lower, upper): + ctx.changeRowBounds(int(pos), float(lo), float(up)) - Returns - ------- - Result - """ - # check for Highs solver compatibility + def _apply_con_coefs( + self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray + ) -> None: + for i in range(rows.size): + ctx.changeCoeff(int(rows[i]), int(cols[i]), float(vals[i])) + + def _apply_obj_linear( + self, ctx: Any, indices: np.ndarray, values: np.ndarray + ) -> None: + ctx.changeColsCost(indices.size, indices, values) + + def _apply_obj_sense(self, ctx: Any, sense: str) -> None: + native = ( + highspy.ObjSense.kMaximize if sense == "max" else highspy.ObjSense.kMinimize + ) + ctx.changeObjectiveSense(native) + + def _build_direct( + self, + explicit_coordinate_names: bool = False, + set_names: bool = True, + log_fn: Path | None = None, + **kwargs: Any, + ) -> None: + model = self.model + assert model is not None if self.solver_options.get("solver") in [ "simplex", "ipm", @@ -848,67 +1619,124 @@ def solve_problem_from_model( "Drop the solver option or use 'choose' to enable quadratic terms / integrality." ) - h = model.to_highspy(explicit_coordinate_names=explicit_coordinate_names) + h = self._build_solver_model( + model, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, + ) self._set_solver_params(h, log_fn) + self.solver_model = h + self.io_api = "direct" + self.sense = model.sense + self._cache_model_labels(model) + + @staticmethod + def _build_solver_model( + model: Model, + explicit_coordinate_names: bool = False, + set_names: bool = True, + ) -> highspy.Highs: + """Build a highspy.Highs instance that mirrors the linopy `model`.""" + if model.variables.sos: + raise NotImplementedError( + "SOS constraints are not supported by the HiGHS direct API. " + "Use io_api='lp' instead." + ) + M = model.matrices + h = highspy.Highs() + h.addVars(len(M.vlabels), M.lb, M.ub) + if ( + len(model.binaries) + + len(model.integers) + + len(list(model.variables.semi_continuous)) + ): + vtypes = M.vtypes + integrality_map = {"C": 0, "B": 1, "I": 1, "S": 2} + int_mask = (vtypes == "B") | (vtypes == "I") | (vtypes == "S") + labels = np.arange(len(vtypes))[int_mask] + integrality = np.array( + [integrality_map[v] for v in vtypes[int_mask]], dtype=np.uint8 + ) + h.changeColsIntegrality(len(labels), labels, integrality) + + c = M.c + h.changeColsCost(len(c), np.arange(len(c), dtype=np.int32), c) + + A = M.A + if A is not None: + A = A.tocsr() + num_cons = A.shape[0] + lower = np.where(M.sense != "<", M.b, -np.inf) + upper = np.where(M.sense != ">", M.b, np.inf) + h.addRows(num_cons, lower, upper, A.nnz, A.indptr, A.indices, A.data) + + if set_names: + print_variables, print_constraints = linopy.io.get_printers_scalar( + model, explicit_coordinate_names=explicit_coordinate_names + ) + lp = h.getLp() + lp.col_names_ = print_variables(M.vlabels) + if len(M.clabels): + lp.row_names_ = print_constraints(M.clabels) + h.passModel(lp) + + Q = M.Q + if Q is not None: + Q = triu(Q).tocsr() + num_vars = Q.shape[0] + h.passHessian(num_vars, Q.nnz, 1, Q.indptr, Q.indices, Q.data) + + if model.objective.sense == "max": + h.changeObjectiveSense(highspy.ObjSense.kMaximize) + + return h + + def _run_direct( + self, + solution_fn: Path | None = None, + log_fn: Path | None = None, + warmstart_fn: Path | None = None, + basis_fn: Path | None = None, + env: Any = None, + **kw: Any, + ) -> Result: return self._solve( - h, - solution_fn, - warmstart_fn, - basis_fn, - model=model, - io_api="direct", - sense=model.sense, + self.solver_model, + solution_fn=solution_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=self.io_api, + sense=self.sense, ) - def solve_problem_from_file( + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the Highs solver. - Reads a linear problem file and passes it to the highs solver. - If the solution is feasible the function returns the - objective, solution and dual constraint variables. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - - Returns - ------- - Result - """ - + problem_fn = self._problem_fn + assert problem_fn is not None problem_fn_ = path_to_string(problem_fn) h = highspy.Highs() self._set_solver_params(h, log_fn) h.readModel(problem_fn_) + self.solver_model = h + self.io_api = read_io_api_from_problem_file(problem_fn) return self._solve( h, solution_fn, warmstart_fn, basis_fn, - io_api=read_io_api_from_problem_file(problem_fn), + io_api=self.io_api, sense=read_sense_from_problem_file(problem_fn), + from_file=True, ) def _set_solver_params( @@ -929,18 +1757,18 @@ def _solve( solution_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - model: Model | None = None, io_api: str | None = None, sense: str | None = None, + from_file: bool = False, ) -> Result: """ - Solve a linear problem from a Highs object. + Solve a linear problem from a HiGHS object. Parameters ---------- h : highspy.Highs - Highs object. + HiGHS object. solution_fn : Path, optional Path to the solution file. log_fn : Path, optional @@ -949,12 +1777,13 @@ def _solve( Path to the warmstart file. basis_fn : Path, optional Path to the basis file. - model : linopy.model, optional - Linopy model for the problem. io_api: str io_api of the problem. For direct API from linopy model this is "direct". sense: str "min" or "max" + from_file: bool + ``True`` when ``h`` was populated via ``readModel`` — HiGHS may have + reordered columns/rows, so values are re-permuted using parsed names. Returns ------- @@ -1006,20 +1835,45 @@ def _solve( def get_solver_solution() -> Solution: objective = h.getObjectiveValue() solution = h.getSolution() - - if model is not None: - sol = pd.Series(solution.col_value, model.matrices.vlabels, dtype=float) - dual = pd.Series(solution.row_dual, model.matrices.clabels, dtype=float) + sol = np.asarray(solution.col_value, dtype=float) + dual = np.asarray(solution.row_dual, dtype=float) + if from_file: + lp = h.getLp() + sol = _solution_from_names(sol, lp.col_names_, self._n_vars) + dual = _solution_from_names(dual, lp.row_names_, self._n_cons) else: - sol = pd.Series(solution.col_value, h.getLp().col_names_, dtype=float) - dual = pd.Series(solution.row_dual, h.getLp().row_names_, dtype=float) - + sol = _solution_from_labels(sol, self._vlabels, self._n_vars) + dual = _solution_from_labels(dual, self._clabels, self._n_cons) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, h) + runtime: float | None = None + mip_gap: float | None = None + dual_bound: float | None = None + with contextlib.suppress(Exception): + runtime = float(h.getRunTime()) + with contextlib.suppress(Exception): + mip_gap = float(h.getInfo().mip_gap) + with contextlib.suppress(Exception): + dual_bound = float(h.getInfo().mip_dual_bound) + + self.io_api = io_api + return self._make_result( + status, + solution, + solver_model=h, + report=SolverReport( + runtime=runtime, mip_gap=mip_gap, dual_bound=dual_bound + ), + ) + + +class _GurobiApplyCtx(NamedTuple): + gm: Any + gvars: list[Any] + gcons: list[Any] class Gurobi(Solver["gurobipy.Env | dict[str, Any] | None"]): @@ -1032,126 +1886,272 @@ class Gurobi(Solver["gurobipy.Env | dict[str, Any] | None"]): options for the given solver """ - def __init__( + display_name: ClassVar[str] = "Gurobi" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.DIRECT_API, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + SolverFeature.IIS_COMPUTATION, + SolverFeature.SOS_CONSTRAINTS, + SolverFeature.INDICATOR_CONSTRAINTS, + SolverFeature.SEMI_CONTINUOUS_VARIABLES, + SolverFeature.SOLVER_ATTRIBUTE_ACCESS, + SolverFeature.MIP_DUAL_BOUND_REPORT, + } + ) + supports_persistent_update: ClassVar[bool] = True + supports_sign_update: ClassVar[bool] = True + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("gurobipy") + + @classmethod + def _license_probe(cls) -> None: + with gurobipy.Env(): + pass + + def _resolve_env(self, env: gurobipy.Env | dict[str, Any] | None) -> gurobipy.Env: + self.close() + self._env_stack = contextlib.ExitStack() + if env is None: + resolved = self._env_stack.enter_context(gurobipy.Env()) + elif isinstance(env, dict): + resolved = self._env_stack.enter_context(gurobipy.Env(params=env)) + else: + resolved = env + self.env = resolved + return resolved + + def _build_direct( self, - **solver_options: Any, + explicit_coordinate_names: bool = False, + env: gurobipy.Env | dict[str, Any] | None = None, + set_names: bool = True, + **kwargs: Any, ) -> None: - super().__init__(**solver_options) + model = self.model + assert model is not None + env_ = self._resolve_env(env) + m = self._build_solver_model( + model, + env=env_, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, + ) + self.solver_model = m + self.io_api = "direct" + self.sense = model.sense + self._cache_model_labels(model) - def solve_problem_from_model( - self, + @staticmethod + def _build_solver_model( model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: gurobipy.Env | dict[str, Any] | None = None, + env: gurobipy.Env | None = None, explicit_coordinate_names: bool = False, - ) -> Result: - """ - Solve a linear problem directly from a linopy model using the Gurobi solver. - Reads a problem file and passes it to the Gurobi solver. - This function communicates with gurobi using the gurobipy package. + set_names: bool = True, + ) -> gurobipy.Model: + """Build a gurobipy.Model that mirrors the linopy `model`.""" + model.constraints.sanitize_missings() + gm = gurobipy.Model(env=env) + + M = model.matrices + + kwargs: dict[str, Any] = {} + if set_names: + print_variables, print_constraints = linopy.io.get_printers_scalar( + model, explicit_coordinate_names=explicit_coordinate_names + ) + kwargs["name"] = print_variables(M.vlabels) + if ( + len(model.binaries.labels) + + len(model.integers.labels) + + len(list(model.variables.semi_continuous)) + ): + kwargs["vtype"] = M.vtypes + x = gm.addMVar(M.vlabels.shape, M.lb, M.ub, **kwargs) + + if model.is_quadratic: + assert M.Q is not None + gm.setObjective(0.5 * x.T @ M.Q @ x + M.c @ x) + else: + gm.setObjective(M.c @ x) + + if model.objective.sense == "max": + gm.ModelSense = -1 + + if M.A is not None: + c = gm.addMConstr(M.A, x, M.sense, M.b) + if set_names: + names = print_constraints(M.clabels) + c.setAttr("ConstrName", names) + + for sos_type, positions, weights in _iter_sos_sets(model): + gm.addSOS(sos_type, x[positions.tolist()].tolist(), weights.tolist()) + + if M.indicator_A is not None: + sense_map = { + "<": gurobipy.GRB.LESS_EQUAL, + ">": gurobipy.GRB.GREATER_EQUAL, + "=": gurobipy.GRB.EQUAL, + } + x_list = x.tolist() + A = M.indicator_A + for i in range(A.shape[0]): + lhs = gurobipy.LinExpr() + start, end = A.indptr[i], A.indptr[i + 1] + for col, coeff in zip(A.indices[start:end], A.data[start:end]): + lhs.add(x_list[int(col)], float(coeff)) + gm.addGenConstrIndicator( + x_list[int(M.indicator_binvar[i])], + bool(M.indicator_binval[i]), + lhs, + sense_map[str(M.indicator_sense[i])], + float(M.indicator_b[i]), + ) - Parameters - ---------- - model : linopy.model - Linopy model for the problem. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : gurobipy.Env or dict, optional - Gurobi environment for the solver, pass env directly or kwargs for creation. - explicit_coordinate_names : bool, optional - Transfer variable and constraint names to the solver (default: False) + gm.update() + return gm + + _GUROBI_VTYPE_MAP: ClassVar[dict[VarKind, str]] = { + VarKind.CONTINUOUS: "C", + VarKind.BINARY: "B", + VarKind.INTEGER: "I", + VarKind.SEMI_CONTINUOUS: "S", + } + _GUROBI_SIGN_MAP: ClassVar[dict[str, str]] = { + short_LESS_EQUAL: "<", + short_GREATER_EQUAL: ">", + EQUAL: "=", + } + _GUROBI_SENSE_MAP: ClassVar[dict[str, int]] = {"min": 1, "max": -1} + + def _apply_begin(self, var_label_index: Any, con_label_index: Any) -> Any: + gm = self.solver_model + gurobi_vars = gm.getVars() + gurobi_cons = gm.getConstrs() + if len(gurobi_vars) != var_label_index.n_active_vars: + raise UnsupportedUpdate("gurobi var count mismatch") + if len(gurobi_cons) != con_label_index.n_active_cons: + raise UnsupportedUpdate("gurobi con count mismatch") + return _GurobiApplyCtx(gm, gurobi_vars, gurobi_cons) + + def _apply_end(self, ctx: Any) -> None: + ctx.gm.update() + + def _apply_var_bounds( + self, ctx: Any, indices: np.ndarray, lower: np.ndarray, upper: np.ndarray + ) -> None: + gm, gvars, _ = ctx + subset = [gvars[int(i)] for i in indices] + gm.setAttr("LB", subset, lower.tolist()) + gm.setAttr("UB", subset, upper.tolist()) - Returns - ------- - Result - """ - with contextlib.ExitStack() as stack: - if env is None: - env_ = stack.enter_context(gurobipy.Env()) - elif isinstance(env, dict): - env_ = stack.enter_context(gurobipy.Env(params=env)) - else: - env_ = env + def _apply_var_types( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray + ) -> None: + gm, gvars, _ = ctx + subset = [gvars[int(p)] for p in positions] + vtypes = [self._GUROBI_VTYPE_MAP[k] for k in kinds] + gm.setAttr("VType", subset, vtypes) - m = model.to_gurobipy( - env=env_, explicit_coordinate_names=explicit_coordinate_names - ) + def _reclamp_binary_bounds( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray + ) -> None: + # Gurobi's VType 'B' natively implies [0, 1]; no bound writes needed. + return None - return self._solve( - m, - solution_fn=solution_fn, - log_fn=log_fn, - warmstart_fn=warmstart_fn, - basis_fn=basis_fn, - io_api="direct", - sense=model.sense, - ) + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + gm, _, gcons = ctx + subset = [gcons[int(r)] for r in diff.con_rhs_indices] + gm.setAttr("RHS", subset, diff.con_rhs_values.tolist()) - def solve_problem_from_file( + def _apply_con_signs( + self, ctx: Any, indices: np.ndarray, signs: np.ndarray + ) -> None: + gm, _, gcons = ctx + senses = [] + for s in signs: + s_str = str(s) + if s_str not in self._GUROBI_SIGN_MAP: + raise UnsupportedUpdate(f"unknown sign {s_str!r}") + senses.append(self._GUROBI_SIGN_MAP[s_str]) + subset = [gcons[int(r)] for r in indices] + gm.setAttr("Sense", subset, senses) + + def _apply_con_coefs( + self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray + ) -> None: + gm, gvars, gcons = ctx + for i in range(rows.size): + gm.chgCoeff(gcons[int(rows[i])], gvars[int(cols[i])], float(vals[i])) + + def _apply_obj_linear( + self, ctx: Any, indices: np.ndarray, values: np.ndarray + ) -> None: + gm, gvars, _ = ctx + subset = [gvars[int(i)] for i in indices] + gm.setAttr("Obj", subset, values.tolist()) + + def _apply_obj_sense(self, ctx: Any, sense: str) -> None: + if sense not in self._GUROBI_SENSE_MAP: + raise UnsupportedUpdate(f"unknown obj sense {sense!r}") + ctx.gm.ModelSense = self._GUROBI_SENSE_MAP[sense] + + def _run_direct( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - env: gurobipy.Env | dict[str, Any] | None = None, + env: Any = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the Gurobi solver. - Reads a problem file and passes it to the Gurobi solver. - This function communicates with gurobi using the gurobipy package. + return self._solve( + self.solver_model, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=self.io_api, + sense=self.sense, + ) - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : gurobipy.Env or dict, optional - Gurobi environment for the solver, pass env directly or kwargs for creation. - - Returns - ------- - Result - """ + def _run_file( + self, + solution_fn: Path | None = None, + log_fn: Path | None = None, + warmstart_fn: Path | None = None, + basis_fn: Path | None = None, + env: gurobipy.Env | dict[str, Any] | None = None, + **kw: Any, + ) -> Result: + problem_fn = self._problem_fn + assert problem_fn is not None sense = read_sense_from_problem_file(problem_fn) io_api = read_io_api_from_problem_file(problem_fn) problem_fn_ = path_to_string(problem_fn) - with contextlib.ExitStack() as stack: - if env is None: - env_ = stack.enter_context(gurobipy.Env()) - elif isinstance(env, dict): - env_ = stack.enter_context(gurobipy.Env(params=env)) - else: - env_ = env - - m = gurobipy.read(problem_fn_, env=env_) + env_ = self._resolve_env(env) + m = gurobipy.read(problem_fn_, env=env_) + self.solver_model = m + self.io_api = io_api - return self._solve( - m, - solution_fn=solution_fn, - log_fn=log_fn, - warmstart_fn=warmstart_fn, - basis_fn=basis_fn, - io_api=io_api, - sense=sense, - ) + return self._solve( + m, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=io_api, + sense=sense, + from_file=True, + ) def _solve( self, @@ -1162,6 +2162,7 @@ def _solve( basis_fn: Path | None, io_api: str | None, sense: str | None, + from_file: bool = False, ) -> Result: """ Solve a linear problem from a Gurobi object. @@ -1239,22 +2240,54 @@ def _solve( def get_solver_solution() -> Solution: objective = m.ObjVal - sol = pd.Series({v.VarName: v.x for v in m.getVars()}, dtype=float) # type: ignore + vars_ = m.getVars() + sol = np.array([v.X for v in vars_], dtype=float) + if from_file: + sol = _solution_from_names( + sol, [v.VarName for v in vars_], self._n_vars + ) + else: + sol = _solution_from_labels(sol, self._vlabels, self._n_vars) try: - dual = pd.Series( - {c.ConstrName: c.Pi for c in m.getConstrs()}, dtype=float - ) + constrs = m.getConstrs() + dual = np.array([c.Pi for c in constrs], dtype=float) + if from_file: + dual = _solution_from_names( + dual, + [c.ConstrName for c in constrs], + self._n_cons, + ) + else: + dual = _solution_from_labels(dual, self._clabels, self._n_cons) except AttributeError: logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + dual = np.array([], dtype=float) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) - solution = solution = maybe_adjust_objective_sign(solution, io_api, sense) + solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + runtime: float | None = None + mip_gap: float | None = None + dual_bound: float | None = None + with contextlib.suppress(Exception): + runtime = float(m.Runtime) + with contextlib.suppress(Exception): + mip_gap = float(m.MIPGap) + with contextlib.suppress(Exception): + dual_bound = float(m.ObjBound) + + self.io_api = io_api + return self._make_result( + status, + solution, + solver_model=m, + report=SolverReport( + runtime=runtime, mip_gap=mip_gap, dual_bound=dual_bound + ), + ) class Cplex(Solver[None]): @@ -1271,60 +2304,35 @@ class Cplex(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "CPLEX" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOS_CONSTRAINTS, + SolverFeature.INDICATOR_CONSTRAINTS, + SolverFeature.SEMI_CONTINUOUS_VARIABLES, + } + ) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - msg = "Direct API not implemented for Cplex" - raise NotImplementedError(msg) + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("cplex") - def solve_problem_from_file( + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the cplex solver. - - This function reads the linear problem file and passes it to the cplex - solver. If the solution is successful it returns variable solutions and - constraint dual values. Cplex must be installed for using this function. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - - Returns - ------- - Result - """ + problem_fn = self._problem_fn + assert problem_fn is not None CONDITION_MAP = { "integer optimal solution": "optimal", "integer optimal, tolerance": "optimal", @@ -1392,25 +2400,30 @@ def get_solver_solution() -> Solution: objective = m.solution.get_objective_value() - solution = pd.Series( - m.solution.get_values(), m.variables.get_names(), dtype=float + solution = _solution_from_names( + np.asarray(m.solution.get_values(), dtype=float), + m.variables.get_names(), + self._n_vars, ) - if is_lp: - dual = pd.Series( - m.solution.get_dual_values(), + try: + dual = _solution_from_names( + np.asarray(m.solution.get_dual_values(), dtype=float), m.linear_constraints.get_names(), - dtype=float, + self._n_cons, ) - else: - logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + except Exception: + logger.warning( + "Dual values not available (e.g. barrier solution without crossover)" + ) + dual = np.array([], dtype=float) return Solution(solution, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) class SCIP(Solver[None]): @@ -1423,58 +2436,33 @@ class SCIP(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "SCIP" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + } + ) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - msg = "Direct API not implemented for SCIP" - raise NotImplementedError(msg) + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("pyscipopt") - def solve_problem_from_file( + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the scip solver. - - This function communicates with scip using the pyscipopt package. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - - Returns - ------- - Result - """ + problem_fn = self._problem_fn + assert problem_fn is not None CONDITION_MAP: dict[str, TerminationCondition] = { # https://github.com/scipopt/scip/blob/b2bac412222296ff2b7f2347bb77d5fc4e05a2a1/src/scip/type_stat.h#L40 "inforunbd": TerminationCondition.infeasible_or_unbounded, @@ -1541,29 +2529,32 @@ def get_solver_solution() -> Solution: vars_to_ignore = {"quadobjvar", "qmatrixvar", "quadobj", "qmatrix"} s = m.getSols()[0] - sol = pd.Series( - {v.name: s[v] for v in m.getVars() if v.name not in vars_to_ignore} + kept_vars = [v for v in m.getVars() if v.name not in vars_to_ignore] + sol = _solution_from_names( + np.array([s[v] for v in kept_vars], dtype=float), + [v.name for v in kept_vars], + self._n_vars, ) cons = m.getConss(False) if len(cons) != 0: - dual = pd.Series( - { - c.name: m.getDualSolVal(c) - for c in cons - if c.name not in vars_to_ignore - } + kept_cons = [c for c in cons if c.name not in vars_to_ignore] + dual = _solution_from_names( + np.array([m.getDualSolVal(c) for c in kept_cons], dtype=float), + [c.name for c in kept_cons], + self._n_cons, ) else: logger.warning("Dual values not available (is this an MILP?)") - dual = pd.Series(dtype=float) + dual = np.array([], dtype=float) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) class Xpress(Solver[None]): @@ -1579,61 +2570,350 @@ class Xpress(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, + display_name: ClassVar[str] = "FICO Xpress" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.DIRECT_API, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + SolverFeature.IIS_COMPUTATION, + SolverFeature.SOS_CONSTRAINTS, + } + ) + supports_persistent_update: ClassVar[bool] = True + supports_sign_update: ClassVar[bool] = True + + _XPRESS_VTYPE_MAP: ClassVar[dict[VarKind, str]] = { + VarKind.CONTINUOUS: "C", + VarKind.BINARY: "B", + VarKind.INTEGER: "I", + VarKind.SEMI_CONTINUOUS: "S", + } + _XPRESS_ROWTYPE_MAP: ClassVar[dict[str, str]] = { + short_LESS_EQUAL: "L", + short_GREATER_EQUAL: "G", + EQUAL: "E", + } + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("xpress") + + def _apply_var_bounds( + self, ctx: Any, indices: np.ndarray, lower: np.ndarray, upper: np.ndarray + ) -> None: + cols = np.concatenate([indices, indices]).astype(np.int64, copy=False) + btypes = ["L"] * indices.size + ["U"] * indices.size + lb = np.where(np.isneginf(lower), -xpress.infinity, lower) + ub = np.where(np.isposinf(upper), xpress.infinity, upper) + vals = np.concatenate([lb, ub]).astype(float, copy=False) + ctx.chgbounds(cols.tolist(), btypes, vals.tolist()) + + def _apply_var_types( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray ) -> None: - super().__init__(**solver_options) + coltypes = [self._XPRESS_VTYPE_MAP[k] for k in kinds] + ctx.chgcoltype(positions.tolist(), coltypes) - def solve_problem_from_model( + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + ctx.chgrhs(_int_list(diff.con_rhs_indices), _float_list(diff.con_rhs_values)) + + def _apply_con_signs( + self, ctx: Any, indices: np.ndarray, signs: np.ndarray + ) -> None: + rowtypes = [] + for s in signs: + s_str = str(s) + if s_str not in self._XPRESS_ROWTYPE_MAP: + raise UnsupportedUpdate(f"unknown sign {s_str!r}") + rowtypes.append(self._XPRESS_ROWTYPE_MAP[s_str]) + ctx.chgrowtype(_int_list(indices), rowtypes) + + def _apply_con_coefs( + self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray + ) -> None: + ctx.chgmcoef(_int_list(rows), _int_list(cols), _float_list(vals)) + + def _apply_obj_linear( + self, ctx: Any, indices: np.ndarray, values: np.ndarray + ) -> None: + ctx.chgobj(_int_list(indices), _float_list(values)) + + def _apply_obj_sense(self, ctx: Any, sense: str) -> None: + if sense == "max": + ctx.chgobjsense(xpress.maximize) + elif sense == "min": + ctx.chgobjsense(xpress.minimize) + else: + raise UnsupportedUpdate(f"unknown obj sense {sense!r}") + + def _build_direct( self, + explicit_coordinate_names: bool = False, + set_names: bool = True, + **kwargs: Any, + ) -> None: + model = self.model + assert model is not None + self.close() + self._env_stack = contextlib.ExitStack() + problem = self._build_solver_model( + model, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, + ) + self._env_stack.enter_context(problem) + self.solver_model = problem + self.io_api = "direct" + self.sense = model.sense + self._cache_model_labels(model) + + @staticmethod + def _build_solver_model( model: Model, + explicit_coordinate_names: bool = False, + set_names: bool = True, + ) -> xpress.problem: + """ + Build an ``xpress.problem`` that mirrors the linopy ``model`` via ``loadproblem``. + + ``loadproblem`` is Xpress' universal native-array entry point loading LP/QP/MIQP + in a single call; see the parameter reference at + https://www.fico.com/fico-xpress-optimization/docs/latest/solver/optimizer/python/HTML/problem.loadproblem.html. + SOS arguments are left ``None`` and sets are added afterwards via ``addSOS`` so + multi-dim ``add_sos_constraints`` can be grouped natively. + """ + model.constraints.sanitize_missings() + problem = xpress.problem() + + M = model.matrices + A = M.A + Q = M.Q + + if A is not None and A.nnz: + if A.format != "csc": + A = A.tocsc() + start = A.indptr.astype(np.int64, copy=False) + rowind = A.indices.astype(np.int64, copy=False) + rowcoef = A.data.astype(float, copy=False) + else: + start = np.zeros(len(M.vlabels) + 1, dtype=np.int64) + rowind = np.empty(0, dtype=np.int64) + rowcoef = np.empty(0, dtype=float) + + lb = np.asarray(M.lb, dtype=float) + ub = np.asarray(M.ub, dtype=float) + np.place(lb, np.isneginf(lb), -xpress.infinity) + np.place(ub, np.isposinf(ub), xpress.infinity) + + rowtype: np.ndarray + rhs: np.ndarray + if len(M.clabels): + sense = M.sense + rowtype = np.full(sense.shape, "E", dtype="U1") + rowtype[sense == "<"] = "L" + rowtype[sense == ">"] = "G" + rhs = np.asarray(M.b, dtype=float) + else: + rowtype = np.empty(0, dtype="U1") + rhs = np.empty(0, dtype=float) + + objqcol1: np.ndarray | None + objqcol2: np.ndarray | None + objqcoef: np.ndarray | None + if Q is not None and Q.nnz: + Qt = Q if Q.format == "coo" else triu(Q, format="coo") # codespell:ignore + mask = Qt.row <= Qt.col + objqcol1 = Qt.row[mask].astype(np.int64, copy=False) + objqcol2 = Qt.col[mask].astype(np.int64, copy=False) + objqcoef = Qt.data[mask].astype(float, copy=False) + else: + objqcol1 = None + objqcol2 = None + objqcoef = None + + vtypes = M.vtypes + integer_mask = (vtypes == "B") | (vtypes == "I") + if integer_mask.any(): + entind = np.flatnonzero(integer_mask).astype(np.int64, copy=False) + coltype = vtypes[entind] + else: + entind = None + coltype = None + + objcoef = np.asarray(M.c, dtype=float) + has_q = objqcol1 is not None + has_int = coltype is not None + base_kwargs: dict[str, Any] = dict( + probname="linopy", + rowtype=rowtype, + rhs=rhs, + rng=None, + objcoef=objcoef, + start=start, + collen=None, + rowind=rowind, + rowcoef=rowcoef, + lb=lb, + ub=ub, + ) + try: # Try new API first (Xpress 9.8+) + if has_q and has_int: + problem.loadMIQP( + **base_kwargs, + objqcol1=objqcol1, + objqcol2=objqcol2, + objqcoef=objqcoef, + coltype=coltype, + entind=entind, + ) + elif has_q: + problem.loadQP( + **base_kwargs, + objqcol1=objqcol1, + objqcol2=objqcol2, + objqcoef=objqcoef, + ) + elif has_int: + problem.loadMIP( + **base_kwargs, + coltype=coltype, + entind=entind, + ) + else: + problem.loadLP(**base_kwargs) + except AttributeError: # Fallback to old API + problem.loadproblem( + probname="linopy", + rowtype=rowtype, + rhs=rhs, + rng=None, + objcoef=objcoef, + start=start, + collen=None, + rowind=rowind, + rowcoef=rowcoef, + lb=lb, + ub=ub, + objqcol1=objqcol1, + objqcol2=objqcol2, + objqcoef=objqcoef, + qrowind=None, + nrowqcoefs=None, + rowqcol1=None, + rowqcol2=None, + rowqcoef=None, + coltype=coltype, + entind=entind, + limit=None, + settype=None, + setstart=None, + setind=None, + refval=None, + ) + + if model.objective.sense == "max": + problem.chgobjsense(xpress.maximize) + + if set_names: + print_variable, print_constraint = linopy.io.get_printers_scalar( + model, explicit_coordinate_names=explicit_coordinate_names + ) + vnames = print_variable(M.vlabels) + if vnames: + try: # Try new API first (Xpress 9.8+) + problem.addNames( + xpress_Namespaces.COLUMN, vnames, 0, len(vnames) - 1 + ) + except AttributeError: # Fallback to old API + problem.addnames( + xpress_Namespaces.COLUMN, vnames, 0, len(vnames) - 1 + ) + cnames = print_constraint(M.clabels) + if cnames: + try: # Try new API first (Xpress 9.8+) + problem.addNames(xpress_Namespaces.ROW, cnames, 0, len(cnames) - 1) + except AttributeError: # Fallback to old API + problem.addnames(xpress_Namespaces.ROW, cnames, 0, len(cnames) - 1) + + for sos_type, positions, weights in _iter_sos_sets(model): + problem.addSOS(positions.tolist(), weights.tolist(), type=sos_type) + + return problem + + @classmethod + def runtime_features(cls) -> frozenset[SolverFeature]: + if _installed_version_in("xpress", ">=9.8.0"): + return frozenset({SolverFeature.GPU_ACCELERATION}) + return frozenset() + + def _run_direct( + self, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, - explicit_coordinate_names: bool = False, + **kw: Any, ) -> Result: - msg = "Direct API not implemented for Xpress" - raise NotImplementedError(msg) + return self._solve( + self.solver_model, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=self.io_api, + sense=self.sense, + ) - def solve_problem_from_file( + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the Xpress solver. + problem_fn = self._problem_fn + assert problem_fn is not None + io_api = read_io_api_from_problem_file(problem_fn) + sense = read_sense_from_problem_file(problem_fn) - This function reads the linear problem file and passes it to - the Xpress solver. If the solution is successful it returns - variable solutions and constraint dual values. The `xpress` module - must be installed for using this function. + self.close() + self._env_stack = contextlib.ExitStack() + m = self._env_stack.enter_context(xpress.problem()) + try: # Try new API first + m.readProb(path_to_string(problem_fn)) + except AttributeError: # Fallback to old API + m.read(path_to_string(problem_fn)) - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver + return self._solve( + m, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=io_api, + sense=sense, + from_file=True, + ) - Returns - ------- - Result - """ + def _solve( + self, + m: xpress.problem, + solution_fn: Path | None, + log_fn: Path | None, + warmstart_fn: Path | None, + basis_fn: Path | None, + io_api: str | None, + sense: str | None, + from_file: bool = False, + ) -> Result: CONDITION_MAP = { xpress.SolStatus.NOTFOUND: "unknown", xpress.SolStatus.OPTIMAL: "optimal", @@ -1642,17 +2922,6 @@ def solve_problem_from_file( xpress.SolStatus.UNBOUNDED: "unbounded", } - io_api = read_io_api_from_problem_file(problem_fn) - sense = read_sense_from_problem_file(problem_fn) - - m = xpress.problem() - - try: # Try new API first - m.readProb(path_to_string(problem_fn)) - except AttributeError: # Fallback to old API - m.read(path_to_string(problem_fn)) - - # Set solver options - new API uses setControl per option, old API accepts dict if self.solver_options is not None: m.setControl(self.solver_options) @@ -1670,7 +2939,6 @@ def solve_problem_from_file( m.optimize() - # if the solver is stopped (timelimit for example), postsolve the problem if m.attributes.solvestatus == xpress.enums.SolveStatus.STOPPED: try: # Try new API first m.postSolve() @@ -1683,7 +2951,7 @@ def solve_problem_from_file( m.writeBasis(path_to_string(basis_fn)) except AttributeError: # Fallback to old API m.writebasis(path_to_string(basis_fn)) - except (xpress.SolverError, xpress.ModelError) as err: + except (xpress.SolverError, xpress.ModelError) as err: # pragma: no cover logger.info("No model basis stored. Raised error: %s", err) if solution_fn is not None: @@ -1692,7 +2960,7 @@ def solve_problem_from_file( m.writeBinSol(path_to_string(solution_fn)) except AttributeError: # Fallback to old API m.writebinsol(path_to_string(solution_fn)) - except (xpress.SolverError, xpress.ModelError) as err: + except (xpress.SolverError, xpress.ModelError) as err: # pragma: no cover logger.info("Unable to save solution file. Raised error: %s", err) condition = m.attributes.solstatus @@ -1703,37 +2971,286 @@ def solve_problem_from_file( def get_solver_solution() -> Solution: objective = m.attributes.objval - try: # Try new API first - var = m.getNameList(xpress_Namespaces.COLUMN, 0, m.attributes.cols - 1) - except AttributeError: # Fallback to old API - var = m.getnamelist(xpress_Namespaces.COLUMN, 0, m.attributes.cols - 1) - sol = pd.Series(m.getSolution(), index=var, dtype=float) + sol_values = np.asarray(m.getSolution(), dtype=float) + if from_file: + sol = _solution_from_names( + sol_values, + [v.name for v in m.getVariable()], + self._n_vars, + ) + else: + sol = _solution_from_labels(sol_values, self._vlabels, self._n_vars) try: - try: # Try new API first - _dual = m.getDuals() - except AttributeError: # Fallback to old API - _dual = m.getDual() - - try: # Try new API first - constraints = m.getNameList( - xpress_Namespaces.ROW, 0, m.attributes.rows - 1 - ) - except AttributeError: # Fallback to old API - constraints = m.getnamelist( - xpress_Namespaces.ROW, 0, m.attributes.rows - 1 - ) - dual = pd.Series(_dual, index=constraints, dtype=float) - except (xpress.SolverError, xpress.ModelError, SystemError): + if m.attributes.rows == 0: + dual = np.array([], dtype=float) + else: + try: # getDuals introduced in 9.5; fallback for 9.4 + dual_values = np.asarray(m.getDuals(), dtype=float) + except AttributeError: + dual_values = np.asarray(m.getDual(), dtype=float) + if from_file: + dual = _solution_from_names( + dual_values, + [c.name for c in m.getConstraint()], + self._n_cons, + ) + else: + dual = _solution_from_labels( + dual_values, self._clabels, self._n_cons + ) + except ( + xpress.SolverError, + xpress.ModelError, + SystemError, + ): # pragma: no cover logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + dual = np.array([], dtype=float) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) + + +KnitroResult = namedtuple( + "KnitroResult", + "reported_runtime mip_relaxation_bnd mip_number_nodes mip_number_solves mip_rel_gap mip_abs_gap abs_feas_error rel_feas_error abs_opt_error rel_opt_error n_vars n_cons n_integer_vars n_continuous_vars", +) + + +class Knitro(Solver[None]): + """ + Solver subclass for the Knitro solver. + + For more information on solver options, see + https://www.artelys.com/app/docs/knitro/3_referenceManual/knitroPythonReference.html + + Attributes + ---------- + **solver_options + options for the given solver + """ + + display_name: ClassVar[str] = "Artelys Knitro" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + SolverFeature.MIP_DUAL_BOUND_REPORT, + } + ) + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("knitro") + + @classmethod + def _license_probe(cls) -> None: + kc = knitro.KN_new() + knitro.KN_free(kc) + + @staticmethod + def _set_option(kc: Any, name: str, value: Any) -> None: + param_id = knitro.KN_get_param_id(kc, name) + + if isinstance(value, bool): + value = int(value) + + if isinstance(value, int): + knitro.KN_set_int_param(kc, param_id, value) + elif isinstance(value, float): + knitro.KN_set_double_param(kc, param_id, value) + elif isinstance(value, str): + knitro.KN_set_char_param(kc, param_id, value) + else: + msg = f"Unsupported Knitro option type for {name!r}: {type(value).__name__}" + raise TypeError(msg) + + @staticmethod + def _extract_values( + kc: Any, + get_count_fn: Callable[..., Any], + get_values_fn: Callable[..., Any], + ) -> np.ndarray: + n = int(get_count_fn(kc)) + if n == 0: + return np.array([], dtype=float) + + try: + # Compatible with KNITRO >= 15 + values = get_values_fn(kc) + except TypeError: + # Fallback for older wrappers requiring explicit indices + values = get_values_fn(kc, list(range(n))) + + return np.asarray(values, dtype=float) + + def _run_file( + self, + solution_fn: Path | None = None, + log_fn: Path | None = None, + warmstart_fn: Path | None = None, + basis_fn: Path | None = None, + env: None = None, + **kw: Any, + ) -> Result: + problem_fn = self._problem_fn + assert problem_fn is not None + CONDITION_MAP: dict[int, TerminationCondition] = { + 0: TerminationCondition.optimal, + -100: TerminationCondition.suboptimal, + -101: TerminationCondition.infeasible, + -102: TerminationCondition.suboptimal, + -200: TerminationCondition.unbounded, + -201: TerminationCondition.infeasible_or_unbounded, + -202: TerminationCondition.iteration_limit, + -203: TerminationCondition.time_limit, + -204: TerminationCondition.terminated_by_limit, + -300: TerminationCondition.unbounded, + -400: TerminationCondition.iteration_limit, + -401: TerminationCondition.time_limit, + -410: TerminationCondition.terminated_by_limit, + -411: TerminationCondition.terminated_by_limit, + } + + READ_OPTIONS: dict[str, str] = {".lp": "l", ".mps": "m"} + + io_api = read_io_api_from_problem_file(problem_fn) + sense = read_sense_from_problem_file(problem_fn) + + suffix = problem_fn.suffix.lower() + if suffix not in READ_OPTIONS: + msg = f"Unsupported problem file format: {suffix}" + raise ValueError(msg) + + kc = knitro.KN_new() + try: + knitro.KN_read_problem( + kc, + path_to_string(problem_fn), + read_options=READ_OPTIONS[suffix], + ) + + if log_fn is not None: + logger.warning("Log file output not implemented for Knitro") + + for k, v in self.solver_options.items(): + self._set_option(kc, k, v) + + ret = int(knitro.KN_solve(kc)) + + reported_runtime: float | None = None + mip_relaxation_bnd: float | None = None + mip_number_nodes: int | None = None + mip_number_solves: int | None = None + mip_rel_gap: float | None = None + mip_abs_gap: float | None = None + abs_feas_error: float | None = None + rel_feas_error: float | None = None + abs_opt_error: float | None = None + rel_opt_error: float | None = None + n_vars: int | None = None + n_cons: int | None = None + n_integer_vars: int | None = None + n_continuous_vars: int | None = None + with contextlib.suppress(Exception): + reported_runtime = float(knitro.KN_get_solve_time_real(kc)) + mip_relaxation_bnd = float(knitro.KN_get_mip_relaxation_bnd(kc)) + mip_number_nodes = int(knitro.KN_get_mip_number_nodes(kc)) + mip_number_solves = int(knitro.KN_get_mip_number_solves(kc)) + mip_rel_gap = float(knitro.KN_get_mip_rel_gap(kc)) + mip_abs_gap = float(knitro.KN_get_mip_abs_gap(kc)) + abs_feas_error = float(knitro.KN_get_abs_feas_error(kc)) + rel_feas_error = float(knitro.KN_get_rel_feas_error(kc)) + abs_opt_error = float(knitro.KN_get_abs_opt_error(kc)) + rel_opt_error = float(knitro.KN_get_rel_opt_error(kc)) + n_vars = int(knitro.KN_get_number_vars(kc)) + n_cons = int(knitro.KN_get_number_cons(kc)) + var_types = list(knitro.KN_get_var_types(kc)) + n_integer_vars = int( + var_types.count(knitro.KN_VARTYPE_INTEGER) + + var_types.count(knitro.KN_VARTYPE_BINARY) + ) + n_continuous_vars = int(var_types.count(knitro.KN_VARTYPE_CONTINUOUS)) + + if ret in CONDITION_MAP: + termination_condition = CONDITION_MAP[ret] + elif ret > 0: + termination_condition = TerminationCondition.internal_solver_error + else: + termination_condition = TerminationCondition.unknown + + status = Status.from_termination_condition(termination_condition) + status.legacy_status = str(ret) + + def get_solver_solution() -> Solution: + objective = float(knitro.KN_get_obj_value(kc)) + + sol = self._extract_values( + kc, + knitro.KN_get_number_vars, + knitro.KN_get_var_primal_values, + ) + n_vars = int(knitro.KN_get_number_vars(kc)) + var_names = [knitro.KN_get_var_names(kc, i) for i in range(n_vars)] + sol = _solution_from_names(sol, var_names, self._n_vars) + + try: + dual = self._extract_values( + kc, + knitro.KN_get_number_cons, + knitro.KN_get_con_dual_values, + ) + n_cons = int(knitro.KN_get_number_cons(kc)) + con_names = [knitro.KN_get_con_names(kc, i) for i in range(n_cons)] + dual = _solution_from_names(dual, con_names, self._n_cons) + except Exception: + logger.warning("Dual values couldn't be parsed") + dual = np.array([], dtype=float) + + return Solution(sol, dual, objective) + + solution = self.safe_get_solution(status=status, func=get_solver_solution) + solution = maybe_adjust_objective_sign(solution, io_api, sense) + + if solution_fn is not None: + solution_fn.parent.mkdir(exist_ok=True) + knitro.KN_write_mps_file(kc, path_to_string(solution_fn)) + + knitro_model = KnitroResult( + reported_runtime=reported_runtime, + mip_relaxation_bnd=mip_relaxation_bnd, + mip_number_nodes=mip_number_nodes, + mip_number_solves=mip_number_solves, + mip_rel_gap=mip_rel_gap, + mip_abs_gap=mip_abs_gap, + abs_feas_error=abs_feas_error, + rel_feas_error=rel_feas_error, + abs_opt_error=abs_opt_error, + rel_opt_error=rel_opt_error, + n_vars=n_vars, + n_cons=n_cons, + n_integer_vars=n_integer_vars, + n_continuous_vars=n_continuous_vars, + ) + self.io_api = io_api + return self._make_result( + status, + solution, + solver_model=knitro_model, + report=SolverReport(runtime=reported_runtime, mip_gap=mip_rel_gap), + ) + finally: + with contextlib.suppress(Exception): + knitro.KN_free(kc) mosek_bas_re = re.compile(r" (XL|XU)\s+([^ \t]+)\s+([^ \t]+)| (LL|UL|BS)\s+([^ \t]+)") @@ -1759,127 +3276,296 @@ class Mosek(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, + display_name: ClassVar[str] = "MOSEK" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.DIRECT_API, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + } + ) + supports_persistent_update: ClassVar[bool] = True + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("mosek") + + @classmethod + def _license_probe(cls) -> None: + with mosek.Env() as env, env.Task(0, 0) as task: + task.optimize() + + def _validate_update(self, diff: ModelDiff) -> None: + super()._validate_update(diff) + if (diff.var_type_kinds == VarKind.SEMI_CONTINUOUS).any(): + raise UnsupportedUpdate("MOSEK does not support semi-continuous variables") + + def _apply_var_bounds( + self, ctx: Any, indices: np.ndarray, lower: np.ndarray, upper: np.ndarray + ) -> None: + for k in range(indices.size): + j = int(indices[k]) + lb = float(lower[k]) + ub = float(upper[k]) + ctx.chgvarbound(j, 1, int(np.isfinite(lb)), lb) + ctx.chgvarbound(j, 0, int(np.isfinite(ub)), ub) + + def _apply_var_types( + self, ctx: Any, positions: np.ndarray, kinds: np.ndarray ) -> None: - super().__init__(**solver_options) + integer_mask = (kinds == VarKind.BINARY) | (kinds == VarKind.INTEGER) + vartypes = np.where( + integer_mask, + mosek.variabletype.type_int, + mosek.variabletype.type_cont, + ).tolist() + ctx.putvartypelist(_int_list(positions, np.int32), vartypes) + + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + lower, upper = diff.con_rhs_as_bounds() + for k, i in enumerate(diff.con_rhs_indices): + lo = float(lower[k]) + up = float(upper[k]) + ctx.chgconbound(int(i), 1, int(np.isfinite(lo)), lo) + ctx.chgconbound(int(i), 0, int(np.isfinite(up)), up) + + def _apply_con_coefs( + self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray + ) -> None: + ctx.putaijlist( + _int_list(rows, np.int32), _int_list(cols, np.int32), _float_list(vals) + ) - def solve_problem_from_model( + def _apply_obj_linear( + self, ctx: Any, indices: np.ndarray, values: np.ndarray + ) -> None: + ctx.putclist(_int_list(indices, np.int32), _float_list(values)) + + def _apply_obj_sense(self, ctx: Any, sense: str) -> None: + if sense == "max": + ctx.putobjsense(mosek.objsense.maximize) + elif sense == "min": + ctx.putobjsense(mosek.objsense.minimize) + else: + raise UnsupportedUpdate(f"unknown obj sense {sense!r}") + + def _run_direct( self, - model: Model, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, + env: Any = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem directly from a linopy model using the MOSEK solver. + return self._solve( + self.solver_model, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=self.io_api, + sense=self.sense, + ) - Parameters - ---------- - model : linopy.model - Linopy model for the problem. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional, deprecated - Deprecated. This parameter is ignored. MOSEK now uses the global - environment automatically. Will be removed in a future version. - explicit_coordinate_names : bool, optional - Transfer variable and constraint names to the solver (default: False) + def _build_direct( + self, + explicit_coordinate_names: bool = False, + set_names: bool = True, + **kwargs: Any, + ) -> None: + model = self.model + assert model is not None + self.close() + self._env_stack = contextlib.ExitStack() + env = self._env_stack.enter_context(mosek.Env()) + task = self._env_stack.enter_context(env.Task(0, 0)) + m = self._build_solver_model( + model, + task, + explicit_coordinate_names=explicit_coordinate_names, + set_names=set_names, + ) + self.solver_model = m + self.io_api = "direct" + self.sense = model.sense + self._cache_model_labels(model) - Returns - ------- - Result - """ + @staticmethod + def _build_solver_model( + model: Model, + task: mosek.Task, + explicit_coordinate_names: bool = False, + set_names: bool = True, + ) -> mosek.Task: + """Populate an empty MOSEK task with the contents of `model`.""" + if model.variables.sos: + raise NotImplementedError("SOS constraints are not supported by MOSEK.") + if model.variables.semi_continuous: + raise NotImplementedError( + "Semi-continuous variables are not supported by MOSEK. " + "Use a solver that supports them (gurobi, cplex, highs)." + ) - if env is not None: - warnings.warn( - "The 'env' parameter in solve_problem_from_model is deprecated and will be " - "removed in a future version. MOSEK now uses the global environment " - "automatically, avoiding unnecessary license checkouts.", - DeprecationWarning, - stacklevel=2, + task.appendvars(model.nvars) + task.appendcons(model.ncons) + + M = model.matrices + + if set_names: + print_variables, print_constraints = linopy.io.get_printers_scalar( + model, explicit_coordinate_names=explicit_coordinate_names + ) + labels = print_variables(M.vlabels) + task.generatevarnames( + np.arange(0, len(labels)), "%0", [len(labels)], None, [0], labels ) - with mosek.Task() as m: - m = model.to_mosek(m, explicit_coordinate_names=explicit_coordinate_names) - return self._solve( - m, - solution_fn=solution_fn, - log_fn=log_fn, - warmstart_fn=warmstart_fn, - basis_fn=basis_fn, - io_api="direct", - sense=model.sense, + bkx = [ + ( + ( + (mosek.boundkey.ra if lb < ub else mosek.boundkey.fx) + if ub < np.inf + else mosek.boundkey.lo + ) + if (lb > -np.inf) + else (mosek.boundkey.up if (ub < np.inf) else mosek.boundkey.fr) ) + for (lb, ub) in zip(M.lb, M.ub) + ] + blx = [b if b > -np.inf else 0.0 for b in M.lb] + bux = [b if b < np.inf else 0.0 for b in M.ub] + task.putvarboundslice(0, model.nvars, bkx, blx, bux) + + if len(model.binaries.labels) + len(model.integers.labels) > 0: + idx = [i for (i, v) in enumerate(M.vtypes) if v in ["B", "I"]] + task.putvartypelist(idx, [mosek.variabletype.type_int] * len(idx)) + + if len(model.constraints) > 0: + if set_names: + names = print_constraints(M.clabels) + for i, n in enumerate(names): + task.putconname(i, n) + bkc = [ + ( + (mosek.boundkey.up if b < np.inf else mosek.boundkey.fr) + if s == "<" + else ( + (mosek.boundkey.lo if b > -np.inf else mosek.boundkey.up) + if s == ">" + else mosek.boundkey.fx + ) + ) + for s, b in zip(M.sense, M.b) + ] + blc = [b if b > -np.inf else 0.0 for b in M.b] + buc = [b if b < np.inf else 0.0 for b in M.b] + if M.A is not None: + A = M.A.tocsr() + task.putarowslice( + 0, model.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data + ) + task.putconboundslice(0, model.ncons, bkc, blc, buc) - def solve_problem_from_file( + if M.Q is not None: + Q = (0.5 * tril(M.Q + M.Q.transpose())).tocoo() + task.putqobj(Q.row, Q.col, Q.data) + task.putclist(list(np.arange(model.nvars)), M.c) + + if model.objective.sense == "max": + task.putobjsense(mosek.objsense.maximize) + else: + task.putobjsense(mosek.objsense.minimize) + return task + + @staticmethod + def _choose_solution(task: mosek.Task) -> mosek.soltype | None: + """ + Pick the Mosek solution with the best status available. + + Mosek may return up to three solutions per task: interior-point + (``soltype.itr``), basic (``soltype.bas``), and integer + (``soltype.itg``). Each carries its own ``solsta``: on a numerically + marginal LP solved with the default IPM+crossover, the interior-point + solver may terminate with ``solsta.dual_infeas_cer`` while crossover + recovers ``solsta.optimal`` for the basic solution. Reading only the + interior-point solution would discard the actual optimum. + + Ranking, best to worst: ``solsta.optimal`` / ``solsta.integer_optimal`` + > any other defined status > undefined. On a tie between ``bas`` and + ``itr`` (e.g. both ``optimal``) we prefer ``itr`` to preserve historical + behaviour. If ``itg`` is defined it always wins, since integer and + continuous solutions do not coexist for a well-posed task. + + Returns ``None`` if no solution is defined at all (e.g. the optimizer + crashed before producing one). + """ + + def _is_defined(soltype: mosek.soltype) -> bool: + try: + return bool(task.solutiondef(soltype)) + except mosek.Error: + return False + + if _is_defined(mosek.soltype.itg): + return mosek.soltype.itg + + optimal_statuses = {mosek.solsta.optimal, mosek.solsta.integer_optimal} + + best: mosek.soltype | None = None + best_score = -1 + # Iterate bas first and only then itr so that on a score tie + # itr wins, preserving the historical default for the common LP case. + for candidate in [mosek.soltype.bas, mosek.soltype.itr]: + if not _is_defined(candidate): + continue + try: + solsta = task.getsolsta(candidate) + except mosek.Error: + continue + score = 1 if solsta in optimal_statuses else 0 + if score >= best_score: + best = candidate + best_score = score + + return best + + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the MOSEK solver. Both mps and - lp files are supported; MPS does not support quadratic terms. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional, deprecated - Deprecated. This parameter is ignored. MOSEK now uses the global - environment automatically. Will be removed in a future version. + problem_fn = self._problem_fn + assert problem_fn is not None + self.close() + self._env_stack = contextlib.ExitStack() + mosek_env = self._env_stack.enter_context(mosek.Env()) + m = self._env_stack.enter_context(mosek_env.Task(0, 0)) + sense = read_sense_from_problem_file(problem_fn) + io_api = read_io_api_from_problem_file(problem_fn) + problem_fn_ = path_to_string(problem_fn) + m.readdata(problem_fn_) + self.solver_model = m + self.io_api = io_api - Returns - ------- - Result - """ - if env is not None: - warnings.warn( - "The 'env' parameter in solve_problem_from_file is deprecated and will be " - "removed in a future version. MOSEK now uses the global environment " - "automatically, avoiding unnecessary license checkouts.", - DeprecationWarning, - stacklevel=2, - ) - with mosek.Task() as m: - # read sense and io_api from problem file - sense = read_sense_from_problem_file(problem_fn) - io_api = read_io_api_from_problem_file(problem_fn) - # for Mosek solver, the path needs to be a string - problem_fn_ = path_to_string(problem_fn) - m.readdata(problem_fn_) - - return self._solve( - m, - solution_fn=solution_fn, - log_fn=log_fn, - warmstart_fn=warmstart_fn, - basis_fn=basis_fn, - io_api=io_api, - sense=sense, - ) + return self._solve( + m, + solution_fn=solution_fn, + log_fn=log_fn, + warmstart_fn=warmstart_fn, + basis_fn=basis_fn, + io_api=io_api, + sense=sense, + from_file=True, + ) def _solve( self, @@ -1890,6 +3576,7 @@ def _solve( basis_fn: Path | None, io_api: str | None, sense: str | None, + from_file: bool = False, ) -> Result: """ Solve a linear problem from a Mosek task object. @@ -2031,25 +3718,25 @@ def _solve( f.write(f" UL {namex}\n") f.write("ENDATA\n") - soltype = None - possible_soltypes = [ - mosek.soltype.bas, - mosek.soltype.itr, - mosek.soltype.itg, - ] - for possible_soltype in possible_soltypes: - try: - if m.solutiondef(possible_soltype): - soltype = possible_soltype - except mosek.Error: - pass + # Inspect both bas and itr (and itg for MILPs) and pick the + # solution with the best status. Reading only the interior-point + # solution may discard a valid crossover optimum. + soltype = Mosek._choose_solution(m) - if solution_fn is not None: + if solution_fn is not None and soltype is not None: try: - m.writesolution(mosek.soltype.bas, path_to_string(solution_fn)) + m.writesolution(soltype, path_to_string(solution_fn)) except mosek.Error as err: logger.info("Unable to save solution file. Raised error: %s", err) + if soltype is None: + condition = "no solution available" + status = Status.from_termination_condition( + TerminationCondition.internal_solver_error + ) + status.legacy_status = condition + return self._make_result(status, None) + condition = str(m.getsolsta(soltype)) termination_condition = CONDITION_MAP.get(condition, condition) status = Status.from_termination_condition(termination_condition) @@ -2058,24 +3745,41 @@ def _solve( def get_solver_solution() -> Solution: objective = m.getprimalobj(soltype) - sol = m.getxx(soltype) - sol = {m.getvarname(i): sol[i] for i in range(m.getnumvar())} - sol = pd.Series(sol, dtype=float) + sol_values = np.asarray(m.getxx(soltype), dtype=float) + if from_file: + sol = _solution_from_names( + sol_values, + [m.getvarname(i) for i in range(m.getnumvar())], + self._n_vars, + ) + else: + sol = _solution_from_labels(sol_values, self._vlabels, self._n_vars) try: - dual = m.gety(soltype) - dual = {m.getconname(i): dual[i] for i in range(m.getnumcon())} - dual = pd.Series(dual, dtype=float) + dual_values = np.asarray(m.gety(soltype), dtype=float) + if from_file: + dual = _solution_from_names( + dual_values, + [m.getconname(i) for i in range(m.getnumcon())], + self._n_cons, + ) + else: + dual = _solution_from_labels( + dual_values, + self._clabels, + self._n_cons, + ) except (mosek.Error, AttributeError): logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + dual = np.array([], dtype=float) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) class COPT(Solver[None]): @@ -2093,56 +3797,38 @@ class COPT(Solver[None]): options for the given solver """ - def __init( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "COPT" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + } + ) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - msg = "Direct API not implemented for COPT" - raise NotImplementedError(msg) + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("coptpy") - def solve_problem_from_file( + @classmethod + def _license_probe(cls) -> None: + env = coptpy.Envr() + env.close() + + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the COPT solver. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - COPT environment for the solver - - Returns - ------- - Result - """ + problem_fn = self._problem_fn + assert problem_fn is not None # conditions: https://guide.coap.online/copt/en-doc/constant.html#chapconst-solstatus CONDITION_MAP = { 0: "unstarted", @@ -2164,59 +3850,71 @@ def solve_problem_from_file( if env is None: env_ = coptpy.Envr() - m = env_.createModel() - - m.read(path_to_string(problem_fn)) + try: + m = env_.createModel() - if log_fn is not None: - m.setLogFile(path_to_string(log_fn)) + m.read(path_to_string(problem_fn)) - for k, v in self.solver_options.items(): - m.setParam(k, v) + if log_fn is not None: + m.setLogFile(path_to_string(log_fn)) - if warmstart_fn is not None: - m.readBasis(path_to_string(warmstart_fn)) + for k, v in self.solver_options.items(): + m.setParam(k, v) - m.solve() + if warmstart_fn is not None: + m.readBasis(path_to_string(warmstart_fn)) - if basis_fn and m.HasBasis: - try: - m.write(path_to_string(basis_fn)) - except coptpy.CoptError as err: - logger.info("No model basis stored. Raised error: %s", err) + m.solve() - if solution_fn: - try: - m.write(path_to_string(solution_fn)) - except coptpy.CoptError as err: - logger.info("No model solution stored. Raised error: %s", err) + if basis_fn and m.HasBasis: + try: + m.write(path_to_string(basis_fn)) + except coptpy.CoptError as err: + logger.warning("No model basis stored. Raised error: %s", err) - # TODO: check if this suffices - condition = m.MipStatus if m.ismip else m.LpStatus - termination_condition = CONDITION_MAP.get(condition, str(condition)) - status = Status.from_termination_condition(termination_condition) - status.legacy_status = str(condition) + if solution_fn: + try: + m.write(path_to_string(solution_fn)) + except coptpy.CoptError as err: + logger.warning("No model solution stored. Raised error: %s", err) - def get_solver_solution() -> Solution: # TODO: check if this suffices - objective = m.BestObj if m.ismip else m.LpObjVal - - sol = pd.Series({v.name: v.x for v in m.getVars()}, dtype=float) - - try: - dual = pd.Series({v.name: v.pi for v in m.getConstrs()}, dtype=float) - except (coptpy.CoptError, AttributeError): - logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + condition = m.MipStatus if m.ismip else m.LpStatus + termination_condition = CONDITION_MAP.get(condition, str(condition)) + status = Status.from_termination_condition(termination_condition) + status.legacy_status = str(condition) + + def get_solver_solution() -> Solution: + # TODO: check if this suffices + objective = m.BestObj if m.ismip else m.LpObjVal + + vars_ = m.getVars() + sol = _solution_from_names( + np.array([v.x for v in vars_], dtype=float), + [v.name for v in vars_], + self._n_vars, + ) - return Solution(sol, dual, objective) + try: + cons = m.getConstrs() + dual = _solution_from_names( + np.array([c.pi for c in cons], dtype=float), + [c.name for c in cons], + self._n_cons, + ) + except (coptpy.CoptError, AttributeError): + logger.warning("Dual values of MILP couldn't be parsed") + dual = np.array([], dtype=float) - solution = self.safe_get_solution(status=status, func=get_solver_solution) - solution = maybe_adjust_objective_sign(solution, io_api, sense) + return Solution(sol, dual, objective) - env_.close() + solution = self.safe_get_solution(status=status, func=get_solver_solution) + solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) + finally: + env_.close() class MindOpt(Solver[None]): @@ -2234,57 +3932,38 @@ class MindOpt(Solver[None]): options for the given solver """ - def __init( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "MindOpt" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.INTEGER_VARIABLES, + SolverFeature.QUADRATIC_OBJECTIVE, + SolverFeature.LP_FILE_NAMES, + SolverFeature.READ_MODEL_FROM_FILE, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + } + ) - def solve_problem_from_model( - self, - model: Model, - solution_fn: Path | None = None, - log_fn: Path | None = None, - warmstart_fn: Path | None = None, - basis_fn: Path | None = None, - env: None = None, - explicit_coordinate_names: bool = False, - ) -> Result: - msg = "Direct API not implemented for MindOpt" - raise NotImplementedError(msg) + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("mindoptpy") - def solve_problem_from_file( + @classmethod + def _license_probe(cls) -> None: + env = mindoptpy.Env() + env.dispose() + + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the MindOpt solver. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - MindOpt environment for the solver - - Returns - ------- - Result - - """ + problem_fn = self._problem_fn + assert problem_fn is not None CONDITION_MAP = { -1: "error", 0: "unknown", @@ -2309,58 +3988,73 @@ def solve_problem_from_file( if env is None: env_ = mindoptpy.Env(path_to_string(log_fn) if log_fn else "") - env_.start() - - m = mindoptpy.read(path_to_string(problem_fn), env_) - - for k, v in self.solver_options.items(): - m.setParam(k, v) - - if warmstart_fn: - try: - m.read(path_to_string(warmstart_fn)) - except mindoptpy.MindoptError as err: - logger.info("Model basis could not be read. Raised error: %s", err) - - m.optimize() + m = None + try: + env_.start() - if basis_fn: - try: - m.write(path_to_string(basis_fn)) - except mindoptpy.MindoptError as err: - logger.info("No model basis stored. Raised error: %s", err) + m = mindoptpy.read(path_to_string(problem_fn), env_) - if solution_fn: - try: - m.write(path_to_string(solution_fn)) - except mindoptpy.MindoptError as err: - logger.info("No model solution stored. Raised error: %s", err) + for k, v in self.solver_options.items(): + m.setParam(k, v) - condition = m.status - termination_condition = CONDITION_MAP.get(condition, condition) - status = Status.from_termination_condition(termination_condition) - status.legacy_status = condition + if warmstart_fn: + try: + m.read(path_to_string(warmstart_fn)) + except mindoptpy.MindoptError as err: + logger.info("Model basis could not be read. Raised error: %s", err) - def get_solver_solution() -> Solution: - objective = m.objval + m.optimize() - sol = pd.Series({v.varname: v.X for v in m.getVars()}, dtype=float) + if basis_fn: + try: + m.write(path_to_string(basis_fn)) + except mindoptpy.MindoptError as err: + logger.info("No model basis stored. Raised error: %s", err) - try: - dual = pd.Series({c.constrname: c.DualSoln for c in m.getConstrs()}) - except (mindoptpy.MindoptError, AttributeError): - logger.warning("Dual values of MILP couldn't be parsed") - dual = pd.Series(dtype=float) + if solution_fn: + try: + m.write(path_to_string(solution_fn)) + except mindoptpy.MindoptError as err: + logger.info("No model solution stored. Raised error: %s", err) + + condition = m.status + termination_condition = CONDITION_MAP.get(condition, condition) + status = Status.from_termination_condition(termination_condition) + status.legacy_status = condition + + def get_solver_solution() -> Solution: + assert m is not None + objective = m.objval + + vars_ = m.getVars() + sol = _solution_from_names( + np.array([v.X for v in vars_], dtype=float), + [v.VarName for v in vars_], + self._n_vars, + ) - return Solution(sol, dual, objective) + try: + cons = m.getConstrs() + dual = _solution_from_names( + np.array([c.DualSoln for c in cons], dtype=float), + [c.ConstrName for c in cons], + self._n_cons, + ) + except (mindoptpy.MindoptError, AttributeError): + logger.warning("Dual values of MILP couldn't be parsed") + dual = np.array([], dtype=float) - solution = self.safe_get_solution(status=status, func=get_solver_solution) - solution = maybe_adjust_objective_sign(solution, io_api, sense) + return Solution(sol, dual, objective) - m.dispose() - env_.dispose() + solution = self.safe_get_solution(status=status, func=get_solver_solution) + solution = maybe_adjust_objective_sign(solution, io_api, sense) - return Result(status, solution, m) + self.io_api = io_api + return self._make_result(status, solution, solver_model=m) + finally: + if m is not None: + m.dispose() + env_.dispose() class PIPS(Solver[None]): @@ -2368,11 +4062,7 @@ class PIPS(Solver[None]): Solver subclass for the PIPS solver. """ - def __init__( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + def __post_init__(self) -> None: msg = "The PIPS solver interface is not yet implemented." raise NotImplementedError(msg) @@ -2397,48 +4087,36 @@ class cuPDLPx(Solver[None]): options for the given solver """ - def __init__( - self, - **solver_options: Any, - ) -> None: - super().__init__(**solver_options) + display_name: ClassVar[str] = "cuPDLPx" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + { + SolverFeature.DIRECT_API, + SolverFeature.GPU_ACCELERATION, + SolverFeature.GPU_ONLY, + SolverFeature.SOLUTION_FILE_NOT_NEEDED, + } + ) - def solve_problem_from_file( + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _has_module("cupdlpx") + + @classmethod + def _license_probe(cls) -> None: + cupdlpx.Model(np.array([0.0]), np.array([[0.0]]), None, None) + + def _run_file( self, - problem_fn: Path, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, env: EnvType | None = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem from a problem file using the solver cuPDLPx. - cuPDLPx does not currently support its own file IO, so this function - reads the problem file using linopy (only support netcf files) and - then passes the model to cuPDLPx for solving. - If the solution is feasible the function returns the - objective, solution and dual constraint variables. - - Parameters - ---------- - problem_fn : Path - Path to the problem file. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - - Returns - ------- - Result - """ + problem_fn = self._problem_fn + assert problem_fn is not None logger.warning( "cuPDLPx doesn't currently support file IO. Building model from file using linopy." ) @@ -2450,8 +4128,9 @@ def solve_problem_from_file( msg = "linopy currently only supports reading models from netcdf files. Try using io_api='direct' instead." raise NotImplementedError(msg) - return self.solve_problem_from_model( - model, + self.model = model + self._build_direct() + return self._run_direct( solution_fn=solution_fn, log_fn=log_fn, warmstart_fn=warmstart_fn, @@ -2459,64 +4138,85 @@ def solve_problem_from_file( env=env, ) - def solve_problem_from_model( + def _build_direct(self, **kwargs: Any) -> None: + model = self.model + assert model is not None + if model.type in ["QP", "MILP"]: + msg = "cuPDLPx does not currently support QP or MILP problems." + raise NotImplementedError(msg) + if kwargs.get("explicit_coordinate_names"): + warnings.warn( + "cuPDLPx does not support named variables/constraints. " + "The explicit_coordinate_names parameter is ignored.", + UserWarning, + stacklevel=2, + ) + cu_model = self._build_solver_model(model) + self.solver_model = cu_model + self.io_api = "direct" + self.sense = model.sense + self._cache_model_labels(model) + + @staticmethod + def _build_solver_model(model: Model) -> cupdlpx.Model: + """Build a cupdlpx.Model that mirrors the linopy `model`.""" + if model.variables.semi_continuous: + raise NotImplementedError( + "Semi-continuous variables are not supported by cuPDLPx. " + "Use a solver that supports them (gurobi, cplex, highs)." + ) + + M = model.matrices + if M.A is None: + raise ValueError("Model has no constraints, cannot export to cuPDLPx.") + A = M.A.tocsr() + lower = np.where( + np.logical_or(np.equal(M.sense, ">"), np.equal(M.sense, "=")), + M.b, + -np.inf, + ) + upper = np.where( + np.logical_or(np.equal(M.sense, "<"), np.equal(M.sense, "=")), + M.b, + np.inf, + ) + + cu_model = cupdlpx.Model( + objective_vector=M.c, + constraint_matrix=A, + constraint_lower_bound=lower, + constraint_upper_bound=upper, + variable_lower_bound=M.lb, + variable_upper_bound=M.ub, + ) + + if model.objective.sense == "max": + cu_model.ModelSense = cupdlpx.PDLP.MAXIMIZE + + return cu_model + + def _run_direct( self, - model: Model, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, basis_fn: Path | None = None, - env: EnvType | None = None, - explicit_coordinate_names: bool = False, + env: Any = None, + **kw: Any, ) -> Result: - """ - Solve a linear problem directly from a linopy model using the solver cuPDLPx. - If the solution is feasible the function returns the - objective, solution and dual constraint variables. - - Parameters - ---------- - model : linopy.model - Linopy model for the problem. - solution_fn : Path, optional - Path to the solution file. - log_fn : Path, optional - Path to the log file. - warmstart_fn : Path, optional - Path to the warmstart file. - basis_fn : Path, optional - Path to the basis file. - env : None, optional - Environment for the solver - explicit_coordinate_names : bool, optional - Transfer variable and constraint names to the solver (default: False) - - Returns - ------- - Result - """ - - if model.type in ["QP", "MILP"]: - msg = "cuPDLPx does not currently support QP or MILP problems." - raise NotImplementedError(msg) - - cu_model = model.to_cupdlpx() - return self._solve( - cu_model, - l_model=model, + self.solver_model, solution_fn=solution_fn, log_fn=log_fn, warmstart_fn=warmstart_fn, basis_fn=basis_fn, - io_api="direct", - sense=model.sense, + io_api=self.io_api, + sense=self.sense, ) def _solve( self, cu_model: cupdlpx.Model, - l_model: Model | None = None, solution_fn: Path | None = None, log_fn: Path | None = None, warmstart_fn: Path | None = None, @@ -2593,23 +4293,31 @@ def _solve( def get_solver_solution() -> Solution: objective = cu_model.ObjVal - - vlabels = None if l_model is None else l_model.matrices.vlabels - clabels = None if l_model is None else l_model.matrices.clabels - - sol = pd.Series(cu_model.X, vlabels, dtype=float) - dual = pd.Series(cu_model.Pi, clabels, dtype=float) + sol = np.asarray(cu_model.X, dtype=float) + dual = np.asarray(cu_model.Pi, dtype=float) if cu_model.ModelSense == cupdlpx.PDLP.MAXIMIZE: - dual *= -1 # flip sign of duals for max problems + dual = -dual + + sol = _solution_from_labels(sol, self._vlabels, self._n_vars) + dual = _solution_from_labels(dual, self._clabels, self._n_cons) return Solution(sol, dual, objective) solution = self.safe_get_solution(status=status, func=get_solver_solution) solution = maybe_adjust_objective_sign(solution, io_api, sense) - # see https://github.com/MIT-Lu-Lab/cuPDLPx/tree/main/python#solution-attributes - return Result(status, solution, cu_model) + runtime: float | None = None + with contextlib.suppress(Exception): + runtime = float(cu_model.Runtime) + + self.io_api = io_api + return self._make_result( + status, + solution, + solver_model=cu_model, + report=SolverReport(runtime=runtime), + ) def _set_solver_params(self, cu_model: cupdlpx.Model) -> None: """ @@ -2620,3 +4328,140 @@ def _set_solver_params(self, cu_model: cupdlpx.Model) -> None: """ for k, v in self.solver_options.items(): cu_model.setParam(k, v) + + +def _solver_class_for(name: str) -> type[Solver] | None: + try: + return globals().get(SolverName(name).name) + except ValueError: + return None + + +QUADRATIC_SOLVERS = [ + n.value + for n in SolverName + if (cls := _solver_class_for(n.value)) is not None + and cls.supports(SolverFeature.QUADRATIC_OBJECTIVE) +] +NO_SOLUTION_FILE_SOLVERS = [ + n.value + for n in SolverName + if (cls := _solver_class_for(n.value)) is not None + and cls.supports(SolverFeature.SOLUTION_FILE_NOT_NEEDED) +] + + +# Defines the iteration order of ``available_solvers`` — the first installed +# entry is the default solver in :meth:`Model.solve`. Matches the historical +# eager-probe order from before lazy availability landed. +_SOLVER_PROBE_ORDER: tuple[str, ...] = ( + "gurobi", + "highs", + "glpk", + "cbc", + "scip", + "cplex", + "xpress", + "knitro", + "mosek", + "mindopt", + "copt", + "cupdlpx", + "pips", +) + + +class _AvailableSolvers(Sequence[str]): + """ + Lazy sequence of installed solver names. + + Probes each solver's :meth:`Solver.is_available` on first access and caches + the result. Membership means the solver's Python package or binary is + importable — it does **not** mean a working license exists. Call + :func:`check_solver_licenses` for an opt-in eager license probe. + + :meth:`refresh` clears the cache (and each per-class ``is_available`` + cache) so the probe re-runs. + """ + + _filter: ClassVar[frozenset[str] | None] = None + + @functools.cached_property + def _names(self) -> list[str]: + names: list[str] = [] + for name in _SOLVER_PROBE_ORDER: + if self._filter is not None and name not in self._filter: + continue + cls = _solver_class_for(name) + if cls is not None and cls.is_available(): + names.append(name) + return names + + def __contains__(self, item: object) -> bool: + return item in self._names + + def __iter__(self) -> Iterator[str]: + return iter(self._names) + + def __len__(self) -> int: + return len(self._names) + + def __getitem__(self, idx: int | slice) -> Any: + return self._names[idx] + + def __repr__(self) -> str: + return repr(self._names) + + def __bool__(self) -> bool: + return bool(self._names) + + def refresh(self) -> None: + self.__dict__.pop("_names", None) + seen: set[int] = set() + for name in _SOLVER_PROBE_ORDER: + cls = _solver_class_for(name) + if cls is None: + continue + fn = cls.__dict__.get("is_available") + if fn is None: + continue + cache_clear = getattr(fn, "cache_clear", None) + if cache_clear is not None and id(fn) not in seen: + cache_clear() + seen.add(id(fn)) + + +class _QuadraticSolvers(_AvailableSolvers): + _filter: ClassVar[frozenset[str] | None] = frozenset(QUADRATIC_SOLVERS) + + +class _LicensedSolvers(_AvailableSolvers): + """Installed solvers whose ``license_status()`` probe currently succeeds.""" + + @functools.cached_property + def _names(self) -> list[str]: + names: list[str] = [] + for name in _SOLVER_PROBE_ORDER: + cls = _solver_class_for(name) + if cls is None or not cls.is_available(): + continue + if cls.license_status().ok: + names.append(name) + return names + + +available_solvers = _AvailableSolvers() +quadratic_solvers = _QuadraticSolvers() +licensed_solvers = _LicensedSolvers() + + +def check_solver_licenses(*names: str) -> dict[str, LicenseStatus]: + """Probe license status for the given solvers, or all installed ones.""" + targets = names or tuple(available_solvers) + out: dict[str, LicenseStatus] = {} + for n in targets: + cls = _solver_class_for(n) + if cls is None: + raise ValueError(f"unknown solver: {n!r}") + out[n] = cls.license_status() + return out diff --git a/linopy/sos_reformulation.py b/linopy/sos_reformulation.py new file mode 100644 index 000000000..4abfb7552 --- /dev/null +++ b/linopy/sos_reformulation.py @@ -0,0 +1,371 @@ +""" +SOS constraint reformulation using Big-M method. + +Converts SOS1/SOS2 constraints to binary + linear constraints for solvers +that don't support them natively. +""" + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal + +import numpy as np +import pandas as pd + +from linopy.constants import SOS_BIG_M_ATTR, SOS_DIM_ATTR, SOS_TYPE_ATTR + +if TYPE_CHECKING: + from xarray import DataArray + + from linopy.model import Model + from linopy.variables import Variable + +logger = logging.getLogger(__name__) + + +@dataclass +class SOSReformulationResult: + """Tracks what was added/changed during SOS reformulation for undo.""" + + reformulated: list[str] = field(default_factory=list) + added_variables: list[str] = field(default_factory=list) + added_constraints: list[str] = field(default_factory=list) + saved_attrs: dict[str, dict] = field(default_factory=dict) + + +def compute_big_m_values(var: Variable) -> DataArray: + """ + Compute Big-M values from variable bounds and custom big_m attribute. + + Uses the tighter of variable upper bound and custom big_m to ensure + the best possible LP relaxation. + + Parameters + ---------- + var : Variable + Variable with bounds (and optionally big_m_upper attr). + + Returns + ------- + DataArray + M_upper for reformulation constraints: x <= M_upper * y + + Raises + ------ + ValueError + If variable has negative lower bounds or infinite upper bounds. + """ + # SOS reformulation requires non-negative variables + if (var.lower < 0).any(): + raise ValueError( + f"Variable '{var.name}' has negative lower bounds. " + "SOS reformulation requires non-negative variables (lower >= 0)." + ) + + big_m_upper = var.attrs.get(SOS_BIG_M_ATTR) + M_upper = var.upper + + if big_m_upper is not None: + M_upper = M_upper.clip(max=big_m_upper) # type: ignore[arg-type] + + # Validate finiteness + if np.isinf(M_upper).any(): + raise ValueError( + f"Variable '{var.name}' has infinite upper bounds. " + "Set finite bounds or specify big_m in add_sos_constraints()." + ) + + return M_upper + + +def reformulate_sos1( + model: Model, var: Variable, prefix: str, M: DataArray | None = None +) -> tuple[list[str], list[str]]: + """ + Reformulate SOS1 constraint as binary + linear constraints. + + For each x[i] with upper bound M[i]: + - Add binary indicator y[i] + - x[i] <= M[i] * y[i] + - sum(y) <= 1 + + Parameters + ---------- + model : Model + Model to add reformulation constraints to. + var : Variable + Variable with SOS1 constraint (must have non-negative lower bounds). + prefix : str + Prefix for naming auxiliary variables and constraints. + M : DataArray, optional + Precomputed Big-M values. Computed from variable bounds if not provided. + + Returns + ------- + tuple[list[str], list[str]] + Names of added variables and constraints. + """ + if M is None: + M = compute_big_m_values(var) + sos_dim = str(var.attrs[SOS_DIM_ATTR]) + name = var.name + + y_name = f"{prefix}{name}_y" + upper_name = f"{prefix}{name}_upper" + card_name = f"{prefix}{name}_card" + + coords = [var.indexes[d] for d in var.dims] + y = model.add_variables(coords=coords, name=y_name, binary=True) + + model.add_constraints(var <= M * y, name=upper_name) + model.add_constraints(y.sum(dim=sos_dim) <= 1, name=card_name) + + return [y_name], [upper_name, card_name] + + +def reformulate_sos2( + model: Model, var: Variable, prefix: str, M: DataArray | None = None +) -> tuple[list[str], list[str]]: + """ + Reformulate SOS2 constraint as binary + linear constraints. + + For ordered x[0..n-1] with upper bounds M[i]: + - Add n-1 binary segment indicators z[i] + - x[0] <= M[0] * z[0] + - x[i] <= M[i] * (z[i-1] + z[i]) for middle elements + - x[n-1] <= M[n-1] * z[n-2] + - sum(z) <= 1 + + Parameters + ---------- + model : Model + Model to add reformulation constraints to. + var : Variable + Variable with SOS2 constraint (must have non-negative lower bounds). + prefix : str + Prefix for naming auxiliary variables and constraints. + M : DataArray, optional + Precomputed Big-M values. Computed from variable bounds if not provided. + + Returns + ------- + tuple[list[str], list[str]] + Names of added variables and constraints. + """ + sos_dim = str(var.attrs[SOS_DIM_ATTR]) + name = var.name + n = var.sizes[sos_dim] + + if n <= 1: + return [], [] + + if M is None: + M = compute_big_m_values(var) + + z_name = f"{prefix}{name}_z" + first_name = f"{prefix}{name}_upper_first" + last_name = f"{prefix}{name}_upper_last" + card_name = f"{prefix}{name}_card" + + z_coords = [ + pd.Index(var.indexes[sos_dim][:-1], name=sos_dim) + if d == sos_dim + else var.indexes[d] + for d in var.dims + ] + z = model.add_variables(coords=z_coords, name=z_name, binary=True) + + x_expr, z_expr = 1 * var, 1 * z + + added_constraints = [first_name] + + model.add_constraints( + x_expr.isel({sos_dim: 0}) <= M.isel({sos_dim: 0}) * z_expr.isel({sos_dim: 0}), + name=first_name, + ) + + if n > 2: + mid_slice = slice(1, n - 1) + x_mid = x_expr.isel({sos_dim: mid_slice}) + M_mid = M.isel({sos_dim: mid_slice}) + + z_left_coords = var.coords[sos_dim].values[: n - 2] + z_right_coords = var.coords[sos_dim].values[1 : n - 1] + + z_left = z_expr.sel({sos_dim: z_left_coords}) + z_right = z_expr.sel({sos_dim: z_right_coords}) + + z_left_aligned = z_left.assign_coords({sos_dim: M_mid.coords[sos_dim].values}) + z_right_aligned = z_right.assign_coords({sos_dim: M_mid.coords[sos_dim].values}) + + mid_name = f"{prefix}{name}_upper_mid" + model.add_constraints( + x_mid <= M_mid * (z_left_aligned + z_right_aligned), + name=mid_name, + ) + added_constraints.append(mid_name) + + model.add_constraints( + x_expr.isel({sos_dim: n - 1}) + <= M.isel({sos_dim: n - 1}) * z_expr.isel({sos_dim: n - 2}), + name=last_name, + ) + added_constraints.extend([last_name, card_name]) + + model.add_constraints(z.sum(dim=sos_dim) <= 1, name=card_name) + + return [z_name], added_constraints + + +def reformulate_sos_constraints( + model: Model, prefix: str = "_sos_reform_" +) -> SOSReformulationResult: + """ + Reformulate SOS constraints as binary + linear constraints. + + This converts SOS1 and SOS2 constraints into equivalent binary variable + formulations using the Big-M method. This allows solving models with SOS + constraints using solvers that don't support them natively (e.g., HiGHS, GLPK). + + Big-M values are determined as follows: + 1. If custom big_m was specified in add_sos_constraints(), use that + 2. Otherwise, use the variable bounds (tightest valid Big-M) + + Note: This permanently mutates the model and returns a token the caller + owns. For a stateful, reversible API use ``model.apply_sos_reformulation()`` + / ``model.undo_sos_reformulation()``; for automatic undo around a single + solve use ``model.solve(reformulate_sos=True)``. + + Parameters + ---------- + model : Model + Model containing SOS constraints to reformulate. + prefix : str, optional + Prefix for auxiliary variables and constraints. Default: "_sos_reform_" + + Returns + ------- + SOSReformulationResult + Tracks what was changed, enabling undo via ``undo_sos_reformulation``. + """ + result = SOSReformulationResult() + + try: + for var_name in list(model.variables.sos): + var = model.variables[var_name] + sos_type = var.attrs[SOS_TYPE_ATTR] + sos_dim = var.attrs[SOS_DIM_ATTR] + + if var.sizes[sos_dim] <= 1: + result.saved_attrs[var_name] = dict(var.attrs) + model.remove_sos_constraints(var) + result.reformulated.append(var_name) + continue + + M = compute_big_m_values(var) + if (M == 0).all(): + result.saved_attrs[var_name] = dict(var.attrs) + model.remove_sos_constraints(var) + result.reformulated.append(var_name) + continue + + result.saved_attrs[var_name] = dict(var.attrs) + + sort_idx = np.argsort(var.coords[sos_dim].values) + if not np.all(sort_idx[:-1] <= sort_idx[1:]): + sorted_var = var.isel({sos_dim: sort_idx}) + M = M.isel({sos_dim: sort_idx}) + else: + sorted_var = var + + if sos_type == 1: + added_vars, added_cons = reformulate_sos1(model, sorted_var, prefix, M) + elif sos_type == 2: + added_vars, added_cons = reformulate_sos2(model, sorted_var, prefix, M) + else: + raise ValueError( + f"Unknown sos_type={sos_type} on variable '{var_name}'" + ) + + result.added_variables.extend(added_vars) + result.added_constraints.extend(added_cons) + + model.remove_sos_constraints(var) + result.reformulated.append(var_name) + except Exception: + undo_sos_reformulation(model, result) + raise + + logger.info(f"Reformulated {len(result.reformulated)} SOS constraint(s)") + return result + + +def undo_sos_reformulation(model: Model, result: SOSReformulationResult) -> None: + """ + Undo a previous SOS reformulation, restoring the model to its original state. + + Parameters + ---------- + model : Model + Model that was reformulated. + result : SOSReformulationResult + Result from ``reformulate_all_sos`` tracking what was added. + """ + objective_value = model.objective._value + + for con_name in result.added_constraints: + if con_name in model.constraints: + model.remove_constraints(con_name) + + for var_name in result.added_variables: + if var_name in model.variables: + model.remove_variables(var_name) + + for var_name, attrs in result.saved_attrs.items(): + if var_name in model.variables: + model.variables[var_name].attrs.update(attrs) + + model.objective._value = objective_value + + +@contextmanager +def sos_reformulation_context( + model: Model, + solver_name: str | None, + reformulate_sos: bool | Literal["auto"], +) -> Iterator[bool]: + """ + Apply SOS reformulation for the duration of the block, then undo. + + Yields whether the reformulation was actually applied, so callers can + branch on it (e.g. to scope a warning suppression). + """ + applied = model._resolve_sos_reformulation(solver_name, reformulate_sos) + if applied: + logger.info(f"Reformulating SOS constraints for solver {solver_name}") + model.apply_sos_reformulation() + try: + yield applied + finally: + if applied: + model.undo_sos_reformulation() + + +@contextmanager +def suppress_serialization_warning(active: bool) -> Iterator[None]: + """Silence the SOS-active-on-serialize UserWarning when ``active`` is True.""" + if not active: + yield + return + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="Serializing a model with an active SOS reformulation", + category=UserWarning, + ) + yield diff --git a/linopy/testing.py b/linopy/testing.py index 0392064ed..310822846 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -1,13 +1,29 @@ from __future__ import annotations +import numpy as np from xarray.testing import assert_equal -from linopy.constraints import Constraint, _con_unwrap +from linopy.constants import TERM_DIM +from linopy.constraints import ConstraintBase, _con_unwrap from linopy.expressions import LinearExpression, QuadraticExpression, _expr_unwrap from linopy.model import Model from linopy.variables import Variable, _var_unwrap +def _sort_by_vars_along_term(expr: LinearExpression) -> LinearExpression: + """Sort a linear expression's terms by variable labels along _term.""" + ds = expr.data + if TERM_DIM not in ds.dims: + return expr + order = np.argsort(ds["vars"].values, axis=-1, kind="stable") + sorted_vars = np.take_along_axis(ds["vars"].values, order, axis=-1) + sorted_coeffs = np.take_along_axis(ds["coeffs"].values, order, axis=-1) + new_ds = ds.copy() + new_ds["vars"] = (ds["vars"].dims, sorted_vars) + new_ds["coeffs"] = (ds["coeffs"].dims, sorted_coeffs) + return LinearExpression(new_ds, expr.model) + + def assert_varequal(a: Variable, b: Variable) -> None: """Assert that two variables are equal.""" return assert_equal(_var_unwrap(a), _var_unwrap(b)) @@ -16,10 +32,18 @@ def assert_varequal(a: Variable, b: Variable) -> None: def assert_linequal( a: LinearExpression | QuadraticExpression, b: LinearExpression | QuadraticExpression ) -> None: - """Assert that two linear expressions are equal.""" + """ + Assert that two linear expressions are semantically equal. + + Terms are sorted by variable labels along _term before comparing, + so expressions with different term orderings but identical mathematical + meaning are considered equal. + """ assert isinstance(a, LinearExpression) assert isinstance(b, LinearExpression) - return assert_equal(_expr_unwrap(a), _expr_unwrap(b)) + a_sorted = _sort_by_vars_along_term(a) + b_sorted = _sort_by_vars_along_term(b) + return assert_equal(_expr_unwrap(a_sorted), _expr_unwrap(b_sorted)) def assert_quadequal( @@ -29,7 +53,7 @@ def assert_quadequal( return assert_equal(_expr_unwrap(a), _expr_unwrap(b)) -def assert_conequal(a: Constraint, b: Constraint, strict: bool = True) -> None: +def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) -> None: """ Assert that two constraints are equal. diff --git a/linopy/types.py b/linopy/types.py index 68e5c3079..6b4cf712d 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -2,7 +2,7 @@ from collections.abc import Hashable, Iterable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, TypeAlias, Union, get_args import numpy import polars as pl @@ -11,7 +11,10 @@ from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import ( + AnonymousScalarConstraint, + ConstraintBase, + ) from linopy.expressions import ( LinearExpression, QuadraticExpression, @@ -19,34 +22,32 @@ ) from linopy.variables import ScalarVariable, Variable -# Type aliases using Union for Python 3.9 compatibility -CoordsLike = Union[ # noqa: UP007 - Sequence[Sequence | Index | DataArray], - Mapping, - DataArrayCoordinates, - DatasetCoordinates, -] -DimsLike = Union[str, Iterable[Hashable]] # noqa: UP007 +CoordsLike: TypeAlias = ( + Sequence[Sequence | Index] | Mapping | DataArrayCoordinates | DatasetCoordinates +) +DimsLike: TypeAlias = str | Iterable[Hashable] -ConstantLike = Union[ # noqa: UP007 - int, - float, - numpy.floating, - numpy.integer, - numpy.ndarray, - DataArray, - Series, - DataFrame, - pl.Series, -] -SignLike = Union[str, numpy.ndarray, DataArray, Series, DataFrame] # noqa: UP007 -VariableLike = Union["ScalarVariable", "Variable"] -ExpressionLike = Union[ - "ScalarLinearExpression", - "LinearExpression", - "QuadraticExpression", +ConstantLike: TypeAlias = ( + int + | float + | numpy.floating + | numpy.integer + | numpy.ndarray + | DataArray + | Series + | DataFrame + | pl.Series +) +CONSTANT_TYPES: tuple[type, ...] = get_args(ConstantLike) +SignLike: TypeAlias = str | numpy.ndarray | DataArray | Series | DataFrame +MaskLike: TypeAlias = numpy.ndarray | DataArray | Series | DataFrame +PathLike: TypeAlias = str | Path + +# These reference types only available under TYPE_CHECKING, so use Union with strings +VariableLike: TypeAlias = Union["ScalarVariable", "Variable"] +ExpressionLike: TypeAlias = Union[ + "ScalarLinearExpression", "LinearExpression", "QuadraticExpression" ] -ConstraintLike = Union["Constraint", "AnonymousScalarConstraint"] -MaskLike = Union[numpy.ndarray, DataArray, Series, DataFrame] # noqa: UP007 +ConstraintLike = Union["ConstraintBase", "AnonymousScalarConstraint"] +LinExprLike = Union["Variable", "LinearExpression"] SideLike = Union[ConstantLike, VariableLike, ExpressionLike] # noqa: UP007 -PathLike = Union[str, Path] # noqa: UP007 diff --git a/linopy/variables.py b/linopy/variables.py index e2570b5de..0e87aec18 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -14,6 +14,7 @@ from typing import ( TYPE_CHECKING, Any, + cast, overload, ) from warnings import warn @@ -26,34 +27,42 @@ from xarray import DataArray, Dataset, broadcast from xarray.core.coordinates import DatasetCoordinates from xarray.core.indexes import Indexes +from xarray.core.types import JoinOptions from xarray.core.utils import Frozen import linopy.expressions as expressions +from linopy.alignment import broadcast_to_coords from linopy.common import ( LabelPositionIndex, LocIndexer, - as_dataarray, + VariableLabelIndex, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, filter_nulls_polars, + format_coord, + format_single_variable, format_string_as_variable_name, generate_indices_for_printout, get_dims_with_index_levels, get_label_position, has_optimized_model, iterate_slices, - print_coord, - print_single_variable, - require_constant, save_join, set_int_index, to_dataframe, to_polars, ) from linopy.config import options -from linopy.constants import HELPER_DIMS, TERM_DIM -from linopy.solver_capabilities import SolverFeature, solver_supports +from linopy.constants import ( + HELPER_DIMS, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, + STASHED_ATTRS, + STASHED_LOWER, + STASHED_UPPER, + TERM_DIM, +) from linopy.types import ( ConstantLike, DimsLike, @@ -195,10 +204,10 @@ def __init__( if "label_range" not in data.attrs: data.assign_attrs(label_range=(data.labels.min(), data.labels.max())) - if "sos_type" in data.attrs or "sos_dim" in data.attrs: - if (sos_type := data.attrs.get("sos_type")) not in (1, 2): + if SOS_TYPE_ATTR in data.attrs or SOS_DIM_ATTR in data.attrs: + if (sos_type := data.attrs.get(SOS_TYPE_ATTR)) not in (1, 2): raise ValueError(f"sos_type must be 1 or 2, got {sos_type}") - if (sos_dim := data.attrs.get("sos_dim")) not in data.dims: + if (sos_dim := data.attrs.get(SOS_DIM_ATTR)) not in data.dims: raise ValueError( f"sos_dim must name a variable dimension, got {sos_dim}" ) @@ -290,9 +299,15 @@ def at(self) -> AtIndexer: @property def loc(self) -> LocIndexer: + """ + Indexing the variable using coordinates. + """ return LocIndexer(self) def to_pandas(self) -> pd.Series: + """ + Convert the variable labels to a pandas Series. + """ return self.labels.to_pandas() def to_linexpr( @@ -313,7 +328,11 @@ def to_linexpr( linopy.LinearExpression Linear expression with the variables and coefficients. """ - coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) + coefficient = broadcast_to_coords( + coefficient, coords=self.coords, dims=self.dims, strict=False + ) + coefficient = coefficient.reindex_like(self.labels, fill_value=0) + coefficient = coefficient.fillna(0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( TERM_DIM, -1 ) @@ -328,8 +347,8 @@ def __repr__(self) -> str: dim_names = self.coord_names dim_sizes = list(self.sizes.values()) masked_entries = (~self.mask).sum().values - sos_type = self.attrs.get("sos_type") - sos_dim = self.attrs.get("sos_dim") + sos_type = self.attrs.get(SOS_TYPE_ATTR) + sos_dim = self.attrs.get(SOS_DIM_ATTR) lines = [] if dims: @@ -342,9 +361,9 @@ def __repr__(self) -> str: ] label = self.labels.values[indices] line = ( - print_coord(coord) + format_coord(coord) + ": " - + print_single_variable(self.model, label) + + format_single_variable(self.model, label) ) lines.append(line) # lines = align_lines_by_delimiter(lines, "∈") @@ -359,7 +378,7 @@ def __repr__(self) -> str: ) else: lines.append( - f"Variable\n{'-' * 8}\n{print_single_variable(self.model, self.labels.item())}" + f"Variable\n{'-' * 8}\n{format_single_variable(self.model, self.labels.item())}" ) return "\n".join(lines) @@ -420,7 +439,9 @@ def __pow__(self, other: int) -> QuadraticExpression: return NotImplemented if other == 2: expr = self.to_linexpr() - return expr._multiply_by_linear_expression(expr) + return cast( + "QuadraticExpression", expr._multiply_by_linear_expression(expr) + ) raise ValueError("Can only raise to the power of 2") @overload @@ -440,7 +461,7 @@ def __matmul__( return self.to_linexpr() @ other def __div__( - self, other: float | int | LinearExpression | Variable + self, other: ConstantLike | LinearExpression | Variable ) -> LinearExpression: """ Divide variables with a coefficient. @@ -451,10 +472,10 @@ 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 + self, coefficient: ConstantLike | LinearExpression | Variable ) -> LinearExpression: """ True divide variables with a coefficient. @@ -525,7 +546,7 @@ def __le__(self, other: SideLike) -> Constraint: def __ge__(self, other: SideLike) -> Constraint: return self.to_linexpr().__ge__(other) - def __eq__(self, other: SideLike) -> Constraint: # type: ignore + def __eq__(self, other: SideLike) -> Constraint: # type: ignore[override] return self.to_linexpr().__eq__(other) def __gt__(self, other: Any) -> NotImplementedType: @@ -541,29 +562,118 @@ 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: JoinOptions | None = None + ) -> LinearExpression | QuadraticExpression: """ 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: JoinOptions | None = None + ) -> LinearExpression | QuadraticExpression: """ 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: JoinOptions | None = None + ) -> LinearExpression | QuadraticExpression: """ 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: JoinOptions | None = None + ) -> LinearExpression | QuadraticExpression: """ 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.to_linexpr().div(other, join=join) + + def le(self, rhs: SideLike, join: JoinOptions | 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.__div__(other) + return self.to_linexpr().le(rhs, join=join) + + def ge(self, rhs: SideLike, join: JoinOptions | 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: JoinOptions | 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: """ @@ -726,15 +836,23 @@ def type(self) -> str: return "Integer Variable" elif self.attrs["binary"]: return "Binary Variable" + elif self.attrs.get("semi_continuous"): + return "Semi-continuous Variable" else: return "Continuous Variable" @property def coord_dims(self) -> tuple[Hashable, ...]: + """ + Get the coordinate dimensions of the variable. + """ return tuple(k for k in self.dims if k not in HELPER_DIMS) @property def coord_sizes(self) -> dict[Hashable, int]: + """ + Get the coordinate sizes of the variable. + """ return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} @property @@ -776,18 +894,18 @@ def upper(self) -> DataArray: return self.data.upper @upper.setter - @require_constant def upper(self, value: ConstantLike) -> None: """ - Set the upper bounds of the variables. - - The function raises an error in case no model is set as a - reference. + Syntactic sugar for :meth:`Variable.update`. Do not add logic + here; mutate via ``update`` so the contract stays single-sourced. """ - value = DataArray(value).broadcast_like(self.upper) - if not set(value.dims).issubset(self.model.variables[self.name].dims): - raise ValueError("Cannot assign new dimensions to existing variable.") - self._data = assign_multiindex_safe(self.data, upper=value) + warn( + "Variable.upper setter is deprecated and will be removed in a " + "future release; use Variable.update(upper=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(upper=value) @property def lower(self) -> DataArray: @@ -800,18 +918,100 @@ def lower(self) -> DataArray: return self.data.lower @lower.setter - @require_constant def lower(self, value: ConstantLike) -> None: """ - Set the lower bounds of the variables. + Syntactic sugar for :meth:`Variable.update`. Do not add logic + here; mutate via ``update`` so the contract stays single-sourced. + """ + warn( + "Variable.lower setter is deprecated and will be removed in a " + "future release; use Variable.update(lower=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + self.update(lower=value) - The function raises an error in case no model is set as a - reference. + def update( + self, + *, + lower: ConstantLike | None = None, + upper: ConstantLike | None = None, + ) -> Variable: """ - value = DataArray(value).broadcast_like(self.lower) - if not set(value.dims).issubset(self.model.variables[self.name].dims): - raise ValueError("Cannot assign new dimensions to existing variable.") - self._data = assign_multiindex_safe(self.data, lower=value) + Update variable bounds in place. + + Canonical mutation API. Validation and coord alignment live here. + Single-attribute setters (`var.lower = …`) forward to this method. + + Parameters + ---------- + lower : ConstantLike, optional + New lower bound. Accepts any constant — scalars, numpy + arrays, pandas Series / DataFrame, xarray DataArray (e.g. + time-varying bounds). Aligned via xarray broadcast against + the variable's existing shape; new dims are rejected. + Decision variables / linear expressions are not accepted. + upper : ConstantLike, optional + New upper bound. Same. + + Returns + ------- + Variable + ``self`` for chaining. + + Raises + ------ + TypeError + If either bound is a Variable / Expression (bounds must be + numeric, not symbolic). + ValueError + If the new bound introduces dimensions not in the variable's + coords, or if the resulting ``lower > upper`` anywhere. + """ + if lower is None and upper is None: + return self + + updates = self._validate_update(lower=lower, upper=upper) + self._data = assign_multiindex_safe(self.data, **updates) + return self + + def _validate_update( + self, + *, + lower: ConstantLike | None = None, + upper: ConstantLike | None = None, + ) -> dict[str, DataArray]: + """ + Validate, broadcast, and cross-check update inputs. + + Returns the broadcasted DataArrays ready for assignment. Raises + before any mutation if any input is wrong. + """ + updates: dict[str, DataArray] = {} + own_dims = self.model.variables[self.name].dims + for name, val, ref in ( + ("lower", lower, self.lower), + ("upper", upper, self.upper), + ): + if val is None: + continue + if not isinstance(val, ConstantLike): + raise TypeError( + f"Variable.update({name}=...) must be a constant; " + f"got {type(val).__name__}." + ) + new_val = DataArray(val).broadcast_like(ref) + if not set(new_val.dims).issubset(own_dims): + raise ValueError("Cannot assign new dimensions to existing variable.") + updates[name] = new_val + + final_lower = updates.get("lower", self.lower) + final_upper = updates.get("upper", self.upper) + if bool((final_lower > final_upper).any()): + raise ValueError( + "Variable.update would leave lower > upper at one or more coordinates." + ) + return updates @property @has_optimized_model @@ -861,9 +1061,11 @@ def get_solver_attribute(self, attr: str) -> DataArray: ------- xr.DataArray """ + from linopy.solver_capabilities import SolverFeature, solver_supports + solver_model = self.model.solver_model if not solver_supports( - self.model.solver_name, SolverFeature.SOLVER_ATTRIBUTE_ACCESS + self.model.solver_name or "", SolverFeature.SOLVER_ATTRIBUTE_ACCESS ): raise NotImplementedError( "Solver attribute getter only supports the Gurobi solver for now." @@ -894,7 +1096,7 @@ def flat(self) -> DataFrame: ------- df : pandas.DataFrame """ - ds = self.data + ds = self.data.drop_vars(STASHED_ATTRS, errors="ignore") def mask_func(data: pd.DataFrame) -> pd.Series: return data["labels"] != -1 @@ -914,7 +1116,8 @@ def to_polars(self) -> pl.DataFrame: ------- pl.DataFrame """ - df = to_polars(self.data) + ds = self.data.drop_vars(STASHED_ATTRS, errors="ignore") + df = to_polars(ds) df = filter_nulls_polars(df) check_has_nulls_polars(df, name=f"{self.type} {self.name}") return df @@ -1020,6 +1223,9 @@ def where( raise ValueError( f"other must be a Variable, ScalarVariable, dict or Dataset, got {type(other)}" ) + if isinstance(_other, dict): + fill: dict[str, float] = {str(k): np.nan for k in self.data} + _other = {**fill, **_other} return self.__class__( self.data.where(cond, _other, **kwargs), self.model, self.name ) @@ -1066,7 +1272,9 @@ def ffill(self, dim: str, limit: None = None) -> Variable: .map(DataArray.ffill, dim=dim, limit=limit) .fillna(self._fill_value) ) - return self.assign_multiindex_safe(labels=data.labels.astype(int)) + return self.assign_multiindex_safe( + labels=data.labels.astype(options["label_dtype"]) + ) def bfill(self, dim: str, limit: None = None) -> Variable: """ @@ -1093,7 +1301,7 @@ def bfill(self, dim: str, limit: None = None) -> Variable: .map(DataArray.bfill, dim=dim, limit=limit) .fillna(self._fill_value) ) - return self.assign(labels=data.labels.astype(int)) + return self.assign(labels=data.labels.astype(options["label_dtype"])) def sanitize(self) -> Variable: """ @@ -1104,10 +1312,25 @@ def sanitize(self) -> Variable: linopy.Variable """ if issubdtype(self.labels.dtype, floating): - return self.assign(labels=self.labels.fillna(-1).astype(int)) + return self.assign( + labels=self.labels.fillna(-1).astype(options["label_dtype"]) + ) return self def equals(self, other: Variable) -> bool: + """ + Check if this Variable is equal to another. + + Parameters + ---------- + other : Variable + The Variable to compare with. + + Returns + ------- + bool + True if the variables have equal labels, False otherwise. + """ return self.labels.equals(other.labels) # Wrapped function which would convert variable to dataarray @@ -1151,6 +1374,163 @@ def equals(self, other: Variable) -> bool: iterate_slices = iterate_slices + def relax(self) -> None: + """ + Relax the integrality of this variable. + + Converts binary or integer variables to continuous. The original type + is stored in the model's ``_relaxed_registry`` so that + :meth:`unrelax` can restore it. + + Semi-continuous variables are not supported and will raise a + ``NotImplementedError``. + + For binary variables, the existing [0, 1] bounds are preserved, + which is the correct LP relaxation. For integer variables, the + existing bounds are preserved as-is. + + If the variable is already continuous, this method is a no-op. + """ + attrs = self.attrs + if attrs.get("semi_continuous"): + msg = ( + f"Relaxation of semi-continuous variable '{self.name}' is not " + "supported. The LP relaxation of a semi-continuous variable " + "requires changing bounds, which is not handled by relax()." + ) + raise NotImplementedError(msg) + + for attr in ("binary", "integer"): + if attrs.get(attr): + self.model._relaxed_registry[self.name] = attr + attrs[attr] = False + return + + def unrelax(self) -> None: + """ + Restore the original integrality type of a relaxed variable. + + Reverses the effect of :meth:`relax`. If the variable was not + previously relaxed, this is a no-op. + """ + registry = self.model._relaxed_registry + original_type = registry.pop(self.name, None) + if original_type is not None: + self.attrs[original_type] = True + + @property + def relaxed(self) -> bool: + """ + Return whether the variable is currently relaxed. + """ + return self.name in self.model._relaxed_registry + + def fix( + self, + value: ConstantLike | None = None, + decimals: int = 8, + overwrite: bool = True, + ) -> None: + """ + Fix the variable to a given value by collapsing its bounds. + + Sets ``lower = upper = value``. + + If no value is given, the current solution value is used. + + A fix value outside the variable's current bounds emits a warning, but + does not cause infeasibilities (the bounds are overridden). Fixing a + binary variable to anything other than 0 or 1 raises. + + Parameters + ---------- + value : float/array_like, optional + Value to fix the variable to. If None, the current solution is used. + decimals : int, optional + Number of decimal places to round continuous variables to. + Integer and binary variables are always rounded to 0 decimal places. + Default is 8. + overwrite : bool, optional + If True (default), re-fix a variable that is already fixed to the + new value (the originally stashed bounds are kept). If False, raise + an error if the variable is already fixed. + """ + if value is None: + try: + value = self.solution + except AttributeError: + msg = ( + f"Cannot fix variable '{self.name}': no solution value " + "available. Solve the model first or provide an explicit " + "value." + ) + raise ValueError(msg) from None + + is_fixed = self.fixed + is_binary = self.attrs["binary"] + is_integer = self.attrs["integer"] + + if is_fixed and not overwrite: + msg = ( + f"Variable '{self.name}' is already fixed. Use " + "overwrite=True to replace the existing fix value." + ) + raise ValueError(msg) + + value = broadcast_to_coords( + value, self.coords, label=f"fix() for variable '{self.name}'" + ) + + if is_binary and not (np.isclose(value, 0) | np.isclose(value, 1)).all(): + msg = ( + f"Cannot fix binary variable '{self.name}' to a value " + "other than 0 or 1." + ) + raise ValueError(msg) + + if is_integer or is_binary: + value = value.round(0) + else: + value = value.round(decimals) + + if is_fixed: + lower, upper = self.data[STASHED_LOWER], self.data[STASHED_UPPER] + else: + lower, upper = self.data.lower, self.data.upper + + if not is_binary and ((value < lower).any() or (value > upper).any()): + warn( + f"Fix values for variable '{self.name}' lie outside its current " + "bounds; the bounds are overridden by the fix value.", + UserWarning, + stacklevel=2, + ) + + if not is_fixed: + self._data = assign_multiindex_safe( + self.data, + **{STASHED_LOWER: lower, STASHED_UPPER: upper}, + ) + + self.update(lower=value, upper=value) + + def unfix(self) -> None: + """ + Unfix the variable, restoring the bounds it had before :meth:`fix`. + """ + if not self.fixed: + return + + self.update(lower=self.data[STASHED_LOWER], upper=self.data[STASHED_UPPER]) + self._data = self.data.drop_vars(STASHED_ATTRS) + + @property + def fixed(self) -> bool: + """ + Return whether the variable is currently fixed. + """ + return all(attr in self.data for attr in STASHED_ATTRS) + class AtIndexer: __slots__ = ("object",) @@ -1183,6 +1563,7 @@ class Variables: data: dict[str, Variable] model: Model _label_position_index: LabelPositionIndex | None = None + _variable_label_index: VariableLabelIndex | None = None dataset_attrs = ["labels", "lower", "upper"] dataset_names = ["Labels", "Lower bounds", "Upper bounds"] @@ -1230,29 +1611,40 @@ def __dir__(self) -> list[str]: ] return base_attributes + formatted_names - def __repr__(self) -> str: - """ - Return a string representation of the linopy model. - """ - r = "linopy.model.Variables" - line = "-" * len(r) - r += f"\n{line}\n" - + def _format_items(self, exclude: set[str] | None = None) -> str: + """Format variable items, optionally excluding names in a group.""" + r = "" + count = 0 for name, ds in self.items(): + if exclude and name in exclude: + continue + count += 1 coords = ( " (" + ", ".join(str(coord) for coord in ds.coords) + ")" if ds.coords else "" ) - if (sos_type := ds.attrs.get("sos_type")) in (1, 2) and ( - sos_dim := ds.attrs.get("sos_dim") + if (sos_type := ds.attrs.get(SOS_TYPE_ATTR)) in (1, 2) and ( + sos_dim := ds.attrs.get(SOS_DIM_ATTR) ): coords += f" - sos{sos_type} on {sos_dim}" + if ds.attrs.get("semi_continuous", False): + coords += " - semi-continuous" r += f" * {name}{coords}\n" - if not len(list(self)): + if count == 0: r += "\n" return r + def __repr__(self) -> str: + """ + Return a string representation of the variables container. + """ + r = "linopy.model.Variables" + line = "-" * len(r) + r += f"\n{line}\n" + r += self._format_items() + return r + def __len__(self) -> int: return self.data.__len__() @@ -1290,6 +1682,15 @@ def _invalidate_label_position_index(self) -> None: """Invalidate the label position index cache.""" if self._label_position_index is not None: self._label_position_index.invalidate() + if self._variable_label_index is not None: + self._variable_label_index.invalidate() + + @property + def label_index(self) -> VariableLabelIndex: + """Index for O(1) label->position mapping and compact vlabels array.""" + if self._variable_label_index is None: + self._variable_label_index = VariableLabelIndex(self) + return self._variable_label_index @property def attrs(self) -> dict[Any, Any]: @@ -1387,7 +1788,23 @@ def continuous(self) -> Variables: { name: self.data[name] for name in self - if not self[name].attrs["integer"] and not self[name].attrs["binary"] + if not self[name].attrs["integer"] + and not self[name].attrs["binary"] + and not self[name].attrs.get("semi_continuous", False) + }, + self.model, + ) + + @property + def semi_continuous(self) -> Variables: + """ + Get all semi-continuous variables. + """ + return self.__class__( + { + name: self.data[name] + for name in self + if self[name].attrs.get("semi_continuous", False) }, self.model, ) @@ -1401,12 +1818,87 @@ def sos(self) -> Variables: { name: self.data[name] for name in self - if self[name].attrs.get("sos_dim") - and self[name].attrs.get("sos_type") in (1, 2) + if self[name].attrs.get(SOS_DIM_ATTR) + and self[name].attrs.get(SOS_TYPE_ATTR) in (1, 2) }, self.model, ) + def fix( + self, + value: int | float | None = None, + decimals: int = 8, + overwrite: bool = True, + ) -> None: + """ + Fix all variables in this container to their solution or a scalar value. + + Delegates to each variable's ``fix()`` method. See + :meth:`Variable.fix` for details. + + Parameters + ---------- + value : int/float, optional + Scalar value to fix all variables to. Only scalar values are + accepted to avoid shape mismatches across differently-shaped + variables. If None, each variable is fixed to its current solution. + decimals : int, optional + Number of decimal places to round continuous variables to. + overwrite : bool, optional + If True, re-fix variables that are already fixed. + """ + for var in self.data.values(): + var.fix(value=value, decimals=decimals, overwrite=overwrite) + + def unfix(self) -> None: + """ + Unfix all variables in this container. + + Delegates to each variable's ``unfix()`` method. + """ + for var in self.data.values(): + var.unfix() + + @property + def fixed(self) -> Variables: + """ + Get all currently fixed variables. + """ + return self.__class__( + {name: self.data[name] for name in self if self[name].fixed}, + self.model, + ) + + def relax(self) -> None: + """ + Relax integrality of all integer/binary variables in this container. + + Delegates to each variable's :meth:`Variable.relax` method. + Semi-continuous variables will raise ``NotImplementedError``. + """ + for var in self.data.values(): + var.relax() + + def unrelax(self) -> None: + """ + Restore integrality of all previously relaxed variables in this + container. + + Delegates to each variable's :meth:`Variable.unrelax` method. + """ + for var in self.data.values(): + var.unrelax() + + @property + def relaxed(self) -> Variables: + """ + Get all currently relaxed variables. + """ + return self.__class__( + {name: self.data[name] for name in self if self[name].relaxed}, + self.model, + ) + @property def solution(self) -> Dataset: """ @@ -1499,17 +1991,36 @@ def get_label_position_with_index( self._label_position_index = LabelPositionIndex(self) return self._label_position_index.find_single_with_index(label) - def print_labels(self, values: list[int]) -> None: + def format_labels(self, values: list[int]) -> str: """ - Print a selection of labels of the variables. + Get a string representation of a selection of variable labels. Parameters ---------- values : list, array-like - One dimensional array of constraint labels. + One dimensional array of variable labels. + + Returns + ------- + str + String representation of the selected variables. + """ + res = [format_single_variable(self.model, v) for v in values] + return "\n".join(res) + + def print_labels(self, values: list[int]) -> None: """ - res = [print_single_variable(self.model, v) for v in values] - print("\n".join(res)) + Print a selection of labels of the variables. + + .. deprecated:: + Use :meth:`format_labels` instead. + """ + warn( + "`Variables.print_labels` is deprecated. Use `Variables.format_labels` instead.", + DeprecationWarning, + stacklevel=2, + ) + print(self.format_labels(values)) @property def flat(self) -> pd.DataFrame: @@ -1525,7 +2036,10 @@ def flat(self) -> pd.DataFrame: """ df = pd.concat([self[k].flat for k in self], ignore_index=True) unique_labels = df.labels.unique() - map_labels = pd.Series(np.arange(len(unique_labels)), index=unique_labels) + map_labels = pd.Series( + np.arange(len(unique_labels), dtype=options["label_dtype"]), + index=unique_labels, + ) df["key"] = df.labels.map(map_labels) return df @@ -1578,7 +2092,7 @@ def __repr__(self) -> str: if self.label == -1: return "ScalarVariable: None" name, coord = self.model.variables.get_label_position(self.label) - coord_string = print_coord(coord) + coord_string = format_coord(coord) return f"ScalarVariable: {name}{coord_string}" @property @@ -1652,7 +2166,7 @@ def __le__(self, other: int | float) -> AnonymousScalarConstraint: def __ge__(self, other: int) -> AnonymousScalarConstraint: return self.to_scalar_linexpr(1).__ge__(other) - def __eq__(self, other: int | float) -> AnonymousScalarConstraint: # type: ignore + def __eq__(self, other: int | float) -> AnonymousScalarConstraint: # type: ignore[override] return self.to_scalar_linexpr(1).__eq__(other) def __gt__(self, other: Any) -> None: diff --git a/pyproject.toml b/pyproject.toml index 52d5e3d5c..8ddbfa076 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,9 +8,12 @@ dynamic = ["version"] description = "Linear optimization with N-D labeled arrays in Python" readme = "README.md" authors = [{ name = "Fabian Hofmann", email = "fabianmarikhofmann@gmail.com" }] +maintainers = [ + { name = "Fabian Hofmann", email = "fabianmarikhofmann@gmail.com" }, + { name = "Felix Bumann", email = "dev@fxbumann.de" }, +] license = { file = "LICENSE" } classifiers = [ - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -23,22 +26,19 @@ classifiers = [ "Operating System :: OS Independent", ] -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ - "numpy; python_version > '3.10'", - "numpy<2; python_version <= '3.10'", + "numpy", "scipy", "bottleneck", "toolz", "numexpr", "xarray>=2024.2.0", "dask>=0.18.0", - "polars", + "polars>=1.31.1", "tqdm", "deprecation", "packaging", - "google-cloud-storage", - "requests", ] [project.urls] @@ -46,20 +46,27 @@ Homepage = "https://github.com/PyPSA/linopy" Source = "https://github.com/PyPSA/linopy" [project.optional-dependencies] +oetc = [ + "google-cloud-storage", + "requests", +] +remote = [ + "paramiko", +] docs = [ "ipython==8.26.0", - "numpydoc==1.7.0", - "sphinx==7.3.7", - "sphinx_rtd_theme==2.0.0", - "sphinx_book_theme==1.1.3", + "numpydoc==1.10.0", + "sphinx_rtd_theme==3.1.0", + "sphinx==9.0.4", + "sphinx_book_theme==1.2.0", "sphinx-copybutton==0.5.2", "nbsphinx==0.9.4", "nbsphinx-link==1.3.0", "docutils<0.21", "numpy<2", - "gurobipy==11.0.2", + "gurobipy>=13.0.0", "ipykernel==6.29.5", - "matplotlib==3.9.1", + "matplotlib==3.11.0", "highspy>=1.7.1", ] dev = [ @@ -73,17 +80,35 @@ dev = [ "types-requests", "gurobipy", "highspy", + "jupyter", +] +# Perf-relevant deps pinned exactly so run-to-run deltas reflect linopy +# changes, not dependency bumps. +benchmarks = [ + "highspy==1.13.1", + "netcdf4==1.7.4", + "numpy==2.4.6", + "scipy==1.17.1", + "xarray==2026.4.0", + "pandas==3.0.3", + "polars==1.41.2", + "dask==2026.6.0", + "pytest==9.1.1", + "pytest-benchmark==5.2.3", + "pytest-memray==1.8.0", + "pytest-codspeed==5.0.3", ] solvers = [ "gurobipy", - "highspy>=1.5.0; python_version < '3.12'", - "highspy>=1.7.1; python_version >= '3.12'", - "cplex; platform_system != 'Darwin' and python_version < '3.12'", + "highspy>=1.5.0,!=1.14.0; python_version < '3.12'", + "highspy>=1.7.1,!=1.14.0; python_version >= '3.12'", + "cplex; platform_system != 'Darwin'", "mosek", - "mindoptpy; python_version < '3.12'", + "mindoptpy", "coptpy!=7.2.1", - "xpress; platform_system != 'Darwin' and python_version < '3.11'", + "xpress; platform_system != 'Darwin'", "pyscipopt; platform_system != 'Darwin'", + "knitro>=15.1.0", # "cupdlpx>=0.1.2", pip package currently unstable ] @@ -99,10 +124,22 @@ version_scheme = "no-guess-dev" [tool.pytest.ini_options] testpaths = ["test"] -norecursedirs = ["dev-scripts", "doc", "examples", "benchmark"] +norecursedirs = ["dev-scripts", "doc", "examples", "benchmark", "benchmarks"] markers = [ "gpu: marks tests as requiring GPU hardware (deselect with '-m \"not gpu\"')", ] +filterwarnings = [ + # Silence our own EvolvingAPIWarning inside the test suite so the + # piecewise tests don't emit 500+ warnings. Users still see them. + # Match by message prefix on the builtin FutureWarning base class + # rather than on ``linopy.EvolvingAPIWarning`` directly — using the + # class reference here would force pytest to import linopy at + # config-parse time, which loads ``linopy.variables`` from + # site-packages and then conflicts with ``--doctest-modules`` + # collection of ``linopy/variables.py`` in the source tree on + # Windows CI. + "ignore:piecewise:FutureWarning", +] [tool.coverage.run] branch = true @@ -112,7 +149,7 @@ omit = ["test/*"] exclude_also = ["if TYPE_CHECKING:"] [tool.mypy] -exclude = ['dev/*', 'examples/*', 'benchmark/*', 'doc/*'] +exclude = ['dev/*', 'examples/*', '^benchmark/', 'doc/*'] ignore_missing_imports = true no_implicit_optional = true warn_unused_ignores = true @@ -156,6 +193,7 @@ ignore = [ 'D101', # Missing docstring in public class 'D102', # Missing docstring in public method 'D103', # Missing docstring in public function + 'D106', # Missing docstring in public nested class 'D107', # Missing docstring in __init__ 'D202', # No blank lines allowed after function docstring 'D203', # 1 blank line required before class docstring diff --git a/test/conftest.py b/test/conftest.py index 3197689ba..067452d2b 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,9 +1,16 @@ """Pytest configuration and fixtures.""" +from __future__ import annotations + import os +from typing import TYPE_CHECKING +import pandas as pd import pytest +if TYPE_CHECKING: + from linopy import Model, Variable + def pytest_addoption(parser: pytest.Parser) -> None: """Add custom command line options.""" @@ -30,13 +37,19 @@ def pytest_configure(config: pytest.Config) -> None: def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item] ) -> None: - """Automatically skip GPU tests unless --run-gpu is passed.""" + """ + Auto-skip GPU-only solvers (e.g. cuPDLPx) unless --run-gpu is passed. + + Solvers that *also* have a CPU mode (e.g. xpress, which carries + ``GPU_ACCELERATION`` from version 9.8) are not skipped — their default + pathway is CPU and they should run in normal CI. + """ if config.getoption("--run-gpu"): return skip_gpu = pytest.mark.skip(reason="need --run-gpu option to run GPU tests") for item in items: - # Check if this is a parametrized test with a GPU solver + # Check if this is a parametrized test with a GPU-only solver if hasattr(item, "callspec") and "solver" in item.callspec.params: solver = item.callspec.params["solver"] # Import here to avoid circular dependency @@ -45,6 +58,46 @@ def pytest_collection_modifyitems( solver_supports, ) - if solver_supports(solver, SolverFeature.GPU_ACCELERATION): + if solver_supports(solver, SolverFeature.GPU_ONLY): item.add_marker(skip_gpu) item.add_marker(pytest.mark.gpu) + + +@pytest.fixture +def m() -> Model: + from linopy import Model + + m = Model() + m.add_variables(pd.Series([0, 0]), 1, name="x") + m.add_variables(4, pd.Series([8, 10]), name="y") + m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]).T, name="z") + m.add_variables(coords=[pd.RangeIndex(20, name="dim_2")], name="v") + idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) + idx.name = "dim_3" + m.add_variables(coords=[idx], name="u") + return m + + +@pytest.fixture +def x(m: Model) -> Variable: + return m.variables["x"] + + +@pytest.fixture +def y(m: Model) -> Variable: + return m.variables["y"] + + +@pytest.fixture +def z(m: Model) -> Variable: + return m.variables["z"] + + +@pytest.fixture +def v(m: Model) -> Variable: + return m.variables["v"] + + +@pytest.fixture +def u(m: Model) -> Variable: + return m.variables["u"] diff --git a/test/remote/test_oetc.py b/test/remote/test_oetc.py index d937e376c..7b2d75f2b 100644 --- a/test/remote/test_oetc.py +++ b/test/remote/test_oetc.py @@ -5,10 +5,11 @@ from unittest.mock import Mock, patch import pytest -import requests -from requests import RequestException -from linopy.remote.oetc import ( +requests = pytest.importorskip("requests") +from requests import RequestException # noqa: E402 + +from linopy.remote.oetc import ( # noqa: E402 AuthenticationResult, ComputeProvider, GcpCredentials, @@ -1391,7 +1392,9 @@ def test_submit_job_success( mock_post.return_value = mock_response # Execute - result = handler_with_auth_setup._submit_job_to_compute_service(input_file_name) + result = handler_with_auth_setup._submit_job_to_compute_service( + input_file_name, "gurobi", {} + ) # Verify request expected_payload = { @@ -1433,7 +1436,9 @@ def test_submit_job_http_error( # Execute and verify exception with pytest.raises(Exception) as exc_info: - handler_with_auth_setup._submit_job_to_compute_service(input_file_name) + handler_with_auth_setup._submit_job_to_compute_service( + input_file_name, "highs", {} + ) assert "Failed to submit job to compute service" in str(exc_info.value) @@ -1451,7 +1456,9 @@ def test_submit_job_missing_uuid_in_response( # Execute and verify exception with pytest.raises(Exception) as exc_info: - handler_with_auth_setup._submit_job_to_compute_service(input_file_name) + handler_with_auth_setup._submit_job_to_compute_service( + input_file_name, "highs", {} + ) assert "Invalid job submission response format: missing field 'uuid'" in str( exc_info.value @@ -1468,7 +1475,9 @@ def test_submit_job_network_error( # Execute and verify exception with pytest.raises(Exception) as exc_info: - handler_with_auth_setup._submit_job_to_compute_service(input_file_name) + handler_with_auth_setup._submit_job_to_compute_service( + input_file_name, "highs", {} + ) assert "Failed to submit job to compute service" in str(exc_info.value) @@ -1567,7 +1576,9 @@ def test_solve_on_oetc_file_upload( "/tmp/linopy-abc123.nc" ) mock_upload.assert_called_once_with("/tmp/linopy-abc123.nc") - mock_submit.assert_called_once_with("uploaded_file.nc.gz") + mock_submit.assert_called_once_with( + "uploaded_file.nc.gz", "highs", {} + ) mock_wait.assert_called_once_with("test-job-uuid") mock_download.assert_called_once_with("result.nc.gz") mock_read_netcdf.assert_called_once_with( @@ -1693,7 +1704,9 @@ def test_solve_on_oetc_with_job_submission( "/tmp/linopy-abc123.nc" ) mock_upload.assert_called_once_with("/tmp/linopy-abc123.nc") - mock_submit.assert_called_once_with(uploaded_file_name) + mock_submit.assert_called_once_with( + uploaded_file_name, "highs", {} + ) mock_wait.assert_called_once_with(job_uuid) mock_download.assert_called_once_with("result.nc.gz") mock_read_netcdf.assert_called_once_with( diff --git a/test/remote/test_oetc_job_polling.py b/test/remote/test_oetc_job_polling.py index 96ec98b4a..4b2681f9a 100644 --- a/test/remote/test_oetc_job_polling.py +++ b/test/remote/test_oetc_job_polling.py @@ -9,9 +9,11 @@ from unittest.mock import Mock, patch import pytest -from requests import RequestException -from linopy.remote.oetc import ( +requests = pytest.importorskip("requests") +from requests import RequestException # noqa: E402 + +from linopy.remote.oetc import ( # noqa: E402 AuthenticationResult, ComputeProvider, OetcCredentials, diff --git a/test/remote/test_ssh.py b/test/remote/test_ssh.py new file mode 100644 index 000000000..c6960c840 --- /dev/null +++ b/test/remote/test_ssh.py @@ -0,0 +1,157 @@ +"""Tests for ``linopy.remote.ssh.RemoteHandler.solve_on_remote``.""" + +from __future__ import annotations + +import warnings +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("paramiko") + +from linopy import Model # noqa: E402 +from linopy.remote.ssh import RemoteHandler # noqa: E402 + + +class _FakeSFTPClient: + """In-memory SFTP stand-in: ``put`` / ``get`` round-trip file bytes.""" + + def __init__(self) -> None: + self.store: dict[str, bytes] = {} + + def open(self, path: str, mode: str) -> Any: + store = self.store + + @contextmanager + def _writer() -> Iterator[Any]: + class _Writer: + def write(self_inner, data: str | bytes) -> None: + store[path] = data.encode() if isinstance(data, str) else data + + yield _Writer() + + return _writer() + + def put(self, local_path: str, remote_path: str) -> None: + with open(local_path, "rb") as fh: + self.store[remote_path] = fh.read() + + def get(self, remote_path: str, local_path: str) -> None: + with open(local_path, "wb") as fh: + fh.write(self.store[remote_path]) + + def remove(self, path: str) -> None: + self.store.pop(path, None) + + +def _make_sos_model() -> Model: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1.0, 2.0, 3.0]), sense="max") + return m + + +@pytest.fixture +def handler() -> RemoteHandler: + """``RemoteHandler`` wired to an in-memory SFTP and a no-op shell.""" + client = MagicMock() + client.invoke_shell.return_value.makefile.return_value = MagicMock() + sftp = _FakeSFTPClient() + client.open_sftp.return_value = sftp + + h = RemoteHandler(hostname="fake", client=client) + # The unsolved model gets put() into sftp.store under model_unsolved_file; + # serve it back as the "solved" model so read_netcdf has something valid. + h.sftp_client = sftp # type: ignore[assignment] + h.execute = MagicMock() # type: ignore[method-assign] + + original_put = sftp.put + + def put_and_mirror(local_path: str, remote_path: str) -> None: + original_put(local_path, remote_path) + if remote_path == h.model_unsolved_file: + sftp.store[h.model_solved_file] = sftp.store[remote_path] + + sftp.put = put_and_mirror # type: ignore[method-assign] + return h + + +class TestSolveOnRemoteSosBracket: + """``solve_on_remote`` must bracket SOS reformulation around transfer.""" + + def test_reformulates_before_transfer_and_restores_after( + self, handler: RemoteHandler + ) -> None: + m = _make_sos_model() + + observed: dict[str, bool] = {} + real_write = handler.write_model_on_remote + + def spy_write(model: Model) -> None: + observed["state_active"] = model._sos_reformulation_state is not None + observed["has_aux_var"] = "_sos_reform_x_y" in model.variables + real_write(model) + + handler.write_model_on_remote = spy_write # type: ignore[method-assign] + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + handler.solve_on_remote(m, reformulate_sos=True, solver_name="highs") + + assert observed["state_active"] is True + assert observed["has_aux_var"] is True + assert not any("active SOS reformulation" in str(w.message) for w in captured) + assert m._sos_reformulation_state is None + assert "_sos_reform_x_y" not in m.variables + assert list(m.variables.sos) == ["x"] + + def test_skips_bracket_when_reformulate_sos_false( + self, handler: RemoteHandler + ) -> None: + m = _make_sos_model() + + observed: dict[str, bool] = {} + real_write = handler.write_model_on_remote + + def spy_write(model: Model) -> None: + observed["state_active"] = model._sos_reformulation_state is not None + real_write(model) + + handler.write_model_on_remote = spy_write # type: ignore[method-assign] + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + handler.solve_on_remote(m, reformulate_sos=False) + + assert observed["state_active"] is False + assert not any("active SOS reformulation" in str(w.message) for w in captured) + assert m._sos_reformulation_state is None + + def test_auto_without_solver_name_raises_on_sos_model( + self, handler: RemoteHandler + ) -> None: + m = _make_sos_model() + with pytest.raises(ValueError, match="requires an explicit `solver_name`"): + handler.solve_on_remote(m, reformulate_sos="auto") + + def test_no_sos_model_passes_through_unchanged( + self, handler: RemoteHandler, tmp_path: Path + ) -> None: + m = Model() + x = m.add_variables(lower=0, upper=1, name="x") + m.add_objective(1.0 * x, sense="max") + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + handler.solve_on_remote(m, reformulate_sos="auto") + + assert m._sos_reformulation_state is None + assert not any("active SOS reformulation" in str(w.message) for w in captured) diff --git a/test/test_alignment.py b/test/test_alignment.py new file mode 100644 index 000000000..73d2cfdb0 --- /dev/null +++ b/test/test_alignment.py @@ -0,0 +1,1043 @@ +#!/usr/bin/env python3 +""" +Tests for linopy.alignment — conversion, broadcasting, and validation of +user input against coordinates. + +Organized by the module's public surface: + +- ``TestAsDataarrayFrom*`` — :func:`as_dataarray` (convert only) +- ``TestCoordsToDict`` — the coords-entry naming rules +- ``TestAddVariablesCoords`` — coords/dims → variable dims (end-to-end) +- ``TestBroadcastToCoords`` — ``broadcast_to_coords(strict=False)`` +- ``TestMultiIndexProjection`` — implicit MI-level projection (values, + deprecation warnings, coverage gaps) — the legacy/v1 fork point +- ``TestStrictMode`` — ``broadcast_to_coords(strict=True)`` +- ``TestValidateAlignment`` — the validation primitive +- ``TestAlign`` — symmetric :func:`align` +""" + +import warnings +from collections.abc import Callable +from typing import Any + +import numpy as np +import pandas as pd +import polars as pl +import pytest +import xarray as xr +from xarray import DataArray +from xarray.testing.assertions import assert_equal + +from linopy import EvolvingAPIWarning, LinearExpression, Model, Variable +from linopy.alignment import ( + _coords_to_dict, + align, + as_dataarray, + broadcast_to_coords, + fill_missing_coords, + validate_alignment, +) +from linopy.testing import assert_linequal, assert_varequal +from linopy.types import CoordsLike + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mi_index() -> pd.MultiIndex: + """Named (level1, level2) MultiIndex backing the stacked dim 'dim_3'.""" + idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) + idx.name = "dim_3" + return idx + + +@pytest.fixture +def mi_coords(mi_index: pd.MultiIndex) -> xr.Coordinates: + """Coordinates of the stacked MultiIndex dim 'dim_3'.""" + return xr.Coordinates.from_pandas_multiindex(mi_index, "dim_3") + + +@pytest.fixture +def by_level1() -> DataArray: + """A constant indexed by level1 only — a partial level set.""" + return DataArray([10.0, 20.0], coords={"level1": [1, 2]}, dims=["level1"]) + + +# --------------------------------------------------------------------------- +# as_dataarray — convert only +# --------------------------------------------------------------------------- + + +class TestAsDataarrayFromPandas: + """Series / DataFrame conversion: pandas axis names vs the dims argument.""" + + @pytest.mark.parametrize( + ("index", "dims", "expected_dim"), + [ + pytest.param([0, 1, 2], None, "dim_0", id="default"), + pytest.param(["a", "b", "c"], ["dim1"], "dim1", id="dims-set"), + pytest.param( + pd.Index(["a", "b", "c"], name="dim1"), [], "dim1", id="dims-given" + ), + pytest.param( + pd.Index(["a", "b", "c"], name="dim1"), + ["other"], + "dim1", + id="pandas-name-has-priority", + ), + pytest.param(["a", "b", "c"], [], "dim_0", id="dims-subset"), + pytest.param( + ["a", "b", "c"], ["dim_a", "other"], "dim_a", id="dims-superset" + ), + ], + ) + def test_series_dim_naming( + self, index: Any, dims: list[str] | None, expected_dim: str + ) -> None: + s = pd.Series([1, 2, 3], index=index) + da = as_dataarray(s, dims=dims) if dims is not None else as_dataarray(s) + assert isinstance(da, DataArray) + assert da.dims == (expected_dim,) + assert list(da.coords[expected_dim].values) == list(s.index) + + @pytest.mark.parametrize( + ("index", "columns", "dims", "expected_dims"), + [ + pytest.param([0, 1], ["A", "B"], None, ("dim_0", "dim_1"), id="default"), + pytest.param( + ["a", "b"], + ["A", "B"], + ("dim1", "dim2"), + ("dim1", "dim2"), + id="dims-set", + ), + pytest.param( + pd.Index(["a", "b"], name="dim1"), + pd.Index(["A", "B"], name="dim2"), + [], + ("dim1", "dim2"), + id="dims-given", + ), + pytest.param( + pd.Index(["a", "b"], name="dim1"), + pd.Index(["A", "B"], name="dim2"), + ["other"], + ("dim1", "dim2"), + id="pandas-name-has-priority", + ), + pytest.param( + ["a", "b"], ["A", "B"], [], ("dim_0", "dim_1"), id="dims-subset" + ), + pytest.param( + ["a", "b"], + ["A", "B"], + ["dim_a", "dim_b", "other"], + ("dim_a", "dim_b"), + id="dims-superset", + ), + ], + ) + def test_dataframe_dim_naming( + self, + index: Any, + columns: Any, + dims: Any, + expected_dims: tuple[str, ...], + ) -> None: + df = pd.DataFrame([[1, 2], [3, 4]], index=index, columns=columns) + da = as_dataarray(df, dims=dims) if dims is not None else as_dataarray(df) + assert isinstance(da, DataArray) + assert da.dims == expected_dims + assert list(da.coords[expected_dims[0]].values) == list(df.index) + assert list(da.coords[expected_dims[1]].values) == list(df.columns) + + @pytest.mark.parametrize( + "coords", + [[["a", "b", "c"]], {"dim_0": ["a", "b", "c"]}], + ids=["list", "dict"], + ) + def test_series_aligned_coords_do_not_warn(self, coords: Any) -> None: + """Coords matching the pandas index are accepted silently — no misalignment warning.""" + s = pd.Series([1, 2, 3], index=["a", "b", "c"]) + with warnings.catch_warnings(): + warnings.simplefilter("error") + da = as_dataarray(s, coords=coords) + assert da.dims == ("dim_0",) + assert list(da.coords["dim_0"].values) == ["a", "b", "c"] + + @pytest.mark.parametrize( + "coords", + [[["a", "b"], ["A", "B"]], {"dim_0": ["a", "b"], "dim_1": ["A", "B"]}], + ids=["list", "dict"], + ) + def test_dataframe_aligned_coords_do_not_warn(self, coords: Any) -> None: + """Coords matching the frame's index/columns are accepted silently.""" + df = pd.DataFrame([[1, 2], [3, 4]], index=["a", "b"], columns=["A", "B"]) + with warnings.catch_warnings(): + warnings.simplefilter("error") + da = as_dataarray(df, coords=coords) + assert da.dims == ("dim_0", "dim_1") + assert list(da.coords["dim_0"].values) == ["a", "b"] + assert list(da.coords["dim_1"].values) == ["A", "B"] + + def test_polars_series(self) -> None: + target_dim = "dim_0" + target_index = [0, 1, 2] + s = pl.Series([1, 2, 3]) + da = as_dataarray(s) + assert isinstance(da, DataArray) + assert da.dims == (target_dim,) + assert list(da.coords[target_dim].values) == target_index + + def test_series_dims_as_bare_string(self) -> None: + """Dims may be a single dim name instead of a list.""" + da = as_dataarray(pd.Series([1, 2, 3]), dims="x") + assert da.dims == ("x",) + + +class TestAsDataarrayFromNumpy: + """ndarray conversion: positional labeling from coords / dims.""" + + arr = np.array([[1, 2], [3, 4]]) + + @pytest.mark.parametrize( + ("coords", "dims", "expected"), + [ + pytest.param( + None, None, {"dim_0": [0, 1], "dim_1": [0, 1]}, id="no-coords-no-dims" + ), + pytest.param( + [["a", "b"], ["A", "B"]], + None, + {"dim_0": ["a", "b"], "dim_1": ["A", "B"]}, + id="coords-list", + ), + pytest.param( + [pd.Index(["a", "b"], name="dim1"), pd.Index(["A", "B"], name="dim2")], + None, + {"dim1": ["a", "b"], "dim2": ["A", "B"]}, + id="coords-named-indexes", + ), + pytest.param( + {"dim_0": ["a", "b"], "dim_2": ["A", "B"]}, + None, + {"dim_0": ["a", "b"], "dim_2": ["A", "B"]}, + id="coords-dict", + ), + pytest.param( + [["a", "b"], ["A", "B"]], + ("dim1", "dim2"), + {"dim1": ["a", "b"], "dim2": ["A", "B"]}, + id="coords-list-and-dims", + ), + pytest.param( + [["a", "b"], ["A", "B"]], + ("dim1", "dim2", "dim3"), + {"dim1": ["a", "b"], "dim2": ["A", "B"]}, + id="dims-superset", + ), + pytest.param( + [["a", "b"], ["A", "B"]], + ["dim0"], + {"dim0": ["a", "b"], "dim_1": ["A", "B"]}, + id="dims-subset", + ), + pytest.param( + [pd.Index(["a", "b"], name="dim1"), pd.Index(["A", "B"], name="dim2")], + ("dim1", "dim2"), + {"dim1": ["a", "b"], "dim2": ["A", "B"]}, + id="named-indexes-and-matching-dims", + ), + pytest.param( + {"dim_0": ["a", "b"], "dim_1": ["A", "B"]}, + ("dim_0", "dim_1"), + {"dim_0": ["a", "b"], "dim_1": ["A", "B"]}, + id="coords-dict-and-matching-dims", + ), + ], + ) + def test_labeling(self, coords: Any, dims: Any, expected: dict[str, list]) -> None: + da = as_dataarray(self.arr, coords=coords, dims=dims) + assert isinstance(da, DataArray) + assert da.dims == tuple(expected) + for dim, values in expected.items(): + assert list(da.coords[dim]) == values + + def test_named_indexes_conflicting_dims_raise(self) -> None: + coords = [pd.Index(["a", "b"], name="dim1"), pd.Index(["A", "B"], name="dim2")] + with pytest.raises(ValueError): + as_dataarray(self.arr, coords=coords, dims=("dim3", "dim4")) + + def test_extra_coord_entries_are_dropped(self) -> None: + """as_dataarray converts only: dims label the axes, extra coord entries are dropped.""" + target_coords = {"dim_0": ["a", "b"], "dim_2": ["A", "B"]} + da = as_dataarray(self.arr, coords=target_coords, dims=("dim_0", "dim_1")) + assert da.dims == ("dim_0", "dim_1") + assert list(da.coords["dim_0"].values) == ["a", "b"] + assert "dim_2" not in da.coords + + def test_dims_as_bare_string(self) -> None: + """Dims may be a single dim name; dict coords are filtered to those dims.""" + da = as_dataarray(np.array([1, 2]), coords={"x": [0, 1], "drop": [9]}, dims="x") + assert da.dims == ("x",) + assert list(da.coords["x"].values) == [0, 1] + assert "drop" not in da.coords + + def test_zero_dim_array_expands_over_dict_coords(self) -> None: + """A 0-d array converts like a scalar, expanding over dict coords.""" + da = as_dataarray(np.array(5.0), coords={"a": [0, 1]}) + assert da.dims == ("a",) + assert da.values.tolist() == [5.0, 5.0] + + +class TestAsDataarrayFromScalar: + """Scalar conversion: numbers expand over coords when given.""" + + @pytest.mark.parametrize( + "num", [1, np.float64(1)], ids=["python-int", "np-float64"] + ) + def test_with_dims_and_coords(self, num: Any) -> None: + da = as_dataarray(num, dims=["dim1"], coords=[["a"]]) + assert isinstance(da, DataArray) + assert da.dims == ("dim1",) + assert list(da.coords["dim1"].values) == ["a"] + + def test_default_dims_coords(self) -> None: + da = as_dataarray(1) + assert isinstance(da, DataArray) + assert da.dims == () + assert da.coords == {} + + def test_with_named_index_coords(self) -> None: + da = as_dataarray(1, coords=[pd.RangeIndex(10, name="a")]) + assert isinstance(da, DataArray) + assert da.dims == ("a",) + assert list(da.coords["a"].values) == list(range(10)) + + +class TestAsDataarrayFromDataArray: + """DataArray inputs pass through; unsupported types raise.""" + + da_in = DataArray( + data=[[1, 2], [3, 4]], + dims=["dim1", "dim2"], + coords={"dim1": ["a", "b"], "dim2": ["A", "B"]}, + ) + + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"dims": ["dim1", "dim2"], "coords": [["a", "b"], ["A", "B"]]}, + id="matching-dims-and-coords", + ), + pytest.param({}, id="default"), + ], + ) + def test_passthrough(self, kwargs: dict[str, Any]) -> None: + da_out = as_dataarray(self.da_in, **kwargs) + assert isinstance(da_out, DataArray) + assert da_out.dims == self.da_in.dims + assert list(da_out.coords["dim1"].values) == list( + self.da_in.coords["dim1"].values + ) + assert list(da_out.coords["dim2"].values) == list( + self.da_in.coords["dim2"].values + ) + + def test_unsupported_type_raises(self) -> None: + with pytest.raises(TypeError): + as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]]) + + def test_fill_missing_coords_rejects_non_xarray(self) -> None: + with pytest.raises( + TypeError, match="Expected xarray.DataArray or xarray.Dataset" + ): + fill_missing_coords([1, 2, 3]) # type: ignore[call-overload] + + def test_does_not_expand_missing_coord_dims(self) -> None: + """as_dataarray converts; only broadcast_to_coords expands missing dims.""" + coords = {"a": [0, 1], "b": [10, 20]} + arr = np.array([1, 2]) + + converted = as_dataarray(arr, coords=coords, dims=["a"]) + assert converted.dims == ("a",) + + broadcast = broadcast_to_coords(arr, coords=coords, dims=["a"], strict=False) + assert broadcast.dims == ("a", "b") + + +class TestAsDataarrayMultiIndexCoords: + """MultiIndex coords inputs: level names must not become extra dims.""" + + station_mi = pd.MultiIndex.from_tuples( + [("a", 1), ("b", 2)], names=["letter", "num"] + ) + + @pytest.mark.parametrize( + ("arr", "expected_values"), + [ + (np.float64(3.0), [3.0, 3.0]), + (3, [3, 3]), + (3.0, [3.0, 3.0]), + (np.array([10.0, 20.0]), [10.0, 20.0]), + ], + ids=["np_number", "python_int", "python_float", "numpy_array"], + ) + def test_input_types(self, arr: object, expected_values: list[float]) -> None: + """Level names in multi-index coords must not be treated as extra dims.""" + source = DataArray( + [1.0, 2.0], coords={"station": self.station_mi}, dims="station" + ) + + da = as_dataarray(arr, coords=source.coords) + + assert da.dims == ("station",) + assert da.shape == (2,) + assert set(da.coords.keys()) == {"station", "letter", "num"} + assert list(da.coords["letter"].values) == ["a", "b"] + assert list(da.coords["num"].values) == [1, 2] + assert da.coords["letter"].dims == ("station",) + assert da.coords["num"].dims == ("station",) + assert list(da.values) == expected_values + + @pytest.mark.parametrize( + "coords_factory", + [ + lambda mi: xr.Coordinates.from_pandas_multiindex(mi, "station"), + lambda mi: {"station": mi}, + lambda mi: ( + DataArray([1.0, 2.0], coords={"station": mi}, dims="station").coords + ), + ], + ids=["xarray_Coordinates", "plain_dict", "dataarray_coords"], + ) + def test_coord_input_forms( + self, coords_factory: Callable[[pd.MultiIndex], CoordsLike] + ) -> None: + """Users may pass a MultiIndex via Coordinates, a dict, or another DataArray's coords.""" + coords = coords_factory(self.station_mi) + + da = as_dataarray(3.0, coords=coords) + + assert da.dims == ("station",) + assert da.shape == (2,) + assert set(da.coords.keys()) == {"station", "letter", "num"} + assert da.coords["letter"].dims == ("station",) + assert da.coords["num"].dims == ("station",) + assert (da.values == 3.0).all() + + def test_explicit_dims_win_over_inference(self) -> None: + """Explicit dims must win over any inference from Coordinates.""" + source = DataArray( + [1.0, 2.0], coords={"station": self.station_mi}, dims="station" + ) + + da = as_dataarray(3.0, coords=source.coords, dims=["station"]) + assert da.dims == ("station",) + assert da.shape == (2,) + assert set(da.coords.keys()) == {"station", "letter", "num"} + + +def _ij_multiindex() -> pd.MultiIndex: + """Unnamed (i, j) MultiIndex used across the coords-entry tests.""" + return pd.MultiIndex.from_product([[0, 1], ["a", "b"]], names=["i", "j"]) + + +def _named_multiindex(name: str = "multi") -> pd.MultiIndex: + """:func:`_ij_multiindex` carrying an overall index name.""" + mi = _ij_multiindex() + mi.name = name + return mi + + +# --------------------------------------------------------------------------- +# _coords_to_dict — the coords-entry naming rules +# --------------------------------------------------------------------------- + + +class TestCoordsToDict: + """ + Executable spec of ``_coords_to_dict``: how each coords-entry form is + named or rejected, parameterized by entry form. The end-to-end dim + assignment these feed lives in :class:`TestAddVariablesCoords`. + """ + + @staticmethod + def _parse(coords: Any, dims: Any = None) -> dict: + return _coords_to_dict(coords, dims=dims) + + @pytest.mark.parametrize( + "coords, dims", + [ + ([("x", [0, 1, 2])], None), + ([pd.Index([0, 1, 2], name="x")], None), + ([pd.Index([0, 1, 2])], ["x"]), + ([[0, 1, 2]], ["x"]), + ([range(3)], ["x"]), + ([np.array([0, 1, 2])], ["x"]), + ], + ids=[ + "tuple", + "named-index", + "unnamed-index+dims", + "list+dims", + "range+dims", + "ndarray+dims", + ], + ) + def test_named_form_parses_to_x(self, coords: Any, dims: Any) -> None: + """Each naming form parses to {"x": [0, 1, 2]} (tuple = xarray form).""" + result = self._parse(coords, dims=dims) + assert set(result) == {"x"} + assert list(result["x"]) == [0, 1, 2] + assert result["x"].name == "x" + + @pytest.mark.parametrize( + "coords, expected", + [ + ( + xr.Coordinates.from_pandas_multiindex(_ij_multiindex(), "stacked"), + {"stacked"}, + ), + ([_named_multiindex()], {"multi"}), + ([("x", [0, 1, 2], {"units": "m"})], {"x"}), + ], + ids=["xarray-coordinates", "named-multiindex", "tuple-with-attrs"], + ) + def test_other_forms_parse_to_expected_names( + self, coords: Any, expected: set + ) -> None: + assert set(self._parse(coords)) == expected + + def test_mapping_returns_shallow_copy(self) -> None: + src = {"x": [0, 1, 2], "y": [10, 20]} + result = self._parse(src) + assert result == src + assert result is not src + + @pytest.mark.parametrize( + "entry", [pd.Index([0, 1, 2]), _ij_multiindex()], ids=["index", "multiindex"] + ) + def test_unnamed_index_named_from_dims_on_a_copy(self, entry: Any) -> None: + result = self._parse([entry], dims=["x"]) + assert result["x"].name == "x" + assert entry.name is None # caller not mutated + + @pytest.mark.parametrize( + "entry", + [[0, 1, 2], range(3), np.array([0, 1, 2]), pd.Index([0, 1, 2])], + ids=["list", "range", "ndarray", "unnamed-index"], + ) + def test_unlabeled_without_dims_is_skipped(self, entry: Any) -> None: + assert self._parse([entry]) == {} + + @pytest.mark.parametrize( + "coords, dims, match", + [ + ([_ij_multiindex()], None, r"MultiIndex.*must have \.name set"), + ([("x",)], None, r"\(dim_name, values\) convention"), + ([(0, 1, 2)], ["x"], r"\(dim_name, values\) convention"), + ([("x", 5)], None, r"with array-like values"), + ( + [DataArray([0, 1, 2], dims=["x"])], + None, + r"coords entries must be pd\.Index", + ), + ([object()], None, r"coords entries must be pd\.Index"), + ], + ids=[ + "unnamed-multiindex", + "tuple-too-short", + "tuple-bare-values", + "tuple-scalar-values", + "dataarray", + "unknown-type", + ], + ) + def test_invalid_entry_raises_typeerror( + self, coords: Any, dims: Any, match: str + ) -> None: + with pytest.raises(TypeError, match=match): + self._parse(coords, dims=dims) + + +# --------------------------------------------------------------------------- +# add_variables — coords / dims map to the variable's dimensions +# --------------------------------------------------------------------------- + + +class TestAddVariablesCoords: + """End-to-end: each coords / dims form sets the variable's dimensions.""" + + @pytest.mark.parametrize( + "coords, dims, expected_dims", + [ + ([("x", [0, 1, 2])], None, ("x",)), + ([pd.Index([0, 1, 2], name="x")], None, ("x",)), + ([pd.Index([0, 1, 2])], ["x"], ("x",)), + ([[0, 1, 2]], ["x"], ("x",)), + ([range(3)], ["x"], ("x",)), + ([np.array([0, 1, 2])], ["x"], ("x",)), + ([[0, 1, 2]], None, ("dim_0",)), + ([range(3)], None, ("dim_0",)), + ([np.array([0, 1, 2])], None, ("dim_0",)), + ([pd.Index([0, 1, 2])], None, ("dim_0",)), + ([("origin", ["a", "b"]), ("dest", ["x", "y"])], None, ("origin", "dest")), + ], + ids=[ + "tuple", + "named-index", + "unnamed-index+dims", + "list+dims", + "range+dims", + "ndarray+dims", + "list", + "range", + "ndarray", + "unnamed-index", + "multiple-tuples", + ], + ) + def test_coords_set_variable_dims( + self, coords: Any, dims: Any, expected_dims: tuple + ) -> None: + m = Model() + v = m.add_variables(lower=0, coords=coords, dims=dims) + assert v.dims == expected_dims + + +# --------------------------------------------------------------------------- +# broadcast_to_coords(strict=False) — broadcast mechanics, mismatches pass +# --------------------------------------------------------------------------- + + +class TestBroadcastToCoords: + """strict=False: dims are made to agree; entry mismatches pass through.""" + + def test_preserves_extra_dims(self) -> None: + """Extra dims in the input are not rejected — they broadcast downstream.""" + arr = DataArray( + [[1, 2], [3, 4], [5, 6]], + dims=["a", "t"], + coords={"a": [0, 1, 2], "t": [10, 20]}, + ) + coords = {"a": [0, 1, 2]} + da = broadcast_to_coords(arr, coords=coords, strict=False) + assert set(da.dims) == {"a", "t"} + assert list(da.coords["t"].values) == [10, 20] + + def test_keeps_disjoint_shared_dim_values(self) -> None: + """Different value sets on a shared dim are passed through (xr.align handles).""" + arr = DataArray([1, 2, 3, 4, 5], dims=["a"], coords={"a": [0, 1, 2, 3, 4]}) + coords = {"a": [2, 3]} + da = broadcast_to_coords(arr, coords=coords, strict=False) + # No exception, no reindex; downstream alignment intersects. + assert list(da.coords["a"].values) == [0, 1, 2, 3, 4] + + def test_extra_coord_entries_broadcast_in(self) -> None: + """Coords is source of truth: extra coord entries broadcast into the result.""" + target_coords = {"dim_0": ["a", "b"], "dim_2": ["A", "B"]} + arr = np.array([[1, 2], [3, 4]]) + da = broadcast_to_coords( + arr, coords=target_coords, dims=("dim_0", "dim_1"), strict=False + ) + # dims labels the positional axes; coords adds dim_2 by broadcast. + assert set(da.dims) == {"dim_0", "dim_1", "dim_2"} + assert list(da.coords["dim_0"].values) == ["a", "b"] + assert list(da.coords["dim_2"].values) == ["A", "B"] + + +# --------------------------------------------------------------------------- +# Implicit MultiIndex-level projection — the legacy/v1 fork point +# --------------------------------------------------------------------------- + + +class TestBroadcastToCoordsMultiIndexProjection: + """ + Inputs indexed by levels of a stacked MultiIndex dim are projected onto it. + + Implicit projection is deprecated (scenario B, #732/#737): it warns under + both modes today and will raise under the v1 convention. Coverage gaps + raise under strict mode. When #717 lands, the deprecation tests here fork + into legacy (warn) and v1 (raise) variants. + """ + + def test_broadcasts_single_level( + self, mi_coords: xr.Coordinates, by_level1: DataArray + ) -> None: + """ + A constant indexed by one MultiIndex level broadcasts across the MI dim. + + PyPSA multi-investment multiplies an expression over a (period, timestep) + 'snapshot' MultiIndex by a weighting indexed only by 'period'. Each level + combination of the MultiIndex must pick up its level's value. + """ + with pytest.warns(EvolvingAPIWarning, match=r"broadcasting level subset"): + da = broadcast_to_coords( + by_level1, coords=mi_coords, dims=["dim_3"], strict=False + ) + + assert da.dims == ("dim_3",) + assert isinstance(da.indexes["dim_3"], pd.MultiIndex) + assert da.sel(dim_3=(1, "a")).item() == 10.0 + assert da.sel(dim_3=(1, "b")).item() == 10.0 + assert da.sel(dim_3=(2, "a")).item() == 20.0 + assert da.sel(dim_3=(2, "b")).item() == 20.0 + + def test_stacks_full_levels(self, mi_coords: xr.Coordinates) -> None: + """ + A constant indexed by all MI level names stacks element-wise into the MI dim. + + PyPSA's storage_weightings is a pandas Series over a (period, timestep) + MultiIndex subset (the last snapshot of each period); it must align onto + the matching level combinations of the 'snapshot' MultiIndex. Combinations + the subset does not cover are left as NaN (broadcast path). + """ + subset = pd.MultiIndex.from_tuples( + [(1, "a"), (2, "b")], names=["level1", "level2"] + ) + weights = pd.Series([10.0, 20.0], index=subset) + + with pytest.warns( + EvolvingAPIWarning, match=r"filling uncovered level combinations" + ): + da = broadcast_to_coords( + weights, coords=mi_coords, dims=["dim_3"], strict=False + ) + + assert da.dims == ("dim_3",) + assert isinstance(da.indexes["dim_3"], pd.MultiIndex) + assert da.sel(dim_3=(1, "a")).item() == 10.0 + assert da.sel(dim_3=(2, "b")).item() == 20.0 + assert np.isnan(da.sel(dim_3=(1, "b")).item()) + assert np.isnan(da.sel(dim_3=(2, "a")).item()) + + def test_full_coverage_is_silent( + self, mi_coords: xr.Coordinates, mi_index: pd.MultiIndex + ) -> None: + """ + Full-level, fully-covering alignment is convention-clean → no warning. + + Aligning an input that reconstructs the whole MultiIndex onto its dim is + equivalent to the input already carrying that dim (future §11), so it must + not emit the EvolvingAPIWarning the partial/gap projections do. + """ + full = pd.Series([1.0, 2.0, 3.0, 4.0], index=mi_index) + + with warnings.catch_warnings(): + warnings.simplefilter("error", EvolvingAPIWarning) + da = broadcast_to_coords( + full, coords=mi_coords, dims=["dim_3"], strict=False + ) + + assert da.dims == ("dim_3",) + assert da.values.tolist() == [1.0, 2.0, 3.0, 4.0] + + def test_expands_missing_mi_dim_keeps_levels(self) -> None: + """ + Broadcasting a missing MultiIndex dim must keep its level coords intact. + + expand_dims drops MultiIndex level coords, leaving a degenerate flat + index that fails to align downstream (PyPSA multi-investment regression). + """ + midx = pd.MultiIndex.from_tuples( + [(2020, 0), (2020, 1), (2030, 0), (2030, 1)], + names=["period", "timestep"], + ) + midx.name = "snapshot" + sc = xr.Coordinates.from_pandas_multiindex(midx, "snapshot") + labels = DataArray( + [[1], [2], [3], [4]], + coords={**sc, "name": ["1"]}, + dims=["snapshot", "name"], + ) + coeff = broadcast_to_coords( + DataArray([1.0], coords={"name": ["1"]}, dims=["name"]), + coords=labels.coords, + dims=labels.dims, + strict=False, + ) + assert set(coeff.xindexes) == {"snapshot", "period", "timestep", "name"} + coeff.reindex_like(labels, fill_value=0) + + def test_ambiguous_level_raises(self) -> None: + """A level name shared by two MI dims cannot be resolved.""" + a = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("shared", "x")) + b = pd.MultiIndex.from_product([[1, 2], ["c", "d"]], names=("shared", "y")) + coords = { + **xr.Coordinates.from_pandas_multiindex(a, "dimA"), + **xr.Coordinates.from_pandas_multiindex(b, "dimB"), + } + arr = DataArray([1.0, 2.0], coords={"shared": [1, 2]}, dims=["shared"]) + + with pytest.raises(ValueError, match=r"shared.*shared by MultiIndex"): + broadcast_to_coords(arr, coords=coords, strict=False) + + def test_missing_level_value_raises(self, mi_coords: xr.Coordinates) -> None: + """A level value absent from the input cannot be broadcast.""" + by_level1 = DataArray([10.0, 20.0], coords={"level1": [1, 9]}, dims=["level1"]) + + with pytest.raises(ValueError, match=r"Cannot align level.*is missing"): + broadcast_to_coords( + by_level1, coords=mi_coords, dims=["dim_3"], strict=False + ) + + def test_unrelated_mi_series_still_unstacks(self) -> None: + """A MI Series whose levels match no coords MI dim keeps unstacking.""" + sub = pd.MultiIndex.from_product([["p", "q"], [1, 2]], names=["foo", "bar"]) + series = pd.Series([1.0, 2.0, 3.0, 4.0], index=sub) + + da = broadcast_to_coords(series, coords={"time": [0, 1, 2]}, strict=False) + + assert set(da.dims) == {"time", "foo", "bar"} + + def test_partially_named_mi_levels(self) -> None: + """A None level name in the MultiIndex is skipped during projection.""" + mi = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", None)) + mi.name = "dim_3" + by_level1 = DataArray([10.0, 20.0], coords={"level1": [1, 2]}, dims=["level1"]) + + with pytest.warns(EvolvingAPIWarning, match=r"broadcasting level subset"): + da = broadcast_to_coords(by_level1, coords={"dim_3": mi}, strict=False) + + assert da.dims == ("dim_3",) + assert da.values.tolist() == [10.0, 10.0, 20.0, 20.0] + + def test_gap_detection_with_extra_dims(self, mi_coords: xr.Coordinates) -> None: + """Gaps are detected per level combination even when the input has extra dims.""" + arr = DataArray( + [[[1.0, np.nan], [2.0, 2.0]], [[3.0, 3.0], [4.0, 4.0]]], + dims=["level1", "level2", "extra"], + coords={"level1": [1, 2], "level2": ["a", "b"], "extra": [0, 1]}, + ) + + with pytest.warns( + EvolvingAPIWarning, match=r"filling uncovered level combinations" + ): + da = broadcast_to_coords( + arr, coords=mi_coords, dims=["dim_3"], strict=False + ) + + assert set(da.dims) == {"dim_3", "extra"} + + def test_strict_gap_error_truncates_long_missing_list(self) -> None: + """More than 5 missing combinations are truncated in the error message.""" + idx = pd.MultiIndex.from_product( + [[1, 2, 3], ["a", "b", "c"]], names=("l1", "l2") + ) + idx.name = "dim_m" + coords = xr.Coordinates.from_pandas_multiindex(idx, "dim_m") + # Diagonal subset: every level value present, 6 of 9 combinations missing. + diagonal = pd.MultiIndex.from_tuples( + [(1, "a"), (2, "b"), (3, "c")], names=["l1", "l2"] + ) + weights = pd.Series([1.0, 2.0, 3.0], index=diagonal) + + with pytest.raises( + ValueError, match=r"no value for 6 level combination.*in total" + ): + broadcast_to_coords(weights, coords, dims=["dim_m"], label="lower bound") + + # --- strict-mode policy on MI projections (deprecation / gaps) --- + + def test_strict_partial_level_warns( + self, mi_coords: xr.Coordinates, by_level1: DataArray + ) -> None: + """ + Per-level bounds broadcast across the MI dim, with the deprecation warning. + + Scenario B (#732 / #737 discussion): implicit MI-level projection is + deprecated everywhere, including the strict (bounds/mask) path, and will + raise under the v1 convention. + """ + with pytest.warns(EvolvingAPIWarning, match=r"broadcasting level subset"): + da = broadcast_to_coords( + by_level1, mi_coords, dims=["dim_3"], label="lower bound" + ) + + assert da.sel(dim_3=(1, "b")).item() == 10.0 + assert da.sel(dim_3=(2, "a")).item() == 20.0 + + def test_strict_rejects_coverage_gap(self, mi_coords: xr.Coordinates) -> None: + """A coverage gap warns on the broadcast rung but raises on the strict rung.""" + subset = pd.MultiIndex.from_tuples( + [(1, "a"), (2, "b")], names=["level1", "level2"] + ) + weights = pd.Series([10.0, 20.0], index=subset) + + with pytest.warns( + EvolvingAPIWarning, match=r"filling uncovered level combinations" + ): + broadcast_to_coords(weights, coords=mi_coords, dims=["dim_3"], strict=False) + + with pytest.raises(ValueError, match=r"no value for .* level combination"): + broadcast_to_coords(weights, mi_coords, dims=["dim_3"], label="lower bound") + + def test_strict_rejects_unnamed_mi_mismatch(self) -> None: + """ + A MultiIndex input with unnamed levels cannot be projected by level name, + so it keeps its own index under the coords dim. The strict rung must still + reject it when its level combinations don't cover coords, just as the + named-level coverage-gap case does. + """ + idx = pd.MultiIndex.from_product([[2020, 2030], ["t1", "t2"]], names=("p", "t")) + idx.name = "snapshot" + coords = xr.Coordinates.from_pandas_multiindex(idx, "snapshot") + sparse_unnamed = pd.Series({(2020, "t1"): 1.0, (2030, "t2"): 2.0}) + + with pytest.raises(ValueError, match=r"MultiIndex for dimension 'snapshot'"): + broadcast_to_coords( + sparse_unnamed, coords, dims=["snapshot"], label="lower bound" + ) + + +# --------------------------------------------------------------------------- +# broadcast_to_coords(strict=True) — the contract +# --------------------------------------------------------------------------- + + +class TestBroadcastToCoordsStrictMode: + """strict=True: anything broadcasting can't resolve raises, naming label.""" + + def test_extra_dims_pass_loose_fail_strict(self) -> None: + """Extra dims pass through the broadcast rung but fail the strict rung.""" + arr = DataArray( + [[1, 2], [3, 4]], dims=["a", "t"], coords={"a": [0, 1], "t": [10, 20]} + ) + coords = {"a": [0, 1]} + + da = broadcast_to_coords(arr, coords=coords, strict=False) + assert set(da.dims) == {"a", "t"} + + with pytest.raises(ValueError, match=r"not declared in coords"): + broadcast_to_coords(arr, coords, label="lower bound") + + def test_requires_label(self) -> None: + """strict=True without label raises: errors must name their subject.""" + with pytest.raises(TypeError, match=r"requires `label`"): + broadcast_to_coords(np.array([1, 2]), {"x": [0, 1]}) # type: ignore[call-overload] + + def test_wraps_conversion_errors(self) -> None: + with pytest.raises(ValueError, match=r"lower bound could not be aligned"): + broadcast_to_coords(np.array([1, 2]), {"x": [0, 1, 2]}, label="lower bound") + + def test_preserves_type_errors(self) -> None: + """Unsupported input types stay TypeError (don't become ValueError).""" + with pytest.raises(TypeError, match=r"lower bound could not be aligned"): + broadcast_to_coords(lambda x: x, {"x": [0, 1, 2]}, label="lower bound") + + def test_does_not_relabel_coords_errors(self) -> None: + """Coords-side TypeError carries its own message, not the value label.""" + mi = pd.MultiIndex.from_product([[0, 1], ["a", "b"]], names=["i", "j"]) + with pytest.raises(TypeError, match=r"MultiIndex.*must have \.name set"): + broadcast_to_coords(np.array([1, 2, 3, 4]), [mi], label="lower bound") + + +# --------------------------------------------------------------------------- +# validate_alignment — the validation primitive +# --------------------------------------------------------------------------- + + +class TestValidateAlignment: + """Raise when arr is incompatible with coords; no-op otherwise.""" + + def test_rejects_extra_dims(self) -> None: + arr = DataArray( + [[1, 2], [3, 4]], dims=["a", "b"], coords={"a": [0, 1], "b": [0, 1]} + ) + with pytest.raises(ValueError, match=r"not declared in coords"): + validate_alignment(arr, {"a": [0, 1]}) + + def test_rejects_value_mismatch(self) -> None: + arr = DataArray([1, 2, 3], dims=["a"], coords={"a": [0, 1, 2]}) + with pytest.raises(ValueError, match="do not match coords"): + validate_alignment(arr, {"a": [10, 20, 30]}) + + def test_allows_subset_dims(self) -> None: + """arr.dims ⊂ coords.dims is fine (broadcasting fills the missing dim).""" + arr = DataArray([1, 2, 3], dims=["a"], coords={"a": [0, 1, 2]}) + validate_alignment(arr, {"a": [0, 1, 2], "b": [10, 20]}) # no raise + + def test_unnamed_coords_and_dims(self) -> None: + """coords=[[...]], dims=[...] enforces the same contract as a named mapping.""" + arr = DataArray([1, 2, 3], dims=["x"], coords={"x": [0, 1, 2]}) + validate_alignment(arr, [[0, 1, 2]], dims=["x"]) # no raise + + bad = DataArray( + [[1, 2], [3, 4]], dims=["x", "y"], coords={"x": [0, 1], "y": [0, 1]} + ) + with pytest.raises(ValueError, match=r"not declared in coords"): + validate_alignment(bad, [[0, 1]], dims=["x"]) + + def test_label_in_error(self) -> None: + arr = DataArray( + [[1, 2], [3, 4]], dims=["a", "b"], coords={"a": [0, 1], "b": [0, 1]} + ) + with pytest.raises(ValueError, match=r"lower bound has dimension\(s\) \['b'\]"): + validate_alignment(arr, {"a": [0, 1]}, label="lower bound") + + +# --------------------------------------------------------------------------- +# align — the symmetric counterpart (wraps xarray.align) +# --------------------------------------------------------------------------- + + +class TestAlign: + """align() conforms multiple linopy / xarray objects to common coords.""" + + def test_inner_join_intersects_coords(self, x: Variable) -> None: + """Default join keeps only the shared coords (x over [0, 1] ∩ alpha over [1, 2]).""" + alpha = xr.DataArray([1, 2], [[1, 2]]) + + x_obs, alpha_obs = align(x, alpha) + + assert isinstance(x_obs, Variable) + assert x_obs.shape == alpha_obs.shape == (1,) + assert_varequal(x_obs, x.loc[[1]]) + + def test_left_join_keeps_left_coords_and_fills(self, x: Variable) -> None: + """join='left' keeps x's coords; the right operand is reindexed with NaN.""" + alpha = xr.DataArray([1, 2], [[1, 2]]) + + x_obs, alpha_obs = align(x, alpha, join="left") + + assert isinstance(x_obs, Variable) + assert x_obs.shape == alpha_obs.shape == (2,) + assert_varequal(x_obs, x) + assert_equal(alpha_obs, DataArray([np.nan, 1], [[0, 1]])) + + def test_inner_join_over_multiindex(self, u: Variable) -> None: + """Inner join intersects MultiIndex coords element-wise across the stacked dim.""" + beta = xr.DataArray( + [1, 2, 3], + [ + ( + "dim_3", + pd.MultiIndex.from_tuples( + [(1, "b"), (2, "b"), (1, "c")], names=["level1", "level2"] + ), + ) + ], + ) + + beta_obs, u_obs = align(beta, u) + + assert isinstance(u_obs, Variable) + assert u_obs.shape == beta_obs.shape == (2,) + assert_varequal(u_obs, u.loc[[(1, "b"), (2, "b")]]) + assert_equal(beta_obs, beta.loc[[(1, "b"), (2, "b")]]) + + def test_aligns_linear_expression(self, x: Variable) -> None: + """A LinearExpression aligns alongside variables, keeping its _term dim.""" + alpha = xr.DataArray([1, 2], [[1, 2]]) + expr = 20 * x + + x_obs, expr_obs, alpha_obs = align(x, expr, alpha) + + assert isinstance(expr_obs, LinearExpression) + assert x_obs.shape == alpha_obs.shape == (1,) + assert expr_obs.shape == (1, 1) # the trailing 1 is the _term dim + assert_linequal(expr_obs, expr.loc[[1]]) diff --git a/test/test_available_solvers.py b/test/test_available_solvers.py new file mode 100644 index 000000000..549d173e4 --- /dev/null +++ b/test/test_available_solvers.py @@ -0,0 +1,240 @@ +"""Tests for the lazy ``available_solvers`` collection and ``license_status``.""" + +from __future__ import annotations + +import subprocess +import sys + +import pytest + +import linopy +from linopy import solvers as solvers_mod +from linopy.solvers import ( + LicenseStatus, + SolverName, + _solver_class_for, + available_solvers, + check_solver_licenses, + quadratic_solvers, +) + + +def test_import_does_not_load_license_managed_packages() -> None: + """ + Importing linopy must not import packages whose ``__init__`` runs license logic. + + Verified in a subprocess so the test isn't fooled by modules other tests + have already imported. + """ + code = ( + "import sys, linopy;" + "loaded = [m for m in ('mindoptpy', 'coptpy') if m in sys.modules];" + "print(','.join(loaded))" + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ) + assert result.stdout.strip() == "" + + +def test_is_available_matches_membership() -> None: + for sn in SolverName: + cls = _solver_class_for(sn.value) + if cls is None: + continue + assert cls.is_available() == (sn.value in available_solvers) + + +def test_available_solvers_caches(monkeypatch: pytest.MonkeyPatch) -> None: + cls = _solver_class_for("highs") + assert cls is not None + counter = {"n": 0} + + def probe() -> bool: + counter["n"] += 1 + return True + + monkeypatch.setattr(cls, "is_available", classmethod(lambda c: probe())) + fresh = solvers_mod._AvailableSolvers() + list(fresh) + list(fresh) + assert counter["n"] == 1 + + +def test_available_solvers_refresh_reprobes() -> None: + fresh = solvers_mod._AvailableSolvers() + first = list(fresh) + fresh.refresh() + second = list(fresh) + assert first == second + + +def test_quadratic_solvers_is_subset_of_available() -> None: + assert set(quadratic_solvers).issubset(set(available_solvers)) + + +def test_license_status_on_uninstalled_solver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("gurobi") + assert cls is not None + monkeypatch.setattr(cls, "is_available", classmethod(lambda c: False)) + probe_called = {"n": 0} + + def _probe() -> None: + probe_called["n"] += 1 + + monkeypatch.setattr(cls, "_license_probe", classmethod(lambda c: _probe())) + status = cls.license_status() + assert status.ok is False + assert status.message == "package not installed" + assert probe_called["n"] == 0 + + +def test_license_status_wraps_probe_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("gurobi") + assert cls is not None + monkeypatch.setattr(cls, "is_available", classmethod(lambda c: True)) + + def _boom() -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(cls, "_license_probe", classmethod(lambda c: _boom())) + status = cls.license_status() + assert status.ok is False + assert "boom" in (status.message or "") + assert bool(status) is False + + +def test_license_status_ok_when_probe_succeeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("highs") + assert cls is not None + monkeypatch.setattr(cls, "is_available", classmethod(lambda c: True)) + monkeypatch.setattr(cls, "_license_probe", classmethod(lambda c: None)) + status = cls.license_status() + assert status.ok is True + assert bool(status) is True + assert isinstance(status, LicenseStatus) + + +def test_check_solver_licenses_rejects_unknown() -> None: + with pytest.raises(ValueError, match="unknown solver"): + check_solver_licenses("not-a-solver") + + +def test_check_solver_licenses_returns_mapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("highs") + assert cls is not None + monkeypatch.setattr(cls, "is_available", classmethod(lambda c: True)) + monkeypatch.setattr(cls, "_license_probe", classmethod(lambda c: None)) + result = check_solver_licenses("highs") + assert set(result) == {"highs"} + assert result["highs"].ok is True + + +def test_available_solvers_reexported_from_top_level() -> None: + assert linopy.available_solvers is available_solvers + + +def test_mosek_license_probe_releases_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("mosek") + assert cls is not None + + events: list[str] = [] + + class _FakeTask: + def __enter__(self) -> _FakeTask: + events.append("task_enter") + return self + + def __exit__(self, *exc: object) -> None: + events.append("task_exit") + + def optimize(self) -> None: + events.append("task_optimize") + + class _FakeEnv: + def __enter__(self) -> _FakeEnv: + events.append("env_enter") + return self + + def __exit__(self, *exc: object) -> None: + events.append("env_exit") + + def Task(self, numcon: int, numvar: int) -> _FakeTask: + events.append(f"env_task({numcon},{numvar})") + return _FakeTask() + + class _FakeMosek: + Env = _FakeEnv + + monkeypatch.setattr(solvers_mod, "mosek", _FakeMosek) + + cls._license_probe() + + assert events == [ + "env_enter", + "env_task(0,0)", + "task_enter", + "task_optimize", + "task_exit", + "env_exit", + ] + + +def test_copt_license_probe_closes_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("copt") + assert cls is not None + + events: list[str] = [] + + class _FakeEnvr: + def __init__(self) -> None: + events.append("envr_init") + + def close(self) -> None: + events.append("envr_close") + + class _FakeCoptpy: + Envr = _FakeEnvr + + monkeypatch.setattr(solvers_mod, "coptpy", _FakeCoptpy) + + cls._license_probe() + + assert events == ["envr_init", "envr_close"] + + +def test_mindopt_license_probe_disposes_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cls = _solver_class_for("mindopt") + assert cls is not None + + events: list[str] = [] + + class _FakeEnv: + def __init__(self) -> None: + events.append("env_init") + + def dispose(self) -> None: + events.append("env_dispose") + + class _FakeMindoptpy: + Env = _FakeEnv + + monkeypatch.setattr(solvers_mod, "mindoptpy", _FakeMindoptpy) + + cls._license_probe() + + assert events == ["env_init", "env_dispose"] diff --git a/test/test_common.py b/test/test_common.py index db2183755..c56b17b17 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -10,430 +10,18 @@ import polars as pl import pytest import xarray as xr -from test_linear_expression import m, u, x # noqa: F401 -from xarray import DataArray -from xarray.testing.assertions import assert_equal -from linopy import LinearExpression, Model, Variable +from linopy import Model from linopy.common import ( - align, - as_dataarray, assign_multiindex_safe, best_int, + coords_from_dataset, + coords_to_dataset_vars, get_dims_with_index_levels, is_constant, iterate_slices, + maybe_group_terms_polars, ) -from linopy.testing import assert_linequal, assert_varequal - - -def test_as_dataarray_with_series_dims_default() -> None: - target_dim = "dim_0" - target_index = [0, 1, 2] - s = pd.Series([1, 2, 3]) - da = as_dataarray(s) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_dims_set() -> None: - target_dim = "dim1" - target_index = ["a", "b", "c"] - s = pd.Series([1, 2, 3], index=target_index) - dims = [target_dim] - da = as_dataarray(s, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_dims_given() -> None: - target_dim = "dim1" - target_index = ["a", "b", "c"] - index = pd.Index(target_index, name=target_dim) - s = pd.Series([1, 2, 3], index=index) - dims: list[str] = [] - da = as_dataarray(s, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_dims_priority() -> None: - """The dimension name from the pandas object should have priority.""" - target_dim = "dim1" - target_index = ["a", "b", "c"] - index = pd.Index(target_index, name=target_dim) - s = pd.Series([1, 2, 3], index=index) - dims = ["other"] - da = as_dataarray(s, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_dims_subset() -> None: - target_dim = "dim_0" - target_index = ["a", "b", "c"] - s = pd.Series([1, 2, 3], index=target_index) - dims: list[str] = [] - da = as_dataarray(s, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_dims_superset() -> None: - target_dim = "dim_a" - target_index = ["a", "b", "c"] - s = pd.Series([1, 2, 3], index=target_index) - dims = [target_dim, "other"] - da = as_dataarray(s, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_override_coords() -> None: - target_dim = "dim_0" - target_index = ["a", "b", "c"] - s = pd.Series([1, 2, 3], index=target_index) - with pytest.warns(UserWarning): - da = as_dataarray(s, coords=[[1, 2, 3]]) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_series_aligned_coords() -> None: - """This should not give out a warning even though coords are given.""" - target_dim = "dim_0" - target_index = ["a", "b", "c"] - s = pd.Series([1, 2, 3], index=target_index) - da = as_dataarray(s, coords=[target_index]) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - da = as_dataarray(s, coords={target_dim: target_index}) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_with_pl_series_dims_default() -> None: - target_dim = "dim_0" - target_index = [0, 1, 2] - s = pl.Series([1, 2, 3]) - da = as_dataarray(s) - assert isinstance(da, DataArray) - assert da.dims == (target_dim,) - assert list(da.coords[target_dim].values) == target_index - - -def test_as_dataarray_dataframe_dims_default() -> None: - target_dims = ("dim_0", "dim_1") - target_index = [0, 1] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - da = as_dataarray(df) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_dims_set() -> None: - target_dims = ("dim1", "dim2") - target_index = ["a", "b"] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - da = as_dataarray(df, dims=target_dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_dims_given() -> None: - target_dims = ("dim1", "dim2") - target_index = ["a", "b"] - target_columns = ["A", "B"] - index = pd.Index(target_index, name=target_dims[0]) - columns = pd.Index(target_columns, name=target_dims[1]) - df = pd.DataFrame([[1, 2], [3, 4]], index=index, columns=columns) - dims: list[str] = [] - da = as_dataarray(df, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_dims_priority() -> None: - """The dimension name from the pandas object should have priority.""" - target_dims = ("dim1", "dim2") - target_index = ["a", "b"] - target_columns = ["A", "B"] - index = pd.Index(target_index, name=target_dims[0]) - columns = pd.Index(target_columns, name=target_dims[1]) - df = pd.DataFrame([[1, 2], [3, 4]], index=index, columns=columns) - dims = ["other"] - da = as_dataarray(df, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_dims_subset() -> None: - target_dims = ("dim_0", "dim_1") - target_index = ["a", "b"] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - dims: list[str] = [] - da = as_dataarray(df, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_dims_superset() -> None: - target_dims = ("dim_a", "dim_b") - target_index = ["a", "b"] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - dims = [*target_dims, "other"] - da = as_dataarray(df, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_override_coords() -> None: - target_dims = ("dim_0", "dim_1") - target_index = ["a", "b"] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - with pytest.warns(UserWarning): - da = as_dataarray(df, coords=[[1, 2], [2, 3]]) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_dataframe_aligned_coords() -> None: - """This should not give out a warning even though coords are given.""" - target_dims = ("dim_0", "dim_1") - target_index = ["a", "b"] - target_columns = ["A", "B"] - df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns) - da = as_dataarray(df, coords=[target_index, target_columns]) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - coords = dict(zip(target_dims, [target_index, target_columns])) - da = as_dataarray(df, coords=coords) - assert isinstance(da, DataArray) - assert da.dims == target_dims - assert list(da.coords[target_dims[0]].values) == target_index - assert list(da.coords[target_dims[1]].values) == target_columns - - -def test_as_dataarray_with_ndarray_no_coords_no_dims() -> None: - target_dims = ("dim_0", "dim_1") - target_coords = [[0, 1], [0, 1]] - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == target_coords[i] - - -def test_as_dataarray_with_ndarray_coords_list_no_dims() -> None: - target_dims = ("dim_0", "dim_1") - target_coords = [["a", "b"], ["A", "B"]] - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == target_coords[i] - - -def test_as_dataarray_with_ndarray_coords_indexes_no_dims() -> None: - target_dims = ("dim1", "dim2") - target_coords = [ - pd.Index(["a", "b"], name="dim1"), - pd.Index(["A", "B"], name="dim2"), - ] - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == list(target_coords[i]) - - -def test_as_dataarray_with_ndarray_coords_dict_set_no_dims() -> None: - """If no dims are given and coords are a dict, the keys of the dict should be used as dims.""" - target_dims = ("dim_0", "dim_2") - target_coords = {"dim_0": ["a", "b"], "dim_2": ["A", "B"]} - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for dim in target_dims: - assert list(da.coords[dim]) == target_coords[dim] - - -def test_as_dataarray_with_ndarray_coords_list_dims() -> None: - target_dims = ("dim1", "dim2") - target_coords = [["a", "b"], ["A", "B"]] - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords, dims=target_dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == target_coords[i] - - -def test_as_dataarray_with_ndarray_coords_list_dims_superset() -> None: - target_dims = ("dim1", "dim2") - target_coords = [["a", "b"], ["A", "B"]] - arr = np.array([[1, 2], [3, 4]]) - dims = [*target_dims, "dim3"] - da = as_dataarray(arr, coords=target_coords, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == target_coords[i] - - -def test_as_dataarray_with_ndarray_coords_list_dims_subset() -> None: - target_dims = ("dim0", "dim_1") - target_coords = [["a", "b"], ["A", "B"]] - arr = np.array([[1, 2], [3, 4]]) - dims = ["dim0"] - da = as_dataarray(arr, coords=target_coords, dims=dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == target_coords[i] - - -def test_as_dataarray_with_ndarray_coords_indexes_dims_aligned() -> None: - target_dims = ("dim1", "dim2") - target_coords = [ - pd.Index(["a", "b"], name="dim1"), - pd.Index(["A", "B"], name="dim2"), - ] - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords, dims=target_dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for i, dim in enumerate(target_dims): - assert list(da.coords[dim]) == list(target_coords[i]) - - -def test_as_dataarray_with_ndarray_coords_indexes_dims_not_aligned() -> None: - target_dims = ("dim3", "dim4") - target_coords = [ - pd.Index(["a", "b"], name="dim1"), - pd.Index(["A", "B"], name="dim2"), - ] - arr = np.array([[1, 2], [3, 4]]) - with pytest.raises(ValueError): - as_dataarray(arr, coords=target_coords, dims=target_dims) - - -def test_as_dataarray_with_ndarray_coords_dict_dims_aligned() -> None: - target_dims = ("dim_0", "dim_1") - target_coords = {"dim_0": ["a", "b"], "dim_1": ["A", "B"]} - arr = np.array([[1, 2], [3, 4]]) - da = as_dataarray(arr, coords=target_coords, dims=target_dims) - assert isinstance(da, DataArray) - assert da.dims == target_dims - for dim in target_dims: - assert list(da.coords[dim]) == target_coords[dim] - - -def test_as_dataarray_with_ndarray_coords_dict_set_dims_not_aligned() -> None: - target_dims = ("dim_0", "dim_1") - target_coords = {"dim_0": ["a", "b"], "dim_2": ["A", "B"]} - arr = np.array([[1, 2], [3, 4]]) - with pytest.raises(ValueError): - as_dataarray(arr, coords=target_coords, dims=target_dims) - - -def test_as_dataarray_with_number() -> None: - num = 1 - da = as_dataarray(num, dims=["dim1"], coords=[["a"]]) - assert isinstance(da, DataArray) - assert da.dims == ("dim1",) - assert list(da.coords["dim1"].values) == ["a"] - - -def test_as_dataarray_with_np_number() -> None: - num = np.float64(1) - da = as_dataarray(num, dims=["dim1"], coords=[["a"]]) - assert isinstance(da, DataArray) - assert da.dims == ("dim1",) - assert list(da.coords["dim1"].values) == ["a"] - - -def test_as_dataarray_with_number_default_dims_coords() -> None: - num = 1 - da = as_dataarray(num) - assert isinstance(da, DataArray) - assert da.dims == () - assert da.coords == {} - - -def test_as_dataarray_with_number_and_coords() -> None: - num = 1 - da = as_dataarray(num, coords=[pd.RangeIndex(10, name="a")]) - assert isinstance(da, DataArray) - assert da.dims == ("a",) - assert list(da.coords["a"].values) == list(range(10)) - - -def test_as_dataarray_with_dataarray() -> None: - da_in = DataArray( - data=[[1, 2], [3, 4]], - dims=["dim1", "dim2"], - coords={"dim1": ["a", "b"], "dim2": ["A", "B"]}, - ) - da_out = as_dataarray(da_in, dims=["dim1", "dim2"], coords=[["a", "b"], ["A", "B"]]) - assert isinstance(da_out, DataArray) - assert da_out.dims == da_in.dims - assert list(da_out.coords["dim1"].values) == list(da_in.coords["dim1"].values) - assert list(da_out.coords["dim2"].values) == list(da_in.coords["dim2"].values) - - -def test_as_dataarray_with_dataarray_default_dims_coords() -> None: - da_in = DataArray( - data=[[1, 2], [3, 4]], - dims=["dim1", "dim2"], - coords={"dim1": ["a", "b"], "dim2": ["A", "B"]}, - ) - da_out = as_dataarray(da_in) - assert isinstance(da_out, DataArray) - assert da_out.dims == da_in.dims - assert list(da_out.coords["dim1"].values) == list(da_in.coords["dim1"].values) - assert list(da_out.coords["dim2"].values) == list(da_in.coords["dim2"].values) - - -def test_as_dataarray_with_unsupported_type() -> None: - with pytest.raises(TypeError): - as_dataarray(lambda x: 1, dims=["dim1"], coords=[["a"]]) def test_best_int() -> None: @@ -486,6 +74,25 @@ def test_assign_multiindex_safe() -> None: assert result["pressure"].equals(data) +def test_coords_dataset_vars_roundtrip_multiindex() -> None: + """MultiIndex and plain coords survive serialization to Dataset vars and back.""" + mi = pd.MultiIndex.from_product( + [[2020, 2030], ["t1", "t2"]], names=("period", "timestep") + ) + mi.name = "snapshot" + plain = pd.Index([1, 2, 3], name="simple") + + ds = xr.Dataset(coords_to_dataset_vars([mi, plain])) + restored = coords_from_dataset(ds, ["snapshot", "simple"]) + + assert isinstance(restored[0], pd.MultiIndex) + assert restored[0].equals(mi) + assert list(restored[0].names) == ["period", "timestep"] + assert restored[0].name == "snapshot" + assert restored[1].equals(plain) + assert restored[1].name == "simple" + + def test_iterate_slices_basic() -> None: ds = xr.Dataset( {"var": (("x", "y"), np.random.rand(10, 10))}, # noqa: NPY002 @@ -671,49 +278,6 @@ 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 - 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) - assert isinstance(x_obs, Variable) - assert x_obs.shape == alpha_obs.shape == (1,) - assert_varequal(x_obs, x.loc[[1]]) - - # left-join - x_obs, alpha_obs = align(x, alpha, join="left") - assert x_obs.shape == alpha_obs.shape == (2,) - assert isinstance(x_obs, Variable) - 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) - assert x_obs.shape == alpha_obs.shape == (1,) - assert expr_obs.shape == (1, 1) # _term dim - assert isinstance(expr_obs, LinearExpression) - assert_linequal(expr_obs, expr.loc[[1]]) - - def test_is_constant() -> None: model = Model() index = pd.Index(range(10), name="t") @@ -737,3 +301,20 @@ def test_is_constant() -> None: ] for cv in constant_values: assert is_constant(cv) + + +def test_maybe_group_terms_polars_no_duplicates() -> None: + """Fast path: distinct (labels, vars) pairs skip group_by.""" + df = pl.DataFrame({"labels": [0, 0], "vars": [1, 2], "coeffs": [3.0, 4.0]}) + result = maybe_group_terms_polars(df) + assert result.shape == (2, 3) + assert result.columns == ["labels", "vars", "coeffs"] + assert result["coeffs"].to_list() == [3.0, 4.0] + + +def test_maybe_group_terms_polars_with_duplicates() -> None: + """Slow path: duplicate (labels, vars) pairs trigger group_by.""" + df = pl.DataFrame({"labels": [0, 0], "vars": [1, 1], "coeffs": [3.0, 4.0]}) + result = maybe_group_terms_polars(df) + assert result.shape == (1, 3) + assert result["coeffs"].to_list() == [7.0] diff --git a/test/test_compatible_arithmetrics.py b/test/test_compatible_arithmetrics.py index 1d1618ba8..edab1ae19 100644 --- a/test/test_compatible_arithmetrics.py +++ b/test/test_compatible_arithmetrics.py @@ -98,13 +98,13 @@ def test_arithmetric_operations_variable(m: Model) -> None: assert_linequal(x + data, x + other_datatype) assert_linequal(x - data, x - other_datatype) assert_linequal(x * data, x * other_datatype) - assert_linequal(x / data, x / other_datatype) # type: ignore - assert_linequal(data * x, other_datatype * x) # type: ignore + assert_linequal(x / data, x / other_datatype) + assert_linequal(data * x, other_datatype * x) # type: ignore[arg-type] assert x.__add__(object()) is NotImplemented assert x.__sub__(object()) is NotImplemented assert x.__mul__(object()) is NotImplemented - assert x.__truediv__(object()) is NotImplemented # type: ignore - assert x.__pow__(object()) is NotImplemented # type: ignore + assert x.__truediv__(object()) is NotImplemented + assert x.__pow__(object()) is NotImplemented # type: ignore[operator] with pytest.raises(ValueError): x.__pow__(3) diff --git a/test/test_constraint.py b/test/test_constraint.py index 35f49ea2b..10febf15e 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -24,7 +24,12 @@ short_LESS_EQUAL, sign_replace_dict, ) -from linopy.constraints import AnonymousScalarConstraint, Constraint, Constraints +from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + ConstraintBase, + Constraints, +) @pytest.fixture @@ -33,7 +38,7 @@ def m() -> Model: x = m.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") m.add_variables(coords=[pd.Index([1, 2, 3], name="second")], name="y") m.add_variables(0, 10, name="z") - m.add_constraints(x >= 0, name="c") + m.add_constraints(x >= 0, name="c", freeze=True) return m @@ -48,19 +53,49 @@ def y(m: Model) -> linopy.Variable: @pytest.fixture -def c(m: Model) -> linopy.constraints.Constraint: +def c(m: Model) -> linopy.constraints.ConstraintBase: return m.constraints["c"] -def test_constraint_repr(c: linopy.constraints.Constraint) -> None: +@pytest.fixture +def mc(m: Model) -> linopy.constraints.Constraint: + return m.constraints["c"].mutable() + + +def test_constraint_repr(c: linopy.constraints.CSRConstraint) -> None: c.__repr__() +def test_constraint_repr_equivalent_to_mutable( + c: linopy.constraints.CSRConstraint, +) -> None: + """Constraint (CSR-backed) and Constraint repr must be identical.""" + frozen = c.freeze() + assert repr(frozen) == repr(c) + + def test_constraints_repr(m: Model) -> None: m.constraints.__repr__() -def test_constraint_name(c: linopy.constraints.Constraint) -> None: +def test_add_constraints_freeze(m: Model, x: linopy.Variable) -> None: + c = m.add_constraints(x >= 1, name="frozen_c", freeze=True) + assert isinstance(c, linopy.constraints.CSRConstraint) + assert isinstance(m.constraints["frozen_c"], linopy.constraints.CSRConstraint) + assert c.ncons == 10 + + +def test_add_constraints_uses_model_freeze_default() -> None: + m = Model(freeze_constraints=True) + x = m.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") + c = m.add_constraints(x >= 1, name="frozen_by_default") + assert isinstance(c, linopy.constraints.CSRConstraint) + assert isinstance( + m.constraints["frozen_by_default"], linopy.constraints.CSRConstraint + ) + + +def test_constraint_name(c: linopy.constraints.CSRConstraint) -> None: assert c.name == "c" @@ -69,13 +104,34 @@ def test_empty_constraints_repr() -> None: Model().constraints.__repr__() +@pytest.mark.parametrize("freeze_constraints", [True, False]) +def test_constraint_handles_empty_rows(freeze_constraints: bool) -> None: + """An empty constraint group must be accepted and solve cleanly.""" + + m = Model(freeze_constraints=freeze_constraints) + x = m.add_variables( + lower=0.0, + coords=[range(3), range(2)], + dims=["time", "product"], + name="x", + ) + empty = x.isel(time=range(1, 1)) + c = m.add_constraints(empty == 0, name="empty") + assert isinstance(c, linopy.constraints.ConstraintBase) + assert c.size == 0 + # Solving a model with only an empty constraint group is also fine. + m.add_objective(x.sum()) + m.solve("highs", io_api="direct", output_flag=False) + assert m.status == "ok" + + def test_cannot_create_constraint_without_variable() -> None: model = linopy.Model() with pytest.raises(ValueError): _ = linopy.LinearExpression(12, model) == linopy.LinearExpression(13, model) -def test_constraints_getter(m: Model, c: linopy.constraints.Constraint) -> None: +def test_constraints_getter(m: Model, c: linopy.constraints.CSRConstraint) -> None: assert c.shape == (10,) assert isinstance(m.constraints[["c"]], Constraints) @@ -228,7 +284,7 @@ def test_constraint_wrapped_methods(x: linopy.Variable, y: linopy.Variable) -> N def test_anonymous_constraint_sel(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con.sel(first=[1, 2]), Constraint) + assert isinstance(con.sel(first=[1, 2]), ConstraintBase) def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) -> None: @@ -236,7 +292,7 @@ def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) con = expr <= 10 con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)}) con = con.swap_dims({"second": "third"}) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert con.coord_dims == ("first", "third") @@ -245,7 +301,7 @@ def test_anonymous_constraint_set_index(x: linopy.Variable, y: linopy.Variable) con = expr <= 10 con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)}) con = con.set_index({"multi": ["second", "third"]}) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert con.coord_dims == ( "first", "multi", @@ -256,13 +312,13 @@ def test_anonymous_constraint_set_index(x: linopy.Variable, y: linopy.Variable) def test_anonymous_constraint_loc(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con.loc[[1, 2]], Constraint) + assert isinstance(con.loc[[1, 2]], ConstraintBase) def test_anonymous_constraint_getitem(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con[1], Constraint) + assert isinstance(con[1], ConstraintBase) def test_constraint_from_rule(m: Model, x: linopy.Variable, y: linopy.Variable) -> None: @@ -271,7 +327,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint: coords = [x.coords["first"], y.coords["second"]] con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert con.lhs.nterm == 2 repr(con) # test repr @@ -286,7 +342,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: coords = [x.coords["first"], y.coords["second"]] con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, Constraint) + assert isinstance(con, ConstraintBase) assert isinstance(con.lhs.vars, xr.DataArray) assert con.lhs.nterm == 2 assert (con.lhs.vars.loc[0, :] == -1).all() @@ -295,152 +351,321 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: def test_constraint_vars_getter( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - assert_equal(c.vars.squeeze(), x.labels) + assert_equal(mc.vars.squeeze(), x.labels) -def test_constraint_coeffs_getter(c: linopy.constraints.Constraint) -> None: - assert (c.coeffs == 1).all() +def test_constraint_coeffs_getter(mc: linopy.constraints.Constraint) -> None: + assert (mc.coeffs == 1).all() -def test_constraint_sign_getter(c: linopy.constraints.Constraint) -> None: +def test_constraint_sign_getter(c: linopy.constraints.CSRConstraint) -> None: assert (c.sign == GREATER_EQUAL).all() -def test_constraint_rhs_getter(c: linopy.constraints.Constraint) -> None: +def test_constraint_rhs_getter(c: linopy.constraints.CSRConstraint) -> None: assert (c.rhs == 0).all() def test_constraint_vars_setter( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - c.vars = x - assert_equal(c.vars, x.labels) + mc.vars = x + assert_equal(mc.vars, x.labels) def test_constraint_vars_setter_with_array( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - c.vars = x.labels - assert_equal(c.vars, x.labels) + """Passing a raw DataArray is deprecated but still works for back-compat.""" + with pytest.warns(FutureWarning, match="DataArray"): + mc.vars = x.labels + assert_equal(mc.vars, x.labels) def test_constraint_vars_setter_invalid( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: with pytest.raises(TypeError): - c.vars = pd.DataFrame(x.labels) + mc.vars = pd.DataFrame(x.labels) -def test_constraint_coeffs_setter(c: linopy.constraints.Constraint) -> None: - c.coeffs = 3 - assert (c.coeffs == 3).all() +def test_constraint_coeffs_setter(mc: linopy.constraints.Constraint) -> None: + mc.coeffs = 3 + assert (mc.coeffs == 3).all() def test_constraint_lhs_setter( - c: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable ) -> None: - c.lhs = x + y - assert c.lhs.nterm == 2 - assert c.vars.notnull().all().item() - assert c.coeffs.notnull().all().item() + mc.lhs = x + y + assert mc.lhs.nterm == 2 + assert mc.vars.notnull().all().item() + assert mc.coeffs.notnull().all().item() def test_constraint_lhs_setter_with_variable( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - c.lhs = x - assert c.lhs.nterm == 1 + mc.lhs = x + assert mc.lhs.nterm == 1 -def test_constraint_lhs_setter_with_constant(c: linopy.constraints.Constraint) -> None: - sizes = c.sizes - c.lhs = 10 - assert (c.rhs == -10).all() - assert c.lhs.nterm == 0 - assert c.sizes["first"] == sizes["first"] +def test_constraint_lhs_setter_with_constant( + mc: linopy.constraints.Constraint, +) -> None: + sizes = mc.sizes + mc.lhs = 10 + assert (mc.rhs == -10).all() + assert mc.lhs.nterm == 0 + assert mc.sizes["first"] == sizes["first"] -def test_constraint_sign_setter(c: linopy.constraints.Constraint) -> None: - c.sign = EQUAL - assert (c.sign == EQUAL).all() +def test_constraint_sign_setter(mc: linopy.constraints.Constraint) -> None: + mc.sign = EQUAL + assert (mc.sign == EQUAL).all() -def test_constraint_sign_setter_alternative(c: linopy.constraints.Constraint) -> None: - c.sign = long_EQUAL - assert (c.sign == EQUAL).all() +def test_constraint_sign_setter_alternative( + mc: linopy.constraints.Constraint, +) -> None: + mc.sign = long_EQUAL + assert (mc.sign == EQUAL).all() -def test_constraint_sign_setter_invalid(c: linopy.constraints.Constraint) -> None: +def test_constraint_sign_setter_invalid( + mc: linopy.constraints.Constraint, +) -> None: # Test that assigning lhs with other type that LinearExpression raises TypeError with pytest.raises(ValueError): - c.sign = "asd" + mc.sign = "asd" + +def test_constraint_rhs_setter(mc: linopy.constraints.Constraint) -> None: + sizes = mc.sizes + mc.rhs = 2 + assert (mc.rhs == 2).all() + assert mc.sizes == sizes -def test_constraint_rhs_setter(c: linopy.constraints.Constraint) -> None: - sizes = c.sizes - c.rhs = 2 # type: ignore - assert (c.rhs == 2).all() - assert c.sizes == sizes + +def test_constraint_update_rhs_and_sign(mc: linopy.constraints.Constraint) -> None: + mc.update(rhs=5, sign=EQUAL) + assert (mc.rhs == 5).all() + assert (mc.sign == EQUAL).all() + + +def test_constraint_update_no_kwargs_is_noop( + mc: linopy.constraints.Constraint, +) -> None: + old_rhs = mc.rhs.copy() + old_sign = mc.sign.copy() + mc.update() + assert (mc.rhs == old_rhs).all() + assert (mc.sign == old_sign).all() + + +def test_constraint_update_rearranges_variable_rhs( + mc: linopy.constraints.Constraint, x: linopy.Variable +) -> None: + """ + Variable / Expression rhs is moved onto lhs; only the constant + part lands on rhs (mirrors add_constraints and the .rhs setter). + """ + mc.update(rhs=x + 3) + assert (mc.rhs == 3).all() + assert mc.lhs.nterm == 2 # original term + the rearranged -x + + +def test_constraint_update_returns_self( + mc: linopy.constraints.Constraint, +) -> None: + out = mc.update(rhs=7) + assert out is mc + + +def test_constraint_update_positional_constraint_expression( + mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable +) -> None: + """``c.update(x + 5 <= 3)`` replaces lhs / sign / rhs in one call.""" + mc.update(x + y <= 7) + assert (mc.rhs == 7).all() + assert (mc.sign == LESS_EQUAL).all() + assert mc.lhs.nterm == 2 + + +def test_constraint_update_positional_rejects_mixing_kwargs( + mc: linopy.constraints.Constraint, x: linopy.Variable +) -> None: + """Positional constraint can't be combined with keyword updates.""" + with pytest.raises(TypeError, match="cannot be combined with keyword"): + mc.update(x <= 3, sign=EQUAL) + + +def test_constraint_update_positional_rejects_non_constraint( + mc: linopy.constraints.Constraint, +) -> None: + """Random objects are rejected with a clear error.""" + with pytest.raises(TypeError, match="must be a ConstraintLike"): + mc.update("not a constraint") # type: ignore + + +def test_constraint_update_lhs_only( + mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable +) -> None: + """lhs= alone replaces the expression; rhs and sign untouched.""" + old_rhs = mc.rhs.copy() + old_sign = mc.sign.copy() + mc.update(lhs=5 * x + 7 * y) + assert (mc.rhs == old_rhs).all() + assert (mc.sign == old_sign).all() + assert mc.lhs.nterm == 2 + + +def test_constraint_update_coeffs_only_keeps_values( + mc: linopy.constraints.Constraint, +) -> None: + """coeffs= alone replaces the coef array element-wise; vars untouched.""" + old_vars = mc.vars.copy() + mc.update(coeffs=mc.coeffs * 10) + assert (mc.vars == old_vars).all() + # original was mc.lhs with leading coeff; *10 → all coeffs *10 + assert mc.coeffs.max() >= 10 + + +def test_constraint_update_lhs_and_sign_together( + mc: linopy.constraints.Constraint, x: linopy.Variable +) -> None: + """Compound updates compose: lhs replacement + sign flip in one call.""" + mc.update(lhs=2 * x, sign=EQUAL) + assert (mc.sign == EQUAL).all() + assert mc.lhs.nterm == 1 + + +def test_constraint_update_lhs_and_coeffs_rejected( + mc: linopy.constraints.Constraint, x: linopy.Variable +) -> None: + """lhs= (full replacement) and coeffs= (partial) are mutually exclusive.""" + with pytest.raises(TypeError, match="lhs.*coeffs.*variables"): + mc.update(lhs=2 * x, coeffs=mc.coeffs * 2) + + +def test_constraint_update_lhs_and_variables_rejected( + mc: linopy.constraints.Constraint, x: linopy.Variable +) -> None: + """lhs= (full replacement) and variables= (partial) are mutually exclusive.""" + with pytest.raises(TypeError, match="lhs.*coeffs.*variables"): + mc.update(lhs=2 * x, variables=mc.vars) def test_constraint_rhs_setter_with_variable( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - c.rhs = x # type: ignore - assert (c.rhs == 0).all() - assert (c.coeffs.isel({c.term_dim: -1}) == -1).all() - assert c.lhs.nterm == 2 + mc.rhs = x + assert (mc.rhs == 0).all() + assert (mc.coeffs.isel({mc.term_dim: -1}) == -1).all() + assert mc.lhs.nterm == 2 def test_constraint_rhs_setter_with_expression( - c: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable ) -> None: - c.rhs = x + y - assert (c.rhs == 0).all() - assert (c.coeffs.isel({c.term_dim: -1}) == -1).all() - assert c.lhs.nterm == 3 + mc.rhs = x + y + assert (mc.rhs == 0).all() + assert (mc.coeffs.isel({mc.term_dim: -1}) == -1).all() + assert mc.lhs.nterm == 3 def test_constraint_rhs_setter_with_expression_and_constant( - c: linopy.constraints.Constraint, x: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: - c.rhs = x + 1 - assert (c.rhs == 1).all() - assert (c.coeffs.sum(c.term_dim) == 0).all() - assert c.lhs.nterm == 2 + mc.rhs = x + 1 + assert (mc.rhs == 1).all() + assert (mc.coeffs.sum(mc.term_dim) == 0).all() + assert mc.lhs.nterm == 2 + + +def test_constraint_rhs_setter_broadcasts_missing_dim() -> None: + """Rhs assignment broadcasts against the constraint coords: missing dims expand.""" + m = Model() + x = m.add_variables( + coords=[pd.RangeIndex(2, name="i"), pd.RangeIndex(3, name="j")], name="x" + ) + con = m.add_constraints(1 * x >= 0, name="con") + + con.rhs = xr.DataArray([1.0, 2.0], dims=["i"], coords={"i": [0, 1]}) + assert dict(con.rhs.sizes) == {"i": 2, "j": 3} + assert (con.rhs.sel(i=1) == 2.0).all() -def test_constraint_labels_setter_invalid(c: linopy.constraints.Constraint) -> None: - # Test that assigning labels raises FrozenInstanceError + +def test_constraint_rhs_setter_projects_multiindex_level() -> None: + """ + Rhs indexed by one MultiIndex level is projected onto the stacked dim. + + Regression: as_expression must convert constants with the broadcast rung + (broadcast_to_coords), not plain conversion — otherwise the level dim + collides with the MI level coord downstream (xarray AlignmentError). + """ + idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) + idx.name = "dim_3" + coords = xr.Coordinates.from_pandas_multiindex(idx, "dim_3") + m = Model() + x = m.add_variables(coords=coords, name="x") + con = m.add_constraints(1 * x >= 0, name="con") + + rhs_by_level = xr.DataArray( + [10.0, 20.0], coords={"level1": [1, 2]}, dims=["level1"] + ) + with pytest.warns(linopy.EvolvingAPIWarning, match="broadcasting level subset"): + con.rhs = rhs_by_level + + assert con.rhs.sel(dim_3=(1, "b")).item() == 10.0 + assert con.rhs.sel(dim_3=(2, "a")).item() == 20.0 + + +def test_constraint_labels_setter_invalid(c: linopy.constraints.CSRConstraint) -> None: + # Test that assigning labels raises AttributeError (Constraint is frozen) with pytest.raises(AttributeError): c.labels = c.labels # type: ignore -def test_constraint_sel(c: linopy.constraints.Constraint) -> None: - assert isinstance(c.sel(first=[1, 2]), Constraint) - assert isinstance(c.isel(first=[1, 2]), Constraint) +def test_constraint_sel(c: linopy.constraints.CSRConstraint) -> None: + assert isinstance(c.mutable().sel(first=[1, 2]), ConstraintBase) + assert isinstance(c.mutable().isel(first=[1, 2]), ConstraintBase) -def test_constraint_flat(c: linopy.constraints.Constraint) -> None: +def test_constraint_flat(c: linopy.constraints.CSRConstraint) -> None: assert isinstance(c.flat, pd.DataFrame) -def test_iterate_slices(c: linopy.constraints.Constraint) -> None: - for i in c.iterate_slices(slice_size=2): - assert isinstance(i, Constraint) - assert c.coord_dims == i.coord_dims +def test_iterate_slices(mc: linopy.constraints.Constraint) -> None: + for i in mc.iterate_slices(slice_size=2): + assert isinstance(i, ConstraintBase) + assert mc.coord_dims == i.coord_dims -def test_constraint_to_polars(c: linopy.constraints.Constraint) -> None: +def test_constraint_to_polars(c: linopy.constraints.CSRConstraint) -> None: assert isinstance(c.to_polars(), pl.DataFrame) +def test_constraint_to_polars_mixed_signs(m: Model, x: linopy.Variable) -> None: + """Test to_polars when a constraint has mixed sign values across dims.""" + # Use Constraint so sign data can be patched + con = m.add_constraints(x >= 0, name="mixed", freeze=False) + # Replace sign data with mixed signs across the first dimension + n = con.sizes["first"] + signs = np.array(["<=" if i % 2 == 0 else ">=" for i in range(n)]) + con.data["sign"] = xr.DataArray(signs, dims=con.data["sign"].dims) + df = con.to_polars() + assert isinstance(df, pl.DataFrame) + assert set(df["sign"].to_list()) == {"<=", ">="} + + def test_constraint_assignment_with_anonymous_constraints( m: Model, x: linopy.Variable, y: linopy.Variable ) -> None: - m.add_constraints(x + y == 0, name="c2") + m.add_constraints(x + y == 0, name="c2", freeze=False) assert m.constraints["c2"].vars.notnull().all() assert m.constraints["c2"].coeffs.notnull().all() @@ -448,10 +673,14 @@ def test_constraint_assignment_with_anonymous_constraints( def test_constraint_assignment_sanitize_zeros( m: Model, x: linopy.Variable, y: linopy.Variable ) -> None: - m.add_constraints(0 * x + y == 0, name="c2") + m.add_constraints(0 * x + y == 0, name="c2", freeze=True) m.constraints.sanitize_zeros() - assert m.constraints["c2"].vars[0, 0, 0].item() == -1 - assert np.isnan(m.constraints["c2"].coeffs[0, 0, 0].item()) + c2 = m.constraints["c2"] + assert c2.nterm == 1 + assert c2.has_variable(y) + assert not c2.has_variable(x) + csr, _ = c2.to_matrix(m.variables.label_index) + assert (csr.data == 1).all() def test_constraint_assignment_with_args( @@ -570,12 +799,15 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: con = Constraint(data, m, "c") expr = m.add_constraints(con) - assert not set(HELPER_DIMS).intersection(set(expr.data.coords)) + assert not set(HELPER_DIMS).intersection(set(expr.coords)) def test_constraint_matrix(m: Model) -> None: - A = m.constraints.to_matrix() - assert A.shape == (10, 14) + # Returns (csr_array, con_labels) — dense: active rows and active-var columns + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (10, n_active_vars) + assert len(con_labels) == 10 def test_constraint_matrix_masked_variables() -> None: @@ -586,54 +818,48 @@ def test_constraint_matrix_masked_variables() -> None: missing. The matrix shoud not be built for constraints which have variables which are missing. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 6) - assert A.shape == (m.ncons, m.nvars) - - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + # Returns dense matrix: active rows only, all active-var columns + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) + assert len(con_labels) == m.ncons def test_constraint_matrix_masked_constraints() -> None: """ Test constraint matrix with missing constraints. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)]) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 11) - assert A.shape == (m.ncons, m.nvars) - - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + # active cons are indices 5-9, which reference vars 5-9 only (all active) + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) + assert len(con_labels) == m.ncons def test_constraint_matrix_masked_constraints_and_variables() -> None: """ - Test constraint matrix with missing constraints. + Test constraint matrix with missing constraints and variables. """ - # now with missing variables m = Model() mask = pd.Series([False] * 5 + [True] * 5) x = m.add_variables(coords=[range(10)], mask=mask) m.add_variables() m.add_constraints(x, EQUAL, 0, mask=mask) - A = m.constraints.to_matrix(filter_missings=True) - assert A.shape == (5, 6) - assert A.shape == (m.ncons, m.nvars) - - A = m.constraints.to_matrix(filter_missings=False) - assert A.shape == m.shape + # both masks align: 5 active cons x all active vars (5 x + 1 scalar) + A, con_labels = m.constraints.to_matrix() + n_active_vars = len(m.variables.label_index.vlabels) + assert A.shape == (m.ncons, n_active_vars) + assert len(con_labels) == m.ncons def test_get_name_by_label() -> None: @@ -660,3 +886,175 @@ def test_constraints_inequalities(m: Model) -> None: def test_constraints_equalities(m: Model) -> None: assert isinstance(m.constraints.equalities, Constraints) + + +def test_freeze_mutable_roundtrip(m: Model) -> None: + frozen = m.constraints["c"] + assert isinstance(frozen, linopy.constraints.CSRConstraint) + mc = frozen.mutable() + assert isinstance(mc, Constraint) + refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) + assert_equal(frozen.labels, refrozen.labels) + assert_equal(frozen.rhs, refrozen.rhs) + assert_equal(frozen.sign, refrozen.sign) + np.testing.assert_array_equal(frozen._csr.toarray(), refrozen._csr.toarray()) + np.testing.assert_array_equal(frozen._con_labels, refrozen._con_labels) + + +def test_freeze_mutable_roundtrip_with_masking() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") + mask = xr.DataArray([True, False, True, False, True], dims=["i"]) + m.add_constraints(x.where(mask) >= 0, name="c", freeze=True) + frozen = m.constraints["c"] + assert isinstance(frozen, linopy.constraints.CSRConstraint) + mc = frozen.mutable() + refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) + assert_equal(frozen.labels, refrozen.labels) + assert_equal(frozen.rhs, refrozen.rhs) + assert frozen.ncons == refrozen.ncons == 3 + + +def test_from_mutable_mixed_signs() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + m.add_constraints(x >= 0, name="mixed", freeze=False) + mc = m.constraints["mixed"] + assert isinstance(mc, Constraint) + mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) + frozen = linopy.constraints.CSRConstraint.from_mutable(mc) + assert isinstance(frozen._sign, np.ndarray) + assert list(frozen._sign) == ["<=", ">=", "<="] + assert_equal(frozen.sign, mc.sign) + + +def test_variable_label_index(m: Model) -> None: + li = m.variables.label_index + assert li.n_active_vars > 0 + assert len(li.vlabels) == li.n_active_vars + assert li.label_to_pos.shape[0] == m._xCounter + for lbl in li.vlabels: + assert li.label_to_pos[lbl] >= 0 + assert (li.label_to_pos[li.vlabels] == np.arange(li.n_active_vars)).all() + + +def test_variable_label_index_invalidation(m: Model) -> None: + li = m.variables.label_index + old_vlabels = li.vlabels.copy() + m.add_variables(name="w") + li.invalidate() + assert len(li.vlabels) > len(old_vlabels) + + +def test_to_matrix_with_rhs(m: Model) -> None: + c = m.constraints["c"] + assert isinstance(c, linopy.constraints.CSRConstraint) + li = m.variables.label_index + csr, con_labels, b, sense = c.to_matrix_with_rhs(li) + assert csr.shape[0] == len(con_labels) + assert csr.shape[0] == len(b) + assert csr.shape[0] == len(sense) + assert all(s in ("<", ">", "=") for s in sense) + np.testing.assert_array_equal(b, c._rhs) + + +def test_to_matrix_with_rhs_mutable(m: Model) -> None: + mc = m.constraints["c"].mutable() + li = m.variables.label_index + csr, con_labels, b, sense = mc.to_matrix_with_rhs(li) + assert csr.shape[0] == len(con_labels) + assert csr.shape[0] == len(b) + assert csr.shape[0] == len(sense) + + +def test_constraint_repr_shows_variable_names(m: Model) -> None: + c = m.constraints["c"] + r = repr(c) + assert "x" in r + + +def test_freeze_mixed_signs_from_rule() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m: Model, i: int) -> AnonymousScalarConstraint: + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="mixed_rule", freeze=True) + assert isinstance(con, linopy.constraints.CSRConstraint) + assert isinstance(con._sign, np.ndarray) + assert con.ncons == 4 + expected_signs = ["=", ">=", "=", ">="] + assert list(con._sign) == expected_signs + np.testing.assert_array_equal(con.sign.values, expected_signs) + + +def test_frozen_lhs_setter_raises() -> None: + m = Model() + time = pd.RangeIndex(5, name="t") + x = m.add_variables(lower=0, coords=[time], name="x") + y = m.add_variables(lower=0, coords=[time], name="y") + con = m.add_constraints(x >= 0, name="c", freeze=True) + assert isinstance(con, linopy.constraints.CSRConstraint) + with pytest.raises(AttributeError, match="read-only"): + con.lhs = 3 * x + 2 * y + + +def test_frozen_rhs_setter_raises() -> None: + m = Model() + time = pd.RangeIndex(5, name="t") + x = m.add_variables(lower=0, coords=[time], name="x") + con = m.add_constraints(x >= 0, name="c", freeze=True) + assert isinstance(con, linopy.constraints.CSRConstraint) + with pytest.raises(AttributeError, match="read-only"): + con.rhs = 10 + + +def test_mixed_sign_to_matrix_with_rhs() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m: Model, i: int) -> AnonymousScalarConstraint: + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="c") + li = m.variables.label_index + csr, con_labels, b, sense = con.to_matrix_with_rhs(li) + assert len(sense) == 4 + assert list(sense) == ["=", ">", "=", ">"] + + +def test_mixed_sign_sanitize_infinities() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + m.add_constraints(x >= 0, name="c", freeze=False) + mc = m.constraints["c"] + assert isinstance(mc, Constraint) + mc._data["sign"] = xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"]) + mc._data["rhs"] = xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"]) + frozen = mc.freeze() + frozen.sanitize_infinities() + assert frozen.ncons == 2 + np.testing.assert_array_equal(frozen._rhs, [1.0, 2.0]) + + +def test_mixed_sign_repr() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + coords = [pd.RangeIndex(4, name="i")] + + def bound(m: Model, i: int) -> AnonymousScalarConstraint: + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + con = m.add_constraints(bound, coords=coords, name="c") + r = repr(con) + assert "≥" in r + assert "=" in r diff --git a/test/test_constraint_coef_dirty.py b/test/test_constraint_coef_dirty.py new file mode 100644 index 000000000..6e32217b6 --- /dev/null +++ b/test/test_constraint_coef_dirty.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import pytest + +from linopy import Model + + +@pytest.fixture +def m_with_c() -> tuple[Model, str]: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 10, coords=[range(3)], name="y") + m.add_constraints(2 * x + y >= 5, name="c") + return m, "c" + + +def test_initial_coef_dirty_false(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + assert m.constraints[name]._coef_dirty is False + + +def test_update_coeffs_sets_dirty(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + c.update(coeffs=c.coeffs * 2) + assert c._coef_dirty is True + + +def test_update_variables_sets_dirty(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + x = m.variables["x"] + c.update(variables=x) + assert c._coef_dirty is True + + +def test_update_lhs_sets_dirty(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + x = m.variables["x"] + c.update(lhs=3 * x) + assert c._coef_dirty is True + + +def test_update_pure_constant_rhs_short_circuits(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + coeffs_buf = c.data["coeffs"].values + vars_buf = c.data["vars"].values + c.update(rhs=9) + assert c._coef_dirty is False + assert c.data["coeffs"].values is coeffs_buf + assert c.data["vars"].values is vars_buf + + +def test_update_rhs_with_variable_sets_dirty(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + x = m.variables["x"] + c.update(rhs=x + 3) + assert c._coef_dirty is True + + +def test_update_sign_does_not_set_dirty(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + c = m.constraints[name] + c.update(sign="<=") + assert c._coef_dirty is False + + +def test_flag_persists_across_container_access(m_with_c: tuple[Model, str]) -> None: + m, name = m_with_c + m.constraints[name].update(coeffs=m.constraints[name].coeffs * 2) + assert m.constraints[name]._coef_dirty is True + + +def test_update_positional_constraint_sets_dirty(m_with_c: tuple[Model, str]) -> None: + """Positional ``c.update(expr <= rhs)`` rewrites lhs and must flip the flag.""" + m, name = m_with_c + c = m.constraints[name] + x = m.variables["x"] + c.update(4 * x >= 7) + assert c._coef_dirty is True + + +def test_update_noop_does_not_set_dirty(m_with_c: tuple[Model, str]) -> None: + """``c.update()`` with no args is a no-op and must not flip the flag.""" + m, name = m_with_c + c = m.constraints[name] + c.update() + assert c._coef_dirty is False diff --git a/test/test_constraint_label_index.py b/test/test_constraint_label_index.py new file mode 100644 index 000000000..3652d5637 --- /dev/null +++ b/test/test_constraint_label_index.py @@ -0,0 +1,86 @@ +import numpy as np +import pandas as pd +import pytest + +import linopy +import linopy.constants +from linopy.constraints import Constraint + + +@pytest.fixture +def model_with_mask() -> linopy.Model: + m = linopy.Model() + coords = pd.Index(range(5), name="i") + mask = np.array([True, False, True, True, False]) + x = m.add_variables(lower=0, coords=[coords], name="x") + y = m.add_variables(lower=0, coords=[coords], name="y") + m.add_constraints(x + y >= 1, name="c_xy", mask=mask) + m.add_constraints(x.sum() + y.sum() <= 100, name="c_sum") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + +def test_clabels_parity_with_matrices(model_with_mask: linopy.Model) -> None: + expected = model_with_mask.matrices.clabels + actual = model_with_mask.constraints.label_index.clabels + np.testing.assert_array_equal(actual, expected) + + +def test_assign_result_does_not_build_matrix( + monkeypatch: pytest.MonkeyPatch, model_with_mask: linopy.Model +) -> None: + calls = {"n": 0} + original = Constraint._matrix_export_data + + def counting(self, label_index): # type: ignore[no-untyped-def] + calls["n"] += 1 + return original(self, label_index) + + monkeypatch.setattr(Constraint, "_matrix_export_data", counting) + + model_with_mask.solve("highs") + + assert model_with_mask.status == "ok" + # one build for solver input is fine; the post-solve mapping must not add more + n_after_solve = calls["n"] + solver = model_with_mask.solver + assert solver is not None + assert solver.status is not None + result = linopy.constants.Result( + status=solver.status, + solution=solver.solution, + solver_model=solver.solver_model, + solver_name=solver.solver_name.value, + report=solver.report, + ) + model_with_mask.assign_result(result) + assert calls["n"] == n_after_solve + + +def test_label_index_invalidated_on_add(model_with_mask: linopy.Model) -> None: + first = model_with_mask.constraints.label_index.clabels.copy() + x = model_with_mask.variables["x"] + model_with_mask.add_constraints(x.sum() >= 0, name="c_extra") + second = model_with_mask.constraints.label_index.clabels + assert len(second) == len(first) + 1 + + +def test_label_index_invalidated_on_remove(model_with_mask: linopy.Model) -> None: + before = len(model_with_mask.constraints.label_index.clabels) + removed = len(model_with_mask.constraints["c_sum"].active_labels()) + model_with_mask.constraints.remove("c_sum") + after = len(model_with_mask.constraints.label_index.clabels) + assert after == before - removed + + +def test_assign_result_correctness_with_mask(model_with_mask: linopy.Model) -> None: + model_with_mask.solve("highs") + assert model_with_mask.status == "ok" + x_sol = model_with_mask.variables["x"].solution.values + y_sol = model_with_mask.variables["y"].solution.values + assert np.isfinite(x_sol).all() + assert np.isfinite(y_sol).all() + dual = model_with_mask.constraints["c_xy"].dual.values + mask = np.array([True, False, True, True, False]) + assert np.isfinite(dual[mask]).all() + assert np.isnan(dual[~mask]).all() diff --git a/test/test_constraints.py b/test/test_constraints.py index cca010e8c..acc41b2e9 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -5,6 +5,8 @@ @author: fabulous """ +from typing import Any + import dask import dask.array.core import numpy as np @@ -12,7 +14,7 @@ import pytest import xarray as xr -from linopy import EQUAL, GREATER_EQUAL, LESS_EQUAL, Model +from linopy import EQUAL, GREATER_EQUAL, LESS_EQUAL, Model, Variable, available_solvers from linopy.testing import assert_conequal # Test model functions @@ -34,10 +36,10 @@ def test_constraint_assignment() -> None: assert "con0" in getattr(m.constraints, attr) assert m.constraints.labels.con0.shape == (10, 10) - assert m.constraints.labels.con0.dtype == int - assert m.constraints.coeffs.con0.dtype in (int, float) - assert m.constraints.vars.con0.dtype in (int, float) - assert m.constraints.rhs.con0.dtype in (int, float) + assert np.issubdtype(m.constraints.labels.con0.dtype, np.integer) + assert np.issubdtype(m.constraints.coeffs.con0.dtype, np.number) + assert np.issubdtype(m.constraints.vars.con0.dtype, np.number) + assert np.issubdtype(m.constraints.rhs.con0.dtype, np.number) assert_conequal(m.constraints.con0, con0) @@ -88,10 +90,10 @@ def test_anonymous_constraint_assignment() -> None: assert "con0" in getattr(m.constraints, attr) assert m.constraints.labels.con0.shape == (10, 10) - assert m.constraints.labels.con0.dtype == int - assert m.constraints.coeffs.con0.dtype in (int, float) - assert m.constraints.vars.con0.dtype in (int, float) - assert m.constraints.rhs.con0.dtype in (int, float) + assert np.issubdtype(m.constraints.labels.con0.dtype, np.integer) + assert np.issubdtype(m.constraints.coeffs.con0.dtype, np.number) + assert np.issubdtype(m.constraints.vars.con0.dtype, np.number) + assert np.issubdtype(m.constraints.rhs.con0.dtype, np.number) def test_constraint_assignment_with_tuples() -> None: @@ -139,6 +141,82 @@ def test_constraint_assignment_with_reindex() -> None: assert (con.coords["dim_0"].values == shuffled_coords).all() +@pytest.mark.parametrize( + "rhs_factory", + [ + pytest.param(lambda m, v: v, id="numpy"), + pytest.param(lambda m, v: xr.DataArray(v, dims=["dim_0"]), id="dataarray"), + pytest.param(lambda m, v: pd.Series(v, index=v), id="series"), + pytest.param( + lambda m, v: m.add_variables(coords=[v]), + id="variable", + ), + pytest.param( + lambda m, v: 2 * m.add_variables(coords=[v]) + 1, + id="linexpr", + ), + ], +) +def test_constraint_rhs_lower_dim(rhs_factory: Any) -> None: + m = Model() + naxis = np.arange(10, dtype=float) + maxis = np.arange(10).astype(str) + x = m.add_variables(coords=[naxis, maxis]) + y = m.add_variables(coords=[naxis, maxis]) + + c = m.add_constraints(x - y >= rhs_factory(m, naxis)) + assert c.shape == (10, 10) + + +@pytest.mark.parametrize( + "rhs_factory", + [ + pytest.param(lambda m: np.ones((5, 3)), id="numpy"), + pytest.param(lambda m: pd.DataFrame(np.ones((5, 3))), id="dataframe"), + ], +) +def test_constraint_rhs_higher_dim_constant_warns( + rhs_factory: Any, caplog: Any +) -> None: + m = Model() + x = m.add_variables(coords=[range(5)], name="x") + + with caplog.at_level("WARNING", logger="linopy.expressions"): + m.add_constraints(x >= rhs_factory(m)) + assert "dimensions" in caplog.text + + +def test_constraint_rhs_higher_dim_dataarray_reindexes() -> None: + """DataArray RHS with extra dims reindexes to expression coords (no raise).""" + m = Model() + x = m.add_variables(coords=[range(5)], name="x") + rhs = xr.DataArray(np.ones((5, 3)), dims=["dim_0", "extra"]) + + c = m.add_constraints(x >= rhs) + assert c.shape == (5, 3) + + +@pytest.mark.parametrize( + "rhs_factory", + [ + pytest.param( + lambda m: m.add_variables(coords=[range(5), range(3)]), + id="variable", + ), + pytest.param( + lambda m: 2 * m.add_variables(coords=[range(5), range(3)]) + 1, + id="linexpr", + ), + ], +) +def test_constraint_rhs_higher_dim_expression(rhs_factory: Any) -> None: + m = Model() + x = m.add_variables(coords=[range(5)], name="x") + + c = m.add_constraints(x >= rhs_factory(m)) + assert c.shape == (5, 3) + + def test_wrong_constraint_assignment_repeated() -> None: # repeated variable assignment is forbidden m: Model = Model() @@ -162,6 +240,50 @@ def test_masked_constraints() -> None: assert (m.constraints.labels.con0[5:10, :] == -1).all() +def test_masked_constraints_broadcast() -> None: + m: Model = Model() + + lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) + upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) + x = m.add_variables(lower, upper) + y = m.add_variables() + + mask = pd.Series([True] * 5 + [False] * 5) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc1", mask=mask) + assert (m.constraints.labels.bc1[0:5, :] != -1).all() + assert (m.constraints.labels.bc1[5:10, :] == -1).all() + + mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"]) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc2", mask=mask2) + assert (m.constraints.labels.bc2[:, 0:5] != -1).all() + assert (m.constraints.labels.bc2[:, 5:10] == -1).all() + + # Pandas Series with named index missing a dim is broadcast to data.coords. + mask_pd = pd.Series( + [True, False, True] + [False] * 7, index=pd.RangeIndex(10, name="dim_0") + ) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc_pd", mask=mask_pd) + assert (m.constraints.labels.bc_pd[[0, 2], :] != -1).all() + assert (m.constraints.labels.bc_pd[[1, 3, 4, 5, 6, 7, 8, 9], :] == -1).all() + + # Mask with sparse coords (subset of data's coords) now raises instead of + # emitting a FutureWarning — the rule from the bounds path applies here too. + mask3 = xr.DataArray( + [True, True, False, False, False], + dims=["dim_0"], + coords={"dim_0": range(5)}, + ) + with pytest.raises( + ValueError, match=r"mask: coordinate values for dimension 'dim_0'" + ): + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc3", mask=mask3) + + # Mask with extra dimension not in data should raise + mask4 = xr.DataArray([True, False], dims=["extra_dim"]) + with pytest.raises(ValueError, match=r"mask has dimension\(s\) \['extra_dim'\]"): + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc4", mask=mask4) + + def test_non_aligned_constraints() -> None: m: Model = Model() @@ -231,3 +353,105 @@ def test_sanitize_infinities() -> None: m.add_constraints(x >= np.inf, name="con_wrong_inf") with pytest.raises(ValueError): m.add_constraints(y <= -np.inf, name="con_wrong_neg_inf") + + +class TestConstraintCoordinateAlignment: + @pytest.fixture(params=["xarray", "pandas_series"], ids=["da", "series"]) + def subset(self, request: Any) -> xr.DataArray | pd.Series: + if request.param == "xarray": + return xr.DataArray([10.0, 30.0], dims=["dim_2"], coords={"dim_2": [1, 3]}) + return pd.Series([10.0, 30.0], index=pd.Index([1, 3], name="dim_2")) + + @pytest.fixture(params=["xarray", "pandas_series"], ids=["da", "series"]) + def superset(self, request: Any) -> xr.DataArray | pd.Series: + if request.param == "xarray": + return xr.DataArray( + np.arange(25, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(25)}, + ) + return pd.Series( + np.arange(25, dtype=float), index=pd.Index(range(25), name="dim_2") + ) + + 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", [LESS_EQUAL, GREATER_EQUAL, EQUAL]) + def test_var_comparison_subset( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == LESS_EQUAL: + con = v <= subset + elif sign == GREATER_EQUAL: + 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()) + + @pytest.mark.parametrize("sign", [LESS_EQUAL, GREATER_EQUAL, EQUAL]) + def test_subset_comparison_var( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == LESS_EQUAL: + con = subset <= v + elif sign == GREATER_EQUAL: + 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) + + @pytest.mark.parametrize("sign", [LESS_EQUAL, GREATER_EQUAL]) + def test_superset_comparison_var( + self, v: Variable, superset: xr.DataArray, sign: str + ) -> None: + if sign == LESS_EQUAL: + 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_constraint_rhs_extra_dims_broadcasts(self, v: Variable) -> None: + rhs = xr.DataArray( + [[1.0, 2.0]], + dims=["extra", "dim_2"], + coords={"dim_2": [0, 1]}, + ) + c = v <= rhs + assert "extra" in c.dims + + def test_subset_constraint_solve_integration(self) -> None: + if not available_solvers: + pytest.skip("No solver available") + solver = "highs" if "highs" in available_solvers else available_solvers[0] + 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=solver) + 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) diff --git a/test/test_dtypes.py b/test/test_dtypes.py new file mode 100644 index 000000000..b30c7eac3 --- /dev/null +++ b/test/test_dtypes.py @@ -0,0 +1,75 @@ +"""Tests for int32 default label dtype.""" + +import numpy as np +import pytest + +from linopy import Model +from linopy.config import options + + +def test_default_label_dtype_is_int32() -> None: + assert options["label_dtype"] == np.int32 + + +def test_variable_labels_are_int32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + assert x.labels.dtype == np.int32 + + +def test_constraint_labels_are_int32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + m.add_constraints(x >= 1, name="c") + assert m.constraints["c"].labels.dtype == np.int32 + + +def test_expression_vars_are_int32() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + expr = 2 * x + 1 + assert expr.vars.dtype == np.int32 + + +@pytest.mark.skipif( + not pytest.importorskip("highspy", reason="highspy not installed"), + reason="highspy not installed", +) +def test_solve_with_int32_labels() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, name="x") + y = m.add_variables(lower=0, upper=10, name="y") + m.add_constraints(x + y <= 15, name="c1") + m.add_objective(x + 2 * y, sense="max") + m.solve("highs") + assert m.objective.value == pytest.approx(25.0) + + +def test_overflow_guard_variables() -> None: + m = Model() + m._xCounter = np.iinfo(np.int32).max - 1 + with pytest.raises(ValueError, match="exceeds the maximum"): + m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + + +def test_overflow_guard_constraints() -> None: + m = Model() + x = m.add_variables(lower=0, upper=1, coords=[range(5)], name="x") + m._cCounter = np.iinfo(np.int32).max - 1 + with pytest.raises(ValueError, match="exceeds the maximum"): + m.add_constraints(x >= 0, name="c") + + +def test_label_dtype_option_int64() -> None: + with options: + options["label_dtype"] = np.int64 + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[range(5)], name="x") + assert x.labels.dtype == np.int64 + expr = 2 * x + 1 + assert expr.vars.dtype == np.int64 + + +def test_label_dtype_rejects_invalid() -> None: + with pytest.raises(ValueError, match="label_dtype must be one of"): + options["label_dtype"] = np.float64 diff --git a/test/test_dualization.py b/test/test_dualization.py new file mode 100644 index 000000000..c819b3fe9 --- /dev/null +++ b/test/test_dualization.py @@ -0,0 +1,608 @@ +"""Tests for linopy/dualization.py - LP dualization.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr +from _pytest.logging import LogCaptureFixture + +from linopy import Model +from linopy.dualization import ( + _build_label_to_flat_index_lookup, + _build_obj_coeff_lookup, + _dual_bounds_from_constraint_signs, + _lookup_flat_indices, + _skip, + _term_slots_for_sorted_flat_indices, + dualize, +) +from linopy.solvers import licensed_solvers + +_lp_solver = next((s for s in ("highs", "glpk", "scip") if s in licensed_solvers), None) +needs_solver = pytest.mark.skipif(_lp_solver is None, reason="No LP solver available") + + +# Structural tests for important internal functions +def test_skip_empty_label_array() -> None: + """Empty label arrays are skipped.""" + labels = xr.DataArray(np.array([], dtype=np.int64), dims=["dim_0"]) + + assert _skip(labels, "variable", "x") + + +def test_skip_fully_masked_label_array() -> None: + """Fully masked label arrays are skipped.""" + labels = xr.DataArray(np.array([-1, -1], dtype=np.int64), dims=["dim_0"]) + + assert _skip(labels, "constraint", "c") + + +def test_build_label_to_flat_index_lookup() -> None: + """Flat labels are mapped to their positions in the flattened label array.""" + labels = np.array([10, -1, 12, 99], dtype=np.int64) + + lookup = _build_label_to_flat_index_lookup(labels) + + assert lookup[10] == 0 + assert lookup[12] == 2 + assert lookup[99] == 3 + assert lookup[11] == -1 + + +def test_build_label_to_flat_index_lookup_all_masked() -> None: + """An all-masked label array produces an empty lookup.""" + labels = np.array([-1, -1], dtype=np.int64) + + lookup = _build_label_to_flat_index_lookup(labels) + + assert lookup.dtype == np.int64 + assert len(lookup) == 0 + + +def test_lookup_flat_indices_bounds_safe() -> None: + """Out-of-range and masked labels safely map to -1.""" + lookup = np.array([-1, -1, 5, -1, 7], dtype=np.int64) + labels = np.array([-1, 0, 2, 4, 99], dtype=np.int64) + + flat_indices = _lookup_flat_indices(labels, lookup) + + np.testing.assert_array_equal( + flat_indices, + np.array([-1, -1, 5, 7, -1], dtype=np.int64), + ) + + +def test_term_slots_for_sorted_flat_indices() -> None: + """Repeated sorted flat indices are assigned increasing term slots.""" + sorted_flat_indices = np.array([2, 2, 2, 5, 5, 9], dtype=np.int64) + + slots = _term_slots_for_sorted_flat_indices(sorted_flat_indices) + + np.testing.assert_array_equal( + slots, + np.array([0, 1, 2, 0, 1, 0], dtype=np.int64), + ) + + +def test_build_obj_coeff_lookup_all_masked() -> None: + """All-masked variable labels produce an empty objective-coefficient lookup.""" + lookup = _build_obj_coeff_lookup( + np.array([-1, -1], dtype=np.int64), + np.array([1.0, 2.0], dtype=np.float64), + ) + + assert lookup.dtype == np.float64 + assert len(lookup) == 0 + + +def test_dual_bounds_from_mixed_constraint_signs_min() -> None: + """Mixed signs produce elementwise dual bounds for a minimization primal.""" + idx = pd.RangeIndex(3, name="i") + labels = xr.DataArray([0, 1, 2], dims=["i"], coords={"i": idx}) + signs = xr.DataArray(["<=", ">=", "="], dims=["i"], coords={"i": idx}) + + lower, upper, valid_sign = _dual_bounds_from_constraint_signs( + signs, + labels, + primal_is_min=True, + ) + + np.testing.assert_allclose(lower.values, np.array([-np.inf, 0.0, -np.inf])) + np.testing.assert_allclose(upper.values, np.array([0.0, np.inf, np.inf])) + np.testing.assert_array_equal(valid_sign.values, np.array([True, True, True])) + + +def test_dual_bounds_from_mixed_constraint_signs_max() -> None: + """Mixed signs produce elementwise dual bounds for a maximization primal.""" + idx = pd.RangeIndex(3, name="i") + labels = xr.DataArray([0, 1, 2], dims=["i"], coords={"i": idx}) + signs = xr.DataArray(["<=", ">=", "="], dims=["i"], coords={"i": idx}) + + lower, upper, valid_sign = _dual_bounds_from_constraint_signs( + signs, + labels, + primal_is_min=False, + ) + + np.testing.assert_allclose(lower.values, np.array([0.0, -np.inf, -np.inf])) + np.testing.assert_allclose(upper.values, np.array([np.inf, 0.0, np.inf])) + np.testing.assert_array_equal(valid_sign.values, np.array([True, True, True])) + + +# Structural tests (no solver required) +def test_dualize_empty_model() -> None: + """Dualizing an empty model returns an empty dual model.""" + m = Model() + m_dual = dualize(m) + assert len(m_dual.variables) == 0 + assert len(m_dual.constraints) == 0 + + +def test_only_lower_bound_lifted_to_dual_variable() -> None: + """A finite lower bound is lifted even when the upper bound is infinite.""" + m = Model() + x = m.add_variables(lower=1, upper=np.inf, name="x") + m.add_objective(x) + + m_dual = dualize(m) + + assert "x-bound-lower" in m_dual.variables + assert "x-bound-upper" not in m_dual.variables + + +def test_only_upper_bound_lifted_to_dual_variable() -> None: + """A finite upper bound is lifted even when the lower bound is infinite.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=3, name="x") + m.add_objective(x) + + m_dual = dualize(m) + + assert "x-bound-lower" not in m_dual.variables + assert "x-bound-upper" in m_dual.variables + + +def test_unbounded_variable_bounds_do_not_create_dual_variables() -> None: + """Infinite variable bounds are not lifted into dual variables.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + m.add_constraints(x == 1, name="c") + m.add_objective(x) + + m_dual = dualize(m) + + assert "x-bound-lower" not in m_dual.variables + assert "x-bound-upper" not in m_dual.variables + + +def test_dualize_model_with_variables_but_no_constraints_or_finite_bounds() -> None: + """A model with variables but no constraints or finite bounds returns an empty dual.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + m.add_objective(x) + + m_dual = dualize(m) + + assert len(m_dual.variables) == 0 + assert len(m_dual.constraints) == 0 + + +def test_variable_bounds_lifted_to_dual_variables() -> None: + """Finite variable bounds are lifted and dualized.""" + m = Model() + x = m.add_variables(lower=1, upper=3, name="x") + m.add_objective(x) + + m_dual = dualize(m) + + assert "x-bound-lower" in m_dual.variables + assert "x-bound-upper" in m_dual.variables + + +def test_dual_variables_named_after_primal_constraints() -> None: + """Dual variables use the names of their corresponding primal constraints.""" + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3)], name="x") + m.add_constraints(x >= 1.0, name="lb_con") + m.add_objective(2.0 * x) + + m_dual = dualize(m) + assert "lb_con" in m_dual.variables + assert "x-bound-lower" in m_dual.variables + + +def test_dual_feasibility_constraints_named_after_primal_variables() -> None: + """Dual-feasibility constraints use the names of the primal variables.""" + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3)], name="x") + m.add_constraints(x >= 1.0, name="lb_con") + m.add_objective(2.0 * x) + + m_dual = dualize(m) + assert "x" in m_dual.constraints + + +def test_dual_objective_sense_flipped_min() -> None: + """A minimization primal produces a maximization dual.""" + m = Model() + x = m.add_variables(lower=0, name="x") + m.add_constraints(x >= 1, name="c") + m.add_objective(x) # min + + m_dual = dualize(m) + assert m_dual.objective.sense == "max" + + +def test_dual_objective_sense_flipped_max() -> None: + """A maximization primal produces a minimization dual.""" + m = Model() + x = m.add_variables(upper=10, name="x") + m.add_constraints(x <= 5, name="c") + m.add_objective(x, sense="max") + + m_dual = dualize(m) + assert m_dual.objective.sense == "min" + + +def test_dual_sign_conventions_min() -> None: + """For a min primal: <= -> dual <= 0, >= -> dual >= 0, = -> dual free.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + m.add_constraints(x == 5, name="eq") + m.add_constraints(x <= 10, name="leq") + m.add_constraints(x >= -10, name="geq") + m.add_objective(x) + + m_dual = dualize(m) + assert m_dual.variables["eq"].lower.item() == -np.inf + assert m_dual.variables["eq"].upper.item() == np.inf + assert m_dual.variables["leq"].upper.item() == 0 + assert m_dual.variables["geq"].lower.item() == 0 + + +def test_dual_sign_conventions_max() -> None: + """For a max primal: <= -> dual >= 0, >= -> dual <= 0.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + m.add_constraints(x <= 10, name="leq") + m.add_constraints(x >= -10, name="geq") + m.add_objective(x, sense="max") + + m_dual = dualize(m) + assert m_dual.variables["leq"].lower.item() == 0 + assert m_dual.variables["geq"].upper.item() == 0 + + +def test_dual_feasibility_rhs_equals_objective_coefficients() -> None: + """The dual feasibility RHS equals the primal objective coefficient.""" + m = Model() + x = m.add_variables( + lower=-np.inf, upper=np.inf, coords=[pd.RangeIndex(4)], name="x" + ) + m.add_constraints(x == np.array([1.0, 2.0, 3.0, 4.0]), name="eq") + c = np.array([10.0, 20.0, 30.0, 40.0]) + m.add_objective(c * x) + + m_dual = dualize(m) + np.testing.assert_allclose(m_dual.constraints["x"].rhs.values, c) + + +def test_dual_multi_constraint_per_variable() -> None: + """A variable appearing in k constraints gets k dual-feasibility terms.""" + m = Model() + n = 3 + x = m.add_variables( + lower=-np.inf, upper=np.inf, coords=[pd.RangeIndex(n)], name="x" + ) + # Two equality constraints, both using x + m.add_constraints(x == 1.0, name="c1") + m.add_constraints(2.0 * x == 2.0, name="c2") + m.add_objective(5.0 * x) + + m_dual = dualize(m) + # dual feas constraint for x: lambda_c1 + 2*lambda_c2 = 5 + con_x = m_dual.constraints["x"] + # Should have 2 terms (one from c1, one from c2) + n_terms = (con_x.vars != -1).sum(dim="_term").max().item() + assert n_terms == 2 + + +def test_dual_with_masked_variable() -> None: + """Partially masked variables only produce constraints for unmasked elements.""" + m = Model() + mask = xr.DataArray([True, False, True], dims=["dim_0"]) + x = m.add_variables(lower=0, coords=[pd.RangeIndex(3)], name="x", mask=mask) + m.add_constraints(x >= 1.0, name="c", mask=mask) + m.add_objective(x) + + m_dual = dualize(m) + assert "x" in m_dual.constraints + assert ( + m_dual.constraints["x"].labels != -1 + ).sum().item() == 2 # only 2 unmasked elements + + +def test_dual_free_unconstrained_zero_cost_variable_no_error() -> None: + """A disconnected zero-cost free variable is skipped.""" + m = Model() + m.add_variables( + lower=-np.inf, upper=np.inf, name="x" + ) # no bounds → no bound constraints + y = m.add_variables(lower=-np.inf, upper=np.inf, name="y") + m.add_constraints(y == 5, name="eq") + m.add_objective(y) # x has no connections at all + + m_dual = dualize(m) # must not raise + assert "y" in m_dual.constraints + assert "x" not in m_dual.constraints # no constraint connections + + +def test_dual_free_unconstrained_variable_with_objective_warns( + caplog: LogCaptureFixture, +) -> None: + """A disconnected variable with nonzero objective coefficient is reported.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + y = m.add_variables(lower=-np.inf, upper=np.inf, name="y") + + m.add_constraints(y == 5, name="eq") + m.add_objective(x + y) + + with caplog.at_level("WARNING"): + m_dual = dualize(m) + + assert "y" in m_dual.constraints + assert "x" not in m_dual.constraints + assert any( + "corresponding dual-feasibility condition is infeasible" in rec.message + for rec in caplog.records + ) + + +def test_dual_multi_constraint_per_variable_coefficients() -> None: + """Dual-feasibility terms keep the correct A-matrix coefficients.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + + m.add_constraints(x == 1.0, name="c1") + m.add_constraints(2.0 * x == 2.0, name="c2") + m.add_objective(5.0 * x) + + m_dual = dualize(m) + coeffs = m_dual.constraints["x"].coeffs.values.ravel() + coeffs = coeffs[m_dual.constraints["x"].vars.values.ravel() != -1] + + np.testing.assert_allclose(np.sort(coeffs), np.array([1.0, 2.0])) + + +def test_dual_mixed_sign_constraint_array_min() -> None: + """Mixed elementwise constraint signs produce elementwise dual bounds.""" + m = Model() + idx = pd.RangeIndex(3, name="i") + x = m.add_variables(lower=-np.inf, upper=np.inf, coords=[idx], name="x") + + signs = xr.DataArray(["<=", ">=", "="], dims=["i"], coords={"i": idx}) + rhs = xr.DataArray([1.0, 2.0, 3.0], dims=["i"], coords={"i": idx}) + + m.add_constraints(x, signs, rhs, name="mixed") + m.add_objective(x.sum()) + + m_dual = dualize(m) + dv = m_dual.variables["mixed"] + + np.testing.assert_allclose( + dv.lower.values, + np.array([-np.inf, 0.0, -np.inf]), + ) + np.testing.assert_allclose( + dv.upper.values, + np.array([0.0, np.inf, np.inf]), + ) + + +def test_dual_mixed_sign_constraint_array_max() -> None: + """Mixed elementwise constraint signs follow maximization dual bounds.""" + m = Model() + idx = pd.RangeIndex(3, name="i") + x = m.add_variables(lower=-np.inf, upper=np.inf, coords=[idx], name="x") + + signs = xr.DataArray(["<=", ">=", "="], dims=["i"], coords={"i": idx}) + rhs = xr.DataArray([1.0, 2.0, 3.0], dims=["i"], coords={"i": idx}) + + m.add_constraints(x, signs, rhs, name="mixed") + m.add_objective(x.sum(), sense="max") + + m_dual = dualize(m) + dv = m_dual.variables["mixed"] + + np.testing.assert_allclose( + dv.lower.values, + np.array([0.0, -np.inf, -np.inf]), + ) + np.testing.assert_allclose( + dv.upper.values, + np.array([np.inf, 0.0, np.inf]), + ) + + +def test_mixed_variable_bounds_lift_only_finite_entries() -> None: + """Mixed finite/infinite bounds are lifted only for finite entries.""" + m = Model() + idx = pd.RangeIndex(3, name="i") + + lower = xr.DataArray([0.0, -np.inf, -np.inf], dims=["i"], coords={"i": idx}) + upper = xr.DataArray([np.inf, 5.0, np.inf], dims=["i"], coords={"i": idx}) + + x = m.add_variables(lower=lower, upper=upper, coords=[idx], name="x") + m.add_objective(x.sum()) + + m_dual = dualize(m) + + lower_labels = m_dual.variables["x-bound-lower"].labels + upper_labels = m_dual.variables["x-bound-upper"].labels + + assert (lower_labels != -1).sum().item() == 1 + assert (upper_labels != -1).sum().item() == 1 + + assert lower_labels.sel(i=0).item() != -1 + assert upper_labels.sel(i=1).item() != -1 + + +# Numerical tests (require solver) +def _solve(model: Model, **kwargs: Any) -> float: + """Solve a model with the available LP solver and return its objective value.""" + assert _lp_solver is not None + model.solve(solver_name=_lp_solver, io_api="lp", **kwargs) + + value = model.objective.value + assert value is not None + return float(value) + + +@needs_solver +def test_strong_duality_simple() -> None: + """Strong duality: primal obj == dual obj at optimality.""" + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_constraints(x + y >= 3, name="c1") + m.add_constraints(x + 2 * y >= 4, name="c2") + m.add_objective(5 * x + 4 * y) + + primal_obj = _solve(m) + dual_obj = _solve(m.dualize()) + assert abs(primal_obj - dual_obj) < 1e-5 + + +@needs_solver +def test_strong_duality_array_variable() -> None: + """Strong duality with array variables.""" + rng = np.random.default_rng(0) + n = 6 + + A = rng.random((4, n)) + 0.1 + b = rng.random(4) + 1.0 + c_obj = rng.random(n) + 0.1 + + m = Model() + x = m.add_variables(lower=0, coords=[pd.RangeIndex(n)], name="x") + + for i in range(4): + lhs: Any = sum(float(A[i, j]) * x[j] for j in range(n)) + m.add_constraints(lhs <= float(b[i]), name=f"c{i}") + + # Bounded because x >= 0 and A x <= b with A > 0. + obj: Any = sum(float(c_obj[j]) * x[j] for j in range(n)) + m.add_objective(obj, sense="max") + + primal_obj = _solve(m) + dual_obj = _solve(m.dualize()) + assert abs(primal_obj - dual_obj) < 1e-4 + + +@needs_solver +def test_strong_duality_mixed_constraint_types() -> None: + """Strong duality with =, <=, >= constraints.""" + m = Model() + x = m.add_variables(lower=-np.inf, upper=np.inf, name="x") + y = m.add_variables(lower=-np.inf, upper=np.inf, name="y") + m.add_constraints(x + y == 5, name="eq") + m.add_constraints(x - y <= 2, name="leq") + m.add_constraints(x >= 0, name="geq") + m.add_objective(2 * x + 3 * y) + + primal_obj = _solve(m) + dual_obj = _solve(m.dualize()) + assert abs(primal_obj - dual_obj) < 1e-5 + + +@needs_solver +def test_strong_duality_maximization() -> None: + """Strong duality for a maximization primal.""" + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_constraints(x + y <= 10, name="c1") + m.add_constraints(x <= 6, name="c2") + m.add_constraints(y <= 8, name="c3") + m.add_objective(3 * x + 5 * y, sense="max") + + primal_obj = _solve(m) + dual_obj = _solve(m.dualize()) + assert abs(primal_obj - dual_obj) < 1e-5 + + +@needs_solver +def test_dualize_mixed_signs_and_mixed_variable_bounds() -> None: + """Dualization handles mixed constraint signs and mixed variable bounds.""" + m = Model() + idx = pd.RangeIndex(3, name="i") + + lower = xr.DataArray([0.0, -np.inf, -np.inf], dims=["i"], coords={"i": idx}) + upper = xr.DataArray([np.inf, 5.0, np.inf], dims=["i"], coords={"i": idx}) + + x = m.add_variables( + lower=lower, + upper=upper, + coords=[idx], + name="x", + ) + + signs = xr.DataArray(["<=", ">=", "="], dims=["i"], coords={"i": idx}) + rhs = xr.DataArray([4.0, 2.0, 3.0], dims=["i"], coords={"i": idx}) + + m.add_constraints(x, signs, rhs, name="mixed") + m.add_objective(x.sum()) + + primal_obj = _solve(m) + + m_dual = dualize(m) + dual_obj = _solve(m_dual) + + assert abs(primal_obj - dual_obj) < 1e-5 + + mixed_dual = m_dual.variables["mixed"] + + np.testing.assert_allclose( + mixed_dual.lower.values, + np.array([-np.inf, 0.0, -np.inf]), + ) + np.testing.assert_allclose( + mixed_dual.upper.values, + np.array([0.0, np.inf, np.inf]), + ) + + assert "x-bound-lower" in m_dual.variables + assert "x-bound-upper" in m_dual.variables + + lower_bound_labels = m_dual.variables["x-bound-lower"].labels + upper_bound_labels = m_dual.variables["x-bound-upper"].labels + + assert (lower_bound_labels != -1).sum().item() == 1 + assert (upper_bound_labels != -1).sum().item() == 1 + + assert lower_bound_labels.sel(i=0).item() != -1 + assert upper_bound_labels.sel(i=1).item() != -1 + + con_x = m_dual.constraints["x"] + + # x[0]: mixed[0] + x-bound-lower[0] = 1 + # x[1]: mixed[1] + x-bound-upper[1] = 1 + # x[2]: mixed[2] = 1 + n_terms = (con_x.vars != -1).sum(dim="_term") + np.testing.assert_array_equal( + n_terms.values, + np.array([2, 2, 1]), + ) + + # The dual values from the standalone dual model should match Linopy's + # primal constraint duals for the original mixed constraint block. + np.testing.assert_allclose( + m.constraints["mixed"].dual.values, + m_dual.variables["mixed"].solution.values, + atol=1e-6, + ) diff --git a/test/test_fix_relax.py b/test/test_fix_relax.py new file mode 100644 index 000000000..18031ea3a --- /dev/null +++ b/test/test_fix_relax.py @@ -0,0 +1,629 @@ +"""Tests for Variable.fix(), Variable.unfix(), and Variable.fixed.""" + +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from xarray import DataArray + +from linopy import Model, Variable +from linopy.types import ConstantLike + + +@pytest.fixture +def model_with_solution() -> Model: + """Create a simple model and simulate a solution.""" + m = Model() + x = m.add_variables(lower=0, upper=10, name="x") + y = m.add_variables(lower=-5, upper=5, coords=[pd.Index([0, 1])], name="y") + z = m.add_variables(binary=True, name="z") + w = m.add_variables(lower=0, upper=100, integer=True, name="w") + + # Simulate solution values + x.solution = 3.14159265 + y.solution = DataArray([2.71828, -1.41421], dims="dim_0") + z.solution = 0.9999999997 + w.solution = 41.9999999998 + m._status = "ok" + m._termination_condition = "optimal" + + return m + + +SCALAR_VALUES: list = [ + pytest.param(5, id="int"), + pytest.param(5.0, id="float"), + pytest.param(np.float64(5.0), id="np.float64"), + pytest.param(np.int64(5), id="np.int64"), + pytest.param(np.array(5.0), id="np.0d-array"), + pytest.param(DataArray(5.0), id="DataArray"), +] + +ARRAY_VALUES: list = [ + pytest.param([2.5, -1.5], id="list"), + pytest.param(np.array([2.5, -1.5]), id="np.array"), + pytest.param(DataArray([2.5, -1.5], dims="dim_0"), id="DataArray"), + pytest.param(pd.Series([2.5, -1.5]), id="pd.Series"), +] + + +class TestVariableFix: + @pytest.mark.parametrize("value", SCALAR_VALUES) + def test_fix_scalar_dtypes(self, model_with_solution: Model, value: object) -> None: + m = model_with_solution + m.variables["x"].fix(value=value) + assert m.variables["x"].fixed + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 5.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 5.0) + + @pytest.mark.parametrize("value", ARRAY_VALUES) + def test_fix_array_dtypes(self, model_with_solution: Model, value: object) -> None: + m = model_with_solution + m.variables["y"].fix(value=value) + assert m.variables["y"].fixed + np.testing.assert_array_almost_equal(m.variables["y"].lower.values, [2.5, -1.5]) + np.testing.assert_array_almost_equal(m.variables["y"].upper.values, [2.5, -1.5]) + + def test_fix_uses_solution(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix() + assert m.variables["x"].fixed + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 3.14159265) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 3.14159265) + + def test_fix_does_not_add_constraint(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + assert len(m.constraints) == 0 + + def test_fix_rounds_binary(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix() + np.testing.assert_equal(m.variables["z"].lower.item(), 1.0) + np.testing.assert_equal(m.variables["z"].upper.item(), 1.0) + + def test_fix_binary_to_zero(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix(value=0) + np.testing.assert_equal(m.variables["z"].lower.item(), 0.0) + np.testing.assert_equal(m.variables["z"].upper.item(), 0.0) + + @pytest.mark.parametrize("value", [5, 0.4, 0.6, -1]) + def test_fix_binary_outside_domain_raises( + self, model_with_solution: Model, value: float + ) -> None: + m = model_with_solution + with pytest.raises(ValueError, match="binary variable"): + m.variables["z"].fix(value=value) + + def test_fix_rounds_integer(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].fix() + np.testing.assert_equal(m.variables["w"].lower.item(), 42.0) + np.testing.assert_equal(m.variables["w"].upper.item(), 42.0) + + def test_fix_rounds_continuous(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(decimals=4) + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 3.1416, decimal=4) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 3.1416, decimal=4) + + def test_fix_above_upper_bound_warns_and_overrides( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + with pytest.warns(UserWarning, match="outside its current"): + m.variables["x"].fix(value=11.0) + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 11.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 11.0) + + def test_fix_below_lower_bound_warns_and_overrides( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + with pytest.warns(UserWarning, match="outside its current"): + m.variables["x"].fix(value=-1.0) + np.testing.assert_almost_equal(m.variables["x"].lower.item(), -1.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), -1.0) + + def test_fix_within_bounds_does_not_warn(self, model_with_solution: Model) -> None: + m = model_with_solution + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + m.variables["x"].fix(value=5.0) + + def test_fix_small_overshoot_rounded_within_bounds( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + m.variables["x"].fix(value=10.0000000001) + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 10.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 10.0) + + def test_fix_raises_if_already_fixed_no_overwrite( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + m.variables["x"].fix(value=3.0) + with pytest.raises(ValueError, match="already fixed"): + m.variables["x"].fix(value=5.0, overwrite=False) + + def test_fix_overwrite_keeps_original_stashed_bounds( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + m.variables["x"].fix(value=3.0) + m.variables["x"].fix(value=5.0, overwrite=True) + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 5.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 5.0) + m.variables["x"].unfix() + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 0.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 10.0) + + def test_fix_multidimensional(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["y"].fix() + assert m.variables["y"].fixed + np.testing.assert_array_almost_equal( + m.variables["y"].lower.values, [2.71828, -1.41421] + ) + np.testing.assert_array_almost_equal( + m.variables["y"].upper.values, [2.71828, -1.41421] + ) + + +class TestVariableUnfix: + def test_unfix_restores_bounds(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + m.variables["x"].unfix() + assert not m.variables["x"].fixed + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 0.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 10.0) + + def test_unfix_restores_multidimensional_bounds( + self, model_with_solution: Model + ) -> None: + m = model_with_solution + m.variables["y"].fix() + m.variables["y"].unfix() + np.testing.assert_array_almost_equal(m.variables["y"].lower.values, [-5, -5]) + np.testing.assert_array_almost_equal(m.variables["y"].upper.values, [5, 5]) + + def test_unfix_restores_binary_bounds(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix() + m.variables["z"].unfix() + np.testing.assert_equal(m.variables["z"].lower.item(), 0.0) + np.testing.assert_equal(m.variables["z"].upper.item(), 1.0) + + def test_unfix_noop_if_not_fixed(self, model_with_solution: Model) -> None: + m = model_with_solution + # Should not raise + m.variables["x"].unfix() + assert not m.variables["x"].fixed + + +class TestFixNoSolution: + def test_fix_without_solution_raises(self) -> None: + m = Model() + m.add_variables(lower=0, upper=10, name="x") + with pytest.raises(ValueError, match="no solution value available"): + m.variables["x"].fix() + + +class TestUnfixDoesNotUnrelaxIndependently: + def test_unfix_on_relaxed_only_variable(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + m.variables["z"].unfix() + assert m.variables["z"].relaxed + assert not m.variables["z"].attrs["binary"] + + +class TestFixThenRelax: + """Test the combined fix() + relax() workflow (fix first, then relax).""" + + def test_fix_then_relax_binary(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix() + m.variables["z"].relax() + assert not m.variables["z"].attrs["binary"] + assert m.variables["z"].fixed + assert m.variables["z"].relaxed + + def test_unfix_does_not_unrelax(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix() + m.variables["z"].relax() + m.variables["z"].unfix() + assert not m.variables["z"].fixed + # unfix restores the original binary bounds regardless of relaxation + np.testing.assert_equal(m.variables["z"].lower.item(), 0.0) + np.testing.assert_equal(m.variables["z"].upper.item(), 1.0) + # relaxation is independent — still in effect + assert m.variables["z"].relaxed + assert not m.variables["z"].attrs["binary"] + # explicit unrelax needed + m.variables["z"].unrelax() + assert m.variables["z"].attrs["binary"] + assert not m.variables["z"].relaxed + + def test_fix_then_relax_integer(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].fix() + m.variables["w"].relax() + assert not m.variables["w"].attrs["integer"] + assert m.variables["w"].fixed + assert m.variables["w"].relaxed + + def test_unfix_does_not_unrelax_integer(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].fix() + m.variables["w"].relax() + m.variables["w"].unfix() + assert not m.variables["w"].fixed + assert m.variables["w"].relaxed + assert not m.variables["w"].attrs["integer"] + + +class TestVariableFixed: + def test_fixed_false_initially(self, model_with_solution: Model) -> None: + m = model_with_solution + assert not m.variables["x"].fixed + + def test_fixed_true_after_fix(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + assert m.variables["x"].fixed + + def test_fixed_false_after_unfix(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + m.variables["x"].unfix() + assert not m.variables["x"].fixed + + +class TestVariablesContainerFixUnfix: + def test_fix_all(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.fix() + for name in m.variables: + assert m.variables[name].fixed + + def test_unfix_all(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.fix() + m.variables.unfix() + for name in m.variables: + assert not m.variables[name].fixed + np.testing.assert_almost_equal(m.variables["x"].lower.item(), 0.0) + np.testing.assert_almost_equal(m.variables["x"].upper.item(), 10.0) + + def test_fix_integers_only(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.integers.fix() + assert m.variables["w"].fixed + assert not m.variables["x"].fixed + + def test_fix_binaries_only(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.binaries.fix() + assert m.variables["z"].fixed + assert not m.variables["x"].fixed + + def test_fixed_returns_container(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + result = m.variables.fixed + assert "x" in result + assert "y" not in result + + def test_fix_then_relax_integers(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.integers.fix() + m.variables.integers.relax() + assert not m.variables["w"].attrs["integer"] + assert m.variables["w"].fixed + m.variables["w"].unfix() + assert not m.variables["w"].attrs["integer"] # still relaxed + m.variables["w"].unrelax() + assert m.variables["w"].attrs["integer"] + + +class TestVariableRelax: + def test_relax_binary(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + assert not m.variables["z"].attrs["binary"] + assert not m.variables["z"].attrs["integer"] + assert m.variables["z"].relaxed + assert m._relaxed_registry["z"] == "binary" + + def test_relax_integer(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].relax() + assert not m.variables["w"].attrs["integer"] + assert not m.variables["w"].attrs["binary"] + assert m.variables["w"].relaxed + assert m._relaxed_registry["w"] == "integer" + + def test_relax_continuous_noop(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].relax() + assert "x" not in m._relaxed_registry + assert not m.variables["x"].relaxed + + def test_relax_semi_continuous_raises(self) -> None: + m = Model() + m.add_variables(lower=1, upper=10, semi_continuous=True, name="sc") + with pytest.raises(NotImplementedError, match="semi-continuous"): + m.variables["sc"].relax() + + def test_unrelax_binary(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + m.variables["z"].unrelax() + assert m.variables["z"].attrs["binary"] + assert not m.variables["z"].relaxed + assert "z" not in m._relaxed_registry + + def test_unrelax_integer(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].relax() + m.variables["w"].unrelax() + assert m.variables["w"].attrs["integer"] + assert not m.variables["w"].relaxed + assert "w" not in m._relaxed_registry + + def test_unrelax_noop_if_not_relaxed(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].unrelax() + assert not m.variables["x"].relaxed + + def test_relax_preserves_binary_bounds(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + assert float(m.variables["z"].lower) == 0.0 + assert float(m.variables["z"].upper) == 1.0 + + def test_relax_preserves_integer_bounds(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["w"].relax() + assert float(m.variables["w"].lower) == 0.0 + assert float(m.variables["w"].upper) == 100.0 + + +class TestVariablesContainerRelax: + def test_relax_all(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.relax() + assert not m.variables["z"].attrs["binary"] + assert not m.variables["w"].attrs["integer"] + assert m.variables["z"].relaxed + assert m.variables["w"].relaxed + # Continuous variables unaffected + assert not m.variables["x"].relaxed + + def test_unrelax_all(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.relax() + m.variables.unrelax() + assert m.variables["z"].attrs["binary"] + assert m.variables["w"].attrs["integer"] + assert not m.variables["z"].relaxed + assert not m.variables["w"].relaxed + + def test_relax_integers_only(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.integers.relax() + assert not m.variables["w"].attrs["integer"] + assert m.variables["w"].relaxed + # Binary should be untouched + assert m.variables["z"].attrs["binary"] + assert not m.variables["z"].relaxed + + def test_relax_binaries_only(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.binaries.relax() + assert not m.variables["z"].attrs["binary"] + assert m.variables["z"].relaxed + # Integer should be untouched + assert m.variables["w"].attrs["integer"] + assert not m.variables["w"].relaxed + + def test_relaxed_returns_container(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + result = m.variables.relaxed + assert "z" in result + assert "x" not in result + + def test_relax_with_semi_continuous_raises(self) -> None: + m = Model() + m.add_variables(lower=0, upper=10, name="x") + m.add_variables(lower=1, upper=10, semi_continuous=True, name="sc") + with pytest.raises(NotImplementedError, match="semi-continuous"): + m.variables.relax() + + def test_relaxed_view_unrelax(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables.relax() + assert len(m.variables.relaxed) == 2 + m.variables.relaxed.unrelax() + assert len(m.variables.relaxed) == 0 + assert m.variables["z"].attrs["binary"] + assert m.variables["w"].attrs["integer"] + + def test_fixed_view_unfix(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + m.variables["z"].fix() + assert len(m.variables.fixed) == 2 + m.variables.fixed.unfix() + assert len(m.variables.fixed) == 0 + + def test_double_relax_is_idempotent(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].relax() + m.variables["z"].relax() + assert m._relaxed_registry["z"] == "binary" + m.variables["z"].unrelax() + assert m.variables["z"].attrs["binary"] + + def test_relax_all_converts_milp_to_lp(self, model_with_solution: Model) -> None: + m = model_with_solution + assert m.type == "MILP" + m.variables.relax() + assert m.type == "LP" + m.variables.unrelax() + assert m.type == "MILP" + + +class TestRemoveVariablesCleansUpFix: + def test_remove_fixed_variable(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + m.remove_variables("x") + assert "x" not in m.variables + + def test_remove_relaxed_variable(self, model_with_solution: Model) -> None: + m = model_with_solution + m.variables["z"].fix() + m.variables["z"].relax() + m.remove_variables("z") + assert "z" not in m._relaxed_registry + assert "z" not in m.variables + + +class TestFixIO: + def test_fixed_bounds_survive_netcdf( + self, model_with_solution: Model, tmp_path: Path + ) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + m.variables["y"].fix() + + path = tmp_path / "model.nc" + m.to_netcdf(path) + + from linopy.io import read_netcdf + + m2 = read_netcdf(path) + assert m2.variables["x"].fixed + assert m2.variables["y"].fixed + np.testing.assert_almost_equal(m2.variables["x"].lower.item(), 5.0) + np.testing.assert_almost_equal(m2.variables["x"].upper.item(), 5.0) + + def test_unfix_after_roundtrip_restores_bounds( + self, model_with_solution: Model, tmp_path: Path + ) -> None: + m = model_with_solution + m.variables["x"].fix(value=5.0) + + path = tmp_path / "model.nc" + m.to_netcdf(path) + + from linopy.io import read_netcdf + + m2 = read_netcdf(path) + m2.variables["x"].unfix() + assert not m2.variables["x"].fixed + np.testing.assert_almost_equal(m2.variables["x"].lower.item(), 0.0) + np.testing.assert_almost_equal(m2.variables["x"].upper.item(), 10.0) + + def test_relaxed_registry_survives_netcdf( + self, model_with_solution: Model, tmp_path: Path + ) -> None: + m = model_with_solution + m.variables["z"].fix() + m.variables["z"].relax() + m.variables["w"].fix() + m.variables["w"].relax() + + path = tmp_path / "model.nc" + m.to_netcdf(path) + + from linopy.io import read_netcdf + + m2 = read_netcdf(path) + assert m2._relaxed_registry == {"z": "binary", "w": "integer"} + assert m2.variables["z"].fixed + assert m2.variables["w"].fixed + + def test_empty_registry_netcdf( + self, model_with_solution: Model, tmp_path: Path + ) -> None: + m = model_with_solution + path = tmp_path / "model.nc" + m.to_netcdf(path) + + from linopy.io import read_netcdf + + m2 = read_netcdf(path) + assert m2._relaxed_registry == {} + + def test_unrelax_after_roundtrip( + self, model_with_solution: Model, tmp_path: Path + ) -> None: + m = model_with_solution + m.variables["z"].relax() + + path = tmp_path / "model.nc" + m.to_netcdf(path) + + from linopy.io import read_netcdf + + m2 = read_netcdf(path) + m2.variables["z"].unrelax() + assert m2.variables["z"].attrs["binary"] + assert "z" not in m2._relaxed_registry + + +TIME = pd.Index([2020, 2030, 2040], name="time") + +ALIGNED_VALUES = [ + pytest.param(5.0, [5.0, 5.0, 5.0], id="scalar-broadcast"), + pytest.param([1.0, 2.0, 3.0], [1.0, 2.0, 3.0], id="list"), + pytest.param(np.array([1.0, 2.0, 3.0]), [1.0, 2.0, 3.0], id="ndarray"), + pytest.param(pd.Series([1.0, 2.0, 3.0], index=TIME), [1.0, 2.0, 3.0], id="series"), + pytest.param( + pd.Series([3.0, 1.0, 2.0], index=pd.Index([2040, 2020, 2030], name="time")), + [1.0, 2.0, 3.0], + id="series-reordered", + ), + pytest.param( + DataArray([1.0, 2.0, 3.0], coords=[TIME]), [1.0, 2.0, 3.0], id="dataarray" + ), +] + + +class TestFixValueAlignment: + """fix() aligns the value to the variable's own coords (broadcast_to_coords).""" + + @pytest.fixture + def variable(self) -> Variable: + m = Model() + m.add_variables(lower=-5, upper=5, coords=[TIME], name="t") + return m.variables["t"] + + @pytest.mark.parametrize("value, expected", ALIGNED_VALUES) + def test_aligns_to_named_dimension( + self, variable: Variable, value: ConstantLike, expected: ConstantLike + ) -> None: + variable.fix(value) + assert variable.lower.dims == ("time",) + np.testing.assert_array_almost_equal(variable.lower.values, expected) + np.testing.assert_array_almost_equal(variable.upper.values, expected) + + def test_unknown_dimension_rejected(self, variable: Variable) -> None: + value = pd.Series([1.0, 2.0], index=pd.Index([0, 1], name="other")) + with pytest.raises(ValueError, match="fix.. for variable 't'"): + variable.fix(value) + + def test_partial_value_rejected(self, variable: Variable) -> None: + value = pd.Series([1.0, 3.0], index=pd.Index([2020, 2040], name="time")) + with pytest.raises(ValueError, match="fix.. for variable 't'"): + variable.fix(value) diff --git a/test/test_indicator_constraints.py b/test/test_indicator_constraints.py new file mode 100644 index 000000000..22962edd9 --- /dev/null +++ b/test/test_indicator_constraints.py @@ -0,0 +1,425 @@ +"""Tests for indicator constraint support.""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +import linopy +from linopy import Model, available_solvers +from linopy.constraints import Constraint, CSRConstraint +from linopy.variables import Variable + +requires_gurobi = pytest.mark.skipif( + "gurobi" not in available_solvers, reason="Gurobi not installed" +) + + +@pytest.fixture +def mbx() -> tuple[Model, Variable, Variable]: + """Model with a scalar binary ``b`` and continuous ``x`` in [0, 10].""" + m = Model() + b = m.add_variables(name="b", binary=True) + x = m.add_variables(lower=0, upper=10, name="x") + return m, b, x + + +@pytest.fixture +def coords_mbx() -> tuple[Model, Variable, Variable, pd.RangeIndex]: + """Model with index-aligned binary ``b`` and continuous ``x`` over a length-3 index.""" + m = Model() + idx = pd.RangeIndex(3, name="i") + b = m.add_variables(coords=[idx], name="b", binary=True) + x = m.add_variables(coords=[idx], lower=0, upper=10, name="x") + return m, b, x, idx + + +class TestConstruction: + """Creating indicator constraints and validating their arguments.""" + + def test_basic_fields(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Indicator constraint is created with correct fields.""" + m, b, x = mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + assert "ic0" in m.indicator_constraints + ic = m.indicator_constraints["ic0"] + for field in ( + "coeffs", + "vars", + "sign", + "rhs", + "binary_var", + "binary_val", + "labels", + ): + assert field in ic + + def test_auto_name(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Omitting ``name`` auto-generates an ``indcon`` name.""" + m, b, x = mbx + con = m.add_indicator_constraints(b, 1, x, "<=", 5) + assert con.name.startswith("indcon") + assert con.name in m.indicator_constraints + + @pytest.mark.parametrize( + ("build_lhs", "kwargs"), + [ + (lambda x: x <= 5, {}), + (lambda x: 2 * x, {"sign": "<=", "rhs": 5}), + (lambda x: 1 * x.at[()] <= 5, {}), + ], + ids=["constraint", "linexpr", "scalarcon"], + ) + def test_valid_lhs( + self, mbx: tuple[Model, Variable, Variable], build_lhs: object, kwargs: dict + ) -> None: + """Each accepted lhs form (constraint, expression, scalar) builds an indicator.""" + m, b, x = mbx + ic = m.add_indicator_constraints(b, 1, build_lhs(x), name="ic0", **kwargs) # type: ignore[operator] + assert ic.is_indicator + assert "ic0" in m.indicator_constraints + + @pytest.mark.parametrize( + ("build_lhs", "kwargs", "exc", "match"), + [ + (lambda x: x <= 5, {"sign": "<=", "rhs": 5}, ValueError, "must be None"), + ( + lambda x: 2 * x, + {}, + ValueError, + "are required when", + ), + (lambda x: x, {}, ValueError, "are required when"), + ( + lambda x: 1 * x.at[()] <= 5, + {"sign": "<=", "rhs": 5}, + ValueError, + "must be None", + ), + (lambda x: "not-a-constraint", {}, TypeError, "must be a LinearExpression"), + ], + ids=[ + "con+sign", + "linexpr-no-sign", + "var-no-sign", + "scalarcon+sign", + "bad-type", + ], + ) + def test_invalid_lhs( + self, + mbx: tuple[Model, Variable, Variable], + build_lhs: object, + kwargs: dict, + exc: type[Exception], + match: str, + ) -> None: + """Invalid lhs / sign / rhs combinations raise with a clear message.""" + m, b, x = mbx + with pytest.raises(exc, match=match): + m.add_indicator_constraints(b, 1, build_lhs(x), name="ic0", **kwargs) # type: ignore[operator] + + def test_non_binary_var_raises(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Non-binary variable raises ValueError.""" + m, _, x = mbx + y = m.add_variables(lower=0, upper=1, name="y") + with pytest.raises(ValueError, match="must be binary"): + m.add_indicator_constraints(y, 1, x, "<=", 5) + + def test_bad_binary_val_raises(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Invalid binary_val raises ValueError.""" + m, b, x = mbx + with pytest.raises(ValueError, match="must be 0 or 1"): + m.add_indicator_constraints(b, 2, x, "<=", 5) + + def test_duplicate_name_raises(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Duplicate name raises ValueError.""" + m, b, x = mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + with pytest.raises(ValueError, match="already assigned"): + m.add_indicator_constraints(b, 0, x, ">=", 0, name="ic0") + + +class TestContainer: + """How indicator constraints live in (and stay separate within) the container.""" + + def test_separate_from_regular(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Indicator constraints are separate from regular constraints.""" + m, b, x = mbx + m.add_constraints(x >= 0, name="regular") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + assert "regular" in m.constraints + assert "regular" in m.constraints.regular + assert "ic0" in m.constraints + assert "ic0" in m.constraints.indicator + assert "ic0" not in m.constraints.regular + assert "ic0" in m.indicator_constraints + + def test_in_unified_container(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Indicator constraint lives in the unified container, not in regular.""" + m, b, x = mbx + m.add_constraints(x >= 0, name="regular") + ic = m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + assert isinstance(ic, Constraint) + assert ic.is_indicator + assert "ic0" in m.constraints + assert "ic0" in m.constraints.indicator + assert "ic0" not in m.constraints.regular + assert "ic0" in m.indicator_constraints + assert "regular" in m.constraints.regular + assert "regular" not in m.constraints.indicator + + def test_remove(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Indicator constraints can be removed.""" + m, b, x = mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + assert "ic0" in m.indicator_constraints + m.remove_indicator_constraints("ic0") + assert "ic0" not in m.indicator_constraints + + def test_regular_constraint_has_no_indicator_fields( + self, mbx: tuple[Model, Variable, Variable] + ) -> None: + """Regular constraints report no indicator metadata, frozen or not.""" + m, _, x = mbx + m.add_constraints(x <= 5, name="c") + con = m.constraints["c"] + assert not con.is_indicator + assert con.binary_var is None + assert con.binary_val is None + + frozen = con.freeze() + assert not frozen.is_indicator + assert frozen.binary_var is None + assert frozen.binary_val is None + + def test_with_coords( + self, coords_mbx: tuple[Model, Variable, Variable, pd.RangeIndex] + ) -> None: + """Indicator constraints work with multi-dimensional coords.""" + m, b, x, idx = coords_mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + assert m.indicator_constraints["ic0"].labels.size == idx.size + + +class TestPersistence: + """Indicator metadata survives freeze, copy, and netCDF round-trips.""" + + def test_freeze_roundtrip(self, mbx: tuple[Model, Variable, Variable]) -> None: + """is_indicator and binary fields survive freeze and freeze->mutable.""" + m, b, x = mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + frozen = m.constraints["ic0"].freeze() + assert isinstance(frozen, CSRConstraint) + assert frozen.is_indicator + assert frozen.binary_var is not None + assert frozen.binary_val == 1 + + mutable = frozen.mutable() + assert isinstance(mutable, Constraint) + assert mutable.is_indicator + assert mutable.binary_var is not None + assert np.all(mutable.binary_val == 1) + + def test_copy_preserves(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Model.copy preserves is_indicator and the binary fields.""" + m, b, x = mbx + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + ic = m.copy().constraints["ic0"] + assert ic.is_indicator + assert ic.binary_var is not None + assert np.all(ic.binary_val == 1) + + @pytest.mark.parametrize("freeze_constraints", [False, True]) + def test_netcdf_roundtrip(self, tmp_path: Path, freeze_constraints: bool) -> None: + """is_indicator and binary fields survive a netCDF round-trip.""" + m = Model(freeze_constraints=freeze_constraints) + b = m.add_variables(name="b", binary=True) + x = m.add_variables(lower=0, upper=10, name="x") + m.add_constraints(x >= 0, name="regular") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + fn = tmp_path / "model.nc" + m.to_netcdf(fn) + ic = linopy.read_netcdf(fn).constraints["ic0"] + + assert ic.is_indicator + assert ic.binary_var is not None + assert np.all(ic.binary_val == 1) + + def test_array_binval_roundtrip(self, tmp_path: Path) -> None: + """A coords-based indicator has per-element binary_val that round-trips.""" + m = Model(freeze_constraints=True) + idx = pd.RangeIndex(3, name="i") + b = m.add_variables(coords=[idx], name="b", binary=True) + x = m.add_variables(coords=[idx], lower=0, upper=10, name="x") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + frozen = m.constraints["ic0"] + assert isinstance(frozen, CSRConstraint) + assert isinstance(frozen.binary_val, np.ndarray) and frozen.binary_val.ndim == 1 + assert "binary_val" in frozen.data + + fn = tmp_path / "model.nc" + m.to_netcdf(fn) + ic = linopy.read_netcdf(fn).constraints["ic0"] + assert ic.is_indicator + assert np.all(ic.binary_val == 1) + + +class TestMatrices: + """Indicator rows are split out of the regular constraint matrix.""" + + def test_matrix_split( + self, coords_mbx: tuple[Model, Variable, Variable, pd.RangeIndex] + ) -> None: + """Regular A excludes indicator rows; indicator arrays carry them.""" + m, b, x, idx = coords_mbx + m.add_constraints(x >= 0, name="regular") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + n = idx.size + assert m.matrices.A is not None + assert m.matrices.indicator_A is not None + assert m.matrices.A.shape[0] == n + assert m.matrices.indicator_A.shape[0] == n + assert len(m.matrices.clabels) == n + assert m.ncons == n + + np.testing.assert_array_equal(m.matrices.indicator_binval, np.full(n, 1)) + np.testing.assert_array_equal(m.matrices.indicator_b, np.full(n, 5.0)) + np.testing.assert_array_equal(m.matrices.indicator_sense, np.full(n, "<")) + assert m.matrices.indicator_binvar.shape == (n,) + + def test_to_matrix_skips_indicator( + self, coords_mbx: tuple[Model, Variable, Variable, pd.RangeIndex] + ) -> None: + """Constraints.to_matrix drops indicator rows; an indicator-only set raises.""" + m, b, x, idx = coords_mbx + m.add_constraints(x <= 5, name="regular") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + A, labels = m.constraints.to_matrix() + assert A.shape[0] == idx.size + + with pytest.raises(ValueError, match="No constraints available"): + m.constraints.indicator.to_matrix() + + +class TestLPFile: + """LP export of indicator constraints.""" + + def test_with_regular_constraints( + self, mbx: tuple[Model, Variable, Variable], tmp_path: Path + ) -> None: + """LP file contains general constraints section.""" + m, b, x = mbx + m.add_constraints(x >= 0, name="dummy") + m.add_objective(x) + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + fn = tmp_path / "test.lp" + m.to_file(fn) + content = fn.read_text() + assert "= 1 ->" in content + label = int(m.indicator_constraints["ic0"].labels.item()) + assert f"ic{label}:" in content + + def test_indicator_only( + self, mbx: tuple[Model, Variable, Variable], tmp_path: Path + ) -> None: + """An LP file with no regular constraints still writes the indicator section.""" + m, b, x = mbx + m.add_objective(x) + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + fn = tmp_path / "test.lp" + m.to_file(fn) + content = fn.read_text() + assert "s.t." in content + assert "= 1 ->" in content + + +class TestSolve: + """Solving models with indicator constraints.""" + + @requires_gurobi + @pytest.mark.parametrize( + ("io_api", "trigger", "expected"), + [ + (None, 1, 5), + (None, 0, 10), + ("direct", 1, 5), + ("direct", 0, 10), + ], + ids=["lp-active", "lp-inactive", "direct-active", "direct-inactive"], + ) + def test_gurobi_enforces_at_trigger( + self, + mbx: tuple[Model, Variable, Variable], + io_api: str | None, + trigger: int, + expected: float, + ) -> None: + """ + The indicator is enforced (x<=5) only when b matches the trigger value. + + With b fixed to 1, an active trigger caps x at 5; an inactive one leaves + x free to its upper bound of 10. + """ + m, b, x = mbx + m.add_constraints(b >= 1, name="fix_b") + m.add_indicator_constraints(b, trigger, x, "<=", 5, name="ic0") + m.add_objective(x, sense="max") + m.solve(solver_name="gurobi", io_api=io_api) + assert m.objective.value is not None + assert np.isclose(m.objective.value, expected, atol=1e-6) + + @requires_gurobi + def test_gurobi_multiple(self) -> None: + """Multiple indicators combine: x <= min(10, 5) = 5.""" + m = Model() + b1 = m.add_variables(name="b1", binary=True) + b2 = m.add_variables(name="b2", binary=True) + x = m.add_variables(lower=0, upper=20, name="x") + m.add_constraints(b1 >= 1, name="fix_b1") + m.add_constraints(b2 >= 1, name="fix_b2") + m.add_indicator_constraints(b1, 1, x, "<=", 10, name="ic1") + m.add_indicator_constraints(b2, 1, x, "<=", 5, name="ic2") + m.add_objective(x, sense="max") + m.solve(solver_name="gurobi") + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5, atol=1e-6) + + def test_unsupported_solver_raises( + self, mbx: tuple[Model, Variable, Variable] + ) -> None: + """Solvers without indicator support raise ValueError.""" + m, b, x = mbx + m.add_constraints(x >= 0, name="dummy") + m.add_objective(x) + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + + for solver in ["glpk", "highs", "mosek", "mindopt"]: + if solver in available_solvers: + with pytest.raises( + ValueError, match="does not support indicator constraints" + ): + m.solve(solver_name=solver) + + @requires_gurobi + def test_mip_has_no_duals(self, mbx: tuple[Model, Variable, Variable]) -> None: + """Indicator rows are skipped when collecting duals; MIPs expose none.""" + m, b, x = mbx + m.add_constraints(b >= 1, name="fix_b") + m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + m.add_objective(x, sense="max") + m.solve(solver_name="gurobi") + with pytest.raises(AttributeError, match="dual"): + _ = m.matrices.dual diff --git a/test/test_infeasibility.py b/test/test_infeasibility.py index 019947898..df6b62739 100644 --- a/test/test_infeasibility.py +++ b/test/test_infeasibility.py @@ -3,6 +3,8 @@ Test infeasibility detection for different solvers. """ +from typing import cast + import pandas as pd import pytest @@ -92,8 +94,9 @@ def test_simple_infeasibility_detection( assert isinstance(labels, list) assert len(labels) > 0 # Should find at least one infeasible constraint - # Test print_infeasibilities (just check it doesn't raise an error) - m.print_infeasibilities() + formatted = m.format_infeasibilities() + assert isinstance(formatted, str) + assert formatted @pytest.mark.parametrize("solver", ["gurobi", "xpress"]) def test_complex_infeasibility_detection( @@ -163,9 +166,8 @@ def test_no_solver_model_error(self, solver: str) -> None: # Solve the model first m.solve(solver_name=solver) - # Manually remove the solver_model to simulate cleanup - m.solver_model = None - m.solver_name = solver # But keep the solver name + assert m.solver is not None + m.solver.solver_model = None # Should raise ValueError since we know it was solved with supported solver with pytest.raises(ValueError, match="No solver model available"): @@ -208,6 +210,7 @@ def test_unsupported_solver_error(self) -> None: x = m.add_variables(name="x") m.add_constraints(x >= 0) m.add_constraints(x <= -1) # Make it infeasible + m.add_objective(1 * x) # Use a solver that doesn't support IIS if "cbc" in available_solvers: @@ -242,3 +245,58 @@ def test_deprecated_method( # Check that it contains constraint labels assert len(subset) > 0 + + @pytest.mark.parametrize("solver", ["gurobi", "xpress"]) + def test_masked_constraint_infeasibility( + self, solver: str, capsys: pytest.CaptureFixture[str] + ) -> None: + """ + Test infeasibility detection with masked constraints. + + This test verifies that the solver correctly maps constraint positions + to constraint labels when constraints are masked (some rows skipped). + The enumeration creates positions [0, 1, 2, ...] that should correspond + to the actual constraint labels which may have gaps like [0, 2, 4, 6]. + """ + if solver not in available_solvers: + pytest.skip(f"{solver} not available") + + m = Model() + + time = pd.RangeIndex(8, name="time") + x = m.add_variables(lower=0, upper=5, coords=[time], name="x") + y = m.add_variables(lower=0, upper=5, coords=[time], name="y") + + # Create a mask that keeps only even time indices (0, 2, 4, 6) + mask = pd.Series([i % 2 == 0 for i in range(len(time))]) + m.add_constraints(x + y >= 10, name="sum_lower", mask=mask) + + mask = pd.Series([False] * (len(time) // 2) + [True] * (len(time) // 2)) + m.add_constraints(x <= 4, name="x_upper", mask=mask) + + m.add_objective(x.sum() + y.sum()) + status, condition = m.solve(solver_name=solver) + + assert status == "warning" + assert "infeasible" in condition + + labels = m.compute_infeasibilities() + assert labels + + positions = [ + cast(tuple[str, dict[str, int]], m.constraints.get_label_position(label)) + for label in labels + ] + grouped_coords: dict[str, set[int]] = {"sum_lower": set(), "x_upper": set()} + for name, coord in positions: + assert name in grouped_coords + grouped_coords[name].add(coord["time"]) + + assert grouped_coords["sum_lower"] + assert grouped_coords["sum_lower"] == grouped_coords["x_upper"] + + print(m.format_infeasibilities()) + output = capsys.readouterr().out + for time_coord in grouped_coords["sum_lower"]: + assert f"sum_lower[{time_coord}]" in output + assert f"x_upper[{time_coord}]" in output diff --git a/test/test_io.py b/test/test_io.py index 4336f29d3..27cba396b 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -5,7 +5,10 @@ @author: fabian """ +import importlib.util +import json import pickle +from collections.abc import Callable from pathlib import Path import numpy as np @@ -18,6 +21,8 @@ from linopy.io import signed_number from linopy.testing import assert_model_equal +HAS_NETCDF4 = importlib.util.find_spec("netCDF4") is not None + @pytest.fixture def model() -> Model: @@ -78,6 +83,48 @@ def test_model_to_netcdf(model: Model, tmp_path: Path) -> None: assert_model_equal(m, p) +def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: + from linopy.constraints import CSRConstraint + + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") + m.add_constraints(x >= 1, name="c", freeze=True) + + assert isinstance(m.constraints["c"], CSRConstraint) + + fn = tmp_path / "test_frozen.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.constraints["c"], CSRConstraint) + assert_model_equal(m, p) + + +def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None: + from linopy.constraints import CSRConstraint + + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + + def bound(m: Model, i: int) -> object: + if i % 2: + return x.at[i] >= i + return x.at[i] == 0.0 + + m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c", freeze=True) + assert isinstance(m.constraints["c"], CSRConstraint) + + fn = tmp_path / "test_mixed_sign.nc" + m.to_netcdf(fn) + p = read_netcdf(fn) + + assert isinstance(p.constraints["c"], CSRConstraint) + import numpy as np + + np.testing.assert_array_equal(m.constraints["c"]._sign, p.constraints["c"]._sign) + assert_model_equal(m, p) + + def test_model_to_netcdf_with_sense(model: Model, tmp_path: Path) -> None: m = model m.objective.sense = "max" @@ -127,11 +174,6 @@ def test_pickle_model(model_with_dash_names: Model, tmp_path: Path) -> None: assert_model_equal(m, p) -# skip it xarray version is 2024.01.0 due to issue https://github.com/pydata/xarray/issues/8628 -@pytest.mark.skipif( - xr.__version__ in ["2024.1.0", "2024.1.1"], - reason="xarray version 2024.1.0 has a bug with MultiIndex deserialize", -) def test_model_to_netcdf_with_multiindex( model_with_multiindex: Model, tmp_path: Path ) -> None: @@ -143,6 +185,57 @@ def test_model_to_netcdf_with_multiindex( assert_model_equal(m, p) +# Regression for https://github.com/PyPSA/linopy/issues/525. +def test_model_to_netcdf_with_multiindex_scipy_engine( + model_with_multiindex: Model, tmp_path: Path +) -> None: + m = model_with_multiindex + fn = tmp_path / "test.nc" + m.to_netcdf(fn, engine="scipy") + + raw_attrs = xr.load_dataset(fn).attrs + multiindex_attrs = {k: v for k, v in raw_attrs.items() if k.endswith("_multiindex")} + assert multiindex_attrs + for k, v in multiindex_attrs.items(): + assert isinstance(v, str), f"{k!r}: {v!r}" + + assert_model_equal(m, read_netcdf(fn)) + + +@pytest.mark.skipif(not HAS_NETCDF4, reason="legacy format requires netCDF4 backend") +def test_read_netcdf_with_multiindex_legacy_list_attr( + model_with_multiindex: Model, tmp_path: Path +) -> None: + # Older linopy stored multiindex names as a Python list (netCDF4-only). + m = model_with_multiindex + fn = tmp_path / "test.nc" + m.to_netcdf(fn, engine="netcdf4") + + ds = xr.load_dataset(fn, engine="netcdf4").load() + ds.attrs = { + k: (json.loads(v) if k.endswith("_multiindex") and isinstance(v, str) else v) + for k, v in ds.attrs.items() + } + fn_legacy = tmp_path / "legacy.nc" + ds.to_netcdf(fn_legacy, engine="netcdf4") + + assert_model_equal(m, read_netcdf(fn_legacy)) + + +def test_read_netcdf_without_version_stamp(model: Model, tmp_path: Path) -> None: + from linopy.io import NETCDF_VERSION_ATTR + + fn = tmp_path / "test.nc" + model.to_netcdf(fn) + + ds = xr.load_dataset(fn).load() + del ds.attrs[NETCDF_VERSION_ATTR] + fn_legacy = tmp_path / "legacy.nc" + ds.to_netcdf(fn_legacy) + + assert_model_equal(model, read_netcdf(fn_legacy)) + + @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed") def test_to_file_lp(model: Model, tmp_path: Path) -> None: import gurobipy @@ -196,12 +289,74 @@ def test_to_file_invalid(model: Model, tmp_path: Path) -> None: @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed") def test_to_gurobipy(model: Model) -> None: - model.to_gurobipy() + gm = model.to_gurobipy() + assert gm.NumVars > 0 + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed") +def test_to_gurobipy_no_names(model: Model) -> None: + m_with = model.to_gurobipy(set_names=True) + m_without = model.to_gurobipy(set_names=False) + names_with = [v.VarName for v in m_with.getVars()] + names_without = [v.VarName for v in m_without.getVars()] + assert names_with != names_without @pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") def test_to_highspy(model: Model) -> None: - model.to_highspy() + h = model.to_highspy() + assert h.getLp().num_col_ > 0 + + +@pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") +def test_to_highspy_no_names(model: Model) -> None: + h = model.to_highspy(set_names=False) + lp = h.getLp() + assert len(lp.col_names_) == 0 + assert len(lp.row_names_) == 0 + + +@pytest.mark.skipif("mosek" not in available_solvers, reason="Mosek not installed") +def test_to_mosek(model: Model) -> None: + task = model.to_mosek() + assert task.getnumvar() > 0 + + +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_to_xpress(model: Model) -> None: + p = model.to_xpress() + assert p.attributes.cols > 0 + assert p.attributes.rows > 0 + + +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_to_xpress_no_names(model: Model) -> None: + p_with = model.to_xpress(set_names=True) + p_without = model.to_xpress(set_names=False) + names_with = [v.name for v in p_with.getVariable()] + names_without = [v.name for v in p_without.getVariable()] + assert names_with != names_without + + +@pytest.mark.skipif("cupdlpx" not in available_solvers, reason="cuPDLPx not installed") +def test_to_cupdlpx(model: Model) -> None: + cu = model.to_cupdlpx() + assert cu is not None + + +def test_model_set_names_in_solver_io_default() -> None: + assert Model().set_names_in_solver_io is True + + +@pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed") +def test_model_set_names_in_solver_io(model: Model) -> None: + model.solve(solver_name="highs", io_api="direct") + expected_obj = model.objective.value + + model.set_names_in_solver_io = False + status, _ = model.solve(solver_name="highs", io_api="direct") + assert status == "ok" + assert model.objective.value == pytest.approx(expected_obj) def test_to_blocks(tmp_path: Path) -> None: @@ -336,3 +491,156 @@ def test_to_file_lp_with_negative_zero_coefficients(tmp_path: Path) -> None: # Verify Gurobi can read it without errors gurobipy.read(str(fn)) + + +def test_to_file_lp_same_sign_constraints(tmp_path: Path) -> None: + """Test LP writing when all constraints have the same sign operator.""" + m = Model() + N = np.arange(5) + x = m.add_variables(coords=[N], name="x") + # All constraints use <= + m.add_constraints(x <= 10, name="upper") + m.add_constraints(x <= 20, name="upper2") + m.add_objective(x.sum()) + + fn = tmp_path / "same_sign.lp" + m.to_file(fn) + content = fn.read_text() + assert "s.t." in content + assert "<=" in content + + +def test_to_file_lp_mixed_sign_constraints(tmp_path: Path) -> None: + """Test LP writing when constraints have different sign operators.""" + m = Model() + N = np.arange(5) + x = m.add_variables(coords=[N], name="x") + # Mix of <= and >= constraints in the same container + m.add_constraints(x <= 10, name="upper") + m.add_constraints(x >= 1, name="lower") + m.add_constraints(2 * x == 8, name="eq") + m.add_objective(x.sum()) + + fn = tmp_path / "mixed_sign.lp" + m.to_file(fn) + content = fn.read_text() + assert "s.t." in content + assert "<=" in content + assert ">=" in content + assert "=" in content + + +class TestLPBinaryBounds: + """LP export honors binary bounds tightened below [0, 1] (#776).""" + + @pytest.fixture + def make_tightened_model(self) -> Callable[[], Model]: + def build() -> Model: + m = Model() + x = m.add_variables( + binary=True, coords=[pd.RangeIndex(4, name="t")], name="x" + ) + x.upper = pd.Series([1, 1, 0, 0], index=pd.RangeIndex(4, name="t")) + m.add_constraints(x.sum() >= 2, name="atleast2") + m.add_objective(-1 * x.sum()) + return m + + return build + + def test_default_bounds_omitted(self, tmp_path: Path) -> None: + """A binary with the implied [0, 1] bounds gets no bounds section.""" + m = Model() + b = m.add_variables(binary=True, coords=[pd.RangeIndex(3, name="t")], name="b") + m.add_constraints(b.sum() >= 1, name="c") + m.add_objective(b.sum()) + + fn = tmp_path / "binary_default.lp" + m.to_file(fn) + assert "bounds" not in fn.read_text() + + def test_tightened_bounds_written( + self, make_tightened_model: Callable[[], Model], tmp_path: Path + ) -> None: + """Per-element bounds tighter than [0, 1] reach the LP `bounds` section.""" + m = make_tightened_model() + fn = tmp_path / "binary_tightened.lp" + m.to_file(fn) + + bounds_section = fn.read_text().split("bounds")[1].split("binary")[0] + for label in m.variables["x"].labels.values[2:]: + assert f"x{label} <= +0.0" in bounds_section + + @pytest.mark.skipif(not available_solvers, reason="No solver installed") + def test_lp_and_direct_agree( + self, make_tightened_model: Callable[[], Model] + ) -> None: + """LP and direct paths see the same feasible set for tightened binaries.""" + solver = available_solvers[0] + + m_direct = make_tightened_model() + m_direct.solve(solver_name=solver, io_api="direct") + + m_lp = make_tightened_model() + m_lp.solve(solver_name=solver, io_api="lp") + + assert m_direct.objective.value == m_lp.objective.value == -2 + + +def test_to_file_lp_frozen_vs_mutable(tmp_path: Path) -> None: + """Test that frozen and mutable constraints produce identical LP output.""" + m_frozen = Model() + N = np.arange(5) + x = m_frozen.add_variables(coords=[N], name="x") + y = m_frozen.add_variables(coords=[N], name="y") + m_frozen.add_constraints(x + y <= 10, name="upper") + m_frozen.add_constraints(x >= 1, name="lower") + m_frozen.add_constraints(2 * x + y == 8, name="eq") + m_frozen.add_objective(x.sum() + 2 * y.sum()) + + m_mutable = Model() + x2 = m_mutable.add_variables(coords=[N], name="x") + y2 = m_mutable.add_variables(coords=[N], name="y") + m_mutable.add_constraints(x2 + y2 <= 10, name="upper", freeze=False) + m_mutable.add_constraints(x2 >= 1, name="lower", freeze=False) + m_mutable.add_constraints(2 * x2 + y2 == 8, name="eq", freeze=False) + m_mutable.add_objective(x2.sum() + 2 * y2.sum()) + + fn_frozen = tmp_path / "frozen.lp" + fn_mutable = tmp_path / "mutable.lp" + m_frozen.to_file(fn_frozen) + m_mutable.to_file(fn_mutable) + + assert fn_frozen.read_text() == fn_mutable.read_text() + + +def test_to_file_lp_frozen_mixed_sign(tmp_path: Path) -> None: + """Test LP writing for frozen constraint with per-row signs.""" + m_frozen = Model() + N = pd.RangeIndex(4, name="i") + x = m_frozen.add_variables(coords=[N], name="x") + + def bound(m: Model, i: int) -> object: + if i % 2: + return x.at[i] >= i + return x.at[i] <= 10 + + m_frozen.add_constraints(bound, coords=[N], name="mixed", freeze=True) + m_frozen.add_objective(x.sum()) + + m_mutable = Model() + x2 = m_mutable.add_variables(coords=[N], name="x") + + def bound2(m: Model, i: int) -> object: + if i % 2: + return x2.at[i] >= i + return x2.at[i] <= 10 + + m_mutable.add_constraints(bound2, coords=[N], name="mixed", freeze=False) + m_mutable.add_objective(x2.sum()) + + fn_frozen = tmp_path / "frozen_mixed.lp" + fn_mutable = tmp_path / "mutable_mixed.lp" + m_frozen.to_file(fn_frozen) + m_mutable.to_file(fn_mutable) + + assert fn_frozen.read_text() == fn_mutable.read_text() diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index a75ace3f7..5ffd7de1e 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -7,60 +7,31 @@ from __future__ import annotations +import warnings +from typing import Any + import numpy as np import pandas as pd import polars as pl import pytest import xarray as xr +from xarray.core.types import JoinOptions from xarray.testing import assert_equal -from linopy import LinearExpression, Model, QuadraticExpression, Variable, merge +from linopy import ( + EvolvingAPIWarning, + LinearExpression, + Model, + QuadraticExpression, + Variable, + merge, +) from linopy.constants import HELPER_DIMS, TERM_DIM from linopy.expressions import ScalarLinearExpression from linopy.testing import assert_linequal, assert_quadequal from linopy.variables import ScalarVariable -@pytest.fixture -def m() -> Model: - m = Model() - - m.add_variables(pd.Series([0, 0]), 1, name="x") - m.add_variables(4, pd.Series([8, 10]), name="y") - m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]).T, name="z") - m.add_variables(coords=[pd.RangeIndex(20, name="dim_2")], name="v") - - idx = pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=("level1", "level2")) - idx.name = "dim_3" - m.add_variables(coords=[idx], name="u") - return m - - -@pytest.fixture -def x(m: Model) -> Variable: - return m.variables["x"] - - -@pytest.fixture -def y(m: Model) -> Variable: - return m.variables["y"] - - -@pytest.fixture -def z(m: Model) -> Variable: - return m.variables["z"] - - -@pytest.fixture -def v(m: Model) -> Variable: - return m.variables["v"] - - -@pytest.fixture -def u(m: Model) -> Variable: - return m.variables["u"] - - def test_empty_linexpr(m: Model) -> None: LinearExpression(None, m) @@ -92,7 +63,7 @@ def test_linexpr_with_helper_dims_as_coords(m: Model) -> None: assert set(HELPER_DIMS).intersection(set(data.coords)) expr = LinearExpression(data, m) - assert not set(HELPER_DIMS).intersection(set(expr.data.coords)) + assert not set(HELPER_DIMS).intersection(set(expr.coords)) def test_linexpr_with_data_without_coords(m: Model) -> None: @@ -325,6 +296,27 @@ def test_linear_expression_multi_indexed(u: Variable) -> None: assert isinstance(expr, LinearExpression) +def test_multiply_expression_by_multiindex_level_constant(u: Variable) -> None: + """ + Expression over a MultiIndex dim times a single-level constant. + + Mirrors PyPSA's ``soc_delta * storage_weightings``: ``u`` is indexed by + the (level1, level2) MultiIndex ``dim_3``; the weighting is indexed only + by ``level1``. The product must not raise, and each ``dim_3`` entry must + take the weight of its ``level1``. + """ + by_level1 = xr.DataArray([10.0, 20.0], coords={"level1": [1, 2]}, dims=["level1"]) + + with pytest.warns(EvolvingAPIWarning, match=r"broadcasting level subset"): + expr = (1 * u) * by_level1 + + coeffs = expr.coeffs.squeeze("_term") + assert coeffs.sel(dim_3=(1, "a")).item() == 10.0 + assert coeffs.sel(dim_3=(1, "b")).item() == 10.0 + assert coeffs.sel(dim_3=(2, "a")).item() == 20.0 + assert coeffs.sel(dim_3=(2, "b")).item() == 20.0 + + def test_linear_expression_with_errors(m: Model, x: Variable) -> None: with pytest.raises(TypeError): x / x @@ -555,6 +547,47 @@ def test_matmul_expr_and_const(x: Variable, y: Variable) -> None: assert_linequal(expr.dot(const), target) +def test_matmul_contracts_only_shared_dims(z: Variable) -> None: + """ + A @ b contracts the genuinely shared dims and keeps the rest. + + ``z`` has dims (dim_0, dim_1); ``b`` has (dim_1, location). Only dim_1 + is shared, so the result must keep dim_0 and location. A conversion that + broadcast ``b`` to ``z``'s coords would expand dim_0 into ``b`` and + contract it away too — collapsing the result to (location,) only. + """ + expr = 1 * z + b = xr.DataArray( + np.ones((3, 2)), + coords={"dim_1": expr.indexes["dim_1"], "location": ["L1", "L2"]}, + dims=["dim_1", "location"], + ) + + res = expr @ b + + assert set(res.coord_dims) == {"dim_0", "location"} + assert_linequal(res, (expr * b).sum("dim_1")) + + +def test_matmul_contracts_all_dims_when_const_covers_them(z: Variable) -> None: + """B covering all of a's dims (and more) contracts a's dims, keeping b's extras.""" + expr = 1 * z # dims (dim_0, dim_1) + b = xr.DataArray( + np.ones((2, 3, 2)), + coords={ + "dim_0": expr.indexes["dim_0"], + "dim_1": expr.indexes["dim_1"], + "location": ["L1", "L2"], + }, + dims=["dim_0", "dim_1", "location"], + ) + + res = expr @ b + + assert set(res.coord_dims) == {"location"} + assert_linequal(res, (expr * b).sum(["dim_0", "dim_1"])) + + def test_matmul_wrong_input(x: Variable, y: Variable, z: Variable) -> None: expr = 10 * x + y + z with pytest.raises(TypeError): @@ -575,6 +608,498 @@ def test_linear_expression_multiplication_invalid( expr / x +class TestCoordinateAlignment: + @pytest.fixture(params=["da", "series"]) + def subset(self, request: Any) -> xr.DataArray | pd.Series: + if request.param == "da": + return xr.DataArray([10.0, 30.0], dims=["dim_2"], coords={"dim_2": [1, 3]}) + return pd.Series([10.0, 30.0], index=pd.Index([1, 3], name="dim_2")) + + @pytest.fixture(params=["da", "series"]) + def superset(self, request: Any) -> xr.DataArray | pd.Series: + if request.param == "da": + return xr.DataArray( + np.arange(25, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(25)}, + ) + return pd.Series( + np.arange(25, dtype=float), index=pd.Index(range(25), name="dim_2") + ) + + @pytest.fixture + def expected_fill(self) -> np.ndarray: + arr = np.zeros(20) + arr[1] = 10.0 + arr[3] = 30.0 + return arr + + @pytest.fixture(params=["xarray", "pandas_series"], ids=["da", "series"]) + def nan_constant(self, request: Any) -> xr.DataArray | pd.Series: + vals = np.arange(20, dtype=float) + vals[0] = np.nan + vals[5] = np.nan + vals[19] = np.nan + if request.param == "xarray": + return xr.DataArray(vals, dims=["dim_2"], coords={"dim_2": range(20)}) + return pd.Series(vals, index=pd.Index(range(20), name="dim_2")) + + class TestSubset: + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_mul_subset_fills_zeros( + self, + v: Variable, + subset: xr.DataArray, + expected_fill: np.ndarray, + operand: str, + ) -> None: + target = v if operand == "var" else 1 * v + result = target * 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("operand", ["var", "expr"]) + def test_add_subset_fills_zeros( + self, + v: Variable, + subset: xr.DataArray, + expected_fill: np.ndarray, + operand: str, + ) -> None: + if operand == "var": + result = v + subset + expected = expected_fill + else: + result = (v + 5) + subset + expected = expected_fill + 5 + 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) + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_sub_subset_fills_negated( + self, + v: Variable, + subset: xr.DataArray, + expected_fill: np.ndarray, + operand: str, + ) -> None: + if operand == "var": + result = v - subset + expected = -expected_fill + else: + result = (v + 5) - subset + expected = 5 - expected_fill + 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) + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_div_subset_inverts_nonzero( + self, v: Variable, subset: xr.DataArray, operand: str + ) -> None: + target = v if operand == "var" else 1 * v + result = target / 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)) + + class TestSuperset: + def test_add_superset_pins_to_lhs_coords( + 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_add_var_commutative(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset + v, v + superset) + + def test_sub_var_commutative(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset - v, -v + superset) + + def test_mul_var_commutative(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset * v, v * superset) + + def test_mul_superset_pins_to_lhs_coords( + 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_div_superset_pins_to_lhs_coords(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() + + class TestDisjoint: + def test_add_disjoint_fills_zeros(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_mul_disjoint_fills_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_div_disjoint_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)) + + class TestCommutativity: + @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)) + + def test_sub_var_anticommutative( + self, v: Variable, subset: xr.DataArray + ) -> None: + assert_linequal(subset - v, -v + subset) + + def test_sub_expr_anticommutative( + self, v: Variable, subset: xr.DataArray + ) -> None: + expr = v + 5 + assert_linequal(subset - expr, -(expr - subset)) + + 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) + + class TestQuadratic: + 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) + + class TestMissingValues: + """ + Same shape as variable but with NaN entries in the constant. + + NaN values are filled with operation-specific neutral elements: + - Addition/subtraction: NaN -> 0 (additive identity) + - Multiplication: NaN -> 0 (zeroes out the variable) + - Division: NaN -> 1 (multiplicative identity, no scaling) + """ + + NAN_POSITIONS = [0, 5, 19] + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_add_nan_filled( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + operand: str, + ) -> None: + base_const = 0.0 if operand == "var" else 5.0 + target = v if operand == "var" else v + 5 + result = target + nan_constant + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.const.values).any() + # At NaN positions, const should be unchanged (added 0) + for i in self.NAN_POSITIONS: + assert result.const.values[i] == base_const + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_sub_nan_filled( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + operand: str, + ) -> None: + base_const = 0.0 if operand == "var" else 5.0 + target = v if operand == "var" else v + 5 + result = target - nan_constant + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.const.values).any() + # At NaN positions, const should be unchanged (subtracted 0) + for i in self.NAN_POSITIONS: + assert result.const.values[i] == base_const + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_mul_nan_filled( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + operand: str, + ) -> None: + target = v if operand == "var" else 1 * v + result = target * nan_constant + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.coeffs.squeeze().values).any() + # At NaN positions, coeffs should be 0 (variable zeroed out) + for i in self.NAN_POSITIONS: + assert result.coeffs.squeeze().values[i] == 0.0 + + @pytest.mark.parametrize("operand", ["var", "expr"]) + def test_div_nan_filled( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + operand: str, + ) -> None: + target = v if operand == "var" else 1 * v + result = target / nan_constant + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.coeffs.squeeze().values).any() + # At NaN positions, coeffs should be unchanged (divided by 1) + original_coeffs = (1 * v).coeffs.squeeze().values + for i in self.NAN_POSITIONS: + assert result.coeffs.squeeze().values[i] == original_coeffs[i] + + def test_add_commutativity( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + ) -> None: + result_a = v + nan_constant + result_b = nan_constant + v + assert not np.isnan(result_a.const.values).any() + assert not np.isnan(result_b.const.values).any() + np.testing.assert_array_equal(result_a.const.values, result_b.const.values) + np.testing.assert_array_equal( + result_a.coeffs.values, result_b.coeffs.values + ) + + def test_mul_commutativity( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + ) -> None: + result_a = v * nan_constant + result_b = nan_constant * v + assert not np.isnan(result_a.coeffs.values).any() + assert not np.isnan(result_b.coeffs.values).any() + np.testing.assert_array_equal( + result_a.coeffs.values, result_b.coeffs.values + ) + + def test_quadexpr_add_nan( + self, + v: Variable, + nan_constant: xr.DataArray | pd.Series, + ) -> None: + qexpr = v * v + result = qexpr + nan_constant + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == 20 + assert not np.isnan(result.const.values).any() + + class TestExpressionWithNaN: + """Test that NaN in expression's own const/coeffs doesn't propagate.""" + + def test_shifted_expr_add_scalar(self, v: Variable) -> None: + expr = (1 * v).shift(dim_2=1) + result = expr + 5 + assert not np.isnan(result.const.values).any() + assert result.const.values[0] == 5.0 + + def test_shifted_expr_mul_scalar(self, v: Variable) -> None: + expr = (1 * v).shift(dim_2=1) + result = expr * 2 + assert not np.isnan(result.coeffs.squeeze().values).any() + assert result.coeffs.squeeze().values[0] == 0.0 + + def test_shifted_expr_add_array(self, v: Variable) -> None: + arr = np.arange(v.sizes["dim_2"], dtype=float) + expr = (1 * v).shift(dim_2=1) + result = expr + arr + assert not np.isnan(result.const.values).any() + assert result.const.values[0] == 0.0 + + def test_shifted_expr_mul_array(self, v: Variable) -> None: + arr = np.arange(v.sizes["dim_2"], dtype=float) + 1 + expr = (1 * v).shift(dim_2=1) + result = expr * arr + assert not np.isnan(result.coeffs.squeeze().values).any() + assert result.coeffs.squeeze().values[0] == 0.0 + + def test_shifted_expr_div_scalar(self, v: Variable) -> None: + expr = (1 * v).shift(dim_2=1) + result = expr / 2 + assert not np.isnan(result.coeffs.squeeze().values).any() + assert result.coeffs.squeeze().values[0] == 0.0 + + def test_shifted_expr_sub_scalar(self, v: Variable) -> None: + expr = (1 * v).shift(dim_2=1) + result = expr - 3 + assert not np.isnan(result.const.values).any() + assert result.const.values[0] == -3.0 + + def test_shifted_expr_div_array(self, v: Variable) -> None: + arr = np.arange(v.sizes["dim_2"], dtype=float) + 1 + expr = (1 * v).shift(dim_2=1) + result = expr / arr + assert not np.isnan(result.coeffs.squeeze().values).any() + assert result.coeffs.squeeze().values[0] == 0.0 + + def test_variable_to_linexpr_nan_coefficient(self, v: Variable) -> None: + nan_coeff = np.ones(v.sizes["dim_2"]) + nan_coeff[0] = np.nan + result = v.to_linexpr(nan_coeff) + assert not np.isnan(result.coeffs.squeeze().values).any() + assert result.coeffs.squeeze().values[0] == 0.0 + + class TestMultiDim: + 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) + + class TestXarrayCompat: + 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_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_expression_inherited_properties(x: Variable, y: Variable) -> None: expr = 10 * x + y assert isinstance(expr.attrs, dict) @@ -632,6 +1157,56 @@ def test_linear_expression_isnull(v: Variable) -> None: assert expr.isnull().sum() == 10 +class TestHasTerms: + """has_terms: true at slots with at least one live term, regardless of the constant.""" + + def test_basic_and_masking(self, v: Variable) -> None: + expr = np.arange(20) * v + assert expr.has_terms.all() + + filter = (expr.coeffs >= 10).any(TERM_DIM) + masked = expr.where(filter) + assert_equal(masked.has_terms, filter.rename("has_terms")) + + def test_ignores_const(self, v: Variable) -> None: + # has_terms differs from isnull() at slots whose constant was revived by + # fillna: no longer null, but still without terms + expr = np.arange(20) * v + filter = (expr.coeffs >= 10).any(TERM_DIM) + masked = expr.where(filter) + assert_equal(masked.isnull(), ~masked.has_terms) + + filled = masked.fillna(0) + assert not filled.isnull().any() + assert_equal(filled.has_terms, filter.rename("has_terms")) + + def test_merge_reindex(self, x: Variable, y: Variable) -> None: + # the nodal-balance pattern: outer merge, then reindex to a superset of + # coordinates; slots beyond the original coordinates carry no terms + lhs = merge([1 * x, 1 * y], join="outer").reindex( + dim_0=pd.RangeIndex(4, name="dim_0") + ) + assert lhs.has_terms.values.tolist() == [True, True, False, False] + + def test_constant_only(self, m: Model) -> None: + expr = LinearExpression(xr.DataArray([1, 2], dims=["dim_0"]), m) + assert expr.nterm == 0 + assert not expr.has_terms.any() + + def test_quadratic(self, v: Variable) -> None: + # linear terms inside a quadratic expression carry one factor == -1; + # they must still count as live terms + quad = v * v + 2 * v + assert quad.has_terms.all() + assert TERM_DIM not in quad.has_terms.dims + + filter = xr.DataArray( + np.arange(20) >= 10, dims="dim_2", coords={"dim_2": range(20)} + ) + masked = quad.where(filter) + assert_equal(masked.has_terms, filter.rename("has_terms")) + + def test_linear_expression_flat(v: Variable) -> None: coeff = np.arange(1, 21) # use non-zero coefficients expr = coeff * v @@ -810,6 +1385,288 @@ def test_linear_expression_groupby_on_same_name_as_target_dim( assert grouped.nterm == 10 +class TestMultiKeyFastPath: + """ + Group a LinearExpression by a list of coordinate names: takes the fast + reindex path and returns one dimension per key, like the xarray fallback. + """ + + @staticmethod + def _expr(period_vals: list, season_vals: list) -> LinearExpression: + n = len(period_vals) + s = pd.RangeIndex(n, name="s") + m = Model() + x = m.add_variables(coords=[s], name="x") + return (1.0 * x).assign_coords( + period=xr.DataArray(period_vals, dims="s", coords={"s": s}, name="period"), + season=xr.DataArray(season_vals, dims="s", coords={"s": s}, name="season"), + ) + + @pytest.mark.parametrize("spelling", [list, tuple], ids=["list", "tuple"]) + def test_matches_fallback(self, spelling: type) -> None: + # the fast path must equal the slow fallback, sparse cells included + expr = self._expr([2020, 2020, 2030, 2030, 2030], list("wswws")) + group = spelling(["period", "season"]) + + fast = expr.groupby(group).sum() + slow = expr.groupby(group).sum(use_fallback=True) + + assert_linequal(fast, slow) + + def test_separate_dims_not_stacked(self) -> None: + # built via a stacked index internally, but returns one dim per key + expr = self._expr([2020, 2020, 2030, 2030], list("wsws")) + + grouped = expr.groupby(["period", "season"]).sum() + + assert {"period", "season"} <= set(grouped.dims) + assert "group" not in grouped.dims + assert not isinstance(grouped.data.indexes.get("period"), pd.MultiIndex) + + def test_sparse_combination_filled(self) -> None: + # (2020, "s") never occurs -> empty term in the grid + expr = self._expr([2020, 2020, 2030, 2030], list("wwws")) + + grouped = expr.groupby(["period", "season"]).sum() + + cell = grouped.sel(period=2020, season="s") + assert (cell.vars == -1).all() + assert cell.coeffs.isnull().all() + + def test_dataframe_grouper_stays_compact(self) -> None: + # the DataFrame grouper keeps the stacked observed-only group dim + expr = self._expr([2020, 2020, 2030, 2030], list("wwws")) + df = expr.data[["period", "season"]].to_dataframe()[["period", "season"]] + + grouped = expr.groupby(df).sum() + + assert "group" in grouped.dims + assert isinstance(grouped.data.indexes["group"], pd.MultiIndex) + assert grouped.sizes["group"] == 3 # observed, not the 2x2=4 grid + + def test_blowup_warns_when_sparse(self) -> None: + # 200 observed combos, 200x200 grid -> nudge toward observed=True + expr = self._expr(list(range(200)), list(range(200))) + + with pytest.warns(UserWarning, match="dense .* grid"): + expr.groupby(["period", "season"]).sum() + + def test_no_warning_when_dense(self) -> None: + expr = self._expr([2020, 2020, 2030, 2030], list("wsws")) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + expr.groupby(["period", "season"]).sum() + + def test_observed_keeps_stacked(self) -> None: + # observed=True skips the unstack: compact stacked MultiIndex, + # identical to the DataFrame grouper output + expr = self._expr([2020, 2020, 2030, 2030], list("wwws")) + df = expr.data[["period", "season"]].to_dataframe()[["period", "season"]] + + grouped = expr.groupby(["period", "season"]).sum(observed=True) + + assert_linequal(grouped, expr.groupby(df).sum()) + assert grouped.sizes["group"] == 3 # observed, not the 2x2=4 grid + + def test_observed_silences_blowup_warning(self) -> None: + expr = self._expr(list(range(200)), list(range(200))) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + grouped = expr.groupby(["period", "season"]).sum(observed=True) + + assert grouped.sizes["group"] == 200 + + def test_observed_with_fallback_raises(self) -> None: + expr = self._expr([2020, 2020], list("ws")) + + with pytest.raises(ValueError, match="observed"): + expr.groupby(["period", "season"]).sum(use_fallback=True, observed=True) + + +class TestGroupbyByAttachedCoordinate: + """ + Group by an attached non-dimension coordinate. + + Asserts grouping against hard-coded ``vars``/``coeffs`` to catch regressions. + """ + + @pytest.fixture + def t(self) -> pd.RangeIndex: + return pd.RangeIndex(4, name="t") + + @pytest.fixture + def period(self, t: pd.RangeIndex) -> xr.DataArray: + return xr.DataArray( + [2020, 2020, 2030, 2030], dims="t", coords={"t": t}, name="period" + ) + + @pytest.fixture + def season(self, t: pd.RangeIndex) -> xr.DataArray: + return xr.DataArray(list("wsws"), dims="t", coords={"t": t}, name="season") + + @pytest.fixture + def expr( + self, t: pd.RangeIndex, period: xr.DataArray, season: xr.DataArray + ) -> LinearExpression: + m = Model() + x = m.add_variables(coords=[t], name="x") + return (2.0 * x).assign_coords(period=period, season=season) + + @pytest.mark.parametrize("use_fallback", [True, False]) + @pytest.mark.parametrize("by", ["name", "dataarray"]) + def test_single_key( + self, + expr: LinearExpression, + period: xr.DataArray, + by: str, + use_fallback: bool, + ) -> None: + group = "period" if by == "name" else period + + grouped = expr.groupby(group).sum(use_fallback=use_fallback) + + assert grouped.data.period.values.tolist() == [2020, 2030] + assert grouped.vars.transpose("period", TERM_DIM).values.tolist() == [ + [0, 1], + [2, 3], + ] + assert grouped.coeffs.transpose("period", TERM_DIM).values.tolist() == [ + [2.0, 2.0], + [2.0, 2.0], + ] + + @pytest.mark.parametrize("spelling", [list, tuple], ids=["list", "tuple"]) + def test_multi_key(self, expr: LinearExpression, spelling: type) -> None: + # A multi-key group always goes through the xarray fallback (a list is + # not a fast-path type), so there is no separate use_fallback case. + group = spelling(["period", "season"]) + + grouped = expr.groupby(group).sum() + + assert dict(grouped.sizes) == {"period": 2, "season": 2, TERM_DIM: 1} + assert grouped.data.period.values.tolist() == [2020, 2030] + assert grouped.data.season.values.tolist() == ["s", "w"] + assert grouped.vars.transpose("period", "season", TERM_DIM).values.tolist() == [ + [[1], [0]], + [[3], [2]], + ] + assert (grouped.coeffs == 2.0).all() + + def test_extra_aux_coord_does_not_change_result( + self, t: pd.RangeIndex, period: xr.DataArray + ) -> None: + # A second auxiliary coord on the grouped dimension must neither break + # the reshape (it raised ``KeyError`` before the fix) nor change the sum. + m = Model() + x = m.add_variables(coords=[t], name="x") + timestep = xr.DataArray( + list("abab"), dims="t", coords={"t": t}, name="timestep" + ) + expr = (2.0 * x).assign_coords(period=period, timestep=timestep) + + grouped = expr.groupby("period").sum() + + assert "timestep" not in grouped.coords + assert grouped.vars.transpose("period", TERM_DIM).values.tolist() == [ + [0, 1], + [2, 3], + ] + assert (grouped.coeffs == 2.0).all() + + @pytest.mark.parametrize("by", ["name", "dataarray"]) + def test_two_dimensional(self, by: str) -> None: + # Grouping one dimension of a 2-D variable by an aux coord must keep the + # other dimension intact and pair up the right variable labels. + m = Model() + snapshot = pd.RangeIndex(4, name="snapshot") + gen = pd.Index(["g1", "g2"], name="gen") + y = m.add_variables(coords=[snapshot, gen], name="y") # labels 0..7 + period = xr.DataArray( + [2020, 2020, 2030, 2030], + dims="snapshot", + coords={"snapshot": snapshot}, + name="period", + ) + expr = (1.0 * y).assign_coords(period=period) + group = "period" if by == "name" else period + + grouped = expr.groupby(group).sum() + + assert grouped.data.period.values.tolist() == [2020, 2030] + assert grouped.data.gen.values.tolist() == ["g1", "g2"] + assert grouped.vars.transpose("period", "gen", TERM_DIM).values.tolist() == [ + [[0, 2], [1, 3]], + [[4, 6], [5, 7]], + ] + assert (grouped.coeffs == 1.0).all() + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_dimension_coordinate_by_name(self, use_fallback: bool) -> None: + # A dimension coordinate may also be grouped by name; it collapses that + # dimension and keeps the other one. + m = Model() + snapshot = pd.RangeIndex(4, name="snapshot") + gen = pd.Index(["g1", "g2"], name="gen") + y = m.add_variables(coords=[snapshot, gen], name="y") # labels 0..7 + + grouped = (1 * y).groupby("gen").sum(use_fallback=use_fallback) + + assert grouped.data.gen.values.tolist() == ["g1", "g2"] + assert grouped.sizes["snapshot"] == 4 + assert grouped.vars.transpose("gen", "snapshot", TERM_DIM).values.tolist() == [ + [[0], [2], [4], [6]], + [[1], [3], [5], [7]], + ] + + @pytest.mark.parametrize("use_fallback", [True, False]) + def test_single_element_list_groups_like_scalar( + self, expr: LinearExpression, use_fallback: bool + ) -> None: + # ``groupby(["period"])`` groups like the scalar key, mirroring xarray. + grouped = expr.groupby(["period"]).sum(use_fallback=use_fallback) + + assert grouped.data.period.values.tolist() == [2020, 2030] + assert grouped.vars.transpose("period", TERM_DIM).values.tolist() == [ + [0, 1], + [2, 3], + ] + assert (grouped.coeffs == 2.0).all() + + def test_multi_key_dataarrays_unsupported( + self, expr: LinearExpression, period: xr.DataArray, season: xr.DataArray + ) -> None: + # Multi-key grouping must be spelled with names; a list of DataArrays + # is unhashable and raises in xarray itself, so linopy mirrors that. + with pytest.raises(TypeError, match="unhashable"): + expr.groupby([period, season]).sum() + + @pytest.mark.parametrize("use_fallback", [True, False]) + @pytest.mark.parametrize( + "level, values, vars_", + [ + ("period", [2020, 2030], [[0, 1, 2], [3, 4, 5]]), + ("timestep", ["t1", "t2", "t3"], [[0, 3], [1, 4], [2, 5]]), + ], + ) + def test_multiindex_level( + self, level: str, values: list, vars_: list, use_fallback: bool + ) -> None: + # Grouping by a level of a real ``MultiIndex`` dimension (the + # pydata/xarray#6836 case, fixed upstream) works through linopy. + m = Model() + mi = pd.MultiIndex.from_product( + [[2020, 2030], ["t1", "t2", "t3"]], names=["period", "timestep"] + ) + x = m.add_variables(coords={"snapshot": mi}, name="x") # labels 0..5 + + grouped = (1 * x).groupby(level).sum(use_fallback=use_fallback) + + assert grouped.data[level].values.tolist() == values + assert grouped.vars.transpose(level, TERM_DIM).values.tolist() == vars_ + + @pytest.mark.parametrize("use_fallback", [True]) def test_linear_expression_groupby_ndim(z: Variable, use_fallback: bool) -> None: # TODO: implement fallback for n-dim groupby, see https://github.com/PyPSA/linopy/issues/299 @@ -1313,3 +2170,450 @@ def test_simplify_partial_cancellation(x: Variable, y: Variable) -> None: assert all(simplified.coeffs.values == 3.0), ( f"Expected coefficient 3.0, got {simplified.coeffs.values}" ) + + +def test_constant_only_expression_mul_dataarray(m: Model) -> None: + const_arr = xr.DataArray([2, 3], dims=["dim_0"]) + const_expr = LinearExpression(const_arr, m) + assert const_expr.is_constant + assert const_expr.nterm == 0 + + data_arr = xr.DataArray([10, 20], dims=["dim_0"]) + expected_const = const_arr * data_arr + + result = const_expr * data_arr + assert isinstance(result, LinearExpression) + assert result.is_constant + assert (result.const == expected_const).all() + + result_rev = data_arr * const_expr + assert isinstance(result_rev, LinearExpression) + assert result_rev.is_constant + assert (result_rev.const == expected_const).all() + + +def test_constant_only_expression_mul_linexpr_with_vars(m: Model, x: Variable) -> None: + const_arr = xr.DataArray([2, 3], dims=["dim_0"]) + const_expr = LinearExpression(const_arr, m) + assert const_expr.is_constant + assert const_expr.nterm == 0 + + expr_with_vars = 1 * x + 5 + expected_coeffs = const_arr + expected_const = const_arr * 5 + + result = const_expr * expr_with_vars + assert isinstance(result, LinearExpression) + assert (result.coeffs == expected_coeffs).all() + assert (result.const == expected_const).all() + + result_rev = expr_with_vars * const_expr + assert isinstance(result_rev, LinearExpression) + assert (result_rev.coeffs == expected_coeffs).all() + assert (result_rev.const == expected_const).all() + + +def test_constant_only_expression_mul_constant_only(m: Model) -> None: + const_arr = xr.DataArray([2, 3], dims=["dim_0"]) + const_arr2 = xr.DataArray([4, 5], dims=["dim_0"]) + const_expr = LinearExpression(const_arr, m) + const_expr2 = LinearExpression(const_arr2, m) + assert const_expr.is_constant + assert const_expr2.is_constant + + expected_const = const_arr * const_arr2 + + result = const_expr * const_expr2 + assert isinstance(result, LinearExpression) + assert result.is_constant + assert (result.const == expected_const).all() + + result_rev = const_expr2 * const_expr + assert isinstance(result_rev, LinearExpression) + assert result_rev.is_constant + assert (result_rev.const == expected_const).all() + + +def test_constant_only_expression_mul_linexpr_with_vars_and_const( + m: Model, x: Variable +) -> None: + const_arr = xr.DataArray([2, 3], dims=["dim_0"]) + const_expr = LinearExpression(const_arr, m) + assert const_expr.is_constant + + expr_with_vars_and_const = 4 * x + 10 + expected_coeffs = const_arr * 4 + expected_const = const_arr * 10 + + result = const_expr * expr_with_vars_and_const + assert isinstance(result, LinearExpression) + assert not result.is_constant + assert (result.coeffs == expected_coeffs).all() + assert (result.const == expected_const).all() + + result_rev = expr_with_vars_and_const * const_expr + assert isinstance(result_rev, LinearExpression) + assert not result_rev.is_constant + assert (result_rev.coeffs == expected_coeffs).all() + assert (result_rev.const == expected_const).all() + + +def test_variable_names() -> None: + m = Model() + time = pd.Index(range(3), name="time") + + a = m.add_variables(name="a", coords=[time]) + b = m.add_variables(name="b", coords=[time]) + + expr = a + b + assert expr.nterm == 2 + assert expr.variable_names == {"a", "b"} + + mask = xr.DataArray(False, coords=[time]) + expr = a + (b * 1).where(mask) + assert expr.nterm == 2 + assert expr.variable_names == {"a"} + + expr = (b * 1).where(mask) + assert expr.nterm == 1 + assert expr.variable_names == set() + + expr = LinearExpression.from_constant(model=m, constant=5) + assert expr.nterm == 0 + assert expr.variable_names == set() + + # Single variable expression + expr = 1 * a + assert expr.variable_names == {"a"} + + # Repeated variable across terms (a + a) + expr = a + a + assert expr.variable_names == {"a"} + + +def test_nterm() -> None: + m = Model() + time = pd.Index(range(3), name="time") + all_false = xr.DataArray(False, coords=[time]) + not_0 = xr.DataArray([False, True, True], coords=[time]) + not_1 = xr.DataArray([True, False, True], coords=[time]) + not_2 = xr.DataArray([True, True, False], coords=[time]) + + a = m.add_variables(name="a", coords=[time]) + b = m.add_variables(name="b", coords=[time]) + c = m.add_variables(name="c", coords=[time]) + + expr = (a.where(not_0) + b.where(not_1) + c.where(not_2)).densify_terms() + assert expr.nterm == 3 + + expr = a + b.where(all_false) + assert expr.nterm == 2 + + expr = expr.simplify() + assert expr.nterm == 1 + + +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"] + + class TestAddition: + 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.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.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.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.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.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.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.indexes["i"]) == [0, 1, 2] + assert (result.const.values == const.values).all() + + 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]}) + joins: list[JoinOptions] = ["override", "outer", "inner"] + for join in joins: + 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] + + class TestSubtraction: + 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.indexes["i"]) == [1, 2] + + 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]) + + class TestMultiplication: + 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.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.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_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") + + class TestDivision: + 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.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.indexes["i"]) == [0, 1, 2, 3] + + 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") + + class TestVariableOperations: + def test_variable_add_join(self, a: Variable, b: Variable) -> None: + result = a.add(b, join="inner") + assert list(result.indexes["i"]) == [1, 2] + + def test_variable_sub_join(self, a: Variable, b: Variable) -> None: + result = a.sub(b, join="inner") + assert list(result.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.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.indexes["i"]) == [1, 2] + + 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_same_shape_add_join_override(self, a: Variable, c: Variable) -> None: + result = a.to_linexpr().add(c.to_linexpr(), join="override") + assert list(result.indexes["i"]) == [0, 1, 2] + + class TestMerge: + def test_merge_join_parameter(self, a: Variable, b: Variable) -> None: + result = merge( + [a.to_linexpr(), b.to_linexpr()], cls=LinearExpression, join="inner" + ) + assert list(result.indexes["i"]) == [1, 2] + + def test_merge_outer_join(self, a: Variable, b: Variable) -> None: + result = merge( + [a.to_linexpr(), b.to_linexpr()], cls=LinearExpression, join="outer" + ) + assert set(result.coords["i"].values) == {0, 1, 2, 3} + + def test_merge_join_left(self, a: Variable, b: Variable) -> None: + result = merge( + [a.to_linexpr(), b.to_linexpr()], cls=LinearExpression, join="left" + ) + assert list(result.indexes["i"]) == [0, 1, 2] + + def test_merge_join_right(self, a: Variable, b: Variable) -> None: + result = merge( + [a.to_linexpr(), b.to_linexpr()], cls=LinearExpression, join="right" + ) + assert list(result.indexes["i"]) == [1, 2, 3] + + class TestValueVerification: + 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_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) + + class TestQuadratic: + 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.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.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.indexes["i"]) == [1, 2, 3] diff --git a/test/test_model.py b/test/test_model.py index c363fe4c1..6b9e31576 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -5,6 +5,8 @@ from __future__ import annotations +import copy as pycopy +import weakref from pathlib import Path from tempfile import gettempdir @@ -12,8 +14,13 @@ import pytest import xarray as xr -from linopy import EQUAL, Model -from linopy.testing import assert_model_equal +from linopy import EQUAL, Model, available_solvers +from linopy.testing import ( + assert_conequal, + assert_equal, + assert_linequal, + assert_model_equal, +) target_shape: tuple[int, int] = (10, 10) @@ -35,6 +42,38 @@ def test_model_solver_dir() -> None: assert m.solver_dir == Path(d) +def test_model_config_defaults() -> None: + m = Model(freeze_constraints=True, set_names_in_solver_io=False) + assert m.freeze_constraints is True + assert m.set_names_in_solver_io is False + + +def test_model_copy_preserves_config() -> None: + m = Model(freeze_constraints=True, set_names_in_solver_io=False) + copied = m.copy() + assert copied.freeze_constraints is True + assert copied.set_names_in_solver_io is False + + +def test_model_is_weakrefable() -> None: + m: Model = Model() + ref = weakref.ref(m) + assert ref() is m + + +def test_model_weakkeydict_use_case() -> None: + # third-party extensions rely on WeakKeyDictionary for per-instance storage + registry: weakref.WeakKeyDictionary[Model, str] = weakref.WeakKeyDictionary() + m: Model = Model() + registry[m] = "extension-state" + assert registry[m] == "extension-state" + del m + import gc + + gc.collect() + assert len(registry) == 0 + + def test_model_variable_getitem() -> None: m = Model() x = m.add_variables(name="x") @@ -87,6 +126,14 @@ def test_objective() -> None: m.objective = m.objective + 3 +def test_solve_without_objective_raises() -> None: + # https://github.com/PyPSA/linopy/issues/668 + m: Model = Model() + m.add_variables(lower=0, upper=10, name="myvar") + with pytest.raises(ValueError, match="No objective has been set"): + m.solve() + + def test_remove_variable() -> None: m: Model = Model() @@ -104,10 +151,11 @@ def test_remove_variable() -> None: assert "x" in m.variables - m.remove_variables("x") + with pytest.warns(UserWarning, match="con0"): + m.remove_variables("x") assert "x" not in m.variables - assert not m.constraints.con0.vars.isin(x.labels).any() + assert "con0" not in m.constraints assert not m.objective.vars.isin(x.labels).any() @@ -163,3 +211,167 @@ def test_assert_model_equal() -> None: m.add_objective(obj) assert_model_equal(m, m) + + +@pytest.fixture(scope="module") +def copy_test_model() -> Model: + """Small representative model used across copy tests.""" + m: Model = Model() + + lower: xr.DataArray = xr.DataArray( + np.zeros((10, 10)), coords=[range(10), range(10)] + ) + upper: xr.DataArray = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) + x = m.add_variables(lower, upper, name="x") + y = m.add_variables(name="y") + + m.add_constraints(1 * x + 10 * y, EQUAL, 0) + m.add_objective((10 * x + 5 * y).sum()) + + return m + + +@pytest.fixture(scope="module") +def solved_copy_test_model(copy_test_model: Model) -> Model: + """Solved representative model used across solved-copy tests.""" + m = copy_test_model.copy(deep=True) + m.solve() + return m + + +def test_model_copy_unsolved(copy_test_model: Model) -> None: + """Copy of unsolved model is structurally equal and independent.""" + m = copy_test_model.copy(deep=True) + c = m.copy(include_solution=False) + + assert_model_equal(m, c) + + # independence: mutating copy does not affect source + c.add_variables(name="z") + assert "z" not in m.variables + + +def test_model_copy_unsolved_with_solution_flag(copy_test_model: Model) -> None: + """Unsolved model with include_solution=True has no extra solve artifacts.""" + m = copy_test_model.copy(deep=True) + + c_include_solution = m.copy(include_solution=True) + c_exclude_solution = m.copy(include_solution=False) + + assert_model_equal(c_include_solution, c_exclude_solution) + assert c_include_solution.status == "initialized" + assert c_include_solution.termination_condition == "" + assert c_include_solution.objective.value is None + + +def test_model_copy_shallow(copy_test_model: Model) -> None: + """Shallow copy has independent wrappers sharing underlying data buffers.""" + m = copy_test_model.copy(deep=True) + c = m.copy(deep=False) + + assert c is not m + assert c.variables is not m.variables + assert c.constraints is not m.constraints + assert c.objective is not m.objective + + # wrappers are distinct, but shallow copy shares payload buffers + c.variables["x"].lower.values[0, 0] = 123.0 + assert m.variables["x"].lower.values[0, 0] == 123.0 + + +def test_model_deepcopy_protocol(copy_test_model: Model) -> None: + """copy.deepcopy(model) dispatches to Model.__deepcopy__ and stays independent.""" + m = copy_test_model.copy(deep=True) + c = pycopy.deepcopy(m) + + assert_model_equal(m, c) + + # Test independence: mutations to copy do not affect source + # 1. Variable mutation: add new variable + c.add_variables(name="z") + assert "z" not in m.variables + + # 2. Variable data mutation (bounds): verify buffers are independent + original_lower = m.variables["x"].lower.values[0, 0].item() + new_lower = 999 + c.variables["x"].lower.values[0, 0] = new_lower + assert c.variables["x"].lower.values[0, 0] == new_lower + assert m.variables["x"].lower.values[0, 0] == original_lower + + # 3. Constraint coefficient mutation: deep copy must not leak back + original_con_coeff = m.constraints["con0"].coeffs.values.flat[0].item() + new_con_coeff = original_con_coeff + 42 + c.constraints["con0"].coeffs.values.flat[0] = new_con_coeff + assert c.constraints["con0"].coeffs.values.flat[0] == new_con_coeff + assert m.constraints["con0"].coeffs.values.flat[0] == original_con_coeff + + # 4. Objective expression coefficient mutation: deep copy must not leak back + original_obj_coeff = m.objective.expression.coeffs.values.flat[0].item() + new_obj_coeff = original_obj_coeff + 20 + c.objective.expression.coeffs.values.flat[0] = new_obj_coeff + assert c.objective.expression.coeffs.values.flat[0] == new_obj_coeff + assert m.objective.expression.coeffs.values.flat[0] == original_obj_coeff + + # 5. Objective sense mutation + original_sense = m.objective.sense + c.objective.sense = "max" + assert c.objective.sense == "max" + assert m.objective.sense == original_sense + + +@pytest.mark.skipif(not available_solvers, reason="No solver installed") +class TestModelCopySolved: + def test_model_deepcopy_protocol_excludes_solution( + self, solved_copy_test_model: Model + ) -> None: + """copy.deepcopy on solved model drops solve state by default.""" + m = solved_copy_test_model + + c = pycopy.deepcopy(m) + + assert c.status == "initialized" + assert c.termination_condition == "" + assert c.objective.value is None + + for v in m.variables: + assert_equal( + c.variables[v].data[c.variables.dataset_attrs], + m.variables[v].data[m.variables.dataset_attrs], + ) + for con in m.constraints: + assert_conequal(c.constraints[con], m.constraints[con], strict=False) + assert_linequal(c.objective.expression, m.objective.expression) + assert c.objective.sense == m.objective.sense + + def test_model_copy_solved_with_solution( + self, solved_copy_test_model: Model + ) -> None: + """Copy with include_solution=True preserves solve state.""" + m = solved_copy_test_model + + c = m.copy(include_solution=True) + assert_model_equal(m, c) + + def test_model_copy_solved_without_solution( + self, solved_copy_test_model: Model + ) -> None: + """Copy with include_solution=False (default) drops solve state but preserves problem structure.""" + m = solved_copy_test_model + + c = m.copy(include_solution=False) + + # solve state is dropped + assert c.status == "initialized" + assert c.termination_condition == "" + assert c.objective.value is None + + # problem structure is preserved — compare only dataset_attrs to exclude solution/dual + for v in m.variables: + assert_equal( + c.variables[v].data[c.variables.dataset_attrs], + m.variables[v].data[m.variables.dataset_attrs], + ) + for con in m.constraints: + assert_conequal(c.constraints[con], m.constraints[con], strict=False) + assert_linequal(c.objective.expression, m.objective.expression) + assert c.objective.sense == m.objective.sense diff --git a/test/test_oetc_settings.py b/test/test_oetc_settings.py new file mode 100644 index 000000000..12deeb66c --- /dev/null +++ b/test/test_oetc_settings.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from linopy.remote.oetc import ( + ComputeProvider, + OetcCredentials, + OetcHandler, + OetcSettings, +) + +REQUIRED_ENV = { + "OETC_EMAIL": "test@example.com", + "OETC_PASSWORD": "secret", + "OETC_NAME": "test-job", + "OETC_AUTH_URL": "https://auth.example.com", + "OETC_ORCHESTRATOR_URL": "https://orch.example.com", +} + + +def _set_required_env(monkeypatch: pytest.MonkeyPatch) -> None: + for k, v in REQUIRED_ENV.items(): + monkeypatch.setenv(k, v) + + +def _clear_oetc_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key in [ + "OETC_EMAIL", + "OETC_PASSWORD", + "OETC_NAME", + "OETC_AUTH_URL", + "OETC_ORCHESTRATOR_URL", + "OETC_CPU_CORES", + "OETC_DISK_SPACE_GB", + "OETC_DELETE_WORKER_ON_ERROR", + ]: + monkeypatch.delenv(key, raising=False) + + +def test_from_env_all_set(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_CPU_CORES", "8") + monkeypatch.setenv("OETC_DISK_SPACE_GB", "20") + monkeypatch.setenv("OETC_DELETE_WORKER_ON_ERROR", "true") + + s = OetcSettings.from_env() + assert s.credentials.email == "test@example.com" + assert s.credentials.password == "secret" + assert s.name == "test-job" + assert s.cpu_cores == 8 + assert s.disk_space_gb == 20 + assert s.compute_provider == ComputeProvider.GCP + assert s.delete_worker_on_error is True + + +def test_from_env_kwargs_override(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + + s = OetcSettings.from_env(email="override@example.com") + assert s.credentials.email == "override@example.com" + + +def test_from_env_missing_required(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + with pytest.raises( + ValueError, + match="OETC_EMAIL.*OETC_PASSWORD.*OETC_NAME.*OETC_AUTH_URL.*OETC_ORCHESTRATOR_URL", + ): + OetcSettings.from_env() + + +def test_from_env_empty_string_required(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + monkeypatch.setenv("OETC_EMAIL", "") + monkeypatch.setenv("OETC_PASSWORD", " ") + monkeypatch.setenv("OETC_NAME", "valid") + monkeypatch.setenv("OETC_AUTH_URL", "https://auth.example.com") + monkeypatch.setenv("OETC_ORCHESTRATOR_URL", "https://orch.example.com") + + with pytest.raises(ValueError, match="OETC_EMAIL.*OETC_PASSWORD"): + OetcSettings.from_env() + + +def test_from_env_partial_kwargs(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + monkeypatch.setenv("OETC_NAME", "env-name") + monkeypatch.setenv("OETC_AUTH_URL", "https://auth.example.com") + monkeypatch.setenv("OETC_ORCHESTRATOR_URL", "https://orch.example.com") + + s = OetcSettings.from_env(email="a@b.com", password="pw") + assert s.credentials.email == "a@b.com" + assert s.name == "env-name" + + +def test_from_env_defaults_applied(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + + s = OetcSettings.from_env() + assert s.solver == "highs" + assert s.solver_options == {} + assert s.cpu_cores == 2 + assert s.disk_space_gb == 10 + assert s.compute_provider == ComputeProvider.GCP + assert s.delete_worker_on_error is False + + +def test_from_env_cpu_cores_valid(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_CPU_CORES", "4") + + assert OetcSettings.from_env().cpu_cores == 4 + + +def test_from_env_cpu_cores_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_CPU_CORES", "abc") + + with pytest.raises(ValueError, match="OETC_CPU_CORES"): + OetcSettings.from_env() + + +@pytest.mark.parametrize("val", ["true", "1", "yes"]) +def test_from_env_bool_true_values(monkeypatch: pytest.MonkeyPatch, val: str) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_DELETE_WORKER_ON_ERROR", val) + + assert OetcSettings.from_env().delete_worker_on_error is True + + +@pytest.mark.parametrize("val", ["false", "0", "no"]) +def test_from_env_bool_false_values(monkeypatch: pytest.MonkeyPatch, val: str) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_DELETE_WORKER_ON_ERROR", val) + + assert OetcSettings.from_env().delete_worker_on_error is False + + +def test_from_env_bool_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_DELETE_WORKER_ON_ERROR", "maybe") + + with pytest.raises(ValueError, match="OETC_DELETE_WORKER_ON_ERROR"): + OetcSettings.from_env() + + +def _make_handler(settings: OetcSettings) -> OetcHandler: + with ( + patch("linopy.remote.oetc._oetc_deps_available", True), + patch.object(OetcHandler, "_OetcHandler__sign_in", return_value=MagicMock()), + patch.object( + OetcHandler, + "_OetcHandler__get_cloud_provider_credentials", + return_value=MagicMock(), + ), + ): + return OetcHandler(settings) + + +def _default_settings(**overrides: Any) -> OetcSettings: + defaults: dict[str, Any] = dict( + credentials=OetcCredentials(email="a@b.com", password="pw"), + name="test", + authentication_server_url="https://auth", + orchestrator_server_url="https://orch", + solver="highs", + solver_options={"TimeLimit": 100}, + ) + defaults.update(overrides) + return OetcSettings(**defaults) + + +def test_solve_on_oetc_mutation_safety() -> None: + settings = _default_settings() + handler = _make_handler(settings) + original_opts = dict(settings.solver_options) + + mock_model = MagicMock() + mock_solved = MagicMock() + mock_solved.objective.value = 42.0 + mock_solved.status = "ok" + + with ( + patch.object(handler, "_upload_file_to_gcp", return_value="file.nc.gz"), + patch.object(handler, "_submit_job_to_compute_service", return_value="uuid"), + patch.object(handler, "wait_and_get_job_data") as mock_wait, + patch.object(handler, "_download_file_from_gcp", return_value="/tmp/sol.nc"), + patch("linopy.read_netcdf", return_value=mock_solved), + patch("os.remove"), + ): + mock_wait.return_value = MagicMock(output_files=["out.nc.gz"]) + + handler.solve_on_oetc(mock_model, Extra=999) + handler.solve_on_oetc(mock_model, Other=1) + + assert settings.solver_options == original_opts + + +def test_solve_on_oetc_solver_name_override() -> None: + settings = _default_settings() + handler = _make_handler(settings) + + mock_model = MagicMock() + mock_solved = MagicMock() + mock_solved.objective.value = 1.0 + mock_solved.status = "ok" + + with ( + patch.object(handler, "_upload_file_to_gcp", return_value="file.nc.gz"), + patch.object( + handler, "_submit_job_to_compute_service", return_value="uuid" + ) as mock_submit, + patch.object(handler, "wait_and_get_job_data") as mock_wait, + patch.object(handler, "_download_file_from_gcp", return_value="/tmp/sol.nc"), + patch("linopy.read_netcdf", return_value=mock_solved), + patch("os.remove"), + ): + mock_wait.return_value = MagicMock(output_files=["out.nc.gz"]) + + handler.solve_on_oetc(mock_model, solver_name="gurobi") + + mock_submit.assert_called_once() + assert mock_submit.call_args[0][1] == "gurobi" + + +def test_solve_on_oetc_solver_options_merge_precedence() -> None: + settings = _default_settings(solver_options={"TimeLimit": 100}) + handler = _make_handler(settings) + + mock_model = MagicMock() + mock_solved = MagicMock() + mock_solved.objective.value = 1.0 + mock_solved.status = "ok" + + with ( + patch.object(handler, "_upload_file_to_gcp", return_value="file.nc.gz"), + patch.object( + handler, "_submit_job_to_compute_service", return_value="uuid" + ) as mock_submit, + patch.object(handler, "wait_and_get_job_data") as mock_wait, + patch.object(handler, "_download_file_from_gcp", return_value="/tmp/sol.nc"), + patch("linopy.read_netcdf", return_value=mock_solved), + patch("os.remove"), + ): + mock_wait.return_value = MagicMock(output_files=["out.nc.gz"]) + + handler.solve_on_oetc(mock_model, TimeLimit=200) + + mock_submit.assert_called_once() + assert mock_submit.call_args[0][2] == {"TimeLimit": 200} + + +def test_solve_on_oetc_solver_name_default_fallback() -> None: + settings = _default_settings(solver="cplex") + handler = _make_handler(settings) + + mock_model = MagicMock() + mock_solved = MagicMock() + mock_solved.objective.value = 1.0 + mock_solved.status = "ok" + + with ( + patch.object(handler, "_upload_file_to_gcp", return_value="file.nc.gz"), + patch.object( + handler, "_submit_job_to_compute_service", return_value="uuid" + ) as mock_submit, + patch.object(handler, "wait_and_get_job_data") as mock_wait, + patch.object(handler, "_download_file_from_gcp", return_value="/tmp/sol.nc"), + patch("linopy.read_netcdf", return_value=mock_solved), + patch("os.remove"), + ): + mock_wait.return_value = MagicMock(output_files=["out.nc.gz"]) + + handler.solve_on_oetc(mock_model) + + mock_submit.assert_called_once() + assert mock_submit.call_args[0][1] == "cplex" + + +def test_from_env_disk_space_gb_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + _clear_oetc_env(monkeypatch) + _set_required_env(monkeypatch) + monkeypatch.setenv("OETC_DISK_SPACE_GB", "abc") + + with pytest.raises(ValueError, match="OETC_DISK_SPACE_GB"): + OetcSettings.from_env() + + +def test_model_solve_forwards_to_oetc() -> None: + from linopy import Model + + m = Model() + x = m.add_variables(lower=0, name="x") + m.add_objective(1 * x) + + handler = MagicMock(spec=OetcHandler) + mock_solved = MagicMock() + mock_solved.status = "ok" + mock_solved.termination_condition = "optimal" + mock_solved.objective.value = 10.0 + mock_solved.variables.items.return_value = [(k, v) for k, v in m.variables.items()] + mock_solved.constraints.items.return_value = [] + for k in m.variables: + mock_solved.variables[k].solution = 0.0 + handler.solve_on_oetc.return_value = mock_solved + + m.solve(solver_name="gurobi", remote=handler, TimeLimit=100) + + handler.solve_on_oetc.assert_called_once_with( + m, solver_name="gurobi", reformulate_sos=False, TimeLimit=100 + ) diff --git a/test/test_optimization.py b/test/test_optimization.py index ff790d6e8..1e771b221 100644 --- a/test/test_optimization.py +++ b/test/test_optimization.py @@ -21,11 +21,16 @@ from linopy.common import to_path from linopy.expressions import LinearExpression from linopy.solver_capabilities import ( - SolverFeature, get_available_solvers_with_feature, solver_supports, ) -from linopy.solvers import _new_highspy_mps_layout, available_solvers, quadratic_solvers +from linopy.solvers import ( + SolverFeature, + _new_highspy_mps_layout, + _solver_class_for, + licensed_solvers, + quadratic_solvers, +) logger = logging.getLogger(__name__) @@ -33,38 +38,42 @@ explicit_coordinate_names = [False, True] -if "highs" in available_solvers: +if "highs" in licensed_solvers: # mps io is only supported via highspy io_apis.append("mps") file_io_solvers = get_available_solvers_with_feature( - SolverFeature.READ_MODEL_FROM_FILE, available_solvers + SolverFeature.READ_MODEL_FROM_FILE, licensed_solvers ) params: list[tuple[str, str, bool]] = list( itertools.product(file_io_solvers, io_apis, explicit_coordinate_names) ) direct_solvers = get_available_solvers_with_feature( - SolverFeature.DIRECT_API, available_solvers + SolverFeature.DIRECT_API, licensed_solvers ) for solver in direct_solvers: params.append((solver, "direct", False)) -if "mosek" in available_solvers: +set_names_direct_solvers = [ + solver for solver in ("highs", "gurobi") if solver in direct_solvers +] + +if "mosek" in licensed_solvers: params.append(("mosek", "lp", False)) params.append(("mosek", "lp", True)) -# Note: Platform-specific solver bugs (e.g., SCIP quadratic on Windows) are now +# Note: Platform-specific solver bugs are now # handled in linopy/solver_capabilities.py by adjusting the registry at import time. feasible_quadratic_solvers: list[str] = list(quadratic_solvers) feasible_mip_solvers: list[str] = get_available_solvers_with_feature( - SolverFeature.INTEGER_VARIABLES, available_solvers + SolverFeature.INTEGER_VARIABLES, licensed_solvers ) gpu_solvers: list[str] = get_available_solvers_with_feature( - SolverFeature.GPU_ACCELERATION, available_solvers + SolverFeature.GPU_ONLY, licensed_solvers ) # set tolerances for solution checking based on solver type (CPU vs. GPU) @@ -75,7 +84,7 @@ def test_print_solvers(capsys: Any) -> None: with capsys.disabled(): print( - f"\ntesting solvers: {', '.join(available_solvers)}\n" + f"\ntesting solvers: {', '.join(licensed_solvers)}\n" f"testing quadratic solvers: {', '.join(feasible_quadratic_solvers)}" ) @@ -298,7 +307,7 @@ def modified_model() -> Model: x = m.add_variables(coords=[lower.index], name="x", binary=True) y = m.add_variables(lower, name="y") - c = m.add_constraints(x + y, GREATER_EQUAL, 10) + c = m.add_constraints(x + y, GREATER_EQUAL, 10, freeze=False) y.lower = 9 c.lhs = 2 * x + y @@ -464,7 +473,7 @@ def test_model_maximization( assert m.objective.sense == "max" assert m.objective.value is None - if solver in ["cbc", "glpk"] and io_api == "mps" and _new_highspy_mps_layout: + if solver in ["cbc", "glpk"] and io_api == "mps" and _new_highspy_mps_layout(): with pytest.raises(ValueError): m.solve( solver, @@ -492,6 +501,22 @@ def test_mock_solve(model_maximization: Model) -> None: assert (x_solution == 0).all() +@pytest.mark.skipif("highs" not in licensed_solvers, reason="HiGHS is not installed") +def test_mock_solve_clears_existing_solver_state(model: Model) -> None: + status, condition = model.solve(solver_name="highs", io_api="direct") + assert status == "ok" + assert model.solver is not None + assert model.solver_model is not None + assert model.solver_name == "highs" + + status, condition = model.solve(solver="some_non_existant_solver", mock_solve=True) + + assert status == "ok" + assert model.solver is None + assert model.solver_model is None + assert model.solver_name is None + + @pytest.mark.parametrize("solver,io_api,explicit_coordinate_names", params) def test_default_settings_chunked( model_chunked: Model, solver: str, io_api: str, explicit_coordinate_names: bool @@ -530,7 +555,7 @@ def test_solver_time_limit_options( "cplex": {"timelimit": 1}, "xpress": {"maxtime": 1}, "highs": {"time_limit": 1}, - "scip": {"limits/time": 1}, + "scip": {"limits/time": 10}, # increase time limit to avoid race condition "mosek": {"MSK_DPAR_OPTIMIZER_MAX_TIME": 1}, "mindopt": {"MaxTime": 1}, "copt": {"TimeLimit": 1}, @@ -672,7 +697,9 @@ def test_infeasible_model( with pytest.warns(DeprecationWarning): model.compute_set_of_infeasible_constraints() model.compute_infeasibilities() - model.print_infeasibilities() + formatted = model.format_infeasibilities() + assert isinstance(formatted, str) + assert formatted else: with pytest.raises((NotImplementedError, ImportError)): model.compute_infeasibilities() @@ -710,6 +737,50 @@ def test_milp_binary_model( ).all() +FIXED_VAR_CASES = [ + pytest.param("continuous", {}, 7.0, 100, id="continuous-lower-raised"), + pytest.param("continuous", {}, 3.0, -100, id="continuous-upper-lowered"), + pytest.param("integer", {"integer": True}, 7.0, 100, id="integer-lower-raised"), + pytest.param("integer", {"integer": True}, 3.0, -100, id="integer-upper-lowered"), + pytest.param("binary", {"binary": True}, 1.0, 100, id="binary-1"), + pytest.param("binary", {"binary": True}, 0.0, -100, id="binary-0"), +] + + +@pytest.mark.parametrize("kind,var_kwargs,fixval,coef", FIXED_VAR_CASES) +@pytest.mark.parametrize( + "solver,io_api,explicit_coordinate_names", + [p for p in params if p[0] not in ["mindopt"]], +) +def test_fixed_variable_is_held( + solver: str, + io_api: str, + explicit_coordinate_names: bool, + kind: str, + var_kwargs: dict, + fixval: float, + coef: float, +) -> None: + if kind in ("integer", "binary") and solver not in feasible_mip_solvers: + pytest.skip(f"{solver} does not support MIP") + + m = Model() + if "binary" in var_kwargs: + v = m.add_variables(name="v", **var_kwargs) + else: + v = m.add_variables(lower=0, upper=10, name="v", **var_kwargs) + x = m.add_variables(lower=0, upper=10, name="x") + m.add_constraints(x >= 2, name="c") + m.add_objective(x + coef * v) + v.fix(fixval) + + status, condition = m.solve( + solver, io_api=io_api, explicit_coordinate_names=explicit_coordinate_names + ) + assert condition == "optimal" + assert float(m.solution.v) == pytest.approx(fixval) + + @pytest.mark.parametrize( "solver,io_api,explicit_coordinate_names", [p for p in params if p[0] not in ["mindopt"]], @@ -748,6 +819,15 @@ def test_milp_model( assert condition == "optimal" assert ((milp_model.solution.y == 9) | (milp_model.solution.x == 0.5)).all() + solver_cls = _solver_class_for(solver) + if solver_cls is not None and solver_cls.supports( + SolverFeature.MIP_DUAL_BOUND_REPORT + ): + assert milp_model.solver is not None + report = milp_model.solver.report + assert report is not None + assert report.dual_bound is not None + @pytest.mark.parametrize( "solver,io_api,explicit_coordinate_names", @@ -986,7 +1066,7 @@ def test_solution_fn_parent_dir_doesnt_exist( assert status == "ok" -@pytest.mark.parametrize("solver", available_solvers) +@pytest.mark.parametrize("solver", licensed_solvers) def test_non_supported_solver_io(model: Model, solver: str) -> None: with pytest.raises(ValueError): model.solve(solver, io_api="non_supported") @@ -1006,6 +1086,60 @@ def test_solver_attribute_getter( assert set(rc) == set(model.variables) +def assert_semantically_equal_direct_solves( + solved_with_names: Model, solved_without_names: Model, solver: str +) -> None: + tol = GPU_SOL_TOL if solver in gpu_solvers else CPU_SOL_TOL + + assert solved_with_names.status == solved_without_names.status + assert ( + solved_with_names.termination_condition + == solved_without_names.termination_condition + ) + assert solved_with_names.objective.value is not None + assert solved_without_names.objective.value is not None + assert solved_with_names.objective.value == pytest.approx( + solved_without_names.objective.value, rel=tol + ) + assert_allclose( + solved_with_names.solution, + solved_without_names.solution, + rtol=tol, + atol=tol, + ) + + dual_with_names = solved_with_names.dual + dual_without_names = solved_without_names.dual + assert set(dual_with_names.data_vars) == set(dual_without_names.data_vars) + if dual_with_names.data_vars: + assert_allclose( + dual_with_names, + dual_without_names, + rtol=tol, + atol=tol, + ) + + +@pytest.mark.parametrize("solver", set_names_direct_solvers) +def test_direct_solve_set_names_semantic_equivalence(model: Model, solver: str) -> None: + model_with_names = model.copy(deep=True) + model_without_names = model.copy(deep=True) + + status_with_names, condition_with_names = model_with_names.solve( + solver_name=solver, io_api="direct", set_names=True + ) + status_without_names, condition_without_names = model_without_names.solve( + solver_name=solver, io_api="direct", set_names=False + ) + + assert status_with_names == "ok" + assert status_without_names == "ok" + assert condition_with_names == condition_without_names + assert_semantically_equal_direct_solves( + model_with_names, model_without_names, solver + ) + + @pytest.mark.parametrize("solver,io_api,explicit_coordinate_names", params) def test_model_resolve( model: Model, solver: str, io_api: str, explicit_coordinate_names: bool @@ -1038,7 +1172,7 @@ def test_solver_classes_from_problem_file( ) -> None: # first test initialization of super class. Should not be possible to initialize with pytest.raises(TypeError): - solvers.Solver() # type: ignore + solvers.Solver() # initialize the solver as object of solver subclass solver_class = getattr(solvers, f"{solvers.SolverName(solver).name}") @@ -1091,6 +1225,70 @@ def test_solver_classes_direct( solver_.solve_problem(model=model) +@pytest.fixture +def auto_mask_variable_model() -> Model: + """Model with auto_mask=True and NaN in variable bounds.""" + m = Model(auto_mask=True) + + x = m.add_variables(lower=0, coords=[range(10)], name="x") + lower = pd.Series([0.0] * 8 + [np.nan, np.nan], range(10)) + y = m.add_variables(lower=lower, name="y") # NaN bounds auto-masked + + m.add_constraints(x + y, GREATER_EQUAL, 10) + m.add_constraints(y, GREATER_EQUAL, 0) + m.add_objective(2 * x + y) + return m + + +@pytest.fixture +def auto_mask_constraint_model() -> Model: + """Model with auto_mask=True and NaN in constraint RHS.""" + m = Model(auto_mask=True) + + x = m.add_variables(lower=0, coords=[range(10)], name="x") + y = m.add_variables(lower=0, coords=[range(10)], name="y") + + rhs = pd.Series([10.0] * 8 + [np.nan, np.nan], range(10)) + m.add_constraints(x + y, GREATER_EQUAL, rhs) # NaN rhs auto-masked + m.add_constraints(x + y, GREATER_EQUAL, 5) + + m.add_objective(2 * x + y) + return m + + +@pytest.mark.parametrize("solver,io_api,explicit_coordinate_names", params) +def test_auto_mask_variable_model( + auto_mask_variable_model: Model, + solver: str, + io_api: str, + explicit_coordinate_names: bool, +) -> None: + """Test that auto_mask=True correctly masks variables with NaN bounds.""" + auto_mask_variable_model.solve( + solver, io_api=io_api, explicit_coordinate_names=explicit_coordinate_names + ) + y = auto_mask_variable_model.variables.y + # Same assertions as test_masked_variable_model + assert y.solution[-2:].isnull().all() + assert y.solution[:-2].notnull().all() + + +@pytest.mark.parametrize("solver,io_api,explicit_coordinate_names", params) +def test_auto_mask_constraint_model( + auto_mask_constraint_model: Model, + solver: str, + io_api: str, + explicit_coordinate_names: bool, +) -> None: + """Test that auto_mask=True correctly masks constraints with NaN RHS.""" + auto_mask_constraint_model.solve( + solver, io_api=io_api, explicit_coordinate_names=explicit_coordinate_names + ) + # Same assertions as test_masked_constraint_model + assert (auto_mask_constraint_model.solution.y[:-2] == 10).all() + assert (auto_mask_constraint_model.solution.y[-2:] == 5).all() + + # def init_model_large(): # m = Model() # time = pd.Index(range(10), name="time") diff --git a/test/test_persistent_apply_update.py b/test/test_persistent_apply_update.py new file mode 100644 index 000000000..41663dab1 --- /dev/null +++ b/test/test_persistent_apply_update.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +from linopy import Model +from linopy.persistent import RebuildReason +from linopy.solvers import Gurobi, Highs, Mosek, Solver, Xpress + +_BACKENDS: dict[str, tuple[type[Solver], dict[str, Any]]] = { + "gurobi": (Gurobi, {"OutputFlag": 0}), + "highs": (Highs, {"output_flag": False}), + "xpress": (Xpress, {"OUTPUTLOG": 0}), + "mosek": (Mosek, {"MSK_IPAR_LOG": 0}), +} + +_SIGN_CHANGE_IN_PLACE: dict[str, bool] = { + "gurobi": True, + "highs": False, + "xpress": True, + "mosek": False, +} + + +def _have(name: str) -> bool: + cls = _BACKENDS[name][0] + if not cls.is_available(): + return False + try: + cls._license_probe() + except Exception: + return False + if name == "xpress": + try: + import xpress + + xpress.problem() + except Exception: + return False + return True + + +SOLVER_PARAMS = [ + pytest.param( + name, + marks=pytest.mark.skipif(not _have(name), reason=f"{name} not installed"), + ) + for name in _BACKENDS +] + + +def _base_model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 10, coords=[range(3)], name="y") + m.add_constraints(x + y >= 4, name="c1") + m.add_constraints(2 * x + y <= 20, name="c2") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + +def _built(solver_name: str, model: Model) -> Solver: + cls, opts = _BACKENDS[solver_name] + s = cls(model=model, io_api="direct", track_updates=True) + s.options = opts + s._build() + return s + + +def _solve(solver: Solver, model: Model) -> float: + result = solver.solve(model, assign=True) + assert result.solution is not None + return float(result.solution.objective) + + +def _obj(model: Model) -> float: + value = model.objective.value + assert value is not None + return float(value) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_var_lb_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + base_obj = _obj(m) + + m.variables["x"].lower.values[...] = 5.0 + obj = _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert s._last_rebuild_reason is None + assert obj > base_obj + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_var_ub_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + m.variables["x"].upper.values[...] = 1.0 + _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_rhs_only_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + base_obj = _obj(m) + + c = m.constraints["c1"] + c.rhs = 8.0 + assert c._coef_dirty is False + obj = _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert obj > base_obj + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_constraint_coef_change_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + base_obj = _obj(m) + + c = m.constraints["c1"] + c.coeffs = c.coeffs * 2 + obj = _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert not np.isclose(obj, base_obj) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_objective_linear_change_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + base_obj = _obj(m) + + x = m.variables["x"] + y = m.variables["y"] + m.objective.expression = 5 * x.sum() + 3 * y.sum() + obj = _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert not np.isclose(obj, base_obj) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_objective_sense_flip_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + min_obj = _obj(m) + + m.objective.sense = "max" + max_obj = _solve(s, m) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert max_obj > min_obj + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_sparsity_change_triggers_rebuild(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + x = m.variables["x"] + m.add_constraints(x <= 5, name="c3") + s.solve(m, assign=True) + assert s._rebuilds == 1 + assert s._last_rebuild_reason in { + RebuildReason.STRUCTURAL_LABELS, + RebuildReason.STRUCTURAL_CONTAINERS, + } + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_cross_model_in_place(solver_name: str) -> None: + m1 = _base_model() + s = _built(solver_name, m1) + s.solve(assign=True) + + m2 = _base_model() + m2.constraints["c1"].rhs = 8.0 + + s.solve(m2, assign=True) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + + cross_obj = _obj(m2) + m3 = _base_model() + m3.constraints["c1"].rhs = 8.0 + s_fresh = _built(solver_name, m3) + s_fresh.solve(assign=True) + assert np.isclose(cross_obj, _obj(m3)) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_sign_flip(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + m.constraints["c1"].sign = "<=" + s.solve(m, assign=True) + if _SIGN_CHANGE_IN_PLACE[solver_name]: + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + else: + assert s._rebuilds == 1 + assert s._last_rebuild_reason is RebuildReason.BACKEND_REJECTED diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py new file mode 100644 index 000000000..bb801ecf2 --- /dev/null +++ b/test/test_persistent_snapshot_buffers.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from linopy import Model +from linopy.persistent import ModelDiff, ModelSnapshot, RebuildReason +from linopy.persistent.snapshot import _extract_con_buffers + + +def _build_permuted_pair() -> tuple[Model, Model]: + m1 = Model() + x1 = m1.add_variables(0, 10, coords=[range(3)], name="x") + y1 = m1.add_variables(0, 5, coords=[range(2)], name="y") + m1.add_constraints(2 * x1 + 3 * y1.sum() >= 4, name="c1") + m1.add_objective(x1.sum()) + + m2 = Model() + x2 = m2.add_variables(0, 10, coords=[range(3)], name="x") + y2 = m2.add_variables(0, 5, coords=[range(2)], name="y") + m2.add_constraints(3 * y2.sum() + 2 * x2 >= 4, name="c1") + m2.add_objective(x2.sum()) + return m1, m2 + + +def test_permuted_term_order_produces_equal_buffers() -> None: + m1, m2 = _build_permuted_pair() + s1 = ModelSnapshot.capture(m1) + s2 = ModelSnapshot.capture(m2) + b1 = s1.con_buffers["c1"] + b2 = s2.con_buffers["c1"] + np.testing.assert_array_equal(b1.indptr, b2.indptr) + np.testing.assert_array_equal(b1.indices, b2.indices) + np.testing.assert_array_equal(b1.data, b2.data) + + +def test_active_labels_match_label_index(baseline_model: Model) -> None: + snap = ModelSnapshot.capture(baseline_model) + expected = baseline_model.constraints.label_index.clabels + concatenated = np.concatenate( + [buf.active_labels for buf in snap.con_buffers.values()] + ) + np.testing.assert_array_equal(concatenated, expected) + + +@pytest.fixture +def baseline_model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 5, coords=[range(2)], name="y") + m.add_constraints(2 * x >= 4, name="c1") + m.add_constraints(x.sum() + y.sum() <= 20, name="c2") + m.add_objective(x.sum()) + return m + + +def test_shape_mismatch_triggers_sparsity_rebuild(baseline_model: Model) -> None: + snap = ModelSnapshot.capture(baseline_model) + x = baseline_model.variables["x"] + y = baseline_model.variables["y"] + baseline_model.constraints["c1"].lhs = 2 * x + 0 * y.sum() + diff = ModelDiff.from_snapshot(snap, baseline_model) + assert diff in { + RebuildReason.SPARSITY, + RebuildReason.STRUCTURAL_LABELS, + } + + +def test_zero_row_container_capture() -> None: + m = Model() + m.add_variables(0, 10, coords=[range(2)], name="x") + m.add_objective(0.0 * m.variables["x"].sum()) + snap = ModelSnapshot.capture(m) + assert snap.con_buffers == {} + diff = ModelDiff.from_snapshot(snap, m) + assert isinstance(diff, ModelDiff) + assert diff.is_empty + + +def test_con_buffers_dtypes(baseline_model: Model) -> None: + snap = ModelSnapshot.capture(baseline_model) + buf = snap.con_buffers["c1"] + assert buf.rhs.dtype == np.float64 + assert buf.sign.dtype == np.dtype("U1") + assert buf.data.dtype == np.float64 + assert np.issubdtype(buf.indices.dtype, np.integer) + assert np.issubdtype(buf.indptr.dtype, np.integer) + + +def test_masked_rows_excluded_from_active_labels() -> None: + m = Model() + x = m.add_variables(0, 10, coords=[range(4)], name="x") + mask = np.array([True, False, True, True]) + m.add_constraints(2 * x >= 1, mask=mask, name="c1") + m.add_objective(x.sum()) + snap = ModelSnapshot.capture(m) + buf = snap.con_buffers["c1"] + assert buf.active_labels.size == 3 + rebuilt = _extract_con_buffers(m.constraints["c1"], m.variables.label_index) + np.testing.assert_array_equal(rebuilt.active_labels, buf.active_labels) + + +def test_csr_capture_deterministic(baseline_model: Model) -> None: + s1 = ModelSnapshot.capture(baseline_model) + s2 = ModelSnapshot.capture(baseline_model) + for name in s1.con_buffers: + b1, b2 = s1.con_buffers[name], s2.con_buffers[name] + np.testing.assert_array_equal(b1.indptr, b2.indptr) + np.testing.assert_array_equal(b1.indices, b2.indices) + np.testing.assert_array_equal(b1.data, b2.data) + + +def test_duplicate_variable_terms_summed() -> None: + m1 = Model() + x1 = m1.add_variables(0, 10, coords=[range(3)], name="x") + m1.add_constraints(2 * x1 + 3 * x1 >= 1, name="c1") + m1.add_objective(x1.sum()) + + m2 = Model() + x2 = m2.add_variables(0, 10, coords=[range(3)], name="x") + m2.add_constraints(5 * x2 >= 1, name="c1") + m2.add_objective(x2.sum()) + + diff = ModelDiff.from_models(m1, m2) + assert isinstance(diff, ModelDiff) + assert diff.is_empty diff --git a/test/test_persistent_snapshot_diff.py b/test/test_persistent_snapshot_diff.py new file mode 100644 index 000000000..8730bf340 --- /dev/null +++ b/test/test_persistent_snapshot_diff.py @@ -0,0 +1,400 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from linopy import Model +from linopy.persistent import ( + ContainerConBuffers, + ContainerVarBuffers, + ModelDiff, + ModelSnapshot, + RebuildReason, + StructuralKey, + VarKind, +) + + +@pytest.fixture +def baseline() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 5, coords=[range(2)], name="y") + m.add_constraints(2 * x + 1 >= 4, name="c1") + m.add_constraints(x.sum() + y.sum() <= 20, name="c2") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + +def test_capture_structural_key(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + assert isinstance(snap, ModelSnapshot) + assert isinstance(snap.structural_key, StructuralKey) + assert snap.structural_key.var_container_names == ("x", "y") + assert snap.structural_key.con_container_names == ("c1", "c2") + np.testing.assert_array_equal( + snap.structural_key.vlabels, baseline.variables.label_index.vlabels + ) + np.testing.assert_array_equal( + snap.structural_key.clabels, baseline.constraints.label_index.clabels + ) + assert isinstance(snap.var_buffers["x"], ContainerVarBuffers) + assert isinstance(snap.con_buffers["c1"], ContainerConBuffers) + + +def test_is_empty_on_unmutated(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert diff.is_empty + + +def test_bounds_only_mutation(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.variables["x"].lower = 1 + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert "x" in diff.changed_variables + assert "y" not in diff.changed_variables + sl = diff.var_slices["x"].bounds + np.testing.assert_array_equal(diff.var_bounds_lower[sl], np.ones(3)) + + +def test_rhs_only_mutation(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.constraints["c1"].rhs = 9 + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert "c1" in diff.changed_constraints + sl = diff.con_slices["c1"] + assert sl.rhs.stop > sl.rhs.start + assert sl.coef.stop == sl.coef.start + + +def test_objective_linear_change(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + x = baseline.variables["x"] + y = baseline.variables["y"] + baseline.add_objective(3 * x.sum() + 2 * y.sum(), overwrite=True) + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert diff.obj_c_indices is not None + assert diff.obj_c_values is not None + + +def test_objective_sense_flip(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.objective.sense = "max" + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert diff.obj_sense == "max" + + +def test_add_constraints_is_structural(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + x = baseline.variables["x"] + baseline.add_constraints(x.sum() <= 99, name="c3") + diff = ModelDiff.from_snapshot(snap, baseline) + assert diff in ( + RebuildReason.STRUCTURAL_LABELS, + RebuildReason.STRUCTURAL_CONTAINERS, + ) + + +def test_remove_variables_is_structural(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.remove_variables("y") + diff = ModelDiff.from_snapshot(snap, baseline) + assert diff in ( + RebuildReason.STRUCTURAL_LABELS, + RebuildReason.STRUCTURAL_CONTAINERS, + ) + + +def test_coef_value_change_same_sparsity(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + c = baseline.constraints["c1"] + c.coeffs = c.coeffs * 3 + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert "c1" in diff.changed_constraints + sl = diff.con_slices["c1"].coef + vals = diff.con_coef_vals[sl] + np.testing.assert_array_equal(vals, np.full(vals.size, 6.0)) + + +def test_coef_changes_across_containers(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + c1 = baseline.constraints["c1"] + c2 = baseline.constraints["c2"] + c1.update(coeffs=c1.coeffs * 3) + c2.update(coeffs=c2.coeffs * 2) + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + sl1 = diff.con_slices["c1"].coef + sl2 = diff.con_slices["c2"].coef + assert diff.n_coef_updates == (sl1.stop - sl1.start) + (sl2.stop - sl2.start) + np.testing.assert_array_equal( + diff.con_coef_vals[sl1], np.full(sl1.stop - sl1.start, 6.0) + ) + np.testing.assert_array_equal( + diff.con_coef_vals[sl2], np.full(sl2.stop - sl2.start, 2.0) + ) + + +def test_coef_sparsity_change(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + x = baseline.variables["x"] + baseline.constraints["c2"].lhs = 2 * x.sum() + diff = ModelDiff.from_snapshot(snap, baseline) + assert diff is RebuildReason.SPARSITY + + +def test_deep_copy_invariant(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.variables["x"].lower.values[...] = 99 + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + assert "x" in diff.changed_variables + + +def test_same_model_false_ignores_dirty_flag(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + c = baseline.constraints["c1"] + c.coeffs = c.coeffs * 5 + c._coef_dirty = False + diff_fast = ModelDiff.from_snapshot(snap, baseline, same_model=True) + assert isinstance(diff_fast, ModelDiff) + fast_coef = diff_fast.con_slices.get("c1") + assert fast_coef is None or fast_coef.coef.stop == fast_coef.coef.start + diff_full = ModelDiff.from_snapshot(snap, baseline, same_model=False) + assert isinstance(diff_full, ModelDiff) + full_coef = diff_full.con_slices["c1"].coef + assert full_coef.stop > full_coef.start + + +def test_from_models_diffs_two_models() -> None: + m1 = Model() + x1 = m1.add_variables(0, 10, coords=[range(3)], name="x") + m1.add_constraints(2 * x1 >= 4, name="c1") + m1.add_objective(x1.sum()) + + m2 = Model() + x2 = m2.add_variables(0, 10, coords=[range(3)], name="x") + m2.add_constraints(2 * x2 >= 7, name="c1") + m2.add_objective(x2.sum()) + + diff = ModelDiff.from_models(m1, m2) + assert isinstance(diff, ModelDiff) + assert "c1" in diff.changed_constraints + sl = diff.con_slices["c1"].rhs + np.testing.assert_array_equal(diff.con_rhs_values[sl], np.full(3, 7.0)) + + +def test_ignore_dims_detects_coord_change() -> None: + m1 = Model() + m1.add_variables(0, 10, coords=[pd.Index([0, 1, 2], name="t")], name="x") + m1.add_constraints(m1.variables["x"] >= 0, name="c1") + m1.add_objective(m1.variables["x"].sum()) + snap = ModelSnapshot.capture(m1) + + m2 = Model() + m2.add_variables(0, 10, coords=[pd.Index([10, 11, 12], name="t")], name="x") + m2.add_constraints(m2.variables["x"] >= 0, name="c1") + m2.add_objective(m2.variables["x"].sum()) + + assert ModelDiff.from_snapshot(snap, m2) is RebuildReason.COORD_REINDEX + assert isinstance(ModelDiff.from_snapshot(snap, m2, ignore_dims={"t"}), ModelDiff) + + +def _assert_snapshot_equal(a: ModelSnapshot, b: ModelSnapshot) -> None: + assert a.structural_key == b.structural_key + assert a.var_buffers.keys() == b.var_buffers.keys() + assert a.con_buffers.keys() == b.con_buffers.keys() + for name, va in a.var_buffers.items(): + vb = b.var_buffers[name] + np.testing.assert_array_equal(va.lower, vb.lower) + np.testing.assert_array_equal(va.upper, vb.upper) + np.testing.assert_array_equal(va.active_labels, vb.active_labels) + assert va.type is vb.type + for name, ca in a.con_buffers.items(): + cb = b.con_buffers[name] + for attr in ("indptr", "indices", "data", "rhs", "sign", "active_labels"): + np.testing.assert_array_equal(getattr(ca, attr), getattr(cb, attr)) + for coords_a, coords_b in ( + (a.var_coords, b.var_coords), + (a.con_coords, b.con_coords), + ): + assert coords_a.keys() == coords_b.keys() + for name in coords_a: + assert coords_a[name].keys() == coords_b[name].keys() + for dim in coords_a[name]: + np.testing.assert_array_equal(coords_a[name][dim], coords_b[name][dim]) + np.testing.assert_array_equal(a.obj_c, b.obj_c) + assert a.obj_quad_present == b.obj_quad_present + assert a.obj_sense == b.obj_sense + + +def test_capture_is_pure(baseline: Model) -> None: + c = baseline.constraints["c1"] + c.update(coeffs=c.coeffs * 2) + assert c._coef_dirty is True + ModelSnapshot.capture(baseline) + assert c._coef_dirty is True + + +@pytest.mark.parametrize( + "mutate", ["none", "rhs", "bounds", "coeffs", "objective", "combined"] +) +def test_diff_snapshot_matches_capture(baseline: Model, mutate: str) -> None: + snap = ModelSnapshot.capture(baseline) + x = baseline.variables["x"] + y = baseline.variables["y"] + if mutate in ("rhs", "combined"): + baseline.constraints["c1"].update(rhs=9) + if mutate in ("bounds", "combined"): + x.update(lower=1) + if mutate in ("coeffs", "combined"): + c2 = baseline.constraints["c2"] + c2.update(coeffs=c2.coeffs * 3) + if mutate in ("objective", "combined"): + baseline.add_objective(3 * x.sum() + 2 * y.sum(), overwrite=True) + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(baseline)) + + +def test_diff_snapshot_matches_capture_under_ignore_dims() -> None: + def build(t0: int) -> Model: + m = Model() + t = pd.Index(range(t0, t0 + 3), name="t") + m.add_variables(0, 10, coords=[t], name="x") + m.add_constraints(m.variables["x"] >= 0, name="c1") + m.add_objective(m.variables["x"].sum()) + return m + + m1, m2 = build(0), build(10) + snap = ModelSnapshot.capture(m1) + diff = ModelDiff.from_snapshot(snap, m2, ignore_dims={"t"}) + assert isinstance(diff, ModelDiff) + _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(m2)) + + +def test_from_models_snapshot_matches_capture() -> None: + def build(rhs: float) -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + m.add_constraints(2 * x >= rhs, name="c1") + m.add_objective(x.sum()) + return m + + m1, m2 = build(4.0), build(7.0) + diff = ModelDiff.from_models(m1, m2) + assert isinstance(diff, ModelDiff) + _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(m2)) + + +@pytest.mark.parametrize( + "kwargs, expected", + [ + ({"binary": True}, VarKind.BINARY), + ({"lower": 0, "upper": 10, "integer": True}, VarKind.INTEGER), + ({"lower": 1, "upper": 10, "semi_continuous": True}, VarKind.SEMI_CONTINUOUS), + ], +) +def test_variable_kind_captured(kwargs: dict, expected: VarKind) -> None: + m = Model() + m.add_variables(coords=[range(2)], name="x", **kwargs) + m.add_objective(m.variables["x"].sum()) + snap = ModelSnapshot.capture(m) + assert snap.var_buffers["x"].type is expected + + +def test_variable_type_change_via_from_models() -> None: + def build(integer: bool) -> Model: + m = Model() + m.add_variables(0, 10, coords=[range(3)], name="x", integer=integer) + m.add_constraints(m.variables["x"] >= 1, name="c1") + m.add_objective(m.variables["x"].sum()) + return m + + diff = ModelDiff.from_models(build(False), build(True)) + assert isinstance(diff, ModelDiff) + sl = diff.var_slices["x"].type + assert sl.stop > sl.start + assert diff.var_type_kinds[sl][0] is VarKind.INTEGER + + +def test_quadratic_objective_triggers_rebuild() -> None: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + m.add_constraints(x >= 1, name="c1") + m.add_objective((x * x).sum()) + snap = ModelSnapshot.capture(m) + x.update(lower=2) + assert ModelDiff.from_snapshot(snap, m) is RebuildReason.QUAD_OBJ + + +def test_variable_count_change_is_structural() -> None: + def build(n: int) -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(n)], name="x") + m.add_constraints(x >= 1, name="c1") + m.add_objective(x.sum()) + return m + + assert ModelDiff.from_models(build(3), build(4)) is RebuildReason.STRUCTURAL_LABELS + + +def test_constraint_count_change_is_structural() -> None: + def build(aggregate: bool) -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + m.add_constraints(x.sum() >= 1 if aggregate else x >= 1, name="c1") + m.add_objective(x.sum()) + return m + + diff = ModelDiff.from_models(build(False), build(True)) + assert diff is RebuildReason.STRUCTURAL_LABELS + + +def test_indices_change_triggers_sparsity() -> None: + def build(on: int) -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(2)], name="x") + m.add_constraints(x.loc[on] >= 1, name="c1") + m.add_objective(x.sum()) + return m + + assert ModelDiff.from_models(build(0), build(1)) is RebuildReason.SPARSITY + + +def test_sign_only_mutation(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + baseline.constraints["c1"].update(sign="<=") + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + sl = diff.con_slices["c1"] + assert sl.sign.stop > sl.sign.start + assert sl.coef.stop == sl.coef.start + + +def test_inspect_and_repr(baseline: Model) -> None: + snap = ModelSnapshot.capture(baseline) + assert repr(ModelDiff.from_snapshot(snap, baseline)) == "ModelDiff(empty)" + + baseline.variables["x"].update(lower=1) + c1 = baseline.constraints["c1"] + c1.update(coeffs=c1.coeffs * 2, rhs=9, sign="<=") + diff = ModelDiff.from_snapshot(snap, baseline) + assert isinstance(diff, ModelDiff) + + var_info = diff.inspect_variable("x") + assert "lower" in var_info and "bounds_indices" in var_info + con_info = diff.inspect_constraint("c1") + assert {"coef_vals", "rhs_values", "sign_values"} <= con_info.keys() + + assert diff.inspect_variable("missing") == {} + assert diff.inspect_constraint("missing") == {} + assert repr(diff).startswith("ModelDiff(") and "empty" not in repr(diff) diff --git a/test/test_persistent_solver_extras.py b/test/test_persistent_solver_extras.py new file mode 100644 index 000000000..3fc28d18d --- /dev/null +++ b/test/test_persistent_solver_extras.py @@ -0,0 +1,466 @@ +from __future__ import annotations + +import pickle +import threading +from typing import Any + +import numpy as np +import pytest + +from linopy import Model +from linopy.persistent import ModelDiff, RebuildReason, UpdatesDisabledError +from linopy.solvers import Gurobi, Highs, Solver + +_BACKENDS: dict[str, tuple[type[Solver], dict[str, Any]]] = { + "gurobi": (Gurobi, {"OutputFlag": 0}), + "highs": (Highs, {"output_flag": False}), +} + + +def _have(name: str) -> bool: + try: + if name == "gurobi": + import gurobipy # noqa: F401 + elif name == "highs": + import highspy # noqa: F401 + return True + except ImportError: + return False + + +SOLVER_PARAMS = [ + pytest.param( + "gurobi", + marks=pytest.mark.skipif(not _have("gurobi"), reason="gurobipy not installed"), + ), + pytest.param( + "highs", + marks=pytest.mark.skipif(not _have("highs"), reason="highspy not installed"), + ), +] + + +def _base_model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 10, coords=[range(3)], name="y") + m.add_constraints(x + y >= 4, name="c1") + m.add_constraints(2 * x + y <= 20, name="c2") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + +def _built(solver_name: str, model: Model) -> Solver: + cls, opts = _BACKENDS[solver_name] + s = cls(model=model, io_api="direct", track_updates=True) + s.options = opts + s._build() + return s + + +def _obj(model: Model) -> float: + value = model.objective.value + assert value is not None + return float(value) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_noop_resolve_increments_in_place(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + first_obj = _obj(m) + + s.solve(m, assign=True) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert np.isclose(_obj(m), first_obj) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_two_consecutive_solves_no_stale_state(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + first_status = s.status + + m.variables["x"].lower.values[...] = 5.0 + s.solve(m, assign=True) + assert s.status is not first_status + assert s.solution is not None + assert np.isclose(float(s.solution.objective), _obj(m)) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_cross_model_scenario_sweep(solver_name: str) -> None: + m1 = _base_model() + m2 = _base_model() + m2.constraints["c1"].rhs = 6.0 + m3 = _base_model() + m3.variables["x"].lower.values[...] = 2.0 + + s = _built(solver_name, m1) + s.solve(assign=True) + obj1 = _obj(m1) + sol1 = m1.solution + + s.solve(m2, assign=True) + s.solve(m3, assign=True) + + assert s._rebuilds == 0 + assert s._in_place_updates >= 2 + + assert m1.objective._value == obj1 + np.testing.assert_array_equal(m1.solution.x.values, sol1.x.values) + assert m2.objective._value is not None + assert m3.objective._value is not None + + for mk in (m2, m3): + fresh = _base_model() + if mk is m2: + fresh.constraints["c1"].rhs = 6.0 + else: + fresh.variables["x"].lower.values[...] = 2.0 + s_fresh = _built(solver_name, fresh) + s_fresh.solve(assign=True) + assert np.isclose(_obj(mk), _obj(fresh)) + s_fresh.close() + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_cross_model_sparsity_change_rebuilds(solver_name: str) -> None: + def build(include_y_in_c1: bool) -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + y = m.add_variables(0, 10, coords=[range(3)], name="y") + if include_y_in_c1: + m.add_constraints(x + y >= 4, name="c1") + else: + m.add_constraints(2 * x >= 4, name="c1") + m.add_constraints(2 * x + y <= 20, name="c2") + m.add_objective(x.sum() + 2 * y.sum()) + return m + + m1 = build(include_y_in_c1=True) + s = _built(solver_name, m1) + s.solve(assign=True) + + m2 = build(include_y_in_c1=False) + + s.solve(m2, assign=True) + assert s._rebuilds == 1 + assert s._last_rebuild_reason in { + RebuildReason.SPARSITY, + RebuildReason.STRUCTURAL_LABELS, + RebuildReason.STRUCTURAL_CONTAINERS, + } + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_cross_model_structural_mismatch_rebuilds(solver_name: str) -> None: + m1 = _base_model() + s = _built(solver_name, m1) + s.solve(assign=True) + + m2 = _base_model() + m2.add_variables(0, 5, coords=[range(3)], name="z") + + s.solve(m2, assign=True) + assert s._rebuilds == 1 + assert s.model is m2 + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_dirty_flag_ignored_across_models(solver_name: str) -> None: + m1 = _base_model() + s = _built(solver_name, m1) + s.solve(assign=True) + + m2 = _base_model() + c = m2.constraints["c1"] + c.coeffs = c.coeffs * 3 + c._coef_dirty = False + + s.solve(m2, assign=True) + assert s._rebuilds == 0 + assert s._in_place_updates == 1 + + fresh = _base_model() + cf = fresh.constraints["c1"] + cf.coeffs = cf.coeffs * 3 + s_fresh = _built(solver_name, fresh) + s_fresh.solve(assign=True) + assert np.isclose(_obj(m2), _obj(fresh)) + s_fresh.close() + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_solver_pickle_round_trip_drops_native(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + state = s.__getstate__() + for key in ("solver_model", "env", "_env_stack", "snapshot", "_lock"): + assert key not in state + + restored = pickle.loads(pickle.dumps(s)) + assert restored.solver_model is None + assert restored.snapshot is None + assert restored._env_stack is None + assert isinstance(restored._lock, type(threading.Lock())) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_model_pickle_round_trip_no_native_handle(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + m2 = pickle.loads(pickle.dumps(m)) + s2 = _built(solver_name, m2) + assert s2.solver_model is not None + s2.solve(assign=True) + assert s2._rebuilds == 0 + assert np.isclose(_obj(m), _obj(m2)) + s2.close() + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_backend_exception_during_apply_rebuilds( + solver_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + c = m.constraints["c1"] + c.coeffs = c.coeffs * 2 + assert c._coef_dirty is True + + def _boom(*args: Any, **kwargs: Any) -> None: + raise RuntimeError("simulated backend failure") + + monkeypatch.setattr(s, "apply_update", _boom) + + dirty_at_rebuild: list[bool] = [] + original_build = s._build + + def _spy_build(**kwargs: Any) -> None: + dirty_at_rebuild.append(m.constraints["c1"]._coef_dirty) + original_build(**kwargs) + + monkeypatch.setattr(s, "_build", _spy_build) + + s.solve(m, assign=True) + assert s._rebuilds == 1 + assert s._last_rebuild_reason is RebuildReason.BACKEND_REJECTED + assert dirty_at_rebuild == [True] + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_concurrent_solves_serialize(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + expected = _obj(m) + + barrier = threading.Barrier(2) + results: list[float] = [] + errors: list[BaseException] = [] + + def _run() -> None: + try: + barrier.wait() + res = s.solve(m, assign=True) + assert res.solution is not None + results.append(float(res.solution.objective)) + except BaseException as e: + errors.append(e) + + threads = [threading.Thread(target=_run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, errors + assert len(results) == 2 + for r in results: + assert np.isclose(r, expected) + + +_SCENARIO_PARAMS = [ + "bound_only", + "rhs_only", + "single_cell_coef", + "multi_row_coef", + "mixed", +] + + +def _apply_scenario(model: Model, scenario: str) -> None: + if scenario == "bound_only": + model.variables["x"].lower.values[...] = 3.0 + elif scenario == "rhs_only": + model.constraints["c1"].rhs = 7.0 + elif scenario == "single_cell_coef": + c = model.constraints["c1"] + new = c.coeffs.copy() + new.values[0, 0] = 5.0 + c.coeffs = new + elif scenario == "multi_row_coef": + c = model.constraints["c2"] + c.coeffs = c.coeffs * 2 + elif scenario == "mixed": + model.variables["x"].lower.values[...] = 1.0 + model.constraints["c1"].rhs = 6.0 + c = model.constraints["c2"] + new = c.coeffs.copy() + new.values[0, 0] = 4.0 + c.coeffs = new + else: + raise ValueError(scenario) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +@pytest.mark.parametrize("scenario", _SCENARIO_PARAMS) +@pytest.mark.parametrize("same_model", [True, False]) +def test_scenario_sweep_in_place( + solver_name: str, scenario: str, same_model: bool +) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + target = m if same_model else _base_model() + _apply_scenario(target, scenario) + s.solve(target, assign=True) + + assert s._rebuilds == 0 + assert s._in_place_updates == 1 + assert s._last_rebuild_reason is None + + fresh = _base_model() + _apply_scenario(fresh, scenario) + s_fresh = _built(solver_name, fresh) + s_fresh.solve(assign=True) + assert np.isclose(_obj(target), _obj(fresh)) + s_fresh.close() + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_disallow_rebuild_raises_on_structural_change(solver_name: str) -> None: + from linopy.persistent import RebuildRequiredError + + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + m2 = _base_model() + m2.add_variables(0, 5, coords=[range(3)], name="z") + + with pytest.raises(RebuildRequiredError): + s.solve(m2, disallow_rebuild=True, assign=True) + assert s._rebuilds == 0 + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_disallow_rebuild_passes_when_update_works(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + s.solve(assign=True) + + m.constraints["c1"].rhs = 6.0 + s.solve(m, disallow_rebuild=True, assign=True) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_solve_without_assign_does_not_mutate_model(solver_name: str) -> None: + m = _base_model() + s = _built(solver_name, m) + + assert m.objective._value is None + s.solve() + assert m.objective._value is None + + s.solve(assign=True) + assert m.objective._value is not None + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_track_updates_false_skips_snapshot(solver_name: str) -> None: + cls, opts = _BACKENDS[solver_name] + m = _base_model() + s = cls(model=m, io_api="direct", track_updates=False) + s.options = opts + s._build() + assert s.snapshot is None + s.solve(assign=True) + assert s.snapshot is None + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_track_updates_false_rejects_resolve_with_model(solver_name: str) -> None: + cls, opts = _BACKENDS[solver_name] + m = _base_model() + s = cls(model=m, io_api="direct", track_updates=False) + s.options = opts + s._build() + s.solve(assign=True) + + m.variables["x"].lower.values[...] = 6.0 + with pytest.raises(UpdatesDisabledError, match="track_updates=False"): + s.solve(m, assign=True) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_track_updates_false_rejects_update(solver_name: str) -> None: + cls, opts = _BACKENDS[solver_name] + m = _base_model() + s = cls(model=m, io_api="direct", track_updates=False) + s.options = opts + s._build() + with pytest.raises(UpdatesDisabledError, match="track_updates=False"): + s.update(m) + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_track_updates_false_cross_instance_resolve(solver_name: str) -> None: + cls, opts = _BACKENDS[solver_name] + m1 = _base_model() + s = cls(model=m1, io_api="direct", track_updates=False) + s.options = opts + s._build() + s.solve(assign=True) + base_obj = _obj(m1) + + m2 = _base_model() + m2.constraints["c1"].rhs = 8.0 + result = s.solve(m2, assign=True) + assert s._in_place_updates == 1 + assert s._rebuilds == 0 + assert s.snapshot is None + assert s.model is m2 + assert result.solution is not None + assert float(result.solution.objective) > base_obj + + +@pytest.mark.parametrize("solver_name", SOLVER_PARAMS) +def test_track_updates_false_cross_instance_update(solver_name: str) -> None: + cls, opts = _BACKENDS[solver_name] + m1 = _base_model() + s = cls(model=m1, io_api="direct", track_updates=False) + s.options = opts + s._build() + s.solve(assign=True) + + m2 = _base_model() + m2.constraints["c1"].rhs = 8.0 + diff = s.update(m2, apply=False) + assert isinstance(diff, ModelDiff) + assert diff.summary()["con_rhs"] == 3 + assert "c1" in diff.changed_constraints + assert s.snapshot is None diff --git a/test/test_persistent_solver_orchestrator.py b/test/test_persistent_solver_orchestrator.py new file mode 100644 index 000000000..9495e9ea0 --- /dev/null +++ b/test/test_persistent_solver_orchestrator.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import pickle +import threading +from typing import Any + +import pytest + +from linopy import Model +from linopy.constants import ( + Result, + Solution, + SolverStatus, + Status, + TerminationCondition, +) +from linopy.persistent import ModelDiff, RebuildReason +from linopy.solvers import Solver, SolverFeature + + +class FakeSolver(Solver[None]): + display_name = "Fake" + features = frozenset({SolverFeature.DIRECT_API}) + accepted_io_apis = frozenset({"direct"}) + supports_persistent_update = False + + @classmethod + def is_available(cls) -> bool: # type: ignore[override] + return True + + @property + def solver_name(self) -> Any: + class _N: + value = "fake" + + return _N() + + def _validate_model(self) -> None: + return None + + def _build_direct(self, **kwargs: Any) -> None: + self.solver_model = object() + + def _run_direct(self, **kwargs: Any) -> Result: + status = Status(SolverStatus.ok, TerminationCondition.optimal) + return Result( + status=status, solution=Solution(objective=0.0), solver_name="fake" + ) + + +@pytest.fixture +def model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + m.add_constraints(2 * x >= 4, name="c1") + m.add_objective(x.sum()) + return m + + +@pytest.fixture +def other_model() -> Model: + m = Model() + x = m.add_variables(0, 10, coords=[range(3)], name="x") + m.add_constraints(2 * x >= 4, name="c1") + m.add_objective(x.sum()) + return m + + +def _built(model: Model) -> FakeSolver: + s = FakeSolver(model=model, io_api="direct", track_updates=True) + s._build() + return s + + +def test_unsupported_falls_through_to_rebuild(model: Model, other_model: Model) -> None: + s = _built(model) + assert s._rebuilds == 0 + s.solve(other_model) + assert s._rebuilds == 1 + assert s._last_rebuild_reason is RebuildReason.BACKEND_REJECTED + assert s.model is other_model + + +def test_update_apply_false_returns_diff(model: Model) -> None: + s = _built(model) + diff = s.update(model, apply=False) + assert isinstance(diff, ModelDiff) + assert s._in_place_updates == 0 + assert s._rebuilds == 0 + + +def test_solve_no_model_still_works(model: Model) -> None: + s = _built(model) + result = s.solve() + assert result.status.status is SolverStatus.ok + + +def test_getstate_drops_native_fields(model: Model) -> None: + s = _built(model) + state = s.__getstate__() + for k in ("solver_model", "env", "_env_stack", "snapshot", "_lock"): + assert k not in state + restored = pickle.loads(pickle.dumps(s)) + assert restored.solver_model is None + assert restored.snapshot is None + + +def test_update_without_snapshot_raises(model: Model) -> None: + s = FakeSolver(model=model, io_api="direct") + with pytest.raises(RuntimeError, match="not been built"): + s.update(model) + + +def test_unmutated_resolve_diff_is_empty(model: Model) -> None: + s = _built(model) + diff = s.update(model, apply=False) + assert isinstance(diff, ModelDiff) + assert diff.is_empty + + +class FakePersistentSolver(FakeSolver): + supports_persistent_update = True + + def apply_update( + self, diff: ModelDiff, var_label_index: Any, con_label_index: Any + ) -> None: + return None + + +def _built_persistent(model: Model) -> FakePersistentSolver: + s = FakePersistentSolver(model=model, io_api="direct", track_updates=True) + s._build() + return s + + +def test_build_clears_coef_dirty(model: Model) -> None: + c = model.constraints["c1"] + c.update(coeffs=c.coeffs * 2) + assert c._coef_dirty is True + _built_persistent(model) + assert c._coef_dirty is False + + +def test_in_place_update_adopts_diff_snapshot(model: Model) -> None: + s = _built_persistent(model) + c = model.constraints["c1"] + c.update(coeffs=c.coeffs * 2) + diff = s.update(model) + assert isinstance(diff, ModelDiff) + assert s.snapshot is diff.snapshot + assert c._coef_dirty is False + rediff = s.update(model, apply=False) + assert isinstance(rediff, ModelDiff) + assert rediff.is_empty + + +def test_update_apply_false_leaves_state_untouched(model: Model) -> None: + s = _built_persistent(model) + snap_before = s.snapshot + c = model.constraints["c1"] + c.update(coeffs=c.coeffs * 2) + diff = s.update(model, apply=False) + assert isinstance(diff, ModelDiff) + assert c._coef_dirty is True + assert s.snapshot is snap_before + + +def test_update_apply_false_does_not_block_running_solve( + model: Model, monkeypatch: pytest.MonkeyPatch +) -> None: + s = _built_persistent(model) + solve_entered = threading.Event() + release_solve = threading.Event() + original_run = s._run_direct + + def _gated_run(**kwargs: Any) -> Result: + solve_entered.set() + assert release_solve.wait(timeout=5) + return original_run(**kwargs) + + monkeypatch.setattr(s, "_run_direct", _gated_run) + + solver_thread = threading.Thread(target=s.solve) + solver_thread.start() + try: + assert solve_entered.wait(timeout=5) + + result: list[ModelDiff | RebuildReason] = [] + preview_thread = threading.Thread( + target=lambda: result.append(s.update(model, apply=False)) + ) + preview_thread.start() + preview_thread.join(timeout=2) + assert not preview_thread.is_alive(), "preview blocked on a running solve" + assert isinstance(result[0], ModelDiff) + finally: + release_solve.set() + solver_thread.join(timeout=5) + + +def test_preview_detects_raw_mutation_apply_skips_it(model: Model) -> None: + """ + Pins the documented preview/apply asymmetry for unsupported raw + ``.values[...]`` coefficient mutations on the build-time model. + """ + s = _built_persistent(model) + c = model.constraints["c1"] + c.coeffs.values[...] = c.coeffs.values * 2 + assert c._coef_dirty is False + + preview = s.update(model, apply=False) + assert isinstance(preview, ModelDiff) + assert "c1" in preview.changed_constraints + + applied = s.update(model) + assert isinstance(applied, ModelDiff) + assert "c1" not in applied.changed_constraints diff --git a/test/test_piecewise_active_fill.py b/test/test_piecewise_active_fill.py new file mode 100644 index 000000000..0af04c8c3 --- /dev/null +++ b/test/test_piecewise_active_fill.py @@ -0,0 +1,222 @@ +""" +Tests for the ``active_fill`` parameter of ``add_piecewise_formulation`` (#796). + +``active_fill`` is a transitional convenience: it pads a partial ``active`` +gate (a subset of the indexed dimension, or a masked gate) to full coverage. +It is slated for removal once the v1 arithmetic semantics (#717) make +``active.reindex(coords).fillna(value)`` correct on its own, so these tests +live in a dedicated module that can be dropped with the parameter. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Literal, TypeAlias + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model, available_solvers, segments +from linopy.piecewise import _resolve_active +from linopy.solver_capabilities import ( + SolverFeature, + get_available_solvers_with_feature, +) + +Method: TypeAlias = Literal["sos2", "incremental", "lp", "auto"] +GateBuilder: TypeAlias = Callable[[Model], Any] + +_any_solvers = [ + s for s in ["highs", "gurobi", "glpk", "cplex"] if s in available_solvers +] +_sos2_solvers = get_available_solvers_with_feature( + SolverFeature.SOS_CONSTRAINTS, available_solvers +) + + +# ``active`` is meaningful only for the committable subset {a, c}; "b" stays +# ungated. The partial-gate shapes below all leave "b" as the gap. +_PWL_GENS = pd.Index(["a", "b", "c"], name="gen") +_COMMITTABLE = pd.Index(["a", "c"], name="gen") + + +def _subset_gate(m: Model) -> Any: + """``active`` indexed over a strict subset of the formulation's dim.""" + return m.add_variables(binary=True, coords=[_COMMITTABLE], name="u") + + +def _masked_gate(m: Model) -> Any: + """``active`` over the full dim but masked where it does not apply.""" + mask = pd.Series([True, False, True], index=_PWL_GENS) + return m.add_variables(binary=True, coords=[_PWL_GENS], name="u", mask=mask) + + +def _full_gate(m: Model) -> Any: + return m.add_variables(binary=True, coords=[_PWL_GENS], name="u") + + +def _scalar_gate(m: Model) -> Any: + return m.add_variables(binary=True, name="u") + + +_PARTIAL_GATES = [ + pytest.param(_subset_gate, id="strict-subset"), + pytest.param(_masked_gate, id="masked"), +] + +# (builder, active_fill, should_raise): partial gates raise unless active_fill +# is set; full/scalar gates are always fine. +_COVERAGE_CASES = [ + pytest.param(_subset_gate, None, True, id="subset-None-raises"), + pytest.param(_masked_gate, None, True, id="masked-None-raises"), + pytest.param(_subset_gate, 1, False, id="subset-fill1-ok"), + pytest.param(_masked_gate, 1, False, id="masked-fill1-ok"), + pytest.param(_subset_gate, 0, False, id="subset-fill0-ok"), + pytest.param(_full_gate, None, False, id="full-ok"), + pytest.param(_scalar_gate, None, False, id="scalar-ok"), +] + + +def _solve_partial_gate( + solver_name: str, + make_active: GateBuilder, + *, + method: Method, + disjunctive: bool = False, +) -> None: + """Fill a partial gate, force the committable units off, demand "b" runs.""" + m = Model() + x = m.add_variables(lower=0, upper=100, coords=[_PWL_GENS], name="x") + y = m.add_variables(lower=0, coords=[_PWL_GENS], name="y") + u = make_active(m) + if disjunctive: + m.add_piecewise_formulation( + (x, segments([[0.0, 50.0], [50.0, 100.0]])), + (y, segments([[0.0, 10.0], [10.0, 50.0]])), + active=u, + active_fill=1, + ) + else: + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=u, + active_fill=1, + method=method, + ) + m.add_constraints(u <= 0, name="force_off") + m.add_constraints(x.sel(gen="b") >= 50, name="demand") + m.add_objective(y.sum(), sense="min") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.sel(gen="a")), 0, atol=1e-4) + np.testing.assert_allclose(float(x.solution.sel(gen="c")), 0, atol=1e-4) + np.testing.assert_allclose(float(x.solution.sel(gen="b")), 50, atol=1e-4) + np.testing.assert_allclose(float(y.solution.sel(gen="b")), 10, atol=1e-4) + + +class TestResolveActiveFill: + """The private ``_resolve_active`` fills gaps with ``active_fill``.""" + + @pytest.mark.parametrize("fill_value", [1, 0]) + @pytest.mark.parametrize("make_active", _PARTIAL_GATES) + def test_fills_gap(self, make_active: GateBuilder, fill_value: int) -> None: + reference = xr.DataArray(np.zeros(len(_PWL_GENS)), coords=[_PWL_GENS]) + gate = _resolve_active(1 * make_active(Model()), reference, fill_value) + assert gate.const.sel(gen="b").item() == fill_value + assert bool((gate.vars.sel(gen="b") < 0).all()) # no variable at "b" + assert bool((gate.vars.sel(gen="a") >= 0).any()) # variable kept at "a" + + +class TestActiveFillValidation: + """``add_piecewise_formulation`` gates a partial ``active`` via ``active_fill``.""" + + @pytest.mark.parametrize("make_active, active_fill, should_raise", _COVERAGE_CASES) + def test_coverage( + self, + make_active: GateBuilder, + active_fill: int | None, + should_raise: bool, + ) -> None: + m = Model() + x = m.add_variables(lower=0, upper=100, coords=[_PWL_GENS], name="x") + y = m.add_variables(lower=0, coords=[_PWL_GENS], name="y") + + def build() -> None: + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=make_active(m), + active_fill=active_fill, + method="incremental", + ) + + if should_raise: + with pytest.raises(ValueError, match="active_fill"): + build() + else: + build() + + def test_active_fill_without_active_raises(self) -> None: + m = Model() + x = m.add_variables(lower=0, upper=100, coords=[_PWL_GENS], name="x") + y = m.add_variables(lower=0, coords=[_PWL_GENS], name="y") + with pytest.raises(ValueError, match="without `active`"): + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active_fill=1, + method="incremental", + ) + + def test_lower_dimensional_active_broadcasts(self) -> None: + """A gate missing an entire dim broadcasts and must not be rejected.""" + ts = pd.Index([0, 1], name="t") + m = Model() + x = m.add_variables(lower=0, upper=100, coords=[_PWL_GENS, ts], name="x") + y = m.add_variables(lower=0, coords=[_PWL_GENS, ts], name="y") + u = m.add_variables(binary=True, coords=[_PWL_GENS], name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), (y, [0, 10, 50]), active=u, method="incremental" + ) + + +@pytest.mark.skipif(len(_any_solvers) == 0, reason="No solver available") +class TestSolverActiveFill: + """End-to-end: ``active_fill`` leaves ungated units free (#796).""" + + @pytest.fixture(params=_any_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + @pytest.mark.parametrize("make_active", _PARTIAL_GATES) + def test_incremental(self, solver_name: str, make_active: GateBuilder) -> None: + _solve_partial_gate(solver_name, make_active, method="incremental") + + +@pytest.mark.skipif(len(_sos2_solvers) == 0, reason="No SOS2-capable solver") +class TestSolverActiveFillSOS2: + @pytest.fixture(params=_sos2_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + @pytest.mark.parametrize("make_active", _PARTIAL_GATES) + @pytest.mark.parametrize( + "method, disjunctive", + [ + pytest.param("sos2", False, id="sos2"), + pytest.param("auto", True, id="disjunctive"), + ], + ) + def test_solves( + self, + solver_name: str, + make_active: GateBuilder, + method: Method, + disjunctive: bool, + ) -> None: + _solve_partial_gate( + solver_name, make_active, method=method, disjunctive=disjunctive + ) diff --git a/test/test_piecewise_constraints.py b/test/test_piecewise_constraints.py new file mode 100644 index 000000000..72b57265d --- /dev/null +++ b/test/test_piecewise_constraints.py @@ -0,0 +1,3108 @@ +"""Tests for the new piecewise linear constraints API.""" + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Callable, Generator +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, TypeAlias + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import ( + Model, + available_solvers, + breakpoints, + segments, + tangent_lines, +) +from linopy.constants import ( + BREAKPOINT_DIM, + LP_PIECE_DIM, + PWL_ACTIVE_BOUND_SUFFIX, + PWL_BINARY_ORDER_SUFFIX, + PWL_CHORD_SUFFIX, + PWL_CONVEX_SUFFIX, + PWL_DELTA_BOUND_SUFFIX, + PWL_DELTA_SUFFIX, + PWL_DOMAIN_HI_SUFFIX, + PWL_DOMAIN_LO_SUFFIX, + PWL_FILL_ORDER_SUFFIX, + PWL_LAMBDA_SUFFIX, + PWL_LINK_SUFFIX, + PWL_ORDER_BINARY_SUFFIX, + PWL_OUTPUT_LINK_SUFFIX, + PWL_SEGMENT_BINARY_SUFFIX, + PWL_SELECT_SUFFIX, + SEGMENT_DIM, +) +from linopy.piecewise import _slopes_to_points +from linopy.solver_capabilities import ( + SolverFeature, + get_available_solvers_with_feature, + solver_supports, +) + +if TYPE_CHECKING: + from linopy.piecewise import BreaksLike, _PwlInputs + +Sign: TypeAlias = Literal["==", "<=", ">="] +Method: TypeAlias = Literal["sos2", "incremental", "lp", "auto"] + +_sos2_solvers = get_available_solvers_with_feature( + SolverFeature.SOS_CONSTRAINTS, available_solvers +) +_sos2_direct_solvers = sorted( + s for s in _sos2_solvers if solver_supports(s, SolverFeature.DIRECT_API) +) +_SOS_PATHS = [ + *[pytest.param(s, "direct", id=f"{s}-direct") for s in _sos2_direct_solvers], + *[pytest.param(s, "lp", id=f"{s}-lp") for s in sorted(_sos2_solvers)], +] +_any_solvers = [ + s for s in ["highs", "gurobi", "glpk", "cplex"] if s in available_solvers +] + +# Solver-output tolerance for solution-value assertions in this file. Matches +# the convention in ``test_piecewise_feasibility.py``. +TOL = 1e-6 + + +# =========================================================================== +# _slopes_to_points (private list utility) +# =========================================================================== + + +class TestSlopesToPointsPrivate: + """ + The list-level slopes→points primitive is private; the public path is + :class:`Slopes`. These tests exist so the math stays under test even + though the helper isn't user-facing. + """ + + def test_basic(self) -> None: + assert _slopes_to_points([0, 1, 2], [1, 2], 0) == [0, 1, 3] + + def test_negative_slopes(self) -> None: + assert _slopes_to_points([0, 10, 20], [-0.5, -1.0], 10) == [10, 5, -5] + + def test_wrong_length_raises(self) -> None: + with pytest.raises(ValueError, match="len\\(slopes\\)"): + _slopes_to_points([0, 1, 2], [1], 0) + + +# =========================================================================== +# breakpoints() factory +# =========================================================================== + + +class TestBreakpointsFactory: + def test_list(self) -> None: + bp = breakpoints([0, 50, 100]) + assert bp.dims == (BREAKPOINT_DIM,) + assert list(bp.values) == [0.0, 50.0, 100.0] + + def test_dict(self) -> None: + bp = breakpoints({"gen1": [0, 50, 100], "gen2": [0, 30]}, dim="generator") + assert set(bp.dims) == {"generator", BREAKPOINT_DIM} + assert bp.sizes[BREAKPOINT_DIM] == 3 + assert np.isnan(bp.sel(generator="gen2").sel({BREAKPOINT_DIM: 2})) + + def test_dict_without_dim_raises(self) -> None: + with pytest.raises(ValueError, match="'dim' is required"): + breakpoints({"a": [0, 50], "b": [0, 30]}) + + def test_slopes_kwargs_removed(self) -> None: + """The slopes mode of ``breakpoints`` was removed in favour of ``Slopes``.""" + with pytest.raises(TypeError): + breakpoints([0, 1], slopes=[1], x_points=[0, 1], y0=0) # type: ignore[call-arg] + + # --- pandas and xarray inputs --- + + def test_series(self) -> None: + bp = breakpoints(pd.Series([0, 50, 100])) + assert bp.dims == (BREAKPOINT_DIM,) + assert list(bp.values) == [0.0, 50.0, 100.0] + + def test_dataframe(self) -> None: + df = pd.DataFrame( + {"gen1": [0, 50, 100], "gen2": [0, 30, np.nan]} + ).T # rows=entities, cols=breakpoints + bp = breakpoints(df, dim="generator") + assert set(bp.dims) == {"generator", BREAKPOINT_DIM} + assert bp.sizes[BREAKPOINT_DIM] == 3 + np.testing.assert_allclose(bp.sel(generator="gen1").values, [0, 50, 100]) + assert np.isnan(bp.sel(generator="gen2").values[2]) + + def test_dataframe_without_dim_raises(self) -> None: + df = pd.DataFrame({"a": [0, 50], "b": [0, 30]}).T + with pytest.raises(ValueError, match="'dim' is required"): + breakpoints(df) + + def test_dataarray_passthrough(self) -> None: + da = xr.DataArray( + [0, 50, 100], + dims=[BREAKPOINT_DIM], + coords={BREAKPOINT_DIM: np.arange(3)}, + ) + bp = breakpoints(da) + xr.testing.assert_equal(bp, da) + + def test_dataarray_missing_dim_raises(self) -> None: + da = xr.DataArray([0, 50, 100], dims=["foo"]) + with pytest.raises(ValueError, match="must have a"): + breakpoints(da) + + +# =========================================================================== +# Slopes class — deferred breakpoint spec +# =========================================================================== + + +class TestSlopesValueType: + """``Slopes`` is a frozen value type with a custom repr.""" + + def test_immutable(self) -> None: + from linopy import Slopes + + s = Slopes([1, 2], y0=0) + with pytest.raises((AttributeError, TypeError)): + s.y0 = 5 # type: ignore[misc] + + @pytest.mark.parametrize( + ("kwargs", "expected"), + [ + pytest.param( + {"values": [1.2, 1.6, 2.15], "y0": 0}, + "Slopes([1.2, 1.6, 2.15], y0=0)", + id="1d_list_defaults_hidden", + ), + pytest.param( + {"values": [np.nan, 1, 2], "y0": 0, "align": "leading"}, + "align='leading'", + id="non_default_align_shown", + ), + pytest.param( + {"values": [1, 2], "y0": 0, "dim": "gen"}, + "dim='gen'", + id="non_default_dim_shown", + ), + ], + ) + def test_repr_renders(self, kwargs: dict[str, Any], expected: str) -> None: + from linopy import Slopes + + r = repr(Slopes(**kwargs)) + if expected.startswith("Slopes("): + assert r == expected + else: + assert expected in r + + def test_repr_truncates_long_sequences(self) -> None: + """Lists/ndarrays over 8 entries must be summarised, not dumped.""" + from linopy import Slopes + + r = repr(Slopes(list(range(50)), y0=0)) + # No 50-element dump — must include the "(50 items)" suffix and + # contain at most a handful of explicit numbers. + assert "(50 items)" in r + assert "..." in r + assert len(r) < 80, f"repr unexpectedly long: {r!r}" + + def test_repr_normalises_numpy_scalars(self) -> None: + """``np.int64`` / ``np.float64`` must render as plain Python numbers.""" + from linopy import Slopes + + r_int = repr(Slopes(np.array([1, 2, 3], dtype=np.int64), y0=0)) + r_float = repr(Slopes(np.array([1.5, 2.5, 3.5]), y0=0)) + # No numpy type prefixes, no surprising precision. + assert "np." not in r_int and "int64" not in r_int + assert r_int == "Slopes([1, 2, 3], y0=0)" + assert r_float == "Slopes([1.5, 2.5, 3.5], y0=0)" + + @pytest.mark.parametrize( + ("a", "b", "expected"), + [ + pytest.param( + {"values": [1, 2], "y0": 0}, + {"values": [1, 2], "y0": 0}, + True, + id="lists_equal", + ), + pytest.param( + {"values": np.array([1, 2]), "y0": 0}, + {"values": np.array([1, 2]), "y0": 0}, + True, + id="ndarrays_equal_no_raise", + ), + pytest.param( + {"values": [1, 2], "y0": 0}, + {"values": [1, 3], "y0": 0}, + False, + id="different_values", + ), + pytest.param( + {"values": [1, 2], "y0": 0}, + {"values": [1, 2], "y0": 5}, + False, + id="different_y0", + ), + pytest.param( + # list and ndarray of same numeric content — list/tuple + # are promoted to ndarray, so they compare equal. + {"values": [1, 2], "y0": 0}, + {"values": np.array([1, 2]), "y0": 0}, + True, + id="list_and_ndarray_same_content", + ), + pytest.param( + # int and float y0 describe the same curve — Real scalars + # coerce numerically. + {"values": [1, 2], "y0": 0}, + {"values": [1, 2], "y0": 0.0}, + True, + id="int_and_float_y0", + ), + pytest.param( + # numpy scalar y0 vs Python float — same numeric value. + {"values": [1, 2], "y0": np.float64(0)}, + {"values": [1, 2], "y0": 0.0}, + True, + id="numpy_scalar_and_float_y0", + ), + pytest.param( + # In-place ``float('nan')`` (not the np.nan singleton) must + # still compare equal — the array-path promotion handles it. + {"values": [float("nan"), 1.0], "y0": 0, "align": "leading"}, + {"values": [float("nan"), 1.0], "y0": 0, "align": "leading"}, + True, + id="float_nan_in_list", + ), + pytest.param( + {"values": [np.nan, 1], "y0": 0, "align": "leading"}, + {"values": [np.nan, 1], "y0": 0, "align": "leading"}, + True, + id="np_nan_in_list", + ), + pytest.param( + {"values": [1, 2], "y0": float("nan")}, + {"values": [1, 2], "y0": float("nan")}, + True, + id="nan_in_scalar_y0", + ), + pytest.param( + {"values": {"a": [1, 2], "b": [3, 4]}, "y0": 0, "dim": "g"}, + {"values": {"a": [1, 2], "b": [3, 4]}, "y0": 0, "dim": "g"}, + True, + id="dict_equal", + ), + pytest.param( + {"values": {"a": [1, 2]}, "y0": 0, "dim": "g"}, + {"values": {"a": [9, 9]}, "y0": 0, "dim": "g"}, + False, + id="dict_different_inner_values", + ), + ], + ) + def test_equality( + self, a: dict[str, Any], b: dict[str, Any], expected: bool + ) -> None: + """Value-equality across the field types accepted by the constructor.""" + from linopy import Slopes + + assert (Slopes(**a) == Slopes(**b)) is expected + + def test_eq_against_non_slopes_returns_notimplemented(self) -> None: + from linopy import Slopes + + # Falls through to bool(False), not raising. + assert (Slopes([1, 2], y0=0) == "not a slopes") is False + assert (Slopes([1, 2], y0=0) == 42) is False + + def test_eq_dataframe_is_order_sensitive(self) -> None: + """``DataFrame.equals`` is order-sensitive — pin the documented caveat.""" + from linopy import Slopes + + df1 = pd.DataFrame({"a": [1, 0.5], "b": [2, 1]}).T + df2 = df1.loc[["b", "a"]] + assert (Slopes(df1, y0=0, dim="g") == Slopes(df2, y0=0, dim="g")) is False + + def test_eq_object_dtype_ndarray_does_not_raise(self) -> None: + """Object/string-dtype ndarrays fall back to plain array_equal.""" + from linopy import Slopes + + a = np.array(["x", "y"], dtype=object) + b = np.array(["x", "y"], dtype=object) + c = np.array(["x", "z"], dtype=object) + # Equal content -> True; different content -> False; neither raises. + assert (Slopes(a, y0=0) == Slopes(b, y0=0)) is True + assert (Slopes(a, y0=0) == Slopes(c, y0=0)) is False + + def test_unhashable(self) -> None: + """ + ``values`` may be a mutable container (list, ndarray, dict), so + ``Slopes`` is intentionally unhashable. Using one as a dict key + or set member must raise rather than silently using identity hash. + """ + from linopy import Slopes + + with pytest.raises(TypeError, match="unhashable"): + {Slopes([1, 2], y0=0): "x"} + + @pytest.mark.parametrize( + ("values", "fragment"), + [ + pytest.param( + pd.DataFrame({"a": [1, 2], "b": [3, 4]}).T, + "", + id="dict", + ), + pytest.param( + np.zeros((20, 5, 30)), + "", + id="multi_dim_ndarray", + ), + ], + ) + def test_repr_summarises_bulky_values( + self, values: BreaksLike, fragment: str + ) -> None: + """Bulky value types must not dump their full content into the repr.""" + from linopy import Slopes + + r = repr(Slopes(values, y0=0, dim="gen")) + assert fragment in r + + +class TestSlopesToBreakpoints1D: + """ + 1D inputs (single shared curve). All callable input types must + resolve to the same DataArray for the same data: slopes [1, 2] over + x = [0, 1, 2] with y0=0 yields y = [0, 1, 3]. + """ + + EXPECTED = [0.0, 1.0, 3.0] + + @pytest.mark.parametrize( + ("slopes_in", "x_in"), + [ + pytest.param([1, 2], [0, 1, 2], id="list-list"), + pytest.param((1, 2), (0, 1, 2), id="tuple-tuple"), + pytest.param(np.array([1, 2]), np.array([0, 1, 2]), id="ndarray-ndarray"), + pytest.param(pd.Series([1, 2]), pd.Series([0, 1, 2]), id="series-series"), + pytest.param([1, 2], np.array([0, 1, 2]), id="list-ndarray-mixed"), + pytest.param( + xr.DataArray([1, 2], dims=[BREAKPOINT_DIM]), + xr.DataArray([0, 1, 2], dims=[BREAKPOINT_DIM]), + id="dataarray-dataarray", + ), + ], + ) + def test_resolves_to_expected_breakpoints( + self, slopes_in: BreaksLike, x_in: BreaksLike + ) -> None: + from linopy import Slopes + + bp = Slopes(slopes_in, y0=0).to_breakpoints(x_in) + assert bp.dims == (BREAKPOINT_DIM,) + np.testing.assert_allclose(bp.values, self.EXPECTED) + + @pytest.mark.parametrize( + ("slopes", "x_pts", "y0", "expected"), + [ + pytest.param([1, 2], [0, 1, 2], 0, [0, 1, 3], id="canonical"), + pytest.param( + [1.2, 1.4, 1.7], + [0, 30, 60, 100], + 0, + [0, 36, 78, 146], + id="non_unit_slopes", + ), + pytest.param([-0.5, -1.0], [0, 10, 20], 10, [10, 5, -5], id="negative"), + pytest.param([1, 2], [0, 1, 2], 5, [5, 6, 8], id="non_zero_y0"), + ], + ) + def test_arithmetic_anchors( + self, + slopes: list[float], + x_pts: list[float], + y0: float, + expected: list[float], + ) -> None: + """Hand-computable cases pinning the slopes→y arithmetic.""" + from linopy import Slopes + + bp = Slopes(slopes, y0=y0).to_breakpoints(x_pts) + np.testing.assert_allclose(bp.values, expected) + + +class TestSlopesToBreakpointsPerEntity: + """ + Per-entity inputs (multiple curves along one entity dim). All input + container types must produce the same per-entity result. + + Reference data: gen=a slopes [1, 0.5] over x=[0, 10, 50] from y0=0 + → [0, 10, 30]; gen=b slopes [2, 1] over x=[0, 20, 80] from y0=10 + → [10, 50, 110]. + """ + + EXPECTED_A = [0.0, 10.0, 30.0] + EXPECTED_B = [10.0, 50.0, 110.0] + + @pytest.mark.parametrize( + ("slopes_in", "x_in"), + [ + pytest.param( + {"a": [1, 0.5], "b": [2, 1]}, + {"a": [0, 10, 50], "b": [0, 20, 80]}, + id="dict-dict", + ), + pytest.param( + pd.DataFrame({"a": [1, 0.5], "b": [2, 1]}).T, + pd.DataFrame({"a": [0, 10, 50], "b": [0, 20, 80]}).T, + id="dataframe-dataframe", + ), + pytest.param( + xr.DataArray( + [[1, 0.5], [2, 1]], + dims=["gen", BREAKPOINT_DIM], + coords={"gen": ["a", "b"], BREAKPOINT_DIM: [0, 1]}, + ), + xr.DataArray( + [[0, 10, 50], [0, 20, 80]], + dims=["gen", BREAKPOINT_DIM], + coords={"gen": ["a", "b"], BREAKPOINT_DIM: [0, 1, 2]}, + ), + id="dataarray-dataarray", + ), + ], + ) + def test_resolves_to_expected_per_entity( + self, slopes_in: BreaksLike, x_in: BreaksLike + ) -> None: + from linopy import Slopes + + bp = Slopes(slopes_in, y0={"a": 0, "b": 10}, dim="gen").to_breakpoints(x_in) + assert "gen" in bp.dims and BREAKPOINT_DIM in bp.dims + np.testing.assert_allclose(bp.sel(gen="a").values, self.EXPECTED_A) + np.testing.assert_allclose(bp.sel(gen="b").values, self.EXPECTED_B) + + def test_shared_x_grid_broadcasts(self) -> None: + """Per-entity slopes against a single shared x grid (1D x_points).""" + from linopy import Slopes + + bp = Slopes( + {"a": [1, 2], "b": [3, 4]}, y0={"a": 0, "b": 0}, dim="gen" + ).to_breakpoints([0, 1, 2]) + np.testing.assert_allclose(bp.sel(gen="a").values, [0, 1, 3]) + np.testing.assert_allclose(bp.sel(gen="b").values, [0, 3, 7]) + + @pytest.mark.parametrize( + ("y0", "id"), + [ + pytest.param(5.0, "scalar"), + pytest.param({"a": 5, "b": 5}, "dict"), + pytest.param(pd.Series({"a": 5, "b": 5}), "series"), + pytest.param( + xr.DataArray([5, 5], dims=["gen"], coords={"gen": ["a", "b"]}), + "dataarray", + ), + ], + ids=lambda x: x if isinstance(x, str) else None, + ) + def test_y0_input_types_broadcast_consistently(self, y0: object, id: str) -> None: + """All accepted ``y0`` shapes resolve to the same per-entity result.""" + from linopy import Slopes + + bp = Slopes({"a": [1, 2], "b": [3, 4]}, y0=y0, dim="gen").to_breakpoints( + {"a": [0, 1, 2], "b": [0, 1, 2]} + ) + np.testing.assert_allclose(bp.sel(gen="a").values, [5, 6, 8]) + np.testing.assert_allclose(bp.sel(gen="b").values, [5, 8, 12]) + + +class TestSlopesToBreakpointsAlignment: + """ + ``align="pieces"`` (n-1 slopes) and ``align="leading"`` (n slopes + with a NaN sentinel in position 0) describe the same curve. They + must produce the same breakpoint DataArray. + """ + + @pytest.mark.parametrize( + ("pieces_input", "leading_input"), + [ + pytest.param([1, 2], [np.nan, 1, 2], id="1d"), + pytest.param( + {"a": [1, 0.5], "b": [2, 1]}, + {"a": [np.nan, 1, 0.5], "b": [np.nan, 2, 1]}, + id="dict_per_entity", + ), + ], + ) + def test_pieces_and_leading_match( + self, pieces_input: BreaksLike, leading_input: BreaksLike + ) -> None: + from linopy import Slopes + + kwargs: dict[str, Any] = {"y0": 0} + if isinstance(pieces_input, dict): + kwargs.update(dim="gen", y0={"a": 0, "b": 10}) + x_pts: BreaksLike = {"a": [0, 10, 50], "b": [0, 20, 80]} + else: + x_pts = [0, 1, 2] + pieces_bp = Slopes(pieces_input, align="pieces", **kwargs).to_breakpoints(x_pts) + leading_bp = Slopes(leading_input, align="leading", **kwargs).to_breakpoints( + x_pts + ) + xr.testing.assert_allclose(pieces_bp, leading_bp) + + def test_leading_ragged_dict(self) -> None: + """``align='leading'`` with ragged per-entity input keeps NaN padding.""" + from linopy import Slopes + + bp = Slopes( + {"a": [np.nan, 1, 0.5], "b": [np.nan, 2]}, + y0={"a": 0, "b": 10}, + dim="gen", + align="leading", + ).to_breakpoints({"a": [0, 10, 50], "b": [0, 20]}) + np.testing.assert_allclose(bp.sel(gen="a").values, [0, 10, 30]) + np.testing.assert_allclose( + bp.sel(gen="b").dropna(BREAKPOINT_DIM).values, [10, 50] + ) + + +class TestSlopesValidationErrors: + """``to_breakpoints`` rejects malformed specs with actionable messages.""" + + @pytest.mark.parametrize( + ("ctor_kwargs", "x_pts", "match"), + [ + pytest.param( + {"values": [1, 2, 3], "y0": 0, "align": "leading"}, + [0, 1, 2], + "first slope", + id="leading_first_not_nan", + ), + pytest.param( + {"values": [1, 2], "y0": {"a": 0}}, + [0, 10, 20], + "scalar float", + id="1d_with_dict_y0", + ), + pytest.param( + {"values": {"a": [1, 2], "b": [3, 4]}, "y0": "bad", "dim": "gen"}, + {"a": [0, 10, 20], "b": [0, 10, 20]}, + "y0", + id="bad_y0_type", + ), + ], + ) + def test_invalid_inputs_raise( + self, + ctor_kwargs: dict[str, Any], + x_pts: BreaksLike, + match: str, + ) -> None: + from linopy import Slopes + + with pytest.raises((TypeError, ValueError), match=match): + Slopes(**ctor_kwargs).to_breakpoints(x_pts) + + +class TestSlopesDispatch: + """Slopes inside ``add_piecewise_formulation`` — sibling resolution.""" + + def test_two_tuple_deferred(self) -> None: + from linopy import Slopes + + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(lower=0, name="fuel") + # Slopes [1.2, 1.4, 1.7] resolved over the borrowed x grid + # [0, 30, 60, 100] -> fuel breakpoints [0, 36, 78, 146]. + # Equality-2-tuple convexity uses pinned_bps[1] as x; with + # increasing dy/dx slopes, the inverse view (power-vs-fuel) is + # concave — that's the label the formulation reports. + f = m.add_piecewise_formulation( + (power, [0, 30, 60, 100]), + (fuel, Slopes([1.2, 1.4, 1.7], y0=0)), + ) + assert f.method in ("sos2", "incremental") + assert f.convexity == "concave" + + def test_slopes_as_bounded_tuple(self) -> None: + from linopy import Slopes + + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=100, name="y") + f = m.add_piecewise_formulation( + (y, Slopes([2, 1, 0.5], y0=0), "<="), # concave + (x, [0, 10, 20, 30]), + ) + assert f.method == "lp" + assert f.convexity == "concave" + + def test_all_slopes_raises(self) -> None: + from linopy import Slopes + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="All tuples are Slopes"): + m.add_piecewise_formulation( + (x, Slopes([1, 2], y0=0)), + (y, Slopes([1, 1], y0=0)), + ) + + def test_multiple_non_slopes_with_slopes_raises(self) -> None: + """ + With Slopes present, two or more non-Slopes tuples is rejected: + each non-Slopes tuple is a y-vector for its own variable, so + there is no canonical x grid for the Slopes to integrate against. + """ + from linopy import Slopes + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + with pytest.raises(ValueError, match="no canonical x grid"): + m.add_piecewise_formulation( + (x, [0, 10, 20, 30]), + (y, [0, 100, 200, 300]), + (z, Slopes([1, 1, 1], y0=0)), + ) + + def test_multiple_slopes_share_x_grid(self) -> None: + """ + Two Slopes tuples plus one non-Slopes — both Slopes resolve against + the same borrowed x grid. Pin via distinct slope sequences so the + two Slopes-derived variables end up with different breakpoint values. + """ + from linopy import Slopes + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + f = m.add_piecewise_formulation( + (x, [0, 10, 20, 30]), + (y, Slopes([1, 1, 1], y0=0)), # → [0, 10, 20, 30] + (z, Slopes([2, 2, 2], y0=0)), # → [0, 20, 40, 60] + ) + # 3-var formulation -> convexity is None. + assert f.convexity is None + assert f.name in m._piecewise_formulations + + def test_slopes_align_leading_in_dispatch(self) -> None: + from linopy import Slopes + + m = Model() + x = m.add_variables(lower=0, upper=2, name="x") + y = m.add_variables(name="y") + f = m.add_piecewise_formulation( + (x, [0, 1, 2]), + (y, Slopes([np.nan, 1, 2], y0=0, align="leading")), + ) + # Resolved bp for y: [0, 1, 3]. As above, the equality-2-tuple + # convention reports the inverse view → concave. + assert f.convexity == "concave" + + +class TestSlopesDispatchEquivalence: + """ + Deferred Slopes dispatch builds the same model as eager breakpoints. + + The wiring tests in :class:`TestSlopesDispatch` verify dispatch attributes + (``method``/``convexity``). These tests pin the *outcome*: the deferred + form must produce a model byte-equal to the eagerly-resolved reference + (same auxiliary variables, same constraint coefficients/RHS). + """ + + def test_two_tuple_matches_eager(self) -> None: + from linopy import Slopes + from linopy.testing import assert_model_equal + + # Slopes([1.2, 1.4, 1.7], y0=0) over [0, 30, 60, 100] resolves to + # fuel breakpoints [0, 36, 78, 146]. + m_eager = Model() + p1 = m_eager.add_variables(lower=0, upper=100, name="power") + f1 = m_eager.add_variables(lower=0, name="fuel") + m_eager.add_piecewise_formulation( + (p1, [0, 30, 60, 100]), (f1, [0, 36, 78, 146]) + ) + + m_deferred = Model() + p2 = m_deferred.add_variables(lower=0, upper=100, name="power") + f2 = m_deferred.add_variables(lower=0, name="fuel") + m_deferred.add_piecewise_formulation( + (p2, [0, 30, 60, 100]), + (f2, Slopes([1.2, 1.4, 1.7], y0=0)), + ) + + assert_model_equal(m_eager, m_deferred) + + def test_multiple_slopes_resolved_breakpoints(self) -> None: + """ + Two Slopes tuples resolve against the same borrowed x grid: + y → [0, 10, 20, 30], z → [0, 20, 40, 60]. + """ + from linopy import Slopes + from linopy.testing import assert_model_equal + + m_eager = Model() + x1 = m_eager.add_variables(lower=0, upper=30, name="x") + y1 = m_eager.add_variables(lower=0, name="y") + z1 = m_eager.add_variables(lower=0, name="z") + m_eager.add_piecewise_formulation( + (x1, [0, 10, 20, 30]), + (y1, [0, 10, 20, 30]), + (z1, [0, 20, 40, 60]), + ) + + m_deferred = Model() + x2 = m_deferred.add_variables(lower=0, upper=30, name="x") + y2 = m_deferred.add_variables(lower=0, name="y") + z2 = m_deferred.add_variables(lower=0, name="z") + m_deferred.add_piecewise_formulation( + (x2, [0, 10, 20, 30]), + (y2, Slopes([1, 1, 1], y0=0)), + (z2, Slopes([2, 2, 2], y0=0)), + ) + + assert_model_equal(m_eager, m_deferred) + + def test_align_leading_matches_eager(self) -> None: + """``align='leading'`` dispatch resolves to bps [0, 1, 3].""" + from linopy import Slopes + from linopy.testing import assert_model_equal + + m_eager = Model() + x1 = m_eager.add_variables(lower=0, upper=2, name="x") + y1 = m_eager.add_variables(name="y") + m_eager.add_piecewise_formulation((x1, [0, 1, 2]), (y1, [0, 1, 3])) + + m_deferred = Model() + x2 = m_deferred.add_variables(lower=0, upper=2, name="x") + y2 = m_deferred.add_variables(name="y") + m_deferred.add_piecewise_formulation( + (x2, [0, 1, 2]), + (y2, Slopes([np.nan, 1, 2], y0=0, align="leading")), + ) + + assert_model_equal(m_eager, m_deferred) + + +# =========================================================================== +# segments() factory +# =========================================================================== + + +class TestSegmentsFactory: + def test_list(self) -> None: + bp = segments([[0, 10], [50, 100]]) + assert set(bp.dims) == {SEGMENT_DIM, BREAKPOINT_DIM} + assert bp.sizes[SEGMENT_DIM] == 2 + assert bp.sizes[BREAKPOINT_DIM] == 2 + + def test_dict(self) -> None: + bp = segments( + {"a": [[0, 10], [50, 100]], "b": [[0, 20], [60, 90]]}, + dim="gen", + ) + assert "gen" in bp.dims + assert SEGMENT_DIM in bp.dims + assert BREAKPOINT_DIM in bp.dims + + def test_ragged(self) -> None: + bp = segments([[0, 5, 10], [50, 100]]) + assert bp.sizes[BREAKPOINT_DIM] == 3 + assert np.isnan(bp.sel({SEGMENT_DIM: 1, BREAKPOINT_DIM: 2})) + + def test_dict_without_dim_raises(self) -> None: + with pytest.raises(ValueError, match="'dim' is required"): + segments({"a": [[0, 10]], "b": [[50, 100]]}) + + def test_dataframe(self) -> None: + df = pd.DataFrame([[0, 10], [50, 100]]) # rows=segments, cols=breakpoints + bp = segments(df) + assert set(bp.dims) == {SEGMENT_DIM, BREAKPOINT_DIM} + assert bp.sizes[SEGMENT_DIM] == 2 + assert bp.sizes[BREAKPOINT_DIM] == 2 + np.testing.assert_allclose(bp.sel({SEGMENT_DIM: 0}).values, [0, 10]) + np.testing.assert_allclose(bp.sel({SEGMENT_DIM: 1}).values, [50, 100]) + + def test_dataarray_passthrough(self) -> None: + da = xr.DataArray( + [[0, 10], [50, 100]], + dims=[SEGMENT_DIM, BREAKPOINT_DIM], + coords={SEGMENT_DIM: [0, 1], BREAKPOINT_DIM: [0, 1]}, + ) + bp = segments(da) + xr.testing.assert_equal(bp, da) + + def test_dataarray_missing_dim_raises(self) -> None: + da_no_seg = xr.DataArray( + [[0, 10], [50, 100]], + dims=["foo", BREAKPOINT_DIM], + ) + with pytest.raises(ValueError, match="must have both"): + segments(da_no_seg) + + da_no_bp = xr.DataArray( + [[0, 10], [50, 100]], + dims=[SEGMENT_DIM, "bar"], + ) + with pytest.raises(ValueError, match="must have both"): + segments(da_no_bp) + + +# =========================================================================== +# Continuous piecewise -- equality +# =========================================================================== + + +class TestContinuousEquality: + def test_sos2(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [5, 2, 20, 80]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_CONVEX_SUFFIX}" in m.constraints + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + # N-var path uses a single stacked link constraint (no separate y_link) + lam = m.variables[f"pwl0{PWL_LAMBDA_SUFFIX}"] + assert lam.attrs.get("sos_type") == 2 + + def test_auto_selects_incremental_for_monotonic(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # Both breakpoint sequences must be monotonic for incremental + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [0, 5, 20, 80]), + ) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + assert f"pwl0{PWL_LAMBDA_SUFFIX}" not in m.variables + + def test_auto_nonmonotonic_falls_back_to_sos2(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # Non-monotonic y-breakpoints force SOS2 + m.add_piecewise_formulation( + (x, [0, 50, 30, 100]), + (y, [5, 20, 15, 80]), + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_DELTA_SUFFIX}" not in m.variables + + def test_multi_dimensional(self) -> None: + m = Model() + gens = pd.Index(["gen_a", "gen_b"], name="generator") + x = m.add_variables(coords=[gens], name="x") + y = m.add_variables(coords=[gens], name="y") + m.add_piecewise_formulation( + ( + x, + breakpoints( + {"gen_a": [0, 10, 50], "gen_b": [0, 20, 80]}, dim="generator" + ), + ), + ( + y, + breakpoints( + {"gen_a": [0, 5, 30], "gen_b": [0, 8, 50]}, dim="generator" + ), + ), + ) + delta = m.variables[f"pwl0{PWL_DELTA_SUFFIX}"] + assert "generator" in delta.dims + + def test_with_slopes(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # slopes=[-0.3, 0.45, 1.2] with y0=5 -> y_points=[5, 2, 20, 80] + # Non-monotonic y-breakpoints, so auto selects SOS2 + from linopy import Slopes + + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, Slopes([-0.3, 0.45, 1.2], y0=5)), + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + + +# =========================================================================== +# Piecewise Envelope +# =========================================================================== + + +class TestTangentLines: + def test_basic_variable(self) -> None: + """Envelope from a Variable produces a LinearExpression with seg dim.""" + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + env = tangent_lines(x, [0, 50, 100], [0, 40, 60]) + assert LP_PIECE_DIM in env.dims + + def test_basic_linexpr(self) -> None: + """Envelope from a LinearExpression works too.""" + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + env = tangent_lines(1 * x, [0, 50, 100], [0, 40, 60]) + assert LP_PIECE_DIM in env.dims + + def test_piece_count(self) -> None: + """Number of pieces = number of breakpoints - 1.""" + m = Model() + x = m.add_variables(name="x") + env = tangent_lines(x, [0, 50, 100], [0, 40, 60]) + assert env.sizes[LP_PIECE_DIM] == 2 + + def test_invalid_x_type_raises(self) -> None: + with pytest.raises(TypeError, match="must be a Variable or LinearExpression"): + tangent_lines(42, [0, 50, 100], [0, 40, 60]) # type: ignore + + def test_concave_le_constraint(self) -> None: + """Using envelope with <= constraint creates regular constraints.""" + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + y = m.add_variables(name="y") + env = tangent_lines(x, [0, 50, 100], [0, 40, 60]) + m.add_constraints(y <= env, name="pwl") + assert "pwl" in m.constraints + + def test_convex_ge_constraint(self) -> None: + """Using envelope with >= constraint creates regular constraints.""" + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + y = m.add_variables(name="y") + env = tangent_lines(x, [0, 50, 100], [0, 10, 60]) + m.add_constraints(y >= env, name="pwl") + assert "pwl" in m.constraints + + def test_dataarray_breakpoints(self) -> None: + """Envelope accepts DataArray breakpoints.""" + m = Model() + x = m.add_variables(name="x") + x_pts = xr.DataArray([0, 50, 100], dims=[BREAKPOINT_DIM]) + y_pts = xr.DataArray([0, 40, 60], dims=[BREAKPOINT_DIM]) + env = tangent_lines(x, x_pts, y_pts) + assert LP_PIECE_DIM in env.dims + + +# =========================================================================== +# Incremental formulation +# =========================================================================== + + +class TestIncremental: + def test_creates_delta_vars(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [5, 10, 20, 80]), + method="incremental", + ) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + delta = m.variables[f"pwl0{PWL_DELTA_SUFFIX}"] + assert delta.labels.sizes[LP_PIECE_DIM] == 3 + assert f"pwl0{PWL_FILL_ORDER_SUFFIX}" in m.constraints + assert f"pwl0{PWL_LAMBDA_SUFFIX}" not in m.variables + + def test_nonmonotonic_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="strictly monotonic"): + m.add_piecewise_formulation( + (x, [0, 50, 30, 100]), + (y, [5, 20, 15, 80]), + method="incremental", + ) + + def test_sos2_nonmonotonic_succeeds(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 50, 30, 100]), + (y, [5, 20, 15, 80]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_DELTA_SUFFIX}" not in m.variables + + def test_two_breakpoints_no_fill(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 100]), + (y, [5, 80]), + method="incremental", + ) + delta = m.variables[f"pwl0{PWL_DELTA_SUFFIX}"] + assert delta.labels.sizes[LP_PIECE_DIM] == 1 + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + # N-var path uses a single stacked link constraint (no separate y_link) + + def test_creates_binary_indicator_vars(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [5, 10, 20, 80]), + method="incremental", + ) + assert f"pwl0{PWL_ORDER_BINARY_SUFFIX}" in m.variables + binary = m.variables[f"pwl0{PWL_ORDER_BINARY_SUFFIX}"] + assert binary.labels.sizes[LP_PIECE_DIM] == 3 + assert f"pwl0{PWL_DELTA_BOUND_SUFFIX}" in m.constraints + + def test_creates_order_constraints(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [5, 10, 20, 80]), + method="incremental", + ) + assert f"pwl0{PWL_BINARY_ORDER_SUFFIX}" in m.constraints + + def test_two_breakpoints_no_order_constraint(self) -> None: + """With only one segment, there's no order constraint needed.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 100]), + (y, [5, 80]), + method="incremental", + ) + assert f"pwl0{PWL_ORDER_BINARY_SUFFIX}" in m.variables + assert f"pwl0{PWL_DELTA_BOUND_SUFFIX}" in m.constraints + assert f"pwl0{PWL_BINARY_ORDER_SUFFIX}" not in m.constraints + + def test_decreasing_monotonic(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [100, 50, 10, 0]), + (y, [80, 20, 5, 2]), + method="incremental", + ) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + + +# =========================================================================== +# Disjunctive piecewise +# =========================================================================== + + +class TestDisjunctive: + def test_equality_creates_binary(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, segments([[0, 10], [50, 100]])), + (y, segments([[0, 5], [20, 80]])), + ) + assert f"pwl0{PWL_SEGMENT_BINARY_SUFFIX}" in m.variables + assert f"pwl0{PWL_SELECT_SUFFIX}" in m.constraints + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_CONVEX_SUFFIX}" in m.constraints + lam = m.variables[f"pwl0{PWL_LAMBDA_SUFFIX}"] + assert lam.attrs.get("sos_type") == 2 + + def test_method_incremental_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="disjunctive"): + m.add_piecewise_formulation( + (x, segments([[0, 10], [50, 100]])), + (y, segments([[0, 5], [20, 80]])), + method="incremental", + ) + + def test_multi_dimensional(self) -> None: + m = Model() + gens = pd.Index(["gen_a", "gen_b"], name="generator") + x = m.add_variables(coords=[gens], name="x") + y = m.add_variables(coords=[gens], name="y") + m.add_piecewise_formulation( + ( + x, + segments( + {"gen_a": [[0, 10], [50, 100]], "gen_b": [[0, 20], [60, 90]]}, + dim="generator", + ), + ), + ( + y, + segments( + {"gen_a": [[0, 5], [20, 80]], "gen_b": [[0, 8], [30, 70]]}, + dim="generator", + ), + ), + ) + binary = m.variables[f"pwl0{PWL_SEGMENT_BINARY_SUFFIX}"] + lam = m.variables[f"pwl0{PWL_LAMBDA_SUFFIX}"] + assert "generator" in binary.dims + assert "generator" in lam.dims + + def test_three_variables(self) -> None: + """Disjunctive with 3 variables creates single link constraint.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + m.add_piecewise_formulation( + (x, segments([[0, 10], [50, 100]])), + (y, segments([[0, 5], [20, 80]])), + (z, segments([[0, 3], [15, 60]])), + ) + assert f"pwl0{PWL_SEGMENT_BINARY_SUFFIX}" in m.variables + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + # Single link constraint with _pwl_var dimension + link = m.constraints[f"pwl0{PWL_LINK_SUFFIX}"] + assert "_pwl_var" in [str(d) for d in link.dims] + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + def test_sign_le_respected_by_solver(self) -> None: + """ + Disjunctive + sign='<=' must actually bound the solved output + (not just structurally wire up the output link). + """ + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=40, name="y") + # Two segments forming a concave profile: (0,0)→(10,20), (10,20)→(20,30) + m.add_piecewise_formulation( + (y, segments([[0.0, 20.0], [20.0, 30.0]]), "<="), + (x, segments([[0.0, 10.0], [10.0, 20.0]])), + ) + m.add_constraints(x == 15) + m.add_objective(-y) # maximise y + m.solve() + # f(15) = 20 + (30-20)*0.5 = 25 + assert m.solution["y"].item() == pytest.approx(25.0, abs=1e-3) + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + @pytest.mark.parametrize( + "x_fix, expected_y", + [ + # Segment 0: (0, 0) → (5, 10), slope 2. At x=2.5, interp y = 5.0. + (2.5, 5.0), + (0.0, 0.0), # segment 0 left edge + (5.0, 10.0), # segment 0 right edge + # Segment 1: (15, 20) → (25, 35), slope 1.5. At x=20, interp y = 27.5. + (20.0, 27.5), + (15.0, 20.0), # segment 1 left edge + (25.0, 35.0), # segment 1 right edge + ], + ) + def test_sign_le_hits_correct_segment( + self, x_fix: float, expected_y: float + ) -> None: + """ + Disjunctive + sign='<=' picks the **right** segment's interpolation. + + With two segments of different slopes, the bound at ``x_fix`` + depends on which segment ``x_fix`` falls in. The solver must + select the binary for that segment and bound ``y`` by *that* + segment's interpolation, not the other. Probes the binary- + select + signed-output-link combination. + """ + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=50, name="y") + m.add_piecewise_formulation( + (y, segments([[0.0, 10.0], [20.0, 35.0]]), "<="), # two slopes: 2 and 1.5 + (x, segments([[0.0, 5.0], [15.0, 25.0]])), + ) + m.add_constraints(x == x_fix) + m.add_objective(-y) + m.solve() + assert m.solution["y"].item() == pytest.approx(expected_y, abs=1e-3) + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + def test_sign_le_in_forbidden_zone_infeasible(self) -> None: + """X in the gap between segments must be infeasible under sign='<='.""" + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=50, name="y") + m.add_piecewise_formulation( + (y, segments([[0.0, 10.0], [20.0, 35.0]]), "<="), + (x, segments([[0.0, 5.0], [15.0, 25.0]])), + ) + m.add_constraints(x == 10.0) # in the gap (5, 15) + m.add_objective(-y) + status, _ = m.solve() + assert status != "ok" + + +# =========================================================================== +# Validation +# =========================================================================== + + +class TestValidation: + def test_wrong_arg_types_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + with pytest.raises(TypeError, match="at least 2"): + m.add_piecewise_formulation((x, [0, 10, 50])) + + def test_invalid_method_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="method must be"): + m.add_piecewise_formulation( + (x, [0, 10, 50]), + (y, [5, 10, 20]), + method="invalid", # type: ignore + ) + + def test_mismatched_breakpoint_sizes_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="same size"): + m.add_piecewise_formulation( + (x, [0, 10, 50]), + (y, [5, 10]), + ) + + def test_non_tuple_arg_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + with pytest.raises(TypeError, match="tuple"): + m.add_piecewise_formulation(x, [0, 10, 50]) # type: ignore + + +# =========================================================================== +# Name generation +# =========================================================================== + + +class TestNameGeneration: + def test_auto_name(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + m.add_piecewise_formulation((x, [0, 10, 50]), (y, [5, 10, 20])) + m.add_piecewise_formulation((x, [0, 20, 80]), (z, [10, 15, 50])) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + assert f"pwl1{PWL_DELTA_SUFFIX}" in m.variables + + def test_custom_name(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50]), + (y, [5, 10, 20]), + name="my_pwl", + ) + assert f"my_pwl{PWL_DELTA_SUFFIX}" in m.variables + assert f"my_pwl{PWL_LINK_SUFFIX}" in m.constraints + # N-var path uses a single stacked link constraint (no separate y_link) + + +# =========================================================================== +# Broadcasting +# =========================================================================== + + +class TestBroadcasting: + def test_broadcast_over_extra_dims(self) -> None: + m = Model() + gens = pd.Index(["gen_a", "gen_b"], name="generator") + times = pd.Index([0, 1, 2], name="time") + x = m.add_variables(coords=[gens, times], name="x") + y = m.add_variables(coords=[gens, times], name="y") + # Points only have generator dim -> broadcast over time + m.add_piecewise_formulation( + ( + x, + breakpoints( + {"gen_a": [0, 10, 50], "gen_b": [0, 20, 80]}, dim="generator" + ), + ), + ( + y, + breakpoints( + {"gen_a": [0, 5, 30], "gen_b": [0, 8, 50]}, dim="generator" + ), + ), + ) + delta = m.variables[f"pwl0{PWL_DELTA_SUFFIX}"] + assert "generator" in delta.dims + assert "time" in delta.dims + + def test_broadcast_points_dim_order_follows_exprs(self) -> None: + """Expanded dims follow the expression dim order, not set ordering.""" + import xarray as xr + + from linopy.piecewise import BREAKPOINT_DIM, _broadcast_points + + m = Model() + coords = [ + pd.Index(["v0", "v1"], name="alpha"), + pd.Index(["w0", "w1"], name="beta"), + pd.Index([0, 1], name="gamma"), + ] + x = m.add_variables(coords=coords, name="x") + points = xr.DataArray([0, 1, 2, 3], dims=[BREAKPOINT_DIM]) + out = _broadcast_points(points, 1 * x) + assert out.dims == ("alpha", "beta", "gamma", BREAKPOINT_DIM) + + +# =========================================================================== +# NaN masking +# =========================================================================== + + +class TestNaNMasking: + def test_nan_masks_lambda_labels(self) -> None: + """NaN in y_points produces masked labels in SOS2 formulation.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + x_pts = xr.DataArray([0, 10, 50, np.nan], dims=[BREAKPOINT_DIM]) + y_pts = xr.DataArray([0, 5, 20, np.nan], dims=[BREAKPOINT_DIM]) + m.add_piecewise_formulation( + (x, x_pts), + (y, y_pts), + method="sos2", + ) + lam = m.variables[f"pwl0{PWL_LAMBDA_SUFFIX}"] + # First 3 should be valid, last masked + assert (lam.labels.isel({BREAKPOINT_DIM: slice(None, 3)}) != -1).all() + assert int(lam.labels.isel({BREAKPOINT_DIM: 3})) == -1 + + @pytest.mark.parametrize("method", ["sos2", "auto"]) + def test_sos2_interior_nan_raises(self, method: Method) -> None: + """SOS2 with interior NaN breakpoints raises ValueError.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + x_pts = xr.DataArray([0, np.nan, 50, 100], dims=[BREAKPOINT_DIM]) + y_pts = xr.DataArray([0, np.nan, 20, 40], dims=[BREAKPOINT_DIM]) + with pytest.raises(ValueError, match="non-trailing NaN"): + m.add_piecewise_formulation( + (x, x_pts), + (y, y_pts), + method=method, + ) + + +# =========================================================================== +# LP file output +# =========================================================================== + + +class TestLPFileOutput: + def test_sos2_equality(self, tmp_path: Path) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0.0, 10.0, 50.0, 100.0]), + (y, [5.0, 2.0, 20.0, 80.0]), + method="sos2", + ) + m.add_objective(y) + fn = tmp_path / "pwl_eq.lp" + m.to_file(fn, io_api="lp") + content = fn.read_text().lower() + assert "sos" in content + assert "s2" in content + + def test_disjunctive_sos2_and_binary(self, tmp_path: Path) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=100) + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, segments([[0.0, 10.0], [50.0, 100.0]])), + (y, segments([[0.0, 5.0], [20.0, 80.0]])), + ) + m.add_objective(y) + fn = tmp_path / "pwl_disj.lp" + m.to_file(fn, io_api="lp") + content = fn.read_text().lower() + assert "s2" in content + assert "binary" in content or "binaries" in content + + +# =========================================================================== +# Solver integration -- SOS2 capable +# =========================================================================== + + +@pytest.mark.skipif(len(_sos2_solvers) == 0, reason="No solver with SOS2 support") +class TestSolverSOS2: + @pytest.fixture(params=_sos2_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + def test_equality_minimize_cost(self, solver_name: str) -> None: + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + cost = m.add_variables(name="cost") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (cost, [0, 10, 50]), + ) + m.add_constraints(x >= 50, name="x_min") + m.add_objective(cost) + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(x.solution.values, 50, atol=1e-4) + np.testing.assert_allclose(cost.solution.values, 10, atol=1e-4) + + def test_equality_maximize_efficiency(self, solver_name: str) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + eff = m.add_variables(name="eff") + m.add_piecewise_formulation( + (power, [0, 25, 50, 75, 100]), + (eff, [0.7, 0.85, 0.95, 0.9, 0.8]), + ) + m.add_objective(eff, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(power.solution.values, 50, atol=1e-4) + np.testing.assert_allclose(eff.solution.values, 0.95, atol=1e-4) + + def test_disjunctive_solve(self, solver_name: str) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, segments([[0.0, 10.0], [50.0, 100.0]])), + (y, segments([[0.0, 5.0], [20.0, 80.0]])), + ) + m.add_constraints(x >= 60, name="x_min") + m.add_objective(y) + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + # x=60 on second segment: y = 20 + (80-20)/(100-50)*(60-50) = 32 + np.testing.assert_allclose(float(x.solution.values), 60, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 32, atol=1e-4) + + +# =========================================================================== +# Solver integration -- Envelope (any solver) +# =========================================================================== + + +@pytest.mark.skipif(len(_any_solvers) == 0, reason="No solver available") +class TestSolverTangentLines: + @pytest.fixture(params=_any_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + def test_concave_le(self, solver_name: str) -> None: + """Y <= concave f(x), maximize y""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + # Concave: [0,0],[50,40],[100,60] + env = tangent_lines(x, [0, 50, 100], [0, 40, 60]) + m.add_constraints(y <= env, name="pwl") + m.add_constraints(x <= 75, name="x_max") + m.add_constraints(x >= 0, name="x_lo") + m.add_objective(y, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + # At x=75: y = 40 + 0.4*(75-50) = 50 + np.testing.assert_allclose(float(x.solution.values), 75, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 50, atol=1e-4) + + def test_convex_ge(self, solver_name: str) -> None: + """Y >= convex f(x), minimize y""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + # Convex: [0,0],[50,10],[100,60] + env = tangent_lines(x, [0, 50, 100], [0, 10, 60]) + m.add_constraints(y >= env, name="pwl") + m.add_constraints(x >= 25, name="x_min") + m.add_objective(y) + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + # At x=25: y = 0.2*25 = 5 + np.testing.assert_allclose(float(x.solution.values), 25, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 5, atol=1e-4) + + def test_slopes_equivalence(self, solver_name: str) -> None: + """Same model with y_points vs slopes produces identical solutions.""" + # Model 1: direct y_points + m1 = Model() + x1 = m1.add_variables(lower=0, upper=100, name="x") + y1 = m1.add_variables(name="y") + env1 = tangent_lines(x1, [0, 50, 100], [0, 40, 60]) + m1.add_constraints(y1 <= env1, name="pwl") + m1.add_constraints(x1 <= 75, name="x_max") + m1.add_constraints(x1 >= 0, name="x_lo") + m1.add_objective(y1, sense="max") + s1, _ = m1.solve(solver_name=solver_name) + + # Model 2: slopes + m2 = Model() + x2 = m2.add_variables(lower=0, upper=100, name="x") + y2 = m2.add_variables(name="y") + from linopy import Slopes + + env2 = tangent_lines( + x2, + [0, 50, 100], + Slopes([0.8, 0.4], y0=0).to_breakpoints([0, 50, 100]), + ) + m2.add_constraints(y2 <= env2, name="pwl") + m2.add_constraints(x2 <= 75, name="x_max") + m2.add_constraints(x2 >= 0, name="x_lo") + m2.add_objective(y2, sense="max") + s2, _ = m2.solve(solver_name=solver_name) + + assert s1 == "ok" + assert s2 == "ok" + np.testing.assert_allclose( + float(y1.solution.values), float(y2.solution.values), atol=1e-4 + ) + + +# =========================================================================== +# Active parameter (commitment binary) +# =========================================================================== + + +class TestActiveParameter: + """Tests for the ``active`` parameter in piecewise constraints.""" + + def test_incremental_creates_active_bound(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [0, 10, 50, 100]), + (y, [5, 10, 20, 80]), + active=u, + method="incremental", + ) + assert f"pwl0{PWL_ACTIVE_BOUND_SUFFIX}" in m.constraints + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + + def test_active_none_is_default(self) -> None: + """Without active, formulation is identical to before.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (x, [0, 10, 50]), + (y, [0, 5, 30]), + method="incremental", + ) + assert f"pwl0{PWL_ACTIVE_BOUND_SUFFIX}" not in m.constraints + + def test_active_with_linear_expression(self) -> None: + """Active can be a LinearExpression, not just a Variable.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=1 * u, + method="incremental", + ) + assert f"pwl0{PWL_ACTIVE_BOUND_SUFFIX}" in m.constraints + + +# =========================================================================== +# Solver integration -- active parameter +# =========================================================================== + + +@pytest.mark.skipif(len(_any_solvers) == 0, reason="No solver available") +class TestSolverActive: + @pytest.fixture(params=_any_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + def test_incremental_active_on(self, solver_name: str) -> None: + """When u=1 (forced on), normal PWL domain is active.""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=u, + method="incremental", + ) + m.add_constraints(u >= 1, name="force_on") + m.add_constraints(x >= 50, name="x_min") + m.add_objective(y) + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.values), 50, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 10, atol=1e-4) + + def test_incremental_active_off(self, solver_name: str) -> None: + """When u=0 (forced off), x and y must be zero.""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=u, + method="incremental", + ) + m.add_constraints(u <= 0, name="force_off") + m.add_objective(y, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.values), 0, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 0, atol=1e-4) + + def test_incremental_nonzero_base_active_off(self, solver_name: str) -> None: + """ + Non-zero base (x0=20, y0=5) with u=0 must still force zero. + + Tests the x0*u / y0*u base term multiplication -- would fail if + base terms aren't multiplied by active. + """ + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [20, 60, 100]), + (y, [5, 20, 50]), + active=u, + method="incremental", + ) + m.add_constraints(u <= 0, name="force_off") + m.add_objective(y, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.values), 0, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 0, atol=1e-4) + + def test_unit_commitment_pattern(self, solver_name: str) -> None: + """Solver decides to commit: verifies correct fuel at operating point.""" + m = Model() + p_min, p_max = 20.0, 100.0 + fuel_at_pmin, fuel_at_pmax = 10.0, 60.0 + + power = m.add_variables(lower=0, upper=p_max, name="power") + fuel = m.add_variables(name="fuel") + u = m.add_variables(binary=True, name="commit") + + m.add_piecewise_formulation( + (power, [p_min, p_max]), + (fuel, [fuel_at_pmin, fuel_at_pmax]), + active=u, + method="incremental", + ) + m.add_constraints(power >= 50, name="demand") + m.add_objective(fuel + 5 * u) + + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(u.solution.values), 1, atol=1e-4) + np.testing.assert_allclose(float(power.solution.values), 50, atol=1e-4) + # fuel = 10 + (60-10)/(100-20) * (50-20) = 28.75 + np.testing.assert_allclose(float(fuel.solution.values), 28.75, atol=1e-4) + + def test_multi_dimensional_solver(self, solver_name: str) -> None: + """Per-entity on/off: gen_a on at x=50, gen_b off at x=0.""" + m = Model() + gens = pd.Index(["a", "b"], name="gen") + x = m.add_variables(lower=0, upper=100, coords=[gens], name="x") + y = m.add_variables(coords=[gens], name="y") + u = m.add_variables(binary=True, coords=[gens], name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=u, + method="incremental", + ) + m.add_constraints(u.sel(gen="a") >= 1, name="a_on") + m.add_constraints(u.sel(gen="b") <= 0, name="b_off") + m.add_constraints(x.sel(gen="a") >= 50, name="a_min") + m.add_objective(y.sum()) + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.sel(gen="a")), 50, atol=1e-4) + np.testing.assert_allclose(float(y.solution.sel(gen="a")), 10, atol=1e-4) + np.testing.assert_allclose(float(x.solution.sel(gen="b")), 0, atol=1e-4) + np.testing.assert_allclose(float(y.solution.sel(gen="b")), 0, atol=1e-4) + + +@pytest.mark.skipif(len(_sos2_solvers) == 0, reason="No SOS2-capable solver") +class TestSolverActiveSOS2: + @pytest.fixture(params=_sos2_solvers) + def solver_name(self, request: pytest.FixtureRequest) -> str: + return request.param + + def test_sos2_active_off(self, solver_name: str) -> None: + """SOS2: u=0 forces sum(lambda)=0, collapsing x=0, y=0.""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, [0, 50, 100]), + (y, [0, 10, 50]), + active=u, + method="sos2", + ) + m.add_constraints(u <= 0, name="force_off") + m.add_objective(y, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.values), 0, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 0, atol=1e-4) + + def test_disjunctive_active_off(self, solver_name: str) -> None: + """Disjunctive: u=0 forces sum(z_k)=0, collapsing x=0, y=0.""" + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + m.add_piecewise_formulation( + (x, segments([[0.0, 10.0], [50.0, 100.0]])), + (y, segments([[0.0, 5.0], [20.0, 80.0]])), + active=u, + ) + m.add_constraints(u <= 0, name="force_off") + m.add_objective(y, sense="max") + status, _ = m.solve(solver_name=solver_name) + assert status == "ok" + np.testing.assert_allclose(float(x.solution.values), 0, atol=1e-4) + np.testing.assert_allclose(float(y.solution.values), 0, atol=1e-4) + + +# =========================================================================== +# N-variable path +# =========================================================================== + + +class TestNVariable: + """Tests for the N-variable tuple-based piecewise constraint API.""" + + def test_sos2_creates_lambda_and_link(self) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(name="fuel") + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + (fuel, [0.0, 20.0, 60.0]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_CONVEX_SUFFIX}" in m.constraints + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + + def test_incremental_creates_delta(self) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(name="fuel") + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + (fuel, [0.0, 20.0, 60.0]), + method="incremental", + ) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + + def test_auto_selects_method(self) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(name="fuel") + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + (fuel, [0.0, 20.0, 60.0]), + ) + # Auto should select incremental for monotonic breakpoints + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + + def test_single_pair_raises(self) -> None: + m = Model() + power = m.add_variables(name="power") + with pytest.raises(TypeError, match="at least 2"): + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + ) + + def test_three_variables(self) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(name="fuel") + heat = m.add_variables(name="heat") + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + (fuel, [0.0, 20.0, 60.0]), + (heat, [0.0, 30.0, 80.0]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + # link constraint should have _pwl_var dimension + link = m.constraints[f"pwl0{PWL_LINK_SUFFIX}"] + assert "_pwl_var" in link.labels.dims + + def test_custom_name(self) -> None: + m = Model() + power = m.add_variables(lower=0, upper=100, name="power") + fuel = m.add_variables(name="fuel") + m.add_piecewise_formulation( + (power, [0.0, 50.0, 100.0]), + (fuel, [0.0, 20.0, 60.0]), + name="chp", + ) + assert f"chp{PWL_DELTA_SUFFIX}" in m.variables + + +# =========================================================================== +# Additional validation and edge-case coverage +# =========================================================================== + + +class TestValidationEdgeCases: + def test_non_1d_sequence_raises(self) -> None: + """breakpoints() with a 2D nested list raises ValueError.""" + with pytest.raises(ValueError, match="1D sequence"): + breakpoints([[1, 2], [3, 4]]) + + def test_breakpoints_no_values_raises(self) -> None: + """breakpoints() with no positional argument raises TypeError.""" + with pytest.raises(TypeError): + breakpoints() # type: ignore[call-arg] + + def test_non_numeric_breakpoint_coords_raises(self) -> None: + """SOS2 with string breakpoint coords raises ValueError.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + x_pts = xr.DataArray( + [0, 10, 50], + dims=[BREAKPOINT_DIM], + coords={BREAKPOINT_DIM: ["a", "b", "c"]}, + ) + y_pts = xr.DataArray( + [0, 5, 20], + dims=[BREAKPOINT_DIM], + coords={BREAKPOINT_DIM: ["a", "b", "c"]}, + ) + with pytest.raises(ValueError, match="numeric coordinates"): + m.add_piecewise_formulation( + (x, x_pts), + (y, y_pts), + method="sos2", + ) + + def test_unordered_sos2_breakpoint_coords_raise(self) -> None: + """SOS2 breakpoint coords define adjacency and must follow data order.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + x_pts = xr.DataArray( + [0, 1, 2], + dims=[BREAKPOINT_DIM], + coords={BREAKPOINT_DIM: [0, 2, 1]}, + ) + y_pts = xr.DataArray( + [0, 100, 0], + dims=[BREAKPOINT_DIM], + coords={BREAKPOINT_DIM: [0, 2, 1]}, + ) + with pytest.raises(ValueError, match="strictly increasing"): + m.add_piecewise_formulation((x, x_pts), (y, y_pts), method="sos2") + + def test_breakpoint_entity_coords_must_match_expression_coords(self) -> None: + """Entity coords on breakpoints must not silently misalign with variables.""" + m = Model() + entities = pd.Index(["a", "b"], name="entity") + x = m.add_variables(coords=[entities], name="x") + y = m.add_variables(coords=[entities], name="y") + x_pts = xr.DataArray( + [[0, 10], [0, 10]], + dims=["entity", BREAKPOINT_DIM], + coords={"entity": ["a", "b"], BREAKPOINT_DIM: [0, 1]}, + ) + y_pts = xr.DataArray( + [[0, 5], [0, 5]], + dims=["entity", BREAKPOINT_DIM], + coords={"entity": ["b", "c"], BREAKPOINT_DIM: [0, 1]}, + ) + with pytest.raises(ValueError, match="coordinates"): + m.add_piecewise_formulation((x, x_pts), (y, y_pts), method="sos2") + + def test_missing_breakpoint_dim_on_second_arg_raises(self) -> None: + """Second breakpoint array missing breakpoint dim raises.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + good = xr.DataArray([0, 10, 50], dims=[BREAKPOINT_DIM]) + bad = xr.DataArray([0, 5, 20], dims=["wrong"]) + with pytest.raises(ValueError, match="missing"): + m.add_piecewise_formulation((x, good), (y, bad)) + + def test_segment_dim_mismatch_raises(self) -> None: + """Segment dim on only one breakpoint array raises.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + x_pts = segments([[0, 10], [50, 100]]) + y_pts = breakpoints([0, 5]) # same breakpoint count but no segment dim + with pytest.raises(ValueError, match="segment dimension"): + m.add_piecewise_formulation((x, x_pts), (y, y_pts)) + + def test_disjunctive_three_pairs(self) -> None: + """Disjunctive with 3 pairs works (N-variable).""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + seg = segments([[0, 10], [50, 100]]) + m.add_piecewise_formulation( + (x, seg), + (y, seg), + (z, seg), + ) + assert f"pwl0{PWL_SEGMENT_BINARY_SUFFIX}" in m.variables + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_LINK_SUFFIX}" in m.constraints + + def test_disjunctive_interior_nan_raises(self) -> None: + """Disjunctive with interior NaN raises ValueError.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # 3 breakpoints per segment, NaN in the middle of segment 0 + x_pts = xr.DataArray( + [[0, np.nan, 10], [50, 75, 100]], + dims=[SEGMENT_DIM, BREAKPOINT_DIM], + ) + y_pts = xr.DataArray( + [[0, np.nan, 5], [20, 50, 80]], + dims=[SEGMENT_DIM, BREAKPOINT_DIM], + ) + with pytest.raises(ValueError, match="non-trailing NaN"): + m.add_piecewise_formulation((x, x_pts), (y, y_pts)) + + def test_expression_name_fallback(self) -> None: + """LinExpr (not Variable) gets numeric name in link coords.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # Non-monotonic so auto picks SOS2 (which creates lambda vars) + m.add_piecewise_formulation( + (1.0 * x, [0, 50, 10]), + (1.0 * y, [0, 20, 5]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + + def test_incremental_with_nan_mask(self) -> None: + """Incremental method with trailing NaN creates masked delta vars.""" + m = Model() + gens = pd.Index(["a", "b"], name="gen") + x = m.add_variables(coords=[gens], name="x") + y = m.add_variables(coords=[gens], name="y") + x_pts = breakpoints({"a": [0, 10, 50], "b": [0, 20]}, dim="gen") + y_pts = breakpoints({"a": [0, 5, 20], "b": [0, 8]}, dim="gen") + m.add_piecewise_formulation( + (x, x_pts), + (y, y_pts), + method="incremental", + ) + delta = m.variables[f"pwl0{PWL_DELTA_SUFFIX}"] + assert delta.labels.shape[0] == 2 # 2 generators + + def test_scalar_coord_dropped(self) -> None: + """Scalar coords on breakpoints are dropped before stacking.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + bp = breakpoints([0, 10, 50]) + bp_with_scalar = bp.assign_coords(extra=42) + m.add_piecewise_formulation( + (x, bp_with_scalar), + (y, [0, 5, 20]), + method="sos2", + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + + +# =========================================================================== +# Sign parameter (inequality bounds) +# =========================================================================== + + +@pytest.fixture +def nan_padded_pwl_model() -> Callable[[Method], Model]: + """Factory: NaN-padded per-entity piecewise model parametrized by method.""" + from linopy.piecewise import breakpoints + + def _build(method: Method) -> Model: + bp_y = pd.DataFrame([[0, 20, 30, 35], [0, 10, 15, np.nan]], index=["a", "b"]) + bp_x = pd.DataFrame([[0, 10, 20, 30], [0, 5, 15, np.nan]], index=["a", "b"]) + + m = Model() + coord = pd.Index(["a", "b"], name="entity") + x = m.add_variables(lower=0, upper=20, coords=[coord], name="x") + y = m.add_variables(lower=0, upper=40, coords=[coord], name="y") + m.add_piecewise_formulation( + (y, breakpoints(bp_y, dim="entity"), "<="), + (x, breakpoints(bp_x, dim="entity")), + method=method, + ) + m.add_constraints(x.sel(entity="b") == 10) + m.add_objective(-y.sel(entity="b")) + return m + + return _build + + +class TestSignParameter: + """Tests for per-tuple sign on add_piecewise_formulation.""" + + def test_default_is_equality(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation((x, [0, 10, 50]), (y, [0, 5, 20])) + # no output_link for equality — single stacked link only + assert f"pwl0{PWL_OUTPUT_LINK_SUFFIX}" not in m.constraints + + def test_invalid_per_tuple_sign_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="sign must be"): + m.add_piecewise_formulation((x, [0, 10], "!"), (y, [0, 5])) # type: ignore + + def test_two_bounded_tuples_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="At most one tuple"): + m.add_piecewise_formulation((x, [0, 10], "<="), (y, [0, 5], ">=")) + + def test_three_tuples_with_inequality_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + with pytest.raises(ValueError, match="3\\+ tuples"): + m.add_piecewise_formulation( + (x, [0, 10], "<="), + (y, [0, 5]), + (z, [0, 1]), + ) + + def test_bounded_tuple_in_second_position(self) -> None: + """User's tuple order is preserved — bounded tuple need not be first.""" + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=40, name="y") + f = m.add_piecewise_formulation( + (x, [0, 10, 20, 30]), + (y, [0, 20, 30, 35], "<="), + ) + # LP fast-path still triggers regardless of tuple position + assert f.method == "lp" + + def test_lp_with_equality_raises(self) -> None: + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="method='lp'"): + m.add_piecewise_formulation((x, [0, 10, 50]), (y, [0, 5, 20]), method="lp") + + def test_auto_picks_lp_for_concave_le(self) -> None: + """Concave curve + sign='<=' + auto → LP tangent lines (no aux vars).""" + m = Model() + power = m.add_variables(lower=0, upper=30, name="power") + fuel = m.add_variables(lower=0, upper=40, name="fuel") + # Concave: slopes 2, 1, 0.5 + m.add_piecewise_formulation( + (fuel, [0, 20, 30, 35], "<="), + (power, [0, 10, 20, 30]), + ) + assert f"pwl0{PWL_CHORD_SUFFIX}" in m.constraints + assert f"pwl0{PWL_DOMAIN_LO_SUFFIX}" in m.constraints + assert f"pwl0{PWL_DOMAIN_HI_SUFFIX}" in m.constraints + # No SOS2 lambdas for LP + assert f"pwl0{PWL_LAMBDA_SUFFIX}" not in m.variables + + def test_auto_picks_lp_for_convex_ge(self) -> None: + """Convex curve + sign='>=' + auto → LP tangent lines.""" + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=100, name="y") + # Convex: slopes 1, 2, 3 + m.add_piecewise_formulation( + (y, [0, 10, 30, 60], ">="), + (x, [0, 10, 20, 30]), + ) + assert f"pwl0{PWL_CHORD_SUFFIX}" in m.constraints + + def test_auto_falls_back_to_sos2_for_nonmonotonic(self) -> None: + """Non-monotonic x + sign='<=' + auto → SOS2 with signed output link.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # Non-monotonic x + m.add_piecewise_formulation( + (y, [0, 5, 2, 20], "<="), + (x, [0, 10, 5, 50]), + ) + assert f"pwl0{PWL_LAMBDA_SUFFIX}" in m.variables + assert f"pwl0{PWL_OUTPUT_LINK_SUFFIX}" in m.constraints + + def test_auto_concave_ge_falls_back_from_lp(self) -> None: + """Concave + sign='>=' is LP-loose → auto must not pick LP.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + f = m.add_piecewise_formulation( + (y, [0, 20, 30, 35], ">="), # concave + (x, [0, 10, 20, 30]), + ) + assert f.method != "lp" # fallback (sos2 or incremental) + + def test_auto_convex_le_falls_back_from_lp(self) -> None: + """Convex + sign='<=' is LP-loose → auto must not pick LP.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + f = m.add_piecewise_formulation( + (y, [0, 10, 30, 60], "<="), # convex + (x, [0, 10, 20, 30]), + ) + assert f.method != "lp" + + def test_lp_concave_ge_raises(self) -> None: + """Explicit LP + sign='>=' on concave curve is loose → raise.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="convex"): + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], ">="), # concave + (x, [0, 10, 20, 30]), + method="lp", + ) + + def test_lp_nonmatching_convexity_raises(self) -> None: + """Explicit LP with sign='<=' on a convex curve → error.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + # Convex curve, sign='<=' mismatch + with pytest.raises(ValueError, match="concave"): + m.add_piecewise_formulation( + (y, [0, 10, 30, 60], "<="), # convex + (x, [0, 10, 20, 30]), + method="lp", + ) + + def test_sos2_sign_le_has_output_link(self) -> None: + """Explicit SOS2 with sign='<=' gets a signed output link.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + method="sos2", + ) + link = m.constraints[f"pwl0{PWL_OUTPUT_LINK_SUFFIX}"] + assert (link.sign == "<=").all().item() + + def test_incremental_sign_le(self) -> None: + """Incremental method honours sign on output link.""" + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + method="incremental", + ) + assert f"pwl0{PWL_DELTA_SUFFIX}" in m.variables + link = m.constraints[f"pwl0{PWL_OUTPUT_LINK_SUFFIX}"] + assert (link.sign == "<=").all().item() + + def test_lp_consistency_with_sos2(self) -> None: + """LP and SOS2 give the same fuel at a fixed power (within domain).""" + x_pts = [0, 10, 20, 30] + y_pts = [0, 20, 30, 35] # concave + + solutions = {} + methods: list[Method] = ["lp", "sos2", "incremental"] + for method in methods: + m = Model() + power = m.add_variables(lower=0, upper=30, name="power") + fuel = m.add_variables(lower=0, upper=40, name="fuel") + m.add_piecewise_formulation( + (fuel, y_pts, "<="), + (power, x_pts), + method=method, + ) + m.add_constraints(power == 15) + m.add_objective(-fuel) # maximize fuel + m.solve() + solutions[method] = float(m.solution["fuel"]) + + # all methods should max out at f(15) = 25 + for method, val in solutions.items(): + assert abs(val - 25.0) < 1e-4, f"{method}: got {val}" + + def test_convexity_invariant_to_x_direction(self) -> None: + """Decreasing x must classify the same curve identically to ascending x.""" + m_asc = Model() + xa = m_asc.add_variables(name="x") + ya = m_asc.add_variables(name="y") + f_asc = m_asc.add_piecewise_formulation( + (ya, [0, 20, 30, 35], ">="), + (xa, [0, 10, 20, 30]), + ) + m_desc = Model() + xd = m_desc.add_variables(name="x") + yd = m_desc.add_variables(name="y") + f_desc = m_desc.add_piecewise_formulation( + (yd, [35, 30, 20, 0], ">="), + (xd, [30, 20, 10, 0]), + ) + assert f_asc.convexity == f_desc.convexity == "concave" + # concave + >= must fall back from LP + assert f_asc.method != "lp" + assert f_desc.method != "lp" + + def test_lp_per_entity_nan_padding( + self, nan_padded_pwl_model: Callable[[Method], Model] + ) -> None: + """ + Per-entity NaN-padded breakpoints with method='lp': padded + segments must be masked out so they don't create spurious + ``y ≤ 0`` constraints (bug-2 regression). + """ + m = nan_padded_pwl_model("lp") + m.solve() + # f_b(10) on chord (5,10)→(15,15) is 12.5 + assert abs(float(m.solution.sel({"entity": "b"})["y"]) - 12.5) < 1e-3 + + @pytest.mark.skipif(not _SOS_PATHS, reason="No SOS-capable solver installed") + @pytest.mark.parametrize(("solver", "io_api"), _SOS_PATHS) + def test_sos2_per_entity_nan_padding( + self, + nan_padded_pwl_model: Callable[[Method], Model], + solver: str, + io_api: str, + ) -> None: + """ + Per-entity NaN-padded breakpoints with method='sos2': the SOS + lambda variable's masked entries must flow through both the + direct API (via label→position resolution) and the LP writer + (via masked-member filtering) so the solve returns the same + answer as ``method='lp'``. Regression for #688. + + Parametrized across every SOS-capable solver × io_api so the + bug surfaces no matter which backend handles the SOS section + (gurobi-lp masked the bug on master by silently dropping + unknown ``x-1`` members; cplex-lp and gurobi-direct surfaced + it as a parse / OOB error). + """ + m = nan_padded_pwl_model("sos2") + m.solve(solver_name=solver, io_api=io_api) + # f_b(10) on chord (5,10)→(15,15) is 12.5 — same oracle as lp variant + assert abs(float(m.solution.sel({"entity": "b"})["y"]) - 12.5) < 1e-3 + + def test_lp_rejects_decreasing_x_concave_ge(self) -> None: + """ + Explicit LP on a concave curve with sign='>=' must raise, even + when x is specified in decreasing order (bug-1 regression). + """ + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.raises(ValueError, match="convex"): + m.add_piecewise_formulation( + (y, [35, 30, 20, 0], ">="), # same concave curve + (x, [30, 20, 10, 0]), # decreasing x + method="lp", + ) + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + @pytest.mark.parametrize("method", ["sos2", "incremental"]) + def test_active_off_with_sign_le_leaves_lower_open(self, method: Method) -> None: + """ + Documents the asymmetry between sign='==' and sign='<=' under + active=0: equality forces y=0, but '<=' only bounds y ≤ 0 — the + lower side still comes from the variable's own bounds. Verified + uniform across sos2 and incremental. A future change to add the + complementary bound automatically should flip this test. + """ + m = Model() + x = m.add_variables(lower=-100, upper=100, name="x") + y = m.add_variables(lower=-100, upper=100, name="y") + active = m.add_variables(binary=True, name="active") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + method=method, + active=active, + ) + m.add_constraints(active == 0) + m.add_objective(y) # minimize y + m.solve() + # y hits its own lower bound (not 0) — matches docstring note. + assert m.solution["y"].item() == pytest.approx(-100.0, abs=1e-6) + # Input x is still pinned to 0 by the equality input link. + assert m.solution["x"].item() == pytest.approx(0.0, abs=1e-6) + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + def test_active_off_with_sign_le_and_lower_zero_pins_output(self) -> None: + """ + Docstring recipe: with ``y.lower = 0`` (the common case for + fuel/cost/heat outputs), the sign='<=' + active=0 asymmetry + disappears — the variable bound combined with y ≤ 0 forces + y = 0 automatically. + """ + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=100, name="y") # the recipe + active = m.add_variables(binary=True, name="active") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + method="sos2", + active=active, + ) + m.add_constraints(active == 0) + m.add_objective(y, sense="max") # try to push y up + m.solve() + assert m.solution["y"].item() == pytest.approx(0.0, abs=1e-6) + + @pytest.mark.skipif(not _sos2_solvers, reason="no SOS2-capable solver available") + def test_active_off_with_sign_le_disjunctive(self) -> None: + """Same asymmetry applies to the disjunctive (segments) path.""" + m = Model() + x = m.add_variables(lower=-100, upper=100, name="x") + y = m.add_variables(lower=-100, upper=100, name="y") + active = m.add_variables(binary=True, name="active") + m.add_piecewise_formulation( + (y, segments([[0.0, 20.0], [20.0, 35.0]]), "<="), + (x, segments([[0.0, 10.0], [10.0, 30.0]])), + active=active, + ) + m.add_constraints(active == 0) + m.add_objective(y) + m.solve() + assert m.solution["y"].item() == pytest.approx(-100.0, abs=1e-6) + assert m.solution["x"].item() == pytest.approx(0.0, abs=1e-6) + + def test_lp_active_explicit_raises(self) -> None: + """ + method='lp' + active is ValueError (silently ignoring active + would produce a wrong model). + """ + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + u = m.add_variables(binary=True, name="u") + with pytest.raises(ValueError, match="active"): + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + method="lp", + active=u, + ) + + def test_lp_accepts_linear_curve(self) -> None: + """ + A linear curve is both convex and concave per detection, so + LP must accept it with either sign and build the formulation. + """ + signs: list[Sign] = ["<=", ">="] + for sign in signs: + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=60, name="y") + f = m.add_piecewise_formulation( + (y, [0, 10, 20, 30], sign), # linear (all slopes = 1) + (x, [0, 10, 20, 30]), + method="lp", + ) + assert f.method == "lp" + assert f.convexity == "linear" + + def test_auto_logs_when_lp_is_skipped( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """ + method='auto' on a non-LP-eligible case emits an INFO log + explaining why LP was passed over. + """ + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with caplog.at_level(logging.INFO, logger="linopy.piecewise"): + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], ">="), # concave + sign='>=' → LP skipped + (x, [0, 10, 20, 30]), + ) + assert "LP not applicable" in caplog.text + + @pytest.mark.skipif(not _any_solvers, reason="no solver available") + def test_lp_domain_bound_infeasible_when_x_out_of_range(self) -> None: + """ + LP's x ∈ [x_min, x_max] domain bound bites — forcing x beyond + the breakpoint range must make the model infeasible. + """ + m = Model() + x = m.add_variables(lower=0, upper=100, name="x") + y = m.add_variables(lower=0, upper=100, name="y") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), # x_max = 30 + method="lp", + ) + m.add_constraints(x >= 50) + m.add_objective(-y) + status, _ = m.solve() + assert status != "ok" + + @pytest.mark.skipif(not _any_solvers, reason="no solver available") + def test_lp_domain_uses_paired_valid_breakpoints(self) -> None: + """A trailing NaN in y must also shrink the LP x-domain.""" + m = Model() + x = m.add_variables(lower=0, upper=2, name="x") + y = m.add_variables(lower=0, upper=10, name="y") + m.add_piecewise_formulation( + (y, [0, 1, np.nan], "<="), + (x, [0, 1, 2]), + method="lp", + ) + m.add_constraints(x == 2) + m.add_objective(-y) + status, _ = m.solve() + assert status != "ok" + + @pytest.mark.skipif(not _any_solvers, reason="no solver available") + def test_lp_matches_sos2_on_multi_dim_variables(self) -> None: + """ + LP with an entity dimension beyond BREAKPOINT_DIM must match + the SOS2 solution per entity. + """ + entities = pd.Index(["a", "b"], name="entity") + bp_x = pd.DataFrame([[0, 10, 20, 30], [0, 10, 20, 30]], index=["a", "b"]) + bp_y = pd.DataFrame([[0, 20, 30, 35], [0, 15, 25, 30]], index=["a", "b"]) + ys: dict[str, xr.DataArray] = {} + methods: list[Method] = ["lp", "sos2"] + for method in methods: + m = Model() + x = m.add_variables(lower=0, upper=30, coords=[entities], name="x") + y = m.add_variables(lower=0, upper=40, coords=[entities], name="y") + m.add_piecewise_formulation( + (y, breakpoints(bp_y, dim="entity"), "<="), + (x, breakpoints(bp_x, dim="entity")), + method=method, + ) + m.add_constraints(x.sel(entity="a") == 15) + m.add_constraints(x.sel(entity="b") == 5) + m.add_objective(-y.sum()) + m.solve() + ys[method] = y.solution + for entity in ["a", "b"]: + assert float(ys["lp"].sel(entity=entity)) == pytest.approx( + float(ys["sos2"].sel(entity=entity)), abs=1e-3 + ) + + @pytest.mark.skipif(not _any_solvers, reason="no solver available") + def test_lp_consistency_with_sos2_both_directions(self) -> None: + """ + Extends test_lp_consistency_with_sos2 to also probe the + minimisation side of y ≤ f(x). + """ + x_pts = [0, 10, 20, 30] + y_pts = [0, 20, 30, 35] # concave + methods: list[Method] = ["lp", "sos2"] + for obj_sign in [-1.0, +1.0]: + sols: dict[str, float] = {} + for method in methods: + m = Model() + p = m.add_variables(lower=0, upper=30, name="p") + f = m.add_variables(lower=0, upper=50, name="f") + m.add_piecewise_formulation((f, y_pts, "<="), (p, x_pts), method=method) + m.add_constraints(p == 15) + m.add_objective(obj_sign * f) + m.solve() + sols[method] = float(m.solution["f"]) + assert sols["lp"] == pytest.approx(sols["sos2"], abs=1e-3) + + +def _bp(values: list[float]) -> xr.DataArray: + """Small helper: plain 1-D breakpoint DataArray for convexity tests.""" + return breakpoints(values) + + +class TestDetectConvexity: + """Direct unit tests for the _detect_convexity classifier.""" + + def test_convex(self) -> None: + from linopy.piecewise import _detect_convexity + + x = _bp([0, 1, 2, 3]) + y = _bp([0, 1, 4, 9]) # y = x^2 + assert _detect_convexity(x, y) == "convex" + + def test_concave(self) -> None: + from linopy.piecewise import _detect_convexity + + x = _bp([0, 1, 2, 3]) + y = _bp([0, 1, 1.5, 1.75]) # diminishing returns + assert _detect_convexity(x, y) == "concave" + + def test_linear_exact(self) -> None: + from linopy.piecewise import _detect_convexity + + x = _bp([0, 1, 2, 3]) + y = _bp([0, 2, 4, 6]) + assert _detect_convexity(x, y) == "linear" + + def test_linear_within_tol(self) -> None: + from linopy.piecewise import _detect_convexity + + # Tiny slope wobble within 1e-10 tolerance + x = _bp([0, 1, 2, 3]) + y = _bp([0, 2.0, 4.0 + 1e-12, 6.0 + 2e-12]) + assert _detect_convexity(x, y) == "linear" + + def test_mixed(self) -> None: + from linopy.piecewise import _detect_convexity + + x = _bp([0, 1, 2, 3, 4]) + y = _bp([0, 1, 4, 5, 4]) # convex then concave + assert _detect_convexity(x, y) == "mixed" + + def test_too_few_points_returns_linear(self) -> None: + from linopy.piecewise import _detect_convexity + + # Only two points — no second difference to examine + x = _bp([0, 1]) + y = _bp([0, 2]) + assert _detect_convexity(x, y) == "linear" + + def test_decreasing_x_matches_ascending(self) -> None: + """Reversing the breakpoint order must not change the label.""" + from linopy.piecewise import _detect_convexity + + # convex + assert _detect_convexity(_bp([0, 1, 2, 3]), _bp([0, 1, 4, 9])) == "convex" + assert _detect_convexity(_bp([3, 2, 1, 0]), _bp([9, 4, 1, 0])) == "convex" + # concave + assert ( + _detect_convexity(_bp([0, 10, 20, 30]), _bp([0, 20, 30, 35])) == "concave" + ) + assert ( + _detect_convexity(_bp([30, 20, 10, 0]), _bp([35, 30, 20, 0])) == "concave" + ) + + def test_trailing_nan_ignored(self) -> None: + from linopy.piecewise import _detect_convexity + + # Concave curve with a trailing NaN padding + x = _bp([0.0, 1.0, 2.0, np.nan]) + y = _bp([0.0, 1.0, 1.5, np.nan]) + assert _detect_convexity(x, y) == "concave" + + def test_multi_entity_same_shape(self) -> None: + from linopy.piecewise import _detect_convexity + + # Both rows convex + bp_x = pd.DataFrame([[0, 1, 2, 3], [0, 1, 2, 3]], index=["a", "b"]) + bp_y = pd.DataFrame([[0, 1, 4, 9], [0, 2, 8, 18]], index=["a", "b"]) + assert ( + _detect_convexity( + breakpoints(bp_x, dim="entity"), + breakpoints(bp_y, dim="entity"), + ) + == "convex" + ) + + def test_multi_entity_mixed_direction(self) -> None: + """Same concave curve, one entity ascending, one descending.""" + from linopy.piecewise import _detect_convexity + + bp_x = pd.DataFrame([[0, 10, 20, 30], [30, 20, 10, 0]], index=["a", "b"]) + bp_y = pd.DataFrame([[0, 20, 30, 35], [35, 30, 20, 0]], index=["a", "b"]) + assert ( + _detect_convexity( + breakpoints(bp_x, dim="entity"), + breakpoints(bp_y, dim="entity"), + ) + == "concave" + ) + + def test_multi_entity_mixed_curvatures(self) -> None: + """One convex, one concave across entities → mixed.""" + from linopy.piecewise import _detect_convexity + + bp_x = pd.DataFrame([[0, 1, 2, 3], [0, 1, 2, 3]], index=["a", "b"]) + bp_y = pd.DataFrame([[0, 1, 4, 9], [0, 1, 1.5, 1.75]], index=["a", "b"]) + assert ( + _detect_convexity( + breakpoints(bp_x, dim="entity"), + breakpoints(bp_y, dim="entity"), + ) + == "mixed" + ) + + +# =========================================================================== +# netCDF round-trip +# =========================================================================== + + +class TestPiecewiseNetCDFRoundtrip: + """ + Each case exercises a different combination of persisted fields: + + * ``equality_2var`` — non-empty ``variable_names`` + ``convexity != None`` + * ``bounded_lp`` — empty ``variable_names`` (LP path) + + ``convexity != None`` (the only path that yields ``method='lp'``) + * ``equality_3var`` — non-empty ``variable_names`` + ``convexity is None`` + (3-var formulations don't classify curvature) + """ + + def _build(self, kind: str) -> tuple[Model, str]: + m = Model() + if kind == "equality_2var": + y = m.add_variables(name="y") + x = m.add_variables(lower=0, upper=30, name="x") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35]), (x, [0, 10, 20, 30]), name="pwl" + ) + elif kind == "bounded_lp": + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=40, name="y") + m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + name="pwl", + ) + elif kind == "equality_3var": + x = m.add_variables(name="x") + y = m.add_variables(name="y") + z = m.add_variables(name="z") + m.add_piecewise_formulation( + (x, [0, 30, 60, 100]), + (y, [0, 40, 85, 160]), + (z, [0, 25, 55, 95]), + name="pwl", + ) + else: + raise ValueError(kind) + return m, "pwl" + + @pytest.mark.parametrize("kind", ["equality_2var", "bounded_lp", "equality_3var"]) + def test_formulation_survives_netcdf(self, tmp_path: Path, kind: str) -> None: + from linopy import read_netcdf + from linopy.piecewise import PiecewiseFormulation + + m, name = self._build(kind) + f = m._piecewise_formulations[name] + + path = tmp_path / "model.nc" + m.to_netcdf(path) + f2 = read_netcdf(path)._piecewise_formulations[name] + + # Compare every slot except the back-reference to the model, so this + # test auto-catches any future field that IO forgets to persist. + fields = [s for s in PiecewiseFormulation.__slots__ if s != "model"] + before = {s: getattr(f, s) for s in fields} + after = {s: getattr(f2, s) for s in fields} + assert before == after + + # The reloaded formulation's properties must work — i.e. the model + # back-reference was rebound and the named members exist. + assert list(f2.variables) == list(f.variables) + assert list(f2.constraints) == list(f.constraints) + + +# =========================================================================== +# PiecewiseFormulation API surface +# =========================================================================== + + +class TestPiecewiseFormulationAPI: + def test_variables_constraints_repr(self) -> None: + m = Model() + y = m.add_variables(name="y") + x = m.add_variables(lower=0, upper=30, name="x") + f = m.add_piecewise_formulation( + (y, [0, 20, 30, 35]), + (x, [0, 10, 20, 30]), + name="pwl", + ) + # Properties return live views from the model. + assert set(f.variables) == set(f.variable_names) + assert set(f.constraints) == set(f.constraint_names) + + # __repr__ at minimum names the formulation, the resolved method, + # and lists each generated variable / constraint by name. + r = repr(f) + assert "pwl" in r + assert f.method in r + for vname in f.variable_names: + assert vname in r + for cname in f.constraint_names: + assert cname in r + + def test_repr_lp_has_no_variables_section_entries(self) -> None: + """LP formulation has zero auxiliary variables; __repr__ must still render.""" + m = Model() + x = m.add_variables(lower=0, upper=30, name="x") + y = m.add_variables(lower=0, upper=40, name="y") + f = m.add_piecewise_formulation( + (y, [0, 20, 30, 35], "<="), + (x, [0, 10, 20, 30]), + name="pwl_lp", + ) + assert f.method == "lp" + assert f.variable_names == [] + r = repr(f) + assert "pwl_lp" in r + assert "lp" in r + + +# =========================================================================== +# _lp_eligibility — each branch's user-facing reason string +# =========================================================================== + + +class TestLPEligibilityReasons: + """ + Pin the diagnostic returned by each branch of ``_lp_eligibility``. + These reason strings flow into both the auto-dispatch INFO log and the + explicit ``method='lp'`` ``ValueError``, so each is a piece of the + user-facing API. + + Direct unit test (rather than via ``add_piecewise_formulation``) so the + branches that the public-API short-circuits hide are still covered. + """ + + @staticmethod + def _make_inputs( + x_pts: list[float], + y_pts: list[float], + sign: str, + n_pinned: int = 1, + ) -> _PwlInputs: + from xarray import DataArray + + from linopy.constants import BREAKPOINT_DIM, EQUAL, sign_replace_dict + from linopy.piecewise import _PwlInputs + + # Normalise per the production code so callers can write "==" instead + # of "=" in the parametrize table. + sign = sign_replace_dict.get(sign, sign) + + m = Model() + x = m.add_variables(name=f"x_{id(x_pts)}") + pinned_exprs = [(x * 1.0)] + pinned_bps = [ + DataArray(x_pts, dims=[BREAKPOINT_DIM], coords={BREAKPOINT_DIM: x_pts}) + ] + pinned_coords = ["x"] + for i in range(1, n_pinned): + xi = m.add_variables(name=f"x{i}_{id(x_pts)}") + pinned_exprs.append(xi * 1.0) + pinned_bps.append(pinned_bps[0]) + pinned_coords.append(f"x{i}") + + if sign == EQUAL: + return _PwlInputs( + pinned_exprs=pinned_exprs, + pinned_bps=pinned_bps, + pinned_coords=pinned_coords, + bounded_expr=None, + bounded_bp=None, + bounded_coord=None, + bounded_sign=EQUAL, + bp_mask=None, + ) + + y = m.add_variables(name=f"y_{id(y_pts)}") + return _PwlInputs( + pinned_exprs=pinned_exprs, + pinned_bps=pinned_bps, + pinned_coords=pinned_coords, + bounded_expr=y * 1.0, + bounded_bp=DataArray( + y_pts, dims=[BREAKPOINT_DIM], coords={BREAKPOINT_DIM: x_pts} + ), + bounded_coord="y", + bounded_sign=sign, + bp_mask=None, + ) + + @pytest.mark.parametrize( + ("kwargs", "active", "fragments"), + [ + pytest.param( + { + "x_pts": [0.0, 10.0], + "y_pts": [0.0, 5.0], + "sign": "==", + "n_pinned": 3, + }, + None, + ("3 expressions", "LP supports only 2"), + id="too_many_tuples", + ), + pytest.param( + { + "x_pts": [0.0, 10.0], + "y_pts": [0.0, 5.0], + "sign": "==", + "n_pinned": 2, + }, + None, + ("all tuples are equality",), + id="all_equality", + ), + pytest.param( + { + "x_pts": [0.0, 10.0, 20.0, 30.0], + "y_pts": [0.0, 20.0, 30.0, 35.0], + "sign": "<=", + }, + "stub", + ("active",), + id="active_present", + ), + pytest.param( + { + "x_pts": [0.0, 20.0, 10.0, 30.0], + "y_pts": [0.0, 20.0, 30.0, 35.0], + "sign": "<=", + }, + None, + ("not strictly monotonic",), + id="non_monotonic_x", + ), + pytest.param( + { + "x_pts": [0.0, 10.0, 20.0, 30.0], + "y_pts": [0.0, 5.0, 15.0, 30.0], + "sign": "<=", + }, + None, + ("sign='<='", "concave"), + id="le_on_convex", + ), + pytest.param( + { + "x_pts": [0.0, 10.0, 20.0, 30.0], + "y_pts": [0.0, 20.0, 30.0, 35.0], + "sign": ">=", + }, + None, + ("sign='>='", "convex"), + id="ge_on_concave", + ), + ], + ) + def test_reason_string( + self, + kwargs: dict[str, Any], + active: object, + fragments: tuple[str, ...], + ) -> None: + from linopy.piecewise import _lp_eligibility + + inputs = self._make_inputs(**kwargs) + # ``active`` carrying the literal "stub" string is a sentinel — the + # eligibility check only looks at ``is None`` vs not, so any non-None + # value triggers the rejection branch. + ok, reason = _lp_eligibility(inputs, active) # type: ignore[arg-type] + assert not ok + for frag in fragments: + assert frag in reason, f"missing {frag!r} in reason: {reason!r}" + + def test_eligible_concave_le_returns_ok(self) -> None: + from linopy.piecewise import _lp_eligibility + + inputs = self._make_inputs( + x_pts=[0.0, 10.0, 20.0, 30.0], + y_pts=[0.0, 20.0, 30.0, 35.0], # concave + sign="<=", + ) + ok, reason = _lp_eligibility(inputs, None) + assert ok + assert reason == "" + + +# =========================================================================== +# EvolvingAPIWarning — fires once per session per entry point +# =========================================================================== + + +class TestEvolvingAPIWarning: + @pytest.fixture(autouse=True) + def _reset_dedup(self) -> Generator[None, None, None]: + """ + Warnings dedup is module-global so order between tests would + otherwise matter. Clear before each test. + """ + from linopy.piecewise import _emitted_evolving_warnings + + _emitted_evolving_warnings.clear() + yield + _emitted_evolving_warnings.clear() + + def test_add_piecewise_formulation_warns_first_call(self) -> None: + from linopy import EvolvingAPIWarning + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with pytest.warns(EvolvingAPIWarning, match="add_piecewise_formulation"): + m.add_piecewise_formulation((x, [0, 10]), (y, [0, 5])) + + def test_add_piecewise_formulation_dedups(self) -> None: + from linopy import EvolvingAPIWarning + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", EvolvingAPIWarning) + m.add_piecewise_formulation((x, [0, 10]), (y, [0, 5])) + m.add_piecewise_formulation((x, [0, 10]), (y, [0, 5])) + m.add_piecewise_formulation((x, [0, 10]), (y, [0, 5])) + evolving = [w for w in caught if issubclass(w.category, EvolvingAPIWarning)] + assert len(evolving) == 1 + + def test_tangent_lines_warns_and_dedups_independently(self) -> None: + from linopy import EvolvingAPIWarning + + m = Model() + x = m.add_variables(lower=0, upper=10, name="x") + x_pts = [0.0, 5.0, 10.0] + y_pts = [0.0, 4.0, 5.0] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", EvolvingAPIWarning) + tangent_lines(x, x_pts, y_pts) + tangent_lines(x, x_pts, y_pts) + evolving = [w for w in caught if issubclass(w.category, EvolvingAPIWarning)] + assert len(evolving) == 1 + assert "tangent_lines" in str(evolving[0].message) + + def test_slopes_construction_warns_and_dedups(self) -> None: + """ + ``Slopes(...)`` is part of the same evolving API surface and emits + on construction so that the standalone ``Slopes(...).to_breakpoints(...)`` + path doesn't silently bypass the signal. Per-key dedup keeps it + quiet for repeated use. + """ + from linopy import EvolvingAPIWarning, Slopes + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", EvolvingAPIWarning) + Slopes([1, 2], y0=0) + Slopes([3, 4], y0=5) + Slopes([1, 1, 1], y0=0, align="leading") + evolving = [w for w in caught if issubclass(w.category, EvolvingAPIWarning)] + assert len(evolving) == 1 + assert "Slopes" in str(evolving[0].message) + + def test_slopes_warning_stacklevel_points_to_user_call(self) -> None: + """ + ``Slopes.__post_init__`` emits via a dataclass-generated ``__init__`` + — ``_warn_evolving_api`` needs ``stacklevel=4`` to skip the helper, + ``__post_init__``, and the synthetic init and land on the actual + user line. + """ + from linopy import EvolvingAPIWarning, Slopes + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", EvolvingAPIWarning) + Slopes([1, 2], y0=0) + evolving = [w for w in caught if issubclass(w.category, EvolvingAPIWarning)] + assert len(evolving) == 1 + assert evolving[0].filename.endswith("test_piecewise_constraints.py") + + def test_warning_stacklevel_points_to_user_call(self) -> None: + """ + ``stacklevel=3`` in ``_warn_evolving_api`` should make the warning + report this test file as the source, not the internal helper. + """ + from linopy import EvolvingAPIWarning + + m = Model() + x = m.add_variables(name="x") + y = m.add_variables(name="y") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", EvolvingAPIWarning) + m.add_piecewise_formulation((x, [0, 10]), (y, [0, 5])) + evolving = [w for w in caught if issubclass(w.category, EvolvingAPIWarning)] + assert len(evolving) == 1 + assert evolving[0].filename.endswith("test_piecewise_constraints.py") diff --git a/test/test_piecewise_feasibility.py b/test/test_piecewise_feasibility.py new file mode 100644 index 000000000..ed5dd49b8 --- /dev/null +++ b/test/test_piecewise_feasibility.py @@ -0,0 +1,352 @@ +""" +Strategic feasibility-region equivalence tests for PWL inequality. + +Stress-tests the documented claim that ``add_piecewise_formulation(sign="<=")`` +(or ``">="``) yields the **same feasible region** for ``(x, y)`` regardless +of which method (``lp`` / ``sos2`` / ``incremental``) dispatches the +formulation, on curves where all three are applicable. + +The strong test is :class:`TestRotatedObjective`: for every rotation +``(α, β)``, the support function ``min α·x + β·y`` under the PWL must match +a vertex-enumeration oracle. Equal support functions across enough +directions imply equal (convex) feasible regions. + +:class:`TestDomainBoundary` and :class:`TestPointwiseInfeasibility` add +targeted sanity checks for cases that rotated objectives don't directly +probe (domain-bound enforcement, numerical precision of the curve bound). + +:class:`TestNVariableInequality` covers 3-variable inequality (LP does not +support it — this is SOS2 vs incremental only) and verifies the split: +bounded first tuple, equality on the rest. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, TypeAlias + +import numpy as np +import pytest + +from linopy import Model, available_solvers +from linopy.solver_capabilities import ( + SolverFeature, + get_available_solvers_with_feature, +) +from linopy.variables import Variable + +Sign: TypeAlias = Literal["<=", ">="] +Method: TypeAlias = Literal["lp", "sos2", "incremental"] + +TOL = 1e-5 +X_LO, X_HI = -100.0, 100.0 +Y_LO, Y_HI = -100.0, 100.0 + +_sos2_solvers = get_available_solvers_with_feature( + SolverFeature.SOS_CONSTRAINTS, available_solvers +) +_any_solvers = [ + s for s in ["highs", "gurobi", "glpk", "cplex"] if s in available_solvers +] + +pytestmark = pytest.mark.skipif( + not (_sos2_solvers and _any_solvers), + reason="need an SOS2-capable LP/MIP solver", +) + + +# --------------------------------------------------------------------------- +# Curve definition + oracle +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Curve: + """A piecewise-linear curve + the sign of the bound it carries.""" + + name: str + x_pts: tuple[float, ...] + y_pts: tuple[float, ...] + sign: Sign + + def f(self, x: float) -> float: + """Linear interpolation of ``y`` at ``x`` (ground truth).""" + return float(np.interp(x, self.x_pts, self.y_pts)) + + def vertices( + self, y_lo: float = Y_LO, y_hi: float = Y_HI + ) -> list[tuple[float, float]]: + """ + Vertices of the feasible polygon — used by the oracle. + + The feasible region for ``sign="<="`` is + ``{(x,y) : x_0 ≤ x ≤ x_n, y_lo ≤ y ≤ f(x)}`` — a polygon whose + vertices are the breakpoints (top edges) plus two bottom corners. + For ``sign=">="`` it is the mirror image clipped to ``y_hi``. + """ + verts = list(zip(self.x_pts, self.y_pts)) + bottom_y = y_lo if self.sign == "<=" else y_hi + verts.append((self.x_pts[0], bottom_y)) + verts.append((self.x_pts[-1], bottom_y)) + return verts + + +CURVES: list[Curve] = [ + Curve("concave-smooth", (0, 1, 2, 3, 4), (0, 1.75, 3, 3.75, 4), "<="), + Curve("concave-shifted", (-2, 0, 5, 10), (-5, 0, 3, 4), "<="), + Curve("convex-steep", (0, 1, 2, 3, 4), (0, 1, 4, 9, 16), ">="), + Curve("linear-lte", (0, 1, 2, 3, 4), (10, 12, 14, 16, 18), "<="), + Curve("linear-gte", (0, 1, 2, 3, 4), (10, 12, 14, 16, 18), ">="), + Curve("two-segment", (0, 10, 20), (0, 15, 20), "<="), +] + + +# --------------------------------------------------------------------------- +# Primitives: build a model, solve, assert infeasibility +# --------------------------------------------------------------------------- + + +def build_model(curve: Curve, method: Method) -> tuple[Model, Variable, Variable]: + """Build a fresh model with bounded x, y linked by the PWL formulation.""" + m = Model() + x = m.add_variables(lower=X_LO, upper=X_HI, name="x") + y = m.add_variables(lower=Y_LO, upper=Y_HI, name="y") + m.add_piecewise_formulation( + (y, list(curve.y_pts), curve.sign), + (x, list(curve.x_pts)), + method=method, + ) + return m, x, y + + +def solve_support( + curve: Curve, method: Method, alpha: float, beta: float +) -> tuple[float, float, float]: + """ + Solve ``min α·x + β·y``; return ``(x_sol, y_sol, objective)``. + + The attained *point* is returned alongside the objective because + the point usually reveals the bug (wrong segment, clipped domain, + etc.) more clearly than the objective value alone. + """ + m, x, y = build_model(curve, method) + m.add_objective(alpha * x + beta * y) + status, _ = m.solve() + assert status == "ok", f"{method}/{curve.name}: solve failed at ({alpha}, {beta})" + x_sol = float(m.solution["x"]) + y_sol = float(m.solution["y"]) + return x_sol, y_sol, alpha * x_sol + beta * y_sol + + +def oracle_support(curve: Curve, alpha: float, beta: float) -> float: + """Ground truth ``min α·x + β·y`` over the feasible polygon (vertex min).""" + return min(alpha * vx + beta * vy for vx, vy in curve.vertices()) + + +def assert_infeasible(m: Model, x: Variable, msg: str) -> None: + """Solve with a trivial objective; any non-'ok' status counts as infeasible.""" + m.add_objective(x) # objective is irrelevant — just needs to be set + status, _ = m.solve() + assert status != "ok", msg + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(params=CURVES, ids=lambda c: c.name) +def curve(request: pytest.FixtureRequest) -> Curve: + return request.param + + +@pytest.fixture(params=["lp", "sos2", "incremental"]) +def method(request: pytest.FixtureRequest) -> Method: + return request.param + + +# --------------------------------------------------------------------------- +# Rotated objective — the strong test +# --------------------------------------------------------------------------- + + +_N_DIRECTIONS = 16 +_DIRECTIONS = [ + pytest.param( + float(np.cos(2 * np.pi * i / _N_DIRECTIONS)), + float(np.sin(2 * np.pi * i / _N_DIRECTIONS)), + id=f"{round(360 * i / _N_DIRECTIONS):03d}deg", + ) + for i in range(_N_DIRECTIONS) +] + + +class TestRotatedObjective: + """ + Support-function equivalence: ``min α·x + β·y`` under the PWL + matches the vertex-enumeration oracle for every direction. + + Equal support functions over a dense enough set of directions imply + equal convex feasible regions — the strongest region-identity check. + """ + + @pytest.mark.parametrize("alpha, beta", _DIRECTIONS) + def test_support_matches_oracle( + self, curve: Curve, method: Method, alpha: float, beta: float + ) -> None: + x_sol, y_sol, got = solve_support(curve, method, alpha, beta) + want = oracle_support(curve, alpha, beta) + assert abs(got - want) < TOL, ( + f"\n curve: {curve.name} sign: {curve.sign} method: {method}" + f"\n direction: (α={alpha:+.3f}, β={beta:+.3f})" + f"\n attained point: (x={x_sol:+.6f}, y={y_sol:+.6f})" + f"\n attained obj: {got:+.6f}" + f"\n oracle obj: {want:+.6f}" + f"\n diff: {got - want:+.3e} (TOL={TOL:.1e})" + ) + + +# --------------------------------------------------------------------------- +# Domain boundary — direct probe that x cannot escape [x_min, x_max] +# --------------------------------------------------------------------------- + + +class TestDomainBoundary: + """ + ``x`` outside ``[x_min, x_max]`` is infeasible under all methods. + + LP enforces this with an explicit constraint; SOS2/incremental enforce + it implicitly via ``sum(λ) = 1`` (or the delta ladder). Worth a direct + probe because the two paths are very different implementations. + """ + + def test_below_x_min(self, curve: Curve, method: Method) -> None: + m, x, _ = build_model(curve, method) + m.add_constraints(x == curve.x_pts[0] - 1.0) + assert_infeasible( + m, x, f"{method}/{curve.name}: x < x_min should be infeasible" + ) + + def test_above_x_max(self, curve: Curve, method: Method) -> None: + m, x, _ = build_model(curve, method) + m.add_constraints(x == curve.x_pts[-1] + 1.0) + assert_infeasible( + m, x, f"{method}/{curve.name}: x > x_max should be infeasible" + ) + + +# --------------------------------------------------------------------------- +# Pointwise infeasibility — sanity check that (x, f(x) ± ε) is excluded +# --------------------------------------------------------------------------- + + +class TestPointwiseInfeasibility: + """ + ``y`` pushed past ``f(x)`` in the sign direction is infeasible. + + Rotated objectives probe extremes; this targeted check makes sure the + curve bound is actually a strict inequality at a representative + interior point (catches 'off by one segment' or NaN-mask bugs that + might accidentally allow a small slack). + """ + + def test_just_past_curve(self, curve: Curve, method: Method) -> None: + x_mid = 0.5 * (curve.x_pts[0] + curve.x_pts[-1]) + fx = curve.f(x_mid) + # nudge y past the bound in the forbidden direction + y_bad = fx + 0.01 if curve.sign == "<=" else fx - 0.01 + m, x, y = build_model(curve, method) + m.add_constraints(x == x_mid) + m.add_constraints(y == y_bad) + assert_infeasible( + m, + x, + f"{method}/{curve.name}: (x={x_mid}, y={y_bad}) beyond " + f"f(x)={fx} in direction {curve.sign} should be infeasible", + ) + + +# --------------------------------------------------------------------------- +# Hand-computed anchors — sanity-check the oracle itself +# --------------------------------------------------------------------------- + + +class TestHandComputedAnchors: + """ + A handful of pinpoint tests with hand-calculable expected values. + + The parameterised tests compare the solver against a vertex-enumeration + oracle — if that oracle or ``np.interp`` ever drifted, the tests could + continue to pass in false agreement with a broken oracle. These + anchors assert *concrete numbers* a reader can verify with a + calculator in ten seconds, so any oracle drift would surface here. + + Every curve below is arithmetically trivial. Each expected value has + a one-line comment showing the arithmetic. + """ + + # y = 2x on [0, 5] — linear, trivial. + LINEAR = Curve("y_eq_2x", (0, 1, 2, 3, 4, 5), (0, 2, 4, 6, 8, 10), "<=") + + # concave: (0,0) (1,1) (2,1.5) (3,1.75) — slopes 1, 0.5, 0.25 (classic + # diminishing returns) + CONCAVE = Curve("dim_returns", (0, 1, 2, 3), (0, 1, 1.5, 1.75), "<=") + + # convex y = x² sampled at 0..3 — slopes 1, 3, 5 + CONVEX = Curve("y_eq_x2", (0, 1, 2, 3), (0, 1, 4, 9), ">=") + + # ---- 2-variable ---------------------------------------------------- + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_linear_at_midsegment(self, method: Method) -> None: + """Y ≤ 2x at x=2.5: max y = 5.0 (halfway between (2, 4) and (3, 6)).""" + m, x, y = build_model(self.LINEAR, method) + m.add_constraints(x == 2.5) + m.add_objective(-y) + m.solve() + assert float(m.solution["y"]) == pytest.approx(5.0, abs=TOL) + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_linear_at_breakpoint(self, method: Method) -> None: + """Y ≤ 2x at x=3 (exact breakpoint): max y = 6.0.""" + m, x, y = build_model(self.LINEAR, method) + m.add_constraints(x == 3.0) + m.add_objective(-y) + m.solve() + assert float(m.solution["y"]) == pytest.approx(6.0, abs=TOL) + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_linear_at_x_min(self, method: Method) -> None: + """Y ≤ 2x at x=0 (domain lower bound): max y = 0.0.""" + m, x, y = build_model(self.LINEAR, method) + m.add_constraints(x == 0.0) + m.add_objective(-y) + m.solve() + assert float(m.solution["y"]) == pytest.approx(0.0, abs=TOL) + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_linear_at_x_max(self, method: Method) -> None: + """Y ≤ 2x at x=5 (domain upper bound): max y = 10.0.""" + m, x, y = build_model(self.LINEAR, method) + m.add_constraints(x == 5.0) + m.add_objective(-y) + m.solve() + assert float(m.solution["y"]) == pytest.approx(10.0, abs=TOL) + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_concave_at_midsegment(self, method: Method) -> None: + """Y ≤ f(x) concave at x=1.5: max y = (1 + 1.5)/2 = 1.25.""" + m, x, y = build_model(self.CONCAVE, method) + m.add_constraints(x == 1.5) + m.add_objective(-y) + m.solve() + assert float(m.solution["y"]) == pytest.approx(1.25, abs=TOL) + + @pytest.mark.parametrize("method", ["lp", "sos2", "incremental"]) + def test_convex_ge_at_midsegment(self, method: Method) -> None: + """Y ≥ f(x) convex at x=1.5: min y = (1 + 4)/2 = 2.5.""" + m, x, y = build_model(self.CONVEX, method) + m.add_constraints(x == 1.5) + m.add_objective(y) # minimise — pushes y against the lower bound (curve) + m.solve() + assert float(m.solution["y"]) == pytest.approx(2.5, abs=TOL) diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index fc1bb25f3..40ab8c41a 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -38,7 +38,7 @@ def test_quadratic_expression_from_variables_multiplication( ) -> None: quad_expr = x * y assert isinstance(quad_expr, QuadraticExpression) - assert quad_expr.data.sizes[FACTOR_DIM] == 2 + assert quad_expr.vars.sizes[FACTOR_DIM] == 2 def test_adding_quadratic_expressions(x: Variable) -> None: @@ -52,7 +52,7 @@ def test_quadratic_expression_from_variables_power(x: Variable) -> None: power_expr = x**2 target: QuadraticExpression = x * x # type: ignore assert isinstance(power_expr, QuadraticExpression) - assert power_expr.data.sizes[FACTOR_DIM] == 2 + assert power_expr.vars.sizes[FACTOR_DIM] == 2 assert_quadequal(power_expr, target) assert_quadequal(x.pow(2), target) @@ -63,7 +63,7 @@ def test_quadratic_expression_from_linexpr_multiplication( mult_expr = (10 * x + y) * y target: QuadraticExpression = 10 * x * y + y * y # type: ignore assert isinstance(mult_expr, QuadraticExpression) - assert mult_expr.data.sizes[FACTOR_DIM] == 2 + assert mult_expr.vars.sizes[FACTOR_DIM] == 2 assert mult_expr.nterm == 2 assert_quadequal(mult_expr, target) @@ -71,7 +71,7 @@ def test_quadratic_expression_from_linexpr_multiplication( def test_quadratic_expression_from_linexpr_power(x: Variable) -> None: expr = (10 * x) ** 2 assert isinstance(expr, QuadraticExpression) - assert expr.data.sizes[FACTOR_DIM] == 2 + assert expr.vars.sizes[FACTOR_DIM] == 2 assert expr.nterm == 1 @@ -79,7 +79,7 @@ def test_quadratic_expression_from_linexpr_with_constant_power(x: Variable) -> N expr = (10 * x + 5) ** 2 target: QuadraticExpression = 100 * x * x + 50 * x + 50 * x + 25 # type: ignore assert isinstance(expr, QuadraticExpression) - assert expr.data.sizes[FACTOR_DIM] == 2 + assert expr.vars.sizes[FACTOR_DIM] == 2 assert expr.nterm == 3 assert_quadequal(expr, target) @@ -90,7 +90,7 @@ def test_quadratic_expression_from_linexpr_with_constant_multiplation( expr: QuadraticExpression = (10 * x + 5) * (y + 5) # type: ignore target: QuadraticExpression = 10 * x * y + 5 * y + 50 * x + 25 # type: ignore assert isinstance(expr, QuadraticExpression) - assert expr.data.sizes[FACTOR_DIM] == 2 + assert expr.vars.sizes[FACTOR_DIM] == 2 assert expr.nterm == 3 assert_quadequal(expr, target) @@ -133,7 +133,7 @@ def test_matmul_with_const(x: Variable) -> None: expr2 = expr @ const assert isinstance(expr2, QuadraticExpression) assert expr2.nterm == 2 - assert expr2.data.sizes[FACTOR_DIM] == 2 + assert expr2.vars.sizes[FACTOR_DIM] == 2 def test_quadratic_expression_dot_and_matmul(x: Variable, y: Variable) -> None: @@ -225,7 +225,7 @@ def test_quadratic_expression_wrong_multiplication(x: Variable, y: Variable) -> def merge_raise_deprecation_warning(x: Variable, y: Variable) -> None: expr: QuadraticExpression = x * y # type: ignore with pytest.warns(DeprecationWarning): - merge(expr, expr) # type: ignore + merge(expr, expr) def test_merge_linear_expression_and_quadratic_expression( @@ -238,11 +238,11 @@ def test_merge_linear_expression_and_quadratic_expression( with pytest.raises(ValueError): merge([linexpr, quadexpr], cls=QuadraticExpression) - new_quad_ex = merge([linexpr.to_quadexpr(), quadexpr]) # type: ignore + new_quad_ex = merge([linexpr.to_quadexpr(), quadexpr]) assert isinstance(new_quad_ex, QuadraticExpression) with pytest.warns(DeprecationWarning): - merge(quadexpr, quadexpr, cls=QuadraticExpression) # type: ignore + merge(quadexpr, quadexpr, cls=QuadraticExpression) quadexpr_2 = linexpr.to_quadexpr() merged_expr = merge([quadexpr_2, quadexpr], cls=QuadraticExpression) @@ -360,3 +360,11 @@ def test_power_of_three(x: Variable) -> None: x**3 with pytest.raises(TypeError): (x * x) * (x * x) + + +def test_variable_names(x: Variable, y: Variable) -> None: + expr = 2 * (x * x) + 3 * y + 1 + assert expr.variable_names == {"x", "y"} + + expr = 2 * (x * x) + 1 + assert expr.variable_names == {"x"} diff --git a/test/test_repr.py b/test/test_repr.py index 9a7af8938..ebe9804c7 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -40,6 +40,7 @@ multiindex = pd.MultiIndex.from_product( [list("asdfhjkg"), list("asdfghj")], names=["level_0", "level_1"] ) +multiindex.name = "multi" g = m.add_variables(coords=[multiindex], name="g") # create linear expression for each variable @@ -183,15 +184,15 @@ def test_print_options(obj: Variable | LinearExpression | Constraint) -> None: obj.print(display_max_rows=20) -def test_print_labels() -> None: - m.variables.print_labels([1, 2, 3]) - m.constraints.print_labels([1, 2, 3]) - m.constraints.print_labels([1, 2, 3], display_max_terms=10) +def test_format_labels() -> None: + assert m.variables.format_labels([1, 2, 3]) + assert m.constraints.format_labels([1, 2, 3]) + assert m.constraints.format_labels([1, 2, 3], display_max_terms=10) def test_label_position_too_high() -> None: with pytest.raises(ValueError): - m.variables.print_labels([1000]) + m.variables.format_labels([1000]) def test_model_repr_empty() -> None: diff --git a/test/test_scalar_constraint.py b/test/test_scalar_constraint.py index cf5b37241..4872fada2 100644 --- a/test/test_scalar_constraint.py +++ b/test/test_scalar_constraint.py @@ -6,7 +6,7 @@ import linopy from linopy import GREATER_EQUAL, Model, Variable -from linopy.constraints import AnonymousScalarConstraint, Constraint +from linopy.constraints import AnonymousScalarConstraint, ConstraintBase @pytest.fixture @@ -32,20 +32,20 @@ def test_anonymous_scalar_constraint_type(x: Variable) -> None: def test_simple_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] >= 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] >= 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_compound_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] + x.at[1] >= 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1] >= 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_explicit_simple_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.ConstraintBase) def test_explicit_compound_constraint_type(m: Model, x: Variable) -> None: - c: Constraint = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.Constraint) + c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.ConstraintBase) diff --git a/test/test_semi_continuous.py b/test/test_semi_continuous.py new file mode 100644 index 000000000..f529c4288 --- /dev/null +++ b/test/test_semi_continuous.py @@ -0,0 +1,180 @@ +"""Tests for semi-continuous variable support.""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from linopy import Model, available_solvers + + +def test_add_semi_continuous_variable() -> None: + """Semi-continuous variable is created with correct attributes.""" + m = Model() + x = m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + assert x.attrs["semi_continuous"] is True + assert not x.attrs["binary"] + assert not x.attrs["integer"] + + +def test_semi_continuous_mutual_exclusivity() -> None: + """Semi-continuous cannot be combined with binary or integer.""" + m = Model() + with pytest.raises(ValueError, match="only be one of"): + m.add_variables(lower=1, upper=10, binary=True, semi_continuous=True) + with pytest.raises(ValueError, match="only be one of"): + m.add_variables(lower=1, upper=10, integer=True, semi_continuous=True) + + +def test_semi_continuous_requires_positive_lb() -> None: + """Semi-continuous variables require a positive lower bound.""" + m = Model() + with pytest.raises(ValueError, match="positive scalar lower bound"): + m.add_variables(lower=-1, upper=10, semi_continuous=True) + with pytest.raises(ValueError, match="positive scalar lower bound"): + m.add_variables(lower=0, upper=10, semi_continuous=True) + + +def test_semi_continuous_collection_property() -> None: + """Variables.semi_continuous filters correctly.""" + m = Model() + m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_variables(lower=0, upper=5, name="y") + m.add_variables(name="z", binary=True) + + assert list(m.variables.semi_continuous) == ["x"] + assert "x" not in m.variables.continuous + assert "y" in m.variables.continuous + assert "z" not in m.variables.continuous + + +def test_semi_continuous_repr() -> None: + """Semi-continuous annotation appears in repr.""" + m = Model() + m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + r = repr(m.variables) + assert "semi-continuous" in r + + +def test_semi_continuous_vtypes() -> None: + """Matrices vtypes returns 'S' for semi-continuous variables.""" + m = Model() + m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_variables(lower=0, upper=5, name="y") + m.add_variables(name="z", binary=True) + # Add a dummy constraint and objective so the model is valid + m.add_constraints(m.variables["y"] >= 0, name="dummy") + m.add_objective(m.variables["y"]) + + vtypes = m.matrices.vtypes + # x is semi-continuous -> "S", y is continuous -> "C", z is binary -> "B" + assert "S" in vtypes + assert "C" in vtypes + assert "B" in vtypes + + +def test_semi_continuous_lp_file(tmp_path: Path) -> None: + """LP file contains semi-continuous section.""" + m = Model() + m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_variables(lower=0, upper=5, name="y") + m.add_constraints(m.variables["y"] >= 0, name="dummy") + m.add_objective(m.variables["y"]) + + fn = tmp_path / "test.lp" + m.to_file(fn) + content = fn.read_text() + assert "semi-continuous" in content + + +def test_semi_continuous_with_coords() -> None: + """Semi-continuous variables work with multi-dimensional coords.""" + m = Model() + idx = pd.RangeIndex(5, name="i") + x = m.add_variables(lower=2, upper=20, coords=[idx], name="x", semi_continuous=True) + assert x.attrs["semi_continuous"] is True + assert list(m.variables.semi_continuous) == ["x"] + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed") +def test_semi_continuous_solve_gurobi() -> None: + """ + Semi-continuous variable solves correctly with Gurobi. + + Maximize x subject to x <= 0.5, x semi-continuous in [1, 10]. + Since x can be 0 or in [1, 10], and x <= 0.5 prevents [1, 10], + the optimal x should be 0. + """ + m = Model() + x = m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_constraints(x <= 0.5, name="ub") + m.add_objective(x, sense="max") + m.solve(solver_name="gurobi") + assert m.objective.value is not None + assert np.isclose(m.objective.value, 0, atol=1e-6) + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed") +def test_semi_continuous_solve_gurobi_active() -> None: + """ + Semi-continuous variable takes value in [lb, ub] when beneficial. + + Maximize x subject to x <= 5, x semi-continuous in [1, 10]. + Optimal x should be 5. + """ + m = Model() + x = m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_constraints(x <= 5, name="ub") + m.add_objective(x, sense="max") + m.solve(solver_name="gurobi") + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5, atol=1e-6) + + +def test_unsupported_solver_raises() -> None: + """Solvers without semi-continuous support raise ValueError.""" + m = Model() + m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_constraints(m.variables["x"] <= 5, name="ub") + m.add_objective(m.variables["x"]) + + for solver in ["glpk", "mosek", "mindopt"]: + if solver in available_solvers: + with pytest.raises(ValueError, match="does not support semi-continuous"): + m.solve(solver_name=solver) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +def test_semi_continuous_solve_highs() -> None: + """ + Semi-continuous variable solves correctly with HiGHS. + + Maximize x subject to x <= 0.5, x semi-continuous in [1, 10]. + Since x can be 0 or in [1, 10], and x <= 0.5 prevents [1, 10], + the optimal x should be 0. + """ + m = Model() + x = m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_constraints(x <= 0.5, name="ub") + m.add_objective(x, sense="max") + m.solve(solver_name="highs") + assert m.objective.value is not None + assert np.isclose(m.objective.value, 0, atol=1e-6) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +def test_semi_continuous_solve_highs_active() -> None: + """ + Semi-continuous variable takes value in [lb, ub] when beneficial with HiGHS. + + Maximize x subject to x <= 5, x semi-continuous in [1, 10]. + Optimal x should be 5. + """ + m = Model() + x = m.add_variables(lower=1, upper=10, name="x", semi_continuous=True) + m.add_constraints(x <= 5, name="ub") + m.add_objective(x, sense="max") + m.solve(solver_name="highs") + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5, atol=1e-6) diff --git a/test/test_solution_lookup.py b/test/test_solution_lookup.py new file mode 100644 index 000000000..3f6475a85 --- /dev/null +++ b/test/test_solution_lookup.py @@ -0,0 +1,54 @@ +import numpy as np +from numpy import nan + +from linopy.common import values_to_lookup_array +from linopy.solvers import _solution_from_names + + +class TestValuesToLookupArray: + def test_basic(self) -> None: + arr = values_to_lookup_array(np.array([10.0, 20.0, 30.0]), np.array([0, 1, 2])) + np.testing.assert_array_equal(arr, [10.0, 20.0, 30.0]) + + def test_negative_labels_skipped(self) -> None: + arr = values_to_lookup_array(np.array([nan, 10.0, 20.0]), np.array([-1, 0, 2])) + assert arr[0] == 10.0 + assert np.isnan(arr[1]) + assert arr[2] == 20.0 + + def test_sparse_labels(self) -> None: + arr = values_to_lookup_array(np.array([5.0, 7.0]), np.array([0, 100])) + assert len(arr) == 101 + assert arr[0] == 5.0 + assert arr[100] == 7.0 + assert np.isnan(arr[50]) + + def test_only_negative_labels(self) -> None: + arr = values_to_lookup_array(np.array([nan]), np.array([-1])) + assert len(arr) == 0 + + def test_explicit_size(self) -> None: + arr = values_to_lookup_array(np.array([5.0, 7.0]), np.array([0, 2]), size=5) + assert len(arr) == 5 + assert arr[0] == 5.0 + assert arr[2] == 7.0 + assert np.isnan(arr[1]) + assert np.isnan(arr[3]) + assert np.isnan(arr[4]) + + +class TestSolutionFromNames: + def test_default_names(self) -> None: + arr = _solution_from_names( + np.array([1.0, 2.0, 3.0]), ["x2", "x0", "x1"], size=4 + ) + np.testing.assert_array_equal(arr[:3], [2.0, 3.0, 1.0]) + assert np.isnan(arr[3]) + + def test_explicit_coordinate_names(self) -> None: + arr = _solution_from_names( + np.array([1.0, 2.0]), ["power[1]#5", "power[0]#3"], size=7 + ) + assert arr[3] == 2.0 + assert arr[5] == 1.0 + assert np.isnan(arr[4]) diff --git a/test/test_solvers.py b/test/test_solvers.py index 7f8b01f2b..3c9272453 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -6,12 +6,213 @@ """ from pathlib import Path +from unittest.mock import MagicMock +import numpy as np import pytest from test_io import model # noqa: F401 -from linopy import Model, solvers -from linopy.solver_capabilities import SolverFeature, solver_supports +from linopy import GREATER_EQUAL, Model, solvers +from linopy.constants import Result, Solution, Status +from linopy.constraints import CSRConstraint +from linopy.solver_capabilities import ( + SOLVER_REGISTRY, + SolverFeature, + SolverInfo, + solver_supports, +) +from linopy.solvers import _installed_version_in + + +@pytest.fixture +def lp_only_solver() -> str: + for name in ("glpk", "cbc"): + if name in solvers.available_solvers: + return name + pytest.skip("Need an LP-only solver (glpk or cbc) installed") + + +@pytest.fixture +def simple_model() -> Model: + m = Model(chunk=None) + x = m.add_variables(name="x") + y = m.add_variables(name="y") + m.add_constraints(2 * x + 6 * y, GREATER_EQUAL, 10) + m.add_constraints(4 * x + 2 * y, GREATER_EQUAL, 3) + m.add_objective(2 * y + x) + return m + + +@pytest.mark.parametrize("solver", sorted(set(solvers.licensed_solvers))) +def test_solver_instance_attached_after_solve(simple_model: Model, solver: str) -> None: + simple_model.solve(solver) + assert isinstance(simple_model.solver, solvers.Solver) + assert simple_model.solver.status is not None + assert simple_model.solver.status.is_ok + assert simple_model.solver.solution is not None + assert simple_model.solver_model is simple_model.solver.solver_model + assert simple_model.solver_name == solver + + +@pytest.mark.parametrize("solver", sorted(set(solvers.licensed_solvers))) +def test_result_carries_solver_name(simple_model: Model, solver: str) -> None: + if not solver_supports(solver, SolverFeature.DIRECT_API): + pytest.skip("Solver does not support direct API.") + instance = solvers.Solver.from_name(solver, simple_model, io_api="direct") + result = instance.solve() + assert result.solver_name == solver + + +@pytest.mark.parametrize("solver", sorted(set(solvers.licensed_solvers))) +def test_from_name_then_solve(simple_model: Model, solver: str) -> None: + if not solver_supports(solver, SolverFeature.DIRECT_API): + pytest.skip("Solver does not support direct API.") + built = solvers.Solver.from_name(solver, simple_model, io_api="direct") + assert built.solver_model is not None + result = built.solve() + simple_model.assign_result(result) + + reference = Model(chunk=None) + rx = reference.add_variables(name="x") + ry = reference.add_variables(name="y") + reference.add_constraints(2 * rx + 6 * ry, GREATER_EQUAL, 10) + reference.add_constraints(4 * rx + 2 * ry, GREATER_EQUAL, 3) + reference.add_objective(2 * ry + rx) + reference.solve(solver, io_api="direct") + + assert simple_model.status == "ok" + assert simple_model.objective.value is not None + assert reference.objective.value is not None + assert np.isclose(simple_model.objective.value, reference.objective.value) + + +@pytest.mark.parametrize("solver", sorted(set(solvers.licensed_solvers))) +def test_from_name_set_names_false(simple_model: Model, solver: str) -> None: + if not solver_supports(solver, SolverFeature.DIRECT_API): + pytest.skip("Solver does not support direct API.") + built = solvers.Solver.from_name( + solver, simple_model, io_api="direct", set_names=False + ) + result = built.solve() + status, condition = simple_model.assign_result(result) + + assert status == "ok" + assert condition == "optimal" + assert simple_model.objective.value == pytest.approx(3.3) + assert float(simple_model.variables["x"].solution) == pytest.approx(-0.1) + assert float(simple_model.variables["y"].solution) == pytest.approx(1.7) + + +def test_from_name_unknown_solver_raises(simple_model: Model) -> None: + with pytest.raises(ValueError, match="unknown solver"): + solvers.Solver.from_name("not_a_real_solver", simple_model, io_api="direct") + + +@pytest.mark.skipif( + "highs" not in set(solvers.licensed_solvers), reason="HiGHS is not installed" +) +def test_from_name_applies_solver_options(simple_model: Model) -> None: + built = solvers.Solver.from_name( + "highs", simple_model, io_api="direct", options={"time_limit": 123} + ) + option_status, time_limit = built.solver_model.getOptionValue("time_limit") + assert str(option_status) == "HighsStatus.kOk" + assert time_limit == 123 + + +@pytest.mark.skipif( + "highs" not in set(solvers.licensed_solvers), reason="HiGHS is not installed" +) +def test_solver_state_compatibility_setters(simple_model: Model) -> None: + simple_model.solver = solvers.Solver.from_name( + "highs", simple_model, io_api="direct" + ) + simple_model.solver_model = None + assert simple_model.solver is None + assert simple_model.solver_model is None + assert simple_model.solver_name is None + + simple_model.solver = solvers.Solver.from_name( + "highs", simple_model, io_api="direct" + ) + simple_model.solver_name = None + assert simple_model.solver is None + assert simple_model.solver_model is None + assert simple_model.solver_name is None + + with pytest.raises(AttributeError, match="managed via model.solver"): + simple_model.solver_model = object() + with pytest.raises(AttributeError, match="managed via model.solver"): + simple_model.solver_name = "highs" + + +def test_assign_result_explicit(simple_model: Model) -> None: + x_labels = simple_model.variables["x"].labels.values + y_labels = simple_model.variables["y"].labels.values + primal = np.full(simple_model._xCounter, np.nan) + primal[int(x_labels)] = 1.5 + primal[int(y_labels)] = 2.0 + solution = Solution(primal=primal, objective=5.5) + result = Result( + status=Status.from_termination_condition("optimal"), + solution=solution, + solver_name="mock", + ) + simple_model.solver = None + simple_model.assign_result(result) + assert simple_model.status == "ok" + assert simple_model.termination_condition == "optimal" + assert simple_model.objective.value == 5.5 + assert float(simple_model.variables["x"].solution) == 1.5 + assert float(simple_model.variables["y"].solution) == 2.0 + + +def test_assign_result_with_csr_constraints_avoids_data_reconstruction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = Model(freeze_constraints=True) + x = m.add_variables(coords=[range(3)], name="x") + m.add_constraints(x >= 0, name="c") + con = m.constraints["c"] + assert isinstance(con, CSRConstraint) + + primal = np.arange(m._xCounter, dtype=float) + dual = np.arange(m._cCounter, dtype=float) + 10 + result = Result( + status=Status.from_termination_condition("optimal"), + solution=Solution(primal=primal, dual=dual, objective=1.0), + solver_name="mock", + ) + + def fail_data(self: CSRConstraint) -> None: + raise AssertionError("CSRConstraint.data was accessed") + + monkeypatch.setattr(CSRConstraint, "data", property(fail_data)) + m.assign_result(result) + + np.testing.assert_array_equal(m.variables["x"].solution.values, primal) + np.testing.assert_array_equal(m.constraints["c"].dual.values, dual) + + +@pytest.mark.skipif( + "gurobi" not in set(solvers.licensed_solvers), reason="Gurobi is not installed" +) +def test_gurobi_env_persists_after_solve(simple_model: Model) -> None: + simple_model.solve("gurobi", io_api="direct") + assert simple_model.solver is not None + assert simple_model.solver.env is not None + assert isinstance(simple_model.solver_model.NumVars, int) + + +@pytest.mark.parametrize("solver", sorted(set(solvers.licensed_solvers))) +def test_solver_close_releases_state(simple_model: Model, solver: str) -> None: + simple_model.solve(solver) + solver_instance = simple_model.solver + assert solver_instance is not None + solver_instance.close() + assert solver_instance.solver_model is None + assert solver_instance.env is None + free_mps_problem = """NAME sample_mip ROWS @@ -45,8 +246,20 @@ ENDATA """ +free_lp_problem = """ +Maximize + z: 3 x + 4 y +Subject To + c1: 2 x + y <= 10 + c2: x + 2 y <= 12 +Bounds + 0 <= x + 0 <= y +End +""" -@pytest.mark.parametrize("solver", set(solvers.available_solvers)) + +@pytest.mark.parametrize("solver", set(solvers.licensed_solvers)) def test_free_mps_solution_parsing(solver: str, tmp_path: Path) -> None: try: solver_enum = solvers.SolverName(solver.lower()) @@ -72,7 +285,89 @@ def test_free_mps_solution_parsing(solver: str, tmp_path: Path) -> None: @pytest.mark.skipif( - "gurobi" not in set(solvers.available_solvers), reason="Gurobi is not installed" + "knitro" not in set(solvers.licensed_solvers), reason="Knitro is not installed" +) +def test_knitro_solver_mps(tmp_path: Path) -> None: + """Test Knitro solver with a simple MPS problem.""" + knitro = solvers.Knitro() + + mps_file = tmp_path / "problem.mps" + mps_file.write_text(free_mps_problem) + sol_file = tmp_path / "solution.sol" + + result = knitro.solve_problem(problem_fn=mps_file, solution_fn=sol_file) + + assert result.status.is_ok + assert result.solution is not None + assert result.solution.objective == 30.0 + + +@pytest.mark.skipif( + "knitro" not in set(solvers.licensed_solvers), reason="Knitro is not installed" +) +def test_knitro_solver_for_lp(tmp_path: Path) -> None: + """Test Knitro solver with a simple LP problem.""" + knitro = solvers.Knitro() + + lp_file = tmp_path / "problem.lp" + lp_file.write_text(free_lp_problem) + sol_file = tmp_path / "solution.sol" + + result = knitro.solve_problem(problem_fn=lp_file, solution_fn=sol_file) + + assert result.status.is_ok + assert result.solution is not None + assert result.solution.objective == pytest.approx(26.666, abs=1e-3) + + +@pytest.mark.skipif( + "knitro" not in set(solvers.licensed_solvers), reason="Knitro is not installed" +) +def test_knitro_solver_with_options(tmp_path: Path) -> None: + """Test Knitro solver with custom options.""" + knitro = solvers.Knitro(options={"maxit": 100, "feastol": 1e-6}) + + mps_file = tmp_path / "problem.mps" + mps_file.write_text(free_mps_problem) + sol_file = tmp_path / "solution.sol" + log_file = tmp_path / "knitro.log" + + result = knitro.solve_problem( + problem_fn=mps_file, solution_fn=sol_file, log_fn=log_file + ) + assert result.status.is_ok + + +@pytest.mark.skipif( + "knitro" not in set(solvers.licensed_solvers), reason="Knitro is not installed" +) +def test_knitro_solver_with_model_raises_error(model: Model) -> None: # noqa: F811 + """Test Knitro solver raises NotImplementedError for model-based solving.""" + knitro = solvers.Knitro() + with pytest.raises( + NotImplementedError, match="Direct API not implemented for knitro" + ): + knitro.solve_problem(model=model) + + +@pytest.mark.skipif( + "knitro" not in set(solvers.licensed_solvers), reason="Knitro is not installed" +) +def test_knitro_solver_no_log(tmp_path: Path) -> None: + """Test Knitro solver without log file.""" + knitro = solvers.Knitro(options={"outlev": 0}) + + mps_file = tmp_path / "problem.mps" + mps_file.write_text(free_mps_problem) + sol_file = tmp_path / "solution.sol" + + result = knitro.solve_problem(problem_fn=mps_file, solution_fn=sol_file) + + assert result.status.is_ok + + +@pytest.mark.skipif( + "gurobi" not in set(solvers.licensed_solvers), reason="Gurobi is not installed" ) def test_gurobi_environment_with_dict(model: Model, tmp_path: Path) -> None: # noqa: F811 gurobi = solvers.Gurobi() @@ -98,7 +393,7 @@ def test_gurobi_environment_with_dict(model: Model, tmp_path: Path) -> None: # @pytest.mark.skipif( - "gurobi" not in set(solvers.available_solvers), reason="Gurobi is not installed" + "gurobi" not in set(solvers.licensed_solvers), reason="Gurobi is not installed" ) def test_gurobi_environment_with_gurobi_env(model: Model, tmp_path: Path) -> None: # noqa: F811 import gurobipy as gp @@ -124,3 +419,273 @@ def test_gurobi_environment_with_gurobi_env(model: Model, tmp_path: Path) -> Non gurobi.solve_problem(model=model, solution_fn=sol_file, env=env) assert result.status.is_ok assert log2_file.exists() + + +@pytest.mark.parametrize( + "solver_cls, feature, expected", + [ + (solvers.Gurobi, SolverFeature.SOS_CONSTRAINTS, True), + (solvers.Gurobi, SolverFeature.GPU_ACCELERATION, False), + (solvers.Highs, SolverFeature.SOS_CONSTRAINTS, False), + (solvers.Highs, SolverFeature.SEMI_CONTINUOUS_VARIABLES, True), + (solvers.CBC, SolverFeature.LP_FILE_NAMES, False), + (solvers.CBC, SolverFeature.INTEGER_VARIABLES, True), + (solvers.cuPDLPx, SolverFeature.DIRECT_API, True), + (solvers.cuPDLPx, SolverFeature.GPU_ACCELERATION, True), + (solvers.cuPDLPx, SolverFeature.GPU_ONLY, True), + (solvers.cuPDLPx, SolverFeature.QUADRATIC_OBJECTIVE, False), + (solvers.Gurobi, SolverFeature.GPU_ONLY, False), + (solvers.Xpress, SolverFeature.GPU_ONLY, False), + (solvers.PIPS, SolverFeature.INTEGER_VARIABLES, False), + ], +) +def test_solver_class_supports_feature( + solver_cls: type[solvers.Solver], feature: SolverFeature, expected: bool +) -> None: + assert solver_cls.supports(feature) is expected + + +def test_solver_instance_supports_matches_class() -> None: + feature = SolverFeature.QUADRATIC_OBJECTIVE + assert solvers.Gurobi.supports(feature) is True + if "gurobi" in solvers.licensed_solvers: + assert solvers.Gurobi().supports(feature) is True + + +@pytest.mark.parametrize("solver_name", [n.value for n in solvers.SolverName]) +def test_capability_shim_round_trips(solver_name: str) -> None: + solver_cls = getattr(solvers, solvers.SolverName(solver_name).name) + for feature in SolverFeature: + assert solver_supports(solver_name, feature) == solver_cls.supports(feature) + + +def test_solver_registry_iter_and_index() -> None: + names = list(SOLVER_REGISTRY) + assert "gurobi" in names + for name in names: + info = SOLVER_REGISTRY[name] + assert isinstance(info, SolverInfo) + assert isinstance(info.features, frozenset) + assert info.name == name + + +@pytest.mark.skipif( + "xpress" not in set(solvers.licensed_solvers), reason="Xpress is not installed" +) +def test_xpress_gpu_feature_reflects_installed_version() -> None: + assert solvers.Xpress.supports( + SolverFeature.GPU_ACCELERATION + ) == _installed_version_in("xpress", ">=9.8.0") + + +class TestValidateModelOnBuild: + """Solver._build() runs solver-feature checks regardless of entry point.""" + + def test_quadratic_without_qp_support_raises(self, lp_only_solver: str) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=10) + m.add_objective(x * x, sense="min") + + with pytest.raises(ValueError, match="does not support quadratic"): + solvers.Solver.from_name(lp_only_solver, m, io_api="lp") + + def test_semi_continuous_without_support_raises(self, lp_only_solver: str) -> None: + m = Model() + x = m.add_variables(name="x", lower=1, upper=10, semi_continuous=True) + m.add_objective(x) + + with pytest.raises(ValueError, match="does not support semi-continuous"): + solvers.Solver.from_name(lp_only_solver, m, io_api="lp") + + @pytest.mark.skipif( + "highs" not in solvers.available_solvers, reason="HiGHS not installed" + ) + def test_solve_without_objective_raises(self) -> None: + m = Model() + m.add_variables(name="x", lower=0, upper=10) + # No objective added — both entry points should raise the same error. + with pytest.raises(ValueError, match="No objective has been set"): + solvers.Solver.from_name("highs", m, io_api="lp").solve() + with pytest.raises(ValueError, match="No objective has been set"): + m.solve("highs") + + +class TestSolverDoesNotMutateModel: + """Solver.from_model() must not mutate model state (sanitize stays Model-level).""" + + @pytest.mark.skipif( + "highs" not in solvers.available_solvers, reason="HiGHS not installed" + ) + def test_from_model_leaves_constraints_untouched(self) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=10) + # Constraint with a near-zero coefficient — would be sanitized away if + # the Solver path were sanitizing on build. + m.add_constraints(1e-12 * x + x >= 0, name="c") + m.add_objective(x) + + before = m.constraints["c"].coeffs.values.copy() + solvers.Solver.from_name("highs", m, io_api="lp") + after = m.constraints["c"].coeffs.values + + assert np.allclose(before, after, equal_nan=True), ( + "Solver.from_model() must not mutate model constraints. " + "Sanitization is a Model-level primitive; call " + "model.constraints.sanitize_zeros() / .sanitize_infinities() " + "explicitly before building." + ) + + +class TestAssignResultWiring: + """assign_result(result, solver=...) populates model.solver.""" + + @pytest.mark.skipif( + "highs" not in solvers.available_solvers, reason="HiGHS not installed" + ) + def test_assign_result_with_solver_wires_model_solver(self) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=10) + m.add_objective(x, sense="min") + + assert m.solver is None + solver = solvers.Solver.from_name("highs", m, io_api="lp") + result = solver.solve() + m.assign_result(result, solver=solver) + + assert m.solver is solver + assert m.solver_model is solver.solver_model + + @pytest.mark.skipif( + "highs" not in solvers.available_solvers, reason="HiGHS not installed" + ) + def test_assign_result_without_solver_kwarg_leaves_solver_unset(self) -> None: + m = Model() + x = m.add_variables(name="x", lower=0, upper=10) + m.add_objective(x, sense="min") + + solver = solvers.Solver.from_name("highs", m, io_api="lp") + result = solver.solve() + m.assign_result(result) # no solver kwarg + + assert m.solver is None + + +mosek_installed = pytest.importorskip("mosek", reason="Mosek is not installed") + + +class TestMosekChooseSolution: + @staticmethod + def _make_task_mock( + *, + bas_solsta: object | None = None, + itr_solsta: object | None = None, + itg_solsta: object | None = None, + ) -> MagicMock: + defined = { + mosek_installed.soltype.bas: bas_solsta, + mosek_installed.soltype.itr: itr_solsta, + mosek_installed.soltype.itg: itg_solsta, + } + task = MagicMock() + task.solutiondef.side_effect = lambda st: defined[st] is not None + task.getsolsta.side_effect = lambda st: defined[st] + return task + + @pytest.mark.parametrize( + "kwargs, expected_soltype", + [ + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.optimal, + itr_solsta=mosek_installed.solsta.dual_infeas_cer, + ), + mosek_installed.soltype.bas, + id="prefers_bas_when_itr_is_farkas", + ), + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.optimal, + itr_solsta=mosek_installed.solsta.optimal, + ), + mosek_installed.soltype.itr, + id="prefers_itr_on_tie", + ), + pytest.param( + dict(itr_solsta=mosek_installed.solsta.optimal), + mosek_installed.soltype.itr, + id="only_itr_defined", + ), + pytest.param( + dict(bas_solsta=mosek_installed.solsta.optimal), + mosek_installed.soltype.bas, + id="only_bas_defined", + ), + pytest.param( + dict(), + None, + id="nothing_defined", + ), + pytest.param( + dict(itg_solsta=mosek_installed.solsta.integer_optimal), + mosek_installed.soltype.itg, + id="itg_for_mip", + ), + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.optimal, + itr_solsta=mosek_installed.solsta.optimal, + itg_solsta=mosek_installed.solsta.integer_optimal, + ), + mosek_installed.soltype.itg, + id="itg_wins_over_bas_itr", + ), + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.unknown, + itr_solsta=mosek_installed.solsta.optimal, + ), + mosek_installed.soltype.itr, + id="optimal_itr_over_unknown_bas", + ), + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.optimal, + itr_solsta=mosek_installed.solsta.unknown, + ), + mosek_installed.soltype.bas, + id="optimal_bas_over_unknown_itr", + ), + pytest.param( + dict( + bas_solsta=mosek_installed.solsta.prim_infeas_cer, + itr_solsta=mosek_installed.solsta.dual_infeas_cer, + ), + mosek_installed.soltype.itr, + id="falls_back_to_itr_when_both_non_optimal", + ), + ], + ) + def test_choose_solution( + self, kwargs: dict[str, object], expected_soltype: object + ) -> None: + task = self._make_task_mock(**kwargs) + assert solvers.Mosek._choose_solution(task) is expected_soltype + + @pytest.mark.skipif( + "mosek" not in set(solvers.licensed_solvers), + reason="Mosek is not licensed", + ) + def test_smoke_lp(self) -> None: + import math + + m = Model() + x = m.add_variables(name="x", lower=0) + m.add_constraints(2 * x >= 10, name="c1") + m.add_objective(x) + + result = solvers.Solver.from_name("mosek", m).solve() + + assert result.status.is_ok + assert result.solution is not None + assert math.isfinite(result.solution.objective) + assert result.solution.objective == pytest.approx(5.0, abs=1e-3) diff --git a/test/test_sos_constraints.py b/test/test_sos_constraints.py index 5d94162ec..8160d5242 100644 --- a/test/test_sos_constraints.py +++ b/test/test_sos_constraints.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd import pytest +import xarray as xr from linopy import Model, available_solvers @@ -137,6 +138,98 @@ def test_sos2_binary_maximize_different_coeffs() -> None: assert np.isclose(m.objective.value, 4) +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_to_xpress_emits_sos_constraints() -> None: + m = Model() + segments = pd.Index([0.0, 0.5, 1.0], name="seg") + var = m.add_variables(coords=[segments], name="lambda") + m.add_sos_constraints(var, sos_type=1, sos_dim="seg") + m.add_objective(var.sum()) + + problem = m.to_xpress() + assert problem.attributes.sets == 1 + + +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_to_xpress_emits_grouped_sos_constraints() -> None: + m = Model() + groups = pd.Index(["a", "b"], name="group") + segments = pd.Index([0.0, 0.5, 1.0], name="seg") + var = m.add_variables(coords=[groups, segments], name="lambda") + m.add_sos_constraints(var, sos_type=1, sos_dim="seg") + m.add_objective(var.sum()) + + problem = m.to_xpress() + assert problem.attributes.sets == len(groups) + + +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_sos2_xpress_direct() -> None: + m = Model() + locations = pd.Index([0, 1, 2], name="locations") + build = m.add_variables(coords=[locations], name="build", binary=True) + m.add_sos_constraints(build, sos_type=2, sos_dim="locations") + m.add_objective(build * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="xpress", io_api="direct") + + assert np.isclose(build.solution.values, [0, 1, 1]).all() + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5) + + +@pytest.mark.skipif("xpress" not in available_solvers, reason="Xpress not installed") +def test_qp_sos1_xpress_direct() -> None: + m = Model() + seg = pd.Index([0, 1, 2], name="seg") + x = m.add_variables(lower=0, upper=10, coords=[seg], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="seg") + m.add_constraints(x.sum() >= 5) + + linear_coeffs = xr.DataArray([0.0, -10.0, 0.0], coords=[seg]) + m.add_objective((x * x).sum() + (linear_coeffs * x).sum(), sense="min") + + m.solve(solver_name="xpress", io_api="direct") + + assert np.isclose(x.solution.values, [0, 5, 0]).all() + assert m.objective.value is not None + assert np.isclose(m.objective.value, -25) + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed") +def test_reformulate_sos_true_reformulates_on_native_solver(tmp_path: Path) -> None: + """ + ``reformulate_sos=True`` must reformulate even when the solver supports SOS. + + Asserted against the artifacts ``reformulate_sos_constraints`` writes into + the LP file (the auxiliary binary + cardinality constraint, no ``sos`` + section). The reformulation is undone after solve, so the model itself + looks unchanged — the LP snapshot is the durable evidence. + """ + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum()) + + problem_fn = tmp_path / "problem.lp" + m.solve( + solver_name="gurobi", + io_api="lp", + reformulate_sos=True, + problem_fn=problem_fn, + keep_files=True, + explicit_coordinate_names=True, + ) + + content = problem_fn.read_text() + # SOS got rewritten to binary + linear: no `sos` section, the auxiliary + # binary indicator and cardinality constraint appear instead. + assert "\nsos\n" not in content + assert "_sos_reform_x_y" in content + assert "_sos_reform_x_card" in content + + def test_unsupported_solver_raises_error() -> None: m = Model() locations = pd.Index([0, 1, 2], name="locations") @@ -150,7 +243,7 @@ def test_unsupported_solver_raises_error() -> None: m.solve(solver_name=solver) -def test_to_highspy_raises_not_implemented() -> None: +def test_to_highspy_raises_when_sos_present() -> None: pytest.importorskip("highspy") m = Model() @@ -158,8 +251,5 @@ def test_to_highspy_raises_not_implemented() -> None: build = m.add_variables(coords=[locations], name="build", binary=True) m.add_sos_constraints(build, sos_type=1, sos_dim="locations") - with pytest.raises( - NotImplementedError, - match="SOS constraints are not supported by the HiGHS direct API", - ): + with pytest.raises(ValueError, match="does not support SOS constraints"): m.to_highspy() diff --git a/test/test_sos_masked.py b/test/test_sos_masked.py new file mode 100644 index 000000000..46906ba14 --- /dev/null +++ b/test/test_sos_masked.py @@ -0,0 +1,338 @@ +""" +Regression coverage for SOS constraints on masked variables (#688). + +The bug being pinned here has two related failure modes: + +1. **Position-vs-label**: direct-API builds (gurobi, xpress) pass linopy variable + labels straight to vendor ``addSOS`` as if they were 0-based column positions + in the active-variable array. They only happen to coincide when no variable + in the model is masked anywhere. + +2. **LP file emits ``x-1``**: the LP writer iterates raw label arrays and emits + names like ``x-1`` for masked SOS entries, which LP parsers either reject + outright or (gurobi LP reader) silently corrupt into wrong SOS sets. + +The fixture asymmetric-coefficient design plus three-layer oracle (status, +objective, element-wise solution) ensures any wrong indexing surfaces as a +visible failure rather than a permutation-equivalent silent pass. +""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import Literal + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model, available_solvers +from linopy.solver_capabilities import SolverFeature, solver_supports + +# --------------------------------------------------------------------------- +# Capability-derived solver / io_api parametrization +# --------------------------------------------------------------------------- + +SOS_DIRECT = sorted( + s + for s in available_solvers + if solver_supports(s, SolverFeature.SOS_CONSTRAINTS) + and solver_supports(s, SolverFeature.DIRECT_API) +) +SOS_FILE = sorted( + s for s in available_solvers if solver_supports(s, SolverFeature.SOS_CONSTRAINTS) +) +SOS_PATHS = [ + *[pytest.param(s, "direct", id=f"{s}-direct") for s in SOS_DIRECT], + *[pytest.param(s, "lp", id=f"{s}-lp") for s in SOS_FILE], +] + +# --------------------------------------------------------------------------- +# Analytical optimum (matches solver semantics: list-position adjacency for SOS2) +# --------------------------------------------------------------------------- + + +def _optimize_sos_set( + active_i: list[int], coefs: dict[int, float], sos_type: int +) -> tuple[float, dict[int, float]]: + """ + Closed-form optimum for one SOS set with binary [0,1] members. + + ``active_i`` is the sorted list of active (unmasked) member indices in the + SOS dimension. ``coefs`` maps each active index to its objective coefficient + (minimization). For SOS2, adjacency is list-position adjacency, matching the + semantics of gurobi/xpress ``addSOS``. + """ + if not active_i: + return 0.0, {} + + best_obj = 0.0 + best_sol: dict[int, float] = {} + + # singletons + for i in active_i: + if coefs[i] < best_obj: + best_obj = coefs[i] + best_sol = {i: 1.0} + + if sos_type == 2: + # adjacent pairs in the (sorted-by-weight) list + for k in range(len(active_i) - 1): + i1, i2 = active_i[k], active_i[k + 1] + pair_obj = coefs[i1] + coefs[i2] + if pair_obj < best_obj: + best_obj = pair_obj + best_sol = {i1: 1.0, i2: 1.0} + + return best_obj, best_sol + + +# --------------------------------------------------------------------------- +# Fixture +# --------------------------------------------------------------------------- + +MaskOnSos = Literal[None, "sos_dim", "non_sos_dim", "both_dims"] + + +@pytest.fixture +def sos_masked_model() -> Callable[..., tuple[Model, float, np.ndarray]]: # noqa: E501 + """ + Factory producing SOS{1,2} models with controllable mask placement. + + Objective coefficients along the SOS dim are ``[-1, -2, -3, -4]``, + asymmetric to break permutation symmetry — wrong indexing then produces an + observably different objective AND solution. + + Returns ``(model, expected_obj, expected_sol)``. ``expected_sol`` is shaped + like ``sos_var.solution`` (with ``NaN`` where the mask removes a slot). + """ + + def _build( + sos_type: Literal[1, 2] = 1, + sos_var_2d: bool = False, + mask_on_sos: MaskOnSos = None, + mask_on_other: bool = False, + ) -> tuple[Model, float, np.ndarray]: + if not sos_var_2d and mask_on_sos in ("non_sos_dim", "both_dims"): + raise ValueError(f"mask_on_sos={mask_on_sos!r} requires sos_var_2d=True") + + m = Model() + + # Optional unrelated masked variable: shifts label->position mapping + # for all subsequent variables, exposing the position-vs-label bug. + if mask_on_other: + ck = pd.Index([0, 1, 2, 3], name="k") + m.add_variables( + lower=0, + upper=1, + coords=[ck], + mask=pd.Series([False, True, True, True], index=ck), + name="other", + ) + + ci = pd.Index([0, 1, 2, 3], name="i") + cj = pd.Index([0, 1], name="j") + + # Construct sos_var mask + if mask_on_sos is None: + sos_mask = None + elif mask_on_sos == "sos_dim": + mask_i = np.array([True, True, False, True]) + if sos_var_2d: + sos_mask = xr.DataArray( + np.broadcast_to(mask_i[:, None], (4, 2)).copy(), + coords=[ci, cj], + dims=["i", "j"], + ) + else: + sos_mask = pd.Series(mask_i, index=ci) + elif mask_on_sos == "non_sos_dim": + assert sos_var_2d + mask_j = np.array([False, True]) + sos_mask = xr.DataArray( + np.broadcast_to(mask_j[None, :], (4, 2)).copy(), + coords=[ci, cj], + dims=["i", "j"], + ) + elif mask_on_sos == "both_dims": + assert sos_var_2d + mask_i = np.array([True, True, False, True]) + mask_j = np.array([False, True]) + combined = mask_i[:, None] & mask_j[None, :] + sos_mask = xr.DataArray(combined, coords=[ci, cj], dims=["i", "j"]) + else: + raise ValueError(f"unknown mask_on_sos={mask_on_sos!r}") + + sos_coords = [ci, cj] if sos_var_2d else [ci] + sos_var = m.add_variables( + lower=0, + upper=1, + coords=sos_coords, + mask=sos_mask, + name="sos_var", + ) + m.add_sos_constraints(sos_var, sos_type=sos_type, sos_dim="i") + + # Asymmetric coefficients along the SOS dim; broadcast across j in 2D + coefs_i = np.array([-1.0, -2.0, -3.0, -4.0]) + if sos_var_2d: + coefs = xr.DataArray( + np.broadcast_to(coefs_i[:, None], (4, 2)).copy(), + coords=[ci, cj], + dims=["i", "j"], + ) + else: + coefs = xr.DataArray(coefs_i, coords=[ci], dims=["i"]) + m.add_objective(sos_var * coefs) + + # ------------------------------------------------------------------ + # Compute expected_obj and expected_sol from the same mask logic + # ------------------------------------------------------------------ + coefs_dict = {i: float(coefs_i[i]) for i in range(4)} + + # active_per_j[j] = sorted list of active i for SOS set at j (or for + # the single 1D set we use j=None as a sentinel) + if sos_var_2d: + # Reconstruct the 2D mask (default to all True if none) + if sos_mask is None: + mask_arr = np.ones((4, 2), dtype=bool) + else: + mask_arr = np.asarray(sos_mask.values, dtype=bool) + active_per_j: dict[int | None, list[int]] = { + j: [i for i in range(4) if mask_arr[i, j]] for j in range(2) + } + else: + if sos_mask is None: + active = list(range(4)) + else: + active = [i for i in range(4) if bool(sos_mask.iloc[i])] + active_per_j = {None: active} + + expected_obj = 0.0 + # Build expected_sol with the right shape and NaN-fill masked slots + if sos_var_2d: + expected_sol: np.ndarray = np.full((4, 2), 0.0) + if sos_mask is not None: + mask_arr = np.asarray(sos_mask.values, dtype=bool) + expected_sol[~mask_arr] = np.nan + else: + expected_sol = np.full(4, 0.0) + if sos_mask is not None: + for i in range(4): + if not bool(sos_mask.iloc[i]): + expected_sol[i] = np.nan + + for j_key, active in active_per_j.items(): + obj_j, sol_j = _optimize_sos_set(active, coefs_dict, sos_type) + expected_obj += obj_j + for i, value in sol_j.items(): + if sos_var_2d: + expected_sol[i, j_key] = value + else: + expected_sol[i] = value + + return m, expected_obj, expected_sol + + return _build + + +# --------------------------------------------------------------------------- +# Test matrix: 11 fixture configs × 2 SOS types × (solver, io_api) +# --------------------------------------------------------------------------- + +# Each entry: (sos_var_2d, mask_on_sos, mask_on_other) +FIXTURE_CONFIGS = [ + pytest.param(False, None, False, id="1d-no_mask"), + pytest.param(False, "sos_dim", False, id="1d-mask_sos"), + pytest.param(False, None, True, id="1d-mask_other"), + pytest.param(False, "sos_dim", True, id="1d-mask_both"), + pytest.param(True, None, False, id="2d-no_mask"), + pytest.param(True, "sos_dim", False, id="2d-mask_sos_dim"), + pytest.param(True, "non_sos_dim", False, id="2d-mask_non_sos_dim"), + pytest.param(True, "both_dims", False, id="2d-mask_both_dims"), + pytest.param(True, "sos_dim", True, id="2d-mask_sos_dim+other"), + pytest.param(True, "non_sos_dim", True, id="2d-mask_non_sos_dim+other"), + pytest.param(True, "both_dims", True, id="2d-mask_both_dims+other"), +] + + +@pytest.mark.skipif(not SOS_PATHS, reason="No SOS-capable solver installed") +@pytest.mark.parametrize("sos_type", [1, 2]) +@pytest.mark.parametrize(("solver", "io_api"), SOS_PATHS) +@pytest.mark.parametrize( + ("sos_var_2d", "mask_on_sos", "mask_on_other"), FIXTURE_CONFIGS +) +def test_sos_with_masked_variables( + sos_masked_model: Callable[..., tuple[Model, float, np.ndarray]], + solver: str, + io_api: str, + sos_type: int, + sos_var_2d: bool, + mask_on_sos: MaskOnSos, + mask_on_other: bool, +) -> None: + """ + Three-oracle test: status + objective + element-wise solution. + + Asymmetric objective + element-wise solution check ensures we catch: + - direct-path OOB raises (status != ok) + - LP parser rejections (status != ok) + - silent SOS-set corruption (objective and/or solution differ) + """ + m, expected_obj, expected_sol = sos_masked_model( + sos_type=sos_type, + sos_var_2d=sos_var_2d, + mask_on_sos=mask_on_sos, + mask_on_other=mask_on_other, + ) + m.solve(solver_name=solver, io_api=io_api) + + # Oracle 1: did the solve succeed? + assert m.status == "ok", ( + f"solver={solver} io_api={io_api} status={m.status!r} " + f"termination={m.termination_condition!r}" + ) + + # Oracle 2: is the objective at the analytical optimum? + assert m.objective.value is not None + assert m.objective.value == pytest.approx(expected_obj, abs=1e-5) + + # Oracle 3: are the right slots at the right values? + actual_sol = m.variables["sos_var"].solution.values + np.testing.assert_allclose( + actual_sol, + expected_sol, + atol=1e-5, + equal_nan=True, + err_msg=( + f"sos_var.solution mismatch for solver={solver} io_api={io_api} " + f"sos_type={sos_type} sos_var_2d={sos_var_2d} " + f"mask_on_sos={mask_on_sos!r} mask_on_other={mask_on_other}" + ), + ) + + +def test_sos_to_file_skips_fully_masked_sos_variable(tmp_path: Path) -> None: + """A fully-masked SOS variable writes no LP ``sos`` set entries.""" + m = Model() + ci = pd.Index([0, 1, 2, 3], name="i") + free = m.add_variables(lower=0, upper=1, name="free") + sos_var = m.add_variables( + lower=0, + upper=1, + coords=[ci], + mask=pd.Series(False, index=ci), + name="sos_var", + ) + m.add_sos_constraints(sos_var, sos_type=1, sos_dim="i") + m.add_objective(free) + + fn = tmp_path / "model.lp" + m.to_file(fn) + lp = fn.read_text() + + assert "x-1" not in lp + sos_section = lp.partition("\nsos\n")[2] + assert "S1 ::" not in sos_section diff --git a/test/test_sos_reformulation.py b/test/test_sos_reformulation.py new file mode 100644 index 000000000..0e9dc9da6 --- /dev/null +++ b/test/test_sos_reformulation.py @@ -0,0 +1,1283 @@ +"""Tests for SOS constraint reformulation.""" + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Callable +from pathlib import Path +from typing import Literal, cast + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model, Variable, available_solvers +from linopy.constants import SOS_TYPE_ATTR +from linopy.remote import RemoteHandler +from linopy.sos_reformulation import ( + compute_big_m_values, + reformulate_sos1, + reformulate_sos2, + reformulate_sos_constraints, + undo_sos_reformulation, +) + + +class TestValidateBounds: + """Tests for bound validation in compute_big_m_values.""" + + def test_finite_bounds_pass(self) -> None: + """Finite non-negative bounds should pass validation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + compute_big_m_values(x) # Should not raise + + def test_infinite_upper_bounds_raise(self) -> None: + """Infinite upper bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + with pytest.raises(ValueError, match="infinite upper bounds"): + compute_big_m_values(x) + + def test_negative_lower_bounds_raise(self) -> None: + """Negative lower bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=-1, upper=1, coords=[idx], name="x") + with pytest.raises(ValueError, match="negative lower bounds"): + compute_big_m_values(x) + + def test_mixed_negative_lower_bounds_raise(self) -> None: + """Mixed finite/negative lower bounds should raise ValueError.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, -1, 0]), + upper=np.array([1, 1, 1]), + coords=[idx], + name="x", + ) + with pytest.raises(ValueError, match="negative lower bounds"): + compute_big_m_values(x) + + +class TestComputeBigM: + """Tests for compute_big_m_values.""" + + def test_positive_bounds(self) -> None: + """Test Big-M computation with positive bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=10, coords=[idx], name="x") + M = compute_big_m_values(x) + assert np.allclose(M.values, [10, 10, 10]) + + def test_varying_bounds(self) -> None: + """Test Big-M computation with varying upper bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, 0, 0]), + upper=np.array([1, 2, 3]), + coords=[idx], + name="x", + ) + M = compute_big_m_values(x) + assert np.allclose(M.values, [1, 2, 3]) + + def test_custom_big_m_scalar(self) -> None: + """Test Big-M uses tighter of custom value and bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + M = compute_big_m_values(x) + # M = min(10, 100) = 10 (custom is tighter) + assert np.allclose(M.values, [10, 10, 10]) + + def test_custom_big_m_allows_infinite_bounds(self) -> None: + """Test that custom big_m allows variables with infinite bounds.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + # Should not raise - custom big_m makes result finite + M = compute_big_m_values(x) + assert np.allclose(M.values, [10, 10, 10]) + + +class TestSOS1Reformulation: + """Tests for SOS1 reformulation.""" + + def test_basic_sos1(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_sos1(m, x, "_test_") + m.remove_sos_constraints(x) + + # Check auxiliary variables and constraints were added + assert "_test_x_y" in m.variables + assert "_test_x_upper" in m.constraints + assert "_test_x_card" in m.constraints + + # Binary variable should have same dimensions + y = m.variables["_test_x_y"] + assert y.dims == x.dims + assert y.sizes == x.sizes + + def test_sos1_multidimensional(self) -> None: + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_sos1(m, x, "_test_") + m.remove_sos_constraints(x) + + # Binary variable should have same dimensions + y = m.variables["_test_x_y"] + assert set(y.dims) == {"i", "j"} + + # Cardinality constraint should have reduced dimensions (summed over i) + card_con = m.constraints["_test_x_card"] + assert "j" in card_con.dims + + +class TestSOS2Reformulation: + """Tests for SOS2 reformulation.""" + + def test_basic_sos2(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Check auxiliary variables and constraints were added + assert "_test_x_z" in m.variables + assert "_test_x_upper_first" in m.constraints + assert "_test_x_upper_last" in m.constraints + assert "_test_x_card" in m.constraints + + # Segment indicators should have n-1 elements + z = m.variables["_test_x_z"] + assert z.sizes["i"] == 2 # n-1 = 3-1 = 2 + + def test_sos2_trivial_single_element(self) -> None: + m = Model() + idx = pd.Index([0], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + + assert "_test_x_z" not in m.variables + + def test_sos2_two_elements(self) -> None: + m = Model() + idx = pd.Index([0, 1], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Should have 1 segment indicator + z = m.variables["_test_x_z"] + assert z.sizes["i"] == 1 + + def test_sos2_with_middle_constraints(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2, 3], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + assert "_test_x_upper_first" in m.constraints + assert "_test_x_upper_mid" in m.constraints + assert "_test_x_upper_last" in m.constraints + + def test_sos2_multidimensional(self) -> None: + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + + reformulate_sos2(m, x, "_test_") + m.remove_sos_constraints(x) + + # Segment indicator should have (n-1) elements in i dimension, same j dimension + z = m.variables["_test_x_z"] + assert set(z.dims) == {"i", "j"} + assert z.sizes["i"] == 2 # n-1 = 3-1 = 2 + assert z.sizes["j"] == 2 + + # Cardinality constraint should have j dimension preserved + card_con = m.constraints["_test_x_card"] + assert "j" in card_con.dims + + +class TestReformulateAllSOS: + """Tests for reformulate_all_sos.""" + + def test_reformulate_single_sos1(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert result.reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 + + def test_reformulate_multiple_sos(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + y = m.add_variables(lower=0, upper=2, coords=[idx], name="y") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_sos_constraints(y, sos_type=2, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert set(result.reformulated) == {"x", "y"} + assert len(list(m.variables.sos)) == 0 + + def test_reformulate_removes_sos_attrs_for_single_element(self) -> None: + m = Model() + idx = pd.Index([0], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert result.reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 + assert len(result.added_variables) == 0 + assert len(result.added_constraints) == 0 + + def test_reformulate_removes_sos_attrs_for_zero_bounds(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=0, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert result.reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 + assert len(result.added_variables) == 0 + assert len(result.added_constraints) == 0 + + def test_reformulate_raises_on_infinite_bounds(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + with pytest.raises(ValueError, match="infinite"): + reformulate_sos_constraints(m) + + def test_reformulate_raises_on_negative_lower_bounds(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=-1, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + with pytest.raises(ValueError, match="negative lower bounds"): + reformulate_sos_constraints(m) + + +class TestModelReformulateSOS: + """Tests for Model.reformulate_sos_constraints method.""" + + def test_reformulate_inplace(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = m.reformulate_sos_constraints() + + assert result.reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 + assert "_sos_reform_x_y" in m.variables + + +class TestApplyUndoSOSReformulation: + """Tests for Model.apply_sos_reformulation / undo_sos_reformulation.""" + + @staticmethod + def _build_sos1_model() -> Model: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + return m + + def test_apply_stashes_state(self) -> None: + m = self._build_sos1_model() + assert m._sos_reformulation_state is None + + m.apply_sos_reformulation() + + assert m._sos_reformulation_state is not None + assert m._sos_reformulation_state.reformulated == ["x"] + assert len(list(m.variables.sos)) == 0 + assert "_sos_reform_x_y" in m.variables + + def test_undo_restores_and_clears_state(self) -> None: + m = self._build_sos1_model() + m.apply_sos_reformulation() + + m.undo_sos_reformulation() + + assert m._sos_reformulation_state is None + assert list(m.variables.sos) == ["x"] + assert "_sos_reform_x_y" not in m.variables + + def test_double_apply_raises(self) -> None: + m = self._build_sos1_model() + m.apply_sos_reformulation() + + with pytest.raises(RuntimeError, match="already been applied"): + m.apply_sos_reformulation() + + def test_undo_without_apply_raises(self) -> None: + m = self._build_sos1_model() + + with pytest.raises(RuntimeError, match="No SOS reformulation"): + m.undo_sos_reformulation() + + @pytest.mark.parametrize( + "copy_fn", + [ + pytest.param(lambda m: m.copy(), id="model.copy()"), + pytest.param(lambda m: __import__("copy").copy(m), id="copy.copy(model)"), + pytest.param( + lambda m: __import__("copy").deepcopy(m), id="copy.deepcopy(model)" + ), + ], + ) + def test_copy_persists_state_and_undo_works_on_copy( + self, copy_fn: Callable[[Model], Model] + ) -> None: + m = self._build_sos1_model() + m.apply_sos_reformulation() + + c = copy_fn(m) + + # State is carried over but is an independent object + assert c._sos_reformulation_state is not None + assert c._sos_reformulation_state is not m._sos_reformulation_state + # Aux vars/cons exist on the copy (they were copied as part of the + # reformulated model state) + assert "_sos_reform_x_y" in c.variables + assert "_sos_reform_x_upper" in c.constraints + assert "_sos_reform_x_card" in c.constraints + # SOS attrs are not on the copy's "x" yet (still in reformulated form) + assert "x" not in list(c.variables.sos) + + # Undo on the copy fully restores the original SOS form + c.undo_sos_reformulation() + assert c._sos_reformulation_state is None + assert list(c.variables.sos) == ["x"] + assert "_sos_reform_x_y" not in c.variables + assert "_sos_reform_x_upper" not in c.constraints + assert "_sos_reform_x_card" not in c.constraints + + # Original is entirely unaffected + assert m._sos_reformulation_state is not None + assert "_sos_reform_x_y" in m.variables + assert len(list(m.variables.sos)) == 0 + + def test_to_netcdf_warns_when_state_active(self, tmp_path: Path) -> None: + m = self._build_sos1_model() + m.apply_sos_reformulation() + + with pytest.warns(UserWarning, match="active SOS reformulation"): + m.to_netcdf(tmp_path / "m.nc") + + # File written despite the warning — the netcdf carries the + # reformulated MILP form. + assert (tmp_path / "m.nc").exists() + + def test_to_netcdf_silent_after_undo(self, tmp_path: Path) -> None: + m = self._build_sos1_model() + m.apply_sos_reformulation() + m.undo_sos_reformulation() + + with warnings.catch_warnings(): + warnings.simplefilter("error") # any warning fails the test + m.to_netcdf(tmp_path / "m.nc") + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestSolverPathSOSCheck: + """Solver._build() must raise on SOS-bearing model with non-SOS solver.""" + + def test_solver_from_name_raises_without_reformulation(self) -> None: + from linopy import solvers + + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + with pytest.raises(ValueError, match="does not support SOS"): + solvers.Solver.from_name("highs", m, io_api="lp") + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestSolveAutoUndoOnFailure: + """Model.solve must auto-undo SOS reformulation when build/solve raises.""" + + def test_state_restored_when_build_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from linopy import solvers + + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + def boom(*args: object, **kwargs: object) -> None: + raise RuntimeError("simulated build failure") + + monkeypatch.setattr(solvers.Solver, "from_name", boom) + + with pytest.raises(RuntimeError, match="simulated build failure"): + m.solve(solver_name="highs", reformulate_sos=True) + + assert m._sos_reformulation_state is None + assert list(m.variables.sos) == ["x"] + assert "_sos_reform_x_y" not in m.variables + + # A subsequent real solve must not hit "already applied" + monkeypatch.undo() + m.solve(solver_name="highs", reformulate_sos=True) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestSolveWithReformulation: + """Tests for solving with SOS reformulation.""" + + @pytest.mark.parametrize( + "solver_name", + [ + pytest.param( + "gurobi", + marks=pytest.mark.skipif( + "gurobi" not in available_solvers, reason="Gurobi not installed" + ), + ), + pytest.param( + "highs", + marks=pytest.mark.skipif( + "highs" not in available_solvers, reason="HiGHS not installed" + ), + ), + ], + ) + def test_reformulate_handles_masked_sos_variables(self, solver_name: str) -> None: + """ + ``reformulate_sos=True`` must handle SOS variables with masked entries. + + Exercises the reformulation pipeline (``apply_sos_reformulation`` → + binary + linking constraints → solve → ``undo``) on a model whose SOS + variable has a masked slot. Parametrized to cover both the native-SOS + case (gurobi: reformulation runs anyway under ``reformulate_sos=True``, + per #689) and the no-native-SOS case (highs: reformulation is the only + way to solve). + """ + m = Model() + coords = pd.Index([0, 1, 2, 3], name="i") + mask = pd.Series([True, True, False, True], index=coords) + var = m.add_variables( + lower=0, upper=1, coords=[coords], mask=mask, name="sos_var" + ) + m.add_sos_constraints(var, sos_type=1, sos_dim="i") + m.add_objective(-var.sum()) + + m.solve(solver_name=solver_name, reformulate_sos=True) + + sol = m.variables["sos_var"].solution.values + # SOS1 over 3 unmasked entries, all in [0, 1], obj = -sum: + # one slot at 1, others at 0, masked stays NaN. + assert m.objective.value is not None + assert np.isclose(m.objective.value, -1.0) + assert np.isnan(sol[2]) + nonzero = np.flatnonzero(~np.isnan(sol) & (sol > 1e-6)) + assert len(nonzero) == 1 + assert np.isclose(sol[nonzero[0]], 1.0) + + def test_sos1_maximize_with_highs(self) -> None: + """Test SOS1 maximize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Should maximize by choosing x[2] = 1 + assert np.isclose(x.solution.values[2], 1, atol=1e-5) + assert np.isclose(x.solution.values[0], 0, atol=1e-5) + assert np.isclose(x.solution.values[1], 0, atol=1e-5) + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_sos1_minimize_with_highs(self) -> None: + """Test SOS1 minimize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([3, 2, 1]), sense="min") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Should minimize to 0 by setting all x = 0 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 0, atol=1e-5) + + def test_sos2_maximize_with_highs(self) -> None: + """Test SOS2 maximize problem with HiGHS using reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # SOS2 allows two adjacent non-zeros, so x[1] and x[2] can both be 1 + # Maximum is 2 + 3 = 5 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 5, atol=1e-5) + # Check that at most two adjacent variables are non-zero + nonzero_count = (np.abs(x.solution.values) > 1e-5).sum() + assert nonzero_count <= 2 + + def test_sos2_different_coefficients(self) -> None: + """Test SOS2 with different coefficients.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x * np.array([2, 1, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # Best is x[1]=1 and x[2]=1 giving 1+3=4 + # or x[0]=1 and x[1]=1 giving 2+1=3 + assert m.objective.value is not None + assert np.isclose(m.objective.value, 4, atol=1e-5) + + def test_reformulate_sos_false_raises_error(self) -> None: + """Test that HiGHS without reformulate_sos raises error.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + with pytest.raises(ValueError, match="does not support SOS"): + m.solve(solver_name="highs", reformulate_sos=False) + + def test_multidimensional_sos1_with_highs(self) -> None: + """Test multi-dimensional SOS1 with HiGHS.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # For each j, at most one x[i, j] can be non-zero + # Maximum is achieved by one non-zero per j column: 2 total + assert m.objective.value is not None + assert np.isclose(m.objective.value, 2, atol=1e-5) + + # Check SOS1 is satisfied for each j + for j in idx_j: + nonzero_count = (np.abs(x.solution.sel(j=j).values) > 1e-5).sum() + assert nonzero_count <= 1 + + def test_multidimensional_sos2_with_highs(self) -> None: + """Test multi-dimensional SOS2 with HiGHS.""" + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # For each j, at most two adjacent x[i, j] can be non-zero + # Maximum is achieved by two adjacent non-zeros per j column: 4 total + assert m.objective.value is not None + assert np.isclose(m.objective.value, 4, atol=1e-5) + + # Check SOS2 is satisfied for each j + for j in idx_j: + sol_j = x.solution.sel(j=j).values + nonzero_indices = np.where(np.abs(sol_j) > 1e-5)[0] + # At most 2 non-zeros + assert len(nonzero_indices) <= 2 + # If 2 non-zeros, they must be adjacent + if len(nonzero_indices) == 2: + assert abs(nonzero_indices[1] - nonzero_indices[0]) == 1 + + +@pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobi not installed") +class TestEquivalenceWithGurobi: + """Tests comparing reformulated solutions with native Gurobi SOS.""" + + def test_sos1_equivalence(self) -> None: + """Test that reformulated SOS1 gives same result as native Gurobi.""" + gurobipy = pytest.importorskip("gurobipy") + + # Native Gurobi solution + m1 = Model() + idx = pd.Index([0, 1, 2], name="i") + x1 = m1.add_variables(lower=0, upper=1, coords=[idx], name="x") + m1.add_sos_constraints(x1, sos_type=1, sos_dim="i") + m1.add_objective(x1 * np.array([1, 2, 3]), sense="max") + + try: + m1.solve(solver_name="gurobi") + except gurobipy.GurobiError as exc: + pytest.skip(f"Gurobi environment unavailable: {exc}") + + # Reformulated solution with HiGHS + m2 = Model() + x2 = m2.add_variables(lower=0, upper=1, coords=[idx], name="x") + m2.add_sos_constraints(x2, sos_type=1, sos_dim="i") + m2.add_objective(x2 * np.array([1, 2, 3]), sense="max") + + if "highs" in available_solvers: + m2.solve(solver_name="highs", reformulate_sos=True) + assert m1.objective.value is not None + assert m2.objective.value is not None + assert np.isclose(m1.objective.value, m2.objective.value, atol=1e-5) + + def test_sos2_equivalence(self) -> None: + """Test that reformulated SOS2 gives same result as native Gurobi.""" + gurobipy = pytest.importorskip("gurobipy") + + # Native Gurobi solution + m1 = Model() + idx = pd.Index([0, 1, 2], name="i") + x1 = m1.add_variables(lower=0, upper=1, coords=[idx], name="x") + m1.add_sos_constraints(x1, sos_type=2, sos_dim="i") + m1.add_objective(x1 * np.array([1, 2, 3]), sense="max") + + try: + m1.solve(solver_name="gurobi") + except gurobipy.GurobiError as exc: + pytest.skip(f"Gurobi environment unavailable: {exc}") + + # Reformulated solution with HiGHS + m2 = Model() + x2 = m2.add_variables(lower=0, upper=1, coords=[idx], name="x") + m2.add_sos_constraints(x2, sos_type=2, sos_dim="i") + m2.add_objective(x2 * np.array([1, 2, 3]), sense="max") + + if "highs" in available_solvers: + m2.solve(solver_name="highs", reformulate_sos=True) + assert m1.objective.value is not None + assert m2.objective.value is not None + assert np.isclose(m1.objective.value, m2.objective.value, atol=1e-5) + + +class TestEdgeCases: + """Tests for edge cases.""" + + def test_preserves_non_sos_variables(self) -> None: + """Test that non-SOS variables are preserved.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_variables(lower=0, upper=2, coords=[idx], name="y") # No SOS + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_sos_constraints(m) + + # y should be unchanged + assert "y" in m.variables + assert SOS_TYPE_ATTR not in m.variables["y"].attrs + + def test_custom_prefix(self) -> None: + """Test custom prefix for reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + reformulate_sos_constraints(m, prefix="_custom_") + + assert "_custom_x_y" in m.variables + assert "_custom_x_upper" in m.constraints + assert "_custom_x_card" in m.constraints + + def test_constraints_with_sos_variables(self) -> None: + """Test that existing constraints with SOS variables work after reformulation.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + y = m.add_variables(lower=0, upper=10, name="y") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + # Add constraint involving SOS variable + m.add_constraints(x.sum() <= y, name="linking") + + # Reformulate + reformulate_sos_constraints(m) + + # Original constraint should still exist + assert "linking" in m.constraints + + def test_float_coordinates(self) -> None: + """Test SOS with float coordinates (common for piecewise linear).""" + m = Model() + breakpoints = pd.Index([0.0, 0.5, 1.0], name="bp") + lambdas = m.add_variables(lower=0, upper=1, coords=[breakpoints], name="lambda") + m.add_sos_constraints(lambdas, sos_type=2, sos_dim="bp") + + reformulate_sos_constraints(m) + + # Should work with float coordinates + assert "_sos_reform_lambda_z" in m.variables + z = m.variables["_sos_reform_lambda_z"] + # Segment indicators have n-1 = 2 elements + assert z.sizes["bp"] == 2 + + def test_custom_big_m_removed_on_remove_sos(self) -> None: + """Test that custom big_m attribute is removed with SOS constraint.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=100, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=10) + + assert "big_m_upper" in x.attrs + + m.remove_sos_constraints(x) + + assert "big_m_upper" not in x.attrs + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestCustomBigM: + """Tests for custom Big-M functionality.""" + + def test_solve_with_custom_big_m(self) -> None: + """Test solving with custom big_m value.""" + m = Model() + idx = pd.Index([0, 1, 2], name="i") + # Large bounds but tight effective constraint + x = m.add_variables(lower=0, upper=1000, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=1) + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + # With big_m=1, maximum should be 3 (x[2]=1) + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_solve_with_infinite_bounds_and_custom_big_m(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=np.inf, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=5) + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 15, atol=1e-5) + + def test_solve_does_not_mutate_model(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + vars_before = set(m.variables) + cons_before = set(m.constraints) + sos_before = list(m.variables.sos) + + m.solve(solver_name="highs", reformulate_sos=True) + + assert set(m.variables) == vars_before + assert set(m.constraints) == cons_before + assert list(m.variables.sos) == sos_before + + def test_solve_twice_with_reformulate_sos(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + obj1 = m.objective.value + + m.solve(solver_name="highs", reformulate_sos=True) + obj2 = m.objective.value + + assert obj1 is not None and obj2 is not None + assert np.isclose(obj1, obj2, atol=1e-5) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestNoSosConstraints: + def test_reformulate_sos_true_with_no_sos(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos=True) + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + +class TestPartialFailure: + def test_partial_failure_rolls_back(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + y = m.add_variables(lower=-1, upper=1, coords=[idx], name="y") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_sos_constraints(y, sos_type=1, sos_dim="i") + + vars_before = set(m.variables) + cons_before = set(m.constraints) + sos_before = list(m.variables.sos) + + with pytest.raises(ValueError, match="negative lower bounds"): + reformulate_sos_constraints(m) + + assert set(m.variables) == vars_before + assert set(m.constraints) == cons_before + assert list(m.variables.sos) == sos_before + + +class TestMixedBounds: + def test_mixed_finite_infinite_with_big_m(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, 0, 0]), + upper=np.array([5, np.inf, 10]), + coords=[idx], + name="x", + ) + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=8) + M = compute_big_m_values(x) + assert np.allclose(M.values, [5, 8, 8]) + + def test_mixed_finite_infinite_without_big_m_raises(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables( + lower=np.array([0, 0, 0]), + upper=np.array([5, np.inf, 10]), + coords=[idx], + name="x", + ) + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + with pytest.raises(ValueError, match="infinite upper bounds"): + compute_big_m_values(x) + + +class TestBigMValidation: + def test_big_m_zero_raises(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + with pytest.raises(ValueError, match="big_m must be positive"): + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=0) + + def test_big_m_negative_raises(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + with pytest.raises(ValueError, match="big_m must be positive"): + m.add_sos_constraints(x, sos_type=1, sos_dim="i", big_m=-5) + + +class TestUndoReformulation: + def test_undo_restores_sos_attrs(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert len(list(m.variables.sos)) == 0 + assert "_sos_reform_x_y" in m.variables + + undo_sos_reformulation(m, result) + + assert list(m.variables.sos) == ["x"] + assert "_sos_reform_x_y" not in m.variables + assert "_sos_reform_x_upper" not in m.constraints + assert "_sos_reform_x_card" not in m.constraints + + def test_double_reformulate_is_noop(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + m.reformulate_sos_constraints() + + result2 = m.reformulate_sos_constraints() + assert result2.reformulated == [] + + def test_undo_restores_skipped_single_element(self) -> None: + m = Model() + idx = pd.Index([0], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert len(list(m.variables.sos)) == 0 + + undo_sos_reformulation(m, result) + + assert list(m.variables.sos) == ["x"] + + def test_undo_restores_skipped_zero_bounds(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=0, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + + result = reformulate_sos_constraints(m) + + assert len(list(m.variables.sos)) == 0 + + undo_sos_reformulation(m, result) + + assert list(m.variables.sos) == ["x"] + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestUnsortedCoords: + def test_sos2_unsorted_coords_matches_sorted(self) -> None: + coeffs = np.array([1, 2, 3]) + + m_sorted = Model() + idx_sorted = pd.Index([1, 2, 3], name="i") + x_sorted = m_sorted.add_variables( + lower=0, upper=1, coords=[idx_sorted], name="x" + ) + m_sorted.add_sos_constraints(x_sorted, sos_type=2, sos_dim="i") + m_sorted.add_objective(x_sorted * coeffs, sense="max") + m_sorted.solve(solver_name="highs", reformulate_sos=True) + + m_unsorted = Model() + idx_unsorted = pd.Index([3, 1, 2], name="i") + x_unsorted = m_unsorted.add_variables( + lower=0, upper=1, coords=[idx_unsorted], name="x" + ) + m_unsorted.add_sos_constraints(x_unsorted, sos_type=2, sos_dim="i") + m_unsorted.add_objective(x_unsorted * coeffs, sense="max") + m_unsorted.solve(solver_name="highs", reformulate_sos=True) + + assert m_sorted.objective.value is not None + assert m_unsorted.objective.value is not None + assert np.isclose( + m_sorted.objective.value, m_unsorted.objective.value, atol=1e-5 + ) + + def test_sos1_unsorted_coords(self) -> None: + m = Model() + idx = pd.Index([3, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + m.solve(solver_name="highs", reformulate_sos=True) + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestAutoReformulation: + """Tests for reformulate_sos='auto' functionality.""" + + @pytest.fixture() + def sos1_model(self) -> tuple[Model, Variable]: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + return m, x + + def test_auto_reformulates_when_solver_lacks_sos( + self, sos1_model: tuple[Model, Variable] + ) -> None: + m, x = sos1_model + m.solve(solver_name="highs", reformulate_sos="auto") + + assert np.isclose(x.solution.values[2], 1, atol=1e-5) + assert np.isclose(x.solution.values[0], 0, atol=1e-5) + assert np.isclose(x.solution.values[1], 0, atol=1e-5) + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_auto_with_sos2(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2, 3], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=2, sos_dim="i") + m.add_objective(x * np.array([10, 1, 1, 10]), sense="max") + + m.solve(solver_name="highs", reformulate_sos="auto") + + assert m.objective.value is not None + nonzero_indices = np.where(np.abs(x.solution.values) > 1e-5)[0] + assert len(nonzero_indices) <= 2 + if len(nonzero_indices) == 2: + assert abs(nonzero_indices[1] - nonzero_indices[0]) == 1 + assert not np.isclose(m.objective.value, 20, atol=1e-5) + + def test_auto_emits_info_no_warning( + self, sos1_model: tuple[Model, Variable], caplog: pytest.LogCaptureFixture + ) -> None: + m, _ = sos1_model + + with caplog.at_level(logging.INFO): + m.solve(solver_name="highs", reformulate_sos="auto") + + assert any("Reformulating SOS" in msg for msg in caplog.messages) + assert not any("supports SOS natively" in msg for msg in caplog.messages) + + @pytest.mark.skipif( + "gurobi" not in available_solvers, reason="Gurobi not installed" + ) + def test_auto_passes_through_native_sos_without_reformulation(self) -> None: + import gurobipy + + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1, 2, 3]), sense="max") + + try: + m.solve(solver_name="gurobi", reformulate_sos="auto") + except gurobipy.GurobiError as exc: + pytest.skip(f"Gurobi environment unavailable: {exc}") + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + assert np.isclose(x.solution.values[2], 1, atol=1e-5) + assert np.isclose(x.solution.values[0], 0, atol=1e-5) + assert np.isclose(x.solution.values[1], 0, atol=1e-5) + + def test_auto_multidimensional_sos1(self) -> None: + m = Model() + idx_i = pd.Index([0, 1, 2], name="i") + idx_j = pd.Index([0, 1], name="j") + x = m.add_variables(lower=0, upper=1, coords=[idx_i, idx_j], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos="auto") + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 2, atol=1e-5) + for j in idx_j: + nonzero_count = (np.abs(x.solution.sel(j=j).values) > 1e-5).sum() + assert nonzero_count <= 1 + + def test_auto_noop_without_sos(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_objective(x.sum(), sense="max") + + m.solve(solver_name="highs", reformulate_sos="auto") + + assert m.objective.value is not None + assert np.isclose(m.objective.value, 3, atol=1e-5) + + def test_invalid_reformulate_sos_value(self) -> None: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x.sum(), sense="max") + + with pytest.raises(ValueError, match="Invalid value for reformulate_sos"): + m.solve(solver_name="highs", reformulate_sos="invalid") # type: ignore[arg-type] + + +class TestResolveSOSReformulation: + """Helper contracts not already exercised end-to-end by ``m.solve(...)``.""" + + @staticmethod + def _sos_model() -> Model: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + return m + + def test_no_sos_short_circuits(self) -> None: + # Fast path: no SOS variables means False regardless of args. + m = Model() + m.add_variables(name="x") + for v in (True, False, "auto"): + assert m._resolve_sos_reformulation(None, v) is False + + def test_true_does_not_consult_solver_name(self) -> None: + # reformulate_sos=True must not require solver_name — no lookup. + assert self._sos_model()._resolve_sos_reformulation(None, True) is True + + def test_auto_with_none_solver_raises(self) -> None: + with pytest.raises(ValueError, match="requires an explicit `solver_name`"): + self._sos_model()._resolve_sos_reformulation(None, "auto") + + +@pytest.mark.skipif("highs" not in available_solvers, reason="HiGHS not installed") +class TestRemoteBracket: + """ + Model.solve(remote=...) must bracket SOS reformulation around the remote + dispatch and suppress the to_netcdf warning that fires inside the helper. + """ + + @staticmethod + def _sos_model() -> Model: + m = Model() + idx = pd.Index([0, 1, 2], name="i") + x = m.add_variables(lower=0, upper=1, coords=[idx], name="x") + m.add_sos_constraints(x, sos_type=1, sos_dim="i") + m.add_objective(x * np.array([1.0, 2.0, 3.0]), sense="max") + return m + + def _fake_handler( + self, observed: dict[str, object], tmp_path: Path + ) -> RemoteHandler: + """ + Non-OetcHandler stand-in with the SSH-shaped `solve_on_remote`. + + Records whether the model arrives in reformulated form, then runs + `model.to_netcdf(...)` and `read_netcdf(...)` (naturally — no + warning recording here, so we can observe at the call-site whether + Model.solve's suppression worked). + """ + from linopy.io import read_netcdf + from linopy.sos_reformulation import ( + sos_reformulation_context, + suppress_serialization_warning, + ) + + class _Handler: + def solve_on_remote( + _self, + model: Model, + *, + reformulate_sos: bool | Literal["auto"] = False, + **kwargs: object, + ) -> Model: + solver_name = kwargs.get("solver_name") + assert solver_name is None or isinstance(solver_name, str) + with sos_reformulation_context( + model, solver_name, reformulate_sos + ) as applied: + observed["state_active"] = ( + model._sos_reformulation_state is not None + ) + observed["solver_name_arg"] = solver_name + with suppress_serialization_warning(active=applied): + model.to_netcdf(tmp_path / "sent.nc") + solved = read_netcdf(tmp_path / "sent.nc") + for _name, var in solved.variables.items(): + arr = np.zeros(var.labels.shape, dtype=float) + var.solution = xr.DataArray(arr, dims=var.labels.dims) + solved.objective.set_value(0.0) + solved.status = "ok" + solved.termination_condition = "optimal" + return solved + + return cast(RemoteHandler, _Handler()) + + def test_remote_brackets_and_suppresses_warning(self, tmp_path: Path) -> None: + m = self._sos_model() + observed: dict[str, object] = {} + handler = self._fake_handler(observed, tmp_path) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + m.solve(solver_name="highs", remote=handler, reformulate_sos=True) + + # Reformulation was active when the handler ran (apply happened + # before the remote dispatch). + assert observed["state_active"] is True + assert observed["solver_name_arg"] == "highs" + + # No "active SOS reformulation" warning escaped Model.solve. + assert not any("active SOS reformulation" in str(w.message) for w in captured) + + # Lifecycle wound down: state cleared, original SOS variable restored. + assert m._sos_reformulation_state is None + assert list(m.variables.sos) == ["x"] + assert "_sos_reform_x_y" not in m.variables + + def test_remote_skips_bracket_when_reformulate_sos_false( + self, tmp_path: Path + ) -> None: + m = self._sos_model() + observed: dict[str, object] = {} + handler = self._fake_handler(observed, tmp_path) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + m.solve(solver_name="highs", remote=handler, reformulate_sos=False) + + # No reformulation happened — model still has the original SOS var + # when the handler sees it, and to_netcdf never warns. + assert observed["state_active"] is False + assert not any("active SOS reformulation" in str(w.message) for w in captured) + assert m._sos_reformulation_state is None + + def test_remote_auto_requires_solver_name_with_sos(self, tmp_path: Path) -> None: + m = self._sos_model() + observed: dict[str, object] = {} + handler = self._fake_handler(observed, tmp_path) + + with pytest.raises(ValueError, match="requires an explicit `solver_name`"): + m.solve(remote=handler, reformulate_sos="auto") diff --git a/test/test_variable.py b/test/test_variable.py index b7aa0491a..8a9000891 100644 --- a/test/test_variable.py +++ b/test/test_variable.py @@ -5,6 +5,8 @@ @author: fabian """ +from typing import Any + import numpy as np import pandas as pd import polars as pl @@ -12,6 +14,7 @@ import xarray as xr import xarray.core.indexes import xarray.core.utils +from xarray import DataArray from xarray.testing import assert_equal import linopy @@ -54,6 +57,21 @@ def test_variable_inherited_properties(x: linopy.Variable) -> None: assert isinstance(x.ndim, int) +def test_variable_type() -> None: + m = Model() + x = m.add_variables(lower=0, upper=10, name="x") + assert x.type == "Continuous Variable" + + b = m.add_variables(binary=True, name="b") + assert b.type == "Binary Variable" + + i = m.add_variables(lower=0, upper=10, integer=True, name="i") + assert i.type == "Integer Variable" + + sc = m.add_variables(lower=1, upper=10, semi_continuous=True, name="sc") + assert sc.type == "Semi-continuous Variable" + + def test_variable_labels(x: linopy.Variable) -> None: isinstance(x.labels, xr.DataArray) @@ -168,6 +186,60 @@ def test_variable_lower_setter_with_array_invalid_dim(x: linopy.Variable) -> Non x.lower = lower +def test_variable_update_bounds(z: linopy.Variable) -> None: + z.update(lower=2, upper=20) + assert z.lower.item() == 2 + assert z.upper.item() == 20 + + +def test_variable_update_lower_only(z: linopy.Variable) -> None: + z.update(lower=3) + assert z.lower.item() == 3 + assert z.upper.item() == 10 # unchanged from fixture default + + +def test_variable_update_no_kwargs_is_noop(z: linopy.Variable) -> None: + old_lower, old_upper = z.lower.item(), z.upper.item() + z.update() + assert z.lower.item() == old_lower + assert z.upper.item() == old_upper + + +def test_variable_update_rejects_inverted_bounds(z: linopy.Variable) -> None: + with pytest.raises(ValueError, match="lower > upper"): + z.update(lower=20, upper=5) + + +def test_variable_update_rejects_non_constant(z: linopy.Variable) -> None: + with pytest.raises(TypeError, match="must be a constant"): + z.update(upper=z) + + +def test_variable_update_returns_self(z: linopy.Variable) -> None: + out = z.update(lower=1) + assert out is z + + +def test_variable_update_array_invalid_dim(x: linopy.Variable) -> None: + with pytest.raises(ValueError): + x.update(lower=pd.Series(range(15, 25))) + + +def test_variable_update_upper_only(z: linopy.Variable) -> None: + """upper= alone changes upper; lower untouched.""" + old_lower = z.lower.copy() + z.update(upper=25) + assert (z.upper == 25).all() + assert (z.lower == old_lower).all() + + +def test_variable_update_with_array(x: linopy.Variable) -> None: + """Array bound that aligns on the variable's coord is accepted.""" + lower = pd.Series(range(10, 20), index=pd.RangeIndex(10, name="first")) + x.update(lower=lower) + np.testing.assert_array_equal(x.lower.values, lower.values) + + def test_variable_sum(x: linopy.Variable) -> None: res = x.sum() assert res.nterm == 10 @@ -204,6 +276,15 @@ def test_variable_where(x: linopy.Variable) -> None: x.where([True] * 4 + [False] * 6, 0) # type: ignore +def test_variable_where_with_solution(x: linopy.Variable) -> None: + x.solution = xr.DataArray(np.arange(10.0), coords=x.labels.coords) + cond = [True] * 4 + [False] * 6 + filtered = x.where(cond) + assert filtered.labels[9] == x._fill_value["labels"] + assert filtered.data["solution"][0] == 0.0 + assert np.isnan(filtered.data["solution"][9]) + + def test_variable_shift(x: linopy.Variable) -> None: x = x.shift(first=3) assert isinstance(x, linopy.variables.Variable) @@ -342,3 +423,525 @@ def test_variable_multiplication(x: linopy.Variable) -> None: assert x.__rmul__(object()) is NotImplemented assert x.__mul__(object()) is NotImplemented + + +class TestAddVariablesBoundsWithCoords: + """Test that add_variables correctly handles all bound types with coords.""" + + SEQ_COORDS = [pd.RangeIndex(3, name="x")] + DICT_COORDS = {"x": [0, 1, 2]} + + @pytest.fixture() + def model(self) -> "Model": + return Model() + + # -- All bound types should work with both coord formats --------------- + + @pytest.mark.parametrize( + "lower", + [ + pytest.param(0, id="scalar"), + pytest.param(np.float64(0), id="np.number"), + pytest.param(np.array(0), id="numpy-0d"), + pytest.param(np.array([0, 0, 0]), id="numpy-1d"), + pytest.param( + pd.Series([0, 0, 0], index=pd.RangeIndex(3, name="x")), id="pandas" + ), + pytest.param([0, 0, 0], id="list"), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + id="dataarray", + ), + pytest.param(DataArray([0, 0, 0], dims=["x"]), id="dataarray-no-coords"), + pytest.param(xr.DataArray(0), id="dataarray-0d"), + ], + ) + @pytest.mark.parametrize( + "coords", + [ + pytest.param([pd.RangeIndex(3, name="x")], id="seq-coords"), + pytest.param({"x": [0, 1, 2]}, id="dict-coords"), + ], + ) + def test_bound_types_with_coords( + self, model: "Model", lower: Any, coords: Any + ) -> None: + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.shape == (3,) + assert var.dims == ("x",) + assert list(var.coords["x"].values) == [0, 1, 2] + + # -- DataArray validation: mismatch and extra dims --------------------- + + @pytest.mark.parametrize( + "coords", + [ + pytest.param([pd.RangeIndex(5, name="x")], id="seq-coords"), + pytest.param({"x": [0, 1, 2, 3, 4]}, id="dict-coords"), + ], + ) + def test_dataarray_coord_mismatch(self, model: "Model", coords: Any) -> None: + lower = DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}) + with pytest.raises(ValueError, match="lower bound.*do not match coords"): + model.add_variables(lower=lower, coords=coords, name="x") + + def test_dataarray_coord_mismatch_upper(self, model: "Model") -> None: + upper = DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]}) + with pytest.raises(ValueError, match="upper bound.*do not match coords"): + model.add_variables(upper=upper, coords=self.SEQ_COORDS, name="x") + + def test_dataarray_extra_dims(self, model: "Model") -> None: + lower = DataArray( + [[1, 2], [3, 4], [5, 6]], dims=["x", "y"], coords={"x": [0, 1, 2]} + ) + with pytest.raises(ValueError, match=r"lower bound has dimension\(s\) \['y'\]"): + model.add_variables(lower=lower, coords=self.DICT_COORDS, name="x") + + def test_mask_extra_dims_with_unnamed_coords_and_dims(self, model: "Model") -> None: + """Mask is validated against coords + dims= like lower/upper.""" + mask = DataArray( + [[True, False], [True, False], [False, True]], + dims=["x", "extra"], + coords={"x": [0, 1, 2]}, + ) + with pytest.raises(ValueError, match=r"mask has dimension\(s\) \['extra'\]"): + model.add_variables( + mask=mask, + coords=[[0, 1, 2]], + dims=["x"], + name="m", + ) + + def test_dataarray_coord_reorder(self, model: "Model") -> None: + """A bound whose coords differ only in order is reindexed to coords.""" + lower = DataArray([3, 1, 2], dims=["x"], coords={"x": ["c", "a", "b"]}) + var = model.add_variables( + lower=lower, coords=[pd.Index(["a", "b", "c"], name="x")], name="x" + ) + assert (var.data.lower == [1, 2, 3]).all() + + def test_positional_bound_aligns_to_coords(self, model: "Model") -> None: + """ + Numpy / unnamed-pandas bounds align to coords positionally, + even when the input's auto-generated coord values would not match. + """ + coords = [pd.Index(list("abc"), name="x")] + # numpy array — no labels at all, positional alignment. + v_np = model.add_variables(upper=np.array([1, 2, 3]), coords=coords, name="np") + assert v_np.dims == ("x",) + assert (v_np.data.upper.sel(x="a") == 1).all() + assert (v_np.data.upper.sel(x="c") == 3).all() + # Unnamed Series — pandas index is auto-generated, ignored in favour + # of coords (positional alignment, principle: coords is source of truth). + v_s = model.add_variables( + upper=pd.Series([10, 20, 30]), coords=coords, name="s" + ) + assert v_s.dims == ("x",) + assert (v_s.data.upper.sel(x="a") == 10).all() + assert (v_s.data.upper.sel(x="c") == 30).all() + # Unnamed DataFrame — both axes positional. + v_df = model.add_variables( + upper=pd.DataFrame([[1, 2], [3, 4], [5, 6]]), + coords=[pd.Index(list("abc"), name="x"), pd.Index(list("xy"), name="y")], + name="df", + ) + assert v_df.dims == ("x", "y") + assert (v_df.data.upper.sel(x="a", y="x") == 1).all() + assert (v_df.data.upper.sel(x="c", y="y") == 6).all() + + def test_positional_bound_wrong_size_raises_clear_error( + self, model: "Model" + ) -> None: + """ + Shape mismatch on positional inputs surfaces as a size error, + not a 'coordinates do not match' error. + """ + coords = [pd.Index(list("abc"), name="x")] + with pytest.raises(ValueError, match=r"upper bound could not be aligned"): + model.add_variables(upper=np.array([1, 2]), coords=coords, name="np_bad") + with pytest.raises(ValueError, match=r"upper bound could not be aligned"): + model.add_variables(upper=pd.Series([1, 2]), coords=coords, name="s_bad") + + def test_unnamed_pd_index_is_size_only(self, model: "Model") -> None: + bound = DataArray([1, 2, 3], dims=["dim_0"]) + var = model.add_variables(upper=bound, coords=[pd.Index([0, 1, 2])], name="x") + assert (var.upper == [1, 2, 3]).all() + + # -- Broadcasting missing dims ----------------------------------------- + + @pytest.mark.parametrize( + "bound", + [ + pytest.param( + DataArray([1, 2, 3], dims=["time"], coords={"time": range(3)}), + id="DataArray", + ), + pytest.param( + pd.Series(index=pd.RangeIndex(3, name="time"), data=[1, 2, 3]), + id="Series", + ), + pytest.param( + pd.DataFrame( + index=pd.RangeIndex(3, name="time"), + columns=pd.Index(["red"], name="colour"), + data=[[1], [2], [3]], + ), + id="DataFrame", + ), + pytest.param( + pd.Series( + index=pd.MultiIndex.from_product( + [pd.RangeIndex(3), ["red"]], names=("time", "colour") + ), + data=[1, 2, 3], + ), + id="Series-multiindex", + ), + pytest.param( + pd.DataFrame( + index=pd.RangeIndex(3, name="time"), + columns=pd.MultiIndex.from_product( + [["a", "b"], ["red"]], names=("space", "colour") + ), + data=[[1, 1], [2, 2], [3, 3]], + ), + id="DataFrame-multicolumns", + ), + pytest.param( + pd.DataFrame( + index=pd.MultiIndex.from_product( + [pd.RangeIndex(3), ["a", "b"]], names=("time", "space") + ), + columns=pd.Index(["red"], name="colour"), + data=[[1], [1], [2], [2], [3], [3]], + ), + id="DataFrame-multiindex", + ), + ], + ) + def test_bound_broadcast_missing_dim( + self, model: "Model", bound: DataArray | pd.Series | pd.DataFrame + ) -> None: + """Pandas / DataArray bounds missing dims are broadcast to coords.""" + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + colour = pd.Index(["red"], name="colour") + var = model.add_variables( + lower=-bound, upper=bound, coords=[time, space, colour], name="x" + ) + assert var.dims == ("time", "space", "colour") + assert var.data.lower.dims == ("time", "space", "colour") + assert var.data.upper.dims == ("time", "space", "colour") + assert var.sizes == {"time": 3, "space": 2, "colour": 1} + assert not var.data.lower.isnull().any() + assert (var.data.lower.sel(space="a", colour="red") == [-1, -2, -3]).all() + assert (var.data.lower.sel(space="b", colour="red") == [-1, -2, -3]).all() + assert (var.data.upper.sel(space="a", colour="red") == [1, 2, 3]).all() + + @pytest.mark.parametrize( + "lower, upper", + [ + pytest.param(0, "da", id="scalar-lower+da-upper"), + pytest.param("da", 1, id="da-lower+scalar-upper"), + pytest.param("da", "da", id="da-lower+da-upper"), + ], + ) + def test_dataarray_broadcast_missing_dim_order( + self, model: "Model", lower: Any, upper: Any + ) -> None: + """Dimension order follows coords, not the type of the bounds (#706).""" + x = pd.Index(["a", "b", "c"], name="x") + y = pd.Index(["X", "Y"], name="y") + full = DataArray( + np.arange(6).reshape(3, 2), coords={"x": x, "y": y}, dims=["x", "y"] + ) + # bounds are DataArrays missing the 'y' dimension + da = full.sum("y") + lower = da if lower == "da" else lower + upper = da if upper == "da" else upper + var = model.add_variables(lower=lower, upper=upper, coords=[x, y], name="x") + assert var.dims == ("x", "y") + assert var.data.lower.dims == ("x", "y") + assert var.data.upper.dims == ("x", "y") + + # -- Special coord formats --------------------------------------------- + + def test_xarray_coordinates_object(self, model: "Model") -> None: + time = pd.RangeIndex(3, name="time") + base = model.add_variables(lower=0, coords=[time], name="base") + lower = DataArray([1, 1, 1], dims=["time"], coords={"time": range(3)}) + var = model.add_variables(lower=lower, coords=base.coords, name="x2") + assert var.shape == (3,) + + # -- Mixed bound type combinations ------------------------------------ + + @pytest.mark.parametrize( + "lower, upper", + [ + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + np.array([1, 1, 1]), + id="da-lower+numpy-upper", + ), + pytest.param( + np.array([0, 0, 0]), + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="numpy-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="da-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + 10, + id="da-lower+scalar-upper", + ), + pytest.param( + 0, + DataArray([1, 1, 1], dims=["x"], coords={"x": [0, 1, 2]}), + id="scalar-lower+da-upper", + ), + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}), + xr.DataArray(10), + id="da-lower+scalar-da-upper", + ), + ], + ) + def test_mixed_bound_types(self, model: "Model", lower: Any, upper: Any) -> None: + var = model.add_variables( + lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" + ) + assert var.shape == (3,) + assert var.dims == ("x",) + assert not var.data.lower.isnull().any() + assert not var.data.upper.isnull().any() + + def test_both_dataarray_different_dim_subsets(self, model: "Model") -> None: + """Lower and upper cover different subsets of dims, both broadcast.""" + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": range(3)}) + upper = DataArray([10, 20], dims=["space"], coords={"space": ["a", "b"]}) + var = model.add_variables( + lower=lower, upper=upper, coords=[time, space], name="x" + ) + assert var.sizes == {"time": 3, "space": 2} + assert not var.data.lower.isnull().any() + assert not var.data.upper.isnull().any() + assert (var.data.upper.sel(time=0) == [10, 20]).all() + + def test_one_dataarray_mismatches_other_ok(self, model: "Model") -> None: + """Only the mismatched bound should raise, regardless of the other.""" + lower = DataArray([0, 0, 0], dims=["x"], coords={"x": [0, 1, 2]}) + upper = DataArray([1, 1], dims=["x"], coords={"x": [10, 20]}) + with pytest.raises(ValueError, match=r"upper bound.*do not match coords"): + model.add_variables( + lower=lower, upper=upper, coords=self.SEQ_COORDS, name="x" + ) + + # -- Coords inferred from bounds (no coords arg) ---------------------- + + @pytest.mark.parametrize( + "lower", + [ + pytest.param( + DataArray([0, 0, 0], dims=["x"], coords={"x": [10, 20, 30]}), + id="dataarray", + ), + pytest.param( + pd.Series([0, 0, 0], index=pd.Index([10, 20, 30], name="x")), + id="pandas", + ), + ], + ) + def test_coords_inferred_from_bounds(self, model: "Model", lower: Any) -> None: + """When coords is None, dims/coords are inferred from the bounds.""" + var = model.add_variables(lower=lower, name="x") + assert var.dims == ("x",) + assert list(var.coords["x"].values) == [10, 20, 30] + + def test_coords_inferred_multidim(self, model: "Model") -> None: + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": [0, 1, 2], "space": ["a", "b"]}, + ) + var = model.add_variables(lower=lower, name="x") + assert set(var.dims) == {"time", "space"} + assert var.sizes == {"time": 3, "space": 2} + + # -- Multi-dimensional coords ----------------------------------------- + + @pytest.mark.parametrize( + "coords", + [ + pytest.param( + [pd.RangeIndex(3, name="time"), pd.Index(["a", "b"], name="space")], + id="seq-coords", + ), + pytest.param( + {"time": [0, 1, 2], "space": ["a", "b"]}, + id="dict-coords", + ), + ], + ) + def test_multidim_coords_with_scalar(self, model: "Model", coords: Any) -> None: + var = model.add_variables(lower=0, upper=1, coords=coords, name="x") + assert set(var.dims) == {"time", "space"} + assert var.sizes == {"time": 3, "space": 2} + + def test_multidim_dataarray_with_coords(self, model: "Model") -> None: + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": [0, 1, 2], "space": ["a", "b"]}, + ) + coords = [pd.RangeIndex(3, name="time"), pd.Index(["a", "b"], name="space")] + var = model.add_variables(lower=lower, coords=coords, name="x") + assert set(var.dims) == {"time", "space"} + assert var.sizes == {"time": 3, "space": 2} + assert not var.data.lower.isnull().any() + + def test_bounds_with_different_dim_order(self, model: "Model") -> None: + """Lower (time, space) and upper (space, time) should align correctly.""" + time = pd.RangeIndex(3, name="time") + space = pd.Index(["a", "b"], name="space") + lower = DataArray( + np.zeros((3, 2)), + dims=["time", "space"], + coords={"time": range(3), "space": ["a", "b"]}, + ) + upper = DataArray( + np.ones((2, 3)), + dims=["space", "time"], + coords={"space": ["a", "b"], "time": range(3)}, + ) + var = model.add_variables( + lower=lower, upper=upper, coords=[time, space], name="x" + ) + assert var.sizes == {"time": 3, "space": 2} + assert (var.data.lower.values == 0).all() + assert (var.data.upper.values == 1).all() + + # -- Reordered coordinates --------------------------------------------- + + def test_reordered_coords_reindexed(self, model: "Model") -> None: + """Same coord values in different order should reindex, not raise.""" + lower = DataArray([10, 20, 30], dims=["x"], coords={"x": ["c", "a", "b"]}) + var = model.add_variables(lower=lower, coords={"x": ["a", "b", "c"]}, name="x") + assert list(var.coords["x"].values) == ["a", "b", "c"] + # Values must follow the reindexed order, not the original + assert list(var.data.lower.values) == [20, 30, 10] + + def test_reordered_coords_different_values_raises(self, model: "Model") -> None: + """Overlapping but not identical coord sets must still raise.""" + lower = DataArray([10, 20], dims=["x"], coords={"x": ["a", "b"]}) + with pytest.raises(ValueError, match=r"lower bound.*do not match coords"): + model.add_variables(lower=lower, coords={"x": ["a", "c"]}, name="x") + + # -- String and datetime coordinates ----------------------------------- + + def test_string_coordinates(self, model: "Model") -> None: + coords = {"region": ["north", "south", "east"]} + lower = DataArray( + [0, 0, 0], + dims=["region"], + coords={"region": ["north", "south", "east"]}, + ) + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.dims == ("region",) + assert list(var.coords["region"].values) == ["north", "south", "east"] + + def test_datetime_coordinates(self, model: "Model") -> None: + dates = pd.date_range("2025-01-01", periods=3) + coords = [dates.rename("time")] + lower = DataArray([0, 0, 0], dims=["time"], coords={"time": dates}) + var = model.add_variables(lower=lower, coords=coords, name="x") + assert var.dims == ("time",) + assert var.shape == (3,) + + def test_string_coords_mismatch(self, model: "Model") -> None: + lower = DataArray( + [0, 0], dims=["region"], coords={"region": ["north", "south"]} + ) + with pytest.raises(ValueError, match=r"lower bound.*do not match coords"): + model.add_variables( + lower=lower, + coords={"region": ["north", "south", "east"]}, + name="x", + ) + + +class TestAddVariablesMultiIndexCoords: + """MultiIndex-specific coord handling in add_variables.""" + + @pytest.fixture + def model(self) -> "Model": + return Model() + + @pytest.fixture + def midx(self) -> pd.MultiIndex: + mi = pd.MultiIndex.from_product([[0, 1], ["a", "b"]], names=("l1", "l2")) + mi.name = "multi" + return mi + + def test_scalar_bounds(self, model: "Model", midx: pd.MultiIndex) -> None: + var = model.add_variables(lower=0, upper=1, coords=[midx], name="x") + assert var.shape == (4,) + assert var.dims == ("multi",) + + def test_dataarray_bound(self, model: "Model", midx: pd.MultiIndex) -> None: + bound = DataArray([1, 2, 3, 4], dims=["multi"], coords={"multi": midx}) + var = model.add_variables(upper=bound, coords=[midx], name="x") + assert var.shape == (4,) + assert (var.data.upper == [1, 2, 3, 4]).all() + + def test_dataarray_bound_broadcast( + self, model: "Model", midx: pd.MultiIndex + ) -> None: + time = pd.Index([10, 20, 30], name="time") + bound = DataArray([1, 2, 3, 4], dims=["multi"], coords={"multi": midx}) + var = model.add_variables( + lower=-bound, upper=bound, coords=[midx, time], name="x" + ) + assert var.dims == ("multi", "time") + assert var.shape == (4, 3) + assert (var.data.upper.sel(time=10) == [1, 2, 3, 4]).all() + + def test_without_name_raises(self, model: "Model") -> None: + midx = pd.MultiIndex.from_product([[0, 1], ["a", "b"]], names=("l1", "l2")) + with pytest.raises(TypeError, match="MultiIndex.*must have .name set"): + model.add_variables(lower=0, upper=1, coords=[midx], name="x") + + def test_mismatched_multiindex_raises( + self, model: "Model", midx: pd.MultiIndex + ) -> None: + other = pd.MultiIndex.from_product([[0, 1], ["x", "y"]], names=("l1", "l2")) + other.name = "multi" + bound = DataArray([1, 2, 3, 4], dims=["multi"], coords={"multi": other}) + with pytest.raises(ValueError, match="MultiIndex.*does not match"): + model.add_variables(upper=bound, coords=[midx], name="x") + + def test_single_level_bound_broadcasts( + self, model: "Model", midx: pd.MultiIndex + ) -> None: + bound = DataArray([5, 6], dims=["l1"], coords={"l1": [0, 1]}) + # Implicit level projection is deprecated (scenario B) — warns until + # the v1 convention makes it an error. + with pytest.warns( + linopy.EvolvingAPIWarning, match=r"broadcasting level subset" + ): + var = model.add_variables(upper=bound, coords=[midx], name="x") + assert var.dims == ("multi",) + assert (var.data.upper == [5, 5, 6, 6]).all() + + def test_incomplete_level_bound_raises( + self, model: "Model", midx: pd.MultiIndex + ) -> None: + subset = pd.MultiIndex.from_tuples([(0, "a"), (1, "b")], names=("l1", "l2")) + bound = pd.Series([1, 2], index=subset) + with pytest.raises(ValueError, match="no value for .* level combination"): + model.add_variables(upper=bound, coords=[midx], name="x") diff --git a/test/test_variable_assignment.py b/test/test_variable_assignment.py index 02da32dfd..a64492bac 100644 --- a/test/test_variable_assignment.py +++ b/test/test_variable_assignment.py @@ -3,6 +3,8 @@ This module aims at testing the correct assignment of variable to the model. """ +from typing import Any + import dask import numpy as np import pandas as pd @@ -248,6 +250,63 @@ def test_variable_assignment_binary_with_error() -> None: m.add_variables(lower=-2, coords=coords, binary=True) +def test_variable_assignment_binary_force_on() -> None: + """A scalar bound defaults the other end: lower=1 forces the binary on.""" + forced_on = Model().add_variables( + binary=True, lower=1, coords=[pd.RangeIndex(4, name="t")] + ) + assert (forced_on.lower.values == 1).all() + assert (forced_on.upper.values == 1).all() + + +@pytest.mark.parametrize( + "upper", + [ + pytest.param([1, 1, 0, 0], id="list"), + pytest.param(np.array([1.0, 1.0, 0.0, 0.0]), id="ndarray"), + pytest.param(pd.Series([1, 1, 0, 0]), id="series"), + pytest.param( + xr.DataArray([1, np.nan, 0, 1], dims="t", coords={"t": range(4)}), + id="dataarray-nan", + ), + ], +) +def test_variable_assignment_binary_array_bounds_ok(upper: Any) -> None: + """0/1 bounds accepted, NaN tolerated (for masking), across containers.""" + Model().add_variables(binary=True, upper=upper, coords=[pd.RangeIndex(4, name="t")]) + + +@pytest.mark.parametrize( + "upper", + [ + pytest.param([1, 1, 2, 0], id="list"), + pytest.param(np.array([0.5, 1.0, 0.0, 1.0]), id="fractional"), + pytest.param(pd.Series([2, 1, 0, 1]), id="series"), + pytest.param( + xr.DataArray([1, np.nan, 2, 0], dims="t", coords={"t": range(4)}), + id="dataarray-nan", + ), + ], +) +def test_variable_assignment_binary_array_bounds_error(upper: Any) -> None: + """A non-0/1 value is rejected, even when NaN is also present.""" + with pytest.raises(ValueError, match="must be 0 or 1"): + Model().add_variables( + binary=True, upper=upper, coords=[pd.RangeIndex(4, name="t")] + ) + + +@pytest.mark.parametrize("bound", [0, 1, 0.0, 1.0]) +def test_variable_assignment_binary_scalar_bound_ok(bound: float) -> None: + Model().add_variables(binary=True, upper=bound, coords=[pd.RangeIndex(2)]) + + +@pytest.mark.parametrize("bound", [0.5, 2, -1]) +def test_variable_assignment_binary_scalar_bound_error(bound: float) -> None: + with pytest.raises(ValueError, match="must be 0 or 1"): + Model().add_variables(binary=True, upper=bound, coords=[pd.RangeIndex(2)]) + + def test_variable_assignment_integer() -> None: m = Model() diff --git a/test/test_variables.py b/test/test_variables.py index 3984b091d..e55ca680f 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -107,6 +107,48 @@ def test_variables_nvars(m: Model) -> None: assert m.variables.nvars == 19 +def test_variables_mask_broadcast() -> None: + m = Model() + + lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) + upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) + + mask = pd.Series([True] * 5 + [False] * 5) + x = m.add_variables(lower, upper, name="x", mask=mask) + assert (x.labels[0:5, :] != -1).all() + assert (x.labels[5:10, :] == -1).all() + + mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"]) + y = m.add_variables(lower, upper, name="y", mask=mask2) + assert (y.labels[:, 0:5] != -1).all() + assert (y.labels[:, 5:10] == -1).all() + + # Pandas Series with named index missing a dim is broadcast to data.coords. + mask_pd = pd.Series( + [True, False, True] + [False] * 7, index=pd.RangeIndex(10, name="dim_0") + ) + v = m.add_variables(lower, upper, name="v", mask=mask_pd) + assert (v.labels[[0, 2], :] != -1).all() + assert (v.labels[[1, 3, 4, 5, 6, 7, 8, 9], :] == -1).all() + + # Mask with sparse coords (subset of data's coords) now raises instead of + # emitting a FutureWarning — the rule from the bounds path applies here too. + mask3 = xr.DataArray( + [True, True, False, False, False], + dims=["dim_0"], + coords={"dim_0": range(5)}, + ) + with pytest.raises( + ValueError, match=r"mask: coordinate values for dimension 'dim_0'" + ): + m.add_variables(lower, upper, name="z", mask=mask3) + + # Mask with extra dimension not in data should raise + mask4 = xr.DataArray([True, False], dims=["extra_dim"]) + with pytest.raises(ValueError, match=r"mask has dimension\(s\) \['extra_dim'\]"): + m.add_variables(lower, upper, name="w", mask=mask4) + + def test_variables_get_name_by_label(m: Model) -> None: assert m.variables.get_name_by_label(4) == "x" assert m.variables.get_name_by_label(12) == "y"