dadi.TwoLocus.numerics API documentation
Description: methods for numerics of two locus diffusion - grids, transition matrices, etc

# Module `dadi.TwoLocus.numerics`

methods for numerics of two locus diffusion - grids, transition matrices, etc

Expand source code

    """
    methods for numerics of two locus diffusion - grids, transition matrices, etc
    """
    import numpy as np
    import dadi
    from numpy import newaxis as nuax
    from scipy.sparse import lil_matrix,identity
    import math

    from dadi.TwoLocus.TLSpectrum_mod import TLSpectrum

    tol = 1e-12

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

        ns: Number of samples

        Returns (D, r2). Both are TLSpectrum objects in which each entry is the
        value of D or 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

    def grid(numpts):
        """
        Default grid is uniform, so that grid points lie directly on the boundaries
        Here, delta_x = 1./numpts
        numpts - number of grid points (less one)
        """
        return np.linspace(0,1,numpts+1)

    def grid_dx(x):
        """
        Given 1D grid x, this returns the 1D grid spacing for x
        """
        return (np.concatenate((np.diff(x),np.array([0]))) + np.concatenate((np.array([0]),np.diff(x))))/2

    def domain(x):
        """
        Array of same dimension as density function discretization
        Has a 1 if grid point lies in domain, 0 if outside domain
        x - 1D grid
        """
        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

    def grid_dx3(x,dx):
        """
        For integrating over the 3D domain (see int3)
        x - 1D grid
        dx - 1D 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

    def int3(DX,U):
        """
        Numerically integrate the density function phi
        phi - density function
        DX - 3D grid spacing for integration weights
        """
        return np.sum(DX*U)

    """
    Injection of new mutations along axes
    """

    def injectA(x,dx,dt,yB,phi,thetaA):
        """
        Injection new derived mutations A onto background of B/b
        x - 1D grid
        dx - 1D grid spacing
        dt - integration time step
        yA - biallelic frequency spectrum integrated by dadi
        phi - density function
        thetaA - scaled mutation rate
        """
        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

    def injectB(x,dx,dt,yA,phi,thetaB):
        """
        Injection new derived mutations B onto background of A/a
        x - 1D grid
        dx - 1D grid spacing
        dt - integration time step
        yB - biallelic frequency spectrum integrated by dadi
        phi - density function
        thetaB - scaled mutation rate
        """
        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

    """
    Transition matrices for forward integration in bulk of domain.
    """

    def transition1(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
        """
        Transition matrix for ADI method along first axis
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = np.zeros((len(x),len(x),3,len(x)))
        for jj in range(len(x)):
            for kk in range(len(x)):
                if len(np.where(U01[:,jj,kk] == 1)[0]) <= 2:
                    continue
                A = np.zeros((len(x),len(x)))
                # drift
                if (jj > 0 or kk > 0) and x[jj] + x[kk] < 1 - tol:
                    V = x*(1-x)/nu
                    for ii in np.where(U01[:,jj,kk] == 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,kk] == 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]) )

                    x2 = x[jj]
                    x3 = x[kk]
                    M = 2*gammaA * x*(1-x-x2)*(hA+(x+x2)*(1-2*hA)) + 2*gammaB * x*(1-x-x3)*(hB+(x+x3)*(1-2*hB)) - rho/2. * (x*(1-x-x2-x3) - x2*x3)
                    for ii in np.where(U01[:,jj,kk] == 1)[0][:-1]:
                        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,kk] == 1)[0][:-1][-1]:
                            A[ii,ii-1] += 1/dx[ii] * ( - M[ii-1] ) / 2
                            A[ii,ii] += 1/dx[ii] * ( - 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

                if jj == 0 and kk == 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]
                    x3 = x[kk]
                    M = 2*gammaA * x*(1-x-x2)*(hA+(x+x2)*(1-2*hA)) + 2*gammaB * x*(1-x-x3)*(hB+(x+x3)*(1-2*hB)) - rho/2. * (x*(1-x-x2-x3) - x2*x3)
                    for ii in range(len(x)):
                        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 == len(x)-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,kk,0,:] = np.concatenate(( np.array([0]), np.diagonal(A,-1) ))
                P[jj,kk,1,:] = np.diagonal(A)
                P[jj,kk,2,:] = np.concatenate(( np.diagonal(A,1), np.array([0]) ))
        return P

    def transition2(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
        """
        Transition matrix for ADI method along second axis
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = np.zeros((len(x),len(x),3,len(x)))
        for ii in range(len(x)):
            for kk in range(len(x)):
                if len(np.where(U01[ii,:,kk] == 1)[0]) <= 2:
                    continue
                A = np.zeros((len(x),len(x)))
                # drift
                if (ii > 0 or kk > 0) and x[ii] + x[kk] < 1 - tol:
                    V = x*(1-x)/nu
                    for jj in np.where(U01[ii,:,kk] == 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,:,kk] == 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]) )

                    x1 = x[ii]
                    x3 = x[kk]
                    M = 2*gammaA * x*(1-x1-x)*(hA+(x1+x)*(1-2*hA)) - 2*gammaB * x*(x1+x3)*(hB+(x1+x3)*(1-2*hB)) + rho/2. * (x1*(1-x1-x-x3) - x*x3)
                    for jj in np.where(U01[ii,:,kk] == 1)[0][:-1]:
                        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,:,kk] == 1)[0][:-1][-1]:
                            A[jj,jj-1] += 1/dx[jj] * ( - M[jj-1] ) / 2
                            A[jj,jj] += 1/dx[jj] * ( - 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

                if ii == 0 and kk == 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]
                    x3 = x[kk]
                    M = 2*gammaA * x*(1-x1-x)*(hA+(x1+x)*(1-2*hA)) - 2*gammaB * x*(x1+x3)*(hB+(x1+x3)*(1-2*hB)) + rho/2. * (x1*(1-x1-x-x3) - x*x3)
                    for jj in range(len(x)):
                        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 == len(x)-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,kk,0,:] = np.concatenate(( np.array([0]), np.diagonal(A,-1) ))
                P[ii,kk,1,:] = np.diagonal(A)
                P[ii,kk,2,:] = np.concatenate(( np.diagonal(A,1), np.array([0]) ))
        return P

    def transition3(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
        """
        Transition matrix for ADI method along third axis
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = np.zeros((len(x),len(x),3,len(x)))
        for ii in range(len(x)):
            for jj in range(len(x)):
                if len(np.where(U01[ii,jj,:] == 1)[0]) <= 2:
                    continue
                A = np.zeros((len(x),len(x)))
                # drift
                if (ii > 0 or jj > 0) and x[ii] + x[jj] < 1 - tol:
                    V = x*(1-x)/nu
                    for kk in np.where(U01[ii,jj,:] == 1)[0][:-1]:
                        if kk == 0:
                            A[kk,kk] =  - 1/(2*dx[kk]) * ( -V[kk]/(x[kk+1]-x[kk]) )
                            A[kk,kk+1] = - 1/(2*dx[kk]) * ( V[kk+1]/(x[kk+1]-x[kk]) )
                        elif kk == np.where(U01[ii,jj,:] == 1)[0][:-1][-1]:
                            A[kk,kk-1] = - 1/(2*dx[kk]) * ( V[kk-1]/(x[kk]-x[kk-1]) )
                            A[kk,kk] = - 1/(2*dx[kk]) * ( -V[kk]/(x[kk]-x[kk-1]) )
                        else:
                            A[kk,kk-1] = - 1/(2*dx[kk]) * ( V[kk-1]/(x[kk]-x[kk-1]) )
                            A[kk,kk] = - 1/(2*dx[kk]) * ( -V[kk]/(x[kk]-x[kk-1]) -V[kk]/(x[kk+1]-x[kk]) )
                            A[kk,kk+1] = - 1/(2*dx[kk]) * ( V[kk+1]/(x[kk+1]-x[kk]) )

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

                if ii == 0 and jj == 0:
                    V = x*(1-x)/nu
                    for kk in range(len(x)):
                        if kk == 0:
                            A[kk,kk] =  - 1/(2*dx[kk]) * ( -V[kk]/(x[kk+1]-x[kk]) )
                            A[kk,kk+1] = - 1/(2*dx[kk]) * ( V[kk+1]/(x[kk+1]-x[kk]) )
                        elif kk == len(x)-1:
                            A[kk,kk-1] = - 1/(2*dx[kk])*2 * ( V[kk-1]/(x[kk]-x[kk-1]) )
                            A[kk,kk] = - 1/(2*dx[kk])*2 * ( -V[kk]/(x[kk]-x[kk-1]) )
                        else:
                            A[kk,kk-1] = - 1/(2*dx[kk]) * ( V[kk-1]/(x[kk]-x[kk-1]) )
                            A[kk,kk] = - 1/(2*dx[kk]) * ( -V[kk]/(x[kk]-x[kk-1]) -V[kk]/(x[kk+1]-x[kk]) )
                            A[kk,kk+1] = - 1/(2*dx[kk]) * ( V[kk+1]/(x[kk+1]-x[kk]) )

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

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

    def transition12(x, dx, U01):
        """
        Covariance discretization for slice along axes 1 and 2
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = []*len(x)
        for kk in range(len(x)):
            C = lil_matrix((len(x)**2,len(x)**2))
            if kk < len(x) - 4:
                for ii in range(len(x)-1)[:-1]:
                    for jj in range(len(x)-1)[:-1]:
                        if U01[ii+2,jj+2,kk] == 1 or U01[ii+1,jj+2,kk] == 1 or U01[ii+2,jj+1,kk] == 1:
                            if ii+1 < len(x)-kk and jj+1 < len(x)-kk and U01[ii+1,jj+1,kk] == 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)-kk and jj-1 >= 0 and U01[ii+1,jj-1,kk] == 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)-kk and U01[ii-1,jj+1,kk] == 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,kk] == 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,kk] == 1 or U01[ii+1,jj,kk] == 1 or U01[ii,jj+1,kk] == 1:
                            if ii+1 < len(x)-kk and jj-1 >= 0 and U01[ii+1,jj-1,kk] == 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)-kk and U01[ii-1,jj+1,kk] == 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,kk] == 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 kk == 0:
                ii = 0
                jj = len(x)-2-kk
                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-kk
                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

            P.append(C)
        return P

    def transition13(x, dx, U01):
        """
        Covariance discretization for slice along axes 1 and 3
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = []*len(x)
        for jj in range(len(x)):
            C = lil_matrix((len(x)**2,len(x)**2))
            if jj < len(x) - 4:
                for ii in range(len(x)-1)[:-1]:
                    for kk in range(len(x)-1)[:-1]:
                        if U01[ii+2,jj,kk+2] == 1 or U01[ii+1,jj,kk+2] == 1 or U01[ii+2,jj,kk+1] == 1:
                            if ii+1 < len(x)-jj and kk+1 < len(x)-jj and U01[ii+1,jj,kk+1] == 1:
                                C[ii*len(x)+kk,(ii+1)*len(x)+(kk+1)] += 1./4 * 1./(dx[ii]*dx[kk]) * (-x[ii+1]*x[kk+1])

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

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

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

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

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

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

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

                ii = len(x)-2-jj
                kk = 0
                C[ii*len(x)+kk,(ii-1)*len(x)+(kk+1)] += -1./4 * 1./(dx[ii]*dx[kk]) * (-x[ii-1]*x[kk+1]) / 2

            P.append(C)
        return P

    def transition23(x, dx, U01):
        """
        Covariance discretization for slice along axes 2 and 3
        """
        ### XXX: We attempt to import cythonized versions of these methods below. Note: changes should be made to
        #        this version and the cython version together.
        P = []*len(x)
        for ii in range(len(x)):
            C = lil_matrix((len(x)**2,len(x)**2))
            if ii < len(x) - 4:
                for jj in range(len(x)-1)[:-1]:
                    for kk in range(len(x)-1)[:-1]:
                        if U01[ii,jj+2,kk+2] == 1 or U01[ii,jj+1,kk+2] == 1 or U01[ii,jj+2,kk+1] == 1:
                            if jj+1 < len(x)-ii and kk+1 < len(x)-ii and U01[ii,jj+1,kk+1] == 1:
                                C[jj*len(x)+kk,(jj+1)*len(x)+(kk+1)] += 1./4 * 1./(dx[jj]*dx[kk]) * (-x[jj+1]*x[kk+1])

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

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

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

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

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

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

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

                jj = len(x)-2-ii
                kk = 0
                C[jj*len(x)+kk,(jj-1)*len(x)+(kk+1)] += -1./4 * 1./(dx[jj]*dx[kk]) * (-x[jj-1]*x[kk+1]) / 2

            P.append(C)
        return P

    """
    Transition matrices and methods for interaction with non-square face.
    Note that this is more complicated than the triallele methods,
    as density can re-enter the domain through recombination; that is,
    the domain boundaries are not necessarily absorbing!
    """

    def phi_to_surf(phi,x):
        """
        non-square  mapped 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
        """
        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

    def surf_to_phi(surf,phi,x):
        """
        move the surf density back to the full phi 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

    def move_density_to_surface(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
        """
        for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection
        x - grid
        dx - grid spacing
        dt - timestep of integration
        nu - relative population size
        """
        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

    def surface_interaction_b(phi,x,Psurf):
        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

    def surface_interaction(phi,x,Psurf):
        """
        density that should be moved to surface
        """
        ### 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))[: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

    def surface_recombination(phi,x,rho,dt):
        """
        recombination pushes density from surface back into interior of domain
        """
        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

    # triallele methods
    # new gammas: \tilde{gamma1} = gammaA(1-x1-x2)/(1-x1) + gammaB(x2)/(1-x1)
    # \tilde{gamma2} = (gammaA-gammaB)(1-x1-x2)/(1-x2) - gammaB(x1)(1-x2)
    # use the cythonized methods from the triallele code
    """
    forward integration
    """

    def advance_adi1(phi,U01,P1,x):
        """
        ADI integration along axis 1 of phi
        phi - density function
        U01 - domain markers
        P1 - ADI transition matrices
        x - grid
        """
        ### 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

    def advance_adi2(phi,U01,P2,x):
        """
        ADI integration along axis 2 of phi
        phi - density function
        U01 - domain markers
        P2 - ADI transition matrices
        x - grid
        """
        ### 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

    def advance_adi3(phi,U01,P3,x):
        """
        ADI integration along axis 3 of phi
        phi - density function
        U01 - domain markers
        P3 - ADI transition matrices
        x - grid
        """
        ### 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

    def advance_adi(phi,U01,P1,P2,P3,x,ii):
        """
        Combined ADI integration of phi
        phi - density function
        U01 - domain markers
        P1/P2/P3 - ADI transition matrices
        x - grid
        """
        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

    def advance_cov12(phi,C12,x):
        """
        Explicit integration of covariance term in the 1/2 axes planes
        phi - density function
        C12 - transition matrices
        x - grid
        """
        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

    def advance_cov13(phi,C13,x):
        """
        Explicit integration of covariance term in the 1/3 axes planes
        phi - density function
        C13 - transition matrices
        x - grid
        """
        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

    def advance_cov23(phi,C23,x):
        """
        Explicit integration of covariance term in the 2/3 axes planes
        phi - density function
        C23 - transition matrices
        x - grid
        """
        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

    def advance_cov(phi,C12,C13,C23,x,ii):
        """
        Combined integration for the covariance terms
        phi - density function
        C12/C13/C23 - transition matrices
        x - grid
        """
        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

    # methods for integration of the surface - see triallele methods for these (cythonized in triallele code for the transition matrices)

    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
        """
        tol = 1e-12
        U01 = np.ones((len(x),len(x)))
        XX = x[:,nuax] + x[nuax,:]
        U01[np.where(XX > 1+tol)] = 0
        return U01

    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
        """
        DXX = dx[:,nuax]*dx[nuax,:]
        for ii in range(len(x)):
            DXX[ii,len(x)-ii-1] *= 1./2
        return DXX

    def int2(DXX,U):
        """
        Integrate the density function over the domain
        DXX - two dimensional grid
        U - density function
        """
        return np.sum(DXX*U)

    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
        x - grid
        dx - grid spacing
        U01 - surface domain marker
        gammaA/B - scaled selection coefficients
        nu - relative population size
        """
        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

    def transition2_surf(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
        """
        adi transition matrix for x2 on surface
        x - grid
        dx - grid spacing
        U01 - surface domain marker
        gammaA/B - scaled selection coefficients
        nu - relative population size
        """
        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

    def transition12_surf(x,dx,U01):
        """
        transition matrix for explicit integration of 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

    def transition1D(x, dx, dt, gamma, nu):
        """
        transition matrix for one dimenional density function
        x - grid
        dx - grid spacing
        dt - timestep of integration
        gamma - scaled selection coefficients
        nu - relative size of population
        """
        P = np.zeros((len(x),len(x)))
        for ii in range(len(x)):
            if ii == 0:
                P[ii,ii] = 1 + dt/nu * 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                P[ii,ii+1] = - dt/nu * 1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii]) + gamma * dt / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
            elif ii == len(x) - 1:
                P[ii,ii-1] = - dt/nu * 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1]) - gamma * dt / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                P[ii,ii] = 1 + dt/nu * 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1])
            else:
                P[ii,ii-1] = - dt/nu * 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1]) - gamma * dt / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                P[ii,ii] = 1 + dt/nu * 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1]) + dt/nu * 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                P[ii,ii+1] = - dt/nu * 1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii]) + gamma * dt / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
        return P

    def advance_surf_adi1(surf,U01surf,P1surf,x):
        """
        Advance the adi method along first axis one of surface domain
        surf - density function along that surface
        U01surf - domain markers along grid
        P1surf - adi transition matrix
        x - uniform grid
        """
        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

    def advance_surf_adi2(surf,U01surf,P2surf,x):
        """
        Advance the adi method along second axis one of surface domain
        surf - density function along that surface
        U01surf - domain markers along grid
        P2surf - adi transition matrix
        x - uniform grid
        """
        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

    def advance_surf_cov(surf,Csurf,x):
        """
        Explicit integration of the covariance term, using scipy's sparse matrix for Csurf
        surf - density function along the surface
        Csurf - transition matrix
        x - grid
        """
        surf = ( Csurf * surf.reshape(len(x)**2)).reshape(len(x),len(x))
        return surf

    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
        """
        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

    def advance1D(u,P):
        """
        Given transition matrix P, use dadi's triagonal solver to integrate in 1D
        u - 1D density function
        P - transition matrix
        """
        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

    def move_surf_to_line(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
        """
        for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection
        x - grid
        dx - grid spacing
        dt - timestep of integration
        gammaA/gammaB - scaled selection coefficients
        nu - relative population size
        """
        #### 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

    def advance_surface(phi,x,P1surf,P2surf,Csurf,Pline,P,U01surf):
        """
        phi - full density function
        x - grid
        P1surf,P2surf - ADI transition matrices for triallele density function
        Csurf - covarance transition matrix
        Pline - transition matrix for diagonal boundary of triallele surface
        P - stores the ammount of density that should be moved from triallele domain (surf) to diagonal line boundary
        """
        # 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

    ####
    # sampling methods

    sample_cache = {}
    def sample_cached(phi, ns, x):
        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

    def quadrinomial(ns,ii,jj,kk):
        """
        returns ns! / (ii! * jj! * kk! * (ns-ii-jj-kk)!)
        """
        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))

    def ln_binomial(n,k):
        return math.lgamma(n+1) - math.lgamma(k+1) - math.lgamma(n-k+1)

    def fold_ancestral(F):
        """
        Takes an unfolded two locus frequency spectrum and folds it (assume don't konw ancestral state)
        """
        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

    def fold_lr(F):
        """
        Takes an unfolded two locus spectrum and folds based on left/right allele
        """
        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

    def extrap_dt_pts(temps):
        # 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]))

    def array_to_spectrum(phi):
        """
        now handled by TLSpectrum_mod.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

    def to_single_locus(phi):
        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

    def mean_r2(F):
        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)

    # gets too big when I need to cache hundreds of different sample sizes
    #projection_cache = {}
    def cached_projection(proj_to, proj_from, hits):
        """
        Coefficients for projection from a larger size to smaller
        proj_to: Number of samples to project down to
        proj_from: Number of samples to project from
        hits: Number of derived alleles projecting from - tuple of (n1,n2,n3)
        """
    #    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

    def project(F_from, proj_to):
        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)

    ## for projection to genotype spectrum
    def pairings(n):
        return math.factorial(n)/(math.factorial(n/2)*2**(n/2))

    def binom(n,k):
        return math.factorial(n)/math.factorial(k)/math.factorial(n-k)

    def multinomial(n,i,j,k):
        if i+j+k != n:
            return "nein"
        return math.factorial(n)/math.factorial(i)/math.factorial(j)/math.factorial(k)

    def genotypes_prob_4(n,colors,hits):
        (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)

    def possible_genotypes_4(n,colors):
        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

    prob_cache = {} # store possible genotypes and their relative probabilities among all possible genotype sets
    def cached_genotype_exact_projection(n,haplotype_counts):
        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]

    def genotype_spectrum_from_F(F):
        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

    def observed_genotype_spectrum_from_F(F):
        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

    def observed_genotype_spectrum_dict_from_F(F):
        """
        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

    from . import projection_genotypes

    genotype_projection_cache = {} # this cache can get very large for n_from large (>~100)
    def projection_cache_Gdict(n_from,n_to,hits):
        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

    def project_Gdict(G,n_from,n_to):
        """ 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

    def misidentification(F,p):
        """
        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)

    def misidentification_genotype_dict(G,p):
        """
        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

    def genotype_exp_data_to_arrays(G_exp,G_data):
        """
        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

    def condition(F,i):
        """
        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

    # Try importing cythonized versions of several slow methods. These imports should overwrite the Python code defined above.
    try:
        from .transition1D import transition1D
        from .transition1 import transition1
        from .transition2 import transition2
        from .transition3 import transition3
        from .transition12 import transition12
        from .transition13 import transition13
        from .transition23 import transition23
        from .surface_interaction import surface_interaction
    except ImportError:
        print("using numpy versions")
        pass

## Functions

`def LD_per_bin(ns)`

  LD statistics per bin for a TLSpectrum.

  ns: Number of samples

  Returns (D, r2). Both are TLSpectrum objects in which each entry is the value of D or r^2 for that combination of haplotypes.

  Expand source code

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

          ns: Number of samples

          Returns (D, r2). Both are TLSpectrum objects in which each entry is the
          value of D or 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

`def advance1D(u, P)`

  Given transition matrix P, use dadi's triagonal solver to integrate in 1D u - 1D density function P - transition matrix

  Expand source code

      def advance1D(u,P):
          """
          Given transition matrix P, use dadi's triagonal solver to integrate in 1D
          u - 1D density function
          P - transition matrix
          """
          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

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

  Combined ADI integration of phi phi - density function U01 - domain markers P1/P2/P3 - ADI transition matrices x - grid

  Expand source code

      def advance_adi(phi,U01,P1,P2,P3,x,ii):
          """
          Combined ADI integration of phi
          phi - density function
          U01 - domain markers
          P1/P2/P3 - ADI transition matrices
          x - grid
          """
          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

`def advance_adi1(phi, U01, P1, x)`

  ADI integration along axis 1 of phi phi - density function U01 - domain markers P1 - ADI transition matrices x - grid

  Expand source code

      def advance_adi1(phi,U01,P1,x):
          """
          ADI integration along axis 1 of phi
          phi - density function
          U01 - domain markers
          P1 - ADI transition matrices
          x - grid
          """
          ### 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

`def advance_adi2(phi, U01, P2, x)`

  ADI integration along axis 2 of phi phi - density function U01 - domain markers P2 - ADI transition matrices x - grid

  Expand source code

      def advance_adi2(phi,U01,P2,x):
          """
          ADI integration along axis 2 of phi
          phi - density function
          U01 - domain markers
          P2 - ADI transition matrices
          x - grid
          """
          ### 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

`def advance_adi3(phi, U01, P3, x)`

  ADI integration along axis 3 of phi phi - density function U01 - domain markers P3 - ADI transition matrices x - grid

  Expand source code

      def advance_adi3(phi,U01,P3,x):
          """
          ADI integration along axis 3 of phi
          phi - density function
          U01 - domain markers
          P3 - ADI transition matrices
          x - grid
          """
          ### 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

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

  Combined integration for the covariance terms phi - density function C12/C13/C23 - transition matrices x - grid

  Expand source code

      def advance_cov(phi,C12,C13,C23,x,ii):
          """
          Combined integration for the covariance terms
          phi - density function
          C12/C13/C23 - transition matrices
          x - grid
          """
          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

`def advance_cov12(phi, C12, x)`

  Explicit integration of covariance term in the 1/2 axes planes phi - density function C12 - transition matrices x - grid

  Expand source code

      def advance_cov12(phi,C12,x):
          """
          Explicit integration of covariance term in the 1/2 axes planes
          phi - density function
          C12 - transition matrices
          x - grid
          """
          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

`def advance_cov13(phi, C13, x)`

  Explicit integration of covariance term in the 1/3 axes planes phi - density function C13 - transition matrices x - grid

  Expand source code

      def advance_cov13(phi,C13,x):
          """
          Explicit integration of covariance term in the 1/3 axes planes
          phi - density function
          C13 - transition matrices
          x - grid
          """
          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

`def advance_cov23(phi, C23, x)`

  Explicit integration of covariance term in the 2/3 axes planes phi - density function C23 - transition matrices x - grid

  Expand source code

      def advance_cov23(phi,C23,x):
          """
          Explicit integration of covariance term in the 2/3 axes planes
          phi - density function
          C23 - transition matrices
          x - grid
          """
          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

`def advance_surf_adi1(surf, U01surf, P1surf, x)`

  Advance the adi method along first axis one of surface domain surf - density function along that surface U01surf - domain markers along grid P1surf - adi transition matrix x - uniform grid

  Expand source code

      def advance_surf_adi1(surf,U01surf,P1surf,x):
          """
          Advance the adi method along first axis one of surface domain
          surf - density function along that surface
          U01surf - domain markers along grid
          P1surf - adi transition matrix
          x - uniform grid
          """
          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

`def advance_surf_adi2(surf, U01surf, P2surf, x)`

  Advance the adi method along second axis one of surface domain surf - density function along that surface U01surf - domain markers along grid P2surf - adi transition matrix x - uniform grid

  Expand source code

      def advance_surf_adi2(surf,U01surf,P2surf,x):
          """
          Advance the adi method along second axis one of surface domain
          surf - density function along that surface
          U01surf - domain markers along grid
          P2surf - adi transition matrix
          x - uniform grid
          """
          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

`def advance_surf_cov(surf, Csurf, x)`

  Explicit integration of the covariance term, using scipy's sparse matrix for Csurf surf - density function along the surface Csurf - transition matrix x - grid

  Expand source code

      def advance_surf_cov(surf,Csurf,x):
          """
          Explicit integration of the covariance term, using scipy's sparse matrix for Csurf
          surf - density function along the surface
          Csurf - transition matrix
          x - grid
          """
          surf = ( Csurf * surf.reshape(len(x)**2)).reshape(len(x),len(x))
          return surf

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

  phi - full density function x - grid P1surf,P2surf - ADI transition matrices for triallele density function Csurf - covarance transition matrix Pline - transition matrix for diagonal boundary of triallele surface P - stores the ammount of density that should be moved from triallele domain (surf) to diagonal line boundary

  Expand source code

      def advance_surface(phi,x,P1surf,P2surf,Csurf,Pline,P,U01surf):
          """
          phi - full density function
          x - grid
          P1surf,P2surf - ADI transition matrices for triallele density function
          Csurf - covarance transition matrix
          Pline - transition matrix for diagonal boundary of triallele surface
          P - stores the ammount of density that should be moved from triallele domain (surf) to diagonal line boundary
          """
          # 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

`def array_to_spectrum(phi)`

  now handled by TLSpectrum_mod.TLSpectrum

  Expand source code

      def array_to_spectrum(phi):
          """
          now handled by TLSpectrum_mod.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

