dadi.Triallele.numerics API documentation

# Module `dadi.Triallele.numerics`

Expand source code

    import numpy as np
    from scipy.sparse import lil_matrix
    from scipy.sparse import identity
    from scipy.sparse.linalg import spsolve
    from scipy.integrate import trapz
    from scipy.special import gamma
    from numpy import newaxis as nuax
    from dadi import Numerics
    from dadi.Spectrum_mod import Spectrum
    import dadi
    import math
    import pickle

    from dadi.Triallele.TriSpectrum_mod import TriSpectrum

    def grid_dx(x):
        """
        We use uniform grids in x, using np.linspace(0,1,numpts)
        Grid spacing Delta, which is halved at the first and last grid points.
        """
        return (np.concatenate((np.diff(x),np.array([0]))) + np.concatenate((np.array([0]),np.diff(x))))/2

    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 domain(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

    ### transition matrices

    def transition1(x, dx, U01, sig1, sig2):
        """
        Implicit transition matrix for the ADI components of the discretization of the diffusion
        Time scaled by 2N, with variance and mean terms x(1-x) and \sigma*x(1-x), resp.
        Store the tridiagonal elements of the matrices, which need to be adjusted by I + dt*P, where I is the identity matrix
        x - grid
        dx - grid spacing
        U01 - domain markers
        sig1/2 - population scaled selection coefficients
        nu - relative population size
        """
        # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
        #      be automatically used instead of this version, if the user has compiled it.
        print("using numpy, not cython")
        PV = np.zeros((len(x),3,len(x)))
        PM = 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)
                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)
                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]) )

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

            A = np.zeros((len(x),len(x)))
            x2 = x[jj]
            sig_new = sig1*(1-x-x2)/(1-x) + (sig1-sig2)*x2/(1-x)
            sig_new[-1] = 0
            M = sig_new*x*(1-x)
            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

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

    def transition2(x, dx, U01, sig1, sig2):
        """
        Implicit transition matrix for the ADI components of the discretization of the diffusion
        Time scaled by 2N, with variance and mean terms x(1-x) and \sigma*x(1-x), resp.
        Store the tridiagonal elements of the matrices, which need to be adjusted by I + dt*P, where I is the identity matrix
        x - grid
        dx - grid spacing
        U01 - domain markers
        sig1/2 - population scaled selection coefficients
        nu - relative population size
        """
        # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
        #      be automatically used instead of this version, if the user has compiled it.
        PV = np.zeros((len(x),3,len(x)))
        PM = 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)
                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)
                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]) )

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

            A = np.zeros((len(x),len(x)))
            x1 = x[ii]
            sig_new = sig2*(1-x-x1)/(1-x) + (sig2-sig1)*x1/(1-x)
            sig_new[-1] = 0
            M = sig_new*x*(1-x)
            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

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

    def transition12(x, dx, U01):
        """
        Transition matrix for the covariance term of the diffusion operator, with term D_{xy} (-x*y*phi)
        As with the ADI components, final transition matrix is given by I + dt/nu*P
        x - grid
        dx - grid spacing
        U01 - domain markers
        """
        # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
        #      be automatically used instead of this version, if the user has compiled it.
        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, sig):
        """
        transition matrix for one dimensional integration
        x - grid
        dx - grid spacing
        dt - timestep for integration
        sig - selection coefficient
        nu - relative population size
        """
        # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
        #      be automatically used instead of this version, if the user has compiled it.
        PV = np.zeros((len(x),len(x)))
        PM = np.zeros((len(x),len(x)))
        for ii in range(len(x)):
            if ii == 0:
                PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                PV[ii,ii+1] = -1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii])
                PM[ii,ii+1] = sig / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
            elif ii == len(x) - 1:
                PV[ii,ii-1] = - 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1])
                PM[ii,ii-1] = sig / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1])
            else:
                PV[ii,ii-1] = - 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1])
                PM[ii,ii-1] = - sig / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1]) + 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                PV[ii,ii+1] = - 1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii])
                PM[ii,ii+1] = sig / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
        return PV,PM

    transition1D_cache = {}
    def cached_transition1D(numpts,sig):
        key = (numpts,sig)
        try:
            return transition1D_cache[key]
        except KeyError:
            pass
        x = np.linspace(0,1,numpts+1)
        dx = grid_dx(x)
        V,M = transition1D(x,dx,sig)
        transition1D_cache[key] = [V,M]
        return transition1D_cache[key]

    def remove_diag_density_weights_nonneutral(x,dt,nu,sig1,sig2):
        """
        Numerically determine the amount of density that should be lost to the diagonal boundary.
        Numerically integrate 1D array with initial point mass at z0, where z0 is the frequency x+y, integrated for time step dt.
        We then check the fraction of density that is lost to z=1.
        If sig1 or sig2 are nonzero, estimate the selection pressure on z as sig = sig1*x/(x+y) + sig2*y/(x+y)
        x - one dimensional grid of domain
        dt - time step of integration
        nu - relative population size
        sig1,sig2 - selection coefficients
        """
        dx = grid_dx(x)
        P = np.zeros((len(x),len(x)))
        for ii in range(len(x)):
            for jj in range(len(x)):
                if x[ii]+x[jj] < 1.0 and x[ii]+x[jj] > .75:
                    sig = sig1*x[ii]/(x[ii]+x[jj]) + sig2*x[jj]/(x[ii]+x[jj])
                    V,M = cached_transition1D(len(x)-1,sig)
                    P1D = np.eye(len(x)) + dt*(V/nu+M)
                    y = np.zeros(len(x))
                    y[ii+jj] = 1./dx[ii+jj]
                    y = advance1D(y,P1D)
                    prob = y[-1]*dx[-1]
                    P[ii,jj] = prob
                    if ii+jj == len(x)-2:
                        if ii==1:
                            V,M = cached_transition1D(len(x)-1,sig1)
                            P1D = np.eye(len(x)) + dt*(V/nu+M)
                            y = np.zeros(len(x))
                            y[ii] = 1./dx[ii]
                            y = advance1D(y,P1D)
                            prob2 = y[0]*dx[0]
                            P[ii,jj] += prob2
                        elif jj==1:
                            V,M = cached_transition1D(len(x)-1,sig2)
                            P1D = np.eye(len(x)) + dt*(V/nu+M)
                            y = np.zeros(len(x))
                            y[jj] = 1./dx[jj]
                            y = advance1D(y,P1D)
                            prob2 = y[0]*dx[0]
                            P[ii,jj] += prob2

        P[:,0] = 0
        P[0,:] = 0
        return P

    def move_density_to_bdry(x,phi,P):
        """
        P tells us how much should be removed, by multiplying phi*P
        Take that density and instead of deleting it, move straight to boundary
        P - stores how much denstity from each grid point should be moved to diagonal
        """
        for ii in range(len(x)):
            for jj in range(len(x)):
                if P[ii,jj] == 0:
                    continue
                else:
                    amnt = P[ii,jj]
                    s = ii+jj
                    if ii == 1 and jj == len(x)-3:
                        phi[ii+1,jj] += phi[ii,jj]*amnt/4. * 2
                        phi[ii,jj+1] += phi[ii,jj]*amnt/4. * 2
                        phi[ii-1,jj] += phi[ii,jj]*amnt/2 * 2
                        phi[ii,jj] *= (1-amnt)
                    elif ii == len(x)-3 and jj == 1:
                        phi[ii+1,jj] += phi[ii,jj]*amnt/4. * 2
                        phi[ii,jj+1] += phi[ii,jj]*amnt/4. * 2
                        phi[ii,jj-1] += phi[ii,jj]*amnt/2 * 2
                        phi[ii,jj] *= (1-amnt)
                    elif (len(x)-1-s) % 2 == 1:
                        # split between two points
                        dist = (len(x)-1-s) // 2
                        phi[ii+dist+1,jj+dist] += phi[ii,jj]*amnt/2. * 2
                        phi[ii+dist,jj+dist+1] += phi[ii,jj]*amnt/2. * 2
                        phi[ii,jj] *= (1-amnt)
                    else:
                        # straight to boundary grid point
                        dist = (len(x)-1-s) // 2
                        phi[ii+dist,jj+dist] += phi[ii,jj]*amnt * 2
                        phi[ii,jj] *= (1-amnt)

        return phi

    ### forward integration methods

    def advance_adi(U,U01,P1,P2,x,ii):
        """
        Integrate the ADI components forward in time, alternating which direction occurs first
        U - density function
        U01 - stores which points are in the domain
        P1,P2 - transition matrices
        ii - count of integration step
        """
        if np.mod(ii,2) == 0:
            for jj in range(len(x)):
                if np.sum(U01[:,jj]) > 1:
                    U[:,jj] = dadi.Integration.tridiag.tridiag(P1[jj,0,:],P1[jj,1,:],P1[jj,2,:],U[:,jj])
            for ii in range(len(x)):
                if np.sum(U01[ii,:]) > 1:
                    U[ii,:] = dadi.Integration.tridiag.tridiag(P2[ii,0,:],P2[ii,1,:],P2[ii,2,:],U[ii,:])
        else:
            for ii in range(len(x)):
                if np.sum(U01[ii,:]) > 1:
                    U[ii,:] = dadi.Integration.tridiag.tridiag(P2[ii,0,:],P2[ii,1,:],P2[ii,2,:],U[ii,:])
            for jj in range(len(x)):
                if np.sum(U01[:,jj]) > 1:
                    U[:,jj] = dadi.Integration.tridiag.tridiag(P1[jj,0,:],P1[jj,1,:],P1[jj,2,:],U[:,jj])
        return U

    def advance_cov(U,C,x,dx):
        """
        Explicit integration of the covariance term, using scipy's sparse matrix for C
        U - density function
        C - transition matrix
        """
        U = ( C * U.reshape(len(x)**2)).reshape(len(x),len(x))
        return U

    def advance1D(u,P):
        """
        tridiag breakdown for integration along the diagonal boundary
        u - density along that diagonal
        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.Integration.tridiag.tridiag(a,b,c,u)
        return u

    def advance_line(x,phi,P):
        """
        Integrate along the diagonal boundary. Density gets fixed along the boundary, and then diffuses along that boundary until being fixed in one of the two corners
        P - one dimensional transition matrix for the diagonal boundary
        """
        u = np.diag(np.fliplr(phi))
        u = advance1D(u,P)
        for ii in range(len(x)):
            phi[ii,len(x) - ii - 1] = u[ii]
        return phi

    ### sampling methods

    sample_cache = {}
    def sample(phi, ns, x):
        """
        Obtain the expected sample frequency spectrum from the density function
        """
        dx = grid_dx(x)
        DXX = grid_dx_2d(x,dx)

        # Assume ns is typically a list of length 1, if not, make it into one.
        try:
            ns = tuple(ns)
        except TypeError:
            ns = (ns,)

        # We cache calculations of several big matrices that will be re-used
        # within and between integrations.
        key = (ns, tuple(x))
        if key not in sample_cache:
            this_cache = {}
            for ii in range(1,ns[0]-1):
                # Create our cache
                this_cache[ii] = (1-x[:,nuax]-x[nuax,:])**ii
                # Somewhat ugly hack to use negative values to store second array
                # to cache.
                this_cache[-ii] = x[nuax,:]**ii
            sample_cache[key] = this_cache
        else:
            this_cache = sample_cache[key]

        #dx = grid_dx(x)
        F = np.zeros((ns[0]+1,ns[0]+1))
        prod_phi = DXX*phi
        for ii in range(len(F)):
            prod_x = prod_phi * x[:,nuax]**ii
            for jj in range(len(F)):
                if ii+jj < ns[0] and ii != 0 and jj != 0:
                    #F[ii,jj] = math.factorial(ns)/(math.factorial(ii)*math.factorial(jj)*math.factorial(ns-ii-jj)) * int2(x, dx, phi*x[:,nuax]**ii*x[nuax,:]**jj*(1-x[:,nuax]-x[nuax,:])**(ns-ii-jj) )
                    F[ii,jj] = trinomial(ns[0],ii,jj) * np.sum(prod_x * this_cache[-jj] * this_cache[ns[0]-ii-jj])
        F = TriSpectrum(F)
        F.folded_major = False
        F.folded_ancestral = False
        F.extrap_x = x[1]
        return F

    def trinomial(ns,ii,jj):
        """
        Return ns!/(ii! * jj! * (ns-ii-jj)!) for large values
        """
        return np.exp(math.lgamma(ns+1) - math.lgamma(ii+1) - math.lgamma(jj+1) - math.lgamma(ns-ii-jj+1))

    ### various methods for spectrum manipulation or other methods needed for data fitting

    def misidentification(F, p):
        """
        Given folded spectrum, and probability p that one of the derived alleles is the actual ancestral allele
        Then refold to return folded spectrum
        """
        F = TriF(np.zeros((len(F),len(F))))
        for ii in range(len(F))[1:-1]:
            for jj in range(len(F))[1:ii+1]:
                if ii+jj < len(F):
                    F[ii,jj] += (1 - p) * F[ii,jj]
                    F[len(F)-1-ii-jj,jj] += p/2. * F[ii,jj]
                    F[ii,len(F)-1-ii-jj] += p/2. * F[ii,jj]

        return F.fold()

    def fold(spectrum):
        """
        Note: this is now handled in the TriSpectrum class
        Given a frequency spectrum over the full domain, fold into a spectrum with major and minor derived alleles
        """
        spectrum = TriSpectrum(spectrum)
        if spectrum.folded_major == True:
            print("error: trying to fold a spectrum that is already folded")
            return spectrum
        else:
            spectrum = (spectrum + np.transpose(spectrum))
            for ii in range(len(spectrum)):
                spectrum[ii,ii] = spectrum[ii,ii]/2
            spectrum.mask[0,:] = True
            spectrum.mask[:,0] = True
            for ii in range(len(spectrum)):
                spectrum.mask[ii,ii+1:] = True
                spectrum.mask[ii,len(spectrum)-1-ii:] = True
            return spectrum

    def univariate_lognormal_pdf(x,sigma,mu):
        """
        Can compare to scipy.stats.lognorm.pdf(x,sigma,0,np.exp(mu))
        """
        return 1./(x*sigma*np.sqrt(2*math.pi)) * np.exp ( -(np.log(x) - mu)**2 / (2*sigma**2) )

    def bivariate_lognormal_pdf(xx, params):
        """
        mu_i = mu_yi and sigma_i = sigma_yi are the associated means and variances of the bivariate normal distr which gets exponentiated
        We assume for our application that mu1=mu2 and sigma1=sigma2, though this isn't necessary for the general bivariate lognormal distribution
        """
        mu1, mu2, sigma1, sigma2, rho = params
        norm = 1./(2 * np.pi * (sigma1*sigma2) * np.sqrt(1-rho**2) * np.outer(xx,xx) )
        q = 1/(1-rho**2) * ( ((np.log(xx[nuax,:])-mu1)/sigma1)**2 - 2*rho*((np.log(xx[nuax,:])-mu1)/sigma1)*((np.log(xx[:,nuax])-mu2)/sigma2) + ((np.log(xx[:,nuax])-mu2)/sigma2)**2 )
        prob = norm * np.exp( -q/2. )

        return prob

    def ms_demo_params_to_dadi(nu_ms,tau_ms):
        """
        convert from ms parameters (which use current pop size) for nu and tau to dadi parameters (which use ancestral pop size)
        """
        return 1./nu_ms,2*tau_ms/nu_ms

    def optimal_sfs_scaling(model,data):
        data = numerics.fold(data)
        model = numerics.fold(data)
        model, data = Numerics.intersect_masks(model, data)
        return data.sum()/model.sum()

    #def tri_spectrum(array):
    #    """
    #    takes array and outputs correctly masked triallele spectrum
    #    """
    #    F = dadi.Spectrum(array)
    #    F.mask[:,0] = True
    #    F.mask[0,:] = True
    #    for ii in range(len(F))[1:]:
    #        F.mask[ii,len(F)-ii-1:] = True
    #    return F

    def fold_ancestral(F):
        """
        Note: this is now handled by the TriSpectrum class
        Don't know ancestral state, so track minor frequencies
        Store spectrum of two minor allele frequencies
        """
        F_new = 0*F
        F_new = TriSpectrum(F_new)
        ns = len(F)-1
        for ii in range(ns):
            for jj in range(ns):
                kk = ns-ii-jj
                if F.mask[ii,jj] == True:
                    continue
                elif ii <= kk and jj <= kk:
                    if ii >= jj:
                        F_new[ii,jj] += F[ii,jj]
                    else:
                        F_new[jj,ii] += F[ii,jj]
                elif ii > kk and jj <= kk:
                    F_new[kk,jj] += F[ii,jj]
                elif ii <= kk and jj > kk:
                    F_new[kk,ii] += F[ii,jj]
                else: # ii > kk and jj > kk
                    if ii >= jj:
                        F_new[jj,kk] += F[ii,jj]
                    else:
                        F_new[ii,kk] += F[ii,jj]
        # mask if not a valid entry for ancestrally folded spectrum
        for ii in range(ns):
            for jj in range(ns):
                kk = ns-ii-jj
                if not (kk>=ii>=jj):
                    F_new.mask[ii,jj] = True
        return F_new

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

    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,n3)
        """
        key = (proj_to, proj_from, hits)
        try:
            return projection_cache[key]
        except KeyError:
            pass

        X1, X2 = hits
        X3 = proj_from - X1 - X2
        proj_weights = np.zeros((proj_to+1,proj_to+1))
        for ii in range(X1+1):
            for jj in range(X2+1):
                kk = proj_to - ii - jj
                if kk > X3 or kk <0:
                    continue
                f = ln_binomial(X1,ii) + ln_binomial(X2,jj) + ln_binomial(X3,kk) - ln_binomial(proj_from,proj_to)
                proj_weights[ii,jj] = 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('sorry, but projection must be to smaller size!')
            return F_from
        else:
            F_proj = np.zeros((proj_to+1,proj_to+1))
            for X1 in range(proj_from):
                for X2 in range(proj_from):
                    if F_from.mask[X1,X2] == False:
                        hits = (X1,X2)
                        proj_weights = cached_projection(proj_to,proj_from,hits)
                        F_proj += proj_weights * F_from[X1,X2]

            return TriSpectrum(F_proj).fold()

    # Try importing cythonized versions of several slow methods. These imports should overwrite the Python code defined above.
    try:
        from transition1 import transition1
        from transition2 import transition2
        from transition12 import transition12
        from transition1D import transition1D
    except ImportError:
        pass

## Functions

`def advance1D(u, P)`

  tridiag breakdown for integration along the diagonal boundary u - density along that diagonal P - transition matrix

  Expand source code

      def advance1D(u,P):
          """
          tridiag breakdown for integration along the diagonal boundary
          u - density along that diagonal
          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.Integration.tridiag.tridiag(a,b,c,u)
          return u

