Skip to content

Commit a6b7103

Browse files
FBumannclaude
andauthored
perf: reduce memory usage in build_model (sparse coefficients + per-effect share constraints) (#595)
* fix: memory issues due to dense large coeficients 1. flixopt/features.py — Added sparse_multiply_sum() function that takes a sparse dict of (group_id, sum_id) -> coefficient instead of a dense DataArray. This avoids ever allocating the massive dense array. 2. flixopt/elements.py — Replaced _coefficients (dense DataArray) and _flow_sign (dense DataArray) with a single _signed_coefficients cached property that returns dict[tuple[str, str], float | xr.DataArray] containing only non-zero signed coefficients. Updated create_linear_constraints to use sparse_multiply_sum instead of sparse_weighted_sum. The dense allocation at line 2385 (np.zeros(n_conv, max_eq, n_flows, *time) ~14.5 GB) is completely eliminated. Memory usage is now proportional to the number of non-zero entries (typically 2-3 flows per converter) rather than the full cartesian product. * fix(effects): avoid massive memory allocation in share variable creation Replace linopy.align(join='outer') with per-contributor accumulation and linopy.merge(dim='contributor'). The old approach reindexed ALL dimensions via xr.where(), allocating ~12.7 GB of dense arrays. Now contributions are split by contributor at registration time and accumulated via linopy addition (cheap for same-shape expressions), then merged along the disjoint contributor dimension. * Switch to per contributor constraints to solve memmory issues * fix(effects): avoid massive memory allocation in share variable creation Replace linopy.align(join='outer') with per-contributor accumulation and individual constraints. The old approach reindexed ALL dimensions via xr.where(), allocating ~12.7 GB of dense arrays. Now contributions are split by contributor at registration time and accumulated via linopy addition (cheap for same-shape expressions). Each contributor gets its own constraint, avoiding any cross-contributor alignment. Reduces effects expression memory from 1.2 GB to 5 MB. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Switch to per contributor constraints to solve memmory issues * perf: improve bus balance to be more memmory efficient * Switch to per effect shares * Firs succesfull drop to 10 GB * Make more readable * Go back to one variable for all shares * ⏺ Instead of adding zero-constraints for uncovered combos, we should just set lower=0, upper=0 on those entries (fix the bounds), or better yet — use a mask on the per-effect constraints and set the variable bounds to 0 for uncovered combos. The simplest fix: create the variable with lower=0, upper=0 by default, then only the covered entries need constraints. * Only create variables needed * _create_share_var went from 1,674ms → 116ms — a 14x speedup! The reindex + + approach is much faster than per-contributor sel + merge * Revert * Revert * 1. effects.py: add_temporal_contribution and add_periodic_contribution now raise ValueError if a DataArray has no effect dimension and no effect= argument is provided. 2. statistics_accessor.py: Early return with empty xr.Dataset() when no contributors are detected, preventing xr.concat from failing on an empty list. --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 701f7af commit a6b7103

9 files changed

Lines changed: 367 additions & 214 deletions

File tree

flixopt/components.py

Lines changed: 45 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -917,25 +917,59 @@ def add_effect_contributions(self, effects_model) -> None:
917917
if inv.effects_per_size is not None:
918918
factors = inv.effects_per_size
919919
size = self.size.sel({dim: factors.coords[dim].values})
920-
effects_model.add_periodic_contribution(size * factors, contributor_dim=dim)
920+
for eid in factors.coords['effect'].values:
921+
f_single = factors.sel(effect=eid, drop=True)
922+
if (f_single == 0).all():
923+
continue
924+
effects_model.add_periodic_contribution(size * f_single, contributor_dim=dim, effect=str(eid))
921925

922926
# Investment/retirement effects
923927
invested = self.invested
924928
if invested is not None:
925-
if (f := inv.effects_of_investment) is not None:
926-
effects_model.add_periodic_contribution(
927-
invested.sel({dim: f.coords[dim].values}) * f, contributor_dim=dim
928-
)
929-
if (f := inv.effects_of_retirement) is not None:
930-
effects_model.add_periodic_contribution(
931-
invested.sel({dim: f.coords[dim].values}) * (-f), contributor_dim=dim
932-
)
929+
if (ff := inv.effects_of_investment) is not None:
930+
for eid in ff.coords['effect'].values:
931+
f_single = ff.sel(effect=eid, drop=True)
932+
if (f_single == 0).all():
933+
continue
934+
effects_model.add_periodic_contribution(
935+
invested.sel({dim: f_single.coords[dim].values}) * f_single,
936+
contributor_dim=dim,
937+
effect=str(eid),
938+
)
939+
if (ff := inv.effects_of_retirement) is not None:
940+
for eid in ff.coords['effect'].values:
941+
f_single = ff.sel(effect=eid, drop=True)
942+
if (f_single == 0).all():
943+
continue
944+
effects_model.add_periodic_contribution(
945+
invested.sel({dim: f_single.coords[dim].values}) * (-f_single),
946+
contributor_dim=dim,
947+
effect=str(eid),
948+
)
933949

934950
# === Constants: mandatory fixed + retirement ===
935951
if inv.effects_of_investment_mandatory is not None:
936-
effects_model.add_periodic_contribution(inv.effects_of_investment_mandatory, contributor_dim=dim)
952+
mandatory = inv.effects_of_investment_mandatory
953+
if 'effect' in mandatory.dims:
954+
for eid in mandatory.coords['effect'].values:
955+
effects_model.add_periodic_contribution(
956+
mandatory.sel(effect=eid, drop=True),
957+
contributor_dim=dim,
958+
effect=str(eid),
959+
)
960+
else:
961+
effects_model.add_periodic_contribution(mandatory, contributor_dim=dim)
937962
if inv.effects_of_retirement_constant is not None:
938-
effects_model.add_periodic_contribution(inv.effects_of_retirement_constant, contributor_dim=dim)
963+
ret_const = inv.effects_of_retirement_constant
964+
if 'effect' in ret_const.dims:
965+
for eid in ret_const.coords['effect'].values:
966+
effects_model.add_periodic_contribution(
967+
ret_const.sel(effect=eid, drop=True),
968+
contributor_dim=dim,
969+
effect=str(eid),
970+
)
971+
else:
972+
effects_model.add_periodic_contribution(ret_const, contributor_dim=dim)
939973

940974
# --- Investment Cached Properties ---
941975

flixopt/effects.py

Lines changed: 110 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -345,9 +345,9 @@ def __init__(self, model: FlowSystemModel, data):
345345
self.share_periodic: linopy.Variable | None = None
346346

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

364-
def add_temporal_contribution(self, defining_expr, contributor_dim: str = 'contributor') -> None:
364+
def add_temporal_contribution(
365+
self,
366+
defining_expr,
367+
contributor_dim: str = 'contributor',
368+
effect: str | None = None,
369+
) -> None:
365370
"""Register contributors for the share|temporal variable.
366371
367372
Args:
368-
defining_expr: Expression with a contributor dimension.
369-
Accepts linopy LinearExpression/Variable or plain xr.DataArray (constants).
373+
defining_expr: Expression with a contributor dimension (no effect dim if effect is given).
370374
contributor_dim: Name of the element dimension to rename to 'contributor'.
375+
effect: If provided, the expression is for this specific effect (no effect dim needed).
371376
"""
372377
if contributor_dim != 'contributor':
373378
defining_expr = defining_expr.rename({contributor_dim: 'contributor'})
374379
if isinstance(defining_expr, xr.DataArray):
380+
if effect is not None:
381+
defining_expr = defining_expr.expand_dims(effect=[effect])
382+
elif 'effect' not in defining_expr.dims:
383+
raise ValueError(
384+
"DataArray contribution must have an 'effect' dimension or an explicit effect= argument."
385+
)
375386
self._temporal_constant_defs.append(defining_expr)
376387
else:
377-
self._temporal_share_defs.append(defining_expr)
388+
self._accumulate_shares(self._temporal_shares, self._as_expression(defining_expr), effect)
378389

379-
def add_periodic_contribution(self, defining_expr, contributor_dim: str = 'contributor') -> None:
390+
def add_periodic_contribution(
391+
self,
392+
defining_expr,
393+
contributor_dim: str = 'contributor',
394+
effect: str | None = None,
395+
) -> None:
380396
"""Register contributors for the share|periodic variable.
381397
382398
Args:
383-
defining_expr: Expression with a contributor dimension.
384-
Accepts linopy LinearExpression/Variable or plain xr.DataArray (constants).
399+
defining_expr: Expression with a contributor dimension (no effect dim if effect is given).
385400
contributor_dim: Name of the element dimension to rename to 'contributor'.
401+
effect: If provided, the expression is for this specific effect (no effect dim needed).
386402
"""
387403
if contributor_dim != 'contributor':
388404
defining_expr = defining_expr.rename({contributor_dim: 'contributor'})
389405
if isinstance(defining_expr, xr.DataArray):
406+
if effect is not None:
407+
defining_expr = defining_expr.expand_dims(effect=[effect])
408+
elif 'effect' not in defining_expr.dims:
409+
raise ValueError(
410+
"DataArray contribution must have an 'effect' dimension or an explicit effect= argument."
411+
)
390412
self._periodic_constant_defs.append(defining_expr)
391413
else:
392-
self._periodic_share_defs.append(defining_expr)
414+
self._accumulate_shares(self._periodic_shares, self._as_expression(defining_expr), effect)
415+
416+
@staticmethod
417+
def _accumulate_shares(
418+
accum: dict[str, list],
419+
expr: linopy.LinearExpression,
420+
effect: str | None = None,
421+
) -> None:
422+
"""Append expression to per-effect list."""
423+
# accum structure: {effect_id: [(expr, contributor_ids), ...]}
424+
if effect is not None:
425+
# Expression has no effect dim — tagged with specific effect
426+
accum.setdefault(effect, []).append(expr)
427+
elif 'effect' in expr.dims:
428+
# Expression has effect dim — split per effect (DataArray sel is cheap)
429+
for eid in expr.data.coords['effect'].values:
430+
eid_str = str(eid)
431+
accum.setdefault(eid_str, []).append(expr.sel(effect=eid, drop=True))
432+
else:
433+
raise ValueError('Expression must have effect dim or effect parameter must be given')
393434

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

545-
# === Create share|temporal variable ===
546-
if self._temporal_share_defs:
547-
self.share_temporal = self._create_share_var(self._temporal_share_defs, 'share|temporal', temporal=True)
586+
# === Create share|temporal variable (one combined with contributor × effect dims) ===
587+
if self._temporal_shares:
588+
self.share_temporal = self._create_share_var(self._temporal_shares, 'share|temporal', temporal=True)
548589
self._eq_per_timestep.lhs -= self.share_temporal.sum('contributor')
549590

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

554-
# === Create share|periodic variable ===
555-
if self._periodic_share_defs:
556-
self.share_periodic = self._create_share_var(self._periodic_share_defs, 'share|periodic', temporal=False)
557-
self._eq_periodic.lhs -= self.share_periodic.sum('contributor').reindex({'effect': self.data.effect_index})
595+
# === Create share|periodic variable (one combined with contributor × effect dims) ===
596+
if self._periodic_shares:
597+
self.share_periodic = self._create_share_var(self._periodic_shares, 'share|periodic', temporal=False)
598+
self._eq_periodic.lhs -= self.share_periodic.sum('contributor')
558599

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

574615
def _create_share_var(
575616
self,
576-
share_defs: list[linopy.LinearExpression],
617+
accum: dict[str, list[linopy.LinearExpression]],
577618
name: str,
578619
temporal: bool,
579620
) -> linopy.Variable:
580-
"""Create a share variable from registered contributor definitions.
621+
"""Create one share variable with (contributor, effect, ...) dims.
622+
623+
accum structure: {effect_id: [expr1, expr2, ...]} where each expr has
624+
(contributor, ...other_dims) dims — no effect dim.
625+
626+
Constraints are added per-effect: var.sel(effect=eid) == merged_for_eid,
627+
which avoids cross-effect alignment.
581628
582-
Aligns all contributor expressions (outer join on contributor dimension),
583-
then sums them to produce a single expression with the full contributor dimension.
629+
Returns:
630+
linopy.Variable with dims (contributor, effect, time/period).
584631
"""
585632
import pandas as pd
586633

587-
# Ensure all share defs have canonical effect order before alignment.
588-
# linopy merge uses join="override" when shapes match, which aligns by
589-
# position not label — mismatched effect order silently shuffles coefficients.
634+
if not accum:
635+
return None
636+
637+
# Collect all contributor IDs across all effects
638+
all_contributor_ids: set[str] = set()
639+
for expr_list in accum.values():
640+
for expr in expr_list:
641+
all_contributor_ids.update(str(c) for c in expr.data.coords['contributor'].values)
642+
643+
contributor_index = pd.Index(sorted(all_contributor_ids), name='contributor')
590644
effect_index = self.data.effect_index
591-
normalized = []
592-
for expr in share_defs:
593-
if 'effect' in expr.dims:
594-
expr_effects = list(expr.data.coords['effect'].values)
595-
if expr_effects != list(effect_index):
596-
expr = linopy.LinearExpression(expr.data.reindex(effect=effect_index), expr.model)
597-
normalized.append(expr)
598-
599-
aligned = linopy.align(*normalized, join='outer', fill_value=0)
600-
combined_expr = sum(aligned[1:], start=aligned[0])
601-
602-
# Extract contributor IDs from the combined expression
603-
all_ids = [str(cid) for cid in combined_expr.data.coords['contributor'].values]
604-
contributor_index = pd.Index(all_ids, name='contributor')
605645
coords = self._share_coords('contributor', contributor_index, temporal=temporal)
606-
var = self.model.add_variables(lower=-np.inf, upper=np.inf, coords=coords, name=name)
607646

608-
self.model.add_constraints(var == combined_expr, name=name)
647+
# Build mask: only create variables for (effect, contributor) combos that have expressions
648+
mask = xr.DataArray(
649+
np.zeros((len(contributor_index), len(effect_index)), dtype=bool),
650+
dims=['contributor', 'effect'],
651+
coords={'contributor': contributor_index, 'effect': effect_index},
652+
)
653+
covered_map: dict[str, list[str]] = {}
654+
for eid, expr_list in accum.items():
655+
cids = set()
656+
for expr in expr_list:
657+
cids.update(str(c) for c in expr.data.coords['contributor'].values)
658+
covered_map[eid] = sorted(cids)
659+
mask.loc[dict(effect=eid, contributor=covered_map[eid])] = True
660+
661+
var = self.model.add_variables(lower=-np.inf, upper=np.inf, coords=coords, name=name, mask=mask)
662+
663+
# Add per-effect constraints (only for covered combos)
664+
for eid, expr_list in accum.items():
665+
contributors = covered_map[eid]
666+
if len(expr_list) == 1:
667+
merged = expr_list[0].reindex(contributor=contributors)
668+
else:
669+
# Reindex all to common contributor set, then sum via linopy.merge (_term addition)
670+
aligned = [e.reindex(contributor=contributors) for e in expr_list]
671+
merged = aligned[0]
672+
for a in aligned[1:]:
673+
merged = merged + a
674+
var_slice = var.sel(effect=eid, contributor=contributors)
675+
self.model.add_constraints(var_slice == merged, name=f'{name}({eid})')
676+
677+
accum.clear()
609678
return var
610679

611680
def get_periodic(self, effect_id: str) -> linopy.Variable:

0 commit comments

Comments
 (0)