Better ordering of coords in repr (pydata#11091)

Author: jsignell

doc/whats-new.rst

@@ -36,6 +36,8 @@ New Features
   of the same name are compared for potential conflicts when performing binary operations. The default for it is
   ``arithmetic_compat='minimal'`` which matches the existing behaviour.
   By `Matthew Willson <https://github.com/mjwillson>`_.
+- Better ordering of coordinates when displaying Xarray objects. (:pull:`11098`).
+  By `Ian Hunt-Isaak <https://github.com/ianhi>`_, `Julia Signell <https://github.com/jsignell>`_.
 
 Breaking Changes
 ~~~~~~~~~~~~~~~~

xarray/core/common.py

@@ -619,9 +619,9 @@ def assign_coords(
         <xarray.Dataset> Size: 360B
         Dimensions:         (x: 2, y: 2, time: 4)
         Coordinates:
-          * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
             lon             (x, y) float64 32B 260.2 260.7 260.2 260.8
             lat             (x, y) float64 32B 42.25 42.21 42.63 42.59
+          * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
             reference_time  datetime64[ns] 8B 2014-09-05
         Dimensions without coordinates: x, y
         Data variables:
@@ -633,9 +633,9 @@ def assign_coords(
         <xarray.Dataset> Size: 360B
         Dimensions:         (x: 2, y: 2, time: 4)
         Coordinates:
-          * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
             lon             (x, y) float64 32B -99.83 -99.32 -99.79 -99.23
             lat             (x, y) float64 32B 42.25 42.21 42.63 42.59
+          * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
             reference_time  datetime64[ns] 8B 2014-09-05
         Dimensions without coordinates: x, y
         Data variables:

xarray/core/dataarray.py

@@ -369,9 +369,9 @@ class DataArray(
            [[22.60070734, 13.78914233, 14.17424919],
             [18.28478802, 16.15234857, 26.63418806]]])
     Coordinates:
-      * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
         lon             (x, y) float64 32B -99.83 -99.32 -99.79 -99.23
         lat             (x, y) float64 32B 42.25 42.21 42.63 42.59
+      * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
         reference_time  datetime64[ns] 8B 2014-09-05
     Dimensions without coordinates: x, y
     Attributes:
@@ -2807,8 +2807,8 @@ def set_index(
                [1., 1., 1.]])
         Coordinates:
           * x        (x) int64 16B 0 1
-          * y        (y) int64 24B 0 1 2
             a        (x) int64 16B 3 4
+          * y        (y) int64 24B 0 1 2
         >>> arr.set_index(x="a")
         <xarray.DataArray (x: 2, y: 3)> Size: 48B
         array([[1., 1., 1.],
@@ -5964,8 +5964,8 @@ def pad(
                [nan, nan, nan, nan]])
         Coordinates:
           * x        (x) float64 32B nan 0.0 1.0 nan
-          * y        (y) int64 32B 10 20 30 40
             z        (x) float64 32B nan 100.0 200.0 nan
+          * y        (y) int64 32B 10 20 30 40
 
         Careful, ``constant_values`` are coerced to the data type of the array which may
         lead to a loss of precision:
@@ -5978,8 +5978,8 @@ def pad(
                [ 1,  1,  1,  1]])
         Coordinates:
           * x        (x) float64 32B nan 0.0 1.0 nan
-          * y        (y) int64 32B 10 20 30 40
             z        (x) float64 32B nan 100.0 200.0 nan
+          * y        (y) int64 32B 10 20 30 40
         """
         ds = self._to_temp_dataset().pad(
             pad_width=pad_width,

xarray/core/dataset.py

@@ -322,10 +322,10 @@ class Dataset(
     <xarray.Dataset> Size: 552B
     Dimensions:         (loc: 2, instrument: 3, time: 4)
     Coordinates:
-      * instrument      (instrument) <U8 96B 'manufac1' 'manufac2' 'manufac3'
-      * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
         lon             (loc) float64 16B -99.83 -99.32
         lat             (loc) float64 16B 42.25 42.21
+      * instrument      (instrument) <U8 96B 'manufac1' 'manufac2' 'manufac3'
+      * time            (time) datetime64[ns] 32B 2014-09-06 ... 2014-09-09
         reference_time  datetime64[ns] 8B 2014-09-05
     Dimensions without coordinates: loc
     Data variables:
@@ -4733,8 +4733,8 @@ def set_index(
         Dimensions:  (x: 2, y: 3)
         Coordinates:
           * x        (x) int64 16B 0 1
-          * y        (y) int64 24B 0 1 2
             a        (x) int64 16B 3 4
+          * y        (y) int64 24B 0 1 2
         Data variables:
             v        (x, y) float64 48B 1.0 1.0 1.0 1.0 1.0 1.0
         >>> ds.set_index(x="a")
@@ -8713,9 +8713,9 @@ def filter_by_attrs(self, **kwargs) -> Self:
         <xarray.Dataset> Size: 192B
         Dimensions:         (x: 2, y: 2, time: 3)
         Coordinates:
-          * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
             lon             (x, y) float64 32B -99.83 -99.32 -99.79 -99.23
             lat             (x, y) float64 32B 42.25 42.21 42.63 42.59
+          * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
             reference_time  datetime64[ns] 8B 2014-09-05
         Dimensions without coordinates: x, y
         Data variables:
@@ -8728,9 +8728,9 @@ def filter_by_attrs(self, **kwargs) -> Self:
         <xarray.Dataset> Size: 288B
         Dimensions:         (x: 2, y: 2, time: 3)
         Coordinates:
-          * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
             lon             (x, y) float64 32B -99.83 -99.32 -99.79 -99.23
             lat             (x, y) float64 32B 42.25 42.21 42.63 42.59
+          * time            (time) datetime64[ns] 24B 2014-09-06 2014-09-07 2014-09-08
             reference_time  datetime64[ns] 8B 2014-09-05
         Dimensions without coordinates: x, y
         Data variables:

xarray/core/formatting.py

@@ -443,12 +443,38 @@ def _mapping_repr(
 )
 
 
+def _coord_sort_key(coord, dims):
+    """Sort key for coordinate ordering.
+
+        Orders by:
+        1. Primary: index of the matching dimension in dataset dims.
+        2. Secondary: dimension coordinates (name == dim) come before non-dimension coordinates
+
+        This groups non-dimension coordinates right after their associated dimension
+    coordinate.
+    """
+    name, var = coord
+
+    # Dimension coordinates come first within their dim section
+    if name in dims:
+        return (dims.index(name), 0)
+
+    # Non-dimension coordinates come second within their dim section
+    # Check the var.dims list in backwards order to put (x, y) after (x) and (y)
+    for d in var.dims[::-1]:
+        if d in dims:
+            return (dims.index(d), 1)
+
+    # Scalar coords or coords with dims not in dataset dims go at the end
+    return (len(dims), 1)
+
+
 def coords_repr(coords: AbstractCoordinates, col_width=None, max_rows=None):
     if col_width is None:
         col_width = _calculate_col_width(coords)
     dims = tuple(coords._data.dims)
     dim_ordered_coords = sorted(
-        coords.items(), key=lambda x: dims.index(x[0]) if x[0] in dims else len(dims)
+        coords.items(), key=functools.partial(_coord_sort_key, dims=dims)
     )
     return _mapping_repr(
         dict(dim_ordered_coords),

xarray/core/formatting_html.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING
 
 from xarray.core.formatting import (
+    _coord_sort_key,
     filter_nondefault_indexes,
     inherited_vars,
     inline_index_repr,
@@ -120,7 +121,7 @@ def summarize_coords(variables) -> str:
     li_items = []
     dims = tuple(variables._data.dims)
     dim_ordered_coords = sorted(
-        variables.items(), key=lambda x: dims.index(x[0]) if x[0] in dims else len(dims)
+        variables.items(), key=partial(_coord_sort_key, dims=dims)
     )
     for k, v in dim_ordered_coords:
         li_content = summarize_variable(k, v, is_index=k in variables.xindexes)

xarray/structure/concat.py

@@ -239,8 +239,8 @@ def concat(
     array([[0, 1, 2],
            [3, 4, 5]])
     Coordinates:
-      * y        (y) int64 24B 10 20 30
         x        (new_dim) <U1 8B 'a' 'b'
+      * y        (y) int64 24B 10 20 30
     Dimensions without coordinates: new_dim
 
     >>> xr.concat(
@@ -253,8 +253,8 @@ def concat(
            [3, 4, 5]])
     Coordinates:
       * new_dim  (new_dim) int64 16B -90 -100
-      * y        (y) int64 24B 10 20 30
         x        (new_dim) <U1 8B 'a' 'b'
+      * y        (y) int64 24B 10 20 30
 
     # Concatenate a scalar variable along a new dimension of the same name with and without creating a new index
 

xarray/tests/test_dataset.py

@@ -294,8 +294,8 @@ def test_repr(self) -> None:
             Coordinates:
               * dim2     (dim2) float64 72B 0.0 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0
               * dim3     (dim3) {data["dim3"].dtype} 40B 'a' 'b' 'c' 'd' 'e' 'f' 'g' 'h' 'i' 'j'
-              * time     (time) datetime64[ns] 160B 2000-01-01 2000-01-02 ... 2000-01-20
                 numbers  (dim3) int64 80B 0 1 2 0 0 1 1 2 2 3
+              * time     (time) datetime64[ns] 160B 2000-01-01 2000-01-02 ... 2000-01-20
             Dimensions without coordinates: dim1
             Data variables:
                 var1     (dim1, dim2) float64 576B -0.9891 -0.3678 1.288 ... -0.2116 0.364
@@ -875,8 +875,8 @@ def test_coords_properties(self) -> None:
             """\
         Coordinates:
           * x        (x) int64 16B -1 -2
-          * y        (y) int64 24B 0 1 2
             a        (x) int64 16B 4 5
+          * y        (y) int64 24B 0 1 2
             b        int64 8B -10"""
         )
         actual = repr(coords)

xarray/tests/test_datatree.py

@@ -644,8 +644,8 @@ def test_properties(self) -> None:
             """\
         Coordinates:
           * x        (x) int64 16B -1 -2
-          * y        (y) int64 24B 0 1 2
             a        (x) int64 16B 4 5
+          * y        (y) int64 24B 0 1 2
             b        int64 8B -10"""
         )
         actual = repr(coords)