Load in necessary packages

# Libraries required for moisture convergence visualization
from datetime import datetime
import numpy as np
import xarray as xr
import xwrf
import glob
import metpy.calc as mpcalc
import math

import matplotlib.pyplot as plt

# Define path to the lasso simulation data
path_shcu_root = "/data/project/ARM_Summer_School_2024_Data/lasso_tutorial/ShCu/untar"  # on Jupyter

#Define LASSO SGP case date and simulation of interest
case_date = datetime(2019, 4, 4) #Options[April 4, 2019; May 6, 2019]
sim_id = 4

#Load in LASSO wrfstat files. These provide 10-minute averages for various metorology variables and diagnostics
ds_stat = xr.open_dataset(f"{path_shcu_root}/{case_date:%Y%m%d}/sim{sim_id:04d}/raw_model/wrfstat_d01_{case_date:%Y-%m-%d_12:00:00}.nc")
#ds_stat

#Load in LASSO-ShCu wrfout data, which is raw simulation output from the Weather Research and Forecasting (WRF) model run in an idealized LES mode.
#Post process using xwrf package
ds = xr.open_mfdataset(f"{path_shcu_root}/{case_date:%Y%m%d}/sim{sim_id:04d}/raw_model/wrfout_d01_*.nc", combine="nested", concat_dim="Time").xwrf.postprocess()

# By default, xarray does not interpret the wrfout/wrfstat time information in a way that attaches 
# it to each variable. Here is a trick to map the time held in XTIME with the Time coordinate 
# associated with each variable.
ds_stat["Time"] = ds_stat["XTIME"]
ds["Time"] = ds["XTIME"]
ds

#Load in u, v, and q data
U10 = ds["U10"]
V10 = ds["V10"]
QVAPOR = ds["QVAPOR"].sel(z=10, method='nearest').sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00")

#U and V have staggered x and y dimensions. The following unstaggers them to align with QVAPOR 
U = ds.U.interp(x_stag=ds.x)
V = ds.V.interp(y_stag=ds.y)

QVAPOR.shape

# We can use xarray's plotting features to get time-labeled plots.
hour_to_plot = 17 #UTC (6hrs after simulation start)

#This line shows the U winds at 10m from the surface at 18UTC
ds["U10"].sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00").plot()

# Calculate the divergence of the flow
# Multiply by the water vapor (QVAPOR) to get the moisture divergenc

div = mpcalc.divergence(QVAPOR*ds["U10"].sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"), QVAPOR*ds["V10"].sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"))

# start figure and set axis
fig, ax = plt.subplots(figsize=(5, 5))

# plot divergence and scale by 1e5
cf = ax.contourf(ds.y, ds.x, div*1e5 , range(-15, 16, 1), cmap=plt.cm.bwr_r) #* 1e5
plt.colorbar(cf, pad=0, aspect=50)
#ax.barbs(ds.y.values, ds.x.values, ds.U10.sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"), ds.V10.sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"), color='black', length=0.5, alpha=0.5)
#ax.set(xlim=(260, 270), ylim=(30, 40))
ax.set_title('Horizontal Moisture Divergence: 10m')
#ax.set_

#plt.show()

#Load in u, v, and q data
#U and V have staggered x and y dimensions. The following unstaggers them to align with QVAPOR 
U = ds.U.interp(x_stag=ds.x)
V = ds.V.interp(y_stag=ds.y)

z = 1000 

QVAPOR = ds["QVAPOR"].sel(z=z,method='nearest').sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00")
QVAPOR.shape

U_at_z = U.sel(z=z,method='nearest').sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00")
V_at_z = V.sel(z=z,method='nearest').sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00")

print(U_at_z.shape)
print(V_at_z.shape)

# Calculate the divergence of the flow
# Multiply by the water vapor (QVAPOR) to get the moisture divergenc

div2 = mpcalc.divergence(QVAPOR*U_at_z, QVAPOR*V_at_z)

# start figure and set axis
fig2, ax = plt.subplots(figsize=(5, 5))

# plot divergence and scale by 1e5
cf = ax.contourf(ds.y, ds.x, div2*1e5 , range(-15, 16, 1), cmap=plt.cm.bwr_r) #* 1e5
plt.colorbar(cf, pad=0, aspect=50)
#ax.barbs(ds.y.values, ds.x.values, ds.U10.sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"), ds.V10.sel(Time=f"{case_date:%Y-%m-%d} {hour_to_plot}:00"), color='black', length=0.5, alpha=0.5)
#ax.set(xlim=(260, 270), ylim=(30, 40))
ax.set_title('Horizontal Moisture Divergence: 1000m')


plt.show()

ds