Skip to content

Instantly share code, notes, and snippets.

@endrebak
Created February 6, 2020 14:35
Show Gist options
  • Save endrebak/ed04633cd949bcbcb9eb01d0579c956c to your computer and use it in GitHub Desktop.
Save endrebak/ed04633cd949bcbcb9eb01d0579c956c to your computer and use it in GitHub Desktop.
#!/usr/bin/env python3
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True
import sys, math, gzip
import numpy as np
import pandas as pd
from time import time
from libc.math cimport exp, fabs
from libc.stdint cimport int32_t, int8_t
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
@cython.cdivision(True)
cpdef calc_covar(haplos, double [::1] autocovars, double theta, int32_t window_size):
cdef:
int i, j, j1, j2, k, pos1, pos2, len_haps, n11, n10, n01, len_g1_int, half_window_size
double len_g1, f11, f1, f2, Ds2, D, nind, thetas, aj1, aj2, rsq
half_window_size = int(window_size / 2)
haps = haplos
records = []
len_g1_int = haps.shape[1]
len_g1 = float(haps.shape[1])
len_haps = len(haps)
assert len_haps >= 2 * half_window_size
thetas = (1-theta)*(1-theta)
# cdef long[:] allpos_view
cdef long[:] g1_view, g2_view
cdef int8_t[:, :] haps_view
cdef double[::1] outvec_view
# allpos_view = allpos.values
haps_view = haps
outvec = np.zeros(len_haps)
outvec_view = outvec
for i in range(0, half_window_size):
for j in range(i):
j1 = i - j
j2 = i + j
if j1 < 0:
break
aj1 = autocovars[j1]
aj2 = autocovars[j2]
n11, n01, n10 = 0, 0, 0
for k in range(len_g1_int):
if haps_view[j1][k] == 1 and haps_view[j2][k] == 1:
n11 += 1
elif haps_view[j1][k] == 0 and haps_view[j2][k] == 1:
n01 += 1
elif haps_view[j1][k] == 1 and haps_view[j2][k] == 0:
n10 += 1
f11 = n11/len_g1
f1 = (n11+n10)/len_g1
f2 = (n11+n01)/len_g1
D = f11 - f1*f2
Ds2 = (thetas*D)
Ds2 = Ds2 * Ds2
rsq = Ds2 / (aj1 * aj2)
if j1 != j2:
outvec_view[i] += (2 * rsq) # times two due to symmetry
else:
outvec_view[i] += rsq # times two due to symmetry
for i in range(half_window_size, len_haps - half_window_size):
for j in range(half_window_size):
j1 = i - j
j2 = i + j
aj1 = autocovars[j1]
aj2 = autocovars[j2]
n11, n01, n10 = 0, 0, 0
for k in range(len_g1_int):
if haps_view[j1][k] == 1 and haps_view[j2][k] == 1:
n11 += 1
elif haps_view[j1][k] == 0 and haps_view[j2][k] == 1:
n01 += 1
elif haps_view[j1][k] == 1 and haps_view[j2][k] == 0:
n10 += 1
f11 = n11/len_g1
f1 = (n11+n10)/len_g1
f2 = (n11+n01)/len_g1
D = f11 - f1*f2
Ds2 = (thetas*D)
Ds2 = Ds2 * Ds2
rsq = Ds2 / (aj1 * aj2)
if j1 != j2:
outvec_view[i] += (2 * rsq) # times two due to symmetry
else:
outvec_view[i] += rsq # times two due to symmetry
for i in range(len_haps - half_window_size, len_haps):
for j in range(i):
j1 = i - j
j2 = i + j
if j2 >= len_haps:
break
aj1 = autocovars[j1]
aj2 = autocovars[j2]
n11, n01, n10 = 0, 0, 0
for k in range(len_g1_int):
if haps_view[j1][k] == 1 and haps_view[j2][k] == 1:
n11 += 1
elif haps_view[j1][k] == 0 and haps_view[j2][k] == 1:
n01 += 1
elif haps_view[j1][k] == 1 and haps_view[j2][k] == 0:
n10 += 1
f11 = n11/len_g1
f1 = (n11+n10)/len_g1
f2 = (n11+n01)/len_g1
D = f11 - f1*f2
Ds2 = (thetas*D)
Ds2 = Ds2 * Ds2
rsq = Ds2 / (aj1 * aj2)
if j1 != j2:
outvec_view[i] += (2 * rsq) # times two due to symmetry
else:
outvec_view[i] += rsq # times two due to symmetry
return outvec
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment