Updraft Tracking

import datetime
import shutil
from pathlib import Path
from six.moves import urllib

import numpy as np
import pandas as pd
import xarray as xr

import matplotlib.pyplot as plt
%matplotlib inline

import glob
import cmweather

import tobac
print('using tobac version', str(tobac.__version__))

file = sorted(glob.glob("/data/project/ARM_Summer_School_2025/radar/dda/multidop/stonybrook/*")) #Star at the end pulls everything that's in that directory
file

ds = xr.open_mfdataset(file) #files[0] just for the first file. Add "mf" before "dataset()" to open multiple files
ds

ds["x"]

ds["vertical_wind_component"].isel(time=0).sel(z=5000).plot(cmap="Spectral_r")

ds = ds.sel(z=slice(0, 8_000))

vert_wind = ds["vertical_wind_component"]

vert_wind

#vert_wind = vert_wind.drop_vars('z') #This is in case we need to drop the vertical axis

#Defining the resolution of our data
dxy = 500
dt = 15*60

# Keyword arguments for feature detection step:
parameters_features={}
parameters_features['position_threshold']='weighted_diff'
parameters_features['sigma_threshold']=0.5
parameters_features['min_distance']=0
parameters_features['sigma_threshold']=1
parameters_features['threshold']=[3, 5, 10] #m/s
parameters_features['n_erosion_threshold']=0
parameters_features['n_min_threshold']=3

# Perform feature detection and save results:
print('start feature detection based on midlevel column maximum vertical velocity')
Features=tobac.feature_detection_multithreshold(vert_wind,dxy, vertical_cord='z', **parameters_features)
print('feature detection performed start saving features')
#Features.to_hdf(savedir / 'Features.h5', 'table')
print('features saved')

Features

# Keyword arguments for linking step:
parameters_linking={}
parameters_linking['method_linking']='predict'
parameters_linking['adaptive_stop']=0.5
parameters_linking['adaptive_step']=0.95
parameters_linking['extrapolate']=0
parameters_linking['order']=1
parameters_linking['subnetwork_size']=5
parameters_linking['memory']=0
parameters_linking['time_cell_min']=dt
parameters_linking['method_linking']='predict'
parameters_linking['v_max']=15

# Perform linking and save results:
Track=tobac.linking_trackpy(Features,vert_wind,dt=dt,dxy=dxy,**parameters_linking)
#Track.to_hdf(savedir / 'Track.h5', 'table')

Track

Track.groupby("cell")["time_cell"].max()

Track = Track[Track["cell"] != -1]

# Updraft lifetimes of tracked cells:
fig_lifetime,ax_lifetime=plt.subplots()
tobac.plot_lifetime_histogram_bar(Track,axes=ax_lifetime,bin_edges=np.arange(0,120,10),density=False,width_bar=8)
ax_lifetime.set_xlabel('lifetime (min)')
ax_lifetime.set_ylabel('counts')

#Segmentation
parameters_segmentation_TWC={}
parameters_segmentation_TWC['method']='watershed'
parameters_segmentation_TWC['threshold']=1 

print('Start segmentation based on total water content')
Mask_TWC, Features_TWC = tobac.segmentation_3D(Features, vert_wind, dxy, **parameters_segmentation_TWC)
print('segmentation TWC performed, start saving results to files')
#Mask_TWC.to_netcdf(savedir / 'Mask_Segmentation_TWC.nc', encoding={"segmentation_mask":{"zlib":True, "complevel":4}})
#Features_TWC.to_hdf(savedir / 'Features_TWC.h5','table')
print('segmentation TWC performed and saved')

Track[Track["frame"]==0]["x"]

ds["vertical_wind_component"].isel(time=0).sel(z=6000).plot(cmap="Spectral_r")
plt.scatter(Track[Track["frame"]==0]["x"], Track[Track["frame"]==0]["y"])

ds["vertical_wind_component"].isel(time=0).max(dim="z").plot(cmap="Spectral_r") #We're going to try to get the wmax. For that, we do .max along the z direction!
plt.scatter(Track[Track["frame"]==0]["x"], Track[Track["frame"]==0]["y"])

