gif gridding debug

chenyangkang · Jun 3, 2024 · e26c34b · e26c34b
1 parent 7f60cf8
commit e26c34b
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/stemflow/utils/plot_gif.py b/stemflow/utils/plot_gif.py
@@ -146,9 +146,15 @@ def animate(i, norm, log_scale=log_scale):
         sub = data[data["Temporal_indexer"] == i].copy()
         temporal_value = np.array(sub[Temporal1].values)[0]
 
-        sub.loc[:, f"{Spatio1}_grid"] = np.digitize(sub[Spatio1], lng_gird, right=True)
-        sub.loc[:, f"{Spatio2}_grid"] = np.digitize(sub[Spatio2], lat_gird, right=False)
-        sub = sub[(sub[f"{Spatio1}_grid"] <= lng_size + 1) & (sub[f"{Spatio2}_grid"] <= lat_size + 1)]
+        g1 = np.digitize(sub[Spatio1], lng_gird, right=True)
+        g1 = np.where(g1 >= np.max(lng_gird), np.max(lng_gird) - 1, g1)
+
+        g2 = np.digitize(sub[Spatio2], lat_gird, right=True)
+        g2 = np.where(g2 >= np.max(lat_gird), np.max(lat_gird) - 1, g2)
+
+        sub.loc[:, f"{Spatio1}_grid"] = g1
+        sub.loc[:, f"{Spatio2}_grid"] = g2
+        sub = sub[(sub[f"{Spatio1}_grid"] <= lng_size - 1) & (sub[f"{Spatio2}_grid"] <= lat_size - 1)]
 
         sub = sub.groupby([f"{Spatio1}_grid", f"{Spatio2}_grid"])[[col]].mean().reset_index(drop=False)