From 376e51d79b55e8d3732befa118d52f4c2ed25087 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Thu, 12 Dec 2024 23:26:11 +0000
Subject: [PATCH] Rejig the flags values

---
 sc2ts/core.py      | 18 ++----------------
 sc2ts/info.py      | 11 ++++-------
 tests/test_info.py |  8 +++++++-
 3 files changed, 13 insertions(+), 24 deletions(-)

diff --git a/sc2ts/core.py b/sc2ts/core.py
index 831d00a..af19903 100644
--- a/sc2ts/core.py
+++ b/sc2ts/core.py
@@ -30,10 +30,8 @@
 NODE_IS_RECOMBINANT = 1 << 23
 NODE_IS_EXACT_MATCH = 1 << 24
 NODE_IS_IMMEDIATE_REVERSION_MARKER = 1 << 25
-NODE_IN_SAMPLE_GROUP = 1 << 26
-NODE_IN_RETROSPECTIVE_SAMPLE_GROUP = 1 << 27
-NODE_IS_REFERENCE = 1 << 28
-NODE_IS_UNCONDITIONALLY_INCLUDED = 1 << 29
+NODE_IS_REFERENCE = 1 << 26
+NODE_IS_UNCONDITIONALLY_INCLUDED = 1 << 27
 
 
 @dataclasses.dataclass(frozen=True)
@@ -77,18 +75,6 @@ class FlagValue:
         "Node is marking the existance of an immediate reversion which "
         "has not been removed for technical reasons",
     ),
-    FlagValue(
-        NODE_IN_SAMPLE_GROUP,
-        "G",
-        "SampleGroup",
-        "Node is a member of a sample group",
-    ),
-    FlagValue(
-        NODE_IN_RETROSPECTIVE_SAMPLE_GROUP,
-        "Q",
-        "RetroSampleGroup",
-        "Node is a member of a retrospective sample group",
-    ),
     FlagValue(
         NODE_IS_REFERENCE,
         "F",
diff --git a/sc2ts/info.py b/sc2ts/info.py
index c3c444c..49e54e6 100644
--- a/sc2ts/info.py
+++ b/sc2ts/info.py
@@ -469,10 +469,7 @@ def node_counts(self):
         pr_nodes = np.sum(self.ts.nodes_flags == core.NODE_IS_REVERSION_PUSH)
         re_nodes = np.sum(self.ts.nodes_flags == core.NODE_IS_RECOMBINANT)
         exact_matches = np.sum((self.ts.nodes_flags & core.NODE_IS_EXACT_MATCH) > 0)
-        sg_nodes = np.sum((self.ts.nodes_flags & core.NODE_IN_SAMPLE_GROUP) > 0)
-        rsg_nodes = np.sum(
-            (self.ts.nodes_flags & core.NODE_IN_RETROSPECTIVE_SAMPLE_GROUP) > 0
-        )
+        u_nodes = np.sum((self.ts.nodes_flags & core.NODE_IS_UNCONDITIONALLY_INCLUDED) > 0)
         immediate_reversion_marker = np.sum(
             (self.ts.nodes_flags & core.NODE_IS_IMMEDIATE_REVERSION_MARKER) > 0
         )
@@ -484,8 +481,7 @@ def node_counts(self):
             "mc": mc_nodes,
             "pr": pr_nodes,
             "re": re_nodes,
-            "sg": sg_nodes,
-            "rsg": rsg_nodes,
+            "u": u_nodes,
             "imr": immediate_reversion_marker,
             "zero_muts": nodes_with_zero_muts,
         }
@@ -734,6 +730,8 @@ def summary(self):
             ),
             ("max_samples_per_day", np.max(self.num_samples_per_day)),
             ("mean_samples_per_day", np.mean(self.num_samples_per_day)),
+            ("sample_groups", len(self.sample_group_nodes)),
+            ("retro_sample_groups", len(self.retro_sample_groups)),
         ]
         df = pd.DataFrame(
             {"property": [d[0] for d in data], "value": [d[1] for d in data]}
@@ -937,7 +935,6 @@ def retro_sample_groups_summary(self):
             data.append(d)
         return pd.DataFrame(data).set_index("group_id")
 
-
     def recombinants_summary(self):
         data = []
         for u in self.recombinants:
diff --git a/tests/test_info.py b/tests/test_info.py
index a8a7667..977db62 100644
--- a/tests/test_info.py
+++ b/tests/test_info.py
@@ -270,13 +270,19 @@ def test_node_summary(self, fx_ti_2020_02_13):
         for u in range(ti.ts.num_nodes):
             d = ti._node_summary(u)
             assert d["node"] == u
-            assert len(d["flags"]) == 10
+            assert len(d["flags"]) == 8
 
     def test_node_report(self, fx_ti_2020_02_13):
         ti = fx_ti_2020_02_13
         report = ti.node_report(strain="SRR11597190")
         assert len(report) > 0
 
+    def test_summary(self, fx_ti_2020_02_15):
+        df = fx_ti_2020_02_15.summary()
+        assert df.loc["samples"].value == 43
+        assert df.loc["sample_groups"].value == 27
+        assert df.loc["retro_sample_groups"].value == 1
+
 
 class TestSampleGroupInfo:
     def test_draw_svg(self, fx_ti_2020_02_13):