w_wind_max = ds["vertical_wind_component"].max(dim="z")

w_wind_max

# Keyword arguments for feature detection step:
parameters_features={}
parameters_features['position_threshold']='weighted_diff'
parameters_features['sigma_threshold']=0.5
parameters_features['min_distance']=0
parameters_features['sigma_threshold']=1
parameters_features['threshold']=[3, 5, 10] #m/s
parameters_features['n_erosion_threshold']=0
parameters_features['n_min_threshold']=3

# Perform feature detection and save results:
print('start feature detection based on midlevel column maximum vertical velocity')
Features=tobac.feature_detection_multithreshold(w_wind_max,dxy, vertical_cord='z', **parameters_features)
print('feature detection performed start saving features')
print('features saved')

Features

# Keyword arguments for linking step:
parameters_linking={}
parameters_linking['method_linking']='predict'
parameters_linking['adaptive_stop']=0.5
parameters_linking['adaptive_step']=0.95
parameters_linking['extrapolate']=0
parameters_linking['order']=1
parameters_linking['subnetwork_size']=5
parameters_linking['memory']=0
parameters_linking['time_cell_min']=dt
parameters_linking['method_linking']='predict'
parameters_linking['v_max']=15

# Perform linking and save results:
Track_w_wind_max = tobac.linking_trackpy(Features,w_wind_max,dt=dt,dxy=dxy,**parameters_linking)
#Track.to_hdf(savedir / 'Track.h5', 'table')

Track_w_wind_max

Track_w_wind_max.groupby("cell")["time_cell"].max()

Track_w_wind_max = Track_w_wind_max[Track_w_wind_max["cell"] != -1]

# Updraft lifetimes of tracked cells:
fig_lifetime,ax_lifetime=plt.subplots()
tobac.plot_lifetime_histogram_bar(Track_w_wind_max,axes=ax_lifetime,bin_edges=np.arange(0,120,10),density=False,width_bar=8)
ax_lifetime.set_xlabel('lifetime (min)')
ax_lifetime.set_ylabel('counts')

#Segmentation
parameters_segmentation_TWC={}
parameters_segmentation_TWC['method']='watershed'
parameters_segmentation_TWC['threshold']=1 

print('Start segmentation based on total water content')
Mask_TWC, Features_TWC = tobac.segmentation_3D(Features, w_wind_max, dxy, **parameters_segmentation_TWC)
print('segmentation TWC performed, start saving results to files')
#Mask_TWC.to_netcdf(savedir / 'Mask_Segmentation_TWC.nc', encoding={"segmentation_mask":{"zlib":True, "complevel":4}})
#Features_TWC.to_hdf(savedir / 'Features_TWC.h5','table')
print('segmentation TWC performed and saved')

Track_w_wind_max[Track_w_wind_max["frame"]==0]["x"]

ds["vertical_wind_component"].isel(time=0).max(dim="z").plot(cmap="Spectral_r") #We're going to try to get the wmax. For that, we do .max along the z direction!
plt.scatter(Track_w_wind_max[Track_w_wind_max["frame"]==0]["x"], Track_w_wind_max[Track_w_wind_max["frame"]==0]["y"])

w_wind_max.isel(time=0).plot()
plt.scatter(Track_w_wind_max[Track_w_wind_max["frame"]==0]["x"], Track_w_wind_max[Track_w_wind_max["frame"]==0]["y"])

w_wind_max.isel(time=1).plot()
plt.scatter(Track_w_wind_max[Track_w_wind_max["frame"]==1]["x"], Track_w_wind_max[Track_w_wind_max["frame"]==1]["y"])

#Setting up a for loop to plot more figures

w_wind_max.isel(time=2).plot()
plt.scatter(Track_w_wind_max[Track_w_wind_max["frame"]==2]["x"], Track_w_wind_max[Track_w_wind_max["frame"]==2]["y"])

single_time = Track_w_wind_max.time.unique()[2]
single_time

