Loading bed picks and layer data

import xopr
from xopr.bedmap import query_bedmap, query_bedmap_catalog, fetch_bedmap

import holoviews as hv
import xarray as xr
import hvplot
import hvplot.xarray
import hvplot.pandas
import geoviews.feature as gf
import cartopy.crs as ccrs
import rioxarray
from tqdm.notebook import tqdm
import numpy as np
import verde as vd
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
from shapely.geometry import box
from datetime import datetime
import time
import warnings
warnings.filterwarnings('ignore')

opr = xopr.OPRConnection(cache_dir='radar_cache')

epsg_3031 = ccrs.Stereographic(central_latitude=-90, true_scale_latitude=-71)
coastline = gf.coastline.options(scale='50m').opts(projection=epsg_3031)
velocity = rioxarray.open_rasterio(
    "https://its-live-data.s3.amazonaws.com/velocity_mosaic/v2/static/cog/ITS_LIVE_velocity_120m_RGI19A_0000_v02_v.tif",
    chunks='auto', overview_level=4, cache=False
).squeeze().drop_vars(['spatial_ref', 'band']).rename('velocity (m/year)')
velocity_map = velocity.hvplot.image(x='x', y='y', cmap='gray_r').opts(clim=(0,500))

region = xopr.geometry.get_antarctic_regions(name=["Vincennes_Bay", "Underwood"], merge_regions=True, simplify_tolerance=100)
region_projected = xopr.geometry.project_geojson(region, source_crs='EPSG:4326', target_crs="EPSG:3031")

region_hv = hv.Polygons([region_projected]).opts(
    color='green',
    line_color='black',
    fill_alpha=0.3)

(velocity_map * coastline * region_hv).opts(aspect='equal')

gdf = opr.query_frames(geometry=region).to_crs('EPSG:3031')
print(f"Found {len(gdf)} radar frames in the selected region.")
gdf.head()

radar_frames_hv = gdf.hvplot(by='collection', hover_cols=['id'])
(velocity_map * coastline * region_hv * radar_frames_hv).opts(aspect='equal', legend_position='top_left')

layer_ds_list = []

for id, frame in tqdm(gdf.iterrows(), total=len(gdf), mininterval=1.0, desc="Loading layers"):
    layers = opr.get_layers(frame)
    bed_layer_name = None
    if 'standard:bottom' in layers: # Generally, the picked bed should be in group "standard" with layer name "bottom"
        bed_layer_name = 'standard:bottom'
    elif ':bottom' in layers:
        bed_layer_name = ':bottom' # But occasionally it seems to be missing the group
    else:
        continue  # No bed layer found
    # Layers are stored in terms of two-way travel time to avoid any questions about travel speed within ice
    # This is different from how BedMap layers are stored, but it does make more sense when the radar data is availble to use twtt
    layer_wgs84 = xopr.radar_util.layer_twtt_to_range(layers[bed_layer_name], layers["standard:surface"], vertical_coordinate='wgs84').rename({'lat': 'Latitude', 'lon': 'Longitude'})
    layer_wgs84 = xopr.geometry.project_dataset(layer_wgs84, target_crs='EPSG:3031')
    layer_wgs84 = layer_wgs84.dropna('slow_time', subset=['wgs84'])
    layer_wgs84['source'] = id
    layer_ds_list.append(layer_wgs84)

bed_merged = xr.concat(layer_ds_list, dim='slow_time')

# Get extent of radar data for plots and bedmap comparison
xlim = (bed_merged.x.min().item(), bed_merged.x.max().item())
ylim = (bed_merged.y.min().item(), bed_merged.y.max().item())

bed_merged

bed_hv = bed_merged.hvplot.scatter(x='x', y='y', c='wgs84', cmap='turbo', s=2).opts(clabel='Bed Elevation WGS84 (m)')
(velocity_map.opts(colorbar=False) * coastline * region_hv * radar_frames_hv * bed_hv).opts(aspect='equal', legend_position='top_left', xlim=xlim, ylim=ylim)

def grid_dataarray(d: xr.DataArray, spacing=1000, aggregation_fns={'median': "median", 'std': 'std', 'count': "count"}):
    """
    Grid a DataArray with x,y coordinates into a regular grid using block aggregation.
    
    Parameters
    ----------
    d : xr.DataArray
        Input DataArray with 'x' and 'y' coordinates
    spacing : float
        Grid spacing in the same units as x,y coordinates
    aggregation_fns : dict
        Dictionary mapping aggregation function names to functions (e.g., {'median': np.median, 'std': np.std})
    
    Returns
    -------
    xr.Dataset
        Dataset with variables named {d.name}_{fn_name} for each aggregation function
    """
    # Get data extent
    x_min = d['x'].min().values
    x_max = d['x'].max().values
    y_min = d['y'].min().values
    y_max = d['y'].max().values
    
    # Extract coordinate and data values
    x_data = d['x'].values
    y_data = d['y'].values
    data_values = d.values
    
    # Create grid coordinates
    grid_x, grid_y = vd.grid_coordinates(
        region=(x_min, x_max, y_min, y_max),
        spacing=spacing
    )
    
    # Dictionary to store gridded results for each aggregation function
    data_vars = {}
    
    for fn_name, fn in aggregation_fns.items():
        # Use Verde's BlockReduce with the specified aggregation function
        gridder = vd.BlockReduce(
            reduction=fn, 
            spacing=spacing, 
            region=(x_min, x_max, y_min, y_max),
            center_coordinates=True
        )
        block_coords, block_values = gridder.filter(
            coordinates=(x_data, y_data), 
            data=data_values
        )
        
        # Initialize grid with NaN
        grid_data = np.full(grid_x.shape, np.nan)
        
        # Vectorized approach: compute indices directly from coordinates
        x_indices = np.floor((block_coords[0] - x_min) / spacing).astype(int)
        y_indices = np.floor((block_coords[1] - y_min) / spacing).astype(int)
        
        for x_idx, y_idx, value in zip(x_indices.flatten(), y_indices.flatten(), block_values.flatten()):
            grid_data[y_idx, x_idx] = value
        
        # Store in dictionary with name pattern
        var_name = f"{d.name}_{fn_name}" if d.name else f"data_{fn_name}"
        data_vars[var_name] = (['y', 'x'], grid_data)
    
    # Create Dataset with all aggregated variables
    return xr.Dataset(
        data_vars=data_vars,
        coords={
            'y': grid_y[:, 0],
            'x': grid_x[0, :]
        }
    )

