How can I do a convolution/rolling() based on physical dimensions (rather than window size)?

[tomchor]

Basically I'd like to calculate a convolution (which could use a Gaussian kernel or a tophat one, which would be equivalent to a rolling mean) over some data that I have.

However, instead of specifying the number of points used for the convolution (or window in the case of rolling().mean()), I'd like to express the kernel as a function of the actual distance/coordinate of the data. The difference is irrelevant in cases where the coordinate is uniformly-spaced, but is important in other cases. For example, taking the DataArray below as an example:

In [11]: import numpy as np

In [12]: import xarray as xr

In [13]: da = xr.DataArray(np.random.randn(30), dims=["x"], coords=dict(x=np.logspace(0, 2, 30)))

In [14]: da
Out[14]: 
<xarray.DataArray (x: 30)>
array([-0.44477985, -1.31223501,  0.06750107, -0.87812805, -1.69879256,
        0.27442237, -1.10400604,  0.2028317 , -1.22788433, -0.50962412,
       -0.98329938,  1.05594705, -1.5718489 , -0.04504426, -1.80016001,
       -0.21303616, -0.60341888, -0.4295976 ,  0.21796551, -0.02088305,
       -0.08396128,  0.52999041, -1.12443824,  0.12394816,  2.90654089,
        1.02446384,  0.33422938, -0.4752087 , -0.45542943,  1.41001154])
Coordinates:
  * x        (x) float64 1.0 1.172 1.374 1.61 1.887 ... 62.1 72.79 85.32 100.0

In [15]: da.x.diff("x")
Out[15]: 
<xarray.DataArray 'x' (x: 29)>
array([ 0.1721023 ,  0.2017215 ,  0.23643823,  0.27712979,  0.32482447,
        0.38072751,  0.44625158,  0.52305251,  0.61307105,  0.71858198,
        0.84225159,  0.98720503,  1.15710528,  1.35624576,  1.58965877,
        1.86324269,  2.18391104,  2.55976715,  3.00030896,  3.51666902,
        4.12189584,  4.83128358,  5.66275859,  6.63733235,  7.7796325 ,
        9.11852513, 10.68784425, 12.5272468 , 14.68321476])
Coordinates:
  * x        (x) float64 1.172 1.374 1.61 1.887 2.212 ... 62.1 72.79 85.32 100.0

If we interpret x as being a distance coordinate in meters, it'd make little physical sense to specify a fixed number of points. Instead it would be much more helpful if I could express my kernel distance directly in meters, and the points used to calculate the convolution in each location would be different. For example, a 10 m wide kernel on the data above would use several points around x=1, but only one close to x=100.

Is that possible to do efficiently the current tools we have?

At the moment my workaround for these cases has been to interpolate the data to a hi-res uniform grid before calling rolling(), but that's proving to be very expensive so I'd like to avoid that.

[dcherian]

xref #3216

