HealPIX regridding

In this notebook, I demonstrate bilinear regridding onto healpy grids in O(10) ms. this is a 3 order of magnitude speed-up compared to what Dale has reported.

Now, lets define a healpix grid with indexing in the XY convention. we convert to NEST indexing in order to use the healpy.pix2ang to get the lat lon coordinates. This operation is near instant.

import matplotlib.pyplot as plt
import numpy as np
import torch

import earth2grid

# level is the resolution
level = 6
hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.XY())
src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64)
regrid = earth2grid.get_regridder(src, hpx)


z = np.cos(np.deg2rad(src.lat)) * np.cos(np.deg2rad(src.lon))


z_torch = torch.as_tensor(z)
z_hpx = regrid(z_torch)

fig, (a, b) = plt.subplots(2, 1)
a.pcolormesh(src.lon, src.lat, z)
a.set_title("Lat Lon Grid")

b.scatter(hpx.lon, hpx.lat, c=z_hpx, s=0.1)
b.set_title("Healpix")

Text(0.5, 1.0, 'Healpix')

nside = 2**level
reshaped = z_hpx.reshape(12, nside, nside)
lat_r = hpx.lat.reshape(12, nside, nside)
lon_r = hpx.lon.reshape(12, nside, nside)

tile = 11
fig, axs = plt.subplots(3, 4, sharex=True, sharey=True)
axs = axs.ravel()

for tile in range(12):
    axs[tile].pcolormesh(lon_r[tile], lat_r[tile], reshaped[tile])

/Users/nbrenowitz/workspace/nvidia/earth2grid/examples/sphinx_gallery/latlon_to_healpix.py:66: UserWarning: The input coordinates to pcolormesh are interpreted as cell centers, but are not monotonically increasing or decreasing. This may lead to incorrectly calculated cell edges, in which case, please supply explicit cell edges to pcolormesh.
  axs[tile].pcolormesh(lon_r[tile], lat_r[tile], reshaped[tile])