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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
def test_unboxing(self): @nb.njit def func(c): pass pd_value = self._pd_value() func(pd_value)