I am not aware of any helper libraries here but my approach would be to construct a new sparse array of weights and use xr.dot(..., optimize=True) (similar to #3937)

This array of weights will be banded, so there's a potential optimization there.

Please do share any code you come up with.

[tomchor]

@dcherian I think you're arguing for the "simple" approach of manually creating a weight function which would act as a mask, no? Something like:

import numpy as np
import xarray as xr

N = 1000
x = np.logspace(0, 2, N) - 1
da = xr.DataArray(np.sin(2*np.pi*x/5),
                  dims=["x"], coords=dict(x=x))

def filter(da, filter_size, dim="x", **kwargs):
    x2 = da[dim].rename({ dim : "x2"})

    aux = (abs(x2 - da[dim]) < filter_size/2).astype(float)
    weights = aux / aux.sum(dim)
    return xr.dot(da, weights, **kwargs).rename(x2=dim)
    
daf = filter(da, 10, dim="x", optimize=True)

Which btw returns:

And, for comparison, the same algorithm but with an evenly spaced x = np.linspace(0, 100):

(The comparison makes me think there's something wrong at the boundaries since the left boundary (with more points) seems be worse than the right one (with fewer points), but at least the middle seems correct...)

But I'm not sure I understand you suggesting sparse arrays. I thought they still aren't supported? Also, I'm pretty unclear on what kind of optimization the optimize flag does and how much optimization I should expect from it since the docs aren't super helpful. I understand it's related to banded arrays, but in this case the size of the diagonal band varies in a non-predictable way, so I'm not sure how optimized things will be in general.

On a more practical note, the figure above seems to suggest that a simple mask-type weight array (i.e. a top-hat kernel) is not appropriate here since there seem to be jumps in the data. Maybe an Gaussian kernel is needed for this approach which would make optimization harder since then the weight array isn't a banded matrix anymore.

EDIT: I think one error in the algorithm above is that the normalization of aux should be done by the integral, not the sum! Same thing with the dot product, which instead of a sum should be an integral.

[dcherian]

think you're arguing for the "simple" approach of manually creating a weight function which would act as a mask, no?

yes! Dang that broadcast trick is really clever. Nice work. I was thinking rolling.construct but this is a lot cleaner.

But I'm not sure I understand you suggesting sparse arrays. I thought they still aren't supported?

https://sparse.pydata.org/en/stable/ should mostly work. See here for some workarounds necessary at the moment

Also, I'm pretty unclear on what kind of optimization the optimize flag does and how much optimization I should expect from it since the docs aren't super helpful.

See https://optimized-einsum.readthedocs.io/en/stable/index.html

I understand it's related to banded arrays, but in this case the size of the diagonal band varies in a non-predictable way, so I'm not sure how optimized things will be in general.

Sorry, this was a distraction. I was thinking of a dask case where a lot of blocks are zeros. And yes it is only "approximately" banded.

[tomchor]

Thanks for the clarifications! The code is now working properly, and I included a couple of bells and whistles. I haven't pursued any optimization yet, but I know what to try now based on opt_einsum (although I'm not sure I can apply that effectively since the proper way to do things is to use integrate()).

Here's a snippet with an example:

import numpy as np
import xarray as xr

N = 100
x = np.logspace(0, 2, N)-1
da = xr.DataArray(np.vstack([np.sin(2*np.pi*x/30)]*N),
                  dims=["y", "x"], coords=dict(x=x, y=x)).chunk(x=1, y=1)

def filter(da, filter_size, dim="x", kernel="gaussian", min_distance=0):
    x2 = da[dim].rename({ dim : "x2"})

    if kernel == "tophat":
        aux = (abs(x2 - da[dim]) < filter_size/2).astype(float)
    elif kernel == "gaussian":
        aux = np.exp(-((x2 - da[dim]) / (filter_size/2))**2)
    else:
        raise(TypeError("Invalid kernel name"))

    weights = aux / aux.integrate(dim)
    da_f = (da * weights).integrate(dim).rename(x2=dim)
    if min_distance: # Exclude edges (easier to do it after the convolution I think)
        da_f = da_f.where((da_f[dim]     - da_f[dim][0] >= min_distance) & \
                          (da_f[dim][-1] - da_f[dim]    >= min_distance), other=np.nan)
    return da_f


da_top = filter(da, 10, kernel="tophat", dim="x", min_distance=10)
da_exp = filter(da, 10, kernel="gaussian", dim="x", min_distance=10)

from matplotlib import pyplot as plt

da.isel(y=0).plot(label="original signal")
da_top.isel(y=0).plot(label="tophat filter")
da_exp.isel(y=0).plot(label="gaussian filter")
plt.legend()

Which produces this:

I haven't had the chance to test this on my actual data to see how efficient it is (server is out this week), but based on my tests the bit (da * weights).integrate("x").rename(x2=dim) is slow when N is large, even though no calculation is done there (since it's lazy due to dask). So I'm not sure what to make of that.

For now, with N=500, these compare surprisingly well with a rolling function (although the rolling window changes the efficiency, so results will probably be different for each application):

In [37]: %time filter(da, Δ, kernel="tophat", dim="x", min_distance=0)
CPU times: user 1.98 s, sys: 0 ns, total: 1.98 s
Wall time: 1.99 s

In [38]: %time filter(da, Δ, kernel="gaussian", dim="x", min_distance=0)
CPU times: user 965 ms, sys: 6.31 ms, total: 971 ms
Wall time: 977 ms

In [39]: %time da.rolling(x=int(Δ/da.x.diff("x").mean()), min_periods=1, center=True).mean()
CPU times: user 1.82 s, sys: 0 ns, total: 1.82 s
Wall time: 1.83 s

[dcherian]

(da * weights).integrate("x").rename(x2=dim) could be rewritten to xr.dot(da, weights, delta_x, optimize=True)?

[tomchor]

Wow, I'm impressed with how much xr.dot(..., optimize=True) actually optimized the result. For my data the processing time for a variable went from about 11 seconds with the un-optimized version above to about 0.6 seconds! My processing time for the same variable was about 2.5 seconds when re-gridding to a uniform spacing and then using .rolling(). Huge improvement! The result also looks better too, even with the tophat kernel.

My 2 cents: a version of this functionality should make it into xarray. Especially since coarse-graining as a means to analyzing energy in different scales is becoming more popular in the ocean+atmosphere community.

Plus I always thought it was odd that the whole point of xarray is to be able to work with labeled physical dimensions but functions like .rolling() can only be used by specifying the number of points instead of a physical scale. @dcherian thoughts?