Using polyfit on large 2d arrays

[tlsw231]

I am trying to use polyfit to find the slope from 2d arrays of long climate timeseries at each grid point. There is a detailed thread elsewhere (#6280) that discusses ways to improve polyval for 3d arrays and improve inefficiencies associated with stacking the spatial dimensions. I would have thought what I am trying here should be simpler as there is only one spatial dimension because this is output from an unstructured grid. The arrays all have shape (nt, ncol), where nt is the time dimension and ncol is the node index. Also, I only want the slope at every grid point, so don't even need polyval.

I suppose I can use rechunker to unchunk along the core dimension but want to know if there are other options.

MVE:

import dask.array as da

def get_linear_slope(da,dim,deg=1):
    p = da.polyfit(dim=dim, deg=deg, skipna=True)
    return p.polyfit_coefficients.sel(degree=1)

# My original data has the same size for nt and but much larger ncol ~ 700000
(nt,ncol)=(237980,100000)

dummy = xr.DataArray(da.random.random([nt,ncol]),dims=['time','ncol'],
                    coords=[range(nt),range(ncol)] ).rename('dummy')

out = get_linear_slope(dummy,'time').compute()

Here is a picture of the dashboard (with 12 cores, 384 gb):

[dcherian]

Can you use skipna=False? that looks a lot more efficient.

I think da.apply_along_axis is forcing a rechunk to a very big block. It would be good to add this to the docstring.

xarray/xarray/core/dask_array_ops.py

Lines 37 to 45 in a37a6b3

	results = da.apply_along_axis(
	nputils._nanpolyfit_1d,
	0,
	rhs,
	lhs_da,
	dtype=float,
	shape=(lhs.shape[1] + 1,),
	rcond=rcond,
	)

More generally, when skipna is True can we instead .fillna(0) and set weights=0 at those points? I think we need to reformulate the problem for skipna=True, this current version only works for a single chunk along dim or for cases where rechunking to a single chunk along dim works well.