cumsum drops index coordinates

[dcherian]

What happened?

cumsum drops index coordinates. Seen in #6525, #3417

What did you expect to happen?

Preserve index coordinates

Minimal Complete Verifiable Example

import xarray as xr

ds = xr.Dataset(
        {"foo": (("x",), [7, 3, 1, 1, 1, 1, 1])},
        coords={"x": [0, 1, 2, 3, 4, 5, 6]},
    )
ds.cumsum("x")

<xarray.Dataset>
Dimensions:  (x: 7)
Dimensions without coordinates: x
Data variables:
    foo      (x) int64 7 10 11 12 13 14 15

Relevant log output

No response

Anything else we need to know?

No response

Environment

xarray main

[Illviljan]

xarray/xarray/core/dataset.py

Lines 5239 to 5240 in 7173bd3

	if not reduce_dims:
	variables[name] = var

If you comment out this if-check so it always add the variable at least the example seems to work.

Not sure how to handle this:

• Add a new keyword argument to reduce? keepcoords or can keepdims be expanded to handle this as well?
• Since cumsum is not really a reduction, should we create a new method instead of reduce? I guess it would look very similar to reduce though.

[dcherian]

Since cumsum is not really a reduction, should we create a new method instead of reduce? I guess it would look very similar to reduce though.

👍 I think we should add cumsum to generate_reductions.py. which maybe should become generate_aggregations.py

[max-sixty]

There was some discussion a while ago about how to handle reduce on things that aren't reductions. There's some overlap with #1618, and that has some links to other issues.

[max-sixty]

Was this closed by #7152?

Edit: No, I think there was a suggestion to change the approach from .reduce to .apply_ufunc here, but it didn't make it in

---

The general reduce issue overlaps with #6528.

[max-sixty]

I had a look at moving this to .apply_ufunc — if I'm thinking about it correctly — it's actually not trivial, since it's not just a gufunc like the .rolling_exp ones from #8114.

Instead, it's a function that we've written which can execute over multiple dimensions

xarray/xarray/core/duck_array_ops.py

Lines 621 to 631 in c9d89e2

	def _nd_cum_func(cum_func, array, axis, **kwargs):
	array = asarray(array)
	if axis is None:
	axis = tuple(range(array.ndim))
	if isinstance(axis, int):
	axis = (axis,)
	out = array
	for ax in axis:
	out = cum_func(out, axis=ax, **kwargs)
	return out

. (not super difficult! But IIUC, less of a lift-and-shift...)

[max-sixty]

Another alternative is to focus on Cumulative — ds.cumulative('lat').sum() works great!

[dcherian]

Can we just redirect cumsum to that then?

[max-sixty]

Can we just redirect cumsum to that then?

I looked at this but it has slightly more features such as skipna...

[jhollowed]

Can we just redirect cumsum to that then?

fwiw, I just tested this with a 33 MB Dataset. I observed that

data.cumsum('time')

is anywhere from 18 to 30x faster than

data.cumulative('time').sum()

[max-sixty]

@jhollowed, is that the first invocation? .cumulative uses functions from numbagg that are JIT compiled with numba, so the first invocation of the session will be much slower¹. But then it should be between a bit faster and much faster if it can parallelize.

If you have a repro then happy to investigate, would be a big deal if that's real.

Footnotes

1. Numba are working on enabling caching on parallel function, which will mean only the first run on the machine, not of the session, will be slow. ↩

[jhollowed]

@max-sixty No, it's not just the first invocation, from what I can tell. The repo in question is a mess, but here's a minimal working example which mimics my use case:

import time
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

data = xr.DataArray(np.random.rand(100,10,1000), dims=('x','y','z'))

cumsum_time = np.zeros(len(data.z))
cumulative_time = np.zeros(len(data.z))

for i in range(100):
    start = time.time()
    data.cumsum('z')
    cumsum_time[i] = time.time() - start

    start = time.time()
    data.cumulative('z').sum()
    cumulative_time[i] = time.time() - start

plt.hist(cumulative_time/cumsum_time, bins=20)
plt.show()

In this example the difference is even more extreme, with a ~50x slowdown by using cumulative().

I tried to also save the return of data.cumulative('z') as a new object prior to the loop, and call only sum() within the loop. This made no qualitative difference to the timing result. I also tried this in a Jupyter notebook, in which the difference was even worse but still qualitatively similar. I am running xarray version 2024.6.0.

If I'm making a mistake or have an obvious misunderstanding please let me know.

This is perhaps veering off-topic for this issue, but long story short, I just ended up using cumsum() and assign_coords() to recover the index coordinate dropped by cumsum(), rather than using the much slower cumulative().sum().

[max-sixty]

Can you share your env? I get something very very different!

The one outlier will be the initial JIT instance; here's a second run without that:

[max-sixty]

FWIW you may also want to use something like %%timeit — the repro above is using a missized array to store the values etc (no stress, it's exactly the sort of mistake I'd have made, but that's why %%timeit is great)

[jhollowed]

@max-sixty Interesting.

timeit.timeit(lambda: data.cumsum('z'), number=100)

and

timeit.timeit(lambda: data.cumulative().sum('z'), number=100)

yield 1.056 and 50.75, respectively.

I'm working on a login node of a cluster with many users, so I'm not totally in control of the context I'm running in. But my conda env looks like this.

[max-sixty]

Ah, OK, neither numbagg nor bottleneck is installed, and so it defaults to the slow routine. I guess that's more of a trap than I realized, with much worse perf than I realized; I should have thought of that initially when seeing your issue. We do at least put that in the .cumulative docstring.

FWIW we're considering including numbagg in the standard xarray release to reduce this confusion, and having a "core" release without any optional dependencies.

In the meantime, you might want to add numbagg to the env. (disclaimer: also a maintainer of numbagg...)

[max-sixty]

Given we now have a reference to numbagg in the .cumulative docstring, will plan to close this.

I thought about whether we should add a warning for .cumulative without numbagg nor bottleneck; I think that's a bit too invasive, and will be screened off if we do the xarray-core package. But open to feedback ofc.

[dcherian]

The core cumsum is still buggy

[max-sixty]

The core cumsum is still buggy

Agree; I closed since it seems unlikely we'll fix. But great if there's a path.

We could deprecate the skipna, and then defer to .cumulative. But requires numbagg...