def test_expand_copy_item_names(self):
a = math.zeros(channel(vector=2))
try:
math.expand(a, channel(vector=3))
self.fail()
except IncompatibleShapes:
pass
b = math.expand(a, channel(vector='x,y'))
self.assertEqual(('x', 'y'), b.vector.item_names)
try:
math.expand(b, channel(vector='a,b'))
self.fail()
except IncompatibleShapes:
pass
def vector_grid(self,
value: Field or Tensor or Number or Geometry or callable = 0.,
type: type = CenteredGrid,
extrapolation: math.Extrapolation or str = 'vector') -> CenteredGrid or StaggeredGrid:
"""
Creates a vector grid matching the resolution and bounds of the domain.
The grid is created from the given `value` which must be one of the following:
* Number (int, float, complex or zero-dimensional tensor): all grid values will be equal to `value`. This has a near-zero memory footprint.
* Field: the given value is resampled to the grid cells of this Domain.
* Tensor with spatial dimensions matcing the domain resolution: grid values will equal `value`.
* Geometry: grid values are determined from the volume overlap between grid cells and geometry. Non-overlapping = 0, fully enclosed grid cell = 1.
* function(location: Tensor) returning one of the above.
The returned grid will have a vector dimension with size equal to the rank of the domain.
Args:
value: constant, Field, Tensor or function specifying the grid values
type: class of Grid to create, must be either CenteredGrid or StaggeredGrid
extrapolation: (optional) grid extrapolation, defaults to Domain.boundaries['vector']
Returns:
Grid of specified type
"""
extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation]
result = type(value, resolution=self.resolution, bounds=self.bounds, extrapolation=extrapolation)
if result.shape.channel_rank == 0:
result = result.with_values(math.expand(result.values, channel(vector=self.rank)))
else:
assert result.shape.get_size('vector') == self.rank
return result
def concat(geometries: tuple or list,
dim: str,
sizes: tuple or list or None = None):
"""
Concatenates multiple geometries of the same type.
Args:
geometries: sequence of `phi.geom.Geometry` objects of the same type
sizes: implicit
dim: dimension to concatenate
Returns:
New `phi.geom.Geometry` object
"""
if all(isinstance(g, type(geometries[0])) for g in geometries):
characteristics = [g.__characteristics__() for g in geometries]
if sizes is not None:
characteristics = [{
key: math.expand(val, dim, size)
for key, val in c.items()
} for c, size in zip(characteristics, sizes)]
new_attributes = {}
for key in characteristics[0].keys():
concatenated = math.concat([c[key] for c in characteristics], dim)
new_attributes[key] = concatenated
return geometries[0].__with__(**new_attributes)
else:
raise NotImplementedError()
def __init__(self,
lower: Tensor or float or int = None,
upper: Tensor or float or int = None,
**size: int or Tensor):
"""
Args:
lower: physical location of lower corner
upper: physical location of upper corner
**size: Upper l
"""
if lower is not None:
self._lower = wrap(lower)
if upper is not None:
self._upper = wrap(upper)
else:
lower = []
upper = []
for item in size.values():
if isinstance(item, (tuple, list)):
assert len(
item
) == 2, f"Box kwargs must be either dim=upper or dim=(lower,upper) but got {item}"
lo, up = item
lower.append(lo)
upper.append(up)
elif item is None:
lower.append(-INF)
upper.append(INF)
else:
lower.append(0)
upper.append(item)
lower = [-INF if l is None else l for l in lower]
upper = [INF if u is None else u for u in upper]
self._upper = math.wrap(upper,
math.channel(vector=tuple(size.keys())))
self._lower = math.wrap(lower,
math.channel(vector=tuple(size.keys())))
vector_shape = self._lower.shape & self._upper.shape
self._lower = math.expand(self._lower, vector_shape)
self._upper = math.expand(self._upper, vector_shape)
if self.size.vector.item_names is None:
warnings.warn(
"Creating a Box without item names prevents certain operations like project()",
DeprecationWarning,
stacklevel=2)
def concat(*fields: SampledField, dim: str):
assert all(isinstance(f, SampledField) for f in fields)
assert all(isinstance(f, type(fields[0])) for f in 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.concat([f.values for f in fields], dim=dim)
return fields[0].with_(values=values)
elif isinstance(fields[0], PointCloud):
elements = geom.concat([f.elements for f in fields],
dim,
sizes=[f.shape.get_size(dim) for f in fields])
values = math.concat([
math.expand(f.values, dim, f.shape.get_size(dim)) for f in fields
], dim)
colors = math.concat(
[math.expand(f.color, dim, f.shape.get_size(dim)) for f in fields],
dim)
return fields[0].with_(elements=elements, values=values, color=colors)
raise NotImplementedError(type(fields[0]))
def test_boolean_mask_batched(self):
for backend in BACKENDS:
with backend:
x = math.expand(math.range(spatial('x'), 4), batch(batch=2)) * math.tensor([1, -1])
mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x'))
selected = math.boolean_mask(x, 'x', mask)
expected_0 = math.tensor([(0, -0), (2, -2)], spatial('x'), channel('vector'))
expected_1 = math.tensor([(1, -1)], spatial('x'), channel('vector'))
math.assert_close(selected.batch[0], expected_0, msg=backend.name)
math.assert_close(selected.batch[1], expected_1, msg=backend.name)
math.assert_close(selected, x.x[mask], msg=backend.name)
def test_scatter_1d(self):
for backend in BACKENDS:
with backend:
base = math.ones(spatial(x=4))
indices = math.wrap([1, 2], instance('points'))
values = math.wrap([11, 12], instance('points'))
updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined')
math.assert_close(updated, [1, 11, 12, 1])
updated = math.scatter(base, indices, values, mode='add', outside_handling='undefined')
math.assert_close(updated, [1, 12, 13, 1])
# with vector dim
indices = math.expand(indices, channel(vector=1))
updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined')
math.assert_close(updated, [1, 11, 12, 1])
def concat(fields: List[SampledFieldType], dim: Shape) -> SampledFieldType:
"""
Concatenates the given `SampledField`s along `dim`.
See Also:
`stack()`.
Args:
fields: List of matching `SampledField` instances.
dim: Concatenation dimension as `Shape`. Size is ignored.
Returns:
`SampledField` matching concatenated fields.
"""
assert all(isinstance(f, SampledField) for f in fields)
assert all(isinstance(f, type(fields[0])) for f in 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.concat([f.values for f in fields], dim)
return fields[0].with_values(values)
elif isinstance(fields[0], PointCloud):
elements = geom.concat([f.elements for f in fields],
dim,
sizes=[f.shape.get_size(dim) for f in fields])
values = math.concat(
[math.expand(f.values, f.shape.only(dim)) for f in fields], dim)
colors = math.concat(
[math.expand(f.color, f.shape.only(dim)) for f in fields], 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 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 test_collapsed_non_uniform_tensor(self):
non_uniform = math.stack(
[math.zeros(spatial(a=2)),
math.ones(spatial(a=3))], batch('b'))
e = math.expand(non_uniform, channel('vector'))
assert e.shape.without('vector') == non_uniform.shape
def grad(_x, _y, df):
return math.flatten(math.expand(df * 0, batch(tmp=2))),
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)
