precip.py

import numpy as np
import xarray as xr

def compute_precipitation_probability(ds_t):
    g = 9.80665
    pressure = ds_t['pressure_level'] * 100
    dp = pressure.diff('pressure_level').values

    q = ds_t['q']  # shape: (pressure_level, lat, lon)
    q1 = q.isel(pressure_level=slice(None, -1))
    q2 = q.isel(pressure_level=slice(1, None))
    q_mid = (q1 + q2) / 2

    dp = dp[:q_mid.sizes['pressure_level']]
    dp_reshaped = dp[:, np.newaxis, np.newaxis]
    pwat = (q_mid * dp_reshaped / g).sum(dim='pressure_level')

    pwat_norm = (pwat - pwat.min()) / (pwat.max() - pwat.min())

    q_low = ds_t['q'].sel(pressure_level=850)
    u = ds_t['u'].sel(pressure_level=850)
    v = ds_t['v'].sel(pressure_level=850)
    dqdx = q_low.differentiate(coord='lon')
    dqdy = q_low.differentiate(coord='lat')
    conv = -(u * dqdx + v * dqdy)

    conv_norm = (conv - conv.min()) / (conv.max() - conv.min())

    ppp = 0.6 * pwat_norm + 0.4 * conv_norm
    return ppp.clip(0, 1)