test for trait consistency

docs/phenotypes.md

@@ -254,7 +254,7 @@ print(f"SLiM: {np.array([x['phenotype'] for x in ind.metadata['per_trait']])}")
 ```
 
 That looks good! To compare we need to account for floating-point error
-(actually the numbers stored by SLiM and computed by us may differ by $10^{-10}$).
+(actually the numbers stored by SLiM and computed by us may differ by `1e-10` or so).
 
 
 ```{code-cell}
@@ -408,3 +408,68 @@ print(f"SLiM: {np.array([x['phenotype'] for x in ind.metadata['per_trait']])}")
 
 More elegant code would pull the trait types out of top-level metadata
 and use additive or multiplicative effects accordingly, etcetera.
+
+(sec_phenotypes_technical_details)=
+
+## Technical details
+
+You're better off using SLiM to calculate phenotypes
+than doing it yourself in python.
+There's lots of corner cases, and
+phenotypes can even be modified by the SLiM script directly.
+But, there's some situations where it's important to know about those corner cases.
+
+First: phenotypes are determined from the various contributions by
+either a sum ("additive" or "logistic") or a product ("multiplicative").
+Logistic traits are then transformed.
+The contributions come from the global `baselineOffset`,
+from individual offsets, and cumulative mutation effects.
+
+The contribution of 0, 1, or 2 copies of a SLiM mutation to a diploid
+is either 0, 2hs, or 2s (additive) or 1, 1+hs, 1+s (multiplicative),
+where h is the dominance coefficient and s is the effect size.
+For a hemizygous individual, the contributions are the same (for 0 and 1 copies),
+but using the hemizygous dominance coefficient for h.
+For an individual without the chromosome at all, there is no effect (obviously).
+
+Usually, that's all we need to know.
+However, there's some more complications that affect the default trait
+(i.e., the trait you get - called `simT` - if you don't explicitly declare any traits),
+or if you have set up a trait without baseline accumulation
+and with mutations that convert to substitutions.
+This can also be important to understand if you remove fixed mutations in python,
+and then want to read the file back into SLiM:
+to keep phenotypes the same, you have to include the effects of any removed mutations
+that count towards the trait in the `baselineOffset`.
+
+Which mutations count?
+Within SLiM, when a Mutation fixes, it *might* convert into a Substitution.
+Whether this occurs is determined by
+the mutation type's `convertToSubstitution` property,
+which is True by default in WF models and False by default in nonWF models.
+Substitutions are no longer used by SLiM for calculating individual traits,
+and so to retain the effects of fixed mutations,
+sometimes scripts will set `convertToSubstitution=F` if it was not already set.
+Another mechanism to retain these effects is the Trait property `accumulateBaseline`,
+which defaults to True.
+If a trait has `accumulateBaseline=T`, then the effects of any Substitutions
+are accumulated in the "baseline offset", and thus contribute to every individual's
+phenotype. Within SLiM, these effects will appear in the `baselineOffset` property.
+(However, these are *not* included in the `baselineOffset` value that we pull from
+top-level metadata of a tree sequence: see below.)
+The net effect of this is that the contributions of any mutations that fix
+will be included in the phenotype values of all individuals
+as long as `convertToSubstitution=F` for their mutation type,
+or `accumulateBaseline=T` for the trait, or both.
+
+However, the tree sequence does not record whether a given mutation was converted to
+a Substitution or not, and also does not record anything about mutation types,
+including their `convertToSubstitution` property.
+(This is because genome structure is set up before .trees files can be loaded.)
+So, when we find a fixed mutation in the tree sequence, we don't know *a priori*
+if this is a Substitution or not.
+If `accumulateBaseline=T` for the trait
+(information that *is* in the top-level metadata),
+then it doesn't matter, mutations will contribute either way.
+The potentially missing information is in the script: are these mutation types
+converting to substitutions or not?

tests/recipe_specs.py

@@ -27,10 +27,23 @@
         "pedigree": True,
         "record_mutations": True,
         "minimal": True,
+        "traits": True,  # just the default
+    },
+    "recipe_nonWF_X.slim": {
+        "nonWF": True,
+        "pedigree": True,
+        "X": True,
+        "minimal": True,
+        "traits": True,
+    },
+    "recipe_nonWF_Y.slim": {"nonWF": True, "pedigree": True, "Y": True, "traits": True},
+    "recipe_nonWF_H.slim": {
+        "nonWF": True,
+        "pedigree": True,
+        "H": True,
+        "minimal": True,
+        "traits": True,
     },
