def __init__(self, shape, line_index=0, n_coils=1, platform="numpy"):
"""Initilize the 'FFT' class.
Parameters
----------
shape: tuple of int
shape of the image (not necessarly a square matrix).
n_coils: int, default 1
Number of coils used to acquire the signal in case of
multiarray receiver coils acquisition. If n_coils > 1,
data shape must be equal to [n_coils, Nx, Ny, NZ]
line_index: int
The index of the column onto the line_axis of the kspace
n_jobs: int, default 1
Number of parallel workers to use for fourier computation
All cores are used if -1
package: str
The plateform on which to run the computation. can be either 'numpy', 'numba', 'cupy'
"""
self.shape = shape
if n_coils <= 0:
n_coils = 1
self.n_coils = n_coils
self._exp_f = np.zeros(shape[1], dtype=complex)
self._exp_b = np.zeros(shape[1], dtype=complex)
self._mask = line_index
if platform == "numba" and n_coils > 1:
self._dft = njit(complex128[:, :](complex128[:, :, :],
complex128[:]),
parallel=True)(numba_njit_dft_multi)
self._idft = njit(complex128[:, :, :](complex128[:, :],
complex128[:]),
parallel=True)(numba_njit_idft_multi)
elif platform == "numba":
self._dft = njit(complex128[:](complex128[:, :], complex128[:]),
parallel=True)(numba_njit_dft_single)
self._idft = njit(complex128[:, :](complex128[:], complex128[:]),
parallel=True)(numba_njit_idft_single)
elif platform == "gufunc":
self._dft = guvectorize(
["complex128[:], complex128[:], complex128[:]"],
"(n),(n)->()",
nopython=True,
target="cpu",
)(numba_dft)
self._idft = guvectorize(
["complex128, complex128[:], complex128[:]"],
"(),(n)->(n)",
target="parallel",
nopython=True,
)(numba_idft)
elif platform == "numpy":
self._dft = np.dot
self._idft = np.multiply.outer
else:
raise NotImplementedError(
f"platform '{platform}' is not supported")
def test_nopython_flag(self):
def foo(A, B):
pass
# nopython = True is fine
guvectorize([void(float32[:], float32[:])], '(x)->(x)', target='cuda',
nopython=True)(foo)
# nopython = False is bad
with self.assertRaises(TypeError) as raises:
guvectorize([void(float32[:], float32[:])], '(x)->(x)',
target='cuda', nopython=False)(foo)
self.assertEqual("nopython flag must be True", str(raises.exception))
def test_invalid_flags(self):
# Check invalid flags
def foo(A, B):
pass
with self.assertRaises(TypeError) as raises:
guvectorize([void(float32[:], float32[:])], '(x)->(x)',
target='cuda', what1=True, ever2=False)(foo)
head = "The following target options are not supported:"
msg = str(raises.exception)
self.assertEqual(msg[:len(head)], head)
items = msg[len(head):].strip().split(',')
items = [i.strip("'\" ") for i in items]
self.assertEqual(set(['what1', 'ever2']), set(items))
def test_nopython_flag(self):
def foo(A, B):
pass
# nopython = True is fine
guvectorize([void(float32[:], float32[:])], '(x)->(x)', target='cuda',
nopython=True)(foo)
# nopython = False is bad
with self.assertRaises(TypeError) as raises:
guvectorize([void(float32[:], float32[:])], '(x)->(x)',
target='cuda', nopython=False)(foo)
self.assertEqual("nopython flag must be True", str(raises.exception))
def test_invalid_flags(self):
# Check invalid flags
def foo(A, B):
pass
with self.assertRaises(TypeError) as raises:
guvectorize([void(float32[:], float32[:])], '(x)->(x)',
target='cuda', what1=True, ever2=False)(foo)
head = "The following target options are not supported:"
msg = str(raises.exception)
self.assertEqual(msg[:len(head)], head)
items = msg[len(head):].strip().split(',')
items = [i.strip("'\" ") for i in items]
self.assertEqual(set(['what1', 'ever2']), set(items))
def make_method_from_factory(fff: FlatFunctionFactory,
vectorize=True,
use_file=False,
debug=False):
mod = compile_factory(fff)
if debug:
print(to_source(mod))
if vectorize:
coredims = [len(v) for k, v in fff.arguments.items()]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(fff.content)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature = None
fun = eval_ast(mod)
from numba import jit, guvectorize
jitted = jit(fun, nopython=True)
if vectorize:
gufun = guvectorize([fty], signature, target='parallel',
nopython=True)(fun)
return jitted, gufun
else:
return jitted
def test_scalar_input_core_type(self):
def pyfunc(inp, n, out):
for i in range(inp.size):
out[i] = n * (inp[i] + 1)
my_gufunc = guvectorize(['int32[:], int32, int32[:]'],
'(n),()->(n)',
target=self.target)(pyfunc)
# test single core loop execution
arr = np.arange(10).astype(np.int32)
got = my_gufunc(arr, 2)
expected = np.zeros_like(got)
pyfunc(arr, 2, expected)
np.testing.assert_equal(got, expected)
# test multiple core loop execution
arr = np.arange(20).astype(np.int32).reshape(10, 2)
got = my_gufunc(arr, 2)
expected = np.zeros_like(got)
for ax in range(expected.shape[0]):
pyfunc(arr[ax], 2, expected[ax])
np.testing.assert_equal(got, expected)
def make_method_from_factory(fff: FlatFunctionFactory, vectorize=True, use_file=False, debug=False):
mod = compile_factory(fff)
if debug:
print(to_source(mod))
if vectorize:
coredims = [len(v) for k, v in fff.arguments.items()]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(fff.content)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature = None
fun = eval_ast(mod)
from numba import jit, guvectorize
jitted = jit(fun, nopython=True)
if vectorize:
gufun = guvectorize(
[fty], signature, target='parallel', nopython=True)(fun)
return jitted, gufun
else:
return jitted
def guvectorize_compute(target: str, *, cache: bool = True):
return nb.guvectorize(
[nb.void(_nb_float[:, :], _nb_float[:], _nb_float, _nb_float[:])],
'(m, p),(p),()->(m)',
nopython=True,
target=target,
cache=cache)
def __call__(self):
sig = ['(f4, f4, f4[:])']
cfunc = guvectorize(sig, '(),()->()', **self._options)(gufunc_foo)
a = b = np.random.random(10).astype(np.float32)
expected = ufunc_foo(a, b)
got = cfunc(a, b)
np.testing.assert_allclose(expected, got)
def test_guvectorize_no_output(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y),(x,y)")(guadd)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
out = numpy.zeros_like(a)
ufunc(a, a, out)
self.assertTrue(numpy.all(a + a == out))
def test_guvectorize_decor(self):
gufunc = guvectorize(
[void(float32[:, :], float32[:, :], float32[:, :])],
'(m,n),(n,p)->(m,p)',
target=self.target)(matmulcore)
self.check_matmul_gufunc(gufunc)
def test_vector():
def fun(s,x,p,out):
out[0] = s[0] + x[0]
out[1] = s[1] + x[1]
s = numpy.random.random((2,))
x = numpy.random.random((2,))
p = numpy.random.random((2,))
out = numpy.zeros((2,))
out1 = numpy.zeros((2,))
from numba import guvectorize, float64, void
gfun = guvectorize(ftylist=[void(float64[:],float64[:],float64[:],float64[:])],
signature='(n),(n),(n)->(n)')(fun)
sfun = standard_function(gfun,2)
fun(s,x,p,out)
sfun(s,x,p,out=out1)
out2 = sfun(s,x,p)
out = sfun(s,x,p,diff=True)
print("OUT")
print(out)
def test_guvectorize_no_output(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y),(x,y)")(guadd)
a = np.arange(10, dtype='int32').reshape(2, 5)
out = np.zeros_like(a)
ufunc(a, a, out)
self.assertPreciseEqual(a + a, out)
def test_guvectorize(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y)->(x,y)")(guadd)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
b = ufunc(a, a)
self.assertTrue(numpy.all(a + a == b))
self.assertEqual(b.dtype, numpy.dtype('int32'))
def test_vector():
def fun(s, x, p, out):
out[0] = s[0] + x[0]
out[1] = s[1] + x[1]
s = numpy.random.random((2, ))
x = numpy.random.random((2, ))
p = numpy.random.random((2, ))
out = numpy.zeros((2, ))
out1 = numpy.zeros((2, ))
from numba import guvectorize, float64, void
gfun = guvectorize(
ftylist=[void(float64[:], float64[:], float64[:], float64[:])],
signature='(n),(n),(n)->(n)')(fun)
sfun = standard_function(gfun, 2)
fun(s, x, p, out)
sfun(s, x, p, out=out1)
out2 = sfun(s, x, p)
out = sfun(s, x, p, diff=True)
print("OUT")
print(out)
def _compile_lookup_ufuncs(cls, target):
# Export lookup tables to global variables so JIT compiling can cache the tables in the binaries
global CHARACTERISTIC, ORDER, EXP, LOG, ZECH_LOG, ZECH_E, ADD_JIT, MULTIPLY_JIT # pylint: disable=global-statement
CHARACTERISTIC = cls.characteristic
ORDER = cls.order
EXP = cls._EXP
LOG = cls._LOG
ZECH_LOG = cls._ZECH_LOG
if cls.characteristic == 2:
ZECH_E = 0
else:
ZECH_E = (cls.order - 1) // 2
kwargs = {"nopython": True, "target": target}
if target == "cuda":
kwargs.pop("nopython")
# JIT-compile add and multiply routines for reference in other routines
ADD_JIT = numba.jit("int64(int64, int64)", nopython=True)(_add_lookup)
MULTIPLY_JIT = numba.jit("int64(int64, int64)", nopython=True)(_multiply_lookup)
# Create numba JIT-compiled ufuncs using the *current* EXP, LOG, and MUL_INV lookup tables
cls._numba_ufunc_add = numba.vectorize(["int64(int64, int64)"], **kwargs)(_add_lookup)
cls._numba_ufunc_subtract = numba.vectorize(["int64(int64, int64)"], **kwargs)(_subtract_lookup)
cls._numba_ufunc_multiply = numba.vectorize(["int64(int64, int64)"], **kwargs)(_multiply_lookup)
cls._numba_ufunc_divide = numba.vectorize(["int64(int64, int64)"], **kwargs)(_divide_lookup)
cls._numba_ufunc_negative = numba.vectorize(["int64(int64)"], **kwargs)(_additive_inverse_lookup)
cls._numba_ufunc_multiple_add = numba.vectorize(["int64(int64, int64)"], **kwargs)(_multiple_add_lookup)
cls._numba_ufunc_power = numba.vectorize(["int64(int64, int64)"], **kwargs)(_power_lookup)
cls._numba_ufunc_log = numba.vectorize(["int64(int64)"], **kwargs)(_log_lookup)
cls._numba_ufunc_poly_eval = numba.guvectorize([(numba.int64[:], numba.int64[:], numba.int64[:])], "(n),(m)->(m)", **kwargs)(_poly_eval_lookup)
def test_guvectorize(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y)->(x,y)")(guadd)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
b = ufunc(a, a)
self.assertTrue(numpy.all(a + a == b))
self.assertEqual(b.dtype, numpy.dtype('int32'))
def test_guvectorize_no_output(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y),(x,y)")(guadd)
a = np.arange(10, dtype='int32').reshape(2, 5)
out = np.zeros_like(a)
ufunc(a, a, out)
self.assertPreciseEqual(a + a, out)
def test_guvectorize_no_output(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y),(x,y)")(guadd)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
out = numpy.zeros_like(a)
ufunc(a, a, out)
self.assertTrue(numpy.all(a + a == out))
def test_guvectorize_error_in_objectmode(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y)->(x,y)",
forceobj=True)(guerror)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
with self.assertRaises(MyException):
ufunc(a, a)
def __call__(self):
sig = ['(f4, f4, f4[:])']
cfunc = guvectorize(sig, '(),()->()', **self._options)(gufunc_foo)
a = b = np.random.random(10).astype(np.float32)
expected = ufunc_foo(a, b)
got = cfunc(a, b)
np.testing.assert_allclose(expected, got)
def test_scalar_input_core_type(self):
def pyfunc(inp, n, out):
for i in range(inp.size):
out[i] = n * (inp[i] + 1)
my_gufunc = guvectorize(['int32[:], int32, int32[:]'],
'(n),()->(n)',
target=self.target)(pyfunc)
# test single core loop execution
arr = np.arange(10).astype(np.int32)
got = my_gufunc(arr, 2)
expected = np.zeros_like(got)
pyfunc(arr, 2, expected)
np.testing.assert_equal(got, expected)
# test multiple core loop execution
arr = np.arange(20).astype(np.int32).reshape(10, 2)
got = my_gufunc(arr, 2)
expected = np.zeros_like(got)
for ax in range(expected.shape[0]):
pyfunc(arr[ax], 2, expected[ax])
np.testing.assert_equal(got, expected)
def gufunc(self):
extra_args = len(inspect.getargspec(self.func).args) - 2
dtype_str = ['void(%s)' % ','.join('%s[:]' % e
for e in d.split(','))
for d in self.signature]
sig = '(n)%s->(n)' % ''.join(',()' for _ in range(extra_args))
vectorize = numba.guvectorize(dtype_str, sig, nopython=True)
return vectorize(self.transformed_func)
def test_guvectorize_scalar_objectmode(self):
"""
Test passing of scalars to object mode gufuncs.
"""
ufunc = guvectorize(["(int32[:,:], int32, int32[:,:])"], "(x,y),()->(x,y)")(guadd_scalar_obj)
a = np.arange(10, dtype="int32").reshape(2, 5)
b = ufunc(a, 3)
self.assertTrue(np.all(a + 3 == b))
def test_is_jitted(self):
def foo(x):
pass
self.assertFalse(is_jitted(foo))
self.assertTrue(is_jitted(njit(foo)))
self.assertFalse(is_jitted(vectorize(foo)))
self.assertFalse(is_jitted(vectorize(parallel=True)(foo)))
self.assertFalse(is_jitted(guvectorize("void(float64[:])", "(m)")(foo)))
def make_wrapper(func):
theargs = inspect.getargspec(func).args
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
def wrapper(*args, **kwargs):
np_arrays = [getattr(args[0], i).values for i in theargs]
ans = vecd_f(*np_arrays)
return ans
return wrapper
def test_guvectorize_scalar_objectmode(self):
"""
Test passing of scalars to object mode gufuncs.
"""
ufunc = guvectorize(['(int32[:,:], int32, int32[:,:])'],
"(x,y),()->(x,y)")(guadd_scalar_obj)
a = np.arange(10, dtype='int32').reshape(2, 5)
b = ufunc(a, 3)
self.assertPreciseEqual(a + 3, b)
def test_guvectorize_scalar_objectmode(self):
"""
Test passing of scalars to object mode gufuncs.
"""
ufunc = guvectorize(["(int32[:,:], int32, int32[:,:])"],
"(x,y),()->(x,y)")(guadd_scalar_obj)
a = np.arange(10, dtype="int32").reshape(2, 5)
b = ufunc(a, 3)
self.assertPreciseEqual(a + 3, b)
def test_guvectorize_scalar_objectmode(self):
"""
Test passing of scalars to object mode gufuncs.
"""
ufunc = guvectorize(['(int32[:,:], int32, int32[:,:])'],
"(x,y),()->(x,y)")(guadd_scalar_obj)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
b = ufunc(a, 3)
self.assertTrue(numpy.all(a + 3 == b))
def check_gufunc_raise(self, **vectorize_args):
f = guvectorize(['int32[:], int32[:], int32[:]'], '(n),()->(n)',
**vectorize_args)(gufunc_foo)
arr = np.array([1, 2, -3, 4], dtype=np.int32)
out = np.zeros_like(arr)
with self.assertRaises(ValueError) as cm:
f(arr, 2, out)
# The gufunc bailed out after the error
self.assertEqual(list(out), [2, 4, 0, 0])
def gufunc(self):
extra_args = len(inspect.getargspec(self.func).args) - 2
dtype_str = [
'void(%s)' % ','.join('%s[:]' % e for e in d.split(','))
for d in self.signature
]
sig = '(n)%s->(n)' % ''.join(',()' for _ in range(extra_args))
vectorize = numba.guvectorize(dtype_str, sig, nopython=True)
return vectorize(self.transformed_func)
def check_gufunc_raise(self, **vectorize_args):
f = guvectorize(['int32[:], int32[:], int32[:]'], '(n),()->(n)',
**vectorize_args)(gufunc_foo)
arr = np.array([1, 2, -3, 4], dtype=np.int32)
out = np.zeros_like(arr)
with self.assertRaises(ValueError) as cm:
f(arr, 2, out)
# The gufunc bailed out after the error
self.assertEqual(list(out), [2, 4, 0, 0])
def make_wrapper(func):
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
@wraps(func)
def wrapper(*args):
arrays = [arg.values for arg in args]
ans = vecd_f(*arrays)
return ans
return wrapper
def make_wrapper(func):
theargs = inspect.getargspec(func).args
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
def wrapper(*args, **kwargs):
np_arrays = [getattr(args[0], i).values for i in theargs]
ans = vecd_f(*np_arrays)
return ans
return wrapper
def make_wrapper(func):
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
@wraps(func)
def wrapper(*args, **kwargs):
# np_arrays = [getattr(args[0], i).values for i in theargs]
arrays = [arg.values for arg in args]
ans = vecd_f(*arrays)
return ans
return wrapper
def make_wrapper(func):
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
@wraps(func)
def wrapper(*args):
arrays = [arg.values for arg in args]
ans = vecd_f(*arrays)
return ans
return wrapper
def make_wrapper(func):
vecd_f = guvectorize(dtype_args, dtype_sig)(func)
@wraps(func)
def wrapper(*args, **kwargs):
# np_arrays = [getattr(args[0], i).values for i in theargs]
arrays = [arg.values for arg in args]
ans = vecd_f(*arrays)
return ans
return wrapper
def rescaleData(data, scale, offset, dtype, clip):
data_out = np.empty_like(data, dtype=dtype)
key = (data.dtype.name, data_out.dtype.name)
func = rescale_functions.get(key)
if func is None:
func = numba.guvectorize([f'{key[0]}[:],f8,f8,f8,f8,{key[1]}[:]'],
'(n),(),(),(),()->(n)',
nopython=True)(rescale_clip_source)
rescale_functions[key] = func
func(data, scale, offset, clip[0], clip[1], out=data_out)
return data_out
def test_guvectorize(self):
def foo(x, out):
out[0] = x + math.sin(x)
x = np.random.random(8).astype(np.float32)
with override_config('DUMP_OPTIMIZED', True):
types = ['(float32, float32[:])']
sig = '()->()'
with captured_stdout() as fast_cap:
fastfoo = guvectorize(types, sig, fastmath=True)(foo)
fastllvm = fast_cap.getvalue()
with captured_stdout() as slow_cap:
slowfoo = guvectorize(types, sig)(foo)
slowllvm = slow_cap.getvalue()
expect = slowfoo(x)
got = fastfoo(x)
np.testing.assert_almost_equal(expect, got)
self.assertIn('fadd fast', fastllvm)
self.assertIn('call fast', fastllvm)
self.assertNotIn('fadd fast', slowllvm)
self.assertNotIn('call fast', slowllvm)
def gu_activate(func, *args, **kwargs):
"""
The guvectorize decorated method, runs in nopython mode and fastmaths which favors speed over precision
since this is a DL library everything is pretty unprecise and fastmaths performs well enough
:param func: a activation method or its derivative method
:param target: choose between | 1: None -> serial | 2: 'parallel' -> parallel | execution method
:return: the decorated activation function
"""
kwargs_ = {k: v for k, v in kwargs.items() if v is not None}
return guvectorize([(float32[:], float32[:]),
(float64[:], float64[:])],
'(n)->(n)', nopython=True, fastmath=True, *args, **kwargs_)(func)
def test_guvectorize(self):
def foo(x, out):
out[0] = x + math.sin(x)
x = np.random.random(8).astype(np.float32)
with override_config('DUMP_OPTIMIZED', True):
types = ['(float32, float32[:])']
sig = '()->()'
with captured_stdout() as fast_cap:
fastfoo = guvectorize(types, sig, fastmath=True)(foo)
fastllvm = fast_cap.getvalue()
with captured_stdout() as slow_cap:
slowfoo = guvectorize(types, sig)(foo)
slowllvm = slow_cap.getvalue()
expect = slowfoo(x)
got = fastfoo(x)
np.testing.assert_almost_equal(expect, got)
self.assertIn('fadd fast', fastllvm)
self.assertIn('call fast', fastllvm)
self.assertNotIn('fadd fast', slowllvm)
self.assertNotIn('call fast', slowllvm)
def gu_adam_optimizer(func, *args, **kwargs):
kwargs_ = {k: v for k, v in kwargs.items() if v is not None}
return guvectorize(
[(float32[:], float32[:], float32[:, :], float32, float32, float32,
float32, float32, float32[:], float32[:, :]),
(float64[:], float64[:], float64[:, :], float64, float64, float64,
float64, float64, float64[:], float64[:, :])],
'(n),(n),(m,n),(),(),(),(),()->(n),(m,n)',
nopython=True,
fastmath=True,
*args,
**kwargs_)(func)
def guvectorize(signature: str, layout: str,
**kwds) -> Callable[[Callable], Any]:
"""Convenience wrapper around :func:`numba.guvectorize`.
Generate signature for all possible data types and set a few healthy
defaults.
:param str signature:
numba signature, containing {T}
:param str layout:
as in :func:`numba.guvectorize`
:param kwds:
passed verbatim to :func:`numba.guvectorize`.
This function changes the default for cache from False to True.
example::
guvectorize("{T}[:], {T}[:]", "(i)->(i)")
Is the same as::
numba.guvectorize([
"float32[:], float32[:]",
"float64[:], float64[:]",
...
], "(i)->(i)", cache=True)
.. note::
Discussing upstream fix; see
`<https://github.com/numba/numba/issues/2936>`_.
"""
DTYPES = [
# uint needs to appear before signed int:
# https://github.com/numba/numba/issues/2934
"uint8",
"uint16",
"uint32",
"uint64",
"int8",
"int16",
"int32",
"int64",
"float32",
"float64",
"complex64",
"complex128",
]
if "{T}" in signature:
signatures = [signature.format(T=dtype) for dtype in DTYPES]
else:
signatures = [signature]
kwds.setdefault("cache", True)
return numba.guvectorize(signatures, layout, **kwds)
def compile_function_ast(equations, symbols, arg_names, output_names=None, funname='anonymous', rhs_only=False,
return_ast=False, print_code=False, definitions=None, vectorize=True, use_file=False):
arguments = OrderedDict()
for an in arg_names:
if an[0] != 'parameters':
t = an[1]
arguments[an[2]] = [(s,t) for s in symbols[an[0]]]
# arguments = [ [ (s,t) for s in symbols[sg]] for sg,t in arg_names if sg != 'parameters']
parameters = [(s,0) for s in symbols['parameters']]
targets = output_names
if targets is not None:
targets = [(s,targets[1]) for s in symbols[targets[0]]]
mod = make_function(equations, arguments, parameters, definitions=definitions, targets=targets, rhs_only=rhs_only, funname=funname)
from dolo.compiler.codegen import to_source
import dolo.config
if dolo.config.debug:
print(to_source(mod))
if vectorize:
from numba import float64, void
coredims = [len(symbols[an[0]]) for an in arg_names]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(equations)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
ftylist = [fty]
else:
signature=None
ftylist=None
if use_file:
fun = eval_ast_with_file(mod, print_code=True)
else:
fun = eval_ast(mod)
from numba import jit, guvectorize
jitted = jit(fun, nopython=True)
if vectorize:
gufun = guvectorize([fty], signature, target='parallel', nopython=True)(fun)
return jitted, gufun
else:
return jitted
return [f,None]
def test_func(nthreads, py_func=False):
def _test_func(acc, buf, local_mask):
set_num_threads(nthreads)
# set threads in parent function
set_num_threads(local_mask[0])
if local_mask[0] < N:
child_func(buf, local_mask[0])
acc[0] += get_num_threads()
buf = np.zeros((M, N), dtype=np.int64)
acc = np.zeros((M, 1), dtype=np.int64)
local_mask = (1 + np.arange(M) % mask).reshape((M, 1))
sig = ["void(int64[:], int64[:, :], int64[:])"]
layout = "(p), (n, m), (p)"
if not py_func:
_test_func = guvectorize(
sig, layout, nopython=True, target="parallel"
)(_test_func)
else:
_test_func = guvectorize(sig, layout, forceobj=True)(_test_func)
_test_func(acc, buf, local_mask)
return acc, buf
def _generate_rowwise_internal(target):
def _rowwise_distance_covariance_sqr_avl_generic_internal(
x, y, unbiased, res):
args = _get_impl_args_compiled(x, y, unbiased)
res[0] = _distance_covariance_sqr_avl_impl_compiled(*args)
return numba.guvectorize(
[(input_array, input_array, boolean, float64[:])],
'(n),(n),()->()',
nopython=True,
cache=True,
target=target)(_rowwise_distance_covariance_sqr_avl_generic_internal)
def target(cls, target):
"""
Retarget the just-in-time compiled `numba` ufuncs.
Parameters
----------
target : str
The `target` keyword argument from :obj:`numba.vectorize`, either `"cpu"`, `"parallel"`, or `"cuda"`.
"""
global CHARACTERISTIC, ADD_JIT, MULTIPLY_JIT # pylint: disable=global-statement
CHARACTERISTIC = cls.characteristic
if target not in ["cpu", "parallel", "cuda"]:
raise ValueError(
f"Valid numba compilation targets are ['cpu', 'parallel', 'cuda'], not {target}"
)
kwargs = {"nopython": True, "target": target}
if target == "cuda":
kwargs.pop("nopython")
cls.ufunc_mode = "calculate"
cls.ufunc_target = target
# JIT-compile add and multiply routines for reference in polynomial evaluation routine
ADD_JIT = numba.jit("int64(int64, int64)",
nopython=True)(_add_calculate)
MULTIPLY_JIT = numba.jit("int64(int64, int64)",
nopython=True)(_multiply_calculate)
# Create numba JIT-compiled ufuncs using the *current* EXP, LOG, and MUL_INV lookup tables
cls._numba_ufunc_add = numba.vectorize(["int64(int64, int64)"],
**kwargs)(_add_calculate)
cls._numba_ufunc_subtract = numba.vectorize(
["int64(int64, int64)"], **kwargs)(_subtract_calculate)
cls._numba_ufunc_multiply = numba.vectorize(
["int64(int64, int64)"], **kwargs)(_multiply_calculate)
cls._numba_ufunc_divide = numba.vectorize(["int64(int64, int64)"],
**kwargs)(_divide_calculate)
cls._numba_ufunc_negative = numba.vectorize(
["int64(int64)"], **kwargs)(_negative_calculate)
cls._numba_ufunc_multiple_add = numba.vectorize(
["int64(int64, int64)"], **kwargs)(_multiple_add_calculate)
cls._numba_ufunc_power = numba.vectorize(["int64(int64, int64)"],
**kwargs)(_power_calculate)
cls._numba_ufunc_log = numba.vectorize(["int64(int64)"],
**kwargs)(_log_calculate)
cls._numba_ufunc_poly_eval = numba.guvectorize(
[(numba.int64[:], numba.int64[:], numba.int64[:])], "(n),(m)->(m)",
**kwargs)(_poly_eval_calculate)
def test_cpu_guvectorize(self):
target = 'cpu'
gufunc = guvectorize([void(float32[:,:], float32[:,:], float32[:,:])],
'(m,n),(n,p)->(m,p)',
target=target)(matmulcore)
matrix_ct = 1001 # an odd number to test thread/block division in CUDA
A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2, 4)
B = np.arange(matrix_ct * 4 * 5, dtype=np.float32).reshape(matrix_ct, 4, 5)
C = gufunc(A, B)
Gold = ut.matrix_multiply(A, B)
self.assertTrue(np.allclose(C, Gold))
def _create_gufunc(self, core_ndim):
# creating compiling gufunc has some significant overhead (~130ms per
# function and number of dimensions to aggregate), so do this in a
# lazy fashion
colons = ','.join(':' for _ in range(core_ndim))
numba_sig = []
for signature in self.signature:
match = re.match('^(\w+)\((\w+)\)$', signature)
if not match:
raise ValueError('invalid signature')
out_dtype, in_dtype = match.groups()
numba_sig.append('void(%s[%s], %s[:])'
% (in_dtype, colons, out_dtype))
gufunc_sig = '(%s)->()' % ','.join(list('abcdefgijk')[:core_ndim])
vectorize = numba.guvectorize(numba_sig, gufunc_sig, nopython=True)
return vectorize(self.transformed_func)
def test_guvectorize_invalid_layout(self):
sigs = ["(int32[:,:], int32[:,:], int32[:,:])"]
# Syntax error
with self.assertRaises(ValueError) as raises:
guvectorize(sigs, ")-:")(guadd)
self.assertIn("bad token in signature", str(raises.exception))
# Output shape can't be inferred from inputs
with self.assertRaises(NameError) as raises:
guvectorize(sigs, "(x,y),(x,y)->(x,z,v)")(guadd)
self.assertEqual(str(raises.exception), "undefined output symbols: v,z")
# Arrow but no outputs
with self.assertRaises(ValueError) as raises:
guvectorize(sigs, "(x,y),(x,y),(x,y)->")(guadd)
def test_guvectorize_identity(self):
args = (['(int32[:,:], int32[:,:], int32[:,:])'], "(x,y),(x,y)->(x,y)")
for identity in self._supported_identities:
ufunc = guvectorize(*args, identity=identity)(guadd)
expected = None if identity == 'reorderable' else identity
self.assertEqual(ufunc.identity, expected)
# Default value is None
ufunc = guvectorize(*args)(guadd)
self.assertIs(ufunc.identity, None)
# Invalid values
with self.assertRaises(ValueError):
guvectorize(*args, identity='none')(add)
with self.assertRaises(ValueError):
guvectorize(*args, identity=2)(add)
def test_columns():
def fun(s,x,out):
out[0] = s[0] + x[0]
out[1] = s[1] + x[1]
from numba import guvectorize, float64, void
gfun = guvectorize(ftylist=[void(float64[:],float64[:],float64[:])],
signature='(n),(n)->(n)')(fun)
N = 5
s = numpy.random.random((N,2,))
x = numpy.random.random((2,))
out = numpy.zeros((N,2,))
out1 = numpy.zeros((N,2,))
sfun = standard_function(gfun,2)
for n in range(N):
fun(s[n,:],x,out[n,:])
sfun(s,x,out=out1)
out2 = sfun(s,x)
# print(out2)
# print(s+x)
print(s+x)
# assert( (abs(out2-s-x).max())<1e-8 )
print(out)
print(out1)
print(out2)
out, out_s, = sfun(s,x,diff=True)
out = X.copy()
arbitrage(m,s,x,y,M,S,X,Y,p,out)
import time
t1 = time.time()
arbitrage(m,s,x,y,M,S,X,Y,p,out)
t2 = time.time()
print("numexpr: {}".format(t2-t1))
from numba import guvectorize, float64
gufun_cpu = guvectorize([(float64[:],)*10],'(a),(b),(c),(d),(a),(b),(c),(d),(e)->(c)', target='cpu')(arbitrage_numba)
gufun_cpu(m,s,x,y,M,S,X,Y,p,out)
t1 = time.time()
gufun_cpu(m,s,x,y,M,S,X,Y,p,out)
t2 = time.time()
print('gufun: {}'.format(t2-t1))
gufun_parallel = guvectorize([(float64[:],)*10],'(a),(b),(c),(d),(a),(b),(c),(d),(e)->(c)', target='parallel')(arbitrage_numba)
gufun_parallel(m,s,x,y,M,S,X,Y,p,out)
t1 = time.time()
gufun_parallel(m,s,x,y,M,S,X,Y,p,out)
t2 = time.time()
print('gufun: {}'.format(t2-t1))
def test_guvectorize_decor(self):
gufunc = guvectorize([void(float32[:,:], float32[:,:], float32[:,:])],
'(m,n),(n,p)->(m,p)',
target=self.target)(matmulcore)
self.check_matmul_gufunc(gufunc)
def compile_function_ast(expressions, symbols, arg_names, output_names=None, funname='anonymous', return_ast=False, print_code=False, definitions=None, vectorize=True, use_file=False):
'''
expressions: list of equations as string
'''
from collections import OrderedDict
table = OrderedDict()
aa = arg_names
if output_names is not None:
aa = arg_names + [output_names]
for a in aa:
symbol_group = a[0]
date = a[1]
an = a[2]
for b in symbols[symbol_group]:
index = symbols[symbol_group].index(b)
table[(b, date)] = (an, index)
table_symbols = {k: (std_date_symbol(*k)) for k in table.keys()}
# standard assignment: i.e. k = s[0]
index = lambda x: Index(Num(x))
# declare symbols
aux_short_names = [e[2] for e in arg_names if e[0]=='auxiliaries']
preamble = []
for k in table: # order it
# k : var, date
arg, pos = table[k]
if not (arg in aux_short_names):
std_name = table_symbols[k]
val = Subscript(value=Name(id=arg, ctx=Load()), slice=index(pos), ctx=Load())
line = Assign(targets=[Name(id=std_name, ctx=Store())], value=val)
if arg != 'out':
preamble.append(line)
body = []
std_dates = StandardizeDates(symbols, aa)
if definitions is not None:
for k,v in definitions.items():
if isinstance(k, str):
lhs = ast.parse(k).body[0].value
if isinstance(v, str):
rhs = ast.parse(v).body[0].value
else:
rhs = v
lhs = std_dates.visit(lhs)
rhs = std_dates.visit(rhs)
vname = lhs.id
line = Assign(targets=[Name(id=vname, ctx=Store())], value=rhs)
preamble.append(line)
outs = []
for i, expr in enumerate(expressions):
expr = ast.parse(expr).body[0].value
# if definitions is not None:
# expr = ReplaceName(defs).visit(expr)
rexpr = std_dates.visit(expr)
rhs = rexpr
if output_names is not None:
varname = symbols[output_names[0]][i]
date = output_names[1]
out_name = table_symbols[(varname, date)]
else:
out_name = 'out_{}'.format(i)
line = Assign(targets=[Name(id=out_name, ctx=Store())], value=rhs)
body.append(line)
line = Assign(targets=[Subscript(value=Name(id='out', ctx=Load()),
slice=index(i), ctx=Store())], value=Name(id=out_name, ctx=Load()))
body.append(line)
arg_names = [e for e in arg_names if e[0]!="auxiliaries"]
args = [e[2] for e in arg_names] + ['out']
if is_python_3:
from ast import arg
f = FunctionDef(name=funname, args=arguments(args=[arg(arg=a) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]),
body=preamble + body, decorator_list=[])
else:
f = FunctionDef(name=funname, args=arguments(args=[Name(id=a, ctx=Param()) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]),
body=preamble + body, decorator_list=[])
mod = Module(body=[f])
mod = ast.fix_missing_locations(mod)
if print_code:
s = "Function {}".format(mod.body[0].name)
print("-" * len(s))
print(s)
print("-" * len(s))
print(to_source(mod))
if vectorize:
from numba import float64, void
coredims = [len(symbols[an[0]]) for an in arg_names]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(expressions)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims)+1)
ftylist = [fty]
else:
signature=None
ftylist=None
if use_file:
fun = eval_ast_with_file(mod, print_code=True)
else:
fun = eval_ast(mod)
jitted = njit(fun)
if vectorize:
gufun = guvectorize([fty], signature, target='parallel', nopython=True)(fun)
return jitted, gufun
else:
return jitted
m[y2, x] -= m[y, x] * c
for y in range(h - 1, 0 - 1, -1): # Backsubstitute
c = m[y, y]
for y2 in range(0, y):
for x in range(w - 1, y - 1, -1):
m[y2, x] -= m[y, x] * m[y2, y] / c
m[y, y] /= c
for x in range(h, w): # Normalize row y
m[y, x] /= c
for y in range(h):
sol[y] = m[y, w - 1]
serial_solve_numba = guvectorize("void(f8[:,:], f8[:])", "(m,n)->(m)")(solve)
from numpy.linalg import solve as linalg_solve
def serial_solve(A, B, diagnose=True):
if diagnose:
sol = zeros_like(B)
for i in range(sol.shape[0]):
try:
sol[i, :] = linalg_solve(A[i, :, :], B[i, :])
except:
def test_guvectorize_objectmode(self):
ufunc = guvectorize(["(int32[:,:], int32[:,:], int32[:,:])"], "(x,y),(x,y)->(x,y)")(guadd_obj)
a = np.arange(10, dtype="int32").reshape(2, 5)
b = ufunc(a, a)
self.assertTrue(np.all(a + a == b))
r : ndarray(float)
Array containing random values in [0, 1).
out : ndarray(float)
Output array.
"""
n = r.shape[0]
r.sort()
out[0] = r[0]
for i in range(1, n):
out[i] = r[i] - r[i-1]
out[n] = 1 - r[n-1]
_probvec_parallel = guvectorize(
['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='parallel'
)(_probvec)
_probvec_cpu = guvectorize(
['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='cpu'
)(_probvec)
@jit
def sample_without_replacement(n, k, num_trials=None, random_state=None):
"""
Randomly choose k integers without replacement from 0, ..., n-1.
Parameters
----------
n : scalar(int)
Number of integers, 0, ..., n-1, to sample from.
def test_guvectorize_error_in_objectmode(self):
ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
"(x,y),(x,y)->(x,y)", forceobj=True)(guerror)
a = numpy.arange(10, dtype='int32').reshape(2, 5)
with self.assertRaises(MyException):
ufunc(a, a)
