def test_assert_close(self):
a = spatial(a=10)
math.assert_close(math.zeros(a), math.zeros(a), math.zeros(a), rel_tolerance=0, abs_tolerance=0)
assert_not_close(math.zeros(a), math.ones(a), rel_tolerance=0, abs_tolerance=0)
for scale in (1, 0.1, 10):
math.assert_close(math.zeros(a), math.ones(a) * scale, rel_tolerance=0, abs_tolerance=scale * 1.001)
math.assert_close(math.zeros(a), math.ones(a) * scale, rel_tolerance=1, abs_tolerance=0)
assert_not_close(math.zeros(a), math.ones(a) * scale, rel_tolerance=0.9, abs_tolerance=0)
assert_not_close(math.zeros(a), math.ones(a) * scale, rel_tolerance=0, abs_tolerance=0.9 * scale)
with math.precision(64):
assert_not_close(math.zeros(a), math.ones(a) * 1e-100, rel_tolerance=0, abs_tolerance=0)
math.assert_close(math.zeros(a), math.ones(a) * 1e-100, rel_tolerance=0, abs_tolerance=1e-15)
def test__periodic_2d_arakawa_poisson_bracket(self):
"""test _periodic_2d_arakawa_poisson_bracket implementation"""
with math.precision(64):
# Define parameters to test
test_params = {
'grid_size': [(4, 4), (32, 32)],
'dx': [0.1, 1],
'gen_func': [
lambda grid_size: np.random.rand(*grid_size).reshape(
grid_size)
]
}
# Generate test cases as the product
test_cases = [
dict(zip(test_params, v))
for v in product(*test_params.values())
]
for params in test_cases:
grid_size = params['grid_size']
d1 = params['gen_func'](grid_size)
d2 = params['gen_func'](grid_size)
dx = params['dx']
padding = extrapolation.PERIODIC
ref = self.arakawa_reference_implementation(
np.pad(d1.copy(), 1, mode='wrap'),
np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1]
d1_tensor = field.CenteredGrid(
values=math.tensor(d1, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
d2_tensor = field.CenteredGrid(
values=math.tensor(d2, names=['x', 'y']),
bounds=geom.Box([0, 0], list(grid_size)),
extrapolation=padding)
val = math._nd._periodic_2d_arakawa_poisson_bracket(
d1_tensor.values, d2_tensor.values, dx)
try:
math.assert_close(ref,
val,
rel_tolerance=1e-14,
abs_tolerance=1e-14)
except BaseException as e:
abs_error = math.abs(val - ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(ref)
print(val)
raise AssertionError(e, params, variation_str)
def test_fft(self):
def get_2d_sine(grid_size, L):
indices = np.array(np.meshgrid(*list(map(range, grid_size))))
phys_coord = indices.T * L / (grid_size[0]) # between [0, L)
x, y = phys_coord.T
d = np.sin(2 * np.pi * x + 1) * np.sin(2 * np.pi * y + 1)
return d
sine_field = get_2d_sine((32, 32), L=2)
fft_ref_tensor = math.wrap(np.fft.fft2(sine_field), 'x,y')
with math.precision(64):
for backend in BACKENDS:
with backend:
sine_tensor = math.tensor(sine_field, 'x,y')
fft_tensor = math.fft(sine_tensor)
math.assert_close(
fft_ref_tensor, fft_tensor, abs_tolerance=1e-12
) # Should usually be more precise. GitHub Actions has larger errors than usual.
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 test_fft(self):
def get_2d_sine(grid_size, L):
indices = np.array(np.meshgrid(*list(map(range, grid_size))))
phys_coord = indices.T * L / (grid_size[0]) # between [0, L)
x, y = phys_coord.T
d = np.sin(2 * np.pi * x + 1) * np.sin(2 * np.pi * y + 1)
return d
sine_field = get_2d_sine((32, 32), L=2)
fft_ref_tensor = math.wrap(np.fft.fft2(sine_field), spatial('x,y'))
with math.precision(64):
for backend in BACKENDS:
if backend.name != 'Jax': # TODO Jax casts to float32 / complex64 on GitHub Actions
with backend:
sine_tensor = math.tensor(sine_field, spatial('x,y'))
fft_tensor = math.fft(sine_tensor)
self.assertEqual(fft_tensor.dtype, math.DType(complex, 128), msg=backend.name)
math.assert_close(fft_ref_tensor, fft_tensor, abs_tolerance=1e-12, msg=backend.name) # Should usually be more precise. GitHub Actions has larger errors than usual.
def test__periodic_2d_arakawa_poisson_bracket(self):
"""test _periodic_2d_arakawa_poisson_bracket implementation"""
with math.precision(64):
# Define parameters to test
test_params = {
'grid_size': [(4, 4), (32, 32)],
'dx': [0.1, 1],
'gen_func': [lambda grid_size: np.random.rand(*grid_size).reshape(grid_size)]
}
# Generate test cases as the product
test_cases = [dict(zip(test_params, v)) for v in product(*test_params.values())]
for params in test_cases:
grid_size = params['grid_size']
d1 = params['gen_func'](grid_size)
d2 = params['gen_func'](grid_size)
dx = params['dx']
padding = extrapolation.PERIODIC
ref = self.arakawa_reference_implementation(np.pad(d1.copy(), 1, mode='wrap'), np.pad(d2.copy(), 1, mode='wrap'), dx)[1:-1, 1:-1]
d1_tensor = field.CenteredGrid(values=math.tensor(d1, spatial('x,y')), bounds=geom.Box([0, 0], list(grid_size)), extrapolation=padding)
d2_tensor = field.CenteredGrid(values=math.tensor(d2, spatial('x,y')), bounds=geom.Box([0, 0], list(grid_size)), extrapolation=padding)
val = math._nd._periodic_2d_arakawa_poisson_bracket(d1_tensor.values, d2_tensor.values, dx)
# try:
math.assert_close(ref, val, rel_tolerance=1e-14, abs_tolerance=1e-14)
from .numpy_reference import Namespace
from phi import math, field
with math.precision(64):
def get_domain_phi(plasma, domain):
return domain.grid(math.fourier_poisson(plasma.omega.values, plasma.dx))
def step_gradient_2d(params, plasma, phi, dt=0):
"""time gradient of model"""
# Diffusion function
def diffuse(arr, N, dx):
for i in range(N):
arr = math.fourier_laplace(arr, dx) # field.laplace(arr)
return arr
def step_gradient_2d(params, plasma, phi, dt=0):
"""time spatial_gradient of model"""
# Diffusion function
def diffuse(arr, N, dx):
for i in range(N):
arr = field.laplace(arr) # math.fourier_laplace(arr, dx)
return arr
# Calculate Gradients
dx_p, dy_p = field.spatial_gradient(phi).vector.unstack()
# Get difference
diff = phi - plasma.density
# Step 2.1: New Omega.
nu = (-1) ** (params["N"] + 1) * params["nu"]