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))