def get_numba_type( aesara_type: Type, layout: str = "A", force_scalar: bool = False, reduce_to_scalar: bool = False, ) -> numba.types.Type: r"""Create a Numba type object for a :class:`Type`. Parameters ---------- aesara_type The :class:`Type` to convert. layout The :class:`numpy.ndarray` layout to use. force_scalar Ignore dimension information and return the corresponding Numba scalar types. reduce_to_scalar Return Numba scalars for zero dimensional :class:`TensorType`\s. """ if isinstance(aesara_type, TensorType): dtype = aesara_type.numpy_dtype numba_dtype = numba.from_dtype(dtype) if force_scalar or (reduce_to_scalar and getattr(aesara_type, "ndim", None) == 0): return numba_dtype return numba.types.Array(numba_dtype, aesara_type.ndim, layout) elif isinstance(aesara_type, ScalarType): dtype = np.dtype(aesara_type.dtype) numba_dtype = numba.from_dtype(dtype) return numba_dtype else: raise NotImplementedError( f"Numba type not implemented for {aesara_type}")
def __init__(self, data_dtype: np.dtype, coords_dtype: np.dtype, ndim: int): assert isinstance(data_dtype, np.dtype) assert isinstance(coords_dtype, np.dtype) self.data_dtype = data_dtype self.coords_dtype = coords_dtype self.ndim = ndim super().__init__(name="COOType[{!r}, {!r}, {!r}]".format( numba.from_dtype(data_dtype), numba.from_dtype(coords_dtype), ndim))
def make_sundials_jac_sparse(self, format='CSR'): # type: ignore jac_sparse = self.make_jac_sparse(format=format) user_dtype = self.user_data_dtype state_dtype = self.state_dtype N_VGetArrayPointer = lib.N_VGetArrayPointer N_VGetLength = lib.N_VGetLength SUNSparseMatrix_Data = lib.SUNSparseMatrix_Data user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.cffi_support.map_type(ffi.typeof('CVLsJacFn')) args = list(func_type.args) args[4] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def jac_dense_wrapper(t, y_, fy_, out_, user_data_, tmp1_, tmp2_, tmp3_): # type: ignore n_vars = N_VGetLength(y_) y_ptr = N_VGetArrayPointer(y_) out_ptr = SUNSparseMatrix_Data(out_) fy_ptr = N_VGetArrayPointer(fy_) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] out = numba.farray(out_ptr, (n_vars, n_vars)) fy = numba.carray(fy_ptr, (n_vars, )) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return jac_dense(out, t, y, fy, user_data) return jac_dense_wrapper
def make_sundials_jac_prod(self): # type: ignore jac_prod = self.make_rhs_jac_prod() user_dtype = self.user_data_dtype state_dtype = self.state_dtype N_VGetArrayPointer = lib.N_VGetArrayPointer N_VGetLength = lib.N_VGetLength_Serial user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.core.typing.cffi_utils.map_type(ffi.typeof('CVLsJacTimesVecFn')) args = list(func_type.args) args[-2] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def jac_prod_wrapper(v_, out_, t, y_, fy_, user_data_, tmp_,): # type: ignore n_vars = N_VGetLength(v_) v_ptr = N_VGetArrayPointer(v_) out_ptr = N_VGetArrayPointer(out_) y_ptr = N_VGetArrayPointer(y_) fy_ptr = N_VGetArrayPointer(fy_) v = numba.carray(v_ptr, (n_vars,)) y = numba.carray(y_ptr, (n_vars,)).view(state_dtype)[0] out = numba.carray(out_ptr, (n_vars,)) fy = numba.carray(fy_ptr, (n_vars,)) user_data = numba.carray(user_data_, (1,), user_dtype)[0] return jac_prod(out, v, t, y, fy, user_data) return jac_prod_wrapper
def make_sundials_rhs(self) -> Any: rhs = self.make_rhs() N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial N_VGetLength_Serial = lib.N_VGetLength_Serial state_dtype = self.state_dtype user_dtype = self.user_data_dtype user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.cffi_support.map_type(ffi.typeof('CVRhsFn')) func_type = func_type.return_type(*(func_type.args[:-1] + (user_ndtype_p, ))) @numba.cfunc(func_type) def rhs_wrapper(t, y_, out_, user_data_): # type: ignore y_ptr = N_VGetArrayPointer_Serial(y_) n_vars = N_VGetLength_Serial(y_) out_ptr = N_VGetArrayPointer_Serial(out_) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] out = numba.carray(out_ptr, (n_vars, )) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return rhs(out, t, y, user_data) return rhs_wrapper
def test_bad_index_npm(self): with self.assertTypingError() as raises: arraytype1 = from_dtype(np.dtype([('x', np.int32), ('y', np.int32)])) arraytype2 = types.Array(types.int32, 2, 'C') compile_isolated(bad_index, (arraytype1, arraytype2), flags=no_pyobj_flags) self.assertIn('unsupported array index type', str(raises.exception))
def test_bad_index_npm(self): with self.assertTypingError() as raises: arraytype1 = from_dtype( np.dtype([('x', np.int32), ('y', np.int32)])) arraytype2 = types.Array(types.int32, 2, 'C') compile_isolated(bad_index, (arraytype1, arraytype2), flags=no_pyobj_flags) self.assertIn('unsupported array index type', str(raises.exception))
def test_record_misaligned(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')]) rec = from_dtype(rec_dtype) # Unlike the CUDA target, this will not generate an error @njit((rec[:],)) def foo(a): for i in range(a.size): a[i].a = a[i].b
def get_numba_type(aesara_type: Type, layout: str = "A", force_scalar: bool = False) -> numba.types.Type: """Create a Numba type object for a ``Type``.""" if isinstance(aesara_type, TensorType): dtype = aesara_type.numpy_dtype numba_dtype = numba.from_dtype(dtype) if force_scalar: return numba_dtype return numba.types.Array(numba_dtype, aesara_type.ndim, layout) elif isinstance(aesara_type, Scalar): dtype = np.dtype(aesara_type.dtype) numba_dtype = numba.from_dtype(dtype) return numba_dtype else: raise NotImplementedError( f"Numba type not implemented for {aesara_type}")
def test_record_misaligned(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')]) rec = from_dtype(rec_dtype) # Unlike the CUDA target, this will not generate an error @njit((rec[:], )) def foo(a): for i in range(a.size): a[i].a = a[i].b
def sample_gen(r,matrixlength,matrixwidth,fiberlength,fiberwidth,flag): Particle = np.dtype([ # sphere (x, y, z) coordinates ('x', 'f8'), ('y', 'f8'), ('z', 'f8'), # sphere Type ('tYpe', 'u1'), # sphere (Vx, Vy, Vz) velocity ('Vx', 'f8'), ('Vy', 'f8'), ('Vz', 'f8'), # sphere (Vx, Vy, Vz) velocity ('ax', 'f8'), ('ay', 'f8'), ('az', 'f8'), #angular velocity ('Wx', 'f8'), ('Wy', 'f8'), ('Wz', 'f8'), ], align=True) Particle_nb = from_dtype(Particle) Points=[] lengthRange=int(matrixlength/r) widthRange=int(matrixwidth/r) for i in range (0,lengthRange+1): for j in range (-widthRange,widthRange+1): for k in range (-widthRange,widthRange+1): if math.sqrt((k*r)**2+(j*r)**2)<=matrixwidth:# and k==0: if math.sqrt((k*r)**2+(j*r)**2)>fiberwidth or i*r<matrixlength-fiberlength: if flag==0: if i==0 or math.sqrt((k*r)**2+(j*r)**2)>=matrixwidth-0.5*r: Points.append([i*r,j*r,k*r,0]) else: Points.append([i*r,j*r,k*r,1]) if flag==1: if i==0: Points.append([i*r,j*r,k*r,0]) else: Points.append([i*r,j*r,k*r,1]) print(1) else: Points.append([i*r,j*r,k*r,2]) # print(2) if i==lengthRange and math.sqrt((k*r)**2+(j*r)**2)<=fiberwidth: Points.append([(i+1)*r,j*r,k*r,3]) Particles=np.empty(len(Points),dtype=Particle_nb) i=0 for x in Points: Particles[i]['x']=x[0] Particles[i]['y']=x[1] Particles[i]['z']=x[2] Particles[i]['tYpe']=x[3] Particles[i]['Vx']=0.0 Particles[i]['Vy']=0.0 Particles[i]['Vz']=0.0 Particles[i]['ax']=0.0 Particles[i]['ay']=0.0 Particles[i]['az']=0.0 Particles[i]['Wx']=0.0 Particles[i]['Wy']=0.0 Particles[i]['Wz']=0.0 i+=1 return Particles
def test_record_alignment_error(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')]) rec = from_dtype(rec_dtype) with self.assertRaises(Exception) as raises: @cuda.jit((rec[:],)) def foo(a): i = cuda.grid(1) a[i].a = a[i].b self.assertTrue('type float64 is not aligned' in str(raises.exception))
def Bnd_propgen(radii,E,Poisson,Strength): BondProp = np.dtype([ ('A', 'f8'), ('I', 'f8'), ('J', 'f8'), ('Sm', 'f8'), ('kr', 'f8'), ('Kn', 'f8'), ('Ks', 'f8'),('R', 'f8'), ('e0', 'f8'), ('ef', 'f8'),('sig', 'f8'),], align=True) BondProp_nb = from_dtype(BondProp) #Bndprop=np.zeros((35,11)) pi=3.14159 typeTable=np.zeros((11,11)) cntr=0 for i in range (0,11): for j in range (i,11): typeTable[i,j]=cntr typeTable[j,i]=cntr cntr+=1 Bndprop=np.zeros((int(cntr),11),dtype=np.float32) BProp=np.empty(int(cntr),dtype=BondProp_nb) for i in range(0,11): for j in range(0,11): Area=pi*((radii[i]/2.0+radii[j]/2.0)/2.0)**2 I=0.25*pi*((radii[i]/2.0+radii[j]/2.0)/2.0)**4 J=0.5*pi*((radii[i]/2.0+radii[j]/2.0)/2.0)**4 EEq=1/((1-Poisson[i]**2)/E[i]+(1-Poisson[j]**2)/E[j]) GEq=1/(((2*(1+Poisson[i]))*(2-Poisson[i]))/E[i]+(2*(1+Poisson[j])*(2-Poisson[j]))/E[j]) Kn=EEq/(radii[i]/2.0+radii[j]/2.0) Kt=6.0*GEq/(radii[i]/2.0+radii[j]/2.0) maxStrength=Strength[i,j,0] Softening=Strength[i,j,1] R=(radii[i]/2+radii[j]/2)/2 ind=int(typeTable[i,j]) Bndprop[ind,0]=Area Bndprop[ind,1]=I Bndprop[ind,2]=J Bndprop[ind,3]=Kn Bndprop[ind,4]=Kt Bndprop[ind,5]=maxStrength Bndprop[ind,6]=Softening Bndprop[ind,7]=R i=0 for x in Bndprop: BProp[i]['A']=x[0] BProp[i]['I']=x[1] BProp[i]['J']=x[2] BProp[i]['Kn']=float(x[3]) BProp[i]['Ks']=float(x[4]) BProp[i]['Sm']=x[5] BProp[i]['kr']=x[6] BProp[i]['R']=x[7] BProp[i]['e0']=x[8] BProp[i]['ef']=x[9] BProp[i]['sig']=x[10] i+=1 return BProp
def sample_genhccpCali(r,matrixlength,matrixwidth,fiberlength,fiberwidth,flag): Particle = np.dtype([ # sphere (x, y, z) coordinates ('x', 'f8'), ('y', 'f8'), ('z', 'f8'), # sphere Type ('tYpe', 'u1'), # sphere (Vx, Vy, Vz) velocity ('Vx', 'f8'), ('Vy', 'f8'), ('Vz', 'f8'), # sphere (Vx, Vy, Vz) velocity ('ax', 'f8'), ('ay', 'f8'), ('az', 'f8'), #angular ('Wx', 'f8'), ('Wy', 'f8'), ('Wz', 'f8'), ], align=True) Particle_nb = from_dtype(Particle) Points=[] Pnt=hccp(r,matrixlength,matrixwidth) Pnt=np.array(Pnt) # lengthRange=np.max(Pnt[:,0])-0.002 lengthRange=matrixlength maxL=np.max(Pnt[:,0]) lengthStart=np.min(Pnt[Pnt[:,0]>0,0]) # print(lengthRange,maxL,lengthStart) # lengthStart=np.min(Pnt[:,0]) # matrixlength=matrixlength+0.002 for p in Pnt: i=p[0] j=p[1] k=p[2] if math.sqrt(k**2+j**2)<=matrixwidth and i<=maxL and i>=lengthStart: # and k==0: if i==lengthStart: Points.append([i,j,k,0]) if i>lengthStart and i < maxL: Points.append([i,j,k,1]) if i ==maxL: Points.append([i,j,k,3]) Particles=np.empty(len(Points),dtype=Particle_nb) i=0 for x in Points: Particles[i]['x']=x[0] Particles[i]['y']=x[1] Particles[i]['z']=x[2] Particles[i]['tYpe']=x[3] Particles[i]['Vx']=0.0 Particles[i]['Vy']=0.0 Particles[i]['Vz']=0.0 Particles[i]['ax']=0.0 Particles[i]['ay']=0.0 Particles[i]['az']=0.0 Particles[i]['Wx']=0.0 Particles[i]['Wy']=0.0 Particles[i]['Wz']=0.0 i+=1 # print(Particles) return Particles
def test_record_alignment_error(self): rec_dtype = np.dtype([("a", "int32"), ("b", "float64")]) rec = from_dtype(rec_dtype) with self.assertRaises(Exception) as raises: @cuda.jit((rec[:], )) def foo(a): i = cuda.grid(1) a[i].a = a[i].b self.assertTrue("type float64 is not aligned" in str(raises.exception))
def test_record_alignment_error(self): rec_dtype = np.dtype([("a", "int32"), ("b", "float64")]) rec = from_dtype(rec_dtype) with self.assertRaises(Exception) as raises: @cuda.jit((rec[:],)) def foo(a): i = cuda.grid(1) a[i].a = a[i].b self.assertTrue("type float64 is not aligned" in str(raises.exception))
def test_record_alignment_error(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')]) rec = from_dtype(rec_dtype) with self.assertRaises(Exception) as raises: @cuda.jit((rec[:], )) def foo(a): i = cuda.grid(1) a[i].a = a[i].b self.assertTrue('type float64 is not aligned' in str(raises.exception))
def test_getitem_key(self): pyfunc = getitem_key cfunc = jit(nopython=True)(pyfunc) for x, i in [(np.array('123'), ()), (np.array(['123']), 0), (np.array(b'123'), ()), (np.array([b'123']), 0)]: d1 = {} d2 = Dict.empty(from_dtype(x.dtype), types.int64) pyfunc(d1, x, i) cfunc(d2, x, i) self.assertEqual(d1, d2) # check for charseq to str conversion: str(d2)
def make_sundials_adjoint_jac_prod(self): # type: ignore jac_prod = self.make_adjoint_jac_prod() user_dtype = self.user_data_dtype state_dtype = self.state_dtype N_VGetArrayPointer = lib.N_VGetArrayPointer N_VGetLength = lib.N_VGetLength_Serial user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.cffi_support.map_type( ffi.typeof('CVLsJacTimesVecFnB')) args = list(func_type.args) args[-2] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def jac_prod_wrapper( vB_, out_, t, y_, yB_, fyB_, user_data_, tmp_, ): # type: ignore n_vars = N_VGetLength(y_) n_varsB = N_VGetLength(vB_) if n_vars != n_varsB: return -1 vB_ptr = N_VGetArrayPointer(vB_) out_ptr = N_VGetArrayPointer(out_) y_ptr = N_VGetArrayPointer(y_) yB_ptr = N_VGetArrayPointer(yB_) fyB_ptr = N_VGetArrayPointer(fyB_) vB = numba.carray(vB_ptr, (n_vars, )) out = numba.carray(out_ptr, (n_vars, )) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] yB = numba.carray(yB_ptr, (n_vars, )) fyB = numba.carray(fyB_ptr, (n_vars, )) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return jac_prod(out, vB, t, y, yB, fyB, user_data) return jac_prod_wrapper
def __init__(self, source, to_range=np.int64): if isinstance(source, pd.Series): m = nb.typed.Dict.empty( key_type=nb.from_dtype(source.index.dtype), value_type=nb.from_dtype(source.dtype), ) for k, v in source.items(): m[k] = v self._in_dtype = source.index.dtype self._out_dtype = source.dtype self._mapper = m elif to_range: m = nb.typed.Dict.empty( key_type=nb.from_dtype(source.dtype), value_type=nb.from_dtype(to_range), ) for v, k in enumerate(source): m[k] = v self._in_dtype = source.dtype self._out_dtype = to_range self._mapper = m else: raise ValueError("invalid input")
def test_do_all_specific_type(modes, typ): from katana.datastructures import InsertBag @do_all_operator() def f(out, i): out[int(i)] = i input = InsertBag[typ]() for i in range(1000): input.push(i) out = np.zeros(1000, dtype=typ) do_all(input, f(out), **modes) assert np.allclose(out, np.array(range(1000))) # Check that the operator was actually compiled for the correct type assert list(f.inspect_llvm().keys())[0][1][0] == from_dtype(np.dtype(typ))
def test_record_alignment(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')], align=True) rec = from_dtype(rec_dtype) @njit((rec[:], )) def foo(a): for i in range(a.size): a[i].a = a[i].b a_recarray = np.recarray(3, dtype=rec_dtype) for i in range(a_recarray.size): a_rec = a_recarray[i] a_rec.a = 0 a_rec.b = (i + 1) * 123 foo(a_recarray) np.testing.assert_equal(a_recarray.a, a_recarray.b)
def test_record_alignment(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')], align=True) rec = from_dtype(rec_dtype) @njit((rec[:],)) def foo(a): for i in range(a.size): a[i].a = a[i].b a_recarray = np.recarray(3, dtype=rec_dtype) for i in range(a_recarray.size): a_rec = a_recarray[i] a_rec.a = 0 a_rec.b = (i + 1) * 123 foo(a_recarray) np.testing.assert_equal(a_recarray.a, a_recarray.b)
def make_sundials_sensitivity_rhs(self): # type: ignore sens_rhs = self.make_sensitivity_rhs() user_dtype = self.user_data_dtype N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial N_VGetLength_Serial = lib.N_VGetLength_Serial state_dtype = self.state_dtype user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.core.typing.cffi_utils.map_type( ffi.typeof('CVSensRhsFn')) args = list(func_type.args) args[-3] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def sens_rhs_wrapper( # type: ignore n_params, t, y_, ydot_, yS_, out_, user_data_, tmp1_, tmp2_): n_vars = N_VGetLength_Serial(y_) y_ptr = N_VGetArrayPointer_Serial(y_) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] user_data = numba.carray(user_data_, (1, ), user_dtype)[0] out_array = user_data.tmp_nparams_nstates yS_array = user_data.tmp2_nparams_nstates for i in range(n_params): yS_i_ptr = N_VGetArrayPointer_Serial(yS_[i]) yS_i = numba.carray(yS_i_ptr, (n_vars, )) yS_array[i, :] = yS_i retcode = sens_rhs(out_array, t, y, yS_array, user_data) if retcode != 0: return retcode for i in range(n_params): out_i_ptr = N_VGetArrayPointer_Serial(out_[i]) out_i = numba.carray(out_i_ptr, (n_vars, )) out_i[:] = out_array[i, :] return retcode return sens_rhs_wrapper
def Prt_propgen_Stiffness(radii,mass,vel,damp,rotation,E,G): ParticleProp = np.dtype([ ('r', 'f8'), ('mass', 'f8'), ('I', 'f8'), ('damp', 'f8'), ('Bx', 'f8'), ('By', 'f8'), ('Bz', 'f8'), ('BCondition', 'f8'), ('BWx', 'f8'), ('BWy', 'f8'), ('BWz', 'f8'), ('BWCondition', 'f8'), ('Kn', 'f8'), ('Ks', 'f8'),], align=True) ParticleProp_nb = from_dtype(ParticleProp) Pprop=np.empty(len(radii),dtype=ParticleProp_nb) Prtprop=[] pi=3.14159 for i in range(0,len(radii)): ms=(4/3)*pi*mass[i]*(radii[i]/2)**3 I=0.4*ms*(radii[i]/2)**2 P=[] P.append(ms) P.append(I) P.append(radii[i]/2) P.append(damp[i]) for l in vel[i]: P.append(l) for k in rotation[i]: P.append(k) P.append(E[i]) P.append(G[i]) Prtprop.append(P) i=0 for x in Prtprop: Pprop[i]['mass']=x[0] Pprop[i]['I']=x[1] Pprop[i]['r']=x[2] Pprop[i]['damp']=x[3] Pprop[i]['BCondition']=x[4] Pprop[i]['Bx']=x[5] Pprop[i]['By']=x[6] Pprop[i]['Bz']=x[7] Pprop[i]['BWCondition']=x[8] Pprop[i]['BWx']=x[9] Pprop[i]['BWy']=x[10] Pprop[i]['BWz']=x[11] Pprop[i]['Kn']=x[12] Pprop[i]['Ks']=x[13] i+=1 return Pprop
def test_record_alignment(self): rec_dtype = np.dtype([('a', 'int32'), ('b', 'float64')], align=True) rec = from_dtype(rec_dtype) @cuda.jit((rec[:], )) def foo(a): i = cuda.grid(1) a[i].a = a[i].b a_recarray = np.recarray(3, dtype=rec_dtype) for i in range(a_recarray.size): a_rec = a_recarray[i] a_rec.a = 0 a_rec.b = (i + 1) * 123 foo[1, 3](a_recarray) self.assertTrue(np.all(a_recarray.a == a_recarray.b))
def test_record_alignment(self): rec_dtype = np.dtype([("a", "int32"), ("b", "float64")], align=True) rec = from_dtype(rec_dtype) @cuda.jit((rec[:],)) def foo(a): i = cuda.grid(1) a[i].a = a[i].b a_recarray = np.recarray(3, dtype=rec_dtype) for i in range(a_recarray.size): a_rec = a_recarray[i] a_rec.a = 0 a_rec.b = (i + 1) * 123 foo[1, 3](a_recarray) self.assertTrue(np.all(a_recarray.a == a_recarray.b))
def get_pred_to_edges(graph): pred_to_edges = List() dtype = graph.edges.dtype numba_type = numba.from_dtype(dtype) for node in range(graph.nodes_number): temp_dict = {} for source, _, edge in graph.in_edges(node): if source in temp_dict: temp_dict[source].append(edge) else: temp_dict[source] = [edge] pred_to_edges_cur = Dict.empty(key_type=numba_type, value_type=numba_type[:]) for source in temp_dict: pred_to_edges_cur[source] = np.array(temp_dict[source], dtype=dtype) pred_to_edges.append(pred_to_edges_cur) return pred_to_edges
def make_sundials_adjoint_rhs(self): # type: ignore user_dtype = self.user_data_dtype adj = self.make_adjoint_rhs() N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial N_VGetLength_Serial = lib.N_VGetLength_Serial state_dtype = self.state_dtype user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.cffi_support.map_type(ffi.typeof('CVRhsFnB')) args = list(func_type.args) args[-1] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def adj_rhs_wrapper(t, y_, yB_, yBdot_, user_data_): # type: ignore n_vars = N_VGetLength_Serial(y_) y_ptr = N_VGetArrayPointer_Serial(y_) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] yB_ptr = N_VGetArrayPointer_Serial(yB_) yB = numba.carray(yB_ptr, (n_vars, )) yBdot_ptr = N_VGetArrayPointer_Serial(yBdot_) yBdot = numba.carray(yBdot_ptr, (n_vars, )) #print(n_vars) #print(N_VGetLength_Serial(yB_)) #print(N_VGetLength_Serial(yBdot_)) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return adj( yBdot, t, y, yB, user_data, ) return adj_rhs_wrapper
def param_lookup(param_dict, default_val, dtype): """ Generate the ufunc lookup(channel, val), which returns a numpy array of values corresponding to various channels that are looked up in the provided param_dict. If there is no key, use default_val instead. """ out_type = from_dtype(np.dtype(dtype)) #convert types to avoid any necessity of casting... param_dict = { types.uint32(k):out_type(v) for k, v in param_dict.items() } default_val = out_type(default_val) @guvectorize(["void(uint32, "+out_type.name+"[:])"], "()->()", forceobj = True) def lookup(channel, val): """Look up a value for the provided channel from a dictionary provided at compile time""" val[0] = param_dict.get(channel, default_val) return lookup
def make_sundials_adjoint_quad_rhs(self): # type: ignore user_dtype = self.user_data_dtype adjoint_quad = self.make_adjoint_quad_rhs() N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial N_VGetLength_Serial = lib.N_VGetLength_Serial state_dtype = self.state_dtype user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.core.typing.cffi_utils.map_type( ffi.typeof('CVQuadRhsFnB')) args = list(func_type.args) args[-1] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def quad_rhs_wrapper(t, y_, yB_, qBdot_, user_data_): # type: ignore n = N_VGetLength_Serial(y_) y_ptr = N_VGetArrayPointer_Serial(y_) y = numba.carray(y_ptr, (n, )).view(state_dtype)[0] yB_ptr = N_VGetArrayPointer_Serial(yB_) n = N_VGetLength_Serial(yB_) yB = numba.carray(yB_ptr, (n, )) qBdot_ptr = N_VGetArrayPointer_Serial(qBdot_) n = N_VGetLength_Serial(qBdot_) qBdot = numba.carray(qBdot_ptr, (n, )) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return adjoint_quad( qBdot, t, y, yB, user_data, ) return quad_rhs_wrapper
def test_do_all_specific_type(modes, typ): from katana.local import InsertBag @do_all_operator() def f(out, i): out[int(i)] = i data = InsertBag[typ]() for i in range(1000): data.push(i) out = np.zeros(1000, dtype=typ) do_all(data, f(out), **modes) assert np.allclose(out, np.array(range(1000))) if not numba.config.DISABLE_JIT: # Check that the operator was actually compiled for the correct type # [0][1][0] = [first overload][second argument][first possible type] # I'm not sure why the last indexing is used. It always seems to be a 1-tuple, but numba makes it. *shrug* assert list(f.inspect_llvm().keys())[0][1][0] == from_dtype( np.dtype(typ))
def inner_dtype_of_form(form): if form is None: return None elif isinstance(form, ak._v2.forms.NumpyForm): return numba.from_dtype( ak._v2.types.numpytype.primitive_to_dtype(form.primitive)) elif isinstance(form, ak._v2.forms.EmptyForm): return numba.types.float64 elif isinstance( form, ( ak._v2.forms.RegularForm, ak._v2.forms.ListForm, ak._v2.forms.ListOffsetForm, ak._v2.forms.IndexedForm, ), ): return inner_dtype_of_form(form.content) elif isinstance( form, ( ak._v2.forms.RecordForm, ak._v2.forms.IndexedOptionForm, ak._v2.forms.ByteMaskedForm, ak._v2.forms.BitMaskedForm, ak._v2.forms.UnmaskedForm, ), ): return None elif isinstance(form, ak._v2.forms.UnionForm): context = numba.core.typing.Context() return context.unify_types( *[inner_dtype_of_form(x) for x in form.contents]) else: raise AssertionError(f"unrecognized Form type: {type(form)}")
def make_sundials_adjoint_jac_dense(self): # type: ignore jac_dense = self.make_adjoint_jac_dense() user_dtype = self.user_data_dtype state_dtype = self.state_dtype N_VGetArrayPointer_Serial = lib.N_VGetArrayPointer_Serial N_VGetLength_Serial = lib.N_VGetLength_Serial SUNDenseMatrix_Data = lib.SUNDenseMatrix_Data user_ndtype = numba.from_dtype(user_dtype) user_ndtype_p = numba.types.CPointer(user_ndtype) func_type = numba.core.typing.cffi_utils.map_type( ffi.typeof('CVLsJacFnB')) args = list(func_type.args) args[5] = user_ndtype_p func_type = func_type.return_type(*args) @numba.cfunc(func_type) def jac_dense_wrapper(t, y_, yB_, fyB_, out_, user_data_, tmp1_, tmp2_, tmp3_): # type: ignore n_vars = N_VGetLength_Serial(y_) n_lamda = N_VGetLength_Serial(yB_) y_ptr = N_VGetArrayPointer_Serial(y_) yB_ptr = N_VGetArrayPointer_Serial(yB_) fyB_ptr = N_VGetArrayPointer_Serial(fyB_) out_ptr = SUNDenseMatrix_Data(out_) y = numba.carray(y_ptr, (n_vars, )).view(state_dtype)[0] yB = numba.carray(yB_ptr, (n_lamda, )) fyB = numba.carray(fyB_ptr, (n_lamda, )) out = numba.farray(out_ptr, (n_lamda, n_lamda)) user_data = numba.carray(user_data_, (1, ), user_dtype)[0] return jac_dense(out, t, y, yB, fyB, user_data) return jac_dense_wrapper
def test_record_access(self): backyard_type = [('statue', np.float64), ('newspaper', np.float64, (6, ))] goose_type = [('garden', np.float64, (12, )), ('town', np.float64, (42, )), ('backyard', backyard_type)] goose_np_type = np.dtype(goose_type, align=True) goose_nb_type = nb.from_dtype(goose_np_type) @cuda.jit def simple_kernel(f): f.garden[0] = 45.0 f.backyard.newspaper[3] = 2.0 f.backyard.newspaper[3] = f.backyard.newspaper[3] + 3.0 item = np.recarray(1, dtype=goose_np_type) simple_kernel[1, 1](item[0]) np.testing.assert_equal(item[0]['garden'][0], 45) np.testing.assert_equal(item[0]['backyard']['newspaper'][3], 5)
def sendBNDS(number): Bond = np.dtype( [ # Bond (d0, d1, d2) values # Bond Damage ('D', 'f8'), ('L', 'f8'), ('X0', 'f8'), # Particles in bond ('P1', 'i4'), ('P2', 'i4'), # Bond max strain #Normal Vector ('Nx', 'f8'), ('Ny', 'f8'), ('Nz', 'f8'), #Bond Type ('tYpe', 'u1'), #Normal strain ('Unx', 'f8'), ('Uny', 'f8'), ('Unz', 'f8'), ('Usx', 'f8'), ('Usy', 'f8'), ('Usz', 'f8'), ('Wnx', 'f8'), ('Wny', 'f8'), ('Wnz', 'f8'), ('Wsx', 'f8'), ('Wsy', 'f8'), ('Wsz', 'f8'), ('MaxK', 'f8'), ], align=True) Bond_nb = from_dtype(Bond) Bnds = np.empty(int(number), dtype=Bond_nb) d_Bnd = cuda.device_array(Bnds.shape[0], dtype=Bond_nb) cuda.to_device(Bnds, to=d_Bnd) return d_Bnd
import numpy as np import numba as nb bags = ["bag1", "bag2", "bag3"] bag_of_word_dtype = np.dtype([(bags[0], np.int32), (bags[1], np.int32), (bags[2], np.int32)]) numba_bag_of_word_dtype = nb.from_dtype(bag_of_word_dtype)
event.all_hits = np.concatenate(hits_per_pulse) if not self.always_find_single_hit: # Remove hits with 0 or negative area (very rare, but possible due to rigid integration bound) # In always-find-single-hit mode (for PMT calibrations) this is undesirable event.all_hits = event.all_hits[event.all_hits['area'] > 0] self.log.debug("Found %d hits in %d pulses" % (len(event.all_hits), len(event.pulses))) else: self.log.warning("Event has no pulses??!") return event @numba.jit(numba.void(numba.float64[:], numba.int64[:, :], numba.from_dtype(datastructure.Hit.get_dtype())[:], numba.float64, numba.int64, numba.float64, numba.int64, numba.int64, numba.int64, numba.float64, numba.int64[:, :]), nopython=True) def build_hits(w, hit_bounds, hits_buffer, adc_to_pe, channel, noise_sigma_pe, dt, start, pulse_i, saturation_threshold, central_bounds): """Populates hits_buffer with properties from hits indicated by hit_bounds. hit_bounds should be a numpy array of (left, right) bounds (inclusive) in w Returns nothing. """ for hit_i in range(len(hit_bounds)): amplitude = -999.9 argmax = -1 area = 0.0
# * http://xoroshiro.di.unimi.it/splitmix64.c # # Splitmix64 is used to generate the initial state of the xoroshiro128+ # generator to ensure that small seeds don't result in predictable output. # **WARNING**: There is a lot of verbose casting in this file to ensure that # NumPy casting conventions (which cast uint64 [op] int32 to float64) don't # turn integers into floats when using these functions in the CUDA simulator. # # There are also no function type signatures to ensure that compilation is # deferred so that import is quick, and Sphinx autodoc works. We are also # using the CPU @jit decorator everywhere to create functions that work as # both CPU and CUDA device functions. xoroshiro128p_dtype = np.dtype([('s0', np.uint64), ('s1', np.uint64)], align=True) xoroshiro128p_type = from_dtype(xoroshiro128p_dtype) @jit def init_xoroshiro128p_state(states, index, seed): '''Use SplitMix64 to generate an xoroshiro128p state from 64-bit seed. This ensures that manually set small seeds don't result in a predictable initial sequence from the random number generator. :type states: 1D array, dtype=xoroshiro128p_dtype :param states: array of RNG states :type index: uint64 :param index: offset in states to update :type seed: int64 :param seed: seed value to use when initializing state
import numba import numpy as np from pax import units, exceptions from pax.datastructure import Hit @numba.jit(numba.int64[:](numba.from_dtype(Hit.get_dtype())[:]), nopython=True) def gaps_between_hits(hits): """Return array of gaps between hits: a hit's 'gap' is the # of samples before that hit free of other hits. The gap of the first hit is 0 by definition. Hits should already be sorted by index of maximum; we'll check this and throw an error if not. """ n_hits = len(hits) gaps = np.zeros(n_hits, dtype=np.int64) if n_hits == 0: return gaps # Keep a running right boundary boundary = hits[0].index_of_maximum for i, hit in enumerate(hits[1:]): gaps[i + 1] = max(0, hit.index_of_maximum - boundary - 1) if hit.index_of_maximum < boundary: raise ValueError("Hits should be sorted by index_of_maximum") boundary = max(hit.index_of_maximum, boundary) return gaps def count_hits_per_channel(peak, config, weights=None): return np.bincount(peak.hits['channel'].astype(np.int16), minlength=config['n_channels'], weights=weights)