diff --git a/dev-scripts/profile_model_memory.py b/dev-scripts/profile_model_memory.py new file mode 100644 index 00000000..a6a5a981 --- /dev/null +++ b/dev-scripts/profile_model_memory.py @@ -0,0 +1,200 @@ +""" +Reusable memory profiling script for linopy model building. + +Run with scalene for line-level memory attribution: + scalene run dev-scripts/profile_model_memory.py --preset medium + +Run standalone for quick peak RSS + timing: + python dev-scripts/profile_model_memory.py --shape 100 100 20 --sparsity 0.05 --n-expr 5 +""" + +import argparse +import json +import resource +import subprocess +import time +from datetime import datetime, timezone + +import numpy as np +import pandas as pd +import xarray as xr + +import linopy + +PRESETS = { + "small": {"shape": (100, 100, 20), "sparsity": 0.2, "n_expr": 5}, + "medium": {"shape": (200, 200, 50), "sparsity": 0.2, "n_expr": 5}, + "large": {"shape": (300, 300, 100), "sparsity": 0.2, "n_expr": 5}, +} + + +def get_git_info(): + """Return current git branch and short SHA.""" + info = {} + try: + info["git_branch"] = ( + subprocess.check_output( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], stderr=subprocess.DEVNULL + ) + .decode() + .strip() + ) + info["git_sha"] = ( + subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL + ) + .decode() + .strip() + ) + except Exception: + info["git_branch"] = "unknown" + info["git_sha"] = "unknown" + return info + + +def make_coords(shape): + return [pd.RangeIndex(s, name=f"d{i}") for i, s in enumerate(shape)] + + +def make_sparse_coeffs(shape, sparsity, rng, dtype=np.float64): + """Return dense array with (1 - sparsity) fraction set to zero, plus bool mask.""" + mask = rng.random(shape) < sparsity + vals = rng.standard_normal(shape).astype(dtype) + vals[~mask] = 0.0 + return vals, mask + + +def build_model(shape, sparsity, n_expr, seed=42): + """Build a linopy model with sparse expressions and constraints.""" + rng = np.random.default_rng(seed) + coords = make_coords(shape) + dims = [c.name for c in coords] + + m = linopy.Model() + + variables = [] + coeff_arrays = [] + masks = [] + + for i in range(n_expr): + v = m.add_variables(lower=0, coords=coords, name=f"x{i}") + variables.append(v) + c, mask = make_sparse_coeffs(shape, sparsity, rng) + coeff_arrays.append(c) + masks.append(mask) + + # Build expressions and sum + exprs = [coeff_arrays[i] * variables[i] for i in range(n_expr)] + total = sum(exprs) + + # Combined mask + combined_mask = np.zeros(shape, dtype=bool) + for mask in masks: + combined_mask |= mask + combined_mask_da = xr.DataArray( + combined_mask, dims=dims, coords={c.name: c for c in coords} + ) + + # Add constraints + m.add_constraints(total <= 1, name="con", mask=combined_mask_da) + + return m + + +def get_peak_rss_mb(): + """Get peak RSS in MB via resource module.""" + ru = resource.getrusage(resource.RUSAGE_SELF) + # On macOS ru_maxrss is in bytes, on Linux it's in KB + import sys + + if sys.platform == "darwin": + return ru.ru_maxrss / 1024**2 + return ru.ru_maxrss / 1024 + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Profile linopy model building memory usage." + ) + parser.add_argument( + "--shape", + type=int, + nargs="+", + help="Dimensions of the problem, e.g. --shape 100 100 20", + ) + parser.add_argument( + "--sparsity", + type=float, + default=0.05, + help="Fraction of non-zero coefficients (default: 0.05)", + ) + parser.add_argument( + "--n-expr", + type=int, + default=5, + help="Number of expressions to build (default: 5)", + ) + parser.add_argument( + "--preset", + choices=PRESETS.keys(), + help="Use a preset configuration (overrides --shape, --sparsity, --n-expr)", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="Path to write JSON results (default: stdout only)", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + + if args.preset: + config = PRESETS[args.preset] + shape = config["shape"] + sparsity = config["sparsity"] + n_expr = config["n_expr"] + elif args.shape: + shape = tuple(args.shape) + sparsity = args.sparsity + n_expr = args.n_expr + else: + config = PRESETS["medium"] + shape = config["shape"] + sparsity = config["sparsity"] + n_expr = config["n_expr"] + + total_elements = int(np.prod(shape)) + print(f"Config: shape={shape} sparsity={sparsity} n_expr={n_expr}") + print(f"Total elements: {total_elements:,}") + print("-" * 60) + + t0 = time.perf_counter() + build_model(shape, sparsity, n_expr) + elapsed = time.perf_counter() - t0 + peak_rss = get_peak_rss_mb() + + print(f"Peak RSS: {peak_rss:.1f} MB") + print(f"Elapsed: {elapsed:.2f} s") + + git_info = get_git_info() + result = { + **git_info, + "timestamp": datetime.now(timezone.utc).isoformat(), + "config": {"shape": list(shape), "sparsity": sparsity, "n_expr": n_expr}, + "peak_rss_mb": round(peak_rss, 1), + "elapsed_s": round(elapsed, 2), + } + + if args.output: + with open(args.output, "w") as f: + json.dump(result, f, indent=2) + print(f"\nResults written to {args.output}") + else: + print(f"\nJSON: {json.dumps(result)}") + + +if __name__ == "__main__": + main() diff --git a/linopy/expressions.py b/linopy/expressions.py index 10e243de..6d002ed8 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1772,6 +1772,354 @@ def from_constant(cls, model: Model, constant: ConstantLike) -> LinearExpression return LinearExpression(const_da, model) +class LazyLinearExpression(LinearExpression): + """ + A LinearExpression that defers merge along _term. + + Instead of concatenating datasets immediately (which triggers dense + outer-join padding in xarray), this class stores a list of compact + per-expression datasets. Materialization into a single dense Dataset + happens lazily on first access to ``.data`` so that all existing code + paths work unchanged. The speedup comes from overriding the ``flat`` + property (the solver hot-path) to iterate parts directly, avoiding the + dense padding entirely. + """ + + __slots__ = ("_parts", "_const_override") + + def __init__( + self, + data_or_parts: Dataset | list[Dataset], + model: Model, + *, + parts: list[Dataset] | None = None, + const: DataArray | None = None, + ) -> None: + # Normal LinearExpression construction (e.g. from exprwrap fallback) + if parts is None and not isinstance(data_or_parts, list): + super().__init__(data_or_parts, model) + self._parts = [] + self._const_override = None + return + + # Lazy construction: store parts, defer materialization + if parts is not None: + part_list = parts + else: + part_list = data_or_parts # type: ignore[assignment] + + object.__setattr__(self, "_parts", part_list) + object.__setattr__(self, "_const_override", const) + object.__setattr__(self, "_model", model) + object.__setattr__(self, "_data", None) # lazy + + @property + def data(self) -> Dataset: + if self._data is None: + object.__setattr__(self, "_data", self._materialize()) + return self._data + + @property + def const(self) -> DataArray: + if self._const_override is not None: + return self._const_override + return self.data.const + + @const.setter + def const(self, value: DataArray) -> None: + object.__setattr__(self, "_const_override", None) + self._data = assign_multiindex_safe(self.data, const=value) + + @property + def is_lazy(self) -> bool: + """True if the expression has not been materialized yet.""" + return self._data is None and len(self._parts) > 0 + + def _materialize(self) -> Dataset: + """Fall back to the standard dense merge.""" + if not self._parts: + raise ValueError("LazyLinearExpression has no parts to materialize") + + # Detect if all parts share the same coordinate space + coord_dims = [ + {k: v for k, v in p.sizes.items() if k not in HELPER_DIMS} + for p in self._parts + ] + override = check_common_keys_values(coord_dims) + + kwargs: dict[str, Any] = { + "coords": "minimal", + "compat": "override", + "join": "override" if override else "outer", + "fill_value": FILL_VALUE, + } + ds = xr.concat(self._parts, TERM_DIM, **kwargs) + + for d in set(HELPER_DIMS) & set(ds.coords): + ds = ds.reset_index(d, drop=True) + + if self._const_override is not None: + ds = assign_multiindex_safe(ds, const=self._const_override) + else: + ds = ds.assign(const=0.0) + + # Normalize: transpose so _term is last, broadcast coord dims + ds = Dataset(ds.transpose(..., TERM_DIM)) + (ds,) = xr.broadcast(ds, exclude=HELPER_DIMS) + + return ds + + @property + def flat(self) -> pd.DataFrame: + """ + Convert the expression to a pandas DataFrame without materializing. + + Each part is converted independently and the results are concatenated, + avoiding the dense outer-join padding entirely. + """ + if not self.is_lazy: + return super().flat + + frames = [] + for part in self._parts: + + def mask_func( + datadict: dict[Hashable, np.ndarray], + ) -> pd.Series: + return (datadict["vars"] != -1) & (datadict["coeffs"] != 0) + + df = to_dataframe(part, mask_func=mask_func) + if len(df): + frames.append(df) + + if not frames: + return pd.DataFrame( + {"coeffs": pd.Series(dtype=float), "vars": pd.Series(dtype=int)} + ) + + df = pd.concat(frames, ignore_index=True) + df = df.groupby("vars", as_index=False).sum() + check_has_nulls(df, name=self.type) + return df + + def __add__( + self, + other: ConstantLike + | VariableLike + | ScalarLinearExpression + | LinearExpression + | QuadraticExpression, + ) -> LinearExpression | QuadraticExpression: + if not self.is_lazy: + return super().__add__(other) + + if isinstance(other, QuadraticExpression): + return other.__add__(self) + + try: + if np.isscalar(other): + const = ( + self._const_override + if self._const_override is not None + else xr.DataArray(0.0) + ) + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=list(self._parts), + const=const + other, + ) + + if isinstance(other, LazyLinearExpression) and other.is_lazy: + # Merge parts lists + const_self = ( + self._const_override + if self._const_override is not None + else xr.DataArray(0.0) + ) + const_other = ( + other._const_override + if other._const_override is not None + else xr.DataArray(0.0) + ) + aligned_self, aligned_other = xr.align( + const_self, const_other, join="outer", fill_value=0 + ) + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=self._parts + other._parts, + const=aligned_self + aligned_other, + ) + + # For LinearExpression/Variable, use merge directly to avoid + # materializing self via coord_dims access in super().__add__ + other = as_expression(other, model=self._model) + return merge([self, other]) + except TypeError: + return NotImplemented + + def __neg__(self) -> LinearExpression: + if not self.is_lazy: + return super().__neg__() + neg_parts = [assign_multiindex_safe(p, coeffs=-p.coeffs) for p in self._parts] + const = ( + self._const_override + if self._const_override is not None + else xr.DataArray(0.0) + ) + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=neg_parts, + const=-const, + ) + + def __sub__( + self, + other: ConstantLike + | VariableLike + | ScalarLinearExpression + | LinearExpression + | QuadraticExpression, + ) -> LinearExpression | QuadraticExpression: + if not self.is_lazy: + return super().__sub__(other) + try: + return self.__add__(-other) + except TypeError: + return NotImplemented + + def _compact(self) -> LazyLinearExpression: + """ + Merge parts that share the same coordinate space. + + Groups parts by their non-helper dimension signature and merges + each group with ``join="override"`` (cheap, no padding). This + keeps the part count low after many chained additions. + """ + if not self.is_lazy or len(self._parts) <= 1: + return self + + from itertools import groupby + + def _coord_key(p: Dataset) -> tuple: + return tuple( + sorted((k, v) for k, v in p.sizes.items() if k not in HELPER_DIMS) + ) + + sorted_parts = sorted(self._parts, key=_coord_key) + merged: list[Dataset] = [] + for _key, group in groupby(sorted_parts, key=_coord_key): + group_list = list(group) + if len(group_list) == 1: + merged.append(group_list[0]) + else: + ds = xr.concat( + group_list, + TERM_DIM, + coords="minimal", + compat="override", + join="override", + fill_value=FILL_VALUE, + ) + for d in set(HELPER_DIMS) & set(ds.coords): + ds = ds.reset_index(d, drop=True) + merged.append(ds) + + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=merged, + const=self._const_override, + ) + + def to_polars(self) -> pl.DataFrame: + """Convert the expression to a polars DataFrame without materializing.""" + if not self.is_lazy: + return super().to_polars() + + frames = [] + for part in self._parts: + df = to_polars(part) + df = filter_nulls_polars(df) + if len(df): + frames.append(df) + + if not frames: + return pl.DataFrame( + { + "vars": pl.Series([], dtype=pl.Int64), + "coeffs": pl.Series([], dtype=pl.Float64), + } + ) + + df = pl.concat(frames) + df = group_terms_polars(df) + check_has_nulls_polars(df, name=self.type) + return df + + def sel(self, *args: Any, **kwargs: Any) -> LazyLinearExpression | LinearExpression: + """Apply sel to each part independently.""" + if not self.is_lazy: + return super().sel(*args, **kwargs) + # Materialize — sel semantics with indexers may not apply cleanly per-part + return LinearExpression.sel(self, *args, **kwargs) + + def rename( + self, *args: Any, **kwargs: Any + ) -> LazyLinearExpression | LinearExpression: + """Apply rename to each part independently, skipping parts that lack the dim.""" + if not self.is_lazy: + return super().rename(*args, **kwargs) + # Build the name mapping + from xarray.core.utils import either_dict_or_kwargs + + name_dict = either_dict_or_kwargs(args[0] if args else {}, kwargs, "rename") + new_parts = [] + for p in self._parts: + # Only rename dims/vars that exist in this part + applicable = {k: v for k, v in name_dict.items() if k in p or k in p.dims} + new_parts.append(p.rename(applicable) if applicable else p) + const = self._const_override + if const is not None: + applicable = {k: v for k, v in name_dict.items() if k in const.dims} + if applicable: + const = const.rename(applicable) + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=new_parts, + const=const, + ) + + def shift( + self, *args: Any, **kwargs: Any + ) -> LazyLinearExpression | LinearExpression: + """Apply shift to each part independently.""" + if not self.is_lazy: + return super().shift(*args, **kwargs) + # Materialize — shift fill behavior needs consistent coord space + return LinearExpression.shift(self, *args, **kwargs) + + def diff(self, dim: str, n: int = 1) -> LazyLinearExpression | LinearExpression: + """Apply diff to each part that contains the dimension.""" + if not self.is_lazy: + return super().diff(dim, n) + new_parts = [] + for p in self._parts: + if dim in p.dims: + new_parts.append(p.diff(dim, n)) + else: + new_parts.append(p) + return LazyLinearExpression( + None, # type: ignore[arg-type] + self._model, + parts=new_parts, + const=self._const_override, + ) + + class QuadraticExpression(BaseExpression): """ A quadratic expression consisting of terms of coefficients and variables. @@ -1875,7 +2223,9 @@ def __sub__(self, other: SideLike) -> QuadraticExpression: return self.assign(const=self.const - other) other = as_expression(other, model=self.model, dims=self.coord_dims) - if type(other) is LinearExpression: + if isinstance(other, LinearExpression) and not isinstance( + other, QuadraticExpression + ): other = other.to_quadexpr() return merge([self, -other], cls=self.__class__) except TypeError: @@ -2094,7 +2444,7 @@ def merge( exprs = [exprs] + list(add_exprs) # type: ignore 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_linear_expression = any(isinstance(e, LinearExpression) for e in exprs) if cls is None: cls = QuadraticExpression if has_quad_expression else LinearExpression @@ -2111,7 +2461,8 @@ def merge( model = exprs[0].model - if cls in linopy_types and dim in HELPER_DIMS: + has_lazy = any(isinstance(e, LazyLinearExpression) and e.is_lazy for e in exprs) + if cls in linopy_types and dim in HELPER_DIMS and not has_lazy: coord_dims = [ {k: v for k, v in e.sizes.items() if k not in HELPER_DIMS} for e in exprs ] @@ -2119,15 +2470,26 @@ def merge( else: override = False - data = [e.data if isinstance(e, linopy_types) else e for e in exprs] - data = [fill_missing_coords(ds, fill_helper_dims=True) for ds in data] + data = [ + e.data + if isinstance(e, linopy_types) + and not (isinstance(e, LazyLinearExpression) and e.is_lazy) + else e + for e in exprs + ] + data = [ + fill_missing_coords(ds, fill_helper_dims=True) + if isinstance(ds, Dataset) + else ds + for ds in data + ] if not kwargs: kwargs = { "coords": "minimal", "compat": "override", } - if cls == LinearExpression: + if issubclass(cls, LinearExpression): kwargs["fill_value"] = FILL_VALUE elif cls == variables.Variable: kwargs["fill_value"] = variables.FILL_VALUE @@ -2138,6 +2500,46 @@ def merge( kwargs.setdefault("join", "outer") if dim == TERM_DIM: + # Use lazy merge for LinearExpression + if issubclass(cls, LinearExpression): + # Flatten any existing LazyLinearExpression parts + parts: list[Dataset] = [] + const_arrays: list[DataArray] = [] + for expr in exprs: + if isinstance(expr, LazyLinearExpression) and expr.is_lazy: + parts.extend(expr._parts) + c = ( + expr._const_override + if expr._const_override is not None + else xr.DataArray(0.0) + ) + const_arrays.append(c) + else: + d = expr.data if isinstance(expr, linopy_types) else expr + d = fill_missing_coords(d, fill_helper_dims=True) + const_arrays.append(d["const"]) + parts.append(d[["coeffs", "vars"]]) + + # Sum constants (cheap — no _term dim involved) + subkwargs_const = ( + {**kwargs, "fill_value": 0} + if kwargs + else { + "coords": "minimal", + "compat": "override", + "join": "outer", + "fill_value": 0, + } + ) + const = xr.concat(const_arrays, "_temp", **subkwargs_const).sum("_temp") + + return LazyLinearExpression( + None, # type: ignore[arg-type] + model, + parts=parts, + const=const, + ) + ds = xr.concat([d[["coeffs", "vars"]] for d in data], dim, **kwargs) subkwargs = {**kwargs, "fill_value": 0} const = xr.concat([d["const"] for d in data], dim, **subkwargs).sum(TERM_DIM) diff --git a/linopy/objective.py b/linopy/objective.py index b1449270..88ab4407 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -232,11 +232,13 @@ def set_value(self, value: float) -> None: @property def is_linear(self) -> bool: - return type(self.expression) is expressions.LinearExpression + return isinstance( + self.expression, expressions.LinearExpression + ) and not isinstance(self.expression, expressions.QuadraticExpression) @property def is_quadratic(self) -> bool: - return type(self.expression) is expressions.QuadraticExpression + return isinstance(self.expression, expressions.QuadraticExpression) def to_matrix(self, *args: Any, **kwargs: Any) -> csc_matrix: """Wrapper for expression.to_matrix"""