def test_svd(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("n", 3)
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_u, vdm_sigma, vdm_vt", "`<builtin>svd`(vdm, n)")
cb("vdm_usigma", "`<builtin>array`(n*n)")
cb("vdm_v", "`<builtin>array`(n*n)")
cb("vdm_usigma[i + j*n]",
"vdm_u[i + j*n] * vdm_sigma[j]",
loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_v[i + j*n]",
"vdm_vt[j + i*n]",
loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_2", "`<builtin>matmul`(vdm_usigma, vdm_vt, n, n)")
cb("diff", "vdm-vdm_2")
cb((), "`<builtin>print`(diff)")
cb.yield_state("diff", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert la.norm(result) < 1e-10
def test_class_preamble():
from dagrt.language import CodeBuilder
with CodeBuilder(name="primary") as cb:
cb.assign("<t>", "<t> + <dt>")
cb.yield_state("f()", "f", 0, "final")
code = create_DAGCode_with_steady_phase(cb.statements)
from dagrt.codegen import PythonCodeGenerator
import dagrt.function_registry as freg
preamble = """
@staticmethod
def f():
return 1
"""
f = freg.Function(identifier="f",
language_to_codegen={"python": lambda *args: "self.f()"})
generator = PythonCodeGenerator(
"PythonMethod",
class_preamble=preamble,
function_registry=freg.base_function_registry.register(f))
class_ = generator.get_class(code)
method = class_(function_map={})
method.set_up(t_start=0, dt_start=1, context={})
events = list(method.run(t_end=1))
assert events
assert isinstance(events[0], class_.StateComputed)
assert events[0].state_component == 1
def test_create_ast():
x = var("<cond>x")
class ComparableNop(Statement):
"""
Works around the fact that __hash__ and comparison isn't defined for
statements, but is necessary for comparison purposes when embedded
in an expression.
"""
def _state(self):
return (self.condition, self.id, self.depends_on)
def __eq__(self, other):
return (self.__class__ == other.__class__
and self._state() == other._state())
def __ne__(self, other):
return not self.__eq__(other)
def __hash__(self):
return hash(self._state())
nop = ComparableNop(condition=x, id="nop", depends_on=())
code = create_DAGCode_with_steady_phase([nop])
ast = create_ast_from_phase(code, "main")
assert ast == IfThen(x, StatementWrapper(nop.copy(condition=True)))
def test_arrays_and_linalg(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("n", "4")
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("identity[i]", "0", loops=[("i", 0, "n*n")])
cb("identity[i*n + i]", "1", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_inverse", "`<builtin>linear_solve`(vdm, identity, n, n)")
cb("myarray", "`<builtin>matmul`(vdm, vdm_inverse, n, n)")
cb((), "`<builtin>print`(myarray)")
cb.yield_state("myarray", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
result = result.reshape(4, 4, order="F")
assert la.norm(result - np.eye(4)) < 1e-10
def test_CodeBuilder_assign(python_method_impl):
with CodeBuilder("phase") as builder:
builder(var("x"), 1)
builder.yield_state(var("x"), "x", 0, "final")
code = create_DAGCode_with_steady_phase(builder.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 1
def test_missing_state_detection():
"""Check that the code generator detects there is a missing state."""
from dagrt.language import CodeBuilder
with CodeBuilder(name="state_1") as cb:
cb.switch_phase("state_2")
code = create_DAGCode_with_steady_phase(statements=cb.statements)
with pytest.raises(CodeGenerationError):
verify_code(code)
def test_arrays_and_looping(python_method_impl):
with CodeBuilder(name="primary") as cb:
cb("myarray", "`<builtin>array`(20)")
cb("myarray[i]", "i", loops=[("i", 0, 20)])
cb.yield_state("myarray[15]", "result", 0, "final")
from utils import execute_and_return_single_result
code = create_DAGCode_with_steady_phase(cb.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 15
def test_CodeBuilder_condition_with_else_not_taken(python_method_impl):
with CodeBuilder("phase") as builder:
builder(var("x"), 1)
with builder.if_(var("x"), "==", 1):
builder(var("x"), 2)
with builder.else_():
builder(var("x"), 3)
builder.yield_state(var("x"), "x", 0, "final")
code = create_DAGCode_with_steady_phase(builder.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == 2
def test_isnan(python_method_impl, value):
cbuild = RawCodeBuilder()
cbuild.add_and_get_ids(
Assign(id="assign_1", assignee="x", assignee_subscript=(),
expression=var("<builtin>isnan")(value)),
YieldState(id="return", time=0, time_id="final",
expression=var("x"), component_id="<state>",
depends_on=["assign_1"]))
cbuild.commit()
code = create_DAGCode_with_steady_phase(cbuild.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert result == np.isnan(value)
def test_elementwise_abs():
with CodeBuilder(name="primary") as cb:
cb("i", "<builtin>array(20)")
cb("i[j]", "-j", loops=(("j", 0, 20), ))
# Test new builtin on an array type.
cb("k", "<builtin>elementwise_abs(i)")
with cb.if_("k[20] > 19"):
cb.raise_(AbsFailure)
with cb.if_("k[20] < 19"):
cb.raise_(AbsFailure)
# Test new builtin on a scalar.
cb("l", "<builtin>elementwise_abs(-20)")
with cb.if_("l > 20"):
cb.raise_(AbsFailure)
with cb.if_("l < 20"):
cb.raise_(AbsFailure)
cb("y", "<func>f(0, <state>ytype)")
cb("<state>ytype", "y")
# Test new builtin on a usertype.
cb("<state>ytype", "<builtin>elementwise_abs(<state>ytype)")
# (We check this in the outer test code)
code = create_DAGCode_with_steady_phase(cb.statements)
rhs_function = "<func>f"
from dagrt.function_registry import (base_function_registry,
register_ode_rhs)
freg = register_ode_rhs(base_function_registry,
"ytype",
identifier=rhs_function,
input_names=("y", ))
freg = freg.register_codegen(
rhs_function, "fortran",
f.CallCode("""
${result} = -2*${y}
"""))
codegen = f.CodeGenerator(
"element_abs_test",
function_registry=freg,
user_type_map={"ytype": f.ArrayType((100, ), f.BuiltinType("real*8"))},
timing_function="second")
code_str = codegen(code)
run_fortran([
("element_abs.f90", code_str),
("test_element_abs.f90", read_file("test_element_abs.f90")),
],
fortran_libraries=["lapack", "blas"])
def test_arrays_and_linalg():
from dagrt.function_registry import base_function_registry as freg
with CodeBuilder(name="primary") as cb:
cb("n", "4")
cb("nodes", "`<builtin>array`(n)")
cb("vdm", "`<builtin>array`(n*n)")
cb("identity", "`<builtin>array`(n*n)")
cb("nodes[i]", "i/n", loops=[("i", 0, "n")])
cb("identity[i]", "0", loops=[("i", 0, "n*n")])
cb("identity[i*n + i]", "1", loops=[("i", 0, "n")])
cb("vdm[j*n + i]", "nodes[i]**j", loops=[("i", 0, "n"), ("j", 0, "n")])
cb("vdm_inverse", "`<builtin>linear_solve`(vdm, identity, n, n)")
cb("myarray", "`<builtin>matmul`(vdm, vdm_inverse, n, n)")
cb("myzero", "myarray - identity")
cb((), "`<builtin>print`(myzero)")
with cb.if_("`<builtin>norm_2`(myzero) > 10**(-8)"):
cb.raise_(MatrixInversionFailure)
code = create_DAGCode_with_steady_phase(cb.statements)
codegen = f.CodeGenerator("arrays",
function_registry=freg,
user_type_map={},
emit_instrumentation=True,
timing_function="second")
code_str = codegen(code)
if 0:
with open("arrays.f90", "wt") as outf:
outf.write(code_str)
run_fortran([
("arrays.f90", code_str),
("test_arrays_and_linalg.f90",
read_file("test_arrays_and_linalg.f90")),
],
fortran_libraries=["lapack", "blas"])
def get_IfThenElse_test_code_and_expected_result():
from dagrt.expression import IfThenElse
with CodeBuilder(name="primary") as cb:
cb(var("c1"), IfThenElse(True, 0, 1))
cb(var("c2"), IfThenElse(False, 0, 1))
cb(var("c3"), IfThenElse(IfThenElse(True, True, False), 0, 1))
cb(var("c4"), IfThenElse(IfThenElse(False, True, False), 0, 1))
cb(var("c5"), IfThenElse(True, IfThenElse(True, 0, 1), 2))
cb(var("c6"), IfThenElse(True, IfThenElse(False, 0, 1), 2))
cb(var("c7"), IfThenElse(False, 0, IfThenElse(True, 1, 2)))
cb(var("c8"), IfThenElse(False, 0, IfThenElse(False, 1, 2)))
cb(var("c9"), 1 + IfThenElse(True, 0, 1))
cb(var("c10"), 1 + IfThenElse(False, 0, 1))
cb.yield_state(tuple(var("c" + str(i)) for i in range(1, 11)),
"result", 0, "final")
code = create_DAGCode_with_steady_phase(cb.statements)
return (code, (0, 1, 0, 1, 0, 1, 1, 2, 1, 2))
def test_basic_codegen():
"""Test whether the code generator returns a working method. The
generated method always returns 0."""
cbuild = RawCodeBuilder()
cbuild.add_and_get_ids(
YieldState(id="return",
time=0,
time_id="final",
expression=0,
component_id="state",
depends_on=[]))
cbuild.commit()
code = create_DAGCode_with_steady_phase(cbuild.statements)
codegen = f.CodeGenerator("simple",
user_type_map={
"state":
f.ArrayType(
(200, ),
f.BuiltinType("real (kind=8)"),
)
})
print(codegen(code))
def test_norm(python_method_impl, order, norm_suffix, test_vector):
def true_norm(x):
if np.isscalar(x):
return abs(x)
return np.linalg.norm(x, ord=order)
cbuild = RawCodeBuilder()
cbuild.add_and_get_ids(
Assign(id="assign_1", assignee="x", assignee_subscript=(),
expression=test_vector),
Assign(id="assign_2", assignee="n", assignee_subscript=(),
expression=(
var("<builtin>norm_%s" % norm_suffix)(var("x"))),
depends_on=["assign_1"]),
YieldState(id="return", time=0, time_id="final",
expression=var("n"), component_id="<state>",
depends_on=["assign_2"]))
cbuild.commit()
code = create_DAGCode_with_steady_phase(cbuild.statements)
result = execute_and_return_single_result(python_method_impl, code)
assert np.allclose(result, true_norm(test_vector))
def test_self_dep_in_loop():
with CodeBuilder(name="primary") as cb:
cb("y", "<state>y")
cb("y",
"<func>f(0, 2*i*<func>f(0, y if i > 2 else 2*y))",
loops=(("i", 0, 5), ))
cb("<state>y", "y")
code = create_DAGCode_with_steady_phase(cb.statements)
rhs_function = "<func>f"
from dagrt.function_registry import (base_function_registry,
register_ode_rhs)
freg = register_ode_rhs(base_function_registry,
"ytype",
identifier=rhs_function,
input_names=("y", ))
freg = freg.register_codegen(
rhs_function, "fortran",
f.CallCode("""
${result} = -2*${y}
"""))
codegen = f.CodeGenerator(
"selfdep",
function_registry=freg,
user_type_map={"ytype": f.ArrayType((100, ), f.BuiltinType("real*8"))},
timing_function="second")
code_str = codegen(code)
run_fortran([
("selfdep.f90", code_str),
("test_selfdep.f90", read_file("test_selfdep.f90")),
],
fortran_libraries=["lapack", "blas"])
