def test_box_product(self):
a = Box(x=4)
b = Box(y=3).shifted(math.wrap(1))
ab = a * b
self.assertEqual(2, ab.spatial_rank)
math.assert_close(ab.size, (4, 3))
math.assert_close(ab.lower, (0, 1))
def test_box_constructor_kwargs(self):
b = Box(x=3.5, y=4)
math.assert_close(b.lower, 0)
math.assert_close(b.upper, (3.5, 4))
b = Box(x=(1, 2), y=None)
math.assert_close(b.lower, (1, -math.INF))
math.assert_close(b.upper, (2, math.INF))
b = Box(x=(None, None))
math.assert_close(b.lower, -math.INF)
math.assert_close(b.upper, math.INF)
def test_plot_scalar_2d_batch(self):
self._test_plot(
CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])))
self._test_plot(CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])),
row_dims='b',
size=(2, 4))
def tensor_as_field(t: Tensor):
"""
Interpret a `Tensor` as a `CenteredGrid` or `PointCloud` depending on its dimensions.
Unlike the `CenteredGrid` constructor, this function will have the values sampled at integer points for each spatial dimension.
Args:
t: `Tensor` with either `spatial` or `instance` dimensions.
Returns:
`CenteredGrid` or `PointCloud`
"""
if spatial(t):
assert not instance(
t
), f"Cannot interpret tensor as Field because it has both spatial and instance dimensions: {t.shape}"
bounds = Box(-0.5, math.wrap(spatial(t), channel('vector')) - 0.5)
return CenteredGrid(t, 0, bounds=bounds)
if instance(t):
assert not spatial(
t
), f"Cannot interpret tensor as Field because it has both spatial and instance dimensions: {t.shape}"
assert 'vector' in t.shape, f"Cannot interpret tensor as PointCloud because it has not vector dimension."
point_count = instance(t).volume
bounds = data_bounds(t)
radius = math.vec_length(
bounds.size) / (1 + point_count**(1 / t.vector.size))
return PointCloud(Sphere(t, radius=radius))
def test_plot_vector_grid_2d(self):
self._test_plot(
CenteredGrid(Noise(vector=2),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1])) * 0.1)
def __init__(self,
resolution: math.Shape = math.EMPTY_SHAPE,
boundaries: dict or tuple or list = OPEN,
bounds: Box = None,
**resolution_):
"""
The Domain specifies the grid resolution, physical size and boundary conditions of a simulation.
It provides convenience methods for creating Grids fitting the domain, e.g. `grid()`, `vector_grid()` and `staggered_grid()`.
Also see the `phi.physics` module documentation at https://tum-pbs.github.io/PhiFlow/Physics.html
Args:
resolution: grid dimensions as Shape or sequence of integers. Alternatively, dimensions can be specified directly as kwargs.
boundaries: specifies the extrapolation modes of grids created from this Domain.
Default materials include OPEN, CLOSED, PERIODIC.
To specify boundary conditions per face of the domain, pass a sequence of boundaries or boundary pairs (lower, upper)., e.g. [CLOSED, (CLOSED, OPEN)].
See https://tum-pbs.github.io/PhiFlow/Physics.html#boundary-conditions .
bounds: physical size of the domain. If not provided, the size is equal to the resolution (unit cubes).
"""
self.resolution: math.Shape = spatial_shape(
resolution) & spatial_shape(resolution_)
""" Grid dimensions as `Shape` object containing spatial dimensions only. """
self.boundaries: dict = _create_boundary_conditions(
boundaries, self.resolution.names)
""" Outer boundary conditions. """
self.bounds: Box = Box(0, math.wrap(
self.resolution, names='vector')) if bounds is None else bounds
""" Physical dimensions of the domain. """
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 test_plot_multi_1d(self):
self._test_plot(
CenteredGrid(
lambda x: math.stack({
'sin': math.sin(x),
'cos': math.cos(x)
}, channel('curves')),
x=100,
bounds=Box(x=2 * math.pi)))
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 test_plot_multiple(self):
grid = CenteredGrid(Noise(batch(b=2)), 0, Box[0:1, 0:1], x=50, y=10)
grid2 = CenteredGrid(grid, 0, Box[0:2, 0:1], x=20, y=50)
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=0.1), bounds=Box(0, [1, 1]))
titles = math.wrap([['b=0', 'b=0', 'points'], ['b=1', 'b=1', '']],
spatial('rows,cols'))
self._test_plot(grid, grid2, cloud, row_dims='b', title=titles)
def vector_potential(self,
value: Field or Tensor or Number or Geometry or callable = 0.,
extrapolation: str or math.Extrapolation = 'scalar',
curl_type=CenteredGrid):
if self.rank == 2 and curl_type == StaggeredGrid:
pot_bounds = Box(self.bounds.lower - 0.5 * self.dx, self.bounds.upper + 0.5 * self.dx)
alt_domain = Domain(self.boundaries, self.resolution + 1, bounds=pot_bounds)
return alt_domain.scalar_grid(value, extrapolation=extrapolation)
raise NotImplementedError()
def __init__(self, elements: Geometry, values: Tensor, extrapolation: float
or math.Extrapolation, resolution: Shape, bounds: Box):
if bounds.size.vector.item_names is None:
with NUMPY:
bounds = bounds.shifted(
math.zeros(channel(vector=spatial(values).names)))
SampledField.__init__(self, elements, values, extrapolation, bounds)
assert values.shape.spatial_rank == elements.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
assert values.shape.spatial_rank == bounds.spatial_rank, f"Spatial dimensions of values ({values.shape}) do not match elements {elements}"
assert values.shape.instance_rank == 0, f"Instance dimensions not supported for grids. Got values with shape {values.shape}"
self._resolution = resolution
def test_overlay(self):
grid = CenteredGrid(Noise(),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1]))
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=.1))
self._test_plot(overlay(grid, grid * (0.1, 0.02), cloud),
title='Overlay')
def test_plot_vector_3d_batched(self):
sphere = CenteredGrid(SoftGeometryMask(
Sphere(x=.5, y=.5, z=.5, radius=.4)),
x=10,
y=10,
z=10,
bounds=Box(x=1, y=1, z=1)) * (.1, 0, 0)
cylinder = CenteredGrid(geom.infinite_cylinder(
x=16, y=16, inf_dim='z', radius=10),
x=32,
y=32,
z=32) * (0, 0, .1)
self._test_plot(sphere, cylinder)
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain(x=5, y=4, bounds=Box(0, [40, 32])),
buoyancy_factor=0.1,
batch_size=centers.shape[0]),
physics=IncompressibleFlow())
world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2))
world.add(Fan(Sphere(center=centers, radius=5), acceleration=[1.0, 0]))
world.step(dt=1.5)
world.step(dt=1.5)
world.step(dt=1.5)
assert not math.close(fluid.density.values, 0)
print()
return fluid.density.values.batch[0], fluid.velocity.values.batch[0]
def expand_staggered(values: Tensor, resolution: Shape,
extrapolation: math.Extrapolation):
""" Add missing spatial dimensions to `values` """
cells = GridCell(
resolution,
Box(
0,
math.wrap((1, ) * resolution.rank,
channel(vector=resolution.names))))
components = values.vector.unstack(resolution.spatial_rank)
tensors = []
for dim, component in zip(resolution.spatial.names, components):
comp_cells = cells.stagger(dim, *extrapolation.valid_outer_faces(dim))
tensors.append(math.expand(component, comp_cells.resolution))
return math.stack(tensors, channel(vector=resolution.names))
def stack(fields, dim: Shape, dim_bounds: Box = None):
"""
Stacks the given `SampledField`s along `dim`.
See Also:
`concat()`.
Args:
fields: List of matching `SampledField` instances.
dim: Stack dimension as `Shape`. Size is ignored.
Returns:
`SampledField` matching stacked fields.
"""
assert all(
isinstance(f, SampledField) for f in fields
), f"All fields must be SampledFields of the same type but got {fields}"
assert all(
isinstance(f, type(fields[0])) for f in fields
), f"All fields must be SampledFields of the same type but got {fields}"
if any(f.extrapolation != fields[0].extrapolation for f in fields):
raise NotImplementedError("Concatenating extrapolations not supported")
if isinstance(fields[0], Grid):
values = math.stack([f.values for f in fields], dim)
if spatial(dim):
if dim_bounds is None:
dim_bounds = Box(**{dim.name: len(fields)})
return type(fields[0])(values,
extrapolation=fields[0].extrapolation,
bounds=fields[0].bounds * dim_bounds)
else:
return fields[0].with_values(values)
elif isinstance(fields[0], PointCloud):
elements = geom.stack([f.elements for f in fields], dim=dim)
values = math.stack([f.values for f in fields], dim=dim)
colors = math.stack([f.color for f in fields], dim=dim)
return PointCloud(elements=elements,
values=values,
color=colors,
extrapolation=fields[0].extrapolation,
add_overlapping=fields[0]._add_overlapping,
bounds=fields[0]._bounds)
raise NotImplementedError(type(fields[0]))
def grid_scatter(self, bounds: Box, resolution: math.Shape, outside_handling: str):
"""
Approximately samples this field on a regular grid using math.scatter().
Args:
outside_handling: `str` passed to `phi.math.scatter()`.
bounds: physical dimensions of the grid
resolution: grid resolution
Returns:
CenteredGrid
"""
closest_index = bounds.global_to_local(self.points) * resolution - 0.5
mode = 'add' if self._add_overlapping else 'mean'
base = math.zeros(resolution)
if isinstance(self.extrapolation, math.extrapolation.ConstantExtrapolation):
base += self.extrapolation.value
scattered = math.scatter(base, closest_index, self.values, mode=mode, outside_handling=outside_handling)
return scattered
def _grid_scatter(self, box: Box, resolution: math.Shape):
"""
Approximately samples this field on a regular grid using math.scatter().
Args:
box: physical dimensions of the grid
resolution: grid resolution
box: Box:
resolution: math.Shape:
Returns:
CenteredGrid
"""
closest_index = math.to_int(math.round(box.global_to_local(self.points) * resolution - 0.5))
if self._add_overlapping:
duplicates_handling = 'add'
else:
duplicates_handling = 'mean'
scattered = math.scatter(closest_index, self.values, resolution, duplicates_handling=duplicates_handling, outside_handling='discard', scatter_dims=('points',))
return scattered
def pad(grid: Grid, widths: int or tuple or list or dict):
if isinstance(widths, int):
widths = {axis: (widths, widths) for axis in grid.shape.spatial.names}
elif isinstance(widths, (tuple, list)):
widths = {
axis: (width if isinstance(width, (tuple, list)) else
(width, width))
for axis, width in zip(grid.shape.spatial.names, widths)
}
else:
assert isinstance(widths, dict)
widths_list = [widths[axis] for axis in grid.shape.spatial.names]
if isinstance(grid, Grid):
data = math.pad(grid.values, widths, grid.extrapolation)
w_lower = math.wrap([w[0] for w in widths_list])
w_upper = math.wrap([w[1] for w in widths_list])
box = Box(grid.box.lower - w_lower * grid.dx,
grid.box.upper + w_upper * grid.dx)
return type(grid)(data, box, grid.extrapolation)
raise NotImplementedError(
f"{type(grid)} not supported. Only Grid instances allowed.")
def pad(grid: GridType, widths: int or tuple or list or dict) -> GridType:
"""
Pads a `Grid` using its extrapolation.
Unlike `phi.math.pad()`, this function also affects the `bounds` of the grid, changing its size and origin depending on `widths`.
Args:
grid: `CenteredGrid` or `StaggeredGrid`
widths: Either `int` or `(lower, upper)` to pad the same number of cells in all spatial dimensions
or `dict` mapping dimension names to `(lower, upper)`.
Returns:
`Grid` of the same type as `grid`
"""
if isinstance(widths, int):
widths = {axis: (widths, widths) for axis in grid.shape.spatial.names}
elif isinstance(widths, (tuple, list)):
widths = {
axis: (width if isinstance(width, (tuple, list)) else
(width, width))
for axis, width in zip(grid.shape.spatial.names, widths)
}
else:
assert isinstance(widths, dict)
widths_list = [widths[axis] for axis in grid.shape.spatial.names]
if isinstance(grid, Grid):
data = math.pad(grid.values, widths, grid.extrapolation)
w_lower = math.wrap([w[0] for w in widths_list])
w_upper = math.wrap([w[1] for w in widths_list])
bounds = Box(grid.box.lower - w_lower * grid.dx,
grid.box.upper + w_upper * grid.dx)
return type(grid)(values=data,
resolution=data.shape.spatial,
bounds=bounds,
extrapolation=grid.extrapolation)
raise NotImplementedError(
f"{type(grid)} not supported. Only Grid instances allowed.")
def test_box_constructor(self):
box = Box(0, (1, 1))
math.assert_close(box.size, 1)
def test_box_eq(self):
self.assertNotEqual(Box(x=1, y=1), Box(x=1))
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 test_box_volume(self):
box = Box(
math.tensor([(0, 0), (1, 1)], batch('boxes'), channel('vector')),
1)
math.assert_close(box.volume, [1, 0])
def test_box_batched(self):
box = Box(
math.tensor([(0, 0), (1, 1)], batch('boxes'), channel('vector')),
1)
self.assertEqual(math.batch(boxes=2), box.shape)
def test_create_grid_int_resolution(self):
g = CenteredGrid(0, 0, Box(x=4, y=3), resolution=10)
self.assertEqual(g.shape, spatial(x=10, y=10))
g = StaggeredGrid(0, 0, Box(x=4, y=3), resolution=10)
self.assertEqual(spatial(g), spatial(x=10, y=10))
def __init__(self,
values: Any,
extrapolation: Any = 0.,
bounds: Box = None,
resolution: int or Shape = None,
**resolution_: int or Tensor):
"""
Args:
values: Values to use for the grid.
Has to be one of the following:
* `phi.geom.Geometry`: sets inside values to 1, outside to 0
* `Field`: resamples the Field to the staggered sample points
* `Number`: uses the value for all sample points
* `tuple` or `list`: interprets the sequence as vector, used for all sample points
* `phi.math.Tensor` compatible with grid dims: uses tensor values as grid values
* Function `values(x)` where `x` is a `phi.math.Tensor` representing the physical location.
The spatial dimensions of the grid will be passed as batch dimensions to the function.
extrapolation: The grid extrapolation determines the value outside the `values` tensor.
Allowed types: `float`, `phi.math.Tensor`, `phi.math.extrapolation.Extrapolation`.
bounds: Physical size and location of the grid as `phi.geom.Box`.
resolution: Grid resolution as purely spatial `phi.math.Shape`.
**resolution_: Spatial dimensions as keyword arguments. Typically either `resolution` or `spatial_dims` are specified.
"""
if resolution is None and not resolution_:
assert isinstance(
values, math.Tensor
), "Grid resolution must be specified when 'values' is not a Tensor."
resolution = values.shape.spatial
bounds = bounds or Box(0, math.wrap(resolution, channel('vector')))
elements = GridCell(resolution, bounds)
else:
if isinstance(resolution, int):
assert not resolution_, "Cannot specify keyword resolution and integer resolution at the same time."
resolution = spatial(
**{
dim: resolution
for dim in bounds.size.shape.get_item_names('vector')
})
resolution = (resolution
or math.EMPTY_SHAPE) & spatial(**resolution_)
bounds = bounds or Box(0, math.wrap(resolution, channel('vector')))
elements = GridCell(resolution, bounds)
if isinstance(values, math.Tensor):
values = math.expand(values, resolution)
elif isinstance(values, Geometry):
values = reduce_sample(HardGeometryMask(values), elements)
elif isinstance(values, Field):
values = reduce_sample(values, elements)
elif callable(values):
values = math.map_s2b(values)(elements.center)
assert isinstance(
values, math.Tensor
), f"values function must return a Tensor but returned {type(values)}"
else:
if isinstance(
values,
(tuple, list)) and len(values) == resolution.rank:
values = math.tensor(values,
channel(vector=resolution.names))
values = math.expand(math.tensor(values), resolution)
if values.dtype.kind not in (float, complex):
values = math.to_float(values)
assert resolution.spatial_rank == bounds.spatial_rank, f"Resolution {resolution} does not match bounds {bounds}"
Grid.__init__(self, elements, values, extrapolation,
values.shape.spatial, bounds)
def __init__(self,
values: Any,
extrapolation: float or math.Extrapolation = 0,
bounds: Box = None,
resolution: Shape = None,
**resolution_: int or Tensor):
"""
Args:
values: Values to use for the grid.
Has to be one of the following:
* `phi.geom.Geometry`: sets inside values to 1, outside to 0
* `Field`: resamples the Field to the staggered sample points
* `Number`: uses the value for all sample points
* `tuple` or `list`: interprets the sequence as vector, used for all sample points
* `phi.math.Tensor` with staggered shape: uses tensor values as grid values.
Must contain a `vector` dimension with each slice consisting of one more element along the dimension they describe.
Use `phi.math.stack()` to manually create this non-uniform tensor.
* Function `values(x)` where `x` is a `phi.math.Tensor` representing the physical location.
The spatial dimensions of the grid will be passed as batch dimensions to the function.
extrapolation: The grid extrapolation determines the value outside the `values` tensor.
Allowed types: `float`, `phi.math.Tensor`, `phi.math.extrapolation.Extrapolation`.
bounds: Physical size and location of the grid.
resolution: Grid resolution as purely spatial `phi.math.Shape`.
**resolution_: Spatial dimensions as keyword arguments. Typically either `resolution` or `spatial_dims` are specified.
"""
if not isinstance(extrapolation, math.Extrapolation):
extrapolation = math.extrapolation.ConstantExtrapolation(
extrapolation)
if resolution is None and not resolution_:
assert isinstance(
values, Tensor
), "Grid resolution must be specified when 'values' is not a Tensor."
any_dim = values.shape.spatial.names[0]
x = values.vector[any_dim]
ext_lower, ext_upper = extrapolation.valid_outer_faces(any_dim)
delta = int(ext_lower) + int(ext_upper) - 1
resolution = x.shape.spatial._replace_single_size(
any_dim,
x.shape.get_size(any_dim) - delta)
bounds = bounds or Box(0, math.wrap(resolution, channel('vector')))
elements = staggered_elements(resolution, bounds, extrapolation)
else:
if isinstance(resolution, int):
assert not resolution_, "Cannot specify keyword resolution and integer resolution at the same time."
resolution = spatial(
**{
dim: resolution
for dim in bounds.size.shape.get_item_names('vector')
})
resolution = (resolution
or math.EMPTY_SHAPE) & spatial(**resolution_)
bounds = bounds or Box(0, math.wrap(resolution, channel('vector')))
elements = staggered_elements(resolution, bounds, extrapolation)
if isinstance(values, math.Tensor):
values = expand_staggered(values, resolution, extrapolation)
elif isinstance(values, Geometry):
values = reduce_sample(HardGeometryMask(values), elements)
elif isinstance(values, Field):
values = reduce_sample(values, elements)
elif callable(values):
values = math.map_s2b(values)(elements.center)
if elements.shape.shape.rank > 1: # Different number of X and Y faces
assert isinstance(
values, TensorStack
), f"values function must return a staggered Tensor but returned {type(values)}"
assert 'staggered_direction' in values.shape
if 'vector' in values.shape:
values = math.stack([
values.staggered_direction[i].vector[i]
for i in range(resolution.rank)
], channel(vector=resolution.names))
else:
values = values.staggered_direction.as_channel('vector')
else:
values = expand_staggered(math.tensor(values), resolution,
extrapolation)
if values.dtype.kind not in (float, complex):
values = math.to_float(values)
assert resolution.spatial_rank == bounds.spatial_rank, f"Resolution {resolution} does not match bounds {bounds}"
Grid.__init__(self, elements, values, extrapolation, resolution,
bounds)
def curl(field: Grid, type: type = CenteredGrid):
""" Computes the finite-difference curl of the give 2D `StaggeredGrid`. """
assert field.spatial_rank in (
2, 3), "curl is only defined in 2 and 3 spatial dimensions."
if isinstance(field, CenteredGrid) and field.spatial_rank == 2:
if 'vector' not in field.shape and type == StaggeredGrid:
# 2D curl of scalar field
grad = math.spatial_gradient(field.values,
dx=field.dx,
difference='forward',
padding=None,
stack_dim=channel('vector'))
result = grad.vector.flip() * (1, -1) # (d/dy, -d/dx)
bounds = Box(field.bounds.lower + 0.5 * field.dx,
field.bounds.upper -
0.5 * field.dx) # lose 1 cell per dimension
return StaggeredGrid(
result,
bounds=bounds,
extrapolation=field.extrapolation.spatial_gradient())
if 'vector' in field.shape and type == CenteredGrid:
# 2D curl of vector field
x, y = field.shape.spatial.names
vy_dx = math.spatial_gradient(field.values.vector[1],
dx=field.dx.vector[0],
padding=field.extrapolation,
dims=x,
stack_dim=None)
vx_dy = math.spatial_gradient(field.values.vector[0],
dx=field.dx.vector[1],
padding=field.extrapolation,
dims=y,
stack_dim=None)
c = vy_dx - vx_dy
return field.with_values(c)
elif isinstance(field, StaggeredGrid) and field.spatial_rank == 2:
if type == CenteredGrid:
for dim in field.resolution.names:
l, u = field.extrapolation.valid_outer_faces(dim)
assert l == u, "periodic extrapolation not yet supported"
values = bake_extrapolation(field).values
x_padded = math.pad(values.vector['x'], {'y': (1, 1)},
field.extrapolation)
y_padded = math.pad(values.vector['y'], {'x': (1, 1)},
field.extrapolation)
vx_dy = math.spatial_gradient(x_padded,
field.dx,
'forward',
None,
dims='y',
stack_dim=None)
vy_dx = math.spatial_gradient(y_padded,
field.dx,
'forward',
None,
dims='x',
stack_dim=None)
result = vx_dy - vy_dx
return CenteredGrid(result,
field.extrapolation.spatial_gradient(),
bounds=field.bounds)
raise NotImplementedError()
