Access and Orthorectify PRISM L1B Radiance

import earthaccess
import numpy as np
import pandas as pd
import geopandas as gpd
import xarray as xr
import rioxarray
import pyproj
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.io.img_tiles as cimgt
import warnings
warnings.filterwarnings('ignore')

xyz = "https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}"
attr = "ESRI"

auth = earthaccess.login()

instrument = "PRISM"
datasets = earthaccess.search_datasets(instrument=instrument)
print(f"{instrument} datasets found: {len(datasets)}\n")
for ds in datasets:
    print(f"  Short name : {ds['umm']['ShortName']}")

shortname = "PRISM_L1B_PACEPAX_2024_2515"
collections = earthaccess.search_datasets(short_name=shortname)

for c in collections:
    temporal = c.get_umm('TemporalExtents')[0]['RangeDateTimes'][0]
    print(f"Title      : {c.get_umm('EntryTitle')}")
    print(f"ShortName  : {c.get_umm('ShortName')}")
    print(f"Temporal   : {temporal['BeginningDateTime']} to {temporal['EndingDateTime']}")
    size = c.get_umm('ArchiveAndDistributionInformation')['FileDistributionInformation'][0]
    print(f"Total size : {size['TotalCollectionFileSize']} {size['TotalCollectionFileSizeUnit']}")

short_name = 'PRISM_L1B_PACEPAX_2024_2515'
granules = earthaccess.search_data(short_name=short_name)
print(f"Granules found: {len(granules)}")

def granule_details(g):
    gname = g.data_links()[0].split('/')[-1]
    sp = gname.removesuffix('.nc').split('_')
    return {
        'file': gname,
        'flight_line': sp[2],
        'scene': sp[-1],
        'start_time': pd.to_datetime(g['umm']['TemporalExtent']['RangeDateTime']['BeginningDateTime']),
        'size_mb': g['size'],
    }

df_granules = pd.DataFrame([granule_details(g) for g in granules])
df_granules.head()

df_granules['date'] = df_granules['start_time'].dt.date
df_granules.groupby('date').agg(
    flight_lines=('flight_line', 'nunique'),
    scenes=('file', 'count'),
    size_gb=('size_mb', lambda s: round(s.sum() / 1000, 1)),
)

# bounding box (west, south, east, north) over Monterey Bay, California
bbox = (-122.2, 36.5, -121.75, 37.05)

granules_mb = earthaccess.search_data(short_name=short_name, 
                                      bounding_box=bbox)
print(f"Granules over Monterey Bay: {len(granules_mb)}")

# convert to geopandas
gdf = gpd.GeoDataFrame(
    [granule_details(g) for g in granules_mb],
    geometry=gpd.GeoSeries(granules_mb, crs=4326),
)
gdf['start_time'] = gdf['start_time'].astype(str)
gdf.groupby('flight_line').size().rename('scenes').to_frame()

# plot granule footprints, colored by flight line
gdf.explore(
    column='flight_line', tiles=xyz, attr=attr, cmap='tab10',
    style_kwds={"weight": 2, "fillOpacity": 0.5},
    tooltip=['file', 'start_time'],
)

granule_name = 'PACEPAX-PRISM-L1B_ER2_20240926181605_R0_019.nc'
granule = [g for g in granules_mb if g.data_links()[0].endswith(granule_name)][0]

granule

f = earthaccess.open([granule])[0]
print(f"Granule opened: {f}")

dt = xr.open_datatree(f, engine='h5netcdf')
dt

root = dt.to_dataset()
rad = dt['radiance'].to_dataset()
glt = dt['geolocation_lookup_table'].to_dataset()

print(dt.attrs['title'])
print(f"Acquisition: {dt.attrs['time_coverage_start']} to {dt.attrs['time_coverage_end']}")
print(f"Radiance shape (wavelength, lines, samples): {rad['radiance'].shape}")
print(f"GLT shape (northing, easting): {glt['line'].shape}")

