def test_noise_scaling():
"""compare the noise strength (in terms of the spectral density of
two different noise sources that should be equivalent)"""
# create a grid
x, w = 2 + 10 * np.random.random(2)
size = np.random.randint(128, 256)
grid = CartesianGrid([[x, x + w]], size, periodic=True)
# colored noise
noise_colored = make_colored_noise(grid.shape, grid.discretization, exponent=2)
# divergence of white noise
shape = (grid.dim,) + grid.shape
div = grid.make_operator("divergence", bc="natural")
def noise_div():
return div(np.random.randn(*shape))
# calculate spectral densities of the two noises
result = []
for noise_func in [noise_colored, noise_div]:
def get_noise():
k, density = spectral_density(data=noise_func(), dx=grid.discretization)
assert k[0] == 0
assert density[0] == pytest.approx(0)
return np.log(density[1]) # log of spectral density
# average spectral density of longest length scale
mean = np.mean([get_noise() for _ in range(64)])
result.append(mean)
np.testing.assert_allclose(*result, rtol=0.5)
def test_unit_rect_grid(periodic):
"""test whether the rectangular grid behaves like a unit grid in special cases"""
dim = random.randrange(1, 4)
shape = np.random.randint(2, 10, size=dim)
g1 = UnitGrid(shape, periodic=periodic)
g2 = CartesianGrid(np.c_[np.zeros(dim), shape], shape, periodic=periodic)
volume = np.prod(shape)
for g in [g1, g2]:
assert g.volume == pytest.approx(volume)
assert g.integrate(1) == pytest.approx(volume)
assert g.make_integrator()(np.ones(shape)) == pytest.approx(volume)
assert g1.dim == g2.dim == dim
np.testing.assert_array_equal(g1.shape, g2.shape)
assert g1.typical_discretization == pytest.approx(
g2.typical_discretization)
for _ in range(10):
p1, p2 = np.random.normal(scale=10, size=(2, dim))
assert g1.distance_real(p1,
p2) == pytest.approx(g2.distance_real(p1, p2))
p0 = np.random.normal(scale=10, size=dim)
np.testing.assert_allclose(g1.polar_coordinates_real(p0),
g2.polar_coordinates_real(p0))
def test_boundary_interpolation_vector():
"""test boundary interpolation"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [3, 3])
field = VectorField.random_normal(grid)
# test boundary interpolation
bndry_val = np.random.randn(2, 3)
for bndry in grid._iter_boundaries():
val = field.get_boundary_values(*bndry, bc={"value": bndry_val})
np.testing.assert_allclose(val, bndry_val)
def test_degenerated_grid():
"""test operators on grids with singular dimensions"""
g1 = CartesianGrid([[0, 1]], 4)
g2 = CartesianGrid([[0, 1], [0, 0.1]], [4, 1], periodic=[False, True])
f1 = ScalarField.random_uniform(g1)
f2 = ScalarField(g2, f1.data.reshape(g2.shape))
res1 = f1.laplace("auto_periodic_neumann").data
res2 = f2.laplace("auto_periodic_neumann").data
np.testing.assert_allclose(res1.flat, res2.flat)
def test_rect_div_grad():
"""compare div grad to laplacian"""
grid = CartesianGrid([[0, 2 * np.pi], [0, 2 * np.pi]], [16, 16], periodic=True)
x, y = grid.cell_coords[..., 0], grid.cell_coords[..., 1]
field = ScalarField(grid, data=np.cos(x) + np.sin(y))
bcs = grid.get_boundary_conditions("auto_periodic_neumann")
a = field.laplace(bcs)
b = field.gradient(bcs).divergence("auto_periodic_curvature")
np.testing.assert_allclose(a.data, -field.data, rtol=0.05, atol=0.01)
np.testing.assert_allclose(b.data, -field.data, rtol=0.05, atol=0.01)
def test_expression_bc_derivative(dim):
"""test boundary conditions that use an expression to calculate the derivative"""
grid = CartesianGrid([[0, 1]] * dim, 4)
bc1 = grid.get_boundary_conditions({"derivative": 0})
bc2 = grid.get_boundary_conditions({"derivative_expression": "0"})
field = ScalarField(grid, 1)
result = field.laplace(bc1)
np.testing.assert_allclose(result.data, 0.0)
result = field.laplace(bc2)
np.testing.assert_allclose(result.data, 0.0)
def test_expression_bc_operator(dim):
"""test boundary conditions that use an expression in an operator"""
grid = CartesianGrid([[0, 1]] * dim, 4)
bc1 = grid.get_boundary_conditions({"value": 1})
bc2 = grid.get_boundary_conditions({"virtual_point": f"2 - value"})
field = ScalarField(grid, 1)
result = field.laplace(bc1)
np.testing.assert_allclose(result.data, 0.0)
result = field.laplace(bc2)
np.testing.assert_allclose(result.data, 0.0)
def test_expression_bc_value(dim):
"""test the boundary conditions that calculate the virtual point directly"""
grid = CartesianGrid([[0, 1]] * dim, 4)
bc1 = grid.get_boundary_conditions({"value": 1})
bc2 = grid.get_boundary_conditions({"value_expression": "1"})
field = ScalarField(grid, 1)
result = field.laplace(bc1)
np.testing.assert_allclose(result.data, 0.0)
result = field.laplace(bc2)
np.testing.assert_allclose(result.data, 0.0)
def test_div_grad_const():
"""compare div grad to laplace operator"""
grid = CartesianGrid([[-1, 1]], 32)
# test constant
y = ScalarField(grid, 3)
for bc in [{"type": "derivative", "value": 0}, {"type": "value", "value": 3}]:
bcs = grid.get_boundary_conditions(bc)
lap = y.laplace(bcs)
divgrad = y.gradient(bcs).divergence("auto_periodic_curvature")
np.testing.assert_allclose(lap.data, np.zeros(32))
np.testing.assert_allclose(divgrad.data, np.zeros(32))
def test_expression_bc_value(dim):
"""test boundary conditions that use an expression to calculate the value"""
grid = CartesianGrid([[0, 1]] * dim, 4)
bc1 = grid.get_boundary_conditions({"value": 1})
bc2 = grid.get_boundary_conditions({"value_expression": "1"})
assert "field" in bc2.get_mathematical_representation("field")
field = ScalarField(grid, 1)
result = field.laplace(bc1)
np.testing.assert_allclose(result.data, 0.0)
result = field.laplace(bc2)
np.testing.assert_allclose(result.data, 0.0)
def test_div_grad_quadratic():
""" compare div grad to laplace operator """
grid = CartesianGrid([[-1, 1]], 32)
x = grid.axes_coords[0]
# test simple quadratic
y = ScalarField(grid, x**2)
bcs = grid.get_boundary_conditions({"type": "derivative", "value": 2})
lap = y.laplace(bcs)
divgrad = y.gradient(bcs).divergence(bcs.differentiated)
np.testing.assert_allclose(lap.data, np.full(32, 2.0))
np.testing.assert_allclose(divgrad.data, np.full(32, 2.0))
def test_pde_poisson_solver_1d():
"""test the poisson solver on 1d grids"""
# solve Laplace's equation
grid = UnitGrid([4])
res = solve_laplace_equation(grid, bc=[{"value": -1}, {"value": 3}])
np.testing.assert_allclose(res.data, grid.axes_coords[0] - 1)
res = solve_laplace_equation(grid, bc=[{"value": -1}, {"derivative": 1}])
np.testing.assert_allclose(res.data, grid.axes_coords[0] - 1)
# test Poisson equation with 2nd Order BC
res = solve_laplace_equation(grid, bc=[{"value": -1}, "extrapolate"])
# solve Poisson's equation
grid = CartesianGrid([[0, 1]], 4)
field = ScalarField(grid, data=1)
res = solve_poisson_equation(field, bc=[{"value": 1}, {"derivative": 1}])
xs = grid.axes_coords[0]
np.testing.assert_allclose(res.data, 1 + 0.5 * xs**2, rtol=1e-2)
# test inconsistent problem
field.data = 1
with pytest.raises(RuntimeError, match="Neumann"):
solve_poisson_equation(field, {"derivative": 0})
def test_vector_gradient_field():
"""test the vector gradient operator"""
grid = CartesianGrid([[0, 2 * np.pi], [0, 2 * np.pi]], [16, 16],
periodic=True)
x, y = grid.cell_coords[..., 0], grid.cell_coords[..., 1]
data = [np.cos(x) + y, np.sin(y) - x]
v = VectorField(grid, data)
t1 = v.gradient("periodic")
assert t1.data.shape == (2, 2, 16, 16)
d00 = -np.sin(x)
d10 = -np.ones(grid.shape)
d01 = np.ones(grid.shape)
d11 = np.cos(y)
t2 = Tensor2Field(grid, np.array([[d00, d01], [d10, d11]]))
np.testing.assert_allclose(t1.data[1:-1, 1:-1],
t2.data[1:-1, 1:-1],
rtol=0.1,
atol=0.1)
v.gradient("auto_periodic_neumann", out=t1)
assert t1.data.shape == (2, 2, 16, 16)
np.testing.assert_allclose(t1.data[1:-1, 1:-1],
t2.data[1:-1, 1:-1],
rtol=0.1,
atol=0.1)
def test_generic_cartesian_grid():
"""test generic cartesian grid functions"""
for dim in (1, 2, 3):
periodic = random.choices([True, False], k=dim)
shape = np.random.randint(2, 8, size=dim)
a = np.random.random(dim)
b = a + np.random.random(dim)
cases = [
UnitGrid(shape, periodic=periodic),
CartesianGrid(np.c_[a, b], shape, periodic=periodic),
]
for grid in cases:
assert grid.dim == dim
dim_axes = len(grid.axes) + len(grid.axes_symmetric)
assert dim_axes == dim
vol = np.prod(grid.discretization) * np.prod(shape)
assert grid.volume == pytest.approx(vol)
assert grid.uniform_cell_volumes
# random points
points = [[np.random.uniform(a[i], b[i]) for i in range(dim)]
for _ in range(10)]
c = grid.point_to_cell(points)
p = grid.cell_to_point(c)
np.testing.assert_array_equal(c, grid.point_to_cell(p))
assert grid.contains_point(grid.get_random_point())
w = 0.499 * (b - a).min()
assert grid.contains_point(grid.get_random_point(w))
assert "laplace" in grid.operators
def test_tensor_symmetrize():
"""test advanced tensor calculations"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [2, 2])
t1 = Tensor2Field(grid)
t1.data[0, 0, :] = 1
t1.data[0, 1, :] = 2
t1.data[1, 0, :] = 3
t1.data[1, 1, :] = 4
# traceless = False
t2 = t1.copy()
t1.symmetrize(make_traceless=False, inplace=True)
tr = t1.trace()
assert np.all(tr.data == 5)
t1_trans = np.swapaxes(t1.data, 0, 1)
np.testing.assert_allclose(t1.data, t1_trans.data)
ts = t1.copy()
ts.symmetrize(make_traceless=False, inplace=True)
np.testing.assert_allclose(t1.data, ts.data)
# traceless = True
t2.symmetrize(make_traceless=True, inplace=True)
tr = t2.trace()
assert np.all(tr.data == 0)
t2_trans = np.swapaxes(t2.data, 0, 1)
np.testing.assert_allclose(t2.data, t2_trans.data)
ts = t2.copy()
ts.symmetrize(make_traceless=True, inplace=True)
np.testing.assert_allclose(t2.data, ts.data)
def test_complex_vectors():
"""test some complex vector fields"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [3, 4])
shape = (2, 2) + grid.shape
numbers = np.random.random(shape) + np.random.random(shape) * 1j
v1 = VectorField(grid, numbers[0])
v2 = VectorField(grid, numbers[1])
assert v1.is_complex and v2.is_complex
for backend in ["numpy", "numba"]:
dot_op = v1.make_dot_operator(backend)
# test complex conjugate
expected = v1.to_scalar("norm_squared").data
np.testing.assert_allclose((v1 @ v1).data, expected)
np.testing.assert_allclose(dot_op(v1.data, v1.data), expected)
# test dot product
res = dot_op(v1.data, v2.data)
for s in (v1 @ v2, (v2 @ v1).conjugate(), v1.dot(v2)):
assert isinstance(s, ScalarField)
assert s.grid is grid
np.testing.assert_allclose(s.data, res)
# test without conjugate
dot_op = v1.make_dot_operator(backend, conjugate=False)
res = v1.dot(v2, conjugate=False)
np.testing.assert_allclose(dot_op(v1.data, v2.data), res.data)
def test_make_derivative(ndim, axis):
""" test the make derivative function """
periodic = random.choice([True, False])
grid = CartesianGrid([[0, 6 * np.pi]] * ndim, 16, periodic=periodic)
field = ScalarField.random_harmonic(grid, modes=1, axis_combination=np.add)
bcs = grid.get_boundary_conditions("natural")
grad = field.gradient(bcs)
for method in ["central", "forward", "backward"]:
msg = f"method={method}, periodic={periodic}"
diff = ops._make_derivative(bcs, axis=axis, method=method)
np.testing.assert_allclose(grad.data[axis],
diff(field.data),
atol=0.1,
rtol=0.1,
err_msg=msg)
def test_make_derivative2(ndim, axis):
"""test the _make_derivative2 function"""
periodic = random.choice([True, False])
grid = CartesianGrid([[0, 6 * np.pi]] * ndim, 16, periodic=periodic)
field = ScalarField.random_harmonic(grid, modes=1, axis_combination=np.add)
bcs = grid.get_boundary_conditions("auto_periodic_neumann")
grad = field.gradient(bcs)[axis]
grad2 = grad.gradient(bcs)[axis]
diff = ops._make_derivative2(grid, axis=axis)
res = field.copy()
res.data[:] = 0
field.set_ghost_cells(bcs)
diff(field._data_full, out=res.data)
np.testing.assert_allclose(grad2.data, res.data, atol=0.1, rtol=0.1)
def test_vector_boundary_conditions():
""" test some boundary conditions of operators of vector fields """
grid = CartesianGrid([[0, 2 * np.pi], [0, 1]], 32, periodic=False)
v_x = np.sin(grid.cell_coords[..., 0])
v_y = grid.cell_coords[..., 1]
vf = VectorField(grid, np.array([v_x, v_y]))
bc_x = [
{
"type": "derivative",
"value": [0, -1]
},
{
"type": "derivative",
"value": [0, 1]
},
]
bc_y = [{
"type": "value",
"value": [0, 0]
}, {
"type": "value",
"value": [1, 1]
}]
tf = vf.gradient(bc=[bc_x, bc_y])
np.testing.assert_allclose(tf[0, 1].data[1:-1, :], 0)
np.testing.assert_allclose(tf[1, 1].data, 1)
def _get_random_grid_bcs(ndim: int, dx="random", periodic="random", rank=0):
""" create a random Cartesian grid with natural bcs """
shape = np.random.randint(2, 5, ndim)
if dx == "random":
dx = np.random.uniform(0.5, 1.5, ndim)
elif dx == "uniform":
dx = np.full(ndim, np.random.uniform(0.5, 1.5))
else:
dx = np.broadcast_to(dx, shape)
if periodic == "random":
periodic = random.choice([True, False])
sizes = [(0, float(s * d)) for s, d in zip(shape, dx)]
grid = CartesianGrid(sizes, shape, periodic=periodic)
return grid.get_boundary_conditions("natural", rank=rank)
def integrate(tfinal, u, u2, udiff, IC):
grid = CartesianGrid([[0, 200]], 200)
field = ScalarField(grid, IC)
storage = MemoryStorage()
eq = libraryPDE(u, u2, udiff) #define PDE with custom parameters
return eq.solve(field,
t_range=tfinal,
dt=0.1,
tracker=storage.tracker(0.1)).data
def test_make_derivative(ndim, axis):
"""test the _make_derivative function"""
periodic = random.choice([True, False])
grid = CartesianGrid([[0, 6 * np.pi]] * ndim, 16, periodic=periodic)
field = ScalarField.random_harmonic(grid, modes=1, axis_combination=np.add)
bcs = grid.get_boundary_conditions("auto_periodic_neumann")
grad = field.gradient(bcs)
for method in ["central", "forward", "backward"]:
msg = f"method={method}, periodic={periodic}"
diff = ops._make_derivative(grid, axis=axis, method=method)
res = field.copy()
res.data[:] = 0
field.set_ghost_cells(bcs)
diff(field._data_full, out=res.data)
np.testing.assert_allclose(
grad.data[axis], res.data, atol=0.1, rtol=0.1, err_msg=msg
)
def iter_grids():
"""generator providing some test grids"""
for periodic in [True, False]:
yield UnitGrid([3], periodic=periodic)
yield UnitGrid([3, 3, 3], periodic=periodic)
yield CartesianGrid([[-1, 2], [0, 3]], [5, 7], periodic=periodic)
yield CylindricalSymGrid(3, [-1, 2], [7, 8], periodic_z=periodic)
yield PolarSymGrid(3, 4)
yield SphericalSymGrid(3, 4)
def test_simple_diffusion_flux_right():
"""test a simple diffusion equation with flux boundary on the right"""
grid = CartesianGrid([[0, 1]], [16])
c = ScalarField.random_uniform(grid, 0, 1)
b_l = {"type": "value", "value": 0}
b_r = {"type": "derivative", "value": 3}
pde = DiffusionPDE(bc=[b_l, b_r])
sol = pde.solve(c, t_range=5, dt=0.001, tracker=None)
np.testing.assert_allclose(sol.data, 3 * grid.axes_coords[0], rtol=5e-3)
def test_tensors_basic():
"""test some tensor calculations"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [3, 4])
t1 = Tensor2Field(grid, np.full((2, 2) + grid.shape, 1))
t2 = Tensor2Field(grid, np.full((2, 2) + grid.shape, 2))
np.testing.assert_allclose(t2.average, [[2, 2], [2, 2]])
assert t1.magnitude == pytest.approx(2)
assert t1["x", "x"] == t1[0, 0]
assert t1["x", 1] == t1[0, "y"] == t1[0, 1]
t1[0, 0] = t1[0, 0]
t3 = t1 + t2
assert t3.grid == grid
np.testing.assert_allclose(t3.data, 3)
t1 += t2
np.testing.assert_allclose(t1.data, 3)
field = Tensor2Field.random_uniform(grid)
trace = field.trace()
assert isinstance(trace, ScalarField)
np.testing.assert_allclose(trace.data, field.data.trace())
t1 = Tensor2Field(grid)
t1[0, 0] = 1
t1[0, 1] = 2
t1[1, 0] = 3
t1[1, 1] = 4
for method, value in [
("min", 1),
("max", 4),
("norm", np.linalg.norm([[1, 2], [3, 4]])),
("squared_sum", 30),
("norm_squared", 30),
("trace", 5),
("invariant1", 5),
("invariant2", -1),
]:
p1 = t1.to_scalar(method)
assert p1.data.shape == grid.shape
np.testing.assert_allclose(p1.data, value)
for idx in ((1, ), (1, 2, 3), (1.5, 2), ("a", "b"), 1.0):
with pytest.raises(IndexError):
t1[idx]
t2 = FieldBase.from_state(t1.attributes, data=t1.data)
assert t1 == t2
assert t1.grid is t2.grid
attrs = Tensor2Field.unserialize_attributes(t1.attributes_serialized)
t2 = FieldBase.from_state(attrs, data=t1.data)
assert t1 == t2
assert t1.grid is not t2.grid
def test_simple_diffusion_value():
"""test a simple diffusion equation with constant boundaries"""
grid = CartesianGrid([[0, 1]], [16])
c = ScalarField.random_uniform(grid, 0, 1)
b_l = {"type": "value", "value": 0}
b_r = {"type": "value", "value": 1}
pde = DiffusionPDE(bc=[b_l, b_r])
sol, info = pde.solve(c, t_range=1, dt=0.001, tracker=None, ret_info=True)
assert isinstance(info, dict)
np.testing.assert_allclose(sol.data, grid.axes_coords[0], rtol=5e-3)
def test_inhomogeneous_bcs():
"""test simulation with inhomogeneous boundary conditions"""
# single coordinate
grid = CartesianGrid([[0, 2 * np.pi], [0, 1]], [32, 2], periodic=[True, False])
state = ScalarField(grid)
pde = DiffusionPDE(bc=["natural", {"type": "value", "value": "sin(x)"}])
sol = pde.solve(state, t_range=1e1, dt=1e-2, tracker=None)
data = sol.get_line_data(extract="project_x")
np.testing.assert_almost_equal(
data["data_y"], 0.9 * np.sin(data["data_x"]), decimal=2
)
# double coordinate
grid = CartesianGrid([[0, 1], [0, 1]], [8, 8], periodic=False)
state = ScalarField(grid)
pde = DiffusionPDE(bc={"type": "value", "value": "x + y"})
sol = pde.solve(state, t_range=1e1, dt=1e-3, tracker=None)
expect = ScalarField.from_expression(grid, "x + y")
np.testing.assert_almost_equal(sol.data, expect.data)
def test_normalize_point(reflect):
"""test normalize_point method for Cartesian Grids"""
grid = CartesianGrid([[1, 3]], [1], periodic=False)
norm_numba = grid.make_normalize_point_compiled(reflect=reflect)
def norm_numba_wrap(x):
y = np.array([x])
norm_numba(y)
return y
if reflect:
values = [(-2, 2), (0, 2), (1, 1), (2, 2), (3, 3), (4, 2), (5, 1), (6, 2)]
else:
values = [(-2, -2), (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]
for norm in [norm_numba_wrap, partial(grid.normalize_point, reflect=reflect)]:
for x, y in values:
assert norm(x) == pytest.approx(y), (norm, x)
def test_spherical_to_cartesian():
""" test conversion of spherical grid to cartesian """
expr_sph = "1 / (1 + r**2)"
expr_cart = expr_sph.replace("r**2", "(x**2 + y**2 + z**2)")
grid_sph = SphericalGrid(7, 16)
pf_sph = ScalarField.from_expression(grid_sph, expression=expr_sph)
grid_cart = CartesianGrid([[-4, 4], [-3.9, 4.1], [-4.1, 3.9]], [16] * 3)
pf_cart1 = pf_sph.interpolate_to_grid(grid_cart)
pf_cart2 = ScalarField.from_expression(grid_cart, expression=expr_cart)
np.testing.assert_allclose(pf_cart1.data, pf_cart2.data, atol=0.1)
def test_polar_to_cartesian():
""" test conversion of polar grid to Cartesian """
expr_pol = "1 / (1 + r**2)"
expr_cart = expr_pol.replace("r**2", "(x**2 + y**2)")
grid_pol = PolarGrid(7, 16)
pf_pol = ScalarField.from_expression(grid_pol, expression=expr_pol)
grid_cart = CartesianGrid([[-4, 4], [-3.9, 4.1]], [16, 16])
pf_cart1 = pf_pol.interpolate_to_grid(grid_cart)
pf_cart2 = ScalarField.from_expression(grid_cart, expression=expr_cart)
np.testing.assert_allclose(pf_cart1.data, pf_cart2.data, atol=0.1)
