Python numba.njit() Examples

The following are 30 code examples of numba.njit(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module numba , or try the search function .
Example #1
Source File: test_numba.py    From fluids with MIT License 7 votes vote down vote up
def test_Colebrook_ignored():
    fd = fluids.numba.Colebrook(1e5, 1e-5)
    assert_close(fd, 0.018043802895063684, rtol=1e-14)


# Not compatible with cache
#@pytest.mark.numba
#@pytest.mark.skipif(numba is None, reason="Numba is missing")
#def test_secant_runs():
#    # Really feel like the kwargs should work in object mode, but it doesn't
#    # Just gets slower
#    @numba.jit
#    def to_solve(x):
#        return sin(x*.3) - .5
#    fluids.numba.secant(to_solve, .3, ytol=1e-10)
#
#@pytest.mark.numba
#@pytest.mark.skipif(numba is None, reason="Numba is missing")
#def test_brenth_runs():
#    @numba.njit
#    def to_solve(x, goal):
#        return sin(x*.3) - goal
#    
#    ans = fluids.numba.brenth(to_solve, .3, 2, args=(.45,))
#    assert_close(ans, 1.555884463490988) 
Example #2
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 6 votes vote down vote up
def applymap(self, apply_func_nb, *args):
        """See `vectorbt.tseries.nb.applymap_nb`.

        Example:
            ```python-repl
            >>> multiply_nb = njit(lambda col, i, a: a ** 2)
            >>> print(df.vbt.tseries.applymap(multiply_nb))
                           a     b    c
            2020-01-01   1.0  25.0  1.0
            2020-01-02   4.0  16.0  4.0
            2020-01-03   9.0   9.0  9.0
            2020-01-04  16.0   4.0  4.0
            2020-01-05  25.0   1.0  1.0
            ```"""
        checks.assert_numba_func(apply_func_nb)

        result = nb.applymap_nb(self.to_2d_array(), apply_func_nb, *args)
        return self.wrap(result) 
Example #3
Source File: test_nearest_neighbors.py    From openTSNE with BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def test_numba_compiled_callable_metric_same_result(self):
        k = 15

        knn_index = self.knn_index("manhattan", random_state=1)
        knn_index.build(self.x1, k=k)
        true_indices_, true_distances_ = knn_index.query(self.x2, k=k)

        @njit(fastmath=True)
        def manhattan(x, y):
            result = 0.0
            for i in range(x.shape[0]):
                result += np.abs(x[i] - y[i])

            return result

        knn_index = self.knn_index(manhattan, random_state=1)
        knn_index.build(self.x1, k=k)
        indices, distances = knn_index.query(self.x2, k=k)
        np.testing.assert_array_equal(
            indices, true_indices_, err_msg="Nearest neighbors do not match"
        )
        np.testing.assert_allclose(
            distances, true_distances_, err_msg="Distances do not match"
        ) 
Example #4
Source File: test_nearest_neighbors.py    From openTSNE with BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def test_numba_compiled_callable_metric_same_result(self):
        k = 15

        knn_index = self.knn_index("manhattan", random_state=1)
        knn_index.build(self.x1, k=k)
        true_indices_, true_distances_ = knn_index.query(self.x2, k=k)

        @njit(fastmath=True)
        def manhattan(x, y):
            result = 0.0
            for i in range(x.shape[0]):
                result += np.abs(x[i] - y[i])

            return result

        knn_index = self.knn_index(manhattan, random_state=1)
        knn_index.build(self.x1, k=k)
        indices, distances = knn_index.query(self.x2, k=k)
        np.testing.assert_array_equal(
            indices, true_indices_, err_msg="Nearest neighbors do not match"
        )
        np.testing.assert_allclose(
            distances, true_distances_, err_msg="Distances do not match"
        ) 
Example #5
Source File: test_numba_extensions.py    From clifford with BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def test_piecewise_construction(self):
        @numba.njit
        def negate(a):
            return cf.MultiVector(a.layout, -a.value)

        n_e1 = negate(e1)
        assert n_e1.layout is e1.layout
        assert n_e1 == -e1

        @numba.njit
        def add(a, b):
            return cf.MultiVector(a.layout, a.value + b.value)

        ab = add(e1, e2)
        assert ab == e1 + e2
        assert ab.layout is e1.layout 
Example #6
Source File: numba_functions.py    From astromodels with BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def ggrb_int_pl(a, b, Ec, Emin, Emax):

    pre = math.pow(a - b, a - b) * math.exp(b - a) / math.pow(Ec, b)

    if b != -2:
        b2 = 2+b

        return pre / (b2) * (math.pow(Emax, b2) - math.pow(Emin, b2))

    else:

        return pre * math.log(Emax/Emin)


# @nb.njit(fastmath=True, cache=True)
# def ggrb_int_cpl(a, Ec, Emin, Emax):

#     # Gammaincc does not support quantities
#     i1 = vec_gammaincc(2 + a, Emin/Ec) * vec_gamma(2 + a)
#     i2 = vec_gammaincc(2 + a, Emax/Ec) * vec_gamma(2 + a)

#     return -Ec * Ec * (i2 - i1) 
Example #7
Source File: hpat_pandas_dataframe_pass.py    From sdc with BSD 2-Clause "Simplified" License 6 votes vote down vote up
def _handle_agg_func(self, in_vars, out_colnames, func_name, lhs, rhs):
        agg_func = get_agg_func(self.state.func_ir, func_name, rhs)
        out_tp_vars = {}

        # sdc.jit() instead of numba.njit() to handle str arrs etc
        agg_func_dis = sdc.jit(agg_func)
        #agg_func_dis = numba.njit(agg_func)
        agg_gb_var = ir.Var(lhs.scope, ir_utils.mk_unique_var("agg_gb"), lhs.loc)
        nodes = [ir.Assign(ir.Global("agg_gb", agg_func_dis, lhs.loc), agg_gb_var, lhs.loc)]
        for out_cname in out_colnames:
            in_var = in_vars[out_cname]

            def to_arr(a, _agg_f):
                b = sdc.hiframes.api.to_arr_from_series(a)
                res = sdc.hiframes.api.init_series(sdc.hiframes.api.agg_typer(b, _agg_f))
            f_block = ir_utils.compile_to_numba_ir(to_arr, {'sdc': sdc, 'numpy': numpy}).blocks.popitem()[1]
            ir_utils.replace_arg_nodes(f_block, [in_var, agg_gb_var])
            nodes += f_block.body[:-3]  # remove none return
            out_tp_vars[out_cname] = nodes[-1].target
        return nodes, agg_func, out_tp_vars 
Example #8
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 6 votes vote down vote up
def apply_and_reduce(self, apply_func_nb, reduce_func_nb, *args, **kwargs):
        """See `vectorbt.tseries.nb.apply_and_reduce_nb`.

        `**kwargs` will be passed to `vectorbt.tseries.common.TSArrayWrapper.wrap_reduced`.

        Example:
            ```python-repl
            >>> greater_nb = njit(lambda col, a: a[a > 2])
            >>> mean_nb = njit(lambda col, a: np.nanmean(a))
            >>> print(df.vbt.tseries.apply_and_reduce(greater_nb, mean_nb))
            a    4.0
            b    4.0
            c    3.0
            dtype: float64
            ```"""
        checks.assert_numba_func(apply_func_nb)
        checks.assert_numba_func(reduce_func_nb)

        result = nb.apply_and_reduce_nb(self.to_2d_array(), apply_func_nb, reduce_func_nb, *args)
        return self.wrap_reduced(result, **kwargs) 
Example #9
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 6 votes vote down vote up
def reduce(self, reduce_func_nb, *args, **kwargs):
        """See `vectorbt.tseries.nb.reduce_nb`.

        `**kwargs` will be passed to `vectorbt.tseries.common.TSArrayWrapper.wrap_reduced`.

        Example:
            ```python-repl
            >>> mean_nb = njit(lambda col, a: np.nanmean(a))
            >>> print(df.vbt.tseries.reduce(mean_nb))
            a    3.0
            b    3.0
            c    1.8
            dtype: float64
            ```"""
        checks.assert_numba_func(reduce_func_nb)

        result = nb.reduce_nb(self.to_2d_array(), reduce_func_nb, *args)
        return self.wrap_reduced(result, **kwargs) 
Example #10
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 6 votes vote down vote up
def reduce_to_array(self, reduce_func_nb, *args, **kwargs):
        """See `vectorbt.tseries.nb.reduce_to_array_nb`.

        `**kwargs` will be passed to `vectorbt.tseries.common.TSArrayWrapper.wrap_reduced`.

        Example:
            ```python-repl
            >>> min_max_nb = njit(lambda col, a: np.array([np.nanmin(a), np.nanmax(a)]))
            >>> print(df.vbt.tseries.reduce_to_array(min_max_nb, index=['min', 'max']))
                   a    b    c
            min  1.0  1.0  1.0
            max  5.0  5.0  3.0
            ```"""
        checks.assert_numba_func(reduce_func_nb)

        result = nb.reduce_to_array_nb(self.to_2d_array(), reduce_func_nb, *args)
        return self.wrap_reduced(result, **kwargs) 
Example #11
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 6 votes vote down vote up
def generate_after(self, choice_func_nb, *args):
        """See `vectorbt.signals.nb.generate_after_nb`.

        Example:
            Fill all space between signals in `signals`:

            ```python-repl
            >>> @njit
            ... def choice_func_nb(col, from_i, to_i):
            ...     return np.arange(from_i, to_i)

            >>> print(signals.vbt.signals.generate_after(choice_func_nb))
                            a      b      c
            2020-01-01  False  False  False
            2020-01-02  False   True  False
            2020-01-03  False  False   True
            2020-01-04   True  False  False
            2020-01-05  False   True  False
            ```"""
        checks.assert_numba_func(choice_func_nb)

        return self.wrap(nb.generate_after_nb(self.to_2d_array(), choice_func_nb, *args)) 
Example #12
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_constructor_categories_set(self):
        @nb.njit
        def func():
            return pd.CategoricalDtype(categories={'b', 'a'}, ordered=True)

        boxed = func()
        assert(boxed == self._pd_dtype()) 
Example #13
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 5 votes vote down vote up
def groupby_apply(self, by, apply_func_nb, *args, on_matrix=False, **kwargs):
        """See `vectorbt.tseries.nb.groupby_apply_nb` and
        `vectorbt.tseries.nb.groupby_apply_matrix_nb` for `on_matrix=True`.

        For `by`, see `pandas.DataFrame.groupby`.

        Example:
            ```python-repl
            >>> mean_nb = njit(lambda col, i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.groupby_apply([1, 1, 2, 2, 3], mean_nb))
                 a    b    c
            1  1.5  4.5  1.5
            2  3.5  2.5  2.5
            3  5.0  1.0  1.0

            >>> mean_matrix_nb = njit(lambda i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.groupby_apply([1, 1, 2, 2, 3],
            ...     mean_matrix_nb, on_matrix=True))
                      a         b         c
            1  2.500000  2.500000  2.500000
            2  2.833333  2.833333  2.833333
            3  2.333333  2.333333  2.333333
            ```"""
        checks.assert_numba_func(apply_func_nb)

        regrouped = self._obj.groupby(by, axis=0, **kwargs)
        groups = Dict()
        for i, (k, v) in enumerate(regrouped.indices.items()):
            groups[i] = np.asarray(v)
        if on_matrix:
            result = nb.groupby_apply_matrix_nb(self.to_2d_array(), groups, apply_func_nb, *args)
        else:
            result = nb.groupby_apply_nb(self.to_2d_array(), groups, apply_func_nb, *args)
        return self.wrap_reduced(result, index=list(regrouped.indices.keys())) 
Example #14
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 5 votes vote down vote up
def resample_apply(self, freq, apply_func_nb, *args, on_matrix=False, **kwargs):
        """See `vectorbt.tseries.nb.groupby_apply_nb` and
        `vectorbt.tseries.nb.groupby_apply_matrix_nb` for `on_matrix=True`.

        For `freq`, see `pandas.DataFrame.resample`.

        Example:
            ```python-repl
            >>> mean_nb = njit(lambda col, i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.resample_apply('2d', mean_nb))
                          a    b    c
            2020-01-01  1.5  4.5  1.5
            2020-01-03  3.5  2.5  2.5
            2020-01-05  5.0  1.0  1.0

            >>> mean_matrix_nb = njit(lambda i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.resample_apply('2d',
            ...     mean_matrix_nb, on_matrix=True))
                               a         b         c
            2020-01-01  2.500000  2.500000  2.500000
            2020-01-03  2.833333  2.833333  2.833333
            2020-01-05  2.333333  2.333333  2.333333
            ```"""
        checks.assert_numba_func(apply_func_nb)

        resampled = self._obj.resample(freq, axis=0, **kwargs)
        groups = Dict()
        for i, (k, v) in enumerate(resampled.indices.items()):
            groups[i] = np.asarray(v)
        if on_matrix:
            result = nb.groupby_apply_matrix_nb(self.to_2d_array(), groups, apply_func_nb, *args)
        else:
            result = nb.groupby_apply_nb(self.to_2d_array(), groups, apply_func_nb, *args)
        result_obj = self.wrap(result, index=list(resampled.indices.keys()))
        resampled_arr = np.full((resampled.ngroups, self.to_2d_array().shape[1]), np.nan)
        resampled_obj = self.wrap(resampled_arr, index=pd.Index(list(resampled.groups.keys()), freq=freq))
        resampled_obj.loc[result_obj.index] = result_obj.values
        return resampled_obj 
Example #15
Source File: nb.py    From vectorbt with GNU General Public License v3.0 5 votes vote down vote up
def generate_nb(shape, choice_func_nb, *args):
    """Create a boolean matrix of `shape` and pick `True` values using `choice_func_nb`.

    `choice_func_nb` must accept index of the current column `col`, index of the start 
    of the range `from_i`, index of the end of the range `to_i`, and `*args`. 
    It must return an array of indices from `[from_i, to_i)` (can be empty).

    !!! note
        Returned indices must be absolute.

    Example:
        ```python-repl
        >>> from numba import njit
        >>> import numpy as np
        >>> from vectorbt.signals.nb import generate_nb

        >>> @njit
        ... def choice_func_nb(col, from_i, to_i):
        ...     return np.array([from_i + col])

        >>> print(generate_nb((5, 3), choice_func_nb))
        [[ True False False]
         [False  True False]
         [False False  True]
         [False False False]
         [False False False]]
        ```"""
    result = np.full(shape, False, dtype=b1)

    for col in range(result.shape[1]):
        idxs = choice_func_nb(col, 0, shape[0], *args)
        result[idxs, col] = True
    return result 
Example #16
Source File: numba_functions.py    From astromodels with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def bb_eval(x, K, kT):

    n = x.shape[0]
    out = np.empty(n)

    for idx in range(n):

        arg = x[idx]/kT
        out[idx] = K * x[idx] * x[idx] / np.expm1(arg)

    return out

# @nb.njit(fastmath=True, cache=True)
# def bbrad_eval(x, K, kT):

#     n = x.shape[0]
#     out = np.empty(n)

#     for idx in range(n):

#         arg = x[idx]/kT
#         out[idx] = K * x[idx] * x[idx] / np.expm1(arg)

#     return out



# band calderone 
Example #17
Source File: utils.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def to_array_overload(A):
    # try regular np.array and return it if it works
    def to_array_impl(A):
        return np.array(A)
    try:
        numba.njit(to_array_impl).get_call_template((A,), {})
        return to_array_impl
    except BaseException:
        pass  # should be handled elsewhere (e.g. Set) 
Example #18
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_attribute_ordered(self):
        @nb.njit
        def func(c):
            return c.ordered

        pd_dtype = self._pd_dtype()
        ordered = func(pd_dtype)
        assert(ordered == pd_dtype.ordered) 
Example #19
Source File: accessors.py    From vectorbt with GNU General Public License v3.0 5 votes vote down vote up
def expanding_apply(self, apply_func_nb, *args, on_matrix=False):
        """See `vectorbt.tseries.nb.expanding_apply_nb` and
        `vectorbt.tseries.nb.expanding_apply_matrix_nb` for `on_matrix=True`.

        Example:
            ```python-repl
            >>> mean_nb = njit(lambda col, i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.expanding_apply(mean_nb))
                          a    b    c
            2020-01-01  1.0  5.0  1.0
            2020-01-02  1.5  4.5  1.5
            2020-01-03  2.0  4.0  2.0
            2020-01-04  2.5  3.5  2.0
            2020-01-05  3.0  3.0  1.8

            >>> mean_matrix_nb = njit(lambda i, a: np.nanmean(a))
            >>> print(df.vbt.tseries.expanding_apply(
            ...     mean_matrix_nb, on_matrix=True))
                               a         b         c
            2020-01-01  2.333333  2.333333  2.333333
            2020-01-02  2.500000  2.500000  2.500000
            2020-01-03  2.666667  2.666667  2.666667
            2020-01-04  2.666667  2.666667  2.666667
            2020-01-05  2.600000  2.600000  2.600000
            ```"""
        checks.assert_numba_func(apply_func_nb)

        if on_matrix:
            result = nb.expanding_apply_matrix_nb(self.to_2d_array(), apply_func_nb, *args)
        else:
            result = nb.expanding_apply_nb(self.to_2d_array(), apply_func_nb, *args)
        return self.wrap(result) 
Example #20
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_constructor_no_categories(self):
        @nb.njit
        def func():
            return pd.CategoricalDtype()

        boxed = func()
        expected = pd.CategoricalDtype(ordered=None)
        assert(boxed == expected)
        assert(boxed.categories == expected.categories)
        assert(boxed.ordered == expected.ordered) 
Example #21
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_constructor_no_order(self):
        @nb.njit
        def func():
            return pd.CategoricalDtype(categories=('b', 'a'))

        boxed = func()
        assert(boxed == self._pd_dtype(ordered=False)) 
Example #22
Source File: test_numba_extensions.py    From clifford with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def test_multivector_shorthand(self):
        @numba.njit
        def double(a):
            return a.layout.MultiVector(a.value*2)

        assert double(e2) == 2 * e2 
Example #23
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_constructor(self):
        @nb.njit
        def func():
            return pd.CategoricalDtype(categories=('b', 'a'), ordered=True)

        boxed = func()
        assert(boxed == self._pd_dtype()) 
Example #24
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_lowering(self):
        pd_dtype = self._pd_dtype()

        @nb.njit
        def func():
            return pd_dtype

        boxed = func()
        assert(boxed == pd_dtype) 
Example #25
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_boxing(self):
        @nb.njit
        def func(c):
            return c

        pd_dtype = self._pd_dtype()
        boxed = func(pd_dtype)
        assert(boxed == pd_dtype) 
Example #26
Source File: test_categoricaldtype.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_unboxing(self):
        @nb.njit
        def func(c):
            pass

        pd_dtype = self._pd_dtype()
        func(pd_dtype) 
Example #27
Source File: test_categorical.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_constructor_values_list(self):
        @nb.njit
        def func():
            return pd.Categorical(values=[1, 2, 3, 2, 1])

        boxed = func()
        assert(boxed.equals(self._pd_value())) 
Example #28
Source File: test_categorical.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_lowering(self):
        pd_value = self._pd_value()

        @nb.njit
        def func():
            return pd_value

        boxed = func()
        assert(boxed.equals(pd_value)) 
Example #29
Source File: test_categorical.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_boxing(self):
        @nb.njit
        def func(c):
            return c

        pd_value = self._pd_value()
        boxed = func(pd_value)
        assert(boxed.equals(pd_value)) 
Example #30
Source File: test_categorical.py    From sdc with BSD 2-Clause "Simplified" License 5 votes vote down vote up
def test_unboxing(self):
        @nb.njit
        def func(c):
            pass

        pd_value = self._pd_value()
        func(pd_value)