def staggered_component_box(resolution, axis, box_like=None):
staggered_box = AABox(0, resolution) if box_like is None else AABox.to_box(
box_like, resolution_hint=resolution)
unit = np.array([
(staggered_box.size[axis] / resolution[axis]) if d == axis else 0
for d in range(len(resolution))
])
box = AABox(staggered_box.lower - unit / 2, staggered_box.upper + unit / 2)
return box
def test_batched_box(self):
mybox = AABox(0,
np.stack(
[np.ones(10), np.linspace(0, 10, 10)], axis=-1))
# 0D indexing
values = mybox.value_at(np.zeros([10, 2]) + [0, 4])
np.testing.assert_equal(values.shape, [10, 1])
np.testing.assert_equal(values[:, 0], [0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
# 1D indexing
values = mybox.value_at(np.zeros([10, 3, 2]) + [0, 4])
np.testing.assert_equal(values.shape, [10, 3, 1])
np.testing.assert_equal(values[:, 0, 0],
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
def padded(self, widths):
data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]],
_pad_mode(self.extrapolation))
w_lower, w_upper = np.transpose(widths)
box = AABox(self.box.lower - w_lower * self.dx,
self.box.upper + w_upper * self.dx)
return self.copied_with(data=data, box=box)
def conv_layer(grid,
filters,
kernel_size,
strides=1,
padding='valid',
activation=None,
use_bias=True,
name=None,
trainable=True,
reuse=None):
assert isinstance(grid, CenteredGrid)
if grid.rank == 2:
result = tf.layers.conv2d(grid.data,
filters,
kernel_size,
strides=strides,
activation=activation,
padding=padding,
name=name,
use_bias=use_bias,
trainable=trainable,
reuse=reuse)
if padding != 'valid':
box = grid.box
else:
w_upper = kernel_size // 2
w_lower = (kernel_size - 1) // 2
box = AABox(grid.box.lower + w_lower * grid.dx,
grid.box.upper - w_upper * grid.dx)
return CenteredGrid(result, box=box, extrapolation=grid.extrapolation)
else:
raise NotImplementedError()
def padded(self, widths):
new_grids = [grid.padded(widths) for grid in self.unstack()]
if isinstance(widths, int):
widths = [[widths, widths]] * self.rank
w_lower, w_upper = np.transpose(widths)
box = AABox(self.box.lower - w_lower * self.dx, self.box.upper + w_upper * self.dx)
return self.copied_with(data=new_grids, box=box)
def padded(self, widths):
extrapolation = self.extrapolation if isinstance(
self.extrapolation, six.string_types) else ['constant'] + list(
self.extrapolation) + ['constant']
data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]],
_pad_mode(extrapolation))
w_lower, w_upper = np.transpose(widths)
box = AABox(self.box.lower - w_lower * self.dx,
self.box.upper + w_upper * self.dx)
return self.copied_with(data=data, box=box)
def padded(self, widths):
if isinstance(widths, int):
widths = [[widths, widths]] * self.rank
data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]],
_pad_mode(self.extrapolation),
constant_values=_pad_value(self.extrapolation_value))
w_lower, w_upper = np.transpose(widths)
box = AABox(self.box.lower - w_lower * self.dx,
self.box.upper + w_upper * self.dx)
return self.copied_with(data=data, box=box)
def data_bounds(field):
assert field.has_points
try:
data = field.points.data
min_vec = math.min(data, axis=tuple(range(len(data.shape) - 1)))
max_vec = math.max(data, axis=tuple(range(len(data.shape) - 1)))
except StaggeredSamplePoints:
boxes = [data_bounds(c) for c in field.unstack()]
min_vec = math.min([b.lower for b in boxes], axis=0)
max_vec = math.max([b.upper for b in boxes], axis=0)
return AABox(min_vec, max_vec)
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain([5, 4], boundaries=CLOSED, box=AABox(0, [40, 32])),
buoyancy_factor=0.1,
batch_size=centers.shape[0]),
physics=IncompressibleFlow(pressure_solver=SparseCG(max_iterations=3)))
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)
print()
return fluid.density.data[0, ...], fluid.velocity.unstack()[0].data[0, ...], fluid.velocity.unstack()[1].data[0, ...]
def padded(self, widths):
extrapolation = self.extrapolation if isinstance(
self.extrapolation, six.string_types) else ['constant'] + list(
self.extrapolation) + ['constant']
data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]],
_pad_mode(extrapolation))
w_lower, w_upper = np.transpose(widths)
box = AABox(self.box.lower - w_lower * self.dx,
self.box.upper + w_upper * self.dx)
return CenteredGrid(data,
box,
extrapolation=self.extrapolation,
name=self.name,
batch_size=self._batch_size)
def box(self, box):
return AABox.to_box(box, resolution_hint=self.resolution)
def box(self, box):
box = AABox.to_box(box, resolution_hint=self.resolution)
assert_same_rank(len(self.data), self.box,
'StaggeredGrid.data does not match box.')
return box
