Analogue to numpy.take?

[krikru]

Hi,

I was just wondering if Iris has any analogue to numpy.take?

I have a cube c and want to do is something similar to c[:, 0, :, :, 0, ...] where some indices should be 0 and the rest of the indices be slice(None) and where the positions of all zero-indices are stored in a set. I want to do this since those zeros represent dimensions in c that always only have one point in them and therefore safely can be discarded. I know I can do this with c.collapsed, but then I need to provide an aggregator, which feels unnecessary when the entire operation basically just amounts to a reshaping of the data.

[krikru]

I solved my problem like this, b.t.w.:

c = c[tuple(0 if dim_idx in dimensions_to_collapse else slice(None) for dim_idx in range(len(c.coords())))]

This corresponds to writing

c[(slice(None), 0, slice(None), slice(None), 0, ...)]

Apparently a tuple of integers or slice objects works to index the cube with too.

[trexfeathers]

Take a look at iris.util.squeeze() for this use-case.

In cases where advanced indexing is needed, we usually construct something similar to your example in https://github.com/SciTools/iris-grib/discussions/302#discussioncomment-3915696.

[krikru]

@trexfeathers Thanks for your answer! That seems like a potentially useful method, but in my case it carries the risk of also removing any dimension that just happens to have length one but doesn't always have length one. In that case I do want to keep the dimension because my code relies on it still being there.

A method that would be more useful to me would be one in which you could choose which dimensions to remove (and that should probably also raise an exception if any of those dimensions would have a length different from one).

[rcomer]

Moving this over to Iris. Hold on to your hats!

[krikru]

Oh, I didn't mean to start this discussion in iris-grib; thanks for moving it!

[rcomer]

Does extract do what you need?

import iris.cube
import iris.coords

foo_coord = iris.coords.DimCoord(1, long_name='foo')
bar_coord = iris.coords.DimCoord(1, long_name='bar')
cube = iris.cube.Cube([[1]])

for dim, coord in enumerate([foo_coord, bar_coord]):
    cube.add_dim_coord(coord, dim)

print(cube)
print(cube.extract(iris.Constraint(bar=1)))

unknown / (unknown)                 (foo: 1; bar: 1)
    Dimension coordinates:
        foo                             x       -
        bar                             -       x
unknown / (unknown)                 (foo: 1)
    Dimension coordinates:
        foo                             x
    Scalar coordinates:
        bar                         1

[krikru]

It seems like it can do what I want to do, but using it requires the extra step of extracting the value of the coordinate for each dimension I want to collapse, which seems like carrying coals to Newcastle in comparison with my solution since I already know that I want to extract the elements with index 0 along that dimension.

[rcomer]

This also works, but is not very intuitive

print(cube.extract(iris.Constraint(bar=lambda c: True)))

unknown / (unknown)                 (foo: 1)
    Dimension coordinates:
        foo                             x
    Scalar coordinates:
        bar                         1

I learned about this pattern in the course of updating one of the gallery examples - #3933