using xarray.apply(np.nansum) with args

Question

I have been trying to apply np.nansum to an xr.Dataset (xarray), but keep on coming up with errors. For a 3D dataset, I try to apply to axis=2. The syntax is not quite clear and I may have misunderstood the documentation, but I have tried:

ds.apply(np.nansum,axis=2)` and `ds.apply(lambda x: np.nansum(x,axis=2))

and get the same error:

cannot set variable 'var' with 2-dimensional data without explicit dimension names. Pass a tuple of (dims, data) instead.

I am guessing this means that it it does not know what dimension names to return to the new dataset object? Any ideas how to fix this?

And does anyone know why and when xarray might implement np.nansum()? Thanks

Answer 1

The problem you're running into here is that nansum returns a numpy ndarray, and not a DataArray, which is what the function passed into apply is supposed to return.

For nansum, you should just use xarray.Dataset.sum, which skips NaNs by default if your data is float.

Answer 2

Jeremy is correct that the built-in sum() method already skips NaN by default. But if you want to supply a custom aggregation function, you can do so with reduce, e.g., ds.reduce(np.nansum, axis=2).