Skip to content

groupby() by the name of a non-dimension coordinate raises (xarray parity) #750

@FBumann

Description

@FBumann

Bug

Grouping a LinearExpression by an existing non-dimension coordinate fails, even though the equivalent xarray.Dataset.groupby works. Both the string form (groupby("name")) and passing the coordinate DataArray raise.

Reproducer

import pandas as pd, xarray as xr, linopy

m = linopy.Model()
snapshots = pd.RangeIndex(4, name="snapshot")
period = xr.DataArray(
    [2020, 2020, 2030, 2030], dims="snapshot",
    coords={"snapshot": snapshots}, name="period",
)
x = m.add_variables(coords=[snapshots], name="x")
expr = (1 * x).assign_coords(period=period)   # attach 'period' as an aux coord

expr.groupby("period").sum()   # ValueError: period already exists as coordinate or variable name.
expr.groupby(period).sum()     # KeyError: 'period'

Both should group by the period coordinate and return a period-dimmed expression, the same way xr.Dataset.groupby("period").sum() does.

Current workaround

You have to detach the coordinate first, and (depending on the form) rename it:

expr.drop_vars("period").groupby(period).sum()                          # works
expr.drop_vars("period").groupby(period.rename("inv_period")).sum()     # also works

Cause

In LinearExpressionGroupby.sum() (linopy/expressions.py):

  1. A string group is not in non_fallback_types, so groupby("period") skips the fast path and hits the xarray fallback, which fails re-attaching the non-dimension coordinate.
  2. The fast path builds names_to_drop from the dimension index only (coords.name), never the aux coordinate being grouped. It then does ds.rename({GROUP_DIM: final_group_name}), which collides with the period coordinate that is still attached → the ValueError.

Fix sketch

In sum(): (a) resolve a string group that names an existing coordinate to self.data[name] so it takes the fast path, and (b) add final_group_name to names_to_drop when it is still an attached coordinate. PR to follow.

Environment

linopy master (a74724f), reproduced in a clean worktree.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions