dask support for lateral filling

[dcherian]

pop-tools/pop_tools/fill.py

Lines 58 to 80 in 8582dd0

	if len(non_lateral_dims) > 0:
	da_in_stack = da_in.stack(non_lateral_dims=non_lateral_dims)
	da_out_stack = xr.full_like(da_in_stack, fill_value=np.nan)
	isvalid_mask_stack = isvalid_mask.stack(non_lateral_dims=non_lateral_dims)
	for i in range(da_in_stack.shape[-1]):
	arr = da_in_stack.data[:, :, i]
	da_out_stack[:, :, i] = lateral_fill_np_array(
	arr, isvalid_mask_stack.data[:, :, i], ltripole, tol, use_sor, rc, max_iter
	)
	da_out = da_out_stack.unstack('non_lateral_dims').transpose(*dims_in)
	else:
	da_out = xr.full_like(da_in, fill_value=np.nan)
	da_out[:, :] = lateral_fill_np_array(
	da_in.data, isvalid_mask.data, ltripole, tol, use_sor, rc, max_iter
	)
	da_out.attrs = attrs
	da_out.encoding = encoding
	for k, da in coords.items():
	da_out[k].attrs = da.attrs

This should be replaced with a call to apply_ufunc(..., dask="parallelized", keep_attrs=True, vectorize=True, ...)