jitclass not working
@@ -7,13 +7,18 @@ import pylab as pl
import sciris as sc
import numba as nb
+spec = [
+ ('n', nb.int32),
+ ('npts', nb.int32),
+
+]
+
+@nb.jitclass(spec)
class NumbaTests(sc.prettyobj):
- def __init__(self, n=400e3, npts=1000, keys=None):
- if keys is None:
- keys = ['a','b']
+ def __init__(self, n=400e3, npts=1000):
self.n = int(n)
self.npts = npts
- self.keys = keys
+ self.keys = ['a','b']
self.states = {key:np.random.random(self.n) for key in self.keys}
self.results = np.zeros(self.npts)
@@ -25,13 +30,17 @@ class NumbaTests(sc.prettyobj):
def b(self):
return self.states['b']
- def next(self, t):
+ def next_mult(self, t):
self.results[t] += (self.a * self.b).sum()
return
+ def next_cond(self, t):
+ self.results[t] += np.logical_and(self.a>0.5, self.b>0.5).sum()
+ return
+
def run(self):
for t in range(self.npts):
- self.next(t)
+ self.next_cond(t)
return
