def test_evaluate_func_scalar():
"""test the evaluate function with scalar fields"""
grid = UnitGrid([2, 4])
field = ScalarField.from_expression(grid, "x")
res1 = evaluate("laplace(a)", {"a": field})
res2 = field.laplace("natural")
np.testing.assert_almost_equal(res1.data, res2.data)
res = evaluate("a - x", {"a": field})
np.testing.assert_almost_equal(res.data, 0)
field_c = ScalarField.from_expression(grid, "1 + 1j")
res = evaluate("c", {"c": field_c})
assert res.dtype == "complex"
np.testing.assert_almost_equal(res.data, field_c.data)
res = evaluate("abs(a)**2", {"a": field_c})
assert res.dtype == "float"
np.testing.assert_allclose(res.data, 2)
# edge case
res = evaluate("sin", {"a": field_c}, consts={"sin": 3.14})
assert res.dtype == "float"
np.testing.assert_allclose(res.data, 3.14)
def test_field_labels():
""" test the FieldCollection.labels property """
grid = UnitGrid([5])
s1 = ScalarField(grid, label="s1")
s2 = ScalarField(grid)
fc = FieldCollection([s1, s2])
assert fc.labels == ["s1", None]
assert len(fc.labels) == 2
assert fc.labels[0] == "s1"
fc.labels = ["a", "b"]
assert fc.labels == ["a", "b"]
fc.labels[0] = "c"
assert fc.labels == ["c", "b"]
assert str(fc.labels) == str(["c", "b"])
assert repr(fc.labels) == repr(["c", "b"])
assert fc.labels[0:1] == ["c"]
assert fc.labels[:] == ["c", "b"]
fc.labels[0:1] = "d"
assert fc.labels == ["d", "b"]
fc.labels[:] = "a"
assert fc.labels == ["a", "a"]
labels = fc.labels[:]
labels[0] = "e"
assert fc.labels == ["a", "a"]
fc = FieldCollection([s1, s2], labels=[None, "b"])
assert fc.labels == [None, "b"]
fc = FieldCollection([s1, s2], labels=["a", "b"])
assert fc.labels == ["a", "b"]
def test_emulsion_plotting():
"""test plotting emulsions"""
# 1d emulsion
e1 = Emulsion([DiffuseDroplet([1], 10, 0.5)] * 2)
with pytest.raises(NotImplementedError):
e1.plot()
# 2d emulsion
field = ScalarField(UnitGrid([10, 10], periodic=True))
es = [
Emulsion([DiffuseDroplet([0, 1], 10, 0.5)] * 2),
Emulsion([droplets.PerturbedDroplet2D([0, 1], 3, 1, [1, 2, 3, 4])]),
]
for e2 in es:
e2.plot()
e2.plot(field=field)
# 3d emulsion
field = ScalarField(UnitGrid([5, 5, 5], periodic=True))
e3 = Emulsion([DiffuseDroplet([0, 1, 2], 10, 0.5)] * 2)
e3.plot()
e3.plot(field=field)
e3 = Emulsion([droplets.PerturbedDroplet3D([0, 1, 2], 3, 1, [1, 2, 3, 4, 5, 6])])
with pytest.raises(NotImplementedError):
e3.plot()
with pytest.raises(NotImplementedError):
Emulsion().plot()
def test_examples_scalar_polar():
"""compare derivatives of scalar fields for polar grids"""
grid = PolarSymGrid(1, 32)
sf = ScalarField.from_expression(grid, "r**3")
# gradient
res = sf.gradient([{"derivative": 0}, {"derivative": 3}])
expect = VectorField.from_expression(grid, ["3 * r**2", 0])
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# gradient squared
expect = ScalarField.from_expression(grid, "9 * r**4")
for c in [True, False]:
res = sf.gradient_squared([{
"derivative": 0
}, {
"derivative": 3
}],
central=c)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# laplace
res = sf.laplace([{"derivative": 0}, {"derivative": 3}])
expect = ScalarField.from_expression(grid, "9 * r")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
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_examples_scalar_cyl():
"""compare derivatives of scalar fields for cylindrical grids"""
grid = CylindricalSymGrid(1, [0, 2 * np.pi], 32)
expr = "r**3 * sin(z)"
sf = ScalarField.from_expression(grid, expr)
bcs = [[{
"derivative": 0
}, {
"value": expr
}], [{
"value": expr
}, {
"value": expr
}]]
# gradient - The coordinates are ordered as (r, z, φ) in py-pde
res = sf.gradient(bcs)
expect = VectorField.from_expression(
grid, ["3 * r**2 * sin(z)", "r**3 * cos(z)", 0])
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# gradient squared
expect = ScalarField.from_expression(
grid, "r**6 * cos(z)**2 + 9 * r**4 * sin(z)**2")
res = sf.gradient_squared(bcs, central=True)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# laplace
bcs[0][1] = {
"curvature": "6 * sin(z)"
} # adjust BC to fit laplacian better
res = sf.laplace(bcs)
expect = ScalarField.from_expression(grid,
"9 * r * sin(z) - r**3 * sin(z)")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
def test_poisson_solver_polar():
""" test the poisson solver on Polar grids """
grid = PolarGrid(4, 8)
for bc_val in ["natural", {"value": 1}]:
bcs = grid.get_boundary_conditions(bc_val)
poisson = grid.get_operator("poisson_solver", bcs)
laplace = grid.get_operator("laplace", bcs)
d = np.random.random(grid.shape)
d -= ScalarField(grid, d).average # balance the right hand side
np.testing.assert_allclose(laplace(poisson(d)),
d,
err_msg=f"bcs = {bc_val}")
grid = PolarGrid([2, 4], 8)
for bc_val in ["natural", {"value": 1}]:
bcs = grid.get_boundary_conditions(bc_val)
poisson = grid.get_operator("poisson_solver", bcs)
laplace = grid.get_operator("laplace", bcs)
d = np.random.random(grid.shape)
d -= ScalarField(grid, d).average # balance the right hand side
np.testing.assert_allclose(laplace(poisson(d)),
d,
err_msg=f"bcs = {bc_val}")
def test_laplacian_field_cyl():
""" test the gradient operator """
grid = CylindricalGrid(2 * np.pi, [0, 2 * np.pi], [8, 16], periodic_z=True)
r, z = grid.cell_coords[..., 0], grid.cell_coords[..., 1]
s = ScalarField(grid, data=np.cos(r) + np.sin(z))
s_lap = s.laplace(bc="natural")
assert s_lap.data.shape == (8, 16)
res = -np.cos(r) - np.sin(r) / r - np.sin(z)
np.testing.assert_allclose(s_lap.data, res, rtol=0.1, atol=0.1)
def test_gradient_field_cyl():
""" test the gradient operator"""
grid = CylindricalGrid(2 * np.pi, [0, 2 * np.pi], [8, 16], periodic_z=True)
r, z = grid.cell_coords[..., 0], grid.cell_coords[..., 1]
s = ScalarField(grid, data=np.cos(r) + np.sin(z))
v = s.gradient(bc="natural")
assert v.data.shape == (3, 8, 16)
np.testing.assert_allclose(v.data[0], -np.sin(r), rtol=0.1, atol=0.1)
np.testing.assert_allclose(v.data[1], np.cos(z), rtol=0.1, atol=0.1)
np.testing.assert_allclose(v.data[2], 0, rtol=0.1, atol=0.1)
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_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_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_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_grid_div_grad_sph():
"""compare div grad to laplacian"""
grid = SphericalSymGrid(2 * np.pi, 16)
field = ScalarField.from_expression(grid, "cos(r)")
a = field.laplace("derivative")
b = field.gradient("derivative").divergence("value")
res = ScalarField.from_expression(grid, "-2 * sin(r) / r - cos(r)")
# do not test the radial boundary points
np.testing.assert_allclose(a.data[1:-1], res.data[1:-1], rtol=0.1, atol=0.1)
np.testing.assert_allclose(b.data[1:-1], res.data[1:-1], rtol=0.1, atol=0.1)
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_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)
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_storage_persistence(compression, tmp_path):
"""test writing to persistent trackers"""
dim = 5
grid = UnitGrid([dim])
field = ScalarField(grid)
path = tmp_path / f"test_storage_persistence_{compression}.hdf5"
# write some data
for write_mode in ["append", "truncate_once", "truncate"]:
storage = FileStorage(path,
info={"a": 1},
write_mode=write_mode,
compression=compression)
# first batch
storage.start_writing(field, info={"b": 2})
field.data = np.arange(dim)
storage.append(field, 0)
field.data = np.arange(dim, 2 * dim)
storage.append(field)
storage.end_writing()
# read first batch
np.testing.assert_array_equal(storage.times, np.arange(2))
np.testing.assert_array_equal(np.ravel(storage.data), np.arange(10))
assert {"a": 1, "b": 2}.items() <= storage.info.items()
# second batch
storage.start_writing(field, info={"c": 3})
field.data = np.arange(2 * dim, 3 * dim)
storage.append(field, 2)
storage.end_writing()
storage.close()
# read the data
storage = FileStorage(path)
if write_mode == "truncate":
np.testing.assert_array_equal(storage.times, np.array([2]))
np.testing.assert_array_equal(np.ravel(storage.data),
np.arange(10, 15))
assert storage.shape == (1, 5)
info = {"c": 3}
assert info.items() <= storage.info.items()
else:
np.testing.assert_array_equal(storage.times, np.arange(0, 3))
np.testing.assert_array_equal(np.ravel(storage.data),
np.arange(0, 15))
assert storage.shape == (3, 5)
info = {"a": 1, "b": 2, "c": 3}
assert info.items() <= storage.info.items()
def test_singular_dimensions_3d(periodic):
"""test grids with singular dimensions"""
dim = np.random.randint(3, 5)
g1 = UnitGrid([dim], periodic=periodic)
g3a = UnitGrid([dim, 1, 1], periodic=periodic)
g3b = UnitGrid([1, 1, dim], periodic=periodic)
field = ScalarField.random_uniform(g1)
expected = field.laplace("auto_periodic_neumann").data
for g in [g3a, g3b]:
f = ScalarField(g, data=field.data.reshape(g.shape))
res = f.laplace("auto_periodic_neumann").data.reshape(g1.shape)
np.testing.assert_allclose(expected, res)
def test_grid_div_grad_cyl():
"""compare div grad to laplacian"""
grid = CylindricalSymGrid(2 * np.pi, (0, 2 * np.pi), (16, 16), periodic_z=True)
field = ScalarField.from_expression(grid, "cos(r) + sin(z)")
bcs = grid.get_boundary_conditions()
a = field.laplace(bcs)
c = field.gradient(bcs)
b = c.divergence(bcs.differentiated)
res = ScalarField.from_expression(grid, "-sin(r)/r - cos(r) - sin(z)")
# do not test the radial boundary points
np.testing.assert_allclose(a.data[1:-1], res.data[1:-1], rtol=0.1, atol=0.05)
np.testing.assert_allclose(b.data[1:-1], res.data[1:-1], rtol=0.1, atol=0.05)
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_cylindrical_to_cartesian():
"""test conversion of cylindrical grid to Cartesian"""
expr_cyl = "cos(z / 2) / (1 + r**2)"
expr_cart = expr_cyl.replace("r**2", "(x**2 + y**2)")
z_range = (-np.pi, 2 * np.pi)
grid_cyl = CylindricalSymGrid(10, z_range, (16, 33))
pf_cyl = ScalarField.from_expression(grid_cyl, expression=expr_cyl)
grid_cart = CartesianGrid([[-7, 7], [-6, 7], z_range], [16, 16, 16])
pf_cart1 = pf_cyl.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_pde_spatial_args():
""" test ScalarFieldExpression without extra dependence """
eq = PDE({"a": "x"})
field = ScalarField(grids.UnitGrid([2]))
rhs = eq.evolution_rate(field)
assert rhs == field.copy(data=[0.5, 1.5])
rhs = eq.make_pde_rhs(field, backend="numba")
np.testing.assert_allclose(rhs(field.data, 0.0), np.array([0.5, 1.5]))
eq = PDE({"a": "x + y"})
with pytest.raises(RuntimeError):
eq.evolution_rate(field)
def test_user_bcs_numba(dim, target):
"""test setting user BCs"""
if dim == 1:
value = 1
elif dim == 2:
value = np.arange(3)
elif dim == 3:
value = np.arange(9).reshape(3, 3)
grid = UnitGrid([3] * dim)
bcs = grid.get_boundary_conditions({"type": "user"})
# use normal method to set BCs
f1 = ScalarField(grid)
f1.set_ghost_cells(bc={target: value})
# use user method to set BCs
f2 = ScalarField(grid)
setter = bcs.make_ghost_cell_setter()
# test whether normal call is a no-op
f2._data_full = 3
setter(f2._data_full)
np.testing.assert_allclose(f2._data_full, 3)
setter(f2._data_full, args={"t": 1})
np.testing.assert_allclose(f2._data_full, 3)
f2._data_full = 0
# test whether calling setter with user data works properly
setter(f2._data_full, args={target: value})
np.testing.assert_allclose(f1._data_full, f2._data_full)
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_collection_plotting():
""" test simple plotting of various fields on various grids """
grid = UnitGrid([5])
s1 = ScalarField(grid, label="s1")
s2 = ScalarField(grid)
fc = FieldCollection([s1, s2])
# test setting different figure sizes
fc.plot(figsize="default")
fc.plot(figsize="auto")
fc.plot(figsize=(3, 3))
# test different arrangements
fc.plot(arrangement="horizontal")
fc.plot(arrangement="vertical")
def test_callback_tracker():
""" test trackers that support a callback """
data = []
def store_mean_data(state):
data.append(state.average)
def get_mean_data(state):
return state.average
grid = UnitGrid([4, 4])
state = ScalarField.random_uniform(grid, 0.2, 0.3)
pde = DiffusionPDE()
data_tracker = trackers.DataTracker(get_mean_data, interval=0.1)
callback_tracker = trackers.CallbackTracker(store_mean_data, interval=0.1)
tracker_list = [data_tracker, callback_tracker]
pde.solve(state,
t_range=0.5,
dt=0.005,
tracker=tracker_list,
backend="numpy")
np.testing.assert_array_equal(data, data_tracker.data)
data = []
def store_time(state, t):
data.append(t)
def get_time(state, t):
return t
grid = UnitGrid([4, 4])
state = ScalarField.random_uniform(grid, 0.2, 0.3)
pde = DiffusionPDE()
data_tracker = trackers.DataTracker(get_time, interval=0.1)
tracker_list = [
trackers.CallbackTracker(store_time, interval=0.1), data_tracker
]
pde.solve(state,
t_range=0.5,
dt=0.005,
tracker=tracker_list,
backend="numpy")
ts = np.arange(0, 0.55, 0.1)
np.testing.assert_allclose(data, ts, atol=1e-2)
np.testing.assert_allclose(data_tracker.data, ts, atol=1e-2)
def test_findiff_cyl():
"""test operator for a simple cylindrical grid. Note that we only
really test the polar symmetry"""
grid = CylindricalSymGrid(1.5, [0, 1], (3, 2), periodic_z=True)
_, r1, r2 = grid.axes_coords[0]
np.testing.assert_array_equal(grid.discretization, np.full(2, 0.5))
s = ScalarField(grid, [[1, 1], [2, 2], [4, 4]])
# test laplace
lap = s.laplace(bc=[{"type": "value", "value": 3}, "periodic"])
y1 = 4 + 3 / r1
y2 = -16
np.testing.assert_allclose(lap.data, [[8, 8], [y1, y1], [y2, y2]])
lap = s.laplace(bc=[{"type": "derivative", "value": 3}, "periodic"])
y2 = -2 + 3.5 / r2
np.testing.assert_allclose(lap.data, [[8, 8], [y1, y1], [y2, y2]])
def test_trackers():
""" test whether simple trackers can be used """
times = []
def store_time(state, t):
times.append(t)
def get_data(state):
return {"integral": state.integral}
devnull = open(os.devnull, "w")
data = trackers.DataTracker(get_data, interval=0.1)
tracker_list = [
trackers.PrintTracker(interval=0.1, stream=devnull),
trackers.CallbackTracker(store_time, interval=0.1),
None, # should be ignored
data,
]
if module_available("matplotlib"):
tracker_list.append(trackers.PlotTracker(interval=0.1, show=False))
grid = UnitGrid([16, 16])
state = ScalarField.random_uniform(grid, 0.2, 0.3)
pde = DiffusionPDE()
pde.solve(state, t_range=1, dt=0.005, tracker=tracker_list)
devnull.close()
assert times == data.times
if module_available("pandas"):
df = data.dataframe
np.testing.assert_allclose(df["time"], times)
np.testing.assert_allclose(df["integral"], state.integral)
