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 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 l1_loss(tensor, batch_norm=True, reduce_batches=True): if isinstance(tensor, StaggeredGrid): tensor = tensor.staggered if reduce_batches: total_loss = math.sum(math.abs(tensor)) else: total_loss = math.sum(math.abs(tensor), axis=list(range(1, len(tensor.shape)))) if batch_norm and reduce_batches: batch_size = math.shape(tensor)[0] return total_loss / math.to_float(batch_size) else: return total_loss
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 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 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 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 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 extrapolate(input_field, valid_mask, voxel_distance=10): """ Create a signed distance field for the grid, where negative signs are fluid cells and positive signs are empty cells. The fluid surface is located at the points where the interpolated value is zero. Then extrapolate the input field into the air cells. :param domain: Domain that can create new Fields :param input_field: Field to be extrapolated :param valid_mask: One dimensional binary mask indicating where fluid is present :param voxel_distance: Optional maximal distance (in number of grid cells) where signed distance should still be calculated / how far should be extrapolated. :return: ext_field: a new Field with extrapolated values, s_distance: tensor containing signed distance field, depending only on the valid_mask """ ext_data = input_field.data dx = input_field.dx if isinstance(input_field, StaggeredGrid): ext_data = input_field.staggered_tensor() valid_mask = math.pad(valid_mask, [[0, 0]] + [[0, 1]] * input_field.rank + [[0, 0]], "constant") dims = range(input_field.rank) # Larger than voxel_distance to be safe. It could start extrapolating velocities from outside voxel_distance into the field. signs = -1 * (2 * valid_mask - 1) s_distance = 2.0 * (voxel_distance + 1) * signs surface_mask = create_surface_mask(valid_mask) # surface_mask == 1 doesn't output a tensor, just a scalar, but >= works. # Initialize the voxel_distance with 0 at the surface # Previously initialized with -0.5*dx, i.e. the cell is completely full (center is 0.5*dx inside the fluid surface). For stability and looks this was changed to 0 * dx, i.e. the cell is only half full. This way small changes to the SDF won't directly change neighbouring empty cells to fluid cells. s_distance = math.where((surface_mask >= 1), -0.0 * math.ones_like(s_distance), s_distance) directions = np.array( list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1))))) # First make a move in every positive direction (StaggeredGrid velocities there are correct, we want to extrapolate these) if isinstance(input_field, StaggeredGrid): for d in directions: if (d <= 0).all(): continue # Shift the field in direction d, compare new distances to old ones. d_slice = tuple([(slice(1, None) if d[i] == -1 else slice(0, -1) if d[i] == 1 else slice(None)) for i in dims]) d_field = math.pad( ext_data, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_field = d_field[(slice(None), ) + d_slice + (slice(None), )] d_dist = math.pad( s_distance, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )] d_dist += np.sqrt((dx * d).dot(dx * d)) * signs if (d.dot(d) == 1) and (d >= 0).all(): # Pure axis direction (1,0,0), (0,1,0), (0,0,1) updates = (math.abs(d_dist) < math.abs(s_distance)) & (surface_mask <= 0) updates_velocity = updates & (signs > 0) ext_data = math.where( math.concat([(math.zeros_like(updates_velocity) if d[i] == 1 else updates_velocity) for i in dims], axis=-1), d_field, ext_data) s_distance = math.where(updates, d_dist, s_distance) else: # Mixed axis direction (1,1,0), (1,1,-1), etc. continue for _ in range(voxel_distance): buffered_distance = 1.0 * s_distance # Create a copy of current voxel_distance. This should not be necessary... for d in directions: if (d == 0).all(): continue # Shift the field in direction d, compare new distances to old ones. d_slice = tuple([(slice(1, None) if d[i] == -1 else slice(0, -1) if d[i] == 1 else slice(None)) for i in dims]) d_field = math.pad( ext_data, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_field = d_field[(slice(None), ) + d_slice + (slice(None), )] d_dist = math.pad( s_distance, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )] d_dist += np.sqrt((dx * d).dot(dx * d)) * signs # We only want to update velocity that is outside of fluid updates = (math.abs(d_dist) < math.abs(buffered_distance)) & (surface_mask <= 0) updates_velocity = updates & (signs > 0) ext_data = math.where( math.concat([updates_velocity] * math.spatial_rank(ext_data), axis=-1), d_field, ext_data) buffered_distance = math.where(updates, d_dist, buffered_distance) s_distance = buffered_distance # Cut off inaccurate values distance_limit = -voxel_distance * (2 * valid_mask - 1) s_distance = math.where( math.abs(s_distance) < voxel_distance, s_distance, distance_limit) if isinstance(input_field, StaggeredGrid): ext_field = input_field.with_data(ext_data) stagger_slice = tuple([slice(0, -1) for i in dims]) s_distance = s_distance[(slice(None), ) + stagger_slice + (slice(None), )] else: ext_field = input_field.copied_with(data=ext_data) return ext_field, s_distance
def abs(self): return StaggeredGrid(math.abs(self.staggered))
def abs(self): return self.with_data(math.abs(self.data))