Merge pull request #213 from benjeffery/fix-quant-cols

Correctly decode tszip cols

Author: benjeffery

tests/test_data_model.py

@@ -1,4 +1,5 @@
 import msprime
+import numpy as np
 import pytest
 import tskit
 import tszip
@@ -9,12 +10,38 @@
 
 
 def test_model(tmpdir):
+    # Generate a tree sequence with populations and migrations
+    N = 1000
+    demography = msprime.Demography()
+    demography.add_population(name="pop1", initial_size=N)
+    demography.add_population(name="pop2", initial_size=N)
+    demography.add_population(name="ancestral", initial_size=N)
+    demography.set_symmetric_migration_rate(["pop1", "pop2"], 0.01)
+    demography.add_population_split(
+        time=1000, derived=["pop1", "pop2"], ancestral="ancestral"
+    )
     ts = msprime.sim_ancestry(
-        recombination_rate=1e-3, samples=10, sequence_length=1_000, random_seed=42
+        samples={"pop1": 5, "pop2": 5},
+        demography=demography,
+        sequence_length=1e4,
+        record_migrations=True,
+        random_seed=42,
     )
-    ts = msprime.sim_mutations(ts, rate=1e-2, random_seed=43)
+    ts = msprime.sim_mutations(ts, rate=1e-8, random_seed=42)
+    assert ts.num_populations > 0
+    assert ts.num_sites > 0
+    assert ts.num_migrations > 0
+    assert ts.num_mutations > 0
+
     tables = ts.tables
     tables.nodes.metadata_schema = tskit.MetadataSchema({"codec": "json"})
+
+    # Give each individual a location
+    indiv_copy = tables.individuals.copy()
+    tables.individuals.clear()
+    for i, ind in enumerate(indiv_copy):
+        tables.individuals.append(ind.replace(location=[i / 2, i + 1]))
+
     ts = tables.tree_sequence()
 
     tszip.compress(ts, tmpdir / "test.tszip")
@@ -25,11 +52,44 @@ def test_model(tmpdir):
     assert tsm.name == "test"
     assert tsm.file_uuid == ts.file_uuid
     assert len(tsm.summary_df) == 9
-    assert len(tsm.edges_df) == ts.num_edges
     assert len(tsm.trees_df) == ts.num_trees
+
+    assert len(tsm.edges_df) == ts.num_edges
+    for col in ["left", "right", "parent", "child"]:
+        assert np.array_equal(tsm.edges_df[col].values, getattr(ts.tables.edges, col))
+
     assert len(tsm.mutations_df) == ts.num_mutations
+    for m1, m2 in zip(ts.mutations(), tsm.mutations_df.to_dict("records")):
+        assert m1.derived_state == m2["derived_state"]
+        assert m1.site == m2["site"]
+        assert m1.node == m2["node"]
+        assert m1.parent == m2["parent"]
+        assert m1.time == m2["time"]
+
     assert len(tsm.nodes_df) == ts.num_nodes
+    for col in ["time", "flags", "population", "individual"]:
+        assert np.array_equal(tsm.nodes_df[col].values, getattr(ts.tables.nodes, col))
+
     assert len(tsm.sites_df) == ts.num_sites
+    for m1, m2 in zip(ts.sites(), tsm.sites_df.to_dict("records")):
+        assert m1.ancestral_state == m2["ancestral_state"]
+        assert m1.position == m2["position"]
+
+    assert len(tsm.individuals_df) == ts.num_individuals
+    for m1, m2 in zip(ts.individuals(), tsm.individuals_df.to_dict("records")):
+        assert m1.flags == m2["flags"]
+        assert np.array_equal(m1.location, m2["location"])
+        assert np.array_equal(m1.parents, m2["parents"])
+
+    assert len(tsm.populations_df) == ts.num_populations
+    for m1, m2 in zip(ts.populations(), tsm.populations_df.to_dict("records")):
+        assert m1.metadata == m2["metadata"]
+
+    assert len(tsm.migrations_df) == ts.num_migrations
+    for col in ["left", "right", "node", "source", "dest", "time"]:
+        assert np.array_equal(
+            tsm.migrations_df[col].values, getattr(ts.tables.migrations, col)
+        )
 
 
 def test_model_errors(tmpdir):

tsbrowse/model.py

@@ -33,9 +33,9 @@ def __init__(self, tsbrowse_path):
         self.ts = tszip.load(tsbrowse_path)
         self.name = tsbrowse_path.stem
         self.full_path = tsbrowse_path
+        ts_tables = self.ts.tables
         for table_name in [
             "edges",
-            "trees",
             "mutations",
             "nodes",
             "sites",
@@ -44,30 +44,35 @@ def __init__(self, tsbrowse_path):
             "migrations",
             "provenances",
         ]:
+            ts_table = getattr(ts_tables, table_name)
             # filter out ragged arrays with offset
             array_names = set(root[table_name].keys())
             ragged_array_names = {
                 "_".join(name.split("_")[:-1])
                 for name in array_names
                 if "offset" in name
             }
-            array_names -= set(ragged_array_names)
             array_names -= {"metadata_schema"}
             array_names -= {f"{name}_offset" for name in ragged_array_names}
-            arrays = {name: root[table_name][name][:] for name in array_names}
-            ragged_array_names -= {"metadata"}
-            for name in ragged_array_names:
-                array = root[table_name][name][:]
-                offsets = root[table_name][f"{name}_offset"][:]
-                arrays[name] = np.array(
-                    [
-                        array[s].tobytes().decode("utf-8")
-                        for s in (
-                            slice(start, end)
-                            for start, end in zip(offsets[:-1], offsets[1:])
-                        )
-                    ]
-                )
+            arrays = {}
+            for name in array_names:
+                if hasattr(ts_table, name):
+                    if name in ragged_array_names:
+                        arrays[name] = [
+                            getattr(row, name) for row in getattr(self.ts, table_name)()
+                        ]
+                    else:
+                        arrays[name] = getattr(ts_table, name)
+                else:
+                    arrays[name] = root[table_name][name][:]
+            df = pd.DataFrame(arrays)
+            df["id"] = df.index
+            setattr(self, f"{table_name}_df", df)
+
+        for table_name in ["trees"]:
+            arrays = {
+                name: root[table_name][name][:] for name in root[table_name].keys()
+            }
             df = pd.DataFrame(arrays)
             df["id"] = df.index
             setattr(self, f"{table_name}_df", df)