Python numba.njit方法代码示例

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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"
        ) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
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)) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
def njit(first=None, *args, **kwargs):
        """Identity JIT, returns unchanged function."""
        def _jit(f):
            return f

        if inspect.isfunction(first):
            return first
        else:
            return _jit

# This function was ported from its Matlab equivalent here:
# http://www.mathworks.com/matlabcentral/fileexchange/23051-vectorized-solar-azimuth-and-elevation-estimation 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
def two_body_mc_en_jit(bond_array_1, c1, etypes1,
                       bond_array_2, c2, etypes2,
                       sig, ls, r_cut, cutoff_func,
                       nspec, spec_mask, bond_mask):
    """Multicomponent two-body energy/energy kernel accelerated with
    Numba's njit decorator."""

    kern = 0

    ls1 = 1 / (2 * ls * ls)
    sig2 = sig * sig

    bc1 = spec_mask[c1]
    bc1n = bc1 * nspec

    for m in range(bond_array_1.shape[0]):
        ri = bond_array_1[m, 0]
        fi, _ = cutoff_func(r_cut, ri, 0)
        e1 = etypes1[m]

        be1 = spec_mask[e1]
        btype = bond_mask[bc1n + be1]

        tls1 = ls1[btype]
        tsig2 = sig2[btype]

        for n in range(bond_array_2.shape[0]):
            e2 = etypes2[n]

            if (c1 == c2 and e1 == e2) or (c1 == e2 and c2 == e1):
                rj = bond_array_2[n, 0]
                fj, _ = cutoff_func(r_cut, rj, 0)
                r11 = ri - rj
                kern += fi * fj * tsig2 * exp(-r11 * r11 * tls1)

    return kern 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
def where_in(x, y, not_in=False):
    """Retrieve the indices of the elements in x which are also in y.

    x and y are assumed to be 1 dimensional arrays.

    :params: not_in: if True, returns the indices of the of the elements in x
    which are not in y.

    """
    # np.isin is not supported in numba. Also: "i in y" raises an error in numba
    # setting njit(parallel=True) slows down the function
    list_y = set(y)
    return np.where(np.array([i in list_y for i in x]) != not_in)[0] 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
def num_nodes_and_leaves(tree):
    n_nodes = 0
    n_leaves = 0
    for i in range(len(tree.children)):
        if tree.children[i][0] < 0:
            n_leaves += 1
            n_nodes += 1
        else:
            n_nodes += 1

    return n_nodes, n_leaves


# #@numba.njit()
# def convert_tree_format(tree):
#     n_nodes, n_leaves = num_nodes_and_leaves(tree)
#     print(n_nodes, n_leaves, len(tree.children))
#     hyperplane_dim = np.max([tree.hyperplanes[i].shape[0] for i in range(len(
#         tree.hyperplanes))])
#
#     hyperplanes = np.zeros((n_nodes, hyperplane_dim), dtype=np.float32)
#     offsets = np.zeros(n_nodes, dtype=np.float32)
#     children = -1 * np.ones((n_nodes, 2), dtype=np.int64)
#     graph_indices = -1 * np.ones((n_leaves, tree.leaf_size), dtype=np.int64)
#     recursive_convert(tree, hyperplanes, offsets, children, graph_indices, 0, 0,
#                       len(tree.children) - 1)
#     return FlatTree(hyperplanes, offsets, children, graph_indices, tree.leaf_size) 

# 需要导入模块: import numba [as 别名]
# 或者: from numba import njit [as 别名]
def transform_lists_to_arrays(module, to_change, __funcs, vec=False, cache_blacklist=set([])):
    if vec:
        conv_fun = numba.vectorize
        extra_args = extra_args_vec
    else:
        conv_fun = numba.njit
        extra_args = extra_args_std
    for s in to_change:
        func = s.split('.')[-1]
        mod = '.'.join(s.split('.')[:-1])
        fake_mod = __funcs[mod]

        try:
            real_mod = getattr(module, mod)
        except:
            real_mod = module
            for s in mod.split('.'):
                real_mod = getattr(real_mod, s)

        orig_func = getattr(real_mod, func)
        source = inspect.getsource(orig_func)
        source = remove_for_numba(source) # do before anything else
        source = return_value_numpy(source)
        source = re.sub(list_mult_expr, numpy_not_list_expr, source)
#        if 'longitude_obliquity_nutation' in s:
#        print(source)
        numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
        new_func = fake_mod.__dict__[func]
        do_cache = caching and func not in cache_blacklist
        obj = conv_fun(cache=do_cache, **extra_args)(new_func)
        __funcs[func] = obj
        fake_mod.__dict__[func] = obj
        obj.__doc__ = ''


#set_signatures = {'Clamond': [numba.float64(numba.float64, numba.float64, numba.boolean),
#                              numba.float64(numba.float64, numba.float64, numba.optional(numba.boolean))
#                              ]
#                    }