WRF-Python Visibility Calculation

[Will H]

Just wanted to share my attempt at pulling out the AFWA visibility diagnostic from fortran into python code.

Please test and see what you think.




import numpy as np
from scipy.ndimage import gaussian_filter
from netCDF4 import Dataset
import wrf

def compute_afwa_visibility_miles(wrf_path):
    """
    AFWA surface visibility diagnostic from a WRF output.
    Returns:
        vis_miles  – visibility in miles
        alpha      – AFWA haze shape parameter
    """

    nc = Dataset(wrf_path)

    # --- Surface fields ---
    psfc = wrf.getvar(nc, "PSFC").values
    t2   = wrf.getvar(nc, "T2").values
    rh2  = wrf.getvar(nc, "rh2").values
    q2   = wrf.getvar(nc, "Q2").values

    # Hydrometeors (lowest model level)
    rain = wrf.getvar(nc, "QRAIN")[0,:,:].values
    snow = wrf.getvar(nc, "QSNOW")[0,:,:].values
    grau = wrf.getvar(nc, "QGRAUP")[0,:,:].values

    # Winds
    u10 = wrf.getvar(nc, "U10").values
    v10 = wrf.getvar(nc, "V10").values
    wind10 = np.sqrt(u10*u10 + v10*v10)

    # Precipitable water
    try:
        pw = wrf.getvar(nc, "pw")
        pw = pw[0,:,:].values if "Time" in pw.dims else pw.values
    except:
        pw = np.full_like(psfc, 20.0)

    # 125 m wind (fallback to 10 m)
    try:
        hgt = wrf.getvar(nc, "height_agl")[0,:,:,:]
        ua  = wrf.getvar(nc, "ua")[0,:,:,:]
        va  = wrf.getvar(nc, "va")[0,:,:,:]
        u125 = wrf.interplevel(ua, hgt, 125.0).values
        v125 = wrf.interplevel(va, hgt, 125.0).values
        wind125 = np.sqrt(u125*u125 + v125*v125)
    except:
        wind125 = wind10

    # -------------------------
    # AFWA VISIBILITY EQUATIONS
    # -------------------------
    R = 287.05
    rho = psfc / (R * t2)

    br = 1.1   * (1000*rho*(rain+grau))**0.75
    bs = 10.36 * (1000*rho*snow)**0.78
    ext  = (br + bs) / 1000.0

    visfactor = 3.912
    vis_hyd = np.where(ext > 0, visfactor / ext, 9.9e5)

    qsafe = np.where(q2 > 0, q2, np.nan)
    mask  = ~np.isnan(qsafe)

    vis_haze = np.full_like(psfc, 9.9e5)
    vis_haze_local = 1500.0 * (105 - rh2[mask]) * (5.0 / np.minimum(1000*qsafe[mask], 5.0))
    vis_haze[mask] = vis_haze_local

    alpha = np.full_like(psfc, 3.6)
    if np.any(mask):
        a_local = (
            0.1
            + pw[mask]/25.0
            + wind125[mask]/3.0
            + (100 - rh2[mask])/10.0
            + 1.0/(1000*qsafe[mask])
        )
        alpha[mask] = np.minimum(a_local, 3.6)

    # Final: min(hydrometeor, haze)
    vis_m = np.where(vis_hyd < vis_haze, vis_hyd, vis_haze)
    vis_m = gaussian_filter(vis_m, sigma=1)

    # Convert to miles
    vis_miles = vis_m * 0.000621371

    return vis_miles, alpha






import numpy as np
from scipy.ndimage import gaussian_filter
from netCDF4 import Dataset
import wrf

def compute_afwa_visibility_km(wrf_path):
    """
    AFWA surface visibility diagnostic from a WRF output.
    Returns:
        vis_km  – visibility in kilometers
        alpha   – AFWA haze shape parameter
    """

    nc = Dataset(wrf_path)

    # --- Surface fields ---
    psfc = wrf.getvar(nc, "PSFC").values
    t2   = wrf.getvar(nc, "T2").values
    rh2  = wrf.getvar(nc, "rh2").values
    q2   = wrf.getvar(nc, "Q2").values

    # Hydrometeors (lowest model level)
    rain = wrf.getvar(nc, "QRAIN")[0,:,:].values
    snow = wrf.getvar(nc, "QSNOW")[0,:,:].values
    grau = wrf.getvar(nc, "QGRAUP")[0,:,:].values

    # Winds
    u10 = wrf.getvar(nc, "U10").values
    v10 = wrf.getvar(nc, "V10").values
    wind10 = np.sqrt(u10*u10 + v10*v10)

    # Precipitable water
    try:
        pw = wrf.getvar(nc, "pw")
        pw = pw[0,:,:].values if "Time" in pw.dims else pw.values
    except:
        pw = np.full_like(psfc, 20.0)

    # 125 m wind (fallback to 10 m)
    try:
        hgt = wrf.getvar(nc, "height_agl")[0,:,:,:]
        ua  = wrf.getvar(nc, "ua")[0,:,:,:]
        va  = wrf.getvar(nc, "va")[0,:,:,:]
        u125 = wrf.interplevel(ua, hgt, 125.0).values
        v125 = wrf.interplevel(va, hgt, 125.0).values
        wind125 = np.sqrt(u125*u125 + v125*v125)
    except:
        wind125 = wind10

    # -------------------------
    # AFWA VISIBILITY EQUATIONS
    # -------------------------
    R = 287.05
    rho = psfc / (R * t2)

    br = 1.1   * (1000*rho*(rain+grau))**0.75
    bs = 10.36 * (1000*rho*snow)**0.78
    ext  = (br + bs) / 1000.0

    visfactor = 3.912
    vis_hyd = np.where(ext > 0, visfactor / ext, 9.9e5)

    qsafe = np.where(q2 > 0, q2, np.nan)
    mask  = ~np.isnan(qsafe)

    vis_haze = np.full_like(psfc, 9.9e5)
    vis_haze_local = 1500.0 * (105 - rh2[mask]) * (5.0 / np.minimum(1000*qsafe[mask], 5.0))
    vis_haze[mask] = vis_haze_local

    alpha = np.full_like(psfc, 3.6)
    if np.any(mask):
        a_local = (
            0.1
            + pw[mask]/25.0
            + wind125[mask]/3.0
            + (100 - rh2[mask])/10.0
            + 1.0/(1000*qsafe[mask])
        )
        alpha[mask] = np.minimum(a_local, 3.6)

    # Final: min(hydrometeor, haze)
    vis_m = np.where(vis_hyd < vis_haze, vis_haze, vis_hyd)
    vis_m = gaussian_filter(vis_m, sigma=1)

    # Convert to km
    vis_km = vis_m / 1000.0

    return vis_km, alpha