def test_default_args_starargs_and_keyonly(self):
spec = [('x', int32), ('y', int32), ('z', int32),
('args', types.UniTuple(int32, 2)), ('a', int32)]
with self.assertRaises(errors.UnsupportedError) as raises:
jitclass(spec)(TestClass2)
msg = "VAR_POSITIONAL argument type unsupported"
self.assertIn(msg, str(raises.exception))
def test_default_args_starargs_and_keyonly(self):
spec = [("x", int32), ("y", int32), ("z", int32),
("args", types.UniTuple(int32, 2)), ("a", int32)]
with self.assertRaises(errors.UnsupportedError) as raises:
jitclass(TestClass2, spec)
msg = "VAR_POSITIONAL argument type unsupported"
self.assertIn(msg, str(raises.exception))
def test_globals(self):
class Mine(object):
constant = 123
def __init__(self):
pass
with self.assertRaises(TypeError) as raises:
jitclass(())(Mine)
self.assertEqual(str(raises.exception),
"class members are not yet supported: constant")
def test_default_args_keyonly(self):
spec = [("x", int32), ("y", int32), ("z", int32), ("a", int32)]
TestClass = jitclass(TestClass1, spec)
tc = TestClass(2, 3)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 3)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 5)
tc = TestClass(y=4, x=2, a=42, z=100)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 100)
self.assertEqual(tc.a, 42)
tc = TestClass(y=4, x=2, a=42)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 42)
tc = TestClass(y=4, x=2)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 5)
def make_jitclass_element(self):
spec = [
('many', types.float64[:]),
('scalar', types.float64),
]
JCItem = jitclass(spec)(Item)
return JCItem
def KalmanODE(wgt_meas, tmin, tmax, n_eval, ode_fun,
mu_state, wgt_state, var_state, z_state=None):
"Create a jitted KalmanODE class."
spec = [
('_wgt_meas', float64[::1, :]),
('n_meas', intc),
('n_state', intc),
('tmin', float64),
('tmax', float64),
('n_eval', intc),
('n_steps', intc),
('_mu_state', float64[::1]),
('_wgt_state', float64[::1, :]),
('_var_state', float64[::1, :]),
('_z_state', float64[::1, :]),
('ode_fun', typeof(ode_fun)),
('mu_state_pred', float64[::1, :]),
('var_state_pred', float64[::1, :, :]),
('mu_state_filt', float64[::1, :]),
('var_state_filt', float64[::1, :, :]),
('var_state_smooth', float64[::1, :]),
('x_meas', float64[::1]),
('mu_meas', float64[::1]),
('var_meas', float64[::1, :]),
('ktv', kalman_type),
('tx_state', float64[::1]),
('twgt_meas', float64[::1, :]),
('tchol_state', float64[::1, :])
]
jcl = jitclass(spec)
kalman_cl = jcl(_KalmanODE)
return kalman_cl(wgt_meas, tmin, tmax, n_eval, ode_fun,
mu_state, wgt_state, var_state, z_state)
def test_default_args_keyonly(self):
spec = [('x', int32), ('y', int32), ('z', int32), ('a', int32)]
TestClass = jitclass(spec)(TestClass1)
tc = TestClass(2, 3)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 3)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 5)
tc = TestClass(y=4, x=2, a=42, z=100)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 100)
self.assertEqual(tc.a, 42)
tc = TestClass(y=4, x=2, a=42)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 42)
tc = TestClass(y=4, x=2)
self.assertEqual(tc.x, 2)
self.assertEqual(tc.y, 4)
self.assertEqual(tc.z, 1)
self.assertEqual(tc.a, 5)
def test_spec_errors(self):
spec1 = [('x', int), ('y', float32[:])]
spec2 = [(1, int32), ('y', float32[:])]
class Test(object):
def __init__(self):
pass
with self.assertRaises(TypeError) as raises:
jitclass(spec1)(Test)
self.assertIn("spec values should be Numba type instances",
str(raises.exception))
with self.assertRaises(TypeError) as raises:
jitclass(spec2)(Test)
self.assertEqual(str(raises.exception),
"spec keys should be strings, got 1")
def test_user_deleter_error(self):
class Foo(object):
def __init__(self):
pass
@property
def value(self):
return 1
@value.deleter
def value(self):
pass
with self.assertRaises(TypeError) as raises:
jitclass([])(Foo)
self.assertEqual(str(raises.exception),
"deleter is not supported: value")
def foo(fun):
"""
Inputs are all those of `_foo`.
From these it creates a spec of all members for `numba.jitclass`.
It then creates the jitted class and returns an instance of it initialized with the given inputs.
"""
spec = [("fun", typeof(fun))]
jcl = jitclass(spec)
foo_cl = jcl(_foo)
return foo_cl(fun)
def test_jitclass(self):
with override_config('DISABLE_JIT', True):
with forbid_codegen():
SimpleJITClass = jitclass(simple_class_spec)(SimpleClass)
obj = SimpleJITClass()
self.assertPreciseEqual(obj.h, 5)
cfunc = jit(nopython=True)(simple_class_user)
self.assertPreciseEqual(cfunc(obj), 5)
def test_name_shadowing_error(self):
class Foo(object):
def __init__(self):
pass
@property
def my_property(self):
pass
def my_method(self):
pass
with self.assertRaises(NameError) as raises:
jitclass([('my_property', int32)])(Foo)
self.assertEqual(str(raises.exception), 'name shadowing: my_property')
with self.assertRaises(NameError) as raises:
jitclass([('my_method', int32)])(Foo)
self.assertEqual(str(raises.exception), 'name shadowing: my_method')
def test_jitclass_function_usecases(self):
spec = OrderedDict(x=float64)
class AnnotatedTest:
x: float
def __init__(self):
self.x = 0
JitTest1 = jitclass(AnnotatedTest)
self.assertIsInstance(JitTest1, JitClassType)
self.assertDictEqual(JitTest1.class_type.struct, spec)
class UnannotatedTest:
def __init__(self):
self.x = 0
JitTest2 = jitclass(UnannotatedTest, spec)
self.assertIsInstance(JitTest2, JitClassType)
self.assertDictEqual(JitTest2.class_type.struct, spec)
def test_annotations(self):
"""
Methods with annotations should compile fine (issue #1911).
"""
from .annotation_usecases import AnnotatedClass
spec = {'x': int32}
cls = jitclass(spec)(AnnotatedClass)
obj = cls(5)
self.assertEqual(obj.x, 5)
self.assertEqual(obj.add(2), 7)
def jit_compile_functions():
# utils_for_numba.USE_JIT = True
utils_for_numba.set_list_type_for_jit()
set_list_type_for_jit()
(
d_int,
d_scalar,
d_iarray,
d_array,
d_matrix,
d_nested,
d_string,
) = initialize_dict_numbafied()
specs_specie = [
("logc", numba.float64),
("logg", numba.float64),
("z", numba.int64),
("phase", numba.int64), # Improve, but See Phase
("name", numba.types.string),
# ('I_factor', numba.float64),
# ('dh_a', numba.float64),
# ('dh_b', numba.float64),
("cond_molar", numba.float64),
("p_int", numba.typeof(d_int)), # scalar parameters
("p_scalar", numba.typeof(d_scalar)), # scalar parameters
("p_iarray", numba.typeof(d_iarray)), # int array parameters
("p_array", numba.typeof(d_array)), # array parameters
("p_matrix", numba.typeof(d_matrix)), # matrix parameters
("d", numba.typeof(d_nested)), # nested dict float parameters
]
# global Specie
core.Specie = jitclass(specs_specie)(core.Specie)
l_species, l_string = create_typed_lists()
l_d_string_float = create_numba_list_of_dict()
type_list_specie = numba.typeof(l_species)
spec_result = [
("c_molal", numba.float64[:]),
("gamma", numba.float64[:]),
("pH", numba.float64),
("I", numba.float64),
(
"DIC",
numba.float64,
), # Only the dissolved Carbon (solid not counted)
("sc", numba.float64),
# ('SI', numba.typeof(SI_Dict)),
# ('SI', numba.float64[:]),
("IAP", numba.float64[:]),
("solid_names", numba.types.List(numba.types.unicode_type)),
("precipitation_molalities", numba.float64[:]),
# ('specie_names', numba.types.List(numba.types.unicode_type)),
("specie_names", numba.typeof(l_string)),
("saturation_index", numba.typeof(d_scalar)),
("preciptation_conc", numba.typeof(d_scalar)),
("ionic_activity_prod", numba.typeof(d_scalar)),
("log_K_solubility", numba.typeof(d_scalar)),
("idx", numba.typeof(d_int)),
("reactions", numba.typeof(l_d_string_float)),
("index_solubility_calculation", numba.int64),
("calculated_solubility", numba.typeof(d_scalar)),
("concentrations", numba.typeof(d_scalar)),
("x", numba.float64[:]), # Numerical solution
("successfull", numba.boolean),
("alkalinity", numba.float64),
]
# global core.SolutionResult
core.SolutionResult = jitclass(spec_result)(core.SolutionResult)
specs_idx = [
("idx", numba.typeof(d_int)),
("s", numba.typeof(d_scalar)),
("a", numba.typeof(d_array)),
("m", numba.typeof(d_matrix)),
]
# global core.IndexController
core.IndexController = jitclass(specs_idx)(core.IndexController)
specs_reacs = [
# ('idx_species__', numba.int64),
# ('stoic_coefs', numba.float64),
# ('constant_T_coefs', numba.float64),
("idx_species", numba.int64[:]),
("stoic_coefs", numba.float64[:]),
# ('stoic_coefs', numba.types.List(numba.int64)),
("idx_reaction_db", numba.int64),
("constant_T_coefs", numba.float64[:]),
("log_K25", numba.float64),
("type", numba.types.unicode_type),
("delta_h", numba.float64),
# ('species_tags', numba.types.List(numba.types.unicode_type)),
("species_tags", numba.typeof(l_string)),
]
# global EqReaction
core.EqReaction = jitclass(specs_reacs)(core.EqReaction)
# global DUMMY_EqREACTION
core.DUMMY_EqREACTION = core.EqReaction(
np.array([-1, -1], dtype=np.int64),
# np.array([-1]),
np.array([np.nan], dtype=np.float64),
-1.0,
np.array([np.nan], dtype=np.float64),
"dummy",
utils_for_numba.create_nb_List([""]),
np.nan,
) # -1 IS DUMMY -> Numba issues
# if os.getenv('NUMBA_DISABLE_JIT') != "1":
l_reactions = numba.typed.List()
l_reactions.append(core.DUMMY_EqREACTION)
spec_mass_balance = [
("idx_species", numba.int64[:]),
("stoic_coefs", numba.float64[:]),
(
"idx_feed",
numba.types.List(numba.types.Tuple((numba.int64, numba.float64))),
),
("use_constant_value", numba.boolean),
("feed_is_unknown", numba.boolean),
]
# global MassBalance
core.MassBalance = jitclass(spec_mass_balance)(core.MassBalance)
# global DUMMY_MB
core.DUMMY_MB = core.MassBalance(np.array([-1]), np.array([-1.0]),
[(-1, -1.0)], False)
specs_eqsys = [
("c_feed", numba.float64),
# Extra variables for equilibrium specific cases: to be more generic is a list of np.ndarray (it has to be an array)
# Arguments in function is necessary for the jacobian, otherwise the jacobian would need to be generated for each change in args
# This field args is just for convenience to have it stored (OR NOT MAY BE REMOVED)
("args", numba.types.List(numba.float64[:])),
("res", numba.float64[:]),
("TK", numba.float64),
# ('idx', Indexes.class_type.instance_type),
("idx_control", core.IndexController.class_type.instance_type),
# ('species', numba.types.List(Specie.class_type.instance_type)),
("species", type_list_specie),
# ('reactions', numba.types.List(EqReaction.class_type.instance_type)),
("reactions", numba.typeof(l_reactions)),
("ionic_strength", numba.float64),
("pH", numba.float64),
("sc", numba.float64),
("molar_conc", numba.float64[:]),
("gamma", numba.float64[:]),
# ('activity_model_type', numba.typeof(TypeActivityCalculation)),
# ('activity_model_type', numba.types.EnumMember(TypeActivityCalculation, numba.int64)),
(
"mass_balances",
numba.types.List(core.MassBalance.class_type.instance_type),
),
(
"mass_balances_known",
numba.types.List(core.MassBalance.class_type.instance_type),
),
("is_there_known_mb", numba.boolean),
(
"dic_idx_coef",
numba.types.List(numba.types.Tuple((numba.int64, numba.float64))),
),
# ('solid_reactions_but_not_equation', numba.types.List(EqReaction.class_type.instance_type)),
("solid_reactions_but_not_equation", numba.typeof(l_reactions)),
("num_of_feeds", numba.int64),
# System creation related
("feed_compounds", numba.typeof(l_string)),
("closing_equation_type", numba.int64),
("element_mass_balance", numba.typeof(l_string)),
("initial_feed_mass_balance", numba.typeof(l_string)),
("fixed_elements", numba.typeof(l_string)),
# ('database_files', numba.typeof(d_string)),
("reactionsStorage", numba.typeof(l_d_string_float)),
("index_solubility_calculation", numba.int64),
(
"fugacity_calculation",
numba.types.unicode_type,
), # TEST: will fail in numba
("allow_precipitation", numba.boolean),
("solid_reactions_in", numba.typeof(l_d_string_float)),
# ('known_tags', numba.typeof(l_string)),
]
# global EquilibriumSystem
core.EquilibriumSystem = jitclass(specs_eqsys)(core.EquilibriumSystem)
# Activity coefficients model compilation
conductivity.solution_conductivity = numba.njit()(
conductivity.solution_conductivity)
act.log10gamma_davies = numba.njit()(act.log10gamma_davies)
act.debye_huckel_constant = numba.njit()(act.debye_huckel_constant)
act.b_dot_equation = numba.njit()(act.b_dot_equation)
act.calc_log_gamma_ideal = numba.njit()(act.calc_log_gamma_ideal)
act.calc_log_gamma_dh_bdot = numba.njit()(act.calc_log_gamma_dh_bdot)
act.bromley_model_ion = numba.njit()(act.bromley_model_ion)
act.calc_bromley_method = numba.njit()(act.calc_bromley_method)
act.calc_sit_method = numba.njit()(act.calc_sit_method)
act.calc_log_gamma_pitzer = numba.njit()(act.calc_log_gamma_pitzer)
act.calc_log_gamma_dh_bdot_mean_activity_neutral = numba.njit()(
act.calc_log_gamma_dh_bdot_mean_activity_neutral)
properties_utils.dieletricconstant_water = numba.njit()(
properties_utils.dieletricconstant_water)
properties_utils.density_water = numba.njit()(
properties_utils.density_water)
core.logK_H = numba.njit()(core.logK_H)
print("End JIT compilation settings")
return
else:
return self.values[sp:ep]
def rowinds(self):
ris = np.zeros(self.nnz, np.intc)
for i in range(self.nrows):
sp, ep = self.row_extent(i)
ris[sp:ep] = i
return ris
_CSR64 = type('_CSR64', _CSR.__bases__, dict(_CSR.__dict__))
_CSR = jitclass({
'nrows': n.intc,
'ncols': n.intc,
'nnz': n.intc,
'rowptrs': n.intc[::1],
'colinds': n.intc[::1],
'values': n.float64[::1]
})(_CSR)
_CSR64 = jitclass({
'nrows': n.intc,
'ncols': n.intc,
'nnz': n.int64,
'rowptrs': n.int64[::1],
'colinds': n.intc[::1],
'values': n.float64[::1]
})(_CSR64)
class CSR:
"""
return x + 5.
class _foo:
def __init__(self, fun):
self.fun = fun
def foo(fun):
"""
Inputs are all those of `_foo`.
From these it creates a spec of all members for `numba.jitclass`.
It then creates the jitted class and returns an instance of it initialized with the given inputs.
"""
spec = [("fun", typeof(fun))]
jcl = jitclass(spec)
foo_cl = jcl(_foo)
return foo_cl(fun)
spec = [("fun", typeof(a))]
bar = jitclass(spec)(_foo)(a)
# breakpoint()
bar = foo(a)
baz = foo(b)
bar.fun(5) - a(5)
baz.fun(10) - b(10)
def create_my_class(value):
cls = jitclass([('value', typeof(value))])(MyClass)
return cls(value)
def bda_mapper(time,
interval,
ant1,
ant2,
uvw,
chan_width,
chan_freq,
max_uvw_dist,
flag_row=None,
max_fov=3.0,
decorrelation=0.98,
time_bin_secs=None,
min_nchan=1):
have_time_bin_secs = not is_numba_type_none(time_bin_secs)
Omitted = numba.types.misc.Omitted
decorr_type = (numba.typeof(decorrelation.value) if isinstance(
decorrelation, Omitted) else decorrelation)
fov_type = (numba.typeof(max_fov.value)
if isinstance(max_fov, Omitted) else max_fov)
# If time_bin_secs is None,
# then we set it to the max of the time dtype
# lower down
time_bin_secs_type = time_bin_secs if have_time_bin_secs else time.dtype
spec = [('tbin', numba.uintp), ('bin_count', numba.uintp),
('bin_flag_count', numba.uintp), ('time_sum', time.dtype),
('interval_sum', interval.dtype), ('rs', numba.uintp),
('re', numba.uintp), ('bin_half_Δψ', uvw.dtype),
('max_lm', fov_type), ('n_max', fov_type),
('decorrelation', decorr_type),
('time_bin_secs', time_bin_secs_type),
('max_chan_freq', chan_freq.dtype), ('max_uvw_dist', max_uvw_dist)]
JitBinner = jitclass(spec)(Binner)
def impl(time,
interval,
ant1,
ant2,
uvw,
chan_width,
chan_freq,
max_uvw_dist,
flag_row=None,
max_fov=3.0,
decorrelation=0.98,
time_bin_secs=None,
min_nchan=1):
# 𝞓 𝝿 𝞇 𝞍 𝝼
if decorrelation < 0.0 or decorrelation > 1.0:
raise ValueError("0.0 <= decorrelation <= 1.0 must hold")
if max_fov <= 0.0 or max_fov > 90.0:
raise ValueError("0.0 < max_fov <= 90.0 must hold")
max_lm = np.deg2rad(max_fov)
ubl, _, bl_inv, _ = unique_baselines(ant1, ant2)
utime, _, time_inv, _ = unique_time(time)
nrow = time.shape[0]
ntime = utime.shape[0]
nbl = ubl.shape[0]
nchan = chan_width.shape[0]
nchan_factors = factors(nchan)
bandwidth = chan_width.sum()
if min_nchan is None:
min_nchan = 1
else:
s = np.searchsorted(nchan_factors, min_nchan, side='left')
min_nchan = max(min_nchan, nchan_factors[s])
if nchan == 0:
raise ValueError("zero channels")
# Create the row lookup
row_lookup = np.full((nbl, ntime), -1, dtype=np.int32)
bin_lookup = np.full((nbl, ntime), -1, dtype=np.int32)
bin_chan_width = np.full((nbl, ntime), 0.0, dtype=chan_width.dtype)
sentinel = np.finfo(time.dtype).max
time_lookup = np.full((nbl, ntime), sentinel, dtype=time.dtype)
interval_lookup = np.full((nbl, ntime), sentinel, dtype=interval.dtype)
# Is the entire bin flagged?
bin_flagged = np.zeros((nbl, ntime), dtype=np.bool_)
bin_chan_map = np.empty((nbl, ntime, nchan), dtype=np.int32)
out_rows = 0
nr_of_time_bins = 0
out_row_chans = 0
def update_lookups(finalised, bl):
"""
Closure which updates lookups for a baseline,
given a binner's finalisation data
"""
# NOTE(sjperkins) Why do scalars need this, but not arrays?
nonlocal out_rows
nonlocal out_row_chans
nonlocal min_nchan
tbin = finalised.tbin
time_lookup[bl, tbin] = finalised.time
interval_lookup[bl, tbin] = finalised.interval
bin_flagged[bl, tbin] = finalised.flag
nchan = max(finalised.nchan, min_nchan)
bin_nchan = chan_width.shape[0] // nchan
bin_chan_width[bl, tbin] = bandwidth / finalised.nchan
# Construct the channel map
for c in range(chan_width.shape[0]):
bin_chan_map[bl, tbin, c] = c // bin_nchan
out_rows += 1
out_row_chans += nchan
for r in range(nrow):
t = time_inv[r]
bl = bl_inv[r]
if row_lookup[bl, t] != -1:
raise ValueError("Duplicate (TIME, ANTENNA1, ANTENNA2)")
row_lookup[bl, t] = r
# If we don't have time_bin_secs
# set it to the maximum floating point value,
# effectively ignoring this limit
if not have_time_bin_secs:
time_bin_secs = np.finfo(time.dtype).max
# This derived from Synthesis & Imaging II (18-31)
# Converts decrease in amplitude into change in phase
dphi = np.sqrt(6. * (1. - decorrelation))
binner = JitBinner(0, 0, max_lm, dphi, time_bin_secs, chan_freq.max())
for bl in range(nbl):
# Reset the binner for this baseline
binner.reset()
# Auto-correlated baseline
auto_corr = ubl[bl, 0] == ubl[bl, 1]
for t in range(ntime):
# Lookup row, continue if non-existent
r = row_lookup[bl, t]
if r == -1:
continue
# Start a new bin
if binner.empty:
binner.start_bin(r, time, interval, flag_row)
# Try add the row to the bin
# If this fails, finalise the current bin and start a new one
elif not binner.add_row(r, auto_corr, time, interval, uvw,
flag_row):
f = binner.finalise_bin(auto_corr, uvw, nchan_factors,
chan_width, chan_freq)
update_lookups(f, bl)
# Post-finalisation, the bin is empty, start a new bin
binner.start_bin(r, time, interval, flag_row)
# Record the time bin associated with this row
bin_lookup[bl, t] = binner.tbin
# Finalise any remaining data in the bin
if not binner.empty:
f = binner.finalise_bin(auto_corr, uvw, nchan_factors,
chan_width, chan_freq)
update_lookups(f, bl)
nr_of_time_bins += binner.tbin
# Mark remaining bins as unoccupied and unflagged
for tbin in range(binner.tbin, ntime):
time_lookup[bl, tbin] = sentinel
bin_flagged[bl, tbin] = False
assert out_rows == nr_of_time_bins
# Flatten the time lookup and argsort it
flat_time = time_lookup.ravel()
argsort = np.argsort(flat_time, kind='mergesort')
inv_argsort = np.empty_like(argsort)
# Generate lookup from flattened (bl, time) to output row
for i, a in enumerate(argsort):
inv_argsort[a] = i
# Generate row offsets
fbin_chan_map = bin_chan_map.reshape((-1, nchan))
offsets = np.zeros(out_rows + 1, dtype=np.uint32)
decorr_chan_width = np.empty(out_rows, dtype=chan_width.dtype)
# NOTE(sjperkins)
# This: out_rows > 0
# does not work here for some strange (numba reason?)
if offsets.shape[0] > 0:
offsets[0] = 0
for r in range(1, out_rows + 1):
prev_bin_chans = fbin_chan_map[argsort[r - 1]].max() + 1
offsets[r] = offsets[r - 1] + prev_bin_chans
# Construct the final row map
row_chan_map = np.full((nrow, nchan), -1, dtype=np.int32)
time_ret = np.full(out_row_chans, -1, dtype=time.dtype)
int_ret = np.full(out_row_chans, -1, dtype=interval.dtype)
chan_width_ret = np.full(out_row_chans, -1, dtype=chan_width.dtype)
# Construct output flag row, if necessary
out_flag_row = (None if flag_row is None else np.empty(
out_row_chans, dtype=flag_row.dtype))
# foreach input row
for in_row in range(time.shape[0]):
# Lookup baseline and time
bl = bl_inv[in_row]
t = time_inv[in_row]
# lookup time bin and output row in inv_argsort
tbin = bin_lookup[bl, t]
bin_time = time_lookup[bl, tbin]
bin_interval = interval_lookup[bl, tbin]
flagged = bin_flagged[bl, tbin]
out_row = inv_argsort[bl * ntime + tbin]
decorr_chan_width[out_row] = bin_chan_width[bl, tbin]
# Should never happen, but check
if out_row >= out_rows:
raise RowMapperError("out_row >= out_rows")
# Handle output row flagging
if flag_row is not None and flag_row[in_row] == 0 and flagged:
raise RowMapperError("Unflagged input row "
"contributing to "
"flagged output row. "
"This should never happen!")
# Set up the row channel map, populate
# time, interval and chan_width
for c in range(nchan):
out_offset = offsets[out_row] + bin_chan_map[bl, tbin, c]
# Should never happen, but check
if out_offset >= out_row_chans:
raise RowMapperError("out_offset >= out_row_chans")
# Set the output row for this input row and channel
row_chan_map[in_row, c] = out_offset
# Broadcast the time and interval to the output row
time_ret[out_offset] = bin_time
int_ret[out_offset] = bin_interval
# Add channel contribution for each row
chan_width_ret[out_offset] += chan_width[c]
if flag_row is not None:
out_flag_row[out_offset] = 1 if flagged else 0
return RowMapOutput(row_chan_map, offsets, decorr_chan_width, time_ret,
int_ret, chan_width_ret, out_flag_row)
return impl
def transform_complete(replaced, __funcs, __all__, normal, vec=False):
cache_blacklist = set(['Stichlmair_flood', 'airmass',
'Spitzglass_high', '_to_solve_Spitzglass_high',
'_to_solve_Spitzglass_low', 'Spitzglass_low',
'Oliphant', '_to_solve_Oliphant',
'P_isothermal_critical_flow', 'P_upstream_isothermal_critical_flow',
'isothermal_gas_err_P1', 'isothermal_gas_err_P2', 'isothermal_gas_err_P2_basis', 'isothermal_gas_err_D', 'isothermal_gas',
'v_terminal', 'differential_pressure_meter_solver', 'err_dp_meter_solver_P1', 'err_dp_meter_solver_D2',
'err_dp_meter_solver_P2', 'err_dp_meter_solver_m', 'V_horiz_spherical', 'V_horiz_torispherical',
'Prandtl_von_Karman_Nikuradse', 'plate_enlargement_factor', 'Stichlmair_wet', 'V_from_h',
'SA_partial_horiz_spherical_head', '_SA_partial_horiz_spherical_head_to_int',
'_SA_partial_horiz_ellipsoidal_head_to_int', '_SA_partial_horiz_ellipsoidal_head_limits', 'SA_partial_horiz_ellipsoidal_head',
'_SA_partial_horiz_guppy_head_to_int', 'SA_partial_horiz_guppy_head', 'SA_partial_horiz_torispherical_head',
'SA_from_h', 'V_tank'])
# cache_blacklist = set([])
if vec:
conv_fun = numba.vectorize
extra_args = extra_args_vec
else:
conv_fun = numba.njit
extra_args = extra_args_std
new_mods = transform_module(normal, __funcs, replaced, vec=vec, cache_blacklist=cache_blacklist)
to_change = ['packed_tower._Stichlmair_flood_f_and_jac',
'packed_tower.Stichlmair_flood', 'compressible.isothermal_gas',
'fittings.Darby3K', 'fittings.Hooper2K', 'geometry.SA_partial_horiz_torispherical_head',
'optional.spa.solar_position', 'optional.spa.longitude_obliquity_nutation',
'optional.spa.transit_sunrise_sunset',
'fittings.bend_rounded_Crane', 'geometry.tank_from_two_specs_err',
'friction.roughness_Farshad',
]
transform_lists_to_arrays(normal_fluids, to_change, __funcs, vec=vec, cache_blacklist=cache_blacklist)
# AvailableMethods will be removed in the future in favor of non-numba only
# calls to method functions
to_change = {}
# to_change['friction.roughness_Farshad'] = 'ID in _Farshad_roughness'
for s, bad_branch in to_change.items():
mod, func = s.split('.')
source = inspect.getsource(getattr(getattr(normal_fluids, mod), func))
fake_mod = __funcs[mod]
source = remove_branch(source, bad_branch)
source = remove_for_numba(source)
numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
new_func = fake_mod.__dict__[func]
obj = conv_fun(cache=caching, **extra_args)(new_func)
__funcs[func] = obj
obj.__doc__ = ''
# Do some classes by hand
PlateExchanger_spec = [(k, float64) for k in ('pitch', 'beta', 'gamma', 'a', 'amplitude', 'wavelength',
'b', 'chevron_angle', 'inclination_angle', 'plate_corrugation_aspect_ratio',
'plate_enlargement_factor', 'D_eq', 'D_hydraulic', 'width', 'length', 'thickness',
'd_port', 'plates', 'length_port', 'A_plate_surface', 'A_heat_transfer',
'A_channel_flow', 'channels', 'channels_per_fluid')]
PlateExchanger_spec.append(('chevron_angles', numba.types.UniTuple(numba.types.float64, 2)))
HelicalCoil_spec = [(k, float64) for k in
('Do', 'Dt', 'Di', 'Do_total', 'N', 'pitch', 'H', 'H_total',
'tube_circumference', 'tube_length', 'surface_area', 'helix_angle',
'curvature', 'total_inlet_area', 'total_volume', 'inner_surface_area',
'inlet_area', 'inner_volume', 'annulus_area', 'annulus_volume')]
ATMOSPHERE_1976_spec = [(k, float64) for k in
('Z', 'dT', 'H', 'T_layer', 'T_increase', 'P_layer', 'H_layer', 'H_above_layer',
'T', 'P', 'rho', 'v_sonic',
'mu', 'k', 'g', 'R')]
# # No string support
# PlateExchanger = jitclass(PlateExchanger_spec)(getattr(__funcs['geometry'], 'PlateExchanger'))
# __funcs['PlateExchanger'] = __funcs['geometry'].PlateExchanger = PlateExchanger
HelicalCoil = jitclass(HelicalCoil_spec)(getattr(__funcs['geometry'], 'HelicalCoil'))
__funcs['HelicalCoil'] = __funcs['geometry'].HelicalCoil = HelicalCoil
ATMOSPHERE_1976 = jitclass(ATMOSPHERE_1976_spec)(getattr(__funcs['atmosphere'], 'ATMOSPHERE_1976'))
__funcs['ATMOSPHERE_1976'] = __funcs['atmosphere'].ATMOSPHERE_1976 = ATMOSPHERE_1976
# Not needed
__funcs['friction'].Colebrook = __funcs['Colebrook'] = __funcs['Clamond']
# Works but 50% slower
#__funcs['geometry']._V_horiz_spherical_toint = __funcs['_V_horiz_spherical_toint'] = cfunc("float64(float64, float64, float64, float64)")(normal_fluids.geometry._V_horiz_spherical_toint)
for mod in new_mods:
mod.__dict__.update(__funcs)
try:
__all__.extend(mod.__all__)
except AttributeError:
pass
def make_jitclass_container(self):
spec = {
'data': types.List(dtype=types.List(types.float64[::1])),
}
JCContainer = jitclass(spec)(Container)
return JCContainer
def create_my_class(value):
cls = jitclass(MyClass, [("value", typeof(value))])
return cls(value)
def transform_complete(replaced, __funcs, __all__, normal, vec=False):
if vec:
conv_fun = numba.vectorize
else:
conv_fun = numba.njit
new_mods = transform_module(normal, __funcs, replaced, vec=vec)
# Do some classes by hand
PlateExchanger_spec = [
('pitch', float64),
('beta', float64),
('gamma', float64),
('a', float64),
('amplitude', float64),
('wavelength', float64),
('b', float64),
('chevron_angle', float64),
('inclination_angle', float64),
('plate_corrugation_aspect_ratio', float64),
('plate_enlargement_factor', float64),
('D_eq', float64),
('D_hydraulic', float64),
('width', float64),
('length', float64),
('thickness', float64),
('d_port', float64),
('plates', float64),
('length_port', float64),
('A_plate_surface', float64),
('A_heat_transfer', float64),
('A_channel_flow', float64),
('channels', float64),
('channels_per_fluid', float64),
]
HelicalCoil_spec = [
(k, float64)
for k in ('Do', 'Dt', 'Di', 'Do_total', 'N', 'pitch', 'H', 'H_total',
'tube_circumference', 'tube_length', 'surface_area',
'helix_angle', 'curvature', 'total_inlet_area',
'total_volume', 'inner_surface_area', 'inlet_area',
'inner_volume', 'annulus_area', 'annulus_volume')
]
ATMOSPHERE_1976_spec = [
(k, float64) for k in ('Z', 'dT', 'H', 'T_layer', 'T_increase',
'P_layer', 'H_layer', 'H_above_layer', 'T', 'P',
'rho', 'v_sonic', 'mu', 'k', 'g', 'R')
]
to_change = [
'packed_tower._Stichlmair_flood_f_and_jac',
'packed_tower.Stichlmair_flood'
]
transform_lists_to_arrays(normal_fluids, to_change, __funcs, vec=vec)
# AvailableMethods will be removed in the future in favor of non-numba only
# calls to method functions
to_change_AvailableMethods = [
'friction.friction_factor_curved', 'friction.friction_factor',
'packed_bed.dP_packed_bed', 'two_phase.two_phase_dP',
'drag.drag_sphere', 'two_phase_voidage.liquid_gas_voidage',
'two_phase_voidage.gas_liquid_viscosity'
]
to_change_full_output = [
'two_phase.Mandhane_Gregory_Aziz_regime',
'two_phase.Taitel_Dukler_regime'
]
to_change = {k: 'AvailableMethods' for k in to_change_AvailableMethods}
to_change.update({k: 'full_output' for k in to_change_full_output})
to_change['fittings.Darby3K'] = 'name in Darby: # NUMBA: DELETE'
to_change['fittings.Hooper2K'] = 'name in Hooper: # NUMBA: DELETE'
to_change['friction.roughness_Farshad'] = 'ID in _Farshad_roughness'
for s, bad_branch in to_change.items():
mod, func = s.split('.')
source = inspect.getsource(getattr(getattr(normal_fluids, mod), func))
fake_mod = __funcs[mod]
source = remove_branch(source, bad_branch)
source = remove_for_numba(source)
numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
new_func = fake_mod.__dict__[func]
obj = conv_fun(cache=caching)(new_func)
__funcs[func] = obj
# globals()[func] = obj
obj.__doc__ = ''
to_change = ['compressible.isothermal_gas']
for s in to_change:
mod, func = s.split('.')
source = inspect.getsource(getattr(getattr(normal_fluids, mod), func))
fake_mod = __funcs[mod]
source = remove_for_numba(source)
numba_exec_cacheable(source, fake_mod.__dict__, fake_mod.__dict__)
new_func = fake_mod.__dict__[func]
obj = conv_fun(cache=caching)(new_func)
__funcs[func] = obj
# globals()[func] = obj
obj.__doc__ = ''
# Almost there but one argument has a variable type
#PlateExchanger = jitclass(PlateExchanger_spec)(getattr(__funcs['geometry'], 'PlateExchanger'))
HelicalCoil = jitclass(HelicalCoil_spec)(getattr(__funcs['geometry'],
'HelicalCoil'))
__funcs['HelicalCoil'] = __funcs['geometry'].HelicalCoil = HelicalCoil
ATMOSPHERE_1976 = jitclass(ATMOSPHERE_1976_spec)(getattr(
__funcs['atmosphere'], 'ATMOSPHERE_1976'))
__funcs['ATMOSPHERE_1976'] = __funcs[
'atmosphere'].ATMOSPHERE_1976 = ATMOSPHERE_1976
# Not needed
__funcs['friction'].Colebrook = __funcs['Colebrook'] = __funcs['Clamond']
#for k in ('flow_meter', 'fittings', 'two_phase', 'friction'):
# __funcs[k].friction_factor = __funcs['friction_factor'] = __funcs['Clamond']
#__funcs['PlateExchanger'] = __funcs['geometry'].PlateExchanger = PlateExchanger
# Works but 50% slower
#__funcs['geometry']._V_horiz_spherical_toint = __funcs['_V_horiz_spherical_toint'] = cfunc("float64(float64, float64, float64, float64)")(normal_fluids.geometry._V_horiz_spherical_toint)
# ex = fluids.numba.geometry.PlateExchanger(amplitude=5E-4, wavelength=3.7E-3, length=1.2, width=.3, d_port=.05, plates=51, thickness=1e-10)
#fluids.numba.geometry.HelicalCoil(Do_total=32.0, H_total=22.0, pitch=5.0, Dt=2.0, Di=1.8)
for mod in new_mods:
mod.__dict__.update(__funcs)
def write_many(self, b, count):
self.data[self.loc:self.loc + count] = b
self.loc += count
def tell(self):
return self.loc
def so_far(self):
""" In write mode, the data we have gathered until now
"""
return self.data[:self.loc]
spec8 = [("data", numba.uint8[:]), ("loc", numba.int64), ("len", numba.int64)]
Numpy8 = jitclass(spec8)(NumpyIO)
spec32 = [("data", numba.uint32[:]), ("loc", numba.int64),
("len", numba.int64)]
Numpy32 = jitclass(spec32)(NumpyIO)
def _assemble_objects(assign, defi, rep, val, dic, d, null, null_val, max_defi,
prev_i):
"""Dremel-assembly of arrays of values into lists
Parameters
----------
assign: array dtype O
To insert lists into
defi: int array
Definition levels, max 3
