diff --git a/.gitignore b/.gitignore
index f6bf4a8..b5f9c63 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,6 +10,7 @@ __pycache__/
.Python
env/
build/
+_build/
develop-eggs/
dist/
downloads/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 161a66d..7ffcd1e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,19 @@
In development
+### Highlights
+
+- Extended `genetic_value` with `level="individual"` (default), `level="node"`,
+ and `level="edge"` to return genetic values for the corresponding entities.
+ {pr}`155`
+- Added `edge_effect` to compute introduced effects on edges {pr}`155`
+
+### Documentation
+
+- Added a worked example relating causal-allele effects,
+ edge effects, and edge, node, and individual genetic values {pr}`155`
+- Clarified that tstrait currently uses a site-mode effect model {pr}`155`
+
## [0.1.2] - 2026-03-03
Maintenance release
diff --git a/README.md b/README.md
index 9c94000..323f064 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
[](https://github.com/tskit-dev/tstrait/blob/main/LICENSE) [](https://pypi.org/project/tstrait/) [](https://pypi.org/project/tstrait/) [](https://github.com/tskit-dev/tstrait/actions/workflows/docs.yml) [](https://github.com/tskit-dev/tstrait/actions/workflows/tests.yml) [](https://codecov.io/gh/tskit-dev/tstrait)
-**tstrait** is a quantitative trait simulator of [tskit](https://tskit.dev) tree sequences. It supports simulation of quantitative traits, sampling genetic effect sizes from various distributions, modelling pleiotropic phenotypes and environmental noise, and computing per‑individual genetic values given the effect sizes of causal mutations.
+**tstrait** is a quantitative trait simulator of [tskit](https://tskit.dev) tree sequences. It supports simulation of quantitative traits, sampling genetic effect sizes from various distributions, modelling pleiotropic phenotypes and environmental noise, and computing per-individual, per-node, and per-edge genetic values given the effect sizes of causal mutations.
Documentation: https://tskit.dev/tstrait/docs/stable/ (latest: https://tskit.dev/tstrait/docs/latest/)
Learn more about tree sequences: https://tskit.dev/learn/
diff --git a/docs/_static/individual_node_edge_alleles.svg b/docs/_static/individual_node_edge_alleles.svg
new file mode 100644
index 0000000..01b1526
--- /dev/null
+++ b/docs/_static/individual_node_edge_alleles.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/_static/individual_node_edge_values.svg b/docs/_static/individual_node_edge_values.svg
new file mode 100644
index 0000000..270a197
--- /dev/null
+++ b/docs/_static/individual_node_edge_values.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/_toc.yml b/docs/_toc.yml
index a394516..8020087 100644
--- a/docs/_toc.yml
+++ b/docs/_toc.yml
@@ -11,6 +11,7 @@ parts:
chapters:
- file: effect-size
- file: genetic
+ - file: individual_node_edge
- file: environment
- file: multi-trait
@@ -24,4 +25,4 @@ parts:
- file: contributing
- file: CHANGELOG
- url: https://tskit.dev/community/
- title: Community
\ No newline at end of file
+ title: Community
diff --git a/docs/effect-size.md b/docs/effect-size.md
index 52483bd..4e9599c 100644
--- a/docs/effect-size.md
+++ b/docs/effect-size.md
@@ -53,10 +53,18 @@ noise is simulated from the following distribution,
where $V_G=Var(X\beta)$ and $h^2$ is the narrow-sense heritability that is defined by the
user.
-The genetic values ($X\beta$) are obtained by simply adding up over all the genomes in each individual,
-regardless of ploidy.
+The genetic values ($X\beta$) are obtained by adding up the effects of causal alleles
+over all the genomes in each individual, regardless of ploidy.
+
+These genetic values correspond to
+the {ref}`**site mode** of tree sequence statistics `.
+That is, contribution to $X\beta$ depends on the allelic state at causal sites,
+not on which mutation event produced the allele.
+Branch and mutation modes are currently not implemented in tstrait
+and may be added in the future.
:::{seealso}
+
- [](genetic_value_doc) for obtaining the genetic value $X\beta$.
- [](environment_noise) for simulating environmental noise $\epsilon$.
- [](numericalise_genotype) for modifying the numericalisation of genotypes.
@@ -89,7 +97,7 @@ model.name
```
tstrait uses this distribution to simulate effect size $\beta$ in the phenotype model described
-in [](eq:phenotype-model).
+in [](eq:phenotype-model).
The following effect size distributions are supported in tstrait, and please refer to links under
**Details** for details on the input and distribution.
@@ -205,7 +213,6 @@ ts.site(0)
```
The details of sites in tree sequences can be found {ref}`here`.
-
(trait_frequency_dependence)=
## Frequency Dependence
@@ -261,7 +268,7 @@ it in your trait simulation.
# trait.sim_trait(ts, num_causal=1000, model=model, random_seed=1)
# also works here
-trait_df = tstrait.sim_trait(ts, num_causal=1000, model=model,
+trait_df = tstrait.sim_trait(ts, num_causal=1000, model=model,
alpha=0, random_seed=1)
plt.scatter(trait_df.allele_freq, trait_df.effect_size)
@@ -277,7 +284,7 @@ we will be simulating effect sizes with `alpha` = -1/2.
```{code-cell}
-trait_df = tstrait.sim_trait(ts, num_causal=1000, model=model,
+trait_df = tstrait.sim_trait(ts, num_causal=1000, model=model,
alpha=-1/2, random_seed=1)
plt.scatter(trait_df.allele_freq, trait_df.effect_size)
@@ -305,7 +312,7 @@ Example:
```{code-cell}
-trait_df = tstrait.sim_trait(ts, causal_sites=[0, 3, 4], model=model,
+trait_df = tstrait.sim_trait(ts, causal_sites=[0, 3, 4], model=model,
alpha=-1/2, random_seed=1)
trait_df
-```
\ No newline at end of file
+```
diff --git a/docs/environment.md b/docs/environment.md
index 5547add..753a6b0 100644
--- a/docs/environment.md
+++ b/docs/environment.md
@@ -159,4 +159,4 @@ sim_result = tstrait.sim_phenotype(
sim_result.phenotype.head()
```
-We see that all values in the `environmental_noise` column are zero.
\ No newline at end of file
+We see that all values in the `environmental_noise` column are zero.
diff --git a/docs/individual_node_edge.md b/docs/individual_node_edge.md
new file mode 100644
index 0000000..5567191
--- /dev/null
+++ b/docs/individual_node_edge.md
@@ -0,0 +1,268 @@
+---
+kernelspec:
+ name: python3
+ display_name: python3
+jupytext:
+ text_representation:
+ extension: .md
+ format_name: myst
+ format_version: '0.13'
+ jupytext_version: 1.13.8
+---
+
+```{eval-rst}
+.. currentmodule:: tstrait
+```
+
+(genetic_individual_node_edge_doc)=
+
+# Individual, node, and edge genetic values
+
+In line with the {ref}`tree sequence data model`,
+this page describes relationships between
+individual, node, and edge genetic values,
+and underlying edge effects arising from mutations.
+
+**Learning Objectives**
+
+After this page, you will be able to:
+
+- Understand how to obtain genetic values in tstrait for individuals, nodes, and edges
+- Understand how causal-allele effects give rise to edge effects and
+ the genetic values.
+
+# Algorithm Overview
+
+The tstrait algorithm for {ref}`individuals' genetic values `
+can also be used to compute node and edge genetic values,
+and there is the related concept of edge effects;
+all following the assumed {ref}`quantitative genetics model `.
+These quantities are related as follows:
+
+- *edge effect* is the sum of effects of causal alleles introduced on the edge by mutations.
+- *edge "genetic" value* is its edge effect plus a sum of edge effects ancestral to the edge.
+- *node "genetic" value* is a sum of "genetic" values of immediately ancestral edges into the node.
+- *individual genetic value* is a sum of "genetic" values of nodes of the individual.
+
+Above we have quoted the term "genetic" when referring to nodes and edges,
+because the term "genetic value" is usually used with individuals.
+Since individuals' genetic value is a sum over the contributions of its nodes,
+the term "genetic value" applies also to nodes, and hence also to immediately ancestral edges of the nodes.
+
+More on terminology.
+Edges are also called branches or haplotypes,
+so we can call
+edge effect as branch effect or haplotype effect and
+edge genetic value as branch genetic value or haplotype genetic value.
+
+# Example
+
+We will demonstrate the concepts of edge effects and
+edge, node, and individual genetic values with
+a tiny example so we can follow the calculations.
+Tree sequence for the example is shown in the following figure.
+It represents two diploid individuals and how their four genomes (nodes)
+are connected with edges to four ancestral nodes and eight mutation events.
+Mutation events and their downstream effects of allele change from ancestral state
+are denoted with a red colour.
+One node is a recombinant,
+meaning that this tree sequence has two local trees.
+We denote the first (second) local tree by shading the nodes' first (second) region with black (white) colour.
+To help tracking IDs across the different elements of a tree sequence
+in the following text, we prepend a letter to the zero-based IDs
+(`iI` for individual I, `nN` for node N, `eE` for edge E, and `mM@S` for mutation M at site S).
+
+```{figure} _static/individual_node_edge_alleles.svg
+:name: fig_individual_node_edge_alleles
+:width: 90%
+:align: center
+:alt: Diagram of a tree sequence with two individuals (`i0` and `i1`), eight nodes (`n0-n7`; `n0-n3` are sample nodes and `n4-n7` are ancestral nodes), eight edges (`e0-e7`; with a "virtual" root edge `e8`), eight mutations (`m0-m7` [each at a different site]; in red colour), and resulting node alleles across ten sites. The first (second) local tree is denoted with black (white) colour.
+
+Diagram of a tree sequence with two individuals (`i0` and `i1`), eight nodes (`n0-n7`; `n0-n3` are sample nodes and `n4-n7` are ancestral nodes), eight edges (`e0-e7`; with a "virtual" root edge `e8`), eight mutations (`m0-m7` [each at a different site]; in red colour), and resulting node alleles across ten sites. The first (second) local tree is denoted with black (white) colour.
+```
+
+We now create tables for the example and
+combined them into a tree sequence object:
+
+```{code-cell}
+import io
+import tskit
+import tstrait
+import pandas as pd
+
+individuals = io.StringIO(
+ """\
+ flags location parents
+ 0 0 -1
+ 0 0 -1
+ """
+)
+
+nodes = io.StringIO(
+ """\
+ id is_sample time individual
+ 0 1 0 0
+ 1 1 0 0
+ 2 1 0 1
+ 3 1 0 1
+ 4 0 1 -1
+ 5 0 1 -1
+ 6 0 10 -1
+ 7 0 1000 -1
+ """
+)
+
+edges = io.StringIO(
+ """\
+ left right parent child
+ 0 10 4 0
+ 0 5 4 1
+ 5 10 5 1
+ 0 10 5 2
+ 0 10 6 3
+ 0 10 6 5
+ 0 10 7 4
+ 0 10 7 6
+ """
+)
+
+sites = io.StringIO(
+ """\
+ position ancestral_state
+ 0 0
+ 1 0
+ 2 0
+ 3 0
+ 4 0
+ 5 0
+ 6 0
+ 7 0
+ 8 0
+ 9 0
+ """
+)
+
+mutations = io.StringIO(
+ """\
+ site node time derived_state parent
+ 0 4 750.0 1 -1
+ 1 3 7.0 1 -1
+ 2 1 0.5 1 -1
+ 4 4 500.0 1 -1
+ 6 6 600.0 1 -1
+ 7 5 9.0 1 -1
+ 8 3 2.0 1 -1
+ 9 1 0.5 1 -1
+ """
+)
+
+ts = tskit.load_text(individuals = individuals,
+ nodes = nodes,
+ edges = edges,
+ sites = sites,
+ mutations = mutations,
+ strict = False)
+print(ts)
+```
+
+We can inspect the tables in the tree sequence and
+IDs of its elements using:
+
+```{code-cell}
+print(ts.tables.individuals)
+print(ts.tables.nodes)
+print(ts.tables.edges)
+print(ts.tables.sites)
+print(ts.tables.mutations)
+```
+
+The resulting alleles across ten sites for the four sample nodes are:
+
+```{code-cell}
+ts.genotype_matrix().T
+```
+
+We now create a data frame with trait information:
+causal site, effect, trait, and allele.
+In this example all eight polymorphic sites are causal sites.
+
+```{code-cell}
+data = {"site_id": [ 0, 1, 2, 4, 6, 7, 8, 9],
+ "effect_size": [ 1, 1, -1, 1, -1, -1, -2, 1],
+ "trait_id": [ 0, 0, 0, 0, 0, 0, 0, 0],
+ "causal_allele": ["1", "1", "1", "1", "1", "1", "1", "1"]}
+trait_df = pd.DataFrame(data)
+print(trait_df)
+```
+
+Combining the tree sequence and trait information
+generates the edge effects, edge genetic values, node genetic values, and
+individuals' genetic values as shown in the figure below.
+We will now demonstrate how these values
+are generated from the tree sequence and trait information.
+
+```{figure} _static/individual_node_edge_values.svg
+:name: fig_individual_node_edge_values
+:width: 90%
+:align: center
+:alt: Diagram of the tree sequence with mutation effects generating edge effects (both in red colour), which in turn generate edge genetic values (in green colour), node genetic values (in blue colour), and individuals' genetic values (in purple colour).
+
+Diagram of the tree sequence with mutation effects generating edge effects (both in red colour), which in turn generate edge genetic values (in green colour), node genetic values (in blue colour), and individuals' genetic values (in purple colour).
+```
+
+We obtain **edge effects** by using the {py:func}`edge_effect` function:
+
+```{code-cell}
+edge_effect_df = tstrait.edge_effect(ts, trait_df)
+print(edge_effect_df)
+```
+
+For example, on edge 6, connecting the parent node 7 and child node 4,
+mutation 0 at site 0 and mutation 3 at site 4 introduce causal alleles,
+each with effect `1` unit, so the joint edge effect is `1 + 1 = 2` units.
+Edge 0, connecting the parent node 4 and child node 0,
+has no mutations, so its edge effect is `0` units.
+
+We obtain **edge genetic values** by using the `level="edge"` argument
+of the {py:func}`genetic_value` function:
+
+```{code-cell}
+edge_value_df = tstrait.genetic_value(ts, trait_df, level="edge")
+print(edge_value_df)
+```
+
+For example, above we saw that edge 6 has an effect of `2` units and
+since the root node 7 has no causal alleles it's genetic value is `0` units,
+so edge 6 has genetic value of `0 + 2 = 2` units.
+Note that both mutations on edge 6 are in interval `[0,5)`.
+This means that when this interval is passed to node 1 via node 4,
+edge 1 inherits genetic value of `2` units,
+but it's effect of `-1` unit changes its genetic value to `2 + (-1) = 1` unit.
+
+We obtain **node genetic values** by using the `level="node"` argument
+of the {py:func}`genetic_value` function:
+
+```{code-cell}
+node_value_df = tstrait.genetic_value(ts, trait_df, level="node")
+print(node_value_df)
+```
+
+For example, above we saw that the root node 7 has genetic value of `0` units
+and that mutations on edge 6 changed it to `2` units for node 4 and
+consequently for node 0 (since there were no causal mutations on edge 0).
+Note that node 1 is recombinant:
+it inherits
+edge 1 from node 4 over `[0, 5)`, with genetic value of `1` unit and
+edge 2 from node 5 over `[5, 10)`, with genetic value of `-1` unit.
+Node 1 genetic value is therefore `1 + (-1) = 0` units.
+
+We obtain **individual genetic values** by using the `level="individual"` argument
+of the {py:func}`genetic_value` function, which is the default level:
+
+```{code-cell}
+individual_value_df = tstrait.genetic_value(ts, trait_df)
+print(individual_value_df)
+```
+
+Individual 0 has nodes 0 and 1, so its genetic value is `2 + 0 = 2` units.
+Individual 1 has nodes 2 and 3, so its genetic value is `-2 + (-2) = -4` units.
diff --git a/docs/intro.md b/docs/intro.md
index 9907f87..7dda87c 100644
--- a/docs/intro.md
+++ b/docs/intro.md
@@ -20,7 +20,9 @@ data structures.
Quantitative trait simulation in tstrait is transparent, and users can control each step in the simulation. Thus,
it is possible for the users to simulate their own environmental noise on top of simulated genetic values,
-or even use their own defined effect sizes and causal sites. The tree sequence data structure is widely used in
+or even use their own defined effect sizes and causal sites.
+Users can also study how tree sequence nodes and edges contribute to individuals' genetic values.
+The tree sequence data structure is widely used in
various population genetic simulation packages, including [SLiM](https://messerlab.org/slim/),
{ref}`msprime`, and {ref}`stdpopsim`; it is therefore easy for
users of these packages to add quantitative traits to their results using tstrait.
@@ -31,7 +33,7 @@ To learn more about tree sequences:
- The [tskit website](https://tskit.dev/) provides [learning materials](https://tskit.dev/learn/) explaining
what tree sequences are, and includes tutorials, publications and videos.
- - The {ref}`PySLiM manual` explains how forward genetic simulation can be create
+- The {ref}`PySLiM manual` explains how forwards-in-time genetic simulation can create
tree sequences.
- The {ref}`msprime manual` details an efficient backward-time genetic simulator that outputs
tree sequences.
diff --git a/tests/test_genetic_value.py b/tests/test_genetic_value.py
index f44785f..715e77a 100644
--- a/tests/test_genetic_value.py
+++ b/tests/test_genetic_value.py
@@ -5,7 +5,6 @@
import tstrait
from tstrait.base import _check_numeric_array
-from tstrait.genetic_value import _GeneticValue
from .data import (
binary_tree,
@@ -21,7 +20,7 @@
@pytest.fixture(scope="class")
def sample_ts():
ts = msprime.sim_ancestry(10, sequence_length=100_000, random_seed=1)
- ts = msprime.sim_mutations(ts, rate=0.01, random_seed=1)
+ ts = msprime.sim_mutations(ts, rate=1e-2, random_seed=1)
return ts
@@ -45,6 +44,20 @@ def sample_trait_model():
return tstrait.trait_model(distribution="normal", mean=0, var=1)
+def individual_genetic_values(tree, site, causal_allele, effect_size):
+ """Return individual values for one causal site via the public API."""
+ trait_df = pd.DataFrame(
+ {
+ "site_id": [site.id],
+ "effect_size": [effect_size],
+ "trait_id": [0],
+ "causal_allele": [causal_allele],
+ }
+ )
+ result = tstrait.genetic_value(ts=tree.tree_sequence, trait_df=trait_df)
+ return result["genetic_value"].to_numpy()
+
+
class TestInput:
"""This test will check that an informative error is raised when the input parameter
does not have an appropriate type or value.
@@ -79,22 +92,149 @@ def test_bad_input(self, sample_ts, sample_df):
df["site_id"] = [2, 0]
tstrait.genetic_value(ts=sample_ts, trait_df=df)
+ @pytest.mark.parametrize("function", [tstrait.genetic_value, tstrait.edge_effect])
@pytest.mark.parametrize("trait_id", [[2, 3], [0, 2]])
- def test_trait_id(self, sample_ts, sample_df, trait_id):
+ def test_trait_id(self, sample_ts, sample_df, function, trait_id):
with pytest.raises(
ValueError, match="trait_id must be consecutive and start from 0"
):
df = sample_df.copy()
df["trait_id"] = trait_id
- tstrait.genetic_value(ts=sample_ts, trait_df=df)
+ function(ts=sample_ts, trait_df=df)
def test_no_individual(self, sample_df):
- ts = msprime.simulate(100, length=10000, mutation_rate=1e-2)
- df = sample_df.copy()
+ ts = msprime.simulate(10, length=100, random_seed=1)
+ # legacy msprime.simulate() creates a tree sequence without individuals
with pytest.raises(
ValueError, match="No individuals in the provided tree sequence dataset"
):
- tstrait.genetic_value(ts=ts, trait_df=df)
+ tstrait.genetic_value(ts=ts, trait_df=sample_df)
+
+ @pytest.mark.parametrize(
+ ("level", "id_column", "size_attribute"),
+ [
+ ("node", "node_id", "num_nodes"),
+ ("edge", "edge_id", "num_edges"),
+ ],
+ )
+ def test_node_or_edge_level_without_individuals(
+ self, sample_df, level, id_column, size_attribute
+ ):
+ ts = msprime.simulate(10, length=100, mutation_rate=1e-2, random_seed=1)
+ result = tstrait.genetic_value(ts=ts, trait_df=sample_df, level=level)
+
+ size = getattr(ts, size_attribute)
+ assert list(result.columns) == ["trait_id", id_column, "genetic_value"]
+ assert len(result) == size
+ np.testing.assert_array_equal(result["trait_id"], np.zeros(size))
+ np.testing.assert_array_equal(result[id_column], np.arange(size))
+
+ def test_invalid_level(self, sample_ts, sample_df):
+ with pytest.raises(
+ ValueError,
+ match="level must be one of 'individual', 'node', or 'edge'",
+ ):
+ tstrait.genetic_value(ts=sample_ts, trait_df=sample_df, level="bla")
+
+ def test_empty_trait_df(self, sample_ts, sample_df):
+ empty_trait_df = sample_df.iloc[0:0]
+ with pytest.raises(ValueError, match="trait_df must contain at least one row"):
+ tstrait.genetic_value(ts=sample_ts, trait_df=empty_trait_df)
+ with pytest.raises(ValueError, match="trait_df must contain at least one row"):
+ tstrait.edge_effect(ts=sample_ts, trait_df=empty_trait_df)
+
+
+class TestEdgeEffect:
+ @pytest.mark.parametrize(
+ ("ancestral_state", "derived_state", "expected_effect"),
+ [
+ pytest.param("A", "T", 2, id="non_causal_to_causal"),
+ pytest.param("T", "A", -2, id="causal_to_non_causal"),
+ pytest.param("A", "G", 0, id="non_causal_to_non_causal"),
+ pytest.param("T", "T", 0, id="causal_to_causal"),
+ ],
+ )
+ def test_causal_allele_state_transition(
+ self, ancestral_state, derived_state, expected_effect
+ ):
+ ts = binary_tree()
+ tables = ts.dump_tables()
+ tables.mutations.clear()
+ tables.sites.clear()
+ site_id = tables.sites.add_row(position=0, ancestral_state=ancestral_state)
+ mutation_node = 4
+ tables.mutations.add_row(
+ site=site_id, node=mutation_node, derived_state=derived_state
+ )
+ ts = tables.tree_sequence()
+
+ trait_df = pd.DataFrame(
+ {
+ "site_id": [site_id],
+ "effect_size": [2],
+ "trait_id": [0],
+ "causal_allele": ["T"],
+ }
+ )
+ result = tstrait.edge_effect(ts=ts, trait_df=trait_df)
+
+ expected = np.zeros(ts.num_edges)
+ expected[ts.first().edge(mutation_node)] = expected_effect
+ np.testing.assert_array_equal(result["effect_size"], expected)
+
+ def test_site_mode_allelic_state_transitions(self):
+ ts = binary_tree()
+ tables = ts.dump_tables()
+ tables.mutations.clear()
+ tables.sites.clear()
+ site_id = tables.sites.add_row(position=0, ancestral_state="A")
+ mutation_id = tables.mutations.add_row(site=site_id, node=4, derived_state="T")
+ tables.mutations.add_row(
+ site=site_id, node=0, derived_state="A", parent=mutation_id
+ )
+ tables.mutations.add_row(site=site_id, node=3, derived_state="T")
+ ts = tables.tree_sequence()
+
+ trait_df = pd.DataFrame(
+ {
+ "site_id": [site_id],
+ "effect_size": [2],
+ "trait_id": [0],
+ "causal_allele": ["T"],
+ }
+ )
+ result = tstrait.edge_effect(ts=ts, trait_df=trait_df)
+
+ tree = ts.first()
+ expected = np.zeros(ts.num_edges)
+ expected[tree.edge(4)] = 2
+ expected[tree.edge(0)] = -2
+ expected[tree.edge(3)] = 2
+ np.testing.assert_array_equal(result["effect_size"], expected)
+
+ def test_mutation_on_root(self):
+ ts = binary_tree()
+ root = ts.first().root
+ tables = ts.dump_tables()
+ tables.mutations.clear()
+ tables.sites.clear()
+ site_id = tables.sites.add_row(position=0, ancestral_state="A")
+ tables.mutations.add_row(site=site_id, node=root, derived_state="T")
+ ts = tables.tree_sequence()
+ trait_df = pd.DataFrame(
+ {
+ "site_id": [site_id],
+ "effect_size": [1],
+ "trait_id": [0],
+ "causal_allele": ["T"],
+ }
+ )
+
+ with pytest.raises(
+ ValueError,
+ match="Cannot assign an edge effect to a mutation on a root node",
+ ):
+ tstrait.edge_effect(ts=ts, trait_df=trait_df)
class TestOutputDim:
@@ -170,10 +310,7 @@ def test_additional_row(self, sample_ts, sample_df):
class TestGenotype:
- """Test the `_individual_genetic_values` and `_obtain_allele_count` method and
- check if they can accurately detect the individual genotype. Afterwards, we will
- check the output of `genetic_value`.
- """
+ """Test that genetic_value accurately detects individual genotypes."""
def test_binary_tree(self):
trait_df = pd.DataFrame(
@@ -187,11 +324,10 @@ def test_binary_tree(self):
ts = binary_tree()
tree = ts.first()
- genetic = _GeneticValue(ts, trait_df)
- g0 = genetic._individual_genetic_values(tree, ts.site(0), "T", 1)
- g1 = genetic._individual_genetic_values(tree, ts.site(1), "T", 2)
- g2 = genetic._individual_genetic_values(tree, ts.site(2), "C", 3)
- g3 = genetic._individual_genetic_values(tree, ts.site(3), "C", 4)
+ g0 = individual_genetic_values(tree, ts.site(0), "T", 1)
+ g1 = individual_genetic_values(tree, ts.site(1), "T", 2)
+ g2 = individual_genetic_values(tree, ts.site(2), "C", 3)
+ g3 = individual_genetic_values(tree, ts.site(3), "C", 4)
np.testing.assert_equal(g0, np.array([1, 0, 2]) * 1)
np.testing.assert_equal(g1, np.array([1, 1, 0]) * 2)
@@ -222,11 +358,10 @@ def test_diff_ind_tree(self):
ts = diff_ind_tree()
tree = ts.first()
- genetic = _GeneticValue(ts, trait_df)
- g0 = genetic._individual_genetic_values(tree, ts.site(0), "T", 1)
- g1 = genetic._individual_genetic_values(tree, ts.site(1), "T", 2)
- g2 = genetic._individual_genetic_values(tree, ts.site(2), "C", 3)
- g3 = genetic._individual_genetic_values(tree, ts.site(3), "C", 4)
+ g0 = individual_genetic_values(tree, ts.site(0), "T", 1)
+ g1 = individual_genetic_values(tree, ts.site(1), "T", 2)
+ g2 = individual_genetic_values(tree, ts.site(2), "C", 3)
+ g3 = individual_genetic_values(tree, ts.site(3), "C", 4)
np.testing.assert_equal(g0, np.array([1, 1, 1]) * 1)
np.testing.assert_equal(g1, np.array([1, 0, 1]) * 2)
@@ -256,9 +391,8 @@ def test_non_binary_tree(self):
ts = non_binary_tree()
tree = ts.first()
- genetic = _GeneticValue(ts, trait_df)
- g0 = genetic._individual_genetic_values(tree, ts.site(0), "T", 1)
- g1 = genetic._individual_genetic_values(tree, ts.site(1), "C", 2)
+ g0 = individual_genetic_values(tree, ts.site(0), "T", 1)
+ g1 = individual_genetic_values(tree, ts.site(1), "C", 2)
np.testing.assert_equal(g0, np.array([0, 1, 2]) * 1)
np.testing.assert_equal(g1, np.array([0, 1, 1]) * 2)
@@ -275,7 +409,7 @@ def test_non_binary_tree(self):
pd.testing.assert_frame_equal(genetic_df, genetic_result, check_dtype=False)
- def test_triploid(self, sample_df):
+ def test_triploid(self):
trait_df = pd.DataFrame(
{
"site_id": [0, 1],
@@ -287,9 +421,8 @@ def test_triploid(self, sample_df):
ts = triploid_tree()
tree = ts.first()
- genetic = _GeneticValue(ts, sample_df)
- g0 = genetic._individual_genetic_values(tree, ts.site(0), "T", 1)
- g1 = genetic._individual_genetic_values(tree, ts.site(1), "C", 2)
+ g0 = individual_genetic_values(tree, ts.site(0), "T", 1)
+ g1 = individual_genetic_values(tree, ts.site(1), "C", 2)
np.testing.assert_equal(g0, np.array([1, 2]) * 1)
np.testing.assert_equal(g1, np.array([1, 1]) * 2)
diff --git a/tests/test_individual_node_edge.py b/tests/test_individual_node_edge.py
new file mode 100644
index 0000000..c61ee52
--- /dev/null
+++ b/tests/test_individual_node_edge.py
@@ -0,0 +1,279 @@
+"""
+Tests for edge effects and edge, node, and individual genetic values.
+
+The tree sequence is based on the recombining ARG in
+docs/individual_node_edge.md. The expected values can be calculated by hand.
+"""
+
+import io
+
+import numpy as np
+import pandas as pd
+import pytest
+import tskit
+
+import tstrait
+from tstrait.genetic_value import _accumulate_edge_values
+
+
+def test_accumulate_edge_values():
+ nodes_genetic_value = np.array([10, 2, 0, 0, -3, 3, 5, 20, -1])
+ node_edges = np.array(
+ [tskit.NULL, 0, tskit.NULL, 1, 1, 1, 2, tskit.NULL, 2], dtype=int
+ )
+
+ observed = _accumulate_edge_values(
+ nodes_genetic_value=nodes_genetic_value,
+ nodes_edge=node_edges,
+ num_nodes=len(nodes_genetic_value),
+ num_edges=3,
+ )
+
+ np.testing.assert_array_equal(observed, [2, 0, 4])
+
+
+@pytest.fixture(scope="module")
+def individual_node_edge_ts():
+ """
+ Return the small recombining tree sequence from the documentation.
+
+ Individuals 0 and 1 are diploid:
+
+ individual 0: nodes 0 and 1
+ individual 1: nodes 2 and 3
+
+ Node 1 has two immediately ancestral edges because of recombination:
+
+ edge 1: [0, 5), parent 4 -> child 1
+ edge 2: [5, 10), parent 5 -> child 1
+ """
+ individuals = io.StringIO(
+ """\
+flags location parents
+ 0 0 -1
+ 0 0 -1
+"""
+ )
+
+ nodes = io.StringIO(
+ """\
+id is_sample time individual
+ 0 1 0 0
+ 1 1 0 0
+ 2 1 0 1
+ 3 1 0 1
+ 4 0 1 -1
+ 5 0 1 -1
+ 6 0 2 -1
+ 7 0 4 -1
+"""
+ )
+
+ edges = io.StringIO(
+ """\
+left right parent child
+ 0 10 4 0
+ 0 5 4 1
+ 5 10 5 1
+ 0 10 5 2
+ 0 10 6 3
+ 0 10 6 5
+ 0 10 7 4
+ 0 10 7 6
+"""
+ )
+
+ sites = io.StringIO(
+ """\
+position ancestral_state
+ 0 0
+ 1 0
+ 2 0
+ 3 0
+ 4 0
+ 5 0
+ 6 0
+ 7 0
+ 8 0
+ 9 0
+"""
+ )
+
+ mutations = io.StringIO(
+ """\
+site node time derived_state parent
+ 0 4 3.0 1 -1
+ 1 3 1.5 1 -1
+ 2 1 0.5 1 -1
+ 4 4 2.0 1 -1
+ 6 6 3.0 1 -1
+ 7 5 1.5 1 -1
+ 8 3 1.0 1 -1
+ 9 1 0.5 1 -1
+"""
+ )
+
+ ts = tskit.load_text(
+ individuals=individuals,
+ nodes=nodes,
+ edges=edges,
+ sites=sites,
+ mutations=mutations,
+ strict=False,
+ )
+
+ # Check if tree sequence is as we expect it to be
+ assert ts.sequence_length == 10
+ assert ts.num_trees == 2
+ assert ts.num_individuals == 2
+ assert ts.num_nodes == 8
+ assert ts.num_edges == 8
+ assert ts.num_sites == 10
+ assert ts.num_mutations == 8
+
+ return ts
+
+
+def make_trait_df(effect_sizes):
+ """Return a causal-allele effect table for the fixed set of sites."""
+ site_ids = [0, 1, 2, 4, 6, 7, 8, 9]
+
+ return pd.DataFrame(
+ {
+ "site_id": site_ids,
+ "effect_size": effect_sizes,
+ "trait_id": [0] * len(site_ids),
+ "causal_allele": ["1"] * len(site_ids),
+ }
+ )
+
+
+def assert_value_frame(
+ observed,
+ id_column,
+ value_column,
+ expected_values,
+):
+ """Check the IDs, columns, and values of a single-trait output."""
+ assert list(observed.columns) == [
+ "trait_id",
+ id_column,
+ value_column,
+ ]
+
+ observed = observed.sort_values(["trait_id", id_column]).reset_index(drop=True)
+
+ expected = pd.DataFrame(
+ {
+ "trait_id": np.zeros(len(expected_values), dtype=int),
+ id_column: np.arange(len(expected_values)),
+ value_column: np.asarray(expected_values, dtype=float),
+ }
+ )
+
+ pd.testing.assert_frame_equal(
+ observed,
+ expected,
+ check_dtype=False,
+ )
+
+
+def test_edge_node_and_individual_genetic_values(individual_node_edge_ts):
+ """
+ Check the complete edge-effect to individual-value calculation.
+
+ The sequence of calculations is:
+
+ causal-allele effects
+ -> edge effects
+ -> edge genetic values
+ -> node genetic values
+ -> individual genetic values
+ """
+ ts = individual_node_edge_ts
+ # See values in docs/individual_node_edge.md.
+ effect_sizes = [1, 1, -1, 1, -1, -1, -2, 1]
+ expected_edge_effects = [0, -1, 1, 0, -1, -1, 2, -1]
+ expected_edge_values = [2, 1, -1, -2, -2, -2, 2, -1]
+ expected_node_values = [2, 0, -2, -2, 2, -2, -1, 0]
+ expected_individual_values = [2, -4]
+
+ trait_df = make_trait_df(effect_sizes)
+
+ edge_effect_df = tstrait.edge_effect(ts, trait_df)
+ edge_value_df = tstrait.genetic_value(
+ ts,
+ trait_df,
+ level="edge",
+ )
+ node_value_df = tstrait.genetic_value(
+ ts,
+ trait_df,
+ level="node",
+ )
+ individual_value_df = tstrait.genetic_value(
+ ts,
+ trait_df,
+ level="individual",
+ )
+
+ assert_value_frame(
+ edge_effect_df,
+ id_column="edge_id",
+ value_column="effect_size",
+ expected_values=expected_edge_effects,
+ )
+ assert_value_frame(
+ edge_value_df,
+ id_column="edge_id",
+ value_column="genetic_value",
+ expected_values=expected_edge_values,
+ )
+ assert_value_frame(
+ node_value_df,
+ id_column="node_id",
+ value_column="genetic_value",
+ expected_values=expected_node_values,
+ )
+ assert_value_frame(
+ individual_value_df,
+ id_column="individual_id",
+ value_column="genetic_value",
+ expected_values=expected_individual_values,
+ )
+
+ # The default level must return individual genetic values.
+ default_value_df = tstrait.genetic_value(ts, trait_df)
+
+ pd.testing.assert_frame_equal(
+ individual_value_df,
+ default_value_df,
+ check_dtype=False,
+ )
+
+ # Independently calculate sample-node values from the genotype matrix.
+ # This does not use tstrait's edge or node accumulation algorithm.
+ effects_by_site = np.zeros(ts.num_sites)
+ for row in trait_df.itertuples(index=False):
+ effects_by_site[int(row.site_id)] += float(row.effect_size)
+
+ sample_values = ts.genotype_matrix().T @ effects_by_site
+
+ np.testing.assert_allclose(
+ sample_values,
+ np.asarray(expected_node_values[:4], dtype=float),
+ )
+
+ # Individual 0 consists of sample nodes 0 and 1.
+ # Individual 1 consists of sample nodes 2 and 3.
+ genotype_based_individual_values = np.array(
+ [
+ sample_values[0] + sample_values[1],
+ sample_values[2] + sample_values[3],
+ ]
+ )
+
+ np.testing.assert_allclose(
+ genotype_based_individual_values,
+ np.asarray(expected_individual_values, dtype=float),
+ )
diff --git a/tstrait/__init__.py b/tstrait/__init__.py
index b072e48..bdb4963 100644
--- a/tstrait/__init__.py
+++ b/tstrait/__init__.py
@@ -23,6 +23,7 @@
) # noreorder
from .genetic_value import (
genetic_value,
+ edge_effect,
normalise_genetic_value,
) # noreorder
from .simulate_environment import (
@@ -44,6 +45,7 @@
"TraitModelNormal",
"TraitModelT",
"genetic_value",
+ "edge_effect",
"normalise_genetic_value",
"sim_env",
]
diff --git a/tstrait/genetic_value.py b/tstrait/genetic_value.py
index d641523..78fca11 100644
--- a/tstrait/genetic_value.py
+++ b/tstrait/genetic_value.py
@@ -16,7 +16,7 @@ def _compute_nodes_genetic_value(
effect_size,
): # pragma: no cover
"""
- Compute the genetic value of each node for the specified set of mutations
+ Compute the node genetic values for the specified set of mutations
encoded in the stack.
"""
genetic_value = np.zeros(num_nodes)
@@ -24,7 +24,7 @@ def _compute_nodes_genetic_value(
parent_node_id = stack.pop()
genetic_value[parent_node_id] = effect_size
child_node_id = left_child_array[parent_node_id]
- while child_node_id != -1:
+ while child_node_id != tskit.NULL:
if not has_mutation[child_node_id]:
stack.append(child_node_id)
child_node_id = right_sib_array[child_node_id]
@@ -36,34 +36,50 @@ def _accumulate_individual_values(
nodes_genetic_value, nodes_individual, num_nodes, num_individuals
): # pragma: no cover
"""
- Accumulate the genetic values by summing their node contributions.
+ Accumulate the individual genetic values by summing their node contributions.
"""
individuals_genetic_value = np.zeros(num_individuals)
for u in range(num_nodes):
ind = nodes_individual[u]
- if ind != -1:
+ if ind != tskit.NULL:
individuals_genetic_value[ind] += nodes_genetic_value[u]
return individuals_genetic_value
-class _GeneticValue:
- """GeneticValue class to compute genetic values of individuals.
+def _accumulate_edge_values(nodes_genetic_value, nodes_edge, num_nodes, num_edges):
+ """
+ Accumulate the edge genetic values by summing their node contributions.
+ """
+ nodes_edge = nodes_edge[:num_nodes]
+ nodes_genetic_value = nodes_genetic_value[:num_nodes]
+ has_edge = nodes_edge != tskit.NULL
+ return np.bincount(
+ nodes_edge[has_edge],
+ weights=nodes_genetic_value[has_edge],
+ minlength=num_edges,
+ )
+
- :param ts: Tree sequence data with mutation
- :type ts: tskit.TreeSequence
- :param trait_df: Dataframe that includes causal site ID, causal allele,
- simulated effect size, and trait ID.
- :type trait_df: pandas.DataFrame
+class _GeneticValue:
+ """
+ GeneticValue class to compute genetic values of individuals, nodes, or edges.
+
+ Parameters
+ ----------
+ ts : tskit.TreeSequence
+ Tree sequence data with mutation
+ trait_df : pandas.DataFrame
+ Dataframe that includes causal site ID, causal allele, simulated effect
+ size, and trait ID.
"""
def __init__(self, ts, trait_df):
self.trait_df = trait_df[["site_id", "effect_size", "trait_id", "causal_allele"]]
self.ts = ts
- def _individual_genetic_values(self, tree, site, causal_allele, effect_size):
+ def _node_genetic_values(self, tree, site, causal_allele, effect_size):
"""
- Returns a numpy array that describes the genetic value of all individuals at
- a particular site.
+ Returns a numpy array with node genetic values.
"""
has_mutation = np.zeros(self.ts.num_nodes + 1, dtype=bool)
state_transitions = {tree.virtual_root: site.ancestral_state}
@@ -86,109 +102,165 @@ def _individual_genetic_values(self, tree, site, causal_allele, effect_size):
num_nodes=self.ts.num_nodes,
effect_size=effect_size,
)
+ return genetic_value
- individuals_genetic_value = _accumulate_individual_values(
- genetic_value,
- self.ts.nodes_individual,
- self.ts.num_nodes,
- self.ts.num_individuals,
- )
- return individuals_genetic_value
-
- def _run(self):
- """Computes genetic values of individuals.
+ def _run(self, level):
+ """
+ Computes genetic values of individuals, nodes, or edges
+ depending on the value of "level"
- :returns: Dataframe with genetic value, individual ID, and trait ID.
- :rtype: pandas.DataFrame
+ Returns
+ -------
+ pandas.DataFrame
+ Dataframe with trait ID, [individual|node|edge] ID, and genetic value.
"""
- num_ind = self.ts.num_individuals
+ ts = self.ts
+ size_map = {
+ "individual": ts.num_individuals,
+ "node": ts.num_nodes,
+ "edge": ts.num_edges,
+ }
+ N = size_map[level]
+
num_trait = np.max(self.trait_df.trait_id) + 1
- genetic_val_array = np.zeros((num_trait, num_ind))
+ genetic_value_table = np.zeros((num_trait, N))
tree = tskit.Tree(self.ts)
for data in self.trait_df.itertuples():
site = self.ts.site(data.site_id)
tree.seek(site.position)
- individual_genetic_value = self._individual_genetic_values(
+ genetic_value = self._node_genetic_values(
tree=tree,
site=site,
causal_allele=data.causal_allele,
effect_size=data.effect_size,
)
- genetic_val_array[data.trait_id, :] += individual_genetic_value
+ if level == "individual":
+ genetic_value = _accumulate_individual_values(
+ genetic_value, ts.nodes_individual, ts.num_nodes, ts.num_individuals
+ )
+ elif level == "edge":
+ genetic_value = _accumulate_edge_values(
+ genetic_value, tree.edge_array, ts.num_nodes, ts.num_edges
+ )
+ genetic_value_table[data.trait_id, :] += genetic_value
df = pd.DataFrame(
{
- "trait_id": np.repeat(np.arange(num_trait), num_ind),
- "individual_id": np.tile(np.arange(num_ind), num_trait),
- "genetic_value": genetic_val_array.flatten(),
+ "trait_id": np.repeat(np.arange(num_trait), N),
+ f"{level}_id": np.tile(np.arange(N), num_trait),
+ "genetic_value": genetic_value_table.flatten(),
}
)
return df
-def genetic_value(ts, trait_df):
+def genetic_value(ts, trait_df, level="individual"):
"""
- Obtains genetic value from a trait dataframe.
+ Compute genetic values for a tree sequence given a trait dataframe.
- :param ts: The tree sequence data that will be used in the quantitative trait
+ Parameters
+ ----------
+ ts : tskit.TreeSequence
+ The tree sequence data that will be used in the quantitative trait
simulation.
- :type ts: tskit.TreeSequence
- :param trait_df: Trait dataframe.
- :type trait_df: pandas.DataFrame
- :returns: Pandas dataframe that includes genetic value of individuals in the
- tree sequence.
- :rtype: pandas.DataFrame
-
- .. seealso::
- :func:`trait_model` Return a trait model, which can be used as `model` input.
-
- :func:`sim_trait` Return a trait dataframe, which can be used as a
- `trait_df` input.
-
- :func:`sim_env` Genetic value dataframe output can be used as an input
- to simulate environmental noise.
-
- .. note::
- The `trait_df` input has some requirements that will be noted below.
-
- 1. Columns
+ trait_df : pandas.DataFrame
+ Trait dataframe. See :ref:`req_trait_df` for column and data
+ requirements.
+ level : {"individual", "node", "edge"}, default "individual"
+ The level (entity) at which genetic values are returned.
+
+ Returns
+ -------
+ pandas.DataFrame
+ Genetic values for each trait and level (entity) in the tree sequence.
+ The dataframe columns are ``trait_id``, the ID column corresponding to
+ ``level`` (``individual_id``, ``node_id``, or ``edge_id``), and
+ ``genetic_value``.
+
+ See Also
+ --------
+ trait_model : Return a trait model, which can be used as `model` input.
+ sim_trait : Return a trait dataframe, which can be used as a `trait_df` input.
+ sim_env : Simulate environmental noise on top of genetic values.
+ edge_effect : Compute effects introduced on each edge.
+
+ Examples
+ --------
+ See :ref:`genetic_value_doc` and :ref:`genetic_individual_node_edge_doc` for
+ worked examples, while :ref:`phenotype_model` describes
+ quantitative genetics model assumptions.
+ """
- The following columns must be included in `trait_df`:
+ ts = _check_instance(ts, "ts", tskit.TreeSequence)
+ valid_levels = ("individual", "node", "edge")
+ if level not in valid_levels:
+ raise ValueError("level must be one of 'individual', 'node', or 'edge'")
+ if level == "individual" and ts.num_individuals == 0:
+ raise ValueError("No individuals in the provided tree sequence dataset")
+ trait_df = _check_dataframe(
+ trait_df, ["site_id", "effect_size", "trait_id", "causal_allele"], "trait_df"
+ )
+ if len(trait_df) == 0:
+ raise ValueError("trait_df must contain at least one row")
+ _check_non_decreasing(trait_df["site_id"], "site_id")
- * **site_id**: Site IDs that have causal allele.
- * **effect_size**: Simulated effect size of causal allele.
- * **causal_allele**: Causal allele.
- * **trait_id**: Trait ID.
+ trait_id = trait_df["trait_id"].unique()
- 2. Data requirements
+ if np.min(trait_id) != 0 or np.max(trait_id) != len(trait_id) - 1:
+ raise ValueError("trait_id must be consecutive and start from 0")
- * Site IDs in **site_id** column must be sorted in an ascending order. Please
- refer to :py:meth:`pandas.DataFrame.sort_values` for details on sorting
- values in a :class:`pandas.DataFrame`.
+ trait_df = trait_df.astype({"trait_id": int})
- * Trait IDs in **trait_id** column must start from zero and be consecutive.
+ genetic = _GeneticValue(ts=ts, trait_df=trait_df)
- The genetic value dataframe contains the following columns:
+ genetic_result = genetic._run(level)
- * **trait_id**: Trait ID.
- * **individual_id**: Individual ID inside the tree sequence input.
- * **genetic_value**: Genetic values that are obtained from the trait
- dataframe.
+ return genetic_result
- .. rubric:: Examples
- See :ref:`genetic_value` for worked examples.
+def edge_effect(ts, trait_df):
+ """
+ Compute effects introduced on each edge for a tree sequence
+ given a trait dataframe.
+
+ Parameters
+ ----------
+ ts : tskit.TreeSequence
+ The tree sequence that will be used in the trait simulation.
+ trait_df : pandas.DataFrame
+ Trait dataframe. See :ref:`req_trait_df` for column and data
+ requirements.
+
+ Returns
+ -------
+ pandas.DataFrame
+ Edge effects for each trait and edge in the tree sequence. The
+ dataframe contains ``trait_id``, ``edge_id``, and ``effect_size``.
+
+ Raises
+ ------
+ ValueError
+ If a causal-allele transition occurs on a root node,
+ which has no immediately ancestral edge to which an edge effect can be assigned.
+
+ See Also
+ --------
+ genetic_value : Compute trait genetic values for individuals, nodes, or edges.
+
+ Examples
+ --------
+ See :ref:`genetic_individual_node_edge_doc` for a worked example,
+ while :ref:`phenotype_model` describes quantitative genetics model assumptions.
"""
-
ts = _check_instance(ts, "ts", tskit.TreeSequence)
- if ts.num_individuals == 0:
- raise ValueError("No individuals in the provided tree sequence dataset")
trait_df = _check_dataframe(
trait_df, ["site_id", "effect_size", "trait_id", "causal_allele"], "trait_df"
)
+ if len(trait_df) == 0:
+ raise ValueError("trait_df must contain at least one row")
_check_non_decreasing(trait_df["site_id"], "site_id")
trait_id = trait_df["trait_id"].unique()
@@ -198,44 +270,83 @@ def genetic_value(ts, trait_df):
trait_df = trait_df.astype({"trait_id": int})
- genetic = _GeneticValue(ts=ts, trait_df=trait_df)
+ N = ts.num_edges
+ num_trait = np.max(trait_df.trait_id) + 1
+ edge_effect_table = np.zeros((num_trait, N))
- genetic_result = genetic._run()
-
- return genetic_result
+ # TODO: This implementation is slow - replace with a numba algorithm using tskit 1.0
+ tree = tskit.Tree(ts)
+ for data in trait_df.itertuples():
+ site = ts.site(data.site_id)
+ tree.seek(site.position)
+ for m in site.mutations:
+ if m.parent == tskit.NULL:
+ state_before_mutation = site.ancestral_state
+ else:
+ state_before_mutation = m.inherited_state
+ had_causal_allele = int(state_before_mutation == data.causal_allele)
+ has_causal_allele = int(m.derived_state == data.causal_allele)
+ state_change = has_causal_allele - had_causal_allele
+ if state_change != 0:
+ e = tree.edge(m.node)
+ if e == tskit.NULL:
+ raise ValueError(
+ "Cannot assign an edge effect to a mutation on a root node"
+ )
+ edge_effect_table[data.trait_id, e] += state_change * data.effect_size
+
+ df = pd.DataFrame(
+ {
+ "trait_id": np.repeat(np.arange(num_trait), N),
+ "edge_id": np.tile(np.arange(N), num_trait),
+ "effect_size": edge_effect_table.flatten(),
+ }
+ )
+ return df
def normalise_genetic_value(genetic_df, mean=0, var=1, ddof=1):
- """Normalise genetic value dataframe.
-
- :param genetic_df: Genetic value dataframe.
- :type genetic_df: pandas.DataFrame
- :param mean: Mean of the resulting genetic value.
- :type mean: float
- :param var: Variance of the resulting genetic value.
- :type var: float
- :param ddof: Delta degrees of freedom. The divisor used in computing the variance
+ """
+ Normalise genetic value dataframe.
+
+ Parameters
+ ----------
+ genetic_df : pandas.DataFrame
+ Genetic value dataframe.
+ mean : float, default 0
+ Mean of the resulting genetic value.
+ var : float, default 1
+ Variance of the resulting genetic value.
+ ddof : int, default 1
+ Delta degrees of freedom. The divisor used in computing the variance
is N - ddof, where N represents the number of elements.
- :type ddof: int
- :returns: Dataframe with normalised genetic value.
- :rtype: pandas.DataFrame
- :raises ValueError: If `var` <= 0.
- .. note::
- The following columns must be included in `genetic_df`:
+ Returns
+ -------
+ pandas.DataFrame
+ Dataframe with normalised genetic value.
+
+ Raises
+ ------
+ ValueError
+ If `var` <= 0.
- * **trait_id**: Trait ID.
- * **individual_id**: Individual ID inside the tree sequence input.
- * **genetic_value**: Simulated genetic values.
+ Notes
+ -----
+ The following columns must be included in `genetic_df`:
- The dataframe output has the following columns:
+ * **trait_id**: Trait ID.
+ * **individual_id**: Individual ID inside the tree sequence input.
+ * **genetic_value**: Simulated genetic values.
- * **trait_id**: Trait ID.
- * **individual_id**: Individual ID inside the tree sequence input.
- * **genetic_value**: Normalised genetic values.
+ The dataframe output has the following columns:
- .. rubric:: Examples
+ * **trait_id**: Trait ID.
+ * **individual_id**: Individual ID inside the tree sequence input.
+ * **genetic_value**: Normalised genetic values.
+ Examples
+ --------
See :ref:`normalise_genetic_value` section for worked examples.
"""
if var <= 0: