Ordering of Coords in Repr

[ianhi]

What is your issue?

Problem

This preferentially placing dimension coordinates first in the order. Which makes a lot of sense. But has a negative side effect for the visualization non-dimension coords that share a dimension wiht a dim coord. e..g here word coord and the other coords sharing the word dim are disconnected:

Or a simple example:

import numpy as np

ds = xr.Dataset(
    {"data": (("a", "b", "c"), np.random.rand(3, 3, 3))},
    coords={"a": [0, 1, 2], "b": [0, 1, 2], "c": [3, 2, 1]},
)

ds = ds.assign_coords({"a_1":('a', [5,6,7])})
ds = ds.assign_coords({"b_1":('b', [5,6,7])})
ds = ds.assign_coords({"c_1":('c', [5,6,7])})
print(ds)

<xarray.Dataset> Size: 360B
Dimensions:  (a: 3, b: 3, c: 3)
Coordinates:
  * a        (a) int64 24B 0 1 2
  * b        (b) int64 24B 0 1 2
  * c        (c) int64 24B 3 2 1
    a_1      (a) int64 24B 5 6 7
    b_1      (b) int64 24B 5 6 7
    c_1      (c) int64 24B 5 6 7
Data variables:
    data     (a, b, c) float64 216B 0.2621 0.7528 0.7416 ... 0.3515 0.1 0.2662

Solution

The order of coords in the repr is currently defined here:

xarray/xarray/core/formatting.py

Lines 450 to 452 in c374754

	dim_ordered_coords = sorted(
	coords.items(), key=lambda x: dims.index(x[0]) if x[0] in dims else len(dims)
	)

We could (possibly optionally) change the sort key to something like this:

def _coord_sort_key(name, var, dims):
    """Sort key for coordinate ordering.

        Orders by:
        1. Primary: index of the first 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.
    """
    # Dimension coordinates sort by their position in dims, come first (0)
    if name in dims:
        return (dims.index(name), 0)

    # Non-dimension coordinates sort by their first dim, come second (1)
    for d in var.dims:
        if d in dims:
            return (dims.index(d), 1)

    # Scalar coords or coords with dims not in dataset dims go at end
    return (len(dims), 1)

which gives

<xarray.Dataset> Size: 360B
Dimensions:  (a: 3, b: 3, c: 3)
Coordinates:
  * a        (a) int64 24B 0 1 2
    a_1      (a) int64 24B 5 6 7
  * b        (b) int64 24B 0 1 2
    b_1      (b) int64 24B 5 6 7
  * c        (c) int64 24B 3 2 1
    c_1      (c) int64 24B 5 6 7
Data variables:
    data     (a, b, c) float64 216B 0.0534 0.7124 0.705 ... 0.5102 0.563 0.1605

or potentially going even farther and indenting the non-dimension coords

[sshekhar563]

Can you please assign this issue to me?

[ianhi]

Hi @sshekhar563 since this is an open-source project that anyone can contribute to we don't explicitly assign issues. Anyone is allowed to work on something and open a PR with their proposed changes - you don't need to ask for permission :)

That said I would be cautious about going straight to opening a PR for this issue. In this case the hard work is not the implementation in code, but deciding what we want the behavior to be. That's why I opened an issue rather than a PR so people can chime in on what the tradeoffs are, and if they think this is a good or bad idea.

If you are looking for something to work on I'd suggest searching for the issues tagged with contrib-help-wanted and or contrib-good-first-issue

[sshekhar563]

@ianhi Thanks for the clarification, that makes sense.

I agree that the main challenge here is deciding on the desired behavior rather than the implementation itself. The proposed ordering that groups non-dimension coordinates right after their associated dimension seems more readable to me, but I understand there may be tradeoffs to consider.

I’m happy to follow the discussion and help test or evaluate different options if that would be useful.

[keewis]

