def wrapper(a, b, out=None, args=None): """wrapper deciding whether the underlying function is called with or without `out`. This uses nb.generated_jit to compile different versions of the same function.""" if isinstance(a, nb.types.Number): # simple scalar call -> do not need to allocate anything raise RuntimeError( "Functions defined with an `out` keyword should not be called " "with scalar quantities") elif isinstance(out, (nb.types.NoneType, nb.types.Omitted)): # function is called without `out` f_jit = register_jitable(**jit_kwargs_inner)(func) if out_shape is None: # we have to obtain the shape of `out` from `a` def f_with_allocated_out(a, b, out=None, args=None): """helper function allocating output array""" return f_jit(a, b, out=np.empty_like(a), args=args) else: # the shape of `out` is given by `out_shape` def f_with_allocated_out(a, b, out=None, args=None): """helper function allocating output array""" out_arr = np.empty(out_shape, dtype=a.dtype) return f_jit(a, b, out=out_arr, args=args) return f_with_allocated_out else: # function is called with `out` argument return func
def _vectorize_operator( make_operator: Callable, grid: CartesianGridBase, *, backend: str = "numba" ) -> OperatorType: """apply an operator to on all dimensions of a vector Args: make_operator (callable): The function that creates the basic operator grid (:class:`~pde.grids.cartesian.CartesianGridBase`): The grid for which the operator is created backend (str): Backend used for calculating the vector gradient operator. Returns: A function that can be applied to an array of values """ dim = grid.dim operator = make_operator(grid, backend=backend) def vectorized_operator(arr: np.ndarray, out: np.ndarray) -> None: """apply vector gradient operator to array `arr`""" for i in range(dim): operator(arr[i], out[i]) if backend == "numba": return register_jitable(vectorized_operator) # type: ignore else: return vectorized_operator
def wrapper(arr, out=None): """wrapper deciding whether the underlying function is called with or without `out`. This uses :func:`nb.generated_jit` to compile different versions of the same function """ if isinstance(arr, nb.types.Number): # simple scalar call -> do not need to allocate anything raise RuntimeError( "Functions defined with an `out` keyword should not be called " "with scalar quantities." ) if not isinstance(arr, nb.types.Array): raise RuntimeError( "Compiled functions need to be called with numpy arrays, not " f"type {arr.__class__.__name__}" ) f_jit = register_jitable(**jit_kwargs_inner)(func) if isinstance(out, (nb.types.NoneType, nb.types.Omitted)): # function is called without `out` if out_shape is None: # we have to obtain the shape of `out` from `arr` def f_with_allocated_out(arr, out): """helper function allocating output array""" return f_jit(arr, out=np.empty_like(arr)) else: # the shape of `out` is given by `out_shape` def f_with_allocated_out(arr, out): """helper function allocating output array""" return f_jit(arr, out=np.empty(out_shape, dtype=arr.dtype)) return f_with_allocated_out else: # function is called with `out` argument if out_shape is None: return f_jit else: def f_with_tested_out(arr, out): """helper function allocating output array""" assert out.shape == out_shape return f_jit(arr, out) return f_with_tested_out
elif np.isnan(indval2.value): # indval1 not a nan but indval2 is so consider indval2 larger return indval2 elif indval2.value > indval1.value: return indval2 elif indval1.value == indval2.value: if indval1.index < indval2.index: return indval1 else: return indval2 return indval1 return max_impl greater_than = register_jitable(lambda a, b: a > b) less_than = register_jitable(lambda a, b: a < b) @register_jitable def min_max_impl(iterable, op): if isinstance(iterable, types.IterableType): def impl(iterable): it = iter(iterable) return_val = next(it) for val in it: if op(val, return_val): return_val = val return return_val
boundary = min(win, length) for i in prange(boundary): arr_range = input_arr[:i + 1] output_arr[i] = apply_minp(arr_range, ddof, minp) for i in prange(boundary, length): arr_range = input_arr[i + 1 - win:i + 1] output_arr[i] = apply_minp(arr_range, ddof, minp) return pandas.Series(output_arr, input_series._index, name=input_series._name) return impl hpat_pandas_rolling_series_median_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_median)) @sdc_register_jitable def pop_corr(x, y, nfinite, result): """Calculate the window sums for corr without old value.""" sum_x, sum_y, sum_xy, sum_xx, sum_yy = result if numpy.isfinite(x) and numpy.isfinite(y): nfinite -= 1 sum_x -= x sum_y -= y sum_xy -= x * y sum_xx -= x * x sum_yy -= y * y return nfinite, (sum_x, sum_y, sum_xy, sum_xx, sum_yy)
def make_jit_quicksort(*args, **kwargs): from numba.extending import register_jitable return make_quicksort_impl((lambda f: register_jitable(f)), *args, **kwargs)
return is_upper(_get_code_point(a, 0)) if l == 0: return False cased = False for idx in range(l): code_point = _get_code_point(a, idx) if is_lower(code_point) or is_title(code_point): return False elif (not cased and is_upper(code_point)): cased = True return cased return impl _always_false = register_jitable(lambda x: False) _ascii_is_upper = register_jitable( _is_upper(_Py_ISLOWER, _Py_ISUPPER, _always_false)) _unicode_is_upper = register_jitable( _is_upper(_PyUnicode_IsLowercase, _PyUnicode_IsUppercase, _PyUnicode_IsTitlecase)) @overload_method(types.UnicodeType, 'isupper') def unicode_isupper(a): """ Implements .isupper() """ def impl(a): if a._is_ascii: return _ascii_is_upper(a)
def make_jit_quicksort(*args, **kwargs): from numba.extending import register_jitable return make_quicksort_impl((lambda f: register_jitable(f)), *args, **kwargs)
lo = 0 hi = n while hi > lo: mid = (lo + hi) >> 1 if func(a[mid], (v)): # mid is too low => go up lo = mid + 1 else: # mid is too high, or is a NaN => go down hi = mid return lo return searchsorted_inner _lt = register_jitable(lambda x, y: x < y) _le = register_jitable(lambda x, y: x <= y) _searchsorted_left = register_jitable(_searchsorted(_lt)) _searchsorted_right = register_jitable(_searchsorted(_le)) @overload(np.searchsorted) def searchsorted(a, v, side='left'): side_val = getattr(side, 'value', side) if side_val == 'left': loop_impl = _searchsorted_left elif side_val == 'right': loop_impl = _searchsorted_right else: raise ValueError("Invalid value given for 'side': %s" % side_val)
for i in prange(boundary): arr_range = input_arr[:i + 1] output_arr[i] = apply_minp(arr_range, ddof, minp) for i in prange(boundary, length): arr_range = input_arr[i + 1 - win:i + 1] output_arr[i] = apply_minp(arr_range, ddof, minp) return pandas.Series(output_arr, input_series._index, name=input_series._name) return impl hpat_pandas_rolling_series_kurt_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_kurt)) hpat_pandas_rolling_series_max_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_max)) hpat_pandas_rolling_series_median_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_median)) hpat_pandas_rolling_series_min_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_min)) hpat_pandas_rolling_series_skew_impl = register_jitable( gen_hpat_pandas_series_rolling_impl(arr_skew)) hpat_pandas_rolling_series_std_impl = register_jitable( gen_hpat_pandas_series_rolling_ddof_impl(arr_std)) hpat_pandas_rolling_series_var_impl = register_jitable( gen_hpat_pandas_series_rolling_ddof_impl(arr_var)) @sdc_register_jitable