Определение функций с условными выражениями в Python

Question

В настоящее время я работаю над кодированием докторской диссертации. В статье рассматривается выборочный дефолт, основанный на работе Arellano (2008). Вдохновлялся блокнотами @quantecon.

Я пытаюсь создать функцию на Python, используя Numba, которая выполняет итерацию нескольких переменных, но одна из них является условной. То есть в зависимости от вероятности состояния природы итерация должна быть разной. Я пытался использовать оператор if/else, но Python указывает на синтаксическую ошибку. Код представляет собой итерацию функции значения для моей докторской диссертации. Вот что я написал:

    @jit(nopython=True)
def solve(model, tol=1e-8, maxiter=10_000):
    """
    Given an Selective_Economy type, this function computes the optimal
    policy and value functions
    """
    # Unpack certain parameters for simplification
    β, σ, r, θ_d, θ_f = model.β, model.σ, model.r, model.θ_d, model.θ_f
    b_d = np.ascontiguousarray(model.b_d)
    b_f = np.ascontiguousarray(model.b_f)
    A, y = np.ascontiguousarray(model.A), np.ascontiguousarray(model.y)
    nb_d, nb_f, ny = b_d.size, b_f.size, y.size

    # Allocate space
    ib_d_star = np.zeros((ny, nb_d), int64)
    ib_f_star = np.zeros((ny, nb_f), int64)
    dom_default_prob = np.zeros((ny, nb_d))
    ext_default_prob = np.zeros((ny, nb_f))
    dom_default_states = np.zeros((ny, nb_d))
    ext_default_states = np.zeros((ny, nb_f))
    tot_default_states = np.zeros((ny, nb_d, nb_f))
    q_f = np.ones((ny, nb_f)) * 0.95
    q_d = np.ones((ny, nb_d)) * 0.95
    Vfd = np.zeros(ny, nb_d)
    Vdd = np.zeros(ny, nb_f)
    Vtd = np.zeros(ny)
    Vc, V, Vupd = np.zeros((ny, nb_d,nb_f)), np.zeros((ny, nb_d,nb_f)), np.zeros((ny, nb_d,nb_f))

    it = 0
    dist = 10.0
    while (it < maxiter) and (dist > tol):

        # Compute expectations used for this iteration
        EV = A@V
        EVfd = A@Vfd
        EVdd = A@Vdd
        EVtd = A@Vtd

        for iy in range(ny):
            # Update value function for default state
            Vfd[iy] = model.bellman_ext_default(iy, EVfd, EV)
            Vdd[iy] = model.bellman_dom_default(iy, EVdd, EV)
            Vtd[iy] = model.bellman_tot_default(iy, EVtd, EV)

            for ib_d, ib_f in range(nb_d), range(nb_f):
                # Update value function for non-default state
                ib_d_star[iy, ib_d] = model.compute_savings_policy(iy, ib_d, q_d, EV)
                ib_f_star[iy, ib_f] = model.compute_savings_policy(iy, ib_f, q_d, EV)
                Vc[iy, ib_d,ib_f] = model.bellman_nondefault(iy, ib_d, ib_f, q_d, q_f, EV, ib_d_star[iy, ib_d], ib_f_star[iy,ib_f])

        # Once value functions are updated, can combine them to get
        # the full value function
        Vfd_compat = np.reshape(np.repeat(Vfd, nb_f), (ny, nb_f))
        Vdd_compat = np.reshape(np.repeat(Vdd, nb_d), (ny, nb_d))
        Vtd_compat = np.reshape(np.repeat(Vtd, nb_d, nb_f), (ny, nb_d, nb_f))
        Vupd[:, :] = np.maximum(Vc, Vfd_compat, Vdd_compat, Vtd_compat)

        # Can also compute default states and update prices
        ext_default_states[:, :] = 1.0 * (Vfd_compat > max(Vc, Vdd, Vtd))
        dom_default_states[:, :] = 1.0 * (Vdd_compat > max(Vc, Vfd, Vtd))
        tot_default_states[:, :] = 1.0 * (Vtd_compat > max(Vc, Vfd, Vdd))
        ext_default_prob[:, :] = A @ ext_default_states
        dom_default_prob[:, :] = A @ dom_default_states

        if dom_default_prob = 0:

            q_f[:, :] = (1 - ext_default_prob) / (1 + r)
        else:
            q_f[:, :] = θ_f * (1 - ext_default_prob) / (1 + r) + (1 - θ_f) (1 - ext_default_prob) / (1 + r)

        if ext_default_prob = 0:

            q_d[:, :] = β * (1 - ext_default_prob) * (1 - dom_default_prob) * u_prime(y[iy]) + ext_default_prob * (1 - dom_default_prob) * u_prime(ext_y[iy]) / u_prime(y[iy])
        else:
            q_d[:, :] = β * θ_f * (1 - dom_default_prob) * u_prime(y[iy]) + (1 - θ_f) * (1 - dom_default_prob) * u_prime(ext_y[iy]) / u_prime(ext_y[iy])

        # Check tolerance etc...
        dist = np.max(np.abs(Vupd - V))
        V[:, :] = Vupd[:, :]
        it += 1

    return V, Vc, Vfd, Vdd, Vtd, ib_d_star, ib_f_star, ext_default_prob, dom_default_prob, ext_default_states, dom_default_states, tot_default_states, q_f, q_d

Ошибка находится в последней части, где я получаю следующее сообщение:

  File "<ipython-input-3-cf2f65d62110>", line 65
if dom_default_prob = 0:
                    SyntaxError: invalid syntax

Если кому поможет, буду Вам признателен!

Спасибо:)

Answer 1

Вы должны использовать

==

чтобы проверить равенство. Один равен= сообщает python присвоить значение правой части левой переменной.

Таким образом, код должен быть

if dom_default_prob == 0: