num_traits and num_chromosomes arguments to annotate; closes #410

CHANGELOG.rst

@@ -27,9 +27,9 @@ https://tskit.dev/pyslim/docs/latest/previous_versions.html
   returned by `pyslim.mutation_metadata(ts)`.
 
 - The SLiM mutation IDs represented by each tskit mutation should no longer be
-  read in from the `derived_state` property, but instead from metadata.
-  (However, SLiM still writes these out in text to the `derived_state`
-  entry as before, so existing code will continue to work.)
+  read in from the `derived_state` property, but instead from the tskit mutation's
+  metadata. (However, SLiM still writes these out in text to the `derived_state`
+  entry as before.)
 
 - Previously, `msprime.sim_mutations` with the `msprime.SLiMMutationModel`
   would record SLiM metadata along with each new mutation. However, msprime
@@ -49,6 +49,9 @@ https://tskit.dev/pyslim/docs/latest/previous_versions.html
 
 **Bug fixes:**
 
+- `pyslim.annotate` now has a `num_chromosomes` argument. Previously it could not be
+  easily used to annotate multichromosome simulations with more than 8 chromosomes.
+
 - In some previous versions, converting files produced by a yet-older version of SLiM
   to the previously-current file version dropped some information from metadata:
   nucleotide values for mutations, and pedigree parent IDs for individuals. This only

pyslim/methods.py

