numerics

Methods for numerics of two-locus diffusion: grids, transition matrices, etc.

LD_per_bin(ns)

Calculate LD statistics per bin for a TLSpectrum.

Parameters:

Name	Type	Description	Default
ns	int	Number of samples.	required

Returns:

Name	Type	Description
D	TLSpectrum object	TLSpectrum object in which each entry is the value of D for that combination of haplotypes.
r2	TLSpectrum object	TLSpectrum object in which each entry is the value of r^2 for that combination of haplotypes.

Source code in dadi/TwoLocus/numerics.py

def LD_per_bin(ns):
    """
    Calculate LD statistics per bin for a TLSpectrum.

    Args:
        ns (int): Number of samples.

    Returns:
        D (TLSpectrum object): TLSpectrum object in which each entry is the value of D for that combination of haplotypes.
        r2 (TLSpectrum object): TLSpectrum object in which each entry is the value of r^2 for that combination of haplotypes.
    """
    temp = np.arange(ns+1, dtype=float)/ns
    # Fancy array arithmetic, to avoid explicity for loops.
    pAB = temp[:,nuax,nuax]
    pAb = temp[nuax,:,nuax]
    paB = temp[nuax,nuax,:]
    pA = pAB + pAb
    pB = pAB + paB
    D = TLSpectrum(pAB - pA*pB)
    r2 = TLSpectrum(D**2/(pA*(1-pA)*pB*(1-pB)))

    return D,r2

advance1D(u, P)

Given transition matrix P, use dadi's triagonal solver to integrate in 1D.

Parameters:

Name	Type	Description	Default
u	array - like	1D density function.	required
P	array - like	Transition matrix.	required

Returns:

Name	Type	Description
u	ndarray	Updated 1D density function.

Source code in dadi/TwoLocus/numerics.py

def advance1D(u,P):
    """
    Given transition matrix P, use dadi's triagonal solver to integrate in 1D.

    Args:
        u (array-like): 1D density function.
        P (array-like): Transition matrix.

    Returns:
        u (ndarray): Updated 1D density function.
    """
    a = np.concatenate((np.array([0]),np.diag(P,-1)))
    b = np.diag(P)
    c = np.concatenate((np.diag(P,1),np.array([0])))
    u = dadi.tridiag.tridiag(a,b,c,u)
    return u

advance_adi(phi, U01, P1, P2, P3, x, ii)

Combined ADI integration of phi.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
U01	array - like	Domain markers.	required
P1	array - like	ADI transition matrices for axis 1.	required
P2	array - like	ADI transition matrices for axis 2.	required
P3	array - like	ADI transition matrices for axis 3.	required
x	array - like	1D grid.	required
ii	int	Current iteration.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_adi(phi,U01,P1,P2,P3,x,ii):
    """
    Combined ADI integration of phi.

    Args:
        phi (array-like): Density function.
        U01 (array-like): Domain markers.
        P1 (array-like): ADI transition matrices for axis 1.
        P2 (array-like): ADI transition matrices for axis 2.
        P3 (array-like): ADI transition matrices for axis 3.
        x (array-like): 1D grid.
        ii (int): Current iteration.

    Returns:
        phi (ndarray): Updated density function.
    """
    if np.mod(ii,3) == 0:
        order = [1,2,3]
    elif np.mod(ii,3) == 1:
        order = [2,3,1]
    else:
        order = [3,1,2]
    for ord in order:
        if ord == 1:
            phi = advance_adi1(phi,U01,P1,x)
        elif ord == 2:
            phi = advance_adi2(phi,U01,P2,x)
        elif ord == 3:
            phi = advance_adi3(phi,U01,P3,x)

    return phi

advance_adi1(phi, U01, P1, x)

ADI integration along axis 1 of phi.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
U01	array - like	Domain markers.	required
P1	array - like	ADI transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_adi1(phi,U01,P1,x):
    """
    ADI integration along axis 1 of phi.

    Args:
        phi (array-like): Density function.
        U01 (array-like): Domain markers.
        P1 (array-like): ADI transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to 
    #        this version and the cython version together.
    for jj in range(len(x)):
        for kk in range(len(x)):
            if np.sum(U01[:,jj,kk]) > 1:
                phi[:,jj,kk] = dadi.tridiag.tridiag(P1[jj,kk,0,:],P1[jj,kk,1,:],P1[jj,kk,2,:],phi[:,jj,kk])
    return phi

advance_adi2(phi, U01, P2, x)

ADI integration along axis 2 of phi.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
U01	array - like	Domain markers.	required
P2	array - like	ADI transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_adi2(phi,U01,P2,x):
    """
    ADI integration along axis 2 of phi.

    Args:
        phi (array-like): Density function.
        U01 (array-like): Domain markers.
        P2 (array-like): ADI transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to 
    #        this version and the cython version together.
    for ii in range(len(x)):
        for kk in range(len(x)):
            if np.sum(U01[ii,:,kk]) > 1:
                phi[ii,:,kk] = dadi.tridiag.tridiag(P2[ii,kk,0,:],P2[ii,kk,1,:],P2[ii,kk,2,:],phi[ii,:,kk])
    return phi

advance_adi3(phi, U01, P3, x)

ADI integration along axis 3 of phi.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
U01	array - like	Domain markers.	required
P3	array - like	ADI transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_adi3(phi,U01,P3,x):
    """
    ADI integration along axis 3 of phi.

    Args:
        phi (array-like): Density function.
        U01 (array-like): Domain markers.
        P3 (array-like): ADI transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to 
    #        this version and the cython version together.
    for ii in range(len(x)):
        for jj in range(len(x)):
            if np.sum(U01[ii,jj,:]) > 1:
                phi[ii,jj,:] = dadi.tridiag.tridiag(P3[ii,jj,0,:],P3[ii,jj,1,:],P3[ii,jj,2,:],phi[ii,jj,:])
    return phi

advance_cov(phi, C12, C13, C23, x, ii)

Combined integration for the covariance terms.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
C12	array - like	Transition matrices for axes 1 and 2.	required
C13	array - like	Transition matrices for axes 1 and 3.	required
C23	array - like	Transition matrices for axes 2 and 3.	required
x	array - like	1D grid.	required
ii	int	Current iteration.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_cov(phi,C12,C13,C23,x,ii):
    """
    Combined integration for the covariance terms.

    Args:
        phi (array-like): Density function.
        C12 (array-like): Transition matrices for axes 1 and 2.
        C13 (array-like): Transition matrices for axes 1 and 3.
        C23 (array-like): Transition matrices for axes 2 and 3.
        x (array-like): 1D grid.
        ii (int): Current iteration.

    Returns:
        phi (ndarray): Updated density function.
    """
    if np.mod(ii,3) == 0:
        order = [1,2,3]
    elif np.mod(ii,3) == 1:
        order = [2,3,1]
    else:
        order = [3,1,2]
    for ord in order:
        if ord == 1:
            phi = advance_cov12(phi,C12,x)
        elif ord == 2:
            phi = advance_cov13(phi,C13,x)
        elif ord == 3:
            phi = advance_cov23(phi,C23,x)
    return phi

advance_cov12(phi, C12, x)

Explicit integration of covariance term in the 1/2 axes planes.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
C12	array - like	Transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_cov12(phi,C12,x):
    """
    Explicit integration of covariance term in the 1/2 axes planes.

    Args:
        phi (array-like): Density function.
        C12 (array-like): Transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    for kk in range(len(x)):
        C = C12[kk]
        phi[:,:,kk] = ( C * phi[:,:,kk].reshape(len(x)**2) ).reshape(len(x),len(x))
    return phi

advance_cov13(phi, C13, x)