moving the dimensions first is a recent addition (#10778), and had been requested a couple of times before (e.g. #712, #4515). Going back is thus not really possible, but adding (yet another) option may be possible. Which means we'd have to figure out whether that's desirable, and if so how that should look: should the option take a set of order names, or also a key function for use with sorted? The latter would support domain-specific orderings.

Indenting the non-dim coords may be hard for coords with multiple dimensions (we could restrict to 1D coords, but that feels kinda arbitrary).

[jsignell]

Yeah given that this is a recent change I wonder if it makes sense to just leave this open for a while and see if it collects 👍🏻 or if people raise duplicates. It could be that the prior behavior (no sort) was actually more intuitive in which case #10778 could just be reverted.

[benbovy]

I don't think @ianhi's solution proposed here is really going back?

IIUC the requests behind #10778 were mainly about showing coordinates in an order that is consistent with the dimensions shown in the repr, which still prevails with the solution proposed here.

#4515 explicitly mentions dimension coordinates but it was opened before the Xarray index refactor (i.e., before it was possible to have non-dimension coordinate with an index), and looking at the example repr shown in that issue I guess the original complain was more related to the no sort at all than specifically about dimension vs. non-dimension coordinates (cc @dcherian).

I actually missed #10778, sorry. I like the change, although what's more annoying to me as a side effect is that coordinates associated with the same index might not always be displayed next to each other anymore:

>>> import xarray as xr

>>> ds = xr.Dataset(coords={"x": [1, 2], "y": [3, 4], "crs": 0})

>>> class CustomIndex(xr.Index):
...     @classmethod
...     def from_variables(cls, variables, options):
...         return cls()
... 

>>> ds.drop_indexes("x").set_xindex(["x", "crs"], CustomIndex)
<xarray.Dataset> Size: 40B
Dimensions:  (x: 2, y: 2)
Coordinates:
  * x        (x) int64 16B 1 2
  * y        (y) int64 16B 3 4
  * crs      int64 8B 0
Data variables:
    *empty*
Indexes:
  ┌ x        CustomIndex
  └ crs

Same for PandasMultiIndex:

>>> import pandas as pd
>>> midx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("foo", "bar"))
>>> ds = xr.Dataset(coords=xr.Coordinates.from_pandas_multiindex(midx, "x"))
>>> ds.assign_coords(y=[1, 2])
<xarray.Dataset> Size: 112B
Dimensions:  (x: 4, y: 2)
Coordinates:
  * x        (x) object 32B MultiIndex
  * y        (y) int64 16B 1 2
  * foo      (x) object 32B 'a' 'a' 'b' 'b'
  * bar      (x) int64 32B 1 2 1 2
Data variables:
    *empty*

---

I'd be +1 for a solution like the one proposed here, and even thinking further:

• primary: show 1-dimensional coordinates first, then n-dimensional coordinates and finally scalar coordinates
• secondary: show 1-dimensional coordinates ordered like the Dataset / DataArray dimensions
• tertiary: show index-coordinates before non-index coordinates ; if a coordinate is associated with a multi-coordinate Xarray index then show the other coordinates associated to that index just below (in the order defined by the index) regardless of their dimensions

This is getting complicated, though... It feels like opening a can of worms. Other users may want to sort coordinates in alphabetical order... Or for a raster DataArray with (y, x) dimension keep showing coordinates ordered as x then y would still make a bit more sense...

I'm wondering if reverting #10778 wouldn't be easier after all, and let users care about the coordinate order.

[jsignell]

I guess I have a mild preference for adding more sorting logic rather than reverting since I haven't seen anyone asking for alphabetical order or anything like that.

I agree that handling custom indexes adds complexity, but the multi-index case at least is covered by Ian's suggestion.

[jsignell]

Anyways I put up a PR (#11091) with a slightly modified version of @ianhi's suggestion. I think it's an improvement over main and an improvement over unsorted.

[benbovy]

Anyways I put up a PR (#11091) with a slightly modified version of @ianhi's suggestion. I think it's an improvement over main and an improvement over unsorted.

Indeed, thanks @jsignell.

Grouping coordinates by index in the repr still seems sensible to me... In fact #11091 should already work with most multi-coordinate indexes that I'm aware of since those are associated with coordinates that either all have the same dimensions or are orthogonal / adjacent dimension coordinates (e.g., PandasMultiIndex, NDPointIndex, RasterIndex). Examples where this may still not work well are rolodex.ForecastIndex or the not yet available CFIntervalIndex (#10296) and UgridIndex.

Anyway, I guess we'll be able to further improve the coordinate ordering for the latter cases in a backward compatible way, if needed.