Naming dimensions in python?

Question

There is something that I would very appreciate, it is the ability to name the dimensions in an array in python. For example I have a numpy array with 3 dimensions and I will regularly have to sum it along a specific dimensions.

So I can do with a ndarray a:

sum(a, axis=2)

if my relevant dimension is the last one, but I want to make it "position independent", i.e. a user can provide any array, as long as he specifies "this dimension is "DI" " (example, for "Dimension of Interest"). So basically I would like to be able to write:

sum(a, axis="DI")

Close to NETCDF, but I don't want to implement a whole netcdf capability.

Answer 1

@M456's idea is clever, but if you have the same naming scheme for several arrays, I think the simpler solution would be just to use a dictionary:

axes = { 'DA': 0, 'DB':1 }
a.sum(axes['DA'])

or even just variables:

DA, DB, DC = range(3)
a.sum(DA)

If it should be your last (or penultimate, etc) axis, just use -1 (or -2, etc.):

a.shape
#(2,3,4)

np.all(a.sum(2) == a.sum(-1))
#True
np.all(a.sum(0) == a.sum(-3))
#True

Answer 2

You can write a thinly wrapped subclass to np.ndarray. But maintaining the correspondence between dimensions and the names can be tricky.

class NamedArray(np.ndarray):
    def __new__(cls, *args, **kwargs):
        obj = np.ndarray(args[0], **kwargs).view(cls)
        return obj

    def __init__(self, *args, **kwargs):
        self.dim_names = None
        if len(args) == 2:
            self.dim_names = args[1]

    def sum(self, *args, **kwargs):
        if (self.dim_names is not None) and (type(kwargs['axis']) == str):
            axis_name = kwargs.pop('axis')
            axis_ind = self.dim_names.index(axis_name)
            kwargs['axis'] = axis_ind
        return super().sum(*args, **kwargs)

#regular ndarray
a = NamedArray([1,2,3], dtype=np.float32)

#ndarray with dimension names
b = NamedArray([1,2,3], ('d1', 'd2', 'd3'), dtype=np.float32)

Edit: Pandas DataFrame nowadays is a pretty close thing to what the OP asked.