Support non-string dimension/variable names

[joshburkart]

Problem description

Currently, it appears that "dimension"/"coordinate" labels must be strings. However, in more rigorous software engineering applications it is often desirable to use something more organized/structured for labels, e.g. enums. I think it would be great if xarray supported this.

Obviously storing to e.g. NetCDF necessitates string-valued field names, so I would think calling str could be appropriate when performing this sort of serialization. This is what pandas seems to do (see below). But I imagine there might be other issues that would need to be resolved to do what I'm suggesting...?

Code sample

import enum

import numpy as np
import pandas as pd
import xarray as xr

class CoordId(enum.Enum):
    LAT = 'lat'
    LON = 'lon'

pd.DataFrame({CoordId.LAT: [1,2,3]}).to_csv()
# Returns: ',CoordId.LAT\n0,1\n1,2\n2,3\n'

xr.DataArray(
    data=np.arange(3 * 2).reshape(3, 2),
    coords={CoordId.LAT: [1, 2, 3], CoordId.LON: [7, 8]},
    dims=[CoordId.LAT, CoordId.LON],
)
# Fails: TypeError: dimension CoordId.LAT is not a string

Output of xr.show_versions()

INSTALLED VERSIONS ------------------ commit: None python: 3.6.5.final.0 python-bits: 64 OS: Linux OS-release: 4.15.0-1010-gcp machine: x86_64 processor: x86_64 byteorder: little LC_ALL: None LANG: C.UTF-8 LOCALE: en_US.UTF-8

xarray: 0.10.7
pandas: 0.23.1
numpy: 1.14.5
scipy: 1.1.0
netCDF4: 1.3.1
h5netcdf: 0.5.0
h5py: 2.7.1
Nio: None
zarr: None
bottleneck: None
cyordereddict: None
dask: None
distributed: None
matplotlib: 2.1.1
cartopy: 0.16.0
seaborn: None
setuptools: 39.2.0
pip: 9.0.1
conda: None
pytest: 3.6.1
IPython: 6.4.0
sphinx: None

[shoyer]

Hi @joshburkart -- thanks for raising this concern.

I agree, it would be nice to support enums (really any hashable value) as dimension names. Our current checks for strings are somewhat inconsistent, e.g., you can actually use these in an xarray.Dataset if you use the fully explicit API for constructing a dataset:

ds = xr.Dataset(
    data_vars={'foo': ((CoordId.LAT, CoordId.LON), np.arange(3 * 2).reshape(3, 2))},
    coords={CoordId.LAT: ((CoordId.LAT,), [1, 2, 3]),
            CoordId.LON: ((CoordId.LON,), [7, 8])},
)

But now if you try to print the dataset, you get an error about sorting:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/IPython/core/formatters.py in __call__(self, obj)
    697                 type_pprinters=self.type_printers,
    698                 deferred_pprinters=self.deferred_printers)
--> 699             printer.pretty(obj)
    700             printer.flush()
    701             return stream.getvalue()

/usr/local/lib/python3.6/dist-packages/IPython/lib/pretty.py in pretty(self, obj)
    396                             if callable(meth):
    397                                 return meth(obj, self, cycle)
--> 398             return _default_pprint(obj, self, cycle)
    399         finally:
    400             self.end_group()

/usr/local/lib/python3.6/dist-packages/IPython/lib/pretty.py in _default_pprint(obj, p, cycle)
    516     if _safe_getattr(klass, '__repr__', None) not in _baseclass_reprs:
    517         # A user-provided repr. Find newlines and replace them with p.break_()
--> 518         _repr_pprint(obj, p, cycle)
    519         return
    520     p.begin_group(1, '<')

/usr/local/lib/python3.6/dist-packages/IPython/lib/pretty.py in _repr_pprint(obj, p, cycle)
    707     """A pprint that just redirects to the normal repr function."""
    708     # Find newlines and replace them with p.break_()
--> 709     output = repr(obj)
    710     for idx,output_line in enumerate(output.splitlines()):
    711         if idx:

/usr/local/lib/python3.6/dist-packages/xarray/core/formatting.py in __repr__(self)
     62 
     63     def __repr__(self):
---> 64         return ensure_valid_repr(self.__unicode__())
     65 
     66 

/usr/local/lib/python3.6/dist-packages/xarray/core/dataset.py in __unicode__(self)
   1188 
   1189     def __unicode__(self):
-> 1190         return formatting.dataset_repr(self)
   1191 
   1192     def info(self, buf=None):

/usr/local/lib/python3.6/dist-packages/xarray/core/formatting.py in dataset_repr(ds)
    415 
    416     dims_start = pretty_print(u'Dimensions:', col_width)
--> 417     summary.append(u'%s(%s)' % (dims_start, dim_summary(ds)))
    418 
    419     if ds.coords:

/usr/local/lib/python3.6/dist-packages/xarray/core/formatting.py in dim_summary(obj)
    324 
    325 def dim_summary(obj):
--> 326     elements = [u'%s: %s' % (k, v) for k, v in obj.sizes.items()]
    327     return u', '.join(elements)
    328 

/usr/local/lib/python3.6/dist-packages/xarray/core/formatting.py in <listcomp>(.0)
    324 
    325 def dim_summary(obj):
