<jupyter_start><jupyter_text>imports<jupyter_code>import scipy.stats as st
import numpy as np
import h5py
import scipy.integrate as integrate
from numba import njit
import toyplot
from tqdm.auto import tqdm
from scipy.special import binom<jupyter_output><empty_output><jupyter_text>define useful functions We'll use these later but don't worry about them now.<jupyter_code># things we'll need:

# because scipy.stats.binom is slow...
def poisson_pmf(k, lam):
    try:
        x = (lam**k) * np.exp(-1*lam) / np.math.factorial(k)
        return(x)
    except:
        return( np.nan )

@njit
def binom_jit(n,k,p,com):
    return(com*p**k*(1-p)**(n-k))

# full binomial function!
def binomial(n,k,p):
    com=binom(n,k)
    return(binom_jit(n,k,p,com))    
    
# get the combinations of different crossovers from n_1 and ali_size that would yield a given n_2
# (that is, the different ways hap1s can switch to hap0 and hap0s can switch to hap1 to give us n_2)
@njit
def get_combs(n_1,ali_size,n_2):
    zs = ali_size - n_1
    arr_ = np.zeros((n_1+1,zs+1),dtype=np.int64)
    for row in range(arr_.shape[0]):
        for col in range(arr_.shape[1]):
            arr_[row,col] = n_1 - row + col
    # top row is the number of n_1 that change to zero, bottom is number of zero that change to 1
    return(np.vstack(np.where(arr_ == n_2)).T)

# calculate the probability of switching to an n_2 value given an n_1 value, p_1, p_2, and aliquot size

def calc_n_2_prob(n_1, n_2, p_1, p_2, ali_size):
    sumprob_n_2 = 0.0
    for i in get_combs(n_1, ali_size, n_2):
        nstay = n_1 - i[0] # nstay is the number of n_1 that stay n_1... i[0] is the number that change!
        nchange = i[1]
        # prob that nstay ones _do_ stay at one
        p_nstay = binomial(k=nstay, n=n_1, p=(1-p_2)*(1-p_1)+(1-p_2)*p_1*(1-p_2)+p_2*p_1*p_2)
        # prob that nchange zeros _do_ change to one
        p_nchange = binomial(k=nchange, n=ali_size-n_1, p=p_2*(1-p_1)+p_2*p_1*(1-p_2)+(1-p_2)*p_1*p_2)
        
        sumprob_n_2 += p_nstay * p_nchange
    return(sumprob_n_2)<jupyter_output><empty_output><jupyter_text>Define a full model  Make a pretend recombination map.<jupyter_code># define an arbitrary recombination map...

# class for making a continuous pdf
class my_pdf(st.rv_continuous):
    def _pdf(self,x):
        return (1+np.cos(12.58*x))  # Normalized over its range, in this case [0,1]
# define recombination map
my_cv = my_pdf(a=0, b=2, name='my_pdf')
# here's our sample recombination map
sns.scatterplot(range(1000), 1+1*np.cos(12.58*np.linspace(0,1,num=1000)+0) );<jupyter_output><empty_output><jupyter_text>Integrating from `alpha` to `beta` along the x-axis gives the bernoulli probability of a single recombination event landing between `alpha` and `beta`.  Sequencing parameters<jupyter_code>np.random.seed(123)
# the number of total aliquots sequenced -- used to calculate expected coverage per locus
nalis = 500

# number of gametes per aliquot
ali_size = 20

# number of loci per gamete that can be sequenced
nloci = 100

# sequencing effort. It's really just the number of random draws drom the pooled aliquots
nreads = 1e7

# from the poisson process of selecting a single locus from all
# pooled loci, which is a tiny probability that happens many times.
exp_cov_per_locus = nreads/nalis/nloci
exp_cov_per_locus<jupyter_output><empty_output><jupyter_text>Biological parameters<jupyter_code># prob of a second crossover for each gamete. This is something defined by me. 
p_1 = 0.2

# first locus location (arbitrary choice for now)
alpha = 0.0

