def test_linear_solve_matrix_tape(self):
y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
with math.SolveTape() as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]],
abs_tolerance=1e-3)
assert len(solves) == 1
assert solves[0] == solves[solve]
math.assert_close(solves[solve].residual.values,
0,
abs_tolerance=1e-3)
assert math.close(solves[solve].iterations, 2) or math.close(
solves[solve].iterations, -1)
with math.SolveTape(record_trajectories=True) as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]],
abs_tolerance=1e-3)
assert solves[solve].x.trajectory.size == 3
math.assert_close(solves[solve].residual.trajectory[-1].values,
0,
abs_tolerance=1e-3)
def plot_solves():
"""
While `plot_solves()` is active, certain performance optimizations and algorithm implementations may be disabled.
"""
from . import math
import pylab
cycle = pylab.rcParams['axes.prop_cycle'].by_key()['color']
with math.SolveTape(record_trajectories=True) as solves:
try:
yield solves
finally:
for i, result in enumerate(solves):
assert isinstance(result, math.SolveInfo)
from phi.math._tensors import disassemble_tree
_, (residual, ) = disassemble_tree(result.residual)
residual_mse = math.mean(math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
residual_mse_max = math.max(
math.sqrt(math.sum(residual**2)),
residual.shape.without('trajectory'))
# residual_mean = math.mean(math.abs(residual), residual.shape.without('trajectory'))
residual_max = math.max(math.abs(residual),
residual.shape.without('trajectory'))
pylab.plot(residual_mse.numpy(),
label=f"{i}: {result.method}",
color=cycle[i % len(cycle)])
pylab.plot(residual_max.numpy(),
'--',
alpha=0.2,
color=cycle[i % len(cycle)])
pylab.plot(residual_mse_max.numpy(),
alpha=0.2,
color=cycle[i % len(cycle)])
print(
f"Solve {i}: {result.method} ({1000 * result.solve_time:.1f} ms)\n"
f"\t{result.solve}\n"
f"\t{result.msg}\n"
f"\tConverged: {result.converged}\n"
f"\tDiverged: {result.diverged}\n"
f"\tIterations: {result.iterations}\n"
f"\tFunction evaulations: {result.function_evaluations.trajectory[-1]}"
)
pylab.yscale('log')
pylab.ylabel("Residual: MSE / max / individual max")
pylab.xlabel("Iteration")
pylab.title(f"Solve Convergence")
pylab.legend(loc='upper right')
pylab.savefig(f"pressure-solvers-FP32.png")
pylab.show()
def test_solve_diverge(self):
y = math.ones(spatial(x=2)) * (1, 2)
x0 = math.zeros(spatial(x=2))
for method in ['CG']:
solve = Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
try:
field.solve_linear(math.jit_compile_linear(math.laplace), y,
solve)
assert False
except Diverged:
pass
with math.SolveTape(record_trajectories=True) as solves:
try:
field.solve_linear(math.jit_compile_linear(math.laplace),
y, solve) # impossible
assert False
except Diverged:
pass
def test_solve_linear_function_batched(self):
y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
x = field.solve_linear(math.jit_compile_linear(field.laplace), y,
solve)
math.assert_close(x.values,
math.wrap([[-1.5, -2, -1.5], [-3, -4, -3]],
channel('vector'), spatial('x')),
abs_tolerance=1e-3)
with math.SolveTape() as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values,
math.wrap([[-1.5, -2, -1.5], [-3, -4, -3]],
channel('vector'), spatial('x')),
abs_tolerance=1e-3)
assert len(solves) == 1
assert solves[0] == solves[solve]
math.assert_close(solves[solve].residual.values,
0,
abs_tolerance=1e-3)
def test_minimize(self):
def loss(x, y):
return math.l2_loss(x - 1) + math.l2_loss(y + 1)
for backend in BACKENDS:
if backend.supports(Backend.functional_gradient):
with backend:
x0 = tensor([0, 0, 0], spatial('x')), tensor([-1, -1, -1],
spatial('y'))
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
x0 = tensor([[0, 0, 0], [1, 1, 1]], batch('batch'),
spatial('x')), tensor(
[[0, 0, 0], [-1, -1, -1]], batch('batch'),
spatial('y'))
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
with math.SolveTape() as solves:
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(solves[0].residual,
0,
abs_tolerance=1e-4)
assert (solves[0].iterations <= (4, 0)).all
assert (solves[0].function_evaluations <= (30, 1)).all
with math.SolveTape(
record_trajectories=True) as trajectories:
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(trajectories[0].residual.trajectory[-1],
0,
abs_tolerance=1e-4)
assert (
trajectories[0].iterations == solves[0].iterations).all
assert trajectories[
0].residual.trajectory.size == trajectories[0].x[
0].trajectory.size
assert trajectories[0].residual.trajectory.size > 1