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_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_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_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_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)
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_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_tensor_div_div_analytical():
"""test double divergence of a tensor field against analytical expression"""
grid = SphericalSymGrid([0.5, 1], 12)
tf = Tensor2Field.from_expression(
grid, [["r**4", 0, 0], [0, "r**3", 0], [0, 0, "r**3"]])
res = tf._apply_operator("tensor_double_divergence", bc="curvature")
expect = ScalarField.from_expression(grid, "2 * r * (15 * r - 4)")
np.testing.assert_allclose(res.data[1:-1], expect.data[1:-1], rtol=0.01)
def test_gradient_squared_polar(r_inner):
"""compare gradient squared operator"""
grid = PolarSymGrid((r_inner, 4 * np.pi), 32)
field = ScalarField.from_expression(grid, "cos(r)")
s1 = field.gradient("auto_periodic_neumann").to_scalar("squared_sum")
s2 = field.gradient_squared("auto_periodic_neumann", central=True)
np.testing.assert_allclose(s1.data, s2.data, rtol=0.1, atol=0.1)
s3 = field.gradient_squared("auto_periodic_neumann", central=False)
np.testing.assert_allclose(s1.data, s3.data, rtol=0.1, atol=0.1)
assert not np.array_equal(s2.data, s3.data)
def test_conservative_laplace_sph():
"""test and compare the two implementation of the laplace operator"""
r_max = 3.14
for r_min in [0, 0.1]:
grid = SphericalSymGrid((r_min, r_max), 8)
f = ScalarField.from_expression(grid, "cos(r)")
res1 = f.laplace("auto_periodic_neumann", conservative=True)
res2 = f.laplace("auto_periodic_neumann", conservative=False)
np.testing.assert_allclose(res1.data, res2.data, rtol=0.5, atol=0.5)
np.testing.assert_allclose(res1.integral, 0, atol=1e-12)
def test_evaluate_func_invalid():
"""test the evaluate function with invalid data"""
grid = UnitGrid([3])
field = ScalarField.from_expression(grid, "x")
with pytest.raises(ValueError):
evaluate("x", {})
with pytest.raises(ValueError):
evaluate("x", {"x": field})
field2 = ScalarField.from_expression(UnitGrid([4]), "x")
with pytest.raises(ValueError):
evaluate("a + b", {"a": field, "b": field2})
with pytest.raises(BCDataError):
evaluate("laplace(a)", {"a": field}, bc=RuntimeError)
with pytest.raises(RuntimeError):
evaluate("a + b", {"a": field})
def test_divergence_field_cyl():
"""test the divergence operator"""
grid = CylindricalSymGrid(2 * np.pi, [0, 2 * np.pi], [16, 32], periodic_z=True)
v = VectorField.from_expression(grid, ["cos(r) + sin(z)**2", "z * cos(r)**2", 0])
s = v.divergence(bc="natural")
assert s.data.shape == grid.shape
res = ScalarField.from_expression(
grid, "cos(r)**2 - sin(r) + (cos(r) + sin(z)**2) / r"
)
np.testing.assert_allclose(
s.data[1:-1, 1:-1], res.data[1:-1, 1:-1], rtol=0.1, atol=0.1
)
def test_grid_laplace():
"""test the cylindrical implementation of the laplace operator"""
grid_cyl = CylindricalSymGrid(6, (0, 4), (4, 4))
grid_cart = CartesianGrid([[-5, 5], [-5, 5], [0, 4]], [10, 10, 4])
a_2d = ScalarField.from_expression(grid_cyl, expression="exp(-5 * r) * cos(z / 3)")
a_3d = a_2d.interpolate_to_grid(grid_cart)
b_2d = a_2d.laplace("natural")
b_3d = a_3d.laplace("natural")
b_2d_3 = b_2d.interpolate_to_grid(grid_cart)
np.testing.assert_allclose(b_2d_3.data, b_3d.data, rtol=0.2, atol=0.2)
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, True])
vf = VectorField.from_expression(grid, ["sin(x)", "0"])
bc_x = [{"derivative": [-1, 0]}, {"derivative": [1, 0]}]
tf = vf.gradient(bc=[bc_x, "periodic"])
res = ScalarField.from_expression(grid, "cos(x)")
np.testing.assert_allclose(tf[0, 0].data, res.data, atol=0.01, rtol=0.01)
np.testing.assert_allclose(tf[0, 1].data, 0)
np.testing.assert_allclose(tf[1, 0].data, 0)
np.testing.assert_allclose(tf[1, 1].data, 0)
def test_conservative_laplace_sph():
""" test and compare the two implementation of the laplace operator """
r_max = 3.14
for r_min in [0, 0.1]:
grid = SphericalGrid((r_min, r_max), 8)
f = ScalarField.from_expression(grid, "cos(r)")
bcs = grid.get_boundary_conditions("natural")
lap1 = ops.make_laplace(bcs, conservative=True)
lap2 = ops.make_laplace(bcs, conservative=False)
np.testing.assert_allclose(lap1(f.data), lap2(f.data), rtol=0.5, atol=0.5)
np.testing.assert_allclose(f.apply(lap1).integral, 0, atol=1e-12)
def test_length_scale_tracker(tmp_path):
"""test the length scale tracker"""
path = tmp_path / "test_length_scale_tracker.json"
grid = CartesianGrid([[0, 10 * np.pi]], 64, periodic=True)
field = ScalarField.from_expression(grid, "sin(2 * x)")
pde = DiffusionPDE()
tracker = LengthScaleTracker(0.05, filename=path)
pde.solve(field, t_range=0.1, backend="numpy", tracker=tracker)
for ls in tracker.length_scales:
assert ls == pytest.approx(np.pi, rel=1e-3)
assert path.stat().st_size > 0 # wrote some result
def test_evaluate_func_vector():
"""test the evaluate function with vector fields"""
grid = UnitGrid([3])
field = ScalarField.from_expression(grid, "x")
vec = VectorField.from_expression(grid, ["x"])
res = evaluate("inner(v, v)", {"v": vec})
assert isinstance(res, ScalarField)
np.testing.assert_almost_equal(res.data, grid.axes_coords[0] ** 2)
res = evaluate("outer(v, v)", {"v": vec})
assert isinstance(res, Tensor2Field)
np.testing.assert_almost_equal(res.data, [[grid.axes_coords[0] ** 2]])
assert isinstance(evaluate("gradient(a)", {"a": field}), VectorField)
def test_examples_vector_sph():
"""compare derivatives of vector fields for spherical grids"""
grid = SphericalSymGrid(1, 32)
# divergence
vf = VectorField.from_expression(grid, ["r**3", 0, "r**2"])
res = vf.divergence([{"derivative": 0}, {"value": 1}])
expect = ScalarField.from_expression(grid, "5 * r**2")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# vector gradient
vf = VectorField.from_expression(grid, ["r**3", 0, 0])
res = vf.gradient([{"derivative": 0}, {"value": [1, 1, 1]}])
expr = [["3 * r**2", 0, 0], [0, "r**2", 0], [0, 0, "r**2"]]
expect = Tensor2Field.from_expression(grid, expr)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
def test_grid_laplace():
"""test the polar implementation of the laplace operator"""
grid_sph = SphericalSymGrid(9, 11)
grid_cart = CartesianGrid([[-5, 5], [-5, 5], [-5, 5]], [12, 10, 11])
a_1d = ScalarField.from_expression(grid_sph, "cos(r)")
a_3d = a_1d.interpolate_to_grid(grid_cart)
b_3d = a_3d.laplace("auto_periodic_neumann")
b_1d = a_1d.laplace("auto_periodic_neumann")
b_1d_3 = b_1d.interpolate_to_grid(grid_cart)
i = slice(1, -1) # do not compare boundary points
np.testing.assert_allclose(
b_1d_3.data[i, i, i], b_3d.data[i, i, i], rtol=0.2, atol=0.2
)
def test_grid_laplace():
""" test the polar implementation of the laplace operator """
grid_sph = PolarGrid(7, 8)
grid_cart = CartesianGrid([[-5, 5], [-5, 5]], [12, 11])
a_1d = ScalarField.from_expression(grid_sph, "cos(r)")
a_2d = a_1d.interpolate_to_grid(grid_cart)
b_2d = a_2d.laplace("natural")
b_1d = a_1d.laplace("natural")
b_1d_2 = b_1d.interpolate_to_grid(grid_cart)
i = slice(1, -1) # do not compare boundary points
np.testing.assert_allclose(b_1d_2.data[i, i],
b_2d.data[i, i],
rtol=0.2,
atol=0.2)
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_examples_vector_polar():
"""compare derivatives of vector fields for polar grids"""
grid = PolarSymGrid(1, 32)
vf = VectorField.from_expression(grid, ["r**3", "r**2"])
# divergence
res = vf.divergence([{"derivative": 0}, {"value": 1}])
expect = ScalarField.from_expression(grid, "4 * r**2")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# # vector Laplacian
# res = vf.laplace([{"derivative": 0}, {"value": 1}])
# expect = VectorField.from_expression(grid, ["8 * r", "3"])
# np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# vector gradient
res = vf.gradient([{"derivative": 0}, {"value": [1, 1]}])
expr = [["3 * r**2", "-r"], ["2 * r", "r**2"]]
expect = Tensor2Field.from_expression(grid, expr)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
def test_examples_vector_cyl():
"""compare derivatives of vector fields for cylindrical grids"""
grid = CylindricalSymGrid(1, [0, 2 * np.pi], 32)
e_r = "r**3 * sin(z)"
e_φ = "r**2 * sin(z)"
e_z = "r**4 * cos(z)"
vf = VectorField.from_expression(grid, [e_r, e_z, e_φ])
bc_r = ({"derivative_normal": 0}, {"value_normal": "r**3 * sin(z)"})
bc_z = {"curvature_normal": "-r**4 * cos(z)"}
bcs = [bc_r, bc_z]
# divergence
res = vf.divergence(bcs)
expect = ScalarField.from_expression(grid,
"4 * r**2 * sin(z) - r**4 * sin(z)")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# vector Laplacian
grid = CylindricalSymGrid(1, [0, 2 * np.pi], 32, periodic_z=True)
vf = VectorField.from_expression(grid, ["r**3 * sin(z)"] * 3)
val_r_outer = np.broadcast_to(6 * np.sin(grid.axes_coords[1]), (3, 32))
bcs = [({"derivative": 0}, {"curvature": val_r_outer}), "periodic"]
res = vf.laplace(bcs)
expr = [
"8 * r * sin(z) - r**3 * sin(z)",
"9 * r * sin(z) - r**3 * sin(z)",
"8 * r * sin(z) - r**3 * sin(z)",
]
expect = VectorField.from_expression(grid, expr)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
# vector gradient
bcs = [({"derivative": 0}, {"curvature": val_r_outer}), "periodic"]
res = vf.gradient(bcs)
expr = [
["3 * r**2 * sin(z)", "r**3 * cos(z)", "-r**2 * sin(z)"],
["3 * r**2 * sin(z)", "r**3 * cos(z)", 0],
["3 * r**2 * sin(z)", "r**3 * cos(z)", "r**2 * sin(z)"],
]
expect = Tensor2Field.from_expression(grid, expr)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
\partial_t \phi = 6 \phi \partial_x \phi - \partial_x^2 \phi
which we implement using a custom PDE class below.
"""
from math import pi
from pde import CartesianGrid, MemoryStorage, PDEBase, ScalarField, plot_kymograph
class KortewegDeVriesPDE(PDEBase):
""" Korteweg-de Vries equation """
def evolution_rate(self, state, t=0):
""" implement the python version of the evolution equation """
assert state.grid.dim == 1 # ensure the state is one-dimensional
grad = state.gradient("natural")[0]
return 6 * state * grad - grad.laplace("natural")
# initialize the equation and the space
grid = CartesianGrid([[0, 2 * pi]], [32], periodic=True)
state = ScalarField.from_expression(grid, "sin(x)")
# solve the equation and store the trajectory
storage = MemoryStorage()
eq = KortewegDeVriesPDE()
eq.solve(state, t_range=3, tracker=storage.tracker(0.1))
# plot the trajectory as a space-time plot
plot_kymograph(storage)
"""
Visualizing a 3d field interactively
====================================
This example demonstrates how to display 3d data interactively using the
`napari <https://napari.org>`_ viewer.
"""
import numpy as np
from pde import CartesianGrid, ScalarField
# create a scalar field with some noise
grid = CartesianGrid([[0, 2 * np.pi]] * 3, 64)
data = ScalarField.from_expression(
grid, "(cos(2 * x) * sin(3 * y) + cos(2 * z))**2")
data += ScalarField.random_normal(grid, std=0.01)
data.label = "3D Field"
data.plot_interactive()
"""
Plotting a scalar field in cylindrical coordinates
==================================================
This example shows how to initialize and visualize the scalar field
:math:`u = \sqrt{z} \, \exp(-r^2)` in cylindrical coordinates.
"""
from pde import CylindricalGrid, ScalarField
grid = CylindricalGrid(radius=3, bounds_z=[0, 4], shape=16)
field = ScalarField.from_expression(grid, "sqrt(z) * exp(-r**2)")
field.plot(title="Scalar field in cylindrical coordinates")
This example implements a complex PDE using the :class:`~pde.pdes.pde.PDE`. We here
chose the `Schrödinger equation <https://en.wikipedia.org/wiki/Schrödinger_equation>`_
without a spatial potential in non-dimensional form:
.. math::
i \partial_t \psi = -\nabla^2 \psi
Note that the example imposes Neumann conditions at the wall, so the wave packet is
expected to reflect off the wall.
"""
from math import sqrt
from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph
grid = CartesianGrid([[0, 20]], 128, periodic=False) # generate grid
# create a (normalized) wave packet with a certain form as an initial condition
initial_state = ScalarField.from_expression(
grid, "exp(I * 5 * x) * exp(-(x - 10)**2)")
initial_state /= sqrt(initial_state.to_scalar("norm_squared").integral.real)
eq = PDE({"ψ": f"I * laplace(ψ)"}) # define the pde
# solve the pde and store intermediate data
storage = MemoryStorage()
eq.solve(initial_state, t_range=2.5, dt=1e-5, tracker=[storage.tracker(0.02)])
# visualize the results as a space-time plot
plot_kymograph(storage, scalar="norm_squared")
"""
Visualizing a scalar field
==========================
This example displays methods for visualizing scalar fields.
"""
import numpy as np
import matplotlib.pyplot as plt
from pde import CylindricalGrid, ScalarField
# create a scalar field with some noise
grid = CylindricalGrid(7, [0, 4 * np.pi], 64)
data = ScalarField.from_expression(grid, 'sin(z) * exp(-r / 3)')
data += 0.05 * ScalarField.random_normal(grid)
# manipulate the field
smoothed = data.smooth() # Gaussian smoothing to get rid of the noise
projected = data.project('r') # integrate along the radial direction
sliced = smoothed.slice({'z': 1}) # slice the smoothed data
# create four plots of the field and the modifications
fig, axes = plt.subplots(nrows=2, ncols=2)
data.plot(ax=axes[0, 0], title='Original field')
smoothed.plot(ax=axes[1, 0], title='Smoothed field')
projected.plot(ax=axes[0, 1], title='Projection on axial coordinate')
sliced.plot(ax=axes[1, 1], title='Slice of smoothed field at $z=1$')
plt.subplots_adjust(hspace=0.8)