def test_pad_tensor(self):
for backend in BACKENDS:
with backend:
a = math.meshgrid(x=4, y=3)
# 0
p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, ZERO)
self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check
math.assert_close(p.x[1:-2].y[:-1], a) # copy inner
math.assert_close(p.x[0], 0)
# 1
p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, ONE)
self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check
math.assert_close(p.x[1:-2].y[:-1], a) # copy inner
math.assert_close(p.x[0], 1)
# periodic
p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, PERIODIC)
self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check
math.assert_close(p.x[1:-2].y[:-1], a) # copy inner
math.assert_close(p.x[0].y[:-1], a.x[-1])
math.assert_close(p.x[-2:].y[:-1], a.x[:2])
# boundary
p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, BOUNDARY)
self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check
math.assert_close(p.x[1:-2].y[:-1], a) # copy inner
math.assert_close(p.x[0].y[:-1], a.x[0])
math.assert_close(p.x[-2:].y[:-1], a.x[-1])
# mixed
p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, combine_sides({'x': PERIODIC, 'y': (ONE, REFLECT)}))
math.print(p)
self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check
math.assert_close(p.x[1:-2].y[:-1], a) # copy inner
math.assert_close(p.x[0].y[:-1], a.x[-1]) # periodic
math.assert_close(p.x[-2:].y[:-1], a.x[:2]) # periodic
def at_faces(self, face_dimension_xyz):
dims = range(self.spatial_rank)
face_dimension_zyx = len(
dims) - face_dimension_xyz - 1 # 0=Z, 1=Y, 2=X, etc.
components = []
for d in dims: # z,y,x
if d == face_dimension_zyx:
components.append(self.staggered[..., len(dims) - d - 1])
else:
# Interpolate other components
vq = self.staggered[..., len(dims) - d - 1]
t = vq
for d2 in dims: # z,y,x
slices1 = [(slice(1, None) if i == d2 else slice(None))
for i in dims]
slices2 = [(slice(-1) if i == d2 else slice(None))
for i in dims]
t = t[[slice(None)] + slices1] + t[[slice(None)] + slices2]
if d2 == d:
t = math.pad(
t, [[0, 0]] + [([0, 1] if i == d2 else [0, 0])
for i in dims]) / 2
else:
t = math.pad(
t, [[0, 0]] + [([1, 0] if i == d2 else [0, 0])
for i in dims]) / 2
components.append(t)
return math.stack(components[::-1], axis=-1)
def test_pad_collapsed(self):
a = math.zeros(b=2, x=10, y=10, batch=10)
p = math.pad(a, {'x': (1, 2)}, ZERO)
self.assertIsInstance(p, CollapsedTensor)
self.assertEqual((10, 2, 13, 10), p.shape.sizes)
p = math.pad(a, {'x': (1, 2)}, PERIODIC)
self.assertIsInstance(p, CollapsedTensor)
self.assertEqual((10, 2, 13, 10), p.shape.sizes)
def stagger(field: CenteredGrid,
face_function: Callable,
extrapolation: math.extrapolation.Extrapolation,
type: type = StaggeredGrid):
"""
Creates a new grid by evaluating `face_function` given two neighbouring cells.
One layer of missing cells is inferred from the extrapolation.
This method returns a Field of type `type` which must be either StaggeredGrid or CenteredGrid.
When returning a StaggeredGrid, the new values are sampled at the faces of neighbouring cells.
When returning a CenteredGrid, the new grid has the same resolution as `field`.
Args:
field: centered grid
face_function: function mapping (value1: Tensor, value2: Tensor) -> center_value: Tensor
extrapolation: extrapolation mode of the returned grid. Has no effect on the values.
type: one of (StaggeredGrid, CenteredGrid)
field: CenteredGrid:
face_function: Callable:
extrapolation: math.extrapolation.Extrapolation:
type: type: (Default value = StaggeredGrid)
Returns:
grid of type matching the `type` argument
"""
all_lower = []
all_upper = []
if type == StaggeredGrid:
for dim in field.shape.spatial.names:
lo_valid, up_valid = extrapolation.valid_outer_faces(dim)
width_lower = {dim: (int(lo_valid), int(up_valid) - 1)}
width_upper = {
dim: (int(lo_valid or up_valid) - 1, int(lo_valid
and up_valid))
}
all_lower.append(
math.pad(field.values, width_lower, field.extrapolation))
all_upper.append(
math.pad(field.values, width_upper, field.extrapolation))
all_upper = math.stack(all_upper, channel('vector'))
all_lower = math.stack(all_lower, channel('vector'))
values = face_function(all_lower, all_upper)
result = StaggeredGrid(values,
bounds=field.bounds,
extrapolation=extrapolation)
assert result.shape.spatial == field.shape.spatial
return result
elif type == CenteredGrid:
left, right = math.shift(field.values, (-1, 1),
padding=field.extrapolation,
stack_dim=channel('vector'))
values = face_function(left, right)
return CenteredGrid(values,
bounds=field.bounds,
extrapolation=extrapolation)
else:
raise ValueError(type)
def linear_function(val):
val = -val
val *= 2
val = math.pad(val, {'x': (2, 0), 'y': (0, 1)}, extrapolation.PERIODIC)
val = val.x[:-2].y[1:] + val.x[2:].y[:-1]
val = math.pad(val, {'x': (0, 0), 'y': (0, 1)}, extrapolation.ZERO)
val = math.pad(val, {'x': (2, 2), 'y': (0, 1)}, extrapolation.BOUNDARY)
# sl = sl.vector[0]
return val
val = val.x[1:4].y[:2]
return math.sum([val, sl], axis=0) - sl
def gradient(scalar_field, padding_mode='replicate'):
assert isinstance(scalar_field, CenteredGrid)
data = scalar_field.data
if data.shape[-1] != 1:
raise ValueError('input must be a scalar field')
tensors = []
for dim in math.spatial_dimensions(data):
upper = math.pad(data, [[0,1] if d == dim else [0,0] for d in math.all_dimensions(data)], padding_mode)
lower = math.pad(data, [[1,0] if d == dim else [0,0] for d in math.all_dimensions(data)], padding_mode)
tensors.append((upper - lower) / scalar_field.dx[dim - 1])
return StaggeredGrid(tensors, scalar_field.box, name='grad(%s)' % scalar_field.name,
batch_size=scalar_field._batch_size)
def bake_extrapolation(grid: GridType) -> GridType:
"""
Pads `grid` with its current extrapolation.
For `StaggeredGrid`s, the resulting grid will have a consistent shape, independent of the original extrapolation.
Args:
grid: `CenteredGrid` or `StaggeredGrid`.
Returns:
Padded grid with extrapolation `phi.math.extrapolation.NONE`.
"""
if grid.extrapolation == math.extrapolation.NONE:
return grid
if isinstance(grid, StaggeredGrid):
values = grid.values.unstack('vector')
padded = []
for dim, value in zip(grid.shape.spatial.names, values):
lower, upper = grid.extrapolation.valid_outer_faces(dim)
padded.append(
math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)},
grid.extrapolation))
return StaggeredGrid(math.stack(padded, channel('vector')),
bounds=grid.bounds,
extrapolation=math.extrapolation.NONE)
elif isinstance(grid, CenteredGrid):
return pad(grid, 1).with_extrapolation(math.extrapolation.NONE)
else:
raise ValueError(f"Not a valid grid: {grid}")
def with_extrapolation(self, extrapolation: math.Extrapolation):
if all(
extrapolation.valid_outer_faces(dim) ==
self.extrapolation.valid_outer_faces(dim)
for dim in self.resolution.names):
return StaggeredGrid(self.values,
extrapolation=extrapolation,
bounds=self.bounds)
else:
values = []
for dim, component in zip(self.shape.spatial.names,
self.values.unstack('vector')):
old_lo, old_hi = [
int(v) for v in self.extrapolation.valid_outer_faces(dim)
]
new_lo, new_hi = [
int(v) for v in extrapolation.valid_outer_faces(dim)
]
widths = (new_lo - old_lo, new_hi - old_hi)
values.append(
math.pad(component, {dim: widths}, self.extrapolation))
values = math.stack(values, channel('vector'))
return StaggeredGrid(values,
extrapolation=extrapolation,
bounds=self.bounds)
def create_surface_mask(liquid_mask):
"""
Computes inner contours of the liquid_mask.
A cell i is flagged 1 if liquid_mask[i] = 1 and it has a non-liquid neighbour.
:param liquid_mask: binary tensor
:return: tensor
"""
# When we create inner contour, we don't want the fluid-wall boundaries to show up as surface, so we should pad with symmetric edge values.
mask = math.pad(liquid_mask, [[0, 0]] +
[[1, 1]] * math.spatial_rank(liquid_mask) + [[0, 0]],
"constant")
dims = range(math.spatial_rank(mask))
bcs = math.zeros_like(liquid_mask)
# Move in every possible direction to assure corners are properly set.
directions = np.array(
list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1)))))
for d in directions:
d_slice = tuple([(slice(2, None) if d[i] == -1 else
slice(0, -2) if d[i] == 1 else slice(1, -1))
for i in dims])
center_slice = tuple([slice(1, -1) for _ in dims])
# Create inner contour of particles
bc_d = math.maximum(mask[(slice(None),) + d_slice + (slice(None),)],
mask[(slice(None),) + center_slice + (slice(None),)]) - \
mask[(slice(None),) + d_slice + (slice(None),)]
bcs = math.maximum(bcs, bc_d)
return bcs
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 _staggered_curl_3d(self):
"""
Calculates the curl operator on a staggered three-dimensional field.
The resulting vector field is a staggered grid.
If the velocities of the vector potential were sampled at the lower faces of a cube, the resulting velocities
are sampled at the centers of the upper edges.
:param vector_potential: three-dimensional vector potential
:return: three-dimensional staggered vector field
"""
kernel = np.zeros((2, 2, 2, 3, 3), np.float32)
derivative = np.array([-1, 1])
# x-component: dz/dy - dy/dz
kernel[0, :, 0, 2, 0] = derivative
kernel[:, 0, 0, 1, 0] = -derivative
# y-component: dx/dz - dz/dx
kernel[:, 0, 0, 0, 1] = derivative
kernel[0, 0, :, 2, 1] = -derivative
# z-component: dy/dx - dx/dy
kernel[0, 0, :, 1, 2] = derivative
kernel[0, :, 0, 0, 2] = -derivative
vector_potential = math.pad(self.staggered,
[[0, 0], [0, 1], [0, 1], [0, 1], [0, 0]],
"SYMMETRIC")
vector_field = math.conv(vector_potential, kernel, padding="VALID")
return StaggeredGrid(vector_field)
def upsample2x(tensor, interpolation="LINEAR"):
if interpolation.lower() != "linear":
raise ValueError("Only linear interpolation supported")
dims = range(spatial_rank(tensor))
vlen = tensor.shape[-1]
spatial_dims = tensor.shape[1:-1]
tensor = math.pad(tensor,
[[0, 0]] + [[1, 1]] * spatial_rank(tensor) + [[0, 0]],
"SYMMETRIC")
for dim in dims:
left_slices_1 = [(slice(2, None) if i == dim else slice(None))
for i in dims]
left_slices_2 = [(slice(1, -1) if i == dim else slice(None))
for i in dims]
right_slices_1 = [(slice(1, -1) if i == dim else slice(None))
for i in dims]
right_slices_2 = [(slice(-2) if i == dim else slice(None))
for i in dims]
left = 0.75 * tensor[[slice(None)] + left_slices_2 +
[slice(None)]] + 0.25 * tensor[
[slice(None)] + left_slices_1 + [slice(None)]]
right = 0.25 * tensor[[slice(None)] + right_slices_2 + [
slice(None)
]] + 0.75 * tensor[[slice(None)] + right_slices_1 + [slice(None)]]
combined = math.stack([right, left], axis=2 + dim)
tensor = math.reshape(combined, [-1] + [
spatial_dims[dim] * 2 if i == dim else tensor.shape[i + 1]
for i in dims
] + [vlen])
return tensor
def linear_function(val):
val = -val
val *= 2
val = math.pad(val, {
'x': (2, 0),
'y': (0, 1)
}, math.extrapolation.PERIODIC)
val = val.x[:-2].y[1:] + val.x[2:].y[:-1]
val = math.pad(val, {
'x': (0, 0),
'y': (0, 1)
}, math.extrapolation.ZERO)
val = math.pad(val, {
'x': (2, 2),
'y': (0, 1)
}, math.extrapolation.BOUNDARY)
return math.sum([val, val], dim='0') - val
def active_tensor(self, extend=0):
"""
Scalar channel encoding active cells as ones and inactive (open/obstacle) as zero.
Active cells are those for which physical constants_dict such as pressure or velocity are calculated.
:param extend: Extend the grid in all directions beyond the grid size specified by the domain
"""
return math.pad(self.active.data, [[0, 0]] + [[extend, extend]] * self.rank + [[0, 0]], "constant")
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 stack_staggered_components(data: Tensor) -> Tensor:
padded = []
for dim, component in zip(data.shape.spatial.names,
data.unstack('vector')):
padded.append(
math.pad(
component,
{d: (0, 1)
for d in data.shape.spatial.without(dim).names},
mode=math.extrapolation.ZERO))
return math.channel_stack(padded, 'vector')
def accessible_tensor(self, extend=0):
"""
Scalar channel encoding cells that are accessible, i.e. not solid, as ones and obstacles as zero.
:param extend: Extend the grid in all directions beyond the grid size specified by the domain
"""
pad_values = struct.map(lambda solid: int(not solid), Material.solid(self.domain.boundaries))
if isinstance(pad_values, (list, tuple)):
pad_values = [0] + list(pad_values) + [0]
result = math.pad(self.accessible.data, [[0,0]] + [[extend, extend]] * self.rank + [[0,0]], constant_values=pad_values)
return result
def _staggered_curl_2d(self):
kernel = np.zeros((2, 2, 1, 2), np.float32)
derivative = np.array([-1, 1])
# x-component: dz/dy
kernel[:, 0, 0, 0] = derivative
# y-component: - dz/dx
kernel[0, :, 0, 1] = -derivative
scalar_potential = math.pad(self.staggered,
[[0, 0], [0, 1], [0, 1], [0, 0]],
"SYMMETRIC")
vector_field = math.conv(scalar_potential, kernel, padding="VALID")
return StaggeredGrid(vector_field)
def _central_diff_nd(field, dims):
field = math.pad(field,
[[0, 0]] + [[1, 1]] * spatial_rank(field) + [[0, 0]],
"symmetric")
df_dq = []
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 = field[[slice(None)] + upper_slices +
[0]] - field[[slice(None)] + lower_slices + [0]]
df_dq.append(diff)
return math.stack(df_dq[::-1], axis=-1)
def _forward_diff_nd(field, dims):
df_dq = []
for dimension in dims:
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 = field[[slice(None)] + upper_slices] - field[[slice(None)] +
lower_slices]
padded = math.pad(diff,
[[0, 0]] + [([0, 1] if i == dimension else [0, 0])
for i in dims])
df_dq.append(padded)
return math.stack(df_dq[::-1], axis=-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 _stagger_sample(self, box, resolution):
"""
Samples this field on a staggered grid.
In addition to sampling, extrapolates the field using an occupancy mask generated from the points.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: StaggeredGrid
"""
resolution = np.array(resolution)
valid_indices = math.to_int(math.floor(self.sample_points))
valid_indices = math.minimum(math.maximum(0, valid_indices), resolution - 1)
# Correct format for math.scatter
valid_indices = batch_indices(valid_indices)
active_mask = math.scatter(self.sample_points, valid_indices, 1, math.concat([[valid_indices.shape[0]], resolution, [1]], axis=-1), duplicates_handling='any')
mask = math.pad(active_mask, [[0, 0]] + [[1, 1]] * self.rank + [[0, 0]], "constant")
if isinstance(self.data, (int, float, np.ndarray)):
values = math.zeros_like(self.sample_points) + self.data
else:
values = self.data
result = []
ones_1d = math.unstack(math.ones_like(values), axis=-1)[0]
staggered_shape = [i + 1 for i in resolution]
dx = box.size / resolution
dims = range(len(resolution))
for d in dims:
staggered_offset = math.stack([(0.5 * dx[i] * ones_1d if i == d else 0.0 * ones_1d) for i in dims], axis=-1)
indices = math.to_int(math.floor(self.sample_points + staggered_offset))
valid_indices = math.maximum(0, math.minimum(indices, resolution))
valid_indices = batch_indices(valid_indices)
values_d = math.expand_dims(math.unstack(values, axis=-1)[d], axis=-1)
result.append(math.scatter(self.sample_points, valid_indices, values_d, [indices.shape[0]] + staggered_shape + [1], duplicates_handling=self.mode))
d_slice = tuple([(slice(0, -2) if i == d else slice(1,-1)) for i in dims])
u_slice = tuple([(slice(2, None) if i == d else slice(1,-1)) for i in dims])
active_mask = math.minimum(mask[(slice(None),) + d_slice + (slice(None),)], active_mask)
active_mask = math.minimum(mask[(slice(None),) + u_slice + (slice(None),)], active_mask)
staggered_tensor_prep = unstack_staggered_tensor(math.concat(result, axis=-1))
grid_values = StaggeredGrid(staggered_tensor_prep)
# Fix values at boundary of liquids (using StaggeredGrid these might not receive a value, so we replace it with a value inside the liquid)
grid_values, _ = extrapolate(grid_values, active_mask, voxel_distance=2)
return grid_values
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 batch_indices(indices):
"""
Reshapes the indices such that, aside from indices, they also contain batch number.
For example the entry (32, 40) as coordinates of batch 2 will become (2, 32, 40).
Transform shape (b, p, d) to (b, p, d+1) where batch size is b, number of particles is p and number of dimensions is d.
"""
batch_size = indices.shape[0]
out_spatial_rank = len(indices.shape) - 2
out_spatial_size = math.shape(indices)[1:-1]
batch_range = math.DYNAMIC_BACKEND.choose_backend(indices).range(batch_size)
batch_ids = math.reshape(batch_range, [batch_size] + [1] * out_spatial_rank)
tile_shape = math.pad(out_spatial_size, [[1,0]], constant_values=1)
batch_ids = math.expand_dims(math.tile(batch_ids, tile_shape), axis=-1)
return math.concat((batch_ids, indices), axis=-1)
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 laplace(tensor, weights=None, padding="symmetric"):
if tensor.shape[-1] != 1:
raise ValueError("Laplace operator requires a scalar field as input")
rank = spatial_rank(tensor)
if padding.lower() != "valid":
tensor = math.pad(tensor, [[0, 0]] + [[1, 1]] * rank + [[0, 0]],
padding)
if weights is not None:
return _weighted_sliced_laplace_nd(tensor, weights)
if rank == 2:
return _conv_laplace_2d(tensor)
elif rank == 3:
return _conv_laplace_3d(tensor)
else:
return _sliced_laplace_nd(tensor)
def from_scalar(scalar_field, axis_forces, padding_mode="constant"):
if scalar_field.shape[-1] != 1:
raise ValueError("Resample requires a scalar field as input")
rank = spatial_rank(scalar_field)
dims = range(rank)
df_dq = []
for dimension in dims: # z,y,x
padded_field = math.pad(scalar_field, [[0, 0]] +
[[1, 1] if i == dimension else [0, 1]
for i in dims] + [[0, 0]], padding_mode)
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]
neighbour_sum = padded_field[[slice(None)] + upper_slices + [slice(None)]] + \
padded_field[[slice(None)] + lower_slices + [slice(None)]]
df_dq.append(axis_forces[dimension] * neighbour_sum * 0.5 / rank)
return StaggeredGrid(math.concat(df_dq[::-1], axis=-1))
def downsample2x(tensor, interpolation="LINEAR"):
if interpolation.lower() != "linear":
raise ValueError("Only linear interpolation supported")
dims = range(spatial_rank(tensor))
tensor = math.pad(tensor,
[[0, 0]] + [([0, 1] if (dim % 2) != 0 else [0, 0])
for dim in tensor.shape[1:-1]] + [[0, 0]],
"SYMMETRIC")
for dimension in dims:
upper_slices = [(slice(1, None, 2) if i == dimension else slice(None))
for i in dims]
lower_slices = [(slice(0, None, 2) if i == dimension else slice(None))
for i in dims]
sum = tensor[[slice(None)] + upper_slices +
[slice(None)]] + tensor[[slice(None)] + lower_slices +
[slice(None)]]
tensor = sum / 2
return tensor
def gradient(scalar_field, padding="symmetric"):
if scalar_field.shape[-1] != 1:
raise ValueError("Gradient requires a scalar field as input")
rank = spatial_rank(scalar_field)
dims = range(rank)
field = math.pad(scalar_field, [[0, 0]] + [[1, 1]] * rank + [[0, 0]],
mode=padding)
df_dq = []
for dimension in dims:
upper_slices = [
(slice(1, None) if i == dimension else slice(1, None))
for i in dims
]
lower_slices = [(slice(-1) if i == dimension else slice(1, None))
for i in dims]
diff = field[[slice(None)] + upper_slices] - field[[slice(None)] +
lower_slices]
df_dq.append(diff)
return StaggeredGrid(math.concat(df_dq[::-1], axis=-1))
