diff --git a/doc/index.rst b/doc/index.rst index e2035ab6b..5efd57149 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -143,6 +143,7 @@ This package is published under MIT license. solve-on-remote solve-on-oetc gpu-acceleration + pips-installation .. toctree:: :hidden: diff --git a/doc/pips-installation.rst b/doc/pips-installation.rst new file mode 100644 index 000000000..fcece5b8e --- /dev/null +++ b/doc/pips-installation.rst @@ -0,0 +1,188 @@ +===================================== +Installing PIPS-IPM++ (block solving) +===================================== + +.. warning:: + + This integration is **experimental** and not tested in CI. It requires an + external C++/MPI build of PIPS-IPM++ and is aimed at contributors and users + with genuinely block-separable models. The Python export/round-trip is + stable; the C++ callback driver is under active development. + +`PIPS-IPM++ `_ is a distributed +interior-point solver (MPI) for doubly-bordered block-diagonal ("arrowhead") +linear programs. linopy can export a model with block structure into the format +PIPS consumes and read the solution back. + +The picture is two-part: **linopy exports** the block files in-process, an +**external PIPS build solves** them via a small callback driver, and linopy +reads the result. PIPS is not a pip-installable package — it must be compiled +against MPI and a linear solver. The instructions below use the fully-open +`MUMPS `_ backend, which needs no proprietary +solver. + +Using the ``pips`` solver +========================= + +Once PIPS-IPM++ is installed and the callback driver is built (see the +:ref:`installation steps ` below), solving is a one-liner. Point +linopy at the driver binary through two environment variables and call +``Model.solve`` with ``solver_name="pips"``: + +.. code-block:: bash + + export PIPS_BINARY=/path/to/build-driver/pips_driver + export PIPS_MPI_RANKS=2 # optional, default 1; must be <= number of blocks + +.. code-block:: python + + m.blocks = ... # assign the block structure (see "When to use it") + m.solve(solver_name="pips") + + m.objective.value # optimal objective + m.solution # primal values, mapped back onto the variables + +The solver exports the model, runs PIPS under ``mpirun`` and reads the primal, +duals and objective back onto ``m`` — the same interface as every other linopy +solver. It is strictly opt-in: without ``PIPS_BINARY`` set, ``pips`` is absent +from ``linopy.available_solvers`` and is never auto-selected, so ordinary +installs are unaffected. + +When to use it +============== + +PIPS only pays off for models that are genuinely **block-separable**: a large +set of local variables/constraints per block, coupled by a comparatively small +set of global (first-stage) variables and linking constraints. Assign the block +structure via ``m.blocks`` (a ``xarray.DataArray`` of block ids over one +dimension); everything independent of that dimension becomes global (block 0). +A model where everything ends up linking defeats the purpose. + +.. _pips-install: + +System dependencies +=================== + +Verified on **Ubuntu 24.04** (cmake 3.28, g++ 13). Install the toolchain and +the MUMPS backend: + +.. code-block:: bash + + sudo apt-get install -y gfortran libopenmpi-dev openmpi-bin \ + libmumps-dev libmetis-dev libscalapack-openmpi-dev \ + libblas-dev liblapack-dev zlib1g-dev + +Notes: + +- Mainline PIPS needs **no Boost**. OpenMP ships with g++ (no package). +- Stick to **one MPI implementation** (OpenMPI here) — mixing OpenMPI and MPICH + breaks the build. +- MUMPS is the default open backend. Panua-PARDISO (≥7.2) or HSL/MA57 are + optional performance backends and are user-supplied at build time; they are + not required to solve. + +Building PIPS-IPM++ and the driver +================================== + +Use the **mainline GitLab repository**, not the stale GitHub mirrors (those lack +the MUMPS backend): + +.. code-block:: bash + + git clone --depth 1 https://gitlab.com/pips-ipmpp/pips-ipmpp.git + +The mainline build tree unconditionally adds the GAMS ``gmspips`` driver, which +needs proprietary GDX sources that are not vendored and are irrelevant to the +callback path. Disable that one subdirectory before configuring: + +.. code-block:: bash + + sed -i 's|^\([[:space:]]*\)add_subdirectory(gams/gmspips)|\1# &|' \ + pips-ipmpp/PIPS-IPM/Drivers/CMakeLists.txt + +Then configure, build and install (MUMPS is auto-detected; ``WITH_MAKETEST=OFF`` +avoids pulling GoogleTest). Mainline exports a clean ``pips-ipmpp`` CMake package +on install, so the callback driver can link against it with +``find_package(pips-ipmpp CONFIG)``: + +.. code-block:: bash + + cmake -S pips-ipmpp -B pips-ipmpp/build \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX="$PWD/pips-install" \ + -DWITH_MAKETEST=OFF + cmake --build pips-ipmpp/build -j"$(nproc)" + cmake --install pips-ipmpp/build + +The callback driver (``pips_driver.cpp``) and a scripted end-to-end build +(``build.sh``, which performs all of the above) live under +``dev-scripts/pips/`` in the linopy repository. + +Exporting and solving a model +============================= + +Export in-process, then hand the directory to the driver: + +.. code-block:: python + + import linopy, xarray as xr, numpy as np, pandas as pd + + m = linopy.Model() + time = pd.Index(range(12), name="time") + m.blocks = xr.DataArray(np.repeat([1, 2], 6), [time]) # 2 local blocks + x = m.add_variables(lower=0, coords=[time], name="x") + g = m.add_variables(lower=0, name="g") # global (block 0) + m.add_constraints(x <= g, name="cap") + m.add_objective(x.sum() + g) + + m.to_pips_files("export-dir") # writes the block files + +.. code-block:: bash + + # N = number of local blocks (here 2); use at most N MPI ranks + mpirun -np 2 ./build-driver/pips_driver export-dir + +.. important:: + + PIPS caps the number of MPI ranks at the **number of local blocks** ``N`` (it + aborts with *"too many MPI processes"* above that). Use ``-np <= N``; PIPS + assigns several blocks per rank when there are fewer ranks than blocks. + +The driver writes the primal, objective and status back into ``export-dir``. +Read them with ``linopy.io.read_pips_solution``. + +Validating the export without PIPS +================================== + +You do not need a PIPS build to check the exporter. ``linopy.io.read_pips_files`` +reconstructs a numerically equivalent, flat model straight from the export +directory; solving it with any ordinary solver reproduces the original optimum: + +.. code-block:: python + + from linopy.io import read_pips_files + + m.solve(solver_name="highs") + m2 = read_pips_files("export-dir") + m2.solve(solver_name="highs") + assert abs(m2.objective.value - m.objective.value) < 1e-6 + +This round-trip runs in linopy's test suite and validates the full arrowhead +carve-up (block partition, linking rows, bound encoding) with no external +dependency. + +Limitations +=========== + +- **LP only** (no MIP/QP through this path yet). +- Requires a manual external build; no conda/pip package yet (a MUMPS-based + conda-forge recipe is feasible and planned). +- Performance depends on genuine block separability and the linear-solver + backend (MUMPS works; PARDISO/MA57 are faster, user-supplied). + +References +========== + +- `PIPS-IPM++ (mainline) `_ +- `Block-structure annotation guide `_ +- `MUMPS `_ diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 29d75df35..92d79d743 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -17,6 +17,11 @@ Upcoming Version * ``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. +*Block-structured solving (PIPS-IPM++)* + +* New ``Model.to_pips_files`` exports a model with block structure (assigned via ``Model.blocks``) into the doubly-bordered block-diagonal ("arrowhead") format consumed by the distributed solver `PIPS-IPM++ `_, deriving the submatrices from linopy's internal CSR. ``linopy.io.read_pips_files`` reconstructs an equivalent flat model (usable to validate the export against any solver with no PIPS build), and ``read_pips_solution`` reads the solver output back. See :doc:`pips-installation`. +* An env-gated ``pips`` solver (``solver_name="pips"``, enabled through the ``PIPS_BINARY`` environment variable) runs the export through PIPS under MPI and maps the primal, duals and objective back onto the model. It stays absent from ``available_solvers`` unless the driver binary is present, so ordinary installs are unaffected. + *Other* * Default internal integer labels to ``int32``, cutting memory ~25% and speeding up model build 10-35%. Models exceeding the int32 maximum (~2.1 billion labels) widen to ``int64`` automatically with a ``UserWarning``; pass ``Model(dtypes={"labels": np.int64})`` upfront to avoid the mid-build upcast (exposed read-only via ``Model.dtypes``). diff --git a/linopy/constraints.py b/linopy/constraints.py index dbd2d2eeb..b95d01542 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -2146,10 +2146,10 @@ def set_blocks(self, block_map: np.ndarray) -> None: Get a dataset of same shape as constraints.labels with block values. Let N be the number of blocks. - The following ciases are considered: + The following cases are considered: - * where are all vars are -1, the block is -1 - * where are all vars are 0, the block is 0 + * where all vars are -1, the block is -1 + * where all vars are 0, the block is 0 * where all vars are n, the block is n * where vars are n or 0 (both present), the block is n * N+1 otherwise @@ -2159,18 +2159,19 @@ def set_blocks(self, block_map: np.ndarray) -> None: 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) + term_dim = constraint.term_dim entries = replace_by_map(constraint.vars, block_map) - not_zero = entries != 0 - not_missing = entries != -1 - for n in range(N + 1): - not_n = entries != n - mask = not_n & not_zero & not_missing - res = res.where(mask.any(constraint.term_dim), n) + is_local = entries > 0 + has_local = is_local.any(term_dim) + has_entry = (entries != -1).any(term_dim) - res = res.where(not_missing.any(constraint.term_dim), -1) - res = res.where(not_zero.any(constraint.term_dim), 0) + max_local = entries.where(is_local, 0).max(term_dim) + min_local = entries.where(is_local, N + 1).min(term_dim) + local = max_local.where(max_local == min_local, N + 1) + + res = xr.where(has_local, local, xr.where(has_entry, 0, -1)) + res = res.astype(block_map.dtype) constraint._data = assign_multiindex_safe(constraint.data, blocks=res) @property diff --git a/linopy/io.py b/linopy/io.py index 462fa5b8f..fa0878ea5 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -11,7 +11,8 @@ import shutil import time import warnings -from collections.abc import Callable, Iterable +from collections.abc import Callable, Hashable, Iterable +from dataclasses import dataclass from importlib.metadata import version from io import BufferedWriter from pathlib import Path @@ -21,18 +22,26 @@ import numpy as np import pandas as pd import polars as pl +import scipy.sparse import xarray as xr from tqdm import tqdm from linopy import solvers -from linopy.common import to_polars -from linopy.constants import CONCAT_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR +from linopy.common import to_dataframe, to_polars +from linopy.constants import ( + CONCAT_DIM, + EQUAL, + GREATER_EQUAL, + SOS_DIM_ATTR, + SOS_TYPE_ATTR, +) from linopy.objective import Objective if TYPE_CHECKING: from cupdlpx import Model as cupdlpxModel from highspy.highs import Highs + from linopy.constraints import Constraint from linopy.model import Model from linopy.variables import Variable @@ -794,12 +803,12 @@ def to_cupdlpx(m: Model) -> cupdlpxModel: return solver.solver_model -def to_block_files(m: Model, fn: Path) -> None: +def to_block_files(m: Model, fn: Path | None) -> None: """ - Write out the linopy model to a block structured output. + Write out the linopy model to a block-structured output for PIPS-IPM++. - Experimental: This function does not support grouping duplicated variables - in one constraint row yet! + Experimental: duplicated variables within one constraint row are not + grouped into a single coefficient. """ if fn is None: fn = Path(TemporaryDirectory(prefix="linopy-problem-", dir=m.solver_dir).name) @@ -818,96 +827,532 @@ def to_block_files(m: Model, fn: Path) -> None: for n in range(N + 2): (path / f"block{n}").mkdir() - raise NotImplementedError("This function is not yet implemented fully.") - - # # Write out variables - # blocks = vars.ravel("blocks", filter_missings=True) - # for key, suffix in zip(["labels", "lower", "upper"], ["x", "xl", "xu"]): - # arr = vars.ravel(key, filter_missings=True) - # for n in tqdm(range(N + 1), desc=f"Write variable {key}"): - # arr[blocks == n].tofile(path / f"block{n}" / suffix, sep="\n") - - # # Write out objective (uses variable blocks from above) - # coeffs = np.zeros(m._xCounter) - # coeffs[np.asarray(m.objective.vars)] = np.asarray(m.objective.coeffs) - # # reorder like non-missing variables - # coeffs = coeffs[vars.ravel("labels", filter_missings=True)] - # for n in tqdm(range(N + 1), desc="Write objective"): - # coeffs[blocks == n].tofile(path / f"block{n}" / "c", sep="\n") - - # # Write out rhs - # blocks = cons.ravel("blocks", filter_missings=True) - # rhs = cons.ravel("rhs", filter_missings=True) - # is_equality = cons.ravel(cons.sign == EQUAL, filter_missings=True) - # is_lower_bound = cons.ravel(cons.sign == GREATER_EQUAL, filter_missings=True) - - # for n in tqdm(range(N + 2), desc="Write RHS"): - # is_blockn = blocks == n - - # rhs[is_blockn & is_equality].tofile(path / f"block{n}" / "b", sep="\n") - - # not_equality = is_blockn & ~is_equality - # is_lower_bound_sub = is_lower_bound[not_equality] - # rhs_sub = rhs[not_equality] - - # lower_bound = np.where(is_lower_bound_sub, rhs_sub, -np.inf) - # lower_bound.tofile(path / f"block{n}" / "dl", sep="\n") - - # upper_bound = np.where(~is_lower_bound_sub, rhs_sub, np.inf) - # upper_bound.tofile(path / f"block{n}" / "du", sep="\n") - - # # Write out constraints - # conblocks = cons.ravel("blocks", "vars", filter_missings=True) - # varblocks = cons.ravel("var_blocks", "vars", filter_missings=True) - # is_equality = cons.ravel(cons.sign == EQUAL, "vars", filter_missings=True) - - # is_varblock_0 = varblocks == 0 - # is_conblock_L = conblocks == N + 1 - - # keys = ["labels", "coeffs", "vars"] - - # def filtered(arr, mask, key): - # """ - # Set coefficients to zero where mask is False, keep others unchanged. - - # PIPS requires this information to set the shape of sub-matrices. - # """ - # assert key in keys - # return np.where(mask, arr, 0) if key == "coeffs" else arr - - # for key, suffix in zip(keys, ["row", "data", "col"]): - # arr = cons.ravel(key, "vars", filter_missings=True) - # for n in tqdm(range(N + 1), desc=f"Write constraint {key}"): - # is_conblock_n = conblocks == n - # is_varblock_n = varblocks == n - - # mask = is_conblock_n & is_equality - # filtered(arr[mask], is_varblock_0[mask], key).tofile( - # path / f"block{n}" / f"A_{suffix}", sep="\n" - # ) - # mask = is_conblock_n & ~is_equality - # filtered(arr[mask], is_varblock_0[mask], key).tofile( - # path / f"block{n}" / f"C_{suffix}", sep="\n" - # ) - - # mask = is_conblock_L & is_equality - # filtered(arr[mask], is_varblock_n[mask], key).tofile( - # path / f"block{n}" / f"BL_{suffix}", sep="\n" - # ) - # mask = is_conblock_L & ~is_equality - # filtered(arr[mask], is_varblock_n[mask], key).tofile( - # path / f"block{n}" / f"DL_{suffix}", sep="\n" - # ) - - # if n: - # mask = is_conblock_n & is_equality - # filtered(arr[mask], is_varblock_n[mask], key).tofile( - # path / f"block{n}" / f"B_{suffix}", sep="\n" - # ) - # mask = is_conblock_n & ~is_equality - # filtered(arr[mask], is_varblock_n[mask], key).tofile( - # path / f"block{n}" / f"D_{suffix}", sep="\n" - # ) + block_map = m.variables.get_blockmap(m.blocks.dtype.type) + + # Write out variables + variables = m.variables.flat + var_blocks = variables.blocks.to_numpy() + for key, suffix in zip(["labels", "lower", "upper"], ["x", "xl", "xu"]): + arr = variables[key].to_numpy() + for n in tqdm(range(N + 1), desc=f"Write variable {key}"): + arr[var_blocks == n].tofile(path / f"block{n}" / suffix, sep="\n") + + # Write out objective, aligned to the non-missing variable order + objective = m.objective.flat + obj_coeffs = objective.groupby("vars")["coeffs"].sum() + coeffs = obj_coeffs.reindex(variables.labels, fill_value=0).to_numpy() + for n in tqdm(range(N + 1), desc="Write objective"): + coeffs[var_blocks == n].tofile(path / f"block{n}" / "c", sep="\n") + + # Flatten constraints to one row per (constraint, term) + def constraint_terms(constraint: Constraint) -> pd.DataFrame: + def mask_func(data: dict[Hashable, np.ndarray]) -> pd.Series: + return (data["vars"] != -1) & (data["labels"] != -1) + + return to_dataframe(constraint.data, mask_func=mask_func) + + terms = pd.concat( + [constraint_terms(m.constraints[k]) for k in m.constraints], + ignore_index=True, + ) + + # Write out rhs, one entry per constraint + per_constraint = terms.drop_duplicates("labels") + con_blocks = per_constraint.blocks.to_numpy() + rhs = per_constraint.rhs.to_numpy() + is_equality = (per_constraint.sign == EQUAL).to_numpy() + is_lower_bound = (per_constraint.sign == GREATER_EQUAL).to_numpy() + + for n in tqdm(range(N + 2), desc="Write RHS"): + is_blockn = con_blocks == n + + rhs[is_blockn & is_equality].tofile(path / f"block{n}" / "b", sep="\n") + + not_equality = is_blockn & ~is_equality + is_lower_bound_sub = is_lower_bound[not_equality] + rhs_sub = rhs[not_equality] + + lower_bound = np.where(is_lower_bound_sub, rhs_sub, -np.inf) + lower_bound.tofile(path / f"block{n}" / "dl", sep="\n") + + upper_bound = np.where(~is_lower_bound_sub, rhs_sub, np.inf) + upper_bound.tofile(path / f"block{n}" / "du", sep="\n") + + # Write out constraints, one entry per (constraint, term) + term_con_blocks = terms.blocks.to_numpy() + term_var_blocks = block_map[terms.vars.to_numpy()] + term_is_equality = (terms.sign == EQUAL).to_numpy() + + is_varblock_0 = term_var_blocks == 0 + is_conblock_L = term_con_blocks == N + 1 + + keys = ["labels", "coeffs", "vars"] + + def filtered(arr: np.ndarray, mask: np.ndarray, key: str) -> np.ndarray: + """ + Zero out coefficients where mask is False, keeping other keys unchanged. + + PIPS requires this information to set the shape of sub-matrices. + """ + return np.where(mask, arr, 0) if key == "coeffs" else arr + + for key, suffix in zip(keys, ["row", "data", "col"]): + arr = terms[key].to_numpy() + for n in tqdm(range(N + 1), desc=f"Write constraint {key}"): + is_conblock_n = term_con_blocks == n + is_varblock_n = term_var_blocks == n + + mask = is_conblock_n & term_is_equality + filtered(arr[mask], is_varblock_0[mask], key).tofile( + path / f"block{n}" / f"A_{suffix}", sep="\n" + ) + mask = is_conblock_n & ~term_is_equality + filtered(arr[mask], is_varblock_0[mask], key).tofile( + path / f"block{n}" / f"C_{suffix}", sep="\n" + ) + + mask = is_conblock_L & term_is_equality + filtered(arr[mask], is_varblock_n[mask], key).tofile( + path / f"block{n}" / f"BL_{suffix}", sep="\n" + ) + mask = is_conblock_L & ~term_is_equality + filtered(arr[mask], is_varblock_n[mask], key).tofile( + path / f"block{n}" / f"DL_{suffix}", sep="\n" + ) + + if n: + mask = is_conblock_n & term_is_equality + filtered(arr[mask], is_varblock_n[mask], key).tofile( + path / f"block{n}" / f"B_{suffix}", sep="\n" + ) + mask = is_conblock_n & ~term_is_equality + filtered(arr[mask], is_varblock_n[mask], key).tofile( + path / f"block{n}" / f"D_{suffix}", sep="\n" + ) + + +@dataclass +class PipsProblem: + A_full: scipy.sparse.csr_array + senses: np.ndarray + rhs: np.ndarray + lb: np.ndarray + ub: np.ndarray + c: np.ndarray + x_labels: np.ndarray + objective_sense: str + objective_offset: float + n_blocks: int + col_offsets: list[int] + + +@dataclass +class PipsSolution: + primal: np.ndarray + labels: np.ndarray + objective: float | None + status: str | None + dual_eq: np.ndarray | None + dual_ineq: np.ndarray | None + + +def _lower_indicators(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + flag = np.isfinite(values).astype(np.int8) + stored = np.where(flag == 1, values, 0.0).astype(np.float64) + return stored, flag + + +def _ineq_indicators( + sense: np.ndarray, rhs: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + is_lower = sense == ">" + iclow = is_lower.astype(np.int8) + icupp = (~is_lower).astype(np.int8) + clow = np.where(is_lower, rhs, 0.0).astype(np.float64) + cupp = np.where(is_lower, 0.0, rhs).astype(np.float64) + return clow, iclow, cupp, icupp + + +def _write(path: Path, arr: np.ndarray, dtype: type) -> None: + arr.astype(dtype).tofile(path, sep="\n") + + +def _read(path: Path, dtype: type) -> np.ndarray: + if not path.exists(): + return np.array([], dtype=dtype) + return np.fromfile(path, dtype=dtype, sep="\n") + + +def _write_csr(directory: Path, name: str, sub: scipy.sparse.csr_array) -> None: + sub = sub.tocsr() + sub.sort_indices() + _write(directory / f"{name}_krow", sub.indptr, np.int32) + _write(directory / f"{name}_jcol", sub.indices, np.int32) + _write(directory / f"{name}_val", sub.data, np.float64) + + +def to_pips_files(m: Model, fn: Path | str | None = None) -> Path: + """ + Export `m` to the PIPS-IPM++ arrowhead block format under `fn`. + + Requires block structure set via `m.blocks`; calls `m.calculate_block_maps()`. + Indicator constraints are skipped. Returns the root directory written. + """ + if m.blocks is None: + raise ValueError("Model does not have blocks defined.") + if fn is None: + fn = Path(TemporaryDirectory(prefix="linopy-pips-", dir=m.solver_dir).name) + path = Path(fn) + if path.exists(): + shutil.rmtree(path) + path.mkdir(parents=True) + + m.calculate_block_maps() + N = int(m.blocks.max()) + for n in range(N + 1): + (path / f"block{n}").mkdir() + + label_index = m.variables.label_index + vlabels = label_index.vlabels + block_map = m.variables.get_blockmap(m.blocks.dtype.type) + col_blocks = block_map[vlabels] + + rows, con_labels_list, b_list, sense_list, row_blocks_list = [], [], [], [], [] + for _, c in m.constraints.items(): + if c.is_indicator: + continue + csr, cl, bb, ss = c.to_matrix_with_rhs(label_index) + mask = c.active_row_mask() + rows.append(csr) + con_labels_list.append(cl) + b_list.append(bb) + sense_list.append(ss) + row_blocks_list.append(c.data["blocks"].values.ravel()[mask]) + A_full = scipy.sparse.vstack(rows, format="csr") + con_labels = np.concatenate(con_labels_list).astype(np.int64) + b_all = np.concatenate(b_list).astype(np.float64) + sense_all = np.concatenate(sense_list).astype("U1") + row_blocks = np.concatenate(row_blocks_list) + + mats = m.matrices + lb, ub, c_obj, vtypes = mats.lb, mats.ub, mats.c, mats.vtypes + + col_masks = [np.flatnonzero(col_blocks == n) for n in range(N + 1)] + block_cols = [len(cm) for cm in col_masks] + col_offsets = np.concatenate([[0], np.cumsum(block_cols)])[:-1].tolist() + n0 = block_cols[0] + + row_eq = sense_all == "=" + row_ineq = ~row_eq + row_diag = [row_blocks == n for n in range(N + 1)] + row_link = row_blocks == N + 1 + eqL = row_eq & row_link + ineqL = row_ineq & row_link + mBL, mDL = int(eqL.sum()), int(ineqL.sum()) + + def carve(rowmask: np.ndarray, cols: np.ndarray) -> scipy.sparse.csr_array: + return A_full[np.flatnonzero(rowmask)][:, cols] + + def write_bounds(directory: Path, cols: np.ndarray) -> None: + xlow, ixlow = _lower_indicators(lb[cols]) + xupp, ixupp = _lower_indicators(ub[cols]) + _write(directory / "col_labels", vlabels[cols], np.int64) + _write(directory / "xlow", xlow, np.float64) + _write(directory / "ixlow", ixlow, np.int8) + _write(directory / "xupp", xupp, np.float64) + _write(directory / "ixupp", ixupp, np.int8) + _write(directory / "c", c_obj[cols], np.float64) + (directory / "vtypes").write_text("\n".join(vtypes[cols].tolist())) + + def write_ineq_rhs( + directory: Path, rowmask: np.ndarray, labels_name: str, suffix: str + ) -> None: + clow, iclow, cupp, icupp = _ineq_indicators(sense_all[rowmask], b_all[rowmask]) + _write(directory / f"clow{suffix}", clow, np.float64) + _write(directory / f"iclow{suffix}", iclow, np.int8) + _write(directory / f"cupp{suffix}", cupp, np.float64) + _write(directory / f"icupp{suffix}", icupp, np.int8) + _write(directory / labels_name, con_labels[rowmask], np.int64) + + manifest_blocks = [] + for k in range(N + 1): + directory = path / f"block{k}" + write_bounds(directory, col_masks[k]) + + eq_diag = row_eq & row_diag[k] + ineq_diag = row_ineq & row_diag[k] + mA_k, mC_k = int(eq_diag.sum()), int(ineq_diag.sum()) + + if mA_k: + _write_csr(directory, "A", carve(eq_diag, col_masks[0])) + if k: + _write_csr(directory, "B", carve(eq_diag, col_masks[k])) + if mC_k: + _write_csr(directory, "C", carve(ineq_diag, col_masks[0])) + if k: + _write_csr(directory, "D", carve(ineq_diag, col_masks[k])) + + _write(directory / "b", b_all[eq_diag], np.float64) + _write(directory / "eq_labels", con_labels[eq_diag], np.int64) + write_ineq_rhs(directory, ineq_diag, "ineq_labels", "") + + link_cols = col_masks[0] if k == 0 else col_masks[k] + if mBL: + _write_csr(directory, "BL", carve(eqL, link_cols)) + if mDL: + _write_csr(directory, "DL", carve(ineqL, link_cols)) + + has = {"A": bool(mA_k), "C": bool(mC_k), "BL": bool(mBL), "DL": bool(mDL)} + if k: + has["B"] = bool(mA_k) + has["D"] = bool(mC_k) + manifest_blocks.append( + { + "block": k, + "n_cols": block_cols[k], + "col_offset": int(col_offsets[k]), + "mA": mA_k, + "mC": mC_k, + "has": has, + } + ) + + directory = path / "block0" + _write(directory / "bL", b_all[eqL], np.float64) + _write(directory / "eqL_labels", con_labels[eqL], np.int64) + write_ineq_rhs(directory, ineqL, "ineqL_labels", "L") + + expr = m.objective.expression + offset = float(expr.const) if np.isfinite(expr.const) else 0.0 + manifest = { + "format": "linopy-pips", + "version": 1, + "n_blocks": N, + "objective_sense": m.objective.sense, + "objective_offset": offset, + "n_global_cols": n0, + "block_cols": block_cols, + "col_offsets": [int(o) for o in col_offsets], + "linking": {"mBL": mBL, "mDL": mDL}, + "blocks": manifest_blocks, + } + (path / "pips.json").write_text(json.dumps(manifest, indent=2)) + return path + + +def read_pips_problem(fn: Path | str) -> PipsProblem: + """Read a PIPS block directory into a flat :class:`PipsProblem`.""" + path = Path(fn) + manifest = json.loads((path / "pips.json").read_text()) + N = manifest["n_blocks"] + block_cols = manifest["block_cols"] + col_offsets = manifest["col_offsets"] + n0 = manifest["n_global_cols"] + mBL, mDL = manifest["linking"]["mBL"], manifest["linking"]["mDL"] + + x_labels, lb_list, ub_list, c_list = [], [], [], [] + for k in range(N + 1): + d = path / f"block{k}" + xlow = _read(d / "xlow", np.float64) + ixlow = _read(d / "ixlow", np.int8) + xupp = _read(d / "xupp", np.float64) + ixupp = _read(d / "ixupp", np.int8) + x_labels.append(_read(d / "col_labels", np.int64)) + lb_list.append(np.where(ixlow == 1, xlow, -np.inf)) + ub_list.append(np.where(ixupp == 1, xupp, np.inf)) + c_list.append(_read(d / "c", np.float64)) + x_labels = np.concatenate(x_labels).astype(np.int64) + lb = np.concatenate(lb_list) + ub = np.concatenate(ub_list) + c = np.concatenate(c_list) + n_active = len(x_labels) + + def read_wide( + d: Path, name: str, nrows: int, ncols: int, base: int + ) -> scipy.sparse.csr_array: + krow = d / f"{name}_krow" + if not krow.exists() or nrows == 0: + return scipy.sparse.csr_array((nrows, n_active)) + indptr = _read(krow, np.int32) + indices = _read(d / f"{name}_jcol", np.int32) + data = _read(d / f"{name}_val", np.float64) + if indptr.size == 0: + indptr = np.zeros(nrows + 1, dtype=np.int32) + sub = scipy.sparse.csr_array((data, indices, indptr), shape=(nrows, ncols)) + coo = sub.tocoo() + return scipy.sparse.csr_array( + (coo.data, (coo.row, coo.col + base)), shape=(nrows, n_active) + ) + + def decode_ineq( + d: Path, suffix: str, nrows: int + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + clow = _read(d / f"clow{suffix}", np.float64) + iclow = _read(d / f"iclow{suffix}", np.int8) + cupp = _read(d / f"cupp{suffix}", np.float64) + icupp = _read(d / f"icupp{suffix}", np.int8) + row_idx, senses, rhs = [], [], [] + for i in range(nrows): + if iclow[i]: + row_idx.append(i) + senses.append(">") + rhs.append(clow[i]) + if icupp[i]: + row_idx.append(i) + senses.append("<") + rhs.append(cupp[i]) + return ( + np.array(row_idx, dtype=int), + np.array(senses, dtype="U1"), + np.array(rhs, dtype=np.float64), + ) + + a_blocks, senses_all, rhs_all = [], [], [] + for k in range(N + 1): + d = path / f"block{k}" + info = manifest["blocks"][k] + mA_k, mC_k = info["mA"], info["mC"] + if mA_k: + wide = read_wide(d, "A", mA_k, n0, 0) + if k: + wide = wide + read_wide(d, "B", mA_k, block_cols[k], col_offsets[k]) + a_blocks.append(wide) + rhs_all.append(_read(d / "b", np.float64)) + senses_all.append(np.full(mA_k, "=", dtype="U1")) + if mC_k: + wide = read_wide(d, "C", mC_k, n0, 0) + if k: + wide = wide + read_wide(d, "D", mC_k, block_cols[k], col_offsets[k]) + row_idx, senses, rhs = decode_ineq(d, "", mC_k) + a_blocks.append(wide[row_idx]) + senses_all.append(senses) + rhs_all.append(rhs) + + if mBL: + d0 = path / "block0" + wide = read_wide(d0, "BL", mBL, n0, 0) + for n in range(1, N + 1): + wide = wide + read_wide( + path / f"block{n}", "BL", mBL, block_cols[n], col_offsets[n] + ) + a_blocks.append(wide) + rhs_all.append(_read(d0 / "bL", np.float64)) + senses_all.append(np.full(mBL, "=", dtype="U1")) + if mDL: + d0 = path / "block0" + wide = read_wide(d0, "DL", mDL, n0, 0) + for n in range(1, N + 1): + wide = wide + read_wide( + path / f"block{n}", "DL", mDL, block_cols[n], col_offsets[n] + ) + row_idx, senses, rhs = decode_ineq(d0, "L", mDL) + a_blocks.append(wide[row_idx]) + senses_all.append(senses) + rhs_all.append(rhs) + + A_full = scipy.sparse.vstack(a_blocks, format="csr") + return PipsProblem( + A_full=A_full, + senses=np.concatenate(senses_all).astype("U1"), + rhs=np.concatenate(rhs_all).astype(np.float64), + lb=lb, + ub=ub, + c=c, + x_labels=x_labels, + objective_sense=manifest["objective_sense"], + objective_offset=manifest["objective_offset"], + n_blocks=N, + col_offsets=col_offsets, + ) + + +def read_pips_files(fn: Path | str) -> Model: + """Read a PIPS block directory back into an equivalent, solvable Model.""" + from linopy.constants import TERM_DIM + from linopy.expressions import LinearExpression + from linopy.model import Model + + prob = read_pips_problem(fn) + m = Model() + lower = xr.DataArray(prob.lb, dims="col") + upper = xr.DataArray(prob.ub, dims="col") + x = m.add_variables(lower=lower, upper=upper, name="x") + x_model_labels = x.labels.values + + A = prob.A_full.tocsr() + A.sort_indices() + nrows = A.shape[0] + counts = np.diff(A.indptr) + max_t = int(counts.max()) if nrows else 1 + fill = np.arange(max_t)[None, :] < counts[:, None] + vars_pad = np.full((nrows, max_t), -1, dtype=np.int64) + coeffs_pad = np.zeros((nrows, max_t), dtype=np.float64) + vars_pad[fill] = x_model_labels[A.indices] + coeffs_pad[fill] = A.data + + data = xr.Dataset( + { + "coeffs": xr.DataArray(coeffs_pad, dims=["row", TERM_DIM]), + "vars": xr.DataArray(vars_pad, dims=["row", TERM_DIM]), + "const": 0.0, + } + ) + lhs = LinearExpression(data, m) + signs = np.where(prob.senses == "=", "=", np.where(prob.senses == "<", "<=", ">=")) + m.add_constraints( + lhs, xr.DataArray(signs, dims="row"), xr.DataArray(prob.rhs, dims="row") + ) + + m.add_objective( + (x * xr.DataArray(prob.c, dims="col")).sum(), sense=prob.objective_sense + ) + return m + + +def read_pips_solution(fn: Path | str, model: Model | None = None) -> PipsSolution: + """Read the solution the C++ driver wrote into `fn` (spec section D).""" + path = Path(fn) + manifest = json.loads((path / "pips.json").read_text()) + N = manifest["n_blocks"] + + primal, labels, y, z = [], [], [], [] + for k in range(N + 1): + d = path / f"block{k}" + primal.append(_read(d / "x_sol", np.float64)) + labels.append(_read(d / "col_labels", np.int64)) + y.append(_read(d / "y_sol", np.float64)) + z.append(_read(d / "z_sol", np.float64)) + y.append(_read(path / "block0" / "yL_sol", np.float64)) + z.append(_read(path / "block0" / "zL_sol", np.float64)) + + primal_arr = np.concatenate(primal).astype(np.float64) + labels_arr = np.concatenate(labels).astype(np.int64) + dual_eq = np.concatenate(y) if any(a.size for a in y) else None + dual_ineq = np.concatenate(z) if any(a.size for a in z) else None + + obj_path = path / "objective" + objective = float(_read(obj_path, np.float64)[0]) if obj_path.exists() else None + status_path = path / "status" + status = status_path.read_text().strip() if status_path.exists() else None + + if model is not None and primal_arr.size: + lookup = np.full(int(labels_arr.max()) + 1, np.nan) + lookup[labels_arr] = primal_arr + for _, var in model.variables.items(): + vl = var.labels.values + vals = np.where(vl == -1, np.nan, lookup[np.where(vl == -1, 0, vl)]) + var.solution = xr.DataArray(vals, var.coords, dims=var.dims) + if objective is not None: + model.objective._value = objective + + return PipsSolution( + primal=primal_arr, + labels=labels_arr, + objective=objective, + status=status, + dual_eq=dual_eq, + dual_ineq=dual_ineq, + ) def non_bool_dict( diff --git a/linopy/model.py b/linopy/model.py index 24594c9cc..5c55fdb72 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -9,6 +9,7 @@ import logging import os import re +import shutil import warnings from collections.abc import Callable, Mapping, Sequence from pathlib import Path @@ -72,6 +73,7 @@ to_highspy, to_mosek, to_netcdf, + to_pips_files, to_xpress, ) from linopy.matrices import MatrixAccessor @@ -2019,7 +2021,7 @@ def solve( 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) + shutil.rmtree(fn) if os.path.isdir(fn) else os.remove(fn) return self.assign_result(result) @@ -2391,4 +2393,6 @@ def reset_solution(self) -> None: to_block_files = to_block_files + to_pips_files = to_pips_files + dualize = dualize diff --git a/linopy/solvers.py b/linopy/solvers.py index 867dc06f6..caede2b0f 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -4110,14 +4110,106 @@ def get_solver_solution() -> Solution: env_.dispose() +PIPS_BINARY_ENV = "PIPS_BINARY" +PIPS_MPI_RANKS_ENV = "PIPS_MPI_RANKS" + + +def _pips_binary() -> str | None: + """Resolve the PIPS-IPM++ driver from ``$PIPS_BINARY``, or None if unset/missing.""" + binary = os.environ.get(PIPS_BINARY_ENV) + return shutil.which(binary) if binary else None + + class PIPS(Solver[None]): """ - Solver subclass for the PIPS solver. + Solver subclass for the distributed PIPS-IPM++ solver. + + The model is exported to the arrowhead block format via + :func:`linopy.io.to_pips_files`, the driver is launched through + ``mpirun -np `` and the solution is read back + with :func:`linopy.io.read_pips_solution`. The driver binary is taken from + ``$PIPS_BINARY`` and the MPI rank count from ``$PIPS_MPI_RANKS`` (default 1). + PIPS only appears in ``available_solvers`` when both the binary and + ``mpirun`` are found. + + Attributes + ---------- + **solver_options + forwarded to the driver as ``-- `` command-line arguments """ - def __post_init__(self) -> None: - msg = "The PIPS solver interface is not yet implemented." - raise NotImplementedError(msg) + display_name: ClassVar[str] = "PIPS-IPM++" + features: ClassVar[frozenset[SolverFeature]] = frozenset( + {SolverFeature.SOLUTION_FILE_NOT_NEEDED} + ) + + @classmethod + @functools.cache + def is_available(cls) -> bool: + return _pips_binary() is not None and shutil.which("mpirun") is not None + + def _build_file(self, **build_kwargs: Any) -> None: + model = self.model + assert model is not None + self._problem_fn = linopy.io.to_pips_files(model) + self.io_api = "blocks" + self._cache_model_labels(model) + + 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: + export_dir = self._problem_fn + assert export_dir is not None + binary = _pips_binary() + if binary is None: + raise RuntimeError( + f"PIPS driver binary not found. Set the {PIPS_BINARY_ENV} " + "environment variable to the PIPS-IPM++ executable." + ) + if shutil.which("mpirun") is None: + raise RuntimeError("`mpirun` not found on PATH; cannot launch PIPS-IPM++.") + + ranks = int(os.environ.get(PIPS_MPI_RANKS_ENV, "1")) + command = ["mpirun", "-np", str(ranks), binary, path_to_string(export_dir)] + for k, v in self.solver_options.items(): + command += [f"--{k}", str(v)] + + if log_fn is not None: + with open(log_fn, "w") as log_f: + proc = sub.run(command, stdout=log_f, stderr=sub.STDOUT, text=True) + else: + proc = sub.run(command, stdout=sub.PIPE, stderr=sub.STDOUT, text=True) + logger.info(proc.stdout) + if proc.returncode != 0: + raise RuntimeError( + f"PIPS driver exited with code {proc.returncode}: {' '.join(command)}" + ) + + pips_solution = linopy.io.read_pips_solution(export_dir) + status = Status.from_termination_condition(pips_solution.status or "unknown") + status.legacy_status = pips_solution.status or "" + + def get_solver_solution() -> Solution: + primal = _solution_from_labels( + pips_solution.primal, pips_solution.labels, self._n_vars + ) + objective = ( + pips_solution.objective + if pips_solution.objective is not None + else np.nan + ) + return Solution(primal, np.array([], dtype=float), objective) + + solution = self.safe_get_solution(status=status, func=get_solver_solution) + + self.io_api = "blocks" + return self._make_result(status, solution, report=SolverReport()) class cuPDLPx(Solver[None]): diff --git a/test/test_io.py b/test/test_io.py index 27cba396b..fa70714c4 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -18,7 +18,12 @@ import xarray as xr from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf -from linopy.io import signed_number +from linopy.io import ( + read_pips_files, + read_pips_problem, + read_pips_solution, + signed_number, +) from linopy.testing import assert_model_equal HAS_NETCDF4 = importlib.util.find_spec("netCDF4") is not None @@ -364,8 +369,8 @@ def test_to_blocks(tmp_path: Path) -> None: lower: pd.Series = pd.Series(range(20)) upper: pd.Series = pd.Series(range(30, 50)) - x = m.add_variables(lower, upper) - y = m.add_variables(lower, upper) + x = m.add_variables(lower, upper, name="x") + y = m.add_variables(lower, upper, name="y") m.add_constraints(x + y, LESS_EQUAL, 10) @@ -373,8 +378,256 @@ def test_to_blocks(tmp_path: Path) -> None: m.blocks = xr.DataArray([1] * 10 + [2] * 10) - with pytest.raises(NotImplementedError): - m.to_block_files(tmp_path) + m.to_block_files(tmp_path) + + assert sorted(p.name for p in tmp_path.iterdir()) == [ + "block0", + "block1", + "block2", + "block3", + ] + + def read(block: int, suffix: str) -> np.ndarray: + return np.fromfile(tmp_path / f"block{block}" / suffix, sep="\n") + + labels = np.concatenate([read(n, "x") for n in range(3)]).astype(int) + assert sorted(labels) == list(range(40)) + + lowers = np.concatenate([read(n, "xl") for n in range(3)]) + upper_bounds = np.concatenate([read(n, "xu") for n in range(3)]) + assert sorted(lowers) == sorted(list(range(20)) * 2) + assert sorted(upper_bounds) == sorted(list(range(30, 50)) * 2) + + coeffs = np.concatenate([read(n, "c") for n in range(3)]) + x_labels = m.variables["x"].labels.values + is_x = np.isin(labels, x_labels) + assert (coeffs[is_x] == 2).all() + assert (coeffs[~is_x] == 3).all() + + +def _pips_time_plant_model(masked: bool) -> Model: + n_time, n_plant, n_blocks = 12, 3, 2 + m = Model() + time = pd.Index(range(n_time), name="time") + plant = pd.Index(range(n_plant), name="plant") + demand = pd.Series(3 + np.sin(np.pi / 24 * time), index=time) + m.blocks = xr.DataArray( + np.repeat(np.arange(1, n_blocks + 1), n_time // n_blocks), [time] + ) + x = m.add_variables(lower=0, coords=[time, plant], name="x") + y = m.add_variables(lower=0, coords=[plant], name="y") + m.add_constraints(y.sum() >= 0, name="total_capacity") + m.add_constraints(y - x >= 0, name="capacity") + m.add_constraints(x.sum("plant") == demand, name="demand") + ramp = (x - x.shift(time=1)).isel(time=slice(1, None)) + m.add_constraints(ramp <= 10, name="ramplimit") + m.add_constraints(x.sum() <= 50, name="co2limit") + if masked: + mask = xr.DataArray(np.arange(n_time) < n_time // 2, [time]) + m.add_constraints(x.sum("plant") <= 20, name="cap", mask=mask) + m.add_objective((2 * x).sum() + y.sum()) + return m + + +def _trivial_pips_model() -> Model: + m = Model() + idx = pd.Index(range(4), name="i") + m.blocks = xr.DataArray([1, 1, 2, 2], [idx]) + x = m.add_variables(lower=0, upper=10, coords=[idx], name="x") + g = m.add_variables(lower=0, upper=5, name="g") + m.add_constraints(x.sum() + g <= 20, name="lim") + m.add_constraints(x >= 1, name="floor") + m.add_objective(x.sum() + 2 * g, sense="max") + return m + + +PIPS_MODELS = { + "time-plant": lambda: _pips_time_plant_model(masked=False), + "time-plant-masked": lambda: _pips_time_plant_model(masked=True), + "trivial": _trivial_pips_model, +} + + +@pytest.mark.parametrize("builder", PIPS_MODELS.values(), ids=PIPS_MODELS.keys()) +def test_pips_problem_matches_matrices( + builder: Callable[[], Model], tmp_path: Path +) -> None: + m = builder() + m.to_pips_files(tmp_path) + prob = read_pips_problem(tmp_path) + mats = m.matrices + + order_prob = np.argsort(prob.x_labels) + order_orig = np.argsort(mats.vlabels) + assert np.array_equal(prob.x_labels[order_prob], mats.vlabels[order_orig]) + assert np.allclose(prob.lb[order_prob], mats.lb[order_orig]) + assert np.allclose(prob.ub[order_prob], mats.ub[order_orig]) + assert np.allclose(prob.c[order_prob], mats.c[order_orig]) + assert prob.A_full.nnz == mats.A.nnz + assert prob.A_full.shape == mats.A.shape + assert sorted(prob.senses.tolist()) == sorted(mats.sense.tolist()) + assert prob.objective_sense == m.objective.sense + + +def _read_pips_block_system( + d: Path, +) -> tuple[ + dict[tuple[int, int], float], dict[int, tuple[str, float]], list[np.ndarray] +]: + n_blocks = json.loads((d / "pips.json").read_text())["n_blocks"] + + def vec(k: int, name: str, dtype: type = np.float64) -> np.ndarray: + path = d / f"block{k}" / name + return ( + np.fromfile(path, dtype=dtype, sep="\n") + if path.exists() + else np.array([], dtype) + ) + + def labels(k: int, name: str) -> np.ndarray: + return vec(k, name, np.int64) + + def csr(k: int, name: str) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None: + if not (d / f"block{k}" / f"{name}_krow").exists(): + return None + return ( + vec(k, f"{name}_krow", np.int32), + vec(k, f"{name}_jcol", np.int32), + vec(k, f"{name}_val"), + ) + + coeffs: dict[tuple[int, int], float] = {} + + def add(triplet, row_labels: np.ndarray, col_labels: np.ndarray) -> None: + if triplet is None: + return + krow, jcol, val = triplet + for i, row in enumerate(row_labels): + for p in range(int(krow[i]), int(krow[i + 1])): + key = (int(row), int(col_labels[int(jcol[p])])) + coeffs[key] = coeffs.get(key, 0.0) + float(val[p]) + + global_cols = labels(0, "col_labels") + eqL, ineqL = labels(0, "eqL_labels"), labels(0, "ineqL_labels") + add(csr(0, "BL"), eqL, global_cols) + add(csr(0, "DL"), ineqL, global_cols) + for k in range(n_blocks + 1): + eqr, ineqr = labels(k, "eq_labels"), labels(k, "ineq_labels") + add(csr(k, "A"), eqr, global_cols) + add(csr(k, "C"), ineqr, global_cols) + if k >= 1: + local = labels(k, "col_labels") + add(csr(k, "B"), eqr, local) + add(csr(k, "D"), ineqr, local) + add(csr(k, "BL"), eqL, local) + add(csr(k, "DL"), ineqL, local) + + rhs: dict[int, tuple[str, float]] = {} + + def add_ineq(row_labels: np.ndarray, suffix: str) -> None: + iclow = vec(0, f"iclow{suffix}", np.int8) + clow, cupp = vec(0, f"clow{suffix}"), vec(0, f"cupp{suffix}") + for i, lbl in enumerate(row_labels): + rhs[int(lbl)] = (">", float(clow[i])) if iclow[i] else ("<", float(cupp[i])) + + for k in range(n_blocks + 1): + for lbl, b in zip(labels(k, "eq_labels"), vec(k, "b")): + rhs[int(lbl)] = ("=", float(b)) + iclow = vec(k, "iclow", np.int8) + clow, cupp = vec(k, "clow"), vec(k, "cupp") + for i, lbl in enumerate(labels(k, "ineq_labels")): + rhs[int(lbl)] = (">", float(clow[i])) if iclow[i] else ("<", float(cupp[i])) + for lbl, b in zip(eqL, vec(0, "bL")): + rhs[int(lbl)] = ("=", float(b)) + add_ineq(ineqL, "L") + + col_partition = [labels(k, "col_labels") for k in range(n_blocks + 1)] + return coeffs, rhs, col_partition + + +@pytest.mark.parametrize("builder", PIPS_MODELS.values(), ids=PIPS_MODELS.keys()) +def test_pips_block_structure_reconstructs_matrix( + builder: Callable[[], Model], tmp_path: Path +) -> None: + m = builder() + m.to_pips_files(tmp_path) + coeffs, rhs, col_partition = _read_pips_block_system(tmp_path) + mats = m.matrices + + all_cols = np.concatenate(col_partition) + assert sorted(all_cols.tolist()) == sorted(mats.vlabels.tolist()) + assert len(all_cols) == len(set(all_cols.tolist())) + + coo = mats.A.tocoo() + canon: dict[tuple[int, int], float] = {} + for r, c, v in zip(coo.row.tolist(), coo.col.tolist(), coo.data.tolist()): + key = (int(mats.clabels[r]), int(mats.vlabels[c])) + canon[key] = canon.get(key, 0.0) + v + assert coeffs.keys() == canon.keys() + for key, value in canon.items(): + assert coeffs[key] == pytest.approx(value) + + canon_rhs = { + int(mats.clabels[i]): (str(mats.sense[i]), float(mats.b[i])) + for i in range(len(mats.clabels)) + } + assert rhs.keys() == canon_rhs.keys() + for key, (sense, b) in canon_rhs.items(): + assert rhs[key][0] == sense + assert rhs[key][1] == pytest.approx(b) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="highs not installed") +@pytest.mark.parametrize("builder", PIPS_MODELS.values(), ids=PIPS_MODELS.keys()) +def test_pips_roundtrip_solve(builder: Callable[[], Model], tmp_path: Path) -> None: + m = builder() + m.to_pips_files(tmp_path) + m2 = read_pips_files(tmp_path) + + assert m2.objective.sense == m.objective.sense + m.solve(solver_name="highs") + m2.solve(solver_name="highs") + assert m2.objective.value == pytest.approx(m.objective.value) + assert set(m2.matrices.vlabels.tolist()) == set(m.matrices.vlabels.tolist()) + + +@pytest.mark.skipif("highs" not in available_solvers, reason="highs not installed") +def test_pips_solution_roundtrip(tmp_path: Path) -> None: + m = _pips_time_plant_model(masked=True) + m.to_pips_files(tmp_path) + m2 = read_pips_files(tmp_path) + m.solve(solver_name="highs") + + lookup = dict(zip(m.matrices.vlabels.tolist(), m.matrices.sol.tolist())) + manifest = json.loads((tmp_path / "pips.json").read_text()) + for k in range(manifest["n_blocks"] + 1): + cl = np.fromfile( + tmp_path / f"block{k}" / "col_labels", dtype=np.int64, sep="\n" + ) + np.array([lookup[int(label)] for label in cl]).tofile( + tmp_path / f"block{k}" / "x_sol", sep="\n" + ) + np.array([m.objective.value]).tofile(tmp_path / "objective", sep="\n") + (tmp_path / "status").write_text("optimal") + + solution = read_pips_solution(tmp_path, model=m2) + assert solution.status == "optimal" + assert solution.objective == pytest.approx(m.objective.value) + assert sorted(solution.primal) == pytest.approx(sorted(m.matrices.sol)) + assert m2.objective.value == pytest.approx(m.objective.value) + + +@pytest.mark.skipif( + "pips" not in available_solvers or "highs" not in available_solvers, + reason="requires the PIPS driver (set PIPS_BINARY) and highs for the reference", +) +@pytest.mark.parametrize("builder", PIPS_MODELS.values(), ids=PIPS_MODELS.keys()) +def test_pips_end_to_end_solve(builder: Callable[[], Model]) -> None: + reference = builder() + reference.solve(solver_name="highs") + m = builder() + assert m.solve(solver_name="pips") == ("ok", "optimal") + assert m.objective.value == pytest.approx(reference.objective.value, abs=1e-3) class TestSignedNumberExpr: