Analysis-Ready Raster Data cubes¶

import itertools

import dask
import dask.array
import icechunk as ic
import numpy as np
import pandas as pd
import xarray as xr

xr.set_options(display_style="text")

import os

# request GDAL to use unauthenticated access
os.environ["GS_NO_SIGN_REQUEST"] = "1"
ds = xr.open_dataset(
    f"{ROOT_URI}/2000/70N_180W.tif",
    # use rasterio to read the TIFF
    engine="rasterio",
    # By default, xarray will replace NoData values with np.nan. Turn that off.
    mask_and_scale=False,
)
ds

def read_tile(
    tile: str, year: int, *, chunks=dict(zip(("band", "y", "x"), DASK_CHUNK_SIZE))
):
    """
    Helper function to read a single tile, specified as a string e.g. '70N_180W'
    for a single year.
    """
    import os

    # unauthenticated access
    os.environ["GS_NO_SIGN_REQUEST"] = "1"

    from rasterio.env import Env

    with Env(GDAL_CACHEMAX=0, VSI_CACHE=False, VSI_CACHE_SIZE=0):
        ds = (
            xr.open_dataset(
                f"{ROOT_URI}/{year}/{tile}.tif",
                engine="rasterio",
                # we must read as dask array so that the expand_dims step is lazy
                chunks=chunks,
                mask_and_scale=False,
                lock=False,
            )
            # Add a nice data variable name
            .rename({"band_data": "lclu"})
            # this is a singleton dimension, drop it
            .squeeze("band", drop=True)
            # add a more useful time dimension
            .expand_dims(year=[year])
        )
    return ds


ds = read_tile(tile="70N_180W", year=2000)
ds

template = xr.Dataset(
    {
        "lclu": (
            ("year", "y", "x"),
            dask.array.zeros(tiling.datacube_shape, chunks=-1, dtype=np.uint8),
            {
                "publication": (
                    "Potapov, P., Hansen, M.C., Pickens, A., Hernandez-Serna, A., Tyukavina, A., Turubanova, S., Zalles, V., Li, X., Khan, A., Stolle, F. "
                    "and Harris, N., 2022. The global 2000-2020 land cover and land use change dataset derived from the Landsat archive: first results. "
                    "Front. Remote Sens. 3: 856903."
                ),
                "doi": "10.3389/frsen",
                "grid_mapping": "spatial_ref",
            },
        )
    },
    coords={"year": ("year", tiling.year), "spatial_ref": ds.spatial_ref},
)
# Update the top-left corner information in the GeoTransform
template.spatial_ref.attrs["GeoTransform"] = (
    f"{tiling.tilex[0]:.1f} 0.00025 0.0 {tiling.tiley[0]:.1f} 0.0 -0.00025"
)

template.attrs["license"] = "CC BY 4.0"
template

session = repo.writable_session("main")
template.to_zarr(
    session.store,
    compute=False,  # important, don't write data values, just the array metadata!
    mode="w",
    encoding={"lclu": {"chunks": ZARR_CHUNK_SIZE, "fill_value": 255}},
    consolidated=False,
)
session.commit("initialize array")

repo = ic.Repository.open(storage)
session = repo.writable_session("main")
update_tile(session=session, tiling=tiling, tile="70N_180W", year=2000)
session.status()

import coiled

# Create software environment
coiled.create_software_environment(
    name="raster-ingest",
    conda={
        "channels": ["conda-forge"],
        "dependencies": [
            "python=3.12",
            "icechunk",
            "xarray",
            "rioxarray",
            "gcsfs",
            "dask",
            "coiled",
        ],
    },
    workspace="earthmover-devs-gcp",
)

cluster = coiled.Cluster(
    n_workers=[10, 200],
    software="raster-ingest",
    name="write-glad",
    workspace="earthmover-devs-gcp",
    idle_timeout="10m",
)

import distributed

client = distributed.Client(cluster)
client

repo = ic.Repository.open(storage)
session = repo.writable_session("main")
with session.allow_pickling():
    tasks = [get_update_tile_tasks_for_year(year) for year in tiling.year]
    tasks = tuple(itertools.chain(*tasks))
    print(f"Issuing {len(tasks)} tasks.")
    sessions = dask.compute(*tasks, scheduler=client)

repo = ic.Repository.open(storage)
list(repo.ancestry(branch="main"))

ds = xr.open_dataset(
    repo.readonly_session("main").store, engine="zarr", consolidated=False
)
ds