Skip to content

perf: reduce memory usage in build_model (sparse coefficients + per-effect share constraints) #595

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FBumann merged 16 commits into from
Feb 1, 2026
Merged

perf: reduce memory usage in build_model (sparse coefficients + per-effect share constraints) #595

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
f6a580b
fix: memory issues due to dense large coeficients
FBumann Jan 31, 2026
5e6c8be
fix(effects): avoid massive memory allocation in share variable creation
FBumann Jan 31, 2026
032d4ea
Switch to per contributor constraints to solve memmory issues
FBumann Jan 31, 2026
a9c2e22
fix(effects): avoid massive memory allocation in share variable creation
FBumann Jan 31, 2026
cdc2da9
Switch to per contributor constraints to solve memmory issues
FBumann Jan 31, 2026
a1d57a2
perf: improve bus balance to be more memmory efficient
FBumann Jan 31, 2026
76442c7
Switch to per effect shares
FBumann Feb 1, 2026
00f92c1
Firs succesfull drop to 10 GB
FBumann Feb 1, 2026
3a61f5b
Make more readable
FBumann Feb 1, 2026
0ef2a02
Go back to one variable for all shares
FBumann Feb 1, 2026
e640104
⏺ Instead of adding zero-constraints for uncovered combos, we should …
FBumann Feb 1, 2026
6e933af
Only create variables needed
FBumann Feb 1, 2026
327de3e
_create_share_var went from 1,674ms → 116ms — a 14x speedup! The rein…
FBumann Feb 1, 2026
47711f8
Revert
FBumann Feb 1, 2026
b30f8c3
Revert
FBumann Feb 1, 2026
35af3d2
1. effects.py: add_temporal_contribution and add_periodic_contribut…
FBumann Feb 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 45 additions & 11 deletions flixopt/components.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -917,25 +917,59 @@ def add_effect_contributions(self, effects_model) -> None:
if inv.effects_per_size is not None:
factors = inv.effects_per_size
size = self.size.sel({dim: factors.coords[dim].values})
effects_model.add_periodic_contribution(size * factors, contributor_dim=dim)
for eid in factors.coords['effect'].values:
f_single = factors.sel(effect=eid, drop=True)
if (f_single == 0).all():
continue
effects_model.add_periodic_contribution(size * f_single, contributor_dim=dim, effect=str(eid))

# Investment/retirement effects
invested = self.invested
if invested is not None:
if (f := inv.effects_of_investment) is not None:
effects_model.add_periodic_contribution(
invested.sel({dim: f.coords[dim].values}) * f, contributor_dim=dim
)
if (f := inv.effects_of_retirement) is not None:
effects_model.add_periodic_contribution(
invested.sel({dim: f.coords[dim].values}) * (-f), contributor_dim=dim
)
if (ff := inv.effects_of_investment) is not None:
for eid in ff.coords['effect'].values:
f_single = ff.sel(effect=eid, drop=True)
if (f_single == 0).all():
continue
effects_model.add_periodic_contribution(
invested.sel({dim: f_single.coords[dim].values}) * f_single,
contributor_dim=dim,
effect=str(eid),
)
if (ff := inv.effects_of_retirement) is not None:
for eid in ff.coords['effect'].values:
f_single = ff.sel(effect=eid, drop=True)
if (f_single == 0).all():
continue
effects_model.add_periodic_contribution(
invested.sel({dim: f_single.coords[dim].values}) * (-f_single),
contributor_dim=dim,
effect=str(eid),
)

# === Constants: mandatory fixed + retirement ===
if inv.effects_of_investment_mandatory is not None:
effects_model.add_periodic_contribution(inv.effects_of_investment_mandatory, contributor_dim=dim)
mandatory = inv.effects_of_investment_mandatory
if 'effect' in mandatory.dims:
for eid in mandatory.coords['effect'].values:
effects_model.add_periodic_contribution(
mandatory.sel(effect=eid, drop=True),
contributor_dim=dim,
effect=str(eid),
)
else:
effects_model.add_periodic_contribution(mandatory, contributor_dim=dim)
if inv.effects_of_retirement_constant is not None:
effects_model.add_periodic_contribution(inv.effects_of_retirement_constant, contributor_dim=dim)
ret_const = inv.effects_of_retirement_constant
if 'effect' in ret_const.dims:
for eid in ret_const.coords['effect'].values:
effects_model.add_periodic_contribution(
ret_const.sel(effect=eid, drop=True),
contributor_dim=dim,
effect=str(eid),
)
else:
effects_model.add_periodic_contribution(ret_const, contributor_dim=dim)

