<jupyter_start><jupyter_text>Take a vector of maf of two population, randomly choose a portion value and sample genotype (for now, fix 0.5)<jupyter_code>import random 
import numpy as np 
import pandas as pd 

from pandas_plink import read_plink1_bin



%load_ext autoreload
%autoreload 2

#read in minor allele frequencies for two population
sas_frq= np.genfromtxt(fname="./sas.frq",  missing_values={"NA"})
irish_frq = np.genfromtxt(fname="./irish.frq")
wb_frq = np.genfromtxt(fname="./wb.frq")
#length of chr22 is 5048
#M is num of snps
#N is num of individuals 
M= 5000
N = 4000
#let local chunk size be 100 snps, thus in total 50 chunks 
np.random.seed(10)
local_assignment = np.random.binomial(size=(4000,50), n=1, p=0.5)
#local_assignment = np.zeros(shape=(4000,50))
#local_assignment[0:2000,:] = 1
#local_assignment = local_assignment.astype(int)
#fix local assignment 
local_assignment = local_assignment.astype(int)
np.savetxt("local.txt", local_assignment, fmt="%u")
C = np.repeat(local_assignment,100, axis=1)
C_prime = 1-C
P1 = np.sum(C, axis=1) /M
P2 = np.sum(C_prime, axis=1)/M

#simulate genotype
def sim_geno(N,M,frq1,frq2):
    num_chunk = 50
    G = np.empty(shape=(N,M))
    for i in range(num_chunk):
        N2 = local_assignment[:,i].sum()
        N1 = N - N2
        for snp in range(100):
            idx = i*100+snp
            sas_frq_cur = frq1[idx]
            irish_frq_cur = frq2[idx]
            sas_sim = np.random.binomial(size=N1, n=2, p=sas_frq_cur)
            irish_sim = np.random.binomial(size=N2, n=2, p=irish_frq_cur)
            G[np.where(local_assignment[:,i]==0),idx ] = sas_sim
            G[np.where(local_assignment[:,i]==1), idx] = irish_sim
    return G

G_sas = sim_geno(N,M,sas_frq, wb_frq)
G_irish = sim_geno(N,M,irish_frq, wb_frq)
#simulation stats sanity check 
#print(G)
#np.mean(G, axis=0)/2
G_sas = G_sas.astype(int)
G_irish = G_irish.astype(int)
np.savetxt("Geno_WB_sas.txt", G_sas, fmt="%u" )
np.savetxt("Geno_WB_irish.txt", G_irish, fmt="%u" )


#sim phenotype 
#N1 is == 0 
#N2 is assignment == 1 
#in sample population specific maf and mean 
def calculate_stats(G,local_assignment):
    mean_hat1 = np.empty(shape=(M,1))
    mean_hat2 = np.empty(shape=(M,1))
    sigma_hat1 = np.empty(shape=(M,1))
    sigma_hat2 = np.empty(shape=(M,1))
    for i in range(50):
        snp_idx = i*100
        N2 = local_assignment[:,i].sum()
        N1 = N - N2
        G1 = G[local_assignment[:,i]==0, snp_idx:(snp_idx+100)] 
        G2 = G[local_assignment[:,i]==1, snp_idx:(snp_idx+100)] 
        mean_hat1[snp_idx:(snp_idx+100),0] = np.mean(G1 , axis=0)
        mean_hat2[snp_idx:(snp_idx+100),0] = np.mean(G2, axis=0)
        sigma_hat1[snp_idx:(snp_idx+100),0] = 1/np.std(G1 , axis=0)
        sigma_hat2[snp_idx:(snp_idx+100),0] = 1/np.std(G2, axis=0)
    return mean_hat1,mean_hat2,sigma_hat1,sigma_hat2

