def grid_sample(self, resolution: math.Shape, size, shape: math.Shape = None): shape = (self._shape if shape is None else shape) & resolution for dim in channel(self._shape): if dim.item_names[0] is None: warnings.warn( f"Please provide item names for Noise dim {dim} using {dim}='x,y,z'", FutureWarning) shape &= channel(**{dim.name: resolution.names}) rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex( random_normal(shape)) # Note: there is no complex32 with math.NUMPY: k = math.fftfreq(resolution) * resolution / math.tensor( size) * math.tensor(self.scale) # in physical units k = math.vec_squared(k) lowest_frequency = 0.1 weight_mask = math.to_float(k > lowest_frequency) # --- Compute 1/k --- k._native[(0, ) * len(k.shape)] = np.inf inv_k = 1 / k inv_k._native[(0, ) * len(k.shape)] = 0 # --- Compute result --- fft = rndj * inv_k**self.smoothness * weight_mask array = math.real(math.ifft(fft)) array /= math.std(array, dim=array.shape.non_batch) array -= math.mean(array, dim=array.shape.non_batch) array = math.to_float(array) return array
def grid_sample(self, resolution, size, batch_size=1, channels=None): channels = channels or self.channels or len(size) shape = (batch_size, ) + tuple(resolution) + (channels, ) rndj = math.to_complex( self.math.random_normal(shape)) + 1j * math.to_complex( self.math.random_normal(shape)) # Note: there is no complex32 k = math.fftfreq( resolution) * resolution / size * self.scale # in physical units k = math.sum(k**2, axis=-1, keepdims=True) lowest_frequency = 0.1 weight_mask = 1 / (1 + math.exp( (lowest_frequency - k) * 1e3)) # High pass filter # --- Compute 1/k --- k[(0, ) * len(k.shape)] = np.inf inv_k = 1 / k inv_k[(0, ) * len(k.shape)] = 0 # --- Compute result --- fft = rndj * inv_k**self.smoothness * weight_mask array = math.real(math.ifft(fft)) array /= math.std(array, axis=tuple(range(1, math.ndims(array))), keepdims=True) array -= math.mean(array, axis=tuple(range(1, math.ndims(array))), keepdims=True) array = math.to_float(array) return array
def diffuse(field, amount, substeps=1): assert isinstance(field, CenteredGrid) if field.extrapolation == 'periodic': frequencies = math.fft(math.to_complex(field.data)) k = math.fftfreq(field.resolution) / field.dx k = math.sum(k**2, axis=-1, keepdims=True) fft_laplace = -(2 * pi)**2 * k diffuse_kernel = math.to_complex(math.exp(fft_laplace * amount)) data = math.ifft(frequencies * diffuse_kernel) data = math.real(data) else: data = field.data for i in range(substeps): data += amount / substeps * field.laplace() return field.with_data(data)
def diffuse(field, amount, substeps=1): u""" Simulate a finite-time diffusion process of the form dF/dt = α · ΔF on a given `Field` F with diffusion coefficient α. If `field` is periodic (set via `extrapolation='periodic'`), diffusion may be simulated in Fourier space. Otherwise, finite differencing is used to approximate the :param field: CenteredGrid, StaggeredGrid or ConstantField :param amount: number of Field, typically α · dt :param substeps: number of iterations to use :return: Field of same type as `field` :rtype: Field """ if isinstance(field, ConstantField): return field if isinstance(field, StaggeredGrid): return struct.map( lambda grid: diffuse(grid, amount, substeps=substeps), field, leaf_condition=lambda x: isinstance(x, CenteredGrid)) assert isinstance( field, CenteredGrid), "Cannot diffuse field of type '%s'" % type(field) if field.extrapolation == 'periodic' and not isinstance(amount, Field): fft_laplace = -(2 * pi)**2 * field.squared_frequencies diffuse_kernel = math.exp(fft_laplace * amount) return math.real( math.ifft(field.fft() * math.to_complex(diffuse_kernel))) else: data = field.data if isinstance(amount, Field): amount = amount.at(field).data else: amount = math.batch_align(amount, 0, data) for i in range(substeps): data += amount / substeps * field.laplace().data return field.with_data(data)
def test_cast(self): for backend in BACKENDS: with backend: x = math.random_uniform(dtype=DType(float, 64)) self.assertEqual(DType(float, 32), math.to_float(x).dtype, msg=backend.name) self.assertEqual(DType(complex, 64), math.to_complex(x).dtype, msg=backend.name) with math.precision(64): self.assertEqual(DType(float, 64), math.to_float(x).dtype, msg=backend.name) self.assertEqual(DType(complex, 128), math.to_complex(x).dtype, msg=backend.name) self.assertEqual(DType(int, 64), math.to_int64(x).dtype, msg=backend.name) self.assertEqual(DType(int, 32), math.to_int32(x).dtype, msg=backend.name) self.assertEqual(DType(float, 16), math.cast(x, DType(float, 16)).dtype, msg=backend.name) self.assertEqual(DType(complex, 128), math.cast(x, DType(complex, 128)).dtype, msg=backend.name) try: math.cast(x, DType(float, 3)) self.fail(msg=backend.name) except KeyError: pass
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 grid_sample(self, resolution: math.Shape, size, shape: math.Shape = None): shape = (self._shape if shape is None else shape).combined(resolution) rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(random_normal(shape)) # Note: there is no complex32 with math.NUMPY_BACKEND: k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units k = math.vec_squared(k) lowest_frequency = 0.1 weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter # --- Compute 1/k --- k.native()[(0,) * len(k.shape)] = np.inf inv_k = 1 / k inv_k.native()[(0,) * len(k.shape)] = 0 # --- Compute result --- fft = rndj * inv_k ** self.smoothness * weight_mask array = math.real(math.ifft(fft)) array /= math.std(array, dim=array.shape.non_batch) array -= math.mean(array, dim=array.shape.non_batch) array = math.to_float(array) return array
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))