# --- Investment Cached Properties ---

Expand Down
151 changes: 110 additions & 41 deletions flixopt/effects.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,9 +345,9 @@ def __init__(self, model: FlowSystemModel, data):
self.share_periodic: linopy.Variable | None = None

# Registered contributions from type models (FlowsModel, StoragesModel, etc.)
# Each entry: a defining_expr with 'contributor' dim
self._temporal_share_defs: list[linopy.LinearExpression] = []
self._periodic_share_defs: list[linopy.LinearExpression] = []
# Per-effect, per-contributor accumulation: effect_id -> {contributor_id -> expr (no effect dim)}
self._temporal_shares: dict[str, dict[str, linopy.LinearExpression]] = {}
self._periodic_shares: dict[str, dict[str, linopy.LinearExpression]] = {}
# Constant (xr.DataArray) contributions with 'contributor' + 'effect' dims
self._temporal_constant_defs: list[xr.DataArray] = []
self._periodic_constant_defs: list[xr.DataArray] = []
Expand All @@ -361,35 +361,76 @@ def effect_index(self):
"""Public access to the effect index for type models."""
return self.data.effect_index

def add_temporal_contribution(self, defining_expr, contributor_dim: str = 'contributor') -> None:
def add_temporal_contribution(
self,
defining_expr,
contributor_dim: str = 'contributor',
effect: str | None = None,
) -> None:
"""Register contributors for the share|temporal variable.

Args:
defining_expr: Expression with a contributor dimension.
Accepts linopy LinearExpression/Variable or plain xr.DataArray (constants).
defining_expr: Expression with a contributor dimension (no effect dim if effect is given).
contributor_dim: Name of the element dimension to rename to 'contributor'.
effect: If provided, the expression is for this specific effect (no effect dim needed).
"""
if contributor_dim != 'contributor':
defining_expr = defining_expr.rename({contributor_dim: 'contributor'})
if isinstance(defining_expr, xr.DataArray):
if effect is not None:
defining_expr = defining_expr.expand_dims(effect=[effect])
elif 'effect' not in defining_expr.dims:
raise ValueError(
"DataArray contribution must have an 'effect' dimension or an explicit effect= argument."
)
self._temporal_constant_defs.append(defining_expr)
else:
self._temporal_share_defs.append(defining_expr)
self._accumulate_shares(self._temporal_shares, self._as_expression(defining_expr), effect)

def add_periodic_contribution(self, defining_expr, contributor_dim: str = 'contributor') -> None:
def add_periodic_contribution(
self,
defining_expr,
contributor_dim: str = 'contributor',
effect: str | None = None,
) -> None:
"""Register contributors for the share|periodic variable.

Args:
defining_expr: Expression with a contributor dimension.
Accepts linopy LinearExpression/Variable or plain xr.DataArray (constants).
defining_expr: Expression with a contributor dimension (no effect dim if effect is given).
contributor_dim: Name of the element dimension to rename to 'contributor'.
effect: If provided, the expression is for this specific effect (no effect dim needed).
"""
if contributor_dim != 'contributor':
defining_expr = defining_expr.rename({contributor_dim: 'contributor'})
if isinstance(defining_expr, xr.DataArray):
if effect is not None:
defining_expr = defining_expr.expand_dims(effect=[effect])
elif 'effect' not in defining_expr.dims:
raise ValueError(
"DataArray contribution must have an 'effect' dimension or an explicit effect= argument."
)
self._periodic_constant_defs.append(defining_expr)
else:
self._periodic_share_defs.append(defining_expr)
self._accumulate_shares(self._periodic_shares, self._as_expression(defining_expr), effect)
Comment on lines 364 to 414
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 1, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