def sim_beta_pheno(pop,G,M,N,arch,local_assignment,C,C_prime):
    sims = np.empty(shape=(4,100, N))
    mean_hat1,mean_hat2,sigma_hat1, sigma_hat2 = calculate_stats(G,local_assignment)
    for t in range(4):
        snp_h2g1= 0.2/M
        snp_h2g2 = 0.3/M
        rho_g = arch[t]
        gamma_g = rho_g * np.sqrt(snp_h2g1)* np.sqrt(snp_h2g2) 
        cov = [[snp_h2g1, gamma_g], [gamma_g, snp_h2g2]]
        mean = [0,0]
        #simulate 100 
        beta = np.random.multivariate_normal(mean, cov, (100,M))
        true_coeff = np.empty(shape=(100,1))
        for i in range(100):
            filename = "./sims/"+pop+"/"+str(rho_g)+"/beta_"+str(i)+".txt"
            true_coeff[i,0] = np.diag(np.corrcoef(beta[i,:,0], beta[i,:,1]), k=1)[0]
            np.savetxt(filename, beta[i,:,:],fmt="%f")
        filename="./sims/"+pop+"/"+str(rho_g)+"/true_rg.txt"
        np.savetxt(filename, true_coeff, fmt="%f")
        #simulate pheno 
        
        pheno=np.zeros(shape= (100,N))
        for i in range(100):
            for j in range(N): 
                temp1 = np.multiply((G[j,:] - mean_hat1.reshape(1,M)),sigma_hat1.reshape(1,M))
                temp1 = np.multiply(temp1, C_prime[j,:])
                temp2 = np.multiply((G[j,:] - mean_hat2.reshape(1,M)),sigma_hat2.reshape(1,M))
                temp2 = np.multiply(temp2, C[j,:])
                temp = temp1.dot(beta[i,:,0].reshape(M,1)) + temp2.dot(beta[i,:,1].reshape(M,1))
                pheno[i,j] = temp 
        e = np.empty(shape=(100,N))
        for i in range(100):
            e[i,:]=np.multiply(P1, np.random.normal(0, 0.8, N)) + np.multiply(P2,np.random.normal(0, 0.7, N))
            filename = "./sims/"+pop+"/"+str(rho_g)+ "/pheno_"+str(i)+".txt"
            pheno[i,] = pheno[i,]+e[i,]
            np.savetxt(filename, np.transpose(pheno[i,]+ e[i,]),fmt="%f" )
        sims[t] = pheno
    return sims
arch_rg=[0, 0.2, 0.5, 0.8]

pheno_sas = sim_beta_pheno("sas",G_sas, M,N,arch_rg, local_assignment, C,C_prime)
pheno_irish = sim_beta_pheno("irish", G_irish, M,N,arch_rg, local_assignment, C,C_prime)
#estimation
def reconstruct_geno(G, C,C_prime):
    mean_hat1,mean_hat2,sigma_hat1, sigma_hat2 = calculate_stats(G,local_assignment)
    X1 = np.multiply((G - mean_hat1.reshape(1,M)).dot(np.diag(sigma_hat1.reshape(M,))),C_prime)
    X2 = np.multiply((G - mean_hat2.reshape(1,M)).dot(np.diag(sigma_hat2.reshape(M,))), C)
    return X1,X2
X1_sas,X2_sas = reconstruct_geno(G_sas, C, C_prime)
X1_irish,X2_irish = reconstruct_geno(G_irish, C,C_prime)
#exact MoM, A matrix 
def construct_K(X1, X2,M):
    K12 = X1.dot(np.transpose(X2))/M
    K11 = X1.dot(np.transpose(X1))/M
    K22 = X2.dot(np.transpose(X2))/M
    return K12,K11,K22
