def distribute_points(density, particles_per_cell=1, distribution='uniform'):
"""
Distribute points according to the distribution specified in density.
:param density: binary tensor
:param particles_per_cell: integer
:param distribution: 'uniform' or 'center'
:return: tensor of shape (batch_size, point_count, rank)
"""
assert distribution in ('center', 'uniform')
index_array = []
batch_size = math.staticshape(density)[0] if math.staticshape(density)[0] is not None else 1
for batch in range(batch_size):
indices = math.where(density[batch, ..., 0] > 0)
indices = math.to_float(indices)
temp = []
for _ in range(particles_per_cell):
if distribution == 'center':
temp.append(indices + 0.5)
elif distribution == 'uniform':
temp.append(indices + math.random_uniform(math.shape(indices)))
index_array.append(math.concat(temp, axis=0))
try:
index_array = math.stack(index_array)
return index_array
except ValueError:
raise ValueError("all arrays in the batch must have the same number of active cells.")
def f(value):
if isinstance(value, field.StaggeredGrid):
shape = math.staticshape(value.staggered_tensor())
return (value._batch_size, ) + shape[1:]
if isinstance(value, field.CenteredGrid):
return value.staticshape.data
else:
return math.staticshape(value)
def gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
gravity = gravity * GLOBAL_AXIS_ORDER.up_vector(rank)
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
return math.to_float(
math.expand_dims([gravity] + [0] * (rank - 1), 0, rank + 1))
else:
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def _expand_axes(data, points, collapse_dimensions=True):
assert math.spatial_rank(data) >= 0
data = math.expand_dims(data, 1, math.spatial_rank(points) - math.spatial_rank(data))
if collapse_dimensions:
return data
else:
points_axes = math.staticshape(points)[1:-1]
data_axes = math.staticshape(data)[1:-1]
for d_points, d_data in zip(points_axes, data_axes):
assert d_points % d_data == 0
tilings = [1] + [d_points // d_data for d_points, d_data in zip(math.staticshape(points)[1:-1], math.staticshape(data)[1:-1])] + [1]
data = math.tile(data, tilings)
return data
def __add__(self, other):
if other is 0:
return self
assert isinstance(other, Gravity), type(other)
if self._batch_size is not None:
assert self._batch_size == other._batch_size
# Add gravity
if math.is_scalar(self.gravity) and math.is_scalar(other.gravity):
return Gravity(self.gravity + other.gravity)
else:
rank = math.staticshape(other.gravity)[-1] if math.is_scalar(self.gravity)\
else math.staticshape(self.gravity)[-1]
sum_tensor = gravity_tensor(self, rank) + gravity_tensor(
other, rank)
return Gravity(sum_tensor)
def _expand_axes(data, points, batch_size=1):
assert math.spatial_rank(data) >= 0
data = math.expand_dims(
data, 1,
math.spatial_rank(points) - math.spatial_rank(data))
points_axes = math.staticshape(points)[1:-1]
data_axes = math.staticshape(data)[1:-1]
for d_points, d_data in zip(points_axes, data_axes):
assert d_points % d_data == 0
tilings = [batch_size or 1] + [
d_points // d_data for d_points, d_data in zip(
math.staticshape(points)[1:-1],
math.staticshape(data)[1:-1])
] + [1]
data = math.tile(data, tilings)
return data
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 __dataop__(self, other, linear_if_scalar, data_operator):
if isinstance(other, StaggeredGrid):
assert self.compatible(
other), 'Fields are not compatible: %s and %s' % (self, other)
data = [
data_operator(c1, c2) for c1, c2 in zip(self.data, other.data)
]
flags = propagate_flags_operation(self.flags + other.flags, False,
self.rank, self.component_count)
elif math.ndims(other) > 0 and math.staticshape(other)[-1] > 1:
other_components = math.unstack(math.as_tensor(other),
axis=-1,
keepdims=True)
data = [
data_operator(c1, c2)
for c1, c2 in zip(self.data, other_components)
]
flags = propagate_flags_operation(self.flags, False, self.rank,
self.component_count)
else:
data = [data_operator(c1, other) for c1 in self.data]
flags = propagate_flags_operation(self.flags, linear_if_scalar,
self.rank, self.component_count)
return self.copied_with(data=np.array(data, dtype=np.object),
flags=flags)
def solve(self, field, domain, guess, enable_backprop):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = math.staticshape(field)[1:-1]
N = int(np.prod(dimensions))
periodic = Material.periodic(domain.domain.boundaries)
if math.choose_backend([field, active_mask,
fluid_mask]).matches_name('SciPy'):
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask,
periodic)
else:
sidx, sorting = sparse_indices(dimensions, periodic)
sval_data = sparse_values(dimensions, active_mask, fluid_mask,
sorting, periodic)
backend = math.choose_backend(field)
sval_data = backend.cast(sval_data, field.dtype)
A = backend.sparse_tensor(indices=sidx,
values=sval_data,
shape=[N, N])
div_vec = math.reshape(field, [-1, int(np.prod(field.shape[1:]))])
if guess is not None:
guess = math.reshape(guess, [-1, int(np.prod(field.shape[1:]))])
def apply_A(pressure):
return math.matmul(A, pressure)
result_vec, iterations = conjugate_gradient(div_vec, apply_A, guess,
self.accuracy,
self.max_iterations,
enable_backprop)
return math.reshape(result_vec, math.shape(field)), iterations
def solve(self, field, domain, guess):
assert isinstance(domain, FluidDomain)
active_mask = domain.active_tensor(extend=1)
fluid_mask = domain.accessible_tensor(extend=1)
dimensions = math.staticshape(field)[1:-1]
N = int(np.prod(dimensions))
periodic = Material.periodic(domain.domain.boundaries)
if math.choose_backend([field, active_mask, fluid_mask]).matches_name('SciPy'):
A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask, periodic)
else:
sidx, sorting = sparse_indices(dimensions, periodic)
sval_data = sparse_values(dimensions, active_mask, fluid_mask, sorting, periodic)
backend = math.choose_backend(field)
sval_data = backend.cast(sval_data, field.dtype)
A = backend.sparse_tensor(indices=sidx, values=sval_data, shape=[N, N])
if self.autodiff:
return sparse_cg(field, A, self.max_iterations, guess, self.accuracy, back_prop=True)
else:
def pressure_gradient(op, grad):
return sparse_cg(grad, A, max_gradient_iterations, None, self.gradient_accuracy)[0]
pressure, iteration = math.with_custom_gradient(sparse_cg,
[field, A, self.max_iterations, guess, self.accuracy],
pressure_gradient, input_index=0, output_index=0,
name_base='scg_pressure_solve')
max_gradient_iterations = iteration if self.max_gradient_iterations == 'mirror' else self.max_gradient_iterations
return pressure, iteration
def solve(self, divergence, domain, guess, enable_backprop):
"""
:param guess: not used in this implementation, Kernel takes the last pressure value for initial_guess
"""
active_mask, accessible_mask = domain.active_tensor(
extend=1), domain.accessible_tensor(extend=1)
# Setup
dimensions = math.staticshape(divergence)[1:-1]
dimensions = dimensions[::
-1] # the custom op needs it in the x,y,z order
dim_array = np.array(dimensions)
dim_product = np.prod(dimensions)
mask_dimensions = dim_array + 2
laplace_matrix = tf.zeros(dim_product * (len(dimensions) * 2 + 1),
dtype=tf.int8)
# Helper variables for CG, make sure new memory is allocated for each variable.
one_vector = tf.ones(dim_product, dtype=tf.float32)
p = tf.zeros_like(divergence, dtype=tf.float32) + 1
z = tf.zeros_like(divergence, dtype=tf.float32) + 2
r = tf.zeros_like(divergence, dtype=tf.float32) + 3
pressure = tf.zeros_like(divergence, dtype=tf.float32) + 4
# Solve
pressure, iteration = pressure_op.pressure_solve(
dimensions, mask_dimensions, active_mask, accessible_mask,
laplace_matrix, divergence, p, r, z, pressure, one_vector,
dim_product, self.accuracy, self.max_iterations)
return pressure, iteration
def tensorshape(tensor):
if tensor is None: return None
default_batched_shape = staticshape(tensor)
if len(default_batched_shape) >= 2:
return (self._batch_size,) + default_batched_shape[1:]
else:
return default_batched_shape
def component_count(self):
if self._rank is None:
return None
try:
example_value = self.sample_at([[0] * self._rank])
return math.staticshape(example_value)[-1]
except NotImplementedError:
return None
def _convert_constant_to_data(value):
if isinstance(value, numbers.Number):
value = math.to_float(math.expand_dims(value))
if isinstance(value, (list, tuple)):
value = math.to_float(numpy.array(value))
if len(math.staticshape(value)) < 2:
value = math.expand_dims(value)
return value
def __init__(self,
center,
size,
wave_vector,
name='wave_packet',
data=1.0,
**kwargs):
rank = math.staticshape(center)[-1]
AnalyticField.__init__(self, **struct.kwargs(locals()))
def resolution(self):
if self.content_type in (struct.VALID, struct.INVALID):
return math.as_tensor(math.staticshape(self.data)[1:-1])
elif self.content_type in (struct.shape, struct.staticshape):
return self.data[1:-1]
else:
raise AssertionError(
'Cannot compute resolution of invalid CenteredGrid (content type = %s)'
% self.content_type)
def __init__(self,
center=(0, 0),
size=0,
mode='RANDOM',
factor=1.0,
name='Seed',
data=1.0,
**kwargs):
rank = math.staticshape(center)[-1]
AnalyticField.__init__(self, **struct.kwargs(locals()))
def sequence_get(self, sequence, name):
if isinstance(sequence, dict):
return sequence[name]
if isinstance(sequence, (tuple, list)):
assert len(sequence) == self.rank
return sequence[self.names.index(name)]
if math.is_tensor(sequence):
assert math.staticshape(sequence) == (self.rank,)
return sequence[self.names.index(name)]
else: # just a constant
return sequence
def step(self, u, dt=1.0, **dependent_states):
assert isinstance(u, CenteredGrid)
grad = u.gradient()
laplace = u.laplace()
laplace2 = laplace.laplace()
du_dt = -laplace - laplace2 - 0.5 * grad**2
result = u + dt * du_dt
result -= math.mean(result.data,
axis=tuple(
range(1, len(math.staticshape(result.data)))),
keepdims=True)
return result.copied_with(age=u.age + dt, name=u.name)
def global_to_child(self, location):
""" Inverse transform """
delta = location - self.center
if math.staticshape(location)[-1] == 2:
angle = -math.batch_align(self.angle, 0, location)
sin = math.sin(angle)
cos = math.cos(angle)
y, x = math.unstack(delta, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
x, y = y, x
rot_x = cos * x - sin * y
rot_y = sin * x + cos * y
rotated = math.stack([rot_y, rot_x], axis=-1)
elif math.staticshape(location)[-1] == 3:
angle = math.batch_align(self.angle, 1, location)
raise NotImplementedError(
'not yet implemented') # ToDo apply angle
else:
raise NotImplementedError('Rotation only supported in 2D and 3D')
final = rotated + self.center
return final
def _component_grid(self, grid, axis):
resolution = list(grid.resolution if isinstance(grid, CenteredGrid) else math.staticshape(grid)[1:-1])
resolution[axis] -= 1
box = staggered_component_box(resolution, axis, self.box)
if isinstance(grid, CenteredGrid):
assert grid.component_count == 1
assert grid.rank == self.rank
assert grid.box == box
if grid.extrapolation != self.extrapolation:
grid = grid.copied_with(extrapolation=self.extrapolation)
else:
grid = CenteredGrid(data=grid, box=box, extrapolation=self.extrapolation, name=_subname(self.name, axis),
batch_size=self._batch_size, flags=propagate_flags_children(self.flags, box.rank, 1))
return grid
def sample_at(self, points):
points_rank = math.spatial_rank(points)
src_rank = math.spatial_rank(self.location)
# --- Expand shapes to format (batch_size, points_dims..., src_dims..., channels) ---
points = math.expand_dims(points, axis=-2, number=src_rank)
src_points = math.expand_dims(self.location,
axis=-2,
number=points_rank)
src_strength = math.expand_dims(self.strength, axis=-1)
src_strength = math.batch_align(src_strength, 0, self.location)
src_strength = math.expand_dims(src_strength,
axis=-1,
number=points_rank)
src_axes = tuple(range(-2, -2 - src_rank, -1))
# --- Compute distances and falloff ---
distances = points - src_points
if self.falloff is not None:
raise NotImplementedError()
# distances_squared = math.sum(distances ** 2, axis=-1, keepdims=True)
# unit_distances = distances / math.sqrt(distances_squared)
# strength = src_strength * math.exp(-distances_squared)
else:
strength = src_strength
# --- Compute velocities ---
if math.staticshape(points)[-1] == 2: # Curl in 2D
dist_1, dist_2 = math.unstack(distances, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
velocity = strength * math.stack([-dist_2, dist_1], axis=-1)
else:
velocity = strength * math.stack([dist_2, -dist_1], axis=-1)
elif math.staticshape(points)[-1] == 3: # Curl in 3D
raise NotImplementedError('not yet implemented')
else:
raise AssertionError(
'Vector product not available in > 3 dimensions')
velocity = math.sum(velocity, axis=src_axes)
return velocity
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 = local_points * math.to_float(self.resolution) - 0.5
if self.extrapolation == 'periodic':
data = math.pad(self.data,
[[0, 0]] + [[0, 1]] * self.rank + [[0, 0]],
mode='wrap')
local_points = local_points % math.to_float(
math.staticshape(self.data)[1:-1])
resampled = math.resample(data,
local_points,
interpolation=self.interpolation)
else:
boundary = 'replicate' if self.extrapolation == 'boundary' else 'zero'
resampled = math.resample(self.data,
local_points,
boundary=boundary,
interpolation=self.interpolation)
return resampled
def cuda_solve_forward(divergence, active_mask, fluid_mask, accuracy,
max_iterations):
# Setup
dimensions = math.staticshape(divergence)[1:-1]
dimensions = dimensions[::-1] # the custom op needs it in the x,y,z order
dim_array = np.array(dimensions)
dim_product = np.prod(dimensions)
mask_dimensions = dim_array + 2
laplace_matrix = tf.zeros(dim_product * (len(dimensions) * 2 + 1),
dtype=tf.int8)
# Helper variables for CG, make sure new memory is allocated for each variable.
one_vector = tf.ones(dim_product, dtype=tf.float32)
p = tf.zeros_like(divergence, dtype=tf.float32) + 1
z = tf.zeros_like(divergence, dtype=tf.float32) + 2
r = tf.zeros_like(divergence, dtype=tf.float32) + 3
pressure = tf.zeros_like(divergence, dtype=tf.float32) + 4
# Solve
pressure, iteration = pressure_op.pressure_solve(
dimensions, mask_dimensions, active_mask, fluid_mask, laplace_matrix,
divergence, p, r, z, pressure, one_vector, dim_product, accuracy,
max_iterations)
return pressure, iteration
def sample(value, domain, batch_size=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:
data = value.sample_at(
CenteredGrid.getpoints(domain.box, domain.resolution).data)
else: # value is constant
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )) + value
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries))
def resolution(self):
return math.as_tensor(math.staticshape(self.data)[1:-1])
class CenteredGrid(Field):
def __init__(self,
data,
box=None,
extrapolation='boundary',
name=None,
**kwargs):
"""Create new CenteredGrid from array like data
:param data: numerical values to be set as values of CenteredGrid (immutable)
:type data: array-like
:param box: numerical values describing the surrounding area of the CenteredGrid, defaults to None
:type box: domain.box, optional
:param extrapolation: set conditions for boundaries, defaults to 'boundary'
:type extrapolation: str, optional
:param name: give CenteredGrid a custom name (immutable), defaults to None
:type name: string, optional
"""
Field.__init__(self, **struct.kwargs(locals()))
self._sample_points = None
@staticmethod
def sample(value, domain, batch_size=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:
data = value.sample_at(
CenteredGrid.getpoints(domain.box, domain.resolution).data)
else: # value is constant
components = math.staticshape(
value)[-1] if math.ndims(value) > 0 else 1
data = math.zeros((batch_size, ) + tuple(domain.resolution) +
(components, )) + value
return CenteredGrid(data,
box=domain.box,
extrapolation=Material.extrapolation_mode(
domain.boundaries))
@struct.variable()
def data(self, data):
if data is None:
return None
if isinstance(data, (tuple, list)):
data = np.array(data) # numbers or objects
while math.ndims(data) < 2:
data = math.expand_dims(data)
return data
data.override(
struct.staticshape, lambda self, data:
(self._batch_size, ) + math.staticshape(data)[1:])
@property
def resolution(self):
return math.as_tensor(math.staticshape(self.data)[1:-1])
@struct.constant(dependencies=Field.data)
def box(self, box):
return AABox.to_box(box, resolution_hint=self.resolution)
@property
def dx(self):
return self.box.size / self.resolution
@property
def rank(self):
return math.spatial_rank(self.data)
@struct.constant(default='boundary')
def extrapolation(self, extrapolation):
if extrapolation is None:
return 'boundary'
assert extrapolation in ('periodic', 'constant',
'boundary') or isinstance(
extrapolation,
(tuple, list)), extrapolation
return collapse(extrapolation)
@struct.constant(default='linear')
def interpolation(self, interpolation):
assert interpolation == 'linear'
return interpolation
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 at(self,
other_field,
collapse_dimensions=True,
force_optimization=False,
return_self_if_compatible=False):
if self.compatible(
other_field
): # and return_self_if_compatible: not applicable for fields with Points
return self
if isinstance(other_field, CenteredGrid) and np.allclose(
self.dx, other_field.dx):
paddings = _required_paddings_transposed(self.box, self.dx,
other_field.box)
if math.sum(paddings) == 0:
origin_in_local = self.box.global_to_local(
other_field.box.lower) * self.resolution
data = _crop_for_interpolation(self.data, origin_in_local,
other_field.resolution)
dimensions = self.resolution != other_field.resolution
dimensions = [
d for d in math.spatial_dimensions(data)
if dimensions[d - 1]
]
data = math.interpolate_linear(data, origin_in_local % 1.0,
dimensions)
return CenteredGrid(data,
other_field.box,
name=self.name,
batch_size=self._batch_size)
elif math.sum(paddings) < 16:
padded = self.padded(np.transpose(paddings).tolist())
return padded.at(other_field, collapse_dimensions,
force_optimization)
return Field.at(self,
other_field,
force_optimization=force_optimization)
@property
def component_count(self):
return self.data.shape[-1]
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
return [
CenteredGrid(math.expand_dims(component),
box=self.box,
name='%s[...,%d]' % (self.name, i),
flags=flags,
batch_size=self._batch_size)
for i, component in enumerate(math.unstack(self.data, -1))
]
@property
def points(self):
if SAMPLE_POINTS in self.flags:
return self
if self._sample_points is None:
self._sample_points = CenteredGrid.getpoints(
self.box, self.resolution)
return self._sample_points
def compatible(self, other_field):
if not other_field.has_points:
return True
if isinstance(other_field, CenteredGrid):
if self.box != other_field.box:
return False
if self.rank != other_field.rank:
return False
for r1, r2 in zip(self.resolution, other_field.resolution):
if r1 != r2 and r2 != 1 and r1 != 1:
return False
return True
else:
return False
def __repr__(self):
if self.is_valid:
return 'Grid[%s(%d), size=%s]' % ('x'.join([
str(r) for r in self.resolution
]), self.component_count, self.box.size)
else:
return struct.Struct.__repr__(self)
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 axis_padded(self, axis, lower, upper):
widths = [[lower, upper] if ax == axis else [0, 0]
for ax in range(self.rank)]
return self.padded(widths)
@staticmethod
def getpoints(box, resolution):
idx_zyx = np.meshgrid(*[
np.linspace(0.5 / dim, 1 - 0.5 / dim, dim) for dim in resolution
],
indexing="ij")
local_coords = math.expand_dims(math.stack(idx_zyx, axis=-1),
0).astype(np.float32)
points = box.local_to_global(local_coords)
return CenteredGrid(points,
box,
name='grid_centers(%s, %s)' % (box, resolution),
flags=[SAMPLE_POINTS])
def laplace(self, physical_units=True, axes=None):
if not physical_units:
data = math.laplace(self.data,
padding=_pad_mode(self.extrapolation),
axes=axes)
else:
if not self.has_cubic_cells:
raise NotImplementedError('Only cubic cells supported.')
laplace = math.laplace(self.data,
padding=_pad_mode(self.extrapolation),
axes=axes)
data = laplace / self.dx[0]**2
extrapolation = map_for_axes(_gradient_extrapolation,
self.extrapolation, axes, self.rank)
return self.copied_with(data=data,
extrapolation=extrapolation,
flags=())
def gradient(self, physical_units=True):
if not physical_units or self.has_cubic_cells:
data = math.gradient(self.data,
dx=np.mean(self.dx),
padding=_pad_mode(self.extrapolation))
return self.copied_with(data=data,
extrapolation=_gradient_extrapolation(
self.extrapolation),
flags=())
else:
raise NotImplementedError('Only cubic cells supported.')
@property
def has_cubic_cells(self):
return np.allclose(self.dx, np.mean(self.dx))
def normalized(self, total, epsilon=1e-5):
if isinstance(total, CenteredGrid):
total = total.data
normalize_data = math.normalize_to(self.data, total, epsilon)
return self.with_data(normalize_data)
def _padded_resample(self, points):
data = self.padded([[1, 1]] * self.rank).data
local_points = self.box.global_to_local(points)
local_points = local_points * math.to_float(
self.resolution) + 0.5 # depends on amount of padding
resampled = math.resample(data,
local_points,
boundary='replicate',
interpolation=self.interpolation)
return resampled
def staggered_grid(tensor, name='manta_staggered'):
tensor = tensor[..., ::-1] # manta: xyz, phiflow: zyx
assert math.staticshape(tensor)[-1] == math.spatial_rank(tensor)
return StaggeredGrid(tensor, name=name)