crs = pyproj.CRS.from_wkt(dt['transverse_mercator'].attrs['crs_wkt'])
pixel_size = float(glt['easting'].diff('easting')[0])
print(f"CRS: {crs.name} (EPSG:{crs.to_epsg()})")
print(f"GLT pixel size: {pixel_size:.1f} m")

wl = rad['wavelength']
fwhm = rad['fwhm']
print(f"Number of bands : {wl.size}")
print(f"Wavelength range : {wl.min():.1f} to {wl.max():.1f} nm")
print(f"Band spacing : {np.diff(wl).mean():.2f} nm")
print(f"FWHM range : {fwhm.min():.2f} to {fwhm.max():.2f} nm")

fwhm.plot.line("g--", marker="o", alpha=0.3)
plt.grid(True, linestyle="--", alpha=0.6)
plt.show()

def stretch(band, lower=2, upper=98):
    """percentile stretch to the 0-1 range"""
    p_low, p_high = np.nanpercentile(band, [lower, upper])
    return np.clip((band - p_low) / (p_high - p_low), 0, 1)

rgb_wl = [640, 550, 460]
rgb_raw = rad['radiance'].sel(wavelength=rgb_wl, method='nearest').load()

rgb_raw_img = np.dstack([stretch(b) for b in rgb_raw.values])

fig, ax = plt.subplots(figsize=(7, 7))
ax.imshow(rgb_raw_img)
ax.set_xlabel('samples (across-track)')
ax.set_ylabel('lines (along-track)')
ax.set_title('PRISM Calibrated Radiance (True Color)')
plt.show()

def orthorectify(data, glt, crs):
    # map a (..., lines, samples) raw-geometry array onto the GLT's UTM grid
    line = glt['line'].values
    sample = glt['sample'].values
    valid = ~np.isnan(line)
    # lines/samples starts with 1
    # subtracting 1 to convert to 0-based index
    # converting both -ve/+ve values to absolute
    li = np.abs(line[valid]).astype(int) - 1   
    si = np.abs(sample[valid]).astype(int) - 1

    data = np.asarray(data)
    out = np.full(data.shape[:-2] + line.shape,
                  np.nan, dtype=np.float32)
    out[..., valid] = data[..., li, si]

    da = xr.DataArray(out, dims=('band', 'y', 'x'),
                      coords={'y': glt['northing'].values, 
                              'x': glt['easting'].values})
    return da.rio.write_crs(crs)

rgb_ortho = orthorectify(rgb_raw.values, glt, crs)
rgb_ortho = rgb_ortho.assign_coords(band=rgb_raw['wavelength'].values)
rgb_ortho

rgb_ortho_img = np.dstack([stretch(b) for b in rgb_ortho.values])
# show pixels outside the swath as white
rgb_ortho_img[np.isnan(rgb_ortho.values[0])] = 1

fig, axes = plt.subplots(1, 2, figsize=(15, 7))
axes[0].imshow(rgb_raw_img)
axes[0].set_title('PRISM Radiance (Unorthorectified) True Color RGB')
axes[0].set_xlabel('Samples')
axes[0].set_ylabel('Lines')

extent = [float(rgb_ortho.x.min()), float(rgb_ortho.x.max()),
          float(rgb_ortho.y.min()), float(rgb_ortho.y.max())]
axes[1].imshow(rgb_ortho_img, extent=extent)
axes[1].set_title(f"PRISM Radiance (Orthorectified {crs.name}) True Color RGB")
axes[1].set_xlabel('Easting (m)')
axes[1].set_ylabel('Northing (m)')
plt.tight_layout()
plt.show()

# add an alpha band for transparent pixels for nan
swath = ~np.isnan(rgb_ortho.values[0])
rgba = np.dstack([rgb_ortho_img, swath.astype(float)])

# UTM projection
proj = ccrs.epsg(crs.to_epsg())
tiles = cimgt.GoogleTiles(url=xyz)

# buffer of 1.5 km around image
pad = 1500
map_extent = [extent[0] - pad, extent[1] + pad, 
              extent[2] - pad, extent[3] + pad]

