test_rasterize.py

#!/usr/bin/env python3

import sys
import logging
logging.basicConfig(format="%(asctime)s %(message)s",
                    level=logging.INFO,
                    stream=sys.stderr)
import numpy as np
import xarray as xr
import pandas as pd
import geopandas as gpd
from xcube.core.geom import rasterize_features
from shapely.geometry import Polygon
import xcube.version

def main():
    logging.info(f"xcube version {get_xcube_version()}")
    
    npolys = int(sys.argv[1])
    logging.info(f"{npolys} polygons")
    ids = range(npolys)
    gdf = gpd.GeoDataFrame(
        {"id": ids, "geometry": [make_polygon(i) for i in ids]},
        crs="EPSG:4326"
    )

    ds = make_dataset()
    logging.info("Rasterizing...")
    rasterize_features(ds, gdf, ["id"], in_place=True)
    logging.info("Done.")


def get_xcube_version():
    if isinstance(xcube.version, str):
        return xcube.version
    else:
        return xcube.version.version
    

def make_polygon(i):
    lon0 = 4.39 + (i // 100) * 0.1
    lat0 = 49.53 + (i % 100) * 0.018 
    return Polygon([[lon0, lat0], [lon0 + 0.05, lat0],
                    [lon0 + 0.05, lat0 + 0.01], [lon0, lat0 + 0.01],
                    [lon0, lat0]])


def make_dataset():
    lat = np.linspace(51.516735, 49.520065, 15360, dtype="float64")
    lon = np.linspace(4.380065, 6.110495, 13312, dtype="float64")
    t = pd.date_range("2019-01-06", periods=1)
    ds = xr.Dataset(
        {
            "v": (("lat", "lon", "time"),
                  np.zeros((len(lat), len(lon), len(t)), dtype="float32"))
        },
        coords={
            "lat": lat,
            "lon": lon,
            "time": t
        }
    )
    return ds

    
if __name__ == "__main__":
    main()