# second locus location (arbitrary choice for now)
beta = 0.32

# prob of a single crossover landing between alpha and beta values
# this is done by integrating between alpha and beta on the recomb map
p_2 = integrate.quad(lambda x: 1+np.cos(12.58*x), alpha, beta)[0]<jupyter_output><empty_output><jupyter_text>Likelihood calculation<jupyter_code># a pretend observation:
n_1_obs = 101 # number of 1-haplotype reads observed at locus Alpha
n_2_obs = 110 # number of 1-haplotype reads observed at locus Beta
# calculating likelihood
total_lik = 0
for n_1 in range(ali_size+1):
    for n_2 in range(ali_size+1):
        # prob of getting n_1 = n_1
        p_n_1 = binomial(k = n_1, n = ali_size, p = 0.5) # .5 prob of each alpha haplo being 1

        # prob of getting n_1_obs if n_1 = n_1?
        p_n_1_obs = poisson_pmf(k=n_1_obs, lam = (n_1/ali_size)*exp_cov_per_locus)

        # prob of n_2 being zero if n_1 is zero?
        p_n_2 = calc_n_2_prob(n_1,n_2,p_1,p_2,ali_size)

        # prob of getting n_2_obs if n_2 = n_2?
        p_n_2_obs = poisson_pmf(k=n_2_obs, lam = (n_2/ali_size)*exp_cov_per_locus)

        # total prob
        if 0 not in [p_n_1,p_n_1_obs,p_n_2,p_n_2_obs]:
            total_lik += np.exp( np.log(p_n_1)+np.log(p_n_1_obs)+np.log(p_n_2)+np.log(p_n_2_obs) )
        else:
            total_lik += 0
total_lik<jupyter_output><empty_output><jupyter_text>Now try to back-infer the recombination rate between two points in our simulated data.  I simulated 500 aliquots, each with 20 gametes, each with 100 loci. 1e7 total reads sequenced.  This was simulated under the recombination map defined at the top of the notebook.<jupyter_code>haps = h5py.File("500ali_20gpa_1e7reads/haplotypes.hdf5",'r')
haps.keys() # aliquots
haps['0'].keys() # loci
np.array(haps['0']['0']) # sequenced haps at this locus<jupyter_output><empty_output><jupyter_text>Try to recover the recombination probability between two points from our simulated data.  1) Loci at an intermediate distance  Choose our two loci, and look at the haplotype numbers at those loci across all 500 aliquots.<jupyter_code># what what is locus Alpha, how many hap1s do we observe there?
locusnumAlpha = 0 # locus location (out of 100)
n_1_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumAlpha)])) for i in range(500)])

# what what is locus Beta, how many hap1s do we observe there?
locusnumBeta = 36 # locus location (out of 100)
n_2_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumBeta)])) for i in range(500)])

# hold the haplotype data at these loci across all aliquots
n1n2 = np.vstack([n_1_obs_set,n_2_obs_set]).T<jupyter_output><empty_output><jupyter_text>Look at the "real" probability of recombination between these two loci.<jupyter_code># what is the true recombination rate between these two loci, used in simulation?
# this is the number we are trying to recover
p_2 = integrate.quad(lambda x: 1+np.cos(12.58*x), locusnumAlpha*0.01, locusnumBeta*0.01)[0]
p_2<jupyter_output><empty_output><jupyter_text>Now, pretending we don't know the real answer, cycle through possible p_2 values, recording the likelihood of each.<jupyter_code># calculate the likelihood of observing n_1_obs hap1 reads at locus 1 and n_2_obs hap1 reads at locus 2
full_set = []
# test across values of p, finding which one yields the highest likelihood.
for p_2 in tqdm(np.linspace(0,1,num=101)):
    set1 = []
    for n_1_obs, n_2_obs in n1n2:
        total_lik = 0
        for n_1 in range(ali_size+1):
            for n_2 in range(ali_size+1):
                # prob of getting n_1 = 0
                p_n_1 = binomial(k = n_1, n = ali_size, p = 0.5) # .5 prob of each alpha haplo being 1

                # prob of getting n_1_obs if n_1 = 0?
                p_n_1_obs = poisson_pmf(k=n_1_obs, lam = (n_1/ali_size)*exp_cov_per_locus)

                # prob of n_2 being zero if n_1 is zero?
                p_n_2 = calc_n_2_prob(n_1,n_2,p_1,p_2,ali_size)

                # prob of getting n_2_obs if n_2 = 0?
                p_n_2_obs = poisson_pmf(k=n_2_obs, lam = (n_2/ali_size)*exp_cov_per_locus)

                # total prob
                if 0 not in [p_n_1,p_n_1_obs,p_n_2,p_n_2_obs]:
                    total_lik += np.exp( np.log(p_n_1)+np.log(p_n_1_obs)+np.log(p_n_2)+np.log(p_n_2_obs) )
                else:
                    total_lik += 0
        set1.append(np.log(total_lik))
    full_set.append(np.array(set1))
    
