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 jit_filter_function(filter_function):
jitted_function = numba.jit(filter_function, nopython=True)
@cfunc(intc(CPointer(float64), intp, CPointer(float64), voidptr))
def wrapped(values_ptr, len_values, result, data):
values = carray(values_ptr, (len_values,), dtype=float64)
result[0] = jitted_function(values)
return 1
return LowLevelCallable(wrapped.ctypes)
def jit_filter_function(filter_function):
"""Decorator for use with scipy.ndimage.generic_filter."""
jitted_function = numba.jit(filter_function, nopython=True)
@cfunc(intc(CPointer(float64), intp, CPointer(float64), voidptr))
def wrapped(values_ptr, len_values, result, data):
values = carray(values_ptr, (len_values, ), dtype=float64)
result[0] = jitted_function(values)
return 1
return LowLevelCallable(wrapped.ctypes)
def jit_geometric_function(geometric_function):
jitted_function = numba.jit(geometric_function, nopython=True)
@cfunc(intc(CPointer(intp), CPointer(float64), intc, intc, voidptr))
def wrapped(output_ptr, input_ptr, output_rank, input_rank, user_data):
output_coords = carray(output_ptr, (output_rank, ), dtype=intp)
input_coords = carray(input_ptr, (output_rank, ), dtype=float64)
jitted_function(output_coords, input_coords)
return 1
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):
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_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 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 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)
import os
PINK = np.array([255, 15, 255]) / 255.
YELLOW = np.array([255, 255, 15]) / 255.
@jit
def tss(a):
# total sum of square difference
total_sum_of_squares = np.sum((a - np.mean(a))**2)
return total_sum_of_squares
@cfunc(intc(CPointer(float64), intp, CPointer(float64), voidptr))
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
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
from numba import carray, cfunc, njit
from numba.types import CPointer, double, intc, void
from pypde.utils import nargs
F_SIG = void(CPointer(double), CPointer(double), CPointer(double), intc)
B_SIG = void(CPointer(double), CPointer(double), intc)
S_SIG = void(CPointer(double), CPointer(double))
def generate_signatures(nFargs, nBargs):
if nFargs == 1:
Fsig = 'double[:](double[:])'
elif nFargs == 2:
Fsig = 'double[:](double[:], intc)'
elif nFargs == 3:
Fsig = 'double[:](double[:], double[:,:], intc)'
if nBargs == 1:
Bsig = 'double[:,:](double[:])'
elif nBargs == 2:
Bsig = 'double[:,:](double[:], intc)'
Ssig = 'double[:](double[:])'
return Fsig, Bsig, Ssig
def generate_cfuncs(F, B, S, ndim, V):