Provide smaller example?

[hyanwong]

For testing purposes, I think it would be useful to provide a small msprime or stdpopsim example that only takes a few seconds to run. At the moment I'm using the example at the end of the example.ipynb (I guess I could turn down the length_multiplier to 0.01)

import phlash.sim

sim_contigs = phlash.sim.stdpopsim_dataset(
    "HomSap",
    "Zigzag_1S14",
    {"generic": 100},
    options=dict(length_multiplier=0.1),
)
test_k = list(sim_contigs["data"])[0]
test_data = sim_contigs["data"][test_k]
train_data = [v for k, v in sim_contigs["data"].items() if k != test_k]
results = phlash.fit(train_data, test_data, truth=sim_contigs["truth"], fold_sfs=False)

For reference, I'm testing Phlash on google colab, so I can demonstrate it to students. It would be helpful not to use all my colab GPU credits/time just for a demo.

Incidentally, when running the phlash.sim.stdpopsim_dataset command above on colab, I get

phlash/sim.py:119: UserWarning: Doubling the number of steps has resulted in different  predictions, please re-run with smaller step sizes to ensure  numerical accuracy.
  c, _ = md.coalescence_rate_trajectory(t, d)

[hyanwong]

Actually, the easiest way to get a trivial example working on colab seems to be to turn down the number of iterations, e.g.

results = phlash.fit(train_data, test_data, truth=sim_contigs["truth"], fold_sfs=False, niter=10)

But I don't know what a reasonable small number of iterations would be to demo phlash to students.

[terhorst]

Another option we might try is increasing the learning rate, which can be passed as a kwarg to mcmc.fit:

phlash/src/phlash/mcmc.py

Line 178 in e5fcf69

opt = optax.amsgrad(learning_rate=options.get("learning_rate", 0.1))

It's set to a conservatively low value to prevent runaway behavior/NaNs, but for a hand-picked example you might be able to tune it a bit higher. I'll play with it and let you know.

[hyanwong]

Thanks, that would be great.