add WrappedTreeSequence properties for top-level metadata (#747)

Make WrappedTreeSequence.ts a property and the stored variable is now _ts

doc/pages/tskit_tools.md

@@ -17,7 +17,41 @@
     .. data:: fwdpy11.tskit_tools.INDIVIDUAL_IS_FIRST_GENERATION
 
     .. autoclass:: fwdpy11.tskit_tools.WrappedTreeSequence
-       :members:
+
+       This class has the following attributes:
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.data
+
+            Return the `data` field from top-level metadata
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.demes_graph
+
+            Return the `demes_graph` field from top-level metadata
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.generation
+
+            Return the `generation` field from top-level metadata
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.model_params
+
+            Return the `model_params` field from top-level metadata
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.seed
+
+            Return the `seed` field from top-level metadata
+
+       .. autoattribute:: fwdpy11.tskit_tools.WrappedTreeSequence.ts
+
+            Access the underlying :class:`tskit.TreeSequence`
+
+       This class has the following methods:
+
+       .. autofunction:: fwdpy11.tskit_tools.WrappedTreeSequence.timepoints_with_individuals
+
+       .. autofunction:: fwdpy11.tskit_tools.WrappedTreeSequence.decode_individual_metadata
+
+       .. autofunction:: fwdpy11.tskit_tools.WrappedTreeSequence.decode_mutation_metadata
+
 ```
 
 

fwdpy11/tskit_tools/trees.py

@@ -39,6 +39,11 @@ class WrappedTreeSequence(object):
     .. versionadded:: 0.15.0
     """
 
+    def _toplevel_metadata_value(self, name):
+        if name in self._ts.metadata:
+            return self._ts.metadata[name]
+        return None
+
     def __init__(self, ts: tskit.TreeSequence):
         found = False
         for row in ts.provenances():
@@ -49,7 +54,7 @@ def __init__(self, ts: tskit.TreeSequence):
         if not found:
             raise ValueError("this tree sequence was not generated by fwdpy11")
 
-        self.ts = ts
+        self._ts = ts
 
     def timepoints_with_individuals(self, *, decode_metadata=False):
         """
@@ -64,29 +69,29 @@ def timepoints_with_individuals(self, *, decode_metadata=False):
         a :class:`list` of :class:`fwdpy11.tskit_tools.DiploidMetadata`.
         """
         # Get rows of the node table where the nodes are in individuals
-        nodes_in_individuals = np.where(self.ts.tables.nodes.individual != tskit.NULL)[
+        nodes_in_individuals = np.where(self._ts.tables.nodes.individual != tskit.NULL)[
             0
         ]
 
         # Get the times
-        node_times = self.ts.tables.nodes.time[nodes_in_individuals]
+        node_times = self._ts.tables.nodes.time[nodes_in_individuals]
 
         unique_node_times = np.unique(node_times)
 
         for utime in unique_node_times[::-1]:
             # Get the node tables rows in individuals at this time
             x = np.where(node_times == utime)
             node_table_rows = nodes_in_individuals[x]
-            assert np.all(self.ts.tables.nodes.time[node_table_rows] == utime)
+            assert np.all(self._ts.tables.nodes.time[node_table_rows] == utime)
 
             # Get the individuals
-            individuals = np.unique(self.ts.tables.nodes.individual[node_table_rows])
+            individuals = np.unique(self._ts.tables.nodes.individual[node_table_rows])
             assert not np.any(individuals == tskit.NULL)
 
             if decode_metadata is True:
                 # now, let's decode the individual metadata for this time slice
                 decoded_individual_metadata = decode_individual_metadata(
-                    self.ts.tables,
+                    self._ts.tables,
                     individuals,
                 )
             else:
@@ -101,7 +106,7 @@ def decode_individual_metadata(
 
         See :func:`fwdpy11.tskit_tools.decode_individual_metadata` for details.
         """
-        return decode_individual_metadata(self.ts.tables, rows)
+        return decode_individual_metadata(self._ts.tables, rows)
 
     def decode_mutation_metadata(
         self, rows: typing.Optional[typing.Union[int, slice]] = None
@@ -111,4 +116,50 @@ def decode_mutation_metadata(
 
         See :func:`fwdpy11.tskit_tools.decode_mutation_metadata` for details.
         """
-        return decode_mutation_metadata(self.ts.tables, rows)
+        return decode_mutation_metadata(self._ts.tables, rows)
+
+    @property
+    def generation(self):
+        return self._toplevel_metadata_value("generation")
+
+    @property
+    def seed(self):
+        return self._toplevel_metadata_value("seed")
+
+    @property
+    def demes_graph(self):
+        rv = self._toplevel_metadata_value("demes_graph")
+        if rv is not None:
+            import demes
+
+            rv = demes.Graph.fromdict(rv)
+
+        return rv
+
+    @property
+    def model_params(self):
+        rv = self._toplevel_metadata_value("model_params")
+        if rv is not None:
+            import fwdpy11
+
+            try:
+                temp = {}
+                for key, value in rv.items():
+                    temp[key] = fwdpy11.ModelParams(**eval(value))
+                rv = temp
+            except:
+                rv = fwdpy11.ModelParams(**eval(rv))
+
+        return rv
+
+    @property
+    def data(self):
+        rv = self._toplevel_metadata_value("data")
+        if rv is None:
+            return rv
+
+        return rv
+
+    @property
+    def ts(self):
+        return self._ts

tests/conftest.py

@@ -16,9 +16,9 @@
 # You should have received a copy of the GNU General Public License
 # along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
 #
-import pytest
-
+import demes
 import fwdpy11
+import pytest
 
 
 @pytest.fixture(scope="function")
@@ -44,3 +44,52 @@ def mslike_pop(request):
         return fwdpy11.DiploidPopulation(N, 1.0)
     except AttributeError:  # NOQA
         return fwdpy11.DiploidPopulation(1000, 1.0)
+
+
+@pytest.fixture
+def gutenkunst():
+    yaml = """
+description: The Gutenkunst et al. (2009) OOA model.
+doi:
+- https://doi.org/10.1371/journal.pgen.1000695
+time_units: years
+generation_time: 25
+
+demes:
+- name: ancestral
+  description: Equilibrium/root population
+  epochs:
+  - {end_time: 220e3, start_size: 7300}
+- name: AMH
+  description: Anatomically modern humans
+  ancestors: [ancestral]
+  epochs:
+  - {end_time: 140e3, start_size: 12300}
+- name: OOA
+  description: Bottleneck out-of-Africa population
+  ancestors: [AMH]
+  epochs:
+  - {end_time: 21.2e3, start_size: 2100}
+- name: YRI
+  description: Yoruba in Ibadan, Nigeria
+  ancestors: [AMH]
+  epochs:
+  - start_size: 12300
+- name: CEU
+  description: Utah Residents (CEPH) with Northern and Western European Ancestry
+  ancestors: [OOA]
+  epochs:
+  - {start_size: 1000, end_size: 29725}
+- name: CHB
+  description: Han Chinese in Beijing, China
+  ancestors: [OOA]
+  epochs:
+  - {start_size: 510, end_size: 54090}
+
+migrations:
+- {demes: [YRI, OOA], rate: 25e-5}
+- {demes: [YRI, CEU], rate: 3e-5}
+- {demes: [YRI, CHB], rate: 1.9e-5}
+- {demes: [CEU, CHB], rate: 9.6e-5}
+"""
+    return demes.loads(yaml)

tests/test_tskit_metadata.py

@@ -48,55 +48,6 @@ def pdict2():
     return pd
 
 
-@pytest.fixture
-def gutenkunst():
-    yaml = """
-description: The Gutenkunst et al. (2009) OOA model.
-doi:
-- https://doi.org/10.1371/journal.pgen.1000695
-time_units: years
-generation_time: 25
-
-demes:
-- name: ancestral
-  description: Equilibrium/root population
-  epochs:
-  - {end_time: 220e3, start_size: 7300}
-- name: AMH
-  description: Anatomically modern humans
-  ancestors: [ancestral]
-  epochs:
-  - {end_time: 140e3, start_size: 12300}
-- name: OOA
-  description: Bottleneck out-of-Africa population
-  ancestors: [AMH]
-  epochs:
-  - {end_time: 21.2e3, start_size: 2100}
-- name: YRI
-  description: Yoruba in Ibadan, Nigeria
-  ancestors: [AMH]
-  epochs:
-  - start_size: 12300
-- name: CEU
-  description: Utah Residents (CEPH) with Northern and Western European Ancestry
-  ancestors: [OOA]
-  epochs:
-  - {start_size: 1000, end_size: 29725}
-- name: CHB
-  description: Han Chinese in Beijing, China
-  ancestors: [OOA]
-  epochs:
-  - {start_size: 510, end_size: 54090}
-
-migrations:
-- {demes: [YRI, OOA], rate: 25e-5}
-- {demes: [YRI, CEU], rate: 3e-5}
-- {demes: [YRI, CHB], rate: 1.9e-5}
-- {demes: [CEU, CHB], rate: 9.6e-5}
-"""
-    return demes.loads(yaml)
-
-
 @pytest.mark.parametrize("pop", [{"N": 100, "genome_length": 1}], indirect=["pop"])
 def test_single_model_params(pop, pdict1):
     mp = fwdpy11.ModelParams(**pdict1)
@@ -165,6 +116,10 @@ def test_user_defined_data(pop):
     ts = pop.dump_tables_to_tskit(data={"mydata": 11})
     assert ts.metadata["data"]["mydata"] == 11
 
+    # Test WrappedTreeSequence propery
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    assert wts.data["mydata"] == 11
+
     class MyType(object):
         def __init__(self, x):
             self.x = x
@@ -175,11 +130,19 @@ def __repr__(self):
     ts = pop.dump_tables_to_tskit(data=str(MyType(x=11)))
     assert eval(ts.metadata["data"]).x == 11
 
+    # Test WrappedTreeSequence property
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    assert eval(wts.data).x == 11
+
 
 @pytest.mark.parametrize("pop", [{"N": 100, "genome_length": 1}], indirect=["pop"])
 def test_seed(pop):
     ts = pop.dump_tables_to_tskit(seed=333)
     assert ts.metadata["seed"] == 333
 
+    # Test WrappedTreeSequence property
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    assert wts.seed == 333
+
     with pytest.raises(ValueError):
         _ = pop.dump_tables_to_tskit(seed=-333)

tests/test_wrapped_tskit_tree_sequence.py

@@ -40,7 +40,7 @@ def simulation():
 
     fwdpy11.evolvets(rng, pop, params, 100, recorder=recorder)
 
-    return pop
+    return pop, params
 
 
 @pytest.mark.parametrize("pop", [{"N": 100, "genome_length": 1}], indirect=["pop"])
@@ -53,10 +53,13 @@ def test_creation(pop):
 
 
 def test_on_simulated_example(simulation):
-    ts = simulation.dump_tables_to_tskit()
+    pop, params = simulation
+    ts = pop.dump_tables_to_tskit(model_params=params)
 
     wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts=ts)
 
+    assert wts.model_params == params
+
     times = []
     for time, nodes, md in wts.timepoints_with_individuals(decode_metadata=True):
         assert np.all(wts.ts.tables.nodes.time[nodes] == time)
@@ -71,6 +74,30 @@ def test_on_simulated_example(simulation):
                 assert wts.ts.tables.nodes.time[n] == time
         times.append(time)
 
-    assert np.all(
-        np.array(times) == np.arange(simulation.generation, dtype=np.float64)[::-1]
-    )
+    assert np.all(np.array(times) == np.arange(pop.generation, dtype=np.float64)[::-1])
+
+    ## Test storing two model param object
+    ts = pop.dump_tables_to_tskit(model_params={"stage1": params, "stage2": params})
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    mp = wts.model_params
+    assert len(mp) == 2
+    for _, value in mp.items():
+        assert value == params
+
+
+def test_generation_property(simulation):
+    pop, _ = simulation
+    assert pop.generation > 0
+    ts = pop.dump_tables_to_tskit()
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts=ts)
+    assert wts.generation == pop.generation
+
+
+def test_demes_graph_property(gutenkunst):
+    pop = fwdpy11.DiploidPopulation(100, 1.0)
+    ts = pop.dump_tables_to_tskit(demes_graph=gutenkunst)
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    assert wts.demes_graph == gutenkunst
+    ts = pop.dump_tables_to_tskit()
+    wts = fwdpy11.tskit_tools.WrappedTreeSequence(ts)
+    assert wts.demes_graph is None