Track_w_wind_max.loc[(Track_w_wind_max.time == single_time) & (Track_w_wind_max.threshold_value == 10)]

def subset_column(ds, time, x, y):
    return ds.sel(time=time, x=x, y=y, method='nearest')

single_cell = Track_w_wind_max.loc[Track_w_wind_max.cell == 42]
ds_list = []
for time in single_cell.time.values:
    row = single_cell.loc[single_cell.time == time]
    ds_list.append(subset_column(ds, time=row.time.values, x=row.x.values, y=row.y.values))

single_cell_ds = xr.concat(ds_list, dim='time')
single_cell_ds
single_cell_ds.vertical_wind_component.max(["y", "x"]).plot(y='z')

# Filter for a single cell (cell 42)
single_cell = Track_w_wind_max.loc[Track_w_wind_max.cell == 42]

# Subset data at each time step using x, y, and time
ds_list = []
for time in single_cell.time.values:
    row = single_cell.loc[single_cell.time == time]
    ds_list.append(subset_column(ds, time=row.time.values, x=row.x.values, y=row.y.values))

# Concatenate all time slices into a single dataset
single_cell_ds = xr.concat(ds_list, dim='time')

# Plot the max vertical wind component over time and height
plot = single_cell_ds.vertical_wind_component.max(["y", "x"]).plot(
    y='z',
    cbar_kwargs={'label': 'Max Vertical Wind (m/s)'}  # custom colorbar label
)

# Add a custom title and axis labels
plt.title("Max Vertical Wind for Cell 42 Over Height and Time", fontsize=14)
plt.xlabel("Time")
plt.ylabel("Height (z)")

# Show plot
plt.show()

ax = plt.subplot(111)
#ds.sel(z=5000).isel(time=2).vertical_wind_component.plot()
ds.sel(z=3000).isel(time=2).radar_network_reflectivity.plot(cmap="Spectral_r", vmin=-20, vmax=70)
Track_w_wind_max.loc[Track_w_wind_max.cell == 42].plot.scatter(x="x", y='y', ax=ax, c="threshold_value", cmap="Reds")

Track_w_wind_max.loc[Track_w_wind_max.cell == 83]

y_line = Track_w_wind_max.groupby('cell')['y'].apply(list).loc[83]

y_line

x_line = Track_w_wind_max.groupby('cell')['x'].apply(list).loc[83]

#ax = plt.subplot(111)
#ds.sel(z=5000).isel(time=2).vertical_wind_component.plot()
#Track_w_wind_max.loc[Track_w_wind_max.cell == 83].plot.scatter(x="x", y='y', ax=ax, c="threshold_value", cmap="Reds")

fig, axs = plt.subplots(1, 3, figsize=(15, 6))
y1 = ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(cmap="Spectral_r", vmin=-20, vmax=70, ax = axs[0])
y2 = ds.sel(z=2000).isel(time=3).radar_network_reflectivity.plot(cmap="Spectral_r", vmin=-20, vmax=70, ax = axs[1])
y3 = ds.sel(z=2000).isel(time=4).radar_network_reflectivity.plot(cmap="Spectral_r", vmin=-20, vmax=70, ax = axs[2])
scatter_data = Track_w_wind_max.loc[Track_w_wind_max.cell == 83]

#axs[0].plot(x_line, y_line)
#axs[1].plot(x_line, y_line)
#axs[2].plot(x_line, y_line)


# Add the scatter and line plot to each subplot
for ax in axs:
    scatter_data.plot.scatter(x="x", y="y", ax=ax, c="threshold_value", cmap="Reds")
    ax.plot(x_line, y_line, color='black')  # Optional: specify color or linestyle

fig, axs = plt.subplots(1, 3, figsize=(15, 6), sharex=True, sharey=True)

# Plotting (as before)
mappable_refl = ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[0], add_colorbar=False)
ds.sel(z=2000).isel(time=3).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[1], add_colorbar=False)
ds.sel(z=2000).isel(time=4).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[2], add_colorbar=False)

