Introduction to GPM-API

import datetime
import gpm
import fsspec
import numpy as np
import ximage  # noqa
import xarray as xr
from xarray.backends.api import open_datatree
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from gpm.utils.geospatial import (
    get_country_extent,
    get_geographic_extent_around_point,
    get_circle_coordinates_around_point,
)
from gpm.gv import volume_matching, compare_maps, calibration_summary

gpm.available_products(product_types="RS")  # research products

gpm.available_products(product_categories="RADAR", product_levels="1B")

gpm.available_products(product_categories="RADAR", product_levels="2A")

# Specify the time period you are interested in
start_time = datetime.datetime.strptime("2023-08-20 22:12:00", "%Y-%m-%d %H:%M:%S")
end_time = datetime.datetime.strptime("2023-08-20 22:13:45", "%Y-%m-%d %H:%M:%S")
# Specify the product and product type
product = "2A-DPR"  # 2A-PR
product_type = "RS"
storage = "GES_DISC"
# Specify the version
version = 7

# Download the data

# gpm.download(
#     product=product,
#     product_type=product_type,
#     version=version,
#     start_time=start_time,
#     end_time=end_time,
#     storage=storage,
#     force_download=False,
#     verbose=True,
#     progress_bar=True,
#     check_integrity=False,
# )

# Load the 2A-DPR dataset

# ds = gpm.open_dataset(
#     product=product,
#     product_type=product_type,
#     version=version,
#     start_time=start_time,
#     end_time=end_time,
# )
# ds

bucket_url = "https://js2.jetstream-cloud.org:8001"
fs = fsspec.filesystem("s3", anon=True, client_kwargs=dict(endpoint_url=bucket_url))
file = fs.get_mapper(
    f"pythia/radar/erad2024/gpm_api/2A.GPM.DPR.V9-20211125.20230820-S213941-E231213.053847.V07B.zarr"
)
ds = xr.open_zarr(file, consolidated=True, chunks={})

# Available variables
print("Available variables: ", list(ds.data_vars))
# Available coordinates
print("Available coordinates: ", list(ds.coords))
# Available dimensions
print("Available dimensions: ", list(ds.dims))
# Spatial dimension:
print("Spatial dimensions: ", ds.gpm.spatial_dimensions)
# Vertical dimension:
print("Vertical dimension: ", ds.gpm.vertical_dimension)

print(ds.gpm.spatial_2d_variables)
ds.gpm.select_spatial_2d_variables()

print(ds.gpm.spatial_3d_variables)
ds.gpm.select_spatial_3d_variables()

print(ds.gpm.frequency_variables)
ds.gpm.select_frequency_variables()

print(ds.gpm.bin_variables)
ds.gpm.select_bin_variables()

variable = "precipRateNearSurface"
da = ds[variable]
print(" Array Class: ", type(da.data))
da

da = da.compute()
print("Array Class: ", type(da.data))
da

print(
    ds.gpm.is_spatial_2d
)  # False because the xarray.Dataset also contains the range and frequency dimensions !
print(
    ds.gpm.is_spatial_3d
)  # False because the xarray.Dataset also contains frequency dimensions !

print(
    ds["zFactorFinal"].isel(range=0).sel(radar_frequency="Ka").gpm.is_spatial_2d
)  # True
print(ds["precipRateNearSurface"].gpm.is_spatial_2d)

ds["zFactorFinal"].sel(radar_frequency="Ka")

start_time = datetime.datetime.strptime("2023-08-20 22:12:00", "%Y-%m-%d %H:%M:%S")
end_time = datetime.datetime.strptime("2023-08-20 22:13:45", "%Y-%m-%d %H:%M:%S")
ds_subset = ds.gpm.sel(time=slice(start_time, end_time))
ds_subset["time"]

print(ds_subset.gpm.start_time)
print(ds_subset.gpm.end_time)

# Subset by gpm_id
start_gpm_id = "53847-2768"
end_gpm_id = "53847-2918"
ds_subset = ds.gpm.sel(gpm_id=slice(start_gpm_id, end_gpm_id))
ds_subset["gpm_id"].data

print(ds.gpm.has_contiguous_scans)
print(ds.gpm.is_regular)

list_slices = ds.gpm.get_slices_contiguous_scans()
print(list_slices)

slc = list_slices[0]
print(slc)
ds_regular = ds.isel(along_track=slc)

ds[variable].gpm.plot_map()

p = ds[variable].gpm.sel(gpm_id=slice(start_gpm_id, end_gpm_id)).gpm.plot_map()
p.axes.set_title(ds[variable].gpm.title(add_timestep=False))

ds[variable].gpm.plot_image()