prob_sums = []
for i in full_set:
    prob_sums.append(np.sum(i[(~np.isnan(i)).astype(int) + (~np.isinf(i)).astype(int) == 2]))

canvas = toyplot.Canvas(width=300, height=300)
axes = canvas.cartesian()
mark = axes.scatterplot(np.linspace(0,1,num=101), np.array(prob_sums))<jupyter_output><empty_output><jupyter_text>Predicted recombination rate (argmax of curve above):<jupyter_code># this should recover something close to our simulated p_2 value
np.linspace(0,1,num=101)[np.argmax(prob_sums)]<jupyter_output><empty_output><jupyter_text>This is pretty close to our "real" recombination probability.  2) Choose a different two loci, this time with a much smaller probability of recombination:<jupyter_code># what what is locus Alpha, how many hap1s do we observe there?
locusnumAlpha = 29 # locus location
n_1_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumAlpha)])) for i in range(500)])

# what what is locus Beta, how many hap1s do we observe there?
locusnumBeta = 38 # locus location
n_2_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumBeta)])) for i in range(500)])

# hold the haplotype data at these loci across all aliquots
n1n2 = np.vstack([n_1_obs_set,n_2_obs_set]).T<jupyter_output><empty_output><jupyter_text>True recombination rate:<jupyter_code># what is the true recombination rate between these two loci, used in simulation?
# this is the number we are trying to recover
p_2 = integrate.quad(lambda x: 1+np.cos(12.58*x), locusnumAlpha*0.01, locusnumBeta*0.01)[0]
p_2<jupyter_output><empty_output><jupyter_text>Back-infer the true recombination rate.<jupyter_code># calculate the likelihood of observing n_1_obs hap1 reads at locus 1 and n_2_obs hap1 reads at locus 2
full_set = []
# test across values of p, finding which one yields the highest likelihood.
for p_2 in tqdm(np.linspace(0,1,num=101)):
    set1 = []
    for n_1_obs, n_2_obs in n1n2:
        total_lik = 0
        for n_1 in range(ali_size+1):
            for n_2 in range(ali_size+1):
                # prob of getting n_1 = 0
                p_n_1 = binomial(k = n_1, n = ali_size, p = 0.5) # .5 prob of each alpha haplo being 1

                # prob of getting n_1_obs if n_1 = 0?
                p_n_1_obs = poisson_pmf(k=n_1_obs, lam = (n_1/ali_size)*exp_cov_per_locus)

                # prob of n_2 being zero if n_1 is zero?
                p_n_2 = calc_n_2_prob(n_1,n_2,p_1,p_2,ali_size)

                # prob of getting n_2_obs if n_2 = 0?
                p_n_2_obs = poisson_pmf(k=n_2_obs, lam = (n_2/ali_size)*exp_cov_per_locus)

                # total prob
                if 0 not in [p_n_1,p_n_1_obs,p_n_2,p_n_2_obs]:
                    total_lik += np.exp( np.log(p_n_1)+np.log(p_n_1_obs)+np.log(p_n_2)+np.log(p_n_2_obs) )
                else:
                    total_lik += 0
        set1.append(np.log(total_lik))
    full_set.append(np.array(set1))
    
