Better handling in TSI

src/cellrank/estimators/terminal_states/_gpcca.py

@@ -544,7 +544,7 @@ def tsi(
         cluster_key: Optional[str] = None,
         **kwargs: Any,
     ) -> float:
-        """Compute terminal state identificiation (TSI) score.
+        """Compute terminal state identification (TSI) score.
 
         Parameters
         ----------
@@ -573,10 +573,15 @@ def tsi(
             if cluster_key is None:
                 raise RuntimeError("`cluster_key` needs to be specified to compute TSI.")
 
+            # copy the `adata` object to avoid overwrites, as the estimator overrides some keys
+            adata = self.adata.copy()
             macrostates = {}
-            for n_states in range(n_macrostates, 0, -1):
-                self.compute_macrostates(n_states=n_states, cluster_key=cluster_key, **kwargs)
-                macrostates[n_states] = self.macrostates.cat.categories
+            for n_states in range(1, n_macrostates):
+                # create a new GPCCA object to avoid unsetting attributes that depend on the macrostates,
+                # e.g. the terminal states
+                g = GPCCA(self.transition_matrix, adata=adata)
+                g = g.compute_macrostates(n_states=n_states, cluster_key=cluster_key, **kwargs)
+                macrostates[n_states] = g.macrostates.cat.categories
 
             max_terminal_states = len(terminal_states)
 

tests/test_plotting.py

@@ -2259,7 +2259,6 @@ def test_scvelo_transition_matrix_projection(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_plot_tsi(self, mc: GPCCA, fpath: str):
-        mc = mc.copy(deep=True)
         terminal_states = ["Neuroblast", "Astrocyte", "Granule mature"]
         cluster_key = "clusters"
         _ = mc.tsi(n_macrostates=3, terminal_states=terminal_states, cluster_key=cluster_key, n_cells=10)