def test_setting_boundary_conditions():
"""test setting some boundary conditions"""
grid = CylindricalSymGrid(1, [0, 1], 3)
b_inner = NeumannBC(grid, 0, upper=False)
assert grid.get_boundary_conditions("natural")[0].low == b_inner
assert grid.get_boundary_conditions({"value": 2})[0].low != b_inner
def test_polar_conversion(periodic):
"""test conversion to polar coordinates"""
grid = CylindricalSymGrid(1, [-1, 1], [5, 5], periodic_z=periodic)
dists = grid.polar_coordinates_real([0, 0, 0])
assert np.all(0.09 <= dists)
assert np.any(dists < 0.11)
assert np.all(dists <= np.sqrt(2))
assert np.any(dists > 0.8 * np.sqrt(2))
def test_cylindrical_grid():
"""test simple cylindrical grid"""
for periodic in [True, False]:
grid = CylindricalSymGrid(4, (-1, 2), (8, 9), periodic_z=periodic)
assert grid.dim == 3
assert grid.numba_type == "f8[:, :]"
assert grid.shape == (8, 9)
assert grid.length == pytest.approx(3)
assert grid.discretization[0] == pytest.approx(0.5)
assert grid.discretization[1] == pytest.approx(1 / 3)
np.testing.assert_array_equal(grid.discretization, np.array([0.5, 1 / 3]))
assert not grid.uniform_cell_volumes
assert grid.volume == pytest.approx(np.pi * 4**2 * 3)
assert grid.volume == pytest.approx(grid.integrate(1))
rs, zs = grid.axes_coords
np.testing.assert_allclose(rs, np.linspace(0.25, 3.75, 8))
np.testing.assert_allclose(zs, np.linspace(-1 + 1 / 6, 2 - 1 / 6, 9))
assert grid.contains_point(grid.get_random_point(coords="cartesian"))
ps = [grid.get_random_point(coords="cartesian") for _ in range(2)]
assert all(grid.contains_point(ps))
ps = grid.get_random_point(coords="cartesian", boundary_distance=1.49)
assert grid.contains_point(ps)
assert "laplace" in grid.operators
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_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 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_laplacian_field_cyl():
"""test the gradient operator"""
grid = CylindricalSymGrid(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_vector_gradient_divergence_field_cyl():
"""test the divergence operator"""
grid = CylindricalSymGrid(2 * np.pi, [0, 2 * np.pi], [8, 16], periodic_z=True)
r, z = grid.cell_coords[..., 0], grid.cell_coords[..., 1]
data = [np.cos(r) + np.sin(z) ** 2, np.cos(r) ** 2 + np.sin(z), np.zeros_like(r)]
v = VectorField(grid, data=data)
t = v.gradient(bc="natural")
assert t.data.shape == (3, 3, 8, 16)
v = t.divergence(bc="natural")
assert v.data.shape == (3, 8, 16)
def test_gradient_field_cyl():
"""test the gradient operator"""
grid = CylindricalSymGrid(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_gradient_squared_cyl():
"""compare gradient squared operator"""
grid = CylindricalSymGrid(2 * np.pi, [0, 2 * np.pi], 64)
field = ScalarField.random_harmonic(grid, modes=1)
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.2, atol=0.2)
s3 = field.gradient_squared("auto_periodic_neumann", central=False)
np.testing.assert_allclose(s1.data, s3.data, rtol=0.2, atol=0.2)
assert not np.array_equal(s2.data, s3.data)
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)
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_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_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_examples_tensor_cyl():
"""compare derivatives of tensorial fields for cylindrical grids"""
grid = CylindricalSymGrid(1, [0, 2 * np.pi], 32, periodic_z=True)
tf = Tensor2Field.from_expression(grid, [["r**3 * sin(z)"] * 3] * 3)
# tensor divergence
rs, zs = grid.axes_coords
val_r_outer = np.broadcast_to(6 * rs * np.sin(zs), (3, 32))
bcs = [({"derivative": 0}, {"curvature": val_r_outer}), "periodic"]
res = tf.divergence(bcs)
expect = VectorField.from_expression(
grid,
[
"r**2 * (r * cos(z) + 3 * sin(z))",
"r**2 * (r * cos(z) + 4 * sin(z))",
"r**2 * (r * cos(z) + 5 * sin(z))",
],
)
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
"""
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 CylindricalSymGrid, ScalarField
grid = CylindricalSymGrid(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")
def test_cylindrical_grid():
"""test simple cylindrical grid"""
for periodic in [True, False]:
grid = CylindricalSymGrid(4, (-1, 2), (8, 9), periodic_z=periodic)
msg = str(grid)
assert grid.dim == 3
assert grid.numba_type == "f8[:, :]"
assert grid.shape == (8, 9)
assert grid.length == pytest.approx(3)
assert grid.discretization[0] == pytest.approx(0.5)
assert grid.discretization[1] == pytest.approx(1 / 3)
np.testing.assert_array_equal(grid.discretization,
np.array([0.5, 1 / 3]))
assert not grid.uniform_cell_volumes
assert grid.volume == pytest.approx(np.pi * 4**2 * 3)
assert grid.volume == pytest.approx(grid.integrate(1))
rs, zs = grid.axes_coords
np.testing.assert_allclose(rs, np.linspace(0.25, 3.75, 8))
np.testing.assert_allclose(zs, np.linspace(-1 + 1 / 6, 2 - 1 / 6, 9))
# random points
c = np.random.randint(8, size=(6, 2))
c1 = grid.point_to_cell(grid.cell_to_point(c))
np.testing.assert_almost_equal(c, c1, err_msg=msg)
assert grid.contains_point(grid.get_random_point())
ps = [grid.get_random_point() for _ in range(2)]
assert all(grid.contains_point(ps))
assert grid.contains_point(grid.get_random_point(1.49))
assert "laplace" in grid.operators
def test_setting_domain_cylindrical():
"""test various versions of settings bcs for cylindrical grids"""
grid = CylindricalSymGrid(1, [0, 1], [2, 2], periodic_z=False)
grid.get_boundary_conditions("auto_periodic_neumann")
grid.get_boundary_conditions(["derivative", "derivative"])
with pytest.raises(ValueError):
grid.get_boundary_conditions(["derivative"])
with pytest.raises(ValueError):
grid.get_boundary_conditions(["derivative"] * 3)
with pytest.raises(RuntimeError):
grid.get_boundary_conditions(["derivative", "periodic"])
grid = CylindricalSymGrid(1, [0, 1], [2, 2], periodic_z=True)
grid.get_boundary_conditions("auto_periodic_neumann")
grid.get_boundary_conditions(["derivative", "periodic"])
with pytest.raises(RuntimeError):
grid.get_boundary_conditions(["derivative", "derivative"])
def main():
"""main routine testing the performance"""
print("Reports calls-per-second (larger is better)")
print(" The `CUSTOM` method implemented by hand is the baseline case.")
print(" The `FLEXIBLE` method is a serial implementation using the "
"boundary conditions from the package.")
print(" The other methods use the functions supplied by the package.\n")
# Cartesian grid with different shapes and boundary conditions
for shape in [(32, 32), (512, 512)]:
for periodic in [True, False]:
grid = UnitGrid(shape, periodic=periodic)
print(grid)
field = ScalarField.random_normal(grid)
bcs = grid.get_boundary_conditions("natural", rank=0)
expected = field.laplace("natural")
for method in [
"CUSTOM", "FLEXIBLE", "OPTIMIZED", "numba", "scipy"
]:
if method == "CUSTOM":
laplace = custom_laplace_2d(shape, periodic=periodic)
elif method == "FLEXIBLE":
laplace = flexible_laplace_2d(bcs)
elif method == "OPTIMIZED":
laplace = optimized_laplace_2d(bcs)
elif method in {"numba", "scipy"}:
laplace = grid.make_operator("laplace",
bc=bcs,
backend=method)
else:
raise ValueError(f"Unknown method `{method}`")
# call once to pre-compile and test result
if method == "OPTIMIZED":
result = laplace(field._data_all)
np.testing.assert_allclose(result, expected.data)
speed = estimate_computation_speed(laplace,
field._data_all)
else:
np.testing.assert_allclose(laplace(field.data),
expected.data)
speed = estimate_computation_speed(laplace, field.data)
print(f"{method:>9s}: {int(speed):>9d}")
print()
# Cylindrical grid with different shapes
for shape in [(32, 64), (512, 512)]:
grid = CylindricalSymGrid(shape[0], [0, shape[1]], shape)
print(f"Cylindrical grid, shape={shape}")
field = ScalarField.random_normal(grid)
bcs = Boundaries.from_data(grid, "derivative")
expected = field.laplace(bcs)
for method in ["CUSTOM", "numba"]:
if method == "CUSTOM":
laplace = custom_laplace_cyl_neumann(shape)
elif method == "numba":
laplace = grid.make_operator("laplace", bc=bcs)
else:
raise ValueError(f"Unknown method `{method}`")
# call once to pre-compile and test result
np.testing.assert_allclose(laplace(field.data), expected.data)
speed = estimate_computation_speed(laplace, field.data)
print(f"{method:>8s}: {int(speed):>9d}")
print()
# Spherical grid with different shapes
for shape in [32, 512]:
grid = SphericalSymGrid(shape, shape)
print(grid)
field = ScalarField.random_normal(grid)
bcs = Boundaries.from_data(grid, "derivative")
for conservative in [True, False]:
laplace = grid.make_operator("laplace",
bcs,
conservative=conservative)
laplace(field.data) # call once to pre-compile
speed = estimate_computation_speed(laplace, field.data)
print(
f" numba (conservative={str(conservative):<5}): {int(speed):>9d}"
)
print()
"""
Visualizing a scalar field
==========================
This example displays methods for visualizing scalar fields.
"""
import matplotlib.pyplot as plt
import numpy as np
from pde import CylindricalSymGrid, ScalarField
# create a scalar field with some noise
grid = CylindricalSymGrid(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)
plt.show()