prob_sums = []
for i in full_set:
    prob_sums.append(np.sum(i[(~np.isnan(i)).astype(int) + (~np.isinf(i)).astype(int) == 2]))

canvas = toyplot.Canvas(width=300, height=300)
axes = canvas.cartesian()
mark = axes.scatterplot(np.linspace(0,1,num=101), np.array(prob_sums))
# this should recover something close to our simulated p_2 value
np.linspace(0,1,num=101)[np.argmax(prob_sums)]<jupyter_output><empty_output><jupyter_text>3) Finally, choose two loci that are far apart with high probability of recombination in-between.<jupyter_code># what what is locus Alpha, how many hap1s do we observe there?
locusnumAlpha = 3 # locus location
n_1_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumAlpha)])) for i in range(500)])

# what what is locus Beta, how many hap1s do we observe there?
locusnumBeta = 98 # locus location
n_2_obs_set = np.array([np.sum(np.array(haps[str(i)][str(locusnumBeta)])) for i in range(500)])

# hold the haplotype data at these loci across all aliquots
n1n2 = np.vstack([n_1_obs_set,n_2_obs_set]).T
# what is the true recombination rate between these two loci, used in simulation?
# this is the number we are trying to recover
p_2 = integrate.quad(lambda x: 1+np.cos(12.58*x), locusnumAlpha*0.01, locusnumBeta*0.01)[0]
p_2
# calculate the likelihood of observing n_1_obs hap1 reads at locus 1 and n_2_obs hap1 reads at locus 2
full_set = []
# test across values of p, finding which one yields the highest likelihood.
for p_2 in tqdm(np.linspace(0,1,num=101)):
    set1 = []
    for n_1_obs, n_2_obs in n1n2:
        total_lik = 0
        for n_1 in range(ali_size+1):
            for n_2 in range(ali_size+1):
                # prob of getting n_1 = 0
                p_n_1 = binomial(k = n_1, n = ali_size, p = 0.5) # .5 prob of each alpha haplo being 1

                # prob of getting n_1_obs if n_1 = 0?
                p_n_1_obs = poisson_pmf(k=n_1_obs, lam = (n_1/ali_size)*exp_cov_per_locus)

                # prob of n_2 being zero if n_1 is zero?
                p_n_2 = calc_n_2_prob(n_1,n_2,p_1,p_2,ali_size)

                # prob of getting n_2_obs if n_2 = 0?
                p_n_2_obs = poisson_pmf(k=n_2_obs, lam = (n_2/ali_size)*exp_cov_per_locus)

                # total prob
                if 0 not in [p_n_1,p_n_1_obs,p_n_2,p_n_2_obs]:
                    total_lik += np.exp( np.log(p_n_1)+np.log(p_n_1_obs)+np.log(p_n_2)+np.log(p_n_2_obs) )
                else:
                    total_lik += 0
        set1.append(np.log(total_lik))
    full_set.append(np.array(set1))
    
prob_sums = []
for i in full_set:
    prob_sums.append(np.sum(i[(~np.isnan(i)).astype(int) + (~np.isinf(i)).astype(int) == 2]))

canvas = toyplot.Canvas(width=300, height=300)
axes = canvas.cartesian()
mark = axes.scatterplot(np.linspace(0,1,num=101), np.array(prob_sums))
# this should recover something close to our simulated p_2 value
np.linspace(0,1,num=101)[np.argmax(prob_sums)]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10655204796648125, 'fraction_numerical': 0.03951455609333237, 'mean_word_length': 3.591708126036484, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '231739', 'license': 'MIT', 'n_tokens_mistral': 5594, 'n_tokens_neox': 5175, 'n_words': 1405, 'path': 'notebooks/pooling-likelihood-pfm.ipynb', 'repo_name': 'eaton-lab/poolparty', 'size': '13843'}