gridded = grid_dataarray(bed_merged['wgs84'], spacing=5000)

gridded_median_hv = hv.Image(gridded, kdims=['x', 'y'], vdims=['wgs84_median', 'wgs84_std', 'wgs84_count']).opts(
    cmap='turbo',
    aspect='equal',
    tools=['hover'],
    colorbar=True,
    clabel='WGS84 Elevation (m)'
)

gridded_std_hv = hv.Image(gridded, kdims=['x', 'y'], vdims=['wgs84_std', 'wgs84_median', 'wgs84_count']).opts(
    cmap='inferno',
    aspect='equal',
    tools=['hover'],
    colorbar=True,
    clabel='Std of WGS84 Elevation (m)'
)

(velocity_map * region_hv * coastline * gridded_median_hv).opts(width=500, aspect='equal', xlim=xlim, ylim=ylim) + \
    (velocity_map * region_hv * coastline * gridded_std_hv).opts(width=500, aspect='equal', xlim=xlim, ylim=ylim)

# Create polygon from radar extent (EPSG:3031) and convert to WGS84 for bedmap query
extent_box_3031 = box(xlim[0], ylim[0], xlim[1], ylim[1])
radar_extent = xopr.geometry.project_geojson(extent_box_3031, source_crs='EPSG:3031', target_crs='EPSG:4326')

# Query the catalog to see what bedmap files match our region
print("Querying STAC catalog for matching bedmap files...")

catalog_items = query_bedmap_catalog(
    geometry=radar_extent,
    collections=['bedmap3']
)

if not catalog_items.empty:
    print(f"\nFound {len(catalog_items)} matching files:")
    for _, row in catalog_items.head(10).iterrows():
        props = row['properties'] if 'properties' in row else {}
        name = props.get('name', row.get('id', 'unknown'))
        row_count = props.get('row_count', 0)
        print(f"  - {name}: {row_count:,} rows")
    if len(catalog_items) > 10:
        print(f"  ... and {len(catalog_items) - 10} more")
else:
    print("No matching files found.")

# Query bedmap data matching the radar data extent for comparison
print("Querying bedmap data from cloud GeoParquet files...")
print(f"Query region bounds: {radar_extent.bounds}")

print("Timing cloud-based query...")
t0 = time.time()

cloud_result = query_bedmap(
    geometry=radar_extent,
    collections=['bedmap3'],
    exclude_geometry=True
)

cloud_time = time.time() - t0
print(f"Cloud query: {cloud_time:.2f}s for {len(cloud_result):,} rows")

if not cloud_result.empty:
    print(f"\nRetrieved {len(cloud_result):,} points from bedmap")
    print("\nColumns available:")
    print(cloud_result.columns.tolist())
    print("\nFirst 5 rows:")
    display(cloud_result.head())
else:
    print("No bedmap data found in query region.")

# Download bedmap STAC catalogs to local cache for faster subsequent queries
# This only needs to be done once - catalog files are cached locally
print("Fetching bedmap catalogs to local cache...")
t0 = time.time()
fetch_bedmap('bedmap3') # Default is all versions
download_time = time.time() - t0
print(f"Local cache: {download_time:.2f}s to download")

# Now time the same query using local cache
print("Timing local cache query...")
t0 = time.time()
local_result = query_bedmap(
    geometry=radar_extent,
    collections=['bedmap3'],
    exclude_geometry=True,
    local_cache=True  # Use locally cached files
)
local_time = time.time() - t0
print(f"Local query: {local_time:.2f}s for {len(local_result):,} rows")

# Compare performance
print(f"\nSpeedup: {cloud_time / local_time:.1f}x faster with local cache")

# Visualize bedmap bed elevation alongside OPR bed picks
# Filter to rows with valid bed elevation for apples-to-apples comparison
bedmap_with_bed = local_result.dropna(subset=['bedrock_altitude (m)'])
print(f"Points with bed elevation: {len(bedmap_with_bed):,} of {len(local_result):,}")

# Create GeoDataFrame for bedmap data
bedmap_gdf = gpd.GeoDataFrame(
    bedmap_with_bed,
    geometry=gpd.points_from_xy(bedmap_with_bed['lon'], bedmap_with_bed['lat']),
    crs='EPSG:4326'
).to_crs('EPSG:3031')

# Extract x/y coordinates for plotting
bedmap_gdf['x'] = bedmap_gdf.geometry.x
bedmap_gdf['y'] = bedmap_gdf.geometry.y

# Plot bedmap bed elevation
bedmap_hv = bedmap_gdf.hvplot.points(
    x='x', y='y', c='bedrock_altitude (m)',
    cmap='turbo', s=5, alpha=0.7
).opts(clabel='Bed Elevation (m)')

# Combine with OPR bed picks
(velocity_map.opts(colorbar=False) * coastline * region_hv * 
 radar_frames_hv * bedmap_hv).opts(
    aspect='equal', legend_position='top_left', xlim=xlim, ylim=ylim,
    title='Bedmap bed elevation in radar extent'
)