safe divide in zscore

matsengrp · Nov 6, 2024 · 85863e9 · 85863e9
1 parent 9bc578a
commit 85863e9
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/phippery/zscore.py b/phippery/zscore.py
@@ -119,14 +119,13 @@ def compute_zscore(
             data = data_df.loc[pid].to_numpy()
             null_means = means_df.loc[ibin].to_numpy()
             null_stds = stds_df.loc[ibin].to_numpy()
-            # TODO This line throws warnings about invalid values, and divide by zero
-            # Compute zs, avoiding division by zero
-            # zs = np.where(null_stds != 0, (data - null_means) / null_stds, 0)
-            zs = (data - null_means) / null_stds
-        zscore_df.loc[pid] = zs
+            # Replace invalid values in null_stds with NaN
+            null_stds = np.where(
+                np.isfinite(null_stds) & (null_stds != 0), null_stds, np.nan
+            )
 
-    # Convert Infs and NaNs to zero
-    zscore_df.replace([np.inf, -np.inf], np.nan, inplace=True)
-    zscore_df.fillna(value=0, inplace=True)
+            # Compute zs safely; np.nan_to_num replaces NaNs with 0 in the result
+            zs = np.nan_to_num((data - null_means) / null_stds, nan=0.0)
+        zscore_df.loc[pid] = zs
 
     return zscore_df, means_df, stds_df