DataArray contributions must carry an effect dimension.
If a DataArray comes in without effect and no effect=... is provided, finalize_shares() later calls .reindex({'effect': ...}), which will error. This can happen for single‑effect constants (e.g., mandatory/retirement contributions). Enforce the requirement or expand with the intended effect to avoid runtime failures.

🛠️ Enforce effect presence 
-        if isinstance(defining_expr, xr.DataArray):
-            if effect is not None:
-                defining_expr = defining_expr.expand_dims(effect=[effect])
-            self._temporal_constant_defs.append(defining_expr)
+        if isinstance(defining_expr, xr.DataArray):
+            if effect is not None:
+                defining_expr = defining_expr.expand_dims(effect=[effect])
+            elif 'effect' not in defining_expr.dims:
+                raise ValueError("DataArray contributions must include an 'effect' dimension or pass effect=")
+            self._temporal_constant_defs.append(defining_expr)
-        if isinstance(defining_expr, xr.DataArray):
-            if effect is not None:
-                defining_expr = defining_expr.expand_dims(effect=[effect])
-            self._periodic_constant_defs.append(defining_expr)
+        if isinstance(defining_expr, xr.DataArray):
+            if effect is not None:
+                defining_expr = defining_expr.expand_dims(effect=[effect])
+            elif 'effect' not in defining_expr.dims:
+                raise ValueError("DataArray contributions must include an 'effect' dimension or pass effect=")
+            self._periodic_constant_defs.append(defining_expr)
🤖 Prompt for AI Agents 
In `@flixopt/effects.py` around lines 364 - 406, The DataArray branches in
add_temporal_contribution and add_periodic_contribution must ensure an 'effect'
dimension is present; if a DataArray lacks the 'effect' dim and no effect
argument is provided, raise a clear ValueError rather than appending a
dimensionless array (which breaks finalize_shares' reindexing). Update
add_temporal_contribution and add_periodic_contribution to check for ('effect'
in defining_expr.dims) when isinstance(defining_expr, xr.DataArray) and if
missing and effect is None raise an error instructing the caller to either pass
effect="..." or supply an 'effect' dimension; if effect is provided, continue to
expand_dims(effect=[effect]) and append to
_temporal_constant_defs/_periodic_constant_defs as before.


@staticmethod
def _accumulate_shares(
accum: dict[str, list],
expr: linopy.LinearExpression,
effect: str | None = None,
) -> None:
"""Append expression to per-effect list."""
# accum structure: {effect_id: [(expr, contributor_ids), ...]}
if effect is not None:
# Expression has no effect dim — tagged with specific effect
accum.setdefault(effect, []).append(expr)
elif 'effect' in expr.dims:
# Expression has effect dim — split per effect (DataArray sel is cheap)
for eid in expr.data.coords['effect'].values:
eid_str = str(eid)
accum.setdefault(eid_str, []).append(expr.sel(effect=eid, drop=True))
else:
raise ValueError('Expression must have effect dim or effect parameter must be given')

def create_variables(self) -> None:
"""Create batched effect variables with 'effect' dimension."""
Expand Down Expand Up @@ -542,19 +583,19 @@ def finalize_shares(self) -> None:
if (sm := self.model._storages_model) is not None:
sm.add_effect_contributions(self)

# === Create share|temporal variable ===
if self._temporal_share_defs:
self.share_temporal = self._create_share_var(self._temporal_share_defs, 'share|temporal', temporal=True)
# === Create share|temporal variable (one combined with contributor × effect dims) ===
if self._temporal_shares:
self.share_temporal = self._create_share_var(self._temporal_shares, 'share|temporal', temporal=True)
self._eq_per_timestep.lhs -= self.share_temporal.sum('contributor')

# === Apply temporal constants directly ===
for const in self._temporal_constant_defs:
self._eq_per_timestep.lhs -= const.sum('contributor').reindex({'effect': self.data.effect_index})

# === Create share|periodic variable ===
if self._periodic_share_defs:
self.share_periodic = self._create_share_var(self._periodic_share_defs, 'share|periodic', temporal=False)
self._eq_periodic.lhs -= self.share_periodic.sum('contributor').reindex({'effect': self.data.effect_index})
# === Create share|periodic variable (one combined with contributor × effect dims) ===
if self._periodic_shares:
self.share_periodic = self._create_share_var(self._periodic_shares, 'share|periodic', temporal=False)
self._eq_periodic.lhs -= self.share_periodic.sum('contributor')

# === Apply periodic constants directly ===
for const in self._periodic_constant_defs:
Expand All @@ -573,39 +614,67 @@ def _share_coords(self, element_dim: str, element_index, temporal: bool = True)

def _create_share_var(
self,
share_defs: list[linopy.LinearExpression],
accum: dict[str, list[linopy.LinearExpression]],
name: str,
temporal: bool,
) -> linopy.Variable:
"""Create a share variable from registered contributor definitions.
"""Create one share variable with (contributor, effect, ...) dims.

accum structure: {effect_id: [expr1, expr2, ...]} where each expr has
(contributor, ...other_dims) dims — no effect dim.

Constraints are added per-effect: var.sel(effect=eid) == merged_for_eid,
which avoids cross-effect alignment.

Aligns all contributor expressions (outer join on contributor dimension),
then sums them to produce a single expression with the full contributor dimension.
Returns:
linopy.Variable with dims (contributor, effect, time/period).
"""
import pandas as pd

# Ensure all share defs have canonical effect order before alignment.
# linopy merge uses join="override" when shapes match, which aligns by
# position not label — mismatched effect order silently shuffles coefficients.
if not accum:
return None

# Collect all contributor IDs across all effects
all_contributor_ids: set[str] = set()
for expr_list in accum.values():
for expr in expr_list:
all_contributor_ids.update(str(c) for c in expr.data.coords['contributor'].values)

contributor_index = pd.Index(sorted(all_contributor_ids), name='contributor')
effect_index = self.data.effect_index
normalized = []
for expr in share_defs:
if 'effect' in expr.dims:
expr_effects = list(expr.data.coords['effect'].values)
if expr_effects != list(effect_index):
expr = linopy.LinearExpression(expr.data.reindex(effect=effect_index), expr.model)
normalized.append(expr)

aligned = linopy.align(*normalized, join='outer', fill_value=0)
combined_expr = sum(aligned[1:], start=aligned[0])

# Extract contributor IDs from the combined expression
all_ids = [str(cid) for cid in combined_expr.data.coords['contributor'].values]
contributor_index = pd.Index(all_ids, name='contributor')
coords = self._share_coords('contributor', contributor_index, temporal=temporal)
var = self.model.add_variables(lower=-np.inf, upper=np.inf, coords=coords, name=name)

self.model.add_constraints(var == combined_expr, name=name)
# Build mask: only create variables for (effect, contributor) combos that have expressions
mask = xr.DataArray(
np.zeros((len(contributor_index), len(effect_index)), dtype=bool),
dims=['contributor', 'effect'],
coords={'contributor': contributor_index, 'effect': effect_index},
)
covered_map: dict[str, list[str]] = {}
for eid, expr_list in accum.items():
cids = set()
for expr in expr_list:
cids.update(str(c) for c in expr.data.coords['contributor'].values)
covered_map[eid] = sorted(cids)
mask.loc[dict(effect=eid, contributor=covered_map[eid])] = True

var = self.model.add_variables(lower=-np.inf, upper=np.inf, coords=coords, name=name, mask=mask)

# Add per-effect constraints (only for covered combos)
for eid, expr_list in accum.items():
contributors = covered_map[eid]
if len(expr_list) == 1:
merged = expr_list[0].reindex(contributor=contributors)
else:
# Reindex all to common contributor set, then sum via linopy.merge (_term addition)
aligned = [e.reindex(contributor=contributors) for e in expr_list]
merged = aligned[0]
for a in aligned[1:]:
merged = merged + a
var_slice = var.sel(effect=eid, contributor=contributors)
self.model.add_constraints(var_slice == merged, name=f'{name}({eid})')

accum.clear()
return var

def get_periodic(self, effect_id: str) -> linopy.Variable:
Expand Down
Loading
Loading