EOPF zarr Performance Comparison - xcube

import requests

from xcube.core.store import new_data_store
from xcube_resampling.utils import reproject_bbox

# for benchmarking
from dataclasses import dataclass
from typing import List
from itertools import product
import pandas as pd
import time

def create_aoi(bbox, reduction):
    """
    Generate a reduced bounding box or centroid based on a reduction factor.
    Helper function to easily create portions of the original bbox around the centroid.

    Parameters:
    - bbox: [min_lon, min_lat, max_lon, max_lat]
    - reduction: float between 0 and 1
        - 0 returns the centroid as (lon, lat)
        - 0 < reduction < 1 returns a scaled bounding box centered at the centroid

    Returns:
    - reduced bounding box list
    """
    if not (0 <= reduction <= 1):
        raise ValueError("Reduction must be between 0 and 1.")

    min_lon, min_lat, max_lon, max_lat = bbox

    # Compute centroid
    centroid_lon = (min_lon + max_lon) / 2
    centroid_lat = (min_lat + max_lat) / 2

    # if reduction == 0:
    #   return (centroid_lon, centroid_lat)

    # Compute reduced bounding box dimensions
    lat_span = (max_lat - min_lat) * reduction
    lon_span = (max_lon - min_lon) * reduction

    return [
        centroid_lon - lon_span / 2,
        centroid_lat - lat_span / 2,
        centroid_lon + lon_span / 2,
        centroid_lat + lat_span / 2,
    ]

# this is used for defining inputs
@dataclass
class BenchmarkConfig:
    data_id: str
    bbox: List[float]
    time_range: List[str]
    spatial_res: int
    crs: str
    variables: List[str]

def access_eopf(cfg: BenchmarkConfig):
    """
    Use the dataclass BenchmarkConfig to read the EOPF zarr.

    Parameters:
    - cfg: BenchmarkConfig object, holding the information to read the data.

    Returns:
    - Loaded dataset as xarray
    """
    return store_zarr.open_data(
        data_id=cfg.data_id,
        bbox=reproject_bbox(
            cfg.bbox, "EPSG:4326", cfg.crs
        ),  # has to be done for xcube # TODO: throws error with dfg.crs = EPSG:4326
        time_range=cfg.time_range,
        spatial_res=cfg.spatial_res,
        crs=cfg.crs,
        variables=cfg.variables,
    ).load()

def access_safe(cfg: BenchmarkConfig):
    """
    Use the dataclass BenchmarkConfig to read the SAFE.

    Parameters:
    - cfg: BenchmarkConfig object, holding the information to read the data.

    Returns:
    - Loaded dataset as xarray
    """
    return store_safe.open_data(
        data_id=cfg.data_id,
        bbox=reproject_bbox(
            cfg.bbox, "EPSG:4326", cfg.crs
        ),  # has to be done for xcube # TODO: throws error with dfg.crs = EPSG:4326
        time_range=cfg.time_range,
        spatial_res=cfg.spatial_res,
        crs=cfg.crs,
        asset_names=[v.upper() for v in cfg.variables],
    ).load()

def benchmark_data_access(configs, access_fn):
    """
    Loops through the given configs using the specified access function.

    Parameters:
    - configs: List of BenchmarkConfigs to loop over
    - access_fn: Access function to use, here either access_eopf or access_safe

    Returns:
    - DataFrame with the access specifics and performance indicators
    """
    results = []

    for cfg in configs:
        print(f"Running: {cfg}")
        start = time.perf_counter()
        ds = access_fn(cfg)
        end = time.perf_counter()

        n_pixels_xy = ds.sizes["x"] * ds.sizes["y"]  # get pixel count
        results.append(
            {
                "data_id": cfg.data_id,
                "bbox": cfg.bbox,
                "time_range": cfg.time_range,
                "spat_res": cfg.spatial_res,
                "crs": cfg.crs,
                "variables": cfg.variables,
                "n_pixels_xy": n_pixels_xy,
                "duration_sec": round(end - start, 4),
            }
        )

    return pd.DataFrame(results)

url = "https://stac.core.eopf.eodc.eu/collections/sentinel-2-l2a/items/S2A_MSIL2A_20200607T101031_N0500_R022_T32TPS_20230502T194134"

response = requests.get(url)
item = response.json()
bbox = item["bbox"]
print(bbox)

crs_native = item["properties"]["proj:code"]  # "EPSG:32632"
print(crs_native)

# define data id
opt_data_id = ["sentinel-2-l2a"]

# define bboxes
# only in lat/lon, reprojection of bbox to chosen crs happens later in code
opt_bbox = [
    # create_aoi(bbox, 256 / 10980), # ml patch approx 256*256
    # create_aoi(bbox, 0.125), # eight of a full scene
    create_aoi(bbox, 0.25),  # quarter of a full scene
]

# define crs
# mandatory in xcube
# if it differs from native crs processing is enforced (reprojection, resampling)
opt_crs = [
    crs_native,
]

# define times
opt_time_range = [
    ["2025-05-01", "2025-05-07"],  # week
]

# define spatial resolution
# everything deviating from native resolution enforces processing (resampling)
opt_spatial_res = [
    10,
]

# define band combinations
# choosing bands with different resolutions enforces processing (resampling)
opt_variables = [
    ["b02"],
    # ["b02", "b04"],
]

configs = [
    BenchmarkConfig(data_id, bbox, time_range, spatial_res, crs, variables)
    for data_id, bbox, time_range, spatial_res, crs, variables in product(
        opt_data_id, opt_bbox, opt_time_range, opt_spatial_res, opt_crs, opt_variables
    )
]
print(f"Number of configs: {len(configs)}")
print(f"Example config: {configs[0]}")

store_zarr = new_data_store("eopf-zarr")

df_benchm_eopf = benchmark_data_access(configs, access_eopf)

df_benchm_eopf["bbox"] = df_benchm_eopf["bbox"].apply(
    lambda b: [round(x, 3) for x in b]
)

df_benchm_eopf

credentials = {
    "key": "XXX",
    "secret": "XXX",
}

store_safe = new_data_store("stac-cdse-ardc", **credentials)

df_benchm_safe = benchmark_data_access(configs, access_safe)

df_benchm_safe["bbox"] = df_benchm_safe["bbox"].apply(
    lambda b: [round(x, 3) for x in b]
)

df_benchm_safe