`def advance_adi(U, U01, P1, P2, x, ii)`

  Integrate the ADI components forward in time, alternating which direction occurs first U - density function U01 - stores which points are in the domain P1,P2 - transition matrices ii - count of integration step

  Expand source code

      def advance_adi(U,U01,P1,P2,x,ii):
          """
          Integrate the ADI components forward in time, alternating which direction occurs first
          U - density function
          U01 - stores which points are in the domain
          P1,P2 - transition matrices
          ii - count of integration step
          """
          if np.mod(ii,2) == 0:
              for jj in range(len(x)):
                  if np.sum(U01[:,jj]) > 1:
                      U[:,jj] = dadi.Integration.tridiag.tridiag(P1[jj,0,:],P1[jj,1,:],P1[jj,2,:],U[:,jj])
              for ii in range(len(x)):
                  if np.sum(U01[ii,:]) > 1:
                      U[ii,:] = dadi.Integration.tridiag.tridiag(P2[ii,0,:],P2[ii,1,:],P2[ii,2,:],U[ii,:])
          else:
              for ii in range(len(x)):
                  if np.sum(U01[ii,:]) > 1:
                      U[ii,:] = dadi.Integration.tridiag.tridiag(P2[ii,0,:],P2[ii,1,:],P2[ii,2,:],U[ii,:])
              for jj in range(len(x)):
                  if np.sum(U01[:,jj]) > 1:
                      U[:,jj] = dadi.Integration.tridiag.tridiag(P1[jj,0,:],P1[jj,1,:],P1[jj,2,:],U[:,jj])
          return U