ds[variable].gpm.sel(gpm_id=slice(start_gpm_id, end_gpm_id)).gpm.plot_image()

ds_subset = ds.gpm.sel(gpm_id=slice(start_gpm_id, end_gpm_id), radar_frequency="Ku")
p = ds_subset["zFactorFinalNearSurface"].gpm.plot_map()
p = ds_subset["zFactorFinalNearSurface"].gpm.plot_map(
    cmap="RdYlBu_r", vmin=15, vmax=45
)  # ex: enable to modify defaults parameters on the fly
p = ds_subset["flagPrecip"].gpm.plot_map()  # ex: defaults to categorical colorbar

plot_kwargs, cbar_kwargs = gpm.get_plot_kwargs("IMERG_Liquid")
ds_subset["precipRateNearSurface"].gpm.plot_map(cbar_kwargs=cbar_kwargs, **plot_kwargs)

ds_subset = ds.gpm.sel(gpm_id=slice(start_gpm_id, end_gpm_id))

# Plot reflectivity at various levels
variable = "zFactorFinal"
da = ds_subset[variable].sel(range=[100, 125, 150, 170])
fc = da.gpm.plot_map(row="radar_frequency", col="range")

# Surface reflectivity at Ku and Ka band
variable = "zFactorFinalNearSurface"
da = ds_subset[variable]
fc = da.gpm.plot_map(col="radar_frequency", col_wrap=2)

ds_subset = ds.gpm.sel(
    gpm_id=slice(start_gpm_id, end_gpm_id), radar_frequency="Ku"
).compute()
ds_subset["zFactorFinal"].isel(cross_track=24).gpm.plot_cross_section(zoom=True)
ds_subset["zFactorFinal"].isel(along_track=30).gpm.plot_cross_section(
    y="height_km", x="horizontal_distance_km", zoom=True
)

ds_transect = ds_subset.isel(cross_track=24)
text_kwargs = {
    "bbox": dict(
        facecolor="white", alpha=0.6, edgecolor="black", boxstyle="round,pad=0.2"
    )
}
p = ds_subset["zFactorFinalNearSurface"].gpm.plot_map()
ds_transect.gpm.plot_transect_line(ax=p.axes, text_kwargs=text_kwargs)

ds_transect = ds_subset.isel(cross_track=24)
print(ds_transect.gpm.cross_section_variables)
ds_transect.gpm.select_cross_section_variables()

da_cross = ds_subset["zFactorFinal"].gpm.extract_transect_between_points(
    start_point=(-118, 32), end_point=(-118, 36), steps=60
)
da_cross.gpm.plot_cross_section(zoom=True)

points = np.ones((100, 2))
points[:, 0] = -118
points[:, 1] = np.linspace(32, 36, 100)
da_cross = ds_subset["zFactorFinal"].gpm.extract_transect_at_points(points=points)
da_cross.gpm.plot_cross_section(zoom=True)

da_cross = ds_subset["zFactorFinal"].gpm.extract_transect_around_point(
    point=(-118, 34), azimuth=0, distance=100_000, steps=100
)  # azimuth [0, 360]
da_cross.gpm.plot_cross_section(zoom=True)

ds.gpm.available_retrievals()

ds_subset = ds.gpm.sel(
    gpm_id=slice(start_gpm_id, end_gpm_id), radar_frequency="Ku"
).compute()

ds_subset["EchoTopHeight40dBZ"] = ds_subset.gpm.retrieve("EchoTopHeight", threshold=40)
ds_subset["EchoTopHeight40dBZ"].gpm.plot_map()

ds_subset["EchoDepthSolid30dBZ"] = ds_subset.gpm.retrieve(
    "EchoDepth", threshold=30, mask_liquid_phase=True
)
ds_subset["EchoDepthSolid30dBZ"].gpm.plot_map()

ds_subset["VIL"] = ds_subset.gpm.retrieve("VIL")
ds_subset["VIL"].gpm.plot_map()

ds_subset["flagPrecipitationType"] = ds_subset.gpm.retrieve(
    "flagPrecipitationType", method="major_rain_type"
)
ds_subset["flagPrecipitationType"].gpm.plot_map()

ds_subset["flagHydroClass"] = ds_subset.gpm.retrieve(
    "flagHydroClass"
)  # this return a 3D array !
ds_subset["flagHydroClass"].isel(cross_track=24).gpm.plot_cross_section()

ds_surface = ds_subset.gpm.slice_range_at_bin(
    bins="binClutterFreeBottom"
)  # or ds.gpm.slice_range_at_bin(bins=ds["binClutterFreeBottom"])
ds_surface["precipRate"].gpm.plot_map()  # precipRate is originally a 3D variable !
ds_surface["zFactorFinal"].gpm.plot_map()  # zFactorFinal is originally a 3D variable !

