def test_step_matrix_sparse():
from leap.step_matrix import StepMatrixFinder, fast_evaluator
from pymbolic import var
component_id = 'y'
code = euler(component_id, show_dag=False)
J = np.diag([-3, -2, -1]) # noqa
def rhs_sym(t, y):
return J.dot(y)
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym},
variables=["<state>" + component_id])
dt = var("<dt>")
mat = finder.get_phase_step_matrix("primary",
shapes={"<state>" + component_id: 3},
sparse=True)
assert mat.shape == (3, 3)
assert mat.indices == [(0, 0), (1, 1), (2, 2)]
true_mat = np.eye(3, dtype=np.object) + dt * J
assert (mat.data == np.diag(true_mat)).all()
eval_mat = fast_evaluator(mat, sparse=True)
eval_result = eval_mat({"<dt>": 1})
assert eval_result.shape == (3, 3)
assert eval_result.indices == [(0, 0), (1, 1), (2, 2)]
assert eval_result.data == [-2, -1, 0]
def test_step_matrix_vector_state(show_matrix=True, show_dag=False):
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
component_id = 'y'
code = euler(component_id, show_dag)
J = np.diag([-3, -2, -1]) # noqa
def rhs_sym(t, y):
return J.dot(y)
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym},
variables=["<state>" + component_id])
mat = finder.get_phase_step_matrix("primary",
shapes={"<state>" + component_id: 3})
if show_matrix:
print('Variables: %s' % finder.variables)
from pytools import indices_in_shape
for i in indices_in_shape(mat.shape):
print(i, mat[i])
# XXX: brittle
dt = var("<dt>")
true_mat = np.eye(3, dtype=np.object) + dt * J
assert (mat == true_mat).all()
def generate_mrab_step_matrix(ngridpoints, coeffs, substep_ratio, filename):
problem = VariableCoeffWaveEquationProblem(ngridpoints, coeffs)
code, rhs_map = make_3_component_multirate_method(problem,
substep_ratio,
code_only=True,
return_rhs_map=True)
from leap.step_matrix import StepMatrixFinder
finder = StepMatrixFinder(code,
function_map=rhs_map,
exclude_variables=["<p>bootstrap_step"])
with timer(f"Constructing MRAB({substep_ratio}) matrix"):
component_sizes = {}
for var in finder.variables:
for i in range(problem.ncomponents):
if "comp%d" % i in var:
component_sizes[var] = problem.component_sizes[i]
break
else:
raise ValueError("cannot infer size of variable: %s" % var)
import pprint
pprint.pprint(component_sizes)
mat = finder.get_phase_step_matrix("primary",
shapes=component_sizes,
sparse=True)
with open(filename, "wb") as outf:
import pickle
pickle.dump(mat, outf)
logging.info("%s: %d nnz, size %s", filename, len(mat.data), mat.shape)
def main():
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
#method = RK4Method("y")
method = AdamsBashforthMethod("y", order=3, static_dt=True)
code = method.generate()
print(code)
def rhs_sym(t, y):
return var("lmbda") * y
finder = StepMatrixFinder(code,
function_map={"<func>y": rhs_sym},
exclude_variables=["<p>step"])
print(finder.get_maxima_expressions("primary"))
mat = finder.get_phase_step_matrix("primary")
print('Variables: %s' % finder.variables)
np.set_printoptions(formatter={"all": str})
print(mat)
tol = 1e-8
from leap.step_matrix import fast_evaluator
evaluate_mat = fast_evaluator(mat)
def is_stable(direction, dt):
smat = evaluate_mat({"<dt>": dt, "lmbda": direction})
eigvals = la.eigvals(smat)
return (np.abs(eigvals) <= 1 + tol).all()
from leap.stability import find_truth_bdry
from functools import partial
prec = 1e-5
print("stable imaginary timestep:",
find_truth_bdry(partial(is_stable, 1j), prec=prec))
print("stable neg real timestep:",
find_truth_bdry(partial(is_stable, -1), prec=prec))
def test_step_matrix_fast_eval():
from leap.step_matrix import StepMatrixFinder, fast_evaluator
component_id = 'y'
code = euler(component_id, show_dag=False)
J = np.diag([-3, -2, -1]) # noqa
def rhs_sym(t, y):
return J.dot(y)
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym},
variables=["<state>" + component_id])
mat = finder.get_phase_step_matrix("primary",
shapes={"<state>" + component_id: 3})
eval_mat = fast_evaluator(mat)
assert (eval_mat({"<dt>": 1}) == np.diag([-2, -1, 0])).all()
def test_step_matrix(method, show_matrix=True, show_dag=False):
component_id = 'y'
code = method.generate()
if show_dag:
from dagrt.language import show_dependency_graph
show_dependency_graph(code)
from dagrt.exec_numpy import NumpyInterpreter
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
# {{{ build matrix
def rhs_sym(t, y):
return var("lambda") * y
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym})
mat = finder.get_phase_step_matrix("primary")
if show_matrix:
print('Variables: %s' % finder.variables)
from pytools import indices_in_shape
for i in indices_in_shape(mat.shape):
print(i, mat[i])
# }}}
dt = 0.1
lambda_ = -0.4
def rhs(t, y):
return lambda_ * y
interp = NumpyInterpreter(code,
function_map={"<func>" + component_id: rhs})
interp.set_up(t_start=0, dt_start=dt, context={component_id: 15})
assert interp.next_phase == "initial"
for event in interp.run_single_step():
pass
assert interp.next_phase == "primary"
start_values = np.array([interp.context[v] for v in finder.variables])
for event in interp.run_single_step():
pass
assert interp.next_phase == "primary"
stop_values = np.array([interp.context[v] for v in finder.variables])
from dagrt.expression import EvaluationMapper
concrete_mat = EvaluationMapper({
"lambda": lambda_,
"<dt>": dt,
}, {})(mat)
stop_values_from_mat = concrete_mat.dot(start_values)
rel_err = (
la.norm(stop_values - stop_values_from_mat) / # noqa: W504
la.norm(stop_values))
assert rel_err < 1e-12
def main():
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
speed_factor = 10
method_name = "Fq"
order = 3
tol = 1e-8
prec = 1e-5
angles = np.linspace(0, 2 * np.pi, 100, endpoint=False)
for step_ratio in [1, 2, 3, 4, 5, 6]:
print("speed factor: %g - step ratio: %g - method: %s "
"- order: %d" % (speed_factor, step_ratio, method_name, order))
method = TwoRateAdamsBashforthMethod(method=method_name,
order=order,
step_ratio=step_ratio,
static_dt=True)
code = method.generate()
finder = StepMatrixFinder(code,
function_map={
"<func>f2f":
lambda t, f, s: var("f2f") * f,
"<func>s2f":
lambda t, f, s: var("s2f") * s,
"<func>f2s":
lambda t, f, s: var("f2s") * f,
"<func>s2s":
lambda t, f, s: var("s2s") * s,
},
exclude_variables=["<p>bootstrap_step"])
mat = finder.get_phase_step_matrix("primary")
if 0:
print('Variables: %s' % finder.variables)
np.set_printoptions(formatter={"all": str})
print(mat)
from leap.step_matrix import fast_evaluator
evaluate = fast_evaluator(mat)
def is_stable(major_eigval, dt):
smat = evaluate({
"<dt>": dt,
"f2f": major_eigval,
"s2f": 1 / speed_factor,
"f2s": 1 / speed_factor,
"s2s": major_eigval * 1 / speed_factor,
})
eigvals = la.eigvals(smat)
return (np.abs(eigvals) <= 1 + tol).all()
from leap.stability import find_truth_bdry
from functools import partial
points = []
for angle in angles:
eigval = np.exp(1j * angle)
max_dt = find_truth_bdry(partial(is_stable, eigval), prec=prec)
stable_fake_eigval = eigval * max_dt
points.append([stable_fake_eigval.real, stable_fake_eigval.imag])
points = np.array(points).T
pt.plot(points[0], points[1], "x", label="steprat: %d" % step_ratio)
pt.legend(loc="best")
pt.grid()
outfile = "mr-stability-diagram.pdf"
pt.savefig(outfile)
print("Output written to %s" % outfile)
def main():
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
speed_factor = 10
step_ratio = 7
method_name = "Fq"
order = 3
print("speed factor: %g - step ratio: %g - method: %s "
"- order: %d" % (speed_factor, step_ratio, method_name, order))
method = TwoRateAdamsBashforthMethodBuilder(method=method_name,
order=order,
step_ratio=step_ratio,
static_dt=True)
code = method.generate()
finder = StepMatrixFinder(code,
function_map={
"<func>f2f": lambda t, f, s: var("f2f") * f,
"<func>s2f": lambda t, f, s: var("s2f") * s,
"<func>f2s": lambda t, f, s: var("f2s") * f,
"<func>s2s": lambda t, f, s: var("s2s") * s,
},
exclude_variables=["<p>bootstrap_step"])
mat = finder.get_phase_step_matrix("primary")
if 0:
print("Variables: %s" % finder.variables)
np.set_printoptions(formatter={"all": str})
print(mat)
tol = 1e-8
from leap.step_matrix import fast_evaluator
evaluate_mat = fast_evaluator(mat)
def is_stable(direction, dt):
smat = evaluate_mat({
"<dt>": dt,
"f2f": direction,
"s2f": 1 / speed_factor,
"f2s": 1 / speed_factor,
"s2s": direction * 1 / speed_factor,
})
eigvals = la.eigvals(smat)
return (np.abs(eigvals) <= 1 + tol).all()
from leap.stability import find_truth_bdry
from functools import partial
prec = 1e-5
print("stable imaginary timestep:",
find_truth_bdry(partial(is_stable, 1j), prec=prec))
print("stable neg real timestep:",
find_truth_bdry(partial(is_stable, -1), prec=prec))
def main():
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
speed_factor = 10
method_name = "Fq"
order = 3
step_ratio = 3
prec = 1e-5
method = TwoRateAdamsBashforthMethod(method=method_name,
order=order,
step_ratio=step_ratio,
static_dt=True)
code = method.generate()
finder = StepMatrixFinder(code,
function_map={
"<func>f2f": lambda t, f, s: var("f2f") * f,
"<func>s2f": lambda t, f, s: var("s2f") * s,
"<func>f2s": lambda t, f, s: var("f2s") * f,
"<func>s2s": lambda t, f, s: var("s2s") * s,
},
exclude_variables=["<p>bootstrap_step"])
mat = finder.get_phase_step_matrix("primary")
left = -3
right = 1
bottom = -4
top = 4
res = 200
points = np.mgrid[left:right:res * 1j, bottom:top:res * 1j]
eigvals = points[0] + 1j * points[1]
eigvals_flat = eigvals.reshape(-1)
from multiprocessing import Pool
pool = Pool()
from leap.step_matrix import fast_evaluator
evaluate_mat = fast_evaluator(mat)
stable_dts_list = pool.map(
partial(process_eigval, evaluate_mat, speed_factor, prec),
eigvals_flat)
max_eigvals = np.zeros(eigvals.shape)
max_eigvals.reshape(-1)[:] = stable_dts_list
pt.title("speed factor: %g - step ratio: %g - method: %s "
"- order: %d" % (speed_factor, step_ratio, method_name, order))
log_max_eigvals = np.log10(1e-15 + max_eigvals.T)
pt.imshow(log_max_eigvals,
extent=(left, right, bottom, top),
cmap="viridis")
pt.colorbar()
pt.contour(eigvals.real, eigvals.imag, log_max_eigvals.T, [0], zorder=10)
pt.gca().set_aspect("equal")
pt.grid()
outfile = "mr-max-eigvals.pdf"
pt.savefig(outfile)
print("Output written to %s" % outfile)