Lots of indepdent loci w/exactly 1 mutation using msprime.

many_indepdendent_loci.py

import demes
import demesdraw
import matplotlib.pyplot as plt
import msprime
import numpy as np
import tskit

dmodel = """
time_units: generations
demes:
- name: ancestor
  epochs:
    - {end_time: 100, start_size: 5000}
- name: d1
  ancestors: [ancestor]
  epochs:
    - {end_time: 0, start_size: 2500}
- name: d2
  ancestors: [ancestor]
  epochs:
    - {end_time: 0, start_size: 2500}
- name: d3
  ancestors: [ancestor]
  epochs:
    - {end_time: 0, start_size: 2500}
- name: d4
  ancestors: [ancestor]
  epochs:
    - {end_time: 0, start_size: 2500}
- name: d5
  ancestors: [ancestor]
  epochs:
    - {end_time: 0, start_size: 2500}
migrations:
  - demes: [d1, d2, d3, d4, d5]
    rate: 1e-4
"""

g = demes.loads(dmodel)

demog = msprime.Demography.from_demes(g)

NSITES = 10000

data = np.zeros(250 * NSITES, dtype=np.int32)
for row, ts in enumerate(
    msprime.sim_ancestry(
        {i + 1: 50 for i in range(5)}, demography=demog, ploidy=1, num_replicates=NSITES
    )
):
    assert ts.num_trees == 1
    assert ts.num_samples == 250
    blens = []
    nodes = []
    t = ts.first()
    p = t.parent_array
    time = ts.tables.nodes.time
    # NOTE: with only 1 tree,
    # all nodes are in the tree.
    # With more than one tree, one must
    # iterate over the nodes in a tree,
    # but that is slower b/c the tskit
    # node orders w/in a tree are built using Python
    for n in range(ts.num_nodes):
        if p[n] != tskit.NULL:
            nodes.append(n)
            blens.append(time[p[n]] - time[n])
    blens = np.array(blens)
    blens = blens / np.sum(blens)
    mnode = np.random.choice(nodes, 1, p=blens)[0]
    data[[(row * 250) + i for i in t.leaves(mnode)]] = 1

nrows = len(data) // NSITES

# Rows are sites, columns are samples
# data = np.array(data).reshape(NSITES, nrows)
data = data.reshape(NSITES, len(data) // NSITES)
print(data[0, :])
print(data[1, :])


#
# a = demesdraw.tubes(g)
# plt.savefig("tubes.png")