Applying a Scikit-learn regression on a lat,lon grid

[ricardobarroslourenco]

Hi everyone.

I have a Dataset with coords (x, y, time), and I want to apply a Scikit-learn Theil-Sen regression on the time component (so the result would be multiple trained models for every lat,lon coordinate (it has a single variable in it). The results would be stored in a separate dimension.

Can you refer to an efficient way to do this? I have thought of using some vectorized form such as apply or map, but in my attempts, I couldn't make it work.

[dcherian]

See https://tutorial.xarray.dev/advanced/apply_ufunc/vectorize_1d.html an example, it does something similar for the interpolation problem

Start here if you are new to apply_ufunc

[ricardobarroslourenco]

@dcherian is it possible to use a function that return more than one variable?

[dcherian]

Yes it should be but I don't think we have a nice example. The tests show some simple versions:

xarray/xarray/tests/test_computation.py

Line 223 in 49ae0f8

def test_apply_two_outputs() -> None:

[ricardobarroslourenco]

Nice! A side question. When I run (in an ideal run) such function as :

var_a, var_b = test_apply_two_outputs() , var_a and var_b would be DataArray or Dataset objects?

[dcherian]

Yes the return type is always an Xarray object.

[jhamman]

apply_ufunc is the way to go if you don't need coordinate information in your computation. If you do, I'd check out map_blocks.

In Scikit-downscale, we ended up using map_blocks in our PointWiseDownscaler.

[ricardobarroslourenco]

I have gone with something working, but quite inefficient:

def vector_theil_sen(Xarray_DataArray):
  X = Xarray_DataArray.year.to_pandas()
  y = Xarray_DataArray.to_pandas()
  model = TheilSenRegressor(random_state=0, verbose=False, max_iter=1000, n_jobs=-1).fit(X.to_numpy().reshape(-1, 1),y.to_numpy())
  model_coefficient = model.coef_
  model_intercept = model.intercept_
  return np.array(pickle.dumps(model)) , model_coefficient[0], model_intercept 

def coefficient_ufunc(x,y):
  data_array_pointer = lai_groupby_year['masked_lai'].sel(x=x,y=y)
  if debugging:
    print('x: '+str(x)+' , y: '+str(y))
  if data_array_pointer.isnull().any():
    rate = np.nan
  else:
    _,rate, _ = vector_theil_sen(data_array_pointer)

  if debugging:
    print(str(rate))
  return rate

result = xr.apply_ufunc(
        coefficient_ufunc,
        lai_groupby_year['masked_lai'].x,
        lai_groupby_year['masked_lai'].y,
        vectorize=True

I have masked the source Dasatet ( refer to masked_lai), but still seems a quite slow run. When debugging is true, the x,y iterations are quite slow (half a second, perhaps) but on the Dask workers cpu timeseries it gets in average 60%. The TheilSen Regressor is with the n_jobs=-1 flag. I wonder if this is causing issues when computing it.

[headtr1ck]

As @jhamman pointed out: use apply_ufunc only if you want to work with e.g. numpy arrays and don't require the DataArray. (In this case simply add the value as a function argument and add it to the apply_ufunc call as well).

If you require the DataArray, use map_blocks instead.