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
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
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
def int_to_float(x):
return types.float64(x) / 2
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 )
