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 total(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
total = math.sum(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
total = math.expand_dims(total, -1)
return total
def normalize(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
global_mean = math.mean(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
global_mean = math.expand_dims(global_mean, -1)
return StaggeredGrid(self.staggered / global_mean)
def approximate_signed_distance(self, location):
"""
Computes the exact distance from location to the closest point on the sphere.
Very close to the sphere center, the distance takes a constant value.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
center = math.batch_align(self.center, 1, location)
radius = math.batch_align(self.radius, 0, location)
distance_squared = math.sum((location - center)**2,
axis=-1,
keepdims=True)
distance_squared = math.maximum(
distance_squared,
radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - radius
def approximate_signed_distance(self, location):
"""
Computes the exact distance from location to the closest point on the sphere.
Very close to the sphere center, the distance takes a constant value.
Args:
location: float tensor of shape (batch_size, ..., rank)
Returns:
float tensor of shape (*location.shape[:-1], 1).
"""
distance_squared = math.vec_squared(location - self.center)
distance_squared = math.maximum(
distance_squared,
self.radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - self.radius
def test_infinite_cylinder(self):
cylinder = geom.infinite_cylinder(x=.5,
y=.5,
radius=.5,
inf_dim=math.spatial('z'))
self.assertEqual(cylinder.spatial_rank, 3)
cylinder = geom.infinite_cylinder(x=.5, y=.5, radius=.5, inf_dim='z')
loc = math.wrap([(0, 0, 0), (.5, .5, 0), (1, 1, 0), (0, 0, 100),
(.5, .5, 100)], math.instance('points'),
math.channel(vector='x,y,z'))
inside = math.wrap([False, True, False, False, True],
math.instance('points'))
math.assert_close(cylinder.lies_inside(loc), inside)
loc = math.wrap([(0, 0, 0), (.5, 1, 0), (1, 1, 0), (0, 0, 100),
(.5, 1, 100)], math.instance('points'),
math.channel(vector='x,y,z'))
corner_distance = math.sqrt(2) / 2 - .5
distance = math.wrap(
[corner_distance, 0, corner_distance, corner_distance, 0],
math.instance('points'))
math.assert_close(cylinder.approximate_signed_distance(loc), distance)
def soft_sqrt(self):
return StaggeredGrid(math.sqrt(math.maximum(self.staggered, 1e-20)))
def normalize_probability(probability_amplitude):
p = math.to_float(abs(probability_amplitude)**2)
P = math.sum(p, math.spatial_dimensions(p), keepdims=True)
return probability_amplitude / math.to_complex(math.sqrt(P))