def test_fourier_laplace_2d_periodic(self):
"""test for convergence of the laplace operator"""
test_params = {
'size': [16, 32, 40],
'L': [1, 2,
3], # NOTE: Cannot test with less than 1 full wavelength
}
test_cases = [
dict(zip(test_params, v)) for v in product(*test_params.values())
]
for params in test_cases:
vec = math.meshgrid(x=params['size'], y=params['size'])
sine_field = math.prod(
math.sin(2 * PI * params['L'] * vec / params['size'] + 1),
'vector')
sin_lap_ref = -2 * (
2 * PI * params['L'] / params['size']
)**2 * sine_field # leading 2 from from x-y cross terms
sin_lap = math.fourier_laplace(sine_field, 1)
try:
math.assert_close(sin_lap,
sin_lap_ref,
rel_tolerance=0,
abs_tolerance=1e-5)
except BaseException as e:
abs_error = math.abs(sin_lap - sin_lap_ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / sin_lap_ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(f"{variation_str}\n{params}")
raise AssertionError(e, f"{variation_str}\n{params}")
def data_bounds(field):
assert field.has_points
try:
data = field.points.data
min_vec = math.min(data, axis=tuple(range(len(data.shape) - 1)))
max_vec = math.max(data, axis=tuple(range(len(data.shape) - 1)))
except StaggeredSamplePoints:
boxes = [data_bounds(c) for c in field.unstack()]
min_vec = math.min([b.lower for b in boxes], axis=0)
max_vec = math.max([b.upper for b in boxes], axis=0)
return AABox(min_vec, max_vec)
def test__periodic_2d_arakawa_poisson_bracket(self):
"""test _periodic_2d_arakawa_poisson_bracket implementation"""
with math.precision(64):
# Define parameters to test
test_params = {
'grid_size': [(4, 4), (32, 32)],
'dx': [0.1, 1],
'gen_func': [
lambda grid_size: np.random.rand(*grid_size).reshape(
grid_size)
]
}
# Generate test cases as the product
test_cases = [
dict(zip(test_params, v))
for v in product(*test_params.values())
]
for params in test_cases:
grid_size = params['grid_size']
d1 = params['gen_func'](grid_size)
d2 = params['gen_func'](grid_size)
dx = params['dx']
padding = extrapolation.PERIODIC
ref = self.arakawa_reference_implementation(
np.pad(d1.copy(), 1, mode='wrap'),
np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1]
d1_tensor = field.CenteredGrid(
values=math.tensor(d1, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
d2_tensor = field.CenteredGrid(
values=math.tensor(d2, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
val = math._nd._periodic_2d_arakawa_poisson_bracket(
d1_tensor.values, d2_tensor.values, dx)
try:
math.assert_close(ref,
val,
rel_tolerance=1e-14,
abs_tolerance=1e-14)
except BaseException as e:
abs_error = math.abs(val - ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(ref)
print(val)
raise AssertionError(e, params, variation_str)
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 _plot_scalar_grid(grid: Grid, title, colorbar, cmap, figsize, same_scale):
batch_size = grid.shape.batch.volume
values = math.join_dimensions(grid.values, grid.shape.channel,
'channel').channel[0]
plt_args = {}
if same_scale:
plt_args['vmin'] = math.min(values).native()
plt_args['vmax'] = math.max(values).native()
b_values = math.join_dimensions(values, grid.shape.batch, 'batch')
fig, axes = plt.subplots(1, batch_size, figsize=figsize)
axes = axes if isinstance(axes, np.ndarray) else [axes]
for b in range(batch_size):
im = axes[b].imshow(b_values.batch[b].numpy('y,x'),
origin='lower',
cmap=cmap,
**plt_args)
if title:
if isinstance(title, str):
sub_title = title
elif title is True:
sub_title = f"{b} of {grid.shape.batch}"
elif isinstance(title, (tuple, list)):
sub_title = title[b]
else:
sub_title = None
if sub_title is not None:
axes[b].set_title(sub_title)
if colorbar:
plt.colorbar(im, ax=axes[b])
plt.tight_layout()
return fig, axes
def extrapolation_helper(elements, t_shift, v_field, mask):
shift = math.ceil(math.max(
math.abs(elements.center - points.center))) - t_shift
t_shift += shift
v_field, mask = extrapolate_valid(v_field, mask, int(shift))
v_field *= accessible
return v_field, mask, t_shift
def sample_at(self, points, collapse_dimensions=True):
if len(self.geometries) == 0:
return _expand_axes(math.zeros([1,1]), points, collapse_dimensions=collapse_dimensions)
if len(self.geometries) == 1:
result = self.geometries[0].value_at(points)
else:
result = math.max([geometry.value_at(points) for geometry in self.geometries], axis=0)
return math.mul(result, self.data)
def value_at(self, points):
if len(self.geometries) == 1:
result = self.geometries[0].value_at(points)
else:
result = math.max(
[geometry.value_at(points) for geometry in self.geometries],
axis=0)
return result
def value_at(self, points, collapse_dimensions=True):
if len(self.geometries) == 1:
result = self.geometries[0].value_at(points)
else:
result = math.max(
[geometry.value_at(points) for geometry in self.geometries],
axis=0)
return result
def bounds(self) -> Box:
if self._bounds is not None:
return self._bounds
else:
from phi.field._field_math import data_bounds
bounds = data_bounds(self.elements.center)
radius = math.max(self.elements.bounding_radius())
return Box(bounds.lower - radius, bounds.upper + radius)
def data_bounds(loc: SampledField or Tensor) -> Box:
if isinstance(loc, SampledField):
loc = loc.points
assert isinstance(
loc,
Tensor), f"loc must be a Tensor or SampledField but got {type(loc)}"
min_vec = math.min(loc, dim=loc.shape.non_batch.non_channel)
max_vec = math.max(loc, dim=loc.shape.non_batch.non_channel)
return Box(min_vec, max_vec)
def approximate_signed_distance(self, location):
"""
Computes the signed L-infinity norm (manhattan distance) from the location to the nearest side of the box.
For an outside location `l` with the closest surface point `s`, the distance is `max(abs(l - s))`.
For inside locations it is `-max(abs(l - s))`.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
lower, upper = math.batch_align([self.lower, self.upper], 1, location)
center = 0.5 * (lower + upper)
extent = upper - lower
distance = math.abs(location - center) - extent * 0.5
return math.max(distance, axis=-1, keepdims=True)
def _choose_solver(resolution, backend):
use_fourier = math.max(resolution) > 64
if backend.precision == 64:
from .fourier import FourierSolver
from .geom import GeometricCG
return FourierSolver() & GeometricCG(accuracy=1e-8) if use_fourier else GeometricCG(accuracy=1e-8)
elif backend.precision == 32 and backend.matches_name('SciPy'):
from .sparse import SparseSciPy
return SparseSciPy()
elif backend.precision == 32:
from .fourier import FourierSolver
from .sparse import SparseCG
return FourierSolver() & SparseCG(accuracy=1e-5) if use_fourier else SparseCG(accuracy=1e-5)
else: # lower precision
from .geom import GeometricCG
return GeometricCG(accuracy=1e-2)
def approximate_signed_distance(self, location):
"""
Computes the signed L-infinity norm (manhattan distance) from the location to the nearest side of the box.
For an outside location `l` with the closest surface point `s`, the distance is `max(abs(l - s))`.
For inside locations it is `-max(abs(l - s))`.
Args:
location: float tensor of shape (batch_size, ..., rank)
Returns:
float tensor of shape (*location.shape[:-1], 1).
"""
center = 0.5 * (self.lower + self.upper)
extent = self.upper - self.lower
distance = math.abs(location - center) - extent * 0.5
return math.max(distance, 'vector')
def mac_cormack(field: GridType,
velocity: Field,
dt: float,
correction_strength=1.0,
integrator=euler) -> GridType:
"""
MacCormack advection uses a forward and backward lookup to determine the first-order error of semi-Lagrangian advection.
It then uses that error estimate to correct the field values.
To avoid overshoots, the resulting value is bounded by the neighbouring grid cells of the backward lookup.
Args:
field: Field to be advected, one of `(CenteredGrid, StaggeredGrid)`
velocity: Vector field, need not be sampled at same locations as `field`.
dt: Time increment
correction_strength: The estimated error is multiplied by this factor before being applied.
The case correction_strength=0 equals semi-lagrangian advection. Set lower than 1.0 to avoid oscillations.
integrator: ODE integrator for solving the movement.
Returns:
Advected field of type `type(field)`
"""
v0 = sample(velocity, field.elements)
points_bwd = integrator(field.elements, velocity, -dt, v0=v0)
points_fwd = integrator(field.elements, velocity, dt, v0=v0)
# Semi-Lagrangian advection
field_semi_la = field.with_values(reduce_sample(field, points_bwd))
# Inverse semi-Lagrangian advection
field_inv_semi_la = field.with_values(
reduce_sample(field_semi_la, points_fwd))
# correction
new_field = field_semi_la + correction_strength * 0.5 * (field -
field_inv_semi_la)
# Address overshoots
limits = field.closest_values(points_bwd)
lower_limit = math.min(
limits, [f'closest_{dim}' for dim in field.shape.spatial.names])
upper_limit = math.max(
limits, [f'closest_{dim}' for dim in field.shape.spatial.names])
values_clamped = math.clip(new_field.values, lower_limit, upper_limit)
return new_field.with_values(values_clamped)
def mac_cormack(field: GridType,
velocity: Field,
dt: float,
correction_strength=1.0) -> GridType:
"""
MacCormack advection uses a forward and backward lookup to determine the first-order error of semi-Lagrangian advection.
It then uses that error estimate to correct the field values.
To avoid overshoots, the resulting value is bounded by the neighbouring grid cells of the backward lookup.
Args:
field: Field to be advected, one of `(CenteredGrid, StaggeredGrid)`
velocity: Vector field, need not be sampled at same locations as `field`.
dt: Time increment
correction_strength: The estimated error is multiplied by this factor before being applied. The case correction_strength=0 equals semi-lagrangian advection. Set lower than 1.0 to avoid oscillations. (Default value = 1.0)
Returns:
Advected field of type `type(field)`
"""
v = velocity.sample_in(field.elements)
x0 = field.points
x_bwd = x0 - v * dt
x_fwd = x0 + v * dt
reduce = x0.shape.non_channel.without(field.shape).names
# Semi-Lagrangian advection
field_semi_la = field.with_(
values=field.sample_at(x_bwd, reduce_channels=reduce))
# Inverse semi-Lagrangian advection
field_inv_semi_la = field.with_(
values=field_semi_la.sample_at(x_fwd, reduce_channels=reduce))
# correction
new_field = field_semi_la + correction_strength * 0.5 * (field -
field_inv_semi_la)
# Address overshoots
limits = field.closest_values(x_bwd, reduce_channels=reduce)
lower_limit = math.min(
limits, [f'closest_{dim}' for dim in field.shape.spatial.names])
upper_limit = math.max(
limits, [f'closest_{dim}' for dim in field.shape.spatial.names])
values_clamped = math.clip(new_field.values, lower_limit, upper_limit)
return new_field.with_(values=values_clamped)
def push(self,
positions: Tensor,
outward: bool = True,
shift_amount: float = 0) -> Tensor:
loc_to_center = positions - self.center
sgn_dist_from_surface = math.abs(loc_to_center) - self.half_size
if outward:
# --- get negative distances (particles are inside) towards the nearest boundary and add shift_amount ---
distances_of_interest = (sgn_dist_from_surface == math.max(
sgn_dist_from_surface, 'vector')) & (sgn_dist_from_surface < 0)
shift = distances_of_interest * (sgn_dist_from_surface -
shift_amount)
else:
shift = (sgn_dist_from_surface + shift_amount) * (
sgn_dist_from_surface > 0
) # get positive distances (particles are outside) and add shift_amount
shift = math.where(
math.abs(shift) > math.abs(loc_to_center),
math.abs(loc_to_center),
shift) # ensure inward shift ends at center
return positions + math.where(loc_to_center < 0, 1, -1) * shift
def plot(field: SampledField, title=False, colorbar=False, cmap='magma', figsize=(12, 5), same_scale=True, **plt_args):
batch_size = field.shape.batch.volume
values = math.join_dimensions(field.values, field.shape.channel, 'channel').channel[0]
fig, axes = plt.subplots(1, batch_size, figsize=figsize)
axes = axes if isinstance(axes, np.ndarray) else [axes]
b_values = math.join_dimensions(values, field.shape.batch, 'batch')
if title:
for b in range(batch_size):
if isinstance(title, str):
sub_title = title
elif title is True:
sub_title = f"{b} of {field.shape.batch}"
elif isinstance(title, (tuple, list)):
sub_title = title[b]
else:
sub_title = None
if sub_title is not None:
axes[b].set_title(sub_title)
# Individual plots
if isinstance(field, Grid) and field.shape.channel.volume == 1:
if same_scale:
plt_args['vmin'] = math.min(values).native()
plt_args['vmax'] = math.max(values).native()
for b in range(batch_size):
im = axes[b].imshow(b_values.batch[b].numpy('y,x'), origin='lower', cmap=cmap, **plt_args)
if colorbar:
plt.colorbar(im, ax=axes[b])
elif isinstance(field, Grid):
if isinstance(field, StaggeredGrid):
field = field.at_centers()
for b in range(batch_size):
x, y = field.points.vector.unstack_spatial('x,y', to_numpy=True)
data = math.join_dimensions(field.values, field.shape.batch, 'batch').batch[b]
u, v = data.vector.unstack_spatial('x,y', to_numpy=True)
axes[b].quiver(x-u/2, y-v/2, u, v)
else:
raise NotImplementedError(f"No figure recipe for {field}")
plt.tight_layout()
return fig, axes
def data_bounds(field: SampledField):
data = field.points
min_vec = math.min(data, dim=data.shape.spatial.names)
max_vec = math.max(data, dim=data.shape.spatial.names)
return Box(min_vec, max_vec)
def bounding_radius(self):
return math.max(self.size, 'vector') * 1.414214
def bounding_radius(self):
return math.max(self.size, axis=-1, keepdims=True) * 1.414214
def _bounding_box(self):
boxes = [bounding_box(g) for g in self.geometries]
lower = math.min([b.lower for b in boxes], dim='0')
upper = math.max([b.upper for b in boxes], dim='0')
return Box(lower, upper)
