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(): 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 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 >>")