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)