From 3eaa65aead8eda3d6fe4377748066ba1863e1439 Mon Sep 17 00:00:00 2001
From: Aaron Meyer <ameyer@ucla.edu>
Date: Wed, 6 Mar 2024 07:51:20 -0800
Subject: [PATCH] Minor fix

---
 lineage/figures/figureS11.py | 3 +--
 lineage/plotTree.py          | 3 +++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/lineage/figures/figureS11.py b/lineage/figures/figureS11.py
index fe9f9d9ad..a8f7b004f 100644
--- a/lineage/figures/figureS11.py
+++ b/lineage/figures/figureS11.py
@@ -6,8 +6,7 @@
 from ..Analyze import Analyze_list
 from ..Lineage_collections import AllLapatinib
 
-num_states = 4
-lapt_tHMMobj_list = Analyze_list(AllLapatinib, num_states)[0]
+lapt_tHMMobj_list = Analyze_list(AllLapatinib, 4, write_states=True)[0]
 
 for i in range(4):
     lapt_tHMMobj_list[i].X = sort_lins(lapt_tHMMobj_list[i])
diff --git a/lineage/plotTree.py b/lineage/plotTree.py
index 37c167fbb..445908dcb 100644
--- a/lineage/plotTree.py
+++ b/lineage/plotTree.py
@@ -89,6 +89,9 @@ def plotLineage(lineage: LineageTree, axes, censor: bool = True, color: bool = T
     for ii, cell in enumerate(lineage.output_lineage):
         cell.state = lineage.states[ii]
 
+        if color:
+            assert cell.state >= 0
+
     root = lineage.output_lineage[0]
     if np.isfinite(root.obs[4]):  # the lineage starts from G1 phase
         if np.isfinite(root.obs[3]):