import xarray as xr, numpy as np
from wrf import getvar
from metpy.calc import brunt_vaisala_frequency as bvf
import metpy.units as units


from netCDF4 import Dataset
from wrf import getvar
import numpy as np

nc      = Dataset("wrfinput_d01")      # <-- your wrfinput path
w_field = getvar(nc, "W", timeidx=0, meta=False)   # derived vertical vel.

print("max |w|  (m s-1):", np.abs(w_field).max())

# --- 3.1  vertical velocity derived from mass divergence ---
w = getvar(ds, "w", timeidx=0)                # diagnostic W (not stored)
print("max|w|  (m s-1):", np.abs(w).max().values)

# --- 3.2  Brunt–Väisälä frequency N (static stability) ---
theta = getvar(ds, "T", timeidx=0)
z     = (ds.PH + ds.PHB)/9.80665
N     = bvf(z.metpy.quantify(), theta.metpy.quantify())
print("min N^2 (s-2):", N.min().to_base_units())