def _mg_solve_forward(divergence, domain, pressure_guess, solvers): fluid_mask = domain.accessible_tensor(extend=1) active_mask = domain.active_tensor(extend=1) if active_mask is not None or fluid_mask is not None: if not np.all([s.supports_continuous_masks for s in solvers[:-1]]): logging.warning( "MultiscaleSolver solver: There are boundary conditions inside the domain but " "not all intermediate solvers support continuous masks") div_lvls = [divergence] act_lvls = [active_mask] fld_lvls = [fluid_mask] for grid_i in range(len(solvers) - 1): div_lvls.insert(0, math.downsample2x(div_lvls[0])) act_lvls.insert( 0, math.downsample2x(act_lvls[0]) if act_lvls[0] is not None else None) fld_lvls.insert( 0, math.downsample2x(fld_lvls[0]) if fld_lvls[0] is not None else None) if pressure_guess is not None: pressure_guess = math.downsample2x(pressure_guess) iter_list = [] for i, div in enumerate(div_lvls): pressure_guess, iteration = solvers[i].solve( div, FluidDomain(act_lvls[i], fld_lvls[i], boundaries), pressure_guess) iter_list.append(iteration) if pressure_guess.shape[1] < divergence.shape[1]: pressure_guess = math.upsample2x( pressure_guess) * 2**math.spatial_rank(divergence) return pressure_guess, iter_list
def downsample2x(grid: Grid) -> GridType: """ Reduces the number of sample points by a factor of 2 in each spatial dimension. The new values are determined via linear interpolation. See Also: `upsample2x()`. Args: grid: `CenteredGrid` or `StaggeredGrid`. Returns: `Grid` of same type as `grid`. """ if isinstance(grid, CenteredGrid): values = math.downsample2x(grid.values, grid.extrapolation) return CenteredGrid(values, bounds=grid.bounds, extrapolation=grid.extrapolation) elif isinstance(grid, StaggeredGrid): values = [] for dim, centered_grid in zip(grid.shape.spatial.names, unstack(grid, 'vector')): odd_discarded = centered_grid.values[{dim: slice(None, None, 2)}] others_interpolated = math.downsample2x( odd_discarded, grid.extrapolation, dims=grid.shape.spatial.without(dim)) values.append(others_interpolated) return StaggeredGrid(math.stack(values, channel('vector')), bounds=grid.bounds, extrapolation=grid.extrapolation) else: raise ValueError(type(grid))
def test_downsample2x(self): meshgrid = math.meshgrid(x=(0, 1, 2, 3), y=(0, -1, -2)) half_size = math.downsample2x(meshgrid, extrapolation.BOUNDARY) math.print(meshgrid, 'Full size') math.print(half_size, 'Half size') math.assert_close(half_size.vector[0], wrap([[0.5, 2.5], [0.5, 2.5]], spatial('y,x'))) math.assert_close(half_size.vector[1], wrap([[-0.5, -0.5], [-2, -2]], spatial('y,x')))
def test_upsample2x(self): meshgrid = math.meshgrid(x=(0, 1, 2, 3), y=(0, -1, -2)) double_size = math.upsample2x(meshgrid, extrapolation.BOUNDARY) same_size = math.downsample2x(double_size) math.print(meshgrid, 'Normal size') math.print(double_size, 'Double size') math.print(same_size, 'Same size') math.assert_close(meshgrid.x[1:-1].y[1:-1], same_size.x[1:-1].y[1:-1])
def downsample2x(grid: Grid) -> GridType: if isinstance(grid, CenteredGrid): values = math.downsample2x(grid.values, grid.extrapolation) return CenteredGrid(values, grid.bounds, grid.extrapolation) elif isinstance(grid, StaggeredGrid): values = [] for dim, centered_grid in zip(grid.shape.spatial.names, grid.unstack()): odd_discarded = centered_grid.values[{dim: slice(None, None, 2)}] others_interpolated = math.downsample2x( odd_discarded, grid.extrapolation, dims=grid.shape.spatial.without(dim)) values.append(others_interpolated) return StaggeredGrid(math.channel_stack(values, 'vector'), grid.bounds, grid.extrapolation) else: raise ValueError(type(grid))
def downsample2x(self): data = [] for axis in range(self.rank): grid = self.unstack()[axis].data grid = grid[tuple([ slice(None, None, 2) if d - 1 == axis else slice(None) for d in range(self.rank + 2) ])] # Discard odd indices along axis grid = math.downsample2x( grid, axes=tuple(filter(lambda ax2: ax2 != axis, range( self.rank)))) # Interpolate values along other axes data.append(grid) return self.with_data(data)