diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index c79fac7..c5a3749 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -11,6 +11,7 @@ jobs: Docs: permissions: contents: read - uses: tskit-dev/.github/.github/workflows/docs.yml@v19 + uses: tskit-dev/.github/.github/workflows/docs.yml@v21 with: install-slim: true + install-slim-branch: "multitrait" diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 8b29520..1fc031f 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -10,4 +10,4 @@ jobs: Lint: permissions: contents: read - uses: tskit-dev/.github/.github/workflows/lint.yml@v19 + uses: tskit-dev/.github/.github/workflows/lint.yml@v21 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index c41d629..223052f 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -12,7 +12,7 @@ permissions: jobs: packaging: name: Python packaging - uses: tskit-dev/.github/.github/workflows/python-packaging.yml@v19 + uses: tskit-dev/.github/.github/workflows/python-packaging.yml@v21 test: name: Python @@ -93,6 +93,7 @@ jobs: git clone https://github.com/messerlab/SLiM.git mkdir -p SLiM/Release cd SLiM/windows_compat/gnulib + git checkout multitrait # <-- note multitrait branch!! touch --date="`date`" aclocal.m4 Makefile.am configure configure.ac config.h.in Makefile.in cd ../.. cd Release @@ -104,9 +105,27 @@ jobs: pip install uv uv sync --locked --group test --no-default-groups - - name: Install SLiM (macOS / Linux) - if: matrix.os == 'macos-latest' || matrix.os == 'ubuntu-24.04' - run: micromamba install slim -y + # UNCOMMENT THIS when the below is commented again + # - name: Install SLiM (macOS / Linux) + # if: matrix.os == 'macos-latest' || matrix.os == 'ubuntu-24.04' + # run: micromamba install slim -y + + - name: Install development SLiM + # This should be COMMENTED OUT for release versions, + # since this builds SLiM from github head. + # Also note that this checks out the multitrait branch!! + if: (matrix.os == 'macos-latest' || matrix.os == 'ubuntu-24.04') && steps.cache.outputs.cache-hit != 'true' + # If we want to re-build slim from a new commit to the slim repo + # we may need to bump the cache key above. + shell: bash -l {0} + run: | + git clone https://github.com/messerlab/SLiM.git + mkdir -p SLiM/Release + cd SLiM/Release + git checkout multitrait + cmake -DCMAKE_BUILD_TYPE=Release .. + make -j 2 + - name: Run tests run: | diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f61487d..a31bb70 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -2,7 +2,52 @@ [1.1.2] - 2026-XX-XX ******************** -In development +**Breaking changes**: + +- The release of SLiM 6.0, changes to metadata (see below) mean that accessing + top-level metadata (e.g., `ts.metadata["SLiM"]`) more than a few times in + a script will take a long time. Scripts that previously ran quickly may take a + prohibitively long. See the documentation for simple changes that fix the problem: + https://tskit.dev/pyslim/docs/latest/previous_versions.html + +- The SLiM tree sequence file version number has changed to 1.0. Use `pyslim.update` + to convert your tree sequence file to this format. + +- Metadata for SLiM's mutations are no longer stored along with the tskit mutations, + because mutation stacking allows each tskit mutation to be associated with more + than one SLiM mutation. Now, metadata for each unique mutation is stored in + top-level metadata, under `ts.metadata["SLiM_mutation_list"]`. The recommended + way to access this information is by obtaining the SLiM ID-to-metadata dict + returned by `pyslim.mutation_metadata(ts)`. + +- Previously, `msprime.sim_mutations` with the `msprime.SLiMMutationModel` + would record SLiM metadata along with each new mutation. However, msprime + does not modify top-level metadata, and so the method `add_mutation_metadata` + should be used after adding SLiM mutations. + +- This is a SLiM change, but top-level metadata is now encoded using the `json+struct` + codec now provided by tskit (so that the mutation metadata is not too large/slow). + +- The top-level and individual metadata schemas now depend on the number of traits + in the model. The methods `slim_tree_sequence_metadata_schema` and + `slim_individual_metadata_schema` can be used to produce correct schema. + +**Bug fixes:** + +- 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 + may have affected users using `pyslim.convert(ts)` in a previous version of pyslim + on a tree sequence `ts` with SLiM file version prior to 0.9. + +**New features**: + +- SLiM now includes in metadata information about the effects of mutations on + quantitative traits, the values of traits for individuals, and the values of + various "tags" defined in SLiM. + +- `default_slim_metadata` can now take additional arguments to modify the returned + values. ******************** [1.1.1] - 2026-03-06 diff --git a/docs/metadata.md b/docs/metadata.md index e4255d5..8332cdb 100644 --- a/docs/metadata.md +++ b/docs/metadata.md @@ -20,7 +20,7 @@ import random random.seed(23) ts = tskit.load("example_sim.trees") -tables = ts.tables +tables = ts.dump_tables() ``` ```{eval-rst} @@ -37,7 +37,7 @@ tables = ts.tables ## Overview SLiM puts SLiM-specific information into the *metadata* for the tree sequence, -as well as for each populations, individuals, nodes and mutations. +as well as for each population, individual, node and mutation. Here is a quick reference to what information is available: see the SLiM manual for the more technical writeup. A good way to get a generic metadata example is with {func}`.default_slim_metadata`. @@ -55,9 +55,26 @@ and `ts.metadata["SLiM"]` contains information about the simulation: - `spatial_dimensionality`: for instance, `""` or `"x"` or `"xy"` (etcetera) - `spatial_periodicity`: whether space wraps around in some directions (same format as dimensionality) - `stage`: the *stage* of the life cycle at which the file was written out (either `"first"`, `"early"`, or `"late"`) +- `name`: the *name* of this species in SLiM +- `this_chromosome`: contains, for the chromosome in SLiM recorded in this tree sequence + * `id`: SLiM's ID + * `index`: the index of this chromosome in the list of chromosomes + * `symbol`: the user-assigned symbol + * `type`: specifies inheritance type, e.g., `"A"` for autosome +- `chromosomes`: (optional) a list of all chromosomes in the simulation, + for each containing the same information as for `this_chromosome` +- `traits`: a list of information for each of the traits: + * `index`: the index of the trait in SLiM + * `name`: the name in SLiM for the trait + * `type`: additive, multiplicative, or logistic + * `baselineOffset`, `baselineAccumulation`: a value added to all traits, and whether the effect of substitutions + accumulate in that value + * `directFitnessEffect`: whether the trait has a direct effect on fitness + * `individualOffsetMean`, `individualOffsetSD`: parameters governing the individual-level offsets + (i.e., "environment" effects) **Populations:** -Information about each SLiM-produced population is written to metatadata. +Information about each SLiM-produced population is written to metadata. The format uses JSON and is extensible, so other keys may be present and some keys may be missing (for instance, there are no spatial bounds in a nonspatial simulation). The metadata may be `None` for populations @@ -68,7 +85,7 @@ that SLiM did not use. The keys that SLiM uses are: - `description`: a string describing the population - `selfing_fraction`, `female_cloning_fraction`, `male_cloning_fraction`, and `sex_ratio`: only present when applicable (e.g., in WF simulations) - `bounds_x0`, `bounds_x1`, `bounds_y0`, `bounds_y1`, `bounds_z0`, and `bounds_z1`: the spatial bounds, when applicable -- `migration_records`: A *list* of entries decribing migration between populations in a WF model. +- `migration_records`: A *list* of entries describing migration between populations in a WF model. **Individuals:** Each individual produced by SLiM contains the following metadata: @@ -81,27 +98,91 @@ Each individual produced by SLiM contains the following metadata: - `subpopulation`: the subpopulation within SLiM the individual was in at the time the file was written out - `sex`: the sex of the individual (either {data}`.INDIVIDUAL_TYPE_FEMALE`, {data}`.INDIVIDUAL_TYPE_MALE`, or {data}`.INDIVIDUAL_TYPE_HERMAPHRODITE`) - `flags`: additional information; currently only recording whether the individual was a "migrant" or not (see the SLiM manual) +- `tag`, `tagF`: the corresponding properties in SLiM: default values returned by pyslim + are the special values that SLiM uses to mean that the values are unset +- `tagL0`, `tagL0_set`, etcetera: again, the corresponding properties in SLiM; + the purpose of `tagLX_set` is to record whether the tag has been set in the simulation +- `per_trait`: a list of information about the trait values for this individual; these are in the same order + as the traits listed in top-level metadata; + * `phenotype`: the trait value + * `offset`: the individual's offset (i.e., the "environmental effect") **Nodes:** Each "node" produced by SLiM (i.e., "genome" within SLiM) has: -- 'slim_id': the unique ID associated with the genome by SLiM -- 'is_null': whether the genome is a "null" genome (in which case it isn't +- `slim_id`: the unique ID associated with the genome by SLiM +- `is_vacant`: records the genome is a "vacant" genome (in which case it isn't really there, so shouldn't have any mutations or relationships in the tree - sequence!) -- 'genome_type': the 'type' of this genome (0 for autosome, 1 for X, 2 for Y) + sequence!) - see [](sec_overview_vacant_nodes) for more explanation **Mutations:** -Each mutation's metadata is a dictionary with a single key, `"mutation_list"`, -whose entry is a *list* of metadata dictionaries corresponding to the mutations that are "stacked", -i.e., all present, in all genomes inheriting from this (tskit) mutation. -So, `ts.mutation(12).metadata["mutation_list"]` is a list, each of whose entries contains: +Prior to SLiM 6.0, mutation metadata was associated with the tskit mutation objects. +Now, this is stored in top-level metadata, under ``ts.metadata["SLiM_mutation_list"]``. +Each entry +- `mutation_id`: the numeric ID of mutation in SLiM - `mutation_type`: the numeric ID of the `MutationType` within SLiM -- `selection_coeff`: the selection coefficient - `subpopulation`: the numeric ID of the subpopulation the mutation occurred in - `slim_time`: the value of `community.tick` when the mutation occurred - `nucleotide`: either `-1` if there is no associated nucleotide, or the numeric code for the nucleotide (see {data}`.NUCLEOTIDES`) +- `per_trait`: a list of information in the same order as the traits in top-level metadata, recording for each: + * `effect_size`: the effect on the trait of this mutation + * `dominance`: its dominance coefficient + * `hemizygous_dominance`: its hemizygous dominance coefficient (see the SLiM manual) +- `padding`: this is simply empty bytes, here for byte-alignment reasons, and is always `None` + + +(sec_metadata_using_top_level)= + +## Using top-level metadata + +If you are going to be using information from top-level metadata, +it is good practice to extract the metadata as a separate python object once +and refer to that object, since otherwise you can incur runtime penalties +for decoding and copying the metadata every time you call `ts.metadata`. +This can be substantial, given the amount of mutation information +in top-level metadata. +For instance, to subtract baseline offsets from individual's trait values, +we might do: +```{code-cell} +md = ts.metadata +traits = md["SLiM"]["traits"] +values = [ + [x['phenotype'] - y["baselineOffset"] for x, y in zip(ind.metadata['per_trait'], traits)] + for ind in ts.individuals() +] +``` +If we instead inserted ``ts.metadata["SLiM"]["traits"]`` directly into the loop, +this would become infeasibly slow. + +In some more detail: +each time python evaluates ``ts.metadata`` (e.g., using ``ts.metadata["SLiM"]``) +a new copy of the metadata dict is decoded and returned. Furthermore, a number +of `pyslim` functions need to look up information from metadata under the hood. +For instance, previously it was acceptable to run +``[pyslim.slim_time(ts, mut.time) for mut in ts.mutations()]``. +However, this could now easily take hours even for moderately-sized simulations. +There are several recommendations for how to mitigate this: + +- If you use information from top-level metadata, make a copy of it + and refer to that copy instead: so, ``ts_metadata = ts.metadata`` + after ``ts = tskit.load(...)`` and then use ``ts_metadata``. However, + be careful that you use the correct metadata object corresponding + to the correct tree sequence object! + +- Use a single pyslim function call rather than many calls to the same one + within a loop. For instance, run: + ``slim_times = pyslim.slim_time(ts, ts.mutations_time)`` and extract + slim times from this vector. Similarly, use {func}`.nodes_vacant` + instead of {func}`.node_is_vacant`. + +- Some pyslim methods will accept a pre-extracted metadata dictionary + as an optional argument. If this is not provided, those methods will + extract the metadata again. The methods that now take a `ts_metadata` argument are: + {func}`.individual_ages`, + {func}`.individual_ages_at`, + {func}`.individuals_alive_at`, and + {func}`.slim_time`. (sec_metadata_tools)= @@ -110,10 +191,20 @@ So, `ts.mutation(12).metadata["mutation_list"]` is a list, each of whose entries The dictionaries describing the schema for these metadata entries are available in `pyslim.slim_metadata_schemas`. -Furthermore, this method may be useful in working with metadata: +Furthermore, these methods may be useful in working with metadata: ```{eval-rst} .. autofunction:: default_slim_metadata + +.. autofunction:: slim_tree_sequence_metadata_schema + +.. autofunction:: slim_individual_metadata_schema + +.. autofunction:: slim_node_metadata_schema + +.. autofunction:: set_tree_sequence_metadata + +.. autofunction:: set_metadata_schemas ``` @@ -125,13 +216,13 @@ see {ref}`tskit's metadata documentation `. ### Top-level metadata The entries of the top-level metadata dict are *read-only*. -So, you might think that +So, although you might think that `tables.metadata["SLiM"]["model_type"] = "nonWF"` would switch the model type, -but this in fact (silently) does nothing. To modify the top-level metadata, -we must (a) work with tables (as tree sequences are immutable, and (b) +this in fact (silently) does nothing. To modify the top-level metadata, +we must (a) work with tables (as tree sequences are immutable), and (b) extract the metadata dict, modify the dict, and copy it back in. -Instead, you should do +So, the code for modifying top-level metadata is: ```{code-cell} md = tables.metadata md["SLiM"]["model_type"] = "nonWF" diff --git a/docs/previous_versions.md b/docs/previous_versions.md index 3c1f74d..ce067fa 100644 --- a/docs/previous_versions.md +++ b/docs/previous_versions.md @@ -16,7 +16,7 @@ kernelspec: import pyslim, tskit, msprime ts = tskit.load("example_sim.trees") -tables = ts.tables +tables = ts.dump_tables() ``` @@ -25,6 +25,83 @@ tables = ts.tables # Migrating from previous versions of pyslim +## 1.2 + +Release 1.2 goes along with SLiM v6, which introduces support for traits. +It also changes the format for storing mutation metadata: now this is stored +in top-level metadata. + +1. Each time python evaluates ``ts.metadata`` (e.g., using ``ts.metadata["SLiM"]``) +a new copy of the metadata dict is decoded and returned. In large SLiM simulations, +this can take seconds, so we should avoid doing it many times. Furthermore, a number +of pyslim functions need to look up information from metadata under the hood. +See [](sec_metadata_using_top_level) for more discussion and examples. +In particular: + + - The method {func}`.node_is_vacant` necessarily uses metadata and acts + only on a single node. This method is now deprecated; + use {func}`.nodes_vacant` instead. + + - Some pyslim methods will accept a pre-extracted metadata dictionary + as an optional ``ts_metadata`` argument; see [](sec_metadata_using_top_level). + Furthermore, {func}`.is_current_version` now accepts top-level metadata directly + as an alternative to the tree sequence. + +2. If you are using `msprime` to generate mutations, you need to use +{func}`.add_mutation_metadata` after generating mutations to add the +information about these that SLiM expects to top-level metadata. +For instance: + +```{code-cell} +next_id = pyslim.next_slim_mutation_id(ts) +ts = pyslim.add_mutation_metadata( + msprime.sim_mutations( + ts, + rate=1e-8, + model=msprime.SLiMMutationModel(type=0, next_id=next_id), + ), + mutation_type=0, +) +``` +Here the ``mutation_type`` argument to {func}`.add_mutation_metadata` +is the important one; the ``type`` argument to ``SLiMMutationModel`` +is now deprecated, and will be effectively ignored. + +3. Instead of looking up metadata for mutations in `mut.metadata`, you need +to pull this information out of top-level metadata using the SLiM ID as a key. +In brief, if `mut` is a mutation, then you should replace +`mut.metadata["mutation_list"][j]` +with `mut_metadata[int(mut.derived_state.split(",")[j])]`, +where `mut_metadata` is the output of {func}`.mutation_metadata`. +For instance, where before you might have done: + +```python +mut = ts.mutation(0) +for k, md in zip(mut.derived_state.split(","), mut.metadata["mutation_list"]): + print(f"SLiM ID: {k}") + print(f"Metadata: {md}") +``` + +Now, you would do: + +```{code-cell} +mut_metadata = pyslim.mutation_metadata(ts) +mut = ts.mutation(0) +for k in mut.derived_state.split(","): + md = mut_metadata[int(k)] + print(f"SLiM ID: {k}") + print(f"Metadata: {md}") +``` + +The function {func}`.mutation_metadata` pulls information out of +`ts.metadata["SLiM_mutation_list"]`. It is useful for two reasons: +first, it puts the information into a dict, so you can look up information +using the SLiM mutation ID instead of searching through the list to find it. +Second, it caches the information: every time you access +`ts.metadata["SLiM_mutation_list"]`, it makes a new, decoded copy +of the entire metadata dictionary. This can be **very slow** if it is done +repeatedly. + ## 1.1 Release 1.1 goes along with SLiM v5, which introduces multichromosome support. diff --git a/docs/python_api.md b/docs/python_api.md index b570cd6..51e96d2 100644 --- a/docs/python_api.md +++ b/docs/python_api.md @@ -18,7 +18,7 @@ from IPython.display import SVG import numpy as np ts = tskit.load("example_sim.trees") -tables = ts.tables +tables = ts.dump_tables() ``` ```{eval-rst} @@ -38,6 +38,7 @@ Here is a quick reference to some of the methods: .. autosummary:: recapitate + mutation_metadata annotate individuals_alive_at individual_ages @@ -47,6 +48,9 @@ Here is a quick reference to some of the methods: has_vacant_samples node_is_vacant slim_time + next_slim_mutation_id + add_mutation_metadata + add_mutation_metadata_tables convert_alleles generate_nucleotides population_size @@ -89,6 +93,11 @@ Here is a quick reference to some of the methods: .. autofunction:: set_slim_state ``` +```{eval-rst} +.. autofunction:: add_mutation_metadata +.. autofunction:: add_mutation_metadata_tables +``` + ## Summarizing tree sequences Additionally, ``pyslim`` contains the following methods: @@ -119,6 +128,10 @@ Additionally, ``pyslim`` contains the following methods: ## Utilities +```{eval-rst} +.. autofunction:: mutation_metadata +``` + ```{eval-rst} .. autofunction:: slim_time ``` @@ -131,24 +144,23 @@ Additionally, ``pyslim`` contains the following methods: .. autofunction:: has_vacant_samples ``` +```{eval-rst} +.. autofunction:: nodes_vacant +``` + ```{eval-rst} .. autofunction:: node_is_vacant ``` +```{eval-rst} +.. autofunction:: is_current_version +``` + ## Metadata -SLiM-specific metadata is made visible to the user by ``.metadata`` properties. -For instance: -```{code-cell} -ts.individual(4).metadata -``` -shows that the fifth individual in the tree sequence was given pedigree ID ``495999`` by SLiM, -had parents with pedigree IDs ``493739`` and ``494784``, -was age 10 at the time that they died (or the simulation ended), -lived in subpopulation 1, -was female (because ``sex`` matches ``pyslim.INDIVIDUAL_TYPE_FEMALE``, below), -and has no additional metadata flags. +SLiM-specific metadata is made visible to the user by ``.metadata`` properties, +described in [](sec_metadata). ### Annotation diff --git a/docs/rapid_adaptation.slim b/docs/rapid_adaptation.slim index ed2ed49..909eeec 100644 --- a/docs/rapid_adaptation.slim +++ b/docs/rapid_adaptation.slim @@ -1,5 +1,5 @@ initialize() { - initializeTreeSeq(); + initializeTreeSeq(timeUnit="generations"); initializeMutationRate(1e-8); initializeMutationType("m1", 0.5, "e", 0.1); initializeGenomicElementType("g1", m1, 1.0); diff --git a/docs/selection.slim b/docs/selection.slim index 3c1860b..b617389 100644 --- a/docs/selection.slim +++ b/docs/selection.slim @@ -1,7 +1,7 @@ initialize() { initializeSLiMModelType("WF"); - initializeTreeSeq(); + initializeTreeSeq(timeUnit="generations"); initializeMutationRate(1e-6); initializeMutationType("m1", 0.5, "e", -0.1); initializeMutationType("m2", 0.5, "e", 0.5); diff --git a/docs/time_units.md b/docs/time_units.md index f997c49..419651b 100644 --- a/docs/time_units.md +++ b/docs/time_units.md @@ -226,7 +226,8 @@ so we expect generation time to go up at first. ```{code-cell} gts = tskit.load("generation_time.trees") -gentimes = gts.metadata["SLiM"]["user_metadata"]["generation_times"] +gts_metadata = gts.metadata +gentimes = gts_metadata["SLiM"]["user_metadata"]["generation_times"] fig, ax = plt.subplots(figsize=(12, 6), dpi=300) ax.set_xlabel("tick") @@ -266,7 +267,6 @@ Furthermore, since we already have mutations up until 100 time units ago, we need to put mutations on only previous to that time. ```{code-cell} -gentimes = gts.metadata["SLiM"]["user_metadata"]["generation_times"] gt = np.mean(gentimes[-50:]) recomb_rate = 1e-8 # per generation Ne = 1000 # generations diff --git a/docs/tutorial.md b/docs/tutorial.md index 35fc967..b7ff051 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -315,8 +315,8 @@ can be done with the {meth}`tskit.TreeSequence.simplify` method: ```{code-cell} import numpy as np rng = np.random.default_rng(seed=3) -alive_inds = pyslim.individuals_alive_at(rts, 0) -keep_indivs = rng.choice(alive_inds, 100, replace=False) +alive_indivs = pyslim.individuals_alive_at(rts, 0) +keep_indivs = rng.choice(alive_indivs, 100, replace=False) keep_nodes = [] for i in keep_indivs: keep_nodes.extend(rts.individual(i).nodes) @@ -364,11 +364,13 @@ This works as follows: ```{code-cell} next_id = pyslim.next_slim_mutation_id(sts) -ts = msprime.sim_mutations( +ts = pyslim.add_mutation_metadata( + msprime.sim_mutations( sts, rate=1e-8, model=msprime.SLiMMutationModel(type=0, next_id=next_id), keep=True, + ) ) print(f"The tree sequence now has {ts.num_mutations} mutations,\n" @@ -466,18 +468,18 @@ simulation (i.e., zero time units ago), we could do: ```{code-cell} orig_ts = tskit.load("example_sim.trees") -alive = pyslim.individuals_alive_at(orig_ts, 0) +alive_indivs = pyslim.individuals_alive_at(orig_ts, 0) -print(f"There are {len(alive)} individuals alive in the final generation.") +print(f"There are {len(alive_indivs)} individuals alive in the final generation.") ``` -Here, ``alive`` is a vector of *individual* IDs, +Here, ``alive_indivs`` is a vector of *individual* IDs, so one way to take a sample of living individuals and write their SNPs to a VCF is: ```{code-cell} rng = np.random.default_rng(seed=1) -keep_indivs = rng.choice(alive, 100, replace=False) +keep_indivs = rng.choice(alive_indivs, 100, replace=False) ts = msprime.sim_mutations(orig_ts, rate=1e-8, random_seed=1) with open("example_snps.vcf", "w") as vcffile: ts.write_vcf(vcffile, individuals=keep_indivs) @@ -517,10 +519,10 @@ For this reason, if at this point we try to extract genotypes for all of the alive individuals, we encounter a (somewhat confusing) error: ```{code-cell} +alive_indivs = pyslim.individuals_alive_at(ts, 0) try: - alive = pyslim.individuals_alive_at(ts, 0) with open("example_snps.vcf", "w") as vcffile: - ts.write_vcf(vcffile, individuals=alive) + ts.write_vcf(vcffile, individuals=alive_indivs) except Exception as e: print ("Error:") print (e) @@ -537,7 +539,7 @@ using {meth}`is_sample() `: ```{code-cell} indivlist = [] -for i in pyslim.individuals_alive_at(ts, 0): +for i in alive_indivs: ind = ts.individual(i) if ts.node(ind.nodes[0]).is_sample(): indivlist.append(i) @@ -571,9 +573,9 @@ we could do: ```{code-cell} orig_ts = tskit.load("migrants.trees") -alive = pyslim.individuals_alive_at(orig_ts, 0) +alive_indivs = pyslim.individuals_alive_at(orig_ts, 0) num_alive = [0 for _ in range(orig_ts.num_populations)] -for i in alive: +for i in alive_indivs: ind = orig_ts.individual(i) ind_population = orig_ts.node(ind.nodes[0]).population num_alive[ind_population] += 1 @@ -652,7 +654,7 @@ print(f"There are {ts.num_mutations} mutations across {ts.num_trees} distinct\n" ## Individual metadata -Each ``Mutation``, ``Population``, ``Node``, and ``Individual``, as well as the tree +Each ``Population``, ``Node``, and ``Individual``, as well as the tree sequence as a whole, carries additional information stored by SLiM in its ``metadata`` property. A fuller description of metadata in general is given in [](sec_metadata), but as a quick introduction, here is the information available @@ -687,7 +689,12 @@ produced by SLiM. This is described in more detail in the SLiM manual, but brief - ``flags`` holds additional information about the individual recorded by SLiM (currently, only whether the individual has migrated or not: see [](sec_constants_and_flags)). - +- the ``tag`` entries contain the correspondly-named "tags" in SLiM, + and for the logical tags ``tagLX``, the ``tagLX_set`` records whether or not + that tag was "set" (as opposed to remaining unset). + The funny values in ``tag`` and ``tagF`` are those special values that SLiM uses to + record that *those* entries were not set either. +- the ``per_trait`` entry is a list of information, one for each trait in the simulation. We can use this metadata in many ways, for example, to create an age distribution by sex: @@ -697,7 +704,8 @@ max_age = max([ind.metadata["age"] for ind in ts.individuals()]) age_table = np.zeros((max_age + 1, 2)) age_labels = { pyslim.INDIVIDUAL_TYPE_FEMALE: 'females', pyslim.INDIVIDUAL_TYPE_MALE: 'males' } -for i in pyslim.individuals_alive_at(ts, 0): +alive_indivs = pyslim.individuals_alive_at(ts, 0) +for i in alive_indivs: ind = ts.individual(i) age_table[ind.metadata["age"], ind.metadata["sex"]] += 1 @@ -728,8 +736,8 @@ This can be done using the numpy arrays returned by {func}`.individual_ages` and `.individuals_population` as follows: ```{code-cell} -alive = pyslim.individuals_alive_at(ts, 0) -adults = alive[pyslim.individual_ages(ts)[alive] > 2] +ages = pyslim.individual_ages(ts) +adults = alive_indivs[ages[alive_indivs] > 2] pops = [ [i for i in adults if ts.individual(i).metadata['subpopulation'] == k] for k in [1, 2] @@ -979,22 +987,40 @@ stored in the mutation metadata. To modify the mutations to be under selection, see [](sec_vignette_coalescent_diversity). ```{code-cell} -ts = msprime.sim_mutations( +ts = pyslim.add_mutation_metadata( + msprime.sim_mutations( ts, rate=1e-8, model=msprime.SLiMMutationModel(type=0), random_seed=9 + ) ) ``` -Now the mutations have SLiM metadata. -For instance, here's the first mutation: +The resulting mutations are in SLiM format. +Now, each `mutation` object in the tree sequence represents +some number of SLiM mutations, whose SLiM IDs are stored in the `derived_state`. +For instance, here's which SLiM mutation(s) the first mutation +in the tree sequence represents: +```{code-cell} +ds = ts.mutation(0).derived_state +print(f"SLiM IDs: {ds}") +``` +To see the information about these, we pull their information out +using {func}`.mutation_metadata`, which provides a dictionary +indexed by the SLiM IDs: ```{code-cell} :tags: ["remove-output"] -ts.mutation(0) +mut_metadata = pyslim.mutation_metadata(ts) +for sid in ds.split(","): + print(mut_metadata[int(sid)]) ``` ```{code-cell} :tags: ["remove-input"] -util.pp(ts.mutation(0)) +for sid in ds.split(","): + util.pp(mut_metadata[int(sid)]) ``` +**Important:** the {func}`.mutation_metadata`-returned dictionary +is indexed by **ints**, not strings, so be sure to convert your +SLiM IDs to ints before looking them up! Finally, we write this out to a file that can be loaded in to SLiM: ```{code-cell} @@ -1053,24 +1079,31 @@ Now, mutations have a ``nucleotide`` property in metadata that is not ``-1``: ```{code-cell} :tags: ["remove-output"] +mut_metadata = pyslim.mutation_metadata(ts) m = ts.mutation(0) +md = [mut_metadata[int(k)] for k in m.derived_state.split(",")] print(m) +for x in md: + print(x) ``` ```{code-cell} :tags: ["remove-input"] util.pp(m) +for x in md: + util.pp(x) ``` We can see which nucleotide is the derived state produced by each mutation - by indexing the {data}`.NUCLEOTIDES` object: +by indexing the {data}`.NUCLEOTIDES` object: ```{code-cell} for k in range(3): m = ts.mutation(k) print(f"Mutation {k}: position {ts.site(m.site).position}, time {m.time}") - for ml in m.metadata['mutation_list']: - print(f" nucleotide: {pyslim.NUCLEOTIDES[ml['nucleotide']]}") + for sid in m.derived_state.split(","): + md = mut_metadata[int(sid)] + print(f" nucleotide: {pyslim.NUCLEOTIDES[md['nucleotide']]}") ``` Here's a script minimally modified from the above to be nucleotide-based: @@ -1108,73 +1141,108 @@ print(f"Number of sites: {ts.num_sites}\n" ``` Note that there are more mutations than sites; -that's because some sites (looks like 24 of them) have multiple mutations. +that's because some sites have multiple mutations. The information about the mutation is put in the mutation's metadata. Here's the first mutation: ```{code-cell} :tags: ["remove-output"] +mut_metadata = pyslim.mutation_metadata(ts) m = ts.mutation(0) +md = [mut_metadata[int(k)] for k in m.derived_state.split(",")] print(m) +for x in md: + print(x) ``` + ```{code-cell} :tags: ["remove-input"] util.pp(m) +for x in md: + util.pp(x) ``` -Here, `m.site` tells us the ID of the *site* on the genome that the mutation occurred at, + +Since we haven't explicitly defined any traits in this simulation, +the only trait is fitness, and the `effect_size` listed under `per_trait` +for this mutation is simply its selection coefficient. +Furthermore, `m.site` tells us the ID of the *site* on the genome that the mutation occurred at, and we can pull up information about that with the `ts.site( )` method: + ```{code-cell} :tags: ["remove-output"] -ts.site(m.site) +s = ts.site(m.site) +md = [ + mut_metadata[int(k)] for m in s.mutations + for k in m.derived_state.split(",") +] +print(s) +for x in md: + print(x) ``` + ```{code-cell} :tags: ["remove-input"] -util.pp(ts.site(m.site)) +util.pp(s) +for x in md: + util.pp(x) ``` + This mutation occurred at position 54 along the genome (from `site.position`) which previously had no mutations (since `site.ancestral_state` is the empty string, `''`) -and was given SLiM mutation ID 1653896 (`m.derived_state`). -The metadata (`m.metadata`, a dict) tells us that -the mutation has selection coefficient 1.5597 and occurred in population 1 in generation 827, -which was 172 generations ago. +and was given SLiM mutation ID 1997358 (`m.derived_state`). +The metadata (`mut_metadata[1997358]`, a dict) tells us that +the mutation has selection coefficient -0.1129 and occurred in population 1 in generation 999, +which was 0 generations ago. This is not a nucleotide model, so the nucleotide entry is `-1`. -Note that `m.time` and `m.metadata['mutation_list'][0]['slim_time']` are in this case redundant: +Note that `m.time` and the `slim_time` entry in metadata are in this case redundant: they contain the same information, but the first is in tskit time (i.e., number of steps before the tree sequence was written out) and the second is using SLiM's internal "tick" counter. -Also note that the mutation's metadata is a *list* of metadata entries. +Also note that each mutation may have associated a *list* of SLiM mutations, +each with their own metadata. That's because of SLiM's mutation stacking feature. We know that some sites have more than one mutation, so to get an example let's pull out one such mutation. -In this case, -`m.metadata['mutation_list']` is a list of length one, -so the mutation was not stacked on top of previous ones. Let's pull out a mutation that was stacked on top of another one: + ```{code-cell} :tags: ["remove-output"] for m in ts.mutations(): if m.parent != tskit.NULL: break +pm = ts.mutation(m.parent) +md = [mut_metadata[int(k)] for k in m.derived_state.split(",")] +pmd = [mut_metadata[int(k)] for k in pm.derived_state.split(",")] + print(m) -print(ts.mutation(m.parent)) +for x in md: + print(x) +print(pm) +for x in pmd: + print(x) ``` + ```{code-cell} :tags: ["remove-input"] util.pp(m) +for x in md: + util.pp(x) util.pp(ts.mutation(m.parent)) +for x in pmd: + util.pp(x) ``` -This mutation (which is `ts.mutation(1020)` in the tree sequence) -was the result of SLiM adding a new mutation of type `m1` and selection coefficient -0.0032 -on top of an existing mutation, also of type `m1` and with selection coefficient 0.3086. -This happened at generation 999 (i.e., at tskit time 0.0 time units ago), -and the older mutation occurred at generation 274 (at tskit time 725 time units ago). -The older mutation has SLiM mutation ID 547531, -and the newer mutation had SLiM mutation ID 1998096, -so the resulting "derived state" is `'1998096,547531'`. +This mutation (which is `ts.mutation(330)` in the tree sequence) +was the result of SLiM adding a new mutation of type `m1` and selection coefficient -0.1547 +on top of an existing mutation, of type `m2` and with (whopping) selection coefficient 1.737. +This happened at generation 998 (i.e., at tskit time 1.0 time units ago), +and the older mutation occurred at generation 83 (at tskit time 916 time units ago). +The older mutation has SLiM mutation ID 1994163, +and the newer mutation had SLiM mutation ID 164833, +so the resulting "derived state" is `'1994163,164833'`. Now that we understand how SLiM mutations are stored in a tree sequence, let's look at the allele frequencies. @@ -1189,12 +1257,12 @@ print(afs.astype('int')) ``` (The `span_normalise=False` argument gives us counts rather than a density per unit length.) -This shows us that there are 4169 alleles that are found among the tree sequence's samples -that are not present in any of our 10 samples, 96 that are present in just one, etcetera. +This shows us that there are 3929 alleles that are found among the tree sequence's samples +that are not present in any of our 10 samples, 585 that are present in just one, etcetera. The surprisingly large number that are near 50% frequency are perhaps positively selected and on their way to fixation: we can check if that's true next. -You may have noticed that the sum of the allele frequency spectrum is 5243, -which is not obviously related to the number of mutations (6044) *or* the number of sites (6020). +You may have noticed that the sum of the allele frequency spectrum is 5029, +which is not obviously related to the number of mutations (5861) *or* the number of sites (5848). That's because each derived allele that is inherited by some but not all of the samples in the tree sequence is counted in the polarised allele frequency spectrum: Fixed mutations, or mutations that were entirely "overwritten" by subsequent mutations, @@ -1206,9 +1274,11 @@ afs_total = 0 for v in ts.variants(): if len(set(v.genotypes)) > 1: afs_total += len(set(v.genotypes) - set([0])) -print(afs_total) +print(afs_total, sum(afs)) ``` +These are equal, verifying our interpretation. + At time of writing, we don't have a built-in ``allele_frequency`` method, so we'll use the following snippet: @@ -1236,7 +1306,8 @@ mut_type = np.zeros(ts.num_sites) for j, s in enumerate(ts.sites()): mt = [] for m in s.mutations: - for md in m.metadata["mutation_list"]: + for sid in m.derived_state.split(","): + md = mut_metadata[int(sid)] mt.append(md["mutation_type"]) if len(set(mt)) > 1: mut_type[j] = 3 @@ -1261,32 +1332,34 @@ print(mut_afs) The first column gives the AFS among these 10 samples for the deleterious alleles, the second for the beneficial mutations; -the third column for the seven sites that had both types of mutation. +the third column for the few sites that had both types of mutation. Interestingly, there are similar numbers of both types of mutation at intermediate frequency: perhaps because beneficial mutations are sweeping linked deleterious alleles along with them. -Many fewer benefical alleles are at low frequency: -3,666 deleterious alleles are not found in our sample of 10 genomes, -while only 486 beneficial alleles are. +Many fewer benefical alleles are at low frequency, however. Finally, let's pull out information on the allele with the largest selection coefficient. ```{code-cell} :tags: ["remove-output"] sel_coeffs = np.array([ - sum(md["selection_coeff"] for md in m.metadata["mutation_list"]) + sum(mut_metadata[int(k)]["per_trait"][0]["effect_size"] + for k in m.derived_state.split(",")) for m in ts.mutations() ]) which_max = np.argmax(sel_coeffs) m = ts.mutation(which_max) +print(f"Max selection coefficient: {sel_coeffs[which_max]} for site {m.site}") ts.site(m.site) ``` + ```{code-cell} :tags: ["remove-input"] +print(f"Max selection coefficient: {sel_coeffs[which_max]} for site {m.site}") util.pp(ts.site(m.site)) ``` -This allele had a whopping selection coefficient of 4.94 -and appeared about halfway through the simulation. +This allele had a whopping selection coefficient of 5.69 +and appeared fairly late in the simulation. Let's find its frequency in the full population: ```{code-cell} @@ -1296,7 +1369,7 @@ print(f"The allele is found in {full_freqs[m.site][0]} copies\n" ``` The allele is above 50% in the population, so it is probably on its way to fixation. -Using its SLiM ID (which is shown in its derived state, ``1616148``), +Using its SLiM ID (which is shown in its derived state, ``305447``), we could reload the tree sequence into SLiM, restart the simulation, and use its ID to track its subsequent progression. @@ -1331,4 +1404,4 @@ Also known as "gotchas". 4. SLiM requires that the two nodes corresponding to the haplosomes of each individual are adjacent in the node table, and are sorted by haplosome ID. SLiM always writes out tree sequences like this, but it is possible to make - tree sequences in python that are leval otherwise but don't satisfy this requirement. + tree sequences in python that are legal otherwise but don't satisfy this requirement. diff --git a/docs/vignette_coalescent_diversity.md b/docs/vignette_coalescent_diversity.md index 589d1bd..43dc4aa 100644 --- a/docs/vignette_coalescent_diversity.md +++ b/docs/vignette_coalescent_diversity.md @@ -166,21 +166,26 @@ mut_map = msprime.RateMap( position=breaks, rate=[0.03e-8, 0.003e-8, 0.03e-8]) mut_model = msprime.SLiMMutationModel(type=2) -ots = msprime.sim_mutations( +ots = pyslim.add_mutation_metadata( + msprime.sim_mutations( ots, rate=mut_map, model=mut_model, keep=True, - random_seed=12) + random_seed=12), + mutation_type=2, +) print(f"The tree sequence now has {ots.num_mutations} mutations, at " f"{ots.num_sites} distinct sites.") ``` -Note the ``type=2`` argument to {class}`msprime.SLiMMutationModel`: +Note the ``type=2`` argument to {func}`.add_mutation_metadata`: this means the mutations will be of type "m2" in SLiM (and, so you must initialize that mutation type in the recipe that loads this tree sequence in). Now, we'll assign selection coefficients. +This is easier than in versions of SLiM before 6.0, +because we simply want to assign each mutation an independent selection coefficient Recall that to accomodate mutation stacking in SLiM, a mutation metadata entry is in fact a *list* of metadata entries, one for each of the SLiM mutations that are stacked at this position. @@ -193,25 +198,12 @@ SLiM mutation ID ``k``. ```{code-cell} rng = np.random.default_rng(seed=1234) -tables = ots.tables -tables.mutations.clear() -mut_map = {} -for m in ots.mutations(): - md_list = m.metadata["mutation_list"] - slim_ids = m.derived_state.split(",") - assert len(slim_ids) == len(md_list) - for sid, md in zip(slim_ids, md_list): - if sid not in mut_map: - mut_map[sid] = rng.exponential(scale=0.04) - md["selection_coeff"] = mut_map[sid] - _ = tables.mutations.append( - m.replace(metadata={"mutation_list": md_list}) - ) - -# check we didn't mess anything up -assert tables.mutations.num_rows == ots.num_mutations -print(f"The selection coefficients range from {min(mut_map.values()):0.2e}") -print(f"to {max(mut_map.values()):0.2e}.") +ts_metadata = ots.metadata +for md in ts_metadata["SLiM_mutation_list"]: + md["per_trait"][0]["effect_size"] = rng.exponential(scale=0.04) + +tables = ots.dump_tables() +tables.metadata = ts_metadata ``` @@ -222,15 +214,15 @@ We can see this with ``tables.metadata``: ```{code-cell} :tags: ['remove-output'] -tables.metadata +tables.metadata["SLiM"] ``` ```{code-cell} :tags: ['remove-input'] -util.pp(tables.metadata) +util.pp(tables.metadata["SLiM"]) ``` -We should edit this to match our planned slimulation -- particularly the ``model_type`` (WF or nonWF) and the ``tick``. +We should edit this to match our planned slimulation - +particularly the ``model_type`` (WF or nonWF) and the ``tick``. The ``tick`` tells SLiM what value to set the tick counter to once this tree sequence is loaded. In principle, it can be set to anything, independently of the times in the tree sequence, @@ -246,7 +238,6 @@ edit the metadata, let's make sure, and then we'll write the tree sequence to a file. ```{code-cell} -ts_metadata = tables.metadata ts_metadata["SLiM"]["model_type"] = "WF" tables.metadata = ts_metadata ots = tables.tree_sequence() @@ -290,7 +281,7 @@ This runs quickly, since it's only 100 generations. First, let's look at what mutations are present. ```{code-cell} ts = tskit.load("vignette_annotated.trees") -num_stacked = np.array([len(m.metadata["mutation_list"]) for m in ts.mutations()]) +num_stacked = np.array([len(m.derived_state.split(",")) for m in ts.mutations()]) init_time = ts.metadata['SLiM']['tick'] old_mut = np.array([m.time > init_time - 1 - 1e-12 for m in ts.mutations()]) assert sum(old_mut) == ots.num_mutations @@ -324,7 +315,9 @@ nodes_by_time = [ts.samples(time=t) for t in times] num_nodes = np.array([len(x) for x in nodes_by_time]) p = ts.sample_count_stat(nodes_by_time, lambda x: x/num_nodes, 2, windows='sites', strict=False, span_normalise=False, polarised=True) -s = np.array([sum([sum([md["selection_coeff"] for md in m.metadata["mutation_list"]]) +mut_metadata = pyslim.mutation_metadata(ts) +s = np.array([sum([sum([mut_metadata[int(k)]["per_trait"][0]["effect_size"] + for k in m.derived_state.split(",")]) for m in site.mutations]) for site in ts.sites()]) ``` @@ -407,12 +400,15 @@ next_id = pyslim.next_slim_mutation_id(ts) neutral_mut_model = msprime.SLiMMutationModel( type=1, next_id=next_id) -mts = msprime.sim_mutations( +mts = pyslim.add_mutation_metadata( + msprime.sim_mutations( ts, rate=neutral_mut_map, model=neutral_mut_model, keep=True, - random_seed=35) + random_seed=35), + mutation_type=1, +) print(f"The tree sequence now has {mts.num_mutations} mutations,") print(f"at {mts.num_sites} distinct sites.") ``` @@ -431,8 +427,9 @@ we'll pull out a tree that had a lot of mutations on it, and print a picture of it, with mutations labeled by their type: ```{code-cell} +mut_metadata = pyslim.mutation_metadata(mts) for t in mts.trees(): - mt = [max([u['mutation_type'] for u in m.metadata['mutation_list']]) for m in t.mutations()] + mt = [max([mut_metadata[int(k)]['mutation_type'] for k in m.derived_state.split(",")]) for m in t.mutations()] if t.num_mutations > 12: break @@ -482,28 +479,36 @@ If you wanted some other arrangement (e.g., to have m1 stack on top of m2), you could go through and modify derived states and metadata appropriately. Let's check there are any sites with stacked mutations of different types in the simulation. -There is indeed one such site: +There are indeed: ```{code-cell} :tags: ['remove-output'] for site in mts.sites(): if len(site.mutations) > 1: - types = [set([md["mutation_type"] for md in mut.metadata["mutation_list"]]) + types = [set([mut_metadata[int(k)]["mutation_type"] for k in mut.derived_state.split(",")]) for mut in site.mutations] if max(map(len, types)) > 1: print(site) + for mut in site.mutations: + print(mut) + for k in mut.derived_state.split(","): + print(mut_metadata[int(k)]) ``` ```{code-cell} :tags: ['remove-input'] for site in mts.sites(): if len(site.mutations) > 1: - types = [set([md["mutation_type"] for md in mut.metadata["mutation_list"]]) + types = [set([mut_metadata[int(k)]["mutation_type"] for k in mut.derived_state.split(",")]) for mut in site.mutations] if max(map(len, types)) > 1: util.pp(site) + for mut in site.mutations: + util.pp(mut) + for k in mut.derived_state.split(","): + util.pp(mut_metadata[int(k)]) ``` -Here, a neutral mutation has been put down on top of a selected mutation, +In each of these, a neutral mutation has been put down on top of a selected mutation, but stacked, so that any samples inheriting either of these mutations carries the selected mutation. For more discussion of how this works, see {class}`msprime.SLiMMutationModel`. diff --git a/docs/vignette_continuing.md b/docs/vignette_continuing.md index 3aff661..80d0536 100644 --- a/docs/vignette_continuing.md +++ b/docs/vignette_continuing.md @@ -81,17 +81,18 @@ so first we check that all the existing mutations are of a different type. rts = pyslim.recapitate(ts, ancestral_Ne=1000, recombination_rate=1e-8, random_seed=6) # check type m0 is not used: -mut_types = set([md['mutation_type'] - for mut in ts.mutations() - for md in mut.metadata['mutation_list']]) +mut_metadata = pyslim.mutation_metadata(rts) +mut_types = set([md['mutation_type'] for md in mut_metadata.values()]) print(f"Keeping {rts.num_mutations} existing mutations of type(s) {mut_types}.") assert 0 not in mut_types # add type m0 mutations next_id = pyslim.next_slim_mutation_id(rts) -rts = msprime.sim_mutations( +rts = pyslim.add_mutation_metadata( + msprime.sim_mutations( rts, rate=1e-8, random_seed=7, keep=True, model=msprime.SLiMMutationModel(type=0, next_id=next_id) + ) ) p = rts.sample_count_stat( @@ -143,7 +144,7 @@ new_ts = msprime.sim_mutations( new_ts, rate=1e-8, random_seed=10, keep=True, model=msprime.SLiMMutationModel(type=0) ) -new_tables = new_ts.tables +new_tables = new_ts.dump_tables() # check that the spurious samples are 20000 and 20001 for n in (20000, 20001): assert n in new_ts.samples() @@ -158,13 +159,11 @@ randomly assign each to a node at the end of the SLiM simulation, and merge them. ```{code-cell} - new_nodes = np.where(new_tables.nodes.time == new_time)[0] print(f"There are {len(new_nodes)} nodes from the start of the new simulation.") -# There are 4425 nodes from the start of the new simulation. slim_nodes = rts.samples(time=0) -assert(len(slim_nodes) == 20000) +assert len(slim_nodes) == 20000 # randomly give new_nodes IDs in rts node_map = np.repeat(tskit.NULL, new_tables.nodes.num_rows) @@ -173,15 +172,15 @@ node_map[new_nodes] = np.random.choice(slim_nodes, len(new_nodes), replace=False # shift times: in nodes and mutations # since tree sequences are not mutable, we do this in the tables directly # also, unmark the nodes at the end of the SLiM simulation as samples -tables = rts.tables +tables = rts.dump_tables() tables.nodes.flags = tables.nodes.flags & ~np.uint32(tskit.NODE_IS_SAMPLE) tables.nodes.time = tables.nodes.time + new_time tables.mutations.time = tables.mutations.time + new_time # merge the two sets of tables tables.union(new_tables, node_map, - add_populations=False, - check_shared_equality=False) + add_populations=False, + check_shared_equality=False) # get back the tree sequence full_ts = tables.tree_sequence() @@ -243,3 +242,12 @@ to be identical in the two tree sequences, so ``union`` by default throws an err We don't expect that in this case, because, for instance, there could be a mutation above one of the terminal nodes in the SLiM tree sequence; this would clearly not be present in the new tree sequence. + +*Note:* sharp-eyed readers will note that the call to ``sim_mutations`` was not wrapped in +{func}`.add_mutation_metadata`. If we wanted to read this tree sequence into SLiM again +we'd need to add mutation metadata for these last mutations. +The easiest place to do this would be a call to {func}`.add_mutation_metadata_tables` +just after the ``union`` +(thus avoiding an extra conversion to tree sequence); +this will add metadata for only those mutations not already recorded. +We've left that step out of the code here for simplicity. diff --git a/docs/vignette_parallel_phylo.md b/docs/vignette_parallel_phylo.md index efa70f2..d75eb92 100644 --- a/docs/vignette_parallel_phylo.md +++ b/docs/vignette_parallel_phylo.md @@ -110,10 +110,7 @@ f.close() Here's the result. Again, don't worry about the details, but you can see that the file encodes the phylogeny through a bunch of ``child : parent`` "rules": -```{code-cell} -:tags: ["hide-input"] -%%bash -cat sims.make +```{literalinclude} sims.make ``` With the makefile in hand, @@ -304,9 +301,7 @@ samples per population. ```{code-cell} rng = np.random.default_rng(seed=123) ind_alive = pyslim.individuals_alive_at(tsu, 0) -# TODO: this will work in the next tskit -# ind_pops = tsu.individuals_population[ind_alive] -ind_pops = np.array([tsu.node(tsu.individual(i).nodes[0]).population for i in ind_alive]) +ind_pops = tsu.individuals_population[ind_alive] subsample_indivs = [ rng.choice(ind_alive[ind_pops == pop_ids[name]], 2) for name in pops diff --git a/docs/vignette_space.md b/docs/vignette_space.md index 3c1d3ea..fc87941 100644 --- a/docs/vignette_space.md +++ b/docs/vignette_space.md @@ -98,13 +98,17 @@ and so individuals may live for more than one time step (even up to age 10, it s Let's check that all these individuals are alive at either (a) today or (b) 1000 time steps ago. ```{code-cell} +slim_ts_md = slim_ts.metadata for t in [0, 1000]: - alive = pyslim.individuals_alive_at(slim_ts, t) + alive = pyslim.individuals_alive_at(slim_ts, t, ts_metadata=slim_ts_md) print(f"There were {len(alive)} individuals alive {t} time steps in the past.") ``` -And, 1242 + 1255 is 2497, the total number of individuals. -So, this all checks out. +We can add these numbers to get the total number of individuals, so this all checks out. +On a separate note: here we used the optional ``ts_metadata`` argument +to {func}`.individuals_alive_at`: see [](sec_metadata_using_top_level) for discussion: +this is not important here because we only run the method twice, but this would be essential +if we had a longer list of times. ## Recapitation and mutation @@ -142,11 +146,13 @@ we would need to pass ``keep_input_roots=True`` to allow recapitation. ```{code-cell} recap_ts = pyslim.recapitate(slim_ts, recombination_rate=1e-8, ancestral_Ne=1000) -ts = msprime.sim_mutations( +ts = pyslim.add_mutation_metadata( + msprime.sim_mutations( recap_ts, rate=1e-8, model=msprime.SLiMMutationModel(type=0), keep=True, + ) ) ts.dump("spatial_sim.recap.trees") diff --git a/pyproject.toml b/pyproject.toml index e86faea..0a43898 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ keywords = ["tree sequences", "tskit"] requires-python = ">=3.11" dependencies = [ "msprime>=1.0.1", - "tskit", + "tskit>=1.0.3", "numpy", ] @@ -56,6 +56,7 @@ test = [ "tskit", "msprime", "pandas", + "frozendict", ] docs = [ diff --git a/pyslim/_version.py b/pyslim/_version.py index 3b480d7..8eea699 100644 --- a/pyslim/_version.py +++ b/pyslim/_version.py @@ -5,6 +5,6 @@ except Exception: pyslim_version = "unknown" -slim_file_version = "0.9" +slim_file_version = "1.0" # other file versions that require no modification -compatible_slim_file_versions = ["0.9"] +compatible_slim_file_versions = ["1.0"] diff --git a/pyslim/methods.py b/pyslim/methods.py index 19c1220..cf84634 100644 --- a/pyslim/methods.py +++ b/pyslim/methods.py @@ -12,6 +12,7 @@ set_metadata_schemas, set_tree_sequence_metadata, ) +from .slim_tree_sequence import mutation_metadata from .util import unique_labels_by_group @@ -43,7 +44,7 @@ def _mark_not_samples(tables, nodes): ) -def _chromosome_index(ts): +def _chromosome_index(ts_metadata): """ For a tree sequence produced by a multichromosome simulation, returns the index of the chromosome whose information is stored in this tree sequence @@ -52,18 +53,19 @@ def _chromosome_index(ts): ``ts.metadata['SLiM']['this_chromosome']``, and provides the index of this chromosome into `ts.metadata['SLiM']['chromosomes']``, if present. - :param tskit.TreeSequence ts: The tree sequence or table collection. + :param dict ts_metadata: The top-level metadata from a tree sequence + or table collection. """ if not ( - isinstance(ts.metadata, dict) - and "SLiM" in ts.metadata - and "this_chromosome" in ts.metadata["SLiM"] + isinstance(ts_metadata, dict) + and "SLiM" in ts_metadata + and "this_chromosome" in ts_metadata["SLiM"] ): raise ValueError( "The tree sequence does not have the necessary " "information in top-level metadata." ) - k = ts.metadata["SLiM"]["this_chromosome"]["index"] + k = ts_metadata["SLiM"]["this_chromosome"]["index"] return k @@ -88,15 +90,17 @@ def _is_chrom_vacant(k, b): return (b >> i & 1) > 0 -def has_vacant_samples(ts): +def has_vacant_samples(ts, _ts_metadata=None): """ Returns whether the tree sequence has vacant sample nodes. See :meth:`remove_vacant`. :param tskit.TreeSequence ts: The tree sequence. """ + if _ts_metadata is None: + _ts_metadata = ts.metadata out = False - k = _chromosome_index(ts) + k = _chromosome_index(_ts_metadata) for n in ts.samples(): md = ts.node(n).metadata if md is not None: @@ -106,8 +110,45 @@ def has_vacant_samples(ts): return out +def nodes_vacant(ts): + """ + Evaluates which nodes in the tree sequence are vacant: returns a boolean + vector whose k-th element is True if the k-th node is labelled as *vacant* + in the node's metadata recorded by SLiM. A vacant node represents a blank + placeholder in SLiM: either a "null haplosome" (used as placeholders for + sex chromosomes and other chromosome types not of consistent ploidy in all + individuals) or simply an unused node for haploid chromosome types. See + :meth:`remove_vacant`. + + + :param tskit.TreeSequence ts: The tree sequence. + :return boolean ndarray: + """ + # not using chrom_index here because we expect people to call this on lots of nodes + k = ts.metadata["SLiM"]["this_chromosome"]["index"] + out = np.array( + [ + node.metadata is not None and _is_chrom_vacant(k, node.metadata["is_vacant"]) + for node in ts.nodes() + ], + dtype="bool", + ) + return out + + def node_is_vacant(ts, node): """ + **DEPRECATED:** use :func:`.nodes_vacant` instead. This function requires + top-level metadata access, which can be costly, so it is much better to do, + for instance: + + .. code-block:: python + + vacant = nodes_vacant(ts) + for node in ts.nodes(): + # instead of node_is_vacant(ts, node), use: + vacant[node.id] + Returns True if the node is labelled as *vacant* in the node's metadata recorded by SLiM. A vacant node represents a blank placeholder in SLiM: either a "null haplosome" (used as placeholders for sex chromosomes and other @@ -117,12 +158,19 @@ def node_is_vacant(ts, node): :param tskit.TreeSequence ts: The tree sequence. :param tskit.Node node: The node object. """ + warnings.warn( + "The node_is_vacant method is deprecated: changes in SLiM v6 " + " means that repeated use of this method will be unreasonably slow, " + "so it will be removed in a future version of pyslim: " + "obtain this information from pyslim.vacant_nodes( ) instead.", + FutureWarning, + ) # not using chrom_index here because we expect people to call this on lots of nodes k = ts.metadata["SLiM"]["this_chromosome"]["index"] return node.metadata is not None and _is_chrom_vacant(k, node.metadata["is_vacant"]) -def _record_vacant_tables(tables): +def _record_vacant_tables(tables, ts_metadata): """ Sets the NODE_IS_VACANT_SAMPLE flag for all vacant, sample nodes. See :meth:`remove_vacant`. @@ -135,7 +183,7 @@ def _record_vacant_tables(tables): "flags are being overwritten; this may mean you've already run " "remove_vacant and so don't need to run it again." ) - k = _chromosome_index(tables) + k = _chromosome_index(ts_metadata) dn = tables.nodes.asdict() dn["flags"] &= ~NODE_IS_VACANT_SAMPLE @@ -156,7 +204,7 @@ def _remove_vacant_sample_flags(tables): tables.nodes.set_columns(**dn) -def remove_vacant(ts): +def remove_vacant(ts, _ts_metadata=None): """ Remove sample flags from all vacant nodes. @@ -177,23 +225,27 @@ def remove_vacant(ts): :param tskit.TreeSequence ts: The tree sequence. """ + if _ts_metadata is None: + _ts_metadata = ts.metadata tables = ts.dump_tables() - remove_vacant_tables(tables) + remove_vacant_tables(tables, _ts_metadata) return tables.tree_sequence() -def remove_vacant_tables(tables): +def remove_vacant_tables(tables, _ts_metadata=None): """ Does the work of :meth:`remove_vacant`, modifying ``tables`` in place. :param tskit.TableCollection tables: The tables underlying a tree sequence. """ - _record_vacant_tables(tables) + if _ts_metadata is None: + _ts_metadata = tables.metadata + _record_vacant_tables(tables, _ts_metadata) is_vacant = np.where(tables.nodes.flags & NODE_IS_VACANT_SAMPLE > 0)[0] _mark_not_samples(tables, is_vacant) -def restore_vacant(ts): +def restore_vacant(ts, _ts_metadata=None): """ The inverse of :meth:`remove_vacant`. @@ -204,19 +256,23 @@ def restore_vacant(ts): :param tskit.TreeSequence ts: The tree sequence. """ + if _ts_metadata is None: + _ts_metadata = ts.metadata tables = ts.dump_tables() - restore_vacant_tables(tables) + restore_vacant_tables(tables, _ts_metadata) return tables.tree_sequence() -def restore_vacant_tables(tables): +def restore_vacant_tables(tables, _ts_metadata=None): """ Does the work of :meth:`restore_vacant`, modifying ``tables`` in place. :param tskit.TableCollection tables: The tables underlying a tree sequence. """ + if _ts_metadata is None: + _ts_metadata = tables.metadata is_vacant = np.where(tables.nodes.flags & NODE_IS_VACANT_SAMPLE > 0)[0] - k = _chromosome_index(tables) + k = _chromosome_index(_ts_metadata) for j in is_vacant: n = tables.nodes[j] if n.metadata is None: @@ -273,14 +329,15 @@ def recapitate(ts, ancestral_Ne=None, *, keep_vacant=False, **kwargs): vacant sample nodes. Default: False. :param dict kwargs: Any other arguments to :func:`msprime.sim_ancestry`. """ - is_current_version(ts, _warn=True) + ts_metadata = ts.metadata + is_current_version(ts_metadata, _warn=True) # we need to ask msprime to *not* simulate from any 'vacant' haplosomes; # which we do by marking these as not samples; note that `initial_state` # can take a TableCollection, not just a TreeSequence - has_vacant = has_vacant_samples(ts) + has_vacant = has_vacant_samples(ts, ts_metadata) if has_vacant: - ts = remove_vacant(ts) + ts = remove_vacant(ts, ts_metadata) if ancestral_Ne is not None: if "demography" in kwargs: @@ -338,11 +395,90 @@ def recapitate(ts, ancestral_Ne=None, *, keep_vacant=False, **kwargs): recap = msprime.sim_ancestry(initial_state=ts, **kwargs) if has_vacant and keep_vacant: - recap = restore_vacant(recap) + recap = restore_vacant(recap, ts_metadata) return recap +def add_mutation_metadata(ts, mutation_type=0, remove_unused=False): + """ + Returns a new tree sequence with default information added to the top-level metadata + for each mutation in the tree sequence for which that information is not already present. + To do this, this method looks for all SLiM IDs that are found in the derived + state of some mutation but are not represented in the top-level metadata + (see :func:`.mutation_metadata`). This function then adds entries to that top-level + metadata with default values (see :func:`.default_slim_metadata`), + except that (a) the ``mutation_type`` can be specified; + and (b) the ``slim_time`` is set using the ``tick`` value in top-level metadata + and the ``time`` of the oldest tskit mutation in which the SLiM mutation occurs. + + :param tskit.TreeSequence ts: The tree sequence to transform. + :param int mutation_type: The numeric ID of the mutation type in SLiM. + :param bool remove_unused: Whether to also remove from metadata information about any + mutations not seen in the derived states of the tree sequence. + """ + tables = ts.dump_tables() + add_mutation_metadata_tables( + tables, mutation_type=mutation_type, remove_unused=remove_unused + ) + return tables.tree_sequence() + + +def add_mutation_metadata_tables(tables, mutation_type=0, remove_unused=False): + """ + Modifies the tables in place to add metadata for any mutations for which it is missing; + see :func:`.add_mutation_metadata`. + + :param tskit.TableCollection tables: The table collection to be modified. + :param int mutation_type: The numeric ID of the mutation type in SLiM. + :param bool remove_unused: Whether to also remove from metadata information about any + mutations not seen in the derived states of the tree sequence. + """ + ts_metadata = tables.metadata + if ( + not isinstance(ts_metadata, dict) + or "SLiM" not in ts_metadata + or "SLiM_mutation_list" not in ts_metadata + ): + raise ValueError( + "Top-level metadata schema is not correct: " + "do you need to run pyslim.annotate()?" + ) + existing_muts = {x["mutation_id"] for x in ts_metadata["SLiM_mutation_list"]} + mut_ids = [ + (int(j), mut.time) + for mut in tables.mutations + for j in mut.derived_state.split(",") + ] + mut_ids.sort() + mut_ids = np.array(mut_ids, dtype="int") # floors times + # remove duplicate IDs, keeping the last (most recent) + keep = np.full(len(mut_ids), True, dtype="bool") + keep[np.where(np.diff(mut_ids[:, 0]) == 0)[0]] = False + mut_ids = mut_ids[keep, :] + mut_ids[:, 1] = slim_time( + tables, mut_ids[:, 1], stage="late", ts_metadata=ts_metadata + ) + # this assumes mutations were added in late(), which is what SLiM does + ts_metadata["SLiM_mutation_list"].extend( + [ + default_slim_metadata( + "mutation_list_entry", + mutation_id=int(j), + mutation_type=mutation_type, + slim_time=int(t), + ) + for j, t in mut_ids + if j not in existing_muts + ] + ) + if remove_unused and len(mut_ids) < len(ts_metadata["SLiM_mutation_list"]): + ts_metadata["SLiM_mutation_list"] = [ + x for x in ts_metadata["SLiM_mutation_list"] if x["mutation_id"] in mut_ids + ] + tables.metadata = ts_metadata + + def convert_alleles(ts): """ Returns a modified tree sequence in which alleles have been replaced by @@ -379,25 +515,28 @@ def convert_alleles(ts): # so we must guess which is the most recent, by choosing the one that # has the largest SLiM time, doesn't appear in the parent list, or has # the lagest SLiM ID. - nuc_inds = tables.mutations.metadata_vector( - ["mutation_list", 0, "nucleotide"], dtype="int" - ) - num_stacked = np.array([len(m.metadata["mutation_list"]) for m in ts.mutations()]) - for k in np.where(num_stacked > 1)[0]: + mut_metadata = mutation_metadata(ts) + mut_ids = np.array([x["mutation_id"] for x in mut_metadata.values()], dtype="int") + alleles = np.array([x["nucleotide"] for x in mut_metadata.values()], dtype="int") + # mut_inds will map from tskit-mutations to slim-mutations + mut_inds = ts.mutations_derived_state.copy() + num_stacked = np.strings.count(mut_inds, ",") + mut_inds[num_stacked > 0] = "-1" + mut_inds = mut_inds.astype("int", copy=False) + for k in np.where(num_stacked > 0)[0]: mut = ts.mutation(k) if mut.parent == tskit.NULL: pids = [] else: pids = ts.mutation(mut.parent).derived_state.split(",") x = [ - (md["slim_time"], i not in pids, int(i), j) - for j, (i, md) in enumerate( - zip(mut.derived_state.split(","), mut.metadata["mutation_list"]) - ) + (mut_metadata[int(i)]["slim_time"], i not in pids, int(i), j) + for j, i in enumerate(mut.derived_state.split(",")) ] x.sort() - j = x[-1][3] - nuc_inds[k] = mut.metadata["mutation_list"][j]["nucleotide"] + mut_inds[k] = x[-1][2] + assert np.all(mut_inds >= 0), "This should not occur: please file a bug report." + nuc_inds = alleles[np.searchsorted(mut_ids, mut_inds)] if np.any(nuc_inds == -1): raise ValueError("All mutations must be nucleotide mutations.") da = np.array(NUCLEOTIDES)[nuc_inds] @@ -405,7 +544,6 @@ def convert_alleles(ts): k = tables.sites.position.astype("int") aa = np.frombuffer(ts.reference_sequence.data.encode("utf-8"), dtype="S1")[k] tables.sites.packset_ancestral_state(aa.tobytes().decode("utf-8")) - return tables.tree_sequence() @@ -461,11 +599,11 @@ def generate_nucleotides(ts, reference_sequence=None, keep=True, seed=None): raise ValueError( "Reference sequence must be a string of A, C, G, and T only." ) - + ts_metadata = ts.metadata + mut_info = mutation_metadata(ts, _ts_metadata=ts_metadata) tables = ts.dump_tables() if reference_sequence is not None: tables.reference_sequence.data = reference_sequence - tables.mutations.clear() sets = [[k for k in range(4) if k != i] for i in range(4)] states = np.full((ts.num_mutations,), -1) k = tables.sites.position.astype("int") @@ -484,9 +622,9 @@ def generate_nucleotides(ts, reference_sequence=None, keep=True, seed=None): pa = states[mut.parent] pds = ts.mutation(mut.parent).derived_state.split(",") this_da = pa - ml = mut.metadata max_time = -np.inf - for i, md in zip(mut.derived_state.split(","), ml["mutation_list"]): + for i in mut.derived_state.split(","): + md = mut_info[int(i)] da = md["nucleotide"] if da == -1 or not keep: if i in muts: @@ -502,22 +640,26 @@ def generate_nucleotides(ts, reference_sequence=None, keep=True, seed=None): this_da = da max_time = md["slim_time"] states[mut.id] = this_da - tables.mutations.append(mut.replace(metadata=ml)) - md = tables.metadata - md["SLiM"]["nucleotide_based"] = True - tables.metadata = md + ts_metadata["SLiM"]["nucleotide_based"] = True + ts_metadata["SLiM_mutation_list"] = list(mut_info.values()) + tables.metadata = ts_metadata return tables.tree_sequence() -def individual_ages(ts): +def individual_ages(ts, ts_metadata=None): """ Returns the ages of all individuals in the tree sequence, extracted from metadata. The result is a array of length equal to the number of individuals, with k-th entry equal to ``ts.individual(k).metadata["age"]``. + :param tskit.TreeSequence ts: The tree sequence. + :param dict ts_metadata: Optionally, the top-level metadata for ``ts``. If + this does not match the actual top-level metadata, incorrect values may result. :return: An array of ages of individuals. """ - if ts.metadata["SLiM"]["model_type"] != "WF": + if ts_metadata is None: + ts_metadata = ts.metadata + if ts_metadata["SLiM"]["model_type"] != "WF": ages = ts.tables.individuals.metadata_vector("age") else: ages = np.zeros(ts.num_individuals, dtype="int") @@ -525,7 +667,13 @@ def individual_ages(ts): def individuals_alive_at( - ts, time, stage="late", remembered_stage=None, population=None, samples_only=False + ts, + time, + stage="late", + remembered_stage=None, + population=None, + samples_only=False, + ts_metadata=None, ): """ Returns an array giving the IDs of all individuals that are known to be @@ -576,26 +724,29 @@ def individuals_alive_at( population(s) with these population ID(s). :param bool samples_only: Whether to return only individuals who have at least one node marked as samples. + :param dict ts_metadata: Optionally, the top-level metadata for ``ts``. If + this does not match the actual top-level metadata, incorrect values may result. """ - is_current_version(ts, _warn=True) if stage not in ("late", "early", "first"): raise ValueError( f"Unknown stage '{stage}': should be either 'first', 'early' or 'late'." ) - + if ts_metadata is None: + ts_metadata = ts.metadata + is_current_version(ts_metadata, _warn=True) if remembered_stage is None: - remembered_stage = ts.metadata["SLiM"]["stage"] + remembered_stage = ts_metadata["SLiM"]["stage"] if remembered_stage not in ("late", "early", "first"): raise ValueError( f"Unknown remembered_stage '{remembered_stage}': " "should be either 'first', 'early' or 'late'." ) - if remembered_stage != ts.metadata["SLiM"]["stage"]: + if remembered_stage != ts_metadata["SLiM"]["stage"]: warnings.warn( f"Provided remembered_stage '{remembered_stage}' does not" " match the stage at which the tree sequence was saved" - f" ('{ts.metadata['SLiM']['stage']}'). This is not necessarily" + f" ('{ts_metadata['SLiM']['stage']}'). This is not necessarily" " an error, but mismatched stages will lead to inconsistencies:" " make sure you know what you're doing." ) @@ -615,12 +766,12 @@ def individuals_alive_at( # let x = 1 if the stage is 'first' or (is 'early' and WF) # and y = 1 if remembered stage is 'late' or (is 'early' and nonWF); # then t = time + x + y - 1 . - is_wf = ts.metadata["SLiM"]["model_type"] == "WF" + is_wf = ts_metadata["SLiM"]["model_type"] == "WF" x = stage == "first" or (stage == "early" and is_wf) y = remembered_stage == "late" or (remembered_stage == "early" and not is_wf) t = time + x + y - 1 birth_times = ts.individuals_time - ages = individual_ages(ts) + ages = individual_ages(ts, ts_metadata) if is_wf: alive_bool = birth_times == t else: @@ -643,7 +794,9 @@ def individuals_alive_at( return np.where(alive_bool)[0] -def individual_ages_at(ts, time, stage="late", remembered_stage="late"): +def individual_ages_at( + ts, time, stage="late", remembered_stage="late", ts_metadata=None +): """ Returns the `ages` of each individual at the corresponding time ago, which will be ``nan`` if the individual is either not born yet or dead. @@ -669,20 +822,24 @@ def individual_ages_at(ts, time, stage="late", remembered_stage="late"): is alive (either "early" or "late"; defaults to "late"). :param str remembered_stage: The stage in the SLiM life cycle during which individuals were Remembered. + :param dict ts_metadata: Optionally, the top-level metadata for ``ts``. If + this does not match the actual top-level metadata, incorrect values may result. """ + if ts_metadata is None: + ts_metadata = ts.metadata ages = np.repeat(np.nan, ts.num_individuals) alive = individuals_alive_at( - ts, time, stage=stage, remembered_stage=remembered_stage + ts, time, stage=stage, remembered_stage=remembered_stage, ts_metadata=ts_metadata ) # to convert individuals_time to number of ticks ago we subtract (y - 1), so - is_wf = ts.metadata["SLiM"]["model_type"] == "WF" + is_wf = ts_metadata["SLiM"]["model_type"] == "WF" y = remembered_stage == "late" or (remembered_stage == "early" and not is_wf) t = time + y - 1 ages[alive] = ts.individuals_time[alive] - t return ages -def slim_time(ts, time, stage="late"): +def slim_time(ts, time, stage="late", ts_metadata=None): """ Converts the given "tskit times" (i.e., in units of time before the end of the simulation) to SLiM times (those recorded by SLiM, usually in units @@ -707,17 +864,26 @@ def slim_time(ts, time, stage="late"): this may not return what you expect. See :ref:`sec_metadata_converting_times` for more discussion. + This method accesses top-level metadata, which may be a costly operation, + so if this method will be called many times, it is recommended to + extract this to a variable (e.g., ``ts_metadata = ts.metadata``) and pass it + to this method (as ``ts_metadata``). However, beware: if ``ts_metadata`` + is not in sync with the actual top-level metadata, incorrect values may result. + :param tskit.TreeSequence ts: A SLiM-compatible TreeSequence. :param numpy.ndarray time: An array of times to be converted. :param str stage: The stage of the SLiM life cycle that the SLiM time should be computed for. + :param dict ts_metadata: Optionally, the top-level metadata for ``ts``. """ - is_current_version(ts, _warn=True) - is_wf = ts.metadata["SLiM"]["model_type"] == "WF" - remembered_stage = ts.metadata["SLiM"]["stage"] + if ts_metadata is None: + ts_metadata = ts.metadata + is_current_version(ts_metadata, _warn=True) + is_wf = ts_metadata["SLiM"]["model_type"] == "WF" + remembered_stage = ts_metadata["SLiM"]["stage"] x = stage == "first" or (stage == "early" and is_wf) y = remembered_stage == "late" or (remembered_stage == "early" and not is_wf) - slim_time = ts.metadata["SLiM"]["tick"] - time + x + y - 1 + slim_time = ts_metadata["SLiM"]["tick"] - time + x + y - 1 return slim_time @@ -920,7 +1086,7 @@ def annotate(ts, **kwargs): :param str reference_sequence: A reference sequence of length equal to ts.sequence_length. :param bool annotate_mutations: Whether to replace mutation metadata - with defaults. (If False, the mutation table is unchanged.) + with defaults. (If False, information about mutations is unchanged.) """ tables = ts.dump_tables() annotate_tables(tables, **kwargs) @@ -965,12 +1131,15 @@ def annotate_tables( top_metadata["tick"] = tick top_metadata["cycle"] = cycle top_metadata["stage"] = stage - set_tree_sequence_metadata(tables, **top_metadata) + md = tables.metadata + 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) _annotate_populations(tables) if annotate_mutations: - _annotate_sites_mutations(tables) + _annotate_sites_mutations(tables, ts_metadata=ts_metadata) if reference_sequence is not None: tables.reference_sequence.data = reference_sequence @@ -1109,46 +1278,46 @@ def _annotate_populations(tables): tables.populations[j] = p.replace(metadata=md) -def _annotate_sites_mutations(tables): +def _annotate_sites_mutations(tables, ts_metadata): """ Adds to a TableCollection the information relevant to mutations required - for SLiM to load in a tree sequence. This means adding to the metadata column - of the Mutation table, It will also + for SLiM to load in a tree sequence. This means adding metadata to the + SLiM_mutation_list in top-level metadata. It will also: - give SLiM IDs to each mutation - replace ancestral states with "" - This will replace any information already in the metadata or derived state - columns of the Mutation table. We set slim_time in metadata so that + This will replace any information already in the metadata, and the derived + state columns of the Mutation table. We set slim_time in metadata so that - tick = floor(tskit time) + slim_time """ - if len(tables.mutations.metadata) > 0: + if ( + isinstance(ts_metadata, dict) + and "SLiM_mutation_list" in ts_metadata + and len(ts_metadata["SLiM_mutation_list"]) > 0 + ): warnings.warn( - "The provided tree sequence already has some mutations with " + "The provided tree sequence already has top-level mutation " "metadata; this metadata will be overwritten." ) num_mutations = tables.mutations.num_rows default_mut = default_slim_metadata("mutation_list_entry") - dsb, dso = tskit.pack_bytes([str(j).encode() for j in range(num_mutations)]) - slim_time = tables.metadata["SLiM"]["tick"] - np.floor(tables.mutations.time).astype( + slim_time = ts_metadata["SLiM"]["tick"] - np.floor(tables.mutations.time).astype( "int" ) - mms = tables.mutations.metadata_schema - mutation_metadata = [ - mms.encode_row( - { - "mutation_list": [ - { - "mutation_type": default_mut["mutation_type"], - "selection_coeff": default_mut["selection_coeff"], - "subpopulation": default_mut["subpopulation"], - "slim_time": st, - "nucleotide": default_mut["nucleotide"], - } - ] - } - ) - for st in slim_time + mutation_list = [ + { + "mutation_id": j, + "mutation_type": default_mut["mutation_type"], + "per_trait": default_mut["per_trait"], + "subpopulation": default_mut["subpopulation"], + "slim_time": int(st), + "nucleotide": default_mut["nucleotide"], + "padding": None, + } + for j, st in enumerate(slim_time) ] - mdb, mdo = tskit.pack_bytes(mutation_metadata) + ts_metadata["SLiM_mutation_list"] = mutation_list + tables.metadata = ts_metadata + dsb, dso = tskit.pack_bytes([str(j).encode() for j in range(num_mutations)]) tables.mutations.set_columns( site=tables.mutations.site, node=tables.mutations.node, @@ -1156,8 +1325,6 @@ def _annotate_sites_mutations(tables): derived_state=dsb, derived_state_offset=dso, parent=tables.mutations.parent, - metadata=mdb, - metadata_offset=mdo, ) tables.sites.set_columns( position=tables.sites.position, diff --git a/pyslim/slim_metadata.py b/pyslim/slim_metadata.py index 0fb4144..e515f55 100644 --- a/pyslim/slim_metadata.py +++ b/pyslim/slim_metadata.py @@ -1,6 +1,8 @@ +import copy import json import warnings +import numpy as np import tskit from ._version import * # noqa F403 @@ -38,227 +40,309 @@ def is_vacant_num_bytes(num_chromosomes): _raw_slim_metadata_schemas = { "tree_sequence": { "$schema": "http://json-schema.org/schema#", - "codec": "json", - "examples": [ - { - "SLiM": { - "file_version": "0.9", - "name": "fox", - "description": "foxes on Catalina island", - "cycle": 123, - "tick": 123, - "model_type": "WF", - "this_chromosome": { - "id": 1, - "index": 0, - "symbol": "1", - "name": "autosome_1", - "type": "A", - }, - "chromosomes": [ - {"id": 1, "symbol": "1", "name": "autosome_1", "type": "A"}, - {"id": 35, "symbol": "MT", "name": "mtDNA", "type": "HF"}, - ], - "nucleotide_based": False, - "separate_sexes": True, - "spatial_dimensionality": "xy", - "spatial_periodicity": "x", + "codec": "json+struct", + "json": { + "codec": "json", + "description": "SLiM schema for JSON top-level metadata.", + "examples": [ + { + "SLiM": { + "chromosomes": [ + {"id": 1, "name": "autosome_1", "symbol": "1", "type": "A"}, + {"id": 35, "name": "mtDNA", "symbol": "MT", "type": "HF"}, + ], + "cycle": 123, + "description": "foxes on Catalina island", + "file_version": "1.0", + "model_type": "WF", + "name": "fox", + "nucleotide_based": False, + "separate_sexes": True, + "spatial_dimensionality": "xy", + "spatial_periodicity": "x", + "this_chromosome": { + "id": 1, + "index": 0, + "name": "autosome_1", + "symbol": "1", + "type": "A", + }, + "tick": 123, + "traits": [ + {"index": 0, "name": "simT", "type": "multiplicative"} + ], + } } - } - ], - "properties": { - "SLiM": { - "description": "Top-level metadata for a SLiM tree sequence, file format version 0.9", - "properties": { - "file_version": { - "description": "The SLiM 'file format version' of this tree sequence.", - "type": "string", - }, - "name": { - "description": "The SLiM species name represented by this tree sequence.", - "type": "string", - }, - "description": { - "description": "A user-configurable description of the species represented by this tree sequence.", - "type": "string", - }, - "cycle": { - "description": "The 'SLiM cycle' counter when this tree sequence was recorded.", - "type": "integer", - }, - "tick": { - "description": "The 'SLiM tick' counter when this tree sequence was recorded.", - "type": "integer", - }, - "model_type": { - "description": "The model type used for the last part of this simulation (WF or nonWF).", - "enum": ["WF", "nonWF"], - "type": "string", - }, - "this_chromosome": { - "description": "The chromosome represented by the tree sequence in this file.", - "properties": { - "id": { - "description": "An integer identifier for the chromosome, unique within this set of tree sequences; often the chromosome number in the organism being represented, such as 1.", - "type": "integer", - }, - "index": { - "description": "The (zero-based) index of this chromosome in the chromosomes metadata array (if present), which should match the information given here.", - "type": "integer", - }, - "symbol": { - "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', - "type": "string", - }, - "name": { - "description": "A user-specified name for the chromosome, such as an accession identifier.", - "type": "string", - }, - "type": { - "description": "The type of chromosome, as specified by SLiM.", - "type": "string", + ], + "properties": { + "SLiM": { + "description": "Top-level metadata for a SLiM tree sequence, file format version 1.0", + "properties": { + "chromosomes": { + "description": "The chromosomes represented by the collection of tree sequences, of which this tree sequence is one member.", + "items": { + "properties": { + "id": { + "description": "An integer identifier for the chromosome, unique within this set of tree sequences; often the chromosome number in the organism being represented, such as 1.", + "type": "integer", + }, + "name": { + "description": "A user-specified name for the chromosome, such as an accession identifier.", + "type": "string", + }, + "symbol": { + "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', + "type": "string", + }, + "type": { + "description": "The type of chromosome, as specified by SLiM.", + "type": "string", + }, + }, + "required": ["id", "symbol", "type"], + "type": "object", }, + "type": "array", }, - "required": ["id", "index", "symbol", "type"], - "type": "object", - }, - "chromosomes": { - "description": "The chromosomes represented by the collection of tree sequences, of which this tree sequence is one member.", - "items": { + "cycle": { + "description": "The 'SLiM cycle' counter when this tree sequence was recorded.", + "type": "integer", + }, + "description": { + "description": "A user-configurable description of the species represented by this tree sequence.", + "type": "string", + }, + "file_version": { + "description": "The SLiM 'file format version' of this tree sequence.", + "type": "string", + }, + "model_type": { + "description": "The model type used for the last part of this simulation (WF or nonWF).", + "enum": ["WF", "nonWF"], + "type": "string", + }, + "name": { + "description": "The SLiM species name represented by this tree sequence.", + "type": "string", + }, + "nucleotide_based": { + "description": "Whether the simulation was nucleotide-based.", + "type": "boolean", + }, + "separate_sexes": { + "description": "Whether the simulation had separate sexes.", + "type": "boolean", + }, + "spatial_dimensionality": { + "description": "The spatial dimensionality of the simulation.", + "enum": ["", "x", "xy", "xyz"], + "type": "string", + }, + "spatial_periodicity": { + "description": "The spatial periodicity of the simulation.", + "enum": ["", "x", "y", "z", "xy", "xz", "yz", "xyz"], + "type": "string", + }, + "stage": { + "description": "The stage of the SLiM life cycle when this tree sequence was recorded.", + "type": "string", + }, + "this_chromosome": { + "description": "The chromosome represented by the tree sequence in this file.", "properties": { "id": { "description": "An integer identifier for the chromosome, unique within this set of tree sequences; often the chromosome number in the organism being represented, such as 1.", "type": "integer", }, - "symbol": { - "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', - "type": "string", + "index": { + "description": "The (zero-based) index of this chromosome in the chromosomes metadata array (if present), which should match the information given here.", + "type": "integer", }, "name": { "description": "A user-specified name for the chromosome, such as an accession identifier.", "type": "string", }, + "symbol": { + "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', + "type": "string", + }, "type": { "description": "The type of chromosome, as specified by SLiM.", "type": "string", }, }, - "required": ["id", "symbol", "type"], + "required": ["id", "index", "symbol", "type"], "type": "object", }, - "type": "array", - }, - "nucleotide_based": { - "description": "Whether the simulation was nucleotide-based.", - "type": "boolean", - }, - "separate_sexes": { - "description": "Whether the simulation had separate sexes.", - "type": "boolean", - }, - "spatial_dimensionality": { - "description": "The spatial dimensionality of the simulation.", - "enum": ["", "x", "xy", "xyz"], - "type": "string", - }, - "spatial_periodicity": { - "description": "The spatial periodicity of the simulation.", - "enum": ["", "x", "y", "z", "xy", "xz", "yz", "xyz"], - "type": "string", - }, - "stage": { - "description": "The stage of the SLiM life cycle when this tree sequence was recorded.", - "type": "string", - }, - }, - "required": [ - "model_type", - "tick", - "file_version", - "spatial_dimensionality", - "spatial_periodicity", - "this_chromosome", - "separate_sexes", - "nucleotide_based", - ], - "type": "object", - } - }, - "required": ["SLiM"], - "type": "object", - }, - "edge": None, - "site": None, - "mutation": { - "$schema": "http://json-schema.org/schema#", - "additionalProperties": False, - "codec": "struct", - "description": "SLiM schema for mutation metadata.", - "examples": [ - { - "mutation_list": [ - { - "mutation_type": 1, - "nucleotide": 3, - "selection_coeff": -0.2, - "slim_time": 243, - "subpopulation": 0, - } - ] - } - ], - "properties": { - "mutation_list": { - "items": { - "additionalProperties": False, - "properties": { - "mutation_type": { - "binaryFormat": "i", - "description": "The index of this mutation's mutationType.", - "index": 1, - "type": "integer", - }, - "nucleotide": { - "binaryFormat": "b", - "description": "The nucleotide for this mutation (0=A , 1=C , 2=G, 3=T, or -1 for none)", - "index": 5, - "type": "integer", - }, - "selection_coeff": { - "binaryFormat": "f", - "description": "This mutation's selection coefficient.", - "index": 2, - "type": "number", - }, - "slim_time": { - "binaryFormat": "i", - "description": "The SLiM tick counter when this mutation occurred.", - "index": 4, + "tick": { + "description": "The 'SLiM tick' counter when this tree sequence was recorded.", "type": "integer", }, - "subpopulation": { - "binaryFormat": "i", - "description": "The ID of the subpopulation this mutation occurred in.", - "index": 3, - "type": "integer", + "traits": { + "description": "The traits defined for this tree sequence; each mutation and individual will have per-trait metadata.", + "items": { + "properties": { + "baselineAccumulation": { + "description": "Whether the baseline offset includes accumulated effects from fixed (substituted) mutations.", + "type": "boolean", + }, + "baselineOffset": { + "description": "The baseline offset of the trait.", + "type": "number", + }, + "directFitnessEffect": { + "description": "Whether the trait's effects are used directly as fitness effects.", + "type": "boolean", + }, + "index": { + "description": "The integer index for the trait; indices must be sequential starting from zero.", + "type": "integer", + }, + "individualOffsetMean": { + "description": "The mean of the trait's individual offset distribution (which might or might not be used).", + "type": "number", + }, + "individualOffsetSD": { + "description": "The standard deviation of the trait's individual offset distribution (which might or might not be used).", + "type": "number", + }, + "name": { + "description": "The string name for the trait.", + "type": "string", + }, + "type": { + "description": "The type of the trait; this must be 'additive', 'multiplicative', or 'logistic'.", + "enum": [ + "additive", + "multiplicative", + "logistic", + ], + "type": "string", + }, + }, + "required": ["index", "name", "type"], + "type": "object", + }, + "type": "array", }, }, "required": [ - "mutation_type", - "selection_coeff", - "subpopulation", - "slim_time", - "nucleotide", + "model_type", + "tick", + "file_version", + "spatial_dimensionality", + "spatial_periodicity", + "this_chromosome", + "separate_sexes", + "nucleotide_based", + "traits", ], "type": "object", - }, - "noLengthEncodingExhaustBuffer": True, - "type": "array", - } + } + }, + "required": ["SLiM"], + "type": "object", + }, + "struct": { + "codec": "struct", + "description": "SLiM schema for binary top-level metadata.", + "properties": { + "SLiM_mutation_list": { + "arrayLengthFormat": "Q", + "items": { + "additionalProperties": False, + "properties": { + "mutation_id": { + "binaryFormat": "q", + "description": "The SLiM mutation ID for this mutation.", + "index": 1, + "type": "integer", + }, + "mutation_type": { + "binaryFormat": "i", + "description": "The id of this mutation's mutationType.", + "index": 2, + "type": "integer", + }, + "nucleotide": { + "binaryFormat": "b", + "description": "The nucleotide for this mutation (0=A , 1=C , 2=G, 3=T, or -1 for none)", + "index": 5, + "type": "integer", + }, + "padding": { + "binaryFormat": "3x", + "description": "Padding bytes for alignment", + "index": 6, + "type": "null", + }, + "per_trait": { + "index": 7, + "items": { + "additionalProperties": False, + "properties": { + "dominance": { + "binaryFormat": "f", + "description": "The dominance coefficient for this trait.", + "index": 2, + "type": "number", + }, + "effect_size": { + "binaryFormat": "f", + "description": "The effect size for this trait.", + "index": 1, + "type": "number", + }, + "hemizygous_dominance": { + "binaryFormat": "f", + "description": "The hemizygous dominance coefficient for this trait.", + "index": 3, + "type": "number", + }, + }, + "required": [ + "dominance", + "effect_size", + "hemizygous_dominance", + ], + "type": "object", + }, + "length": 1, # NOTE this may need to be changed to match the number of traits! + "type": "array", + }, + "slim_time": { + "binaryFormat": "i", + "description": "The SLiM tick counter when this mutation occurred.", + "index": 4, + "type": "integer", + }, + "subpopulation": { + "binaryFormat": "i", + "description": "The ID of the subpopulation this mutation occurred in.", + "index": 3, + "type": "integer", + }, + }, + "required": [ + "mutation_id", + "mutation_type", + "slim_time", + "subpopulation", + "nucleotide", + "per_trait", + ], + "type": "object", + }, + "type": "array", + } + }, + "required": ["SLiM_mutation_list"], + "type": "object", }, - "required": ["mutation_list"], - "type": "object", }, + "edge": None, + "site": None, + "mutation": None, "node": { "$schema": "http://json-schema.org/schema#", "additionalProperties": False, @@ -269,12 +353,12 @@ def is_vacant_num_bytes(num_chromosomes): "slim_id": { "binaryFormat": "q", "description": "The 'pedigree ID' of the haplosomes associated with this node in SLiM.", - "index": 0, + "index": 1, "type": "integer", }, "is_vacant": { "description": "A vector of byte (uint8_t) values, with each bit representing whether the node represents a vacant position, either unused or a null haplosome (1), or a non-null haplosome (0), in the corresponding chromosome. This field encodes vacancy for all of the chromosomes in the model, not just the chromosome represented in this file (so that the node table is identical across all chromosomes for a multi-chromosome model). Each chromosome receives one bit here; there are two node table entries per individual, used for the two haplosomes of every chromosome, so only one bit is needed in each entry (making two bits total per chromosome, across the two node table entries). The least significant bit of the first byte is used first (for one haplosome of the first chromosome); the most significant bit of the last byte is used last. The number of bytes present in this field is indicated by this schema's 'binaryFormat' field, which is variable (!), and can also be deduced from the number of chromosomes in the model as given in the top-level 'chromosomes' metadata key, which should always be present if this metadata is present.", - "index": 1, + "index": 2, "type": "array", "length": 1, # MAY NEED TO BE CHANGED (in SLiM code is "%d") "items": {"type": "number", "binaryFormat": "B"}, @@ -285,8 +369,8 @@ def is_vacant_num_bytes(num_chromosomes): }, "individual": { "$schema": "http://json-schema.org/schema#", - "additionalProperties": False, "codec": "struct", + "type": "object", "description": "SLiM schema for individual metadata.", "examples": [ { @@ -295,64 +379,170 @@ def is_vacant_num_bytes(num_chromosomes): "pedigree_id": 123, "pedigree_p1": 12, "pedigree_p2": 23, + "per_trait": [{"offset": 1.0, "phenotype": 1.1}], "sex": 0, "subpopulation": 0, + "tag": 1, + "tagF": 5.5, + "tagL0_set": True, + "tagL0": True, + "tagL1_set": True, + "tagL1": False, + "tagL2_set": False, + "tagL2": False, + "tagL3_set": False, + "tagL3": False, + "tagL4_set": False, + "tagL4": False, } ], "flags": { "SLIM_INDIVIDUAL_METADATA_MIGRATED": { - "description": "Whether this individual was a migrant, either in the tick when the tree sequence " - "was written out (if the individual was alive then), or in the tick of the last time " - "they were Remembered (if not).", + "description": "Whether this individual was a migrant, either in the tick when the tree sequence was written out (if the individual was alive then), or in the tick of the last time they were Remembered (if not).", "value": 1, } }, "properties": { - "age": { - "binaryFormat": "i", - "description": "The age of this individual, either when the tree sequence was written out " - "(if the individual was alive then), or the last time they were Remembered (if not).", - "index": 4, - "type": "integer", - }, - "flags": { - "binaryFormat": "I", - "description": "Other information about the individual: see 'flags'.", - "index": 7, - "type": "integer", - }, "pedigree_id": { - "binaryFormat": "q", - "description": "The 'pedigree ID' of this individual in SLiM.", "index": 1, "type": "integer", + "binaryFormat": "q", + "description": "The 'pedigree ID' of this individual in SLiM.", }, "pedigree_p1": { - "binaryFormat": "q", - "description": "The 'pedigree ID' of this individual's first parent in SLiM.", "index": 2, "type": "integer", + "binaryFormat": "q", + "description": "The 'pedigree ID' of this individual's first parent in SLiM.", }, "pedigree_p2": { + "index": 3, + "type": "integer", "binaryFormat": "q", "description": "The 'pedigree ID' of this individual's second parent in SLiM.", - "index": 3, + }, + "age": { + "index": 4, "type": "integer", + "binaryFormat": "i", + "description": "The age of this individual, either when the tree sequence was written out (if the individual was alive then), or the last time they were Remembered (if not).", + }, + "subpopulation": { + "index": 5, + "type": "integer", + "binaryFormat": "i", + "description": "The ID of the subpopulation the individual was part of, either when the tree sequence was written out (if the individual was alive then), or the last time they were Remembered (if not).", }, "sex": { + "index": 6, + "type": "integer", "binaryFormat": "i", "description": "The sex of the individual (0 for female, 1 for male, -1 for hermaphrodite).", - "index": 6, + }, + "flags": { + "index": 7, "type": "integer", + "binaryFormat": "I", + "description": "Other information about the individual: see 'flags'.", }, - "subpopulation": { - "binaryFormat": "i", - "description": "The ID of the subpopulation the individual was part of, either when the tree sequence " - "was written out (if the individual was alive then), or the last time they were Remembered (if not).", - "index": 5, + "tag": { + "index": 8, "type": "integer", + "binaryFormat": "q", + "description": "The `tag` property of this individual; INT64_MIN if unset.", + }, + "tagF": { + "index": 9, + "type": "number", + "binaryFormat": "d", + "description": "The `tagF` property of this individual; -DBL_MAX if unset.", + }, + "tagL0_set": { + "index": 10, + "type": "boolean", + "binaryFormat": "?", + "description": "A flag indicating whether the `tagL0` property is set; if false, accessing `tagL0` is invalid.", + }, + "tagL0": { + "index": 11, + "type": "boolean", + "binaryFormat": "?", + "description": "The `tagL0` property of this individual; only valid if `tagL0_set` is true.", + }, + "tagL1_set": { + "index": 12, + "type": "boolean", + "binaryFormat": "?", + "description": "A flag indicating whether the `tagL1` property is set; if false, accessing `tagL1` is invalid.", + }, + "tagL1": { + "index": 13, + "type": "boolean", + "binaryFormat": "?", + "description": "The `tagL1` property of this individual; only valid if `tagL1_set` is true.", + }, + "tagL2_set": { + "index": 14, + "type": "boolean", + "binaryFormat": "?", + "description": "A flag indicating whether the `tagL2` property is set; if false, accessing `tagL2` is invalid.", + }, + "tagL2": { + "index": 15, + "type": "boolean", + "binaryFormat": "?", + "description": "The `tagL2` property of this individual; only valid if `tagL2_set` is true.", + }, + "tagL3_set": { + "index": 16, + "type": "boolean", + "binaryFormat": "?", + "description": "A flag indicating whether the `tagL3` property is set; if false, accessing `tagL3` is invalid.", + }, + "tagL3": { + "index": 17, + "type": "boolean", + "binaryFormat": "?", + "description": "The `tagL3` property of this individual; only valid if `tagL3_set` is true.", + }, + "tagL4_set": { + "index": 18, + "type": "boolean", + "binaryFormat": "?", + "description": "A flag indicating whether the `tagL4` property is set; if false, accessing `tagL4` is invalid.", + }, + "tagL4": { + "index": 19, + "type": "boolean", + "binaryFormat": "?", + "description": "The `tagL4` property of this individual; only valid if `tagL4_set` is true.", + }, + "per_trait": { + "index": 20, + "type": "array", + "length": 1, # MAY NEED TO BE CHANGED (in SLiM code is "%d") + "items": { + "additionalProperties": False, + "properties": { + "phenotype": { + "index": 1, + "type": "number", + "binaryFormat": "d", + "description": "The phenotype for this trait.", + }, + "offset": { + "index": 2, + "type": "number", + "binaryFormat": "d", + "description": "The individual offset for this trait.", + }, + }, + "required": ["offset", "phenotype"], + "type": "object", + }, }, }, + "additionalProperties": False, "required": [ "pedigree_id", "pedigree_p1", @@ -360,9 +550,21 @@ def is_vacant_num_bytes(num_chromosomes): "age", "subpopulation", "sex", + "tag", + "tagF", + "tagL0_set", + "tagL0", + "tagL1_set", + "tagL1", + "tagL2_set", + "tagL2", + "tagL3_set", + "tagL3", + "tagL4_set", + "tagL4", "flags", + "per_trait", ], - "type": "object", }, "population": { "$schema": "http://json-schema.org/schema#", @@ -464,6 +666,45 @@ def is_vacant_num_bytes(num_chromosomes): } +def slim_tree_sequence_metadata_schema(num_traits=1): + """ + The top-level metadata schema depends on the number of traits, and + {data}`.slim_metadata_schemas` + returns the schema for a single-trait simulation. This function + returns the correct schema for a simulation with arbitrary number of + traits. (The resulting schemas only differ in the + "length" of the "per_trait" property of + ``schema["properties"]["SLiM_mutation_list"]["items"]``). + + :param int num_traits: The number of traits in the model. + :return tskit.MetadataSchema: The metadata schema to be used + in the node table. + """ + schema = _raw_slim_metadata_schemas["tree_sequence"] + schema["struct"]["properties"]["SLiM_mutation_list"]["items"]["properties"][ + "per_trait" + ]["length"] = num_traits + return tskit.MetadataSchema(schema) + + +def slim_individual_metadata_schema(num_traits=1): + """ + The individual metadata schema depends on the number of traits, and + {data}`.slim_metadata_schemas` + returns the schema for a single-trait simulation. This function + returns the correct schema for a simulation with arbitrary number of + traits. (The resulting schemas only differ in the + "length" of the "per_trait" property.) + + :param int num_traits: The number of traits in the model. + :return tskit.MetadataSchema: The metadata schema to be used + in the node table. + """ + schema = _raw_slim_metadata_schemas["individual"] + schema["properties"]["per_trait"]["length"] = num_traits + return tskit.MetadataSchema(schema) + + def slim_node_metadata_schema(num_chromosomes=1): """ Unlike other schema, the node metadata schema depends on the number of @@ -510,12 +751,15 @@ def slim_node_metadata_schema(num_chromosomes=1): """ -def default_slim_metadata(name, num_chromosomes=1): +def default_slim_metadata(name, num_chromosomes=1, num_traits=1, **kwargs): """ Returns default metadata of type ``name``, where ``name`` is one of "tree_sequence", "edge", "site", "mutation", "mutation_list_entry", "node", "individual", or "population". + Additional kwargs are used to update the resulting metadata + (without validity checking). + :param str name: The type of metadata requested. :rtype dict: """ @@ -540,21 +784,26 @@ def default_slim_metadata(name, num_chromosomes=1): "type": "A", }, "chromosomes": [{"id": 1, "index": 0, "symbol": "A", "type": "A"}], - } + "traits": [{"index": 0, "name": "simT", "type": "multiplicative"}], + }, + "SLiM_mutation_list": [], } elif name == "edge": out = None elif name == "site": out = None elif name == "mutation": - out = {"mutation_list": []} + out = None elif name == "mutation_list_entry": out = { + "mutation_id": 0, "mutation_type": 0, - "selection_coeff": 0.0, "subpopulation": tskit.NULL, "slim_time": 0, "nucleotide": -1, + "per_trait": num_traits + * [{"effect_size": 0.0, "dominance": 0.5, "hemizygous_dominance": 1.0}], + "padding": None, } elif name == "node": out = { @@ -570,6 +819,19 @@ def default_slim_metadata(name, num_chromosomes=1): "flags": 0, "pedigree_p1": tskit.NULL, "pedigree_p2": tskit.NULL, + "tag": np.iinfo(np.int64).min, + "tagF": np.finfo(np.float64).min, + "tagL0_set": False, + "tagL0": False, + "tagL1_set": False, + "tagL1": False, + "tagL2_set": False, + "tagL2": False, + "tagL3_set": False, + "tagL3": False, + "tagL4_set": False, + "tagL4": False, + "per_trait": num_traits * [{"phenotype": np.nan, "offset": 1.0}], } elif name == "population": out = { @@ -594,6 +856,8 @@ def default_slim_metadata(name, num_chromosomes=1): "'edge', 'site', 'mutation', 'mutation_list_entry', 'node', " "'individual', or 'population'." ) + if out is not None: + out.update(kwargs) return out @@ -606,6 +870,7 @@ def set_tree_sequence_metadata( tables, model_type, tick, + *, cycle=None, spatial_dimensionality="", spatial_periodicity="", @@ -618,35 +883,58 @@ def set_tree_sequence_metadata( chromosomes=None, file_version=None, set_table_schemas=True, + traits=None, + SLiM_mutation_list=None, ): if file_version is None: file_version = slim_file_version - if isinstance(tables.metadata, bytes): - if len(tables.metadata) > 0: + if traits is None: + traits = [{"index": 0, "name": "simT", "type": "multiplicative"}] + num_traits = len(traits) + schema_dict = slim_tree_sequence_metadata_schema(num_traits).schema + old_schema_dict = tables.metadata_schema.schema + old_json_schema_dict = {} + old_struct_schema_dict = {} + tmd = tables.metadata + if isinstance(tmd, bytes): + if len(tmd) > 0: raise ValueError( "Tree sequence has top-level metadata but no schema: this is a problem " "since pyslim is trying to add to the metadata." ) - schema_dict = slim_metadata_schemas["tree_sequence"].schema metadata_dict = {} else: # we need to keep other keys in the metadata (and schema) if there are any - schema_dict = tables.metadata_schema.schema metadata_dict = tables.metadata + if old_schema_dict["codec"] == "json": + old_json_schema_dict = tables.metadata_schema.schema + else: + assert old_schema_dict["codec"] == "json+struct", ( + "You are using an unexpected codec; " + "please raise an issue on pyslim if " + "you need this functionality." + ) + old_json_schema_dict = tables.metadata_schema.schema["json"] + old_struct_schema_dict = tables.metadata_schema.schema["struct"] if cycle is None: cycle = tick - defaults = default_slim_metadata("tree_sequence") + if chromosomes is None: + num_chromosomes = 1 + else: + num_chromosomes = len(chromosomes) + defaults = default_slim_metadata( + "tree_sequence", num_chromosomes=num_chromosomes, num_traits=num_traits + ) if this_chromosome is None: this_chromosome = defaults["SLiM"]["this_chromosome"] if chromosomes is None: chromosomes = defaults["SLiM"]["chromosomes"] - assert schema_dict["codec"] == "json" - assert schema_dict["type"] == "object" - if "properties" not in schema_dict: - schema_dict["properties"] = {} - schema_dict["properties"]["SLiM"] = slim_metadata_schemas["tree_sequence"].schema[ - "properties" - ]["SLiM"] + if "properties" in old_json_schema_dict: + schema_dict["json"]["properties"].update(old_json_schema_dict["properties"]) + if "properties" in old_struct_schema_dict: + schema_dict["struct"]["properties"].update(old_struct_schema_dict["properties"]) + if SLiM_mutation_list is None: + SLiM_mutation_list = [] tables.metadata_schema = tskit.MetadataSchema(schema_dict) metadata_dict["SLiM"] = { "model_type": model_type, @@ -662,16 +950,19 @@ def set_tree_sequence_metadata( "description": description, "this_chromosome": this_chromosome, "chromosomes": chromosomes, + "traits": traits, } + metadata_dict["SLiM_mutation_list"] = SLiM_mutation_list tables.metadata = metadata_dict + return metadata_dict -def set_metadata_schemas(tables, num_chromosomes=1): +def set_metadata_schemas(tables, num_chromosomes=1, num_traits=1): tables.edges.metadata_schema = slim_metadata_schemas["edge"] tables.sites.metadata_schema = slim_metadata_schemas["site"] tables.mutations.metadata_schema = slim_metadata_schemas["mutation"] tables.nodes.metadata_schema = slim_node_metadata_schema(num_chromosomes) - tables.individuals.metadata_schema = slim_metadata_schemas["individual"] + tables.individuals.metadata_schema = slim_individual_metadata_schema(num_traits) tables.populations.metadata_schema = slim_metadata_schemas["population"] @@ -683,6 +974,161 @@ def _old_metadata_schema(name, file_version): # Returns a metadata schema *if the format has changed*, # and None otherwise. ms = None + if name == "tree_sequence" and file_version == "0.9": + pre_1_0_tree_sequence = { + "$schema": "http://json-schema.org/schema#", + "codec": "json", + "examples": [ + { + "SLiM": { + "file_version": "0.9", + "name": "fox", + "description": "foxes on Catalina island", + "cycle": 123, + "tick": 123, + "model_type": "WF", + "this_chromosome": { + "id": 1, + "index": 0, + "symbol": "1", + "name": "autosome_1", + "type": "A", + }, + "chromosomes": [ + {"id": 1, "symbol": "1", "name": "autosome_1", "type": "A"}, + {"id": 35, "symbol": "MT", "name": "mtDNA", "type": "HF"}, + ], + "nucleotide_based": False, + "separate_sexes": True, + "spatial_dimensionality": "xy", + "spatial_periodicity": "x", + } + } + ], + "properties": { + "SLiM": { + "description": "Top-level metadata for a SLiM tree sequence, file format version 0.9", + "properties": { + "file_version": { + "description": "The SLiM 'file format version' of this tree sequence.", + "type": "string", + }, + "name": { + "description": "The SLiM species name represented by this tree sequence.", + "type": "string", + }, + "description": { + "description": "A user-configurable description of the species represented by this tree sequence.", + "type": "string", + }, + "cycle": { + "description": "The 'SLiM cycle' counter when this tree sequence was recorded.", + "type": "integer", + }, + "tick": { + "description": "The 'SLiM tick' counter when this tree sequence was recorded.", + "type": "integer", + }, + "model_type": { + "description": "The model type used for the last part of this simulation (WF or nonWF).", + "enum": ["WF", "nonWF"], + "type": "string", + }, + "this_chromosome": { + "description": "The chromosome represented by the tree sequence in this file.", + "properties": { + "id": { + "description": "An integer identifier for the chromosome, unique within this set of tree sequences; often the chromosome number in the organism being represented, such as 1.", + "type": "integer", + }, + "index": { + "description": "The (zero-based) index of this chromosome in the chromosomes metadata array (if present), which should match the information given here.", + "type": "integer", + }, + "symbol": { + "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', + "type": "string", + }, + "name": { + "description": "A user-specified name for the chromosome, such as an accession identifier.", + "type": "string", + }, + "type": { + "description": "The type of chromosome, as specified by SLiM.", + "type": "string", + }, + }, + "required": ["id", "index", "symbol", "type"], + "type": "object", + }, + "chromosomes": { + "description": "The chromosomes represented by the collection of tree sequences, of which this tree sequence is one member.", + "items": { + "properties": { + "id": { + "description": "An integer identifier for the chromosome, unique within this set of tree sequences; often the chromosome number in the organism being represented, such as 1.", + "type": "integer", + }, + "symbol": { + "description": 'A short string symbol for the chromosome, unique within this set of tree sequences, such as "1" or "MT".', + "type": "string", + }, + "name": { + "description": "A user-specified name for the chromosome, such as an accession identifier.", + "type": "string", + }, + "type": { + "description": "The type of chromosome, as specified by SLiM.", + "type": "string", + }, + }, + "required": ["id", "symbol", "type"], + "type": "object", + }, + "type": "array", + }, + "nucleotide_based": { + "description": "Whether the simulation was nucleotide-based.", + "type": "boolean", + }, + "separate_sexes": { + "description": "Whether the simulation had separate sexes.", + "type": "boolean", + }, + "spatial_dimensionality": { + "description": "The spatial dimensionality of the simulation.", + "enum": ["", "x", "xy", "xyz"], + "type": "string", + }, + "spatial_periodicity": { + "description": "The spatial periodicity of the simulation.", + "enum": ["", "x", "y", "z", "xy", "xz", "yz", "xyz"], + "type": "string", + }, + "stage": { + "description": "The stage of the SLiM life cycle when this tree sequence was recorded.", + "type": "string", + }, + }, + "required": [ + "model_type", + "tick", + "file_version", + "spatial_dimensionality", + "spatial_periodicity", + "this_chromosome", + "separate_sexes", + "nucleotide_based", + ], + "type": "object", + } + }, + "required": ["SLiM"], + "type": "object", + } + + ms = pre_1_0_tree_sequence + if name == "tree_sequence" and file_version == "0.8": pre_0_9_tree_sequence = { "$schema": "http://json-schema.org/schema#", @@ -963,6 +1409,90 @@ def _old_metadata_schema(name, file_version): } ms = pre_0_7_population + if name == "individual" and file_version in ["0.7", "0.8", "0.9"]: + pre_1_0_individual = { + "$schema": "http://json-schema.org/schema#", + "additionalProperties": False, + "codec": "struct", + "description": "SLiM schema for individual metadata.", + "examples": [ + { + "age": -1, + "flags": 0, + "pedigree_id": 123, + "pedigree_p1": 12, + "pedigree_p2": 23, + "sex": 0, + "subpopulation": 0, + } + ], + "flags": { + "SLIM_INDIVIDUAL_METADATA_MIGRATED": { + "description": "Whether this individual was a migrant, either in the tick when the tree sequence " + "was written out (if the individual was alive then), or in the tick of the last time " + "they were Remembered (if not).", + "value": 1, + } + }, + "properties": { + "age": { + "binaryFormat": "i", + "description": "The age of this individual, either when the tree sequence was written out " + "(if the individual was alive then), or the last time they were Remembered (if not).", + "index": 4, + "type": "integer", + }, + "flags": { + "binaryFormat": "I", + "description": "Other information about the individual: see 'flags'.", + "index": 7, + "type": "integer", + }, + "pedigree_id": { + "binaryFormat": "q", + "description": "The 'pedigree ID' of this individual in SLiM.", + "index": 1, + "type": "integer", + }, + "pedigree_p1": { + "binaryFormat": "q", + "description": "The 'pedigree ID' of this individual's first parent in SLiM.", + "index": 2, + "type": "integer", + }, + "pedigree_p2": { + "binaryFormat": "q", + "description": "The 'pedigree ID' of this individual's second parent in SLiM.", + "index": 3, + "type": "integer", + }, + "sex": { + "binaryFormat": "i", + "description": "The sex of the individual (0 for female, 1 for male, -1 for hermaphrodite).", + "index": 6, + "type": "integer", + }, + "subpopulation": { + "binaryFormat": "i", + "description": "The ID of the subpopulation the individual was part of, either when the tree sequence " + "was written out (if the individual was alive then), or the last time they were Remembered (if not).", + "index": 5, + "type": "integer", + }, + }, + "required": [ + "pedigree_id", + "pedigree_p1", + "pedigree_p2", + "age", + "subpopulation", + "sex", + "flags", + ], + "type": "object", + } + ms = pre_1_0_individual + if name == "individual" and file_version in [ "0.1", "0.2", @@ -1019,6 +1549,87 @@ def _old_metadata_schema(name, file_version): } ms = pre_0_7_individual + if name == "mutation" and file_version in [ + "0.3", + "0.4", + "0.5", + "0.6", + "0.7", + "0.8", + "0.9", + ]: + mutation_pre_1_0 = { + "$schema": "http://json-schema.org/schema#", + "additionalProperties": False, + "codec": "struct", + "description": "SLiM schema for mutation metadata.", + "examples": [ + { + "mutation_list": [ + { + "mutation_type": 1, + "nucleotide": 3, + "selection_coeff": -0.2, + "slim_time": 243, + "subpopulation": 0, + } + ] + } + ], + "properties": { + "mutation_list": { + "items": { + "additionalProperties": False, + "properties": { + "mutation_type": { + "binaryFormat": "i", + "description": "The index of this mutation's mutationType.", + "index": 1, + "type": "integer", + }, + "nucleotide": { + "binaryFormat": "b", + "description": "The nucleotide for this mutation (0=A , 1=C , 2=G, 3=T, or -1 for none)", + "index": 5, + "type": "integer", + }, + "selection_coeff": { + "binaryFormat": "f", + "description": "This mutation's selection coefficient.", + "index": 2, + "type": "number", + }, + "slim_time": { + "binaryFormat": "i", + "description": "The SLiM tick counter when this mutation occurred.", + "index": 4, + "type": "integer", + }, + "subpopulation": { + "binaryFormat": "i", + "description": "The ID of the subpopulation this mutation occurred in.", + "index": 3, + "type": "integer", + }, + }, + "required": [ + "mutation_type", + "selection_coeff", + "subpopulation", + "slim_time", + "nucleotide", + ], + "type": "object", + }, + "noLengthEncodingExhaustBuffer": True, + "type": "array", + } + }, + "required": ["mutation_list"], + "type": "object", + } + ms = mutation_pre_1_0 + if name == "mutation" and file_version in ["0.1", "0.2"]: mutation_pre_0_3 = { "$schema": "http://json-schema.org/schema#", @@ -1072,6 +1683,33 @@ def _old_metadata_schema(name, file_version): } ms = mutation_pre_0_3 + if name == "node" and file_version == "0.9": + node_0_9 = { + "$schema": "http://json-schema.org/schema#", + "additionalProperties": False, + "codec": "struct", + "description": "SLiM schema for node metadata.", + "examples": [{"slim_id": 123, "is_vacant": 0}], + "properties": { + "slim_id": { + "binaryFormat": "q", + "description": "The 'pedigree ID' of the haplosomes associated with this node in SLiM.", + "index": 0, + "type": "integer", + }, + "is_vacant": { + "description": "A vector of byte (uint8_t) values, with each bit representing whether the node represents a vacant position, either unused or a null haplosome (1), or a non-null haplosome (0), in the corresponding chromosome. This field encodes vacancy for all of the chromosomes in the model, not just the chromosome represented in this file (so that the node table is identical across all chromosomes for a multi-chromosome model). Each chromosome receives one bit here; there are two node table entries per individual, used for the two haplosomes of every chromosome, so only one bit is needed in each entry (making two bits total per chromosome, across the two node table entries). The least significant bit of the first byte is used first (for one haplosome of the first chromosome); the most significant bit of the last byte is used last. The number of bytes present in this field is indicated by this schema's 'binaryFormat' field, which is variable (!), and can also be deduced from the number of chromosomes in the model as given in the top-level 'chromosomes' metadata key, which should always be present if this metadata is present.", + "index": 1, + "type": "array", + "length": 1, # MAY NEED TO BE CHANGED (in SLiM code is "%d") + "items": {"type": "number", "binaryFormat": "B"}, + }, + }, + "required": ["slim_id", "is_vacant"], + "type": ["object", "null"], + } + ms = node_0_9 + if name == "node" and file_version in [ "0.1", "0.2", @@ -1119,24 +1757,66 @@ def _old_metadata_schema(name, file_version): return ms +def _make_mutation_list(mutations, file_version): + # Prior to 1.0, mutation metadata was a list of entries like: + # {'mutation_type': 1, 'selection_coeff': -0.1, + # 'subpopulation': 1, 'slim_time': 5, 'nucleotide': -1} + # with mutation id stored in the derived state. + # + # As of 1.0, this lives in the top-level ts.metadata['SLiM_mutation_list'], + # with entries like + # {'mutation_id': 87, 'mutation_type': 1, 'subpopulation': 1, 'slim_time': 5, + # 'nucleotide': -1, 'padding': None, + # 'per_trait': [{'effect_size': -0.1, 'dominance': 0.5, + # 'hemizygous_dominance': 1.0}]} + if mutations.metadata_schema == tskit.MetadataSchema(None): + mutations.metadata_schema = _old_metadata_schema("mutation", file_version) + mutation_list = [] + for mut in mutations: + for sid, md in zip(mut.derived_state.split(","), mut.metadata["mutation_list"]): + if "nucleotide" not in md: + md["nucleotide"] = -1 + md["mutation_id"] = int(sid) + md["per_trait"] = [ + { + "effect_size": md["selection_coeff"], + "dominance": 0.5, # WE DON'T KNOW THIS + "hemizygous_dominance": 1.0, # OR THIS + }, + ] + del md["selection_coeff"] + md["padding"] = None + mutation_list.append(md) + return mutation_list + + def is_current_version(ts, _warn=False): """ - Tests whether the tree sequence or table collection provided is the current - SLiM file format or not. If not, use `pyslim.update( )` to bring it up to - date. + Tests whether the metadata provided is the current SLiM file format or not. + If not, use `pyslim.update( )` to bring it up to date. - :param TreeSequence ts: The tree sequence or table collection. + This method may be provided either a TreeSequence or TableCollection directly, + or the metadata from one of these. The latter is useful because + accessing top-level metadata can be a costly operation. + + :param dict ts: Either the top-level metadata of a tree sequence, + or a TreeSequence or TableCollection that carries this metadata. :return bool: Whether the tree sequence is the current version. """ + if ( + isinstance(ts, tskit.TreeSequence) + or isinstance(ts, tskit.TableCollection) + or isinstance(ts, tskit.ImmutableTableCollection) + ): + ts = ts.metadata out = ( - isinstance(ts.metadata, dict) - and ("SLiM" in ts.metadata) - and (ts.metadata["SLiM"]["file_version"] == slim_file_version) + isinstance(ts, dict) + and ("SLiM" in ts) + and (ts["SLiM"]["file_version"] == slim_file_version) ) if _warn and not out: warnings.warn( - "This tree sequence is not the current SLiM format, " - "so some operations may not work. " + "This tree sequence is not the current SLiM format. " "Use `pyslim.update( )` to update the tree sequence." ) return out @@ -1161,10 +1841,13 @@ def update_tables(tables): """ # First we ensure we can find the file format version number # in top-level metadata. Then we proceed to fix up the tables as necessary. - if not (isinstance(tables.metadata, dict) and "SLiM" in tables.metadata): + md = tables.metadata + if not (isinstance(md, dict) and "SLiM" in md): # Old versions kept information in provenance, not top-level metadata. # Note this uses defaults on keys not present in provenance, # 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"] prov = None file_version = "unknown" @@ -1192,25 +1875,38 @@ def update_tables(tables): values[k] = record["slim"][k] except: raise ValueError("Failed to obtain metadata from provenance.") - set_tree_sequence_metadata(tables, **values) + md = set_tree_sequence_metadata(tables, **values) - file_version = tables.metadata["SLiM"]["file_version"] + file_version = md["SLiM"]["file_version"] if file_version != slim_file_version: warnings.warn( - "This is a version {} SLiM tree sequence.".format(file_version) - + " If you write this out to a file, " - + "it will be converted to version {}.".format(slim_file_version) + f"This is a version {file_version} SLiM tree sequence. " + "If you write this out to a file, " + f"it will be converted to version {slim_file_version}." ) - # the only tables to have metadata schema changed thus far - # are nodes, populations, individuals, mutations, and top-level: old_schema = _old_metadata_schema("tree_sequence", file_version) if old_schema is not None: - md = tables.metadata - new_schema = slim_metadata_schemas["tree_sequence"] - new_properties = new_schema.asdict()["properties"]["SLiM"]["required"] + assert ( + "struct" not in old_schema.schema + or "SLiM_mutation_list" not in old_schema.schema["struct"]["properties"] + ) + # we should get the number of traits from the metadata, + # but for old file versions, this won't be present + assert ( + "json" not in old_schema.schema + or "traits" not in old_schema.schema["json"]["properties"] + ) + md["SLiM_mutation_list"] = _make_mutation_list( + tables.mutations, file_version + ) + num_traits = 1 + new_schema = slim_tree_sequence_metadata_schema(num_traits=num_traits) + new_properties = new_schema.asdict()["json"]["properties"]["SLiM"][ + "required" + ] tables.metadata_schema = new_schema - defaults = default_slim_metadata("tree_sequence") + defaults = default_slim_metadata("tree_sequence", num_traits=num_traits) for k in new_properties: if k not in md["SLiM"]: if k == "tick": @@ -1226,63 +1922,71 @@ def update_tables(tables): tables.nodes.clear() if nodes.metadata_schema == tskit.MetadataSchema(None): nodes.metadata_schema = old_schema - assert "chromosomes" not in tables.metadata - # if chromosomes was in metadata - # we should use its length to get num_chroms, - # but old file versions did not have this. - num_chroms = 1 + if "chromosomes" not in md["SLiM"]: + num_chroms = 1 + else: + num_chroms = len(md["SLiM"]["chromosomes"]) new_schema = slim_node_metadata_schema(num_chroms) tables.nodes.metadata_schema = new_schema - not_vacant = [0] # single chromosome - yes_vacant = [1] - gt = None - for n in nodes: - md = n.metadata - if len(md) > 0: - md["is_vacant"] = yes_vacant if md["is_null"] else not_vacant - if not md["is_null"]: - if gt is None: - gt = md["genome_type"] - else: - assert md["genome_type"] == gt, ( - "Inconsistent tables: " - f"mismatching genome types {gt} and " - f"{md['genome_type']} in node metadata." - ) - del md["is_null"] - del md["genome_type"] - tables.nodes.append(n.replace(metadata=md)) - # flags for node genome type pre-0.9: - # confusingly and sub-optimally, these were redundant: - # all non-null nodes in the same sim would have the same genome type - assert gt is not None - GENOME_TYPE_AUTOSOME = 0 - GENOME_TYPE_X = 1 - GENOME_TYPE_Y = 2 - top_md = tables.metadata - i = top_md["SLiM"]["this_chromosome"]["index"] - # 'chromosomes' is not required so won't be inserted, - # so we won't update it also - assert "chromosomes" not in top_md["SLiM"]["this_chromosome"], ( - "" - "This is an unexpected result: if you hit this, " - "please file a bug at " - "https://github.com/tskit-dev/pyslim." - ) - if gt == GENOME_TYPE_X: - top_md["SLiM"]["this_chromosome"]["type"] = "X" - top_md["SLiM"]["this_chromosome"]["symbol"] = "X" - # if "chromosomes" in tables.metadata['SLiM']: - # top_md['SLiM']['chromosomes'][i]['type'] = "X" - # top_md['SLiM']['chromosomes'][i]['symbol'] = "X" - elif gt == GENOME_TYPE_Y: - top_md["SLiM"]["this_chromosome"]["type"] = "-Y" - top_md["SLiM"]["this_chromosome"]["symbol"] = "Y" - # top_md['SLiM']['chromosomes'][i]['type'] = "Y" - # top_md['SLiM']['chromosomes'][i]['symbol'] = "Y" + new_node_schema = new_schema + if file_version in ("0.1", "0.2", "0.3", "0.4", "0.5", "0.6", "0.7", "0.8"): + # 0.8->0.9 switched from is_null to is_vacant, + # and moved chromosome type from node metadata to top-level + not_vacant = [0] # single chromosome + yes_vacant = [1] + gt = None + for n in nodes: + md = n.metadata + if len(md) > 0: + md["is_vacant"] = yes_vacant if md["is_null"] else not_vacant + if not md["is_null"]: + if gt is None: + gt = md["genome_type"] + else: + assert md["genome_type"] == gt, ( + "Inconsistent tables: " + f"mismatching genome types {gt} and " + f"{md['genome_type']} in node metadata." + ) + del md["is_null"] + del md["genome_type"] + tables.nodes.append(n.replace(metadata=md)) + # flags for node genome type pre-0.9: + # confusingly and sub-optimally, these were redundant: + # all non-null nodes in the same sim would have the same genome type + GENOME_TYPE_AUTOSOME = 0 + GENOME_TYPE_X = 1 + GENOME_TYPE_Y = 2 + top_md = tables.metadata + i = top_md["SLiM"]["this_chromosome"]["index"] + # 'chromosomes' is not required so won't be inserted, + # so we won't update it also + assert "chromosomes" not in top_md["SLiM"]["this_chromosome"], ( + "" + "This is an unexpected result: if you hit this, " + "please file a bug at " + "https://github.com/tskit-dev/pyslim." + ) + if gt == GENOME_TYPE_X: + top_md["SLiM"]["this_chromosome"]["type"] = "X" + top_md["SLiM"]["this_chromosome"]["symbol"] = "X" + # if "chromosomes" in tables.metadata['SLiM']: + # top_md['SLiM']['chromosomes'][i]['type'] = "X" + # top_md['SLiM']['chromosomes'][i]['symbol'] = "X" + elif gt == GENOME_TYPE_Y: + top_md["SLiM"]["this_chromosome"]["type"] = "-Y" + top_md["SLiM"]["this_chromosome"]["symbol"] = "Y" + # top_md['SLiM']['chromosomes'][i]['type'] = "Y" + # top_md['SLiM']['chromosomes'][i]['symbol'] = "Y" + else: + assert gt == GENOME_TYPE_AUTOSOME + tables.metadata = top_md else: - assert gt == GENOME_TYPE_AUTOSOME - tables.metadata = top_md + assert file_version == "0.9" + # just needs recoding (and doesn't really need that, we just + # changed the index of some entries in the schema) + for n in nodes: + tables.nodes.append(n) old_schema = _old_metadata_schema("population", file_version) if old_schema is not None: @@ -1302,15 +2006,33 @@ def update_tables(tables): tables.individuals.clear() if inds.metadata_schema == tskit.MetadataSchema(None): inds.metadata_schema = old_schema - new_schema = slim_metadata_schemas["individual"] + num_traits = len(tables.metadata["SLiM"]["traits"]) + new_schema = slim_individual_metadata_schema(num_traits=num_traits) tables.individuals.metadata_schema = new_schema - defaults = default_slim_metadata("individual") - d = {} - for k in ["pedigree_p1", "pedigree_p2"]: - d[k] = defaults[k] + # new(er) additions are pedigree_pX in 0.7 + # and per_trait in 1.0 + defaults = default_slim_metadata("individual", num_traits=num_traits) for ind in inds: md = ind.metadata - md.update(d) + for k in [ + "pedigree_p1", + "pedigree_p2", + "tag", + "tagF", + "tagL0", + "tagL0_set", + "tagL1", + "tagL1_set", + "tagL2", + "tagL2_set", + "tagL3", + "tagL3_set", + "tagL4", + "tagL4_set", + ]: + md.setdefault(k, defaults[k]) + if "per_trait" not in md: + md["per_trait"] = copy.deepcopy(defaults["per_trait"]) tables.individuals.append(ind.replace(metadata=md)) old_schema = _old_metadata_schema("mutation", file_version) @@ -1321,10 +2043,8 @@ def update_tables(tables): muts.metadata_schema = old_schema tables.mutations.metadata_schema = slim_metadata_schemas["mutation"] for mut in muts: - md = mut.metadata - for ml in md["mutation_list"]: - ml["nucleotide"] = -1 - tables.mutations.append(mut.replace(metadata=md)) + # drop metadata: it should have been copied into top-level above + tables.mutations.append(mut.replace(metadata=None)) if file_version == "0.1": # shift times @@ -1364,7 +2084,6 @@ def update_tables(tables): tskit.validate_provenance(new_record) tables.provenances.add_row(json.dumps(new_record)) - set_metadata_schemas(tables) md = tables.metadata md["SLiM"]["file_version"] = slim_file_version tables.metadata = md diff --git a/pyslim/slim_tree_sequence.py b/pyslim/slim_tree_sequence.py index 95de033..f6bc9c0 100644 --- a/pyslim/slim_tree_sequence.py +++ b/pyslim/slim_tree_sequence.py @@ -10,6 +10,45 @@ def load(*args, **kwargs): raise RuntimeError("This method has been removed: use tskit.load( ) instead.") +def mutation_metadata(ts, check=True, _ts_metadata=None): + """ + Returns a dictionary whose keys are the numeric SLiM IDs of mutations, + and whose values are metadata entries for those mutations. + *Note:* this is indexed by integers, not strings, so if you obtain SLiM IDs + from something like ``mut.derived_state.split(",")``, you must convert + the result to integers before looking up metadata! + + This is a simple extraction function that places the list of metadata entries + stored in ``ts.metadata["SLiM_mutation_list"]`` in a dictionary + indexed by SLiM ID. It is recommended to extract this information once + and use the result in script, because calling this function many times + (or, even just referring to ``ts.metadata`` many times) + can slow down scripts considerably. + + :param tskit.TreeSequence ts: The tree sequence. + + :returns dict: A dictionary of metadata entries, indexed by SLiM ID + and in sorted order by SLiM ID. + """ + if _ts_metadata is None: + _ts_metadata = ts.metadata + # Note that dictionaries preserve insertion order + ml = _ts_metadata["SLiM_mutation_list"] + ml.sort(key=lambda x: x["mutation_id"]) + out = {mut["mutation_id"]: mut for mut in ml} + if check: + ids = {int(j) for x in ts.mutations_derived_state for j in x.split(",")} + for k in ids: + if k not in out: + raise ValueError( + "Top-level mutation metadata is missing " + f"information for mutation ID {k}: " + "do you need to run " + "pyslim.add_mutation_metadata(ts)?" + ) + return out + + def mutation_at(ts, node, position, time=None): """ Finds the mutation present in the genome of ``node`` at ``position``, @@ -59,7 +98,7 @@ def mutation_at(ts, node, position, time=None): return out -def nucleotide_at(ts, node, position, time=None): +def nucleotide_at(ts, node, position, time=None, mut_metadata=None): """ Finds the nucleotide present in the genome of ``node`` at ``position``. Warning: if ``node`` is not actually in the tree sequence (e.g., not @@ -69,20 +108,34 @@ def nucleotide_at(ts, node, position, time=None): at ``position`` inherited by ``node`` that occurred at or before ``time`` ago. + This method uses a dictionary of mutation metadata, computed by + :meth:`mut_metadata`. This step can be expensive if there are + many mutations, so this can be pre-computed and passed in as + ``mutations``. If not provided, it will be computed. + :param int node: The index of a node in the tree sequence. :param float position: A position along the genome. :param int time: The time ago that we want the nucleotide, or None, in which case the ``time`` of ``node`` is used. + :param dict mut_metadata: If provided, a dictionary mapping + mutation ID to metadata, as returned by ``pyslim.mutation_metadata(ts)``. :returns: Index of the nucleotide in ``NUCLEOTIDES`` (0=A, 1=C, 2=G, 3=T). """ if not ts.has_reference_sequence(): raise ValueError("This tree sequence has no reference sequence.") + if mut_metadata is None: + mut_metadata = mutation_metadata(ts) mut_id = mutation_at(ts, node, position, time) if mut_id == tskit.NULL: out = NUCLEOTIDES.index(ts.reference_sequence.data[int(position)]) else: mut = ts.mutation(mut_id) - k = np.argmax([u["slim_time"] for u in mut.metadata["mutation_list"]]) - out = mut.metadata["mutation_list"][k]["nucleotide"] + _, k = max( + [ + (mut_metadata[int(j)]["slim_time"], int(j)) + for j in mut.derived_state.split(",") + ] + ) + out = mut_metadata[k]["nucleotide"] return out diff --git a/tests/__init__.py b/tests/__init__.py index 0996d1a..376b6ce 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -20,6 +20,25 @@ class PyslimTestCase: Base class for test cases in pyslim. """ + def assert_indiv_metadata_equal(self, a, b): + # individual metadata can have nan's in it, thus this function + assert isinstance(a, dict) + assert type(a) == type(b) + assert a.keys() == b.keys() + for k in a: + if k != "per_trait": + assert a[k] == b[k] + apt = a["per_trait"] + bpt = b["per_trait"] + assert len(apt) == len(bpt) + for x, y in zip(apt, bpt): + assert x.keys() == y.keys() + for k in x: + if k == "phenotype": + assert (np.isnan(x[k]) and np.isnan(y[k])) or (x[k] == y[k]) + else: + assert x[k] == y[k] + def verify_haplotype_equality(self, ts, slim_ts): assert ts.num_sites == slim_ts.num_sites for j, v1, v2 in zip(range(ts.num_sites), ts.variants(), slim_ts.variants()): @@ -53,9 +72,8 @@ def assertMetadataEqual(self, t1, t2): assert t1.metadata_schema == t2.metadata_schema assert t1.metadata == t2.metadata # and now check the underlying bytes - # TODO: use the public interface if https://github.com/tskit-dev/tskit/issues/832 happens - md1 = t1._ll_tables.metadata - md2 = t2._ll_tables.metadata + md1 = t1.metadata_bytes + md2 = t2.metadata_bytes assert md1 == md2 def verify_trees_equal(self, ts1, ts2): @@ -76,10 +94,6 @@ def verify_trees_equal(self, ts1, ts2): if n.metadata is not None: map2[n.metadata["slim_id"]] = j assert set(map1.keys()) == set(map2.keys()) - print(ts1) - print(map1) - print(ts2) - print(map2) sids = list(map1.keys()) for sid in sids: n1 = ts1.node(map1[sid]) @@ -88,22 +102,13 @@ def verify_trees_equal(self, ts1, ts2): assert n1.metadata == n2.metadata i1 = ts1.individual(n1.individual) i2 = ts2.individual(n2.individual) - if i1.metadata != i2.metadata: - print("i1: ", i1.metadata) - print("i2: ", i2.metadata) - assert i1.metadata == i2.metadata + self.assert_indiv_metadata_equal(i1.metadata, i2.metadata) for _ in range(10): pos = random.uniform(0, ts1.sequence_length) t1 = ts1.at(pos) t2 = ts2.at(pos) for _ in range(10): a, b = random.choices(sids, k=2) - print(a, b, map1[a], map1[b], map2[a], map2[b]) - print(t1) - print("a", t1.time(map1[a])) - print("b", t1.time(map1[b])) - print("1", t1.tmrca(map1[a], map1[b])) - print("2", t2.tmrca(map2[a], map2[b])) assert t1.tmrca(map1[a], map1[b]) == t2.tmrca(map2[a], map2[b]) def assertTableCollectionsEqual( diff --git a/tests/conftest.py b/tests/conftest.py index 0a3fa87..6895b5e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ import pytest import tskit from filelock import FileLock +from frozendict import frozendict from .recipe_specs import recipe_specs @@ -73,7 +74,7 @@ def load_ts(self, path): c, e = os.path.splitext(cfile) if e == ".trees": out[c] = tskit.load(os.path.join(path, cfile)) - return out + return frozendict(out) def __init__(self, out_dir): # Note: the 'key' below cannot match a 'key' in recipe_specs diff --git a/tests/recipe_specs.py b/tests/recipe_specs.py index 283daff..d31c2f6 100644 --- a/tests/recipe_specs.py +++ b/tests/recipe_specs.py @@ -12,11 +12,16 @@ # retained: has retained individuals # multipop: has more than one population # multichrom: has more than one chromosome +# traits: has more than just the usual trait +# no_simplify: does not run simplify when writing out +# old_mutations: uses addMutation to add back in some previously lost mutations +# record_mutations: whether mutations and substitutions and reference sequence +# are in top-level metadata # long: kinda big # (chromosome type) # All files are of the form `tests/test_recipes/{key}` recipe_specs = { - "recipe_nonWF.slim": {"nonWF": True, "pedigree": True}, + "recipe_nonWF.slim": {"nonWF": True, "pedigree": True, "record_mutations": True}, "recipe_nonWF_X.slim": {"nonWF": True, "pedigree": True, "X": True}, "recipe_nonWF_Y.slim": {"nonWF": True, "pedigree": True, "Y": True}, "recipe_nonWF_H.slim": {"nonWF": True, "pedigree": True, "H": True}, @@ -78,14 +83,23 @@ }, "recipe_long_nonWF.slim": {"nonWF": True, "long": True}, "recipe_old_nonWF.slim": {"nonWF": True, "remembered_first": True}, - "recipe_WF.slim": {"WF": True, "pedigree": True}, + "recipe_WF.slim": {"WF": True, "pedigree": True, "record_mutations": True}, + "recipe_no_simplify.slim": {"WF": True, "no_simplify": True}, "recipe_long_WF.slim": {"WF": True, "long": True}, "recipe_WF_migration.slim": {"WF": True, "pedigree": True, "multipop": True}, - "recipe_nucleotides_WF.slim": {"WF": True, "pedigree": True, "nucleotides": True}, + "recipe_nucleotides_WF.slim": { + "WF": True, + "pedigree": True, + "nucleotides": True, + "record_mutations": True, + "refseq": True, + }, "recipe_nucleotides_nonWF.slim": { "nonWF": True, "pedigree": True, "nucleotides": True, + "record_mutations": True, + "refseq": True, }, "recipe_nucleotides_plus_others.slim": { "WF": True, @@ -177,7 +191,10 @@ "adds_mutations": True, "nucleotides": True, "non-nucleotides": True, + "record_mutations": True, + "refseq": True, }, + "recipe_adds_old_muts.slim": {"WF": True, "old_mutations": True}, "recipe_many_chromosomes.slim": { "nonWF": True, "pedigree": True, @@ -189,6 +206,15 @@ "multichrom": True, "H-": True, }, + "recipe_with_traits.slim": { + "WF": True, + "traits": True, + "multichrom": True, + "begun_late": True, + "X": True, + "Y": True, + "H": True, + }, } for x in ("first", "early", "late"): diff --git a/tests/test_annotation.py b/tests/test_annotation.py index 6336bc4..ead0a04 100644 --- a/tests/test_annotation.py +++ b/tests/test_annotation.py @@ -15,17 +15,33 @@ import pyslim import tests -from .recipe_specs import restarted_recipe_eq +from .recipe_specs import recipe_eq, restarted_recipe_eq def mutcontext(ts): - if ts.num_mutations > 0: + md = ts.metadata + if ( + isinstance(md, dict) + and "SLiM_mutation_list" in md + and len(md["SLiM_mutation_list"]) > 0 + ): handler = pytest.warns(Warning, match="already has.*metadata") else: handler = contextlib.nullcontext() return handler +def canonicalise_tables(tables): + md = tables.metadata + num_muts = len(md["SLiM_mutation_list"]) + md["SLiM_mutation_list"] = list( + pyslim.mutation_metadata(tables, check=False).values() + ) + assert num_muts == len(md["SLiM_mutation_list"]) + tables.metadata = md + tables.canonicalise() + + def verify_slim_restart_equality(in_ts_dict, out_ts_dict, check_prov=True): """ Check for equality, in everything but the last provenance. @@ -37,9 +53,9 @@ def verify_slim_restart_equality(in_ts_dict, out_ts_dict, check_prov=True): if check_prov: assert in_ts.num_provenances + 1 == out_ts.num_provenances in_tables = in_ts.dump_tables() - in_tables.canonicalise() + canonicalise_tables(in_tables) out_tables = out_ts.dump_tables() - out_tables.canonicalise() + canonicalise_tables(out_tables) in_tables.assert_equals(out_tables, ignore_provenance=True) @@ -103,13 +119,14 @@ def verify_defaults(self, ts): Verify the default values have been entered into metadata. """ do_pops = [False for _ in ts.populations()] - for m in ts.mutations(): - md = m.metadata - assert isinstance(md["mutation_list"], list) - for mdl in md["mutation_list"]: - assert mdl["mutation_type"] == 0 - assert mdl["selection_coeff"] == 0.0 - assert mdl["subpopulation"] == tskit.NULL + for mdl in ts.metadata["SLiM_mutation_list"]: + assert mdl["mutation_type"] == 0 + assert mdl["per_trait"] == [ + {"effect_size": 0.0, "dominance": 0.5, "hemizygous_dominance": 1.0} + ] + assert mdl["subpopulation"] == tskit.NULL + # assert mdl["slim_time"] == 0 # this is set to something meaningful + assert mdl["nucleotide"] == -1 for n in ts.nodes(): md = n.metadata if not n.is_sample(): @@ -220,17 +237,17 @@ def verify_remapping(self, ts, rts, subpop_map): assert m == rm assert ts.num_mutations == rts.num_mutations + mut_info = pyslim.mutation_metadata(ts) + rmut_info = pyslim.mutation_metadata(rts) + assert len(mut_info) == len(rmut_info) # mutations may have changed order - tsm = {m.derived_state: m for m in ts.mutations()} - rtsm = {m.derived_state: m for m in rts.mutations()} + tsm = {m.derived_state for m in ts.mutations()} + rtsm = {m.derived_state for m in rts.mutations()} for x in tsm: assert x in rtsm - m = tsm[x] - rm = rtsm[x] - md = m.metadata - rmd = rm.metadata - assert len(md["mutation_list"]) == len(rmd["mutation_list"]) - for x, y in zip(md["mutation_list"], rmd["mutation_list"]): + for sid in x.split(","): + x = mut_info[int(sid)].copy() + y = rmut_info[int(sid)] x["subpopulation"] = fwd_map[x["subpopulation"]] assert x == y @@ -239,7 +256,7 @@ def verify_remapping(self, ts, rts, subpop_map): md = i.metadata rmd = ri.metadata md["subpopulation"] = fwd_map[md["subpopulation"]] - assert md == rmd + self.assert_indiv_metadata_equal(md, rmd) def test_annotate_errors(self, helper_functions): for ts in helper_functions.get_msprime_examples(): @@ -279,6 +296,11 @@ def test_warns_overwriting_mutations(self, helper_functions): ts, rate=1, random_seed=12, model=msprime.SLiMMutationModel(type=1) ) assert ts.num_mutations > 0 + t = ts.dump_tables() + t.metadata_schema = pyslim.slim_metadata_schemas["tree_sequence"] + t.metadata = pyslim.default_slim_metadata("tree_sequence") + ts = t.tree_sequence() + ts = pyslim.add_mutation_metadata(ts) with pytest.warns(Warning, match="already has.*metadata"): slim_ts = pyslim.annotate(ts, model_type="WF", tick=1) @@ -335,45 +357,35 @@ def test_just_simulate(self, helper_functions, tmp_path): def test_basic_annotation(self, helper_functions, tmp_path): for ts in helper_functions.get_msprime_examples(): - for do_mutations in [False, True]: - tick = 4 - cycle = 1 - stage = "late" - if do_mutations: - handler = mutcontext(ts) - else: - handler = contextlib.nullcontext() - with handler: - slim_ts = pyslim.annotate( - ts, - model_type="WF", - tick=tick, - cycle=cycle, - stage=stage, - annotate_mutations=do_mutations, - ) - assert slim_ts.metadata["SLiM"]["model_type"] == "WF" - assert slim_ts.metadata["SLiM"]["tick"] == tick - assert slim_ts.metadata["SLiM"]["cycle"] == cycle - assert slim_ts.metadata["SLiM"]["stage"] == stage - assert ( - slim_ts.metadata["SLiM"]["file_version"] == pyslim.slim_file_version - ) - self.verify_annotated_tables( - ts, slim_ts, check_alleles=(not do_mutations) + tick = 4 + cycle = 1 + stage = "late" + with mutcontext(ts): + slim_ts = pyslim.annotate( + ts, + model_type="WF", + tick=tick, + cycle=cycle, + stage=stage, + annotate_mutations=True, ) - self.verify_annotated_trees(ts, slim_ts) - if not do_mutations: - self.verify_haplotype_equality(ts, slim_ts) - self.verify_defaults(slim_ts) - self.verify_provenance(slim_ts) - # try loading this into SLiM - loaded_ts = helper_functions.run_msprime_restart( - {"default": slim_ts}, tmp_path, multichrom=False, WF=True - )["default"] - self.verify_annotated_tables(loaded_ts, slim_ts) - self.verify_annotated_trees(loaded_ts, slim_ts) - self.verify_haplotype_equality(loaded_ts, slim_ts) + slim_ts_md = slim_ts.metadata["SLiM"] + assert slim_ts_md["model_type"] == "WF" + assert slim_ts_md["tick"] == tick + assert slim_ts_md["cycle"] == cycle + assert slim_ts_md["stage"] == stage + assert slim_ts_md["file_version"] == pyslim.slim_file_version + self.verify_annotated_tables(ts, slim_ts, check_alleles=False) + self.verify_annotated_trees(ts, slim_ts) + self.verify_defaults(slim_ts) + self.verify_provenance(slim_ts) + # try loading this into SLiM + loaded_ts = helper_functions.run_msprime_restart( + {"default": slim_ts}, tmp_path, multichrom=False, WF=True + )["default"] + self.verify_annotated_tables(loaded_ts, slim_ts) + self.verify_annotated_trees(loaded_ts, slim_ts) + self.verify_haplotype_equality(loaded_ts, slim_ts) def test_annotate_refseq(self): ts = msprime.sim_ancestry(2, sequence_length=10, random_seed=77) @@ -447,22 +459,22 @@ def test_annotate_nodes(self, helper_functions): # not testing SLiM because needs annotation of indivs to make sense def test_annotate_mutations(self, helper_functions): + # test workflow of annotating and editing + rng = np.random.default_rng(seed=123) for ts in helper_functions.get_msprime_examples(): with mutcontext(ts): slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1) tables = slim_ts.dump_tables() - metadata = [m.metadata for m in tables.mutations] - selcoefs = [random.uniform(0, 1) for _ in metadata] + md = tables.metadata + metadata = md["SLiM_mutation_list"] + # these are stored as "f" meaning "float" (not "double") + selcoefs = rng.uniform(size=len(metadata)).astype("float32") for j in range(len(metadata)): - metadata[j]["mutation_list"][0]["selection_coeff"] = selcoefs[j] - ms = tables.mutations.metadata_schema - tables.mutations.packset_metadata( - [ms.validate_and_encode_row(r) for r in metadata] - ) + metadata[j]["per_trait"][0]["effect_size"] = selcoefs[j] + tables.metadata = md new_ts = tables.tree_sequence() - for j, x in enumerate(new_ts.mutations()): - md = x.metadata - assert np.isclose(md["mutation_list"][0]["selection_coeff"], selcoefs[j]) + for j, md in enumerate(new_ts.metadata["SLiM_mutation_list"]): + assert md["per_trait"][0]["effect_size"] == selcoefs[j] def test_dont_annotate_mutations(self): # Test the option to not overwrite mutation annotations @@ -828,11 +840,14 @@ def test_remapping(self, helper_functions, tmp_path): random_seed=455, ) ts = pyslim.annotate(ts, model_type="WF", tick=1) - ts = msprime.sim_mutations( - ts, - rate=1e-2, - random_seed=9, - model=msprime.SLiMMutationModel(type=1), + ts = pyslim.add_mutation_metadata( + msprime.sim_mutations( + ts, + rate=1e-2, + random_seed=9, + model=msprime.SLiMMutationModel(type=1), + ), + mutation_type=1, ) assert ts.num_mutations > 0 for subpop_map in ( @@ -947,8 +962,8 @@ def test_reload_annotate(self, restart_name, recipe, helper_functions, tmp_path) in_ts = {} for chrom, ts in recipe["ts"].items(): tables = ts.dump_tables() - metadata = [m.metadata for m in tables.mutations] - has_nucleotides = tables.metadata["SLiM"]["nucleotide_based"] + metadata = tables.metadata + has_nucleotides = metadata["SLiM"]["nucleotide_based"] if has_nucleotides: nucs = [random.choice([0, 1, 2, 3]) for _ in metadata] refseq = "".join( @@ -957,17 +972,13 @@ def test_reload_annotate(self, restart_name, recipe, helper_functions, tmp_path) k=int(ts.sequence_length), ), ) - for n, md in zip(nucs, metadata): - for m in md["mutation_list"]: - m["nucleotide"] = n + for n, md in zip(nucs, metadata["SLiM_mutation_list"]): + md["nucleotide"] = n tables.reference_sequence.data = refseq - for md in metadata: - for m in md["mutation_list"]: - m["selection_coeff"] = random.random() - ms = tables.mutations.metadata_schema - tables.mutations.packset_metadata( - [ms.validate_and_encode_row(r) for r in metadata] - ) + for md in metadata["SLiM_mutation_list"]: + for x in md["per_trait"]: + x["effect_size"] = random.random() + tables.metadata = metadata in_ts[chrom] = tables.tree_sequence() # put it through SLiM (which just reads in and writes out) out_ts = helper_functions.run_slim_restart( @@ -1046,3 +1057,136 @@ def test_restarts_and_runs_simplified( it = in_ts[k] ot = out_ts[k] assert ot.metadata["SLiM"]["tick"] >= it.metadata["SLiM"]["tick"] + + +class TestAddMutationMetadata(tests.PyslimTestCase): + def test_add_mutation_metadata_errors(self): + ts = msprime.sim_ancestry(10, random_seed=5) + ts = msprime.sim_mutations( + ts, rate=5, random_seed=3, model=msprime.SLiMMutationModel(type=0) + ) + # bad metadata schema + t = ts.dump_tables() + t.metadata_schema = tskit.MetadataSchema(None) + bad_ts = t.tree_sequence() + with pytest.raises(ValueError, match="metadata schema is not"): + _ = pyslim.add_mutation_metadata(bad_ts) + t.metadata_schema = tskit.MetadataSchema.permissive_json() + t.metadata = {} + bad_ts = t.tree_sequence() + with pytest.raises(ValueError, match="metadata schema is not"): + _ = pyslim.add_mutation_metadata(bad_ts) + t.metadata["SLiM"] = pyslim.default_slim_metadata("tree_sequence")["SLiM"] + bad_ts = t.tree_sequence() + with pytest.raises(ValueError, match="metadata schema is not"): + _ = pyslim.add_mutation_metadata(bad_ts) + + def test_add_mutation_metadata_mutation_type(self): + ts = msprime.sim_ancestry(10, random_seed=5) + ts = msprime.sim_mutations( + ts, rate=5, random_seed=3, model=msprime.SLiMMutationModel(type=0) + ) + t = ts.dump_tables() + t.metadata_schema = pyslim.slim_metadata_schemas["tree_sequence"] + t.metadata = pyslim.default_slim_metadata("tree_sequence") + ts = t.tree_sequence() + for k in (0, 1, 5): + new_ts = pyslim.add_mutation_metadata(ts, mutation_type=k) + for x in new_ts.metadata["SLiM_mutation_list"]: + assert x["mutation_type"] == k + + @pytest.mark.parametrize( + "recipe", + recipe_eq( + exclude=( + "long", + "no_simplify", + "multichrom", + "everyone", + "init_mutated", + "old_mutations", + ) + ), + indirect=True, + ) + def test_add_mutation_metadata(self, recipe): + for _, ts in recipe["ts"].items(): + tables = ts.dump_tables() + md = tables.metadata + md["SLiM_mutation_list"] = [] + tables.metadata = md + new_ts = pyslim.add_mutation_metadata( + tables.tree_sequence(), + ) + old_metadata = pyslim.mutation_metadata(ts) + new_metadata = pyslim.mutation_metadata(new_ts) + assert len(old_metadata) == len(new_metadata) + for k in old_metadata: + assert 0 == new_metadata[k]["mutation_type"] # default + assert old_metadata[k]["slim_time"] == new_metadata[k]["slim_time"] + assert old_metadata[k]["mutation_id"] == new_metadata[k]["mutation_id"] + + @pytest.mark.parametrize( + "recipe", + recipe_eq( + exclude=( + "long", + "no_simplify", + "multichrom", + "everyone", + "init_mutated", + "old_mutations", + ) + ), + indirect=True, + ) + def test_add_mutation_metadata_keeps(self, recipe): + for _, ts in recipe["ts"].items(): + if ts.num_mutations > 15: + tables = ts.dump_tables() + md = tables.metadata + metadata = md["SLiM_mutation_list"] + del metadata[:2] + del metadata[10:] + kept_ids = {m["mutation_id"] for m in metadata} + md["SLiM_mutation_list"] = metadata + tables.metadata = md + new_ts = pyslim.add_mutation_metadata( + tables.tree_sequence(), + ) + old_metadata = pyslim.mutation_metadata(ts) + new_metadata = pyslim.mutation_metadata(new_ts) + assert len(old_metadata) == len(new_metadata) + for k in old_metadata: + if k in kept_ids: + assert old_metadata[k] == new_metadata[k] + else: + assert 0 == new_metadata[k]["mutation_type"] # default + assert ( + old_metadata[k]["slim_time"] == new_metadata[k]["slim_time"] + ) + assert ( + old_metadata[k]["mutation_id"] + == new_metadata[k]["mutation_id"] + ) + + def test_add_mutation_metadata_removes(self): + ts = msprime.sim_ancestry(10, random_seed=5) + ts = msprime.sim_mutations( + ts, rate=5, random_seed=3, model=msprime.SLiMMutationModel(type=0) + ) + ts = pyslim.add_mutation_metadata(pyslim.annotate(ts, model_type="WF", tick=1)) + sts = ts.simplify([0, 1]) + assert sts.num_mutations < ts.num_mutations + assert ts.metadata["SLiM_mutation_list"] == sts.metadata["SLiM_mutation_list"] + nsts = pyslim.add_mutation_metadata(sts, remove_unused=True) + mut_info = pyslim.mutation_metadata(ts) + nmut_info = pyslim.mutation_metadata(nsts) + mut_ids = np.unique( + [int(k) for mut in nsts.mutations() for k in mut.derived_state.split(",")] + ) + assert len(mut_ids) == len(nsts.metadata["SLiM_mutation_list"]) + for k in mut_ids: + assert k in nmut_info + assert k in mut_info + assert mut_info[k] == nmut_info[k] diff --git a/tests/test_metadata.py b/tests/test_metadata.py index b0a4433..e472e8f 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -12,6 +12,20 @@ from .recipe_specs import recipe_eq +def assert_nan_equal(a, b): + if isinstance(a, float): + assert (a == b) or (np.isnan(a) and np.isnan(b)) + elif isinstance(a, dict) and isinstance(b, dict): + assert a.keys() == b.keys() + for k in a: + assert_nan_equal(a[k], b[k]) + elif isinstance(a, list) and isinstance(b, list): + for x, y in zip(a, b, strict=True): + assert_nan_equal(x, y) + else: + assert a == b + + class TestMetadataSchemas(tests.PyslimTestCase): def validate_table_metadata(self, table): ms = table.metadata_schema @@ -46,35 +60,47 @@ def test_default_metadata(self): schema = pyslim.slim_metadata_schemas[k] entry = pyslim.default_slim_metadata(k) sd = schema.asdict() - if sd is not None: - for p in sd["properties"]: - assert p in entry - encoded = schema.validate_and_encode_row(entry) - decoded = schema.decode_row(encoded) - if entry is None: - assert decoded is None + if k != "tree_sequence": + if sd is not None: + for p in sd["properties"]: + assert p in entry + encoded = schema.validate_and_encode_row(entry) + decoded = schema.decode_row(encoded) + if entry is None: + assert decoded is None + else: + # some defaults have nans, which are not equal + assert_nan_equal(entry, decoded) else: + assert k == "tree_sequence" + for p in sd["json"]["properties"]: + assert p in entry + encoded = schema.validate_and_encode_row(entry) + decoded = schema.decode_row(encoded) + assert entry == decoded + entry["SLiM_mutation_list"].append( + pyslim.default_slim_metadata("mutation_list_entry") + ) + encoded = schema.validate_and_encode_row(entry) + decoded = schema.decode_row(encoded) + assert entry == decoded + entry["SLiM_mutation_list"].append( + pyslim.default_slim_metadata("mutation_list_entry") + ) + encoded = schema.validate_and_encode_row(entry) + decoded = schema.decode_row(encoded) assert entry == decoded - schema = pyslim.slim_metadata_schemas["mutation"] - entry = pyslim.default_slim_metadata("mutation") - entry["mutation_list"].append( - pyslim.default_slim_metadata("mutation_list_entry") - ) - encoded = schema.validate_and_encode_row(entry) - decoded = schema.decode_row(encoded) - assert entry == decoded - entry["mutation_list"].append( - pyslim.default_slim_metadata("mutation_list_entry") - ) - encoded = schema.validate_and_encode_row(entry) - decoded = schema.decode_row(encoded) - assert entry == decoded def test_slim_metadata_schema_equality(self, recipe): num_chromosomes = len(recipe["ts"]) for ts in recipe["ts"].values(): t = ts.dump_tables() - assert t.metadata_schema == pyslim.slim_metadata_schemas["tree_sequence"] + num_traits = len(t.metadata["SLiM"]["traits"]) + ts_schema = pyslim.slim_metadata_schemas["tree_sequence"].asdict() + ts_schema["struct"]["properties"]["SLiM_mutation_list"]["items"][ + "properties" + ]["per_trait"]["length"] = num_traits + assert t.metadata_schema.asdict() == ts_schema assert t.edges.metadata_schema == pyslim.slim_metadata_schemas["edge"] assert t.sites.metadata_schema == pyslim.slim_metadata_schemas["site"] assert ( @@ -85,10 +111,9 @@ def test_slim_metadata_schema_equality(self, recipe): (num_chromosomes + 7) / 8 ) assert t.nodes.metadata_schema.asdict() == node_schema - assert ( - t.individuals.metadata_schema - == pyslim.slim_metadata_schemas["individual"] - ) + ind_schema = pyslim.slim_metadata_schemas["individual"].asdict() + ind_schema["properties"]["per_trait"]["length"] = num_traits + assert t.individuals.metadata_schema.asdict() == ind_schema assert ( t.populations.metadata_schema == pyslim.slim_metadata_schemas["population"] @@ -117,18 +142,25 @@ class TestTreeSequenceMetadata(tests.PyslimTestCase): def validate_slim_metadata(self, t): # t could be tables or a tree sequence schema = t.metadata_schema.schema - assert "SLiM" in schema["properties"] - assert "SLiM" in t.metadata + assert schema["codec"] == "json+struct" + assert "SLiM" in schema["json"]["properties"] + tmd = t.metadata + assert "SLiM" in tmd for k in pyslim.default_slim_metadata("tree_sequence")["SLiM"]: - assert k in schema["properties"]["SLiM"]["properties"] - assert k in t.metadata["SLiM"] + assert k in schema["json"]["properties"]["SLiM"]["properties"] + assert k in tmd["SLiM"] + sml = schema["struct"]["properties"] + assert "SLiM_mutation_list" in sml + for k in pyslim.default_slim_metadata("mutation_list_entry"): + assert k in sml["SLiM_mutation_list"]["items"]["properties"] def validate_model_type(self, tsdict, model_type): for _, ts in tsdict.items(): - assert ts.metadata["SLiM"]["file_version"] == pyslim.slim_file_version - assert ts.metadata["SLiM"]["model_type"] == model_type - assert ts.metadata["SLiM"]["tick"] > 0 - assert ts.metadata["SLiM"]["tick"] >= np.max(ts.tables.nodes.time) + md = ts.metadata + assert md["SLiM"]["file_version"] == pyslim.slim_file_version + assert md["SLiM"]["model_type"] == model_type + assert md["SLiM"]["tick"] > 0 + assert md["SLiM"]["tick"] >= np.max(ts.tables.nodes.time) @pytest.mark.parametrize("recipe", arbitrary_recipe, indirect=True) def test_set_tree_sequence_metadata_errors(self, recipe): @@ -156,19 +188,59 @@ def test_set_tree_sequence_metadata_keeps(self, recipe): tables.metadata = dummy_metadata pyslim.set_tree_sequence_metadata(tables, "nonWF", 0) schema = tables.metadata_schema.schema + tmd = tables.metadata for k in dummy_metadata: if len(x) > 0: - assert k in schema["properties"] - assert k in tables.metadata - assert tables.metadata[k] == dummy_metadata[k] + assert k in schema["json"]["properties"] + assert k in tmd + assert tmd[k] == dummy_metadata[k] self.validate_slim_metadata(tables) - assert tables.metadata["SLiM"]["model_type"] == "nonWF" - assert tables.metadata["SLiM"]["tick"] == 0 + assert tmd["SLiM"]["model_type"] == "nonWF" + assert tmd["SLiM"]["tick"] == 0 + + @pytest.mark.parametrize("recipe", arbitrary_recipe, indirect=True) + def test_set_tree_sequence_metadata_keeps_struct(self, recipe): + # make sure doesn't overwrite other stuff + ts = list(recipe["ts"].values())[0] + json_props = {"num": {"type": "number"}} + struct_props = {"binnum": {"type": "integer", "binaryFormat": "i", "default": 0}} + schema_dict = { + "codec": "json+struct", + "type": "object", + "json": {"codec": "json", "type": "object", "properties": json_props}, + "struct": {"struct": "json", "type": "object", "properties": struct_props}, + } + dummy_schema = tskit.MetadataSchema(schema_dict) + dummy_metadata = {"num": 12, "binnum": 42} + tables = ts.dump_tables() + tables.metadata_schema = dummy_schema + tables.metadata = dummy_metadata + pyslim.set_tree_sequence_metadata(tables, "nonWF", 0) + schema = tables.metadata_schema.schema + tmd = tables.metadata + for k in dummy_metadata: + if k in json_props: + assert schema["json"]["properties"][k] == json_props[k] + else: + assert schema["struct"]["properties"][k] == struct_props[k] + assert k in tmd + assert tmd[k] == dummy_metadata[k] + self.validate_slim_metadata(tables) + assert tmd["SLiM"]["model_type"] == "nonWF" + assert tmd["SLiM"]["tick"] == 0 @pytest.mark.parametrize("recipe", arbitrary_recipe, indirect=True) def test_set_tree_sequence_metadata(self, recipe): ts = list(recipe["ts"].values())[0] tables = ts.dump_tables() + chroms = [ + {"id": 1, "name": "autosome_1", "symbol": "1", "type": "A", "index": 0}, + {"id": 35, "name": "mtDNA", "symbol": "MT", "type": "HF", "index": 1}, + ] + traits = [ + {"index": 0, "name": "theTrait", "type": "additive"}, + {"index": 1, "name": "perfectness", "type": "multiplicative"}, + ] pyslim.set_tree_sequence_metadata( tables, "WF", @@ -179,16 +251,23 @@ def test_set_tree_sequence_metadata(self, recipe): spatial_periodicity="y", separate_sexes=False, nucleotide_based=True, + this_chromosome=chroms[1], + chromosomes=chroms, + traits=traits, ) self.validate_slim_metadata(tables) - assert tables.metadata["SLiM"]["model_type"] == "WF" - assert tables.metadata["SLiM"]["tick"] == 99 - assert tables.metadata["SLiM"]["cycle"] == 40 - assert tables.metadata["SLiM"]["stage"] == "early" - assert tables.metadata["SLiM"]["spatial_dimensionality"] == "xy" - assert tables.metadata["SLiM"]["spatial_periodicity"] == "y" - assert tables.metadata["SLiM"]["separate_sexes"] == False - assert tables.metadata["SLiM"]["nucleotide_based"] == True + tmd = tables.metadata + assert tmd["SLiM"]["model_type"] == "WF" + assert tmd["SLiM"]["tick"] == 99 + assert tmd["SLiM"]["cycle"] == 40 + assert tmd["SLiM"]["stage"] == "early" + assert tmd["SLiM"]["spatial_dimensionality"] == "xy" + assert tmd["SLiM"]["spatial_periodicity"] == "y" + assert tmd["SLiM"]["separate_sexes"] == False + assert tmd["SLiM"]["nucleotide_based"] == True + assert tmd["SLiM"]["chromosomes"] == chroms + assert tmd["SLiM"]["this_chromosome"] == chroms[1] + assert tmd["SLiM"]["traits"] == traits @pytest.mark.parametrize("recipe", recipe_eq("WF"), indirect=True) def test_WF_model_type(self, recipe): @@ -199,12 +278,12 @@ def test_nonWF_model_type(self, recipe): self.validate_model_type(recipe["ts"], "nonWF") @pytest.mark.parametrize( - "recipe", recipe_eq(exclude=["user_metadata", "multichrom"]), indirect=True + "recipe", + recipe_eq(exclude=["user_metadata", "multichrom", "record_mutations"]), + indirect=True, ) def test_recover_metadata(self, recipe): # msprime <=0.7.5 discards metadata, but we can recover it from provenance - # HOWEVER: multichromosome information is not saved - # but this is not something we need to maintain any more. for _, ts in recipe["ts"].items(): tables = ts.dump_tables() tables.metadata_schema = tskit.MetadataSchema(None) @@ -212,13 +291,12 @@ def test_recover_metadata(self, recipe): pyslim.update_tables(tables) md = tables.metadata assert "SLiM" in md - for k in ts.metadata["SLiM"]: - if k in ("chromosomes", "this_chromosome"): - continue # TODO: see https://github.com/MesserLab/SLiM/issues/520 + tsmd = ts.metadata["SLiM"] + for k in tsmd: assert k in md["SLiM"] # slim does not write out empty descriptions - if k != "description" or ts.metadata["SLiM"][k] != "": - assert ts.metadata["SLiM"][k] == md["SLiM"][k] + if k != "description" or tsmd[k] != "": + assert tsmd[k] == md["SLiM"][k] @pytest.mark.parametrize( "recipe", recipe_eq("recipe_with_metadata.slim"), indirect=True @@ -267,10 +345,9 @@ def test_nucleotides(self, recipe): -1, 0, 1, 2, or 3. """ for _, ts in recipe["ts"].items(): - for mut in ts.mutations(): - for u in mut.metadata["mutation_list"]: - assert u["nucleotide"] >= -1 - assert u["nucleotide"] <= 3 + for u in ts.metadata["SLiM_mutation_list"]: + assert u["nucleotide"] >= -1 + assert u["nucleotide"] <= 3 class TestMultichrom(tests.PyslimTestCase): diff --git a/tests/test_provenance.py b/tests/test_provenance.py index d95b275..71568b2 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -122,6 +122,18 @@ old_provenance_examples = [_slim_v3_0_example, _slim_v3_1_example, _slim_v3_3_1_example] +def yield_ts(path): + out = {} + if os.path.isfile(path): + yield tskit.load(path) + elif os.path.isdir(path): + chroms = os.listdir(path) + for cfile in os.listdir(path): + _, e = os.path.splitext(cfile) + if e == ".trees": + yield tskit.load(os.path.join(path, cfile)) + + class TestProvenance(tests.PyslimTestCase): script_dir = os.path.dirname(os.path.realpath(__file__)) @@ -183,6 +195,18 @@ def get_0_8_slim_examples(self): ]: yield tskit.load(filename) + def get_0_9_slim_examples(self): + for filename in [ + os.path.join(self.script_dir, "test_recipes", "recipe_WF.v5.2.trees"), + os.path.join(self.script_dir, "test_recipes", "recipe_WF_X.v5.2.trees"), + os.path.join(self.script_dir, "test_recipes", "recipe_WF_Y.v5.2.trees"), + os.path.join(self.script_dir, "test_recipes", "recipe_nonWF.v5.2.trees"), + os.path.join( + self.script_dir, "test_recipes", "recipe_WF_many_chromosomes.v5.2.trees" + ), + ]: + yield from yield_ts(filename) + def get_mixed_slim_examples(self): for filename in [ os.path.join( @@ -209,6 +233,115 @@ def verify_upgrade(self, ts): for x in t: _ = ms.validate_and_encode_row(x.metadata) + def verify_consistency(self, ts, pts, file_version): + # Check for stuff we know should be copied over verbatim + # 0.1-0.4 we had no metadata schemas; we could pull those old ones + # from slim_metadata.py but we're not + self.verify_top_level_consistency(ts, pts, file_version) + self.verify_nodes_consistency(ts, pts, file_version) + self.verify_edges_consistency(ts, pts, file_version) + self.verify_sites_consistency(ts, pts, file_version) + self.verify_mutations_consistency(ts, pts, file_version) + self.verify_individuals_consistency(ts, pts, file_version) + self.verify_populations_consistency(ts, pts, file_version) + + def verify_top_level_consistency(self, ts, pts, file_version): + # 0.1-0.7: + # model_type, generation, spatial_dimesionality, spatial_periodicity, + # separate_sexes, nucleotide_based + # 0.8: + # changed generation to tick + # 0.9: + # added this_chromosome + # 1.0: + # added traits + if file_version not in ("0.1", "0.2", "0.3", "0.4"): + md = ts.metadata["SLiM"] + pmd = pts.metadata["SLiM"] + for k in ( + "model_type", + "spatial_dimensionality", + "spatial_periodicity", + "separate_sexes", + "nucleotide_based", + ): + assert md[k] == pmd[k] + k = pk = "tick" + if file_version in ("0.5", "0.6", "0.7"): + k = "generation" + assert md[k] == pmd[pk] + k = "this_chromosome" + if file_version == "0.9": + assert md[k] == pmd[k] + + def verify_edges_consistency(self, ts, pts, file_version): + # no metadata + ts.tables.edges.assert_equals(pts.tables.edges, ignore_metadata=True) + + def verify_sites_consistency(self, ts, pts, file_version): + # no metadata + ts.tables.sites.assert_equals(pts.tables.sites, ignore_metadata=True) + + def verify_mutations_consistency(self, ts, pts, file_version): + ts.tables.mutations.assert_equals(pts.tables.mutations, ignore_metadata=True) + # As of 1.0, metadata moved to top level + if file_version not in ("0.1", "0.2", "0.3", "0.4"): + ptsmd = pts.metadata + num_traits = len(ptsmd["SLiM"]["traits"]) + mut_info = {x["mutation_id"]: x for x in ptsmd["SLiM_mutation_list"]} + for mut in ts.mutations(): + for sid, md in zip( + mut.derived_state.split(","), mut.metadata["mutation_list"] + ): + assert int(sid) in mut_info + mi = mut_info[int(sid)] + for k in ("mutation_type", "subpopulation", "slim_time"): + assert mi[k] == md[k] + assert len(mi["per_trait"]) == 1 + if "nucleotide" in md: + assert mi["nucleotide"] == md["nucleotide"] + assert len(mi["per_trait"]) == num_traits + # we're only converting from single-trait slim so far + assert num_traits == 1 + assert mi["per_trait"][0]["effect_size"] == md["selection_coeff"] + + def verify_individuals_consistency(self, ts, pts, file_version): + ts.tables.individuals.assert_equals(pts.tables.individuals, ignore_metadata=True) + # This has: + # pedigree_id, age, subpopulation, sex, flags + # Starting in 0.7 also: + # pedigree_p1, pedigree_p2 + # Starting in 1.0: per_trait + for a, b in zip(ts.individuals(), pts.individuals()): + if file_version not in ("0.1", "0.2", "0.3", "0.4"): + for k in ("pedigree_id", "age", "subpopulation", "sex", "flags"): + assert a.metadata[k] == b.metadata[k] + if file_version not in ("0.5", "0.6"): + for k in ("pedigree_p1", "pedigree_p2"): + assert a.metadata[k] == b.metadata[k] + + def verify_nodes_consistency(self, ts, pts, file_version): + # 0.1-0.8: had slim_id, is_null, genome_type + # 0.9: removed genome_type + # and changed is_null to is_vacant + # 1.0: same as 0.9 but changed some indexes + if file_version != "0.1": + # 0.1 had a shift in time we're not checking here + ts.tables.nodes.assert_equals(pts.tables.nodes, ignore_metadata=True) + for a, b in zip(ts.nodes(), pts.nodes()): + if file_version in ("0.5", "0.6", "0.7", "0.8"): + assert a.metadata["slim_id"] == b.metadata["slim_id"] + elif file_version == "0.9": + assert a.metadata == b.metadata + + def verify_populations_consistency(self, ts, pts, file_version): + ts.tables.populations.assert_equals(pts.tables.populations, ignore_metadata=True) + # This has a whole bunch of things, none of which are required. + for a, b in zip(ts.individuals(), pts.individuals()): + if file_version not in ("0.1", "0.2", "0.3", "0.4"): + for k in a.metadata: + assert a.metadata[k] == b.metadata[k] + def test_convert_0_1_files(self): for ts in self.get_0_1_slim_examples(): assert not pyslim.is_current_version(ts) @@ -216,24 +349,16 @@ def test_convert_0_1_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.1") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert record["model_type"] == pts.metadata["SLiM"]["model_type"] - assert record["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_2_files(self): for ts in self.get_0_2_slim_examples(): @@ -242,26 +367,16 @@ def test_convert_0_2_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.2") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_3_files(self): for ts in self.get_0_3_slim_examples(): @@ -270,26 +385,16 @@ def test_convert_0_3_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.3") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_4_files(self): # Note that with version 0.5 and above, we *don't* get information from @@ -300,26 +405,16 @@ def test_convert_0_4_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.4") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_5_files(self): for ts in self.get_0_5_slim_examples(): @@ -328,26 +423,16 @@ def test_convert_0_5_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.5") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_6_files(self): for ts in self.get_0_6_slim_examples(): @@ -356,26 +441,16 @@ def test_convert_0_6_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.6") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_7_files(self): for ts in self.get_0_7_slim_examples(): @@ -384,26 +459,16 @@ def test_convert_0_7_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.7") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] - assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) - samples = list(ts.samples()) - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] def test_convert_0_8_files(self): for ts in self.get_0_8_slim_examples(): @@ -412,18 +477,17 @@ def test_convert_0_8_files(self): pts = pyslim.update(ts) assert pyslim.is_current_version(pts) self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.8") assert ts.num_provenances == 1 assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["tick"] == pts.metadata["SLiM"]["tick"] + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["tick"] == ptsmd["SLiM"]["tick"] assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) samples = list(ts.samples()) genome_type = None for n in samples: @@ -432,7 +496,7 @@ def test_convert_0_8_files(self): genome_type = md["genome_type"] break assert genome_type is not None - chromosome_type = pts.metadata["SLiM"]["this_chromosome"]["type"] + chromosome_type = ptsmd["SLiM"]["this_chromosome"]["type"] GENOME_TYPE_AUTOSOME = 0 GENOME_TYPE_X = 1 GENOME_TYPE_Y = 2 @@ -442,13 +506,24 @@ def test_convert_0_8_files(self): assert chromosome_type == "X" elif genome_type == GENOME_TYPE_Y: assert chromosome_type == "-Y" - t = ts.first() - pt = pts.first() - for _ in range(20): - u = random.sample(samples, 1)[0] - assert t.parent(u) == pt.parent(u) - if t.parent(u) != tskit.NULL: - assert t.branch_length(u) == pt.branch_length(u) + + def test_convert_0_9_files(self): + for ts in self.get_0_9_slim_examples(): + assert not pyslim.is_current_version(ts) + with pytest.warns(Warning): + pts = pyslim.update(ts) + assert pyslim.is_current_version(pts) + self.verify_upgrade(pts) + self.verify_consistency(ts, pts, file_version="0.9") + assert ts.num_provenances == 1 + assert pts.num_provenances == 2 + assert ts.provenance(0).record == pts.provenance(0).record + record = json.loads(ts.provenance(0).record) + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["tick"] == ptsmd["SLiM"]["tick"] def test_convert_mixed_files(self): for ts in self.get_mixed_slim_examples(): @@ -461,14 +536,13 @@ def test_convert_mixed_files(self): assert pts.num_provenances == 2 assert ts.provenance(0).record == pts.provenance(0).record record = json.loads(ts.provenance(0).record) - assert isinstance(pts.metadata, dict) - assert "SLiM" in pts.metadata - assert ( - record["parameters"]["model_type"] == pts.metadata["SLiM"]["model_type"] - ) - assert record["slim"]["generation"] == pts.metadata["SLiM"]["tick"] + ptsmd = pts.metadata + assert isinstance(ptsmd, dict) + assert "SLiM" in ptsmd + assert record["parameters"]["model_type"] == ptsmd["SLiM"]["model_type"] + assert record["slim"]["generation"] == ptsmd["SLiM"]["tick"] assert list(ts.samples()) == list(pts.samples()) - assert np.array_equal(ts.tables.nodes.flags, pts.tables.nodes.flags) + assert np.array_equal(ts.tables.nodes.flags, pts.nodes_flags) samples = list(ts.samples()) t = ts.first() pt = pts.first() diff --git a/tests/test_recipes/make_v3_tests.sh b/tests/test_recipes/make_old_file_versions.sh similarity index 80% rename from tests/test_recipes/make_v3_tests.sh rename to tests/test_recipes/make_old_file_versions.sh index eb88af9..ab6b48f 100644 --- a/tests/test_recipes/make_v3_tests.sh +++ b/tests/test_recipes/make_old_file_versions.sh @@ -95,4 +95,20 @@ $SLIMDIR/slim recipe_WF_X.slim && mv out.trees recipe_WF_X.${TAG}.trees $SLIMDIR/slim recipe_WF_Y.slim && mv out.trees recipe_WF_Y.${TAG}.trees git add -f recipe_nonWF.${TAG}.trees recipe_WF.${TAG}.trees recipe_WF_X.${TAG}.trees recipe_WF_Y.${TAG}.trees +# To make the v5.2 files: + +TAG=v5.2 +git checkout $TAG +mkdir -p build_$TAG && cd build_$TAG +cmake .. && make +SLIMDIR=$(pwd) +cd ../.. +$SLIMDIR/slim recipe_nonWF.slim && mv out.trees recipe_nonWF.${TAG}.trees +$SLIMDIR/slim recipe_WF.slim && mv out.trees recipe_WF.${TAG}.trees +$SLIMDIR/slim recipe_WF_X.slim && mv out.trees recipe_WF_X.${TAG}.trees +$SLIMDIR/slim recipe_WF_Y.slim && mv out.trees recipe_WF_Y.${TAG}.trees +$SLIMDIR/slim recipe_all_the_chromosome_types.slim && mv out.trees recipe_all_the_chromosome_types.${TAG}.trees +git add -f recipe_nonWF.${TAG}.trees recipe_WF.${TAG}.trees recipe_WF_X.${TAG}.trees recipe_WF_Y.${TAG}.trees recipe_all_the_chromosome_types.${TAG}.trees + + diff --git a/tests/test_recipes/recipe_WF.slim b/tests/test_recipes/recipe_WF.slim index 2503236..1170250 100644 --- a/tests/test_recipes/recipe_WF.slim +++ b/tests/test_recipes/recipe_WF.slim @@ -10,14 +10,49 @@ initialize() initializeGenomicElementType("g1", m1, 1.0); initializeGenomicElement(g1, 0, 99); initializeRecombinationRate(1e-2); + defineGlobal("MD", Dictionary()); } 1 early() { sim.addSubpop("p1", 10); } +// MUTATION/GENOTYPE INFO +1 first() { // mutation information + MD.setValue("mutations", Dictionary()); +} +mutation() { + muts = MD.getValue("mutations"); + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + return T; +} +10 late() { + subs = Dictionary(); + for (mut in sim.substitutions) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "fixationTick", mut.fixationTick + ); + subs.setValue(asString(mut.id), m); + } + MD.setValue("substitutions", subs); +} + +// OUTPUT/FINISH 10 late() { - sim.treeSeqOutput(TREES_FILE); + sim.treeSeqOutput(TREES_FILE, metadata=MD); catn("Done."); sim.simulationFinished(); } diff --git a/tests/test_recipes/recipe_WF.v5.2.trees b/tests/test_recipes/recipe_WF.v5.2.trees new file mode 100644 index 0000000..936f1cd Binary files /dev/null and b/tests/test_recipes/recipe_WF.v5.2.trees differ diff --git a/tests/test_recipes/recipe_WF_X.v5.2.trees b/tests/test_recipes/recipe_WF_X.v5.2.trees new file mode 100644 index 0000000..d09e122 Binary files /dev/null and b/tests/test_recipes/recipe_WF_X.v5.2.trees differ diff --git a/tests/test_recipes/recipe_WF_Y.v5.2.trees b/tests/test_recipes/recipe_WF_Y.v5.2.trees new file mode 100644 index 0000000..e1f3309 Binary files /dev/null and b/tests/test_recipes/recipe_WF_Y.v5.2.trees differ diff --git a/tests/test_recipes/recipe_adds_old_muts.slim b/tests/test_recipes/recipe_adds_old_muts.slim new file mode 100644 index 0000000..a148a6b --- /dev/null +++ b/tests/test_recipes/recipe_adds_old_muts.slim @@ -0,0 +1,30 @@ +initialize() +{ + setSeed(23); + if (!exists("TREES_FILE")) defineGlobal("TREES_FILE", "out.trees"); + initializeSLiMOptions(keepPedigrees=T); + initializeTreeSeq(timeUnit="generations"); + initializeMutationRate(1e-2); + initializeMutationType("m1", 0.5, "f", -0.1); + initializeGenomicElementType("g1", m1, 1.0); + initializeGenomicElement(g1, 0, 99); + initializeRecombinationRate(1e-2); + // record mutations and add them back at the end + defineGlobal("MUTS", c()); +} + +1 early() { + sim.addSubpop("p1", 10); +} + +early() { + defineGlobal("MUTS", c(MUTS, sim.mutations)); +} + +10 late() { + ind = p1.individuals[0]; + ind.haploidGenome1.addMutations(MUTS); + sim.treeSeqOutput(TREES_FILE); + catn("Done."); + sim.simulationFinished(); +} diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_A.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_A.trees new file mode 100644 index 0000000..3df791e Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_A.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_FL.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_FL.trees new file mode 100644 index 0000000..3d3fe25 Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_FL.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_H.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_H.trees new file mode 100644 index 0000000..a0fdca8 Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_H.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HF.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HF.trees new file mode 100644 index 0000000..473023a Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HF.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HM.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HM.trees new file mode 100644 index 0000000..1fb4b3f Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_HM.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_ML.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_ML.trees new file mode 100644 index 0000000..09bd52f Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_ML.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_W.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_W.trees new file mode 100644 index 0000000..740a592 Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_W.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_X.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_X.trees new file mode 100644 index 0000000..b70be3f Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_X.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Y.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Y.trees new file mode 100644 index 0000000..78604af Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Y.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Z.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Z.trees new file mode 100644 index 0000000..252e442 Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_Z.trees differ diff --git a/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_nY.trees b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_nY.trees new file mode 100644 index 0000000..b599982 Binary files /dev/null and b/tests/test_recipes/recipe_all_the_chromosome_types.v5.2.trees/chromosome_nY.trees differ diff --git a/tests/test_recipes/recipe_chromosomes_adds_muts.slim b/tests/test_recipes/recipe_chromosomes_adds_muts.slim index e53e733..2ffd329 100644 --- a/tests/test_recipes/recipe_chromosomes_adds_muts.slim +++ b/tests/test_recipes/recipe_chromosomes_adds_muts.slim @@ -20,10 +20,11 @@ initialize() { for (id in ids, symbol in symbols, type in types) { initializeChromosome(id, length, type, symbol); - initializeAncestralNucleotides(paste0(rep("A", length))); + initializeAncestralNucleotides(randomNucleotides(length)); initializeRecombinationRate(1e-5); initializeGenomicElement(g1, 0, length-1); } + defineGlobal("MD", Dictionary()); } 1 early() { sim.addSubpop("p1", 10); @@ -34,13 +35,68 @@ initialize() { inds = p1.individuals; haps = inds.haplosomesForChromosomes(chrom, includeNulls=F); if (length(haps) > 0) { - sample(haps, 1 + asInteger(length(haps)/2)).addNewDrawnMutation(m2, rdunif(10, 0, chrom.length-1)); + mutl = sample(haps, 1 + asInteger(length(haps)/2)).addNewDrawnMutation(m2, rdunif(10, 0, chrom.length-1)); + muts = MD.getValue("mutations"); + for (mut in mutl) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + } } } } + +// MUTATION/GENOTYPE INFO +1 first() { // reference sequences + refseqs = Dictionary(); + for (chrom in sim.chromosomes) { + refseqs.setValue("chr" + chrom.id, chrom.ancestralNucleotides()); + } + MD.setValue("reference_sequence", refseqs); +} + +1 first() { // mutation information + MD.setValue("mutations", Dictionary()); +} +mutation() { + muts = MD.getValue("mutations"); + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + return T; +} +10 late() { + subs = Dictionary(); + for (mut in sim.substitutions) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "fixationTick", mut.fixationTick + ); + subs.setValue(asString(mut.id), m); + } + MD.setValue("substitutions", subs); +} + +// OUTPUT/FINISH 10 late() { - sim.treeSeqOutput(TREES_FILE); + sim.treeSeqOutput(TREES_FILE, metadata=MD); catn("Done."); sim.simulationFinished(); } diff --git a/tests/test_recipes/recipe_no_simplify.slim b/tests/test_recipes/recipe_no_simplify.slim new file mode 100644 index 0000000..ceb49b0 --- /dev/null +++ b/tests/test_recipes/recipe_no_simplify.slim @@ -0,0 +1,22 @@ +initialize() +{ + setSeed(23); + if (!exists("TREES_FILE")) defineGlobal("TREES_FILE", "out.trees"); + initializeSLiMOptions(keepPedigrees=T); + initializeTreeSeq(simplificationRatio=INF, timeUnit="generations"); + initializeMutationRate(1e-2); + initializeMutationType("m1", 0.5, "f", -0.1); + initializeGenomicElementType("g1", m1, 1.0); + initializeGenomicElement(g1, 0, 99); + initializeRecombinationRate(1e-2); +} + +1 early() { + sim.addSubpop("p1", 10); +} + +10 late() { + sim.treeSeqOutput(TREES_FILE, simplify=F); + catn("Done."); + sim.simulationFinished(); +} diff --git a/tests/test_recipes/recipe_nonWF.slim b/tests/test_recipes/recipe_nonWF.slim index 72727ad..7174a85 100644 --- a/tests/test_recipes/recipe_nonWF.slim +++ b/tests/test_recipes/recipe_nonWF.slim @@ -12,6 +12,7 @@ initialize() initializeGenomicElement(g1, 0, 99); initializeRecombinationRate(1e-2); defineConstant("K", 10); + defineGlobal("MD", Dictionary()); } reproduction() { @@ -26,8 +27,43 @@ early() { p1.fitnessScaling = K / p1.individualCount; } + +// MUTATION/GENOTYPE INFO +1 first() { // mutation information + MD.setValue("mutations", Dictionary()); +} +mutation() { + muts = MD.getValue("mutations"); + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + return T; +} +10 late() { + subs = Dictionary(); + for (mut in sim.substitutions) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "fixationTick", mut.fixationTick + ); + subs.setValue(asString(mut.id), m); + } + MD.setValue("substitutions", subs); +} + +// OUTPUT/FINISH 10 late() { - sim.treeSeqOutput(TREES_FILE); + sim.treeSeqOutput(TREES_FILE, metadata=MD); catn("Done."); sim.simulationFinished(); } diff --git a/tests/test_recipes/recipe_nonWF.v5.2.trees b/tests/test_recipes/recipe_nonWF.v5.2.trees new file mode 100644 index 0000000..137d2b0 Binary files /dev/null and b/tests/test_recipes/recipe_nonWF.v5.2.trees differ diff --git a/tests/test_recipes/recipe_nucleotides_WF.slim b/tests/test_recipes/recipe_nucleotides_WF.slim index 29baf77..32b30e9 100644 --- a/tests/test_recipes/recipe_nucleotides_WF.slim +++ b/tests/test_recipes/recipe_nucleotides_WF.slim @@ -12,14 +12,57 @@ initialize() { initializeGenomicElementType("g1", m1, 1.0, mmJukesCantor(4e-2)); initializeGenomicElement(g1, 0, L-1); initializeRecombinationRate(1e-2); + defineGlobal("MD", Dictionary()); } 1 early() { sim.addSubpop("p1", 10); } +// MUTATION/GENOTYPE INFO +1 first() { // reference sequences + refseqs = Dictionary(); + for (chrom in sim.chromosomes) { + refseqs.setValue("chr" + chrom.id, chrom.ancestralNucleotides()); + } + MD.setValue("reference_sequence", refseqs); +} + +1 first() { // mutation information + MD.setValue("mutations", Dictionary()); +} +mutation() { + muts = MD.getValue("mutations"); + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + return T; +} +10 late() { + subs = Dictionary(); + for (mut in sim.substitutions) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "fixationTick", mut.fixationTick + ); + subs.setValue(asString(mut.id), m); + } + MD.setValue("substitutions", subs); +} + +// OUTPUT/FINISH 10 late() { - sim.treeSeqOutput(TREES_FILE); + sim.treeSeqOutput(TREES_FILE, metadata=MD); catn("Done."); sim.simulationFinished(); } diff --git a/tests/test_recipes/recipe_nucleotides_nonWF.slim b/tests/test_recipes/recipe_nucleotides_nonWF.slim index 0ae0dc0..5896984 100644 --- a/tests/test_recipes/recipe_nucleotides_nonWF.slim +++ b/tests/test_recipes/recipe_nucleotides_nonWF.slim @@ -12,6 +12,7 @@ initialize() { initializeGenomicElement(g1, 0, L-1); initializeRecombinationRate(1e-2); defineConstant("K", 10); + defineGlobal("MD", Dictionary()); } reproduction() { @@ -26,8 +27,51 @@ early() { p1.fitnessScaling = K / p1.individualCount; } + +// MUTATION/GENOTYPE INFO +1 first() { // reference sequences + refseqs = Dictionary(); + for (chrom in sim.chromosomes) { + refseqs.setValue("chr" + chrom.id, chrom.ancestralNucleotides()); + } + MD.setValue("reference_sequence", refseqs); +} + +1 first() { // mutation information + MD.setValue("mutations", Dictionary()); +} +mutation() { + muts = MD.getValue("mutations"); + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "originTick", mut.originTick + ); + muts.setValue(asString(mut.id), m); + return T; +} +10 late() { + subs = Dictionary(); + for (mut in sim.substitutions) { + nuc = mut.mutationType.nucleotideBased ? mut.nucleotide else "N"; + m = Dictionary( + "chromosome_id", mut.chromosome.id, + "mutationType", mut.mutationType.id, + "nucleotide", nuc, + "position", mut.position, + "fixationTick", mut.fixationTick + ); + subs.setValue(asString(mut.id), m); + } + MD.setValue("substitutions", subs); +} + +// OUTPUT/FINISH 10 late() { - sim.treeSeqOutput(TREES_FILE); + sim.treeSeqOutput(TREES_FILE, metadata=MD); catn("Done."); sim.simulationFinished(); } diff --git a/tests/test_recipes/recipe_with_traits.slim b/tests/test_recipes/recipe_with_traits.slim new file mode 100644 index 0000000..72602d2 --- /dev/null +++ b/tests/test_recipes/recipe_with_traits.slim @@ -0,0 +1,101 @@ +initialize() { + setSeed(23); + if (!exists("TREES_FILE")) defineGlobal("TREES_FILE", "out.trees"); + initializeSLiMOptions(keepPedigrees=T); + initializeTreeSeq(timeUnit="generations"); + defineConstant("I1", 5.0); + defineConstant("I2", -5.0); + defineConstant("OPT1", 10.0); + defineConstant("OPT2", 10.0); + defineConstant("SD1", 2.0); + defineConstant("SD2", 2.0); + + initializeSex(); + + popgen1T = initializeTrait("popgen1T", "m", 1.0, 0.0, 0.01, directFitnessEffect=T); + popgen2T = initializeTrait("popgen2T", "m", 1.0, 0.0, 0.01, directFitnessEffect=T); + n1T = initializeTrait("n1T", "m", directFitnessEffect=T); + n2T = initializeTrait("n2T", "m", directFitnessEffect=F); + + quant1T = initializeTrait("quant1T", "a", I1, 0.0, 0.01, directFitnessEffect=F); + quant2T = initializeTrait("quant2T", "a", I2, 0.0, 0.01, directFitnessEffect=F); + n3T = initializeTrait("n3T", "a", directFitnessEffect=F, baselineAccumulation=F); + + + logistic1T = initializeTrait("logistic1T", "l", 0.0, 0.01, 0.01, directFitnessEffect=T); + initializeMutationType("m1", 0.4, "f", 0.0); + initializeMutationType("m2", 0.4, "e", 0.05); + m2.setEffectSizeDistributionForTrait(c(n1T, n2T), "f", 0.0); + m2.setEffectSizeDistributionForTrait(c(quant1T, quant2T), "n", 0.0, 0.1); + m2.setEffectSizeDistributionForTrait(c(logistic1T), "n", -0.05, 0.1); + + initializeMutationType("m3", 0.4, "g", -0.05, 1.0); + m3.setEffectSizeDistributionForTrait(c(n1T, n2T), "f", 0.0); + m3.setEffectSizeDistributionForTrait(c(quant1T, quant2T), "n", 0.0, 0.1); + m3.setEffectSizeDistributionForTrait(c(logistic1T), "n", -0.05, 0.1); + + c(m2,m3).setEffectSizeDistributionForTrait(n3T, "n", -5.0, 0.5); + + c(m2,m3).setDefaultDominanceForTrait(c(popgen2T, quant2T), NAN); + + c(m2,m3).logMutationData(T, trait=NULL, effectSize=T, dominance=T); + + initializeGenomicElementType("g1", m1, 1.0); + initializeGenomicElementType("g2", 1:3, c(3, 1, 2)); + + ids = 1:5; + symbols = c(1, 2, "X", "Y", "MT"); + lengths = rdunif(5, 1e7, 2e7); + types = c("A", "A", "X", "Y", "H"); + names = c("A1", "A2", "X", "Y", "MT"); + + for (id in ids, symbol in symbols, length in lengths, type in types, name in names) + { + initializeChromosome(id, length, type, symbol, name); + initializeMutationRate(1e-7); + initializeRecombinationRate(1e-8); + + if (id == 1) + initializeGenomicElement(g1); // autosome 1 is pure-neutral, using only m1 + else + initializeGenomicElement(g2); // autosome 2 is a mix, using m1 / m2 / m3 + } +} + +mutation(m2) { + // set random dominance effects for the popgen1T and quant1T and logistic1TDominance traits + // other effects are generated as specified by the mutation type DES + mut.popgen1TDominance = runif(1); + mut.quant1TDominance = runif(1); + mut.logistic1TDominance = runif(1); + return T; +} +mutation(m3) { + // set random dominance effects for the popgen1T and quant1T and logistic1TDominance traits + // other effects are generated as specified by the mutation type DES + mut.popgen1TDominance = runif(1); + mut.quant1TDominance = runif(1); + mut.logistic1TDominance = runif(1); + return T; +} + +1 late() { + sim.addSubpop("p1", 20); +} + +1: late() { + inds = sim.subpopulations.individuals; + sim.demandPhenotype(NULL, c(sim.quant1T, sim.quant2T)); + phenotypes_q1 = inds.quant1T; + phenotypes_q2 = inds.quant2T; + fitnessEffect_q1 = dnorm(phenotypes_q1, OPT1, SD1) / dnorm(0.0, 0.0, SD1); + fitnessEffect_q2 = dnorm(phenotypes_q2, OPT2, SD2) / dnorm(0.0, 0.0, SD2); + inds.fitnessScaling = fitnessEffect_q1 * fitnessEffect_q2; +} + + +10 late() { + sim.treeSeqOutput(TREES_FILE); + catn("Done."); + sim.simulationFinished(); +} diff --git a/tests/test_tree_sequence.py b/tests/test_tree_sequence.py index 57f9b0e..18a9611 100644 --- a/tests/test_tree_sequence.py +++ b/tests/test_tree_sequence.py @@ -2,6 +2,7 @@ Test cases for tree sequences. """ +import copy import json import random import sys @@ -17,6 +18,18 @@ from .recipe_specs import recipe_eq, restarted_recipe_eq +def run_with_ts_metadata(f, ts_metadata, *args, **kwargs): + # check for equality in a method with and without passing in + # the ts_metadata argument + a = f(*args, **kwargs) + new_kwargs = copy.deepcopy(kwargs) + new_kwargs["ts_metadata"] = ts_metadata + b = f(*args, **new_kwargs) + assert len(a) == len(b) + np.testing.assert_equal(a, b) + return a + + def mutations_above(ts, node, pos): for s in ts.sites(): if s.position == pos: @@ -42,26 +55,65 @@ def naive_mutation_at(ts, node, pos, time=None): return mut_id +def verify_mutation_metadata(ts): + # Verify that all derived states are properly accounted for + # in mutation metadata. + mdl = ts.metadata["SLiM_mutation_list"] + mut_info = {str(mut["mutation_id"]): mut for mut in mdl} + assert len(mut_info) == len(mdl) + for mut in ts.mutations(): + for j in mut.derived_state.split(","): + assert j in mut_info + + +class TestMutationMetadata(tests.PyslimTestCase): + @pytest.mark.parametrize("recipe", recipe_eq("multichrom"), indirect=True) + def test_mutation_IDs_unique(self, recipe): + ids = set() + for _, ts in recipe["ts"].items(): + mut_info = pyslim.mutation_metadata(ts) + new_ids = set(mut_info.keys()) + assert len(ids.intersection(new_ids)) == 0 + ids = ids.union(new_ids) + + def test_mutation_metadata(self, recipe): + # test that mutation metadata is properly present + for _, ts in recipe["ts"].items(): + verify_mutation_metadata(ts) + + @pytest.mark.parametrize("recipe", [next(recipe_eq())], indirect=True) + def test_check(self, recipe): + for _, ts in recipe["ts"].items(): + assert ts.num_mutations > 5 + t = ts.dump_tables() + md = t.metadata + del md["SLiM_mutation_list"][5:] + t.metadata = md + ts = t.tree_sequence() + mut_info = pyslim.mutation_metadata(ts, check=False) + assert len(mut_info) == 5 + with pytest.raises(ValueError, match="missing information for mutation"): + _ = pyslim.mutation_metadata(ts) + + class TestSlimTime(tests.PyslimTestCase): # Tests for slim_time() - @pytest.mark.parametrize("recipe", recipe_eq(exclude="long"), indirect=True) + @pytest.mark.parametrize( + "recipe", recipe_eq(exclude=["long", "old_mutations"]), indirect=True + ) def test_slim_time(self, recipe): for _, ts in recipe["ts"].items(): - if "init_mutated" not in recipe: - for mut in ts.mutations(): - mut_time = max( - [x["slim_time"] for x in mut.metadata["mutation_list"]] - ) - assert mut_time == pyslim.slim_time(ts, mut.time) + muts = pyslim.mutation_metadata(ts) # the mutations in "init_mutated" examples have mutations that are *added* # in *early*, and so their times match in that stage. - else: - for mut in ts.mutations(): - mut_time = max( - [x["slim_time"] for x in mut.metadata["mutation_list"]] - ) - assert mut_time == pyslim.slim_time(ts, mut.time, stage="early") + stage = "early" if "init_mutated" in recipe else None + slim_times = pyslim.slim_time(ts, ts.mutations_time, stage=stage) + for t, mut in zip(slim_times, ts.mutations()): + mut_time = max( + [muts[int(j)]["slim_time"] for j in mut.derived_state.split(",")] + ) + assert mut_time == t class TestNextMutationID(tests.PyslimTestCase): @@ -99,12 +151,15 @@ def test_reload_slim(self, recipe, helper_functions, tmp_path): ) next_id = pyslim.next_slim_mutation_id(rts) T = max(1, rts.segregating_sites(mode="branch", span_normalise=False)) - mts = msprime.sim_mutations( - rts, - rate=max(6e-4, 10 / T), - keep=True, - model=msprime.SLiMMutationModel(type=1, next_id=next_id), - random_seed=135, + mts = pyslim.add_mutation_metadata( + msprime.sim_mutations( + rts, + rate=max(6e-4, 10 / T), + keep=True, + model=msprime.SLiMMutationModel(type=1, next_id=next_id), + random_seed=135, + ), + mutation_type=1, ) assert mts.num_mutations > rts.num_mutations recapped[chrom] = mts @@ -127,11 +182,16 @@ def test_reload_slim(self, recipe, helper_functions, tmp_path): assert chrom in recipe["ts"] assert pyslim.next_slim_mutation_id(mts) == pyslim.next_slim_mutation_id(ts) assert ts.num_mutations == recapped[chrom].num_mutations - a = ts.metadata - a["SLiM"].pop("user_metadata", None) - b = recipe["ts"][chrom].metadata - b["SLiM"].pop("user_metadata", None) - assert a == b + ots = recipe["ts"][chrom] + assert ts.metadata["SLiM"] == ots.metadata["SLiM"] + mut_info = pyslim.mutation_metadata(ts) + omut_info = pyslim.mutation_metadata(ots) + assert len(mut_info) == len(ts.metadata["SLiM_mutation_list"]) + assert len(omut_info) == len(ots.metadata["SLiM_mutation_list"]) + # we've added new mutations but originals should all be there + for k in omut_info: + assert k in mut_info + assert omut_info[k] == mut_info[k] def test_invalid_derived_state(self): ts = msprime.sim_ancestry( @@ -156,10 +216,12 @@ class TestRecapitate(tests.PyslimTestCase): """ def check_recap_consistency(self, ts, recap, with_ancestral_Ne=True): - assert ts.metadata["SLiM"]["tick"] == recap.metadata["SLiM"]["tick"] - assert ts.metadata["SLiM"]["cycle"] == recap.metadata["SLiM"]["cycle"] - assert ts.metadata["SLiM"]["stage"] == recap.metadata["SLiM"]["stage"] - assert ts.metadata["SLiM"]["name"] == recap.metadata["SLiM"]["name"] + tsmd = ts.metadata + remd = recap.metadata + assert tsmd["SLiM"]["tick"] == remd["SLiM"]["tick"] + assert tsmd["SLiM"]["cycle"] == remd["SLiM"]["cycle"] + assert tsmd["SLiM"]["stage"] == remd["SLiM"]["stage"] + assert tsmd["SLiM"]["name"] == remd["SLiM"]["name"] assert all(tree.num_roots == 1 for tree in recap.trees()) assert ts.has_reference_sequence() == recap.has_reference_sequence() if ts.has_reference_sequence(): @@ -264,6 +326,7 @@ def test_unique_names(self): assert names[0] == "ancestral" assert names[-2] == "ancestral_ancestral" + @pytest.mark.parametrize("recipe", recipe_eq(exclude="no_simplify"), indirect=True) def test_recapitation(self, recipe): for _, ts in recipe["ts"].items(): recomb_rate = 1.0 / ts.sequence_length @@ -282,7 +345,9 @@ def test_recapitation(self, recipe): assert t.num_roots == 1 assert recap.node(t.root).time >= old_root_time - @pytest.mark.parametrize("recipe", recipe_eq(exclude="long"), indirect=True) + @pytest.mark.parametrize( + "recipe", recipe_eq(exclude=["long", "no_simplify"]), indirect=True + ) def test_with_recomb_map(self, recipe): for _, ts in recipe["ts"].items(): recomb_rate = 1.0 / ts.sequence_length @@ -333,9 +398,10 @@ def test_first_gen_nodes(self, recipe): # (note this will fail if some populations were started at different # times than others or if the tick has been changed) for _, ts in recipe["ts"].items(): - root_time = ts.metadata["SLiM"]["tick"] - is_wf = ts.metadata["SLiM"]["model_type"] == "WF" - remembered_stage = ts.metadata["SLiM"]["stage"] + tsmd = ts.metadata + root_time = tsmd["SLiM"]["tick"] + is_wf = tsmd["SLiM"]["model_type"] == "WF" + remembered_stage = tsmd["SLiM"]["stage"] if (not is_wf) or (remembered_stage != "late"): root_time -= 1 if (not is_wf) and ("begun_first" in recipe): @@ -344,12 +410,10 @@ def test_first_gen_nodes(self, recipe): root_time -= 1 if is_wf and ("begun_late" in recipe): root_time -= 1 + vacant = pyslim.nodes_vacant(ts) for t in ts.trees(): for u in t.roots: - assert ( - pyslim.node_is_vacant(ts, ts.node(u)) - or ts.node(u).time == root_time - ) + assert vacant[u] or ts.node(u).time == root_time class TestIndividualAges(tests.PyslimTestCase): @@ -384,14 +448,23 @@ def test_mismatched_remembered_stage(self, recipe): def test_population(self, recipe): for _, ts in recipe["ts"].items(): individual_populations = ts.individuals_population - all_inds = pyslim.individuals_alive_at(ts, 0) + ts_metadata = ts.metadata + all_inds = run_with_ts_metadata( + pyslim.individuals_alive_at, ts_metadata, ts, 0 + ) assert len(all_inds) > 0 for p in range(ts.num_populations): - sub_inds = pyslim.individuals_alive_at(ts, 0, population=p) + sub_inds = pyslim.individuals_alive_at( + ts, 0, population=p, ts_metadata=ts_metadata + ) assert set(sub_inds) == set(all_inds[individual_populations == p]) - sub_inds = pyslim.individuals_alive_at(ts, 0, population=[p]) + sub_inds = pyslim.individuals_alive_at( + ts, 0, population=[p], ts_metadata=ts_metadata + ) assert set(sub_inds) == set(all_inds[individual_populations == p]) - sub_inds = pyslim.individuals_alive_at(ts, 0, population=np.arange(p)) + sub_inds = pyslim.individuals_alive_at( + ts, 0, population=np.arange(p), ts_metadata=ts_metadata + ) assert set(sub_inds) == set(all_inds[individual_populations != p]) @pytest.mark.parametrize( @@ -399,9 +472,14 @@ def test_population(self, recipe): ) def test_samples_only(self, recipe): for _, ts in recipe["ts"].items(): - all_inds = pyslim.individuals_alive_at(ts, 0) + ts_metadata = ts.metadata + all_inds = run_with_ts_metadata( + pyslim.individuals_alive_at, ts_metadata, ts, 0 + ) assert set(all_inds) == set( - pyslim.individuals_alive_at(ts, 0, samples_only=False) + pyslim.individuals_alive_at( + ts, 0, samples_only=False, ts_metadata=ts_metadata + ) ) sub_inds = np.random.choice( all_inds, size=min(len(all_inds), 4), replace=False @@ -440,18 +518,19 @@ def test_after_simplify(self, recipe): @pytest.mark.parametrize("recipe", recipe_eq("pedigree"), indirect=True) def test_ages(self, recipe): for _, ts in recipe["ts"].items(): + ts_metadata = ts.metadata info = recipe["info"] remembered_stage = "late" if "remembered_first" in recipe: remembered_stage = "first" elif "remembered_early" in recipe: remembered_stage = "early" - assert remembered_stage == ts.metadata["SLiM"]["stage"] - max_time_ago = ts.metadata["SLiM"]["tick"] + assert remembered_stage == ts_metadata["SLiM"]["stage"] + max_time_ago = ts_metadata["SLiM"]["tick"] if remembered_stage in ("first", "early"): max_time_ago -= 1 for time in range(0, max_time_ago): - slim_tick = ts.metadata["SLiM"]["tick"] - time + slim_tick = ts_metadata["SLiM"]["tick"] - time check_stages = ("first", "early", "late") if time == 0: if remembered_stage == "first": @@ -470,10 +549,18 @@ def test_ages(self, recipe): check_stages = ("late",) for stage in check_stages: alive = pyslim.individuals_alive_at( - ts, time, stage=stage, remembered_stage=remembered_stage + ts, + time, + stage=stage, + remembered_stage=remembered_stage, + ts_metadata=ts_metadata, ) ages = pyslim.individual_ages_at( - ts, time, stage=stage, remembered_stage=remembered_stage + ts, + time, + stage=stage, + remembered_stage=remembered_stage, + ts_metadata=ts_metadata, ) for ind in ts.individuals(): ind_time = ts.node(ind.nodes[0]).time @@ -489,7 +576,7 @@ def test_ages(self, recipe): assert slim_alive == pyslim_alive if slim_alive: slim_age = info[slim_id]["age"][(slim_tick, stage)] - if ts.metadata["SLiM"]["model_type"] == "WF": + if ts_metadata["SLiM"]["model_type"] == "WF": # SLiM records -1 but we return 0 in late and 1 in early slim_age = 0 + (stage in ("first", "early")) assert ages[ind.id] == slim_age @@ -595,9 +682,10 @@ def test_post_simplify(self, recipe): for _, ts in recipe["ts"].items(): rng = np.random.default_rng(seed=3) individual_times = ts.individuals_time + md_tick = ts.metadata["SLiM"]["tick"] keep_indivs = rng.choice( # assumes tick hasn't been changed - np.where(individual_times < ts.metadata["SLiM"]["tick"] - 1)[0], + np.where(individual_times < md_tick - 1)[0], size=30, replace=False, ) @@ -694,14 +782,66 @@ def test_pedigree_parents(self, recipe): gfolks = [] for a in set(info[sid]["parents"]) - set(ts_p): gfolks.extend(info[a]["parents"]) - print("===== ", hasp, ind) - print("sid: ", sid, "ts_p: ", ts_p) - print("slim_p: ", slim_p) - print(gfolks) + # print("===== ", hasp, ind) + # print("sid: ", sid, "ts_p: ", ts_p) + # print("slim_p: ", slim_p) + # print(gfolks) for a in set(ts_p) - set(slim_p): assert a in gfolks +class TestMutationConsistency(tests.PyslimTestCase): + """ + Test for consistency between what SLiM has written down in top-level metadata + and what's in the tree sequence + """ + + @pytest.mark.parametrize("recipe", recipe_eq("refseq"), indirect=True) + def test_reference_sequence_consistency(self, recipe): + for n, ts in recipe["ts"].items(): + tsmd = ts.metadata + chrom_id = tsmd["SLiM"]["this_chromosome"]["id"] + assert ts.has_reference_sequence() + ref = list( + tsmd["SLiM"]["user_metadata"]["reference_sequence"][0][f"chr{chrom_id}"][ + 0 + ] + ) + ts_ref = ts.reference_sequence.data + subs = [ + x[0] + for x in tsmd["SLiM"]["user_metadata"]["substitutions"][0].values() + if x[0]["chromosome_id"][0] == chrom_id + ] + subs.sort(key=lambda x: (x["position"][0], x["fixationTick"][0])) + for s in subs: + nuc = s["nucleotide"][0] + if nuc != "N": + ref[s["position"][0]] = nuc + assert ts_ref == "".join(ref) + + @pytest.mark.parametrize("recipe", recipe_eq("record_mutations"), indirect=True) + def test_mutation_consistency(self, recipe): + for n, ts in recipe["ts"].items(): + tsmd = ts.metadata + chrom_id = tsmd["SLiM"]["this_chromosome"]["id"] + # this is just making these things not lists, mostly + debug_info = { + int(k): {x: y[0] for x, y in v[0].items()} + for k, v in tsmd["SLiM"]["user_metadata"]["mutations"][0].items() + } + mut_info = pyslim.mutation_metadata(ts) + for mut in ts.mutations(): + for k in mut.derived_state.split(","): + k = int(k) + assert k in debug_info or mut_info[k]["mutation_id"] == 2 + assert k in mut_info + assert debug_info[k]["chromosome_id"] == chrom_id + assert debug_info[k]["position"] == ts.site(mut.site).position + assert debug_info[k]["mutationType"] == mut_info[k]["mutation_type"] + assert debug_info[k]["originTick"] == mut_info[k]["slim_time"] + + class TestReferenceSequence(tests.PyslimTestCase): """ Test for operations involving the reference sequence @@ -749,6 +889,7 @@ def test_nucleotide_at_errors(self, recipe): def test_mutation_at(self, recipe): rng = random.Random(42) for _, ts in recipe["ts"].items(): + L = int(min(50000, ts.sequence_length)) for _ in range(min(10, ts.num_sites)): site = rng.choice(ts.sites()) pos = site.position @@ -766,10 +907,8 @@ def test_mutation_at(self, recipe): a = pyslim.mutation_at(ts, node, pos, time=time) b = naive_mutation_at(ts, node, pos, time=time) assert a == b - for _ in range(min(10, int(ts.sequence_length - ts.num_sites))): - pos = rng.choice( - list(set(range(int(ts.sequence_length))) - set(ts.sites_position)) - ) + for _ in range(min(10, int(L - ts.num_sites))): + pos = rng.choice(list(set(range(L)) - set(ts.sites_position))) tree = ts.at(pos) for _ in range(10): node = rng.randint(0, ts.num_nodes - 1) @@ -778,33 +917,63 @@ def test_mutation_at(self, recipe): for time in [None, ts.node(node).time, ut]: assert naive_mutation_at(ts, node, pos, time=time) == -1 + @pytest.mark.parametrize("recipe", recipe_eq("nucleotides"), indirect=True) def test_nucleotide_at(self, recipe): random.seed(42) for _, ts in recipe["ts"].items(): if ts.num_mutations > 0: + mut_metadata = pyslim.mutation_metadata(ts) mut_md = ts.mutation(0).metadata - has_nucleotides = mut_md["mutation_list"][0]["nucleotide"] >= 0 - if has_nucleotides: - assert ts.has_reference_sequence() - assert len(ts.reference_sequence.data) == ts.sequence_length - for _ in range(100): - node = random.randint(0, ts.num_nodes - 1) - pos = random.randint(0, int(ts.sequence_length) - 1) - tree = ts.at(pos) - parent = tree.parent(node) - a = pyslim.nucleotide_at(ts, node, pos) - if parent == tskit.NULL: - nuc = ts.reference_sequence.data[int(pos)] - assert a == pyslim.NUCLEOTIDES.index(nuc) - else: - b = pyslim.nucleotide_at(ts, parent, pos) - c = pyslim.nucleotide_at(ts, node, pos, ts.node(parent).time) - assert b == c - for k in np.where(node == ts.tables.mutations.node)[0]: - mut = ts.mutation(k) - if ts.site(mut.site).position == pos: - b = mut.metadata["mutation_list"][0]["nucleotide"] - assert a == b + tsmd = ts.metadata + # check we've got nucleotide mutations + nucs = np.array([x["nucleotide"] for x in tsmd["SLiM_mutation_list"]]) + assert np.sum(nucs >= 0) > 1 + mut_info = { + str(mut["mutation_id"]): mut for mut in tsmd["SLiM_mutation_list"] + } + assert ts.has_reference_sequence() + assert len(ts.reference_sequence.data) == ts.sequence_length + for _ in range(100): + node = random.randint(0, ts.num_nodes - 1) + pos = random.randint(0, int(ts.sequence_length) - 1) + tree = ts.at(pos) + parent = tree.parent(node) + a = pyslim.nucleotide_at(ts, node, pos) + if parent == tskit.NULL: + nuc = ts.reference_sequence.data[int(pos)] + assert a == pyslim.NUCLEOTIDES.index(nuc) + else: + b = pyslim.nucleotide_at( + ts, parent, pos, mut_metadata=mut_metadata + ) + c = pyslim.nucleotide_at( + ts, + node, + pos, + ts.node(parent).time, + mut_metadata=mut_metadata, + ) + assert b == c + for k in np.where(node == ts.tables.mutations.node)[0]: + mut = ts.mutation(k) + if ts.site(mut.site).position == pos: + b = mut_info[mut.derived_state.split(",")[0]][ + "nucleotide" + ] + assert a == b + + @pytest.mark.parametrize("recipe", [next(recipe_eq("nucleotides"))], indirect=True) + def test_nucleotide_at_without_mut_metadata(self, recipe): + random.seed(23) + for _, ts in recipe["ts"].items(): + assert ts.num_mutations > 0 + mut_metadata = pyslim.mutation_metadata(ts) + for _ in range(100): + node = random.randint(0, ts.num_nodes - 1) + pos = random.randint(0, int(ts.sequence_length) - 1) + a = pyslim.nucleotide_at(ts, node, pos) + b = pyslim.nucleotide_at(ts, node, pos, mut_metadata=mut_metadata) + assert a == b @pytest.mark.parametrize("recipe", recipe_eq("mutation_spectrum"), indirect=True) def test_nucleotide_spectrum(self, recipe): @@ -814,6 +983,7 @@ def test_nucleotide_spectrum(self, recipe): # access to the parental genome, so if two adjacent mutations # occur in the same meiosis then each will not know about the other. for _, ts in recipe["ts"].items(): + mut_info = pyslim.mutation_metadata(ts) mutation_spectrum = recipe["mutation_info"] M = { a + b + c + "," + d: 0 @@ -827,7 +997,7 @@ def test_nucleotide_spectrum(self, recipe): pos = ts.site(mut.site).position if pos > 0 and pos < ts.sequence_length - 1: nmuts += 1 - mut_list = mut.metadata["mutation_list"] + mut_list = [mut_info[int(k)] for k in mut.derived_state.split(",")] k = np.argmax([u["slim_time"] for u in mut_list]) derived_nuc = mut_list[k]["nucleotide"] left_nuc = pyslim.nucleotide_at( @@ -847,13 +1017,8 @@ def test_nucleotide_spectrum(self, recipe): ) key = context + "," + pyslim.NUCLEOTIDES[derived_nuc] M[key] += 1 - if key == "ACA,T" or key == "CCA,T": - print(key, pos, mut.node, mut.time) assert sum([M[k] for k in M]) == nmuts assert sum([mutation_spectrum[k][0] for k in mutation_spectrum]) == nmuts - for k in M: - if M[k] != mutation_spectrum[k][0]: - print(k, M[k], mutation_spectrum[k]) for k in M: assert len(mutation_spectrum[k]) == 1 assert M[k] == mutation_spectrum[k][0] @@ -869,12 +1034,12 @@ def last_slim_mutations(self, ts): # (slim id, slim mutation metadata) of the slim mutation that is the # *most recent* one of any possibly stacked mutations. Note that it # is possible that this is ambiguous. + mut_info = pyslim.mutation_metadata(ts) for mut in ts.mutations(): slim_muts = { k: v - for k, v in zip( - mut.derived_state.split(","), mut.metadata["mutation_list"] - ) + for k, v in mut_info.items() + if str(k) in mut.derived_state.split(",") } if mut.parent == tskit.NULL: parent_slim_ids = [] @@ -938,10 +1103,7 @@ def scramble_mutations(self, ts): for m in ts.mutations(): a = np.array(m.derived_state.split(",")) ii = rng.permutation(len(a)) - ml = [m.metadata["mutation_list"][i] for i in ii] - t.mutations.append( - m.replace(derived_state=",".join(a[ii]), metadata={"mutation_list": ml}) - ) + t.mutations.append(m.replace(derived_state=",".join(a[ii]))) t.compute_mutation_parents() return t.tree_sequence() @@ -956,6 +1118,7 @@ def test_convert_alleles_errors(self): ts, model=msprime.SLiMMutationModel(type=1), rate=0.1, random_seed=23 ) assert mts.num_mutations > 0 + mts = pyslim.add_mutation_metadata(mts) mtt = mts.dump_tables() mtt.reference_sequence.data = "A" * int(mts.sequence_length) mts = mtt.tree_sequence() @@ -967,6 +1130,7 @@ def test_convert_alleles_errors(self): ) def test_convert_alleles(self, recipe): for _, ts in recipe["ts"].items(): + verify_mutation_metadata(ts) cts = pyslim.convert_alleles(ts) self.verify_converted_nucleotides(ts, cts) @@ -1021,18 +1185,19 @@ def test_generate_nucleotides_errors(self): def verify_generate_nucleotides(self, ts, check_transitions=False): # if check_transitions is True, verify that derived states differ # from parental states - which we try to do but is not guaranteed, - # for instance, if keep=True or in other weird situations. + # for instance, if keep=True, there was more than one mutation in + # single generation, or in other weird situations. assert ts.metadata["SLiM"]["nucleotide_based"] assert len(ts.reference_sequence.data) == ts.sequence_length + mut_info = pyslim.mutation_metadata(ts) muts = {} ts_muts = { j: v["nucleotide"] for j, (_, v) in enumerate(self.last_slim_mutations(ts)) } for mut in ts.mutations(): aa = ts.reference_sequence.data[int(ts.site(mut.site).position)] - for i, md in zip( - mut.derived_state.split(","), mut.metadata["mutation_list"] - ): + for i in mut.derived_state.split(","): + md = mut_info[int(i)] nuc = md["nucleotide"] assert nuc in [0, 1, 2, 3] if i in muts: @@ -1042,9 +1207,14 @@ def verify_generate_nucleotides(self, ts, check_transitions=False): if mut.parent == tskit.NULL: assert pyslim.NUCLEOTIDES[nuc] != aa else: - if ts.mutation(mut.parent).derived_state != mut.derived_state: - assert ts_muts[mut.parent] != ts_muts[mut.id] - + mp = ts.mutation(mut.parent) + if mp.derived_state != mut.derived_state: + assert (ts_muts[mut.parent] != ts_muts[mut.id]) or ( + len(mut.derived_state.split(",")) + > 1 + len(mp.derived_state.split(",")) + ) + + @pytest.mark.parametrize("recipe", recipe_eq(exclude="old_mutations"), indirect=True) def test_generate_nucleotides(self, recipe): for _, ts in recipe["ts"].items(): nts = pyslim.generate_nucleotides(ts, keep=False, seed=5) @@ -1061,8 +1231,10 @@ def test_generate_nucleotides_refseq(self): random_seed=10, ) ts = pyslim.annotate(ts, model_type="nonWF", tick=1) - mts = msprime.sim_mutations( - ts, model=msprime.SLiMMutationModel(type=1), rate=0.5, random_seed=23 + mts = pyslim.add_mutation_metadata( + msprime.sim_mutations( + ts, model=msprime.SLiMMutationModel(type=1), rate=0.5, random_seed=23 + ) ) refseq = "A" * int(mts.sequence_length) nts = pyslim.generate_nucleotides(mts, reference_sequence=refseq, seed=6) @@ -1072,35 +1244,42 @@ def test_generate_nucleotides_refseq(self): def test_generate_nucleotides_keep(self): ts = msprime.sim_ancestry(4, sequence_length=10, population_size=10) ts = pyslim.annotate(ts, model_type="nonWF", tick=1) - mts1 = msprime.sim_mutations( - ts, model=msprime.SLiMMutationModel(type=1), rate=0.1, random_seed=23 + mts1 = pyslim.add_mutation_metadata( + msprime.sim_mutations( + ts, model=msprime.SLiMMutationModel(type=1), rate=0.1, random_seed=23 + ) ) - mts1.dump("out.trees") nts1 = pyslim.generate_nucleotides(mts1, seed=10, keep=False) assert nts1.num_mutations > 0 self.verify_generate_nucleotides(nts1, check_transitions=False) - mts2 = msprime.sim_mutations( - nts1, - model=msprime.SLiMMutationModel( - type=2, - next_id=nts1.num_mutations, - ), - rate=0.1, - random_seed=24, + mut_info1 = { + str(mut["mutation_id"]): mut for mut in nts1.metadata["SLiM_mutation_list"] + } + mts2 = pyslim.add_mutation_metadata( + msprime.sim_mutations( + nts1, + model=msprime.SLiMMutationModel( + type=2, + next_id=nts1.num_mutations, + ), + rate=0.1, + random_seed=24, + ) ) # keep defaults to True nts2 = pyslim.generate_nucleotides(mts2, seed=12) assert nts2.num_mutations > nts1.num_mutations + mut_info2 = { + str(mut["mutation_id"]): mut for mut in nts2.metadata["SLiM_mutation_list"] + } muts1 = {} for mut in nts1.mutations(): - for i, md in zip( - mut.derived_state.split(","), mut.metadata["mutation_list"] - ): + for i in mut.derived_state.split(","): + md = mut_info1[i] muts1[i] = md["nucleotide"] for mut in nts2.mutations(): - for i, md in zip( - mut.derived_state.split(","), mut.metadata["mutation_list"] - ): + for i in mut.derived_state.split(","): + md = mut_info2[i] if md["mutation_type"] == 1: assert i in muts1 assert muts1[i] == md["nucleotide"] @@ -1160,8 +1339,9 @@ def get_vacant_samples(self, ts): def verify_remove_vacant(self, ts, rts): vacant_samples = self.get_vacant_samples(ts) - for node in rts.nodes(): - assert not (pyslim.node_is_vacant(rts, node) and (node.is_sample() == 1)) + vacant = pyslim.nodes_vacant(ts) + for v, node in zip(vacant, rts.nodes()): + assert not (v and (node.is_sample() == 1)) assert (node.id in vacant_samples) == ( node.flags & pyslim.NODE_IS_VACANT_SAMPLE > 0 ) @@ -1262,14 +1442,23 @@ def test_has_vacant_msprime(self): tables.nodes.metadata_schema = pyslim.slim_metadata_schemas["node"] assert not pyslim.has_vacant_samples(tables.tree_sequence()) - def test_node_is_vacant(self, recipe): + def test_nodes_vacant(self, recipe): + np.random.seed(123) num_chromosomes = len(recipe["ts"]) for _, ts in recipe["ts"].items(): + test_nodes = np.random.choice( + np.arange(ts.num_nodes), size=min(1, ts.num_nodes), replace=False + ) k = ts.metadata["SLiM"]["this_chromosome"]["index"] - for node in ts.nodes(): + vacant = pyslim.nodes_vacant(ts) + for pv, node in zip(vacant, ts.nodes()): v = self.vacancy_values(node) isv = v is not None and v[k] - assert isv == pyslim.node_is_vacant(ts, node) + assert pv == isv + if node.id in test_nodes: + # node_is_vacant is kinda slow so don't test all of them + with pytest.warns(FutureWarning, match="deprecated"): + assert isv == pyslim.node_is_vacant(ts, node) for j in range(num_chromosomes, len(v)): assert not v[j] @@ -1394,6 +1583,16 @@ def test_no_change(self, restart_name, recipe, helper_functions, tmp_path): ) for chrom, ts in recipe["ts"].items(): self.verify_reset(ts, out_ts[chrom]) + # again!! + out_out_ts = helper_functions.run_slim_restart( + out_ts, + restart_name, + tmp_path, + "multichrom" in recipe, + WF="WF" in recipe, + ) + for chrom, ts in recipe["ts"].items(): + self.verify_reset(ts, out_out_ts[chrom]) @pytest.mark.parametrize( "restart_name, recipe", restarted_recipe_eq("no_op"), indirect=["recipe"] @@ -1425,16 +1624,20 @@ def test_set_individuals( ): in_ts = {} ts = list(recipe["ts"].values())[0] + tsmd = ts.metadata assert ( - "user_metadata" in ts.metadata["SLiM"] - and "reset_tick" in ts.metadata["SLiM"]["user_metadata"] + "user_metadata" in tsmd["SLiM"] + and "reset_tick" in tsmd["SLiM"]["user_metadata"] ), "Simulation not set up for this test." - reset_tick = ts.metadata["SLiM"]["user_metadata"]["reset_tick"][0] + reset_tick = tsmd["SLiM"]["user_metadata"]["reset_tick"][0] if time is None: - for time in range(ts.metadata["SLiM"]["tick"] + 1): - if pyslim.slim_time(ts, time) == reset_tick: - break - individuals = pyslim.individuals_alive_at(ts, time)[:num_indivs] + pytimes = run_with_ts_metadata( + pyslim.slim_time, tsmd, ts, np.arange(tsmd["SLiM"]["tick"] + 1) + ) + time = np.searchsorted(pytimes, reset_tick) + individuals = pyslim.individuals_alive_at(ts, time, ts_metadata=tsmd)[ + :num_indivs + ] for chrom, ts in recipe["ts"].items(): in_ts[chrom] = pyslim.set_slim_state(ts, time=time, individuals=individuals) out_ts = helper_functions.run_slim_restart( diff --git a/uv.lock b/uv.lock index 68ae021..fbf98b3 100644 --- a/uv.lock +++ b/uv.lock @@ -645,6 +645,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/4e/ce75a57ff3aebf6fc1f4e9d508b8e5810618a33d900ad6c19eb30b290b97/fonttools-4.61.1-py3-none-any.whl", hash = "sha256:17d2bf5d541add43822bcf0c43d7d847b160c9bb01d15d5007d84e2217aaa371", size = 1148996, upload-time = "2025-12-12T17:31:21.03Z" }, ] +[[package]] +name = "frozendict" +version = "2.4.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/90/b2/2a3d1374b7780999d3184e171e25439a8358c47b481f68be883c14086b4c/frozendict-2.4.7.tar.gz", hash = "sha256:e478fb2a1391a56c8a6e10cc97c4a9002b410ecd1ac28c18d780661762e271bd", size = 317082, upload-time = "2025-11-11T22:40:14.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl", hash = "sha256:972af65924ea25cf5b4d9326d549e69a9a4918d8a76a9d3a7cd174d98b237550", size = 16264, upload-time = "2025-11-11T22:40:12.836Z" }, +] + [[package]] name = "greenlet" version = "3.3.2" @@ -1844,6 +1853,7 @@ dependencies = [ [package.dev-dependencies] dev = [ { name = "filelock" }, + { name = "frozendict" }, { name = "jupyter-book" }, { name = "matplotlib" }, { name = "msprime" }, @@ -1879,6 +1889,7 @@ packaging = [ ] test = [ { name = "filelock" }, + { name = "frozendict" }, { name = "msprime" }, { name = "pandas" }, { name = "pytest" }, @@ -1891,12 +1902,13 @@ test = [ requires-dist = [ { name = "msprime", specifier = ">=1.0.1" }, { name = "numpy" }, - { name = "tskit" }, + { name = "tskit", specifier = ">=1.0.3" }, ] [package.metadata.requires-dev] dev = [ { name = "filelock" }, + { name = "frozendict" }, { name = "jupyter-book", specifier = "<2" }, { name = "matplotlib" }, { name = "msprime" }, @@ -1932,6 +1944,7 @@ packaging = [ ] test = [ { name = "filelock" }, + { name = "frozendict" }, { name = "msprime" }, { name = "pandas" }, { name = "pytest" }, @@ -2798,23 +2811,38 @@ wheels = [ [[package]] name = "tskit" -version = "0.6.4" +version = "1.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonschema" }, { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/94/95/2c2d8bdaae4a3948181de68d1fac0569d9c937a42f7ccfcf097f9a428721/tskit-0.6.4.tar.gz", hash = "sha256:bdac1bb7e3ae3d1f562ec191b5d840156e082dd2adc6af7c41b170c4fb1be792", size = 874772, upload-time = "2025-05-21T18:18:18.343Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ed/28/547acef423709fad5b70bbc68332c1fa0a4c1887ec4e76cb93a433a95bbb/tskit-0.6.4-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:510fd219f2c6d5d669e178a3ea1e4cd60d7fc0b40a50fae4859a6a8f35d2394a", size = 748024, upload-time = "2025-05-21T18:17:54.919Z" }, - { url = "https://files.pythonhosted.org/packages/88/0d/5816ee9ac9708f2e1def387e36b2d0763bdea2267db3ff4a55c17f079e85/tskit-0.6.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b631d38c352b618c3ede33515cf97ffcbc24e3a4252b6979a80b1770434a534", size = 1322885, upload-time = "2025-05-21T18:17:56.958Z" }, - { url = "https://files.pythonhosted.org/packages/bd/db/b98964916b3f9c603f8ad92045db0c050ca4699a976051acdc1368085c3c/tskit-0.6.4-cp311-cp311-win_amd64.whl", hash = "sha256:e121226092816a1e36b2835aa0348e652c09bdf0604a936f63c4fd0d64ca6422", size = 470677, upload-time = "2025-05-21T18:17:58.584Z" }, - { url = "https://files.pythonhosted.org/packages/b4/cf/d2a1c6a6ad29b16310b60bf08dd5b1d1e4e0d23819c49284662373c9fc54/tskit-0.6.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:8fd5a0c94f302f5fe69a9f7a662b8ddf8219e4dc7325a1990b279afd084bb649", size = 749765, upload-time = "2025-05-21T18:18:00.234Z" }, - { url = "https://files.pythonhosted.org/packages/80/46/1630514e8a9f97a8f75e805a7872fb3083119d7e421309e05e99d5f5ca66/tskit-0.6.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f086e648a624004343882ca57f9365d5005cc2e472113b61874c9403c14be272", size = 1324515, upload-time = "2025-05-21T18:18:02.461Z" }, - { url = "https://files.pythonhosted.org/packages/07/44/2116904f37ffe1db0e675b4c70b70cdaf7ab38b474cd546f6ff9d46de4b9/tskit-0.6.4-cp312-cp312-win_amd64.whl", hash = "sha256:507eee5b20c5e47202d90a70143c601aef3b3ec70321b5251c2e6d896b1e0722", size = 470047, upload-time = "2025-05-21T18:18:04.185Z" }, - { url = "https://files.pythonhosted.org/packages/c5/d7/5ee3b118281f2a7032f5d38fca0ee7e11bf7cf1823d2dbfa5a991ccf3497/tskit-0.6.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:66f07b76f18ad576a7585b6e4d9e46d58994e3df1dd2d1808d32d294354695cf", size = 749773, upload-time = "2025-05-21T18:18:06.48Z" }, - { url = "https://files.pythonhosted.org/packages/f2/6d/6368c2ffbccb4d7ee8d82096f1ad0b3cb0d639b958c408589cceff09600e/tskit-0.6.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bfb40717382cf47b8844c1f73d73c1e7ca7f63e6a506e7f859f0980667677de2", size = 1324529, upload-time = "2025-05-21T18:18:08.674Z" }, - { url = "https://files.pythonhosted.org/packages/ee/b9/5b092c5b409cad0005b560622b34f9fd38fe0945939974671306fedef30a/tskit-0.6.4-cp313-cp313-win_amd64.whl", hash = "sha256:c6b1f1b22e5d55a906ee33bca33447500df0cf5b12ec20926773ca1e3c3c8931", size = 470097, upload-time = "2025-05-21T18:18:11.022Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/22/f5/c23333b3ffd86cbe79520bbee85f050f815c16a625b014dbacbe486fb6c2/tskit-1.0.3.tar.gz", hash = "sha256:8a305b3bcf4145688be17962f6ea2ac20aa65db440ed0c02b49f9b4a3c00836b", size = 933435, upload-time = "2026-05-14T18:20:16.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/6c/7fdf85e62204319d77da52d2647e1e65c6dc94660a27d7f202fd134f42c9/tskit-1.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a5d5f5e0d0ad53ac554563c76b91c5f2bb721d021b7c761606a9a72fe1e49d74", size = 526356, upload-time = "2026-05-14T18:19:40.482Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ce/4ff9ba98eca28992250fa8eab25db5bc4a3c7b3af4ad46652f7d1b5b22c3/tskit-1.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af978aa415f04fbabd613ad2933526110b4767e8e64f319ec03cefca9e2242bd", size = 496008, upload-time = "2026-05-14T18:19:42.332Z" }, + { url = "https://files.pythonhosted.org/packages/2f/3d/d4d8567a2bca51df46f0c7ad706678b106ce2ba094d8584d8829a24451c0/tskit-1.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f509b7b271dd80a68ee3f6bfdff614e3b8ac989dcc44dd8d36c25c8181443b79", size = 1395725, upload-time = "2026-05-14T18:19:43.83Z" }, + { url = "https://files.pythonhosted.org/packages/4c/f3/d5ce78eeeda65c4673966d00567f4166b9500843ce953c0dad927a2200c3/tskit-1.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0722a4c1fc9c287c96a2141ae3587df6defed2a049cca7ab90479958c613835f", size = 1377087, upload-time = "2026-05-14T18:19:45.238Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7c/9e53fb9f71bf89eb8327cb2a859d4c9764dc447a42d3407f7ec7c7a25fb8/tskit-1.0.3-cp311-cp311-win32.whl", hash = "sha256:3901c9fc02497e2c7e0ee5ca36ce5b65749f9c630e5ace1daa990963d8f5a63d", size = 449670, upload-time = "2026-05-14T18:19:46.641Z" }, + { url = "https://files.pythonhosted.org/packages/72/88/a80beb0adfb8ceade30b1913010c56e6bd4f837ef2f75aeb76222377d148/tskit-1.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:0d374768d422e941f8ad6ffdba85d6ca1585467532f6f4945b171f18d6613285", size = 496418, upload-time = "2026-05-14T18:19:47.948Z" }, + { url = "https://files.pythonhosted.org/packages/56/1f/d6726034289911b326e74b8b5339c9d1ec6fe46bf203253ae8d90b146a89/tskit-1.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3d035946888c9da5eefb93ca99c2ae3ad2cb7cb64bc6037588fe285f775ee67a", size = 528485, upload-time = "2026-05-14T18:19:49.299Z" }, + { url = "https://files.pythonhosted.org/packages/ad/ed/7fa3f43bc0d7569b75bc6a5c03bc14bd2e9de12a3f6a2f8b3251d505226a/tskit-1.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:31aee6a8b6d9a0bdae06f0b186fb7aad54d53f34050c0ee72a6caa560506950f", size = 496860, upload-time = "2026-05-14T18:19:50.716Z" }, + { url = "https://files.pythonhosted.org/packages/28/b1/fb828d0479be0f178e9f33b4483003a421b069e6331dc5f35066dc3029e4/tskit-1.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fb07ead500c7e88c69677469b4f489c8997ad944103102b04d724c855b4b980b", size = 1398919, upload-time = "2026-05-14T18:19:52.1Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3e/1c7adaae07a00ce009e692a6d171c7a0882d1ed9753017b78a7db7ae7d72/tskit-1.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:241192b12d310e49b21bfae524c76768a15d0ed3f7b950d6ee09dce1528b8973", size = 1378470, upload-time = "2026-05-14T18:19:53.881Z" }, + { url = "https://files.pythonhosted.org/packages/71/a1/10d12e9b7908775cf97ad9e6c3d94b4b970940c74a161cc9aaf540fa83c0/tskit-1.0.3-cp312-cp312-win32.whl", hash = "sha256:3e11c9fe328b27c0e4a3ff64a0fed972d2cd48a743c03c8bc339352664322ee8", size = 450010, upload-time = "2026-05-14T18:19:55.419Z" }, + { url = "https://files.pythonhosted.org/packages/2c/b7/d1842bec89cd4993958165d6da774e85ad79b9426baca571c39457a1b6b5/tskit-1.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:66c83c7f971bc160fd29ca6169e1ee31023a5283ca88936df3eed1e568eee209", size = 496000, upload-time = "2026-05-14T18:19:56.75Z" }, + { url = "https://files.pythonhosted.org/packages/e6/10/aa8cc65060669b911946649cd1802f29640998c5b6ca78c4c845fb3c473d/tskit-1.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d52e2ad0eb79e3b21da8a236d0f5eb888f9b18569e48f4b84a3875523e80ee22", size = 528459, upload-time = "2026-05-14T18:19:58.156Z" }, + { url = "https://files.pythonhosted.org/packages/eb/42/34d981123a335d80e3030e9b0a81a41f654c9cf1c5db0ac630d6b20d8bfd/tskit-1.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5bd1744a020a518e5835b85971fea14ca827d76f778650eb37e3f44c6db0a711", size = 496838, upload-time = "2026-05-14T18:19:59.909Z" }, + { url = "https://files.pythonhosted.org/packages/65/9c/09b072c01b8e58a3137011851d4761daf957f2f86447bc2a2c25e0d5aa4a/tskit-1.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:405c20fad5256edd6abd1f8b514c5bc653d808102f509c4512c9b174507dbe84", size = 1399080, upload-time = "2026-05-14T18:20:01.353Z" }, + { url = "https://files.pythonhosted.org/packages/d0/bb/789b91f7a2deca4d50f649b685e2b3dca4b3358e36bf7a5b2b84709bb8d5/tskit-1.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:bc1c97f00ce8e79460aa9ff3b36987daf70f2dc267e36a32db77018dbefc822b", size = 1378631, upload-time = "2026-05-14T18:20:03.273Z" }, + { url = "https://files.pythonhosted.org/packages/c7/8b/3b2b19246b2ba3159456e4068887e505d7a2ec857ecd7ae7606d6c129fda/tskit-1.0.3-cp313-cp313-win32.whl", hash = "sha256:53444ef3b2fd7ea9ca863e562c33402e46a55e645887dd4d983686b607e3f788", size = 449996, upload-time = "2026-05-14T18:20:04.788Z" }, + { url = "https://files.pythonhosted.org/packages/23/fc/53ef7ad8b3f2183584af81b0ff046a660c3164530a1f768770dbd4d60cb7/tskit-1.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:721ae7f02730ed91233a6da2a9d4b59632d1512877cd0d6fe948ea5d9ddc17e9", size = 496006, upload-time = "2026-05-14T18:20:06.577Z" }, + { url = "https://files.pythonhosted.org/packages/66/b9/91194216907ebf278d7a58288f3e64881e11a0eedcca7125a49e0bbdafdd/tskit-1.0.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:62a45d5e57f5a9181eb113c0dd478783ef3c4829856c0831ea2eba2cdf00a8de", size = 528594, upload-time = "2026-05-14T18:20:08.032Z" }, + { url = "https://files.pythonhosted.org/packages/6e/2c/0a888f4f140922c7a3945dbca5e444fa8052a620c784359673fe9eb1eda7/tskit-1.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8df716837129652fc82d095d80fd4e23a728213a4b18de9233f3aebb073b60b4", size = 496866, upload-time = "2026-05-14T18:20:09.34Z" }, + { url = "https://files.pythonhosted.org/packages/2f/ce/56aa69b768734ede0b8b36c5ea2ce842d50c786b74cfefdfc036a1fd29da/tskit-1.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:151eefbbae9a339f3ea64980937fa89d6ac9fb371dba880a9b30791a461753ad", size = 1397616, upload-time = "2026-05-14T18:20:10.762Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5a/b05eab4a4a32bce6a51f991033ed460a836aafc1f406eebc624831558d9b/tskit-1.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2170cacd8ade47fd8f1510aca78549e09e9ba875c25496f0ccfb755e09c0b06b", size = 1377477, upload-time = "2026-05-14T18:20:12.649Z" }, + { url = "https://files.pythonhosted.org/packages/ea/c5/116419617468dc06428e507001d0a7b212116e9e8d336bd8c05a6531b2a0/tskit-1.0.3-cp314-cp314-win32.whl", hash = "sha256:24f8fffc1e1dd4154654e91bb7bac22078b3cadd9837e5c490225aa5c313b819", size = 455586, upload-time = "2026-05-14T18:20:14.165Z" }, + { url = "https://files.pythonhosted.org/packages/6e/7f/61818ed922629178abb07e45439d0af3bb5867b629fe392e9864df58ce5c/tskit-1.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:cbe791757cbd4d060a9c06c270bbcacc932f60ce1d77e98dcea35c931d0b0aa9", size = 505409, upload-time = "2026-05-14T18:20:15.527Z" }, ] [[package]]