ds_isothermal = ds_subset.gpm.slice_range_at_temperature(
    273.15
)  # use by default the airTemperature variable
ds_3000 = ds_subset.gpm.slice_range_at_height(3000)
ds_max_z = ds_subset.gpm.slice_range_at_max_value(variable="zFactorFinal")
ds_min_z = ds_subset.gpm.slice_range_at_min_value(variable="zFactorFinal")
ds_at_z_30 = ds_subset.gpm.slice_range_at_value(variable="zFactorFinal", value=30)

ds_valid = ds_subset.gpm.subset_range_with_valid_data(variable="zFactorFinal")
ds_intense = ds_subset.gpm.subset_range_where_values(
    variable="precipRate", vmin=200, vmax=300
)  # mm/hr
print("Range with valid data:", ds_valid["range"].data)
print("Range with intense precipitation:", ds_intense["range"].data)

ds_below_melting_layer = ds_subset.gpm.extract_dataset_below_bin(
    "binBBBottom"
)  # the first new range index corresponds to original binBBBottom
ds_above_melting_layer = ds_subset.gpm.extract_dataset_above_bin(
    "binBBTop"
)  # the last new range index corresponds to original binBBTop

ds_rain = ds_subset.gpm.mask_above_bin("binZeroDeg")
ds_snow = ds_subset.gpm.mask_below_bin("binZeroDeg")
ds_masked_melting_layer = ds_subset.gpm.mask_between_bins(
    bottom_bins="binBBBottom", top_bins="binBBTop"
)

# Define the variable to display
variable = "precipRateNearSurface"

# Crop by country
list_isel_dict = ds.gpm.get_crop_slices_by_country("United States")
print(list_isel_dict)

# Crop by extent
extent = get_country_extent("United States")
list_isel_dict = ds.gpm.get_crop_slices_by_extent(extent)
print(list_isel_dict)

# Plot the swaths crossing the country
for isel_dict in list_isel_dict:
    da_subset = ds[variable].isel(isel_dict)
    slice_title = da_subset.gpm.title(add_timestep=True)
    p = da_subset.gpm.plot_map()
    p.axes.set_extent(extent)
    p.axes.set_title(label=slice_title)

# Define the variable to display
variable = "precipRateNearSurface"

# Crop around a point (i.e. radar location)
lon = -117.0418
lat = 32.9190
distance = 200_000  # 200 km
list_isel_dict = ds.gpm.get_crop_slices_around_point(
    lon=lon, lat=lat, distance=distance
)
print(list_isel_dict)

extent = get_geographic_extent_around_point(lon=lon, lat=lat, distance=distance)
list_isel_dict = ds.gpm.get_crop_slices_by_extent(extent)
print(list_isel_dict)

# Define ROI coordinates
circle_lons, circle_lats = get_circle_coordinates_around_point(
    lon, lat, radius=distance, num_vertices=360
)

# Plot the swaths crossing the ROI
for isel_dict in list_isel_dict:
    da_subset = ds[variable].isel(isel_dict)
    slice_title = da_subset.gpm.title(add_timestep=True)
    p = da_subset.gpm.plot_map()
    p.axes.plot(circle_lons, circle_lats, "r-", transform=ccrs.Geodetic())
    p.axes.scatter(lon, lat, c="black", marker="X", s=100, transform=ccrs.Geodetic())
    p.axes.set_extent(extent)

# Alternatives if working with a single granule:

# ds_subset = ds.gpm.crop_by_extent(extent)
# ds_subset = ds.gpm.crop_by_country("United States")

# Define the variable to display
variable = "precipRateNearSurface"

# Crop around a point (i.e. radar location)
lon = -117.0418
lat = 32.9190
distance = 200_000  # 200 km
list_isel_dict = ds.gpm.get_crop_slices_around_point(
    lon=lon, lat=lat, distance=distance
)
print(list_isel_dict)

extent = get_geographic_extent_around_point(lon=lon, lat=lat, distance=distance)
list_isel_dict = ds.gpm.get_crop_slices_by_extent(extent)
print(list_isel_dict)

# Define ROI coordinates
circle_lons, circle_lats = get_circle_coordinates_around_point(
    lon, lat, radius=distance, num_vertices=360
)

# Plot the swaths crossing the ROI
for isel_dict in list_isel_dict:
    da_subset = ds[variable].isel(isel_dict)
    slice_title = da_subset.gpm.title(add_timestep=True)
    p = da_subset.gpm.plot_map()
    p.axes.plot(circle_lons, circle_lats, "r-", transform=ccrs.Geodetic())
    p.axes.scatter(lon, lat, c="black", marker="X", s=100, transform=ccrs.Geodetic())
    p.axes.set_extent(extent)