@@ -1125,12 +1125,18 @@ def annotate(
     stage="early",
     reference_sequence=None,
     annotate_mutations=True,
+    num_chromosomes=1,
+    num_traits=1,
 ):
     """
     Takes a tree sequence (as produced by msprime, for instance), and adds in the
     information necessary for SLiM to use it as an initial state, filling in
     mostly default values. Returns a :class:`tskit.TreeSequence`.
 
+    This method sets up a tree sequence for a given number of chromosomes and/or traits
+    (since both of these affect metadata schemas, this is important),
+    but any information about these will need to be added after the fact.
+
     :param tskit.TreeSequence ts: A :class:`tskit.TreeSequence`.
     :param str model_type: SLiM model type: either "WF" or "nonWF".
     :param int tick: What tick number in SLiM correponds to
@@ -1143,6 +1149,8 @@ def annotate(
         equal to ts.sequence_length.
     :param bool annotate_mutations: Whether to replace mutation metadata
         with defaults. (If False, information about mutations is unchanged.)
+    :param int num_chromosomes: The number of chromosomes.
+    :param int num_traits: The number of traits.
     """
     tables = ts.dump_tables()
     annotate_tables(
@@ -1153,6 +1161,8 @@ def annotate(
         stage=stage,
         reference_sequence=reference_sequence,
         annotate_mutations=annotate_mutations,
+        num_chromosomes=num_chromosomes,
+        num_traits=num_traits,
     )
     return tables.tree_sequence()
 
@@ -1165,6 +1175,8 @@ def annotate_tables(
     stage="early",
     reference_sequence=None,
     annotate_mutations=True,
+    num_chromosomes=1,
+    num_traits=1,
 ):
     """
     Does the work of :func:`annotate`, but modifies the tables in place: so,
@@ -1190,7 +1202,9 @@ def annotate_tables(
             "but must be for loading into SLiM: generate mutations with "
             "sim_mutations(..., discrete_genome=True), not simulate()."
         )
-    top_metadata = default_slim_metadata("tree_sequence")["SLiM"]
+    top_metadata = default_slim_metadata(
+        "tree_sequence", num_chromosomes=num_chromosomes, num_traits=num_traits
+    )["SLiM"]
     top_metadata["model_type"] = model_type
     top_metadata["tick"] = tick
     top_metadata["cycle"] = cycle
@@ -1199,11 +1213,13 @@ def annotate_tables(
     if isinstance(md, dict) and "SLiM_mutation_list" in md:
         top_metadata["SLiM_mutation_list"] = md["SLiM_mutation_list"]
     ts_metadata = set_tree_sequence_metadata(tables, **top_metadata)
-    set_metadata_schemas(tables)
-    _annotate_nodes_individuals(tables, age=default_ages)
+    set_metadata_schemas(tables, num_chromosomes=num_chromosomes, num_traits=num_traits)
+    _annotate_nodes_individuals(
+        tables, age=default_ages, num_chromosomes=num_chromosomes, num_traits=num_traits
+    )
     _annotate_populations(tables)
     if annotate_mutations:
-        _annotate_sites_mutations(tables, ts_metadata=ts_metadata)
+        _annotate_sites_mutations(tables, ts_metadata=ts_metadata, num_traits=num_traits)
     if reference_sequence is not None:
         tables.reference_sequence.data = reference_sequence
 
@@ -1235,7 +1251,7 @@ def next_slim_mutation_id(ts):
     return max_id + 1
 
 
-def _annotate_nodes_individuals(tables, age):
+def _annotate_nodes_individuals(tables, age, num_chromosomes=1, num_traits=1):
     """
     Adds to a TableCollection the information relevant to individuals required
     for SLiM to load in a tree sequence, that is found in Node and Individual
@@ -1276,7 +1292,7 @@ def _annotate_nodes_individuals(tables, age):
             else:
                 ind_population[i] = n.population
                 ind_slim_id[i] = 1
-            md = default_slim_metadata("node")
+            md = default_slim_metadata("node", num_chromosomes=num_chromosomes)
             md["slim_id"] = nid
             nid += 1
         else:
@@ -1294,7 +1310,7 @@ def _annotate_nodes_individuals(tables, age):
     ind_flags = tables.individuals.flags
     for j, ind in enumerate(tables.individuals):
         if slim_ind[j]:
-            md = default_slim_metadata("individual")
+            md = default_slim_metadata("individual", num_traits=num_traits)
             md["pedigree_id"] = int(ind_slim_id[j])
             md["subpopulation"] = int(ind_population[j])
             md["age"] = age
@@ -1340,7 +1356,7 @@ def _annotate_populations(tables):
                 tables.populations[j] = p.replace(metadata=md)
 
 
-def _annotate_sites_mutations(tables, ts_metadata):
+def _annotate_sites_mutations(tables, ts_metadata, num_traits=1):
     """
     Adds to a TableCollection the information relevant to mutations required
     for SLiM to load in a tree sequence. This means adding metadata to the
@@ -1361,10 +1377,9 @@ def _annotate_sites_mutations(tables, ts_metadata):
             "metadata; this metadata will be overwritten."
         )
     num_mutations = tables.mutations.num_rows
-    default_mut = default_slim_metadata("mutation_list_entry")
-    slim_time = ts_metadata["SLiM"]["tick"] - np.floor(tables.mutations.time).astype(
-        "int"
-    )
+    default_mut = default_slim_metadata("mutation_list_entry", num_traits=num_traits)
+    t = ts_metadata["SLiM"]["tick"]
+    slim_time = t - np.floor(tables.mutations.time).astype("int")
     mutation_list = [
         {
             "mutation_id": j,

pyslim/slim_metadata.py

@@ -1885,7 +1885,9 @@ def update_tables(tables):
         # which prior to 0.5 was everything but generation and model_type.
         # Recovering from provenance has also been useful for operations
         # that discard metadata (eg as msprime did prior to 0.7.5).
-        values = default_slim_metadata("tree_sequence")["SLiM"]
+        values = default_slim_metadata("tree_sequence", num_chromosomes=1, num_traits=1)[
+            "SLiM"
+        ]
         prov = None
         file_version = "unknown"
         # use only the last SLiM provenance
@@ -1943,7 +1945,9 @@ def update_tables(tables):
                 "required"
             ]
             tables.metadata_schema = new_schema
-            defaults = default_slim_metadata("tree_sequence", num_traits=num_traits)
+            defaults = default_slim_metadata(
+                "tree_sequence", num_traits=num_traits, num_chromosomes=1
+            )
             for k in new_properties:
                 if k not in md["SLiM"]:
                     if k == "tick":

tests/test_annotation.py

@@ -1012,6 +1012,47 @@ def test_reload_annotate(self, restart_name, recipe, helper_functions, tmp_path)
         # check for equality, in everything but the last provenance
         verify_slim_restart_equality(in_ts, out_ts, phenotypes=False)
 
+    def test_annotate_chromosomes(self):
+        ts = msprime.sim_ancestry(
+            2,
+            population_size=2,
+            sequence_length=1,
+            recombination_rate=0.001,
+            random_seed=123,
+            ploidy=2,
+        )
+        for n in (1, 8, 23):
+            ats = pyslim.annotate(ts, model_type="nonWF", tick=1, num_chromosomes=n)
+            ms = pyslim.slim_node_metadata_schema(num_chromosomes=n)
+            assert ms == ats.tables.nodes.metadata_schema
+            k = pyslim.is_vacant_num_bytes(n)
+            v = ats.node(0).metadata["is_vacant"]
+            assert len(v) == k
+
+    def test_annotate_traits(self):
+        ts = msprime.sim_ancestry(
+            2,
+            population_size=2,
+            sequence_length=1,
+            recombination_rate=0.001,
+            random_seed=123,
+            ploidy=2,
+        )
+        for n in (1, 8, 23):
+            ats = pyslim.annotate(ts, model_type="nonWF", tick=1, num_traits=n)
+            mts = pyslim.add_mutation_metadata(
+                msprime.sim_mutations(ats, model=msprime.SLiMv6MutationModel(), rate=1.0)
+            )
+            assert mts.num_mutations > 0
+            ms = pyslim.slim_tree_sequence_metadata_schema(num_traits=n)
+            assert ms == mts.metadata_schema
+            md = mts.metadata["SLiM"]["traits"]
+            assert len(md) == n
+            ms = pyslim.slim_individual_metadata_schema(num_traits=n)
+            assert ms == mts.tables.individuals.metadata_schema
+            md = mts.metadata["SLiM_mutation_list"][0]
+            assert len(md["per_trait"]) == n
+
 
 class TestReload(tests.PyslimTestCase):
     """