-    "recipe_nonWF_X.slim": {"nonWF": True, "pedigree": True, "X": True, "minimal": True},
-    "recipe_nonWF_Y.slim": {"nonWF": True, "pedigree": True, "Y": True},
-    "recipe_nonWF_H.slim": {"nonWF": True, "pedigree": True, "H": True, "minimal": True},
     "recipe_WF_X.slim": {
         "WF": True,
         "pedigree": True,
@@ -219,11 +232,24 @@
         "multichrom": True,
         "H-": True,
     },
-    "recipe_with_traits.slim": {
+    "recipe_with_traits_single_chrom.slim": {
         "WF": True,
         "traits": True,
+        "begun_late": True,
+    },
+    "recipe_with_traits_simple.slim": {
+        "WF": True,
+        "traits": True,
+        "begun_late": True,
         "multichrom": True,
+        "X": True,
+        "H": True,
+    },
+    "recipe_with_traits.slim": {
+        "WF": True,
+        "traits": True,
         "begun_late": True,
+        "multichrom": True,
         "X": True,
         "Y": True,
         "H": True,

tests/test_metadata.py

@@ -2,10 +2,13 @@
 Test cases for the metadata reading/writing of pyslim.
 """
 
+import dataclasses
 import sys
+from collections import Counter
 
 import numpy as np
 import pytest
+import scipy.sparse as sparse
 import tskit
 
 import pyslim
@@ -424,3 +427,241 @@ def test_chromosome_types(self, recipe):
         for chrom in chroms.keys():
             assert chrom in chrom_list_d
             assert chrom_list_d[chrom] == chrom_info[chrom]
+
+
+class TraitCalculator:
+    """
+    Calculating phenotypes by hand from metadata takes a good bit of
+    internal state, especially if it's from a multi-chromosome simulation,
+    so here's a class that does it.  See below for how to use this.
+    """
+
+    def __init__(self):
+        self.ts_metadata = None
+        self.slim_phenotypes = None
+        self.offsets = None
+
+    def add_ts(self, ts):
+        ts_metadata = ts.metadata
+        if self.ts_metadata is not None:
+            assert ts_metadata["SLiM"]["traits"] == self.ts_metadata["SLiM"]["traits"]
+        self.ts_metadata = ts_metadata
+        self.ts = ts
+        chrom_type = self.ts_metadata["SLiM"]["this_chromosome"]["type"]
+        self.haploid_chromosome = chrom_type in ("H", "HF", "HM")
+        self.mut_metadata = pyslim.mutation_metadata(ts)
+        self.nodes_vacant = pyslim.nodes_vacant(ts)
+        self.num_traits = len(self.ts_metadata["SLiM"]["traits"])
+        self.types = [m["type"] for m in self.ts_metadata["SLiM"]["traits"]]
+        self.accumulates = [
+            m["baselineAccumulation"] for m in self.ts_metadata["SLiM"]["traits"]
+        ]
+        # to properly do traits in the case where baselineAccumulation=F and there are
+        # mutation types with convertToSubstitution=T then we need to store this in metadata:
+        # (so, this code won't work if this is not present for such traits, which are probably
+        # only the default trait!)
+        if (
+            "user_metadata" in ts_metadata["SLiM"]
+            and "mutation_types" in ts_metadata["SLiM"]["user_metadata"]
+        ):
+            self.mutation_types = {
+                k: v[0]["convertToSubstitution"][0]
+                for k, v in ts_metadata["SLiM"]["user_metadata"]["mutation_types"][
+                    0
+                ].items()
+            }
+        else:
+            # guess T if there just a single trait without baseline accumulation
+            # whose name is simT and this is a WF model
+            # since that describes "default trait"
+            convert = (
+                self.ts_metadata["SLiM"]["model_type"] == "WF"
+                and len(self.ts_metadata["SLiM"]["traits"]) == 1
+                and self.ts_metadata["SLiM"]["traits"][0]["name"] == "simT"
+                and self.ts_metadata["SLiM"]["traits"][0]["baselineAccumulation"]
+                == False
+            )
+            print("convert!", convert)
+            mtypes = set(md["mutation_type"] for md in self.mut_metadata.values())
+            self.mutation_types = {f"m{k}": convert for k in mtypes}
+        sp = np.column_stack(
+            [
+                self.ts.tables.individuals.metadata_vector(["per_trait", j, "phenotype"])
+                for j in range(self.num_traits)
+            ]
+        )
+        if self.slim_phenotypes is None:
+            self.slim_phenotypes = sp
+        assert np.all(self.slim_phenotypes == sp)
+        # initialize phenotypes where we'll accumulate changes
+        self.phenotypes = np.full(
+            self.slim_phenotypes.shape,
+            [1.0 if t == "multiplicative" else 0.0 for t in self.types],
+            dtype="float",
+        )
+        offsets = np.column_stack(
+            [
+                self.ts.tables.individuals.metadata_vector(["per_trait", j, "offset"])
+                for j in range(self.num_traits)
+            ]
+        )
+        if self.offsets is None:
+            self.offsets = offsets
+            # only do offsets if this is the first ts!
+            self.do_offsets()
+        else:
+            assert np.all(offsets == self.offsets)
+        self.get_frequencies()
+        self.do_effects()
+
+    def transform(self):
+        for k, t in enumerate(self.types):
+            if t == "logistic":
+                self.phenotypes[:, k] = 1 / (1 + np.exp(-self.phenotypes[:, k]))
+
+    def is_converted(self, a):
+        # is allele a from a convert-to-substitution mutation type?
+        # if we haven't recorded whether or not it's converted, assume T
+        md = self.mut_metadata[int(a)]
+        return self.mutation_types[f"m{md['mutation_type']}"]
+
+    def get_frequencies(self):
+        # Fixed mutations count unless they are converted to substitutions
+        # and baseline accumulation is off, so we need to look up both these things:
+        # (frequency, whether they're converted to substitutions) for each mutation.
+        self.frequencies = {}
+        for v in self.ts.variants(isolated_as_missing=False):
+            freqs = {}
+            for ds, f in v.counts().items():
+                for a in ds.split(","):
+                    if a != "":
+                        if a not in freqs:
+                            convert = self.is_converted(a)
+                            freqs[a] = [0, convert]
+                        freqs[a][0] += f
+            self.frequencies[v.site.id] = freqs
+
+    def do_offsets(self):
+        for k, t in enumerate(self.types):
+            bO = self.ts_metadata["SLiM"]["traits"][k]["baselineOffsetFromUser"]
+            if t == "multiplicative":
+                self.phenotypes[:, k] *= bO * self.offsets[:, k]
+            else:
+                self.phenotypes[:, k] += bO + self.offsets[:, k]
+
+    def do_effects(self):
+        for ind in self.ts.individuals():
+            ploidy = np.sum(~self.nodes_vacant[ind.nodes])
+            if ploidy == 0:
+                continue
+            else:
+                hemizygous = ploidy == 1
+                for k, t in enumerate(self.types):
+                    if t == "multiplicative":
+                        g = self.multiplicative_effect(ind, k, hemizygous)
+                        self.phenotypes[ind.id, k] *= g
+                    else:
+                        g = self.additive_effect(ind, k, hemizygous)
+                        self.phenotypes[ind.id, k] += g
+
+    def skip_mut(self, m, freqs, trait_id):
+        if m == "":
+            out = True
+        else:
+            f, convert = freqs[m]
+            out = f == 0 or (
+                f == self.ts.num_samples and convert and not self.accumulates[trait_id]
+            )
+        return out
+
+    def additive_effect(self, ind, trait_id, hemizygous):
+        out = 0.0
+        for v in self.ts.variants(samples=ind.nodes, isolated_as_missing=False):
+            freqs = self.frequencies[v.site.id]
+            a = ",".join([v.alleles[g] for g in v.genotypes])
+            muts = Counter(a.split(","))
+            for m in muts:
+                if self.skip_mut(m, freqs, trait_id):
+                    continue
+                md = self.mut_metadata[int(m)]["per_trait"][trait_id]
+                s = md["effect_size"]
+                if muts[m] == 2:
+                    h = 1
+                else:
+                    assert muts[m] == 1
+                    if self.haploid_chromosome:
+                        h = 1
+                    else:
+                        if hemizygous:
+                            h = md["hemizygous_dominance"]
+                        else:
+                            h = md["dominance"]
+                        if np.isnan(h):
+                            h = 1 / 2
+                out += 2 * h * s
+        return out
+
+    def multiplicative_effect(self, ind, trait_id, hemizygous):
+        out = 1.0
+        for v in self.ts.variants(samples=ind.nodes, isolated_as_missing=False):
+            freqs = self.frequencies[v.site.id]
+            a = ",".join([v.alleles[g] for g in v.genotypes])
+            muts = Counter(a.split(","))
+            for m in muts:
+                if self.skip_mut(m, freqs, trait_id):
+                    continue
+                assert muts[m] > 0 and muts[m] <= 2
+                md = self.mut_metadata[int(m)]["per_trait"][trait_id]
+                s = md["effect_size"]
+                if muts[m] == 2:
+                    h = 1
+                else:
+                    assert muts[m] == 1
+                    if self.haploid_chromosome:
+                        h = 1
+                    elif hemizygous:
+                        h = md["hemizygous_dominance"]
+                    else:
+                        h = md["dominance"]
+                    if np.isnan(h):
+                        # "independent dominance occurs when (1+hs)(1+hs) equals 1+s,
+                        # which occurs when h=(sqrt(1+s)−1)/s"
+                        h = (np.sqrt(1 + s) - 1) / s if s != 0 else 0
+                out *= 1 + h * s
+        return out
+
+
+class TestTraits(tests.PyslimTestCase):
+    @pytest.mark.parametrize("recipe", recipe_eq("traits"), indirect=True)
+    def test_traits_consistency(self, recipe):
+        trait_md = None
+        offsets = None
+        slim_phenotypes = None
+        tc = TraitCalculator()
+        for ts in recipe["ts"].values():
+            assert ts.num_mutations > 0
+            ts_metadata = ts.metadata
+            # top-level info about traits
+            if trait_md is None:
+                trait_md = ts_metadata["SLiM"]["traits"]
+            assert trait_md == ts_metadata["SLiM"]["traits"]
+            # now compute them
+            tc.add_ts(ts)
+        # phenotypes as computed by us
+        tc.transform()
+        print([ind.metadata["sex"] for ind in ts.individuals()])
+        for k in range(len(tc.types)):
+            print(tc.types[k])
+            print(
+                np.column_stack(
+                    [
+                        tc.phenotypes[:, k],
+                        tc.slim_phenotypes[:, k],
+                        tc.phenotypes[:, k] - tc.slim_phenotypes[:, k],
+                    ]
+                )
+            )
+        print(np.max(tc.phenotypes - tc.slim_phenotypes))
+        print(np.min(tc.phenotypes - tc.slim_phenotypes))
+        print(np.sum(tc.phenotypes - tc.slim_phenotypes))
+        assert np.allclose(tc.phenotypes, tc.slim_phenotypes)

tests/test_recipes/recipe_WF_H.slim

@@ -29,6 +29,7 @@ initialize()
 }
 
 10 late() {
+    sim.demandPhenotype(sim.subpopulations);
     sim.treeSeqOutput(TREES_FILE);
     catn("Done.");
     sim.simulationFinished();