def expm1(a, b):
print((numba.typeof(a)))
print((numba.typeof(np.expm1(a))))
# result = a**2 + b**2
# print "... :)"
# print np.expm1(result), "..."
return np.expm1(a**2) + b
def check_scal(scal):
x = 4
y = 5
cres = compile_isolated(pyfunc, (typeof(scal), typeof(x), typeof(y)))
expected = pyfunc(scal, x, y)
got = cres.entry_point(scal, x, y)
self.assertPreciseEqual(got, expected)
def test_use_dtype(self):
# Test using the dtype attribute inside the Numba function itself
b = np.empty(1, dtype=np.int16)
pyfunc = use_dtype
cfunc = self.get_cfunc(pyfunc, (typeof(self.a), typeof(b)))
expected = pyfunc(self.a, b)
self.assertPreciseEqual(cfunc(self.a, b), expected)
def test_array_return(self):
a = numpy.arange(10)
i = 2
at, it = typeof(a), typeof(i)
cres = compile_isolated(array_return, (at, it))
cfunc = cres.entry_point
self.assertIs(a, cfunc(a, i))
def test_array_ndenumerate_2d(self):
arr = np.arange(12).reshape(4, 3)
arrty = typeof(arr)
self.assertEqual(arrty.ndim, 2)
self.assertEqual(arrty.layout, 'C')
self.assertTrue(arr.flags.c_contiguous)
# Test with C-contiguous array
self.check_array_ndenumerate_sum(arr, arrty)
# Test with Fortran-contiguous array
arr = arr.transpose()
self.assertFalse(arr.flags.c_contiguous)
self.assertTrue(arr.flags.f_contiguous)
arrty = typeof(arr)
self.assertEqual(arrty.layout, 'F')
self.check_array_ndenumerate_sum(arr, arrty)
# Test with non-contiguous array
arr = arr[::2]
self.assertFalse(arr.flags.c_contiguous)
self.assertFalse(arr.flags.f_contiguous)
arrty = typeof(arr)
self.assertEqual(arrty.layout, 'A')
self.check_array_ndenumerate_sum(arr, arrty)
# Boolean array
arr = np.bool_([1, 0, 0, 1]).reshape((2, 2))
self.check_array_ndenumerate_sum(arr, typeof(arr))
def test_record_type_equiv(self):
rec_dt = np.dtype([('a', np.int32), ('b', np.float32)])
rec_ty = typeof(rec_dt)
art1 = rec_ty[::1]
arr = np.zeros(5, dtype=rec_dt)
art2 = typeof(arr)
self.assertEqual(art2.dtype.dtype, rec_ty)
self.assertEqual(art1, art2)
def test_typeof_numba2(arg):
"""
>>> test_typeof_numba2(10)
(10+0j)
"""
x = 1 + 2j
arg = numba.typeof(x)(arg)
return numba.typeof(arg)(arg)
def test_np_ndindex_array(self):
func = np_ndindex_array
arr = np.arange(12, dtype=np.int32)
self.check_array_unary(arr, typeof(arr), func)
arr = arr.reshape((4, 3))
self.check_array_unary(arr, typeof(arr), func)
arr = arr.reshape((2, 2, 3))
self.check_array_unary(arr, typeof(arr), func)
def run_array_index_test(self, func):
A1 = np.arange(6).reshape(2,3)
A2 = A1.copy()
i = 0
pfunc = self.compile_parallel(func, (numba.typeof(A1), numba.typeof(i)))
func(A1, i)
pfunc(A2, i)
np.testing.assert_array_equal(A1, A2)
def check(arr, ind):
arrty = typeof(arr)
cr = self.ccache.compile(pyfunc, (arrty, typeof(ind), arrty.dtype))
# Use np.copy() to keep the layout
expected = np.copy(arr)
got = np.copy(arr)
pyfunc(expected, ind, 123)
cr.entry_point(got, ind, 123)
self.assertPreciseEqual(got, expected)
def test_sinc(self):
"""
Tests the sinc() function.
This test is purely to assert numerical computations are correct.
"""
# Ignore sign of zeros, this will need masking depending on numpy
# version once the fix to numpy complex division is in upstream
# See: https://github.com/numpy/numpy/pull/6699
isoz = True
# Testing sinc(1.) leads to sin(pi)/pi, which is below machine
# precision in practice on most machines. Small floating point
# differences in sin() etc. may lead to large differences in the result
# that are at a range that is inaccessible using standard width
# floating point representations.
# e.g. Assume float64 type.
# sin(pi) ~= 1e-16, but should be zero
# sin(pi)/pi ~= 1e-17, should be zero, error carried from above
# float64 has log10(2^53)~=15.9 digits of precision and the magnitude
# change in the alg is > 16 digits (1.0...0 -> 0.0...0),
# so comparison via ULP is invalid.
# We therefore opt to assume that values under machine precision are
# equal in this case.
tol = "eps"
pyfunc = sinc
def check(x_types, x_values, **kwargs):
self.run_unary(pyfunc, x_types, x_values,
ignore_sign_on_zero=isoz, abs_tol=tol,
**kwargs)
# real domain scalar context
x_values = [1., -1., 0.0, -0.0, 0.5, -0.5, 5, -5, 5e-21, -5e-21]
x_types = [types.float32, types.float64] * (len(x_values) // 2)
check(x_types, x_values)
# real domain vector context
x_values = [np.array(x_values, dtype=np.float64)]
x_types = [typeof(v) for v in x_values]
check(x_types, x_values)
# complex domain scalar context
x_values = [1.+0j, -1+0j, 0.0+0.0j, -0.0+0.0j, 0+1j, 0-1j, 0.5+0.0j,
-0.5+0.0j, 0.5+0.5j, -0.5-0.5j, 5+5j, -5-5j,
# the following are to test sin(x)/x for small x
5e-21+0j, -5e-21+0j, 5e-21j, +(0-5e-21j)
]
x_types = [types.complex64, types.complex128] * (len(x_values) // 2)
check(x_types, x_values, ulps=2)
# complex domain vector context
x_values = [np.array(x_values, dtype=np.complex128)]
x_types = [typeof(v) for v in x_values]
check(x_types, x_values, ulps=2)
def check_scal(scal):
x = 4
y = 5
np.random.shuffle(_types)
x = _types[0](4)
y = _types[1](5)
cres = compile_isolated(pyfunc, (typeof(scal), typeof(x), typeof(y)))
expected = pyfunc(scal, x, y)
got = cres.entry_point(scal, x, y)
self.assertPreciseEqual(got, expected)
def check_array_flat(self, arr):
out = np.zeros(arr.size, dtype=arr.dtype)
nb_out = out.copy()
cres = compile_isolated(array_flat, [typeof(arr), typeof(out)])
cfunc = cres.entry_point
array_flat(arr, out)
cfunc(arr, nb_out)
self.assertTrue(np.all(out == nb_out), (out, nb_out))
def test_simple_literal(self):
@njit
def foo():
d = Dict()
d[123] = 321
return d
k, v = 123, 321
d = foo()
self.assertEqual(dict(d), {k: v})
self.assertEqual(typeof(d).key_type, typeof(k))
self.assertEqual(typeof(d).value_type, typeof(v))
def check_arr(arr):
x = np.zeros_like(arr, dtype=np.float64)
y = np.copy(x)
x.fill(4)
y.fill(9)
cres = compile_isolated(pyfunc, (typeof(arr), typeof(x), typeof(y)))
expected = pyfunc(arr, x, y)
got = cres.entry_point(arr, x, y)
# Contiguity of result varies accross Numpy versions, only
# check contents.
self.assertEqual(got.dtype, expected.dtype)
np.testing.assert_array_equal(got, expected)
def test_simple_args(self):
@njit
def foo(k, v):
d = Dict()
d[k] = v
return d
k, v = 123, 321
d = foo(k, v)
self.assertEqual(dict(d), {k: v})
self.assertEqual(typeof(d).key_type, typeof(k))
self.assertEqual(typeof(d).value_type, typeof(v))
def test_simple_upcast(self):
@njit
def foo(k, v, w):
d = Dict()
d[k] = v
d[k] = w
return d
k, v, w = 123, 32.1, 321
d = foo(k, v, w)
self.assertEqual(dict(d), {k: w})
self.assertEqual(typeof(d).key_type, typeof(k))
self.assertEqual(typeof(d).value_type, typeof(v))
def test_pow_usecase(self):
args = [
(2, 3),
(2.0, 3),
(2, 3.0),
(2j, 3.0j),
]
for x, y in args:
cres = compile_isolated(pow_usecase, (typeof(x), typeof(y)),
flags=no_pyobj_flags)
r = cres.entry_point(x, y)
self.assertPreciseEqual(r, pow_usecase(x, y))
def check_array_flat(self, arr, arrty=None):
out = np.zeros(arr.size, dtype=arr.dtype)
nb_out = out.copy()
if arrty is None:
arrty = typeof(arr)
cres = compile_isolated(array_flat, [arrty, typeof(out)])
cfunc = cres.entry_point
array_flat(arr, out)
cfunc(arr, nb_out)
self.assertPreciseEqual(out, nb_out)
def check_arr(arr):
cres = compile_isolated(pyfunc, (typeof(arr),))
expected = pyfunc(arr)
# NOTE: Numpy 1.9 returns readonly arrays for multidimensional
# arrays. Workaround this by copying the results.
expected = [a.copy() for a in expected]
self.assertPreciseEqual(cres.entry_point(arr), expected)
def test_array_flat_3d(self):
arr = np.arange(50).reshape(5, 2, 5)
arrty = typeof(arr)
self.assertEqual(arrty.ndim, 3)
out = np.zeros(arr.size, dtype=arr.dtype)
nb_out = out.copy()
cres = compile_isolated(array_flat, [arrty, typeof(out)])
cfunc = cres.entry_point
array_flat(arr, out)
cfunc(arr, nb_out)
self.assertTrue(np.all(out == nb_out))
def test_inspect_cfg(self):
# Exercise the .inspect_cfg(). These are minimal tests and do not fully
# check the correctness of the function.
@jit
def foo(the_array):
return the_array.sum()
# Generate 3 overloads
a1 = np.ones(1)
a2 = np.ones((1, 1))
a3 = np.ones((1, 1, 1))
foo(a1)
foo(a2)
foo(a3)
# Call inspect_cfg() without arguments
cfgs = foo.inspect_cfg()
# Correct count of overloads
self.assertEqual(len(cfgs), 3)
# Makes sure all the signatures are correct
[s1, s2, s3] = cfgs.keys()
self.assertEqual(set([s1, s2, s3]),
set(map(lambda x: (typeof(x),), [a1, a2, a3])))
for cfg in cfgs.values():
self._check_cfg_display(cfg)
self.assertEqual(len(list(cfgs.values())), 3)
# Call inspect_cfg(signature)
cfg = foo.inspect_cfg(signature=foo.signatures[0])
self._check_cfg_display(cfg)
def run_comparative(compare_func, test_array):
arrty = typeof(test_array)
cres = compile_isolated(compare_func, [arrty])
numpy_result = compare_func(test_array)
numba_result = cres.entry_point(test_array)
return numpy_result, numba_result
def check_arr(arr):
cres = compile_isolated(pyfunc, (typeof(arr),))
expected = pyfunc(arr)
got = cres.entry_point(arr)
self.assertPreciseEqual(expected, got)
if check_sameness:
self.assertEqual(is_same(expected, arr), is_same(got, arr))
def test_c04(self):
def bar(x):
f = x
f[0][0] = 10
return f
r = [[np.arange(3)]]
with self.assertRaises(errors.TypingError) as raises:
self.compile_and_test(bar, r)
self.assertIn(
"invalid setitem with value of {} to element of {}".format(
typeof(10),
typeof(r[0][0]),
),
str(raises.exception),
)
def test_simple_expr(self):
'''
Using a simple array expression, verify that rewriting is taking
place, and is fusing loops.
'''
A = np.linspace(0,1,10)
X = np.linspace(2,1,10)
Y = np.linspace(1,2,10)
arg_tys = [typeof(arg) for arg in (A, X, Y)]
control_pipeline, nb_axy_0, test_pipeline, nb_axy_1 = \
self._compile_function(axy, arg_tys)
control_pipeline2 = RewritesTester.mk_no_rw_pipeline(arg_tys)
cres_2 = control_pipeline2.compile_extra(ax2)
nb_ctl = cres_2.entry_point
expected = nb_axy_0(A, X, Y)
actual = nb_axy_1(A, X, Y)
control = nb_ctl(A, X, Y)
np.testing.assert_array_equal(expected, actual)
np.testing.assert_array_equal(control, actual)
ir0 = control_pipeline.interp.blocks
ir1 = test_pipeline.interp.blocks
ir2 = control_pipeline2.interp.blocks
self.assertEqual(len(ir0), len(ir1))
self.assertEqual(len(ir0), len(ir2))
# The rewritten IR should be smaller than the original.
self.assertGreater(len(ir0[0].body), len(ir1[0].body))
self.assertEqual(len(ir0[0].body), len(ir2[0].body))
def test_ufunc_and_dufunc_calls(self):
'''
Verify that ufunc and DUFunc calls are being properly included in
array expressions.
'''
A = np.random.random(10)
B = np.random.random(10)
arg_tys = [typeof(arg) for arg in (A, B)]
vaxy_descr = vaxy._dispatcher.targetdescr
control_pipeline = RewritesTester.mk_no_rw_pipeline(
arg_tys,
typing_context=vaxy_descr.typing_context,
target_context=vaxy_descr.target_context)
cres_0 = control_pipeline.compile_extra(call_stuff)
nb_call_stuff_0 = cres_0.entry_point
test_pipeline = RewritesTester.mk_pipeline(
arg_tys,
typing_context=vaxy_descr.typing_context,
target_context=vaxy_descr.target_context)
cres_1 = test_pipeline.compile_extra(call_stuff)
nb_call_stuff_1 = cres_1.entry_point
expected = call_stuff(A, B)
control = nb_call_stuff_0(A, B)
actual = nb_call_stuff_1(A, B)
np.testing.assert_array_almost_equal(expected, control)
np.testing.assert_array_almost_equal(expected, actual)
self._assert_total_rewrite(control_pipeline.interp.blocks,
test_pipeline.interp.blocks)
def test_alias_ctypes(self):
# use xxnrm2 to test call a C function with ctypes
from numba.targets.linalg import _BLAS
xxnrm2 = _BLAS().numba_xxnrm2(types.float64)
def remove_dead_xxnrm2(rhs, lives, call_list):
if call_list == [xxnrm2]:
return rhs.args[4].name not in lives
return False
# adding this handler has no-op effect since this function won't match
# anything else but it's a bit cleaner to save the state and recover
old_remove_handlers = remove_call_handlers[:]
remove_call_handlers.append(remove_dead_xxnrm2)
def func(ret):
a = np.ones(4)
xxnrm2(100, 4, a.ctypes, 1, ret.ctypes)
A1 = np.zeros(1)
A2 = A1.copy()
try:
pfunc = self.compile_parallel(func, (numba.typeof(A1),))
numba.njit(func)(A1)
pfunc(A2)
finally:
# recover global state
remove_call_handlers[:] = old_remove_handlers
self.assertEqual(A1[0], A2[0])
def test_a_is_not_b_array(self):
pyfunc = a_is_not_b
ary = numpy.arange(2)
aryty = typeof(ary)
cres = compile_isolated(pyfunc, [aryty, aryty])
cfunc = cres.entry_point
self.assertTrue(cfunc(ary, ary))
def compile(self, func, *args, **kwargs):
assert not kwargs
sig = tuple([numba.typeof(x) for x in args])
std = compile_isolated(func, sig, flags=self.flags)
fast = compile_isolated(func, sig, flags=self.fastflags)
return std, fast
def get(cls, device: Device, targets: List[Element],
bg_gases: Optional[List[BackgroundGas]] = None,
options: ModelOptions = DEFAULT_MODEL_OPTIONS) -> AdvancedModel:
# Types needed by numba
_T_BG_GAS = numba.typeof(BackgroundGas.get("He", 1e-8))
_T_F8_ARRAY = numba.float64[:]
bg_gases = bg_gases or []
if not bg_gases:
bg_gases_nblist = numba.typed.List.empty_list(_T_BG_GAS)
else:
bg_gases_nblist = numba.typed.List(bg_gases)
# Determine array bounds for different targets in state vector
lb = np.zeros(len(targets), dtype=np.int32)
ub = np.zeros(len(targets), dtype=np.int32)
offset = 0
for i, trgt in enumerate(targets):
lb[i] = offset
ub[i] = lb[i] + trgt.z + 1
offset = ub[i]
# Compute total number of charge states for all targets (len of "n" or "kT" state vector)
nq = int(ub[-1])
# Define vectors listing the charge state and mass for each state
q = np.zeros(nq, dtype=np.int32)
a = np.zeros(nq, dtype=np.int32)
for i, trgt in enumerate(targets):
q[lb[i]:ub[i]] = np.arange(trgt.z + 1, dtype=np.int32)
a[lb[i]:ub[i]] = np.full(trgt.z + 1, trgt.a, dtype=np.int32)
# Initialise cross section vectors
_eixs = np.zeros(nq)
_rrxs = np.zeros(nq)
_drxs = np.zeros(nq)
for i, trgt in enumerate(targets):
_eixs[lb[i]:ub[i]] = xs.eixs_vec(trgt, device.e_kin)
_rrxs[lb[i]:ub[i]] = xs.rrxs_vec(trgt, device.e_kin)
_drxs[lb[i]:ub[i]] = xs.drxs_vec(trgt, device.e_kin, device.fwhm)
# Precompute CX cross sections (invariant)
_cxxs_bggas = numba.typed.List.empty_list(_T_F8_ARRAY)
_cxxs_trgts = numba.typed.List.empty_list(_T_F8_ARRAY)
for gas in bg_gases:
_cxxs_bggas.append(xs.cxxs(q, gas.ip))
for trgt in targets:
_cxxs_trgts.append(xs.cxxs(q, trgt.ip))
return cls(
device=device,
targets=numba.typed.List(targets),
bg_gases=bg_gases_nblist,
options=options,
lb=lb,
ub=ub,
nq=nq,
q=q,
a=a,
eixs=_eixs,
rrxs=_rrxs,
drxs=_drxs,
cxxs_bggas=_cxxs_bggas,
cxxs_trgts=_cxxs_trgts
)
traces (np.ndarray): 2d array (Ntrials x Ntime) of evidence traces
"""
# generate storage objects for data/traces
data, rProb, traces = gen_ddm_storage_objects(parameters, ntrials,
deadline)
# simulate ntrials w/ DDM and fill data & traces array
_sim_ddm_trials_(parameters, data, rProb, traces)
# filter data/traces and return as pd.DataFrame
df, traces = clean_output(data, traces, deadline=deadline)
return df, traces
@jit(nb.typeof((1.0, 1.0))(float64[:], float64[:], float64[:]), nopython=True)
def sim_ddm(parameters, rProb, trace):
""" single trial simulation of the DDM (discrete time / random walk)
::Arguments::
parameters: 1d array (Nparams) of DDM parameters
rProb: 1d array (Ntimesteps) of random floats between 0 and 1
trace: 1d array (Ntimesteps) for storing the evidence trace
::Returns::
RT (float): the time that evidence crossed one of the boundaries
choice: 1 if evidence terminated at upper bound, 0 if lower bound
"""
# extract parameters
a, tr, v, z, si, dx, dt = parameters
def assert_matches(arr, val, exact):
expect = conv(arr, val)
got = jit_conv(arr, val)
self.assertPreciseEqual(expect, exact)
self.assertPreciseEqual(typeof(expect), typeof(got))
self.assertPreciseEqual(expect, got)
ufc_form00 = dolfinx.jit.ffcx_jit(a00)
kernel00 = ufc_form00.create_cell_integral(-1).tabulate_tensor
ufc_form01 = dolfinx.jit.ffcx_jit(a01)
kernel01 = ufc_form01.create_cell_integral(-1).tabulate_tensor
ufc_form10 = dolfinx.jit.ffcx_jit(a10)
kernel10 = ufc_form10.create_cell_integral(-1).tabulate_tensor
ffi = cffi.FFI()
cffi_support.register_type(ffi.typeof('double _Complex'),
numba.types.complex128)
c_signature = numba.types.void(
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.typeof(PETSc.ScalarType())),
numba.types.CPointer(numba.types.double),
numba.types.CPointer(numba.types.int32),
numba.types.CPointer(numba.types.uint8), numba.types.uint32)
@numba.cfunc(c_signature, nopython=True)
def tabulate_condensed_tensor_A(A_,
w_,
c_,
coords_,
entity_local_index,
permutation=ffi.NULL,
cell_permutation_info=0):
def typeof(val, c):
return RegularArrayType(numba.typeof(val.content),
numba.typeof(val.identities),
util.dict2parameters(val.parameters))
def check_arr(arr):
cres = compile_isolated(pyfunc, (typeof(arr), ))
expected = pyfunc(arr)
expected = [a.copy() for a in expected]
self.assertPreciseEqual(cres.entry_point(arr), expected)
def _initialise_list(self, item):
lsttype = types.ListType(typeof(item))
self._list_type, self._opaque = self._parse_arg(lsttype)
def run(shape):
cres = self.ccache.compile(pyfunc, (typeof(shape), ))
return cres.entry_point(shape)
def run(buf):
cres = self.ccache.compile(pyfunc, (typeof(buf), ))
return cres.entry_point(buf)
import numba
def foo(x):
for k in x:
print("k:", k)
for a, b in x[k]:
print(a, b)
if sys.argv[-1] == "-c":
# numba compile
print("compiling")
from numba.pycc import CC
cc = CC('my_module')
cc.export(
'foo', 'void(DictType(int64,ListType(UniTuple(int64,2))))')(foo)
cc.compile()
exit()
else:
x = numba.typed.Dict()
x[1] = numba.typed.List([(1, 2), (3, 4)])
x[100] = numba.typed.List([(100, 200)])
if sys.argv[-1] == "-p":
print(numba.typeof(x))
# => DictType[int64,ListType[UniTuple(int64 x 2)]]
else:
from my_module import foo
foo(x)
def create_my_class(value):
cls = jitclass([('value', typeof(value))])(MyClass)
return cls(value)
def _initialise_dict(self, key, value):
dcttype = types.DictType(typeof(key), typeof(value))
self._dict_type, self._opaque = self._parse_arg(dcttype)
def __init__(self, size):
self.arr = np.zeros(size, dtype=float)
self.type = nb.typeof(self.arr)
def interpolate_compiled(compiled_expression, target_func):
"""Compiled interpolation
Interpolates UFL expression into target function using FFCx
code generation and compilation.
Note
----
Works only for affine-mapped point-evaluation finite elements, e.g.
lagrange/discontinuous lagrange of arbitrary order.
"""
kernel = compiled_expression.module[0].tabulate_tensor_float64
# Register complex types
cffi_support.register_type(ffi.typeof('double _Complex'),
numba.types.complex128)
cffi_support.register_type(ffi.typeof('float _Complex'),
numba.types.complex64)
# Unpack mesh and dofmap data
mesh = target_func.function_space.mesh
geom_dofmap = mesh.geometry.dofmap.array
geom_pos = mesh.geometry.dofmap.offsets
geom = mesh.geometry.x
gdim = mesh.geometry.dim
dofmap = target_func.function_space.dofmap.list.array
# Prepare coefficients and their dofmaps
# Global vectors and dofmaps are prepared here, local are
# fetched inside hot cell-loop
# Number of coefficients in ffcx-processed ufl form
num_coeffs = compiled_expression.module[0].num_coefficients
# Positions of ffcx-preprocessed coefficients in original form
cpos = compiled_expression.module[0].original_coefficient_positions
coeffs = ufl.algorithms.analysis.extract_coefficients(compiled_expression.expr)
coeffs_dofmaps = List.empty_list(numba.types.Array(numba.typeof(dofmap[0]), 1, "C", readonly=True))
coeffs_vectors = List.empty_list(numba.types.Array(numba.typeof(PETSc.ScalarType()), 1, "C", readonly=True))
coeffs_bs = List.empty_list(numba.types.int_)
for i in range(num_coeffs):
coeffs_dofmaps.append(coeffs[cpos[i]].function_space.dofmap.list.array)
coeffs_vectors.append(np.asarray(coeffs[cpos[i]].vector))
coeffs_bs.append(coeffs[cpos[i]].function_space.dofmap.bs)
coeffs_sizes = np.asarray([coeff.function_space.element.space_dimension
for coeff in coeffs], dtype=np.int_)
# Prepare and pack constants
constants = ufl.algorithms.analysis.extract_constants(compiled_expression.expr)
constants_vector = np.array([], dtype=PETSc.ScalarType())
if len(constants) > 0:
constants_vector = np.hstack([c.value.flatten() for c in constants])
value_size = int(np.product(compiled_expression.expr.ufl_shape))
basix_space_dim = compiled_expression.ffcx_element.dim
space_dim = target_func.function_space.element.space_dimension
dofmap_bs = target_func.function_space.dofmap.bs
element_bs = target_func.function_space.dofmap.dof_layout.block_size
# Prepare mapping of subelements into the parent finite element
# This mapping stores also how dofs are collapsed when symmetry to a TensorElement
# is applied
if hasattr(compiled_expression.target_el, "flattened_sub_element_mapping"):
subel_map = np.array(compiled_expression.target_el.flattened_sub_element_mapping())
else:
subel_map = np.array(range(value_size))
num_coeffs = len(coeffs_vectors)
with target_func.vector.localForm() as b:
b.set(0.0)
assemble_vector_ufc(np.asarray(b), kernel, (geom_dofmap, geom_pos, geom), dofmap,
coeffs_vectors, coeffs_dofmaps, coeffs_bs, constants_vector,
coeffs_sizes, gdim, basix_space_dim, space_dim,
value_size, subel_map, dofmap_bs, element_bs)
for x in range(0, row):
for y in range(0, col):
if (img[x, y] == 0):
xVal += x
yVal += y
n += 1.0
xVal /= n
yVal /= n
return [np.int64(np.round(xVal)), np.int64(np.round(yVal))]
@nb.jit(
nb.typeof(np.matrix([[], []],
np.uint8))(nb.typeof(np.matrix([[], []], np.uint8)),
nb.typeof([1, 1]), nb.float64))
def createLine(img, centroid, alpha):
img2 = np.copy(img)
row, col = img.shape
xVal = np.float64(centroid[0])
yVal = np.float64(centroid[1])
distance = -1
distanceTemp = 0
while ((xVal >= 0 and xVal < row) and (yVal >= 0 and yVal < col)):
#print(xVal,yVal)
img2[np.int64(np.floor(xVal)), np.int64(np.floor(yVal))] = 100
if (img[np.int64(np.floor(xVal)), np.int64(np.floor(yVal))] == 0):
import numba
import numpy as np
from numba import float64, int64
from numba import generated_jit, njit
import ast
from numba.extending import overload
from numba.types.containers import Tuple, UniTuple
# from math import max, min
####################
# Dimension helper #
####################
t_coord = numba.typeof((2.3,2.4,1)) # type of an evenly spaced dimension
t_array = numba.typeof(np.array([4.0, 3.9])) # type of an unevenly spaced dimension
@njit
def clamp(x,a,b):
return min(b,max(a,x))
# returns the index of a 1d point along a 1d dimension
@generated_jit(nopython=True)
def get_index(gc, x):
if gc == t_coord:
# regular coordinate
def fun(gc,x):
T_inf = 8.2934
T_sup = 36.0729
elif self.stage == 'pupa':
a = 0.0003102
T_inf = 11.9433
T_sup = 40
else:
return
mean_rate = a * self.T * (self.T - T_inf) * np.sqrt(T_sup - self.T)
mean = 0.2789 * density**0.2789 / mean_rate
self.stage_length = np.random.gamma(mean * 10, 0.1)
# In[22]:
mosquito_type = typeof(Mosquito('egg', 25, 0))
mosquito_list_type = typeof([Mosquito('egg', 25, 0)])
specs = [('egg_numbers', types.ListType(int64)),
('larva_numbers', types.ListType(int64)),
('pupa_numbers', types.ListType(int64)),
('naive_numbers', types.ListType(int64)),
('adult_numbers', types.ListType(int64)),
('total_numbers', types.ListType(int64)),
('Temps', types.ListType(float64)), ('T', float64),
('mosquitos', mosquito_list_type)]
@jitclass(specs)
class Population:
def __init__(self, egg_no, larva_no, pupa_no, naive_no, adult_no, Temps):
def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval,
whereval, advanced):
if len(wheretpe.types) == 0:
return arrayval
headtpe = wheretpe.types[0]
tailtpe = numba.types.Tuple(wheretpe.types[1:])
headval = numba.cgutils.unpack_tuple(builder, whereval)[0]
tailval = context.make_tuple(
builder, tailtpe,
numba.cgutils.unpack_tuple(builder, whereval)[1:])
proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context,
builder,
value=arrayval)
leng = util.arraylen(context,
builder,
arraytpe,
arrayval,
totpe=numba.int64)
if isinstance(headtpe, numba.types.Integer):
assert advanced is None
nextcarry = util.newindex64(context, builder, numba.int64, leng)
util.call(context, builder,
cpu.kernels.awkward_regulararray_getitem_next_at_64,
(util.arrayptr(context, builder, util.index64tpe, nextcarry),
util.cast(context, builder, headtpe, numba.int64,
headval), leng, proxyin.size),
"in {0}, indexing error".format(arraytpe.shortname))
nextcontenttpe = arraytpe.contenttpe.carry()
nextcontentval = arraytpe.contenttpe.lower_carry(
context, builder, arraytpe.contenttpe, util.index64tpe,
proxyin.content, nextcarry)
return nextcontenttpe.lower_getitem_next(context, builder,
nextcontenttpe, tailtpe,
nextcontentval, tailval,
advanced)
elif isinstance(headtpe, numba.types.SliceType):
proxyslicein = context.make_helper(builder, headtpe, value=headval)
numba.targets.slicing.guard_invalid_slice(context, builder, headtpe,
proxyslicein)
numba.targets.slicing.fix_slice(
builder, proxyslicein,
util.cast(context, builder, numba.int64, numba.intp, proxyin.size))
nextsize = util.cast(
context, builder, numba.intp, numba.int64,
numba.targets.slicing.get_slice_length(builder, proxyslicein))
nextcarry = util.newindex64(context, builder, numba.int64,
builder.mul(leng, nextsize))
util.call(context, builder,
cpu.kernels.awkward_regulararray_getitem_next_range_64,
(util.arrayptr(context, builder, util.index64tpe, nextcarry),
util.cast(context, builder, numba.intp, numba.int64,
proxyslicein.start),
util.cast(context, builder, numba.intp, numba.int64,
proxyslicein.step), leng, proxyin.size, nextsize),
"in {0}, indexing error".format(arraytpe.shortname))
nextcontenttpe = arraytpe.contenttpe.carry()
nextcontentval = arraytpe.contenttpe.lower_carry(
context, builder, arraytpe.contenttpe, util.index64tpe,
proxyin.content, nextcarry)
if advanced is None:
outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False)
outcontentval = nextcontenttpe.lower_getitem_next(
context, builder, nextcontenttpe, tailtpe, nextcontentval,
tailval, advanced)
else:
nextadvanced = util.newindex64(context, builder, numba.int64,
builder.mul(leng, nextsize))
util.call(
context, builder, cpu.kernels.
awkward_regulararray_getitem_next_range_spreadadvanced_64,
(util.arrayptr(context, builder, util.index64tpe,
nextadvanced),
util.arrayptr(context, builder, util.index64tpe,
advanced), leng, nextsize),
"in {0}, indexing error".format(arraytpe.shortname))
outcontenttpe = nextcontenttpe.getitem_next(tailtpe, True)
outcontentval = nextcontenttpe.lower_getitem_next(
context, builder, nextcontenttpe, tailtpe, nextcontentval,
tailval, nextadvanced)
outtpe = RegularArrayType(outcontenttpe, arraytpe.identitiestpe,
arraytpe.parameters)
proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder)
proxyout.content = outcontentval
proxyout.size = nextsize
if arraytpe.identitiestpe != numba.none:
proxyout.identities = proxyin.identities
return proxyout._getvalue()
elif isinstance(headtpe, numba.types.StringLiteral):
nexttpe = arraytpe.getitem_str(headtpe.literal_value)
nextval = lower_getitem_str(context, builder,
nexttpe(arraytpe,
headtpe), (arrayval, headval))
return lower_getitem_next(context, builder, nexttpe, tailtpe, nextval,
tailval, advanced)
elif isinstance(headtpe, numba.types.EllipsisType):
raise NotImplementedError("RegularArray.getitem_next(ellipsis)")
elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))):
raise NotImplementedError("RegularArray.getitem_next(newaxis)")
elif isinstance(headtpe, numba.types.Array):
if headtpe.ndim != 1:
raise NotImplementedError("array.ndim != 1")
flathead = numba.targets.arrayobj.array_ravel(context, builder,
util.index64tpe(headtpe),
(headval, ))
lenflathead = util.arraylen(context,
builder,
util.index64tpe,
flathead,
totpe=numba.int64)
regular_flathead = util.newindex64(context, builder, numba.int64,
lenflathead)
util.call(
context, builder,
cpu.kernels.awkward_regulararray_getitem_next_array_regularize_64,
(util.arrayptr(context, builder, util.index64tpe,
regular_flathead),
util.arrayptr(context, builder, util.index64tpe,
flathead), lenflathead, proxyin.size),
"in {0}, indexing error".format(arraytpe.shortname))
if advanced is None:
lencarry = builder.mul(leng, lenflathead)
nextcarry = util.newindex64(context, builder, numba.int64,
lencarry)
nextadvanced = util.newindex64(context, builder, numba.int64,
lencarry)
util.call(
context, builder,
cpu.kernels.awkward_regulararray_getitem_next_array_64,
(util.arrayptr(context, builder, util.index64tpe, nextcarry),
util.arrayptr(context, builder, util.index64tpe,
nextadvanced),
util.arrayptr(
context, builder, util.index64tpe,
regular_flathead), leng, lenflathead, proxyin.size),
"in {0}, indexing error".format(arraytpe.shortname))
nexttpe = arraytpe.contenttpe.carry()
nextval = arraytpe.contenttpe.lower_carry(context, builder,
arraytpe.contenttpe,
util.index64tpe,
proxyin.content,
nextcarry)
contenttpe = nexttpe.getitem_next(tailtpe, True)
contentval = nexttpe.lower_getitem_next(context, builder, nexttpe,
tailtpe, nextval, tailval,
nextadvanced)
outtpe = RegularArrayType(contenttpe, arraytpe.identitiestpe,
arraytpe.parameters)
proxyout = numba.cgutils.create_struct_proxy(outtpe)(context,
builder)
proxyout.content = contentval
proxyout.size = lenflathead
if outtpe.identitiestpe != numba.none:
proxyout.identities = awkward1._numba.identities.lower_getitem_any(
context, builder, outtpe.identitiestpe, util.index64tpe,
proxyin.identities, flathead)
return proxyout._getvalue()
else:
nextcarry = util.newindex64(context, builder, numba.int64, leng)
nextadvanced = util.newindex64(context, builder, numba.int64, leng)
util.call(
context, builder, cpu.kernels.
awkward_regulararray_getitem_next_array_advanced_64,
(util.arrayptr(context, builder, util.index64tpe, nextcarry),
util.arrayptr(context, builder, util.index64tpe,
nextadvanced),
util.arrayptr(context, builder, util.index64tpe, advanced),
util.arrayptr(
context, builder, util.index64tpe,
regular_flathead), leng, lenflathead, proxyin.size),
"in {0}, indexing error".format(arraytpe.shortname))
nexttpe = arraytpe.contenttpe.carry()
nextval = arraytpe.contenttpe.lower_carry(context, builder,
arraytpe.contenttpe,
util.index64tpe,
proxyin.content,
nextcarry)
outtpe = nexttpe.getitem_next(tailtpe, True)
return nexttpe.lower_getitem_next(context, builder, nexttpe,
tailtpe, nextval, tailval,
nextadvanced)
else:
raise AssertionError(headtpe)
locale.setlocale(locale.LC_NUMERIC, "C")
# Used for a floating point "nearly zero" comparison
EPS = 1e-8
INT32_MIN = np.iinfo(np.int32).min + 1
INT32_MAX = np.iinfo(np.int32).max - 1
FlatTree = namedtuple(
"FlatTree", ["hyperplanes", "offsets", "children", "indices", "leaf_size"]
)
dense_hyperplane_type = numba.float32[::1]
sparse_hyperplane_type = numba.float64[:, ::1]
offset_type = numba.float64
children_type = numba.typeof((np.int32(-1), np.int32(-1)))
point_indices_type = numba.int32[::1]
@numba.njit(
numba.types.Tuple(
(numba.int32[::1], numba.int32[::1], dense_hyperplane_type, offset_type)
)(numba.float32[:, ::1], numba.int32[::1], numba.int64[::1]),
locals={
"n_left": numba.uint32,
"n_right": numba.uint32,
"hyperplane_vector": numba.float32[::1],
"hyperplane_offset": numba.float32,
"margin": numba.float32,
"d": numba.uint32,
"i": numba.uint32,
def check(a, b, expected):
cres = compile_isolated(pyfunc, (typeof(a), typeof(b)))
self.assertPreciseEqual(cres.entry_point(a, b),
(expected, not expected))
def get_conn_flattened(self, x, sections, pad=False):
batch_size = x.shape[0]
N = x.shape[1] // 2
n_jops = len(self.jump_operators)
assert sections.shape[0] == batch_size
# Separate row and column inputs
xr, xc = x[:, 0:N], x[:, N:2 * N]
# Compute all flattened connections of each term
sections_r = np.empty(batch_size, dtype=np.int64)
sections_c = np.empty(batch_size, dtype=np.int64)
xr_prime, mels_r = self._Hnh.get_conn_flattened(xr, sections_r)
xc_prime, mels_c = self._Hnh.get_conn_flattened(xc, sections_c)
if pad:
# if else to accomodate for batches of 1 element, because
# sections don't start from 0-index...
# TODO: if we ever switch to 0-indexing, remove this.
if batch_size > 1:
max_conns_r = np.max(np.diff(sections_r))
max_conns_c = np.max(np.diff(sections_c))
else:
max_conns_r = sections_r[0]
max_conns_c = sections_c[0]
max_conns_Lrc = 0
# Must type those lists otherwise, if they are empty, numba
# cannot infer their type
L_xrps = List.empty_list(numba.typeof(x.dtype)[:, :])
L_xcps = List.empty_list(numba.typeof(x.dtype)[:, :])
L_mel_rs = List.empty_list(numba.typeof(self.dtype())[:])
L_mel_cs = List.empty_list(numba.typeof(self.dtype())[:])
sections_Lr = np.empty(batch_size * n_jops, dtype=np.int32)
sections_Lc = np.empty(batch_size * n_jops, dtype=np.int32)
for (i, L) in enumerate(self._jump_ops):
L_xrp, L_mel_r = L.get_conn_flattened(
xr, sections_Lr[i * batch_size:(i + 1) * batch_size])
L_xcp, L_mel_c = L.get_conn_flattened(
xc, sections_Lc[i * batch_size:(i + 1) * batch_size])
L_xrps.append(L_xrp)
L_xcps.append(L_xcp)
L_mel_rs.append(L_mel_r)
L_mel_cs.append(L_mel_c)
if pad:
if batch_size > 1:
max_lr = np.max(
np.diff(sections_Lr[i * batch_size:(i + 1) *
batch_size]))
max_lc = np.max(
np.diff(sections_Lc[i * batch_size:(i + 1) *
batch_size]))
else:
max_lr = sections_Lr[i * batch_size]
max_lc = sections_Lc[i * batch_size]
max_conns_Lrc += max_lr * max_lc
# compose everything again
if self._xprime_f.shape[0] < self._max_conn_size * batch_size:
# refcheck=False because otherwise this errors when testing
self._xprime_f.resize(self._max_conn_size * batch_size,
self.hilbert.size,
refcheck=False)
self._mels_f.resize(self._max_conn_size * batch_size,
refcheck=False)
if pad:
pad = max_conns_Lrc + max_conns_r + max_conns_c
else:
pad = 0
self._xprime_f[:] = 0
self._mels_f[:] = 0
return self._get_conn_flattened_kernel(
self._xprime_f,
self._mels_f,
sections,
np.asarray(xr),
np.asarray(xc),
sections_r,
sections_c,
xr_prime,
mels_r,
xc_prime,
mels_c,
L_xrps,
L_xcps,
L_mel_rs,
L_mel_cs,
sections_Lr,
sections_Lc,
n_jops,
batch_size,
N,
pad,
)
def check_arr(arr):
cres = compile_isolated(pyfunc, (typeof(arr), ))
expected = pyfunc(arr)
got = cres.entry_point(arr)
self.assertPreciseEqual(expected, got)
self.assertEqual(is_same(expected, arr), is_same(got, arr))
def test_typeof_type(arg):
"""
>>> test_typeof_type(int_)
meta(int)
"""
return numba.typeof(arg)
using_numba = True
try:
from numba import b1, f8, i2, i4, njit, typeof, u2
except ImportError:
using_numba = False
# replace numba functionality with "transparent" implementations
from timezonefinder._numba_replacements import b1, f8, i2, i4, njit, typeof, u2
from timezonefinder.configs import (
COORD2INT_FACTOR,
INT2COORD_FACTOR,
OCEAN_TIMEZONE_PREFIX,
)
dtype_3float_tuple = typeof((1.0, 1.0, 1.0))
dtype_2float_tuple = typeof((1.0, 1.0))
dtype_2int_tuple = typeof((1, 1))
# @cc.export('inside_polygon', 'b1(i4, i4, i4[:, :])')
@njit(b1(i4, i4, i4[:, :]), cache=True)
def inside_polygon(x, y, coordinates):
"""
Implementing the ray casting point in polygon test algorithm
cf. https://en.wikipedia.org/wiki/Point_in_polygon#Ray_casting_algorithm
:param x:
:param y:
:param coordinates: a polygon represented by a list containing two lists (x and y coordinates):
[ [x1,x2,x3...], [y1,y2,y3...]]
those lists are actually numpy arrays which are being read directly from a binary file
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
numba.types.BaseTuple)
def lower_getitem_tuple(context, builder, sig, args):
return content.lower_getitem_tuple(context, builder, sig, args)
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
numba.types.Array)
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
numba.types.List)
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
numba.types.ArrayCompatible)
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
numba.types.EllipsisType)
@numba.extending.lower_builtin(operator.getitem, ListArrayType,
type(numba.typeof(numpy.newaxis)))
def lower_getitem_other(context, builder, sig, args):
return content.lower_getitem_other(context, builder, sig, args)
def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval,
whereval, advanced):
import awkward1._numba.array.listoffsetarray
import awkward1._numba.array.regulararray
if len(wheretpe.types) == 0:
return arrayval
headtpe = wheretpe.types[0]
tailtpe = numba.types.Tuple(wheretpe.types[1:])
headval = numba.cgutils.unpack_tuple(builder, whereval)[0]
def run(arr, dtype):
pyfunc = make_array_astype(dtype)
cres = self.ccache.compile(pyfunc, (typeof(arr), ))
return cres.entry_point(arr)
def typeof(val, c):
return ListArrayType(numba.typeof(numpy.asarray(val.starts)),
numba.typeof(numpy.asarray(val.stops)),
numba.typeof(val.content),
numba.typeof(val.identities),
util.dict2parameters(val.parameters))
def is_model(agents, model):
"""Test if agent if type same type as model name
Args:
agents (numpy.ndarray):
model (str):
Returns:
bool:
"""
return hash(agents.dtype) == hash(AgentModelToType[model])
@numba.jit(void(typeof(agent_type_three_circle)[:]),
nopython=True,
nogil=True,
cache=True)
def shoulders(agents):
"""Positions of the center of mass, left- and right shoulders.
Args:
agents (ndarray):
Numpy array of datatype ``dtype=agent_type_three_circle``.
"""
for agent in agents:
tangent = rotate270(unit_vector(agent['orientation']))
offset = tangent * agent['r_ts']
agent['position_ls'][:] = agent['position'] - offset
agent['position_rs'][:] = agent['position'] + offset
def bar(x):
if isinstance(typeof(x), types.Float):
return x + 1234
else:
return x + 1
def create_my_class(value):
cls = jitclass(MyClass, [("value", typeof(value))])
return cls(value)