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