larray-project / larray

N-dimensional labelled arrays in Python
https://larray.readthedocs.io/
GNU General Public License v3.0
8 stars 6 forks source link

expand functionality of axis.i to allow expressions on indices #1007

Open gdementen opened 2 years ago

gdementen commented 2 years ago

Unsure it is possible, but I just wanted to write code like this:

>>> arr = ndtest((3, 4))
>>> num_values_per_row = stack({'a0': 2, 'a1': 1, 'a2': 3}, 'a')
>>> arr.b.i < num_values_per_row
False
>>> arr.b.i[:] < num_values_per_row
False

The ultimate goal in this case was to be able to "select" a different number on items per row (because this is not allowed using the slice/getitem syntax), for example to aggregate them:

>>> where(arr.b.i < num_values_per_row, arr, 0).sum()
0

The workaround is to use .ignore_labels() but this does not seem very intuitive/discoverable to me.

>>> arr.b.ignore_labels() < num_values_per_row
b*\a     a0     a1     a2
   0   True   True   True
   1   True  False   True
   2  False  False   True
   3  False  False  False
>>> where(arr.b.ignore_labels() < num_values_per_row, arr, 0).sum()
32