def test_prob3numba(ignore_fails=False, define_as_ref=False):
"""Run all unit test cases for prob3numba code"""
# Pull first test case to test calling `propagate_array`
tc_name, tc = next(iter(TEST_CASES.items()))
tc_ = deepcopy(tc)
logging.info(
"Testing call and return shape of `propagate_array` with test case '%s'",
tc_name,
)
# Test simple broadcasting over `nubars` and `energies` where both have
# same shape, as this is the "typical" use case
input_shape = (4, 5)
# Without broadcasting, a single probability matrix is 3x3
prob_array_shape = (3, 3)
# Broadcasted shape
out_shape = input_shape + prob_array_shape
nubars = np.full(shape=input_shape, fill_value=tc_["nubar"], dtype=IX)
energies = np.full(shape=input_shape, fill_value=tc_["energy"], dtype=FX)
# Fill with NaN to ensure all elements are assinged a value
probabilities = SmartArray(np.full(shape=out_shape, fill_value=np.nan, dtype=FX))
propagate_array(
SmartArray(tc_["dm"].astype(FX)).get(WHERE),
SmartArray(tc_["pmns"].astype(CX)).get(WHERE),
SmartArray(tc_["mat_pot"].astype(CX)).get(WHERE),
SmartArray(nubars).get(WHERE),
SmartArray(energies).get(WHERE),
SmartArray(tc_["layer_densities"].astype(FX)).get(WHERE),
SmartArray(tc_["layer_distances"].astype(FX)).get(WHERE),
# output:
probabilities.get(WHERE),
)
probabilities.mark_changed(WHERE)
probabilities = probabilities.get("host")
# Check that all probability matrices have no NaNs and are equal to one
# another
ref_probs = probabilities[0, 0]
for i in range(input_shape[0]):
for j in range(input_shape[1]):
probs = probabilities[i, j]
assert np.all(np.isfinite(probs))
assert np.all(probs == ref_probs)
# Run all test cases
for tc_name, tc in TEST_CASES.items():
run_test_case(
tc_name, tc, ignore_fails=ignore_fails, define_as_ref=define_as_ref
)
def test_transpose(self):
# To verify non-redundant data movement run this test with NUMBA_TRACE=1
a = SmartArray(np.arange(16, dtype=float).reshape(4, 4))
b = SmartArray(where='gpu', shape=(4, 4), dtype=float)
c = SmartArray(where='gpu', shape=(4, 4), dtype=float)
event("initialization done")
transpose(a, b)
event("checkpoint")
transpose(b, c)
event("done")
self.assertTrue((c.get('host') == a.get('host')).all())
def test_clear_matrix():
"""Unit tests of `clear_matrix` and `clear_matrix_guf`"""
A = SmartArray(np.ones((4, 3), dtype=FTYPE))
clear_matrix_guf(A.get(WHERE), A.get(WHERE))
A.mark_changed(WHERE)
test = A.get()
ref = np.zeros((4, 3), dtype=FTYPE)
assert np.array_equal(test, ref), f"test:\n{test}\n!= ref:\n{ref}"
logging.info("<< PASS : test_clear_matrix >>")
def test_transpose(self):
# To verify non-redundant data movement run this test with NUMBA_TRACE=1
a = SmartArray(np.arange(16, dtype=float).reshape(4,4))
b = SmartArray(where='gpu', shape=(4,4), dtype=float)
c = SmartArray(where='gpu', shape=(4,4), dtype=float)
event("initialization done")
transpose(a, b)
event("checkpoint")
transpose(b, c)
event("done")
self.assertTrue((c.get('host') == a.get('host')).all())
def test_conjugate():
"""Unit tests of `conjugate` and `conjugate_guf`"""
A = SmartArray((np.linspace(1, 12, 12) +
1j * np.linspace(21, 32, 12)).reshape(4, 3).astype(CX))
B = SmartArray(np.ones((4, 3), dtype=CX))
conjugate_guf(A.get(WHERE), B.get(WHERE))
B.mark_changed(WHERE)
test = B.get()
ref = A.get().conj()
assert np.allclose(test, ref,
**ALLCLOSE_KW), f"test:\n{test}\n!= ref:\n{ref}"
A = SmartArray(np.linspace(1, 12, 12, dtype=FX).reshape(3, 4))
B = SmartArray(np.ones((3, 4), dtype=FX))
conjugate_guf(A.get(WHERE), B.get(WHERE))
B.mark_changed(WHERE)
test = B.get()
ref = A.get().conj()
assert np.allclose(test, ref,
**ALLCLOSE_KW), f"test:\n{test}\n!= ref:\n{ref}"
logging.info("<< PASS : test_conjugate >>")
def test_matrix_dot_matrix():
"""Unit tests of `matrix_dot_matrix` and `matrix_dot_matrix_guf`"""
A = SmartArray(np.linspace(1, 12, 12, dtype=FTYPE).reshape(3, 4))
B = SmartArray(np.linspace(1, 12, 12, dtype=FTYPE).reshape(4, 3))
C = SmartArray(np.ones((3, 3), dtype=FTYPE))
matrix_dot_matrix_guf(A.get(WHERE), B.get(WHERE), C.get(WHERE))
C.mark_changed(WHERE)
test = C.get()
ref = np.dot(A, B).astype(FTYPE)
assert np.allclose(test, ref,
**ALLCLOSE_KW), f"test:\n{test}\n!= ref:\n{ref}"
logging.info("<< PASS : test_matrix_dot_matrix >>")
def test_matrix_dot_vector():
"""Unit tests of `matrix_dot_vector` and `matrix_dot_vector_guf`"""
A = SmartArray(np.linspace(1, 12, 12, dtype=FTYPE).reshape(4, 3))
v = SmartArray(np.linspace(1, 3, 3, dtype=FTYPE))
w = SmartArray(np.ones(4, dtype=FTYPE))
matrix_dot_vector_guf(A.get(WHERE), v.get(WHERE), w.get(WHERE))
w.mark_changed(WHERE)
test = w.get()
ref = np.dot(A, v).astype(FTYPE)
assert np.allclose(test, ref,
**ALLCLOSE_KW), f"test:\n{test}\n!= ref:\n{ref}"
logging.info("<< PASS : test_matrix_dot_vector >>")
def test_imul_and_scale():
"""Unit tests for function ``imul_and_scale``"""
a = SmartArray(np.linspace(0, 1, 1000, dtype=FTYPE))
out = SmartArray(np.ones_like(a))
imul_and_scale(vals=a, scale=10.0, out=out)
assert np.allclose(out.get("host"), np.linspace(0, 10, 1000, dtype=FTYPE))
logging.info("<< PASS : test_multiply_and_scale >>")
def add_array_data(self, key, data):
"""
Parameters
----------
key : string
identifier
data : ndarray
"""
if isinstance(data, np.ndarray):
data = SmartArray(data)
if self.array_length is None:
self.array_length = data.get('host').shape[0]
assert data.get('host').shape[0] == self.array_length
self.array_data[key] = data
def test_astype(self):
a = SmartArray(np.int32([42, 8, -5]))
aa = a.astype(np.float64)
self.assertIsInstance(aa, SmartArray)
# verify that SmartArray.astype() operates like ndarray.astype()...
self.assertPreciseEqual(aa.get('host'), a.get('host').astype(np.float64))
# ...and that both actually yield the expected dtype.
self.assertPreciseEqual(aa.get('host').dtype.type, np.float64)
self.assertIs(aa.dtype.type, np.float64)
def test_astype(self):
a = SmartArray(np.int32([42, 8, -5]))
aa = a.astype(np.float64)
self.assertIsInstance(aa, SmartArray)
# verify that SmartArray.astype() operates like ndarray.astype()...
self.assertPreciseEqual(aa.get('host'),
a.get('host').astype(np.float64))
# ...and that both actually yield the expected dtype.
self.assertPreciseEqual(aa.get('host').dtype.type, np.float64)
self.assertIs(aa.dtype.type, np.float64)
def test():
from numba import SmartArray
a = np.linspace(0, 1, 1000, dtype=FTYPE)
a = SmartArray(a)
out = np.ones_like(a)
out = SmartArray(out)
multiply_and_scale(10., a, out)
assert np.allclose(out.get('host'), np.linspace(0, 10, 1000, dtype=FTYPE))
def main():
print("ftype=", ftype)
# hist arrays
mix = np.ones((3, 3), dtype=np.float64)
n = 1000000
inp = np.arange(3 * n, dtype=np.int32).reshape(n, 3)
out = np.ones((n), dtype=np.int32)
inp = SmartArray(inp)
out = SmartArray(out)
start_t = time.time()
sum_row(mix, 42.0 + 2j, inp.get(WHERE), out=out.get(WHERE))
end_t = time.time()
print("took %.5f" % (end_t - start_t))
start_t = time.time()
sum_row(mix, 42.0 + 2j, inp.get(WHERE), out=out.get(WHERE))
end_t = time.time()
print("took %.5f" % (end_t - start_t))
out.mark_changed(WHERE)
print(out.get("host"))
def test_ufunc(self):
a = SmartArray(np.int32([42, 8, -5]))
cfunc = jit(nopython=True)(npyufunc_usecase)
aa = cfunc(a)
self.assertIsInstance(aa, SmartArray)
self.assertPreciseEqual(aa.get('host'), np.cos(np.sin(a.get('host'))))
def test_find_index():
"""Unit tests for `find_index` function.
Correctness is defined as producing the same histogram as numpy.histogramdd
by using the output of `find_index` (ignoring underflow and overflow values).
Additionally, -1 should be returned if a value is below the range
(underflow) or is nan, and num_bins should be returned for a value above
the range (overflow).
"""
# Negative, positive, integer, non-integer, binary-unrepresentable (0.1) edges
basic_bin_edges = [-1, -0.5, -0.1, 0, 0.1, 0.5, 1, 2, 3, 4]
for basic_bin_edges in [
# Negative, positive, integer, non-integer, binary-unrepresentable (0.1) edges
[-1, -0.5, -0.1, 0, 0.1, 0.5, 1, 2, 3, 4],
# A single infinite bin: [-np.inf, np.inf]
[],
# Half-infinite bins (lower or upper edge) & [-inf, .1, +inf]
[0.1],
# Single bin with finite edges & +/-inf-edge(s)-added variants
[-0.1, 0.1],
]:
# Bin edges from above, w/ and w/o +/-inf as left and/or right edges
for le, re in [(None, None), (-np.inf, None), (None, np.inf),
(-np.inf, np.inf)]:
bin_edges = deepcopy(basic_bin_edges)
if le is not None:
bin_edges = [le] + bin_edges
if re is not None:
bin_edges = bin_edges + [re]
if len(bin_edges) < 2:
continue
logging.debug('bin_edges being tested: %s', bin_edges)
bin_edges = SmartArray(np.array(bin_edges, dtype=FTYPE))
num_bins = len(bin_edges) - 1
underflow_idx = -1
overflow_idx = num_bins
#
# Construct test values to try out
#
non_finite_vals = [-np.inf, +np.inf, np.nan]
# Values within bins (i.e., not on edges)
inbin_vals = []
for idx in range(len(bin_edges) - 1):
lower_be = bin_edges[idx]
upper_be = bin_edges[idx + 1]
if np.isfinite(lower_be):
if np.isfinite(upper_be):
inbin_val = (lower_be + upper_be) / 2
else:
inbin_val = lower_be + 10.5
else:
if np.isfinite(upper_be):
inbin_val = upper_be - 10.5
else:
inbin_val = 10.5
inbin_vals.append(inbin_val)
# Values above/below bin edges by one unit of floating point
# accuracy
eps = np.finfo(FTYPE).eps # pylint: disable=no-member
below_edges_vals = [FTYPE((1 - eps) * be) for be in bin_edges]
above_edges_vals = [FTYPE((1 + eps) * be) for be in bin_edges]
test_vals = np.concatenate([
non_finite_vals,
bin_edges,
inbin_vals,
below_edges_vals,
above_edges_vals,
])
logging.trace('test_vals = %s', test_vals)
#
# Run tests
#
for val in test_vals:
val = FTYPE(val)
np_histvals, _ = np.histogramdd([val],
np.atleast_2d(bin_edges))
nonzero_indices = np.nonzero(np_histvals)[
0] # select first & only dim
if np.isnan(val):
assert len(nonzero_indices) == 0
expected_idx = underflow_idx
elif val < bin_edges[0]:
assert len(nonzero_indices) == 0
expected_idx = underflow_idx
elif val > bin_edges[-1]:
assert len(nonzero_indices) == 0
expected_idx = overflow_idx
else:
assert len(nonzero_indices) == 1
expected_idx = nonzero_indices[0]
if TARGET == 'cpu':
found_idx = find_index(val, bin_edges)
elif TARGET == 'cuda':
found_idx_ary = SmartArray(np.zeros(1, dtype=np.int))
find_index_cuda(
SmartArray(np.array([val], dtype=FTYPE)).get('gpu'),
bin_edges.get('gpu'),
found_idx_ary,
)
found_idx = found_idx_ary.get()[0]
else:
raise NotImplementedError(f"TARGET='{TARGET}'")
if found_idx != expected_idx:
msg = 'val={}, edges={}: Expected idx={}, found idx={}'.format(
val, bin_edges, expected_idx, found_idx)
assert False, msg
logging.info('<< PASS : test_find_index >>')
def lookup(sample, flat_hist, binning):
"""The inverse of histograming: Extract the histogram values at `sample`
points.
Parameters
----------
sample : num_dims list of length-num_samples SmartArrays
Points at which to find histogram's values
flat_hist : SmartArray
Histogram values
binning : num_dims MultiDimBinning
Histogram's binning
Returns
-------
hist_vals : len-num_samples SmartArray
Notes
-----
Only handles 2d and 3d right now
"""
#print(binning)
assert binning.num_dims in [2, 3], 'can only do 2d and 3d at the moment'
bin_edges = [edges.magnitude for edges in binning.bin_edges]
# TODO: directly return smart array
if flat_hist.ndim == 1:
#print 'looking up 1D'
hist_vals = SmartArray(np.zeros_like(sample[0]))
if binning.num_dims == 2:
lookup_vectorized_2d(
sample[0].get(WHERE),
sample[1].get(WHERE),
flat_hist.get(WHERE),
bin_edges[0],
bin_edges[1],
out=hist_vals.get(WHERE),
)
elif binning.num_dims == 3:
lookup_vectorized_3d(
sample[0].get(WHERE),
sample[1].get(WHERE),
sample[2].get(WHERE),
flat_hist.get(WHERE),
bin_edges[0],
bin_edges[1],
bin_edges[2],
out=hist_vals.get(WHERE),
)
elif flat_hist.ndim == 2:
#print 'looking up ND'
hist_vals = SmartArray(
np.zeros((sample[0].size, flat_hist.shape[1]), dtype=FTYPE))
if binning.num_dims == 2:
lookup_vectorized_2d_arrays(
sample[0].get(WHERE),
sample[1].get(WHERE),
flat_hist.get(WHERE),
bin_edges[0],
bin_edges[1],
out=hist_vals.get(WHERE),
)
elif binning.num_dims == 3:
lookup_vectorized_3d_arrays(
sample[0].get(WHERE),
sample[1].get(WHERE),
sample[2].get(WHERE),
flat_hist.get(WHERE),
bin_edges[0],
bin_edges[1],
bin_edges[2],
out=hist_vals.get(WHERE),
)
else:
raise NotImplementedError()
hist_vals.mark_changed(WHERE)
return hist_vals
def execute_func(func, func_kw):
"""Run `func` with *func_kw.values() where `outputs` specify names in
`func_kw` taken to be outputs of the function; for these, mark changed.
Retrieve both input and output values as Numpy arrays on the host and
aggregate together in a single dict before returning.
Parameters
----------
func : numba CPUDispatcher or CUDADispatcher
func_kw : OrderedDict
Returns
-------
ret_dict : OrderedDict
Keys are arg names and vals are type-"correct" values; all arrays are
converted to host Numpy arrays
"""
py_func = func.py_func
func_name = ".".join([getmodule(py_func).__name__, py_func.__name__])
arg_names = list(signature(py_func).parameters.keys())
if hasattr(func, "signatures"):
arg_types = func.signatures[0]
else:
arg_types = func.compiled.argument_types
# Convert types; wrap arrays with SmartArray and place on device (if necessary)
missing = set(arg_names).difference(func_kw.keys())
excess = set(func_kw.keys()).difference(arg_names)
if missing or excess:
msgs = []
if missing:
msgs.append(f"missing kwargs {missing}")
if excess:
msgs.append(f"excess kwargs {excess}")
raise KeyError(f"{func_name}:" + ", ".join(msgs))
typed_args = OrderedDict()
for arg_name, arg_type in zip(arg_names, arg_types):
val = func_kw[arg_name]
if arg_type.name.startswith("array"):
arg_val = SmartArray(val.astype(arg_type.dtype.key))
arg_val = arg_val.get("host")
else:
arg_val = arg_type(val)
typed_args[arg_name] = arg_val
# Call the host function with typed args
try:
func(*list(typed_args.values()))
except Exception:
logging.error("Failed running `%s` with args %s", func_name, str(typed_args))
raise
# All arrays converted to Numpy host arrays
ret_dict = OrderedDict()
for key, val in typed_args.items():
if isinstance(val, SmartArray):
val.mark_changed(WHERE)
val = val.get("host")
ret_dict[key] = val
return ret_dict
def lookup_indices(sample, binning):
"""Lookup (flattened) bin index for sample points.
Parameters
----------
sample : length-M_dimensions sequence of length-N_events SmartArrays
All smart arrays must have the same lengths; corresponding elements of
the arrays are the coordinates of an event in the dimensions each array
represents.
binning : pisa.core.binning.MultiDimBinning or convertible thereto
`binning` is passed to instantiate ``MultiDimBinning``, so e.g., a
pisa.core.binning.OneDimBinning is valid to pass as `binning`
Returns
-------
indices : length-N_events SmartArray
One for each event the index of the histogram in which it falls into
Notes
-----
this method works for 1d, 2d and 3d histogram only
"""
# Convert non-MultiDimBinning objects into MultiDimBinning if possible;
# if this fails, an error will result, as it should
binning = MultiDimBinning(binning)
if len(sample) != binning.num_dims:
raise ValueError(
f"`binning` has {binning.num_dims} dimension(s), but `sample`"
f"contains {len(sample)} arrays (so represents {len(sample)}"
f" dimensions)")
lookup_funcs = {
1: lookup_indices_vectorized_1d,
2: lookup_indices_vectorized_2d,
3: lookup_indices_vectorized_3d,
}
if binning.num_dims not in lookup_funcs:
raise NotImplementedError(
"binning must have num_dims in {}; got {}".format(
sorted(lookup_funcs.keys()), binning.num_dims))
lookup_func = lookup_funcs[binning.num_dims]
lookup_func_args = (
[a.get(WHERE) for a in sample] +
[SmartArray(dim.edge_magnitudes).get(WHERE) for dim in binning])
logging.trace("lookup_func_args = {}".format(lookup_func_args))
# Create an array to store the results
indices = SmartArray(np.empty_like(sample[0], dtype=np.int64))
# Perform the lookup
lookup_func(*lookup_func_args, out=indices.get(WHERE))
indices.mark_changed(WHERE)
return indices
def test_ufunc(self):
a = SmartArray(np.int32([42, 8, -5]))
cfunc = jit(nopython=True)(npyufunc_usecase)
aa = cfunc(a)
self.assertIsInstance(aa, SmartArray)
self.assertPreciseEqual(aa.get('host'), np.cos(np.sin(a.get('host'))))
def lookup(sample, flat_hist, binning):
'''
the inverse of histograming
Paramters
--------
sample : list of SmartArrays
flat_hist : SmartArray
binning : PISA MultiDimBinning
Notes
-----
this is only a 2d method right now
'''
#print(binning)
assert binning.num_dims in [2,3], 'can only do 2d and 3d at the moment'
bin_edges = [edges.magnitude for edges in binning.bin_edges]
# todo: directly return smart array
if flat_hist.ndim == 1:
#print 'looking up 1D'
array = SmartArray(np.zeros_like(sample[0]))
if binning.num_dims == 2:
lookup_vectorized_2d(sample[0].get(WHERE), sample[1].get(WHERE), flat_hist.get(WHERE), bin_edges[0], bin_edges[1], out=array.get(WHERE))
elif binning.num_dims == 3:
lookup_vectorized_3d(sample[0].get(WHERE), sample[1].get(WHERE), sample[2].get(WHERE), flat_hist.get(WHERE), bin_edges[0], bin_edges[1], bin_edges[2], out=array.get(WHERE))
elif flat_hist.ndim == 2:
#print 'looking up ND'
array = SmartArray(np.zeros((sample[0].size, flat_hist.shape[1]), dtype=FTYPE))
if binning.num_dims == 2:
lookup_vectorized_2d_arrays(sample[0].get(WHERE), sample[1].get(WHERE), flat_hist.get(WHERE), bin_edges[0], bin_edges[1], out=array.get(WHERE))
elif binning.num_dims == 3:
lookup_vectorized_3d_arrays(sample[0].get(WHERE), sample[1].get(WHERE), sample[2].get(WHERE), flat_hist.get(WHERE), bin_edges[0], bin_edges[1], bin_edges[2], out=array.get(WHERE))
else:
raise NotImplementedError()
array.mark_changed(WHERE)
return array
def test_lookup_indices():
"""Unit tests for `lookup_indices` function"""
#
# Test a variety of points.
# Points falling exactly on the bound are included in the
#
n_evts = 100
x = np.array([-5, 0.5, 1.5, 7.0, 6.5, 8.0, 6.5], dtype=FTYPE)
y = np.array([-5, 0.5, 1.5, 1.5, 3.0, 1.5, 2.5], dtype=FTYPE)
z = np.array([-5, 0.5, 1.5, 1.5, 0.5, 6.0, 0.5], dtype=FTYPE)
w = np.ones(n_evts, dtype=FTYPE)
x = SmartArray(x)
y = SmartArray(y)
z = SmartArray(z)
w = SmartArray(w)
binning_x = OneDimBinning(name="x", num_bins=7, is_lin=True, domain=[0, 7])
binning_y = OneDimBinning(name="y", num_bins=4, is_lin=True, domain=[0, 4])
binning_z = OneDimBinning(name="z", num_bins=2, is_lin=True, domain=[0, 2])
binning_1d = binning_x
binning_2d = binning_x * binning_y
binning_3d = binning_x * binning_y * binning_z
# 1D case: check that each event falls into its predicted bin
#
# All values higher or equal to the last bin edges are assigned an index of zero
#
logging.trace("TEST 1D:")
logging.trace("Total number of bins: {}".format(7))
logging.trace("array in 1D: {}".format(x.get()))
logging.trace("Binning: {}".format(binning_1d.bin_edges[0]))
indices = lookup_indices([x], binning_1d)
logging.trace("indices of each array element: {}".format(indices.get()))
logging.trace("*********************************")
test = indices.get()
ref = np.array([-1, 0, 1, 6, 6, 7, 6])
assert np.array_equal(test, ref), "test={} != ref={}".format(test, ref)
# 2D case:
#
# The binning edges are flattened as follows:
# [(x=0, y=0), (x=0, y=1), (x=1, y=0), ...]
#
logging.trace("TEST 2D:")
logging.trace("Total number of bins: {}".format(7 * 4))
logging.trace("array in 2D: {}".format(list(zip(x.get(), y.get()))))
logging.trace("Binning: {}".format(binning_2d.bin_edges))
indices = lookup_indices([x, y], binning_2d)
logging.trace("indices of each array element: {}".format(indices.get()))
logging.trace("*********************************")
test = indices.get()
ref = np.array([-1, 0, 5, 25, 27, 28, 26])
assert np.array_equal(test, ref), "test={} != ref={}".format(test, ref)
# 3D case:
#
# the binning edges are flattened as follows:
# [(x=0, y=0, z=0), (x=0, y=0, z=1), (x=0, y=1, z=0)...]
#
logging.trace("TEST 3D:")
logging.trace("Total number of bins: {}".format(7 * 4 * 2))
logging.trace("array in 3D: {}".format(list(zip(x.get(), y.get(),
z.get()))))
logging.trace("Binning: {}".format(binning_3d.bin_edges))
indices = lookup_indices([x, y, z], binning_3d)
logging.trace("indices of each array element: {}".format(indices.get()))
logging.trace("*********************************")
test = indices.get()
ref = np.array([-1, 0, 11, 51, 54, 56, 52])
assert np.array_equal(test, ref), "test={} != ref={}".format(test, ref)
logging.info("<< PASS : test_lookup_indices >>")
