Getting traits into slim_engine.py

[jeffspence]

Work in progress (not ready for primetime yet)

This is where @roshnipatel and I have been working on getting traits implemented in slim_engine.

This is what we've done:

• in slim_engine.py:
  • we simulate effect sizes for mutations from the DME using mutation callbacks
  • we simulate environments in late blocks
  • add environment to phenotype and then potentially transform
  • apply fitness functions in fitness callbacks
  • Checks for consistency between contig and TraitsModel
  • Checks for consistency between TraitsModel and demographic model
• Updated Environment and FitnessFunction in traits.py to specify the times and populations to which they apply (which is then used in slim_engine.py)
• We've implemented tests for the changes in traits.py, and have some hacky, non-functional tests for a couple of slim_engine.py things.

Here are our TODOs:

• We currently implement everything in late() blocks or fitnessEffect() code blocks. This is inconsistent with the ethos of the rest of slim_engine.py which uses registerLateEvent() to programmatically register demographic events based on tables of values defined at the beginning of the SLiM script.
  • This also causes some issues for us because registered events always happen after late() code blocks. This is a problem when new populations are created by population splits, because we want to access traits in populations that don't exist yet in their first generation. For now we've switch population splits to registerEarlyEvent()s but this is probably bad.
• Tests for things in slim_engine.
• We implemented a way to check that Environments and FitnessFunctions are consistent with the demography, but it would be cool to be able to print out a "decorated" demography somewhere, where it prints out simultaneously the demography and fitness functions and environments. @andrewkern -- would you be interested in tackling this?
• We have a bunch of TODOs sprinkled throughout, for things like implementing more fitness functions (we only implemented gaussian) or environmental distributions (we only implemented MVN)
• Testing and writing metadata.
• Update environments to use SLiM's individual offsets instead of doing it all ourselves.

Other things:

