def test_llvm_1_to_1_mapping_parameter():
llvm_program = LLVMBuilder(initial_values, variable_names, STATES, DERIVATIVES)
llvm_program.add_mapping(["oscillator1.mechanics.b"], ["oscillator1.mechanics.x_dot"])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.1, 2.2, 2.3]))) == np.array([5, 8, 9.])
def test_llvm_idx_write():
llvm_program = LLVMBuilder(initial_values, variable_names, STATES, DERIVATIVES)
llvm_program.add_mapping(["oscillator1.mechanics.b"],
["oscillator1.mechanics.x_dot", "oscillator1.mechanics.y_dot"])
diff, var_func, var_write = llvm_program.generate(filename)
var_write(100, 4)
assert approx(var_func()) == np.array([1.0, 2.0, 3.0, 4., 100., 6., 7., 8., 9.])
assert approx(diff(np.array([2.6, 2.2, 2.3]))) == np.array([100, 100., 9.])
def test_llvm_1_function_and_mapping_unordered_vars():
llvm_program = LLVMBuilder(initial_values, variable_distributed, STATES, DERIVATIVES)
llvm_program.add_external_function(eval_llvm, eval_llvm_signature, number_of_args=4, target_ids=[2, 3])
llvm_program.add_call(eval_llvm.__qualname__,
["oscillator1.mechanics.x", "oscillator1.mechanics.y",
"oscillator1.mechanics.a", "oscillator1.mechanics.y_dot"], target_ids=[2, 3])
llvm_program.add_mapping(args=["oscillator1.mechanics.a"],
targets=["oscillator1.mechanics.x_dot"])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.1, 2.2, 2.3]))) == np.array([50., -50., 4.])
assert approx(diff(np.array([2.3, 2.2, 2.1]))) == np.array([-100, 100, 4.])
def test_llvm_2_function_and_mappings():
llvm_program = LLVMBuilder(initial_values, variable_names, STATES, DERIVATIVES)
llvm_program.add_external_function(eval_llvm, eval_llvm_signature, number_of_args=4, target_ids=[2, 3])
llvm_program.add_call(eval_llvm.__qualname__, ["oscillator1.mechanics.b", "oscillator1.mechanics.y",
"oscillator1.mechanics.a", "oscillator1.mechanics.y_dot"],
target_ids=[2, 3])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.1, 2.2, 2.3]))) == np.array([7., 100., 9.])
assert approx(var_func()) == np.array([2.1, 2.2, 2.3, -100., 5., 6., 7., 100., 9.])
def test_llvm_loop_mix():
llvm_program = LLVMBuilder(initial_values, variable_names, STATES, DERIVATIVES)
llvm_program.add_external_function(eval_llvm_mix, eval_llvm_mix_signature, number_of_args=4, target_ids=[1, 3])
llvm_program.add_set_call(eval_llvm_mix.__qualname__, [
["oscillator1.mechanics.y", "oscillator1.mechanics.z", "oscillator1.mechanics.a", "oscillator1.mechanics.b"],
["oscillator1.mechanics.c", "oscillator1.mechanics.x_dot", "oscillator1.mechanics.y_dot",
"oscillator1.mechanics.z_dot"]],
targets_ids=[1, 3])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.6, 2.2, 2.3]))) == np.array([50, 8, -50])
assert approx(var_func()) == np.array([2.6, 2.2, 2.3, 4, -50., 6., 50, 8., -50.])
def test_llvm_loop_seq():
llvm_program = LLVMBuilder(initial_values, variable_names, STATES, DERIVATIVES)
llvm_program.add_external_function(eval_llvm, eval_llvm_signature, number_of_args=4, target_ids=[2, 3])
llvm_program.add_set_call(eval_llvm.__qualname__, [
["oscillator1.mechanics.y", "oscillator1.mechanics.z", "oscillator1.mechanics.a", "oscillator1.mechanics.b"],
["oscillator1.mechanics.c", "oscillator1.mechanics.x_dot", "oscillator1.mechanics.y_dot",
"oscillator1.mechanics.z_dot"]],
targets_ids=[2, 3])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.6, 2.2, 2.3]))) == np.array([7, 50., -50])
##Note that state is not changed. States can only be changed by the solver
assert approx(var_func()) == np.array([2.6, 2.2, 2.3, 50, -50., 6., 7., 50., -50.])
def __init__(self,
filename,
equation_graph,
scope_variables,
equations,
scoped_equations,
temporary_variables,
system_tag="",
use_llvm=True,
imports=None):
self.filename = filename
self.imports = imports
self.system_tag = system_tag
self.scope_variables = scope_variables
self.set_variables = {}
self.states = []
self.deriv = []
for ix, (sv_id, sv) in enumerate(self.scope_variables.items()):
full_tag = d_u(sv.id)
if sv.type == VariableType.STATE:
self.states.append(full_tag)
elif sv.type == VariableType.DERIVATIVE:
self.deriv.append(full_tag)
for k, var in temporary_variables.items():
if var.type == VariableType.TMP_PARAMETER_SET:
self.set_variables.update({k: var})
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k in var.set_var.variables:
tail.update({k: v})
new_sv.update({
var.tmp_vars[v.set_var_ix].id:
var.tmp_vars[v.set_var_ix]
})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
if var.type == VariableType.TMP_PARAMETER:
new_sv = {}
tail = {}
for k, v in self.scope_variables.items():
if k == var.scope_var_id:
tail.update({k: v})
new_sv.update({var.id: var})
else:
new_sv.update({k: v})
self.scope_variables = dict(new_sv, **tail)
self.scoped_equations = scoped_equations
self.temporary_variables = temporary_variables
self.values_order = {}
self.scope_var_node = {}
self.scalar_variables = {}
self._parse_variables()
# Sort the graph topologically to start generating code
self.topo_sorted_nodes = equation_graph.topological_nodes()
self.equation_graph = equation_graph.clean()
self.number_of_states = len(self.states)
self.number_of_derivatives = len(self.deriv)
# Initialize llvm builder - will be a list of intermediate llvm instructions to be lowered in generate
self.llvm = use_llvm
if self.llvm:
self.generated_program = LLVMBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
else:
self.generated_program = ASTBuilder(np.ascontiguousarray(
[x.value for x in self.scope_variables.values()],
dtype=np.float64),
self.values_order,
self.states,
self.deriv,
system_tag=self.system_tag)
self.mod_body = []
# Create a kernel of assignments and calls
self.eq_vardefs = {}
# Loop over equation functions and generate code for each equation.
self._parse_equations(equations)
self.all_targeted = []
self.all_read = []
self.all_targeted_set_vars = []
self.all_read_set_vars = []
def test_llvm_unordered_vars():
llvm_program = LLVMBuilder(initial_values, variable_distributed, STATES, DERIVATIVES)
llvm_program.add_external_function(eval_llvm, eval_llvm_signature, number_of_args=4, target_ids=[2, 3])
diff, var_func, _ = llvm_program.generate(filename)
assert approx(diff(np.array([2.1, 2.2, 2.3]))) == np.array([2., 3., 4.])
assert approx(var_func()) == np.array([1., 2., 3., 4., 2.1, 6., 7., 2.2, 2.3])