# plot
fig = plt.figure(figsize=(10, 10))
ax = plt.axes(projection=proj)
ax.set_extent(map_extent, crs=proj)
ax.add_image(tiles, 12)
ax.imshow(rgba, extent=extent, origin='upper', transform=proj, zorder=2)
gl = ax.gridlines(draw_labels=True, alpha=0.3)
gl.top_labels = gl.right_labels = False
ax.set_title(f'PRISM Corrected Radiance True Color RGB\n{granule_name}')
plt.show()

# points of interest (latitude, longitude)
points = {
    'Open water': (36.625, -121.87),
    'Beach sand': (36.623, -121.847),
    'Vegetation': (36.565, -121.850),
    'Urban': (36.612, -121.848),
}

lat = root['lat'].values
lon = root['lon'].values

spectra = {}
for name, (plat, plon) in points.items():
    # nearest pixel 
    dist = (lat - plat)**2 + ((lon - plon) * np.cos(np.radians(plat)))**2
    i, j = np.unravel_index(np.nanargmin(dist), dist.shape)
    spectra[name] = rad['radiance'].isel(lines=i, samples=j).values
    print(f"{name:<12} line={i:<4} sample={j:<4}")

units = rad['radiance'].attrs['units']
spec_df = pd.DataFrame(spectra)
spec_df['Wavelength']=wl
spec_df.set_index('Wavelength', inplace=True)
spec_df.plot(marker="o", alpha=0.3, figsize=(10, 5),
            title='PRISM L1B Calibrated Radiance Spectra',
            ylabel=f"Radiance ({units})")
plt.grid(linestyle="--", alpha=0.5) 
plt.show()

out_tif = f"{granule_name}_RGB_ortho.tif"
rgb_ortho.rio.write_nodata(np.nan).rio.to_raster(out_tif, driver='COG')
print(f"Saved: {out_tif}")

# PRISM acquisition time
prism_time = pd.to_datetime(dt.attrs['time_coverage_start'])
prism_date = prism_time.strftime('%Y-%m-%d')

# bounding box (west, south, east, north) of the four points
pt_lats = [p[0] for p in points.values()]
pt_lons = [p[1] for p in points.values()]
pts_bbox = (min(pt_lons), min(pt_lats), max(pt_lons), max(pt_lats))

pace_granules = earthaccess.search_data(short_name='PACE_OCI_L2_SFREFL', version='3.1',
                                        bounding_box=pts_bbox,
                                        temporal=(prism_date, prism_date))
for g in pace_granules:
    pace_time = pd.to_datetime(g['umm']['TemporalExtent']['RangeDateTime']['BeginningDateTime'])
    print(f"{g['umm']['GranuleUR']} => {(pace_time - prism_time).total_seconds() / 3600:+.1f} hours from PRISM")

pace_f = earthaccess.open(pace_granules[:1])[0]
pace_dt = xr.open_datatree(pace_f, engine='h5netcdf')
pace_dt

pace_ds = xr.merge(
    (
        pace_dt.ds,
        pace_dt["geophysical_data"].to_dataset(),
        pace_dt["sensor_band_parameters"].coords,
        pace_dt["navigation_data"].ds.set_coords(("longitude", "latitude")).coords,
    )
)

# make 2D lat/lons selectable
pace_ds = pace_ds.set_xindex(("latitude", "longitude"), xr.indexes.NDPointIndex)

# add row, column numbers for each pixel to track pixel index
pace_ds = pace_ds.assign_coords(
    line_idx=("number_of_lines", np.arange(pace_ds.sizes["number_of_lines"])),
    pixel_idx=("pixels_per_line", np.arange(pace_ds.sizes["pixels_per_line"]))
)

pace_ds

print(f"Swath size (lines, pixels): {pace_ds["latitude"].shape}, bands: {pace_ds["wavelength_3d"].size}" \
      f" ({pace_ds["wavelength_3d"].min():.0f}-{pace_ds["wavelength_3d"].max():.0f} nm)")

flag_names = pace_ds["l2_flags"].attrs['flag_meanings'].split()
flag_bits = pace_ds["l2_flags"].attrs['flag_masks']

