Merge pull request #736 from Starlitnightly/master

Enhance _snapshot method in StructureWatcher to capture detailed colu…

Author: Starlitnightly

dynamo/prediction/fate.py

@@ -253,14 +253,18 @@ def _fate(
         t_stack, prediction_stack = np.hstack(t), np.hstack(prediction)
         n_cell, n_feature = init_states.shape
 
-        t_len = int(len(t_stack) / n_cell)
-        avg = np.zeros((n_feature, t_len))
+        if len(t_stack) > 0 and len(prediction_stack) > 0:
+            t_len = int(len(t_stack) / n_cell)
+            avg = np.zeros((n_feature, t_len))
 
-        for i in range(t_len):
-            avg[:, i] = np.mean(prediction_stack[:, np.arange(n_cell) * t_len + i], 1)
+            for i in range(t_len):
+                avg[:, i] = np.mean(prediction_stack[:, np.arange(n_cell) * t_len + i], 1)
 
-        prediction = [avg]
-        t = [np.sort(np.unique(t))]
+            prediction = [avg]
+            t = [np.sort(np.unique(t))]
+        else:
+            # If stack is empty (e.g. failed integrations), keep original or handle as needed
+            pass
 
     return t, prediction
 

dynamo/tools/_track.py

@@ -77,10 +77,24 @@ def _get_column_type(self, series):
 
     def _snapshot(self, adata):
         """Captures keys, types, and shapes for comparison."""
+        obs_types = {}
+        for c in adata.obs.columns:
+            try:
+                obs_types[c] = self._get_column_type(adata.obs[c])
+            except Exception:
+                obs_types[c] = "unknown (error)"
+
+        var_types = {}
+        for c in adata.var.columns:
+            try:
+                var_types[c] = self._get_column_type(adata.var[c])
+            except Exception:
+                var_types[c] = "unknown (error)"        
+
         snapshot = {
             "shape": adata.shape,
-            "obs": {c: self._get_column_type(adata.obs[c]) for c in adata.obs.columns},
-            "var": {c: self._get_column_type(adata.var[c]) for c in adata.var.columns},
+            "obs": obs_types,
+            "var": var_types,
             "uns": set(adata.uns.keys()),
             # For complex slots, store a dict of {key: description_string}
             "obsm": {k: self._get_type_desc(v) for k, v in adata.obsm.items()},