def generate_current_integrand( width: float, field_to_k_factor: float, current_distribution: np.ndarray, phase_distribution: np.ndarray, ): """Integrand to compute the current through a junction at a given field. """ step = width / len(current_distribution) @cfunc(float64(intc, CPointer(float64))) def real_current_integrand(n, args): """Cfunc to be used with quad to calculate the current from the distribution. """ pos, field = args[0], args[1] x = int(pos // step) return current_distribution[x] * np.cos( (phase_distribution[x] + field_to_k_factor * pos) * field) @cfunc(float64(intc, CPointer(float64))) def imag_current_integrand(n, args): """Cfunc to be used with quad to calculate the current from the distribution. """ pos, field = args[0], args[1] x = int(pos // step) return current_distribution[x] * np.sin( (phase_distribution[x] + field_to_k_factor * pos) * field) return ( LowLevelCallable(real_current_integrand.ctypes), LowLevelCallable(imag_current_integrand.ctypes), )
def _fpext(tyctx, val): def impl(cgctx, builder, signature, args): val = args[0] return builder.fpext(val, lc.Type.double()) sig = types.float64(types.float32) return sig, impl
def jit_integrand_function(integrand_function): jitted_function = numba.jit(integrand_function, nopython=True) @cfunc(float64(intc, CPointer(float64))) def wrapped(n, xx): values = carray(xx,n) return jitted_function(values) return LowLevelCallable(wrapped.ctypes)
def jit_integrand_function(integrand_function): """ decorator function to compile numerical integration routine This function is used to build AMPS functions as pre-compiled routines for faster evalution. From the StackOverflow answer: https://stackoverflow.com/questions/49683653/how-to-pass-additional-parameters-to-numba-cfunc-passed-as-lowlevelcallable-to-s We are using the function signature: double func(int n, double *xx) Where n is the length of xx array, xx[0] is the variable to be integrated over. The rest is the extra arguments that would normally be passed in 'args' of integrator 'scipy.integrate.quad' Usage: @jit_integrand_function def f(x, *args): a = args[0] b = args[1] return np.exp(-a*x/b) quad(f, 0, pi, args=(a,b)) """ jitted_function = jit(integrand_function, nopython=True) @cfunc(float64(intc, CPointer(float64))) def wrapped(n, xx): """ """ return jitted_function(xx[0], xx[1], xx[2], xx[3]) return LowLevelCallable(wrapped.ctypes)
def jit_integrand_function(integrand_function): """Based on https://stackoverflow.com/a/49732825/4779220""" jitted_function = numba.jit(integrand_function, nopython=True, nogil=True) @numba.cfunc(nut.float64(nut.intc, nut.CPointer(nut.float64))) def wrapped(n, xx): # TODO: nicer way to not hard code number of args? `*carray()` may not expand correctly return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5]) return LowLevelCallable(wrapped.ctypes)
def jit_psf5args(func): """A complicated piece of code used by numba to compile the PSF. This one works for `gaussianWithConstantSkyPsf()` """ jitted_function = njit(func) @cfunc(float64(intc, CPointer(float64))) def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6]) return LowLevelCallable(wrapped.ctypes)
def get_pointcloud(path): # converting point cloud data to numba type ans = typed.List() with open(path, 'rb') as f: plydata = plyfile.PlyData.read(f) plydata = np.asarray(plydata.elements[1].data) for i in range(len(plydata)): raw = typed.List() [raw.append(types.float64(n)) for n in plydata[i]] ans.append(raw) return ans
def nbtss(values_ptr, len_values, result, data): # total sum of square difference (C-implementation for speedup) values = carray(values_ptr, (len_values, ), dtype=float64) sum = 0.0 for v in values: sum += v mean = sum / float64(len_values) result[0] = 0 for v in values: result[0] += (v - mean)**2 return 1
def _fma(typing_context, x, y, z): """Compute x * y + z using a fused multiply add.""" sig = float64(float64, float64, float64) def codegen(context, builder, signature, args): ty = args[0].type mod = builder.module fnty = ir.types.FunctionType(ty, [ty, ty, ty]) fn = mod.declare_intrinsic('llvm.fma', [ty], fnty) ret = builder.call(fn, args) return ret return sig, codegen
def _integrand_function(integrand_function): """Wrap `integrand_function` as a `LowLevelCallable` to be used with quad. `integrand_function` has to have the signature (float, complex) -> float. This speeds up integration by removing call overhead. However only float arguments can be passed to and from the function. """ @numba.cfunc(nb_t.float64(nb_t.intc, nb_t.CPointer(nb_t.float64))) def wrapped(__, xx): return integrand_function(xx[0], xx[1] + xx[2]) return LowLevelCallable(wrapped.ctypes)
def jit_integrand(integrand_function): jitted_function = numba.jit(integrand_function, nopython=True) no_args = len(inspect.getfullargspec(integrand_function).args) wrapped = None if no_args == 4: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3]) elif no_args == 5: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4]) elif no_args == 6: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5]) elif no_args == 7: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6]) elif no_args == 8: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7]) elif no_args == 9: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8]) elif no_args == 10: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8], xx[9]) elif no_args == 11: def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8], xx[9], xx[10]) cf = cfunc(float64(intc, CPointer(float64))) return LowLevelCallable(cf(wrapped).ctypes)
def test_basic5(self): a = 1 @njit def foo1(x): return x + 1 @njit def foo2(x): return x + 2 @njit def bar1(x): return x / 10 @njit def bar2(x): return x / 1000 tup = (foo1, foo2) tup_bar = (bar1, bar2) int_int_fc = types.FunctionType(types.int64(types.int64, )) flt_flt_fc = types.FunctionType(types.float64(types.float64, )) @njit((types.UniTuple(int_int_fc, 2), types.UniTuple(flt_flt_fc, 2))) def bar(fcs, ffs): x = 0 for i in range(2): x += fcs[i](a) for fn in ffs: x += fn(a) return x got = bar(tup, tup_bar) expected = foo1(a) + foo2(a) + bar1(a) + bar2(a) self.assertEqual(got, expected)
@jit_integrand_function def s_tpf_integrand(θ, *args): μ = args[0] σ = args[1] n = args[2] x1 = 2.0 * np.pi * np.arange(-n, n + 1.0) + θ x2 = 2.0 * np.pi * np.arange(-n, n + 1.0) - θ arg1 = (x1 - μ) / σ arg2 = (x2 - μ) / σ return (np.sin(θ)**6 * (np.sum(np.exp(-0.5 * arg1 * arg1)) + np.sum(np.exp(-0.5 * arg2 * arg2))) * np.sin(θ)) # cheap integration routine for approximating solutions @vectorize([float64(float64, float64)]) def rough_shgratio(μ, σ): θ = np.linspace(0.0, np.pi, 100) W = np.zeros_like(θ) for n in range(-2, 3): x1 = 2.0 * np.pi * n + θ x2 = 2.0 * np.pi * n - θ arg1 = (x1 - μ) / σ arg2 = (x2 - μ) / σ W += np.exp(-0.5 * arg1 * arg1) + np.exp(-0.5 * arg2 * arg2) _P = np.sum(np.cos(θ)**3 * W * np.sin(θ)) _S = np.sum(np.sin(θ)**2 * np.cos(θ) * W * np.sin(θ)) return 1.88 * 4.0 * (_P / _S)**2
def factorial(x): """ Numba-styled implementation of factorial calculation. :param x: (int) Required. :returns n: (int). """ n = 1 for i in range(2, x + 1): n *= i return n @jit(types.float64(types.int64, types.float64), nopython=True) def pois_survival_partial(x, mu): """ Implements the partial derivative of the poisson CDF with respect to lambda. :param x: (int) Required. The value in which the CDF is calculated from. Density is integers to the right of this number. :param mu: (float) Required. The mean at which to calculate the density with. :returns val: (float) Required. The derivative of the cumulative density. """ ret_val = 0 for i in range(x + 1): summation = ((i * (mu**(i - 1))) - mu**i) / factorial(i)
References ---------- [1] Knuth, "The Art of Computer Programming, Volume II" [2] Holin et al., "Polynomial and Rational Function Evaluation", http://www.boost.org/doc/libs/1_61_0/libs/math/doc/html/math_toolkit/roots/rational.html """ import numpy as np from numba import njit from numba.types import complex128, float64, intc, Array from .fma import _fma @njit(float64(Array(float64, 1, "C", readonly=True), float64)) def _devalpoly(coeffs, x): """Evaluate a polynomial using Horner's method.""" res = coeffs[0] for j in range(1, len(coeffs)): res = _fma(res, x, coeffs[j]) return res @njit(complex128(Array(float64, 1, "C", readonly=True), complex128)) def _cevalpoly(coeffs, z): """Evaluate a polynomial with real coefficients at a complex point. Uses equation (3) in section 4.6.4 of [1]. Note that it is more
Fluxe in (cm^2 s sr GeV)^-1 """ phi0 = 1.62e-4 * 1e-4 # (cm^2 s sr GeV)^-1 eb = 914. # GeV gamma1 = 3.09 gamma2 = 3.92 delta = 0.1 return phi0 * (100. / e)**gamma1 * \ (1. + (eb/e)**((gamma1 - gamma2) / delta))**(-delta) phi_e_bg_dampe_jit = jit(phi_e_bg_dampe, nopython=True) @cfunc(float64(float64)) def phi_e_bg_dampe_cfunc(e): return phi_e_bg_dampe_jit(e) phi_e_bg_dampe_llc = LowLevelCallable(phi_e_bg_dampe_cfunc.ctypes) def phi_e_bg_alt(e): """Background model for the e+ e- flux from Ge et al, arXiv:1712.02744 Note ----- The data in several bins is excluded from the background fit. See Fig. 1 in the reference above.
################################## # Useful collections of signatures ################################## _unary_bool = [nt.boolean(nt.boolean)] _unary_int = [ nt.uint8(nt.uint8), nt.int8(nt.int8), nt.uint16(nt.uint16), nt.int16(nt.int16), nt.uint32(nt.uint32), nt.int32(nt.int32), nt.uint64(nt.uint64), nt.int64(nt.int64) ] _unary_float = [nt.float32(nt.float32), nt.float64(nt.float64)] _unary_all = _unary_bool + _unary_int + _unary_float _binary_bool = [nt.boolean(nt.boolean, nt.boolean)] _binary_int = [ nt.uint8(nt.uint8, nt.uint8), nt.int8(nt.int8, nt.int8), nt.uint16(nt.uint16, nt.uint16), nt.int16(nt.int16, nt.int16), nt.uint32(nt.uint32, nt.uint32), nt.int32(nt.int32, nt.int32), nt.uint64(nt.uint64, nt.uint64), nt.int64(nt.int64, nt.int64) ] _binary_float = [ nt.float32(nt.float32, nt.float32),
import warnings from numba.targets.imputils import implement, Registry from numba import types from numba.itanium_mangler import mangle from .hsaimpl import _declare_function registry = Registry() register = registry.register # ----------------------------------------------------------------------------- _unary_b_f = types.int32(types.float32) _unary_b_d = types.int32(types.float64) _unary_f_f = types.float32(types.float32) _unary_d_d = types.float64(types.float64) _binary_f_ff = types.float32(types.float32, types.float32) _binary_d_dd = types.float64(types.float64, types.float64) function_descriptors = { 'isnan': (_unary_b_f, _unary_b_d), 'isinf': (_unary_b_f, _unary_b_d), 'ceil': (_unary_f_f, _unary_d_d), 'floor': (_unary_f_f, _unary_d_d), 'fabs': (_unary_f_f, _unary_d_d), 'sqrt': (_unary_f_f, _unary_d_d), 'exp': (_unary_f_f, _unary_d_d), 'expm1': (_unary_f_f, _unary_d_d), 'log': (_unary_f_f, _unary_d_d), 'log10': (_unary_f_f, _unary_d_d), 'log1p': (_unary_f_f, _unary_d_d),
@njit(cache=True, fastmath=checks.fastmath, nogil=True) def _find_edge_idx(node_edge_map: Dict, edge_data: np.ndarray, start_nd_idx: int, end_nd_idx: int) -> int: """ Finds an edge from start and end nodes """ # iterate the start node's edges for edge_idx in node_edge_map[start_nd_idx]: # check whether the edge's out node matches the target node if edge_data[edge_idx, 1] == end_nd_idx: return edge_idx node_close_func_proto = types.FunctionType( types.float64(types.float64, types.float64, types.float64, types.float64)) # node density @njit(cache=True, fastmath=checks.fastmath, nogil=True) def node_density(to_short_dist, to_simpl_dist, beta, cycles): return 1.0 # return float explicitly # node farness @njit(cache=True, fastmath=checks.fastmath, nogil=True) def node_farness(to_short_dist, to_simpl_dist, beta, cycles): return to_short_dist # node cycles
Function lngamma(), lncombination(), hypergeometric_probability(), were originally written by Oyvind Langsrud: Oyvind Langsrud Copyright (C) : All right reserved. Contact Oyvind Langsrud for permission. Adapted to Cython version by: Haibao Tang, Brent Pedersen """ from numba import jit, prange from numba.types import int64, float64 import numpy as np from math import log, exp, lgamma @jit(float64(int64), nopython=True, parallel=True, cache=True) def _naive_lnfactorial(n): acc = 0.0 for i in prange(2, n + 1): acc += log(i) return acc # Tabulated ln n! for n \in [0, 1023] _lnfactorials1 = np.zeros(1024) for i in range(1024): _lnfactorials1[i] = _naive_lnfactorial(i) # Logarithm of n! with algorithmic approximation @jit(float64(int64), nopython=True, cache=True)
# Authors: Stephane Gaiffas <*****@*****.**> # License: BSD 3 clause import numpy as np from numpy import sin, sign, pi, abs, sqrt import matplotlib.pyplot as plt from numba import vectorize from numba.types import float32, float64 # Examples of signals originally made by David L. Donoho. @vectorize([float32(float32), float64(float64)], nopython=True) def heavisine(x): """Computes the "heavisine" signal. Parameters ---------- x : `numpy.array`, shape=(n_samples,) Inputs values Returns ------- output : `numpy.array`, shape=(n_samples,) The value of the signal at given inputs Notes ----- Inputs are supposed to belong to [0, 1] and must have dtype `float32` or `float64` """ return 4 * sin(4 * pi * x) - sign(x - 0.3) - sign(0.72 - x)
def digit_count(integer): if integer <= 0: return 0 else: return int64(floor(log10(float64(integer))) + 1)
def jit_integrand(integrand_function): jitted_function = decorator_util.jit(nopython=True, cache=True)(integrand_function) no_args = len(inspect.getfullargspec(integrand_function).args) wrapped = None if no_args == 4: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3]) elif no_args == 5: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4]) elif no_args == 6: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5]) elif no_args == 7: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6]) elif no_args == 8: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7]) elif no_args == 9: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8]) elif no_args == 10: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8], xx[9]) elif no_args == 11: # noinspection PyUnusedLocal def wrapped(n, xx): return jitted_function( xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8], xx[9], xx[10], ) cf = cfunc(float64(intc, CPointer(float64))) return LowLevelCallable(cf(wrapped).ctypes)
k = 0 for k in range(len(cumsum)): c = cumsum[k] if random_number <= c: break # put back to model cdk[d, k] += 1 cd[d] += 1 ckn[k, n] += 1 ck[k] += 1 return k @jit(float64(int32, int32, int32, float64[:], float64[:], float64, float64, int32[:, :], int32[:], int32[:, :], int32[:]), nopython=True) def _nb_ll(D, N, K, alpha, beta, N_beta, K_alpha, cdk, cd, ckn, ck): temp_sum = 0 for b in beta: temp_sum += math.lgamma(b) ll = K * (math.lgamma(N_beta) - temp_sum) for k in range(K): for n in range(N): ll += math.lgamma(ckn[k, n] + beta[n]) ll -= math.lgamma(ck[k] + N_beta) temp_sum = 0 for a in alpha: temp_sum += math.lgamma(a)
def int_to_float(x): return types.float64(x) / 2
import warnings from numba.targets.imputils import Registry from numba import types from numba.itanium_mangler import mangle from .hsaimpl import _declare_function registry = Registry() lower = registry.lower # ----------------------------------------------------------------------------- _unary_b_f = types.int32(types.float32) _unary_b_d = types.int32(types.float64) _unary_f_f = types.float32(types.float32) _unary_d_d = types.float64(types.float64) _binary_f_ff = types.float32(types.float32, types.float32) _binary_d_dd = types.float64(types.float64, types.float64) function_descriptors = { 'isnan': (_unary_b_f, _unary_b_d), 'isinf': (_unary_b_f, _unary_b_d), 'ceil': (_unary_f_f, _unary_d_d), 'floor': (_unary_f_f, _unary_d_d), 'fabs': (_unary_f_f, _unary_d_d), 'sqrt': (_unary_f_f, _unary_d_d), 'exp': (_unary_f_f, _unary_d_d), 'expm1': (_unary_f_f, _unary_d_d),
lam = D0 * e0 / (b0 * (1. - delta)) * \ ((e0 / e)**(1. - delta) - (e0 / mx)**(1. - delta)) fact_e = 1. / (b * (4. * np.pi * lam)**1.5) # Term with purely radial dependence r_term = r**2 * np.exp(-(r * kpc_to_cm)**2 / (4. * lam)) # Term from performing theta integral K_term = K_full(r, d, rs, rhos, gamma) return fact_const * fact_e * K_term * (r_term * kpc_to_cm**3) dphi_e_dr_jit = jit(dphi_e_dr, nopython=True) @cfunc(float64(intc, CPointer(float64))) def dphi_e_dr_cfunc(n, xx): return dphi_e_dr_jit(xx[0], xx[1], xx[2], xx[3], xx[4], xx[5], xx[6], xx[7], xx[8]) dphi_e_dr_llc = LowLevelCallable(dphi_e_dr_cfunc.ctypes) # J-factor integrand def dJ_dr(r, th_max, d, rs, rhos, gamma): # Numerically stable expression for solid angle subtended by target dOmega = 4 * np.pi * np.sin(0.5 * th_max)**2 th_term = K(r, 0, d, rs, rhos, gamma) - K(r, th_max, d, rs, rhos, gamma) return 2 * np.pi / dOmega * th_term
cumsum[i] = total k = 0 for k in range(len(cumsum)): c = cumsum[k] if random_number <= c: break # put back to model cdk[d, k] += 1 cd[d] += 1 ckn[k, n] += 1 ck[k] += 1 return k @jit(float64(int32, int32, int32, float64[:], float64[:], float64, float64, int32[:, :], int32[:], int32[:, :], int32[:]), nopython=True) def _nb_ll(D, N, K, alpha, beta, N_beta, K_alpha, cdk, cd, ckn, ck): temp_sum = 0 for b in beta: temp_sum += math.lgamma(b) ll = K * ( math.lgamma(N_beta) - temp_sum ) for k in range(K): for n in range(N): ll += math.lgamma(ckn[k, n]+beta[n]) ll -= math.lgamma(ck[k] + N_beta) temp_sum = 0 for a in alpha: temp_sum += math.lgamma(a) ll += D * ( math.lgamma(K_alpha) - temp_sum )