Python Misc.ensure_1arg_func方法代码示例

# 需要导入模块: import Misc [as 别名]
# 或者: from Misc import ensure_1arg_func [as 别名]
def three_pops(phi, xx, T, nu1=1, nu2=1, nu3=1,
               m12=0, m13=0, m21=0, m23=0, m31=0, m32=0,
               gamma1=0, gamma2=0, gamma3=0, h1=0.5, h2=0.5, h3=0.5,
               theta0=1, initial_t=0, frozen1=False, frozen2=False,
               frozen3=False):
    """
    Integrate a 3-dimensional phi foward.

    phi: Initial 3-dimensional phi
    xx: 1-dimensional grid upon (0,1) overwhich phi is defined. It is assumed
        that this grid is used in all dimensions.

    nu's, gamma's, m's, and theta0 may be functions of time.
    nu1,nu2,nu3: Population sizes
    gamma1,gamma2,gamma3: Selection coefficients on *all* segregating alleles
    h1,h2,h3: Dominance coefficients. h = 0.5 corresponds to genic selection.
    m12,m13,m21,m23,m31,m32: Migration rates. Note that m12 is the rate 
                             *into 1 from 2*.
    theta0: Propotional to ancestral size. Typically constant.

    T: Time at which to halt integration
    initial_t: Time at which to start integration. (Note that this only matters
               if one of the demographic parameters is a function of time.)

    Note: Generalizing to different grids in different phi directions is
          straightforward. The tricky part will be later doing the extrapolation
          correctly.
    """
    phi = phi.copy()

    if T - initial_t == 0:
        return phi
    elif T - initial_t < 0:
        raise ValueError('Final integration time T (%f) is less than '
                         'intial_time (%f). Integration cannot be run '
                         'backwards.' % (T, initial_t))


    if (frozen1 and (m12 != 0 or m21 != 0 or m13 !=0 or m31 != 0))\
       or (frozen2 and (m12 != 0 or m21 != 0 or m23 !=0 or m32 != 0))\
       or (frozen3 and (m13 != 0 or m31 != 0 or m23 !=0 or m32 != 0)):
        raise ValueError('Population cannot be frozen and have non-zero '
                         'migration to or from it.')

    vars_to_check = [nu1,nu2,nu3,m12,m13,m21,m23,m31,m32,gamma1,gamma2,
                     gamma3,h1,h2,h3,theta0]
    if numpy.all([numpy.isscalar(var) for var in vars_to_check]):
        return _three_pops_const_params(phi, xx, T, nu1, nu2, nu3, 
                                        m12, m13, m21, m23, m31, m32, 
                                        gamma1, gamma2, gamma3, h1, h2, h3,
                                        theta0, initial_t,
                                        frozen1, frozen2, frozen3)
    zz = yy = xx

    nu1_f = Misc.ensure_1arg_func(nu1)
    nu2_f = Misc.ensure_1arg_func(nu2)
    nu3_f = Misc.ensure_1arg_func(nu3)
    m12_f = Misc.ensure_1arg_func(m12)
    m13_f = Misc.ensure_1arg_func(m13)
    m21_f = Misc.ensure_1arg_func(m21)
    m23_f = Misc.ensure_1arg_func(m23)
    m31_f = Misc.ensure_1arg_func(m31)
    m32_f = Misc.ensure_1arg_func(m32)
    gamma1_f = Misc.ensure_1arg_func(gamma1)
    gamma2_f = Misc.ensure_1arg_func(gamma2)
    gamma3_f = Misc.ensure_1arg_func(gamma3)
    h1_f = Misc.ensure_1arg_func(h1)
    h2_f = Misc.ensure_1arg_func(h2)
    h3_f = Misc.ensure_1arg_func(h3)
    theta0_f = Misc.ensure_1arg_func(theta0)

    current_t = initial_t
    nu1,nu2,nu3 = nu1_f(current_t), nu2_f(current_t), nu3_f(current_t)
    m12,m13 = m12_f(current_t), m13_f(current_t)
    m21,m23 = m21_f(current_t), m23_f(current_t)
    m31,m32 = m31_f(current_t), m32_f(current_t)
    gamma1,gamma2 = gamma1_f(current_t), gamma2_f(current_t)
    gamma3 = gamma3_f(current_t)
    h1,h2,h3 = h1_f(current_t), h2_f(current_t), h3_f(current_t)
    dx,dy,dz = numpy.diff(xx),numpy.diff(yy),numpy.diff(zz)
    dt = min(_compute_dt(dx,nu1,[m12,m13],gamma1,h1),
             _compute_dt(dy,nu2,[m21,m23],gamma2,h2),
             _compute_dt(dz,nu3,[m31,m32],gamma3,h3))
    steps = 0
    while current_t < T:
        dt = min(_compute_dt(dx,nu1,[m12,m13],gamma1,h1),
                 _compute_dt(dy,nu2,[m21,m23],gamma2,h2),
                 _compute_dt(dz,nu3,[m31,m32],gamma3,h3))
        this_dt = min(dt, T - current_t)

        next_t = current_t + this_dt
        steps += 1

        nu1,nu2,nu3 = nu1_f(next_t), nu2_f(next_t), nu3_f(next_t)
        m12,m13 = m12_f(next_t), m13_f(next_t)
        m21,m23 = m21_f(next_t), m23_f(next_t)
        m31,m32 = m31_f(next_t), m32_f(next_t)
        gamma1,gamma2 = gamma1_f(next_t), gamma2_f(next_t)
        gamma3 = gamma3_f(next_t)
        h1,h2,h3 = h1_f(next_t), h2_f(next_t), h3_f(next_t)
