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')