# Scatter and lines
scatter_data = Track_w_wind_max.loc[Track_w_wind_max.cell == 83]
sc = axs[0].scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
for ax in axs:
    ax.scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
    ax.plot(x_line, y_line, color='black')
    ax.set_xlabel('')  # Remove individual x labels
    ax.set_ylabel('')  # Remove individual y labels

# Shared axis labels
fig.supxlabel("East West Distance (m)", fontsize=14)
fig.supylabel("North South Distance (m)", fontsize=14)

# Shared colorbars
cbar1 = fig.colorbar(mappable_refl, ax=axs, orientation='vertical', fraction=0.02, pad=0.04)
cbar1.set_label('Reflectivity (dBZ)')
cbar2 = fig.colorbar(sc, ax=axs, orientation='vertical', fraction=0.02, pad=0.01)
cbar2.set_label('Vertical Wind Threshold Value (m/s)')

fig, axs = plt.subplots(1, 3, figsize=(15, 6), sharex=True, sharey=True)
fig.suptitle("Radar Reflectivity and Wind Thresholds for Cell 83", fontsize=16, y=1.02)

# Reflectivity plots without auto-labels
mappable_refl = ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[0], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=3).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[1], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=4).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[2], add_colorbar=False, add_labels=False)

# Custom titles (example using time values)
axs[0].set_title(str(ds.time.values[2])[:16], fontsize=10)
axs[1].set_title(str(ds.time.values[3])[:16], fontsize=10)
axs[2].set_title(str(ds.time.values[4])[:16], fontsize=10)

# Scatter and lines
scatter_data = Track_w_wind_max.loc[Track_w_wind_max.cell == 83]
sc = axs[0].scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
for ax in axs:
    ax.scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
    ax.plot(x_line, y_line, color='black')
    ax.set_xlabel('')
    ax.set_ylabel('')

# Shared axis labels
fig.supxlabel("East West Distance (m)", fontsize=14)
fig.supylabel("North South Distance (m)", fontsize=14)

# Shared colorbars
cbar1 = fig.colorbar(mappable_refl, ax=axs, orientation='vertical', fraction=0.02, pad=0.04)
cbar1.set_label('Reflectivity (dBZ)')

cbar2 = fig.colorbar(sc, ax=axs, orientation='vertical', fraction=0.02, pad=0.01)
cbar2.set_label('Vertical Wind Threshold Value (m/s)')

plt.savefig('cell_83_graph.pdf', dpi=300)
plt.savefig("cell_83_graph.png", dpi=300, bbox_inches='tight')

single_cell = Track_w_wind_max.loc[Track_w_wind_max.cell == 83]
ds_list = []
for time in single_cell.time.values:
    row = single_cell.loc[single_cell.time == time]
    ds_list.append(subset_column(ds, time=row.time.values, x=row.x.values, y=row.y.values))

single_cell_ds = xr.concat(ds_list, dim='time')
single_cell_ds
single_cell_ds.vertical_wind_component.max(["y", "x"]).plot(y='z')

# Filter for cell 83
single_cell = Track_w_wind_max.loc[Track_w_wind_max.cell == 83]

# Subset data at each time step
ds_list = []
for time in single_cell.time.values:
    row = single_cell.loc[single_cell.time == time]
    ds_list.append(subset_column(ds, time=row.time.values, x=row.x.values, y=row.y.values))

# Concatenate all time slices
single_cell_ds = xr.concat(ds_list, dim='time')

# Plot max vertical wind component with custom colorbar label
plot = single_cell_ds.vertical_wind_component.max(["y", "x"]).plot(
    y='z',
    cbar_kwargs={'label': 'Max Vertical Wind (m/s)'}
)

# Add figure title and axes labels
plt.title("Max Vertical Wind for Cell 83 Over Height and Time", fontsize=14)
plt.xlabel("Time")
plt.ylabel("Height (z)")

# Show plot

plt.savefig("cell_83_secondgraph.png", dpi=300)

