def test_spherical_grid():
"""test simple spherical grid"""
grid = SphericalSymGrid(4, 8)
assert grid.dim == 3
assert grid.numba_type == "f8[:]"
assert grid.shape == (8, )
assert not grid.has_hole
assert grid.discretization[0] == pytest.approx(0.5)
assert not grid.uniform_cell_volumes
np.testing.assert_array_equal(grid.discretization, np.array([0.5]))
assert grid.volume == pytest.approx(4 / 3 * np.pi * 4**3)
assert grid.volume == pytest.approx(grid.integrate(1))
np.testing.assert_allclose(grid.axes_coords[0], np.linspace(0.25, 3.75, 8))
a = grid.make_operator("laplace",
"auto_periodic_neumann")(np.random.random(8))
assert a.shape == (8, )
assert np.all(np.isfinite(a))
# random points
c = np.random.randint(8, size=(6, 1))
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())
assert grid.contains_point(grid.get_random_point(3.99))
assert "laplace" in grid.operators
def test_spherical_annulus():
"""test simple spherical grid with a hole"""
grid = SphericalSymGrid((2, 4), 8)
assert grid.dim == 3
assert grid.numba_type == "f8[:]"
assert grid.shape == (8,)
assert grid.has_hole
assert grid.discretization[0] == pytest.approx(0.25)
assert not grid.uniform_cell_volumes
np.testing.assert_array_equal(grid.discretization, np.array([0.25]))
assert grid.volume == pytest.approx(4 / 3 * np.pi * (4**3 - 2**3))
assert grid.volume == pytest.approx(grid.integrate(1))
assert grid.radius == (2, 4)
np.testing.assert_allclose(grid.axes_coords[0], np.linspace(2.125, 3.875, 8))
a = grid.make_operator("laplace", "auto_periodic_neumann")(np.random.random(8))
assert a.shape == (8,)
assert np.all(np.isfinite(a))
assert grid.contains_point(grid.get_random_point(coords="cartesian"))
p = grid.get_random_point(boundary_distance=1.99, coords="cartesian")
assert grid.contains_point(p)
# test boundary points
np.testing.assert_equal(grid._boundary_coordinates(0, False), np.array([2]))
np.testing.assert_equal(grid._boundary_coordinates(0, True), np.array([4]))
def test_poisson_solver_spherical():
"""test the poisson solver on Polar grids"""
for grid in [SphericalSymGrid(4, 8), SphericalSymGrid([2, 4], 8)]:
for bc_val in ["auto_periodic_neumann", {"value": 1}]:
bcs = grid.get_boundary_conditions(bc_val)
d = ScalarField.random_uniform(grid)
d -= d.average # balance the right hand side
sol = solve_poisson_equation(d, bcs)
test = sol.laplace(bcs)
np.testing.assert_allclose(
test.data, d.data, err_msg=f"bcs={bc_val}, grid={grid}", rtol=1e-6
)
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_gradient_squared(r_inner):
"""compare gradient squared operator"""
grid = SphericalSymGrid((r_inner, 5), 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.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_examples_tensor_sph():
"""compare derivatives of tensorial fields for spherical grids"""
grid = SphericalSymGrid(1, 32)
tf = Tensor2Field.from_expression(grid, [["r**3"] * 3] * 3)
tfd = tf.data
tfd[0, 1] = tfd[1, 1] = tfd[1, 2] = tfd[2, 1] = tfd[2, 2] = 0
# tensor divergence
res = tf.divergence([{"derivative_normal": 0}, {"value_normal": [1, 1, 1]}])
expect = VectorField.from_expression(grid, ["5 * r**2", "5 * r**2", "6 * r**2"])
np.testing.assert_allclose(res.data, expect.data, 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 = SphericalSymGrid(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_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_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_examples_scalar_sph():
"""compare derivatives of scalar fields for spherical grids"""
grid = SphericalSymGrid(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, 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}, {"value": 1}], 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, "12 * r")
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
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()
""" Spherically symmetric PDE ========================= This example illustrates how to solve a PDE in a spherically symmetric geometry. """ from pde import DiffusionPDE, ScalarField, SphericalSymGrid grid = SphericalSymGrid(radius=[1, 5], shape=128) # generate grid state = ScalarField.random_uniform(grid) # generate initial condition eq = DiffusionPDE(0.1) # define the PDE result = eq.solve(state, t_range=0.1, dt=0.001) result.plot(kind="image")