#.........这里部分代码省略.........

# 需要导入模块: import Misc [as 别名]
# 或者: from Misc import ensure_1arg_func [as 别名]
def one_pop(phi, xx, T, nu=1, gamma=0, h=0.5, theta0=1.0, initial_t=0, 
            frozen=False, beta=1):
    """
    Integrate a 1-dimensional phi foward.

    phi: Initial 1-dimensional phi
    xx: Grid upon (0,1) overwhich phi is defined.

    nu, gamma, and theta0 may be functions of time.
    nu: Population size
    gamma: Selection coefficient on *all* segregating alleles
    h: Dominance coefficient. h = 0.5 corresponds to genic selection. 
       Heterozygotes have fitness 1+2sh and homozygotes have fitness 1+2s.
    theta0: Propotional to ancestral size. Typically constant.
    beta: Breeding ratio, beta=Nf/Nm.

    T: Time at which to halt integration
    initial_t: Time at which to start integration. (Note that this only matters
               if one of the demographic parameters is a function of time.)

    frozen: If True, population is 'frozen' so that it does not change.
            In the one_pop case, this is equivalent to not running the
            integration at all.
    """
    phi = phi.copy()

    # For a one population integration, freezing means just not integrating.
    if frozen:
        return phi

    if T - initial_t == 0:
        return phi
    elif T - initial_t < 0:
        raise ValueError('Final integration time T (%f) is less than '
                         'intial_time (%f). Integration cannot be run '
                         'backwards.' % (T, initial_t))

    vars_to_check = (nu, gamma, h, theta0, beta)
    if numpy.all([numpy.isscalar(var) for var in vars_to_check]):
        return _one_pop_const_params(phi, xx, T, nu, gamma, h, theta0, 
                                     initial_t, beta)

    nu_f = Misc.ensure_1arg_func(nu)
    gamma_f = Misc.ensure_1arg_func(gamma)
    h_f = Misc.ensure_1arg_func(h)
    theta0_f = Misc.ensure_1arg_func(theta0)
    beta_f = Misc.ensure_1arg_func(beta)

    current_t = initial_t
    nu, gamma, h = nu_f(current_t), gamma_f(current_t), h_f(current_t)
    beta = beta_f(current_t)
    dx = numpy.diff(xx)
    while current_t < T:
        dt = _compute_dt(dx,nu,[0],gamma,h)
        this_dt = min(dt, T - current_t)

        # Because this is an implicit method, I need the *next* time's params.
        # So there's a little inconsistency here, in that I'm estimating dt
        # using the last timepoints nu,gamma,h.
        next_t = current_t + this_dt
        nu, gamma, h = nu_f(next_t), gamma_f(next_t), h_f(next_t)
        beta = beta_f(next_t)
        theta0 = theta0_f(next_t)

        if numpy.any(numpy.less([T,nu,theta0], 0)):
            raise ValueError('A time, population size, migration rate, or '
                             'theta0 is < 0. Has the model been mis-specified?')
        if numpy.any(numpy.equal([nu], 0)):
            raise ValueError('A population size is 0. Has the model been '
                             'mis-specified?')

        _inject_mutations_1D(phi, this_dt, xx, theta0)
        # Do each step in C, since it will be faster to compute the a,b,c
        # matrices there.
        phi = int_c.implicit_1Dx(phi, xx, nu, gamma, h, beta, this_dt, 
                                 use_delj_trick=use_delj_trick)
        current_t = next_t
    return phi