Python numba.guvectorize() Examples

def tensor_outer_at(vtarget, **kwargs):
    @guvectorize(['void(float64[:], float64[:, :], float64[:, :], intp[:], intp[:], float64[:, :])'],
                 '(),(i,m),(j,n),(),()->(m,n)',
                 target=vtarget, nopython=True, **kwargs)
    def tensor_outer_wrapped(val, v, w, i, j, res):
        r1 = v.shape[1]
        r2 = w.shape[1]
        for m in range(r1):
            for n in range(r2):
                res[m, n] = val[0] * v[i[0], m] * w[j[0], n]
    return tensor_outer_wrapped 

def calculate_value_functions_and_flow_utilities(
    wage, nonpec, continuation_value, draw, delta, value_function, flow_utility
):
    """Calculate the choice-specific value functions and flow utilities.

    To apply :func:`aggregate_keane_wolpin_utility` to arrays with arbitrary dimensions,
    this function uses :func:`numba.guvectorize`. One cannot use :func:`numba.vectorize`
    because it does not support multiple return values.

    See also
    --------
    aggregate_keane_wolpin_utility

    """
    value_function[0], flow_utility[0] = aggregate_keane_wolpin_utility(
        wage, nonpec, continuation_value, draw, delta
    ) 

def calculate_conditional_draws(
    base_draws, updated_mean, updated_chols, chol_index, max_log_float, conditional_draw
):
    """Calculate the conditional draws from base draws, updated means and updated chols.

    We need to pass ``max_log_float`` to the function, because the global variables
    ``MAX_LOG_FLOAT`` cannot be used directly withing the guvectorize.

    Parameters
    ----------
    base_draws : np.ndarray
        iid standard normal draws
    updated_mean : np.ndarray
        conditional mean, given the observed shock. Contains the observed shock in the
        corresponding position.
    updated_chols : np.ndarray
        cholesky factor of conditional covariance, given the observed shock. If there is
        no measurement error, it contains a zero column and row at the position of the
        observed shock.
    chol_index : float
        index of the relevant updated cholesky factor
    max_log_float : float
        Value at which numbers soon to be exponentiated are clipped.

    Returns
    -------
    conditional draws : np.ndarray
        draws from the conditional distribution of the shocks.

    """
    n_draws, n_choices = base_draws.shape
    n_wages = len(updated_chols) - 1

    for d in range(n_draws):
        for i in range(n_choices):
            cd = updated_mean[i]
            for j in range(i + 1):
                cd += base_draws[d, j] * updated_chols[chol_index, i, j]
            if i < n_wages:
                if cd > max_log_float:
                    cd = max_log_float
                cd = np.exp(cd)
            conditional_draw[d, i] = cd 

def _crps_ensemble_gufunc(observation, forecasts, weights, result):
    # beware: forecasts are assumed sorted in NumPy's sort order

    # we index the 0th element to get the scalar value from this 0d array:
    # http://numba.pydata.org/numba-doc/0.18.2/user/vectorize.html#the-guvectorize-decorator
    obs = observation[0]

    if np.isnan(obs):
        result[0] = np.nan
        return

    total_weight = 0.0
    for n, weight in enumerate(weights):
        if np.isnan(forecasts[n]):
            # NumPy sorts NaN to the end
            break
        if not weight >= 0:
            # this catches NaN weights
            result[0] = np.nan
            return
        total_weight += weight

    obs_cdf = 0
    forecast_cdf = 0
    prev_forecast = 0
    integral = 0

    for n, forecast in enumerate(forecasts):
        if np.isnan(forecast):
            # NumPy sorts NaN to the end
            if n == 0:
                integral = np.nan
            # reset for the sake of the conditional below
            forecast = prev_forecast
            break

        if obs_cdf == 0 and obs < forecast:
            integral += (obs - prev_forecast) * forecast_cdf ** 2
            integral += (forecast - obs) * (forecast_cdf - 1) ** 2
            obs_cdf = 1
        else:
            integral += ((forecast - prev_forecast)
                         * (forecast_cdf - obs_cdf) ** 2)

        forecast_cdf += weights[n] / total_weight
        prev_forecast = forecast

    if obs_cdf == 0:
        # forecast can be undefined here if the loop body is never executed
        # (because forecasts have size 0), but don't worry about that because
        # we want to raise an error in that case, anyways
        integral += obs - forecast

    result[0] = integral