def test_sin(self): for backend in BACKENDS: with backend: math.assert_close(math.sin(math.zeros(spatial(x=4))), 0, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.sin(math.tensor(math.PI / 2)), 1, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.sin(math.tensor(math.PI)), 0, abs_tolerance=1e-6, msg=backend.name) math.assert_close(math.sin(math.tensor(math.PI * 3 / 2)), -1, abs_tolerance=1e-6, msg=backend.name)
def test_fourier_laplace_2d_periodic(self): """test for convergence of the laplace operator""" test_params = { 'size': [16, 32, 40], 'L': [1, 2, 3], # NOTE: Cannot test with less than 1 full wavelength } test_cases = [ dict(zip(test_params, v)) for v in product(*test_params.values()) ] for params in test_cases: vec = math.meshgrid(x=params['size'], y=params['size']) sine_field = math.prod( math.sin(2 * PI * params['L'] * vec / params['size'] + 1), 'vector') sin_lap_ref = -2 * ( 2 * PI * params['L'] / params['size'] )**2 * sine_field # leading 2 from from x-y cross terms sin_lap = math.fourier_laplace(sine_field, 1) try: math.assert_close(sin_lap, sin_lap_ref, rel_tolerance=0, abs_tolerance=1e-5) except BaseException as e: abs_error = math.abs(sin_lap - sin_lap_ref) max_abs_error = math.max(abs_error) max_rel_error = math.max(math.abs(abs_error / sin_lap_ref)) variation_str = "\n".join([ f"max_absolute_error: {max_abs_error}", f"max_relative_error: {max_rel_error}", ]) print(f"{variation_str}\n{params}") raise AssertionError(e, f"{variation_str}\n{params}")
def sample_at(self, x): phase_offset = math.batch_align_scalar(self.phase_offset, 0, x) k = math.batch_align(self.k, 1, x) data = math.batch_align_scalar(self.data, 0, x) spatial_phase = math.sum(k * x, -1, keepdims=True) wave = math.sin(spatial_phase + phase_offset) * data return math.cast(wave, self.dtype)
def test_Dict(self): d1 = math.Dict(a=1, b=math.ones(), c=math.ones(spatial(x=3))) math.assert_close(d1 * 2, d1 + d1, 2 * d1, 2 / d1) math.assert_close(0 + d1, d1, d1 - 0, abs(d1), round(d1)) math.assert_close(-d1, 0 - d1) math.assert_close(d1 // 2, d1 * 0, d1 % 1) math.assert_close(d1 / 2, d1 * 0.5, 0.5 * d1) math.assert_close(math.sin(d1 * 0), d1 * 0)
def sample_at(self, x): phase_offset = math.batch_align_scalar(self.phase_offset, 0, x) k = math.batch_align(self.k, 1, x) data = math.batch_align(self.data, 1, x) spatial_phase = math.sum(k * x, -1, keepdims=True) result = math.sin( math.to_float(spatial_phase + phase_offset)) * math.to_float(data) return result
def _rotate(self, location): sin = math.sin(self.angle) cos = math.cos(self.angle) y, x = location.vector.unstack() if GLOBAL_AXIS_ORDER.is_x_first: x, y = y, x rot_x = cos * x - sin * y rot_y = sin * x + cos * y return math.channel_stack([rot_y, rot_x], 'vector')
def test_plot_multi_1d(self): self._test_plot( CenteredGrid( lambda x: math.stack({ 'sin': math.sin(x), 'cos': math.cos(x) }, channel('curves')), x=100, bounds=Box(x=2 * math.pi)))
def test_fourier_laplace_2d_periodic(self): """test for convergence of the laplace operator""" test_params = { 'size': [16, 32, 40], 'L': [1, 2, 3], # NOTE: Cannot test with less than 1 full wavelength } test_cases = [dict(zip(test_params, v)) for v in product(*test_params.values())] for params in test_cases: vec = math.meshgrid(x=params['size'], y=params['size']) sine_field = math.prod(math.sin(2 * PI * params['L'] * vec / params['size'] + 1), 'vector') sin_lap_ref = - 2 * (2 * PI * params['L'] / params['size']) ** 2 * sine_field # leading 2 from from x-y cross terms sin_lap = math.fourier_laplace(sine_field, 1) # try: math.assert_close(sin_lap, sin_lap_ref, rel_tolerance=0, abs_tolerance=1e-5)
def global_to_child(self, location): """ Inverse transform """ delta = location - self.center if math.staticshape(location)[-1] == 2: angle = -math.batch_align(self.angle, 0, location) sin = math.sin(angle) cos = math.cos(angle) y, x = math.unstack(delta, axis=-1) if GLOBAL_AXIS_ORDER.is_x_first: x, y = y, x rot_x = cos * x - sin * y rot_y = sin * x + cos * y rotated = math.stack([rot_y, rot_x], axis=-1) elif math.staticshape(location)[-1] == 3: angle = math.batch_align(self.angle, 1, location) raise NotImplementedError( 'not yet implemented') # ToDo apply angle else: raise NotImplementedError('Rotation only supported in 2D and 3D') final = rotated + self.center return final
def test_plot_1d(self): self._test_plot( CenteredGrid(lambda x: math.sin(x), x=100, bounds=Box(x=2 * math.pi)))