Update precip.py

precip.py

@@ -1,41 +1,31 @@
 import numpy as np
 import xarray as xr
 
-def compute_precipitation_probability(ds):
-    g = 9.80665  # gravity [m/s^2]
+def compute_precipitation_probability(ds_t):
+    g = 9.80665
+    pressure = ds_t['pressure_level'] * 100
+    dp = pressure.diff('pressure_level').values
 
-    # Get pressure levels and compute pressure differences (Pa)
-    pressure = ds['pressure_level'] * 100  # Convert hPa to Pa
-    dp = pressure.diff('pressure_level').values  # shape (12,)
-
-    # Compute precipitable water (PWAT)
-    q = ds['q']  # shape: (time, pressure_level, lat, lon)
+    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  # shape: (time, 12, lat, lon)
-
-    # Align dp shape with q_mid
-    dp = dp[:q_mid.sizes['pressure_level']]  # Ensure correct length (12,)
-    dp_reshaped = dp[:, np.newaxis, np.newaxis]  # shape (12, 1, 1)
+    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')
 
-    # Normalize PWAT
     pwat_norm = (pwat - pwat.min()) / (pwat.max() - pwat.min())
 
-    # Approximate low-level moisture convergence at 850 hPa
-    q_low = q.sel(pressure_level=850)
-    u = ds['u'].sel(pressure_level=850)
-    v = ds['v'].sel(pressure_level=850)
+    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)
 
-    # Normalize convergence
     conv_norm = (conv - conv.min()) / (conv.max() - conv.min())
 
-    # Combine into precipitation potential proxy (PPP)
     ppp = 0.6 * pwat_norm + 0.4 * conv_norm
-    ppp = ppp.clip(0, 1)
+    return ppp.clip(0, 1)
 
-    return ppp