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 _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 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 unstack(field: Field, dim: str) -> tuple:
"""
Unstack `field` along one of its dimensions.
The dimension can be batch, spatial or channel.
Args:
field: `Field` to unstack.
dim: name of the dimension to unstack, must be part of `self.shape`
Returns:
`tuple` of `Fields`. The returned fields may be of different types than `field`.
"""
size = field.shape.get_size(dim)
if isinstance(size, Tensor):
size = math.min(size) # unstack StaggeredGrid along x or y
return tuple([field[{dim: i}] for i in range(size)])
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 spatial_gradient at sphere center
distance = math.sqrt(distance_squared)
return math.min(distance - self.radius,
self.shape.instance) # union for instance dimensions
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
distance = math.max(distance, 'vector')
distance = math.min(
distance, self.shape.instance) # union for instance dimensions
return distance
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 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_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)
def approximate_signed_distance(self, location):
return math.min([
geometry.approximate_signed_distance(location)
for geometry in self.geometries
],
dim='0')