def __init__(self):
TFModel.__init__(self, "Burger CFE", model_scope_name='CFE', data_validation_fraction=0.1)
sim = self.sim = TFFluidSimulation([32], 'open', 100)
u_true = [sim.placeholder(name="u%d" % i) for i in range(n+1)]
u_reconstructed = self.u_reconstructed = [u_true[i] if i == 0 or i == n else None for i in range(n+1)]
for i in range(1,n):
with self.model_scope():
u_reconstructed[i] = sm_resnet(u_reconstructed[i-1], u_true[-1], np.stack([(i-1)/32.0]*sim.batch_size).astype(np.float32))
force_losses = []
forces = []
for i in range(n):
force = u_reconstructed[i+1] - sim_step(u_reconstructed[i])
force_losses.append(l2_loss(force))
forces.append(l1_loss(force, reduce_batches=False))
force_loss = math.add(force_losses) * self.editable_float("Force_Loss_Scale", 1e-1)
self.minimizer("DiffPhys", force_loss)
self.total_force = math.add(forces)
for i in range(n+1):
field = BatchSelect(lambda len, i=i: range(i, len-n+i), "u")
self.database.add("u%d"%i, augment.AxisFlip(2, field))
self.database.put_scenes(scenes("~/data/control/forced-burgerclash"), range(33), logf=self.info)
self.finalize_setup(u_true)
self.frame = EditableInt("Frame", 0, (0,n))
self.add_field("Ground Truth", lambda: self.view_batch("u%d"%self.frame))
self.add_field("Prediction", lambda: self.view(u_reconstructed[self.frame]))
def _weighted_sliced_laplace_nd(tensor, weights):
if tensor.shape[-1] != 1:
raise ValueError("Laplace operator requires a scalar field as input")
dims = range(spatial_rank(tensor))
components = []
for dimension in dims:
center_slices = [(slice(1, -1) if i == dimension else slice(1, -1))
for i in dims]
upper_slices = [(slice(2, None) if i == dimension else slice(1, -1))
for i in dims]
lower_slices = [(slice(-2) if i == dimension else slice(1, -1))
for i in dims]
lower_weights = weights[[slice(None)] + lower_slices + [slice(
None)]] * weights[[slice(None)] + center_slices + [slice(None)]]
upper_weights = weights[[slice(None)] + upper_slices + [slice(
None)]] * weights[[slice(None)] + center_slices + [slice(None)]]
center_weights = -lower_weights - upper_weights
lower_values = tensor[[slice(None)] + lower_slices + [slice(None)]]
upper_values = tensor[[slice(None)] + upper_slices + [slice(None)]]
center_values = tensor[[slice(None)] + center_slices + [slice(None)]]
diff = upper_values * upper_weights + lower_values * lower_weights + center_values * center_weights
components.append(diff)
return math.add(components)
def sample(value, domain, batch_size=None, name=None):
assert isinstance(domain, Domain)
if isinstance(value, Field):
assert_same_rank(
value.rank, domain.rank,
'rank of value (%s) does not match domain (%s)' %
(value.rank, domain.rank))
if isinstance(value,
CenteredGrid) and value.box == domain.box and np.all(
value.resolution == domain.resolution):
data = value.data
else:
point_field = CenteredGrid.getpoints(domain.box,
domain.resolution)
point_field._batch_size = batch_size
data = value.at(point_field).data
else: # value is constant
if callable(value):
x = CenteredGrid.getpoints(
domain.box, domain.resolution).copied_with(
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
value = value(x)
return value
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.add(
math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )), value)
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries),
name=name)
def total(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
total = math.sum(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
total = math.expand_dims(total, -1)
return total
def normalize(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
global_mean = math.mean(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
global_mean = math.expand_dims(global_mean, -1)
return StaggeredGrid(self.staggered / global_mean)
def divergence(self, physical_units=True):
components = []
for dim, field in enumerate(self.data):
grad = math.axis_gradient(field.data, dim)
if physical_units:
grad /= self.dx[dim]
components.append(grad)
data = math.add(components)
return CenteredGrid(data,
self.box,
name='div(%s)' % self.name,
batch_size=self._batch_size)
def divergence(self):
dims = range(self.spatial_rank)
components = []
for dimension in dims:
comp = self.spatial_rank - dimension - 1
upper_slices = [(slice(1, None) if i == dimension else slice(-1))
for i in dims]
lower_slices = [(slice(-1) if i == dimension else slice(-1))
for i in dims]
diff = self.staggered[[slice(None)] + upper_slices + [comp]] - \
self.staggered[[slice(None)] + lower_slices + [comp]]
components.append(diff)
return math.expand_dims(math.add(components), -1)
def _central_divergence_nd(tensor):
rank = spatial_rank(tensor)
dims = range(rank)
components = []
tensor = math.pad(tensor, [[0, 0]] + [[1, 1]] * rank + [[0, 0]])
for dimension in dims:
upper_slices = [(slice(2, None) if i == dimension else slice(1, -1))
for i in dims]
lower_slices = [(slice(-2) if i == dimension else slice(1, -1))
for i in dims]
diff = tensor[[slice(None)] + upper_slices + [rank - dimension - 1]] - \
tensor[[slice(None)] + lower_slices + [rank - dimension - 1]]
components.append(diff)
return math.expand_dims(math.add(components), -1)
def _sliced_laplace_nd(tensor):
# Laplace code for n dimensions
dims = range(spatial_rank(tensor))
components = []
for dimension in dims:
center_slices = [(slice(1, -1) if i == dimension else slice(1, -1))
for i in dims]
upper_slices = [(slice(2, None) if i == dimension else slice(1, -1))
for i in dims]
lower_slices = [(slice(-2) if i == dimension else slice(1, -1))
for i in dims]
diff = tensor[[slice(None)] + upper_slices + [slice(None)]] \
+ tensor[[slice(None)] + lower_slices + [slice(None)]] \
- 2 * tensor[[slice(None)] + center_slices + [slice(None)]]
components.append(diff)
return math.add(components)
def _forward_divergence_nd(field):
rank = spatial_rank(field)
dims = range(rank)
components = []
for dimension in dims:
vq = field[..., rank - dimension - 1]
upper_slices = [(slice(1, None) if i == dimension else slice(None))
for i in dims]
lower_slices = [(slice(-1) if i == dimension else slice(None))
for i in dims]
diff = vq[[slice(None)] + upper_slices] - vq[[slice(None)] +
lower_slices]
padded = math.pad(diff,
[[0, 0]] + [([0, 1] if i == dimension else [0, 0])
for i in dims])
components.append(padded)
return math.expand_dims(math.add(components), -1)
def __add__(self, other):
return self.__dataop__(other, False, lambda d1, d2: math.add(d1, d2))
