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.