def test_direct_initializers(self): np.testing.assert_equal(math.zeros([1, 16]), np.zeros([1, 16])) self.assertEqual(math.zeros([1, 16]).dtype, np.float32) np.testing.assert_equal(math.ones([1, 16, 1]), np.ones([1, 16, 1])) np.testing.assert_equal(math.zeros_like(math.ones([1, 16, 1])), np.zeros([1, 16, 1])) np.testing.assert_equal(math.randn([1, 4]).shape, [1, 4]) self.assertEqual(math.randn([1, 4]).dtype, np.float32)
def test_fluid_initializers(self): def typetest(fluid): self.assertIsInstance(fluid, Fluid) self.assertIsInstance(fluid.velocity, StaggeredGrid) numpy.testing.assert_equal(fluid.density.data.shape, [1, 4, 4, 1]) numpy.testing.assert_equal(fluid.velocity.resolution, [4, 4]) numpy.testing.assert_equal(fluid.velocity.data[0].resolution, [5, 4]) typetest(Fluid(Domain([4, 4]), density=0.0, velocity=0.0)) typetest(Fluid(Domain([4, 4]), density=1.0, velocity=1.0)) typetest(Fluid(Domain([4, 4]), density=0, velocity=math.zeros)) typetest( Fluid(Domain([4, 4]), density=lambda s: math.randn(s), velocity=lambda s: math.randn(s))) typetest( Fluid(Domain([4, 4]), density=numpy.zeros([1, 4, 4, 1]), velocity=numpy.zeros([1, 5, 5, 2]))) typetest( Fluid(Domain([4, 4]), density=numpy.zeros([1, 4, 4, 1]), velocity=numpy.zeros([1, 5, 5, 2]))) typetest(Fluid(Domain([4, 4])))
def grid_sample(self, resolution, size, batch_size=1, dtype=np.float32): shape = (batch_size,) + tuple(resolution) + (self.channels,) rndj = math.randn(shape, dtype) + 1j * math.randn(shape, dtype) 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)).astype(dtype) 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) return array