Convolution of xarray objects

Convolution of xarray objects#

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import xdggs
from xdggs.healpix import HealpixInfo

import healpix_convolution.xarray as hc
from healpix_convolution.plotting import xr_plot_healpix as plot_healpix

rng = np.random.default_rng(seed=0)

Global convolution#

We first have to define the data we want to convolve:

grid_info = HealpixInfo(level=3, indexing_scheme="nested")

cell_ids = np.arange(12 * 4**grid_info.level)
data1 = rng.normal(size=cell_ids.shape)
data2 = rng.normal(size=cell_ids.shape)

ds = xr.Dataset(
    {"data1": ("cells", data1), "data2": ("cells", data2)},
    coords={"cell_ids": ("cells", cell_ids, grid_info.to_dict())},
).pipe(xdggs.decode)
ds

<xarray.Dataset> Size: 18kB
Dimensions:   (cells: 768)
Coordinates:
  * cell_ids  (cells) int64 6kB 0 1 2 3 4 5 6 7 ... 761 762 763 764 765 766 767
Dimensions without coordinates: cells
Data variables:
    data1     (cells) float64 6kB 0.1257 -0.1321 0.6404 ... 0.008861 -0.1973
    data2     (cells) float64 6kB -0.7557 0.5313 0.7384 ... -0.4298 1.775 -1.508
Indexes:
    cell_ids  HealpixIndex(level=3, indexing_scheme=nested, kind=pandas)

And configure the kernel:

kernel = hc.kernels.gaussian_kernel(
    ds["cell_ids"], sigma=0.5, kernel_size=5, weights_threshold=1e-8
)
kernel

<xarray.DataArray (output_cells: 768, input_cells: 768)> Size: 456kB
<COO: shape=(768, 768), dtype=float64, nnz=18984, fill_value=0.0>
Coordinates:
    output_cell_ids  (output_cells) int64 6kB 0 1 2 3 4 ... 763 764 765 766 767
    input_cell_ids   (input_cells) int64 6kB 0 1 2 3 4 5 ... 763 764 765 766 767
Dimensions without coordinates: output_cells, input_cells
Attributes:
    kernel_type:  gaussian
    method:       continuous
    sigma:        0.5
    ring:         2
    kernel_size:  5

The convolution is then just:

convolved = hc.convolve(ds, kernel).pipe(xdggs.decode)
convolved

<xarray.Dataset> Size: 18kB
Dimensions:   (cells: 768)
Coordinates:
  * cell_ids  (cells) int64 6kB 0 1 2 3 4 5 6 7 ... 761 762 763 764 765 766 767
Dimensions without coordinates: cells
Data variables:
    data1     (cells) float64 6kB -0.1657 -0.05277 -0.3286 ... 0.02172 0.01102
    data2     (cells) float64 6kB -0.04611 -0.09172 -0.05156 ... 0.01582 -0.2399
Indexes:
    cell_ids  HealpixIndex(level=3, indexing_scheme=nested, kind=pandas)

And we can compare the two variables:

fig, axes = plt.subplots(
    nrows=2, ncols=2, figsize=(14, 9), subplot_kw={"projection": ccrs.Robinson()}
)
plot_healpix(ds["data1"], ax=axes[0, 0], title="original1", features=["coastlines"])
plot_healpix(
    convolved["data1"], ax=axes[0, 1], title="convolved1", features=["coastlines"]
)
plot_healpix(ds["data2"], ax=axes[1, 0], title="original2", features=["coastlines"])
plot_healpix(
    convolved["data2"], ax=axes[1, 1], title="convolved2", features=["coastlines"]
)
fig.subplots_adjust(hspace=-0.15, wspace=0.1)
../_images/xarray_4_0.png

Regional convolution#

Similarly to above, we also first have to define the data:

grid_info = HealpixInfo(level=4, indexing_scheme="nested")

cell_ids = 4 * 4**grid_info.level + np.arange(4**grid_info.level)
data1 = rng.normal(size=cell_ids.shape)
data2 = rng.normal(size=cell_ids.shape)

