how to remove polynomial fit over cftime dimension?

[aaronspring]

I want to use xr.polyfit to remove a polynomial from a cftime dimension. On an integer coordinate this works. The problem is the x in bx^2 as cftimeindex * int is not defined.

def rm_poly(ds, dim, deg=2, **kwargs):
    """Remove degree polynomial along dimension dim from ds.
    Uses http://xarray.pydata.org/en/latest/generated/xarray.Dataset.polyfit.html"""
    fit = ds.polyfit(dim,deg=deg, **kwargs)
    fit = fit.rename({v: v.replace("_polyfit_coefficients","") for v in fit.data_vars})
    for deg in fit.degree:
        ds = ds - fit.sel(degree=deg,drop=True)*ds[dim]**deg
    if 'degree' in ds.coords:
        del ds.coords['degree']
    return ds

ds = xr.tutorial.load_dataset('rasm')
rm_poly(ds, dim='time', deg=2)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-85-f240578a0039> in <module>
----> 1 my_rm_poly(ds, 'time')

<ipython-input-62-e311d59a439a> in my_rm_poly(ds, dim, order)
      3     fit = fit.rename({v: v.replace("_polyfit_coefficients","") for v in fit.data_vars})
      4     for deg in fit.degree:
----> 5         ds = ds - fit.sel(degree=deg,drop=True)*ds[dim]**deg
      6     if 'degree' in ds.coords:
      7         del ds.coords['degree']

~/anaconda3/envs/climpred-dev/lib/python3.8/site-packages/xarray/core/dataarray.py in func(self, other)
   2989 
   2990             variable = (
-> 2991                 f(self.variable, other_variable)
   2992                 if not reflexive
   2993                 else f(other_variable, self.variable)

~/anaconda3/envs/climpred-dev/lib/python3.8/site-packages/xarray/core/variable.py in func(self, other)
   2299             with np.errstate(all="ignore"):
   2300                 new_data = (
-> 2301                     f(self_data, other_data)
   2302                     if not reflexive
   2303                     else f(other_data, self_data)

TypeError: unsupported operand type(s) for ** or pow(): 'cftime._cftime.DatetimeNoLeap' and 'int'

One workaround is to push the cftimeindex into an int index, which would even work for evenly separated coords.
Even better would be converting cftimeindex back into an int with cftime.date2num.

# works with cftime index
def rm_poly(ds, dim, deg=2, **kwargs):
    """Remove degree polynomial along dimension dim from ds.
    Uses http://xarray.pydata.org/en/latest/generated/xarray.Dataset.polyfit.html"""
    if isinstance(ds[dim].to_index(), xr.coding.cftimeindex.CFTimeIndex):
        was_cftime=True
        import cftime
        old_coords = ds.coords[dim]
        units = "days since 1900-01-01"
        ds[dim] = cftime.date2num(ds[dim], units)
        print('transform ds[dim]')
    else:
        was_cftime=False
    fit = ds.polyfit(dim,deg=deg, **kwargs)
    fit = fit.rename({v: v.replace("_polyfit_coefficients","") for v in fit.data_vars})
    for deg in fit.degree:
        ds = ds - fit.sel(degree=deg,drop=True)*ds[dim]**deg
    if 'degree' in ds.coords:
        del ds.coords['degree']
    if was_cftime:
        ds[dim]=old_coords
    return ds

My question: Is there a nicer way of doing so?

[dcherian]

polyfit is using xarray.core.missing.get_clean_interp_index which in turn uses datetime_to_numeric to get a numeric coordinate. Seems like we should return that coordinate along with the fit coefficients? We could also support passing datetime_unit and offset down to that function.

[spencerkclark]

I think xarray.polyval handles this gracefully. A bonus is that you also don't need to write the logic to evaluate the polynomials yourself. The one catch is that it only operates on DataArrays of coefficients, but that is straightforward to work around:

def rm_poly(ds, dim, deg=2, **kwargs):
    """Remove degree polynomial along dimension dim from ds."""
    coefficients = ds.polyfit(dim, deg=deg, **kwargs)
    coord = ds[dim]
    fits = []
    for v in coefficients:
        name = v.replace("_polyfit_coefficients", "")
        fit = xr.polyval(coord, coefficients[v]).rename(name)
        fits.append(fit)
    fits = xr.merge(fits)
    return ds - fits

[aaronspring]

Thanks. xr.polyval is not in the menu on the left, therefore I didnt see it.

Updated the function to allow xr.Datasets, xr.DataArrays and also xr.Dataset.map(rm_poly):

def rm_poly(ds, dim='time', deg=2, **kwargs):
    """Remove degree polynomial along dimension dim from ds."""
    coefficients = ds.polyfit(dim, deg=deg, **kwargs)
    coord = ds[dim]
    fits = []
    if isinstance(ds,xr.Dataset):
        for v in coefficients:
            name = v.replace("_polyfit_coefficients", "")
            fit = xr.polyval(coord, coefficients[v]).rename(name)
            fits.append(fit)
        fits = xr.merge(fits)
    elif isinstance(ds,xr.DataArray):
        name = ds.name
        v=list(coefficients.data_vars)[0]
        fits = xr.polyval(coord, coefficients[v]).rename(name)
    ds_rm_poly =  ds - fits
    return ds_rm_poly

[spencerkclark]

I think it is in the "Top-level functions" section of the API reference (since it is not a method of Dataset or DataArray). I agree though it is a little hard to discover. I think it would help if polyval was listed in the "See also" sections of Dataset.polyfit and DataArray.polyfit. Would you be interested in making a PR for that?

[aaronspring]

#5020

[spencerkclark]

Thanks!