feat: show multiindex levels in repr & improve multiindexed coords repr

Author: FabianHofmann

doc/release_notes.rst

@@ -6,6 +6,7 @@ Upcoming Version
 
 * When writing out an LP file, large variables and constraints are now chunked to avoid memory issues. This is especially useful for large models with constraints with many terms. The chunk size can be set with the `slice_size` argument in the `solve` function.
 * Constraints which of the form `<= infinity` and `>= -infinity` are now automatically filtered out when solving. The `solve` function now has a new argument `sanitize_infinities` to control this feature. Default is set to `True`.
+* The representation of linopy objects with multiindexed coordinates was improved to be more readable.
 
 Version 0.3.15
 --------------

linopy/common.py

@@ -613,6 +613,33 @@ def align_lines_by_delimiter(lines: list[str], delimiter: str | list[str]):
     return formatted_lines
 
 
+def get_dims_with_index_levels(
+    ds: Dataset, dims: list[Hashable] | tuple[Hashable, ...] | None = None
+) -> list[str]:
+    """
+    Get the dimensions of a Dataset with their index levels.
+
+    Example usage with a dataset that has:
+    - regular dimension 'time'
+    - multi-indexed dimension 'station' with levels ['country', 'city']
+    The output would be: ['time', 'station (country, city)']
+    """
+    dims_with_levels = []
+    if dims is None:
+        dims = list(ds.dims)
+
+    for dim in dims:
+        if isinstance(ds.indexes[dim], pd.MultiIndex):
+            # For multi-indexed dimensions, format as "dim (level0, level1, ...)"
+            names = ds.indexes[dim].names
+            dims_with_levels.append(f"{dim} ({', '.join(names)})")
+        else:
+            # For regular dimensions, just add the dimension name
+            dims_with_levels.append(str(dim))
+
+    return dims_with_levels
+
+
 def get_label_position(
     obj, values: int | np.ndarray
 ) -> (
@@ -658,10 +685,42 @@ def find_single(value: int) -> tuple[str, dict] | tuple[None, None]:
         raise ValueError("Array's with more than two dimensions is not supported")
 
 
-def print_coord(coord):
-    if isinstance(coord, dict):
-        coord = coord.values()
-    return "[" + ", ".join([str(c) for c in coord]) + "]" if len(coord) else ""
+def print_coord(coord: dict | Iterable) -> str:
+    """
+    Format coordinates into a string representation.
+
+    Args:
+        coord: Dictionary or iterable containing coordinate values.
+              Values can be numbers, strings, or nested iterables.
+
+    Returns:
+        Formatted string representation of coordinates in brackets,
+        with nested coordinates grouped in parentheses.
+
+    Examples:
+        >>> print_coord({"x": 1, "y": 2})
+        '[1, 2]'
+        >>> print_coord([1, 2, 3])
+        '[1, 2, 3]'
+        >>> print_coord([(1, 2), (3, 4)])
+        '[(1, 2), (3, 4)]'
+    """
+    # Handle empty input
+    if not coord:
+        return ""
+
+    # Extract values if input is dictionary
+    values = coord.values() if isinstance(coord, dict) else coord
+
+    # Convert each coordinate component to string
+    formatted = []
+    for value in values:
+        if isinstance(value, (list, tuple)):
+            formatted.append(f"({', '.join(str(x) for x in value)})")
+        else:
+            formatted.append(str(value))
+
+    return f"[{', '.join(formatted)}]"
 
 
 def print_single_variable(model, label):

linopy/constraints.py

@@ -6,7 +6,7 @@
 
 import functools
 import warnings
-from collections.abc import ItemsView, Iterator
+from collections.abc import Hashable, ItemsView, Iterator
 from dataclasses import dataclass
 from itertools import product
 from typing import (
@@ -35,6 +35,7 @@
     filter_nulls_polars,
     format_string_as_variable_name,
     generate_indices_for_printout,
+    get_dims_with_index_levels,
     get_label_position,
     group_terms_polars,
     has_optimized_model,
@@ -255,13 +256,20 @@ def name(self):
         return self.attrs["name"]
 
     @property
-    def coord_dims(self):
+    def coord_dims(self) -> tuple[Hashable, ...]:
         return tuple(k for k in self.dims if k not in HELPER_DIMS)
 
     @property
-    def coord_sizes(self):
+    def coord_sizes(self) -> dict[Hashable, int]:
         return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS}
 
+    @property
+    def coord_names(self) -> list[Hashable]:
+        """
+        Get the names of the coordinates.
+        """
+        return get_dims_with_index_levels(self.data, self.coord_dims)
+
     @property
     def is_assigned(self):
         return self._assigned
@@ -273,6 +281,7 @@ def __repr__(self):
         max_lines = options["display_max_rows"]
         dims = list(self.coord_sizes.keys())
         ndim = len(dims)
+        dim_names = self.coord_names
         dim_sizes = list(self.coord_sizes.values())
         size = np.prod(dim_sizes)  # that the number of theoretical printouts
         masked_entries = (~self.mask).sum().values if self.mask is not None else 0
@@ -297,16 +306,16 @@ def __repr__(self):
                         )
                         sign = SIGNS_pretty[self.sign.values[indices]]
                         rhs = self.rhs.values[indices]