Explicit integration of covariance term in the 1/3 axes planes.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
C13	array - like	Transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_cov13(phi,C13,x):
    """
    Explicit integration of covariance term in the 1/3 axes planes.

    Args:
        phi (array-like): Density function.
        C13 (array-like): Transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    for jj in range(len(x)):
        C = C13[jj]
        phi[:,jj,:] = ( C * phi[:,jj,:].reshape(len(x)**2) ).reshape(len(x),len(x))
    return phi

advance_cov23(phi, C23, x)

Explicit integration of covariance term in the 2/3 axes planes.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
C23	array - like	Transition matrices.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def advance_cov23(phi,C23,x):
    """
    Explicit integration of covariance term in the 2/3 axes planes.

    Args:
        phi (array-like): Density function.
        C23 (array-like): Transition matrices.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated density function.
    """
    for ii in range(len(x)):
        C = C23[ii]
        phi[ii,:,:] = ( C * phi[ii,:,:].reshape(len(x)**2) ).reshape(len(x),len(x))
    return phi

advance_surf_adi1(surf, U01surf, P1surf, x)

Advance the ADI method along first axis one of surface domain.

Parameters:

Name	Type	Description	Default
surf	array - like	Density function along that surface.	required
U01surf	array - like	Domain markers along grid.	required
P1surf	array - like	ADI transition matrix.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
surf	ndarray	Updated surface density function.

Source code in dadi/TwoLocus/numerics.py

def advance_surf_adi1(surf,U01surf,P1surf,x):
    """
    Advance the ADI method along first axis one of surface domain.

    Args:
        surf (array-like): Density function along that surface.
        U01surf (array-like): Domain markers along grid.
        P1surf (array-like): ADI transition matrix.
        x (array-like): 1D grid.

    Returns:
        surf (ndarray): Updated surface density function.
    """
    for jj in range(len(x)):
        if np.sum(U01surf[:,jj]) > 1:
            surf[:,jj] = dadi.tridiag.tridiag(P1surf[jj,0,:],P1surf[jj,1,:],P1surf[jj,2,:],surf[:,jj])
    return surf

advance_surf_adi2(surf, U01surf, P2surf, x)

Advance the ADI method along second axis one of surface domain.

Parameters:

Name	Type	Description	Default
surf	array - like	Density function along that surface.	required
U01surf	array - like	Domain markers along grid.	required
P2surf	array - like	ADI transition matrix.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
surf	ndarray	Updated surface density function.

Source code in dadi/TwoLocus/numerics.py

def advance_surf_adi2(surf,U01surf,P2surf,x):
    """
    Advance the ADI method along second axis one of surface domain.

    Args:
        surf (array-like): Density function along that surface.
        U01surf (array-like): Domain markers along grid.
        P2surf (array-like): ADI transition matrix.
        x (array-like): 1D grid.

    Returns:
        surf (ndarray): Updated surface density function.
    """
    for ii in range(len(x)):
        if np.sum(U01surf[ii,:]) > 1:
            surf[ii,:] = dadi.tridiag.tridiag(P2surf[ii,0,:],P2surf[ii,1,:],P2surf[ii,2,:],surf[ii,:])
    return surf

advance_surf_cov(surf, Csurf, x)

Explicit integration of the covariance term, using scipy's sparse matrix for Csurf.

Parameters:

Name	Type	Description	Default
surf	array - like	Density function along the surface.	required
Csurf	lil_matrix	Transition matrix.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
surf	ndarray	Updated surface density function.

Source code in dadi/TwoLocus/numerics.py

def advance_surf_cov(surf,Csurf,x):
    """
    Explicit integration of the covariance term, using scipy's sparse matrix for Csurf.

    Args:
        surf (array-like): Density function along the surface.
        Csurf (lil_matrix): Transition matrix.
        x (array-like): 1D grid.

    Returns:
        surf (ndarray): Updated surface density function.
    """
    surf = ( Csurf * surf.reshape(len(x)**2)).reshape(len(x),len(x))
    return surf

advance_surface(phi, x, P1surf, P2surf, Csurf, Pline, P, U01surf)

Advance the surface density function.

Parameters:

Name	Type	Description	Default
phi	array - like	Full density function.	required
x	array - like	1D grid.	required
P1surf	array - like	ADI transition matrices for triallele density function along axis 1.	required
P2surf	array - like	ADI transition matrices for triallele density function along axis 2.	required
Csurf	lil_matrix	Covariance transition matrix.	required
Pline	array - like	Transition matrix for diagonal boundary of triallele surface.	required
P	array - like	Amount of density to be moved from triallele domain (surf) to diagonal line boundary.	required
U01surf	array - like	Domain markers for the surface.	required

Returns:

Name	Type	Description
phi	ndarray	Updated full density function.

Source code in dadi/TwoLocus/numerics.py

def advance_surface(phi,x,P1surf,P2surf,Csurf,Pline,P,U01surf):
    """
    Advance the surface density function.

    Args:
        phi (array-like): Full density function.
        x (array-like): 1D grid.
        P1surf (array-like): ADI transition matrices for triallele density function along axis 1.
        P2surf (array-like): ADI transition matrices for triallele density function along axis 2.
        Csurf (lil_matrix): Covariance transition matrix.
        Pline (array-like): Transition matrix for diagonal boundary of triallele surface.
        P (array-like): Amount of density to be moved from triallele domain (surf) to diagonal line boundary.
        U01surf (array-like): Domain markers for the surface.

    Returns:
        phi (ndarray): Updated full density function.
    """
    # create the triallele domain of the nonsquare surface of the full density function
    surf = phi_to_surf(phi,x) # done

    surf = advance_surf_adi1(surf,U01surf,P1surf,x)
    surf = advance_surf_adi2(surf,U01surf,P2surf,x)
    surf = advance_surf_cov(surf,Csurf,x)
    surf = move_surf_density_to_bdry(x,surf,P)
    #surf = advance1D(x,surf,Pline) # this is an absorbing edge, so we'll just leave it alone - density lost!

    phi = surf_to_phi(surf,phi,x) # done
    return phi

array_to_spectrum(phi)

Convert an array to a TLSpectrum.

Parameters:

Name	Type	Description	Default
phi	array - like	Array to be converted.	required

Returns:

Name	Type	Description
phi	TLSpectrum	Converted TLSpectrum.

Source code in dadi/TwoLocus/numerics.py

def array_to_spectrum(phi):
    """
    Convert an array to a TLSpectrum.

    Args:
        phi (array-like): Array to be converted.

    Returns:
        phi (TLSpectrum): Converted TLSpectrum.
    """
    ns = len(phi)-1
    phi = dadi.Spectrum(phi)
    phi = dadi.Spectrum(phi)
    phi.mask[0,0,0] = True
    phi.mask[0,:,0] = True
    phi.mask[0,0,:] = True
    for ii in range(len(phi)):
        for jj in range(len(phi)):
            for kk in range(len(phi)):
                if ii+jj+kk > ns:
                    phi.mask[ii,jj,kk] = True

    for ii in range(len(phi)):
        phi.mask[ii,ns-ii,0] = True
        phi.mask[ii,0,ns-ii] = True

    return phi

binom(n, k)

Compute the binomial coefficient.

Parameters:

Name	Type	Description	Default
n	int	Number of trials.	required
k	int	Number of successes.	required

Returns:

Name	Type	Description
int	int	Binomial coefficient.

Source code in dadi/TwoLocus/numerics.py

def binom(n,k):
    """
    Compute the binomial coefficient.

    Args:
        n (int): Number of trials.
        k (int): Number of successes.

    Returns:
        int (int): Binomial coefficient.
    """
    return math.factorial(n)/math.factorial(k)/math.factorial(n-k)

cached_genotype_exact_projection(n, haplotype_counts)

Cache the probabilities of observing specific sets of genotypes for a given set of haplotype counts.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
haplotype_counts	tuple	Number of individuals with each haplotype.	required

Returns:

Name	Type	Description
prob_cache	dict	Dictionary of probabilities for each set of genotypes.

Source code in dadi/TwoLocus/numerics.py

def cached_genotype_exact_projection(n,haplotype_counts):
    """
    Cache the probabilities of observing specific sets of genotypes for a given set of haplotype counts.

    Args:
        n (int): Number of individuals.
        haplotype_counts (tuple): Number of individuals with each haplotype.

    Returns:
        prob_cache (dict): Dictionary of probabilities for each set of genotypes.
    """
    key = (n,haplotype_counts)
    try:
        return prob_cache[key]
    except KeyError:
        pass

    gl = possible_genotypes_4(n,haplotype_counts)
    tot = float(pairings(n))
    probs = []
    prob_cache.setdefault(key,{})
    for hits in gl:
        prob_cache[key][hits] = float(genotypes_prob_4(n,haplotype_counts,hits))/tot
    return prob_cache[key]

cached_projection(proj_to, proj_from, hits)

Coefficients for projection from a larger size to smaller.

Parameters:

Name	Type	Description	Default
proj_to	int	Number of samples to project down to.	required
proj_from	int	Number of samples to project from.	required
hits	tuple	Number of derived alleles projecting from - tuple of (n1,n2,n3).	required

Returns:

Name	Type	Description
proj_weights	ndarray	Projection coefficients.

Source code in dadi/TwoLocus/numerics.py

def cached_projection(proj_to, proj_from, hits):
    """
    Coefficients for projection from a larger size to smaller.

    Args:
        proj_to (int): Number of samples to project down to.
        proj_from (int): Number of samples to project from.
        hits (tuple): Number of derived alleles projecting from - tuple of (n1,n2,n3).

    Returns:
        proj_weights (ndarray): Projection coefficients.
    """
#    key = (proj_to, proj_from, hits)
#    try:
#        return projection_cache[key]
#    except KeyError:
#        pass

    X1, X2, X3 = hits
    X4 = proj_from - X1 - X2 - X3
    proj_weights = np.zeros((proj_to+1,proj_to+1,proj_to+1))
    for ii in range(X1+1):
        for jj in range(X2+1):
            for kk in range(X3+1):
                ll = proj_to - ii - jj - kk
                if ll > X4 or ll <0:
                    continue
                f = ln_binomial(X1,ii) + ln_binomial(X2,jj) + ln_binomial(X3,kk) + ln_binomial(X4,ll) - ln_binomial(proj_from,proj_to)
                proj_weights[ii,jj,kk] = np.exp(f)


#    projection_cache[key] = proj_weights
    return proj_weights

condition(F, i)

Condition on seeing nA = i in a two-locus spectrum.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required
i	int	Number of derived alleles for the first locus.	required

Returns:

Name	Type	Description
Fcond	ndarray	Conditioned two-locus spectrum.

condition on seeing nA = i Fcond[i,j]: i - number AB, j - number of aB i ranges from 0 to i, j ranges from 0 to n-i

Source code in dadi/TwoLocus/numerics.py

def condition(F,i):
    """
    Condition on seeing nA = i in a two-locus spectrum.

    Args:
        F (TLSpectrum): Two-locus spectrum.
        i (int): Number of derived alleles for the first locus.

    Returns:
        Fcond (ndarray): Conditioned two-locus spectrum.

    condition on seeing nA = i
    Fcond[i,j]: i - number AB, j - number of aB
    i ranges from 0 to i, j ranges from 0 to n-i
    """
    ns = len(F)-1
    Fcond = np.zeros((i+1,ns-i+1))
    for ii in range(ns):
        for jj in range(ns):
            for kk in range(ns):
                if F.mask[ii,jj,kk] == False:
                    if ii+jj == i:
                        Fcond[ii,kk] = F[ii,jj,kk]
    return Fcond

domain(x)

Create an array indicating whether grid points lie inside the domain.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required

Returns:

Name	Type	Description
U01	ndarray	Array with 1 for points inside the domain and 0 for points outside.

Source code in dadi/TwoLocus/numerics.py

def domain(x):
    """
    Create an array indicating whether grid points lie inside the domain.

    Args:
        x (array-like): 1D grid.

    Returns:
        U01 (ndarray): Array with 1 for points inside the domain and 0 for points outside.
    """
    tol = 1e-12
    U01 = np.ones((len(x),len(x),len(x)))
    XX = x[:,nuax,nuax] + x[nuax,:,nuax] + x[nuax,nuax,:]
    U01[np.where(XX > 1+tol)] = 0
    return U01

domain_surf(x)

Constructs a matrix with the same dimension as the density function discretization, a 1 indicates that the corresponding point is inside the triangular domain or on the boundary, while a 0 indicates that point falls outside the domain.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required

Returns:

Name	Type	Description
U01	ndarray	Domain markers for the surface.

Source code in dadi/TwoLocus/numerics.py

def domain_surf(x):
    """
    Constructs a matrix with the same dimension as the density function discretization, a 1 indicates that the corresponding point is inside the triangular domain or on the boundary, while a 0 indicates that point falls outside the domain.

    Args:
        x (array-like): 1D grid.

    Returns:
        U01 (ndarray): Domain markers for the surface.
    """
    tol = 1e-12
    U01 = np.ones((len(x),len(x)))
    XX = x[:,nuax] + x[nuax,:]
    U01[np.where(XX > 1+tol)] = 0
    return U01

extrap_dt_pts(temps)

Extrapolate data points for different time steps and grid points.

Parameters:

Name	Type	Description	Default
temps	dict	Dictionary of data points for different time steps and grid points.	required

Returns:

Name	Type	Description
ndarray	ndarray	Extrapolated data points.

Source code in dadi/TwoLocus/numerics.py

def extrap_dt_pts(temps):
    """
    Extrapolate data points for different time steps and grid points.

    Args:
        temps (dict): Dictionary of data points for different time steps and grid points.

    Returns:
        ndarray (ndarray): Extrapolated data points.
    """
    # in form of temps[dt][numpts]
    dts = sorted(temps.keys())[::-1]
    gridpts = sorted(temps[dts[0]].keys())
    F0 = dadi.Numerics.quadratic_extrap((temps[dts[0]][gridpts[0]],temps[dts[0]][gridpts[1]],temps[dts[0]][gridpts[2]]),(1./gridpts[0],1./gridpts[1],1./gridpts[2]))
    F1 = dadi.Numerics.quadratic_extrap((temps[dts[1]][gridpts[0]],temps[dts[1]][gridpts[1]],temps[dts[1]][gridpts[2]]),(1./gridpts[0],1./gridpts[1],1./gridpts[2]))
    F2 = dadi.Numerics.quadratic_extrap((temps[dts[2]][gridpts[0]],temps[dts[2]][gridpts[1]],temps[dts[2]][gridpts[2]]),(1./gridpts[0],1./gridpts[1],1./gridpts[2]))

    return dadi.Numerics.quadratic_extrap((F0,F1,F2),(dts[0],dts[1],dts[2]))

fold_ancestral(F)

Fold an unfolded two-locus frequency spectrum (assume don't know ancestral state).

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Unfolded two-locus frequency spectrum.	required

Returns:

Name	Type	Description
F	TLSpectrum	Folded two-locus frequency spectrum.

Source code in dadi/TwoLocus/numerics.py

def fold_ancestral(F):
    """
    Fold an unfolded two-locus frequency spectrum (assume don't know ancestral state).

    Args:
        F (TLSpectrum): Unfolded two-locus frequency spectrum.

    Returns:
        F (TLSpectrum): Folded two-locus frequency spectrum.
    """
    ns = len(F[:,0,0]) - 1
    for ii in range(ns+1):
        for jj in range(ns+1):
            for kk in range(ns+1):
                if F.mask[ii,jj,kk]:
                    continue
                p = ii + jj
                q = ii + kk
                if p > ns/2 and q > ns/2:
                    # Switch A/a and B/b, so AB becomes ab, Ab becomes aB, etc
                    F[ns-ii-jj-kk,kk,jj] += F[ii,jj,kk]
                    F.mask[ii,jj,kk] = True
                elif p > ns/2:
                    # Switch A/a, so AB -> aB, Ab -> ab, aB -> AB, and ab -> Ab
                    F[kk,ns-ii-jj-kk,ii] += F[ii,jj,kk]
                    F.mask[ii,jj,kk] = True
                elif q > ns/2:
                    # Switch B/b, so AB -> Ab, Ab -> AB, aB -> ab, and ab -> aB
                    F[jj,ii,ns-ii-jj-kk] += F[ii,jj,kk]
                    F.mask[ii,jj,kk] = True

    return F

fold_lr(F)

Fold an unfolded two-locus spectrum based on left/right allele.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Unfolded two-locus frequency spectrum.	required

Returns:

Name	Type	Description
F	TLSpectrum	Folded two-locus frequency spectrum.

Source code in dadi/TwoLocus/numerics.py

def fold_lr(F):
    """
    Fold an unfolded two-locus spectrum based on left/right allele.

    Args:
        F (TLSpectrum): Unfolded two-locus frequency spectrum.

    Returns:
        F (TLSpectrum): Folded two-locus frequency spectrum.
    """
    ns = len(F[:,0,0]) - 1
    for ii in range(ns+1):
        for jj in range(ns+1):
            for kk in range(ns+1):
                if F.mask[ii,jj,kk]:
                    continue
                if kk > jj:
                    F[ii,kk,jj] += F[ii,jj,kk]
                    F[ii,jj,kk] = 0
                    F.mask[ii,jj,kk] = True
    return F

genotype_exp_data_to_arrays(G_exp, G_data)

Convert dictionaries of expected and observed genotype spectra to arrays.

Parameters:

Name	Type	Description	Default
G_exp	dict	Dictionary of expected genotype spectra.	required
G_data	dict	Dictionary of observed genotype spectra.	required

Returns:

Name	Type	Description
e	array	Array of expected and observed genotype spectra.
d	array	Array of expected and observed genotype spectra.

return 1D spectra, to perform likelihood calcs so given two dicts, one for genotype freq expectations, one for data, returns spectrum arrays of data that can be passed to inference methods

Source code in dadi/TwoLocus/numerics.py

def genotype_exp_data_to_arrays(G_exp,G_data):
    """
    Convert dictionaries of expected and observed genotype spectra to arrays.

    Args:
        G_exp (dict): Dictionary of expected genotype spectra.
        G_data (dict): Dictionary of observed genotype spectra.

    Returns:
        e (array): Array of expected and observed genotype spectra.
        d (array): Array of expected and observed genotype spectra.

    return 1D spectra, to perform likelihood calcs
    so given two dicts, one for genotype freq expectations, one for data, returns spectrum arrays of data that can be passed to inference methods
    """
    n_exp = G_exp.keys()[0]
    n_data = G_data.keys()[0]
    if n_exp != n_data:
        return 'must have same number of sampled individuals'
    exp_arr = np.zeros(len(G_exp[n_exp]))
    data_arr = np.zeros(len(G_exp[n_exp]))
    gl = G_exp[n_exp].keys()
    for ii in range(len(gl)):
        g = gl[ii]
        exp_arr[ii] = G_exp[n_exp][g]
        try:
            data_arr[ii] = G_data[n_exp][g]
        except KeyError:
            pass

    e = dadi.Spectrum(np.concatenate((np.array([0]),exp_arr,np.array([0]))))
    d = dadi.Spectrum(np.concatenate((np.array([0]),data_arr,np.array([0]))))
    return e,d

genotype_spectrum_from_F(F)

Convert a two-locus spectrum to a genotype spectrum.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required

Returns:

Name	Type	Description
G	ndarray	Genotype spectrum.

Source code in dadi/TwoLocus/numerics.py

def genotype_spectrum_from_F(F):
    """
    Convert a two-locus spectrum to a genotype spectrum.

    Args:
        F (TLSpectrum): Two-locus spectrum.

    Returns:
        G (ndarray): Genotype spectrum.
    """
    n = len(F)-1
    ng = n/2
    tot = pairings(n)
    G = np.zeros((ng+1,ng+1,ng+1,ng+1,ng+1,ng+1,ng+1,ng+1,ng+1))
    for ii in range(n+1):
        for jj in range(n+1):
            for kk in range(n+1):
                if F[ii,jj,kk] > 0:# and F.mask[ii,jj,kk] == False:
                    nAB,nAb,naB,nab = ii,jj,kk,n-ii-jj-kk
                    gen_probs = cached_genotype_exact_projection(n,(nAB,nAb,naB,nab))
                    for gens in gen_probs.keys():
                        g1,g2,g3,g4,g5,g6,g7,g8,g9,g10 = gens
                        prob = gen_probs[gens]
                        G[g1,g2,g3,g4,g5,g6,g7,g8,g9] += F[nAB,nAb,naB]*prob
    return G

genotypes_prob_4(n, colors, hits)

Compute the probability of observing a specific set of genotypes.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
colors	tuple	Number of individuals with each haplotype.	required
hits	tuple	Number of individuals with each genotype.	required

Returns:

Name	Type	Description
float	float	Probability of observing the specific set of genotypes.

Source code in dadi/TwoLocus/numerics.py

def genotypes_prob_4(n,colors,hits):
    """
    Compute the probability of observing a specific set of genotypes.

    Args:
        n (int): Number of individuals.
        colors (tuple): Number of individuals with each haplotype.
        hits (tuple): Number of individuals with each genotype.

    Returns:
        float (float): Probability of observing the specific set of genotypes.
    """
    (nR,nG,nB,nY) = colors
    (nRR,nGG,nBB,nYY,nRG,nRB,nRY,nGB,nGY,nBY) = hits
    if sum(hits) != n/2 or sum(colors) != n: return "nuh uh"
    return pairings(nR-nRG-nRB-nRY)*pairings(nG-nRG-nGB-nGY)*pairings(nB-nRB-nGB-nBY)*pairings(nY-nRY-nGY-nBY)*binom(nR,nRG+nRB+nRY)*multinomial(nRG+nRB+nRY,nRG,nRB,nRY)*binom(nG,nRG+nGB+nGY)*multinomial(nRG+nGB+nGY,nRG,nGB,nGY)*binom(nB,nRB+nGB+nBY)*multinomial(nRB+nGB+nBY,nRB,nGB,nBY)*binom(nY,nRY+nGY+nBY)*multinomial(nRY+nGY+nBY,nRY,nGY,nBY)*math.factorial(nRG)*math.factorial(nRB)*math.factorial(nRY)*math.factorial(nGB)*math.factorial(nGY)*math.factorial(nBY)

grid(numpts)

Generate a default uniform grid where grid points lie directly on the boundaries.

Parameters:

Name	Type	Description	Default
numpts	int	Number of grid points (less one).	required

Returns:

Name	Type	Description
ndarray	ndarray	Uniform grid points.

Source code in dadi/TwoLocus/numerics.py

def grid(numpts):
    """
    Generate a default uniform grid where grid points lie directly on the boundaries.

    Args:
        numpts (int): Number of grid points (less one).

    Returns:
        ndarray (ndarray): Uniform grid points.
    """
    return np.linspace(0,1,numpts+1)

grid_dx(x)

Compute the 1D grid spacing for a given grid.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required

Returns:

Name	Type	Description
ndarray	ndarray	1D grid spacing.

Source code in dadi/TwoLocus/numerics.py

def grid_dx(x):
    """
    Compute the 1D grid spacing for a given grid.

    Args:
        x (array-like): 1D grid.

    Returns:
        ndarray (ndarray): 1D grid spacing.
    """
    return (np.concatenate((np.diff(x),np.array([0]))) + np.concatenate((np.array([0]),np.diff(x))))/2

grid_dx3(x, dx)

Compute 3D grid spacing for integration weights.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required

Returns:

Name	Type	Description
DX	ndarray	3D grid spacing.

Source code in dadi/TwoLocus/numerics.py

def grid_dx3(x,dx):
    """
    Compute 3D grid spacing for integration weights.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.

    Returns:
        DX (ndarray): 3D grid spacing.
    """
    DX = dx[:,nuax,nuax]*dx[nuax,:,nuax]*dx[nuax,nuax,:]
    for ii in range(len(x)):
        for jj in range(len(x)):
            DX[ii,jj,len(x)-ii-jj-1] *= 1./2
    return DX

grid_dx_2d(x, dx)

The two dimensional grid spacing over the domain. Grid points lie along the diagonal boundary, and Delta for those points is halved.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required

Returns:

Name	Type	Description
DXX	ndarray	2D grid spacing.

Source code in dadi/TwoLocus/numerics.py

def grid_dx_2d(x,dx):
    """
    The two dimensional grid spacing over the domain.
    Grid points lie along the diagonal boundary, and Delta for those points is halved.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.

    Returns:
        DXX (ndarray): 2D grid spacing.
    """
    DXX = dx[:,nuax]*dx[nuax,:]
    for ii in range(len(x)):
        DXX[ii,len(x)-ii-1] *= 1./2
    return DXX

injectA(x, dx, dt, yB, phi, thetaA)

Inject new derived mutations A onto the background of B/b.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
dt	float	Integration time step.	required
yB	Spectrum	Biallelic frequency spectrum integrated by dadi.	required
phi	array - like	Density function.	required
thetaA	float	Scaled mutation rate.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def injectA(x,dx,dt,yB,phi,thetaA):
    """
    Inject new derived mutations A onto the background of B/b.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        dt (float): Integration time step.
        yB (Spectrum): Biallelic frequency spectrum integrated by dadi.
        phi (array-like): Density function.
        thetaA (float): Scaled mutation rate.

    Returns:
        phi (ndarray): Updated density function.
    """
    phi[0,1,1:-1] += dt/dx[1] / x[1]**2 * yB[1:-1] * (1-x[1:-1]) * thetaA/2.
    phi[1,0,1:-1] += dt/dx[1] / x[1]**2 * yB[1:-1] * x[1:-1] * thetaA/2.
    phi[1,0,0] += dt/dx[1] * 2 / x[1]**2 * yB[1] * x[1] * thetaA/2. 
    return phi

injectB(x, dx, dt, yA, phi, thetaB)

Inject new derived mutations B onto the background of A/a.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
dt	float	Integration time step.	required
yA	Spectrum	Biallelic frequency spectrum integrated by dadi.	required
phi	array - like	Density function.	required
thetaB	float	Scaled mutation rate.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def injectB(x,dx,dt,yA,phi,thetaB):
    """
    Inject new derived mutations B onto the background of A/a.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        dt (float): Integration time step.
        yA (Spectrum): Biallelic frequency spectrum integrated by dadi.
        phi (array-like): Density function.
        thetaB (float): Scaled mutation rate.

    Returns:
        phi (ndarray): Updated density function.
    """
    phi[0,1:-1,1] += dt/dx[1] / x[1]**2 * yA[1:-1] * (1-x[1:-1]) * thetaB/2.
    phi[1,1:-1,0] += dt/dx[1] / x[1]**2 * yA[1:-1] * x[1:-1] * thetaB/2.
    phi[1,0,0] += dt/dx[1] * 2 / x[1]**2 * yA[1] * x[1] * thetaB/2. 
    return phi

int2(DXX, U)

Integrate the density function over the domain.

Parameters:

Name	Type	Description	Default
DXX	array - like	Two dimensional grid.	required
U	array - like	Density function.	required

Returns:

Name	Type	Description
float	float	Result of the integration.

Source code in dadi/TwoLocus/numerics.py

def int2(DXX,U):
    """
    Integrate the density function over the domain.

    Args:
        DXX (array-like): Two dimensional grid.
        U (array-like): Density function.

    Returns:
        float (float): Result of the integration.
    """
    return np.sum(DXX*U)

int3(DX, U)

Numerically integrate the density function over a 3D domain.

Parameters:

Name	Type	Description	Default
DX	array - like	3D grid spacing for integration weights.	required
U	array - like	Density function.	required

Returns:

Name	Type	Description
float	float	Result of the integration.

Source code in dadi/TwoLocus/numerics.py

def int3(DX,U):
    """
    Numerically integrate the density function over a 3D domain.

    Args:
        DX (array-like): 3D grid spacing for integration weights.
        U (array-like): Density function.

    Returns:
        float (float): Result of the integration.
    """
    return np.sum(DX*U)

ln_binomial(n, k)

Compute the natural logarithm of the binomial coefficient.

Parameters:

Name	Type	Description	Default
n	int	Number of trials.	required
k	int	Number of successes.	required

Returns:

Name	Type	Description
float	float	Natural logarithm of the binomial coefficient.

Source code in dadi/TwoLocus/numerics.py

def ln_binomial(n,k):
    """
    Compute the natural logarithm of the binomial coefficient.

    Args:
        n (int): Number of trials.
        k (int): Number of successes.

    Returns:
        float (float): Natural logarithm of the binomial coefficient.
    """
    return math.lgamma(n+1) - math.lgamma(k+1) - math.lgamma(n-k+1)

mean_r2(F)

Compute the mean r^2 value for a two-locus spectrum.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required

Returns:

Name	Type	Description
float	float	Mean r^2 value.

Source code in dadi/TwoLocus/numerics.py

def mean_r2(F):
    """
    Compute the mean r^2 value for a two-locus spectrum.

    Args:
        F (TLSpectrum): Two-locus spectrum.

    Returns:
        float (float): Mean r^2 value.
    """
    ns = len(F)-1
    tot = np.sum(F)
    r2s = []
    weights = []
    for ii in range(ns+1):
        for jj in range(ns+1-ii):
            for kk in range(ns+1-ii-jj):
                if F.mask[ii,jj,kk]:
                    continue
                p = (ii + jj)/float(ns)
                q = (ii + kk)/float(ns)
                pAB = ii/float(ns)
                D = pAB - p*q
                r2s.append(D**2/(p*(1-p)*q*(1-q)))
                weights.append(F[ii,jj,kk])
    return np.sum(np.array(r2s)*np.array(weights)/tot)

misidentification(F, p)

Handle misidentification of ancestral state in a two-locus spectrum.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required
p	float	Probability of misidentification.	required

Returns:

Name	Type	Description
TLSpectrum	TLSpectrum	Updated two-locus spectrum with misidentification.

with probability p, ancestral state is misidentified with prob p(1-p) A -> a but B correct with prob p(1-p) B -> b but A correct with prob p^2 A -> a and B -> b

Source code in dadi/TwoLocus/numerics.py

def misidentification(F,p):
    """
    Handle misidentification of ancestral state in a two-locus spectrum.

    Args:
        F (TLSpectrum): Two-locus spectrum.
        p (float): Probability of misidentification.

    Returns:
        TLSpectrum (TLSpectrum): Updated two-locus spectrum with misidentification.

    with probability p, ancestral state is misidentified
    with prob p(1-p) A -> a but B correct
    with prob p(1-p) B -> b but A correct
    with prob p^2 A -> a and B -> b
    """
    F_new = np.zeros(np.shape(F))
    ns = len(F) - 1
    for ii in range(len(F)):
        for jj in range(len(F)):
            for kk in range(len(F)):
                if F.mask[ii,jj,kk] == True:
                    continue
                ll = ns - ii - jj - kk
                weight = F[ii,jj,kk]
                F_new[ii,jj,kk] += weight * (1 - 2*p + p**2)
                F_new[kk,ll,ii] += weight * p * (1-p)
                F_new[jj,ii,ll] += weight * p * (1-p)
                F_new[ll,kk,jj] += weight * p**2
    return TLSpectrum(F_new)

misidentification_genotype_dict(G, p)

Handle misidentification of ancestral state in a dictionary of genotype spectra.

Parameters:

Name	Type	Description	Default
G	dict	Dictionary of genotype spectra.	required
p	float	Probability of misidentification.	required

Returns:

Name	Type	Description
G_new	dict	Updated dictionary of genotype spectra with misidentification.

same misidentification probabilities as misidentification(F,p) with prob p(1-p), AA -> aa, aa -> AA, Aa stays the same, B/b stay same with prob p(1-p),

Source code in dadi/TwoLocus/numerics.py

def misidentification_genotype_dict(G,p):
    """
    Handle misidentification of ancestral state in a dictionary of genotype spectra.

    Args:
        G (dict): Dictionary of genotype spectra.
        p (float): Probability of misidentification.

    Returns:
        G_new (dict): Updated dictionary of genotype spectra with misidentification.

    same misidentification probabilities as misidentification(F,p)
    with prob p(1-p), AA -> aa, aa -> AA, Aa stays the same, B/b stay same
    with prob p(1-p), 
    """
    G_new = {}
    n = G.keys()[0] # first key of G is number of individuals in sample
    G_new.setdefault(n,{})
    for genotypes in G[n].keys():
        G_new[n].setdefault(genotypes,0.0)
    for genotypes in G[n].keys():
        weight = G[n][genotypes]
        g1,g2,g3,g4,g5,g6,g7,g8 = genotypes
        g9 = n-g1-g2-g3-g4-g5-g6-g7-g8
        G_new[n][(g1,g2,g3,g4,g5,g6,g7,g8)] += weight * (1 - 2*p + p**2)
        G_new[n][(g7,g8,g9,g4,g5,g6,g1,g2)] += weight * p * (1-p)
        G_new[n][(g3,g2,g1,g6,g5,g4,g9,g8)] += weight * p * (1-p)
        G_new[n][(g9,g8,g7,g6,g5,g4,g3,g2)] += weight * p**2
    return G_new

move_density_to_surface(x, dx, dt, gammaA, gammaB, nu, hA=1.0 / 2, hB=1.0 / 2)

For each point in the full domain, compute how much should be lost to the nonsquare surface boundary due to diffusion/selection.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
dt	float	Timestep of integration.	required
gammaA	float	Scaled selection coefficient for A.	required
gammaB	float	Scaled selection coefficient for B.	required
nu	float	Relative population size.	required
hA	float	Dominance coefficient for A. Defaults to 1./2.	1.0 / 2
hB	float	Dominance coefficient for B. Defaults to 1./2.	1.0 / 2

Returns:

Name	Type	Description
Psurf	ndarray	Amount of density to be moved to the surface.

Source code in dadi/TwoLocus/numerics.py

def move_density_to_surface(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
    """
    For each point in the full domain, compute how much should be lost to the nonsquare surface boundary due to diffusion/selection.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        dt (float): Timestep of integration.
        gammaA (float): Scaled selection coefficient for A.
        gammaB (float): Scaled selection coefficient for B.
        nu (float): Relative population size.
        hA (float, optional): Dominance coefficient for A. Defaults to 1./2.
        hB (float, optional): Dominance coefficient for B. Defaults to 1./2.

    Returns:
        Psurf (ndarray): Amount of density to be moved to the surface.
    """
    Psurf = np.zeros((len(x),len(x),len(x)))
    if gammaA == 0 and gammaB == 0:
        P1D = transition1D(x,dx,dt,0,nu)
        weights = {}
        for ii in range(len(x))[1:-1]:
            y = np.zeros(len(x))
            y[ii] = 1./dx[ii]
            y = advance1D(y,P1D)
            weights[ii] = y[-1] * dx[-1]
        for ii in range(len(x))[:len(x)-1]:
            for jj in range(len(x))[:len(x)-1-ii]:
                for kk in range(len(x))[:len(x)-1-ii-jj]:
                    if ii+jj+kk == 0:
                        continue
                    Psurf[ii,jj,kk] = weights[ii+jj+kk]
    else:
        for ii in range(len(x))[:len(x)-1]:
            for jj in range(len(x))[:len(x)-1-ii]:
                for kk in range(len(x))[:len(x)-1-ii-jj]:
                    if ii+jj+kk == 0:
                        continue
                    p = x[ii]+x[jj]
                    q = x[ii]+x[kk]
                    #gamma = ((gammaA+gammaB)*x[ii] + gammaA*x[jj] + gammaB*x[kk]) / (x[ii] + x[jj] + x[kk])
                    ### XXX: why is this divided by (x[ii]+x[jj]+x[kk])?? because it's an average, sort of
                    gamma = ((gammaA+gammaB)*x[ii]**2 + 2*(gammaA+hB*gammaB)*x[ii]*x[jj] + 2*(hA*gammaA+gammaB)*x[ii]*x[kk] + (gammaA)*x[jj]**2 + 2*(hA*gammaA+hB*gammaB)*x[jj]*x[kk] + (gammaB)*x[kk]*2 + 2*(hA*gammaA+hB*gammaB)*x[ii]*(1-x[ii]-x[jj]-x[kk]) + 2*(hA*gammaA)*x[jj]*(1-x[ii]-x[jj]-x[kk]) + 2*(hB*gammaB)*x[kk]*(1-x[ii]-x[jj]-x[kk])) / (x[ii]+x[jj]+x[kk])
                    P1D = transition1D(x,dx,dt,gamma,nu)
                    y = np.zeros(len(x))
                    y[ii+jj+kk] = 1./dx[ii+jj+kk]
                    y = advance1D(y,P1D)
                    Psurf[ii,jj,kk] = y[-1]*dx[-1]
    return Psurf

move_surf_density_to_bdry(x, surf, P)

This is an absorbing boundary, so we'll just remove that density and not worry about integrating along that diagonal.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
surf	array - like	Surface density function.	required
P	array - like	Amount of density to be moved to the boundary.	required

Returns:

Name	Type	Description
surf	ndarray	Updated surface density function.

Source code in dadi/TwoLocus/numerics.py

def move_surf_density_to_bdry(x,surf,P):
    """
    This is an absorbing boundary, so we'll just remove that density and not worry about integrating along that diagonal.

    Args:
        x (array-like): 1D grid.
        surf (array-like): Surface density function.
        P (array-like): Amount of density to be moved to the boundary.

    Returns:
        surf (ndarray): Updated surface density function.
    """
    for ii in range(len(x))[1:]:
        for jj in range(len(x))[1:]:
            if x[ii] + x[jj] < 1:
                surf[ii,jj] *= (1-P[ii,jj])
    return surf

move_surf_to_line(x, dx, dt, gammaA, gammaB, nu, hA=1.0 / 2, hB=1.0 / 2)

For each point in the full domain, compute how much should be lost to the nonsquare surface boundary due to diffusion/selection.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
dt	float	Timestep of integration.	required
gammaA	float	Scaled selection coefficient for A.	required
gammaB	float	Scaled selection coefficient for B.	required
nu	float	Relative population size.	required
hA	float	Dominance coefficient for A. Defaults to 1./2.	1.0 / 2
hB	float	Dominance coefficient for B. Defaults to 1./2.	1.0 / 2

Returns:

Name	Type	Description
Psurf	ndarray	Amount of density to be moved to the surface.

Source code in dadi/TwoLocus/numerics.py

def move_surf_to_line(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
    """
    For each point in the full domain, compute how much should be lost to the nonsquare surface boundary due to diffusion/selection.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        dt (float): Timestep of integration.
        gammaA (float): Scaled selection coefficient for A.
        gammaB (float): Scaled selection coefficient for B.
        nu (float): Relative population size.
        hA (float, optional): Dominance coefficient for A. Defaults to 1./2.
        hB (float, optional): Dominance coefficient for B. Defaults to 1./2.

    Returns:
        Psurf (ndarray): Amount of density to be moved to the surface.
    """
    #### XXXX: 11/8/16: assuming that gammaB = 0
    Psurf = np.zeros((len(x),len(x)))

    for ii in range(len(x))[1:len(x)-1]:
        for jj in range(len(x))[1:len(x)-1-ii]:
            #gamma = ((gammaA+gammaB)*x[ii] + gammaA*x[jj]) / (x[ii] + x[jj]) ### this should also account for dominance
            gamma = (gammaA*x[ii]**2 + 2*(gammaA)*x[ii]*x[jj] + 2*(hA*gammaA)*x[ii]*(1-x[ii]-x[jj]) + gammaA*x[jj]**2 + 2*(hA*gammaA)*x[jj]*(1-x[ii]-x[jj])) / (x[ii]+x[jj])
            P1D = transition1D(x,dx,dt,gamma,nu)
            y = np.zeros(len(x))
            y[ii+jj] = 1./dx[ii+jj]
            y = advance1D(y,P1D)
            Psurf[ii,jj] = y[-1]*dx[-1]
    return Psurf

multinomial(n, i, j, k)

Compute the multinomial coefficient.

Parameters:

Name	Type	Description	Default
n	int	Number of trials.	required
i	int	Number of successes for the first category.	required
j	int	Number of successes for the second category.	required
k	int	Number of successes for the third category.	required

Returns:

Name	Type	Description
int	int	Multinomial coefficient.

Source code in dadi/TwoLocus/numerics.py

def multinomial(n,i,j,k):
    """
    Compute the multinomial coefficient.

    Args:
        n (int): Number of trials.
        i (int): Number of successes for the first category.
        j (int): Number of successes for the second category.
        k (int): Number of successes for the third category.

    Returns:
        int (int): Multinomial coefficient.
    """
    if i+j+k != n:
        return "nein"
    return math.factorial(n)/math.factorial(i)/math.factorial(j)/math.factorial(k)

observed_genotype_spectrum_dict_from_F(F)

Convert a two-locus spectrum to a dictionary of observed genotype spectra.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required

Returns:

Name	Type	Description
Gdict	dict	Dictionary of observed genotype spectra.

Gdict has keys n (num individuals in sample), observed genotypes

Source code in dadi/TwoLocus/numerics.py

def observed_genotype_spectrum_dict_from_F(F):
    """
    Convert a two-locus spectrum to a dictionary of observed genotype spectra.

    Args:
        F (TLSpectrum): Two-locus spectrum.

    Returns:
        Gdict (dict): Dictionary of observed genotype spectra.

    Gdict has keys n (num individuals in sample), observed genotypes
    """
    n = len(F)-1
    ng = n/2
    Gdict = {}
    Gdict.setdefault(ng,{})
    for ii in range(n+1):
        for jj in range(n+1):
            for kk in range(n+1):
                if F.mask[ii,jj,kk] == False:
                    nAB,nAb,naB,nab = ii,jj,kk,n-ii-jj-kk
                    gen_probs = cached_genotype_exact_projection(n,(nAB,nAb,naB,nab))
                    for gens in gen_probs.keys():
                        g1,g2,g3,g4,g5,g6,g7,g8,g9,g10 = gens
                        o1,o2,o3,o4,o5,o6,o7,o8,o9 = g1,g5,g2,g6,g7+g8,g9,g3,g10,g4
                        prob = gen_probs[gens]
                        try:
                            Gdict[ng][(o1,o2,o3,o4,o5,o6,o7,o8)] += F[nAB,nAb,naB]*prob
                        except KeyError:
                            Gdict[ng][(o1,o2,o3,o4,o5,o6,o7,o8)] = F[nAB,nAb,naB]*prob
    return Gdict

observed_genotype_spectrum_from_F(F)

Convert a two-locus spectrum to an observed genotype spectrum.

Parameters:

Name	Type	Description	Default
F	TLSpectrum	Two-locus spectrum.	required

Returns:

Name	Type	Description
G	ndarray	Observed genotype spectrum.

Source code in dadi/TwoLocus/numerics.py

def observed_genotype_spectrum_from_F(F):
    """
    Convert a two-locus spectrum to an observed genotype spectrum.

    Args:
        F (TLSpectrum): Two-locus spectrum.

    Returns:
        G (ndarray): Observed genotype spectrum.
    """
    n = len(F)-1
    ng = n/2
    G = np.zeros((ng+1,ng+1,ng+1,ng+1,ng+1,ng+1,ng+1,ng+1))
    for ii in range(n+1):
        for jj in range(n+1):
            for kk in range(n+1):
                if F[ii,jj,kk] > 0:# and F.mask[ii,jj,kk] == False:
                    nAB,nAb,naB,nab = ii,jj,kk,n-ii-jj-kk
                    gen_probs = cached_genotype_exact_projection(n,(nAB,nAb,naB,nab))
                    for gens in gen_probs.keys():
                        g1,g2,g3,g4,g5,g6,g7,g8,g9,g10 = gens
                        o1,o2,o3,o4,o5,o6,o7,o8,o9 = g1,g5,g2,g6,g7+g8,g9,g3,g10,g4
                        prob = gen_probs[gens]
                        G[o1,o2,o3,o4,o5,o6,o7,o8] += F[nAB,nAb,naB]*prob
    return G

pairings(n)

Compute the number of pairings for n individuals.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required

Returns:

Name	Type	Description
int	int	Number of pairings.

Source code in dadi/TwoLocus/numerics.py

def pairings(n):
    """
    Compute the number of pairings for n individuals.

    Args:
        n (int): Number of individuals.

    Returns:
        int (int): Number of pairings.
    """
    return math.factorial(n)/(math.factorial(n/2)*2**(n/2))

phi_to_surf(phi, x)

Map non-square to triallelic domain.

Note that many of the surface interaction methods could be cythonized for increase in speed. If x2 or x3 is lost, then only types AB and Ab or AB and aB are left, so either A or B have fixed, and that state is absorbing (no recombination can send you back to the interior). Thus, make sure that x2 lost or x3 lost is the diagonal boundary of the surface domain, so we don't have to worry about that density anymore. Here, we have x3 lost.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
surf	ndarray	Surface density function.

Source code in dadi/TwoLocus/numerics.py

def phi_to_surf(phi,x):
    """
    Map non-square to triallelic domain.

    Note that many of the surface interaction methods could be cythonized for increase in speed.
    If x2 or x3 is lost, then only types AB and Ab or AB and aB are left, so either A or B have fixed, and that state is absorbing (no recombination can send you back to the interior).
    Thus, make sure that x2 lost or x3 lost is the diagonal boundary of the surface domain, so we don't have to worry about that density anymore.
    Here, we have x3 lost.

    Args:
        phi (array-like): Density function.
        x (array-like): 1D grid.

    Returns:
        surf (ndarray): Surface density function.
    """
    surf = np.zeros((len(x),len(x)))
    for ii in range(len(x)): # loop through x1 indices
        for jj in range(len(x)): # loop through x2 indices
            kk = len(x)-1 - ii - jj
            surf[ii,jj] = phi[ii,jj,kk]
    return surf

possible_genotypes_4(n, colors)

Generate all possible sets of genotypes for a given set of haplotype counts.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
colors	tuple	Number of individuals with each haplotype.	required

Returns:

Name	Type	Description
gl	list	List of all possible sets of genotypes.

Source code in dadi/TwoLocus/numerics.py

def possible_genotypes_4(n,colors):
    """
    Generate all possible sets of genotypes for a given set of haplotype counts.

    Args:
        n (int): Number of individuals.
        colors (tuple): Number of individuals with each haplotype.

    Returns:
        gl (list): List of all possible sets of genotypes.
    """
    nR,nG,nB,nY = colors
    gl = []
    for pure_red in range(nR/2+1):
        for pure_green in range(nG/2+1):
            for pure_blue in range(nB/2+1):
                for pure_yellow in range(nY/2+1):
                    mixed_red = nR-2*pure_red
                    mixed_green = nG-2*pure_green
                    mixed_blue = nB-2*pure_blue
                    mixed_yellow = nY-2*pure_yellow
                    for mixRG in range(min(mixed_red,mixed_green)+1):
                        for mixRB in range(min(mixed_red-mixRG,mixed_blue)+1):
                            for mixRY in range(mixed_red-mixRG-mixRB,min(mixed_red-mixRG-mixRB,mixed_yellow)+1):
                                for mixGB in range(min(mixed_green-mixRG,mixed_blue-mixRB)+1):
                                    for mixGY in range(mixed_green-mixRG-mixGB,min(mixed_green-mixRG-mixGB,mixed_yellow-mixRY)+1):
                                        mixBY1 = mixed_blue-mixRB-mixGB
                                        mixBY2 = mixed_yellow-mixRY-mixGY
                                        if mixBY1 == mixBY2:
                                            gl.append((pure_red,pure_green,pure_blue,pure_yellow,mixRG,mixRB,mixRY,mixGB,mixGY,mixBY1))
    return gl

project(F_from, proj_to)

Project a two-locus spectrum to a smaller sample size.

Parameters:

Name	Type	Description	Default
F_from	TLSpectrum	Two-locus spectrum to project from.	required
proj_to	int	Number of samples to project down to.	required

Returns:

Name	Type	Description
TLSpectrum	TLSpectrum	Projected two-locus spectrum.

Source code in dadi/TwoLocus/numerics.py

def project(F_from, proj_to):
    """
    Project a two-locus spectrum to a smaller sample size.

    Args:
        F_from (TLSpectrum): Two-locus spectrum to project from.
        proj_to (int): Number of samples to project down to.

    Returns:
        TLSpectrum (TLSpectrum): Projected two-locus spectrum.
    """
    proj_from = len(F_from)-1
    if proj_to == proj_from:
        return F_from
    elif proj_to > proj_from:
        print('nope!')
        return F_from
    else:
        F_proj = np.zeros((proj_to+1,proj_to+1,proj_to+1))
        for X1 in range(proj_from):
            for X2 in range(proj_from):
                for X3 in range(proj_from):
                    if F_from.mask[X1,X2,X3] == False:
                        hits = (X1,X2,X3)
                        proj_weights = cached_projection(proj_to,proj_from,hits)
                        F_proj += proj_weights * F_from[X1,X2,X3]

        return TLSpectrum(F_proj)

project_Gdict(G, n_from, n_to)

Project a dictionary of genotype spectra to a smaller sample size.

Parameters:

Name	Type	Description	Default
G	dict	Dictionary of genotype spectra.	required
n_from	int	Number of individuals in the original sample.	required
n_to	int	Number of individuals in the projected sample.	required

Returns:

Name	Type	Description
G_to	dict	Dictionary of projected genotype spectra.

n_from and n_to are the individual counts, not haplotype counts (so ns/2)

Source code in dadi/TwoLocus/numerics.py

def project_Gdict(G,n_from,n_to):
    """
    Project a dictionary of genotype spectra to a smaller sample size.

    Args:
        G (dict): Dictionary of genotype spectra.
        n_from (int): Number of individuals in the original sample.
        n_to (int): Number of individuals in the projected sample.

    Returns:
        G_to (dict): Dictionary of projected genotype spectra.

    n_from and n_to are the individual counts, not haplotype counts (so ns/2)
    """
    G_to = {}
    G_to.setdefault(n_to,{})
    for genotypes in G[n_from].keys():
        weights = projection_cache_Gdict(n_from,n_to,genotypes)
        for projected_gens in weights.keys():
            try:
                G_to[n_to][projected_gens] += weights[projected_gens]*G[genotypes]
            except KeyError:
                G_to[n_to][projected_gens] = weights[projected_gens]*G[genotypes]
    return G_to

projection_cache_Gdict(n_from, n_to, hits)

Cache the projection weights for genotype spectra.

Parameters:

Name	Type	Description	Default
n_from	int	Number of individuals in the original sample.	required
n_to	int	Number of individuals in the projected sample.	required
hits	tuple	Number of individuals with each genotype in the original sample.	required

Returns:

Name	Type	Description
weights_to	dict	Dictionary of projection weights for each set of genotypes.

Source code in dadi/TwoLocus/numerics.py

def projection_cache_Gdict(n_from,n_to,hits):
    """
    Cache the projection weights for genotype spectra.

    Args:
        n_from (int): Number of individuals in the original sample.
        n_to (int): Number of individuals in the projected sample.
        hits (tuple): Number of individuals with each genotype in the original sample.

    Returns:
        weights_to (dict): Dictionary of projection weights for each set of genotypes.
    """
    key = (n_from,n_to,hits)
    try:
        return genotype_projection_cache[key]
    except KeyError:
        pass

#    g1,g2,g3,g4,g5,g6,g7,g8 = hits
#    g9 = n_from - g1 - g2 - g3 - g4 - g5 - g6 - g7 - g8
#    weights_to = {}
#    for o1 in range(0,g1+1):
#        for o2 in range(0,g2+1):
#            for o3 in range(0,g3+1):
#                for o4 in range(0,g4+1):
#                    for o5 in range(0,g5+1):
#                        for o6 in range(0,g6+1):
#                            for o7 in range(0,g7+1):
#                                for o8 in range(0,g8+1):
#                                    o9 = n_to-o1-o2-o3-o4-o5-o6-o7-o8
#                                    if o9 < 0 or o9 > g9:
#                                        continue
#                                    else:
#                                        weight = np.exp(ln_binomial(g1,o1) + ln_binomial(g2,o2) + ln_binomial(g3,o3) + ln_binomial(g4,o4) + ln_binomial(g5,o5) + ln_binomial(g6,o6) + ln_binomial(g7,o7) + ln_binomial(g8,o8) + ln_binomial(g9,o9) - ln_binomial(n_from,n_to))
#                                        weights_to[(o1,o2,o3,o4,o5,o6,o7,o8)] = weight
#

    weights_to = projection_genotypes.projection_genotypes(n_from,n_to,hits)
    genotype_projection_cache[key] = weights_to
    return weights_to

quadrinomial(ns, ii, jj, kk)

Compute the quadrinomial coefficient.

Parameters:

Name	Type	Description	Default
ns	int	Number of samples.	required
ii	int	Number of derived alleles for the first locus.	required
jj	int	Number of derived alleles for the second locus.	required
kk	int	Number of derived alleles for the third locus.	required

Returns:

Name	Type	Description
float	float	Quadrinomial coefficient.

Source code in dadi/TwoLocus/numerics.py

def quadrinomial(ns,ii,jj,kk):
    """
    Compute the quadrinomial coefficient.

    Args:
        ns (int): Number of samples.
        ii (int): Number of derived alleles for the first locus.
        jj (int): Number of derived alleles for the second locus.
        kk (int): Number of derived alleles for the third locus.

    Returns:
        float (float): Quadrinomial coefficient.
    """
    return np.exp(math.lgamma(ns+1) - math.lgamma(ii+1) - math.lgamma(jj+1) - math.lgamma(kk+1) - math.lgamma(ns-ii-jj-kk+1))

sample_cached(phi, ns, x)

Sample from the cached density function.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
ns	int or tuple	Number of samples.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
F	TLSpectrum	Sampled spectrum.

Source code in dadi/TwoLocus/numerics.py

def sample_cached(phi, ns, x):
    """
    Sample from the cached density function.

    Args:
        phi (array-like): Density function.
        ns (int or tuple): Number of samples.
        x (array-like): 1D grid.

    Returns:
        F (TLSpectrum): Sampled spectrum.
    """
    dx = grid_dx(x)
    dx3 = grid_dx3(x,dx)

    if type(ns) == int:
        ns = (ns,)
    else:
        if len(ns) == 1:
            ns = tuple(ns)
        else:
            ns = (ns[0],)

    # cache calculations of several large matrices
    # cache x**ii, x**jj, and x**kk 
    # just cache x[:,nuax,nuax]**ii, and then later use np.swapaxes
    key = (ns, tuple(x))
    if key not in sample_cache:
        this_cache = {}
        for ii in range(ns[0]+1):
            this_cache[ii] = x[:,nuax,nuax]**ii
            this_cache[-ii] = (1 - x[:,nuax,nuax] - x[nuax,:,nuax] - x[nuax,nuax,:]) ** ii
        sample_cache[key] = this_cache
    else:
        this_cache = sample_cache[key]

    F = np.zeros((ns[0]+1,ns[0]+1,ns[0]+1))
    prod_phi = dx3*phi
    for ii in range(len(F)):
        prod_x = prod_phi * this_cache[ii]
        for jj in range(len(F)):
            prod_y = prod_x * np.swapaxes(this_cache[jj],0,1)
            for kk in range(len(F)):
                if ii+jj+kk <= ns[0]:
                    F[ii,jj,kk] = quadrinomial(ns[0],ii,jj,kk) * np.sum(prod_y * np.swapaxes(this_cache[kk],0,2) * this_cache[-(ns[0]-ii-jj-kk)])

    F = TLSpectrum(F)
    return F

surf_to_phi(surf, phi, x)

Move the surface density back to the full phi density function.

Parameters:

Name	Type	Description	Default
surf	array - like	Surface density function.	required
phi	array - like	Full density function.	required
x	array - like	1D grid.	required

Returns:

Name	Type	Description
phi	ndarray	Updated full density function.

Source code in dadi/TwoLocus/numerics.py

def surf_to_phi(surf,phi,x):
    """
    Move the surface density back to the full phi density function.

    Args:
        surf (array-like): Surface density function.
        phi (array-like): Full density function.
        x (array-like): 1D grid.

    Returns:
        phi (ndarray): Updated full density function.
    """
    for ii in range(len(x)): # loop through x1 indices
        for jj in range(len(x)): # loop through x2 indices
            kk = len(x)-1 - ii - jj
            phi[ii,jj,kk] = surf[ii,jj]
    return phi

surface_interaction_b(phi, x, Psurf)

Handle surface interaction for boundary conditions.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
x	array - like	1D grid.	required
Psurf	array - like	Amount of density to be moved to the surface.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def surface_interaction_b(phi,x,Psurf):
    """
    Handle surface interaction for boundary conditions.

    Args:
        phi (array-like): Density function.
        x (array-like): 1D grid.
        Psurf (array-like): Amount of density to be moved to the surface.

    Returns:
        phi (ndarray): Updated density function.
    """
    for ii in range(len(x))[:len(x)-1]:
        for jj in range(len(x))[:len(x)-1-ii]:
            for kk in range(len(x))[:len(x)-1-ii-jj]:
                if ii == 0 and jj == 0:
                    continue
                if ii == 0 and kk == 0:
                    continue
                if jj == 0 and kk == 0:
                    continue
                if ii == 0:
                    amnt = Psurf[ii,jj,kk]
                    dist = (len(x)-1 - (jj+kk))//2
                    rmdr = (len(x)-1 - (jj+kk))%2
                    if rmdr == 0:
                        phi[ii,jj+dist,kk+dist] += amnt * phi[ii,jj,kk] * 2
                        phi[ii,jj,kk] *= (1-amnt)
                    elif rmdr == 1:
                        phi[ii,jj+dist+1,kk+dist] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii,jj+dist,kk+dist+1] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii,jj,kk] *= (1-amnt)
                elif jj == 0:
                    amnt = Psurf[ii,jj,kk]
                    dist = (len(x)-1 - (ii+kk))//2
                    rmdr = (len(x)-1 - (ii+kk))%2
                    if rmdr == 0:
                        phi[ii+dist,jj,kk+dist] += amnt * phi[ii,jj,kk] * 2
                        phi[ii,jj,kk] *= (1-amnt)
                    elif rmdr == 1:
                        phi[ii+dist+1,jj,kk+dist] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii+dist,jj,kk+dist+1] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii,jj,kk] *= (1-amnt)
                elif kk == 0:
                    amnt = Psurf[ii,jj,kk]
                    dist = (len(x)-1 - (ii+jj))//2
                    rmdr = (len(x)-1 - (ii+jj))%2
                    if rmdr == 0:
                        phi[ii+dist,jj+dist,kk] += amnt * phi[ii,jj,kk] * 2
                        phi[ii,jj,kk] *= (1-amnt)
                    elif rmdr == 1:
                        phi[ii+dist+1,jj+dist,kk] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii+dist,jj+dist+1,kk] += amnt * phi[ii,jj,kk] * 2/2.
                        phi[ii,jj,kk] *= (1-amnt)
                else:
                    amnt = Psurf[ii,jj,kk]
                    dist = (len(x)-1 - (ii+jj+kk))//3
                    rmdr = (len(x)-1 - (ii+jj+kk))%3
                    if rmdr == 0: # in line with boundary grid point
                        phi[ii+dist,jj+dist,kk+dist] += amnt * phi[ii,jj,kk] * 2
                        phi[ii,jj,kk] *= (1-amnt)
                    elif rmdr == 1:
                        phi[ii+dist+1,jj+dist,kk+dist] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii+dist,jj+dist+1,kk+dist] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii+dist,jj+dist,kk+dist+1] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii,jj,kk] *= (1-amnt)
                    elif rmdr == 2:
                        phi[ii+dist+1,jj+dist+1,kk+dist] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii+dist+1,jj+dist,kk+dist+1] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii+dist,jj+dist+1,kk+dist+1] += amnt * phi[ii,jj,kk] * 2/3.
                        phi[ii,jj,kk] *= (1-amnt)

    return phi

surface_recombination(phi, x, rho, dt)

Handle recombination that pushes density from surface back into interior of domain.

Parameters:

Name	Type	Description	Default
phi	array - like	Density function.	required
x	array - like	1D grid.	required
rho	float	Recombination rate.	required
dt	float	Timestep of integration.	required

Returns:

Name	Type	Description
phi	ndarray	Updated density function.

Source code in dadi/TwoLocus/numerics.py

def surface_recombination(phi,x,rho,dt):
    """
    Handle recombination that pushes density from surface back into interior of domain.

    Args:
        phi (array-like): Density function.
        x (array-like): 1D grid.
        rho (float): Recombination rate.
        dt (float): Timestep of integration.

    Returns:
        phi (ndarray): Updated density function.
    """
    if rho == 0:
        return phi
    else:
        for ii in range(len(x)):
            for jj in range(len(x))[1:]:
                kk = len(x)-1 - ii - jj
                if kk > 0:
                    D = -x[jj]*x[kk] # -x2*x3, since x4=0
                    if jj-1 > 0:
                        phi[ii,jj-1,kk] += 1./2 * (-D) * dt * phi[ii,jj,kk] * rho/4. / x[1]
                    else:
                        phi[ii,jj-1,kk] += (-D) * dt * phi[ii,jj,kk] * rho/4. / x[1]
                    if kk-1 > 0:
                        phi[ii,jj,kk-1] += 1./2 * (-D) * dt * phi[ii,jj,kk] * rho/4. / x[1]
                    else:
                        phi[ii,jj,kk-1] += (-D) * dt * phi[ii,jj,kk] * rho/4. / x[1]

                    phi[ii,jj,kk] -= 2 * (-D) * dt * phi[ii,jj,kk] * rho/4. / x[1]
    return phi

to_single_locus(phi)

Convert a two-locus spectrum to single-locus spectra.

Parameters:

Name	Type	Description	Default
phi	TLSpectrum	Two-locus spectrum.	required

Returns:

Name	Type	Description
fsA	Spectrum	Single-locus spectra.
fsB	Spectrum	Single-locus spectra.

Source code in dadi/TwoLocus/numerics.py

def to_single_locus(phi):
    """
    Convert a two-locus spectrum to single-locus spectra.

    Args:
        phi (TLSpectrum): Two-locus spectrum.

    Returns:
        fsA (Spectrum): Single-locus spectra.
        fsB (Spectrum): Single-locus spectra.
    """
    ns = len(phi)-1
    fsA = dadi.Spectrum(np.zeros(ns+1))
    fsB = dadi.Spectrum(np.zeros(ns+1))
    for ii in range(ns+1):
        for jj in range(ns+1-ii):
            for kk in range(ns+1-ii-jj):
                if phi.mask[ii,jj,kk]:
                    continue
                p = ii + jj
                q = ii + kk
                fsA[p] += phi[ii,jj,kk]
                fsB[q] += phi[ii,jj,kk]
    return fsA,fsB

transition12_surf(x, dx, U01)

Transition matrix for explicit integration of mixed derivative term along surface.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
U01	array - like	Surface domain marker.	required

Returns:

Name	Type	Description
C	ndarray	Transition matrix for mixed derivative term along surface.

Source code in dadi/TwoLocus/numerics.py

def transition12_surf(x,dx,U01):
    """
    Transition matrix for explicit integration of mixed derivative term along surface.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        U01 (array-like): Surface domain marker.

    Returns:
        C (ndarray): Transition matrix for mixed derivative term along surface.
    """
    C = lil_matrix((len(x)**2,len(x)**2))
    for ii in range(len(x)-1)[:-1]:
        for jj in range(len(x)-1)[:-1]:
            if U01[ii+2,jj+2] == 1 or U01[ii+1,jj+2] == 1 or U01[ii+2,jj+1] == 1:
                if ii+1 < len(x) and jj+1 < len(x) and U01[ii+1,jj+1] == 1:
                    C[ii*len(x)+jj,(ii+1)*len(x)+(jj+1)] += 1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii+1]*x[jj+1])

                if ii+1 < len(x) and jj-1 >= 0 and U01[ii+1,jj-1] == 1:
                    C[ii*len(x)+jj,(ii+1)*len(x)+(jj-1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii+1]*x[jj-1])

                if ii-1 >= 0 and jj+1 < len(x) and U01[ii-1,jj+1] == 1:
                    C[ii*len(x)+jj,(ii-1)*len(x)+(jj+1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii-1]*x[jj+1])

                if ii-1 >= 0 and jj-1 >= 0 and U01[ii-1,jj-1] == 1:
                    C[ii*len(x)+jj,(ii-1)*len(x)+(jj-1)] += 1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii-1]*x[jj-1])

            elif U01[ii+1,jj+1] == 1 or U01[ii+1,jj] == 1 or U01[ii,jj+1] == 1:
                if ii+1 < len(x) and jj-1 >= 0 and U01[ii+1,jj-1] == 1:
                    C[ii*len(x)+jj,(ii+1)*len(x)+(jj-1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii+1]*x[jj-1]) / 2

                if ii-1 >= 0 and jj+1 < len(x) and U01[ii-1,jj+1] == 1:
                    C[ii*len(x)+jj,(ii-1)*len(x)+(jj+1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii-1]*x[jj+1]) / 2

                if ii-1 >= 0 and jj-1 >= 0 and U01[ii-1,jj-1] == 1:
                    C[ii*len(x)+jj,(ii-1)*len(x)+(jj-1)] += 1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii-1]*x[jj-1])
    ii = 0
    jj = len(x)-2
    C[ii*len(x)+jj,(ii+1)*len(x)+(jj-1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii+1]*x[jj-1]) / 2
    ii = len(x)-2
    jj = 0
    C[ii*len(x)+jj,(ii-1)*len(x)+(jj+1)] += -1./4 * 1./(dx[ii]*dx[jj]) * (-x[ii-1]*x[jj+1]) / 2
    return C

transition1_surf(x, dx, U01, gammaA, gammaB, rho, nu, hA=0.5, hB=0.5)

ADI transition matrix for x1 on surface.

See cythonized versions of these in the dadi/Triallele code.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
U01	array - like	Surface domain marker.	required
gammaA	float	Scaled selection coefficient for A.	required
gammaB	float	Scaled selection coefficient for B.	required
rho	float	Recombination rate.	required
nu	float	Relative population size.	required
hA	float	Dominance coefficient for A. Defaults to .5.	0.5
hB	float	Dominance coefficient for B. Defaults to .5.	0.5

Returns:

Name	Type	Description
P	ndarray	ADI transition matrix for x1 on surface.

Source code in dadi/TwoLocus/numerics.py

def transition1_surf(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
    """
    ADI transition matrix for x1 on surface.

    See cythonized versions of these in the dadi/Triallele code.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        U01 (array-like): Surface domain marker.
        gammaA (float): Scaled selection coefficient for A.
        gammaB (float): Scaled selection coefficient for B.
        rho (float): Recombination rate.
        nu (float): Relative population size.
        hA (float, optional): Dominance coefficient for A. Defaults to .5.
        hB (float, optional): Dominance coefficient for B. Defaults to .5.

    Returns:
        P (ndarray): ADI transition matrix for x1 on surface.
    """
    P = np.zeros((len(x),3,len(x)))
    for jj in range(len(x)):
        A = np.zeros((len(x),len(x)))
        if jj > 0:
            V = x*(1-x)/nu
            V[np.max(np.where(U01[:,jj] == 1))] = 0
            for ii in np.where(U01[:,jj] == 1)[0][:-1]:
                if ii == 0:
                    A[ii,ii] =  - 1/(2*dx[ii]) * ( -V[ii]/(x[ii+1]-x[ii]) )
                    A[ii,ii+1] = - 1/(2*dx[ii]) * ( V[ii+1]/(x[ii+1]-x[ii]) )
                elif ii == np.where(U01[:,jj] == 1)[0][:-1][-1]:
                    A[ii,ii-1] = - 1/(2*dx[ii]) * ( V[ii-1]/(x[ii]-x[ii-1]) )
                    A[ii,ii] = - 1/(2*dx[ii]) * ( -V[ii]/(x[ii]-x[ii-1]) )
                else:
                    A[ii,ii-1] = - 1/(2*dx[ii]) * ( V[ii-1]/(x[ii]-x[ii-1]) )
                    A[ii,ii] = - 1/(2*dx[ii]) * ( -V[ii]/(x[ii]-x[ii-1]) -V[ii]/(x[ii+1]-x[ii]) )
                    A[ii,ii+1] = - 1/(2*dx[ii]) * ( V[ii+1]/(x[ii+1]-x[ii]) )
        if jj == 0:
            V = x*(1-x)/nu
            for ii in range(len(x)):
                if ii == 0:
                    A[ii,ii] =  - 1/(2*dx[ii]) * ( -V[ii]/(x[ii+1]-x[ii]) )
                    A[ii,ii+1] = - 1/(2*dx[ii]) * ( V[ii+1]/(x[ii+1]-x[ii]) )
                elif ii == len(x)-1:
                    A[ii,ii-1] = - 1/(2*dx[ii])*2 * ( V[ii-1]/(x[ii]-x[ii-1]) )
                    A[ii,ii] = - 1/(2*dx[ii])*2 * ( -V[ii]/(x[ii]-x[ii-1]) )
                else:
                    A[ii,ii-1] = - 1/(2*dx[ii]) * ( V[ii-1]/(x[ii]-x[ii-1]) )
                    A[ii,ii] = - 1/(2*dx[ii]) * ( -V[ii]/(x[ii]-x[ii-1]) -V[ii]/(x[ii+1]-x[ii]) )
                    A[ii,ii+1] = - 1/(2*dx[ii]) * ( V[ii+1]/(x[ii+1]-x[ii]) )

        x2 = x[jj]
        M = 2*gammaA * x*(1-x-x2)*(hA+(x+x2)*(1-2*hA)) + 2*gammaB * x*x2*(hB+(1-x2)*(1-2*hB)) - rho/2. * (-x2*(1-x-x2)) / 2.
        #M = gammaA * x * (1-x-x2) + gammaB * x * x2 - rho/2. * (-x2*(1-x-x2))
        M[np.where(U01[:,jj] == 1)[0][-1]] = 0
        for ii in np.where(U01[:,jj] == 1)[0]:
            if ii == 0:
                A[ii,ii] += 1/dx[ii] * ( M[ii] ) / 2
                A[ii,ii+1] += 1/dx[ii] * ( M[ii+1] ) / 2
            elif ii == np.where(U01[:,jj] == 1)[0][-1]:
                A[ii,ii-1] += 1/dx[ii] * 2 * ( - M[ii-1] ) / 2
                A[ii,ii] += 1/dx[ii] * 2 * ( - M[ii] ) / 2
            else:
                A[ii,ii-1] += 1/dx[ii] * ( - M[ii-1] ) / 2
                A[ii,ii] += 0 #1/dx[ii] * ( M[ii] - M[ii-1] ) / 2
                A[ii,ii+1] += 1/dx[ii] * ( M[ii+1] ) / 2

        P[jj,0,:] = np.concatenate(( np.array([0]), np.diagonal(A,-1) ))
        P[jj,1,:] = np.diagonal(A)
        P[jj,2,:] = np.concatenate(( np.diagonal(A,1), np.array([0]) ))
    return P

transition2_surf(x, dx, U01, gammaA, gammaB, rho, nu, hA=0.5, hB=0.5)

ADI transition matrix for x2 on surface.

Parameters:

Name	Type	Description	Default
x	array - like	1D grid.	required
dx	array - like	1D grid spacing.	required
U01	array - like	Surface domain marker.	required
gammaA	float	Scaled selection coefficient for A.	required
gammaB	float	Scaled selection coefficient for B.	required
rho	float	Recombination rate.	required
nu	float	Relative population size.	required
hA	float	Dominance coefficient for A. Defaults to .5.	0.5
hB	float	Dominance coefficient for B. Defaults to .5.	0.5

Returns:

Name	Type	Description
P	ndarray	ADI transition matrix for x2 on surface.

Source code in dadi/TwoLocus/numerics.py

def transition2_surf(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
    """
    ADI transition matrix for x2 on surface.

    Args:
        x (array-like): 1D grid.
        dx (array-like): 1D grid spacing.
        U01 (array-like): Surface domain marker.
        gammaA (float): Scaled selection coefficient for A.
        gammaB (float): Scaled selection coefficient for B.
        rho (float): Recombination rate.
        nu (float): Relative population size.
        hA (float, optional): Dominance coefficient for A. Defaults to .5.
        hB (float, optional): Dominance coefficient for B. Defaults to .5.

    Returns:
        P (ndarray): ADI transition matrix for x2 on surface.
    """
    P = np.zeros((len(x),3,len(x)))
    for ii in range(len(x)):
        A = np.zeros((len(x),len(x)))
        if ii > 0:
            V = x*(1-x)/nu
            V[np.max(np.where(U01[ii,:] == 1))] = 0
            for jj in np.where(U01[ii,:] == 1)[0][:-1]:
                if jj == 0:
                    A[jj,jj] =  - 1/(2*dx[jj]) * ( -V[jj]/(x[jj+1]-x[jj]) )
                    A[jj,jj+1] = - 1/(2*dx[jj]) * ( V[jj+1]/(x[jj+1]-x[jj]) )
                elif jj == np.where(U01[ii,:] == 1)[0][:-1][-1]:
                    A[jj,jj-1] = - 1/(2*dx[jj]) * ( V[jj-1]/(x[jj]-x[jj-1]) )
                    A[jj,jj] = - 1/(2*dx[jj]) * ( -V[jj]/(x[jj]-x[jj-1]) )
                else:
                    A[jj,jj-1] = - 1/(2*dx[jj]) * ( V[jj-1]/(x[jj]-x[jj-1]) )
                    A[jj,jj] = - 1/(2*dx[jj]) * ( -V[jj]/(x[jj]-x[jj-1]) -V[jj]/(x[jj+1]-x[jj]) )
                    A[jj,jj+1] = - 1/(2*dx[jj]) * ( V[jj+1]/(x[jj+1]-x[jj]) )
        if ii == 0:
            V = x*(1-x)/nu
            for jj in range(len(x)):
                if jj == 0:
                    A[jj,jj] =  - 1/(2*dx[jj]) * ( -V[jj]/(x[jj+1]-x[jj]) )
                    A[jj,jj+1] = - 1/(2*dx[jj]) * ( V[jj+1]/(x[jj+1]-x[jj]) )
                elif jj == len(x)-1:
                    A[jj,jj-1] = - 1/(2*dx[jj])*2 * ( V[jj-1]/(x[jj]-x[jj-1]) )
                    A[jj,jj] = - 1/(2*dx[jj])*2 * ( -V[jj]/(x[jj]-x[jj-1]) )
                else:
                    A[jj,jj-1] = - 1/(2*dx[jj]) * ( V[jj-1]/(x[jj]-x[jj-1]) )
                    A[jj,jj] = - 1/(2*dx[jj]) * ( -V[jj]/(x[jj]-x[jj-1]) -V[jj]/(x[jj+1]-x[jj]) )
                    A[jj,jj+1] = - 1/(2*dx[jj]) * ( V[jj+1]/(x[jj+1]-x[jj]) )

        x1 = x[ii]
        M = 2*gammaA * x*(1-x1-x)*(hA+(x+x1)*(1-2*hA)) - 2*gammaB * x*(1-x) + rho/2. * (-x * (1-x1-x)) / 2.
        #M = gammaA * x * (1-x1-x) - gammaB * x * (1-x) + rho/2. * (-x * (1-x1-x)) / 2. ### note the (/2.) term at the end, since x3 also would decrease by this amount - we separately handle reentry into domain interior
        M[np.where(U01[ii,:] == 1)[0][-1]] = 0
        for jj in np.where(U01[ii,:] == 1)[0]:
            if jj == 0:
                A[jj,jj] += 1/dx[jj] * ( M[jj] ) / 2
                A[jj,jj+1] += 1/dx[jj] * ( M[jj+1] ) / 2
            elif jj == np.where(U01[ii,:] == 1)[0][-1]:
                A[jj,jj-1] += 1/dx[jj] * 2 * ( - M[jj-1] ) / 2
                A[jj,jj] += 1/dx[jj] * 2 * ( - M[jj] ) / 2
            else:
                A[jj,jj-1] += 1/dx[jj] * ( - M[jj-1] ) / 2
                A[jj,jj] += 0 # 1/dx[jj] * ( M[jj] - M[jj-1] ) / 2
                A[jj,jj+1] += 1/dx[jj] * ( M[jj+1] ) / 2
        P[ii,0,:] = np.concatenate(( np.array([0]), np.diagonal(A,-1) ))
        P[ii,1,:] = np.diagonal(A)
        P[ii,2,:] = np.concatenate(( np.diagonal(A,1), np.array([0]) ))
    return P