`def advance_cov(U, C, x, dx)`

  Explicit integration of the covariance term, using scipy's sparse matrix for C U - density function C - transition matrix

  Expand source code

      def advance_cov(U,C,x,dx):
          """
          Explicit integration of the covariance term, using scipy's sparse matrix for C
          U - density function
          C - transition matrix
          """
          U = ( C * U.reshape(len(x)**2)).reshape(len(x),len(x))
          return U

`def advance_line(x, phi, P)`

  Integrate along the diagonal boundary. Density gets fixed along the boundary, and then diffuses along that boundary until being fixed in one of the two corners P - one dimensional transition matrix for the diagonal boundary

  Expand source code

      def advance_line(x,phi,P):
          """
          Integrate along the diagonal boundary. Density gets fixed along the boundary, and then diffuses along that boundary until being fixed in one of the two corners
          P - one dimensional transition matrix for the diagonal boundary
          """
          u = np.diag(np.fliplr(phi))
          u = advance1D(u,P)
          for ii in range(len(x)):
              phi[ii,len(x) - ii - 1] = u[ii]
          return phi

`def bivariate_lognormal_pdf(xx, params)`

  mu_i = mu_yi and sigma_i = sigma_yi are the associated means and variances of the bivariate normal distr which gets exponentiated We assume for our application that mu1=mu2 and sigma1=sigma2, though this isn't necessary for the general bivariate lognormal distribution

  Expand source code

      def bivariate_lognormal_pdf(xx, params):
          """
          mu_i = mu_yi and sigma_i = sigma_yi are the associated means and variances of the bivariate normal distr which gets exponentiated
          We assume for our application that mu1=mu2 and sigma1=sigma2, though this isn't necessary for the general bivariate lognormal distribution
          """
          mu1, mu2, sigma1, sigma2, rho = params
          norm = 1./(2 * np.pi * (sigma1*sigma2) * np.sqrt(1-rho**2) * np.outer(xx,xx) )
          q = 1/(1-rho**2) * ( ((np.log(xx[nuax,:])-mu1)/sigma1)**2 - 2*rho*((np.log(xx[nuax,:])-mu1)/sigma1)*((np.log(xx[:,nuax])-mu2)/sigma2) + ((np.log(xx[:,nuax])-mu2)/sigma2)**2 )
          prob = norm * np.exp( -q/2. )

          return prob

`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,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,n3)
          """
          key = (proj_to, proj_from, hits)
          try:
              return projection_cache[key]
          except KeyError:
              pass

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

          projection_cache[key] = proj_weights
          return proj_weights

`def cached_transition1D(numpts, sig)`
  Expand source code

      def cached_transition1D(numpts,sig):
          key = (numpts,sig)
          try:
              return transition1D_cache[key]
          except KeyError:
              pass
          x = np.linspace(0,1,numpts+1)
          dx = grid_dx(x)
          V,M = transition1D(x,dx,sig)
          transition1D_cache[key] = [V,M]
          return transition1D_cache[key]

`def domain(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(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 fold(spectrum)`

  Note: this is now handled in the TriSpectrum class Given a frequency spectrum over the full domain, fold into a spectrum with major and minor derived alleles

  Expand source code

      def fold(spectrum):
          """
          Note: this is now handled in the TriSpectrum class
          Given a frequency spectrum over the full domain, fold into a spectrum with major and minor derived alleles
          """
          spectrum = TriSpectrum(spectrum)
          if spectrum.folded_major == True:
              print("error: trying to fold a spectrum that is already folded")
              return spectrum
          else:
              spectrum = (spectrum + np.transpose(spectrum))
              for ii in range(len(spectrum)):
                  spectrum[ii,ii] = spectrum[ii,ii]/2
              spectrum.mask[0,:] = True
              spectrum.mask[:,0] = True
              for ii in range(len(spectrum)):
                  spectrum.mask[ii,ii+1:] = True
                  spectrum.mask[ii,len(spectrum)-1-ii:] = True
              return spectrum

`def fold_ancestral(F)`

  Note: this is now handled by the TriSpectrum class Don't know ancestral state, so track minor frequencies Store spectrum of two minor allele frequencies

  Expand source code

      def fold_ancestral(F):
          """
          Note: this is now handled by the TriSpectrum class
          Don't know ancestral state, so track minor frequencies
          Store spectrum of two minor allele frequencies
          """
          F_new = 0*F
          F_new = TriSpectrum(F_new)
          ns = len(F)-1
          for ii in range(ns):
              for jj in range(ns):
                  kk = ns-ii-jj
                  if F.mask[ii,jj] == True:
                      continue
                  elif ii <= kk and jj <= kk:
                      if ii >= jj:
                          F_new[ii,jj] += F[ii,jj]
                      else:
                          F_new[jj,ii] += F[ii,jj]
                  elif ii > kk and jj <= kk:
                      F_new[kk,jj] += F[ii,jj]
                  elif ii <= kk and jj > kk:
                      F_new[kk,ii] += F[ii,jj]
                  else: # ii > kk and jj > kk
                      if ii >= jj:
                          F_new[jj,kk] += F[ii,jj]
                      else:
                          F_new[ii,kk] += F[ii,jj]
          # mask if not a valid entry for ancestrally folded spectrum
          for ii in range(ns):
              for jj in range(ns):
                  kk = ns-ii-jj
                  if not (kk>=ii>=jj):
                      F_new.mask[ii,jj] = True
          return F_new

`def grid_dx(x)`

  We use uniform grids in x, using np.linspace(0,1,numpts) Grid spacing Delta, which is halved at the first and last grid points.

  Expand source code

      def grid_dx(x):
          """
          We use uniform grids in x, using np.linspace(0,1,numpts)
          Grid spacing Delta, which is halved at the first and last grid points.
          """
          return (np.concatenate((np.diff(x),np.array([0]))) + np.concatenate((np.array([0]),np.diff(x))))/2

`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 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 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 misidentification(F, p)`

  Given folded spectrum, and probability p that one of the derived alleles is the actual ancestral allele Then refold to return folded spectrum

  Expand source code

      def misidentification(F, p):
          """
          Given folded spectrum, and probability p that one of the derived alleles is the actual ancestral allele
          Then refold to return folded spectrum
          """
          F = TriF(np.zeros((len(F),len(F))))
          for ii in range(len(F))[1:-1]:
              for jj in range(len(F))[1:ii+1]:
                  if ii+jj < len(F):
                      F[ii,jj] += (1 - p) * F[ii,jj]
                      F[len(F)-1-ii-jj,jj] += p/2. * F[ii,jj]
                      F[ii,len(F)-1-ii-jj] += p/2. * F[ii,jj]

          return F.fold()

`def move_density_to_bdry(x, phi, P)`

  P tells us how much should be removed, by multiplying phi*P Take that density and instead of deleting it, move straight to boundary P - stores how much denstity from each grid point should be moved to diagonal

  Expand source code

      def move_density_to_bdry(x,phi,P):
          """
          P tells us how much should be removed, by multiplying phi*P
          Take that density and instead of deleting it, move straight to boundary
          P - stores how much denstity from each grid point should be moved to diagonal
          """
          for ii in range(len(x)):
              for jj in range(len(x)):
                  if P[ii,jj] == 0:
                      continue
                  else:
                      amnt = P[ii,jj]
                      s = ii+jj
                      if ii == 1 and jj == len(x)-3:
                          phi[ii+1,jj] += phi[ii,jj]*amnt/4. * 2
                          phi[ii,jj+1] += phi[ii,jj]*amnt/4. * 2
                          phi[ii-1,jj] += phi[ii,jj]*amnt/2 * 2
                          phi[ii,jj] *= (1-amnt)
                      elif ii == len(x)-3 and jj == 1:
                          phi[ii+1,jj] += phi[ii,jj]*amnt/4. * 2
                          phi[ii,jj+1] += phi[ii,jj]*amnt/4. * 2
                          phi[ii,jj-1] += phi[ii,jj]*amnt/2 * 2
                          phi[ii,jj] *= (1-amnt)
                      elif (len(x)-1-s) % 2 == 1:
                          # split between two points
                          dist = (len(x)-1-s) // 2
                          phi[ii+dist+1,jj+dist] += phi[ii,jj]*amnt/2. * 2
                          phi[ii+dist,jj+dist+1] += phi[ii,jj]*amnt/2. * 2
                          phi[ii,jj] *= (1-amnt)
                      else:
                          # straight to boundary grid point
                          dist = (len(x)-1-s) // 2
                          phi[ii+dist,jj+dist] += phi[ii,jj]*amnt * 2
                          phi[ii,jj] *= (1-amnt)

          return phi

`def ms_demo_params_to_dadi(nu_ms, tau_ms)`

  convert from ms parameters (which use current pop size) for nu and tau to dadi parameters (which use ancestral pop size)

  Expand source code

      def ms_demo_params_to_dadi(nu_ms,tau_ms):
          """
          convert from ms parameters (which use current pop size) for nu and tau to dadi parameters (which use ancestral pop size)
          """
          return 1./nu_ms,2*tau_ms/nu_ms

`def optimal_sfs_scaling(model, data)`
  Expand source code

      def optimal_sfs_scaling(model,data):
          data = numerics.fold(data)
          model = numerics.fold(data)
          model, data = Numerics.intersect_masks(model, data)
          return data.sum()/model.sum()

`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('sorry, but projection must be to smaller size!')
              return F_from
          else:
              F_proj = np.zeros((proj_to+1,proj_to+1))
              for X1 in range(proj_from):
                  for X2 in range(proj_from):
                      if F_from.mask[X1,X2] == False:
                          hits = (X1,X2)
                          proj_weights = cached_projection(proj_to,proj_from,hits)
                          F_proj += proj_weights * F_from[X1,X2]

              return TriSpectrum(F_proj).fold()

`def remove_diag_density_weights_nonneutral(x, dt, nu, sig1, sig2)`

  Numerically determine the amount of density that should be lost to the diagonal boundary. Numerically integrate 1D array with initial point mass at z0, where z0 is the frequency x+y, integrated for time step dt. We then check the fraction of density that is lost to z=1. If sig1 or sig2 are nonzero, estimate the selection pressure on z as sig = sig1x/(x+y) + sig2y/(x+y) x - one dimensional grid of domain dt - time step of integration nu - relative population size sig1,sig2 - selection coefficients

  Expand source code

      def remove_diag_density_weights_nonneutral(x,dt,nu,sig1,sig2):
          """
          Numerically determine the amount of density that should be lost to the diagonal boundary.
          Numerically integrate 1D array with initial point mass at z0, where z0 is the frequency x+y, integrated for time step dt.
          We then check the fraction of density that is lost to z=1.
          If sig1 or sig2 are nonzero, estimate the selection pressure on z as sig = sig1*x/(x+y) + sig2*y/(x+y)
          x - one dimensional grid of domain
          dt - time step of integration
          nu - relative population size
          sig1,sig2 - selection coefficients
          """
          dx = grid_dx(x)
          P = np.zeros((len(x),len(x)))
          for ii in range(len(x)):
              for jj in range(len(x)):
                  if x[ii]+x[jj] < 1.0 and x[ii]+x[jj] > .75:
                      sig = sig1*x[ii]/(x[ii]+x[jj]) + sig2*x[jj]/(x[ii]+x[jj])
                      V,M = cached_transition1D(len(x)-1,sig)
                      P1D = np.eye(len(x)) + dt*(V/nu+M)
                      y = np.zeros(len(x))
                      y[ii+jj] = 1./dx[ii+jj]
                      y = advance1D(y,P1D)
                      prob = y[-1]*dx[-1]
                      P[ii,jj] = prob
                      if ii+jj == len(x)-2:
                          if ii==1:
                              V,M = cached_transition1D(len(x)-1,sig1)
                              P1D = np.eye(len(x)) + dt*(V/nu+M)
                              y = np.zeros(len(x))
                              y[ii] = 1./dx[ii]
                              y = advance1D(y,P1D)
                              prob2 = y[0]*dx[0]
                              P[ii,jj] += prob2
                          elif jj==1:
                              V,M = cached_transition1D(len(x)-1,sig2)
                              P1D = np.eye(len(x)) + dt*(V/nu+M)
                              y = np.zeros(len(x))
                              y[jj] = 1./dx[jj]
                              y = advance1D(y,P1D)
                              prob2 = y[0]*dx[0]
                              P[ii,jj] += prob2

          P[:,0] = 0
          P[0,:] = 0
          return P

`def sample(phi, ns, x)`

  Obtain the expected sample frequency spectrum from the density function

  Expand source code

      def sample(phi, ns, x):
          """
          Obtain the expected sample frequency spectrum from the density function
          """
          dx = grid_dx(x)
          DXX = grid_dx_2d(x,dx)

          # Assume ns is typically a list of length 1, if not, make it into one.
          try:
              ns = tuple(ns)
          except TypeError:
              ns = (ns,)

          # We cache calculations of several big matrices that will be re-used
          # within and between integrations.
          key = (ns, tuple(x))
          if key not in sample_cache:
              this_cache = {}
              for ii in range(1,ns[0]-1):
                  # Create our cache
                  this_cache[ii] = (1-x[:,nuax]-x[nuax,:])**ii
                  # Somewhat ugly hack to use negative values to store second array
                  # to cache.
                  this_cache[-ii] = x[nuax,:]**ii
              sample_cache[key] = this_cache
          else:
              this_cache = sample_cache[key]

          #dx = grid_dx(x)
          F = np.zeros((ns[0]+1,ns[0]+1))
          prod_phi = DXX*phi
          for ii in range(len(F)):
              prod_x = prod_phi * x[:,nuax]**ii
              for jj in range(len(F)):
                  if ii+jj < ns[0] and ii != 0 and jj != 0:
                      #F[ii,jj] = math.factorial(ns)/(math.factorial(ii)*math.factorial(jj)*math.factorial(ns-ii-jj)) * int2(x, dx, phi*x[:,nuax]**ii*x[nuax,:]**jj*(1-x[:,nuax]-x[nuax,:])**(ns-ii-jj) )
                      F[ii,jj] = trinomial(ns[0],ii,jj) * np.sum(prod_x * this_cache[-jj] * this_cache[ns[0]-ii-jj])
          F = TriSpectrum(F)
          F.folded_major = False
          F.folded_ancestral = False
          F.extrap_x = x[1]
          return F

`def transition1(x, dx, U01, sig1, sig2)`

  Implicit transition matrix for the ADI components of the discretization of the diffusion Time scaled by 2N, with variance and mean terms x(1-x) and \sigmax(1-x), resp. Store the tridiagonal elements of the matrices, which need to be adjusted by I + dtP, where I is the identity matrix x - grid dx - grid spacing U01 - domain markers sig1/2 - population scaled selection coefficients nu - relative population size

  Expand source code

      def transition1(x, dx, U01, sig1, sig2):
          """
          Implicit transition matrix for the ADI components of the discretization of the diffusion
          Time scaled by 2N, with variance and mean terms x(1-x) and \sigma*x(1-x), resp.
          Store the tridiagonal elements of the matrices, which need to be adjusted by I + dt*P, where I is the identity matrix
          x - grid
          dx - grid spacing
          U01 - domain markers
          sig1/2 - population scaled selection coefficients
          nu - relative population size
          """
          # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
          #      be automatically used instead of this version, if the user has compiled it.
          print("using numpy, not cython")
          PV = np.zeros((len(x),3,len(x)))
          PM = 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)
                  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)
                  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]) )

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

              A = np.zeros((len(x),len(x)))
              x2 = x[jj]
              sig_new = sig1*(1-x-x2)/(1-x) + (sig1-sig2)*x2/(1-x)
              sig_new[-1] = 0
              M = sig_new*x*(1-x)
              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

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