-                        line = f"{print_coord(coord)}: {expr} {sign} {rhs}"
+                        line = print_coord(coord) + f": {expr} {sign} {rhs}"
                     else:
-                        line = f"{print_coord(coord)}: None"
+                        line = print_coord(coord) + ": None"
                     lines.append(line)
             lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values()))
 
-            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dims, dim_sizes))
+            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes))
             mask_str = f" - {masked_entries} masked entries" if masked_entries else ""
             underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4)
-            lines.insert(0, f"{header_string} ({shape_str}){mask_str}:\n{underscore}")
+            lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}")
         elif size == 1:
             expr = print_single_expression(self.coeffs, self.vars, 0, self.model)
             lines.append(

linopy/expressions.py

@@ -49,10 +49,12 @@
     filter_nulls_polars,
     forward_as_properties,
     generate_indices_for_printout,
+    get_dims_with_index_levels,
     get_index_map,
     group_terms_polars,
     has_optimized_model,
     iterate_slices,
+    print_coord,
     print_single_expression,
     to_dataframe,
     to_polars,
@@ -404,7 +406,8 @@ def __repr__(self) -> str:
         Print the expression arrays.
         """
         max_lines = options["display_max_rows"]
-        dims = list(self.coord_sizes.keys())
+        dims = list(self.coord_sizes)
+        dim_names = self.coord_names
         ndim = len(dims)
         dim_sizes = list(self.coord_sizes.values())
         size = np.prod(dim_sizes)  # that the number of theoretical printouts
@@ -418,26 +421,26 @@ def __repr__(self) -> str:
                 if indices is None:
                     lines.append("\t\t...")
                 else:
-                    coord_values = ", ".join(
-                        str(self.data.indexes[dims[i]][ind])
-                        for i, ind in enumerate(indices)
-                    )
+                    coord = [
+                        self.data.indexes[dims[i]][ind] for i, ind in enumerate(indices)
+                    ]
                     if self.mask is None or self.mask.values[indices]:
                         expr = print_single_expression(
                             self.coeffs.values[indices],
                             self.vars.values[indices],
                             self.const.values[indices],
                             self.model,
                         )
-                        line = f"[{coord_values}]: {expr}"
+
+                        line = print_coord(coord) + f": {expr}"
                     else:
-                        line = f"[{coord_values}]: None"
+                        line = print_coord(coord) + ": None"
                     lines.append(line)
 
-            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dims, dim_sizes))
+            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes))
             mask_str = f" - {masked_entries} masked entries" if masked_entries else ""
             underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4)
-            lines.insert(0, f"{header_string} ({shape_str}){mask_str}:\n{underscore}")
+            lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}")
         elif size == 1:
             expr = print_single_expression(
                 self.coeffs, self.vars, self.const, self.model
@@ -735,9 +738,13 @@ def coord_dims(self) -> tuple[Hashable, ...]:
         return tuple(k for k in self.dims if k not in HELPER_DIMS)
 
     @property
-    def coord_sizes(self) -> dict[str, int]:
+    def coord_sizes(self) -> dict[Hashable, int]:
         return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS}
 
+    @property
+    def coord_names(self) -> list[Hashable]:
+        return get_dims_with_index_levels(self.data, self.coord_dims)
+
     @property
     def vars(self):
         return self.data.vars

linopy/variables.py

@@ -39,6 +39,7 @@
     filter_nulls_polars,
     format_string_as_variable_name,
     generate_indices_for_printout,
+    get_dims_with_index_levels,
     get_label_position,
     has_optimized_model,
     is_constant,
@@ -50,7 +51,7 @@
     to_polars,
 )
 from linopy.config import options
-from linopy.constants import TERM_DIM
+from linopy.constants import HELPER_DIMS, TERM_DIM
 from linopy.solvers import set_int_index
 from linopy.types import NotImplementedType
 
@@ -321,7 +322,8 @@ def __repr__(self) -> str:
         Print the variable arrays.
         """
         max_lines = options["display_max_rows"]
-        dims = list(self.sizes.keys())
+        dims = list(self.sizes)
+        dim_names = self.coord_names
         dim_sizes = list(self.sizes.values())
         masked_entries = (~self.mask).sum().values
         lines = []
@@ -343,7 +345,7 @@ def __repr__(self) -> str:
                     lines.append(line)
             # lines = align_lines_by_delimiter(lines, "∈")
 
-            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dims, dim_sizes))
+            shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes))
             mask_str = f" - {masked_entries} masked entries" if masked_entries else ""
             lines.insert(
                 0,
@@ -663,6 +665,21 @@ def type(self):
         else:
             return "Continuous Variable"
 
+    @property
+    def coord_dims(self) -> tuple[Hashable, ...]:
+        return tuple(k for k in self.dims if k not in HELPER_DIMS)
+
+    @property
+    def coord_sizes(self) -> dict[Hashable, int]:
+        return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS}
+
+    @property
+    def coord_names(self) -> list[Hashable]:
+        """
+        Get the names of the coordinates.
+        """
+        return get_dims_with_index_levels(self.data, self.coord_dims)
+
     @property
     def range(self) -> tuple[int, int]:
         """
@@ -1089,8 +1106,6 @@ def __getitem__(self, keys: Any) -> ScalarVariable:
         keys = keys if isinstance(keys, tuple) else (keys,)
         object = self.object
 
-        if not all(map(pd.api.types.is_scalar, keys)):
-            raise ValueError("Only scalar keys are allowed.")
         # return single scalar
         if not object.labels.ndim:
             return ScalarVariable(object.labels.item(), object.model)

test/test_common.py

@@ -15,6 +15,7 @@
     as_dataarray,
     assign_multiindex_safe,
     best_int,
+    get_dims_with_index_levels,
     iterate_slices,
 )
 
@@ -529,3 +530,77 @@ def test_iterate_slices_no_slice_dims():
     for s in slices:
         assert isinstance(s, xr.Dataset)
         assert set(s.dims) == set(ds.dims)
+
+
+def test_get_dims_with_index_levels():
+    # Create test data
+
+    # Case 1: Simple dataset with regular dimensions
+    ds1 = xr.Dataset(
+        {"temp": (("time", "lat"), np.random.rand(3, 2))},  # noqa: NPY002
+        coords={"time": pd.date_range("2024-01-01", periods=3), "lat": [0, 1]},
+    )
+
+    # Case 2: Dataset with a multi-index dimension
+    stations_index = pd.MultiIndex.from_product(
+        [["USA", "Canada"], ["NYC", "Toronto"]], names=["country", "city"]
+    )
+    stations_coords = xr.Coordinates.from_pandas_multiindex(stations_index, "station")
+    ds2 = xr.Dataset(
+        {"temp": (("time", "station"), np.random.rand(3, 4))},  # noqa: NPY002
+        coords={"time": pd.date_range("2024-01-01", periods=3), **stations_coords},
+    )
+
+    # Case 3: Dataset with unnamed multi-index levels
+    unnamed_stations_index = pd.MultiIndex.from_product(
+        [["USA", "Canada"], ["NYC", "Toronto"]]
+    )
+    unnamed_stations_coords = xr.Coordinates.from_pandas_multiindex(
+        unnamed_stations_index, "station"
+    )
+    ds3 = xr.Dataset(
+        {"temp": (("time", "station"), np.random.rand(3, 4))},  # noqa: NPY002
+        coords={
+            "time": pd.date_range("2024-01-01", periods=3),
+            **unnamed_stations_coords,
+        },
+    )
+
+    # Case 4: Dataset with multiple multi-indexed dimensions
+    locations_index = pd.MultiIndex.from_product(
+        [["North", "South"], ["A", "B"]], names=["region", "site"]
+    )
+    locations_coords = xr.Coordinates.from_pandas_multiindex(
+        locations_index, "location"
+    )
+
+    ds4 = xr.Dataset(
+        {"temp": (("time", "station", "location"), np.random.rand(2, 4, 4))},  # noqa: NPY002
+        coords={
+            "time": pd.date_range("2024-01-01", periods=2),
+            **stations_coords,
+            **locations_coords,
+        },
+    )
+
+    # Run tests
+
+    # Test case 1: Regular dimensions
+    assert get_dims_with_index_levels(ds1) == ["time", "lat"]
+
+    # Test case 2: Named multi-index
+    assert get_dims_with_index_levels(ds2) == ["time", "station (country, city)"]
+
+    # Test case 3: Unnamed multi-index
+    assert get_dims_with_index_levels(ds3) == [
+        "time",
+        "station (station_level_0, station_level_1)",
+    ]
+
+    # Test case 4: Multiple multi-indices
+    expected = ["time", "station (country, city)", "location (region, site)"]
+    assert get_dims_with_index_levels(ds4) == expected
+
+    # Test case 5: Empty dataset
+    ds5 = xr.Dataset()
+    assert get_dims_with_index_levels(ds5) == []