def unstack_staggered_tensor(tensor):
tensors = math.unstack(tensor, -1)
result = []
for i, dim in enumerate(math.spatial_dimensions(tensor)):
slices = [slice(None, -1) if d != dim else slice(None) for d in math.spatial_dimensions(tensor)]
result.append(math.expand_dims(tensors[i][tuple([slice(None)]+slices)], -1))
return result
def at(self, other_field):
if self.compatible(other_field):
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)
return Field.at(self, other_field)
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 step(self, state, dt=1.0, potentials=(), obstacles=()):
if len(potentials) == 0:
potential = 0
else:
potential = math.zeros_like(math.real(
state.amplitude)) # for the moment, allow only real potentials
for pot in potentials:
potential = effect_applied(pot, potential, dt)
potential = potential.data
amplitude = state.amplitude.data
# Rotate by potential
rotation = math.exp(1j * math.to_complex(potential * dt))
amplitude = amplitude * rotation
# Move by rotating in Fourier space
amplitude_fft = math.fft(amplitude)
laplace = math.fftfreq(state.resolution, mode='square')
amplitude_fft *= math.exp(-1j * (2 * np.pi)**2 * math.to_complex(dt) *
laplace / (2 * state.mass))
amplitude = math.ifft(amplitude_fft)
obstacle_mask = union_mask([
obstacle.geometry for obstacle in obstacles
]).at(state.amplitude).data
amplitude *= 1 - obstacle_mask
normalized = False
symmetric = False
if not symmetric:
boundary_mask = math.zeros(
state.domain.centered_shape(1, batch_size=1)).data
boundary_mask[[slice(None)] + [
slice(self.margin, -self.margin)
for i in math.spatial_dimensions(boundary_mask)
] + [slice(None)]] = 1
amplitude *= boundary_mask
if len(obstacles) > 0 or not symmetric:
amplitude = normalize_probability(amplitude)
normalized = True
return state.copied_with(amplitude=amplitude)
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)
def normalize_probability(probability_amplitude):
p = math.to_float(abs(probability_amplitude)**2)
P = math.sum(p, math.spatial_dimensions(p), keepdims=True)
return probability_amplitude / math.to_complex(math.sqrt(P))