• For testing purposes I added generate_slim.py and generate_slim.sh which make test_script.slim (all in the stdpopsim/ directory. These are just temporary and I'm using them to look at the SLiM code generated by slim_engine.py. Once the PR is ready, they'll get removed.

[mergify]

Tick the box to add this pull request to the merge queue (same as @mergifyio queue).

• Queue this pull request

[roshnipatel]

Also, we'd love input on certain decisions we've made about how generation times get interpreted since we newly allow users to specify generation times for both environments and fitness functions.

In general, the decisions can be summarized as:

• (Positive) infinity is interpreted as the start of burn-in
• Burn-in is scheduled to end at the earliest (finite) time specified in the demographic model, the environments, and the fitness functions
• If no population is specified, the condition (environment or fitness function) applies to all active populations; if no time is specified, the condition applies to all times.
• Finite intervals are required to be fully contained within the intervals on which a population exists; infinite intervals are interpreted more generously.

For some concrete examples, suppose we have a demographic model that begins 5,000 generations ago.

• The interval [0, inf) describes an event that occurs throughout burn-in and continues to the present.
• The interval [0, 5000) describes an event that begins at the start of the demographic model (but NOT during burn-in) and continues to the present.
• The interval [0, 6000) describes an event that begins 1,000 generations before the demographic model and continues to the present; burn-in is then scheduled to end 6,000 generations ago (rather than 5,000 generations ago, as would be implied by the demographic model).

Also suppose that in this demographic model, pop B splits 3,000 generations ago, i.e. only exists on the interval [0, 3000).

• With population_list = [popB], specifying time_intervals as [0, inf) and [0, 3000) has the same result: the condition will apply to population B in its entirety.
• With population_list = [popB], specifying time_intervals as [0, 4000) will error -- this is a finite interval, so it is required to fully overlap the time(s) that pop B exists.

[andrewkern]

i'm a bit worried about this implementation choice. here subpopulations are referred to by index, so e.g., SLiM will go and pull subpop '2'. But if the number of active populations is changing over time, this could be an issue.

elsewhere I see that pop_ids (strings) are used, for example in the fitness function generator. I think that's the thing to do here too?

Like line 1831 could use f'p{pop_id}.individuals'?

[andrewkern]

the TODO note is right here-- I think this will fail if the population isn't active, like with models with ancestral populations not present at sampling time. I think the fix is to set time=None and then msprime will sample each population at its default sampling time. I think....

[andrewkern]

actually i went and tested this and I'm wrong about time=None -- that doesn't help.

instead i think we need to use demography.debug().epochs to walk through each to see which population is around

[andrewkern]

i think we want to check that interval[0] is finite -- this would be a user error that hopefully we wouldn't see, but stuff would blow up if the user handed [inf, inf)

also we should check for zero length intervals like [4, 4) -- those would slip by, but probably not affect things?

[andrewkern]

generally this looks really good! I like the decisions you've outlined above @roshnipatel. I've placed some comments on some nitpicks I saw along the way, but haven't done a complete review of the code. Happy to do that whenever

[andrewkern]

and yeah -- i could take a swing at the decorated demography thing @jeffspence. are you thinking a function that produces a table a la msprime's demography debugger?

[jeffspence]

@andrewkern yeah! something like that would be great. Just something showing the populations and times of events (both demographic and traits related).

[bhaller]

Hi folks! A few superficial comments, not having looked at the underlying code:

• A while back, we discussed the terminology of "DME" versus "DES", and my (perhaps erroneous) recollection is that you guys decided to follow SLiM in using "DES", for consistency; I think that'd be a good idea. See section 1.7 of the draft SLiM manual for terminological discussion; it'd probably be good for stdpopsim to match SLiM's terminology on this stuff, to try to minimize confusion.

• You write "add environment to phenotype" as something you've done, and you have "Update environments to use SLiM's individual offsets instead of doing it all ourselves" but it's crossed out. What's the situation there? If the way SLiM does individual offsets is not fit for its purpose, and so you have to roll your own equivalent mechanism for some reason, I'd like to hear about that.

• You write "we simulate effect sizes for mutations from the DME using mutation callbacks"; is this because you are supporting distributions that SLiM doesn't support built-in? Or because you want to have correlations between effects of different traits? Again, if SLiM's built-in facilities are not sufficient I'd like to know why. If there's a distribution you want added, just ask. :-> (On the other hand, if it is for reasons of correlated effect sizes, then yeah, SLiM doesn't support that built-in.)

• You write "apply fitness functions in fitness callbacks"; doing this in a vectorized fashion and putting the results into fitnessScaling would probably be more efficient (as in the examples in chapter 13 of the draft manual, I think), but perhaps this is not a huge deal since the callback(s) would only be called once per individual. If you've got a clear reason for your design, then don't worry about it, I doubt the performance hit will be large for most models.

• You write "This also causes some issues for us because registered events always happen after late() code blocks". Yes, hmm. SLiM's scheduling internals are not very exposed right now. I could conceivably add a parameter to the various registerX() methods that would let you specify whether to add the new X to the start or to the end of the queue, if that would be useful...? Adding more flexibility than a simple start/end choice would be more complex, and probably would probably be overkill but could be discussed.

That's all I see for right now. Great to see the progress! :-> You guys are kind of the guinea pigs for SLiM 6, since I haven't managed to get the SLiM 6 beta out yet; basically nobody is using it yet except you. So feedback on anything that seems non-optimal would be super useful. Thanks!

[jeffspence]

@bhaller thanks for all of this!! We are very happy to be guinea pigs, and everything has been going pretty smoothly so far.

A while back, we discussed the terminology of "DME" versus "DES", and my (perhaps erroneous) recollection is that you guys decided to follow SLiM in using "DES", for consistency; I think that'd be a good idea. See section 1.7 of the draft SLiM manual for terminological discussion; it'd probably be good for stdpopsim to match SLiM's terminology on this stuff, to try to minimize confusion.

I believe that we're discussing that in this issue: #1830 . I'm inclined to stick with DME for now here, but am open to changing it in another PR.

You write "add environment to phenotype" as something you've done, and you have "Update environments to use SLiM's individual offsets instead of doing it all ourselves" but it's crossed out. What's the situation there? If the way SLiM does individual offsets is not fit for its purpose, and so you have to roll your own equivalent mechanism for some reason, I'd like to hear about that.

Crossed out because I committed some changes to switch everything over to using the individual offsets. We rolled our own because I didn't read the SLiM documentation carefully enough, and then we switched it over when @petrelharp pointed us to the offsets. I think that the existing offset mechanism in SLiM works perfectly for us.

You write "we simulate effect sizes for mutations from the DME using mutation callbacks"; is this because you are supporting distributions that SLiM doesn't support built-in? Or because you want to have correlations between effects of different traits? Again, if SLiM's built-in facilities are not sufficient I'd like to know why. If there's a distribution you want added, just ask. :-> (On the other hand, if it is for reasons of correlated effect sizes, then yeah, SLiM doesn't support that built-in.)

Yes, sorry -- I was a bit sloppy writing this up. We use the SLiM built-ins for mutations that only affect one trait and mutation callbacks for mutations that affect multiple traits.

You write "apply fitness functions in fitness callbacks"; doing this in a vectorized fashion and putting the results into fitnessScaling would probably be more efficient (as in the examples in chapter 13 of the draft manual, I think), but perhaps this is not a huge deal since the callback(s) would only be called once per individual. If you've got a clear reason for your design, then don't worry about it, I doubt the performance hit will be large for most models.

Thanks for pointing this out. I'll take a look and see if we can switch to using fitnessScaling

You write "This also causes some issues for us because registered events always happen after late() code blocks". Yes, hmm. SLiM's scheduling internals are not very exposed right now. I could conceivably add a parameter to the various registerX() methods that would let you specify whether to add the new X to the start or to the end of the queue, if that would be useful...? Adding more flexibility than a simple start/end choice would be more complex, and probably would probably be overkill but could be discussed.

I think we settled on switching over all of the traits stuff to also use various registerX() methods to be consistent with the way we are handling demography. So at least for now, I think that this will be a non-issue once we've switched over to doing that, but we will let you know if we end up needing this kind of functionality.

Thanks again!!