diff --git a/linopy/expressions.py b/linopy/expressions.py index 10e243de..a5b995c2 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2122,6 +2122,40 @@ def merge( 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] + # When using join='override', xr.concat places values positionally instead of + # aligning by label. We need to reindex datasets that have the same coordinate + # values but in a different order to ensure proper alignment. + if override and len(data) > 1: + reference = data[0] + aligned_data = [reference] + for ds_item in data[1:]: + reindex_dims = {} + for dim_name in reference.dims: + if dim_name in HELPER_DIMS or dim_name not in ds_item.dims: + continue + if dim_name not in reference.coords or dim_name not in ds_item.coords: + continue # pragma: no cover + ref_coord = reference.coords[dim_name].values + ds_coord = ds_item.coords[dim_name].values + # Check: same length, same set of values, but different order + if len(ref_coord) == len(ds_coord) and not np.array_equal( + ref_coord, ds_coord + ): + try: + same_values = set(ref_coord) == set(ds_coord) + except TypeError: # pragma: no cover + # Unhashable types - convert to strings for comparison + same_values = {str(v) for v in ref_coord} == { + str(v) for v in ds_coord + } + if same_values: + reindex_dims[dim_name] = reference.coords[dim_name] + if reindex_dims: + aligned_data.append(ds_item.reindex(reindex_dims)) + else: + aligned_data.append(ds_item) + data = aligned_data + if not kwargs: kwargs = { "coords": "minimal", diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index a75ace3f..f2226f81 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -1194,6 +1194,33 @@ def test_merge(x: Variable, y: Variable, z: Variable) -> None: merge(expr1, expr2) +def test_merge_with_override_and_reordered_coords(m: Model) -> None: + """Test merge with join='override' when coordinates have same values but different order.""" + import pandas as pd + + # Create variables with same coordinate values but different order + coords_a = pd.Index(["x", "y", "z"], name="dim_0") + coords_b = pd.Index(["z", "x", "y"], name="dim_0") # Same values, different order + + v1 = m.add_variables(coords=[coords_a], name="v1") + v2 = m.add_variables(coords=[coords_b], name="v2") + + expr1 = 1 * v1 + expr2 = 2 * v2 + + # Merging along _term (default) triggers the override logic because + # both expressions have the same dimension sizes + res = merge([expr1, expr2], cls=LinearExpression) + + # Verify that the coordinates match the first expression's order + assert list(res.coords["dim_0"].values) == ["x", "y", "z"] + # The result should have 2 terms (one from each expression) + assert res.nterm == 2 + # Verify the coefficients are correctly aligned (not mismatched due to positional concat) + assert res.sel(dim_0="x").coeffs.values.tolist() == [1.0, 2.0] + assert res.sel(dim_0="z").coeffs.values.tolist() == [1.0, 2.0] + + def test_linear_expression_outer_sum(x: Variable, y: Variable) -> None: expr = x + y expr2: LinearExpression = sum([x, y]) # type: ignore