`def binom(n, k)`
  Expand source code

      def binom(n,k):
          return math.factorial(n)/math.factorial(k)/math.factorial(n-k)

`def cached_genotype_exact_projection(n, haplotype_counts)`
  Expand source code

      def cached_genotype_exact_projection(n,haplotype_counts):
          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]

`def cached_projection(proj_to, proj_from, hits)`

  Coefficients for projection from a larger size to smaller proj_to: Number of samples to project down to proj_from: Number of samples to project from hits: Number of derived alleles projecting from - tuple of (n1,n2,n3)

  Expand source code

      def cached_projection(proj_to, proj_from, hits):
          """
          Coefficients for projection from a larger size to smaller
          proj_to: Number of samples to project down to
          proj_from: Number of samples to project from
          hits: Number of derived alleles projecting from - tuple of (n1,n2,n3)
          """
      #    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

`def condition(F, i)`

  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

  Expand source code

      def condition(F,i):
          """
          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

`def domain(x)`

  Array of same dimension as density function discretization Has a 1 if grid point lies in domain, 0 if outside domain x - 1D grid

  Expand source code

      def domain(x):
          """
          Array of same dimension as density function discretization
          Has a 1 if grid point lies in domain, 0 if outside domain
          x - 1D grid
          """
          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

`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

  Expand source code

      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
          """
          tol = 1e-12
          U01 = np.ones((len(x),len(x)))
          XX = x[:,nuax] + x[nuax,:]
          U01[np.where(XX > 1+tol)] = 0
          return U01

