Regridding ClimateDT to

import os

del os.environ["PROJ_DATA"]
del os.environ["ESMFMKFILE"]
del os.environ["_CONDA_SET_ESMFMKFILE"]

import earthkit.data
import earthkit.regrid
from earthkit.plots.interactive import Chart
from polytope.api import Client
import xesmf

import xdggs
import dask.array as da

import numpy as np
import healpix_geo.nested
import xarray as xr

# Set True if you want to make a live request for the data, or false if you want to use the cached grib file
LIVE_REQUEST = False

pressure_levels = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 70, 50, 30, 20, 10, 5, 1]
request = {
    "activity": "scenariomip",
    "class": "d1",
    "dataset": "climate-dt",
    "experiment": "ssp3-7.0",
    "generation": "1",
    "levtype": "pl",
    "date": "20210101",
    "model": "ifs-nemo",
    "expver": "0001",
    "param": "130",
    "realization": "1",
    "resolution": "high",
    "stream": "clte",
    "type": "fc",
    "time": "/".join(f"{n * 6:02d}00" for n in range(4)),
    "levelist": "/".join(str(_) for _ in pressure_levels),
}
request

data_file = "climate-dt-earthkit.grib"
if LIVE_REQUEST:
    data = earthkit.data.from_source("polytope", "destination-earth", request, address="polytope.lumi.apps.dte.destination-earth.eu", stream=False)
    data.to_target("file", data_file)
else:
    data = earthkit.data.from_source("file", data_file)

ds = data.to_xarray(chunks={"values": 4**8, "forecast_reference_time": -1, "level": -1})
ds

decoded = (
    ds.assign_coords(cell_ids=("values", da.arange(ds.sizes["values"], dtype="u8")))
    .dggs.decode({"grid_name": "healpix", "level": 10, "indexing_scheme": "nested"}, index_kind="moc")
)
decoded

def spherical_to_ellipsoidal(ds):
    source_grid = ds[["latitude", "longitude"]].load()
    grid_info = ds.dggs.grid_info

    cell_ids = np.arange(12 * 4**grid_info.level, dtype="uint64")
    longitude, latitude = healpix_geo.nested.healpix_to_lonlat(
        cell_ids, depth=grid_info.level, ellipsoid="WGS84"
    )
    target_grid = xr.Dataset(
        coords={
            "cell_ids": ("cells", cell_ids),
            "longitude": ("cells", longitude),
            "latitude": ("cells", latitude),
        }
    )

    regridder = xesmf.Regridder(
        source_grid,
        target_grid,
        method="nearest_s2d",
        locstream_in=True,
        locstream_out=True,
        periodic=True,
    )

    return regridder.regrid_dataset(ds, keep_attrs=True).dggs.decode(grid_info)

%%time
regridded = decoded.pipe(spherical_to_ellipsoidal)
regridded

def as_cf(ds):
    grid_info = ds.dggs.grid_info
    metadata = {
        "grid_mapping_name": "healpix",
        "refinement_level": grid_info.level,
        "indexing_scheme": grid_info.indexing_scheme,
        "reference_body": "WGS84",
        "refinement_ratio": 4,
    }

    new = ds.assign_coords(crs=xr.Variable((), np.int8(0), metadata))
    new["cell_ids"].attrs = {"standard_name": "healpix", "units": "1"}

    var_attrs = {"coordinates": ds.dggs._name, "grid_mapping": "crs"}

    for name, var in new.variables.items():
        if name in {"cell_ids"}:
            continue

        var.attrs |= var_attrs

    return new

cf = as_cf(regridded)
cf["t"].attrs.pop("_earthkit", None)
cf

zarr_url = f"{data_file.removesuffix('.grib')}.zarr"
cf.to_zarr(zarr_url, mode="w")

reloaded = xr.open_datatree(zarr_url, engine="zarr", chunks={})
reloaded