`def transition12(x, dx, U01)`

  Transition matrix for the covariance term of the diffusion operator, with term D_{xy} (-xyphi) As with the ADI components, final transition matrix is given by I + dt/nu*P x - grid dx - grid spacing U01 - domain markers

  Expand source code

      def transition12(x, dx, U01):
          """
          Transition matrix for the covariance term of the diffusion operator, with term D_{xy} (-x*y*phi)
          As with the ADI components, final transition matrix is given by I + dt/nu*P
          x - grid
          dx - grid spacing
          U01 - domain markers
          """
          # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
          #      be automatically used instead of this version, if the user has compiled it.
          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, sig)`

  transition matrix for one dimensional integration x - grid dx - grid spacing dt - timestep for integration sig - selection coefficient nu - relative population size

  Expand source code

      def transition1D(x, dx, sig):
          """
          transition matrix for one dimensional integration
          x - grid
          dx - grid spacing
          dt - timestep for integration
          sig - selection coefficient
          nu - relative population size
          """
          # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
          #      be automatically used instead of this version, if the user has compiled it.
          PV = np.zeros((len(x),len(x)))
          PM = np.zeros((len(x),len(x)))
          for ii in range(len(x)):
              if ii == 0:
                  PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                  PV[ii,ii+1] = -1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii])
                  PM[ii,ii+1] = sig / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
              elif ii == len(x) - 1:
                  PV[ii,ii-1] = - 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1])
                  PM[ii,ii-1] = sig / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                  PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1])
              else:
                  PV[ii,ii-1] = - 1./2 * 1./dx[ii] * (x[ii-1]*(1-x[ii-1]))/(x[ii] - x[ii-1])
                  PM[ii,ii-1] = - sig / 2 / dx[ii] * x[ii-1] * (1-x[ii-1])
                  PV[ii,ii] = 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii] - x[ii-1]) + 1./2 * 1./dx[ii] * (x[ii]*(1-x[ii]))/(x[ii+1] - x[ii])
                  PV[ii,ii+1] = - 1./2 * 1./dx[ii] * (x[ii+1]*(1-x[ii+1]))/(x[ii+1] - x[ii])
                  PM[ii,ii+1] = sig / 2 / dx[ii] * x[ii+1] * (1-x[ii+1])
          return PV,PM

`def transition2(x, dx, U01, sig1, sig2)`

  Implicit transition matrix for the ADI components of the discretization of the diffusion Time scaled by 2N, with variance and mean terms x(1-x) and \sigmax(1-x), resp. Store the tridiagonal elements of the matrices, which need to be adjusted by I + dtP, where I is the identity matrix x - grid dx - grid spacing U01 - domain markers sig1/2 - population scaled selection coefficients nu - relative population size

  Expand source code

      def transition2(x, dx, U01, sig1, sig2):
          """
          Implicit transition matrix for the ADI components of the discretization of the diffusion
          Time scaled by 2N, with variance and mean terms x(1-x) and \sigma*x(1-x), resp.
          Store the tridiagonal elements of the matrices, which need to be adjusted by I + dt*P, where I is the identity matrix
          x - grid
          dx - grid spacing
          U01 - domain markers
          sig1/2 - population scaled selection coefficients
          nu - relative population size
          """
          # XXX: Note that this function has been Cythonized. Below, an attempt is made to import the Cythonized version so that it will
          #      be automatically used instead of this version, if the user has compiled it.
          PV = np.zeros((len(x),3,len(x)))
          PM = 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)
                  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)
                  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]) )

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

              A = np.zeros((len(x),len(x)))
              x1 = x[ii]
              sig_new = sig2*(1-x-x1)/(1-x) + (sig2-sig1)*x1/(1-x)
              sig_new[-1] = 0
              M = sig_new*x*(1-x)
              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

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

`def trinomial(ns, ii, jj)`

  Return ns!/(ii! * jj! * (ns-ii-jj)!) for large values

  Expand source code

      def trinomial(ns,ii,jj):
          """
          Return ns!/(ii! * jj! * (ns-ii-jj)!) for large values
          """
          return np.exp(math.lgamma(ns+1) - math.lgamma(ii+1) - math.lgamma(jj+1) - math.lgamma(ns-ii-jj+1))

`def univariate_lognormal_pdf(x, sigma, mu)`

  Can compare to scipy.stats.lognorm.pdf(x,sigma,0,np.exp(mu))

  Expand source code

      def univariate_lognormal_pdf(x,sigma,mu):
          """
          Can compare to scipy.stats.lognorm.pdf(x,sigma,0,np.exp(mu))
          """
          return 1./(x*sigma*np.sqrt(2*math.pi)) * np.exp ( -(np.log(x) - mu)**2 / (2*sigma**2) )