def decode_flags(value):
    """returns the names of the flags set in an l2_flag."""
    return [n for n, b in zip(flag_names, flag_bits) if int(value) & int(b)]

pace_spectra = {}
rows = []
for name, (plat, plon) in points.items():
    # select a point from the full PACE dataset
    pace_point = pace_ds.sel(latitude=plat, longitude=plon, method="nearest")

    # distance (km) from the point to each OCI pixel center
    dist = np.hypot((pace_point["latitude"] - plat) * 111.32,
                    (pace_point["longitude"] - plon) * 111.32 * np.cos(np.radians(plat)))
    
    pace_spectra[name] = pace_point["rhos"].values
    
    rows.append({'location': name, 
                 'line': pace_point.line_idx.item(), 
                 'pixel': pace_point.pixel_idx.item(),
                 'distance_km': round(float(dist.values), 2),
                 'flags': ', '.join(decode_flags(pace_point["l2_flags"].values))})
pd.DataFrame(rows)

# OCI swath window around the points
d0 = np.hypot((pace_ds["latitude"] - pt_lats[0]) * 111.32,
              (pace_ds["longitude"] - pt_lons[0]) * 111.32 * np.cos(np.radians(pace_ds["latitude"])))
i0, j0 = np.unravel_index(np.nanargmin(d0), d0.shape)
win = (slice(i0 - 15, i0 + 15),
       slice(j0 - 15, j0 + 15))
bands = [int(np.abs(pace_ds["wavelength_3d"].values - w).argmin()) for w in [645, 555, 465]]
pace_rgb = np.clip(pace_ds["rhos"][win[0], win[1], bands].values / 0.15, 0, 1)

fig, ax = plt.subplots(figsize=(8, 8))
ax.pcolormesh(pace_ds["longitude"][win], pace_ds["latitude"][win],
              np.nan_to_num(pace_rgb, nan=1), shading='nearest')
for (name, (plat, plon)), c in zip(points.items(), plt.cm.tab10.colors):
    ax.plot(plon, plat, 'o', color=c, mec='white', ms=10, label=name)
# PRISM scene extent
lat_p, lon_p = root['lat'].values, root['lon'].values
ax.add_patch(plt.Rectangle((np.nanmin(lon_p), np.nanmin(lat_p)),
                           np.nanmax(lon_p) - np.nanmin(lon_p),
                           np.nanmax(lat_p) - np.nanmin(lat_p),
                           fill=False, ec='k', lw=1.5, label='PRISM scene extent'))
ax.set_xlim(np.nanmin(lon_p) - 0.05, np.nanmax(lon_p) + 0.05)
ax.set_ylim(np.nanmin(lat_p) - 0.05, np.nanmax(lat_p) + 0.05)
ax.set_aspect(1 / np.cos(np.radians(np.mean(pt_lats))))
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.set_title(f"PACE OCI L2 surface reflectance, true color\n{pace_granules[0]['umm']['GranuleUR'][19:]}")
ax.legend(loc='lower left')
plt.show()

pace_df = pd.DataFrame(pace_spectra, index=pd.Index(pace_ds["wavelength_3d"].values, name='Wavelength'))
# keep only the OCI bands within the PRISM wavelength range
pace_df = pace_df[(pace_df.index >= float(wl.min())) & (pace_df.index <= float(wl.max()))]

fig, axes = plt.subplots(1, 2, figsize=(16, 5))
spec_df.plot(ax=axes[0], marker='o', ms=3, alpha=0.5,
             title='PRISM L1B radiance (~18 m pixels)', ylabel=f"Radiance ({units})")
pace_df.plot(ax=axes[1], marker='o', ms=3, alpha=0.5,
             title='PACE OCI L2 surface reflectance (~1.2 km pixels)', ylabel='Surface reflectance')
for ax in axes:
    ax.set_xlim(float(wl.min()), float(wl.max()))
    ax.grid(linestyle='--', alpha=0.5)
plt.tight_layout()
plt.show()