def construct_A(K12,K11,K22, P1, P2,M):
    A = np.empty(shape=(5,5))
    A[0,0] = np.trace(K12.dot(K12.transpose()))
    A[0,1] = np.trace(K12.dot(K11))
    A[0,2] = np.trace(K12.dot(K22))
    A[0,3]  = np.trace(K12.dot(np.diag(P1)))
    A[0,4] = np.trace(K12.dot(np.diag(P2)))

    A[1,0] = A[0,1]
    A[1,1] = np.trace(K11.dot(K11))
    A[1,2] = np.trace(K11.dot(K22))
    A[1,3]= np.trace(K11.dot(np.diag(P1)))
    A[1,4]= np.trace(K11.dot(np.diag(P2)))

    A[2,0] = A[0,2]
    A[2,1] = A[1,2]
    A[2,2] = np.trace(K22.dot(K22))
    A[2,3]= np.trace(K22.dot(np.diag(P1)))
    A[2,4] = np.trace(K22.dot(np.diag(P2)))

    A[3,0] = np.trace(K12.dot(np.diag(P1)))
    A[3,1] = np.trace(K11.dot(np.diag(P1)))
    A[3,2] = np.trace(K22.dot(np.diag(P1)))
    A[3,3] = np.sum(np.multiply(P1,P1))
    A[3,4] = np.sum(np.multiply(P1,P2))

    A[4,0] = np.trace(K12.dot(np.diag(P2)))
    A[4,1] = np.trace(K11.dot(np.diag(P2)))
    A[4,2] = np.trace(K22.dot(np.diag(P2)))
    A[4,3] = np.sum(np.multiply(P1,P2))
    A[4,4] = np.sum(np.multiply(P2,P2))
    return A

K12_sas,K11_sas,K22_sas = construct_K(X1_sas,X2_sas, M)
K12_irish,K11_irish,K22_irish = construct_K(X1_irish, X2_irish, M)

A_sas = construct_A(K12_sas,K11_sas,K22_sas ,P1,P2,M)
A_irish = construct_A(K12_irish, K11_irish,K22_irish, P1, P2, M)
np.savetxt("A_sas.txt", A_sas,fmt="%f" )
np.savetxt("A_irish.txt", A_irish, fmt="%f")
print(A_sas)
print(A_irish)
print(A)
def construct_b(K12,K11,K22, P1,P2, M,y):
    b=np.empty(shape=(5,1))

    b[0,0] = np.trace(np.transpose(y).dot(K12.dot(y))) 
    b[1,0] = np.trace(np.transpose(y).dot(K11.dot(y)))
    b[2,0] = np.trace(np.transpose(y).dot(K22.dot(y)))
    b[3,0] = np.transpose(y).dot(np.multiply(P1.reshape(N,1), y))
    b[4,0] = np.transpose(y).dot(np.multiply(P2.reshape(N,1), y))
    return b 


result_sas = np.empty(shape=(4,100))
result_irish = np.empty(shape=(4,100))

for i in range(4):
    pheno = pheno_sas[i]
    for t in range(100):
        if(t%50 ==0):
            print(i,t)
        y = pheno[t,:].transpose().reshape(N,1)
        y = y-np.mean(y)
        y = y/ np.std(y)
        b = construct_b(K12_sas,K11_sas,K22_sas,P1,P2, M, y)
        x =np.linalg.solve(A_sas, b)
        result_sas[i,t]= x[0]/np.sqrt(x[1])/ np.sqrt(x[2])

        
        
for i in range(4):
    pheno = pheno_irish[i]
    for t in range(100):
        if(t%50 ==0):
            print(i,t)
        y = pheno[t,:].transpose().reshape(N,1)
        y = y-np.mean(y)
        y = y/ np.std(y)
        b = construct_b(K12_irish,K11_irish,K22_irish,P1,P2, M, y)
        x =np.linalg.solve(A_irish, b)
        result_irish[i,t]= x[0]/np.sqrt(x[1])/ np.sqrt(x[2])

print(np.mean(result_sas,axis=1))
print(np.std(result_sas, axis=1))
print(np.mean(result_irish, axis=1))
print(np.std(result_irish, axis=1))
#visualize
import pandas as pd
pd_sas = pd.DataFrame(result_sas.transpose())
pd_irish = pd.DataFrame(result_irish.transpose())
pd_sas.columns = ["0","0.2","0.5","0.8"]
pd_irish.columns = ["0","0.2","0.5","0.8"]
print(pd_sas)
import matplotlib.pyplot as plt