ds = xr.Dataset(
    {"data1": ("cells", data1), "data2": ("cells", data2)},
    coords={"cell_ids": ("cells", cell_ids, grid_info.to_dict())},
).pipe(xdggs.decode)
ds

<xarray.Dataset> Size: 6kB
Dimensions:   (cells: 256)
Coordinates:
  * cell_ids  (cells) int64 2kB 1024 1025 1026 1027 1028 ... 1276 1277 1278 1279
Dimensions without coordinates: cells
Data variables:
    data1     (cells) float64 2kB -0.03965 1.384 -0.1077 ... -1.03 -1.511 -2.234
    data2     (cells) float64 2kB -1.176 -0.8523 -0.05254 ... 0.1866 0.3077
Indexes:
    cell_ids  HealpixIndex(level=4, indexing_scheme=nested, kind=pandas)

and we can define the kernel:

kernel = hc.kernels.gaussian_kernel(
    ds["cell_ids"], sigma=0.5, kernel_size=5, weights_threshold=1e-8
)
kernel

<xarray.DataArray (output_cells: 256, input_cells: 392)> Size: 153kB
<COO: shape=(256, 392), dtype=float64, nnz=6382, fill_value=0.0>
Coordinates:
    output_cell_ids  (output_cells) int64 2kB 1024 1025 1026 ... 1277 1278 1279
    input_cell_ids   (input_cells) int64 3kB 0 1 2 3 8 ... 3068 3069 3070 3071
Dimensions without coordinates: output_cells, input_cells
Attributes:
    kernel_type:  gaussian
    method:       continuous
    sigma:        0.5
    ring:         2
    kernel_size:  5

However, since we now have boundaries we have to pad the data before convolving:

padding = hc.pad(
    ds["cell_ids"],
    ring=kernel.attrs["ring"],
    mode="constant",
    constant_value=0,
)
padded_ds = padding.apply(ds)
padded_ds

<xarray.Dataset> Size: 9kB
Dimensions:   (cells: 392)
Coordinates:
  * cell_ids  (cells) int64 3kB 0 1 2 3 8 9 10 ... 3062 3063 3068 3069 3070 3071
Dimensions without coordinates: cells
Data variables:
    data1     (cells) float64 3kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
    data2     (cells) float64 3kB 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0
Indexes:
    cell_ids  HealpixIndex(level=4, indexing_scheme=nested, kind=pandas)

After that, convolving works as before (but on the padded dataset):

convolved = hc.convolve(padded_ds, kernel).pipe(xdggs.decode)
convolved

<xarray.Dataset> Size: 6kB
Dimensions:   (cells: 256)
Coordinates:
  * cell_ids  (cells) int64 2kB 1024 1025 1026 1027 1028 ... 1276 1277 1278 1279
Dimensions without coordinates: cells
Data variables:
    data1     (cells) float64 2kB 0.1892 0.2312 0.1216 ... -0.182 -0.2131
    data2     (cells) float64 2kB -0.1219 -0.02719 0.006069 ... 0.02166 0.03977
Indexes:
    cell_ids  HealpixIndex(level=4, indexing_scheme=nested, kind=pandas)

Again, we can look at the result:

fig, axes = plt.subplots(
    nrows=2, ncols=2, figsize=(14, 9), subplot_kw={"projection": ccrs.Robinson()}
)
plot_healpix(
    ds["data1"], ax=axes[0, 0], title="original1", features=["coastlines"], xsize=1000
)
plot_healpix(
    convolved["data1"],
    ax=axes[0, 1],
    title="convolved1",
    features=["coastlines"],
    xsize=1000,
)
plot_healpix(
    ds["data2"], ax=axes[1, 0], title="original2", features=["coastlines"], xsize=1000
)
plot_healpix(
    convolved["data2"],
    ax=axes[1, 1],
    title="convolved2",
    features=["coastlines"],
    xsize=1000,
)
fig.subplots_adjust(hspace=0.1, wspace=-0.55)
../_images/xarray_9_0.png
Contents