# Alternatives if working with a single granule:

# ds_subset = ds.gpm.crop_around_point(lon=lon, lat=lat, distance=distance)

print(da_subset.gpm.extent(padding=0.1))  # expanding
print(da_subset.gpm.extent(padding=-0.1))  # shrinking
print(da_subset.gpm.extent(size=0.5))
print(da_subset.gpm.extent(size=0))  # centroid

# Retrieve labeled xarray object
label_name = "label"
ds = ds.ximage.label(
    variable="precipRateNearSurface",
    min_value_threshold=1,
    min_area_threshold=5,
    footprint=5,  # assign same label to precipitating areas 5 pixels apart
    sort_by="area",  # "maximum", "minimum", <custom_function>
    sort_decreasing=True,
    label_name=label_name,
)
# Plot full label array
ds[label_name].ximage.plot_labels()

gpm.plot_labels(ds[label_name].isel(along_track=slice(2700, 3500)))

# Define the patch generator
da_patch_gen = ds["precipRateNearSurface"].ximage.label_patches(
    label_name=label_name,
    patch_size=(49, 49),
    # Output options
    n_patches=4,
    # Patch extraction Options
    padding=0,
    centered_on="max",
    # Tiling/Sliding Options
    debug=False,
    verbose=False,
)

# # Retrieve list of patches
list_label_patches = list(da_patch_gen)
list_da = [da for label, da in list_label_patches]

# Display patches
gpm.plot_patches(list_label_patches)

list_da_without_coords = [da.drop_vars(["lon", "lat"]) for da in list_da]
da_patch = xr.concat(list_da_without_coords, dim="patch")
da_patch.isel(patch=slice(0, 4)).gpm.plot_image(col="patch", col_wrap=2)

bucket_url = "https://js2.jetstream-cloud.org:8001"
fs = fsspec.filesystem("s3", anon=True, client_kwargs=dict(endpoint_url=bucket_url))
file = fs.get_mapper(f"pythia/radar/erad2024/gpm_api/KNKX20230820_221341_V06.zarr")
dt_gr = open_datatree(file, engine="zarr", consolidated=True, chunks={})
display(dt_gr)

sweep_idx = 0
sweep_group = f"sweep_{sweep_idx}"  # GR sweep (elevation to be used)
ds_gr = dt_gr[sweep_group].to_dataset().compute()

ds_gr["DBZH"].where(ds_gr["DBZH"] > -10).xradar_dev.plot_map()
ds_gr["DBZV"] = (ds_gr["DBZH"].gpm.idecibel / ds_gr["ZDR"].gpm.idecibel).gpm.decibel
ds_gr["DBZV"].where(ds_gr["DBZV"] > -10).xradar_dev.plot_map()
# ds_gr["ZDR"].where(ds_gr["DBZV"] > -10).xradar_dev.plot_map()

radar_band = "S"
beamwidth_gr = 1
z_min_threshold_gr = 0
z_min_threshold_sr = 10

gdf_match = volume_matching(
    ds_gr=ds_gr,
    ds_sr=ds,  # optional
    z_variable_gr="DBZH",
    radar_band=radar_band,
    beamwidth_gr=beamwidth_gr,
    z_min_threshold_gr=z_min_threshold_gr,
    z_min_threshold_sr=z_min_threshold_sr,
    min_gr_range=0,
    max_gr_range=150_000,
    # gr_sensitivity_thresholds=None,
    # sr_sensitivity_thresholds=None,
    download_sr=False,  # require NASA accounts and internet connection !
    display_quicklook=True,
)
display(gdf_match)

sr_z_column = f"SR_zFactorFinal_{radar_band}_mean"
gr_z_column = "GR_Z_mean"

# Compare reflectivities
fig = compare_maps(
    gdf_match,
    sr_column=sr_z_column,
    gr_column=gr_z_column,
    sr_label="SR Reflectivity (dBz)",
    gr_label="GR Reflectivity (dBz)",
    cmap="Spectral_r",
    unified_color_scale=True,
    vmin=15,
    # vmax=40
)
fig.tight_layout()

# Display calibration summary
fig = calibration_summary(
    df=gdf_match,
    gr_z_column=gr_z_column,
    sr_z_column=sr_z_column,
    # Histogram options
    bin_width=2,
    # Scatterplot options
    # hue_column="GR_gate_volume_sum",
    hue_column="SR_fraction_clutter",
    # gr_range=[15, 50]
    # sr_range=[15, 50]
    marker="+",
    cmap="Spectral",
)
fig.tight_layout()