From c0cac104ef3f2242568d09b54c014771bddd2290 Mon Sep 17 00:00:00 2001 From: peter Date: Mon, 14 Sep 2026 23:25:07 -0700 Subject: [PATCH] num_traits and num_chromosomes arguments to annotate; closes #410 --- CHANGELOG.rst | 9 ++++++--- pyslim/methods.py | 39 ++++++++++++++++++++++++++------------ pyslim/slim_metadata.py | 8 ++++++-- tests/test_annotation.py | 41 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 80 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2582e33..e1ef107 100644 --- a/CHANGELOG.rst +++ b/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 diff --git a/pyslim/methods.py b/pyslim/methods.py index e5627d4..ff71edd 100644 --- a/pyslim/methods.py +++ b/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, diff --git a/pyslim/slim_metadata.py b/pyslim/slim_metadata.py index dbebed8..49fa203 100644 --- a/pyslim/slim_metadata.py +++ b/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": diff --git a/tests/test_annotation.py b/tests/test_annotation.py index 0383890..95eb0ab 100644 --- a/tests/test_annotation.py +++ b/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): """