single_cell_ds.vertical_wind_component.max(dim=["x", "y", "z"]).plot()

ax = plt.subplot(111)
#ds.sel(z=5000).isel(time=2).vertical_wind_component.plot()
ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(cmap="Spectral_r", vmin=-20, vmax=70)
Track_w_wind_max.loc[Track_w_wind_max.cell == 83].plot. scatter(x="x", y='y', ax=ax, c="threshold_value", cmap="Reds")

Track_w_wind_max.loc[Track_w_wind_max.cell == 42]

y_line2 = Track_w_wind_max.groupby('cell')['y'].apply(list).loc[42]
x_line2 = Track_w_wind_max.groupby('cell')['x'].apply(list).loc[42]

fig, axs = plt.subplots(1, 5, figsize=(15, 6), sharex=True, sharey=True)

# Plotting (as before)
mappable_refl = ds.sel(z=2000).isel(time=0).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[0], add_colorbar=False)
ds.sel(z=2000).isel(time=1).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[1], add_colorbar=False)
ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[2], add_colorbar=False)
ds.sel(z=2000).isel(time=3).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[3], add_colorbar=False)
ds.sel(z=2000).isel(time=4).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[4], add_colorbar=False)

# Scatter and lines
scatter_data = Track_w_wind_max.loc[Track_w_wind_max.cell == 42]
sc = axs[0].scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
for ax in axs:
    ax.scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
    ax.plot(x_line2, y_line2, color='black')
    ax.set_xlabel('')  # Remove individual x labels
    ax.set_ylabel('')  # Remove individual y labels

# Shared axis labels
fig.supxlabel("East West Distance (m)", fontsize=14)
fig.supylabel("North South Distance (m)", fontsize=14)

# Shared colorbars
cbar1 = fig.colorbar(mappable_refl, ax=axs, orientation='vertical', fraction=0.02, pad=0.04)
cbar1.set_label('Reflectivity (dBZ)')
cbar2 = fig.colorbar(sc, ax=axs, orientation='vertical', fraction=0.02, pad=0.01)
cbar2.set_label('Vertical Wind Threshold Value (m/s)')

fig, axs = plt.subplots(1, 5, figsize=(15, 6), sharex=True, sharey=True)
fig.suptitle("Radar Reflectivity and Wind Thresholds for Cell 42", fontsize=16, y=1.02)

# Reflectivity plots with suppressed auto-labels
mappable_refl = ds.sel(z=2000).isel(time=0).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[0], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=1).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[1], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=2).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[2], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=3).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[3], add_colorbar=False, add_labels=False)
ds.sel(z=2000).isel(time=4).radar_network_reflectivity.plot(
    cmap="Spectral_r", vmin=-20, vmax=70, ax=axs[4], add_colorbar=False, add_labels=False)

# Add smaller custom titles using time values
for i in range(5):
    time_str = str(ds.time.values[i])[:16]  # Shorten for readability
    axs[i].set_title(f"Time: {time_str}", fontsize=10)

# Scatter and lines
scatter_data = Track_w_wind_max.loc[Track_w_wind_max.cell == 42]
sc = axs[0].scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
for ax in axs:
    ax.scatter(scatter_data["x"], scatter_data["y"], c=scatter_data["threshold_value"], cmap="Reds")
    ax.plot(x_line2, y_line2, color='black')
    ax.set_xlabel('')
    ax.set_ylabel('')

# Shared axis labels
fig.supxlabel("East West Distance (m)", fontsize=14)
fig.supylabel("North South Distance (m)", fontsize=14)

# Shared colorbars
cbar1 = fig.colorbar(mappable_refl, ax=axs, orientation='vertical', fraction=0.02, pad=0.04)
cbar1.set_label('Reflectivity (dBZ)')

cbar2 = fig.colorbar(sc, ax=axs, orientation='vertical', fraction=0.02, pad=0.01)
cbar2.set_label('Vertical Wind Threshold Value (m/s)')
plt.savefig("cell_42_graph.png", dpi=300, bbox_inches='tight')

x_line2

y_line2