Best way to run many replicates while drawing parameters from a prior distribution?

[petrelharp]

Let's say I want to run a model across a range of parameter values (eg for ABC inference). Modifying an example from the manual:

modern_size, ancestral_size, split_time = 300, 1000, 500
demography = msprime.Demography()
demography.add_population(name="popA", initial_size=modern_size)
demography.add_population(name="popB", initial_size=modern_size)
demography.add_population(name="ancestral", initial_size=ancestral_size)
demography.add_population_split(time=split_time, derived=["popA", "popB"], ancestral="ancestral")

let's say I want to run this model and extract statistics either across a grid of parameter values, or by drawing the parameters from a prior. For instance:

ts = msprime.sim_ancestry(
        {"popA": 100, "popB": 100},
        demography=demography,
        sequence_length=1e6,
        recombination_rate=1e-8
)
ts = msprime.sim_mutations(ts, rate=1e-8)
pop_names = [p.metadata['name'] for p in ts.populations()]
sample_sets = [ts.samples(pop_names.index(n)) for n in ["popA", "popB"]]
ts.divergence(
        sample_sets,
        indexes=[(0, 0), (0, 1), (1, 1)]
)

What's the best way to do this?

[petrelharp]

One option would be to make the code a function that takes the parameters and outputs the statistics of interest:

def run_sim(modern_size, ancestral_size, split_time):
    demography = msprime.Demography()
    demography.add_population(name="popA", initial_size=modern_size)
    demography.add_population(name="popB", initial_size=modern_size)
    demography.add_population(name="ancestral", initial_size=ancestral_size)
    demography.add_population_split(time=split_time, derived=["popA", "popB"], ancestral="ancestral")

    ts = msprime.sim_ancestry(
            {"popA": 100, "popB": 100},
            demography=demography,
            sequence_length=1e6,
            recombination_rate=1e-8
    )
    ts = msprime.sim_mutations(ts, rate=1e-8)
    pop_names = [p.metadata['name'] for p in ts.populations()]
    sample_sets = [ts.samples(pop_names.index(n)) for n in ["popA", "popB"]]
    out = ts.divergence(
            sample_sets,
            indexes=[(0, 0), (0, 1), (1, 1)]
    )
    return out

and then do e.g.

nreps = 100
df = pd.DataFrame()
for n in ["modern_size", "ancestral_size", "split_time", "piA", "piB", "dxy"]:
    df[n] = np.repeat(-1, nreps)

rng = np.random.default_rng()
for j in range(100):
    modern_size = rng.exponential(1) * 300
    ancestral_size = rng.exponential(1) * 1000
    split_time = rng.exponential(1) * 500
    piA, dxy, piB = run_sim(modern_size, ancestral_size, split_time)
    df['modern_size'][j] = modern_size
    df['ancestral_size'][j] = ancestral_size
    df['split_time'][j] = split_time
    df['piA'][j] = piA
    df['piB'][j] = piB
    df['dxy'][j] = dxy

[eldonb]

Maybe not exactly an answer to the question but I have used the "multiprocessing" python library as follows:

as @petrelharp suggested define a function to simulate your model and compute the statistics, e.g.

def run_sims( input_alpha ) : 
    ## model and data and  statistics 

where input_a is the parameter value I am computing over ( note, can be multi.dimensional) :
then define the grid. e.g.

## grid of alpha values
valpha = numpy.arange(1.01, 2.0,  0.01)

the set it up as follows:

## prepare for running  the  simulator
## by spreading the alpha grid over cores
if __name__ == '__main__':
    ## 'processes' is number of available cores
    with Pool(processes = 8) as p:
        p.map(run_sims, valpha)

this works on a laptop/desktop with several cpus with the imports

import msprime
import tskit
##import scipy.special
import numpy
##import math
##import sys
from multiprocessing import Pool

[grahamgower]

The wrapper function pattern is also appropriate when doing optimisation over the parameters, eg. using scipy.optimize.minimize(). WARNING: I didn't check that this actually works for the given example.

def run_sims(modern_size, ancestral_size, split_time):
    # As above.
    ...

def loss(sim_out, observed=3.14152):
    # E.g. minimise squared distance to empirical data.
    return (sim_out - observed)**2

def opt_func(x):
    modern_size, ancestral_size, split_time = x
    out = run_sims(modern_size, ancestral_size, split_time)
    return loss(out)


import scipy.optimize
x0 = 300, 1000, 500  # initial guess for params
res = scipy.optimize.minimize(opt_func, x0)
if res.success:
    # Best fit params are now in `res.x`. Note: this might be a local optima!
    print(res.x)
else:
    print("optimisation failed")