Creating mosaic from DEM data with stackstac

[glaroc]

Hi, I'm trying to create a mosaic of DEM data at 1km resolution that I then want to save to disk. The mosaic resulting from the code below is mostly empty, with a few blocks where there is data. I've tried with alos-dem and cop-dem-glo-90. Any hints on how I can fix this?

import numpy as np
import xarray as xr
import rasterio.features
import stackstac
import pystac_client
import pystac
import planetary_computer
from rasterio.enums import Resampling
import rioxarray
import dask
from dask_gateway import GatewayCluster
catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

cluster = GatewayCluster() 
cluster.scale(24)
client = cluster.get_client()
print(cluster.dashboard_link)

aoi = {
    "type": "Polygon",
    "coordinates": [
        [
            [-100, 30],
            [-100, 66],
            [-47, 66],
            [-47, 30],
            [-100, 30]
        ]
    ],
}

search = catalog.search(
    #collections=["cop-dem-glo-90"], intersects=aoi
    collections=["alos-dem"], intersects=aoi
)
items = list(search.get_items())
signed_items = [planetary_computer.sign(item).to_dict() for item in items]

data = stackstac.stack(
        signed_items,
        assets=["data"],  
        chunksize=512,
        resolution=1000,
        epsg=32198,
        resampling=Resampling.average,
        bounds=[-2009488, -715776, 1401061, 2597757],
      ).drop('band').squeeze()

mosaic=stackstac.mosaic(data).persist()

mosaic.drop('band').plot.imshow(vmin=0, vmax=1850, figsize=(12,12))
#mosaic.compute().rio.to_raster("/home/jovyan/dem-ql.tif")

[TomAugspurger]

Sorry for the delay! I was sidetracked by gjoseph92/stackstac#113, and forgot to return to this answer.

One opening question: can you confirm that the desired bounding box is correct? If I plot the bounds on a map (yellow) with the STAC Item footprints (blue), I get

The bounding box in red is the bounding box for the item footprints, and is what I use below.

Next, you might want to confirm the expected data types. The source data are int16. If we use

bounds = rasterio.warp.transform_bounds(4326, epsg, *df.unary_union.bounds)

data = stackstac.stack(
        signed_items,
        assets=["data"],
        chunksize=2000,
        resolution=1000,
        epsg=epsg,
        resampling=Resampling.average,
        dtype="int16",
        fill_value=0,
        bounds=bounds,
      ).drop('band').squeeze()
data

We get something that should behave a bit better in terms of memory. Notice that I also increased chunksize to 2000. https://docs.dask.org/en/stable/array-best-practices.html#select-a-good-chunk-size applies here.

Next, we hit gjoseph92/stackstac#92. I'll put up a PR to fix that shortly, but in the meantime, we can define our own mosaic:

from typing import Hashable, Sequence, Union
import numpy as np
import xarray as xr


# TODO `fill_value`s besides NaN
def _mosaic(arr, axis, reverse: bool = False, fill_value=np.nan):
    ax_length = arr.shape[axis]

    # "normal" means last -> first, "reversed" means first -> last,
    # so we apply `reversed` in the opposite case you'd expect
    indices = iter(range(ax_length) if reverse else reversed(range(ax_length)))
    out = np.take(arr, next(indices), axis=axis)

    for i in indices:
        layer = np.take(arr, i, axis=axis)
        if np.isnan(fill_value):
            out = np.where(np.isnan(out), layer, out)
        else:
            out = np.where(out == fill_value, layer, out)
    return out


def mosaic(
    arr: xr.DataArray,
    dim: Union[None, Hashable, Sequence[Hashable]] = None,
    axis: Union[None, int, Sequence[int]] = 0,
    reverse: bool = False,
    fill_value=np.nan
):
    return arr.reduce(_mosaic, dim=dim, axis=axis, keep_attrs=True, reverse=reverse, fill_value=fill_value)

Which we can use as

result = mosaic(data, fill_value=0)
result

But we still run into memory issues :/ I'll post an update later with some more investigation. You might want to keep an eye on gjoseph92/stackstac#66. There might be some performance issues with stackstac.mosaic.

[glaroc]

Thank you for the investigation! I'll try to keep exploring for a solution on my side too. For the bounding box, you are right, I was experimenting with different bounds and forgot to update them.

[gjoseph92]

FYI, I just released stackstac 0.3.1, which fixes gjoseph92/stackstac#113 and gjoseph92/stackstac#92. So now you don't have to use the custom mosaic function at least!

For the memory issues, I'd be interested in hearing more about what's going on. I thought we'd concluded in gjoseph92/stackstac#66 (comment) that memory with mosaic seemed okay. But the mosaic implementation definitely isn't as efficient as it could be.

[glaroc]

I just tested this again with 0.3.1 on the Planetary Comptuer hub and it is still not working. If I set dtype='uint16' and fill_value=0, I get the error:

ValueError: Cannot use nodata=nan when mosaicing a uint16 array, since nan cannot exist in the array.

If I remove the dtype and fill_value options, it runs and then it seems to get into memory issues and I get a Killedworker error.

[gjoseph92]

If I set dtype='uint16' and fill_value=0

I assume you mean in the stackstac.stack call? Like data = stackstac.stack(..., dtype='uint16', fill_value=0)?

You just need to call stackstac.mosaic(data, nodata=0). I'll make the error message clearer, but the default nodata value for mosaic is NaN.

[glaroc]

Thank you. I thought that .mosaic() was pulling the fill_value from stac(). This fixes it. However, even with dtype='uint16', I still get memory issues and the job crashes with a Killedworker error.

[gjoseph92]

Yeah, I wish it could, but there's no xarray convention for how to track/store the fill_value. See gjoseph92/stackstac#92.

[gjoseph92]

it seems to get into memory issues and I get a Killedworker error

Try setting chunksize=(-1, ...) (like chunksize=(-1, 1, 2000, 2000)). This will make the chunks into "spatial skewers" and greatly reduce the size of the graph. It's not perfect but it might help. See gjoseph92/stackstac#116 (comment).

[glaroc]

This gives me an error saying that the chunksize and the shape have to have the same number of dimensions.
Got chunks=(-1,1, 2000, 2000), shape=(4993, 7936)

If I put chunksize(-1,2000), I still get the memory issue.

[gjoseph92]

Are you using the latest version of stackstac?

[glaroc]

I'm running this on the Planetary Computer hub which has stackstac v.0.3.1.

[gjoseph92]

Ah, you'll want 0.4.0. Unfortunately you'd also need to get this installed on the PC dask workers, maybe using a PipInstall plugin?

@TomAugspurger any progress on getting a newer version into the PC base image?

[TomAugspurger]

I'm hoping to update the images in ~2-3 weeks.