Build a sample datacube from two datasets

from pathlib import Path

import geopandas as gpd
import numpy as np
import rioxarray
import xarray as xr
import matplotlib.pyplot as plt
from dask.diagnostics import ProgressBar
from rasterio.enums import Resampling
from rioxarray.merge import merge_arrays


FINAL_CUBE = "alboran_daily_demo_cube.zarr"

INPUT_PATH1 = "https://s3.waw4-1.cloudferro.com/EarthCODE/OSCAssets/ocean_datasets/sssd.zarr"
INPUT_PATH2 = "https://s3.waw4-1.cloudferro.com/EarthCODE/OSCAssets/ocean_datasets/careheat_mhw_2d.zarr"


source_dataset1 = xr.open_zarr(
    INPUT_PATH1, chunks={}, decode_coords="all",
)
source_dataset2 = xr.open_zarr(
    INPUT_PATH2, chunks={}, decode_coords="all",
)

# 1. Define bounding box for the Mediterranean Sea
LON_MIN, LON_MAX = -6.0, 37.0   # From the Strait of Gibraltar to the Levantine coast
LAT_MIN, LAT_MAX = 30.0, 46.0   # From the North African coast to the northern Adriatic
TARGET_RESOLUTION_DEG = 1 / 12
TIME_TARGET = "2019-06-01"

# 2. Calculate the number of cells for longitude and latitude separately
n_cells_lon = int(round((LON_MAX - LON_MIN) / TARGET_RESOLUTION_DEG))
n_cells_lat = int(round((LAT_MAX - LAT_MIN) / TARGET_RESOLUTION_DEG))

# 3. Create 1D coordinate arrays (cell centers)
target_lon = LON_MIN + (np.arange(n_cells_lon) + 0.5) * TARGET_RESOLUTION_DEG
target_lat = LAT_MIN + (np.arange(n_cells_lat) + 0.5) * TARGET_RESOLUTION_DEG

source1 = source_dataset1[["sos"]].sel(
    time=slice(TIME_TARGET, TIME_TARGET), depth=0, drop=True,
    lon=slice(LON_MIN - 0.25, LON_MAX + 0.25),
    lat=slice(LAT_MIN - 0.25, LAT_MAX + 0.25),
)

source2 = source_dataset2[["category"]].sel(
    time=slice(TIME_TARGET, TIME_TARGET), drop=True,
    lon=slice(LON_MIN - 0.25, LON_MAX + 0.25),
    lat=slice(LAT_MIN - 0.25, LAT_MAX + 0.25),
)

sss = source1.interp(lon=target_lon, lat=target_lat, method="nearest")
caheat_cat = source2.interp(lon=target_lon, lat=target_lat, method="nearest")

# align the time coordinates
sss['time'] = caheat_cat['time']

cube = xr.merge([sss, caheat_cat])
cube = cube.rio.write_crs("EPSG:4326").rio.set_spatial_dims(x_dim="lon", y_dim="lat")
cube

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10,5))
cube.sos.isel(time=0).plot(ax=ax[0])
cube.category.isel(time=0).plot(ax=ax[1])

left, bottom, right, top = cube.rio.bounds()
WC_BASE = "https://esa-worldcover.s3.eu-central-1.amazonaws.com"
tiles = gpd.read_file(f"{WC_BASE}/v100/2020/esa_worldcover_2020_grid.geojson").cx[left:right, bottom:top]
wc = merge_arrays([
    rioxarray.open_rasterio(
        f"{WC_BASE}/v200/2021/map/ESA_WorldCover_10m_2021_v200_{tile}_Map.tif",
        overview_level=4,
    )
    for tile in sorted(tiles.ll_tile)
], nodata=0).squeeze("band", drop=True)

water = wc.isin([0, 80]).astype("uint8").rio.write_nodata(255)
water_mask = water.rio.reproject_match(cube, resampling=Resampling.min, nodata=255)
water_mask = water_mask.where(water_mask != 255, 1)
water_mask = water_mask.rename({"x": "lon", "y": "lat"}).rename("water_mask")
cube["sos"] = cube.sos.where(water_mask == 1)
cube["water_mask"] = water_mask
cube.sos.isel(time=0).plot()

cube.attrs = {
    "title": "Alboran Sea salinity demonstration cube",
    "crs": "EPSG:4326",
    "target_resolution_degrees": TARGET_RESOLUTION_DEG,
    "source_store1": INPUT_PATH1,
    "source_store2": INPUT_PATH2,
    "interpolation_method": "nearest",
}

# Clear the source storage settings before writing the new grid.
for name in cube.variables:
    cube[name].encoding = {}
for name in cube.data_vars:
    cube[name].attrs["grid_mapping"] = "spatial_ref"

with ProgressBar():
    cube.chunk({"time": 1, "lat": -1, "lon": -1}).to_zarr(
        FINAL_CUBE, mode="w", consolidated=True, zarr_format=3,
    )

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature


combined_ds = xr.open_zarr(FINAL_CUBE, consolidated=True, decode_coords="all")
ds_plot = combined_ds.isel(time=0)

fig, axes = plt.subplots(
    nrows=1, ncols=2, 
    figsize=(16, 6), 
    subplot_kw={'projection': ccrs.PlateCarree()}
)


# ---------------------------------------------------------
# Plot 1: Sea Surface Salinity (sos)
# ---------------------------------------------------------
ds_plot['sos'].plot(
    ax=axes[0],
    transform=ccrs.PlateCarree(),
    cmap='viridis',
    cbar_kwargs={'label': 'Sea Surface Salinity (.0001)', 'orientation': 'horizontal'}
)
axes[0].set_title('Sea Surface Salinity (sos)')

# ---------------------------------------------------------
# Plot 2: Extreme Event Category (category)
# ---------------------------------------------------------
ds_plot['category'].plot(
    ax=axes[1],
    transform=ccrs.PlateCarree(),
    cmap='Reds', # A sequential colormap is usually good for categorical/extreme data
    cbar_kwargs={'label': 'Extreme Event Category', 'orientation': 'horizontal'}
)
axes[1].set_title('Extreme Event Category')

for ax in axes:
    # Add land mask (zorder=1 to ensure it draws on top of any NaN ocean artifacts, or underneath data)
    ax.add_feature(cfeature.LAND, facecolor='lightgray', zorder=1)
    
    # Add coastlines
    ax.coastlines(resolution='50m', color='black', linewidth=0.8, zorder=2)
    
    # Restrict the plot extent to your Mediterranean bounds
    # Format: [lon_min, lon_max, lat_min, lat_max]
    ax.set_extent([-6.0, 37.0, 30.0, 46.0], crs=ccrs.PlateCarree())
    
plt.tight_layout()
plt.show()