`def extrap_dt_pts(temps)`
  Expand source code

      def extrap_dt_pts(temps):
          # 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]))

`def fold_ancestral(F)`

  Takes an unfolded two locus frequency spectrum and folds it (assume don't konw ancestral state)

  Expand source code

      def fold_ancestral(F):
          """
          Takes an unfolded two locus frequency spectrum and folds it (assume don't konw ancestral state)
          """
          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

`def fold_lr(F)`

  Takes an unfolded two locus spectrum and folds based on left/right allele

  Expand source code

      def fold_lr(F):
          """
          Takes an unfolded two locus spectrum and folds based on left/right allele
          """
          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

`def genotype_exp_data_to_arrays(G_exp, G_data)`

  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

  Expand source code

      def genotype_exp_data_to_arrays(G_exp,G_data):
          """
          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

`def genotype_spectrum_from_F(F)`
  Expand source code

      def genotype_spectrum_from_F(F):
          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

`def genotypes_prob_4(n, colors, hits)`
  Expand source code

      def genotypes_prob_4(n,colors,hits):
          (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)

`def grid(numpts)`

  Default grid is uniform, so that grid points lie directly on the boundaries Here, delta_x = 1./numpts numpts - number of grid points (less one)

  Expand source code

      def grid(numpts):
          """
          Default grid is uniform, so that grid points lie directly on the boundaries
          Here, delta_x = 1./numpts
          numpts - number of grid points (less one)
          """
          return np.linspace(0,1,numpts+1)

`def grid_dx(x)`

  Given 1D grid x, this returns the 1D grid spacing for x

  Expand source code

      def grid_dx(x):
          """
          Given 1D grid x, this returns the 1D grid spacing for x
          """
          return (np.concatenate((np.diff(x),np.array([0]))) + np.concatenate((np.array([0]),np.diff(x))))/2

`def grid_dx3(x, dx)`

  For integrating over the 3D domain (see int3) x - 1D grid dx - 1D grid spacing

  Expand source code

      def grid_dx3(x,dx):
          """
          For integrating over the 3D domain (see int3)
          x - 1D grid
          dx - 1D 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

`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

  Expand source code

      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
          """
          DXX = dx[:,nuax]*dx[nuax,:]
          for ii in range(len(x)):
              DXX[ii,len(x)-ii-1] *= 1./2
          return DXX

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

  Injection new derived mutations A onto background of B/b x - 1D grid dx - 1D grid spacing dt - integration time step yA - biallelic frequency spectrum integrated by dadi phi - density function thetaA - scaled mutation rate

  Expand source code

      def injectA(x,dx,dt,yB,phi,thetaA):
          """
          Injection new derived mutations A onto background of B/b
          x - 1D grid
          dx - 1D grid spacing
          dt - integration time step
          yA - biallelic frequency spectrum integrated by dadi
          phi - density function
          thetaA - scaled mutation rate
          """
          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

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

  Injection new derived mutations B onto background of A/a x - 1D grid dx - 1D grid spacing dt - integration time step yB - biallelic frequency spectrum integrated by dadi phi - density function thetaB - scaled mutation rate

  Expand source code

      def injectB(x,dx,dt,yA,phi,thetaB):
          """
          Injection new derived mutations B onto background of A/a
          x - 1D grid
          dx - 1D grid spacing
          dt - integration time step
          yB - biallelic frequency spectrum integrated by dadi
          phi - density function
          thetaB - scaled mutation rate
          """
          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

`def int2(DXX, U)`

  Integrate the density function over the domain DXX - two dimensional grid U - density function

  Expand source code

      def int2(DXX,U):
          """
          Integrate the density function over the domain
          DXX - two dimensional grid
          U - density function
          """
          return np.sum(DXX*U)

`def int3(DX, U)`

  Numerically integrate the density function phi phi - density function DX - 3D grid spacing for integration weights

  Expand source code

      def int3(DX,U):
          """
          Numerically integrate the density function phi
          phi - density function
          DX - 3D grid spacing for integration weights
          """
          return np.sum(DX*U)

`def ln_binomial(n, k)`
  Expand source code

      def ln_binomial(n,k):
          return math.lgamma(n+1) - math.lgamma(k+1) - math.lgamma(n-k+1)

`def mean_r2(F)`
  Expand source code

      def mean_r2(F):
          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)

`def misidentification(F, p)`

  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

  Expand source code

      def misidentification(F,p):
          """
          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)

`def misidentification_genotype_dict(G, p)`

  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),

  Expand source code

      def misidentification_genotype_dict(G,p):
          """
          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

`def move_density_to_surface(x, dx, dt, gammaA, gammaB, nu, hA=0.5, hB=0.5)`

  for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection x - grid dx - grid spacing dt - timestep of integration nu - relative population size

  Expand source code

      def move_density_to_surface(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
          """
          for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection
          x - grid
          dx - grid spacing
          dt - timestep of integration
          nu - relative population size
          """
          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

`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

  Expand source code

      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
          """
          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

`def move_surf_to_line(x, dx, dt, gammaA, gammaB, nu, hA=0.5, hB=0.5)`

  for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection x - grid dx - grid spacing dt - timestep of integration gammaA/gammaB - scaled selection coefficients nu - relative population size

  Expand source code

      def move_surf_to_line(x,dx,dt,gammaA,gammaB,nu,hA=1./2,hB=1./2):
          """
          for each point in the full domain, how much should be lost to the nonsquare surface boundary due to diffusion/selection
          x - grid
          dx - grid spacing
          dt - timestep of integration
          gammaA/gammaB - scaled selection coefficients
          nu - relative population size
          """
          #### 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

`def multinomial(n, i, j, k)`
  Expand source code

      def multinomial(n,i,j,k):
          if i+j+k != n:
              return "nein"
          return math.factorial(n)/math.factorial(i)/math.factorial(j)/math.factorial(k)

`def observed_genotype_spectrum_dict_from_F(F)`

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

  Expand source code

      def observed_genotype_spectrum_dict_from_F(F):
          """
          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

`def observed_genotype_spectrum_from_F(F)`
  Expand source code

      def observed_genotype_spectrum_from_F(F):
          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

`def pairings(n)`
  Expand source code

      def pairings(n):
          return math.factorial(n)/(math.factorial(n/2)*2**(n/2))

`def phi_to_surf(phi, x)`

  non-square mapped 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

  Expand source code

      def phi_to_surf(phi,x):
          """
          non-square  mapped 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
          """
          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

`def possible_genotypes_4(n, colors)`
  Expand source code

      def possible_genotypes_4(n,colors):
          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

`def project(F_from, proj_to)`
  Expand source code

      def project(F_from, proj_to):
          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)

