How to use the numba.prange function in numba

To help you get started, we’ve selected a few numba examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github pmgbergen / porepy / src / porepy / numerics / contact_mechanics / contact_conditions.py View on Github external
def _isclose(a: np.ndarray, b: np.ndarray, rtol: float = 1e-05, atol: float = 1e-08, equal_nan=False) -> np.ndarray:
    """ Numba-friendly implementation of np.isclose
    Assume a, b are 1d-like-arrays

    Returns
    -------
    res : np.ndarray (1d-array)

    See: https://numpy.org/doc/stable/reference/generated/numpy.isclose.html
    """
    a, b = np.atleast_2d(a), np.atleast_2d(b)
    assert a.shape[0] == b.shape[0] == 1
    res = np.full(a.shape[1], False)
    for i in prange(a.shape[1]):
        _a, _b = a[0, i], b[0, i]
        if np.isnan(_a) and np.isnan(_b) and equal_nan:
            res[i] = True
        else:
            res[i] = np.abs(_a - _b) <= atol + rtol * np.abs(_b)
    return res
github yzhao062 / pyod / pyod / models / sod.py View on Github external
def _snn_imp(ind, ref_set_):
    """Internal function for fast snn calculation

    Parameters
    ----------
    ind : int
        Indices return by kNN.

    ref_set_ : int, optional (default=10)
        specifies the number of shared nearest neighbors to create the
        reference set. Note that ref_set must be smaller than n_neighbors.

    """
    n = ind.shape[0]
    _count = np.zeros(shape=(n, ref_set_), dtype=nb.uint16)
    for i in nb.prange(n):
        temp = np.empty(n, dtype=nb.int16)
        test_element_set = set(ind[i])
        for j in nb.prange(n):
            temp[j] = len(set(ind[j]).intersection(test_element_set))
        temp[i] = np.iinfo(np.uint16).max
        _count[i] = np.argsort(temp)[::-1][1:ref_set_ + 1]

    return _count
github alessandro-zomparelli / tissue / numba_functions.py View on Github external
def numba_set_ab(a, b, brush):
        n = len(a)
        _brush = np.full(n, brush[0]) if len(brush) == 1 else brush
        for i in prange(len(b)):
            b[i] += _brush[i]
            if b[i] < 0: b[i] = 0
            elif b[i] > 1: b[i] = 1
            if a[i] < 0: a[i] = 0
            elif a[i] > 1: a[i] = 1
github ogrisel / pygbm / pygbm / predictor.py View on Github external
def _predict_from_numeric_data(nodes, numeric_data, out):
    for i in prange(numeric_data.shape[0]):
        out[i] = _predict_one_from_numeric_data(nodes, numeric_data[i])
github mpes-kit / mpes / mpes / beamtime.py View on Github external
:Parameters:
        U, V : numpy.ndarray (float32)
            3D matrices, U is the original, V is the modified version of U.
        lsh, rsh, ush, dsh : int
            Pixel shifts along the four primary directions (left, right, up, down).

    :Return:
        V : numpy.ndarray (float32)
            Modified 3D matrix after averaging.
    """

    a, x, y = V.shape

    for i in nb.prange(ush, x-dsh):
        for j in nb.prange(lsh, y-rsh):
            for ia in nb.prange(a):

                # Average over nearby points perpendicular to the specified axis
                V[ia, i, j] = np.nanmean(U[ia, i-ush:i+dsh, j-lsh:j+rsh])

    return V
github Blueqat / Blueqat / blueqat / backends / _numba_backend_impl.py View on Github external
def _u3gate(qubits, n_qubits, target, theta, phi, lambd):
    theta *= 0.5
    cos = math.cos(theta)
    sin = math.sin(theta)
    expadd = cmath.exp((phi + lambd) * 0.5j)
    expsub = cmath.exp((phi - lambd) * 0.5j)
    a = expadd.conjugate() * cos
    b = -expsub.conjugate() * sin
    c = expsub * sin
    d = expadd * cos
    lower_mask = (1 << _QSMask(target)) - 1
    for i in prange(1 << (_QSMask(n_qubits) - 1)):
        i0 = _shifted(lower_mask, i)
        t = qubits[i0]
        u = qubits[i0 + (1 << target)]
        qubits[i0] = a * t + b * u
        qubits[i0 + (1 << target)] = c * t + d * u
github IntelPython / sdc / sdc / functions / numpy_like.py View on Github external
def sdc_fillna_str_impl(self, inplace=False, value=None):
                n = len(self)
                num_chars = 0
                # get total chars in new array
                for i in prange(n):
                    s = self[i]
                    if sdc.hiframes.api.isna(self, i):
                        num_chars += len(value)
                    else:
                        num_chars += len(s)

                filled_data = pre_alloc_string_array(n, num_chars)
                for i in prange(n):
                    if sdc.hiframes.api.isna(self, i):
                        filled_data[i] = value
                    else:
                        filled_data[i] = self[i]
                return filled_data
github asarnow / pyem / pyem / geom / quat_numba.py View on Github external
def cdistdq(q1, q2, d):
    relq = np.zeros((4,), dtype=np.complex128)
    for i in numba.prange(d.shape[0]):
        for j in range(d.shape[1]):
            relq = dqtimes_sca(dqconj_sca(q1[i]), q2[j])
            theta, dax, l, m = dq2sc(relq)
            p0 = cross3_sca(l, m)
            r2 = np.sum(p0 ** 2)
            v = np.sqrt(dax ** 2 + theta ** 2 * r2)
            d[i, j] = v
    return d
github eamid / trimap / trimap / trimap_.py View on Github external
n_outliers: Number of outlier points

    Output
    ------

    triplets: Sampled triplets
    """
    n, n_neighbors = nbrs.shape
    triplets = np.empty((n * n_inliers * n_outliers, 3), dtype=np.int32)
    for i in numba.prange(n):
        sort_indices = np.argsort(-P[i])
        for j in numba.prange(n_inliers):
            sim = nbrs[i][sort_indices[j + 1]]
            samples = rejection_sample(n_outliers, n, sort_indices[: j + 2])
            for k in numba.prange(n_outliers):
                index = i * n_inliers * n_outliers + j * n_outliers + k
                out = samples[k]
                triplets[index][0] = i
                triplets[index][1] = sim
                triplets[index][2] = out
    return triplets