plot1=pd_sas.plot(kind='box', title="SAS vs. White British")
plot1.set_xlabel("rg")
plot1.axhline(y=0, color='r', linestyle='--')
plot1.axhline(y=0.2, color='r', linestyle='--')
plot1.axhline(y=0.5, color='r', linestyle='--')
plot1.axhline(y=0.8, color='r', linestyle='--')


plt.savefig("SAS.png")
plot2=pd_irish.plot(kind='box' , title="Irish vs. White British")
plot2.set_xlabel("rg")
plot2.axhline(y=0, color='r', linestyle='--')
plot2.axhline(y=0.2, color='r', linestyle='--')
plot2.axhline(y=0.5, color='r', linestyle='--')
plot2.axhline(y=0.8, color='r', linestyle='--')
plt.savefig("Irish.png")<jupyter_output>0       0.2       0.5       0.8
0  -0.017261  0.287470  0.484131  0.910238
1   0.061375  0.130979  0.602139  0.778789
2   0.042696  0.286032  0.463373  0.869830
3   0.054035  0.098172  0.622906  0.998762
4   0.022520  0.253287  0.557028  0.975334
..       ...       ...       ...       ...
95 -0.032808  0.302131  0.460798  0.880148
96 -0.039601  0.207648  0.655785  0.803323
97 -0.069687  0.059003  0.711403  0.851855
98 -0.021666  0.222264  0.454218  0.861964
99 -0.095664  0.146033  0.418604  1.076518

[100 rows x 4 columns]<jupyter_text>Transethnic setting: use same solution, where for each sample the snps are all from one population<jupyter_code>G_trans = np.empty(shape=(N,M))
for i in range(M):
    sas_frq_cur = sas_frq[i]
    irish_frq_cur = irish_frq[i]
    sas_sim = np.random.binomial(size=N1, n=2, p=sas_frq_cur)
    irish_sim = np.random.binomial(size=N2, n=2, p=irish_frq_cur)
    G_trans[0:N1,i] = sas_sim
    G_trans[N1:N,i] = irish_sim

maf_hat1_trans = np.mean(G_trans[0:N1,:], axis=0)/2
maf_hat2_trans = np.mean(G_trans[N1:N,:], axis=0)/2
mean_hat1_trans= np.mean(G_trans[0:N1,:],axis=0)
mean_hat2_trans = np.mean(G_trans[N1:N,:],axis=0)
sigma_hat1_trans =1/ np.std(G_trans[0:N1,:],axis=0)
sigma_hat2_trans = 1/np.std(G_trans[N1:N,:],axis=0)

#beta use the same simulations
e_trans = np.empty(shape=(100,N))
for t in range(100):
    e_trans[t,0:N1] = np.random.normal(0,0.8,N1).reshape(1,N1)
    e_trans[t,N1:N] = np.random.normal(0,0.7,N2).reshape(1,N2)


pheno_trans = np.zeros(shape=(100, N))

for t in range(100):
    if(t % 50 ==0):
        print(t)
    for i in range(N1):
       
        temp1 = np.multiply((G_trans[i,:] - mean_hat1_trans.reshape(1,M)),sigma_hat1_trans.reshape(1,M))
        temp = temp1.dot(beta[t,:,0].reshape(M,1))
        pheno_trans[t,i] = temp 
    for i in range(N1, N): 
        temp1 = np.multiply((G_trans[i,:] - mean_hat2_trans.reshape(1,M)),sigma_hat2_trans.reshape(1,M))
        temp = temp1.dot(beta[t,:,1].reshape(M,1))
        pheno_trans[t,i] = temp 
y = pheno_trans+e_trans 

C_trans = np.row_stack((np.ones(shape=(N1,M)), np.zeros(shape=(N2, M))))
C_prime_trans = 1-C_trans
P1_trans = np.sum(C_trans, axis=1) /M
P2_trans = np.sum(C_prime_trans, axis=1)/M
X1_trans = np.multiply((G_trans - mean_hat1_trans.reshape(1,M)).dot(np.diag(sigma_hat1_trans)),C_trans)
X2_trans = np.multiply((G_trans - mean_hat2_trans.reshape(1,M)).dot(np.diag(sigma_hat2_trans)), C_prime_trans)
C_trans.shape
#exact MoM

K12 = X1_trans.dot(np.transpose(X2_trans))/M
K11 = X1_trans.dot(np.transpose(X1_trans))/M
K22 = X2_trans.dot(np.transpose(X2_trans))/M

A = np.empty(shape=(5,5))
A[0,0] = np.trace(K12.dot(K12.transpose()))
A[0,1] = np.trace(K12.dot(K11))
A[0,2] = np.trace(K12.dot(K22))
A[0,3]  = np.trace(K12.dot(np.diag(P1_trans)))
A[0,4] = np.trace(K12.dot(np.diag(P2_trans)))

A[1,0] = A[0,1]
A[1,1] = np.trace(K11.dot(K11))
A[1,2] = np.trace(K11.dot(K22))
A[1,3]= np.trace(K11.dot(np.diag(P1_trans)))
A[1,4]= np.trace(K11.dot(np.diag(P2_trans)))

A[2,0] = A[0,2]
A[2,1] = A[1,2]
A[2,2] = np.trace(K22.dot(K22))
A[2,3]= np.trace(K22.dot(np.diag(P1_trans)))
A[2,4] = np.trace(K22.dot(np.diag(P2_trans)))

A[3,0] = np.trace(K12.dot(np.diag(P1_trans)))
A[3,1] = np.trace(K11.dot(np.diag(P1_trans)))
A[3,2] = np.trace(K22.dot(np.diag(P1_trans)))
A[3,3] = np.sum(np.multiply(P1_trans,P1_trans))
A[3,4] = np.sum(np.multiply(P1_trans,P2_trans))

A[4,0] = np.trace(K12.dot(np.diag(P2_trans)))
A[4,1] = np.trace(K11.dot(np.diag(P2_trans)))
A[4,2] = np.trace(K22.dot(np.diag(P2_trans)))
A[4,3] = np.sum(np.multiply(P1_trans,P2_trans))
A[4,4] = np.sum(np.multiply(P2_trans,P2_trans))

result= np.empty(shape=(100,1))
for t in range(100):
    y_cur = y[t,:].reshape(N,1)
    y_cur =y_cur - np.mean(y_cur)
    y_cur = y_cur/ np.std(y_cur)
    b = np.empty(shape=(5,1))
    b[0,0] = np.trace(np.transpose(y_cur).dot(K12.dot(y_cur))) 
    b[1,0] = np.trace(np.transpose(y_cur).dot(K11.dot(y_cur)))
    b[2,0] = np.trace(np.transpose(y_cur).dot(K22.dot(y_cur)))
    b[3,0] = np.transpose(y_cur).dot(np.multiply(P1_trans.reshape(N,1), y_cur))
    b[4,0] = np.transpose(y_cur).dot(np.multiply(P2_trans.reshape(N,1), y_cur))
    
    x=np.linalg.solve(A,b)
    result[t,0] = x[0]/np.sqrt(x[1])/ np.sqrt(x[2])

print(np.mean(result))
print(np.std(result))<jupyter_output>0.8248982501878389
0.15238340355450353
{'fraction_non_alphanumeric': 0.19796879599646747, 'fraction_numerical': 0.08581100971445393, 'mean_word_length': 3.7381450488145047, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '337341', 'license': 'MIT', 'n_tokens_mistral': 6976, 'n_tokens_neox': 6512, 'n_words': 900, 'path': 'Mixture.ipynb', 'repo_name': 'ariel-wu/SCORE_helper', 'size': '13588'}