Benchmark against GRAPP BOLT-LMM#
This tutorial compares the GRG-backed evolutionary model with GRAPP’s
bolt_lmm_inf [DeHaas et al., 2026] on the forward-simulated populations from
Explicit forward simulation with SLiM. It asks whether the evolutionary model
reproduces the distribution of genic variance across allele frequencies, and
how much fitting time it requires on the same GRG-backed data.
The shared simulation uses:
2,000 diploid individuals and a
10N = 20,000generation burn-in;L = 10^6bases andV_S = 2N = 4,000, henceW_S = V_S/(2N) = 1;sigma_a^2 = sigma_b^2 = 1and the simplifiedrho^2 = 1model; andresidual variance
sigma_e^2 = 0.4.
The ten deterministic replicates (seeds 812–821) are shared between
the two tutorials. Each has its own phenotype and full GRG. For fitting, the
tree sequence is split into two blocks so that GRAPP’s leave-one-chromosome-out
calibration has a block to leave out. Both blocks retain all 2,000 individuals
and use the Genotype Representation Graph data structure
[DeHaas et al., 2025], including coalescent counts for GRAPP’s XTX
traversal.
Timings cover fitting only: operator setup, variance-component estimation, and
each method’s calibration or trace work. They exclude the shared SLiM run,
tskit simplification, and GRG conversion. Both methods use 15 stochastic
trials/probes and a 5e-4 CG tolerance. Times will vary by machine; the
example prints the values observed locally.
Cumulative genic variance#
For a variant with minor allele frequency x, the genic contribution is
2*x*(1-x)*E[beta^2 | x]. The single panel below follows the layout and MAF
bins of Figure A11 in stab1_genetics_template.pdf. It shows the realized
SLiM contribution, the configured evolutionary prior, evo-lmm’s fitted prior,
and GRAPP BOLT-LMM’s global standardized-GRM allocation.
The GRAPP curve is intentionally described as a global-GRM allocation. BOLT-LMM
assigns the fitted sigma_g2 equally across its M eligible standardized
markers (sigma_g2/M per marker). In raw-dosage effect notation this is
E[beta^2 | x] = sigma_g2/(M*2*x*(1-x)); after multiplying by
2*x*(1-x), each marker contributes sigma_g2/M. Its cumulative curve is
therefore a neutral frequency-independent reference. evo-lmm instead uses the
frequency-dependent simplified prior
sigma_b^2/(1 + 2*tau*x*(1-x)).
Each point is the mean across the ten forward replicates, and error bars show
the sample standard deviation. The runtime annotation uses the same summary.
Show the benchmark implementation
1"""Benchmark evo-lmm against GRAPP's GRG-backed BOLT-LMM implementation."""
2
3from __future__ import annotations
4
5from dataclasses import dataclass
6from pathlib import Path
7import sys
8import tempfile
9import time
10from typing import Any
11
12import matplotlib.pyplot as plt
13import numpy as np
14import pygrgl
15import tskit
16
17from evo_lmm import (
18 EvolutionaryLmmOps,
19 SimplifiedPrior,
20 fit_reml,
21 sample_allele_frequencies,
22)
23from evo_lmm.grapp_backend import wrap_grg
24
25
26# The forward tutorial is the source of truth for the simulation configuration.
27# Sphinx's plot directive executes this file without defining __file__, so add
28# the tutorial directory explicitly in that execution mode.
29if "__file__" in globals():
30 _TUTORIAL_DIRECTORY = Path(__file__).resolve().parent
31else:
32 _TUTORIAL_DIRECTORY = Path("docs/tutorials").resolve()
33if str(_TUTORIAL_DIRECTORY) not in sys.path:
34 sys.path.insert(0, str(_TUTORIAL_DIRECTORY))
35
36from slim_forward_simplified import ( # noqa: E402
37 N_GENERATIONS,
38 N_INDIVIDUALS,
39 RESIDUAL_VARIANCE,
40 SEED,
41 SIGMA_A2,
42 TRUE_TAU,
43 V_S,
44 W_S,
45 mutation_effects,
46 run_slim,
47)
48
49
50MAF_THRESHOLDS = np.array([0.001, 0.01, 0.1, 0.2, 0.3, 0.4, 0.5])
51N_FORWARD_BLOCKS = 2
52# Match GRAPP's numerical work budget. At N=2,000 its automatic Monte Carlo
53# rule selects 15 trials and its public BOLT driver defaults to this CG
54# tolerance. Using 64 probes and 1e-8 here measures a different accuracy
55# target rather than the overhead of the evolutionary kernel.
56TRACE_PROBES = 15
57CG_TOL = 5e-4
58
59
60@dataclass(frozen=True)
61class BenchmarkData:
62 """One forward-simulation data set and its two physical GRG blocks."""
63
64 tree_sequence: tskit.TreeSequence
65 full_grg: Any
66 full_frequencies: np.ndarray
67 alpha: np.ndarray
68 phenotype: np.ndarray
69 blocks: tuple[tuple[str, Any], ...]
70 block_frequencies: dict[str, np.ndarray]
71 bolt_blocks: tuple[tuple[str, Any], ...]
72 bolt_block_frequencies: dict[str, np.ndarray]
73
74
75@dataclass(frozen=True)
76class BenchmarkReplicateResult:
77 """One fit pair, runtime pair, and data set used in the benchmark."""
78
79 seed: int
80 data: BenchmarkData
81 evo_fit: Any
82 bolt_fit: Any
83 bolt_stats: list[Any]
84 evo_seconds: float
85 bolt_seconds: float
86
87
88@dataclass(frozen=True)
89class BenchmarkResult:
90 """Aggregate benchmark result over one or more forward replicates."""
91
92 replicates: tuple[BenchmarkReplicateResult, ...]
93
94 def __post_init__(self) -> None:
95 if not self.replicates:
96 raise ValueError("benchmark requires at least one replicate")
97
98 @property
99 def seed(self) -> int:
100 """Return the first seed for backwards-compatible single-run callers."""
101
102 return self.replicates[0].seed
103
104 @property
105 def data(self) -> BenchmarkData:
106 """Return the first data set for backwards-compatible callers."""
107
108 return self.replicates[0].data
109
110 @property
111 def evo_fit(self) -> Any:
112 return self.replicates[0].evo_fit
113
114 @property
115 def bolt_fit(self) -> Any:
116 return self.replicates[0].bolt_fit
117
118 @property
119 def bolt_stats(self) -> list[Any]:
120 return self.replicates[0].bolt_stats
121
122 @property
123 def evo_seconds(self) -> float:
124 return float(np.mean([result.evo_seconds for result in self.replicates]))
125
126 @property
127 def bolt_seconds(self) -> float:
128 return float(np.mean([result.bolt_seconds for result in self.replicates]))
129
130 @property
131 def n_replicates(self) -> int:
132 return len(self.replicates)
133
134 def runtime_summary(self) -> dict[str, tuple[float, float]]:
135 """Return mean and sample standard deviation of each fit runtime."""
136
137 evo = np.asarray([result.evo_seconds for result in self.replicates])
138 bolt = np.asarray([result.bolt_seconds for result in self.replicates])
139 ddof = 1 if self.n_replicates > 1 else 0
140 return {
141 "evo-lmm": (float(np.mean(evo)), float(np.std(evo, ddof=ddof))),
142 "GRAPP BOLT-LMM": (float(np.mean(bolt)), float(np.std(bolt, ddof=ddof))),
143 }
144
145
146def _split_forward_tree_sequence(
147 tree_sequence: tskit.TreeSequence,
148 output_directory: Path,
149) -> tuple[tuple[str, Any], ...]:
150 """Convert two physical halves of one tree sequence to coalescent GRGs.
151
152 GRAPP's top-level BOLT-LMM driver performs leave-one-chromosome-out
153 calibration. Two blocks are therefore needed even though the simulation
154 has one continuous chromosome. The blocks contain disjoint variants from
155 the same tree sequence and preserve the original individual sample set.
156 """
157
158 midpoint = float(tree_sequence.sequence_length) / N_FORWARD_BLOCKS
159 intervals = ((0.0, midpoint), (midpoint, float(tree_sequence.sequence_length)))
160 blocks: list[tuple[str, Any]] = []
161 for index, (left, right) in enumerate(intervals, start=1):
162 block = tree_sequence.keep_intervals([[left, right]], simplify=True).trim()
163 if block.num_mutations == 0:
164 raise RuntimeError(f"forward block {index} contains no mutations")
165 block_path = output_directory / f"slim_forward_block_{index}.trees"
166 block.dump(str(block_path))
167 # GRAPP's XTX traversal requires coalescent counts in the GRG.
168 grg = pygrgl.grg_from_trees(str(block_path), compute_coals=True)
169 blocks.append((f"block-{index}", grg))
170 return tuple(blocks)
171
172
173def _make_bolt_blocks(
174 blocks: tuple[tuple[str, Any], ...],
175 block_frequencies: dict[str, np.ndarray],
176 output_directory: Path,
177) -> tuple[tuple[tuple[str, Any], ...], dict[str, np.ndarray]]:
178 """Remove fixed columns before GRAPP's standardized-GRM calibration."""
179
180 bolt_blocks: list[tuple[str, Any]] = []
181 bolt_frequencies: dict[str, np.ndarray] = {}
182 for label, grg in blocks:
183 frequencies = block_frequencies[label]
184 selected = np.flatnonzero(frequencies * (1.0 - frequencies) > 0.0)
185 if selected.size == 0:
186 raise RuntimeError(f"forward block {label} has no segregating variants")
187 bolt_path = output_directory / f"{label}.segregating.grg"
188 if not pygrgl.save_subset(
189 grg,
190 str(bolt_path),
191 pygrgl.TraversalDirection.DOWN,
192 selected.tolist(),
193 ):
194 raise RuntimeError(f"could not write segregating GRG for {label}")
195 bolt_grg = pygrgl.load_immutable_grg(str(bolt_path))
196 bolt_blocks.append((label, bolt_grg))
197 bolt_frequencies[label] = sample_allele_frequencies(bolt_grg)
198 return tuple(bolt_blocks), bolt_frequencies
199
200
201def _build_benchmark_data(output_directory: Path, seed: int) -> BenchmarkData:
202 """Generate one benchmark data set into a persistent directory."""
203
204 output_directory = Path(output_directory).resolve()
205 output_directory.mkdir(parents=True, exist_ok=True)
206 tree_path = run_slim(output_directory, seed)
207 recorded = tskit.load(str(tree_path))
208 alpha = mutation_effects(recorded)
209 tree_sequence = recorded.simplify(filter_sites=False)
210 if alpha.size != tree_sequence.num_mutations:
211 raise ValueError(
212 "SLiM effect order changed during simplification: "
213 f"{alpha.size} effects for {tree_sequence.num_mutations} mutations"
214 )
215
216 full_path = output_directory / "slim_forward.simplified.trees"
217 tree_sequence.dump(str(full_path))
218 full_grg = pygrgl.grg_from_trees(str(full_path), compute_coals=True)
219 full_frequencies = sample_allele_frequencies(full_grg)
220 if alpha.size != full_grg.num_mutations:
221 raise ValueError(
222 "SLiM effect order does not match GRG mutation order: "
223 f"{alpha.size} effects for {full_grg.num_mutations} mutations"
224 )
225
226 prior = SimplifiedPrior(sigma_b2=SIGMA_A2, tau=TRUE_TAU)
227 full_ops = EvolutionaryLmmOps(
228 full_grg,
229 frequencies=full_frequencies,
230 model="simplified",
231 )
232 genetic_value = full_ops.apply_model_x(alpha)
233 rng = np.random.default_rng(seed + 1)
234 phenotype = genetic_value + rng.normal(
235 0.0,
236 np.sqrt(RESIDUAL_VARIANCE),
237 size=full_ops.n,
238 )
239
240 blocks = _split_forward_tree_sequence(tree_sequence, output_directory)
241 block_frequencies = {
242 label: sample_allele_frequencies(grg) for label, grg in blocks
243 }
244 bolt_blocks, bolt_block_frequencies = _make_bolt_blocks(
245 blocks,
246 block_frequencies,
247 output_directory,
248 )
249 concatenated_frequencies = np.concatenate(
250 [block_frequencies[label] for label, _ in blocks]
251 )
252 np.testing.assert_allclose(concatenated_frequencies, full_frequencies)
253 np.save(output_directory / "alpha.npy", alpha)
254 np.save(output_directory / "phenotype.npy", phenotype)
255 np.save(output_directory / "full.frequencies.npy", full_frequencies)
256 for label, frequencies in block_frequencies.items():
257 np.save(output_directory / f"{label}.frequencies.npy", frequencies)
258 for label, frequencies in bolt_block_frequencies.items():
259 np.save(output_directory / f"{label}.segregating.frequencies.npy", frequencies)
260 (output_directory / "seed.txt").write_text(f"{seed}\n", encoding="utf-8")
261
262 return BenchmarkData(
263 tree_sequence=tree_sequence,
264 full_grg=full_grg,
265 full_frequencies=full_frequencies,
266 alpha=alpha,
267 phenotype=phenotype,
268 blocks=blocks,
269 block_frequencies=block_frequencies,
270 bolt_blocks=bolt_blocks,
271 bolt_block_frequencies=bolt_block_frequencies,
272 )
273
274
275def prepare_benchmark_data(output_directory: Path, seed: int = SEED) -> BenchmarkData:
276 """Run SLiM once and persist all data required by subsequent fit runs."""
277
278 return _build_benchmark_data(Path(output_directory), seed)
279
280
281def load_benchmark_data(output_directory: Path) -> BenchmarkData:
282 """Load a previously prepared benchmark without rerunning SLiM."""
283
284 directory = Path(output_directory)
285 if not (directory / "slim_forward.simplified.trees").exists():
286 raise FileNotFoundError(
287 f"benchmark data are missing in {directory}; run prepare_bolt_benchmark.py first"
288 )
289 tree_sequence = tskit.load(str(directory / "slim_forward.simplified.trees"))
290 full_grg = pygrgl.grg_from_trees(
291 str(directory / "slim_forward.simplified.trees"), compute_coals=True
292 )
293 labels = ("block-1", "block-2")
294 blocks = tuple(
295 (
296 label,
297 pygrgl.grg_from_trees(
298 str(directory / f"slim_forward_{label.replace('-', '_')}.trees"),
299 compute_coals=True,
300 ),
301 )
302 for label in labels
303 )
304 bolt_blocks = tuple(
305 (
306 label,
307 pygrgl.load_immutable_grg(str(directory / f"{label}.segregating.grg")),
308 )
309 for label in labels
310 )
311 block_frequencies = {
312 label: np.load(directory / f"{label}.frequencies.npy") for label in labels
313 }
314 bolt_block_frequencies = {
315 label: np.load(directory / f"{label}.segregating.frequencies.npy")
316 for label in labels
317 }
318 return BenchmarkData(
319 tree_sequence=tree_sequence,
320 full_grg=full_grg,
321 full_frequencies=np.load(directory / "full.frequencies.npy"),
322 alpha=np.load(directory / "alpha.npy"),
323 phenotype=np.load(directory / "phenotype.npy"),
324 blocks=blocks,
325 block_frequencies=block_frequencies,
326 bolt_blocks=bolt_blocks,
327 bolt_block_frequencies=bolt_block_frequencies,
328 )
329
330
331def benchmark_data_from_forward_result(
332 forward_result: dict,
333 output_directory: Path,
334) -> BenchmarkData:
335 """Build benchmark blocks from a previously simulated forward replicate.
336
337 ``slim_forward_simplified.run_replicates`` returns the simplified tree
338 sequence, GRG, frequencies, effects, and phenotype used by its two-panel
339 figure. This adapter reuses those objects and only performs the physical
340 block split needed by GRAPP; it never invokes SLiM.
341 """
342
343 output_directory = Path(output_directory).resolve()
344 output_directory.mkdir(parents=True, exist_ok=True)
345 tree_sequence = forward_result["tree_sequence"]
346 full_grg = forward_result["grg"]
347 full_frequencies = np.asarray(forward_result["frequencies"], dtype=np.float64)
348 alpha = np.asarray(forward_result["alpha"], dtype=np.float64)
349 phenotype = np.asarray(forward_result["phenotype"], dtype=np.float64)
350 if alpha.size != full_grg.num_mutations:
351 raise ValueError(
352 "forward effect order does not match GRG mutation order: "
353 f"{alpha.size} effects for {full_grg.num_mutations} mutations"
354 )
355 blocks = _split_forward_tree_sequence(tree_sequence, output_directory)
356 block_frequencies = {
357 label: sample_allele_frequencies(grg) for label, grg in blocks
358 }
359 bolt_blocks, bolt_block_frequencies = _make_bolt_blocks(
360 blocks,
361 block_frequencies,
362 output_directory,
363 )
364 concatenated_frequencies = np.concatenate(
365 [block_frequencies[label] for label, _ in blocks]
366 )
367 np.testing.assert_allclose(concatenated_frequencies, full_frequencies)
368 return BenchmarkData(
369 tree_sequence=tree_sequence,
370 full_grg=full_grg,
371 full_frequencies=full_frequencies,
372 alpha=alpha,
373 phenotype=phenotype,
374 blocks=blocks,
375 block_frequencies=block_frequencies,
376 bolt_blocks=bolt_blocks,
377 bolt_block_frequencies=bolt_block_frequencies,
378 )
379
380
381def simulate_forward_data(seed: int = SEED) -> BenchmarkData:
382 """Run the simulation in a temporary directory for in-process callers."""
383
384 with tempfile.TemporaryDirectory(prefix="evo_lmm_bolt_benchmark_") as directory:
385 return _build_benchmark_data(Path(directory), seed)
386
387
388def _fit_benchmark_data(seed: int, data: BenchmarkData) -> BenchmarkReplicateResult:
389 """Fit both methods for one prepared data set and time only the fits."""
390
391 initial = SimplifiedPrior(sigma_b2=SIGMA_A2, tau=TRUE_TAU)
392
393 start = time.perf_counter()
394 evo_ops = EvolutionaryLmmOps(
395 data.blocks,
396 frequencies=data.block_frequencies,
397 model="simplified",
398 )
399 evo_fit = fit_reml(
400 evo_ops,
401 data.phenotype,
402 initial=initial,
403 trace_probes=TRACE_PROBES,
404 # GRAPP's secant search makes at most seven variance-component
405 # evaluations (two initial points plus five updates). Allow one extra
406 # AI-REML update while keeping optimizer work comparable.
407 max_iter=8,
408 cg_tol=CG_TOL,
409 seed=seed + 2,
410 trace_method="hutchinson",
411 )
412 evo_seconds = time.perf_counter() - start
413
414 # Import GRAPP's public BOLT-LMM driver here so the data-generation helper
415 # remains usable for inspecting the shared data without fitting.
416 from grapp.assoc.bolt_inf_core import CovariateBasis
417 from grapp.assoc.bolt_lmm import bolt_lmm_inf
418
419 grapp_blocks = [(label, wrap_grg(grg)) for label, grg in data.bolt_blocks]
420 covariates = CovariateBasis.intercept_only(N_INDIVIDUALS)
421 start = time.perf_counter()
422 bolt_fit, _calibration, _residuals, bolt_stats = bolt_lmm_inf(
423 grapp_blocks,
424 data.phenotype,
425 covariates,
426 mc_trials=TRACE_PROBES,
427 cg_tol=CG_TOL,
428 seed=seed + 2,
429 threads=1,
430 batched_apply_x=True,
431 )
432 bolt_seconds = time.perf_counter() - start
433
434 return BenchmarkReplicateResult(
435 seed=seed,
436 data=data,
437 evo_fit=evo_fit,
438 bolt_fit=bolt_fit,
439 bolt_stats=bolt_stats,
440 evo_seconds=evo_seconds,
441 bolt_seconds=bolt_seconds,
442 )
443
444
445def run_benchmark(
446 seed: int = SEED,
447 *,
448 data_directory: Path | None = None,
449 forward_results: list[dict] | None = None,
450) -> BenchmarkResult:
451 """Fit evo-lmm and GRAPP over one or more prepared replicates.
452
453 ``forward_results`` is the output of the two-panel SLiM tutorial and is
454 preferred for the documentation benchmark: it reuses those simulations
455 and phenotypes while fitting each method with the benchmark's matched
456 numerical budget. ``data_directory`` remains a single-replicate fallback
457 for the fit-only command and backwards-compatible callers.
458 """
459
460 if forward_results is not None and data_directory is not None:
461 raise ValueError("pass either forward_results or data_directory, not both")
462
463 if forward_results is None:
464 data = (
465 load_benchmark_data(data_directory)
466 if data_directory is not None
467 else simulate_forward_data(seed)
468 )
469 return BenchmarkResult((_fit_benchmark_data(seed, data),))
470
471 if not forward_results:
472 raise ValueError("forward_results must contain at least one replicate")
473 with tempfile.TemporaryDirectory(prefix="evo_lmm_benchmark_blocks_") as directory:
474 root = Path(directory)
475 fitted = []
476 for index, forward_result in enumerate(forward_results):
477 replicate_seed = int(forward_result.get("seed", seed + index))
478 data = benchmark_data_from_forward_result(
479 forward_result,
480 root / f"seed_{replicate_seed}",
481 )
482 fitted.append(_fit_benchmark_data(replicate_seed, data))
483 return BenchmarkResult(tuple(fitted))
484
485
486def _cumulative_by_maf(
487 maf: np.ndarray,
488 contributions: np.ndarray,
489 *,
490 eligible: np.ndarray | None = None,
491) -> np.ndarray:
492 """Sum per-variant genic contributions below each MAF threshold."""
493
494 maf = np.asarray(maf, dtype=np.float64)
495 contributions = np.asarray(contributions, dtype=np.float64)
496 if maf.shape != contributions.shape:
497 raise ValueError("maf and contributions must have matching shapes")
498 if eligible is not None:
499 eligible = np.asarray(eligible, dtype=bool)
500 if eligible.shape != maf.shape:
501 raise ValueError("eligible must match maf")
502 maf = maf[eligible]
503 contributions = contributions[eligible]
504 order = np.argsort(maf, kind="stable")
505 sorted_maf = maf[order]
506 sorted_contributions = contributions[order]
507 cumulative = np.cumsum(sorted_contributions)
508 positions = np.searchsorted(sorted_maf, MAF_THRESHOLDS, side="right")
509 return np.where(positions > 0, cumulative[np.maximum(positions - 1, 0)], 0.0)
510
511
512def _bolt_curve(result: BenchmarkReplicateResult) -> np.ndarray:
513 """Return GRAPP's global-GRM cumulative variance allocation.
514
515 BOLT-LMM uses a standardized global GRM, so its fitted genetic scale is
516 allocated equally across model SNPs: each eligible marker contributes
517 ``sigma_g2 / M``. This is the neutral reference curve for this comparison,
518 rather than the frequency-dependent evolutionary prior used by evo-lmm.
519 """
520
521 frequencies = np.concatenate(
522 [np.asarray(stats.a1freq, dtype=np.float64) for stats in result.bolt_stats]
523 )
524 se = np.concatenate([np.asarray(stats.se, dtype=np.float64) for stats in result.bolt_stats])
525 eligible = np.isfinite(se)
526 marker_count = int(np.count_nonzero(eligible))
527 if marker_count == 0:
528 raise RuntimeError("GRAPP BOLT-LMM returned no model markers")
529 marker_contribution = np.zeros(frequencies.size, dtype=np.float64)
530 marker_contribution[eligible] = result.bolt_fit.sigma_g2 / marker_count
531 maf = np.minimum(frequencies, 1.0 - frequencies)
532 return _cumulative_by_maf(maf, marker_contribution, eligible=eligible)
533
534
535def _replicate_curves(result: BenchmarkReplicateResult) -> dict[str, np.ndarray]:
536 """Compute the four cumulative genic-variance curves for one replicate."""
537
538 data = result.data
539 maf = np.minimum(data.full_frequencies, 1.0 - data.full_frequencies)
540 q = data.full_frequencies * (1.0 - data.full_frequencies)
541 segregating = q > 0.0
542 genic_factor = 2.0 * q
543 configured_prior = SimplifiedPrior(sigma_b2=SIGMA_A2, tau=TRUE_TAU)
544 return {
545 "SLiM realization": _cumulative_by_maf(
546 maf,
547 np.square(data.alpha) * genic_factor,
548 eligible=segregating,
549 ),
550 "configured evolutionary prior": _cumulative_by_maf(
551 maf,
552 configured_prior.effect_variances(data.full_frequencies) * genic_factor,
553 eligible=segregating,
554 ),
555 "evo-lmm fitted prior": _cumulative_by_maf(
556 maf,
557 result.evo_fit.prior.effect_variances(data.full_frequencies) * genic_factor,
558 eligible=segregating,
559 ),
560 "GRAPP BOLT-LMM": _bolt_curve(result),
561 }
562
563
564def make_summary(result: BenchmarkResult) -> plt.Figure:
565 """Plot mean cumulative genic variance with replicate variation bars."""
566
567 curve_values = {
568 label: np.stack([_replicate_curves(rep)[label] for rep in result.replicates])
569 for label in (
570 "SLiM realization",
571 "configured evolutionary prior",
572 "evo-lmm fitted prior",
573 "GRAPP BOLT-LMM",
574 )
575 }
576 means = {label: np.mean(values, axis=0) for label, values in curve_values.items()}
577 ddof = 1 if result.n_replicates > 1 else 0
578 variations = {
579 label: np.std(values, axis=0, ddof=ddof)
580 for label, values in curve_values.items()
581 }
582
583 figure, axis = plt.subplots(figsize=(6.4, 4.2), constrained_layout=True)
584 x = np.arange(MAF_THRESHOLDS.size)
585 axis.errorbar(
586 x,
587 means["SLiM realization"],
588 yerr=variations["SLiM realization"],
589 color="tab:green",
590 marker="o",
591 linewidth=1.8,
592 capsize=3,
593 label="SLiM realization",
594 )
595 axis.errorbar(
596 x,
597 means["configured evolutionary prior"],
598 yerr=variations["configured evolutionary prior"],
599 color="black",
600 linewidth=1.4,
601 capsize=3,
602 label=r"configured evolutionary prior ($\sigma_a^2=1$, $W_S=1$)",
603 )
604 axis.errorbar(
605 x,
606 means["evo-lmm fitted prior"],
607 yerr=variations["evo-lmm fitted prior"],
608 color="tab:red",
609 linestyle="--",
610 marker="s",
611 linewidth=1.8,
612 capsize=3,
613 label="evo-lmm fitted prior",
614 )
615 axis.errorbar(
616 x,
617 means["GRAPP BOLT-LMM"],
618 yerr=variations["GRAPP BOLT-LMM"],
619 color="tab:blue",
620 linestyle=":",
621 marker="^",
622 linewidth=2.0,
623 capsize=3,
624 label=r"GRAPP BOLT-LMM (global GRM)",
625 )
626 axis.set_xticks(x, [f"{threshold:g}" for threshold in MAF_THRESHOLDS], rotation=35, ha="right")
627 axis.set_xlabel("cumulative MAF bin")
628 axis.set_ylabel("cumulative genic variance")
629 axis.set_title(
630 "Cumulative genic variance across MAF bins\n"
631 rf"$N={N_INDIVIDUALS}$, $V_S={V_S:g}$, $W_S={W_S:g}$, "
632 rf"$\rho^2=1$, {result.n_replicates} replicates"
633 )
634 axis.grid(axis="y", linestyle=":", alpha=0.55)
635 axis.set_ylim(bottom=0.0)
636 axis.legend(loc="upper left", fontsize=8)
637 axis.text(
638 0.98,
639 0.06,
640 "runtime (mean $\\pm$ SD)\n"
641 f"evo-lmm: {result.runtime_summary()['evo-lmm'][0]:.2f} "
642 f"$\\pm$ {result.runtime_summary()['evo-lmm'][1]:.2f} s\n"
643 f"GRAPP BOLT-LMM: {result.runtime_summary()['GRAPP BOLT-LMM'][0]:.2f} "
644 f"$\\pm$ {result.runtime_summary()['GRAPP BOLT-LMM'][1]:.2f} s",
645 transform=axis.transAxes,
646 ha="right",
647 va="bottom",
648 fontsize=8,
649 bbox={"facecolor": "white", "edgecolor": "0.7", "alpha": 0.85},
650 )
651 return figure
652
653
654if __name__ == "__main__":
655 forward_artifacts = Path("docs/_artifacts/forward_replicates")
656 artifact_directory = Path("docs/_artifacts/bolt_seed_812")
657 if forward_artifacts.exists():
658 from slim_forward_simplified import load_simulation_replicates
659
660 benchmark = run_benchmark(
661 forward_results=load_simulation_replicates(forward_artifacts),
662 )
663 else:
664 benchmark = run_benchmark(
665 data_directory=artifact_directory if artifact_directory.exists() else None,
666 )
667 print(
668 f"replicates={benchmark.n_replicates} "
669 f"mutations_first={benchmark.data.full_grg.num_mutations} "
670 f"evo_lmm_seconds={benchmark.evo_seconds:.6f} "
671 f"grapp_bolt_lmm_seconds={benchmark.bolt_seconds:.6f}"
672 )
673 print(
674 f"evo_lmm_sigma_b2={benchmark.evo_fit.prior.sigma_b2:.6g} "
675 f"evo_lmm_tau={benchmark.evo_fit.prior.tau:.6g} "
676 f"grapp_sigma_g2={benchmark.bolt_fit.sigma_g2:.6g}"
677 )
678 make_summary(benchmark)
679 if "agg" not in plt.get_backend().lower():
680 plt.show()
Generate the ten forward-simulation replicates and their GRGs:
uv run python docs/tutorials/prepare_bolt_benchmark.py
Then rerun the ten-replicate fitting comparison without repeating SLiM:
uv run python docs/tutorials/fit_bolt_benchmark.py
The fit-only script reports each replicate together with the mean and sample standard deviation of both runtimes.
The current ten-replicate fit run reports:
Method |
Mean fit seconds |
Sample SD seconds |
|---|---|---|
evo-lmm (warm Hutchinson, 15 vectors) |
27.24 |
6.93 |
GRAPP BOLT-LMM |
10.67 |
0.96 |