`def project_Gdict(G, n_from, n_to)`

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

  Expand source code

      def project_Gdict(G,n_from,n_to):
          """ 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

`def projection_cache_Gdict(n_from, n_to, hits)`
  Expand source code

      def projection_cache_Gdict(n_from,n_to,hits):
          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

`def quadrinomial(ns, ii, jj, kk)`

  returns ns! / (ii! * jj! * kk! * (ns-ii-jj-kk)!)

  Expand source code

      def quadrinomial(ns,ii,jj,kk):
          """
          returns ns! / (ii! * jj! * kk! * (ns-ii-jj-kk)!)
          """
          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))

`def sample_cached(phi, ns, x)`
  Expand source code

      def sample_cached(phi, ns, x):
          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

`def surf_to_phi(surf, phi, x)`

  move the surf density back to the full phi density function

  Expand source code

      def surf_to_phi(surf,phi,x):
          """
          move the surf density back to the full phi 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

`def surface_interaction_b(phi, x, Psurf)`
  Expand source code

      def surface_interaction_b(phi,x,Psurf):
          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

`def surface_recombination(phi, x, rho, dt)`

  recombination pushes density from surface back into interior of domain

  Expand source code

      def surface_recombination(phi,x,rho,dt):
          """
          recombination pushes density from surface back into interior of domain
          """
          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

`def to_single_locus(phi)`
  Expand source code

      def to_single_locus(phi):
          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

`def transition12_surf(x, dx, U01)`

  transition matrix for explicit integration of mixed derivative term along surface

  Expand source code

      def transition12_surf(x,dx,U01):
          """
          transition matrix for explicit integration of 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

`def 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 x - grid dx - grid spacing U01 - surface domain marker gammaA/B - scaled selection coefficients nu - relative population size

  Expand source code

      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
          x - grid
          dx - grid spacing
          U01 - surface domain marker
          gammaA/B - scaled selection coefficients
          nu - relative population size
          """
          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

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

  adi transition matrix for x2 on surface x - grid dx - grid spacing U01 - surface domain marker gammaA/B - scaled selection coefficients nu - relative population size

  Expand source code

      def transition2_surf(x,dx,U01,gammaA,gammaB,rho,nu,hA=.5,hB=.5):
          """
          adi transition matrix for x2 on surface
          x - grid
          dx - grid spacing
          U01 - surface domain marker
          gammaA/B - scaled selection coefficients
          nu - relative population size
          """
          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