def _weighted_sliced_laplace_nd(tensor, weights):
if tensor.shape[-1] != 1:
raise ValueError('Laplace operator requires a scalar channel as input')
dims = range(math.spatial_rank(tensor))
components = []
for dimension in dims:
lower_weights, center_weights, upper_weights = _dim_shifted(
weights, dimension, (-1, 0, 1), diminish_others=(1, 1))
lower_values, center_values, upper_values = _dim_shifted(
tensor, dimension, (-1, 0, 1), diminish_others=(1, 1))
diff = math.mul(
upper_values, upper_weights * center_weights) + math.mul(
lower_values, lower_weights * center_weights) + math.mul(
center_values, -lower_weights - upper_weights)
components.append(diff)
return math.sum(components, 0)
def sample_at(self, points, collapse_dimensions=True):
if len(self.geometries) == 0:
return _expand_axes(math.zeros([1,1]), points, collapse_dimensions=collapse_dimensions)
if len(self.geometries) == 1:
result = self.geometries[0].value_at(points)
else:
result = math.max([geometry.value_at(points) for geometry in self.geometries], axis=0)
return math.mul(result, self.data)
def sample_at(self, points, collapse_dimensions=True):
if not isinstance(self.extrapolation, six.string_types):
return self._padded_resample(points)
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
resampled = math.resample(self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
interpolation=self.interpolation)
return resampled
def sample_at(self, points):
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
resampled = math.resample(self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
interpolation=self.interpolation,
constant_values=_pad_value(
self.extrapolation_value))
return resampled
def general_sample_at(self, points, reduce):
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
result = general_grid_sample_nd(
self.data,
local_points,
boundary=_pad_mode(self.extrapolation),
constant_values=_pad_value(self.extrapolation_value),
math=math.choose_backend([self.data, points]),
reduce=reduce)
return result
def effect_applied(effect, field, dt):
effect_field = effect.field.at(field)
if effect._mode == GROW:
return field + math.mul(effect_field, dt)
elif effect._mode == ADD:
return field + effect_field
elif effect._mode == FIX:
assert effect.bounds is not None
mask = GeometryMask([effect.bounds]).at(field)
return field * (1 - mask) + effect_field * mask
else:
raise ValueError('Invalid mode: %s' % effect.mode)
def effect_applied(effect, field, dt):
effect_field = effect.field.at(field)
if effect._mode == GROW:
return field + math.mul(effect_field, dt)
elif effect._mode == ADD:
return field + effect_field
elif effect._mode == FIX:
assert effect.bounds is not None
inside = mask([effect.bounds]).at(field)
return math.where(inside, effect_field, field)
else:
raise ValueError('Invalid mode: %s' % effect.mode)
def _weighted_sliced_laplace_nd(tensor, weights):
if tensor.shape[-1] != 1:
raise ValueError('Laplace operator requires a scalar channel as input')
dims = range(math.spatial_rank(tensor))
components = []
for dimension in dims:
center_slices = tuple([(slice(1, -1) if i == dimension else slice(1,-1)) for i in dims])
upper_slices = tuple([(slice(2, None) if i == dimension else slice(1,-1)) for i in dims])
lower_slices = tuple([(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 = math.mul(upper_values, upper_weights) + \
math.mul(lower_values, lower_weights) + \
math.mul(center_values, center_weights)
components.append(diff)
return math.sum(components, 0)
def from_scalar(scalar_field, axis_forces, name=None):
assert isinstance(scalar_field, CenteredGrid)
assert scalar_field.component_count == 1, 'channel must be scalar but has %d components' % scalar_field.component_count
tensors = []
for axis in range(scalar_field.rank):
force = axis_forces[axis] if isinstance(axis_forces, (list, tuple)) else axis_forces[...,axis]
if isinstance(force, Number) and force == 0:
dims = list(math.staticshape(scalar_field.data))
dims[axis+1] += 1
tensors.append(math.zeros(dims, math.dtype(scalar_field.data)))
else:
upper = scalar_field.axis_padded(axis, 0, 1).data
lower = scalar_field.axis_padded(axis, 1, 0).data
tensors.append(math.mul((upper + lower) / 2, force))
return StaggeredGrid(tensors, scalar_field.box, name=name, batch_size=scalar_field._batch_size)
def sample_at(self, points, collapse_dimensions=True):
if not isinstance(self.extrapolation, six.string_types):
return self._padded_resample(points)
local_points = self.box.global_to_local(points)
local_points = math.mul(local_points, math.to_float(
self.resolution)) - 0.5
boundary = {
'periodic': 'circular',
'boundary': 'replicate',
'constant': 'constant'
}[self.extrapolation]
resampled = math.resample(self.data,
local_points,
boundary=boundary,
interpolation=self.interpolation)
return resampled
def __mul__(self, other):
return self.__dataop__(other, True, lambda d1, d2: math.mul(d1, d2))
