flowchart TD A[TE genotype matrix + flanking SNPs<br/>per-window image] --> B[Two-branch CNN<br/>state channel + continuous channel] B --> C[Per-window embeddings] C --> D[attention pooling/deepsets<br/>variable number of TEs K] D --> E[Neural posterior estimator<br/>normalising flow, sbi] E --> F[DFE parameters<br/>gamma shape, mean s,<br/>burst timing/duration/amplitude,<br/>insertion rate]
Inferring the fitness effects and dynamics of transposable elements with simulation-based deep learning
An informal progress report
Overview
A long-standing goal of my work has been to obtain more precise estimates of the fitness effects and dynamics of transposable elements (TEs). TEs are known to be eco-responsive, to be at least mildly deleterious, to undergo bursts of transposition and to display insertion preferences. These properties make them substantially harder to model than SNPs. A promising avenue of research probably consists in contrasting patterns of TE age, frequency, and genome-wide distribution to obtain quantitative estimates of the parameters above.
This informal post summarises a suite of tools I have been developing towards that goal, and the directions I am currently pursuing. It is perhaps a bit naive, and I still need it to be reviewed by experts in the field of deep learning and population genetics, but the first results are encouraging enough from my perspective that I think they may be shared with the community. I hope this can stimulate more, better work on this, and maybe illustrate the challenges and excitement of taking advantage of some of the most powerful methods in population genomics for what remains a “niche” question, the population dynamics of TEs.
Potential of Simulation-Based Inference for the population dynamics of TEs.
A workhorse of population genomics is Approximate Bayesian Computation (ABC): by comparing observed summary statistics with simulations, one retains the simulations closest to the data and recovers posterior distributions for the parameters of interest. This works well for SNP data. TEs are more delicate, for three reasons.
First, TEs are far less abundant than SNPs along the genome. If one wishes to simulate TEs in their genomic context as realistically as possible, it may require simulating megabases of host’s sequence to obtain tens of polymorphic TEs for a given TE lineage. Second, the insertion rate can vary through time, for instance during a burst of transposition, which produces cohorts of similarly aged elements. A recent burst can generate an excess of singletons in the allele frequency spectrum that has nothing to do with selection, yet would ordinarily be read as such (Bourgeois and Boissinot 2019). Third, simulating selection against TEs requires forward-in-time simulators such as SLiM, which are slow and computationally costly (cf infra).
So far I had only compared simulations with observations in a qualitative or semi-quantitative way (Bourgeois et al. 2020; Horvath et al. 2024). That past work suggests that the abundance of TEs along the genome, together with the contrast between their frequency and their age, should carry enough signal to disentangle transposition activity, preferential insertion, and selection. An age-of-allele test of neutrality has also been proposed a few years ago already (Horvath et al. 2022), and implemented in a Bayesian framework that showed promise (Blumenstiel et al. 2014). Using ABC (or any simulation-based inference methodology) to compare TE simulations with observations should therefore be doable to jointly infer selection, insertion bias, and time-varying transposition dynamics.
However, ABC requires a lot of simulations of rather long regions to obtain enough simulated TEs, and the choice of relevant summary statistics is non-trivial (although abcrf can handle many summary statistics). There is also the issue of scale: one can simulate a fraction of the genome, retrieve summary statistics for the corresponding subset of simulated TEs, and infer on observed TEs drawn from the entire genome ; but there is a risk of distortion between simulations and observations: a subset of the genome is not always enough to cover the entire variance in recombination rates or gene content, assuming that we want to incorporate this in our simulations. We can of course pick carefully regions that are representative, or infer on genomic subsets with a similar distribution of recombination rate and gene densities as the windows used for training, but this may become quickly cumbersome and unpractical. The development of new deep learning methods may provide a more elegant and efficient way to retrieve parameters while learning to generalize to new configurations and number of observations in the observed data, while keeping the number of simulations relatively low (in the thousands range).
Deep learning: motivation and requirements
Two recent papers recently drew my attention : popgen-npe (Min et al. 2026) and an attention-based approach of Sanchez and colleagues (Sanchez et al. 2026). Both aim to make deep learning more accessible and more readily generalisable to genome-wide data.
An ideal pipeline would:
- take a (phased or unphased, polarised or not) VCF of TEs and SNPs
- extract features relevant to estimate TE age, frequency, genotypes, and where available distance to genes or other functional features and local recombination rate
- compare these features with simulations spanning a broad range of selection coefficients, proportions of strongly versus mildly deleterious insertions, transposition bursts, and insertion preferences and selection that varies along the genome. An advantage of the
popgen-npeimplementation is the ability to obtain an estimate of uncertainties around parameter estimation using neural posterior estimates. I wanted to incorporate this in the pipeline. - be as computationally efficient as possible. Whole-genome SLiM simulations are slow even when only deleterious markers and a few thousand segregating sites are modelled, so the method should learn from a representative subset of the genome and generalise to the whole. Of particular interest in the Sanchez et al. attention-based paper was the ability for the algorithm to weigh demographic inference obtained from distinct different genomic windows to produce a consensus demographic history through a voting scheme. In the preprint, they train the model on 200 x 2Mb windows but infer using the entire human genome (nearly ten times more windows). It seemed like a valuable way of using on TEs drawn from a genome of 1Gb a model trained on, for example, 20 x 5Mb windows (100Mb).
- optionally support iterative refinement in the spirit of ABC-MCMC or the sequential schemes used in tools such as DILS (Fraïsse et al. 2021): resampling from a broad initial posterior and refining, without overfitting
- remain inspectable. Deep learning is often called a black box, and coming from classical ABC I wanted a method whose behaviour I could interrogate: what the summary statistics look like, how the simulations behave, and how well they fit the observed data. A colleague of mine, François Sabot, told me about how some deep-learning methods designed to distinguish between images of huskies and other dog breeds were using the snow background as a distinguishing feature. Having the option to implement even a coarse ABC inference seemed useful to me from that perspective.
Working with AI
Large language models have become common and debated as research tools over the past few years. I had previously used Claude (Opus 4.7 to 5.0 here) for auxiliary tasks, mainly producing figures, improving my written English, and improving on Python, so I began a conversation to build my understanding of deep learning. This developed into what I later learned is called “vibe coding”: I give instructions, then collect, check and test the code. Python is fairly easy to reverse-engineer, so catching errors is easier than writing code from scratch.
What really made coding these trials much easier was the assistance AI provided for coding functions wrapping the simulators and the inference pipelines together. In my experience, this has always been the most difficult part of bioinformatics pipeline, converting the output of one (simulation) program into the input for another (inference) portion. It also really helped documenting the pipeline, and to make it adapted to a HPC environment.
On the more dangerous side, the AI agent is not particularly creative, and the speed at which it produces and edit code creates a perverse incentive to always complexify the method before properly testing what is already there. It is also biased towards more common languages, which is fine since Python is the main language used for the deep learning part. However, the SLiM code needed several corrections, especially since some functions in Eidos look like, but are not exactly the same, as their R equivalent, which generated a few issues during the first trials (for example, setting the gamma distribution in SLiM/Eidos involves specifying the shape and the mean, while in R one specifies the shape and the rate). Using summary statistics and ABC approaches was also useful to identify bugs (produced by the AI) in the scripts: for example, the benchmarking simulator only sampled windows with lower recombination rates in an earlier version. Interestingly, the benchmark was still decent, although a PCA on summary statistics clearly revealed the drift of benchmark simulations compared to the ones used for training.
What is shown below is the result of roughly three months of work (first prompts beginning of May 2026) to obtain a functional pipeline, starting by building on the popgen-npe architecture (Min et al. 2026) (itself a wrapper around the excellent sbi (Tejero-Cantero et al. 2020)) and on ideas from dnadna (Sanchez et al. 2026) about using attention to generalize inference to an arbitrary number of TEs beyond the range used in training. At this stage, I haven’t tested this framework on biological data, and all the parts of the code dealing with polarization, masking or phasing are yet to be tested. If you think this work may bring something to your research, or if you have any comment, feel free to reach out.
The hardest part was not, perhaps, the deep learning itself, for which there are a lot of excellent tutorials now. Reading the wiki pages for the sbi python library helped me a lot to understand the basic concepts around flows or embeddings. I could also count on a few lectures I could find. For example, I can only recommend Fumagalli, Perez and Jay’s lecture on deep learning here: https://github.com/mfumagalli/ImaGene/blob/master/Tutorials/Workshop/lecture.pdf
A typical simulation-based deep learning pipeline is organized this way (again, popgen-npe is really useful to understand the process, which will be familiar to ABC practitioners):
A simulator produces training data. I use SLiM v4.3 (Haller and Messer 2023)here, and intend to upgrade to v5. I also use msprime (Baumdicker et al. 2022) as a fast, rough simulator whose results can be recalibrated using a few SLiM simulations (see below).
A processor extracts the relevant features from the simulations, be it an image of the alignment (CNN), information about local recombination rate, gene density, or frequency of TEs.
A network architecture/flow can be then specified, that will take these features as input, and can be trained to transform the initial features into a set that can be used to run either regression (what we do here) or classification tasks
Once trained, we can benchmark on an independent set of simulations. In our case, we benchmark over a range of TEs or 5 Mb genomic windows.
If the algorithm seems to perform reasonably well, we can then infer based on actual data.
Overall, the modelling decisions, the diagnostics, and the interpretation are mine, AI accelerated implementation and helped me reason through alternatives in a domain that I did not master. I typically worked from my preliminary knowledge of ABC to identify equivalents in deep learning. For example, by asking whether there was an equivalent for improving posterior estimates through an enrichment of simulations around observed summary statistics (like in ABC-MCMC), the agent identified Sequential Neural Posterior estimate as a rough equivalent.
From images to sets of TEs.
The initial design used images of the TE genotype matrix together with flanking regions, the CNN model below. This is close to the model in popgen-npe, which in turn follows the exchangeable CNN model in Chan et al (Chan et al. 2018). I first chose this mostly because it seemed elegant to have an algorithm that could directly extract relevant information and associations between genomic features and TEs from the alignment.
Its drawbacks are essentially practical. Training is slow and memory-hungry, and the posteriors it returns are rather broad (a behaviour I still need to investigate). The tests below also use small effective population sizes. Production settings would be larger, but even then, modelling Drosophila populations (Ne on the order of two million) directly in SLiM is prohibitive and would require rescaling, reducing effective sizes and times while inflating mutation, insertion, and selection rates by the same factor. There is an excellent literature on doing this correctly (Marsh, Kaushik, and Johri 2026), which cautions against using windows larger than ~50-200 kb under two-orders-of-magnitude rescaling (it actually depends on the number of crossing-over events per simulated sequence/genome, see their Figure 1). In that regime the CNN approach, which needs a reasonably long genomic window, may be poorly suited.
The alternative, which is also more amenable to ABC, is to compute summary statistics on individual TEs, obtainable from much shorter, independent windows, and treat them as an unordered set of observations. We can however keep simulating genomic blocks, which may help keeping Linkage Disequilibrium between TEs, in the same way one can do block-bootstrap in population genetic inference to preserve correlations in the data. Here, I kept the same design throughout, with 5Mb windows, and sampled TEs from these windows, but the approach would also work when simulating many shorter windows.
The current pipeline extracts, for the SNPs flanking each TE, proxies for haplotype extension, local Tajima’s D for the carrier and non-carrier haplotype groups, local heterozygosity, and related quantities. These can be aggregated for ABC (the whole population of TEs is summarised: mean and quantiles of the recombination rate at insertion sites, of frequency, of nearby heterozygosity, and so on), which also mitigates the black-box problem, since I can inspect what the simulations produce and how well they match the data. The same per-TE summaries can equally be fed to a multi-layer perceptron (MLP).
A key point is that we never simulate an entire genome. The attention-based work by Sanchez et al. suggests a way to improve inference as more TEs are added without simulating their full genomic context. The design I settled on applies attention gates to the features produced by the CNN or the MLP, trains on a variable number of TEs so that differences in variance are less of an issue later on during inference, and extrapolates at inference to the observed data. This appears to work (see below): adding more TEs, even beyond the number of windows used in training continues to improve inference, most noticeably for the selection coefficient.
Technically this is not a transformer: the gates are computed from each TE’s own features, with no attention between TEs. It belongs to the Deep Sets family (Zaheer et al. 2018), using gated attention pooling in the sense of attention-based multiple-instance learning. Each TE contributes independently, adding more does not perturb those already there.
There are three main strategies that I tried, which all give roughly similar results, though there seems to be a slight advantage in using softmax attention gates, compared to sigmoid attention, or deepsets, which does not weigh individual TEs, but examines the distribution of per-TE features before feeding to the trunk. I also tried (mostly for myself) to provide an explanation to all this, which can be found as an annex.
Architectures
The pipeline implements six models, crossing three input/encoder strategies with two parameter sets. The two parameter sets are the DFE model, a gamma distribution of fitness effects (shape and mean \(|s|\)), the burst parameters (timing, duration, amplitude), and the global insertion rate ; and the context model, which additionally estimates how selection and insertion vary with local recombination through two slopes, \(\alpha_r\) (selection) and \(\gamma_r\) (insertion).
The parameterisation of the recombination dependence is
\[ s_{\text{local}}(p) = \exp\!\big(\log s_{\text{global}} + \alpha_r \, \phi_r(p)\big), \qquad \lambda_{\text{local}}(p) = \lambda_{\text{global}} \, \exp\!\big(\gamma_r \, \phi_r(p)\big), \]
where \(\phi_r(p)\) is a normalised log local recombination rate. Note that \(s_{\text{global}}\) and \(s_{\text{local}}\) are positive in the equation, but are then turned negative (purifying selection).
In the context model, the \(\lambda_{\text{global}}\) that we report in the tests below is the insertion rate expected for a specific median value of local recombination rate (RECOMB_TYPICAL in the config_base.py file), not the average insertion rate for the whole genome.
Note that we could very well use linear or quadratic functions, or use another feature (maybe chromatin accessibility, or gene density). I use recombination since it is usually correlated to gene density and chromatin accessibility, and is in itself a likely determinant of TE fitness due to ectopic recombination (Kent, Uzunović, and Wright 2017). Editing this should not require too much work, and most of the pipeline is already in place since I include distance to nearest gene as a feature used by the algorithm.
We fix the burst duration over 100 generations before rescaling. This is mostly because a 100 generations burst with a 100X increase from the baseline rate is not easy to distinguish from a 500 generations burst with a 20X increase at this scale, and I wanted to start testing relatively simple models. Note however that this is also a parameter that can be estimated in the current pipeline.
The three strategies are: (1) the CNN on genotype-matrix images with flanks; (2) normalised summary statistics aggregated across all TEs, fed to ABC; and (3) per-TE summary statistics treated as an unordered set, encoded by a shared MLP with attention gates. All results below use the same attention scheme (softmax gates, see Annex 2), which seems to produce slightly better results, but differences between the three (softmax, sigmoid, deepsets) are rather small.
(a) CNN + attention + DFE
(b) Normalised sumstats + ABC / neuralnet + DFE
flowchart TD
A[Per-TE summary statistics] --> B[Aggregate ACROSS TEs<br/>mean, sd, quantiles per feature]
B --> C{ABC-RF /-neuralnet / rejection}
C --> F[DFE parameters<br/>gamma shape, mean s,<br/>burst, insertion rate]
(c) Per-TE sumstats + MLP + attention + DFE
flowchart TD A[Per-TE summary statistics<br/>unordered set] --> B[Shared per-TE MLP encoder] B --> C[attention gates/deepsets + quantile pooling<br/>variable number of TEs K] C --> D[Neural posterior estimator<br/>normalising flow, sbi] D --> E[DFE parameters<br/>gamma shape, mean s,<br/>burst, insertion rate]
(d) CNN + attention + context
flowchart TD A[TE genotype matrix + flanking SNPs<br/>+ local recombination map] --> B[Two-branch CNN<br/>state + continuous, recomb-aware] B --> C[Per-window embeddings] C --> D[attention pooling/deepsets<br/>variable number of TEs K] D --> E[Neural posterior estimator<br/>normalising flow, sbi] E --> F[Context parameters<br/>DFE + burst + insertion rate,<br/>plus alpha_r selection slope<br/>and gamma_r insertion slope]
(e) Normalised sumstats + ABC / neuralnet + context
flowchart TD
A[Per-TE summary statistics<br/>incl. local recombination] --> B[Aggregate ACROSS TEs<br/>mean, sd, quantiles; recomb-stratified]
B --> C{ABC-RF /-neuralnet / rejection}
C--> D[Context parameters<br/>+ alpha_r, gamma_r]
(f) Per-TE sumstats + MLP + attention + context
flowchart TD A[Per-TE summary statistics<br/>unordered set, incl. local recombination] --> B[Shared per-TE MLP encoder] B --> C[Attention gates/deepsets + quantile pooling<br/>+ recombination summary features] C --> D[Neural posterior estimator<br/>normalising flow, sbi] D --> E[Context parameters<br/>DFE + burst + insertion rate,<br/>plus alpha_r and gamma_r]
Simulation set-up
The preliminary tests use a population of constant effective size with background selection at coding sites, driven by a recombination and gene map from a pilot project on the argan tree (Sideroxylon spinosum). The map was obtained using pyrho (Spence and Song 2019). The intent is a realistic picture of how gene density and recombination co-vary. This remains exploratory, and I will move to more complex/realistic demographic models later. To speed up the 5000 forward simulations that are used here, all quantities are rescaled by a factor of 10: effective sizes and times are divided by ten, and mutation, recombination, insertion, and selection rates are multiplied by ten. The parameter combinations are summarised below.
| Parameter | Symbol | Unscaled | Rescaled (10×) |
|---|---|---|---|
| Effective population size | \(N_e\) | 5000 diploids | 500 diploids |
| SNP mutation rate | \(\mu\) | \(1.25\times10^{-8}\) | \(1.25\times10^{-7}\) |
| Simulated duration | \(T\) | 50 000 gen | 5000 gen (\(=10\,N_e\)) |
| Window size | \(L\) | 5 Mb | — |
| Windows per replicate | — | 20 | — |
| Recombination / gene map | — | argan (S. spinosum). Recombination on average of the same order of magnitude as mutation rate. | — |
| Background sel. shape | BG_DEL_SHAPE |
0.3 | 0.3 |
| Background sel. mean \(s\) | BG_DEL_MEAN_S |
0.01 | 0.1 |
| Background sel. mutation rate | BG_DEL_MU |
\(0.85\times10^{-7}\) | \(0.85\times10^{-6}\) |
At the moment, tests are conducted on a set of 50 diploid individuals (100 haplotypes). Future tests will explore the robustness of methods to low sample size and lack of phase.
Under these background-selection settings roughly two-thirds of new coding mutations are deleterious. Simulations run for ten rescaled \(N_e\) generations to approach equilibrium. This is of course extremely simplified, a proper benchmark would use more realistic values for \(N_e\) and test various demographies.
The two recombination-dependence slopes of the context model, \(\alpha_r\) and \(\gamma_r\), may use a visual intuition. The figure below shows how the local selection coefficient and the local insertion intensity vary across four orders of magnitude of recombination, for two extreme (opposite-sign) values of each slope.
Code
#| label: fig-recomb-scaling
#| fig-cap: "Local selection coefficient (left) and local insertion intensity (right) as a function of the local recombination rate, for two extreme slope values. alpha_r governs the selection panel and gamma_r the insertion panel. A positive slope concentrates the effect in high-recombination regions; a negative slope does the opposite. phi_r is the normalised log local recombination rate, clipped to [-2, 2]."
#| warning: false
#| message: false
#| fig-width: 9
#| fig-height: 4
library(ggplot2)
library(dplyr)
library(tidyr)
library(knitr)
knitr::opts_chunk$set(
cache = FALSE
)
## Normalised recombination feature, matching the pipeline:
## phi_r(r) = clip( log10(r / RECOMB_TYPICAL) / RECOMB_LOG_SCALE, -2, 2 )
RECOMB_TYPICAL <- 1e-7 # "typical" recombination rate -> phi_r = 0
RECOMB_LOG_SCALE <- 1 # one log10 unit -> one unit of phi_r
feature_r <- function(r) pmin(pmax(log10(r / RECOMB_TYPICAL) / RECOMB_LOG_SCALE, -2), 2)
## Global baselines (illustrative)
log_s_global <- log(0.01) # mean selection coefficient s = 0.01
lambda_global <- 1 # baseline insertion intensity (relative units)
## Recombination axis and two extreme slopes
r <- 10^seq(-9, -5, length.out = 300)
extremes <- c(-2, 2)
df <- expand.grid(r = r, slope = extremes) |>
mutate(
phi = feature_r(r),
Selection = exp(log_s_global + slope * phi), # s_local (alpha_r = slope)
Insertion = lambda_global * exp(slope * phi) # lambda_local (gamma_r = slope)
) |>
pivot_longer(c(Selection, Insertion),
names_to = "quantity", values_to = "value") |>
mutate(
quantity = factor(quantity, levels = c("Selection", "Insertion")),
slope = factor(slope, labels = c("slope = -1", "slope = +1"))
)
ggplot(df, aes(r, value, colour = slope, linetype = slope)) +
geom_line(linewidth = 1) +
facet_wrap(~ quantity, scales = "free_y",
labeller = as_labeller(c(Selection = "Selection (alpha_r)",
Insertion = "Insertion (gamma_r)"))) +
scale_x_log10() +
scale_y_log10() +
scale_colour_manual(values = c("#B3452F", "#378ADD")) +
labs(x = "Local recombination rate (per bp per generation)",
y = "Local value (log scale)",
colour = NULL, linetype = NULL) +
theme_minimal(base_size = 12) +
theme(legend.position = "bottom")
We can also show where along our argan tree genome TEs would fall or be counterselected. For the context model, this is something useful to check.
Code
# Model (matching te_dfe/config_base.py):
# phi_r(p) = clip( log10(r(p) / RECOMB_TYPICAL) / RECOMB_LOG_SCALE, -2, 2 )
# lambda_local = lambda_global * exp( gamma_r * phi_r ) # insertion intensity
# s_local = exp( log_s_global + alpha_r * phi_r ) # selection against TEs
#
#Colour is on a log scale because both quantities span ~3000-fold across the
# map at |slope| = 2 (exp(+/-4)).
library(ggplot2)
library(dplyr)
library(tidyr)
## ---- CONFIG: mirror te_dfe/config_base.py -----------------------------------
RECOMB_MAP_PATH <- "recomb_map.txt" # HapMap: Chromosome Position(bp) Rate(cM/Mb) Map(cM)
CHROMS <- NULL # NULL = all chromosomes in the file
RECOMB_TYPICAL <- 1e-7 # r giving phi_r = 0
RECOMB_LOG_SCALE <- 1 # one log10 unit of r = one unit of phi_r
LOG_S_GLOBAL <- log(0.01) # baseline selection coefficient (s = 0.01)
LAMBDA_GLOBAL <- 1 # baseline insertion intensity (relative units)
BIN_BP <- 1e5 # colour-bar resolution
GAMMA_VALUES <- c(-2, 2) # extreme insertion slopes (Figure 1)
ALPHA_VALUES <- c(-2, 2) # extreme selection slopes (Figure 2)
feature_r <- function(r) {
pmin(pmax(log10(r / RECOMB_TYPICAL) / RECOMB_LOG_SCALE, -2), 2)
}
## ---- Read the HapMap map ----------------------------------------------------
## Rate is cM/Mb; Morgan/bp = rate * 1e-8 (same conversion as
## te_dfe/demography.py::load_hapmap_recomb).
raw <- read.table(RECOMB_MAP_PATH, header = FALSE, fill = TRUE,
stringsAsFactors = FALSE,
col.names = c("chrom", "pos", "rate_cM_Mb", "map_cM"))
raw <- raw[!is.na(suppressWarnings(as.numeric(raw$pos))), ] # drop header row
raw$pos <- as.numeric(raw$pos)
raw$rate_cM_Mb <- as.numeric(raw$rate_cM_Mb)
raw$chrom <- as.character(raw$chrom)
if (!is.null(CHROMS)) raw <- raw[raw$chrom %in% CHROMS, ]
## Natural order: chr2 before chr10 (a plain sort would invert them).
uc <- unique(raw$chrom)
chrom_num <- suppressWarnings(as.numeric(gsub("[^0-9]", "", uc)))
chrom_lev <- uc[order(ifelse(is.na(chrom_num), Inf, chrom_num), uc)]
## Step segments, computed within each chromosome so none straddles a boundary.
bins <- raw |>
mutate(chrom = factor(chrom, levels = chrom_lev)) |>
arrange(chrom, pos) |>
group_by(chrom) |>
mutate(seg_start = lag(pos, default = 0),
seg_end = pos,
seg_len = seg_end - seg_start,
r = rate_cM_Mb * 1e-8) |> # cM/Mb -> Morgan/bp
ungroup() |>
filter(seg_len > 0, is.finite(r)) |>
## bin to BIN_BP: length-weighted mean rate per bin
mutate(bin = floor(((seg_start + seg_end) / 2) / BIN_BP)) |>
group_by(chrom, bin) |>
summarise(r = sum(r * seg_len) / sum(seg_len), .groups = "drop") |>
mutate(mid_mb = (bin + 0.5) * BIN_BP / 1e6,
phi = feature_r(r))
cat(sprintf("map: %d chromosomes, %d bins, phi_r in [%.2f, %.2f]\n",
nlevels(droplevels(bins$chrom)), nrow(bins),
min(bins$phi), max(bins$phi)))map: 11 chromosomes, 5496 bins, phi_r in [-2.00, 0.92]Code
## ---- Shared plot skeleton ---------------------------------------------------
chrom_bars <- function(df, fill_name, title, subtitle, option = "viridis") {
ggplot(df, aes(x = mid_mb, y = chrom, fill = value)) +
geom_tile(height = 0.7, width = BIN_BP / 1e6) +
facet_wrap(~ panel, ncol = 1) +
scale_fill_viridis_c(trans = "log10", option = option, name = fill_name) +
scale_y_discrete(limits = rev(levels(df$chrom))) +
scale_x_continuous(expand = c(0.01, 0)) +
labs(x = "Position (Mb)", y = NULL, title = title, subtitle = subtitle) +
theme_minimal(base_size = 11) +
theme(panel.grid = element_blank(),
legend.position = "bottom",
legend.key.width = unit(2.5, "cm"),
strip.text = element_text(face = "bold"))
}
## ---- Figure 1: insertion intensity (gamma_r) --------------------------------
fig1_df <- tidyr::crossing(bins, gamma_r = GAMMA_VALUES) |>
mutate(value = LAMBDA_GLOBAL * exp(gamma_r * phi),
panel = paste0("gamma_r = ", sprintf("%+g", gamma_r)))
fig1 <- chrom_bars(
fig1_df,
fill_name = "Insertion intensity (relative to baseline)",
title = "Figure 1 — Where TEs insert preferentially",
subtitle = "lambda_local = lambda_global * exp(gamma_r * phi_r). Bright = TEs inserted more often.",
option = "viridis"
)
print(fig1)
Code
## ---- Figure 2: selection coefficient (alpha_r) ------------------------------
fig2_df <- tidyr::crossing(bins, alpha_r = ALPHA_VALUES) |>
mutate(value = exp(LOG_S_GLOBAL + alpha_r * phi),
panel = paste0("alpha_r = ", sprintf("%+g", alpha_r)))
fig2 <- chrom_bars(
fig2_df,
fill_name = "Selection coefficient s against TEs",
title = "Figure 2 — Where TEs are counter-selected",
subtitle = "s_local = exp(log_s_global + alpha_r * phi_r). Bright = selection removes TEs more efficiently.",
option = "magma"
)
print(fig2)
Code
## ggsave("fig1_insertion_gamma_r.png", fig1, width = 10, height = 6, dpi = 150)
## ggsave("fig2_selection_alpha_r.png", fig2, width = 10, height = 6, dpi = 150)Arguably, a model with strong selection against TEs in regions of lower recombination is less likely, but I was interested in checking whether \(\alpha_r\) and \(\gamma_r\) were indeed distinguishable.
Comparing the approaches
For each of the two models I report recovery of the target parameters against their true values, and I split between two sets of simulations: those with an intense transposition burst (> 100X the baseline), and a lower-intensity burst (< 100X). This is mostly because bursts that are not intense and old are difficult to distinguish from a flat influx of TEs. Note that we are using a rather low number of simulations, particularly for ABC inference (5000 simulations of 20 x 5Mb windows for training).
We can extract summary statistics about Tajima’s D, heterozygosity quantiles, haplotype extension for each simulation. The full list of sumstats is in annex. There are many summary statistics, but we can simplify them into a PCA and run that into a classical ABC framework. This can be a useful comparison to more advanced deep learning techniques, and makes it easier to explore the behaviour of our statistics depending on the parameters. It is also very helpful to spot bugs in the model, or determine whether we approximate our observed data well.
Training was done by randomly subsampling between 20 and 1000 TEs (capped at 1000 TEs collected over 20 windows), to expose the algorithm to varying amounts of variance. Of course, some statistics still reflect the original number of TEs in the dataset (for example, higher transposition rates give more TEs and therefore higher TE density values in flanking windows). We benchmark by running the trained algorithm on a set of 200 independent simulations of 200 x 5Mb windows that were not used for training. This allows testing the robustness of our inference pipelines to a number of TEs that is never encountered during training (ten times more). Note that 10,000 polymorphic TEs may be on the high side for any given lineage (let’s say ATHILA in plants, or any consensus of interest), although such numbers may be reached by pooling TEs into broader families (for example all Ty1/Copia together, although this may not make much sense biologically in some cases). The point is that the method is trained on 1000, which may be a more realistic number for (abundant) lineages in most genomes.
(a) Normalised sumstats + ABC/network + DFE
We can obtain the distribution of per-TE summary statistics using this python script, which runs on the datasets generated for each of the SLiM simulations (see also the git repository and the annex for more commands and explanations on the simulation + training commands):
Code
python 07_export_sumstat_summaries.py --model simple --dataset <path_to_collated_training_dataset> --out summaries_simple.csvWe also obtain the same summary statistics automatically for each 200 benchmarking datasets (themselves containing 200 x 5 Mb windows spanning a range of gene density and recombination rates) when we run the benchmarking scripts after training the deep learning algorithm. So we can now use ABC to examine our summary stats and compare.
Code
# ABC benchmark on the per-TE summary statistics.
library(abc)
library(HDInterval)
library(dplyr)
library(ggplot2)
library(patchwork)
library(readr); library(dplyr); library(stringr)
library(FactoMineR); library(data.table)
base="Results_final_blogpost/DFE_model_sumstats/"
ABC_METHOD <- "neuralnet" # "rejection" is ~10x faster for a first pass but not very efficient
ABC_TOL <- 0.05
CRED <- 0.90 # can be set to any interval deemed interesting, can be a way to check different slices of the posterior distribution. More of a sanity check.
NPCS <- 10
All <- na.omit(read.csv(file.path(base,"../summaries_simple.csv"), header = TRUE, fill = TRUE))
parameters <- All[, 1:6]
##mean= alpha/beta with beta the rate and alpha the shape. so rate is shape/mean
parameters$p_strong <- pgamma(100, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(1e-2, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
parameters$p_weak <- pgamma(1e-2, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(1e-3, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
parameters$p_nn <-pgamma(1e-3, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(0, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
train_df <- All[, 8:ncol(All)]
# drop zero-variance / constant columns and column about the number of TEs per simulation (which is clearly different here by one order of magnitude)
train_df <- train_df %>% dplyr::select(where(~ if (is.numeric(.)) var(., use = "complete.obs") != 0 else length(unique(.)) > 1))
res.pca <- PCA(train_df, graph = FALSE, ncp = 10)
output_coord <- data.frame(res.pca$ind$coord)
output_coord$Category <- ifelse(exp(parameters$log_mean_s) > 0.01,
"strong sel", "weaker sel")
pnames <- colnames(parameters)
P <- length(pnames)
train_pcs <- output_coord[, paste0("Dim.", seq_len(NPCS))]
# --- PCA scatter of training vs benchmark, with the CORRECT projection ---
observed_full <- read.csv(file.path(base, "benchmark_NTE10000_abc_summaries.csv"), header = TRUE)
observed_full <- na.omit(observed_full) # drop rows first -> alignment holds
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # same order
observed_df <- as.data.frame(predict(res.pca, newdata = obs_sumstat)$coord)
observed_df$Category <- "Observed"
print(
ggplot(rbind(output_coord, observed_df),
aes(Dim.1, Dim.2, color = Category)) +
geom_point(size = 3, alpha = 0.8) + theme_minimal() +
scale_color_brewer(palette = "Set1"))
The good thing is that the “observed” summary statistics fall within the cloud of training simulations, so using normalized statistics works reasonably well and is not too sensitive to the change in the number of TEs, although it would be more rigorous to compare simulations producing the same number (so 20 windows against 20 windows, not 200).
We can dig a bit more and examine the correlation of summary statistics with distinct parameters.
Code
library(dplyr); library(ggplot2)
param_cols <- colnames(All[, 1:6])
target <- "log_lambda" # parameter to screen
method <- "spearman" # "spearman" (robust) or "pearson"
partial <- TRUE # control for the other parameters?
top_n <- 50
stopifnot(target %in% names(All))
# stats = everything that isn't a parameter (n_valid_te kept as a candidate stat)
stat_cols <- setdiff(names(All), param_cols)
stat_cols <- stat_cols[sapply(All[stat_cols], is.numeric)]
v <- sapply(All[stat_cols], var, na.rm = TRUE) # drop constant / all-NA cols
stat_cols <- stat_cols[is.finite(v) & v > 0]
y <- All[[target]]
if (!partial) {
cors <- sapply(stat_cols, function(s)
suppressWarnings(cor(y, All[[s]], method = method, use = "pairwise.complete.obs")))
} else {
others <- setdiff(param_cols, target)
others <- others[sapply(All[others], function(x) var(x, na.rm = TRUE) > 0)] # drop fixed params
ctrl <- as.matrix(All[others])
resid_of <- function(z) residuals(lm(z ~ ctrl, na.action = na.exclude))
ry <- resid_of(y)
cors <- sapply(stat_cols, function(s)
suppressWarnings(cor(ry, resid_of(All[[s]]), method = method, use = "pairwise.complete.obs")))
}
corr_tbl <- tibble(stat = stat_cols, cor = as.numeric(cors)) |>
filter(is.finite(cor)) |> arrange(desc(abs(cor)))
print(corr_tbl, n = top_n)# A tibble: 129 × 2
stat cor
<chr> <dbl>
1 n_valid_te 0.707
2 local_te_density__mean 0.650
3 local_te_density__q50 0.648
4 local_te_density__q75 0.645
5 local_te_density__q25 0.644
6 local_te_density__q90 0.642
7 local_te_density__q10 0.639
8 local_te_density__max 0.636
9 local_te_density__sd -0.540
10 local_te_density__min 0.216
11 tajD_carrier__max 0.210
12 he_carrier_sd__max 0.204
13 tajD_carrier__min -0.192
14 he_noncarrier_sd__max 0.189
15 tajD_noncarrier__min -0.188
16 local_recomb__max 0.180
17 tajD_noncarrier__max 0.180
18 local_snp_density__min -0.179
19 he_noncarrier_mean__max 0.176
20 hapext_noncarrier__min -0.172
21 n_flank_snps_norm__min -0.168
22 local_snp_density__q90 -0.165
23 he_carrier_mean__max 0.163
24 local_snp_density__q75 -0.160
25 hapext_noncarrier__max 0.149
26 n_flank_snps_norm__max 0.142
27 local_snp_density__q50 -0.133
28 hapext_carrier__max 0.132
29 te_freq__max 0.126
30 he_noncarrier_sd__min -0.122
31 local_snp_density__mean -0.111
32 he_noncarrier_mean__min -0.109
33 n_flank_snps_norm__q75 -0.103
34 n_flank_snps_norm__mean -0.0932
35 n_flank_snps_norm__q50 -0.0839
36 n_flank_snps_norm__q25 -0.0790
37 local_snp_density__q25 -0.0787
38 n_flank_snps_norm__q10 -0.0749
39 hapext_carrier__q75 -0.0718
40 local_snp_density__q10 -0.0672
41 n_flank_snps_norm__q90 -0.0632
42 he_noncarrier_mean__sd 0.0607
43 hapext_carrier__q90 -0.0605
44 local_recomb__q25 0.0587
45 tajD_noncarrier__sd 0.0544
46 he_noncarrier_sd__sd 0.0535
47 hapext_carrier__mean -0.0524
48 he_noncarrier_mean__q10 -0.0504
49 tajD_noncarrier__q10 -0.0492
50 hapext_carrier__q50 -0.0483
# ℹ 79 more rowsCode
plot_tbl <- corr_tbl |> slice_head(n = top_n) |>
mutate(stat = factor(stat, levels = rev(stat)))
ggplot(plot_tbl, aes(cor, stat, color = cor > 0)) +
geom_segment(aes(x = 0, xend = cor, yend = stat), linewidth = 0.6) +
geom_point(size = 2.5) +
geom_vline(xintercept = 0, linewidth = 0.3) +
scale_color_manual(values = c(`TRUE`="#2c7fb8", `FALSE`="#de2d26"), guide = "none") +
labs(x = sprintf("%s %scorrelation with %s", method,
if (partial) "partial " else "", target),
y = NULL, title = sprintf("Stats most associated with %s", target)) +
theme_minimal(base_size = 11)
Code
library(dplyr); library(tidyr); library(ggplot2)
target <- "log_mean_s"
which_stats <- c("local_recomb__mean","local_recomb__q90","local_snp_density__mean","local_te_density__mean") # <- your picks
smoother <- "loess" # "loess" reveals saturation; "lm" for a straight line
df <- All |>
dplyr::select(all_of(c(target, which_stats))) |>
pivot_longer(-all_of(target), names_to = "stat", values_to = "value")
# parameter on X, stat on Y: shows how each stat RESPONDS across the parameter
# range, so a stat that stops moving at high gamma_r (your r in [1,2] flattening)
# shows up directly as the curve going flat.
ggplot(df, aes(log10(exp(.data[[target]])), value)) +
geom_point(alpha = 0.25, size = 0.7) +
geom_smooth(method = smoother, se = TRUE, linewidth = 0.8, color = "#2c7fb8") +
facet_wrap(~ stat, scales = "free_y") +
labs(x = "log10(mean_s)", y = "summary statistic",
title = sprintf("Summary statistics vs %s", target)) +
theme_minimal(base_size = 11)
Here we can check that, for example, \(\lambda\) is strongly correlated with the number of TEs for each simulation, which is what we expect. We also clearly see the effect of increasing selection coefficients on (normalized) TE density in 250kb flanks, as well as SNP density (although the effect is much more subtle) due to linked selection. We are not working with the context model, so no clear effect on the local recombination rate for TEs is observed, apart maybe for very high selection rates, which might be due to selection being more efficient in regions of higher recombination.
We summarized all the summary statistics through a PCA, and will use coordinates on the first 10 PC axes to run ABC with neural networks:
Code
###estimate the discretized categories for the DFE: weak is s < 1/(2Ne_rescaled) = -0.001, weak is between -0.001 and -0.01, strong is lower than -0.01
if (file.exists(file.path(base, "abc_DFE_recovery_values.csv")) ) {
abc_df <- read.csv(file.path(base, "abc_DFE_recovery_values.csv"))
} else {
observed_full$p_strong <- pgamma(100, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(1e-2, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
observed_full$p_weak <- pgamma(1e-2, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(1e-3, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
observed_full$p_nn <-pgamma(1e-3, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(0, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
stopifnot(all(pnames %in% colnames(observed_full)))
stopifnot(all(colnames(train_df) %in% colnames(observed_full)))
obs_true <- observed_full[, pnames, drop = FALSE] # (n_obs x P) true params
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # SAME cols + order as PCA fit
obs_pcs <- observed_df[, seq_len(NPCS)]
stopifnot(nrow(obs_true) == nrow(obs_pcs))
rows <- vector("list", nrow(obs_pcs) * P)
k <- 1
n_fail <- 0
for (i in seq_len(nrow(obs_pcs))) {
res <- tryCatch(
abc(target = obs_pcs[i, ],
param = parameters,
sumstat = train_pcs,
tol = ABC_TOL,
method = ABC_METHOD),
error = function(e) { message(sprintf(" sim %d failed: %s", i - 1L, conditionMessage(e))); NULL })
if (is.null(res)) { n_fail <- n_fail + 1; next }
draws <- if (ABC_METHOD == "rejection") res$unadj.values else res$adj.values
draws <- as.matrix(draws) # (n_accepted x P)
for (j in seq_len(P)) {
dj <- draws[, j]
truej <- obs_true[i, j]
hpd <- HDInterval::hdi(dj, credMass = CRED) # 90% highest-density interval
rows[[k]] <- data.frame(
sim_idx = i - 1L, # 0-indexed, like the CSVs
param_name = pnames[j],
true_value = truej,
posterior_mean = mean(dj),
posterior_median = median(dj),
hpd_lo = hpd[1],
hpd_hi = hpd[2],
rank = sum(dj < truej), # for SBC
n_draws = length(dj),
stringsAsFactors = FALSE)
k <- k + 1
}
if (i %% 25 == 0 || i == nrow(obs_pcs))
cat(sprintf(" %d/%d done\n", i, nrow(obs_pcs)))
}
abc_df <- do.call(rbind, rows[seq_len(k - 1)])
if (n_fail) cat(sprintf(" %d benchmark sims failed and were skipped\n", n_fail))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
RMSE = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_lo & true_value <= hpd_hi), # at CRED
sbc_p = if (has_sbc)
tryCatch(suppressWarnings(
ks.test(rank / n_draws, "punif")$p.value),
error = function(e) NA_real_)
else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics <- bench_metrics(abc_df)
#cat(sprintf("\nABC benchmark (method=%s, tol=%.3f, coverage at %d%% HPD)\n", ABC_METHOD, ABC_TOL, round(100 * CRED)))
#print(as.data.frame(abc_metrics), digits = 4)
#same artifacts the Python pipeline writes, for archiving / side-by-side
write.csv(abc_df, file.path(base, "abc_DFE_recovery_values.csv"), row.names = FALSE)
write.csv(abc_metrics, file.path(base, "abc_DFE_benchmark_summary.csv"), row.names = FALSE)
}
# log10 conversion mirrors the DL figure: natural-log 'log_*' -> 'log10_*'.
abc_plot <- abc_df %>%
mutate(
is_log = grepl("^log_", param_name),
true_value = ifelse(is_log, true_value / log(10), true_value),
posterior_mean = ifelse(is_log, posterior_mean / log(10), posterior_mean),
hpd_lo = ifelse(is_log, hpd_lo / log(10), hpd_lo),
hpd_hi = ifelse(is_log, hpd_hi / log(10), hpd_hi),
param_name = ifelse(is_log, sub("^log_", "log10_", param_name), param_name)) %>%
dplyr::select(-is_log) %>%
rename(hpd_10pct = hpd_lo, hpd_90pct = hpd_hi) %>%
filter(!param_name %in% c("log10_duration", "log10_gamma_shape"))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
RMSE = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_10pct & true_value <= hpd_90pct), # at CRED
# sbc_p = if (has_sbc)
# tryCatch(suppressWarnings(
# ks.test(rank / n_draws, "punif")$p.value),
# error = function(e) NA_real_)
# else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics_log10 <- bench_metrics(abc_plot)
print(abc_metrics_log10)# A tibble: 7 × 5
param_name RMSE bias coverage n_sims
<chr> <dbl> <dbl> <dbl> <int>
1 log10_lambda 0.109 0.00874 0.805 200
2 log10_mean_s 0.519 -0.219 0.755 200
3 log10_multiplier 0.509 -0.0425 0.885 200
4 log10_t_start 0.439 -0.00707 0.87 200
5 p_nn 0.195 0.0738 0.7 200
6 p_strong 0.107 -0.0430 0.69 200
7 p_weak 0.175 -0.0288 0.72 200Code
# strong vs weak burst, keyed by the TRUE log10_multiplier (>= 2 = strong)
burst <- abc_plot %>%
filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
abc_plot <- left_join(abc_plot, burst, by = "sim_idx")
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = hpd_10pct, ymax = hpd_90pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme(legend.position = "none") + ggtitle(ttl)
A <- mk(filter(abc_plot, burst == "strong"), "Strong burst (ABC)")
B <- mk(filter(abc_plot, burst == "weak"), "Weak burst (ABC)")
A
Code
B
This is not too bad. Even with 5000 simulations, which is very low for ABC, we recover the expected trends. When selection is too low (> -0.001 after rescaling), it becomes undistinguishable from neutrality, and inferred s plateaus. Parameters about the burst are reasonably well recovered. For selection, strong bursts are associated with an underestimation of selection, but strong purifying selection can be distinguished.
We can check whether using the same number of TEs (roughly 1000) between training simulations and pseudo observations improves things:
Code
library(abc)
library(HDInterval)
library(dplyr)
library(ggplot2)
library(patchwork)
library(readr); library(dplyr); library(stringr)
library(FactoMineR); library(data.table)
base="Results_final_blogpost/DFE_model_sumstats/"
ABC_METHOD <- "neuralnet" # "rejection" is ~10x faster for a first pass but not very efficient
ABC_TOL <- 0.05
CRED <- 0.90 # can be set to any interval deemed interesting, can be a way to check different slices of the posterior distribution. More of a sanity check.
NPCS <- 10
All <- na.omit(read.csv(file.path(base,"../summaries_simple.csv"), header = TRUE, fill = TRUE))
parameters <- All[, 1:6]
##mean= alpha/beta with beta the rate and alpha the shape. so rate is shape/mean
parameters$p_strong <- pgamma(100, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(1e-2, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
parameters$p_weak <- pgamma(1e-2, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(1e-3, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
parameters$p_nn <-pgamma(1e-3, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s)) - pgamma(0, shape = exp(parameters$log_gamma_shape), rate = exp(parameters$log_gamma_shape)/exp(parameters$log_mean_s))
train_df <- All[, 8:ncol(All)]
# drop zero-variance / constant columns and column about the number of TEs per simulation (which is clearly different here by one order of magnitude)
train_df <- train_df %>% dplyr::select(where(~ if (is.numeric(.)) var(., use = "complete.obs") != 0 else length(unique(.)) > 1))
res.pca <- PCA(train_df, graph = FALSE, ncp = 10)
output_coord <- data.frame(res.pca$ind$coord)
output_coord$Category <- ifelse(exp(parameters$log_mean_s) > 0.01,
"strong sel", "weaker sel")
pnames <- colnames(parameters)
P <- length(pnames)
train_pcs <- output_coord[, paste0("Dim.", seq_len(NPCS))]
# --- PCA scatter of training vs benchmark, with the CORRECT projection ---
observed_full <- read.csv(file.path(base, "benchmark_NTE1000_abc_summaries.csv"), header = TRUE)
observed_full <- na.omit(observed_full) # drop rows first -> alignment holds
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # same order
observed_df <- as.data.frame(predict(res.pca, newdata = obs_sumstat)$coord)
observed_df$Category <- "Observed"
print(
ggplot(rbind(output_coord, observed_df),
aes(Dim.1, Dim.2, color = Category)) +
geom_point(size = 3, alpha = 0.8) + theme_minimal() +
scale_color_brewer(palette = "Set1"))
Code
if (file.exists(file.path(base, "abc_DFE_recovery_values_1000TEs.csv")) ) {
abc_df <- read.csv(file.path(base, "abc_DFE_recovery_values_1000TEs.csv"))
} else {
observed_full$p_strong <- pgamma(100, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(1e-2, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
observed_full$p_weak <- pgamma(1e-2, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(1e-3, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
observed_full$p_nn <-pgamma(1e-3, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s)) - pgamma(0, shape = exp(observed_full$log_gamma_shape), rate = exp(observed_full$log_gamma_shape)/exp(observed_full$log_mean_s))
stopifnot(all(pnames %in% colnames(observed_full)))
stopifnot(all(colnames(train_df) %in% colnames(observed_full)))
obs_true <- observed_full[, pnames, drop = FALSE] # (n_obs x P) true params
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # SAME cols + order as PCA fit
obs_pcs <- observed_df[, seq_len(NPCS)]
stopifnot(nrow(obs_true) == nrow(obs_pcs))
rows <- vector("list", nrow(obs_pcs) * P)
k <- 1
n_fail <- 0
for (i in seq_len(nrow(obs_pcs))) {
res <- tryCatch(
abc(target = obs_pcs[i, ],
param = parameters,
sumstat = train_pcs,
tol = ABC_TOL,
method = ABC_METHOD),
error = function(e) { message(sprintf(" sim %d failed: %s", i - 1L, conditionMessage(e))); NULL })
if (is.null(res)) { n_fail <- n_fail + 1; next }
draws <- if (ABC_METHOD == "rejection") res$unadj.values else res$adj.values
draws <- as.matrix(draws) # (n_accepted x P)
for (j in seq_len(P)) {
dj <- draws[, j]
truej <- obs_true[i, j]
hpd <- HDInterval::hdi(dj, credMass = CRED) # 90% highest-density interval
rows[[k]] <- data.frame(
sim_idx = i - 1L, # 0-indexed, like the CSVs
param_name = pnames[j],
true_value = truej,
posterior_mean = mean(dj),
posterior_median = median(dj),
hpd_lo = hpd[1],
hpd_hi = hpd[2],
rank = sum(dj < truej), # for SBC
n_draws = length(dj),
stringsAsFactors = FALSE)
k <- k + 1
}
if (i %% 25 == 0 || i == nrow(obs_pcs))
cat(sprintf(" %d/%d done\n", i, nrow(obs_pcs)))
}
abc_df <- do.call(rbind, rows[seq_len(k - 1)])
if (n_fail) cat(sprintf(" %d benchmark sims failed and were skipped\n", n_fail))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
RMSE = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_lo & true_value <= hpd_hi), # at CRED
sbc_p = if (has_sbc)
tryCatch(suppressWarnings(
ks.test(rank / n_draws, "punif")$p.value),
error = function(e) NA_real_)
else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics <- bench_metrics(abc_df)
#cat(sprintf("\nABC benchmark (method=%s, tol=%.3f, coverage at %d%% HPD)\n",ABC_METHOD, ABC_TOL, round(100 * CRED)))
#print(as.data.frame(abc_metrics), digits = 4)
#same artifacts the Python pipeline writes, for archiving / side-by-side
write.csv(abc_df, file.path(base, "abc_DFE_recovery_values_1000TEs.csv"), row.names = FALSE)
write.csv(abc_metrics, file.path(base, "abc_DFE_benchmark_summary_1000TEs.csv"), row.names = FALSE)
}
# log10 conversion mirrors the DL figure: natural-log 'log_*' -> 'log10_*'.
abc_plot <- abc_df %>%
mutate(
is_log = grepl("^log_", param_name),
true_value = ifelse(is_log, true_value / log(10), true_value),
posterior_mean = ifelse(is_log, posterior_mean / log(10), posterior_mean),
hpd_lo = ifelse(is_log, hpd_lo / log(10), hpd_lo),
hpd_hi = ifelse(is_log, hpd_hi / log(10), hpd_hi),
param_name = ifelse(is_log, sub("^log_", "log10_", param_name), param_name)) %>%
dplyr::select(-is_log) %>%
rename(hpd_10pct = hpd_lo, hpd_90pct = hpd_hi) %>%
filter(!param_name %in% c("log10_duration", "log10_gamma_shape"))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
RMSE = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_10pct & true_value <= hpd_90pct), # at CRED
# sbc_p = if (has_sbc)
# tryCatch(suppressWarnings(
# ks.test(rank / n_draws, "punif")$p.value),
# error = function(e) NA_real_)
# else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics_log10 <- bench_metrics(abc_plot)
print(abc_metrics_log10)# A tibble: 7 × 5
param_name RMSE bias coverage n_sims
<chr> <dbl> <dbl> <dbl> <int>
1 log10_lambda 0.108 -0.0191 0.81 200
2 log10_mean_s 0.495 -0.135 0.78 200
3 log10_multiplier 0.506 -0.0100 0.895 200
4 log10_t_start 0.443 -0.00654 0.87 200
5 p_nn 0.179 0.0345 0.8 200
6 p_strong 0.0967 -0.0215 0.825 200
7 p_weak 0.172 -0.0131 0.76 200Code
# strong vs weak burst, keyed by the TRUE log10_multiplier (>= 2 = strong)
burst <- abc_plot %>%
filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
abc_plot <- left_join(abc_plot, burst, by = "sim_idx")
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = hpd_10pct, ymax = hpd_90pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme(legend.position = "none") + ggtitle(ttl)
A <- mk(filter(abc_plot, burst == "strong"), "Strong burst (ABC)")
B <- mk(filter(abc_plot, burst == "weak"), "Weak burst (ABC)")
A
Code
B
It does look like there is an improvement. This is perhaps not surprising. The distributions of our summary statistics shifts when including more TEs, creating a mismatch between observations and simulations. A possible way to circumvent this would be to use ABC with random forests, which may be less sensitive to this issue. However, this method does not produce joint posterior distributions.
(b) Per-TE sumstats + MLP + attention + DFE
We can now look at the results obtained with our deep learning inference method, using a (very) large number of attention heads. I am not showing all the tests here, but the benefits of more attention heads seems to be maximal when going from 3 to 6 attention heads, then becomes negligible. We look at inference of a 200 x 5 Mb windows benchmark.
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/DFE_model_sumstats_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures//benchmark_NTE10000/"
if (file.exists("combined_1.csv") )
{combined=read.csv("combined_1.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# ---- recompute 50% CIs from the raw posterior draws ----
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
df_plot <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
df_plot$is_log_param <- NULL
# ---- plot original variables (shape and mean): 80% HPD as the error bar ----
#ggplot(df_plot, aes(x = true_value, y = posterior_mean, color = param_name)) +
# geom_point(size = 3) +
# geom_errorbar(aes(ymin = hpd_10pct, ymax = hpd_90pct), width = 0.1, alpha = 0.3) +
# geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
# facet_wrap(~ param_name, scales = "free", ncol = 2) +
# labs(
# x = "True Value (log10 scale for log_* parameters)",
# y = "Posterior Mean (log10 scale for log_* parameters)",
# title = "True Value vs. Posterior Mean — 50% CI"
# ) +
# theme_minimal() +
# theme(
# panel.grid.major = element_blank(),
# panel.grid.minor = element_blank(),
# legend.position = "none"
# )
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
bin_samples <- d$f[["bin_samples"]] # (n_sims, n_post, 3)
true_bins <- d$f[["true_bins"]] # (n_sims, 3)
bin_names <- as.character(d$f[["bin_names"]]) # e.g. c("p_nn","p_w","p_strong")
dimsb <- dim(bin_samples)
n_sims <- dimsb[1]; n_bins <- dimsb[3]
# Build a tidy table: point estimate + 50% (25/75) + 90% (5/95) CIs per sim/bin.
bin_rows <- vector("list", n_sims * n_bins)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_bins)) {
draws <- bin_samples[i, , j]
q <- quantile(draws, c(0.05, 0.25, 0.75, 0.95), names = FALSE)
bin_rows[[k]] <- data.frame(
sim_idx = i - 1L, # match CSV convention; see note
bin_name = bin_names[j],
true_value = true_bins[i, j],
posterior_mean = mean(draws),
ci_5pct = q[1],
hpd_10pct = q[2],
hpd_90pct = q[3],
ci_95pct = q[4]
)
k <- k + 1
}
}
bins_df <- do.call(rbind, bin_rows)
# Order the facets neutral -> weak -> strong for readability.
bins_df$bin_name <- factor(bins_df$bin_name,
levels = c("p_nn", "p_w", "p_strong"))
common <- c("sim_idx", "param_name", "true_value", "posterior_mean",
"ci_5pct", "ci_95pct")
bins_for_join <- bins_df %>%
dplyr::rename(param_name = bin_name) %>%
dplyr::select(all_of(common))
combined <- bind_rows(df_plot %>% dplyr::select(all_of(common)), bins_for_join)
###getting rid of duration and gamma shape, poorly identified and not that important.
combined=subset(combined,combined$param_name != "log10_duration")
combined=subset(combined,combined$param_name != "log10_gamma_shape")
write.csv(combined,file = "combined_1.csv")
}
###RMSE
rmse_tbl <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl)# A tibble: 7 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 log10_lambda 0.0562 0.00148 0.0410 200
2 log10_mean_s 0.427 -0.0812 0.320 200
3 log10_multiplier 0.453 -0.0499 0.315 200
4 log10_t_start 0.399 -0.0206 0.270 200
5 p_nn 0.151 0.0357 0.115 200
6 p_strong 0.0868 -0.0207 0.0606 200
7 p_w 0.147 -0.0151 0.110 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
C=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin =ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (MLP)")
D=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (MLP)")
C
Code
D
Results are also decent, and there is an improvement compared to the ABC/neuralnet approach, especially when looking at the fit on a ~10,000 TEs benchmark. The main advantage is that we do not need to obtain summary statistics about the distribution of individual TE summary statistics, and the method seems more robust to deployment on a large number of TEs.
What happens if we apply the model trained on 1000 TEs to a lower or a higher number of TEs during benchmarking? We can look for example at RMSE:
Code
table=read.csv("Results_final_blogpost/DFE_model_sumstats_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures/benchmark_NTE_sweep_summary.csv")
ggplot(table, aes(n_te, rmse, color = param)) +
geom_point(size = 3) +
facet_wrap(~ param, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("RMSE over increasing number of TEs included in benchmark")
As expected, the lower the number of TEs we use, the higher the error rate. An interesting aspect however is that including more TEs does not harm, and even helps, recovering parameters, apart for a slight increase for the shape parameter. The average selection coefficient in particular seems to benefit a lot from including more TEs. This is a valuable aspect of the method, although I hope to find ways to further improve this behaviour. In particular, how low can we go during training: would it be reasonable to train only on 100 TEs when simulations are costly, while deploying on 1000 or 10,000 TEs, taking advantage of the increased information?
(c) CNN + attention + DFE
This shows the result of training a model on an image summarizing 60 TE genotypes per genomic window, genotypes at their flanks, and a set of 60 random SNPs collected within the same genomic window. I also add channels providing information about the (genetic and physical) distance of each flanking marker to the focal TE, as well as local SNP and TE density. We can plot what the image and channels for a given window look like with a python script:
Code
python inspect_sim.py simulations_CNN_context/sim_000001.npz
Note that here, the insertion rate is low, and we have to pad the missing TEs and flanks (in grey). Some flanking SNPs are also missing since we require a minimal distance between them and any TE to minimize linkage. The original idea was to get a SNP contrast to the TEs, which would be closer to neutrality and display a stable mutation rate.
There are many things that I could probably test to improve this. For example, I partition the image so that TEs are on the left side, and their flanks then follow, the idea being that the algorithm can learn to focus its attention to the distinct, constant parts of the image to extract relevant features. I am not fully convinced that this would be superior to a more classical representation where flanking SNPs actually flank TEs.
Below are the results for inference with two attention heads. Note that this version of the algorithm takes about six hours to train, and consumes a lot of GPU resource. The fact that it does not really perform better than the previous summary statistics method makes it the lowest ranked in my opinion. I keep it for future possible improvement. Here too we can look at benchmark on 200 x 5 Mb windows.
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/DFE_model_CNN/benchmark_K200//"
if (file.exists("combined_2.csv") )
{combined=read.csv("combined_2.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# ---- recompute 50% CIs from the raw posterior draws ----
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
df_plot <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
df_plot$is_log_param <- NULL
# ---- plot original variables (shape and mean): 80% HPD as the error bar ----
#ggplot(df_plot, aes(x = true_value, y = posterior_mean, color = param_name)) +
# geom_point(size = 3) +
# geom_errorbar(aes(ymin = hpd_10pct, ymax = hpd_90pct), width = 0.1, alpha = 0.3) +
# geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
# facet_wrap(~ param_name, scales = "free", ncol = 2) +
# labs(
# x = "True Value (log10 scale for log_* parameters)",
# y = "Posterior Mean (log10 scale for log_* parameters)",
# title = "True Value vs. Posterior Mean — 50% CI"
# ) +
# theme_minimal() +
# theme(
# panel.grid.major = element_blank(),
# panel.grid.minor = element_blank(),
# legend.position = "none"
# )
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
bin_samples <- d$f[["bin_samples"]] # (n_sims, n_post, 3)
true_bins <- d$f[["true_bins"]] # (n_sims, 3)
bin_names <- as.character(d$f[["bin_names"]]) # e.g. c("p_nn","p_w","p_strong")
dimsb <- dim(bin_samples)
n_sims <- dimsb[1]; n_bins <- dimsb[3]
# Build a tidy table: point estimate + 50% (25/75) + 90% (5/95) CIs per sim/bin.
bin_rows <- vector("list", n_sims * n_bins)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_bins)) {
draws <- bin_samples[i, , j]
q <- quantile(draws, c(0.05, 0.25, 0.75, 0.95), names = FALSE)
bin_rows[[k]] <- data.frame(
sim_idx = i - 1L, # match CSV convention; see note
bin_name = bin_names[j],
true_value = true_bins[i, j],
posterior_mean = mean(draws),
ci_5pct = q[1],
hpd_10pct = q[2],
hpd_90pct = q[3],
ci_95pct = q[4]
)
k <- k + 1
}
}
bins_df <- do.call(rbind, bin_rows)
# Order the facets neutral -> weak -> strong for readability.
bins_df$bin_name <- factor(bins_df$bin_name,
levels = c("p_nn", "p_w", "p_strong"))
common <- c("sim_idx", "param_name", "true_value", "posterior_mean",
"ci_5pct", "ci_95pct")
bins_for_join <- bins_df %>%
dplyr::rename(param_name = bin_name) %>%
dplyr::select(all_of(common))
combined <- bind_rows(df_plot %>% dplyr::select(all_of(common)), bins_for_join)
###getting rid of duration and gamma shape, poorly identified and not that important.
combined=subset(combined,combined$param_name != "log10_duration")
combined=subset(combined,combined$param_name != "log10_gamma_shape")
write.csv(combined,file="combined_2.csv") }
###RMSE
rmse_tbl <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl)# A tibble: 7 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 log10_lambda 0.103 0.00675 0.0754 200
2 log10_mean_s 0.504 0.0406 0.385 200
3 log10_multiplier 0.514 -0.0318 0.401 200
4 log10_t_start 0.452 -0.0179 0.335 200
5 p_nn 0.177 -0.0161 0.141 200
6 p_strong 0.109 0.0154 0.0827 200
7 p_w 0.167 0.000778 0.133 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
E=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin =ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (CNN)")
F=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (CNN)")
E
Code
F
And we can inspect whether deploying the model trained with 20 windows (~1000 TEs) does benefit, even slightly, from the increased number of observations:
Code
table=read.csv("Results_final_blogpost/DFE_model_CNN/benchmark_K_sweep_summary.csv")
ggplot(table, aes(K, rmse, color = param)) +
geom_point(size = 3) +
facet_wrap(~ param, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("RMSE over increasing number of TEs included in benchmark")
Here again, using more windows helps, except for the gamma shape, although variation in RMSE remains extremely low for this challenging parameter. The improvement is particularly pronounced for the average selection coefficient.
(d) Normalised sumstats + ABC/neuralnet + context
Again, the pipeline gives us summary statistics that we can use to check that simulations are behaving in the way we expect. For example, we can examine whether changing the strength of selection with recombination rate behaves as expected.
Code
library(dplyr)
library(ggplot2)
base="./Results_final_blogpost/"
All <- read.csv(file.path(base,"summaries_context.csv"), header = TRUE, fill = TRUE)
param_cols <- c("log_lambda","log_s_global","alpha_r","gamma_r","log10_t_start","log_duration","log_multiplier")
target <- "alpha_r" # parameter to screen
method <- "spearman" # "spearman" (robust) or "pearson"
partial <- TRUE # control for the other parameters?
top_n <- 50
stopifnot(target %in% names(All))
# stats = everything that isn't a parameter (n_valid_te kept as a candidate stat)
stat_cols <- setdiff(names(All), param_cols)
stat_cols <- stat_cols[sapply(All[stat_cols], is.numeric)]
v <- sapply(All[stat_cols], var, na.rm = TRUE) # drop constant / all-NA cols
stat_cols <- stat_cols[is.finite(v) & v > 0]
y <- All[[target]]
if (!partial) {
cors <- sapply(stat_cols, function(s)
suppressWarnings(cor(y, All[[s]], method = method, use = "pairwise.complete.obs")))
} else {
others <- setdiff(param_cols, target)
others <- others[sapply(All[others], function(x) var(x, na.rm = TRUE) > 0)] # drop fixed params
ctrl <- as.matrix(All[others])
resid_of <- function(z) residuals(lm(z ~ ctrl, na.action = na.exclude))
ry <- resid_of(y)
cors <- sapply(stat_cols, function(s)
suppressWarnings(cor(ry, resid_of(All[[s]]), method = method, use = "pairwise.complete.obs")))
}
corr_tbl <- tibble(stat = stat_cols, cor = as.numeric(cors)) |>
filter(is.finite(cor)) |> arrange(desc(abs(cor)))
#print(corr_tbl, n = top_n)
plot_tbl <- corr_tbl |> slice_head(n = top_n) |>
mutate(stat = factor(stat, levels = rev(stat)))
ggplot(plot_tbl, aes(cor, stat, color = cor > 0)) +
geom_segment(aes(x = 0, xend = cor, yend = stat), linewidth = 0.6) +
geom_point(size = 2.5) +
geom_vline(xintercept = 0, linewidth = 0.3) +
scale_color_manual(values = c(`TRUE`="#2c7fb8", `FALSE`="#de2d26"), guide = "none") +
labs(x = sprintf("%s %scorrelation with %s", method,
if (partial) "partial " else "", target),
y = NULL, title = sprintf("Stats most associated with %s", target)) +
theme_minimal(base_size = 11)
Code
library(dplyr); library(tidyr); library(ggplot2)
target <- "alpha_r"
which_stats <- c("local_recomb__mean","local_recomb__q90","local_snp_density__mean","local_te_density__mean")
smoother <- "loess" # "loess" reveals saturation; "lm" for a straight line
df <- All |>
dplyr::select(all_of(c(target, which_stats))) |>
pivot_longer(-all_of(target), names_to = "stat", values_to = "value")
# parameter on X, stat on Y: shows how each stat RESPONDS across the parameter
# range, so a stat that stops moving at high gamma_r (your r in [1,2] flattening)
# shows up directly as the curve going flat.
ggplot(df, aes(.data[[target]], value)) +
geom_point(alpha = 0.25, size = 0.7) +
geom_smooth(method = smoother, se = TRUE, linewidth = 0.8, color = "#2c7fb8") +
facet_wrap(~ stat, scales = "free_y") +
labs(x = target, y = "summary statistic",
title = sprintf("Summary statistics vs %s", target)) +
theme_minimal(base_size = 11)
We can then use an ABC approach to infer parameters, contrasting our pseudo-observed data (200 x 5Mb windows, about ~10,000 TEs) with the ones used for training.
Code
library(abc)
library(HDInterval)
library(dplyr)
library(ggplot2)
library(patchwork)
library(readr); library(dplyr); library(stringr)
library(FactoMineR); library(data.table)
base="./Results_final_blogpost/"
ABC_METHOD <- "neuralnet" #"rejection" is ~10x faster for a first pass but is far less efficient
ABC_TOL <- 0.05
CRED <- 0.90 # credible level, to line up with the DL 80% HPDs. Mostly to show this is possible to show different confidence intervals.
NPCS <- 10
All <- read.csv(file.path(base,"summaries_context.csv"), header = TRUE, fill = TRUE)
parameters <- All[, 1:7]
train_df <- All[, 9:ncol(All)] #exclude n_valid_te, not used by the network.
# drop zero-variance / constant columns
train_df <- train_df[,1:(ncol(train_df)-1)] %>% dplyr::select(where(~ if (is.numeric(.)) var(., use = "complete.obs") != 0 else length(unique(.)) > 1))
##bug to fix in distance to nearest gene, correlation coeff not reported
res.pca <- PCA(train_df, graph = FALSE, ncp = 10)
output_coord <- data.frame(res.pca$ind$coord)
output_coord$Category="Simulations"
pnames <- colnames(parameters)
P <- length(pnames)
train_pcs <- output_coord[, paste0("Dim.", seq_len(NPCS))]
# --- PCA scatter of training vs benchmark, with the CORRECT projection ---
observed_full <- read.csv(file.path(base, "context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures/benchmark_NTE10000_abc_summaries.csv"), header = TRUE)
observed_full <- na.omit(observed_full) # drop rows first -> alignment holds
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # same order
observed_df <- as.data.frame(predict(res.pca, newdata = obs_sumstat)$coord)
observed_df$Category <- "Observed"
print(
ggplot(rbind(output_coord, observed_df),
aes(Dim.1, Dim.2, color = Category)) +
geom_point(size = 3, alpha = 0.8) + theme_minimal() +
scale_color_brewer(palette = "Set1"))
Code
print(
ggplot(rbind(output_coord, observed_df),
aes(Dim.2, Dim.3, color = Category)) +
geom_point(size = 3, alpha = 0.8) + theme_minimal() +
scale_color_brewer(palette = "Set1"))
Code
if (file.exists(file.path(base, "abc_recovery_values_context.csv")) ) {
abc_df <- read.csv(file.path(base, "abc_recovery_values_context.csv"))
} else {
stopifnot(all(pnames %in% colnames(observed_full)))
stopifnot(all(colnames(train_df) %in% colnames(observed_full)))
obs_true <- observed_full[, pnames, drop = FALSE] # (n_obs x P) true params
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # SAME cols + order as PCA fit
obs_pcs <- observed_df[, seq_len(NPCS)]
stopifnot(nrow(obs_true) == nrow(obs_pcs))
rows <- vector("list", nrow(obs_pcs) * P)
k <- 1
n_fail <- 0
for (i in seq_len(nrow(obs_pcs))) {
res <- tryCatch(
abc(target = obs_pcs[i, ],
param = parameters,
sumstat = train_pcs,
tol = ABC_TOL,
method = ABC_METHOD),
error = function(e) { message(sprintf(" sim %d failed: %s", i - 1L, conditionMessage(e))); NULL })
if (is.null(res)) { n_fail <- n_fail + 1; next }
draws <- if (ABC_METHOD == "rejection") res$unadj.values else res$adj.values
draws <- as.matrix(draws) # (n_accepted x P)
for (j in seq_len(P)) {
dj <- draws[, j]
truej <- obs_true[i, j]
hpd <- HDInterval::hdi(dj, credMass = CRED) # 80% highest-density interval
rows[[k]] <- data.frame(
sim_idx = i - 1L, # 0-indexed, like the CSVs
param_name = pnames[j],
true_value = truej,
posterior_mean = mean(dj),
posterior_median = median(dj),
hpd_lo = hpd[1],
hpd_hi = hpd[2],
rank = sum(dj < truej), # for SBC
n_draws = length(dj),
stringsAsFactors = FALSE)
k <- k + 1
}
if (i %% 25 == 0 || i == nrow(obs_pcs))
cat(sprintf(" %d/%d done\n", i, nrow(obs_pcs)))
}
abc_df <- do.call(rbind, rows[seq_len(k - 1)])
if (n_fail) cat(sprintf(" %d benchmark sims failed and were skipped\n", n_fail))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
RMSE = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_lo & true_value <= hpd_hi), # at CRED
# sbc_p = if (has_sbc)
# tryCatch(suppressWarnings(
# ks.test(rank / n_draws, "punif")$p.value),
# error = function(e) NA_real_)
# else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics <- bench_metrics(abc_df)
#cat(sprintf("\nABC benchmark (method=%s, tol=%.3f, coverage at %d%% HPD)\n",
# ABC_METHOD, ABC_TOL, round(100 * CRED)))
#print(as.data.frame(abc_metrics), digits = 4)
# same artifacts the Python pipeline writes, for archiving / side-by-side
write.csv(abc_df, file.path(base, "abc_recovery_values_context.csv"), row.names = FALSE)
write.csv(abc_metrics, file.path(base, "abc_benchmark_summary_context.csv"), row.names = FALSE)
}
# ---- final graph: true vs posterior mean, 80% HPD bars, strong/weak burst ----
# log10 conversion mirrors the DL figure: natural-log 'log_*' -> 'log10_*'.
abc_plot <- abc_df %>%
mutate(
is_log = grepl("^log_", param_name),
true_value = ifelse(is_log, true_value / log(10), true_value),
posterior_mean = ifelse(is_log, posterior_mean / log(10), posterior_mean),
hpd_lo = ifelse(is_log, hpd_lo / log(10), hpd_lo),
hpd_hi = ifelse(is_log, hpd_hi / log(10), hpd_hi),
param_name = ifelse(is_log, sub("^log_", "log10_", param_name), param_name)) %>%
dplyr::select(-is_log) %>%
rename(hpd_10pct = hpd_lo, hpd_90pct = hpd_hi) %>%
filter(!param_name %in% c("log10_duration", "log10_gamma_shape"))
bench_metrics <- function(d) {
has_sbc <- all(c("rank", "n_draws") %in% names(d))
d %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
coverage = mean(true_value >= hpd_10pct & true_value <= hpd_90pct), # at CRED
# sbc_p = if (has_sbc)
# tryCatch(suppressWarnings(
# ks.test(rank / n_draws, "punif")$p.value),
# error = function(e) NA_real_)
# else NA_real_,
n_sims = dplyr::n(),
.groups = "drop")
}
abc_metrics_log10 <- bench_metrics(abc_plot)
print(abc_metrics_log10)# A tibble: 6 × 5
param_name rmse bias coverage n_sims
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.731 -0.209 0.855 200
2 gamma_r 0.385 -0.113 0.9 200
3 log10_lambda 0.181 0.0828 0.54 200
4 log10_multiplier 0.591 0.103 0.74 200
5 log10_s_global 0.317 -0.149 0.745 200
6 log10_t_start 0.522 -0.119 0.745 200Code
# strong vs weak burst, keyed by the TRUE log10_multiplier (>= 2 = strong)
burst <- abc_plot %>%
filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
abc_plot <- left_join(abc_plot, burst, by = "sim_idx")
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = hpd_10pct, ymax = hpd_90pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme(legend.position = "none") + ggtitle(ttl)
G <- mk(filter(abc_plot, burst == "strong"), "Strong burst (ABC)")
H <- mk(filter(abc_plot, burst == "weak"), "Weak burst (ABC)")
G
Code
H
(e) Per-TE sumstats + MLP + attention + context
Here we can directly plot the output of the pipeline. We show the results for 10,000 TEs collected over 200 x 5Mb simulated windows (so 1 Gb of fake genome).
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures//benchmark_NTE10000/" ##this one is based on 10 000 TEs, sumstats, context, extrapolation, 6 attention heads.
if (file.exists("combined_3.csv") )
{combined=read.csv("combined_3.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# recompute 80% CIs from the raw posterior draws, just to show how it can be done. posterior_samples.npz contains everything needed to examine joint posterios if needed.
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
combined <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
combined$is_log_param <- NULL
###getting rid of duration fixed.
combined=subset(combined,combined$param_name != "log10_duration")
write.csv(combined,file="combined_3.csv")}
###RMSE
rmse_tbl_sumstat_context <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl_sumstat_context)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.467 0.0478 0.332 200
2 gamma_r 0.172 -0.00766 0.128 200
3 log10_lambda 0.0753 -0.00381 0.0526 200
4 log10_multiplier 0.505 0.0482 0.384 200
5 log10_s_global 0.158 -0.0170 0.111 200
6 log10_t_start 0.440 0.0231 0.323 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
Z=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (MLN)")
W=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (MLN)")
Z
Code
W
The reader may have noticed that in the case of strong bursts, estimates of negative \(\alpha_r\) are often poor across all methods. This is however explained by the fact that these simulations also have a very high or very low basal selective coefficient (most TEs are effectively neutral or strongly counterselected across the genome), which reduces the contrast.
Code
poor_fit_sims=subset(strong_burst,strong_burst$param_name=="alpha_r" & strong_burst$true_value < (-1) & strong_burst$posterior_mean > (-0.5))$sim_idx
poor_fit_sims_tbl=subset(strong_burst,strong_burst$sim_idx %in% poor_fit_sims & strong_burst$param_name=="log10_s_global")
better_fit_sims=subset(strong_burst,strong_burst$param_name=="alpha_r" & strong_burst$true_value < (-1) & strong_burst$posterior_mean < (-0.5))$sim_idx
better_fit_sims_tbl=subset(strong_burst,strong_burst$sim_idx %in% better_fit_sims & strong_burst$param_name=="log10_s_global")
library(ggplot2)
poor_fit_sims_tbl$fit <- "Poor fit"
better_fit_sims_tbl$fit <- "Better fit"
plot_df <- rbind(
poor_fit_sims_tbl[, c("true_value", "fit")],
better_fit_sims_tbl[, c("true_value", "fit")]
)
ggplot(plot_df, aes(x = true_value, y = fit,col=fit)) +
geom_jitter(
height = 0.12,
width = 0,
size = 3,
alpha = 0.7
) +
theme_classic() +
labs(
x = expression(log[10](s)),
y = NULL
)
Again, given that this configuration (negative \(\alpha_r\) means strong counter-selection in regions of low recombination) is less likely for TEs, this is probably not a major issue from a biological perspective.
Here too we can check whether adding more TEs (beyond training values) helps inference:
Code
table=read.csv("Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures/benchmark_NTE_sweep_summary.csv")
ggplot(table, aes(n_te, rmse, color = param)) +
geom_point(size = 3) +
facet_wrap(~ param, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("RMSE over increasing number of TEs included in inference")
(f) CNN + attention + context
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "/home/yannbourgeois/Desktop/CurrentWORK/TE_DatePalm/Trials_estimate_DFE/scripts_deep_learning/latest/Results_final_blogpost/context_model_CNN/benchmark_K200/"
if (file.exists("combined_4.csv") )
{combined=read.csv("combined_4.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
#bins=read.csv(file.path(base,"binned_props_available.csv"))
df = merge(parameters, sim_properties, by = "sim_idx")
# ---- recompute 50% CIs from the raw posterior draws ----
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
df_plot <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
df_plot$is_log_param <- NULL
###getting rid of duration and gamma shape, poorly identified and not that important.
combined=subset(df_plot,df_plot$param_name != "log10_duration")
write.csv(combined,file="combined_4.csv")
}
###RMSE
rmse_tbl <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.672 0.204 0.544 200
2 gamma_r 0.332 0.101 0.251 200
3 log10_lambda 0.112 0.0144 0.0818 200
4 log10_multiplier 0.630 0.199 0.516 200
5 log10_s_global 0.191 0.0135 0.148 200
6 log10_t_start 0.504 0.115 0.391 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
X=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst")
Y=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst")
X
Code
Y
This is not too bad, but the method struggles more with burst parameters. The average selection coefficient remains well estimated, but its spread and association with recombination are not as well estimated. Maybe changing the image representation to better reflect the correlation between local recombination and TE flanking diversity would help.
A faster route with msprime. The potential of multi-fidelity (MF-NPE).
A purely neutral version of burst detection should be implementable with coalescent-based msprime as the simulator, which would make simulation dramatically faster. Mutations can be overlaid on msprime tree sequences at rates that vary over time and along the genome; the one thing msprime cannot model natively is selection, which is why SLiM has been necessary so far. If we assume a neutral model, msprime is enough to build a method that recovers every parameter relating to insertion rate and preference.
But we can go further. If we can find a way to crudely approximate the effect of purifying selection with msprime, there may be a way to i) train a model on these crude simulations, and ii) recalibrate this model using (ideally few) more principled SLiM simulations. This approach has been proposed under the name “multi-fidelity simulation-based inference” (Krouglova et al. 2026).
Selection can be approximated at TE sites. I originally thought of following the idea of Xue and colleagues (Xue et al. 2018), by starting from the inferred demography and adding an expansion parameter \(\nu\) that mimics purifying selection: purifying selection holds deleterious elements rare, which distorts the frequency spectrum in the same way as a population expansion. Making that expansion recombination-dependent reproduces the \(\alpha_r\) effect. A small SLiM grid might then be enough to calibrate the mapping from \(\nu\) to a true selection coefficient and to map where the proxy is valid.
Approximating selection with msprime, moments, and tree sequences.
Ultimately, after a few tests, I decided against it (although the option remains in script 01_simulate_msprime.py). This produces simulations that are too distinct from the SLiM ones to be very useful, and the equivalence between \(\nu\) and selection is not straightforward. Fortunately, there is a more elegant way. The way selection increases the probability for the counter-selected allele to be removed can be obtained, building on diffusion equations in population genetics (Poisson Random Field from Sawyer and Hartl (Sawyer and Hartl 1992)). The moments package (Jouganous et al. 2017) can provide such an estimate through its function moments.LinearSystem_1D.steady_state_1D(). The idea here is to estimate the neutral spectrum, the one expected under a given selection coefficients, and obtain the relative probability for an allele to remain at the observed frequency or to be pruned at any given site. We then remove TE variants randomly in the lineage that carries them with a probability depending on \(\gamma\) = 4 \(N_e\)\(s\). This provides a much closer approximation of the effect of selection on individual TEs, but at this stage the simulator does not simulate the effects of linked selection. This, however, produces simulations that are much closer to SLiM, making retraining with a few SLiM simulations much more efficient.
Below is a comparison of 10,000 msprime simulations and 5000 SLiM simulations, using the same prior range. We project the msprime simulations in the PCA space obtained for SLiM data.
Code
library(abc)
library(HDInterval)
library(dplyr)
library(ggplot2)
library(patchwork)
library(readr); library(dplyr); library(stringr)
library(FactoMineR); library(data.table)
base="./Results_final_blogpost/"
NPCS <- 10
All <- read.csv(file.path(base,"summaries_context.csv"), header = TRUE, fill = TRUE)
parameters <- All[, 1:7]
train_df <- All[, 9:ncol(All)] #exclude n_valid_te, not used by the network.
# drop zero-variance / constant columns
train_df <- train_df[,1:(ncol(train_df)-1)] %>% dplyr::select(where(~ if (is.numeric(.)) var(., use = "complete.obs") != 0 else length(unique(.)) > 1))
##bug to fix in distance to nearest gene, correlation coeff not reported
res.pca <- PCA(train_df, graph = FALSE, ncp = 10)
output_coord <- data.frame(res.pca$ind$coord)
output_coord$Category="SLiM sims"
pnames <- colnames(parameters)
P <- length(pnames)
train_pcs <- output_coord[, paste0("Dim.", seq_len(NPCS))]
# --- PCA scatter of training vs benchmark, with the CORRECT projection ---
observed_full <- read.csv(file.path(base, "summaries_msprime_context.csv"), header = TRUE)
observed_full <- na.omit(observed_full) # drop rows first -> alignment holds
obs_sumstat <- observed_full[, colnames(train_df), drop = FALSE] # same order
observed_df <- as.data.frame(predict(res.pca, newdata = obs_sumstat)$coord)
observed_df$Category <- "msprime sims"
ggplot(rbind(output_coord, observed_df),
aes(Dim.1, Dim.2, color = Category)) +
geom_point(size = 3, alpha = 0.1) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~Category)
Code
ggplot(rbind(output_coord, observed_df),
aes(Dim.2, Dim.3, color = Category)) +
geom_point(size = 3, alpha = 0.1) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~Category)
It looks like the coordinates for the first PC axis are strongly flattened for msprime compared to SLiM, although this mostly affects the . We can check what summary statistics are the most associated with that axis:
Code
res.pca$var$coord[, "Dim.1"] te_freq__mean te_freq__sd
0.20365541 0.08783517
te_freq__max te_freq__q10
-0.03562001 0.02551889
te_freq__q25 te_freq__q50
0.08139632 0.17608374
te_freq__q75 te_freq__q90
0.25965823 0.20825512
he_carrier_mean__mean he_carrier_mean__sd
0.50531603 0.51785685
he_carrier_mean__max he_carrier_mean__q25
0.25152574 0.05103595
he_carrier_mean__q50 he_carrier_mean__q75
0.37628931 0.48653271
he_carrier_mean__q90 he_carrier_sd__mean
0.50223166 0.49603071
he_carrier_sd__sd he_carrier_sd__max
0.49561530 0.17076584
he_carrier_sd__q25 he_carrier_sd__q50
0.08189305 0.44114470
he_carrier_sd__q75 he_carrier_sd__q90
0.45454206 0.44098297
he_noncarrier_mean__mean he_noncarrier_mean__sd
0.82391035 -0.47008512
he_noncarrier_mean__min he_noncarrier_mean__max
0.16522567 -0.19925830
he_noncarrier_mean__q10 he_noncarrier_mean__q25
0.74982517 0.83211035
he_noncarrier_mean__q50 he_noncarrier_mean__q75
0.84467264 0.80367785
he_noncarrier_mean__q90 he_noncarrier_sd__mean
0.69209640 0.91035535
he_noncarrier_sd__sd he_noncarrier_sd__min
-0.76645944 0.22913945
he_noncarrier_sd__max he_noncarrier_sd__q10
-0.10670945 0.87704314
he_noncarrier_sd__q25 he_noncarrier_sd__q50
0.89375373 0.88433921
he_noncarrier_sd__q75 he_noncarrier_sd__q90
0.82849914 0.65797443
tajD_carrier__mean tajD_carrier__sd
0.19809948 0.38050461
tajD_carrier__min tajD_carrier__max
-0.01901368 0.17688589
tajD_carrier__q10 tajD_carrier__q25
-0.26280659 -0.10405063
tajD_carrier__q50 tajD_carrier__q75
0.07540060 0.11293254
tajD_carrier__q90 tajD_noncarrier__mean
0.54586934 0.94263678
tajD_noncarrier__sd tajD_noncarrier__min
-0.51055878 0.47846097
tajD_noncarrier__max tajD_noncarrier__q10
-0.07515928 0.90507151
tajD_noncarrier__q25 tajD_noncarrier__q50
0.93198649 0.92539490
tajD_noncarrier__q75 tajD_noncarrier__q90
0.90436298 0.81296171
hapext_carrier__mean hapext_carrier__sd
0.63992096 0.41544583
hapext_carrier__max hapext_carrier__q10
0.01503397 0.05440268
hapext_carrier__q25 hapext_carrier__q50
0.37384060 0.42076577
hapext_carrier__q75 hapext_carrier__q90
0.72882237 0.81552468
hapext_noncarrier__mean hapext_noncarrier__sd
0.53202910 0.19938856
hapext_noncarrier__min hapext_noncarrier__max
0.17825871 0.10893459
hapext_noncarrier__q10 hapext_noncarrier__q25
0.66214581 0.62000493
hapext_noncarrier__q50 hapext_noncarrier__q75
0.54059836 0.45514117
hapext_noncarrier__q90 local_recomb__mean
0.39106756 0.86955580
local_recomb__sd local_recomb__min
-0.58452105 0.20794338
local_recomb__max local_recomb__q10
0.60325373 0.73671886
local_recomb__q25 local_recomb__q50
0.82992599 0.85281703
local_recomb__q75 local_recomb__q90
0.81781543 0.80702213
local_snp_density__mean local_snp_density__sd
0.94860088 -0.82780767
local_snp_density__min local_snp_density__max
0.66603395 0.21326907
local_snp_density__q10 local_snp_density__q25
0.88965480 0.91635268
local_snp_density__q50 local_snp_density__q75
0.87795456 0.77625500
local_snp_density__q90 local_te_density__mean
0.73734313 -0.40820625
local_te_density__sd local_te_density__min
-0.70434076 -0.16136981
local_te_density__max local_te_density__q10
-0.47970222 -0.17513231
local_te_density__q25 local_te_density__q50
-0.30032936 -0.44282697
local_te_density__q75 local_te_density__q90
-0.48628783 -0.49439098
n_flank_snps_norm__mean n_flank_snps_norm__sd
0.97685337 -0.54305590
n_flank_snps_norm__min n_flank_snps_norm__max
0.72400938 0.07327834
n_flank_snps_norm__q10 n_flank_snps_norm__q25
0.92859363 0.95691727
n_flank_snps_norm__q50 n_flank_snps_norm__q75
0.96542320 0.93571916
n_flank_snps_norm__q90 valid_tajD__mean
0.88788622 0.47577884
valid_tajD__sd valid_tajD__max
0.39131191 0.03524574
valid_tajD__q50 valid_tajD__q75
0.19722961 0.42740221
valid_tajD__q90 valid_hapext__mean
0.28613439 0.31193876
valid_hapext__sd valid_hapext__q10
-0.31644176 0.03852499
valid_hapext__q25 valid_hapext__q50
0.34667218 0.18412505
valid_hapext__q75 corr_te_freq_local_recomb
0.08804898 0.15132337
corr_local_te_density_local_recomb
0.80665361 Code
most_strongly_correlated=names(sort(abs(res.pca$var$coord[, "Dim.1"]), decreasing = TRUE))[1:10]
res.pca$var$coord[most_strongly_correlated, "Dim.1"]n_flank_snps_norm__mean n_flank_snps_norm__q50 n_flank_snps_norm__q25
0.9768534 0.9654232 0.9569173
local_snp_density__mean tajD_noncarrier__mean n_flank_snps_norm__q75
0.9486009 0.9426368 0.9357192
tajD_noncarrier__q25 n_flank_snps_norm__q10 tajD_noncarrier__q50
0.9319865 0.9285936 0.9253949
local_snp_density__q25
0.9163527 We see mostly positive correlations with flanking SNP diversity and allele frequency spectra. This makes sense: selection and background selection is properly simulated in SLiM, and results in a drop in diversity around the counterselected region which is less pronounced in neutral simulations. However, we can check that the two simulators produce the same trends. For example, we can inspect the way selection correlates with PC projections:
Code
total=rbind(cbind(parameters,output_coord),cbind(observed_full[,1:7],observed_df))
ggplot(total,aes(log_s_global, Dim.1, color = Category)) +
geom_point(size = 3, alpha = 0.8) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~Category)
We can inspect further that parameters controlling TE insertion rates and selection, such as s, \(\gamma_r\) and \(\alpha_r\) produce what we expect with both simulator. Do they alter TE frequency? The average recombination rate of regions where TEs are found?
Code
All$simulator="SLiM"
observed_full$simulator="msprime"
total2=rbind(All,observed_full)
ggplot(total2,aes(log_s_global, te_freq__mean, color = simulator)) +
geom_point(size = 3, alpha = 0.3) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~simulator)
Code
ggplot(total2,aes(alpha_r, local_recomb__mean, color = simulator)) +
geom_point(size = 3, alpha = 0.3) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~simulator)
Code
ggplot(total2,aes(gamma_r, local_recomb__mean, color = simulator)) +
geom_point(size = 3, alpha = 0.3) + theme_minimal() +
scale_color_brewer(palette = "Set1") + facet_wrap(.~simulator)
As expected, stronger selection leads to reduced diversity and lower TE frequency.
Multifidelity, results.
Below is an example of the multi-fidelity pipeline. We run 10,000 msprime simulations using our selection proxy, then fit on benchmark SLiM simulations. We also recalibrate using 100, 250, or 500 SLiM simulations for training. At this stage, I only implement this for the context model, but it should be straightforward to adapt it to the DFE model.
Code
mkdir -p logs
REP_PER_TASK=10 sbatch --array=0-999%800 slurm/msprime_sproxy_sims.sh ###5Mb windows, times 20
SHARD_DIR=sim_shards_msprime_sproxy/ REP_PER_TASK=10 sbatch --array=0-999%800 slurm/msprime_sproxy_sims.sh ###simulating windows, and parallelize by having blocks of windows with the same scenario running on distinct nodes/cpus.
SHARD_DIR=sim_shards_msprime_sproxy OUT=simulations_msprime/dataset sbatch slurm/msprime_sproxy_collate.sh ###collation of windows simulated in parallel for the same set of parameters.
MSPRIME_DATASET=simulations_msprime/dataset SLIM_DATASET=results_test_12headsattention_simple_with_recomb_bins_corr_coeff_384dim_plus_quantiles27pc75pc_on_tefeatures/dataset \
BUDGETS="100 250 500 1000 2500" TE_DFE_MODEL="context" MODEL="context" AGGREGATOR=attention HEADS=12 OUT_ROOT=mfnpe bash ./run_mfnpe.sh
###here we specify where the SLiM simulations are, and how many we use to recalibrate the msprime simulations.We can examine how the crude msprime-based model fares when benchmarked against more realistic SLiM pseudo-observations. Note that we use the same parameters as for the previous context + sumstats MLP (section (e)). Firstly, let’s examine how the model trained with msprime simulations fares on a benchmark of 200 benchmark simulations (each corresponding to 20 x 5Mb windows, as before):
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures_MFNPE_only_msprime_training///benchmark_NTE10000/" ##this one is based on 10 000 TEs, sumstats, context, extrapolation, 6 attention heads.
if (file.exists("combined_5.csv") )
{combined=read.csv("combined_5.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# recompute 80% CIs from the raw posterior draws, just to show how it can be done. posterior_samples.npz contains everything needed to examine joint posterios if needed.
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
combined <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
combined$is_log_param <- NULL
###getting rid of duration fixed.
combined=subset(combined,combined$param_name != "log10_duration")
write.csv(combined,file="combined_5.csv")
}
###RMSE
rmse_tbl_msprime <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl_msprime)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.716 0.302 0.479 200
2 gamma_r 0.277 0.0183 0.161 200
3 log10_lambda 0.0804 -0.0189 0.0594 200
4 log10_multiplier 0.601 -0.0574 0.454 200
5 log10_s_global 0.381 -0.301 0.322 200
6 log10_t_start 0.493 0.103 0.386 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
I=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (only msprime)")
J=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (only msprime)")
I
Code
J
Now, we recalibrate by retraining with 100 SLiM simulations, using a lower (5 x lower) learning rate:
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures_MFNPE_finetune_n100///benchmark_NTE10000/" ##this one is based on 10 000 TEs, sumstats, context, extrapolation, 6 attention heads.
if (file.exists("combined_6.csv") )
{combined=read.csv("combined_6.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# recompute 80% CIs from the raw posterior draws, just to show how it can be done. posterior_samples.npz contains everything needed to examine joint posterios if needed.
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
combined <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
combined$is_log_param <- NULL
###getting rid of duration fixed.
combined=subset(combined,combined$param_name != "log10_duration")
write.csv(combined,file="combined_6.csv")}
###RMSE
rmse_tbl_100 <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl_100)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.613 0.0260 0.418 200
2 gamma_r 0.268 0.00625 0.154 200
3 log10_lambda 0.0880 0.00111 0.0658 200
4 log10_multiplier 0.601 -0.0752 0.450 200
5 log10_s_global 0.176 -0.0203 0.122 200
6 log10_t_start 0.486 0.00743 0.370 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
K=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (100 SLiM)")
L=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (100 SLiM)")
K
Code
L
250 SLiM simulations :
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures_MFNPE_finetune_n250//benchmark_NTE10000/" ##this one is based on 10 000 TEs, sumstats, context, extrapolation, 6 attention heads.
if (file.exists("combined_7.csv") )
{combined=read.csv("combined_7.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# recompute 80% CIs from the raw posterior draws, just to show how it can be done. posterior_samples.npz contains everything needed to examine joint posterios if needed.
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
combined <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
combined$is_log_param <- NULL
###getting rid of duration fixed.
combined=subset(combined,combined$param_name != "log10_duration")
write.csv(combined,file="combined_7.csv")}
###RMSE
rmse_tbl_250 <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl_250)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.586 0.0350 0.385 200
2 gamma_r 0.252 0.0196 0.144 200
3 log10_lambda 0.0760 0.00374 0.0543 200
4 log10_multiplier 0.585 -0.0701 0.430 200
5 log10_s_global 0.163 -0.000340 0.113 200
6 log10_t_start 0.480 0.0550 0.357 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
M=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (250 SLiM)")
N=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (250 SLiM)")
M
Code
N
500 SLiM simulations :
Code
library(reticulate)
library(dplyr)
library(ggplot2)
library(patchwork)
base <- "Results_final_blogpost/context_model_sumstats_12heads_attention_strict_384dim_plus_quantiles27pc75pc_on_tefeatures_MFNPE_finetune_n500///benchmark_NTE10000/" ##this one is based on 10 000 TEs, sumstats, context, extrapolation, 6 attention heads.
if (file.exists("combined_8.csv") )
{combined=read.csv("combined_8.csv")} else {
parameters <- read.csv(file.path(base, "recovery_values.csv"), h = TRUE)
sim_properties <- read.csv(file.path(base, "sim_summary.csv"), h = TRUE)
df <- merge(parameters, sim_properties, by = "sim_idx")
# recompute 80% CIs from the raw posterior draws, just to show how it can be done. posterior_samples.npz contains everything needed to examine joint posterios if needed.
np <- import("numpy")
d <- np$load(file.path(base, "posterior_samples.npz"), allow_pickle = TRUE)
samples <- d$f[["samples"]] # (n_sims, n_post, n_params)
param_names <- as.character(d$f[["param_names"]])
dims <- dim(samples); n_sims <- dims[1]; n_par <- dims[3]
# HPD 80% -> long format
ci80_list <- vector("list", n_sims * n_par)
k <- 1
for (i in seq_len(n_sims)) {
for (j in seq_len(n_par)) {
#qs <- quantile(samples[i, , j], c(0.25, 0.75), names = FALSE) #for quantiles
qs <- HDInterval::hdi(samples[i, , j], credMass = 0.80) #for Highest posterior density
ci80_list[[k]] <- data.frame(
sim_idx = i - 1L, # npz is 0-indexed like the CSV
param_name = param_names[j],
hpd_10pct = qs[1],
hpd_90pct = qs[2]
)
k <- k + 1
}
}
ci80 <- do.call(rbind, ci80_list)
# IMPORTANT: match sim_idx convention. If recovery_values.csv sim_idx does NOT
# start at 0, drop the "- 1L" above. Check: range(parameters$sim_idx).
df <- merge(df, ci80, by = c("sim_idx", "param_name"))
# ---- log10 conversion (now also covering the new 50% columns) ----
combined <- df %>%
mutate(
is_log_param = grepl("^log_", param_name),
true_value = ifelse(is_log_param, true_value / log(10), true_value),
posterior_mean = ifelse(is_log_param, posterior_mean / log(10), posterior_mean),
posterior_median = ifelse(is_log_param, posterior_median / log(10), posterior_median),
ci_5pct = ifelse(is_log_param, ci_5pct / log(10), ci_5pct),
ci_95pct = ifelse(is_log_param, ci_95pct / log(10), ci_95pct),
hpd_10pct = ifelse(is_log_param, hpd_10pct / log(10), hpd_10pct),
hpd_90pct = ifelse(is_log_param, hpd_90pct / log(10), hpd_90pct),
param_name = ifelse(is_log_param, sub("^log_", "log10_", param_name), param_name)
)
combined$is_log_param <- NULL
###getting rid of duration fixed.
combined=subset(combined,combined$param_name != "log10_duration")
write.csv(combined,file="combined_8.csv")}
###RMSE
rmse_tbl_500 <- combined %>%
filter(is.finite(true_value), is.finite(posterior_mean)) %>%
group_by(param_name) %>%
summarise(
rmse = sqrt(mean((posterior_mean - true_value)^2)),
bias = mean(posterior_mean - true_value),
mae = mean(abs(posterior_mean - true_value)),
n = n(),
.groups = "drop"
) %>%
arrange(param_name)
print(rmse_tbl_500)# A tibble: 6 × 5
param_name rmse bias mae n
<chr> <dbl> <dbl> <dbl> <int>
1 alpha_r 0.553 0.0406 0.359 200
2 gamma_r 0.240 0.0149 0.139 200
3 log10_lambda 0.0781 -0.000935 0.0538 200
4 log10_multiplier 0.564 -0.0420 0.412 200
5 log10_s_global 0.164 0.0156 0.113 200
6 log10_t_start 0.481 0.0428 0.358 200Code
##To check behaviour under high and low burst regimes. Selection harder to correctly infer if strong bursts (expected).
strong_burst=subset(combined,combined$sim_idx %in% subset(combined,(combined$param_name=="log10_multiplier" & combined$true_value>=2) )$sim_idx)
weak_burst=subset(combined,combined$sim_idx %in% subset(combined, (combined$param_name=="log10_multiplier" & combined$true_value < 2))$sim_idx)
O=ggplot(strong_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none") + ggtitle("Strong burst (500 SLiM)")
P=ggplot(weak_burst, aes(true_value, posterior_mean, color = param_name)) +
geom_point(size = 3) +
geom_errorbar(aes(ymin = ci_5pct, ymax = ci_95pct), width = 0.1, alpha = 0.3) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme( legend.position = "none")+ ggtitle("Weak burst (500 SLiM)")
#O + P + plot_annotation(tag_levels = 'A')
O
Code
P
This is already fairly decent. We systematically underestimate the strength of selection, and this seems to be due to the lack of drop in diversity near TEs under selection. Note that we still do better than ABC based on RMSE metrics.
Code
all_rmse=rbind(abc_metrics_log10[,1:2],rmse_tbl_sumstat_context[,1:2],rmse_tbl_msprime[,1:2],rmse_tbl_100[,1:2],rmse_tbl_250[,1:2],rmse_tbl_500[,1:2])
all_rmse$method=c(rep("ABC",6),rep("5000 SLiM, MLP",6),rep("0 SLiM + msprime + MLP + MF",6),rep("100 SLiM, MLP + MF",6),rep("250 SLiM, MLP + MF",6),rep("500 SLiM, MLP + MF",6))
ggplot(all_rmse, aes(method, rmse, color = method)) +
geom_point(size = 3) +
facet_wrap(~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme(axis.text.x = element_blank()) + ggtitle("RMSE for different methods")
An extremely fast solution based on diffusion theory: fitting age-dependent SFS with time-varying \(\theta\) on bins of TEs age.
A promising approach to estimate selection on TEs consists in using age-of-allele tests of neutrality (Horvath et al. 2022; Blumenstiel et al. 2014): by contrasting the expected TE frequency given its age with observation, we can derive the selective coefficient. As I developed the approximate msprime simulator, I realized that we already have the theoretical framework to fit allele frequency spectra for distinct age bins depending on selection coefficients. Since most of the available methods dealing with diffusion are implemented in moments/\(\partial\)a\(\partial\)i (Jouganous et al. 2017; Gutenkunst et al. 2009), I checked the API. I found something interesting about the function moments.Demes.SFS (see here: https://momentsld.github.io/moments/_modules/moments/Demes/Demes.html#SFS ).
This function provides the option to output SFS for distinct age bins (mutation_time_windows option). This means that, coupled with a likelihood optimizer, we should be able to infer the different parameters through likelihood. The idea is not new. For example, Christelle Fraïsse developed an implementation of \(\partial\)a\(\partial\)i that simulates a mixture of SFS (one with no migration, the other with migration for example) to model heterogeneity in linked selection or migration rates along the genome (Christe et al. 2016). The same idea was later ported into DILS (Fraïsse et al. 2021). Here, I reuse the same idea, but pool over different age bins.
An advantage of this method is that it does not rely on slow simulations. It only requires the age of TEs (obtained for example through GEVA, as was done recently in Brachypodium or maize (Liu et al. 2025) ), and their frequencies. The fit is extremely fast (we already have the demography fixed, which helps). We can fit a DFE model with the same parameters as in the deep learning implementation. Here I fit using 13 bins of age, and estimate uncertainties by refitting bootstrapped TEs. There is an option to implement block-bootstrap, and one can also use Godambe approximation to estimate confidence intervals, all typical options in \(\partial\)a\(\partial\)i /moments.
Code
###We can control the bounds with EXTRA options. We also estimate uncertainties by doing a block-bootstrap over the 20x5Mb genomic windows simulated with SLiM. Note that we are using the datasets used for training. We also use a dataset with no background selection here. We use 13 age edges here. In general, more edges/slices seems to be beneficial for the algorithm, but requires enough TEs in each. We also fix the duration of the burst at 10 generations for the time being.
EXTRA="--bound log_mean_s=-12,-0.69 --bound log_gamma_shape=-2.3,2.3 --uncerts-bootstrap-refit 100" AGE_JITTER=0.3 TREES_ROOT=/shared/projects/datepalm/test_DFE_deeplearning/June_versions/te_dfe_merged/simulations_wo_BGS/trees AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 BOOT_BLOCK_COL=window FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 sbatch --array=0-200%200 slurm/dfe_aged_sweep.sh
###This is a collection script to gather all results in a single file.
python collect_fits.py dfe_aged_fits -o dfe_aged_fits/summary_final.txtI try to make the most of the diagnostic plots available in moments, so we can for example examine TEs age, their SFS for each age bin, and the fit moments applies. Below is an example for an average selection coefficient of approx. -0.02, and a recent burst 50 generations ago multiplying by a factor of 10x the baseline insertion rate.
Let’s now see the results if we use the true age of TEs to build our SFS slices (inferred directly from the tree sequences generated during the SLiM simulations):
Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
Attaching package: 'purrr'The following object is masked from 'package:data.table':
transposeThe following object is masked from 'package:locfit':
noneCode
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fits_true_ages/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
Q <- mk(filter(plt, burst == "strong"), "Strong burst")
R <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (Q / R) + plot_annotation(tag_levels = "A")
Q
Code
R
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.180 -0.0147 201
2 log_gamma_shape 1.68 0.334 201
3 log_mean_s 2.35 -0.361 201
4 log_gamma_shape 1.68 0.334 201
5 log10_t_start 0.512 -0.0933 201
6 log_multiplier 1.25 0.0636 201This is fairly decent, although it looks like some simulations are difficult to fit, hitting the bounds in our search ranges. I may have to play further with bins of TE ages, optimizers and search bounds.
Now, we can introduce some noise in the allele age estimation. This will give us an idea on whether the algorithm will struggle in real conditions. I approximate this by jittering the allele age using a lognormal approximation lognormal(0, jitter_factor). Below is the shift for a jitter of 0.1:
Code
library(ggplot2)
sigma <- 0.1
n_te <- 2000
max_age <- 500
set.seed(1)
true_age <- runif(n_te, 0, max_age)
jittered <- true_age * rlnorm(n_te, meanlog = 0, sdlog = sigma)
band <- data.frame(true_age = c(1, max_age))
band$lo <- band$true_age * exp(-sigma)
band$hi <- band$true_age * exp( sigma)
ggplot(data.frame(true_age, jittered), aes(true_age, jittered)) +
geom_ribbon(data = band, aes(x = true_age, ymin = lo, ymax = hi),
inherit.aes = FALSE, fill = "steelblue", alpha = 0.15) +
geom_point(alpha = 0.25, size = 0.9, color = "steelblue") +
geom_abline(slope = 1, intercept = 0, color = "red", linetype = "dashed") +
scale_x_log10(limits = c(1, max_age)) +
scale_y_log10() +
annotation_logticks(sides = "bl") +
coord_equal() +
labs(x = "true age (generations)", y = "jittered age (generations)") +
theme_minimal()Warning: Removed 7 rows containing missing values or values outside the scale range
(`geom_point()`).
And for a jitter factor of 0.3:
Code
sigma <- 0.3
n_te <- 2000
max_age <- 500
set.seed(1)
true_age <- runif(n_te, 0, max_age)
jittered <- true_age * rlnorm(n_te, meanlog = 0, sdlog = sigma)
band <- data.frame(true_age = c(1, max_age))
band$lo <- band$true_age * exp(-sigma)
band$hi <- band$true_age * exp( sigma)
ggplot(data.frame(true_age, jittered), aes(true_age, jittered)) +
geom_ribbon(data = band, aes(x = true_age, ymin = lo, ymax = hi),
inherit.aes = FALSE, fill = "steelblue", alpha = 0.15) +
geom_point(alpha = 0.25, size = 0.9, color = "steelblue") +
geom_abline(slope = 1, intercept = 0, color = "red", linetype = "dashed") +
scale_x_log10(limits = c(1, max_age)) +
scale_y_log10() +
annotation_logticks(sides = "bl") +
coord_equal() +
labs(x = "true age (generations)", y = "jittered age (generations)") +
theme_minimal()Warning: Removed 7 rows containing missing values or values outside the scale range
(`geom_point()`).
However, this is probably not something entirely realistic. For example, GEVA tends to perform better for intermediate allele ages but not so much for very recent or very old alleles. I suspect that there are two levers that could be actioned here to determine whether the method is biased by this with biological data: use several age estimators (tsinfer, Relate, ARGWeaver are possible options), and use different time bins to determine which parameters seem robustly inferred.
Below are the results for a jitter of 0.1:
Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fits_jit0.1//summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
S <- mk(filter(plt, burst == "strong"), "Strong burst")
T <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (S/T) + plot_annotation(tag_levels = "A")
S
Code
T
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, #na.rm = TRUE)),
# bias = mean(fit_value - true_value, na.rm = #TRUE),
# n = sum(is.finite(fit_value) & is.finite#(true_value)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.187 0.00748 201
2 log_gamma_shape 1.74 0.573 201
3 log_mean_s 2.21 -0.371 201
4 log_gamma_shape 1.74 0.573 201
5 log10_t_start 0.501 -0.118 201
6 log_multiplier 1.27 -0.0368 201And now with a jitter of 0.3:
Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fits_jit0.3/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
U <- mk(filter(plt, burst == "strong"), "Strong burst")
V <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (U / V) + plot_annotation(tag_levels = "A")
U
Code
V
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.640 0.288 201
2 log_gamma_shape 1.81 0.800 201
3 log_mean_s 2.10 -0.0172 201
4 log_gamma_shape 1.81 0.800 201
5 log10_t_start 0.512 -0.0964 201
6 log_multiplier 1.39 -0.467 201As suspected, the quality of the inference drops when we introduce noise in TEs age. The effect is especially pronounced in the case of a large burst. However, parameters linked to TE activity remain rather robustly well estimated.
Although this method is very fast (hundreds of fits are done in a few minutes for this flat demography), I also need to determine its robustness to changes in population size. It is also probably wise to check it against (slow-ish) SLiM simulations under the assumed demography. I also include a diffusion simulator (again, just reusing the moments pipeline) in the script, which may be a faster way to address robustness.
Injecting an estimated error rate in the model
We may be able to improve the quality of inference above by explicitly adding a jitter error parameter to the model, and estimating it from the data. The idea consists in moving SNPs from their own age bin to others with a probability depending on the error rate. This is a fairly limited addition to the model and may help improving our inference.
Code
TREES_ROOT=/shared/projects/datepalm/test_DFE_deeplearning/June_versions/te_dfe_merged/simulations_wo_BGS/trees \
AGE_JITTER=0.3 AGE_ERROR_SIGMA=0.3 FIT_AGE_ERROR=1 \
AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 \
FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 BOOTSTRAP=100 BOOT_BLOCK_COL=window \
EXTRA="--bound log_mean_s=-9.21,-0.69 --bound log_gamma_shape=-2.3,2.3 --n-s-bins 80" OUT_DIR=dfe_aged_fit_with_error_fit \
sbatch --array=0-200%200 slurm/dfe_aged_sweep.sh
###This is a collection script to gather all results in a single file.
python collect_fits.py dfe_aged_fit_with_error_fit -o dfe_aged_fit_with_error_fit/summary_final.txtWe can inspect whether we manage to recover the error rate and the parameter values:
Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fit_with_error_fit/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
AA <- mk(filter(plt, burst == "strong"), "Strong burst")
AB <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (AA/ AB) + plot_annotation(tag_levels = "A")
AA
Code
AB
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse1 <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse1)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.171 -0.0189 201
2 log_gamma_shape 1.48 0.211 201
3 log_mean_s 1.71 -0.0440 201
4 log_gamma_shape 1.48 0.211 201
5 log10_t_start 0.551 -0.131 201
6 log_multiplier 1.24 0.0474 201Code
raw_rmse1$correction="inferred"
raw_rmse1$contrast="demographic model"We simulated with an error factor of 0.3, do we actually manage to infer it correctly?
Code
# grep age_error_sigma dfe_aged_fit_with_error_fit/*fit* | awk '{print $3}' > dfe_aged_fit_with_error_fit/error_rates_estimated.txt
error_rates_est=read.table("Results_final_blogpost/dfe_aged_fit_with_error_fit//error_rates_estimated.txt")
hist(error_rates_est$V1,breaks=100,xlab="Error rate (estimated)",main="Estimate of the jitter error rate with a starting value of 0.3")
abline(v=0.3,lty=2,col="red")
We provide the true value as a starting parameter, but this sort of error rate can also be reasonably estimated from simulations, or retrieved from the original articles describing the methods, which can help the algorithm.
There are many cases (15%) here where the algorithm does not identify any error, so it may be better to fix the error rate if already known. However, it seems also likely that the algorithm uses this new parameter to freely adjust misfits between expectations and observations. It is probably wise to check the results obtained with a fixed and freely estimated error rate.
Let’s examine what happens if we provide the jitter parameter as a fixed parameter in the model:
Code
TREES_ROOT=/shared/projects/datepalm/test_DFE_deeplearning/June_versions/te_dfe_merged/simulations_wo_BGS/trees \
AGE_JITTER=0.3 AGE_ERROR_SIGMA=0.3 FIT_AGE_ERROR=0 \
AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 \
FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 BOOTSTRAP=100 BOOT_BLOCK_COL=window \
EXTRA="--bound log_mean_s=-9.21,-0.69 --bound log_gamma_shape=-2.3,2.3" OUT_DIR=dfe_aged_fit_without_error_fit \
sbatch --array=0-200%200 slurm/dfe_aged_v0.2.sh
###Note the FIT_AGE_ERROR=0
###This is a collection script to gather all results in a single file.
python collect_fits.py dfe_aged_fit_without_error_fit -o dfe_aged_fit_without_error_fit/summary_final.txt
##The name of the folder is misleading, there is no fit but the error rate is provided and fixed in the model.Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fit_without_error_fit/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
AA <- mk(filter(plt, burst == "strong"), "Strong burst")
AB <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (AA/ AB) + plot_annotation(tag_levels = "A")
AA
Code
AB
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse2 <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse2)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.194 -0.0447 201
2 log_gamma_shape 1.47 -0.0117 201
3 log_mean_s 1.86 -0.0987 201
4 log_gamma_shape 1.47 -0.0117 201
5 log10_t_start 0.563 -0.178 201
6 log_multiplier 1.19 0.106 201Code
raw_rmse2$correction="fixed"
raw_rmse2$contrast="demographic model"Calibrate using random SNPs to fit errors and deviations between expected and observed neutral spectra
Most of the methods inferring the DFE actually compare a set of selected variants against a set of putatively neutral variants. I do not think that it is straightforward to compare age-binned empirical spectra to derive 2\(N_e\)\(s\), but one can use the SNPs to estimate the mismatch between the theoretical expectation from moments based on demography, and propagate the error into the inference. This may further help improving the fit. We can use a set of calibrating SNPs along with their age, obtained along with the TEs. It may also help correcting for demographic misspecification, making the method more robust: even if the demography is very different from what is simulated, we can partly correct for it. Note that moments needs the age bins defined in units of 2 \(N_0\) generations (diploid \(N_0\)). This could be defined before any fit simply by estimating 4\(N_0\)\(\mu\) from the SNP data, or the user can keep using the demographic trajectory in the MSMC format that is used by all the tools in the current pipeline, knowing that the harmonic mean will be used to define bins of age and obtain the demographic units for selection and insertion rates.
Code
TREES_ROOT=/shared/projects/datepalm/test_DFE_deeplearning/June_versions/te_dfe_merged/simulations_wo_BGS/trees \
AGE_JITTER=0.3 AGE_ERROR_SIGMA=0.3 FIT_AGE_ERROR=1 \
NEUTRAL=1 NEUTRAL_CALIBRATION=full MAX_SNPS=20000 \
AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 \
FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 BOOTSTRAP=100 BOOT_BLOCK_COL=window \
EXTRA="--bound log_mean_s=-9.21,-0.69 --bound log_gamma_shape=-2.3,2.3 --n-s-bins 80" OUT_DIR=dfe_aged_fit_SNPcal_with_error_fit \
sbatch --array=0-200%200 slurm/dfe_aged_sweep.sh
###This is a collection script to gather all results in a single file.
python collect_fits.py dfe_aged_fit_SNPcal_with_error_fit -o dfe_aged_fit_SNPcal_with_error_fit/summary_final.txtNow, we can inspect the fit:
Code
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fit_SNPcal_with_error_fit/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
AC <- mk(filter(plt, burst == "strong"), "Strong burst")
AD <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (AC/ AD) + plot_annotation(tag_levels = "A")
AC
Code
AD
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse3 <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse3)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.154 0.00378 201
2 log_gamma_shape 1.56 0.467 201
3 log_mean_s 1.57 0.0687 201
4 log_gamma_shape 1.56 0.467 201
5 log10_t_start 0.521 0.00853 201
6 log_multiplier 1.40 0.192 201Code
raw_rmse3$correction="inferred"
raw_rmse3$contrast="SNPs"Now, we fix the error rate on the age estimate, assuming it is known:
Code
TREES_ROOT=/shared/projects/datepalm/test_DFE_deeplearning/June_versions/te_dfe_merged/simulations_wo_BGS/trees \
AGE_JITTER=0.3 AGE_ERROR_SIGMA=0.3 FIT_AGE_ERROR=0 \
NEUTRAL=1 NEUTRAL_CALIBRATION=full MAX_SNPS=20000 \
AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 \
FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 BOOTSTRAP=100 BOOT_BLOCK_COL=window \
EXTRA="--bound log_mean_s=-9.21,-0.69 --bound log_gamma_shape=-2.3,2.3" OUT_DIR=dfe_aged_fit_SNPcal_without_error_fit \
sbatch --array=0-200%200 slurm/dfe_aged_v0.2.sh
###This is a collection script to gather all results in a single file.
python collect_fits.py dfe_aged_fit_SNPcal_without_error_fit -o dfe_aged_fit_SNPcal_without_error_fit/summary_final.txtCode
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fit_SNPcal_without_error_fit/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
# return an all-NA column of the right length when a column is absent
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
n <- nrow(d); d$sim_idx <- seq_len(n)
# class edges MUST match collect_fits.py --s-class-edges (default 1e-3,1e-2)
S_EDGES <- c(0, 1e-3, 1e-2, 100)
CLASSES <- c("p_nn", "p_weak", "p_strong")
props <- function(log_shape, log_mean_s) {
shape <- exp(log_shape); rate <- exp(log_shape) / exp(log_mean_s)
cdf <- sapply(S_EDGES, function(s) pgamma(s, shape = shape, rate = rate))
m <- matrix(cdf, nrow = length(shape))
setNames(as.data.frame(t(apply(m, 1, diff))), CLASSES)
}
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
rows <- list()
# (1) proportions: prefer the collector's point columns (p_*_fit / p_*_true) if
# present, else compute them in R. CIs come from the collector (percentile of the
# refit draws) via ci_lo_p_* / ci_hi_p_*.
for (which in c("fit", "true")) {
if (all(paste0(CLASSES, "_", which) %in% names(d))) {
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- num(d[[paste0(cl, "_", which)]])
} else {
pr <- props(num(d[[paste0(which, "_log_gamma_shape")]]),
num(d[[paste0(which, "_log_mean_s")]]))
for (cl in CLASSES) d[[paste0(cl, "_", which, "_calc")]] <- pr[[cl]]
}
}
for (cl in CLASSES) {
rows[[length(rows) + 1]] <- tibble(
sim_idx = d$sim_idx, param_name = cl,
true_value = d[[paste0(cl, "_true_calc")]],
fit_value = d[[paste0(cl, "_fit_calc")]],
ci_lo = col(paste0("ci_lo_", cl), n),
ci_hi = col(paste0("ci_hi_", cl), n),
reliable = d$reliable)
}
# (2) direct parameters
add_param <- function(name, fit_col, true_col, cilo_col, cihi_col, to_log10) {
tv <- col(true_col, n); fv <- col(fit_col, n)
lo <- col(cilo_col %||% "", n); hi <- col(cihi_col %||% "", n)
if (to_log10) { f <- log(10); tv<-tv/f; fv<-fv/f; lo<-lo/f; hi<-hi/f }
rows[[length(rows) + 1]] <<- tibble(
sim_idx = d$sim_idx, param_name = name,
true_value = tv, fit_value = fv, ci_lo = lo, ci_hi = hi, reliable = d$reliable)
}
`%||%` <- function(a, b) if (is.null(a)) b else a
# log_lambda: profiled out, so no CI; still plot the point estimate.
add_param("log10_lambda", "fit_log_lambda", "true_log_lambda", NULL, NULL, TRUE)
# log10_mean_s CI: mean = shape*scale is monotone in scale at fixed shape, so the
# scale-endpoint propagation is valid here (unlike the proportions).
have_scale_ci <- has("ci_lo_log_scale") && has("ci_hi_log_scale")
if (has("fit_log_mean_s")) {
lm_lo <- lm_hi <- rep(NA_real_, n)
if (have_scale_ci) {
shp <- num(d$fit_log_gamma_shape)
lm_lo <- shp + num(d$ci_lo_log_scale); lm_hi <- shp + num(d$ci_hi_log_scale)
}
d$._mslo <- lm_lo; d$._mshi <- lm_hi
add_param("log10_mean_s", "fit_log_mean_s", "true_log_mean_s", "._mslo", "._mshi", TRUE)
}
add_param("log10_t_start", "fit_log10_t_start", "true_log10_t_start",
"ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE)
add_param("log10_multiplier", "fit_log_multiplier", "true_log_multiplier",
"ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE)
plt <- bind_rows(rows)
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name,
levels = c(CLASSES, "log10_lambda", "log10_mean_s", "log10_t_start", "log10_multiplier"))
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
AE <- mk(filter(plt, burst == "strong"), "Strong burst")
AF <- mk(filter(plt, burst == "weak"), "Weak burst")
p <- (AE/ AF) + plot_annotation(tag_levels = "A")
AE
Code
AF
Code
#rmse_tbl <- plt %>%
# group_by(param_name) %>%
# summarise(
# rmse = sqrt(mean((fit_value - true_value)^2, na.rm = #TRUE)),
# bias = mean(fit_value - true_value, na.rm = TRUE),
# n = sum(is.finite(fit_value) & is.finite(true_value#)),
# .groups = "drop")
#print(rmse_tbl)
raw_params <- c("log_lambda","log_gamma_shape", "log_mean_s", "log_gamma_shape","log10_t_start","log_multiplier") # whichever you want
raw_rmse4 <- purrr::map_dfr(raw_params, function(p) {
tv <- col(paste0("true_", p), n)
fv <- col(paste0("fit_", p), n)
tibble(param_name = p,
rmse = sqrt(mean((fv - tv)^2, na.rm = TRUE)),
bias = mean(fv - tv, na.rm = TRUE),
n = sum(is.finite(fv) & is.finite(tv)))
})
print(raw_rmse4)# A tibble: 6 × 4
param_name rmse bias n
<chr> <dbl> <dbl> <int>
1 log_lambda 0.151 -0.0353 201
2 log_gamma_shape 1.47 0.283 201
3 log_mean_s 1.48 0.199 201
4 log_gamma_shape 1.47 0.283 201
5 log10_t_start 0.544 -0.148 201
6 log_multiplier 1.23 0.201 201Code
raw_rmse4$correction="fixed"
raw_rmse4$contrast="SNPs"This is a bit better.
We can perform a rough check of how inference changes depending on whether we provide a demographic history or SNPs to calibrate, and whether we fix or estimate an error rate:
Code
error_rates=rbind(raw_rmse1,raw_rmse2,raw_rmse3,raw_rmse4)
ggplot(error_rates, aes(paste(contrast,correction), rmse, color = paste(contrast,correction))) +
geom_point(size = 3) +
facet_wrap( ~ param_name, scales = "free", ncol = 2) +
theme_minimal() + theme(axis.text.x = element_blank()) + ggtitle("RMSE for different methods")
So for this pipeline, I would recommend providing a fixed parameter for the age-estimate error rate, and use SNPs as a neutral contrast for the SFS instead of the theoretical one. However this may change once we incorporate unstable demography to our tests. It seems useful to carefully compare consistency between inference with SNPs v. demographic scenario inferred from SNPs.
Expansion to the context model
For the context model, we further discretize our SFS into discrete bins of age x recombination, and fit a single s coefficient in each recombination bin, the value of s varying with \(\alpha_r\), while the rate at which TEs are injected in the recombination bin depends on the burst parameters, but also \(\gamma_r\). This is probably where we reach the limits of the model in terms of resolution and the number of bins that can be obtained will strongly depend on the number of TEs. Here is a test using five recombination bins and 13 bins of age. We also provide the error rate as a fixed parameter, which is set at a high level (jitter of 0.3 lognormal).
Code
MODEL_KIND=context RECOMB=auto RECOMB_BINS=5 NEUTRAL=1 \
TREES_ROOT=simulations_CNN_context/trees AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 AGE_JITTER=0.3 AGE_ERROR_SIGMA=0.3 FIT_AGE_ERROR=0 \
FIX="log_duration=2.3026" EXTRA="--n-s-bins 80" N_STARTS=100 N_ROUNDS=10 OUT_DIR=dfe_aged_fit_context_SNPcal_with_error_fit \
sbatch --array=0-200%200 slurm/dfe_aged_sweep.sh
python collect_fits.py dfe_aged_fit_context_SNPcal_with_error_fit/ -o dfe_aged_fit_context_SNPcal_with_error_fit/summary_final.txtCode
library(ggplot2); library(dplyr); library(tidyr); library(purrr);library(patchwork)
#table=read.csv("Results_final_blogpost/test_moments.csv")
d <- read.delim("Results_final_blogpost/dfe_aged_fit_context_SNPcal_without_error_fit/summary_final.txt", sep = "\t", header = TRUE,
na.strings = c("NA", "", "nan", "NaN"),
stringsAsFactors = FALSE, check.names = TRUE)
num <- function(x) suppressWarnings(as.numeric(x))
has <- function(nm) nm %in% names(d)
col <- function(nm, n) if (nm %in% names(d)) num(d[[nm]]) else rep(NA_real_, n)
# keep context rows only, if the table mixes both models
if (has("model")) d <- d[is.na(d$model) | d$model == "context", ]
n <- nrow(d); d$sim_idx <- seq_len(n)
pinned <- if (has("pinned")) d$pinned else rep(NA, n)
hnpd <- if (has("hessian_not_pd")) num(d$hessian_not_pd) else rep(0, n)
d$reliable <- (is.na(pinned) | pinned == "") & (is.na(hnpd) | hnpd == 0)
# name, fit col, true col, CI cols, convert natural-log -> log10?
specs <- tibble::tribble(
~param, ~fit, ~true, ~lo, ~hi, ~to_log10,
"log10_lambda", "fit_log_lambda", "true_log_lambda", NA, NA, TRUE,
"log10_s_global", "fit_log_s_global", "true_log_s_global", "ci_lo_log_s_global", "ci_hi_log_s_global", TRUE,
"alpha_r", "fit_alpha_r", "true_alpha_r", "ci_lo_alpha_r", "ci_hi_alpha_r", FALSE,
"gamma_r", "fit_gamma_r", "true_gamma_r", "ci_lo_gamma_r", "ci_hi_gamma_r", FALSE,
"log10_t_start", "fit_log10_t_start", "true_log10_t_start", "ci_lo_log10_t_start", "ci_hi_log10_t_start", FALSE,
"log10_multiplier", "fit_log_multiplier", "true_log_multiplier", "ci_lo_log_multiplier", "ci_hi_log_multiplier", TRUE
)
plt <- pmap_dfr(specs, function(param, fit, true, lo, hi, to_log10) {
tv <- col(true, n); fv <- col(fit, n)
l <- if (is.na(lo)) rep(NA_real_, n) else col(lo, n)
h <- if (is.na(hi)) rep(NA_real_, n) else col(hi, n)
if (to_log10) { f <- log(10); tv <- tv/f; fv <- fv/f; l <- l/f; h <- h/f }
tibble(sim_idx = d$sim_idx, param_name = param,
true_value = tv, fit_value = fv, ci_lo = l, ci_hi = h,
reliable = d$reliable)
})
burst <- plt %>% filter(param_name == "log10_multiplier") %>%
transmute(sim_idx, burst = ifelse(true_value >= 2, "strong", "weak"))
plt <- left_join(plt, burst, by = "sim_idx") %>%
filter(is.finite(true_value), is.finite(fit_value))
plt$param_name <- factor(plt$param_name, levels = specs$param)
mk <- function(dd, ttl)
ggplot(dd, aes(true_value, fit_value)) +
geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red") +
geom_errorbar(aes(ymin = ci_lo, ymax = ci_hi, color = param_name),
width = 0, alpha = 0.3, na.rm = TRUE) +
geom_point(aes(color = param_name), size = 2.2) +
facet_wrap(~ param_name, scales = "free", ncol = 3) +
labs(x = "true", y = "fitted") + guides(color = "none") +
theme_minimal() + theme(legend.position = "bottom") + ggtitle(ttl)
BA <- mk(filter(plt, burst == "strong"), "Strong burst (context model)")
BB <- mk(filter(plt, burst == "weak"), "Weak burst (context model)")
BA
Code
BB
Code
# RMSE / bias per parameter, overall and by burst class
rmse_tbl <- plt %>% group_by(param_name) %>%
summarise(rmse = sqrt(mean((fit_value - true_value)^2, na.rm = TRUE)),
bias = mean(fit_value - true_value, na.rm = TRUE),
n = n(), .groups = "drop")
print(rmse_tbl)# A tibble: 6 × 4
param_name rmse bias n
<fct> <dbl> <dbl> <int>
1 log10_lambda 0.177 -0.143 201
2 log10_s_global 0.250 -0.0118 201
3 alpha_r 0.819 0.253 201
4 gamma_r 0.347 0.121 201
5 log10_t_start 0.656 -0.237 201
6 log10_multiplier 0.679 0.0179 201Code
print(plt %>% group_by(param_name, burst) %>%
summarise(rmse = sqrt(mean((fit_value - true_value)^2, na.rm = TRUE)),
bias = mean(fit_value - true_value, na.rm = TRUE),
n = n(), .groups = "drop"))# A tibble: 12 × 5
param_name burst rmse bias n
<fct> <chr> <dbl> <dbl> <int>
1 log10_lambda strong 0.189 -0.144 79
2 log10_lambda weak 0.169 -0.143 122
3 log10_s_global strong 0.360 -0.0381 79
4 log10_s_global weak 0.139 0.00514 122
5 alpha_r strong 1.10 0.555 79
6 alpha_r weak 0.572 0.0570 122
7 gamma_r strong 0.436 0.185 79
8 gamma_r weak 0.274 0.0796 122
9 log10_t_start strong 0.496 -0.0635 79
10 log10_t_start weak 0.742 -0.349 122
11 log10_multiplier strong 0.652 -0.0514 79
12 log10_multiplier weak 0.697 0.0628 122Overall, this is not as good as the deep learning pipeline to infer \(\alpha_r\), but it is otherwise doing rather well. \(\lambda\) is systematically underestimated, and I am not sure why, but this is likely due to the fact that this quantity is the rate expected in regions with a RECOMB_TYPICAL recombination rate, not the weighed average of the influx across the whole genome. I will keep investigating.
A general observation across all these tests is that we generally infer burst and insertion rate variation with a rather high precision, while selection and its variance are the hardest to infer. This is perhaps not surprising, given that selection acts mostly on the shape of the spectrum while insertion parameters are more directly linked to the number of segregating TE insertions. Given the speed of the method, and the fact that it also requires fewer SLiM simulations to check its behaviour, I think it should definitely be used as a first analysis before starting with more time-consuming simulations. It also produces plots showing the age distribution of TEs and the spectrum for each age bin, which may already provide a few ideas to the user, if they wish to refine the model or better define their priors in simulation-based inference.
First conclusions and ongoing work
All methods show decent performance at this stage, and I am optimistic that we can already address a few interesting questions about TE evolutionary dynamics with them. I will leave the CNN option in the scripts for now, but it does not appear to be the most efficient. The summary statistics option is much faster, and has the advantage of pushing the user to inspect the results from simulations. It also provides a possible ABC-path which may be useful if one can afford enough simulations. I am particularly interested in the multi-fidelity and sequential NPE framework. It is essentially domain adaptation/transfer learning, so it can be easily modified to recalibrate an initial model. It provides a way to make these pipelines more efficient.
The diffusion-based method relies on moments and shows a rather good performance, and I would strongly advise using it if the user can estimate allele age (and the error rate that comes with this estimate) with some reliability. This will probably require rather large cohorts, for statistical phasing and age estimate. A possible difficulty will lie in the definition of age bins: a lower number means a loss in resolution over burst parameters, and simulations will probably still be required to assess whether estimates can be trusted (maybe through some sort of parametric bootstrap estimate).
I intend to include an option in the deep learning framework to use allele age as another summary statistic, which I strongly suspect would make the method more precise. This would require incorporating an error parameter in the simulations too, to make them more realistic.
Possible extensions
- Length-dependent selection. Allow the selection coefficient to differ between short and long elements. In the meantime this can be approximated manually: inferring parameters for short, near-neutral LINEs, then fitting a model with priors on activity or insertion preference narrowed for long, counter-selected LINEs.
- Element life-cycle dynamics. Model conversion from full-length LTR elements to solo LTRs, and possibly DNA-transposon excision. At the moment, the models are likely applicable to DNA transposons, assuming their rate of removal/excision remains sufficiently low to consider the current hypothesis of no back-mutation acceptable. Note that after transposition, the transposon copy on the second chromosome, if present, can refill the empty space through recombination-driven DNA repair, so excision rate is higher than actual allele loss.
- Using the flanking SNPs. At present the random flanking SNPs are largely unused. I do not show it here, but I do not get better estimates when including a random subset of SNPs alongside TEs in the CNN models. However, they could calibrate the analysis and help identify the demographic history when the user does not supply one, and may also help account for background selection.
- Gene conversion, which can matter under rescaling, though it is not available for many systems.
Already implemented but untested are VCF conversion to a format that can be used to run inference on actual data, masking of genomic regions with poor quality, unphased inference, unpolarized SNPs for CNN, and mimicking patterns of missingness in real data.
There is also a 06_snpe_refine.pyscript that could be used to perform SNPE inference, but has not been benchmarked yet.
On positive selection
Although not incorporated here, an extension to positive selection could contribute to the debate over the role of TEs in adaptation. In particular, it could test the idea that positively selected TEs arise through bursts of transposition: if there are more positively selected elements than expected around a burst, relative to before or after it, that may be evidence that TEs contribute to host adaptation in the manner McClintock anticipated. Adding this parameter and its dynamics to the SLiM model and the inference pipeline should be straightforward. I suspect the moments-based pipeline may also be adequate at identifing any excess of positively selected, high frequency alleles for any given age-bin. But more complex models are also dangerous since they add parameters that the algorithm can use to fit noise in the data.
On scaling with SLiM
Rescaling is one reason I want to avoid simulating long regions in future benchmarks (Marsh, Kaushik, and Johri 2026). In Drosophila, with \(N_e\) on the order of two million diploids, one might rescale down to a factor of 100 (so simulations have a \(N_e\) of 20,000), but doing so cleanly means simulating shorter windows to avoid spurious correlations and the accumulated effects of background selection, if the user chooses to include many deleterious variants, including, or in addition to, TEs. This may matter less when only TEs are simulated as deleterious, given their lower density relative to functional SNPs, but realistic simulation is precisely the point of using SLiM.
Annex: installation and a few HPC commands
A few commands
Given how heavy simulations can be, I assume that the user will work on a cluster. I myself am working on the Institut Français de Bioinformatique cluster (IFB). Here is a possible way to install the pipeline, using uv:
Code
# Install uv:
# Tell the installer where to place the binary if you don't have rights:
INSTALL_PATH="/local/install/path/" ##replace local/install/path by your own personal path where you can install uv
export UV_INSTALL_DIR="${INSTALL_PATH}/uv-bin"
#Disable the installer's prompt to modify your shell rc files
export UV_NO_MODIFY_PATH=1
export XDG_CONFIG_HOME="${INSTALL_PATH}/.config"
# Run the installer
curl -LsSf https://astral.sh/uv/install.sh | sh
# /shared/projects/datepalm/test_DFE_deeplearning/activate.sh
cat > ${INSTALL_PATH}/activate.sh << 'EOF'
# Project environment for te_dfe_burst on the HPC cluster
# Source this from your shell or from SLURM scripts:
# source /shared/projects/datepalm/test_DFE_deeplearning/activate.sh
export UV_PROJECT_DIR="/shared/projects/datepalm/test_DFE_deeplearning"
export PATH="${UV_PROJECT_DIR}/uv-bin:$PATH"
export UV_CACHE_DIR="${UV_PROJECT_DIR}/.uv_cache"
export UV_PYTHON_INSTALL_DIR="${UV_PROJECT_DIR}/.uv_python"
# Redirect uv's config away from ~/.config (which is quota-limited)
export XDG_CONFIG_HOME="${UV_PROJECT_DIR}/.config"
mkdir -p "${XDG_CONFIG_HOME}"
EOF
source ${INSTALL_PATH}/activate.sh
uv venv --python 3.11
source .venv/bin/activate
uv pip install msprime pyslim tskit torch sbi joblib scipy matplotlib numpy
uv pip install moments-popgen ##latest addition
#Download popgen-npe:
cd ${INSTALL_PATH}
git clone https://github.com/kr-colab/popgen-npe.git
##then
module load slim/4.3 ##I have SLiM preinstalled on my cluster. You can also install it locally, in which case you will need to add it to your PATH. The scripts work with version 4.3 at the moment, but should be compatible with version 5.xx, the main change being in the SLiM scripts where sim.chromosome would become sim.chromosomes with an S at the end.
source ${INSTALL_PATH}/.venv/bin/activate
source ${INSTALL_PATH}/activate.sh
export PYTHONPATH="${INSTALL_PATH}/popgen-npe:$PYTHONPATH"Here is how we can run the procedure step by step for the DFE model with summary statistics. For the context model, once can switch the environment variables TE_DFE_MODEL="simple" MODEL="simple" to TE_DFE_MODEL="context" MODEL="context". Obviously, a clean version of the pipeline should use automated scripts, for example with snakemake. Since this project is still in development, I keep working sequentially.
Code
source ${INSTALL_PATH}/.venv/bin/activate
source ${INSTALL_PATH}/activate.sh
export PYTHONPATH="${INSTALL_PATH}/popgen-npe:$PYTHONPATH"
###Simulate under the DFE (simple) model, keep the tree sequences. The number and length of genomic windows per simulation is controlled in config_base.py. In case you do not have enough tmp space, you can specify a path to a tmp directory (TE_DFE_TMPDIR)
TE_DFE_TMPDIR=/shared/projects/datepalm/tmp_scratch TE_DFE_MODEL="simple" MODEL="simple" KEEP_TREES=1 N_MIN=20 N_MAX=1000 SEEDS_PER_TASK=10 MAX_CONCURRENT=999 SIM_DIR=/shared/projects/datepalm/simulations_trees_sumstats_simple sbatch slurm/01_simulate.sh
###If you want to reprocess simulations stored as tree sequences (for example if we add new summary statistics and need to update), this can be done also:
TREES_DIR=/shared/projects/datepalm/simulations_trees_sumstats_simple/trees OUT_DIR=/shared/projects/datepalm/simulations_trees_sumstats_simple_renewed/ SEEDS_PER_TASK=5 \
sbatch --array=0-995%990 slurm/01b_reprocess.sh
###Here we group all simulations into a single file that will be stored in results_sumstat_simple/ This allows faster access of the training algorithm to the training sims.
TE_DFE_MODEL="simple" MODEL="simple" KEEP_TREES=1 N_MIN=20 N_MAX=1000 SEEDS_PER_TASK=10 MAX_CONCURRENT=999 SIM_DIR=/shared/projects/datepalm/simulations_trees_sumstats_simple COLLATED_TRAINING_FOLDER="results_sumstat_simple" sbatch slurm/02_collate.sh
TE_DFE_MODEL="simple" MODEL="simple" KEEP_TREES=1 N_MIN=20 N_MAX=1000 SEEDS_PER_TASK=10 MAX_CONCURRENT=999 AGGREGATOR=attention COLLATED_TRAINING_FOLDER="results_sumstat_simple" RESULTS_DIR="results_sumstat_simple" N_HEADS=12 BATCH_SIZE=32 sbatch slurm/03_train.sh
###Model 200 pseudo-observed benchmark simulations, each made of 200 windows of 5Mb (1Gb in total).
MODEL="simple" TE_DFE_MODEL="simple" N_SCENARIOS=200 N_WINDOWS=200 CHUNK_SIZE=50 SHARD_DIR=shards_simple sbatch slurm/05a_benchmark.sh
###We parallelized before by splitting the same scenario across distinct SLURM jobs. Now we gather all windows from the same scenario together.
MODEL="simple" TE_DFE_MODEL="simple" N_SCENARIOS=200 N_WINDOWS=200 CHUNK_SIZE=50 SHARD_DIR=shards_simple TEST_DIR=benchmark_simple sbatch slurm/05a_collate.sh
###Fitting the trained model on benchmark
MODEL="simple" TE_DFE_MODEL="simple" N_SCENARIOS=200 TEST_DIR=benchmark_simple RESULTS_DIR=results_sumstat_simple/ sbatch slurm/05d_sweep_benchmark_sumstats.sh
###Benchmarking results and the trained model files are found in results_CNN_simple/
###reserve 05c_sweep_benchmark.sh for the matrix representation. For example:
MODEL="simple" TE_DFE_MODEL="simple" N_SCENARIOS=200 RESULTS_DIR=results_CNN_simple/ TEST_DIR=benchmark_CNN_simple/ sbatch slurm/05c_sweep_benchmark.shThere are a few helper functions that can help checking the content of any given simulation/trained model:
Code
##checking whether a sim is under the expected model (DFE, context) and representation (sumstats, CNN...)
python inspect_dataset.py simulations/sim_000001.npz
##check the properties of the model: execute this in python
import pickle; print(pickle.load(open("results_CNN_simple/training_meta.pkl","rb")))We can also train using msprime simulations, then test whether adding an increasing number of SLiM simulations improves our algorithm. At the moment this is only implemented for the context model:
Code
###For msprime we also parallelize heavily, like for benchmarking sims.
SHARD_DIR=sim_shards_msprime_sproxy/ REP_PER_TASK=10 sbatch --array=0-999%800 slurm/msprime_sproxy_sims.sh
SHARD_DIR=sim_shards_msprime_sproxy OUT=simulations_msprime/dataset sbatch slurm/msprime_sproxy_collate.sh
###BUDGETS specifies how many SLiM simulations (stored in this example in results_sumstat_context/) should be used to retrain the pure msprime model.
MSPRIME_DATASET=simulations_msprime/dataset SLIM_DATASET=results_sumstat_context/dataset \
BUDGETS="100 250 500 1000 2500" TE_DFE_MODEL="context" MODEL="context" AGGREGATOR=attention HEADS=12 OUT_ROOT=mfnpe bash ./run_mfnpe.shWe also have wrapper scripts to fit the DFE model with the moments implementation.
Code
###In this example we reextract TE frequencies and their age from the tree sequence files stored in /shared/projects/datepalm/simulations_wo_BGS/trees, alter the actual age with a jitter of 0.3, and use 13 bins of age, whose edges are entered as 0,10,25,50,75,100,150,200,250,300,350,400,500. We estimate uncertainties using block bootstrap over our simulated genomic windows and refit the model to each bootstrapped dataset. We try 100 independent combination of starting parameters, refine the best likelihood estimates starting from the previous estimate with the highest likelihood, over 5 rounds. We fix the duration of the burst at exp(2.3026) = 10 generations.
EXTRA="--bound log_mean_s=-12,-0.69 --bound log_gamma_shape=-2.3,2.3 --uncerts-bootstrap-refit 100" AGE_JITTER=0.3 TREES_ROOT=/shared/projects/datepalm/simulations_wo_BGS/trees AGE_EDGES=0,10,25,50,75,100,150,200,250,300,350,400,500 BOOT_BLOCK_COL=window FIX="log_duration=2.3026" N_STARTS=100 N_ROUNDS=5 sbatch --array=0-200%200 slurm/dfe_aged_sweep.sh
###This script will parse the individual reports for each of the 200 training simulations and produce a summary file we can examine in R for benchmarking:
python collect_fits.py dfe_aged_fits -o dfe_aged_fits/summary_final.txtSummary statistics are stored for each benchmark dataset, with TEs being subsampled at the benchmark value. These datasets can be compared to training data as illustrated above. We can obtain these summary statistics when using the sumstats representation, like this:
Code
module load slim/4.3
source ${INSTALL_PATH}/.venv/bin/activate
source ${INSTALL_PATH}/activate.sh
export PYTHONPATH="${INSTALL_PATH}/popgen-npe:$PYTHONPATH"
###These four calls are needed if you used uv to install the required programs.
###statistics for DFE training model
python 07_export_sumstat_summaries.py --model simple --dataset results_sumstat_simple/dataset --out summaries_simple.csv
###statistics for context training simulations
python 07_export_sumstat_summaries.py --model context --dataset results_sumstat_context/dataset --out summaries_context.csv
###statistics for benchmarking data (pseudo observed data for ABC benchmarking)
python 07b_collate_benchmark.py --in-dir benchmark_simple/ --out benchmark_simple/observations
python 07_export_sumstat_summaries.py --dataset benchmark_simple/observations --out observed_summaries.csvOf course, this is still very dense and complex, I am working on making the whole pipeline more accessible, but hopefully this should already give an idea of the procedure.
Priors and parameters
Most of the parameters are stored in the config_base.py script under the te_dfe/ folder. This is still at the draft stage, and many options may disappear in the future. There are also a few remnants from previous trials and dead code. It needs to be cleaned.
Code
"""
config_base.py — configuration shared by BOTH the simple and context models.
This holds everything that does not depend on which model is selected: input
paths, the demographic/simulation setup, the per-window genotype-matrix layout,
the base continuous channels, the SLiM rescaling factor, the realism options
(background selection, coverage mask, dead-pixel missingness, SFS folding,
phasing), training-stage feature filtering, and the network/training defaults.
The model overlays (config_simple.py, config_context.py) import everything from
here with `from .config_base import *` and then add or override the
model-specific pieces (parameter set, prior ranges, continuous-channel count).
config.py dispatches to one overlay based on the TE_DFE_MODEL environment
variable (set by each pipeline script from its --model flag; CLI overrides the
default). See config.py for the dispatch.
The realism knobs live here because they apply identically to both models — a
single source of truth so the two never drift.
"""
import numpy as np
# ---------------------------------------------------------------------------
# External input files
# ---------------------------------------------------------------------------
MSMC_PATH = "data/msmc_output.txt" ###we input the demography using the MSMC format, which is the coalescence rate over time. See here for an explanation and how it relates to Ne: https://github.com/stschiff/msmc-tools/blob/master/msmc-tutorial/guide.md
MU = 1.25e-7 # per-bp per-gen mutation rate; must match MSMC --mu
RECOMB_MAP_PATH = "data/recomb_map.txt"
ANNOTATION_PATH = "data/genes.gtf"
ANNOTATION_FMT = "gtf"
CHROMOSOMES = None
# ---------------------------------------------------------------------------
# Simulation
# ---------------------------------------------------------------------------
SIM_LENGTH = 5000
NEUTRAL_MU = MU
N_INDIVIDUALS = 50
N_HAPS = N_INDIVIDUALS * 2
WINDOW_SIZE = 5_000_000
MAX_INSERTIONS_PER_WINDOW = 5_000
# Burst timing prior (shared by both models).
LOG10_T_START_RANGE = (1.5, 3.5)
LOG_DURATION_RANGE = (np.log(3.0), np.log(100.0))
LOG_MULTIPLIER_RANGE = (np.log(1.0), np.log(1000.0))
# ---------------------------------------------------------------------------
# SLiM rescaling factor Q. Q=1 is exact
# ---------------------------------------------------------------------------
SLIM_Q = 1.0 ###THIS DOES NOT WORK, ONE NEEDS TO RESCALE THE INPUT FILES/PARAMETERS MANUALLY AT THIS STAGE.
# ---------------------------------------------------------------------------
# Phasing mode (shared). PHASED=True -> one row per haplotype, 4 state channels.
# PHASED=False -> collapse to per-individual diploid dosage, 6 state channels.
# Switching requires retraining (input shape changes).
# ---------------------------------------------------------------------------
PHASED = True
N_INPUT_ROWS = N_HAPS if PHASED else N_INDIVIDUALS
N_STATE_CHANNELS = 4 if PHASED else 6
# ---------------------------------------------------------------------------
# Individual subsampling (robustness testing: simulate a large panel once, then
# process/train at a smaller sample size without re-simulating).
# ---------------------------------------------------------------------------
# Default OFF. When OFF, the tree's sample count MUST equal N_HAPS or the
# processor raises (catches config/tree mismatches early). When ON, the
# processor draws N_INDIVIDUALS individuals from a tree that has at least that
# many, KEEPING each individual's two haplotype nodes together (via the tskit
# individuals table that SLiM/pyslim populate) — so the unphased pairing is
# correct and any within-individual structure (e.g. selfing) is preserved.
# Selection is the first N_INDIVIDUALS by default (unbiased for unstructured
# sims, no seed needed); set SUBSAMPLE_SEED to an int for a seeded RANDOM draw
# (use this if the sim carries population/spatial structure in sample order).
SUBSAMPLE_INDIVIDUALS = False
SUBSAMPLE_SEED = None
# ---------------------------------------------------------------------------
# SNP polarization (SFS folding). False=polarized (0=anc,1=der),
# True=folded (0=major,1=minor). SNPs only; TEs never folded.
# ---------------------------------------------------------------------------
FOLD_SFS = False
# ---------------------------------------------------------------------------
# Background selection (optional nuisance; not inferred). Uses the GTF to simulate a fraction of deleterious non-synonymous variants.
# ---------------------------------------------------------------------------
BG_SELECTION = True
BG_USE_GTF = True
BG_DEL_SHAPE = 0.3
BG_DEL_MEAN_S = 0.01
BG_DEL_MU = 0.85e-7 ##say about 2/3 of new mutations are deleterious. Approx.
BG_DEL_DOMINANCE = 0.5
BG_MEAN_EXON_LEN = 150
BG_SD_EXON_LEN = 60
BG_MEAN_INTEREXON = 1500
BG_SD_INTEREXON = 1000
# ---------------------------------------------------------------------------
# Coverage mask (optional, BED-driven). MASK_AS_COVARIATE changes N_COVARIATES.
# ---------------------------------------------------------------------------
COVERAGE_MASK_BED = None
MASK_AS_COVARIATE = False
DROP_FULLY_MASKED_WINDOWS = False
MIN_CALLABLE_FRACTION = 0.5
# ---------------------------------------------------------------------------
# Dead-pixel masking (per-individual missingness; training augmentation).
# None | scalar | list[len N_INDIVIDUALS] | path. For real-data inference: None. NOT TESTED YET. USEFUL FOR MODELING MISSING DATA
# ---------------------------------------------------------------------------
MISSINGNESS = None
# ---------------------------------------------------------------------------
# Training-stage feature filtering (train-without, for fair comparison).
# Mirrors the 05b ablation vocabulary; applied via te_dfe/filtering.py so train
# and benchmark masking are identical. This is a way to test what the model actually uses during training.
# ---------------------------------------------------------------------------
TRAIN_ZERO_CHANNELS = ()
TRAIN_ZERO_REGIONS = ()
TRAIN_ZERO_STATES = ()
# ---------------------------------------------------------------------------
# Covariates: [gene_density, recomb_scaled, te_density, snp_density] (+masked
# fraction when MASK_AS_COVARIATE). Window-level summary covariates.
# ---------------------------------------------------------------------------
N_COVARIATES = 5 if MASK_AS_COVARIATE else 4
RECOMB_SCALE = 1e-7
# ---------------------------------------------------------------------------
# Per-window genotype-matrix layout (three regions on the site axis). Only for CNN representation
# ---------------------------------------------------------------------------
# THREE PRIMITIVES drive the whole layout; every other size is DERIVED from
# them with the same formulas the processor uses, so config and processor can
# never disagree (a past source of drift):
# N_TE_SITES number of TE anchors = max TEs represented per window
# N_RANDOM_BLOCKS number of random background blocks
# N_FLANK_PER_SIDE flanking SNPs taken on EACH side of a TE (and of each
# random-block anchor); 12/side -> 24 flank slots per TE.
N_TE_SITES = 60 # TE anchors (max TEs/window; >this -> subsample)
N_RANDOM_BLOCKS = 60
N_FLANK_PER_SIDE = 12 # 12 flanking SNPs on EACH side (24 SNPs total + 1 TE/block)
# Derived (do not hand-edit): a TE's flank region is 2*N_FLANK_PER_SIDE slots,
# a random block is its anchor + 2*N_FLANK_PER_SIDE flanks.
BLOCK_SIZE = 1 + 2 * N_FLANK_PER_SIDE # random-block width
N_TE_FLANK_SITES = N_TE_SITES * 2 * N_FLANK_PER_SIDE # = anchors * 24
N_RANDOM_BLOCK_SITES = N_RANDOM_BLOCKS * BLOCK_SIZE
N_SITES = N_TE_SITES + N_TE_FLANK_SITES + N_RANDOM_BLOCK_SITES
MAF_MIN = 0.0
SNP_SAMPLING_STRATEGY = "random"
FLANK_WINDOW_BP = 15_000
MAX_FLANK_DISTANCE_BP = 15_000 #vestigial TODO: remove
MIN_TE_FLANK_DIST_BP = 50
MIN_RANDOM_ANCHOR_DIST_BP = 15_000
# State codes in the genotype matrix.
ANCESTRAL_STATE = 0
SNP_STATE = 1
TE_STATE = 2
PADDING_STATE = 3
# ---------------------------------------------------------------------------
# Base continuous channels (per site): phys pos, gen pos, TE density, SNP
# density. Models extend this: the simple model adds a recomb channel; the
# context model adds dist-to-gene + log-recomb. N_CONTINUOUS_CHANNELS is set in
# the overlay, since it differs by model.
# ---------------------------------------------------------------------------
N_POSITION_CHANNELS = 2
N_DENSITY_CHANNELS = 2
CM_NORMALIZER = 0.5
LOCAL_DENSITY_BP = 250_000
DENSITY_LOG_NORMALIZER = 50.0
# This option is to avoid edge effects. It drops TEs within this many bp of a window edge from the FOCAL (emitted) set,
# where the local-density estimate is noisiest even after length-normalization.
# Excluded TEs STILL count as neighbours for interior TEs' density, so this only
# trims which TEs are analysed, not the counts. 0 = keep all TEs (rely on the
# length-normalized rescale alone). Set to LOCAL_DENSITY_BP for the fully
# edge-clean set (drops ~2*EDGE_EXCLUDE_BP/window_size of TEs, e.g. ~10% at 5 Mb).
EDGE_EXCLUDE_BP = 250_000 #this is the safest way to avoid edge effects. to edit if TE density is too low, and to adjust depending on LOCAL_DENSITY_BP and window_size
#TO CHECK: DOES THIS ALSO AFFECT THE NUMBER OF TEs OUTPUT FOR SPECTRA-BASED APPROACHES (MOMENTS). COULD ARTIFICIALLY REDUCE LAMBDA ESTIMATE.
# Recomb-rate channel normalization (shared formula; used by simple ch4 and
# context ch5). log10(r / RECOMB_TYPICAL) / RECOMB_LOG_SCALE, clipped [-2,2].
RECOMB_TYPICAL = 1e-7
RECOMB_LOG_SCALE = 1.0
RECOMB_MIN_FOR_LOG = 1e-12
# Distance-to-gene channel normalization (context model).
DIST_TO_GENE_NORMALIZER = 1e6
# ---------------------------------------------------------------------------
# Network / training defaults (shared).
# ---------------------------------------------------------------------------
DROPOUT = 0.1
N_TRAIN_WINDOWS = 20
# ---------------------------------------------------------------------------
# Inferred-parameter subset (marginalize-out unwanted outputs).
# ---------------------------------------------------------------------------
# By default the network infers ALL of the model's parameters (PARAM_NAMES in
# the overlay). To infer only a SUBSET, set INFER_PARAMS to a list of the names
# you want as outputs; the rest are MARGINALIZED — they still vary in the
# simulations (the sims and existing datasets are untouched), but they are
# dropped from the theta vector the flow sees, so the network learns the
# posterior over the kept parameters with the others integrated out.
#
# INFER_PARAMS = None -> infer all (default)
# INFER_PARAMS = ["log_multiplier"] -> infer only the burst multiplier
# INFER_PARAMS = ["log10_t_start", "log_duration", "log_multiplier"] -> burst only
#
# Names must be a subset of the active model's PARAM_NAMES. The kept order
# follows PARAM_NAMES (not the order you list), so the column mapping is
# canonical and stable. training_meta records the kept names + indices so
# inference/benchmarking slice the same columns; a checkpoint trained with a
# given INFER_PARAMS must be used with the same subset.
INFER_PARAMS = None ##Honestly not extremely useful, remnant of previous trials.
# ---------------------------------------------------------------------------
# Fixed parameters (remove a degree of freedom from the generative model).
# ---------------------------------------------------------------------------
# Unlike INFER_PARAMS (which marginalizes — the parameter still varies in the
# sims, just isn't output), FIX_PARAMS HOLDS a parameter at a constant in the
# simulations, so it is removed from the generative model entirely. This is the
# lever that can genuinely CONCENTRATE the remaining parameters' posteriors when
# they are confounded with the fixed one — at the cost of asserting you know its
# value. A fixed parameter is automatically excluded from the inferred outputs
# (you cannot infer a constant).
# Values are given in the parameter's NATIVE (transformed) units — the same
# units as the theta vector and the *_RANGE bounds. To fix burst duration at
# 30 generations, set log_duration to log(30):
#
# import numpy as np
# FIX_PARAMS = {"log_duration": np.log(30.0)} # duration fixed at 30 gen
# FIX_PARAMS = {"log10_t_start": 2.5} # t_start fixed at 10^2.5
#
# Names must be in the active model's PARAM_NAMES and must not also appear in
# INFER_PARAMS. training_meta records the fix so inference/benchmarking use the
# identical generative model.
FIX_PARAMS = {"log_duration": np.log(10.0)}
# ---------------------------------------------------------------------------
# Observation representation: how a window/locus is turned into network input.
# ---------------------------------------------------------------------------
# "matrix" : the per-window genotype matrix (state one-hot + continuous
# channels) consumed by the exchangeable CNN. Default.
# "sumstat" : a per-TE summary-statistic vector (te_dfe/sumstats.py) consumed
# by a per-TE MLP + the same permutation-invariant aggregator. The
# set of per-TE vectors generalizes to any number of TEs (train on
# N, deploy on more) the same way window K-extrapolation works.
# The two representations share the simulator, prior, flow, training loop, and
# the FIX/INFER machinery; they differ only in the processor output and the
# embedding net.
REPRESENTATION = "sumstat"
# Per-TE summary-statistic vector length (sumstat representation). Derived from
# te_dfe.sumstats so the two stay in lockstep.
def _n_sumstats_per_te():
try:
from .sumstats import N_SUMSTATS_PER_TE
return int(N_SUMSTATS_PER_TE)
except Exception:
return 14
N_SUMSTATS_PER_TE = _n_sumstats_per_te()
# Max TEs per locus used at TRAINING (the set is variable-size; this caps the
# padded tensor width, analogous to N_TRAIN_WINDOWS). Deployment can exceed it
# via the aggregator's cardinality extrapolation.
N_TE_TRAIN_MAX = 10000 #this one should be set to the maximum number of TEs that will be used (algorithm masks empty entries for training and benchmarking). If you benchmark using 10,000 TEs but set this number at 1000 during training, rise it to 10,000 before benchmarking. I need to edit this so it does not fail at benchmarking. Or keep everything at 10,000 (easier, just wasteful since it will pad most of the training simulations)
# ---------------------------------------------------------------------------
# ADDED (OFF by default): recombination-decile O/E enrichment of TEs. THIS IS NOT TESTED YET (06/08/2026), AND PROBABLY NOT NECESSARY.
# Availability-normalized replacement for the old 3-bin TE recomb proportions.
# O = TE proportion per recomb decile (numerator, where TEs ARE);
# E = genome-bp proportion per recomb decile (exposure, from the recomb map);
# feature = log((O+eps)/(E+eps)) per decile -> a per-sim GLOBAL covariate fed on
# the existing global_cov rail (te_dfe/enrichment.py). Supplies the empty-region
# exposure the per-TE view discards -> targets high gamma_r.
#
# Edges are FIXED genome-wide deciles (computed ONCE from the recomb map, shared
# by every sim and the observed data); the DENOMINATOR is each sim's own realized
# windows (so window-sampling variation divides out). See enrichment.py.
ENRICH_OE_DECILES = False # master switch. When True, wire the O/E vector
# into global_cov (see the 4 integration blocks).
OE_N_BINS = 10 # number of recombination deciles
OE_PSEUDOCOUNT = 1e-3 # eps in log((O+eps)/(E+eps)); keeps thin bins finite
OE_EXPOSURE_COMPANION = False # also append log10 bp/bin so the trunk can
# down-weight thin deciles (adds OE_N_BINS more
# global_cov dims)
# The decile edges are computed at startup, once, e.g.:
# from te_dfe.enrichment import genome_decile_edges
# RECOMB_DECILE_EDGES = genome_decile_edges(map_positions, map_rates, OE_N_BINS)
# and cached; they are data-derived (not a static constant) so they live wherever
# the recomb map is first loaded (WindowSampler / processor init).
# --- Global sequence-length covariate (sumstat representation) -------------
# When ON, the total sequence length the TEs were pooled from is fed to the
# network as a global covariate, so abundance is usable in a GENOME-SIZE-
# INVARIANT way (N per unit sequence = density) rather than via raw set size.
# REQUIRES training sims with a VARIABLE window count (so the covariate varies
# and is decorrelated from lambda); with a constant window count it is a
# constant and teaches nothing. Default OFF -> no behavior change. THIS IS NOT TESTED YET (06/08/2026), AND PROBABLY NOT NECESSARY SINCE WE ALSO USE LOCAL TE DENSITY.
SUMSTAT_USE_SEQLEN_COV = False
N_SUMSTAT_GLOBAL_COV = 1 if SUMSTAT_USE_SEQLEN_COV else 0
# The covariate value fed to the net is (log10(total_bp) - REF) / SCALE, kept
# ~O(1). REF ~ log10(20 windows * 5 Mb = 1e8) = 8; SCALE spans the range of
# genome/window-count sizes seen in training.
SEQLEN_LOG_REF = 8.0
SEQLEN_LOG_SCALE = 1.0
def seqlen_covariate(total_bp):
"""Normalized global covariate from a total sequence length in bp."""
import numpy as _np
total_bp = max(float(total_bp), 1.0)
return (_np.log10(total_bp) - SEQLEN_LOG_REF) / SEQLEN_LOG_SCALE
# Max flanking SNPs per TE used for the per-TE statistics (upper bound; fewer is
# fine when missing data / low diversity reduce the count). Matches the matrix
# flank layout for consistency.
SUMSTAT_FLANK_MAX = N_TE_FLANK_SITES // 2
# Number of TEs sampled per replicate for the sumstat set during training
# (deploy/extrapolate to more at inference) If I get it right, max number of TEs that are sampled at inference/benchmarking.
SUMSTAT_N_TE_TRAIN = 10000 ##this one is dead
# DFE selection-class thresholds (|2 Ne_h s|): nearly-neutral / weak / strong.
SEL_THRESHOLD_NN = 1.0
SEL_THRESHOLD_STRONG = 10.0We can modify the parameters that are specific to the DFE model here:
Code
"""
config_simple.py — simple-model overlay (6-parameter gamma-DFE + burst).
Inherits all shared infrastructure and realism options from config_base, then
adds the simple model's parameter set, prior ranges, and continuous-channel
layout. Selected when TE_DFE_MODEL=simple (the default).
Inferred parameter vector (6):
0 log_lambda baseline per-bp per-gen insertion rate (log)
1 log_gamma_shape shape of the gamma DFE over |s| (log)
2 log_mean_s mean |s| of the gamma DFE (log)
3 log10_t_start burst start, generations ago (log10)
4 log_duration burst duration in generations (log)
5 log_multiplier burst insertion-rate multiplier, >=1x (log)
The raw gamma (shape, mean_s) is weakly identified; report the post-hoc
DFE-bin proportions (p_nn / p_w / p_strong via the gamma CDF), which are much
better identified. See prior.py.
"""
from .config_base import * # noqa: F401,F403
import numpy as np
MODEL_NAME = "simple"
# --- Prior ranges -----------------------------------------------------------
LOG_LAMBDA_RANGE = (np.log(5e-10), np.log(2e-9))
LOG_GAMMA_SHAPE_RANGE = (np.log(0.1), np.log(2.0))
LOG_MEAN_S_RANGE = (np.log(1e-4), np.log(0.1))
# (burst ranges inherited from base: LOG10_T_START_RANGE etc.)
# --- Parameter bookkeeping --------------------------------------------------
N_PARAMS = 6
PARAM_NAMES = [
"log_lambda", "log_gamma_shape", "log_mean_s",
"log10_t_start", "log_duration", "log_multiplier",
]
PARAM_LABELS = [
"log lambda (baseline insertion rate)",
"log gamma shape", "log mean |s|",
"log10 t_start (gen ago)", "log duration (gen)", "log multiplier",
]
# --- Continuous channels ----------------------------------------------------
# Base 4 (phys, gen, TE-dens, SNP-dens) plus an optional local-recomb channel
# (ch4). Enabling it changes the input width and invalidates checkpoints.
USE_RECOMB_CHANNEL = True
N_RECOMB_CHANNELS = 1 if USE_RECOMB_CHANNEL else 0
N_CONTINUOUS_CHANNELS = (N_POSITION_CHANNELS + N_DENSITY_CHANNELS
+ N_RECOMB_CHANNELS)
N_CNN_INPUT_CHANNELS = N_STATE_CHANNELS + N_CONTINUOUS_CHANNELSAnd we can modify the parameters for the context-based model here:
Code
"""
config_context.py — context-model overlay (context-dependent DFE + insertion
preference). Selected when TE_DFE_MODEL=context.
Inherits shared infrastructure and realism from config_base, then adds the
context model's parameter set, prior ranges, and the two extra continuous
channels (dist-to-gene, log-recomb) that drive the context regressions.
Generative model (recombination-only default):
s_local(p) = exp(log_s_global + alpha_r * feature_r(p))
lambda_local(p) = lambda_global * exp(gamma_r * feature_r(p))
where feature_r(p) is the normalized log local recombination rate. The
simple-model limit is alpha_* = gamma_* = 0.
CONTEXT_PARAMS selects which context coefficients are FREE (inferred); any
omitted coefficient is pinned to 0 in the simulator and dropped from theta, so
the flow's parameter space stays non-degenerate.
"""
from .config_base import * # noqa: F401,F403
import numpy as np
MODEL_NAME = "context"
# Prior ranges
LOG_LAMBDA_RANGE = (np.log(5e-10), np.log(2e-9))
LOG_S_GLOBAL_RANGE = (np.log(3e-4), np.log(0.1))
ALPHA_D_RANGE = (-2.0, 2.0)
ALPHA_R_RANGE = (-2.0, 2.0)
GAMMA_D_RANGE = (-2.0, 2.0)
GAMMA_R_RANGE = (-2.0, 2.0)
# (burst ranges inherited from base.)
# --- Parameter bookkeeping (configurable context set) -----------------------
# Recombination-only (default): ["alpha_r","gamma_r"] -> 7 params
# Distance-only: ["alpha_d","gamma_d"] -> 7 params
# Full: ["alpha_d","alpha_r","gamma_d","gamma_r"] -> 9
# Constant-DFE (no context): [] -> 5 params
CONTEXT_PARAMS = ["alpha_r", "gamma_r"]
_ALL_CONTEXT_ORDER = ["alpha_d", "alpha_r", "gamma_d", "gamma_r"]
_ACTIVE_CONTEXT = [c for c in _ALL_CONTEXT_ORDER if c in CONTEXT_PARAMS]
_CONTEXT_LABELS = {
"alpha_d": "alpha_d (selection slope on log dist-to-gene)",
"alpha_r": "alpha_r (selection slope on log local recomb)",
"gamma_d": "gamma_d (insertion preference on log dist-to-gene)",
"gamma_r": "gamma_r (insertion preference on log local recomb)",
}
PARAM_NAMES = (
["log_lambda", "log_s_global"]
+ list(_ACTIVE_CONTEXT)
+ ["log10_t_start", "log_duration", "log_multiplier"]
)
PARAM_LABELS = (
["log lambda (baseline insertion rate at zero context)",
"log s_global (baseline |s| at zero context)"]
+ [_CONTEXT_LABELS[c] for c in _ACTIVE_CONTEXT]
+ ["log10 t_start (gen ago)", "log duration (gen)", "log multiplier"]
)
N_PARAMS = len(PARAM_NAMES)
# --- Continuous channels ----------------------------------------------------
# Base 4 + 2 context channels: ch4 = dist-to-gene, ch5 = log-recomb (feature_r).
# These are ALWAYS present in the context model (feature_r drives the context
# regression even when a coefficient is pinned to 0).
N_CONTEXT_CHANNELS = 2
N_CONTINUOUS_CHANNELS = (N_POSITION_CHANNELS + N_DENSITY_CHANNELS
+ N_CONTEXT_CHANNELS)
N_CNN_INPUT_CHANNELS = N_STATE_CHANNELS + N_CONTINUOUS_CHANNELSAnnex: more details about the network.
Rationale for the architecture
A normalizing flow needs a fixed-size input vector. Every simulation yields a different number of TEs. On top of that, the TEs have no meaningful order, so shuffling them must not change the answer. This choice simplifies the way we can deal with variable number of TEs between training and fitting.
Those two constraints determine everything below:
a shared per-TE encoder, applied identically to every TE, so the rows stay comparable;
a pooling step that collapses a variable number of rows into a fixed number of columns, in a way that ignores order.
1. Input for the network (summary statistics)
Each TE contributes one row of 16 numbers. The columns are fixed at simulation time by FEATURE_NAMES, and are the same for every model:
# te_dfe/sumstats.py
FEATURE_NAMES = [
"te_freq", # 0
"he_carrier_mean", "he_carrier_sd", # 1, 2
"he_noncarrier_mean","he_noncarrier_sd", # 3, 4
"tajD_carrier", "tajD_noncarrier", # 5, 6
"hapext_carrier", "hapext_noncarrier", # 7, 8
"dist_to_gene", # 9
"local_recomb", # 10
"local_snp_density", "local_te_density", # 11, 12
"n_flank_snps_norm", # 13
"valid_tajD", "valid_hapext", # 14, 15
]The names of the summary statistics are relatively self-explanatory: we extract the frequency of the focal TE, the heterozygosity in its flanks (whose length is defined in te_dfe/config_base.py), the Tajima’s D for the haplotypes carrying the TE and those that do not, the extension of the haplotype for the TE-carrying haplotypes (considered as the derived allele) and the non-carrier haplotypes, distance to gene, the local recombination rate, the local SNP and TE density (in larger windows of 250kb in the default, also defined in te_dfe/config_base.py), the normalized number of SNPs in the flanks, and two booleans indicating whether there are enough variants to compute Tajima’s and haplotype extension. This is a fairly simple set of statistics, and it should be straightforward and valuable to add more, although performance already seems reasonably accurate with this set.
On disk a simulation is stored padded to a fixed K with a companion mask, so feats is (K, 16) and mask is (K,) with 1 for a real TE and 0 for padding. Every operation downstream is mask-aware: padded rows must never contribute.
2. Embeddings
The encoder is a two-layer MLP applied to one TE at a time. The same weights are used for every TE, which is exactly what makes the rows poolable later.
# te_dfe/network.py — TESumstatNet.__init__
self.encoder = nn.Sequential(
nn.Linear(n_features, per_te_dim), nn.ELU(),
nn.Linear(per_te_dim, per_te_dim), nn.ELU(),
)
self.enc_norm = nn.LayerNorm(per_te_dim) if use_embedding_norm else None
# TESumstatNet.per_te_embeddings
def per_te_embeddings(self, feats):
"""feats: (B, K, n_features) -> (B, K, per_te_dim)."""
emb = self.encoder(feats)
if self.enc_norm is not None:
emb = self.enc_norm(emb)
return embnn.Linear(16, 384) holds 384 separate sets of weights. Each set produces one number from the 16 inputs, so one TE’s 16 statistics become 384 numbers. Those 384 numbers are the embedding, a re-description of that TE in coordinates the network chose during training, rather than ones that were initially designed.
Having 384 outputs may be a lot compared to 16 inputs. The extra width gives the network room to pull apart things the raw features entangle. “Rare” and “rare AND in low recombination” can occupy separate coordinates, which a 16-dimensional representation has no space to do. Note however that I set this number rather high. Halving this value and using 6 attention heads gives similar results.
Knobs: per_te_dim (the 384), use_embedding_norm (the LayerNorm).
3. What is ELU.
ELU is the only nonlinear step in the encoder, and without it the whole thing would collapse. Two stacked linear layers are mathematically just one linear layer, so the encoder could represent nothing but weighted sums of the 16 inputs.
The nonlinearity is what lets it represent an interaction: a TE at frequency 0.02 meaning something different in a recombination desert than in a hotspot. That is the sort of structure alpha_r encodes for example, and we want to catch it.
ELU specifically, rather than ReLU, keeps a small gradient for negative inputs instead of switching off entirely, which may make training better behaved.
4. Pooling
Every pooling operation runs down a column, across TEs, for each of the 384 coordinates independently. Each produces one row of 384 numbers regardless of how many TEs went in.
# te_dfe/network.py — the two pooling helpers
def _masked_mean_per_dim(x, mask):
"""(B, K, D) -> (B, D): plain mean over kept TEs."""
m = (mask > 0).float().unsqueeze(-1)
return (x * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0)
def _masked_quantile_per_dim(x, mask, q):
"""(B, K, D) -> (B, D): per-dimension q-quantile over kept TEs.
Padded TEs are pushed to -inf so they never enter a quantile; the rank
is taken per row from K = mask.sum, so the result is a function of the
kept set only."""
B, W, D = x.shape
keep = mask > 0
xm = x.masked_fill(~keep.unsqueeze(-1), float("-inf"))
xs, _ = torch.sort(xm, dim=1, descending=True)
K = keep.sum(dim=1)
rank = torch.floor((1.0 - q) * (K - 1).clamp(min=0)).long()
idx = rank.view(B, 1, 1).expand(B, 1, D)
return xs.gather(1, idx).squeeze(1)Both are order-free and both work for any K TEs. This is why a network trained on 200 TEs could be run on 10,000.
Knobs: pool_mean, pool_quantiles, pool_proj_dim (an optional projection applied before pooling, to cap the width).
5. Inspecting the distribution of embeddings
Two sets of TEs can share a mean and be completely different. The set that matters here is the second one: two groups of TEs behaving differently, which is what strong alpha_r produces when selection varies with local recombination.
A mean-only pool maps both sets to the same number. Adding quantiles separates them immediately. During testing, it did seem to improve the ability of the model to identify hard parameters linked to selection, such as \(\alpha_r\).
6. Aggregators
Alongside the fixed statistics, the aggregator computes a learned weighted mean. The three kinds differ in how the weights are formed:
# te_dfe/aggregators.py (structure)
# attention : w = softmax(score) -> competitive, sum to 1
# sigmoid_attn : g = sigmoid(score) -> independent, then normalised
# deepsets : no learned weights -> mean + max + std, three rows# te_dfe/network.py — TESumstatNet.aggregate (pooling section)
ctx, weights = self.aggregator(emb, feats, mask) # weighted MEAN over TEs
if self.context_norm is not None:
ctx = self.context_norm(ctx)
if self.pool_quantiles or self.pool_mean:
m = mask if mask is not None else torch.ones_like(emb[:, :, 0])
pe = self.pool_proj(emb) if self.pool_proj is not None else emb
parts = []
if self.pool_mean:
parts.append(_masked_mean_per_dim(pe, m))
parts += [_masked_quantile_per_dim(pe, m, q) for q in self.pool_quantiles]
ctx = torch.cat([ctx] + parts, dim=-1)
Important
Note the interaction: if pool_mean=True, the trunk already receives the plain mean. A sigmoid_attn aggregator whose gates stay near uniform then reproduces that mean and contributes little, whereas softmax differentiates readily and supplies a genuinely different view. This might be why attention slightly outperforms sigmoid_attn in the preliminary runs.
Weighing embeddings
Each attention head has its own softmax and its own slice of the embedding. With n_attention_heads = 3 and per_te_dim = 384, head 1 weights the TEs one way and pools dimensions 1 to 128, head 2 weights them differently and pools 129 to 256, and so on; the results are concatenated back to 384. So heads are parallel specialists, and this is why embed_dim must be divisible by n_heads.
The spread of the weights differs by an order of magnitude. The softmax head above ranges from 0.03 to 0.30, a ratio of 10. The sigmoid gate ranges from 0.13 to 0.20, a ratio of 1.5, because the gates are computed independently and sit near 0.5 until the scores separate strongly. Deepsets is exactly uniform for its mean, and compensates by emitting max and std as well.
Softmax or sigmoid: what differs
The one-line version is that softmax performs a relative comparison with a fixed voting budget, while sigmoid rates each TE independently, and then compares.
softmax: w_i = exp(s_i) / sum_j exp(s_j) every weight depends on every score
sigmoid: g_i = sigmoid(s_i) each TE judged alone
w_i = g_i / sum_j g_j normalised only afterwardsSoftmax hands out a fixed budget of 100 votes, so giving more to one candidate necessarily takes votes from the others and the only thing that matters is who beats whom. Sigmoid sends each candidate to a judge who never sees the others and scores them against an “absolute standard”, converting to percentages only at the end. If every candidate strikes the judge as “pretty good” they all land near the same score and come out nearly equal, even when one is clearly best.
That has a formal signature. Add a constant to every score and softmax is unchanged, because only differences survive normalisation. Do the same to sigmoid and every gate moves. A sigmoid gate therefore carries an absolute meaning, “how relevant is this TE”, while a softmax weight carries only a relative one.
7. Additional information about TE distribution
Three quantities are computed from the raw feature columns, and are not used for the embedding. I experimented with adding them because some statistics can make it difficult for the network to reconstruct from pooled embeddings but remain trivial to compute directly.
# te_dfe/network.py — TESumstatNet.aggregate (extras section)
if self.recomb_tail_quantiles: # upper tail of local_recomb
r = feats[:, :, self.recomb_idx]
tails = torch.stack([_masked_quantile(r, m, q)
for q in self.recomb_tail_quantiles], dim=-1)
ctx = torch.cat([ctx, tails], dim=-1)
if self.recomb_bin_edges: # TE fractions in low/typ/high rec
props = _masked_bin_proportions(r, m, lo, hi) # (B, 3)
ctx = torch.cat([ctx, props], dim=-1)
if self.corr_feature_pairs: # cor(te_freq, local_recomb)
corrs = torch.stack([_masked_correlation(feats[:, :, i], feats[:, :, j], m)
for (i, j) in self.corr_idx_pairs], dim=-1)
ctx = torch.cat([ctx, corrs], dim=-1)The correlation is the most direct: cor(te_freq, local_recomb) within a simulation is a single scalar reading of how TE frequency tracks recombination, which is alpha_r’s definition.
Knobs: recomb_tail_quantiles, recomb_bin_edges, corr_feature_pairs, n_global_cov.
8. Concatenation and injecting to the trunk
All the pooled statistics and extras are concatenated, and the width of that concatenation is trunk_in:
pooled_dim = pool_proj_dim if pool_proj_dim else per_te_dim
n_pool_stats = len(self.pool_quantiles) + (1 if self.pool_mean else 0)
trunk_in = (agg_out + n_global_cov
+ len(self.recomb_tail_quantiles)
+ (3 if self.recomb_bin_edges else 0)
+ len(self.corr_feature_pairs)
+ pooled_dim * n_pool_stats)
self.trunk = nn.Sequential(
nn.Linear(trunk_in, output_dim), nn.ELU(), nn.Dropout(0.1),
nn.Linear(output_dim, output_dim), nn.ELU(),
)The trunk compresses to output_dim numbers, and the resulting vector is what the normalizing flow conditions on when it estimates the posterior over parameters.
9. Knob reference
| knob | where it acts | effect |
|---|---|---|
per_te_dim |
encoder | width of each TE’s embedding |
use_embedding_norm |
encoder, context | LayerNorm on embeddings and pooled context |
aggregator_kind |
aggregator | attention / sigmoid_attn / deepsets |
n_attention_heads |
aggregator | used by attention and sigmoid_attn; deepsets warns |
pool_mean |
pooling | adds the plain unweighted mean |
pool_quantiles |
pooling | adds one row per quantile; the dispersion signal |
pool_proj_dim |
pooling | projects before pooling to cap width |
recomb_tail_quantiles |
trunk extras | tail of local_recomb; tracks gamma_r |
recomb_bin_edges |
trunk extras | TE fractions per recombination class |
corr_feature_pairs |
trunk extras | direct alpha_r readout |
n_global_cov |
trunk extras | e.g. total sequence length |
output_dim |
trunk | width of the vector handed to the flow |
Caution
Neither 03_train.py nor load_core_network passes per_te_dim, the pooling arguments, or the trunk-extra tuples. They come from the defaults in network.py, so editing that file silently changes what every existing checkpoint means. trunk_in is the observable consequence, and inspect_network_arch.py recovers the configuration from a saved checkpoint’s tensor shapes. I need to expose these as configuration parameters once the scripts are cleaner.
References
Baumdicker, Franz, Gertjan Bisschop, Daniel Goldstein, Graham Gower, Aaron P Ragsdale, Georgia Tsambos, Sha Zhu, et al. 2022. “Efficient Ancestry and Mutation Simulation with Msprime 1.0.” Genetics 220 (3): iyab229. https://doi.org/10.1093/genetics/iyab229.
Blumenstiel, Justin P., Xi Chen, Miaomiao He, and Casey M. Bergman. 2014. “An Age-of-Allele Test of Neutrality for Transposable Element Insertions.” Genetics 196 (2): 523–38. https://doi.org/10.1534/genetics.113.158147.
Bourgeois, Yann, and Stephane Boissinot. 2019. “On the Population Dynamics of Junk: A Review on the Population Genomics of Transposable Elements.” Genes 10 (6): 419. https://doi.org/10.3390/genes10060419.
Bourgeois, Yann, Robert Ruggiero, Imtiyaz Hariyani, and Stéphane Boissinot. 2020. “Disentangling the Determinants of Transposable Elements Dynamics in Vertebrate Genomes Using Empirical Evidences and Simulations.” PloS Genetics, 1–32. https://doi.org/10.1101/2020.04.12.037754.
Chan, Jeffrey, Valerio Perrone, Jeffrey P Spence, Paul A Jenkins, Sara Mathieson, and Yun S Song. 2018. “A Likelihood-Free Inference Framework for Population Genetic Data Using Exchangeable Neural Networks.” In Advances in Neural Information Processing Systems, 31:3432–42. Curran Associates, Inc.
Christe, Camille, Kai N. Stolting, Margot Paris, Christelle Fraisse, Nicolas Bierne, and Christian Lexer. 2016. “Adaptive Evolution and Segregating Load Contribute to the Genomic Landscape of Divergence in Two Tree Species Connected by Episodic Gene Flow.” Molecular Ecology 26 (1): 59–76. https://doi.org/10.1111/mec.13765.
Fraïsse, Christelle, Iva Popovic, Clément Mazoyer, Bruno Spataro, Stéphane Delmotte, Jonathan Romiguier, Étienne Loire, et al. 2021. “DILS: Demographic Inferences with Linked Selection by Using ABC.” Molecular Ecology Resources 21 (8): 2629–44. https://doi.org/10.1111/1755-0998.13323.
Gutenkunst, Ryan N., Ryan D. Hernandez, Scott H. Williamson, and Carlos D. Bustamante. 2009. “Inferring the Joint Demographic History of Multiple Populations from Multidimensional SNP Frequency Data.” PLoS Genetics 5 (10). https://doi.org/10.1371/journal.pgen.1000695.
Haller, Benjamin C., and Philipp W. Messer. 2023. “SLiM 4: Multispecies Eco-Evolutionary Modeling.” The American Naturalist 201 (5): E127–39. https://doi.org/10.1086/723601.
Horvath, Robert, Mitra Menon, Michelle Stitzer, and Jeffrey Ross-Ibarra. 2022. “Controlling for Variable Transposition Rate with an Age-Adjusted Site Frequency Spectrum.” Genome Biology and Evolution 14 (2): 1–12. https://doi.org/10.1093/gbe/evac016.
Horvath, Robert, Nikolaos Minadakis, Yann Bourgeois, and Anne C. Roulin. 2024. “The Evolution of Transposable Elements in Brachypodium Distachyon Is Governed by Purifying Selection, While Neutral and Adaptive Processes Play a Minor Role.” eLife 12 (February). https://doi.org/10.7554/eLife.93284.2.
Jouganous, Julien, Will Long, Aaron P. Ragsdale, and Simon Gravel. 2017. “Inferring the Joint Demographic History of Multiple Populations: Beyond the Diffusion Approximation.” Genetics 206 (3): 1549–67. https://doi.org/10.1534/genetics.117.200493.
Kent, Tyler V., Jasmina Uzunović, and Stephen I. Wright. 2017. “Coevolution Between Transposable Elements and Recombination.” Philosophical Transactions of the Royal Society B: Biological Sciences 372 (1736). https://doi.org/10.1098/rstb.2016.0458.
Krouglova, Anastasia N., Hayden R. Johnson, Basile Confavreux, Michael Deistler, and Pedro J. Gonçalves. 2026. “Multifidelity Simulation-Based Inference for Computationally Expensive Simulators.” https://arxiv.org/abs/2502.08416.
Liu, Beibei, Manisha Munasinghe, Regina A. Fairbanks, Candice N. Hirsch, and Jeffrey Ross-Ibarra. 2025. “Genome-Wide Selection on Transposable Elements in Maize.” Evolutionary Biology. https://doi.org/10.1101/2025.09.16.676665.
Marsh, Jacob I, Sachin Kaushik, and Parul Johri. 2026. “Effects of Rescaling Forward-in-Time Population Genetic Simulations.” Genetics 232 (2): iyaf263. https://doi.org/10.1093/genetics/iyaf263.
Min, Jiseon, Yuxin Ning, Nathaniel S Pope, Franz Baumdicker, and Andrew D Kern. 2026. “Neural Posterior Estimation for Population Genetics.” Edited by J Novembre. GENETICS 233 (3): iyag107. https://doi.org/10.1093/genetics/iyag107.
Sanchez, Théophile, Pierre Jobic, Cyril Regan, Paul Verdu, Guillaume Charpiat, and Flora Jay. 2026. “Attentive-SPIDNA: Attention-Based Neural Networks for Population Genetics.” Evolutionary Biology. https://doi.org/10.64898/2026.04.15.718687.
Sawyer, S A, and D L Hartl. 1992. “Population Genetics of Polymorphism and Divergence.” Genetics 132 (4): 1161–76. https://doi.org/10.1093/genetics/132.4.1161.
Spence, Jeffrey P., and Yun S. Song. 2019. “Inference and Analysis of Population-Specific Fine-Scale Recombination Maps Across 26 Diverse Human Populations.” Science Advances 5 (10): eaaw9206. https://doi.org/10.1126/sciadv.aaw9206.
Tejero-Cantero, Alvaro, Jan Boelts, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pedro J. Gonçalves, David S. Greenberg, and Jakob H. Macke. 2020. “Sbi: A Toolkit for Simulation-Based Inference.” Journal of Open Source Software 5 (52): 2505. https://doi.org/10.21105/joss.02505.
Xue, Alexander T., Robert P. Ruggiero, Michael J. Hickerson, and Stéphane Boissinot. 2018. “Differential Effect of Selection Against LINE Retrotransposons Among Vertebrates Inferred from Whole-Genome Data and Demographic Modeling.” Genome Biology and Evolution 10 (5): 1265–81. https://doi.org/10.1093/gbe/evy083.
Zaheer, Manzil, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. 2018. “Deep Sets.” https://arxiv.org/abs/1703.06114.