Add branch and mutation mode for genetic values

[gregorgorjanc]

In terms of the tskit stats mode, tstrait implements the site version. Would it be helpful to also implement branch mode and mutation mode (see tskit-dev/tskit#3299)?

We should think about the API for this - specifically on how to prepare trait information data frame and how that will be used downstream. Current setup works with causal sites and causal alleles. branch mode should work with, well, branches - so, will all branches have an effect or just branches with mutations at causal sites or branches in the data frame!? mutation mode should work with mutations, I guess with those with effect given the data frame!!

@petrelharp @jeromekelleher do you have any thoughts about this?

@iozkaraca might be able to take this extension forwards.

[jeromekelleher]

I think it's worth sketching out a proposal here on small example @gregorgorjanc so that we can see what the different options are. What would the dataframes look like in the different modes?

I'm not really up to speed on this currently so not in a position to make any proposals, but happy to review and comment.

[gregorgorjanc]

Thinking a bit more about this, the API for tstrait.genetic_value() would be:

tstrait.genetic_value(
    ts,
    trait_df,
    mode="site",          # site (default for backward compatibility), mutation, or branch
    level="individual",   # individual, node, or edge
)

• We should also think what to do with the sim_trait() - how to pass arguments to that level for the mode
• The below discussion is about the API and not yet about the underlying algorithm, which would come later once we have a clear plan and API

Then the data frame trait_df would be specific to each mode.

For mode="site" (current behaviour) we require:

• trait_id
• site_id
• causal_allele
• effect_size

For mode="mutation" we would require:

• trait_id
• mutation_id
• effect_size
• anything else?

These would be generated from the trait models like site effects. We should discuss what to do with reverse mutations, say A-->T gets an effect of 3 units, do we then have T-->A gets an effect of -3 units or do we draw completely different effect? One could also start thinking of codon specific effects - not sure how to address this as we would have to have information of the local DNA sequence. Thoughts?

Once we have these effects, we can generate genetic values for individuals, nodes, or edges that inherited the mutations and their effects.

For mode="branch" we would require:

• trait_id
• edge_id
• effect_size
• anything else? position?

These would be generated from the trait models like mutation effects. In principle this could be without the position, BUT say that the branch has an effect of +1 unit due to mutation on the first half of the haplotype represented by the edge. If second half of this haplotype is passed then onwards, creating a new edge, the starting value of that edge would be 0 units (since we did not inherit the bit with the mutation) and then we would add any new increment on this new edge (so stating value + new effect gives the edge genetic value, while the effect is just the innovation/change that arises on this new edge). I wonder if we need this one at all - it feels to me it's exactly? like the mutation mode.

Once we have these effects, we can generate genetic values for individuals, nodes, or edges that inherited the edger and their effects.

I would appreciate discussion around this topic @petrelharp @nspope @iozkaraca @hanbin973 @terhorst