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)