Merge pull request #507 from jeromekelleher/update-node-date

Update node date

Author: jeromekelleher

sc2ts/stats.py

@@ -44,22 +44,30 @@ def node_data(ts, inheritance_stats=True):
 
     md = ts.nodes_metadata
     cols = {k: md[k].astype(str) for k in md.dtype.names}
+    dtype = {k: pd.StringDtype() for k in md.dtype.names}
     flags = ts.nodes_flags
     cols["node_id"] = np.arange(ts.num_nodes)
+    dtype["node_id"] = "int"
     cols["is_sample"] = (flags & tskit.NODE_IS_SAMPLE) > 0
+    dtype["is_sample"] = "bool"
     cols["is_recombinant"] = (flags & core.NODE_IS_RECOMBINANT) > 0
+    dtype["is_recombinant"] = "bool"
     # Are other flags useful of just debug info? Lets leave them out
     # for now.
     cols["num_mutations"] = np.bincount(ts.mutations_node, minlength=ts.num_nodes)
+    dtype["num_mutations"] = "int"
     # This is the same as is_recombinant but less obvious
     # cols["num_parents"] = np.bincount(ts.edges_child,
     #         minlength=ts.num_edges)
 
     if inheritance_stats:
         counter = jit.count(ts)
         cols["max_descendant_samples"] = counter.nodes_max_descendant_samples
-    cols["date"] = convert_date(ts, ts.nodes_time)
-    return pd.DataFrame(cols)
+        dtype["max_descendant_samples"] = "int"
+    if "time_zero_date" in ts.metadata:
+        cols["date"] = convert_date(ts, ts.nodes_time)
+        # Let Pandas infer the dtype of this to get the appropriate date type
+    return pd.DataFrame(cols).astype(dtype)
 
 
 def mutation_data(ts, inheritance_stats=True):
@@ -77,10 +85,15 @@ def mutation_data(ts, inheritance_stats=True):
     cols["node"] = ts.mutations_node
     cols["inherited_state"] = inherited_state
     cols["derived_state"] = derived_state
-    cols["date"] = convert_date(ts, ts.mutations_time)
+    if "time_zero_date" in ts.metadata:
+        cols["date"] = convert_date(ts, ts.mutations_time)
     if inheritance_stats:
         counter = jit.count(ts)
         cols["num_descendants"] = counter.mutations_num_descendants
         cols["num_inheritors"] = counter.mutations_num_inheritors
 
-    return pd.DataFrame(cols)
+    dtype = {k: "int" for k in cols if k != "date"}
+    dtype["inherited_state"] = pd.StringDtype()
+    dtype["derived_state"] = pd.StringDtype()
+
+    return pd.DataFrame(cols).astype(dtype)

tests/test_info.py

@@ -29,7 +29,12 @@ def fx_ti_2020_02_15(fx_ts_map):
 @pytest.fixture
 def fx_ts_min_2020_02_15(fx_ts_map):
     ts = fx_ts_map["2020-02-15"]
-    return sc2ts.minimise_metadata(ts)
+    field_mapping = {
+        "strain": "sample_id",
+        "Viridian_pangolin": "pango",
+        "Viridian_scorpio": "scorpio",
+    }
+    return sc2ts.minimise_metadata(ts, field_mapping)
 
 
 @pytest.fixture
@@ -55,6 +60,7 @@ def test_copying_table(self, fx_ti_recombinant_example_1):
         assert "TestChild" in ct_via_ti
         assert ct_via_ti == ct_via_ts
 
+
 class TestTreeInfo:
     def test_tree_info_values(self, fx_ti_2020_02_13):
         ti = fx_ti_2020_02_13
@@ -229,18 +235,30 @@ class TestDataFuncs:
 
     def test_example_node(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
         ts = fx_ts_min_2020_02_15
-        df = sc2ts.node_data(fx_ts_min_2020_02_15)
         ti = fx_ti_2020_02_15
+        df = sc2ts.node_data(ts)
         assert df.shape[0] == ti.ts.num_nodes
         nt.assert_array_equal(ti.nodes_num_mutations, df["num_mutations"])
         nt.assert_array_equal(np.arange(ti.ts.num_nodes), df["node_id"])
         nt.assert_array_equal(
             ti.nodes_max_descendant_samples, df["max_descendant_samples"]
         )
+        print(ti.nodes_date.dtype)
+        print(df["date"].dtype)
         nt.assert_array_equal(ti.nodes_date, df["date"])
         assert list(np.where(df["is_recombinant"])[0]) == list(ti.recombinants)
         assert list(np.where(df["is_sample"])[0]) == list(ts.samples())
 
+    def test_example_node_no_date(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
+        ts = fx_ts_min_2020_02_15
+        df1 = sc2ts.node_data(ts)
+        tables = ts.dump_tables()
+        tables.metadata = {}
+        df2 = sc2ts.node_data(tables.tree_sequence())
+        assert set(df1) == set(df2) | {"date"}
+        for col in df2:
+            nt.assert_array_equal(df1[col].values, df2[col].values)
+
     def test_example_mutation(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
         ts = fx_ts_min_2020_02_15
         df = sc2ts.mutation_data(fx_ts_min_2020_02_15)
@@ -253,3 +271,13 @@ def test_example_mutation(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
         nt.assert_array_equal(ti.mutations_inherited_state, df["inherited_state"])
         nt.assert_array_equal(ts.mutations_node, df["node"])
         nt.assert_array_equal(ts.mutations_parent, df["parent"])
+
+    def test_example_mutation_no_date(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
+        ts = fx_ts_min_2020_02_15
+        df1 = sc2ts.mutation_data(ts)
+        tables = ts.dump_tables()
+        tables.metadata = {}
+        df2 = sc2ts.mutation_data(tables.tree_sequence())
+        assert set(df1) == set(df2) | {"date"}
+        for col in df2:
+            nt.assert_array_equal(df1[col].values, df2[col].values)