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_binned_data(self, key, data, flat=True):
''' add data to binned_data
key : string
data : PISA Map or (array, binning)-tuple
flat : bool
is the data already flattened (i.e. the binning dimesnions unrolled)
'''
# TODO: logic to not copy back and forth
if isinstance(data, Map):
flat_array = data.hist.ravel()
self.binned_data[key] = (SmartArray(flat_array), data.binning)
elif isinstance(data, Sequence) and len(data) == 2:
binning, array = data
assert isinstance(binning, MultiDimBinning)
if isinstance(array, SmartArray):
array = array.get('host')
if flat:
flat_array = array
else:
# first dimesnions must match
assert array.shape[:binning.num_dims] == binning.shape
#flat_shape = [-1] + [d for d in array.shape[binning.num_dims-1:-1]]
flat_shape = [binning.size, -1]
#print(flat_shape)
flat_array = array.reshape(flat_shape)
if not isinstance(flat_array, SmartArray):
flat_array = SmartArray(flat_array.astype(FTYPE))
self.binned_data[key] = (binning, flat_array)
else:
raise TypeError('unknown dataformat')
def get_hist_gpu(sample, weights, binning, apply_weights=True):
# ToDo:
# * make for d > 3
if binning.num_dims in [2, 3]:
bin_edges = [edges.magnitude for edges in binning.bin_edges]
if len(weights.shape) > 1:
# so we have arrays
flat_hist = SmartArray(np.zeros((binning.size, weights.shape[1]), dtype=FTYPE))
arrays = True
else:
flat_hist = SmartArray(np.zeros(binning.size, dtype=FTYPE))
arrays = False
size = weights.shape[0]
d_bin_edges_x = cuda.to_device(bin_edges[0])
d_bin_edges_y = cuda.to_device(bin_edges[1])
if binning.num_dims == 2:
if arrays:
histogram_2d_kernel_arrays[(size+511)//512, 512](sample[0].get('gpu'),
sample[1].get('gpu'),
flat_hist,
d_bin_edges_x,
d_bin_edges_y,
weights.get('gpu'),
apply_weights)
else:
histogram_2d_kernel[(size+511)//512, 512](sample[0].get('gpu'),
sample[1].get('gpu'),
flat_hist,
d_bin_edges_x,
d_bin_edges_y,
weights.get('gpu'),
apply_weights)
elif binning.num_dims == 3:
d_bin_edges_z = cuda.to_device(bin_edges[2])
if arrays:
histogram_3d_kernel_arrays[(size+511)//512, 512](sample[0].get('gpu'),
sample[1].get('gpu'),
sample[2].get('gpu'),
flat_hist,
d_bin_edges_x,
d_bin_edges_y,
d_bin_edges_z,
weights.get('gpu'),
apply_weights)
else:
histogram_3d_kernel[(size+511)//512, 512](sample[0].get('gpu'),
sample[1].get('gpu'),
sample[2].get('gpu'),
flat_hist,
d_bin_edges_x,
d_bin_edges_y,
d_bin_edges_z,
weights.get('gpu'),
apply_weights)
return flat_hist
else:
raise NotImplementedError('Other dimesnions that 2 and 3 on the GPU not supported right now')
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_interface(self):
# show that the SmartArray type supports all ndarray operations transparently
a = np.arange(16).reshape(4, 4)
ap = SmartArray(a)
ap[:, :] = 1
ref = SmartArray(np.ones(dtype=ap.dtype, shape=(4, 4)))
eq = ap == ref
self.assertIsInstance(eq, SmartArray)
self.assertTrue(eq.all())
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 test_copy_matrix():
"""Unit tests of `copy_matrix` and `copy_matrix_guf`"""
A = SmartArray(np.ones((3, 3), dtype=FTYPE))
B = SmartArray(np.zeros((3, 3), dtype=FTYPE))
copy_matrix_guf(A.get(WHERE), B.get(WHERE))
B.mark_changed(WHERE)
test = B.get()
ref = A.get()
assert np.array_equal(test, ref), f"test:\n{test}\n!= ref:\n{ref}"
logging.info("<< PASS : test_copy_matrix >>")
def test_getitem(self):
a = SmartArray(np.int32([42, 8, -5]))
cfunc = jit(nopython=True)(getitem_usecase)
self.assertPreciseEqual(cfunc(a, 1), 8)
aa = cfunc(a, slice(1, None))
self.assertIsInstance(aa, SmartArray)
self.assertEqual(list(aa), [8, -5])
def binned_to_binned(self, key, new_binning):
'''
resample a binned key into a different binning
Parameters
----------
key : str
new_binning : MultiDimBinning
the new binning
'''
logging.debug('Resampling %s' % (key))
old_binning, hist = self.binned_data[key]
sample = [
self.get_binned_data(name, old_binning)
for name in old_binning.names
]
new_sample = [
SmartArray(self.unroll_binning(name, new_binning))
for name in new_binning.names
]
hist = resample(hist, sample, old_binning, new_sample, new_binning)
self.add_binned_data(key, (new_binning, hist))
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 get_binned_data(self, key, out_binning=None):
"""Get data array from binned data:
if the key is a binning dimensions, then unroll the binning
otherwise return the corresponding flattened array
"""
if out_binning is not None:
# check if key is binning dimension
if key in out_binning.names:
return self.unroll_binning(key, out_binning)
binning, data = self.binned_data[key]
if out_binning is not None:
if not binning == out_binning:
logging.warning('Automatically re-beinning data %s'%key)
sample = [SmartArray(self.unroll_binning(name, binning)) for name in binning.names]
new_sample = [SmartArray(self.unroll_binning(name, out_binning)) for name in out_binning.names]
return resample(data, sample, binning, new_sample, out_binning)
return data
def test_histogram():
"""Unit tests for `histogram` function.
Correctness is defined as matching the histogram produced by
numpy.histogramdd.
"""
all_num_bins = [2, 3, 4]
n_evts = 10000
rand = np.random.RandomState(seed=0)
weights = SmartArray(rand.rand(n_evts).astype(FTYPE))
binning = []
sample = []
for num_dims, num_bins in enumerate(all_num_bins, start=1):
binning.append(
OneDimBinning(
name=f'dim{num_dims - 1}',
num_bins=num_bins,
is_lin=True,
domain=[0, num_bins],
))
sample.append(SmartArray(rand.rand(n_evts).astype(FTYPE) * num_bins))
if TARGET == "cuda" and num_dims == 1:
continue
bin_edges = [b.edge_magnitudes for b in binning]
test = histogram(sample, weights, binning, averaged=False).get()
ref, _ = np.histogramdd(sample=sample, bins=bin_edges, weights=weights)
ref = ref.astype(FTYPE).ravel()
assert recursiveEquality(test, ref), f'\ntest:\n{test}\n\nref:\n{ref}'
test_avg = histogram(sample, weights, binning, averaged=True).get()
ref_counts, _ = np.histogramdd(sample=sample,
bins=bin_edges,
weights=None)
ref_counts = ref_counts.astype(FTYPE).ravel()
ref_avg = (ref / ref_counts).astype(FTYPE)
assert recursiveEquality(test_avg, ref_avg), \
f'\ntest_avg:\n{test_avg}\n\nref_avg:\n{ref_avg}'
logging.info('<< PASS : test_histogram >>')
def test_smartarray(self):
# tests deprecation of SmartArray
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("ignore", category=NumbaWarning)
warnings.simplefilter("always", category=NumbaDeprecationWarning)
SmartArray(np.zeros(1))
self.assertEqual(len(w), 1)
self.assertEqual(w[0].category, NumbaDeprecationWarning)
warn_msg = str(w[0].message)
msg = "SmartArray is deprecated"
self.assertIn(msg, warn_msg)
self.assertIn("http://numba.pydata.org", warn_msg)
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.host() == a.host()).all())
def test_histogram():
n_evts = 100
x = np.arange(n_evts, dtype=FTYPE)
y = np.arange(n_evts, dtype=FTYPE)
w = np.ones(n_evts, dtype=FTYPE)
#w *= np.random.rand(n_evts)
x = SmartArray(x)
y = SmartArray(y)
w = SmartArray(w)
binning_x = OneDimBinning(name='x', num_bins=10, is_lin=True, domain=[0, 100])
binning_y = OneDimBinning(name='y', num_bins=10, is_lin=True, domain=[0, 100])
binning = MultiDimBinning([binning_x, binning_y])
sample = [x, y]
weights = w
averaged = True
histo = histogram(sample, weights, binning, averaged)
assert np.array_equal(histo.reshape(10, 10), np.zeros(shape=(10, 10)))
def get_hist_np(sample, weights, binning, apply_weights=True):
'''helper function for numoy historams'''
bin_edges = [edges.magnitude for edges in binning.bin_edges]
sample = [s.get('host') for s in sample]
weights = weights.get('host')
if weights.ndim == 2:
# that means it's 1-dim data instead of scalars
hists = []
for i in range(weights.shape[1]):
w = weights[:, i] if apply_weights else None
hist, _ = np.histogramdd(sample=sample, weights=w, bins=bin_edges)
hists.append(hist.ravel())
flat_hist = np.stack(hists, axis=1)
else:
w = weights if apply_weights else None
hist, _ = np.histogramdd(sample=sample, weights=w, bins=bin_edges)
flat_hist = hist.ravel()
return SmartArray(flat_hist.astype(FTYPE))
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 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 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. + 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. + 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 elastic_gpu_optimized_1D(p, q):
n = p.size
# Load data
min_energy_values = np.full((n, n), np.inf, dtype=np.float32)
path_nodes = np.zeros((n, n, 2), dtype=np.int16)
min_energy_values[1][1] = integrate(p, q, 0, 1, 0, 1, 1 / (n - 1), 0, 0)
p = SmartArray(p)
q = SmartArray(q)
min_energy_values = SmartArray(min_energy_values)
path_nodes = SmartArray(path_nodes)
m = SmartArray(np.full(1, 0, dtype=np.float32))
# t_global_mem = cuda.const.array_like(t)
# p_global_mem = cuda.const.array_like(p)
# q_global_mem = cuda.const.array_like(q)
# flag1_global_mem = cuda.to_device(flag1)
# flag2_global_mem = cuda.to_device(flag2)
# Set the number of threads in a block
threadsperblock = (TPB, TPB)
# Calculate the number of thread blocks in the grid
blockspergrid = ((p.size + (threadsperblock[0] - 1)) // threadsperblock[0],
(p.size + (threadsperblock[1] - 1)) // threadsperblock[1])
# blockspergrid = (1, 1)
# relax_edges
prev = np.inf
for i in range(2 * n - 3):
relax[blockspergrid, threadsperblock](p, q, min_energy_values, m)
if prev != np.inf and prev == m[0]:
break
prev = m[0]
# print(min_energy_values.__array__())
relax_final[blockspergrid, threadsperblock](p, q, min_energy_values,
path_nodes)
# Print the result
gamma_interval = 1 / (n - 1)
i = n - 1
j = n - 1
min_energy_values[i][j] = integrate(p, q, 0, i, 0, j, gamma_interval, 0, 0)
k = i - width
if k <= 0:
k = 0
minimum = min_energy_values[i][j]
while k < i:
l = j - width
if l <= 0:
l = 0
while l < j:
e = min_energy_values[k, l] + integrate(p, q, k, i, l, j,
gamma_interval, 0, 0)
if e < minimum:
minimum = e
path_nodes[i][j][0] = k
path_nodes[i][j][1] = l
l = l + 1
k = k + 1
min_energy_values[i][j] = minimum
path = np.zeros(n, dtype=np.float64)
# !! Interpolate
path_indices = np.zeros((n, 2), dtype=np.int16)
path_indices[0][0] = n - 1
path_indices[0][1] = n - 1
i = 0
while path_indices[i][
0] != 0 or path_indices[i][1] != 0 and i + 1 < path.size:
result = path_nodes[path_indices[i][0]][path_indices[i][1]]
path_indices[i + 1][0] = result[0]
path_indices[i + 1][1] = result[1]
i = i + 1
i = 0
previous = 1
previousIndex_domain = n - 1
previousIndex_gamma = n - 1
path[path_indices[0][0]] = gamma_interval * path_indices[0][1]
while i < path_indices.size // 2 and previousIndex_domain != 0:
path[path_indices[i][0]] = gamma_interval * path_indices[i][1]
if previousIndex_domain - path_indices[i][0] > 1:
j = 0
val = (gamma_interval * (previousIndex_gamma - path_indices[i][1])) / \
(gamma_interval * previousIndex_domain - gamma_interval * path_indices[i][0])
while j < previousIndex_domain - path_indices[i][0]:
path[previousIndex_domain -
j] = previous - (gamma_interval * previousIndex_domain -
gamma_interval *
(previousIndex_domain - j)) * val
j = j + 1
previousIndex_domain = path_indices[i][0]
previousIndex_gamma = path_indices[i][1]
previous = gamma_interval * path_indices[i][1]
i = i + 1
return path
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'))))
for j in range(template_map.shape[1]):
for k in range(template_map.shape[2]):
template_map[i, j, k]['index'] = index_table[i, j, k]
template_map[i, j, k]['weight'] = table_3d[i, j, k]
return template_map
fname = os.path.join(args.dir, 'cl%s/ckv_table.npy' % (args.cluster_idx))
outname = os.path.join(args.dir,
'cl%s/ckv_template_map.npy' % (args.cluster_idx))
chiname = os.path.join(args.dir,
'cl%s/template_chi2s.npy' % (args.cluster_idx))
templates = np.load(os.path.join(args.dir, 'ckv_dir_templates.npy'))
templates = templates.astype(np.float32)
templates = SmartArray(templates)
if os.path.isfile(outname):
if args.overwrite:
print('overwritting existing file')
else:
print('file exists, abort')
sys.exit()
print('table cluster %s' % (args.cluster_idx))
table_5d = np.load(fname)
print('table loaded')
table_3d = np.sum(table_5d, axis=(3, 4))
# normalize the tables such that all directionality maps sum to 1
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_identity(self):
# make sure unboxing and boxing works.
a = SmartArray(np.arange(3))
cfunc = jit(nopython=True)(identity)
self.assertIs(cfunc(a), a)
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_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 >>")
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 unroll_binning(key, binning):
grid = binning.meshgrid(entity='weighted_centers', attach_units=False)
return SmartArray(grid[binning.index(key)].ravel())
def test_shape(self):
a = SmartArray(np.arange(3))
cfunc = jit(nopython=True)(shape_usecase)
self.assertPreciseEqual(cfunc(a), (3, ))
def test_len(self):
a = SmartArray(np.arange(3))
cfunc = jit(nopython=True)(len_usecase)
self.assertPreciseEqual(cfunc(a), 3)