--> 326     elements = [u'%s: %s' % (k, v) for k, v in obj.sizes.items()]
    327     return u', '.join(elements)
    328 

/usr/lib/python3.6/_collections_abc.py in __iter__(self)
    741 
    742     def __iter__(self):
--> 743         for key in self._mapping:
    744             yield (key, self._mapping[key])
    745 

/usr/local/lib/python3.6/dist-packages/xarray/core/utils.py in __iter__(self)
    311 
    312     def __iter__(self):
--> 313         return iter(self.mapping)
    314 
    315     def __len__(self):

/usr/local/lib/python3.6/dist-packages/xarray/core/utils.py in __iter__(self)
    347 
    348     def __iter__(self):
--> 349         return iter(sorted(self.mapping))
    350 
    351     def __len__(self):

TypeError: '<' not supported between instances of 'CoordId' and 'CoordId'

I would be open to PRs to improve the situation.

[joshburkart]

Thanks @shoyer. I'll see if I can take a look in the near future...

[shoyer]

It would be an improvement even to clearly document the requirements for dimension/variable names.

I suspect we don't actually need them to be sortable, though we do using sorting as part of the current repr() for some xarray objecs. This is mostly to ensure reproducible displays across multiple loads/runs of a file, but it's increasingly less relevant now than Python's dict is ordered by default (since Python 3.6).

[shoyer]

Another choice would be to intentionally simplify xarray's data model and not allow anything other than strings for variable/dimension names.

[joshburkart]

Some options that come to mind:

1. Allow any object with a __str__ method to be supplied as a variable/dimension label, but then delegate all internal sorting/printing/etc. logic to str(label).
2. Just implicitly run str on everything a user tries to input as a label (both when creating an xarray object and when accessing fields from an existing object), so that only strings are used internally.
3. Put this on the user, and only allow calling xarray objects/methods/etc. with labels already strings, as you suggested @shoyer.

I dunno. Whatever the maintainers think is best? (3) seems least complex on the xarray side, but (1) or (2) might be more convenient for users.

[shoyer]

(2) (1) seems like a pretty decent option to me. It's compatible with Python's duck-typing philosophy, and we don't really need string variable/dimension names for anything other than various serialization formats like netCDF. So the full requirement for names would be "Hashable, can be coerced with str() and not None" (we use None as a sentinel value to indicate the lack of a name in xarray).

CC @pydata/xarray in case anyone else has opinions.

[joshburkart]

Just to clarify @shoyer, you said (2) sounds best to you, but your other comments (e.g. duck typing, requiring hashable) seem to describe (1)...? Slightly confused...

[shoyer]

Oops, I was reading a little too quickly. I did indeed mean (1) above. The one thing I would emphasize is that we don't actually want to check for something like hasattr(dim, '__str__') if possible, but rather just call str(dim). (Though I guess Python's object type defines a default __str__ method, so pretty much everything will pass that test.)

[ttung]

We're using xarray in a project that is encouraging use of python typing, and we too would like to use enums as data dimension names. How do you feel about using a base class that data dimension classes need to subclass?

Here's a really simple proof-of-concept (though not very thorough, as it would certainly fail serialization): ttung@8e623eb

In [1]: import xarray as xr

In [2]: import numpy as np

In [5]: from enum import Enum

In [6]: class A(xr.core.dataarray.DimensionBase, Enum):
   ...:     X = "abc"
   ...:     Y = "def"
   ...:     Z = "ghi"
   ...:     

In [7]: a = xr.DataArray(np.random.randint(0, 255, size=(4, 3, 5)), dims=[A.X, A.Y, A.Z])

In [8]: a[A.X]
Out[8]: 
<xarray.DataArray <A.X: 'abc'> (A.X: 4)>
array([0, 1, 2, 3])
Dimensions without coordinates: A.X

In [9]: a.max(A.X)
Out[9]: 
<xarray.DataArray (A.Y: 3, A.Z: 5)>
array([[254, 226, 181, 191, 233],
       [139, 195, 212, 167, 169],
       [191, 241, 199, 174, 208]])
Dimensions without coordinates: A.Y, A.Z

In [10]: 

[shoyer]

Most of the places in the code where we do the isinstance(obj, string) checks are where we allow passing in a single string as a convenient shortcut to a list of names. So I'm not sure it's really essential to allow flexible types there.

That said, checking explicitly for strings wasn't take a careful API choice. I could see a case for replacing all these "as sequence" casts with something more generic, e.g., based on checking explicitly for more generic scalar types. Certainly enums should be scalars.

If possible, I would rather stick to duck typing for any requirements we put on names. Base classes don't feel terribly Pythonic.

[ttung]

The problem with generic scalar types is that it wouldn't work after serialization and deserialization (which would presumably go to strings). My suggestion has the advantage of being able to create a __eq__ method in the base class that would match both the object itself or its string equivalent, so that one could use the scalar type even after ser/deser. I disagree that base classes aren't very pythonic.

However, I think (1)/(2) are both reasonable solution (in fact, they seem to be identical except for when you call str). It has its warts, as even a mutable sequence would pass muster. :)

If that's the direction you'd like to see the project go towards, I'd be happy to take a stab at it.