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 test_findiff_cyl():
"""test operator for a simple cylindrical grid. Note that we only
really test the polar symmetry"""
grid = CylindricalGrid(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]])
v = VectorField(grid,
[[[1, 1], [2, 2], [4, 4]], [[0, 0]] * 3, [[0, 0]] * 3])
# test gradient
grad = s.gradient(bc=["value", "periodic"])
np.testing.assert_allclose(grad.data[0], [[1, 1], [3, 3], [-6, -6]])
grad = s.gradient(bc=["derivative", "periodic"])
np.testing.assert_allclose(grad.data[0], [[1, 1], [3, 3], [2, 2]])
# test divergence
div = v.divergence(bc=["value", "periodic"])
y1 = 3 + 2 / r1
y2 = -6 + 4 / r2
np.testing.assert_allclose(div.data, [[5, 5], [y1, y1], [y2, y2]])
div = v.divergence(bc=["derivative", "periodic"])
y2 = 2 + 4 / r2
np.testing.assert_allclose(div.data, [[5, 5], [y1, y1], [y2, y2]])
# 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_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_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_examples_tensor_polar():
"""compare derivatives of tensorial fields for polar grids"""
grid = PolarSymGrid(1, 32)
tf = Tensor2Field.from_expression(grid, [["r**3"] * 2] * 2)
# tensor divergence
res = tf.divergence([{"derivative": 0}, {"normal_value": [1, 1]}])
expect = VectorField.from_expression(grid, ["3 * r**2", "5 * r**2"])
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
res = tf.divergence([{"derivative": 0}, {"value": np.ones((2, 2))}])
expect = VectorField.from_expression(grid, ["3 * r**2", "5 * r**2"])
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
def test_vector_plotting_2d(transpose):
"""test plotting of 2d vector fields"""
grid = UnitGrid([3, 4])
field = VectorField.random_uniform(grid, 0.1, 0.9)
for method in ["quiver", "streamplot"]:
ref = field.plot(method=method, transpose=transpose)
field._update_plot(ref)
# test sub-sampling
grid = UnitGrid([32, 15])
field = VectorField.random_uniform(grid, 0.1, 0.9)
field.get_vector_data(transpose=transpose, max_points=7)
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_divergence_field_cyl():
""" test the divergence 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]
data = [
np.cos(r) + np.sin(z)**2,
np.cos(r)**2 + np.sin(z),
np.zeros_like(r)
]
v = VectorField(grid, data=data)
s = v.divergence(bc="natural")
assert s.data.shape == (8, 16)
res = np.cos(z) - np.sin(r) + (np.cos(r) + np.sin(z)**2) / r
np.testing.assert_allclose(s.data, res, rtol=0.1, atol=0.1)
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_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_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_interactive_collection_plotting():
""" test the interactive plotting """
grid = UnitGrid([3, 3])
sf = ScalarField.random_uniform(grid, 0.1, 0.9)
vf = VectorField.random_uniform(grid, 0.1, 0.9)
field = FieldCollection([sf, vf])
field.plot_interactive(viewer_args={"show": False, "close": True})
def test_divergence():
"""test the divergence 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)
s1 = v.divergence("auto_periodic_neumann")
assert s1.data.shape == (16, 16)
div = np.cos(y) - np.sin(x)
np.testing.assert_allclose(s1.data, div, rtol=0.1, atol=0.1)
v.divergence("auto_periodic_neumann", out=s1)
assert s1.data.shape == (16, 16)
np.testing.assert_allclose(s1.data, div, rtol=0.1, atol=0.1)
def test_vector_laplace_cart(ndim):
"""test different vector laplace operators"""
bcs = _get_random_grid_bcs(ndim, dx="uniform", periodic="random", rank=1)
field = VectorField.random_uniform(bcs.grid)
res1 = field.laplace(bcs, backend="scipy").data
res2 = field.laplace(bcs, backend="numba").data
assert res1.shape == (ndim,) + bcs.grid.shape
np.testing.assert_allclose(res1, res2)
def test_plotting_2d():
""" test plotting of 2d vector fields """
grid = UnitGrid([3, 3])
field = VectorField.random_uniform(grid, 0.1, 0.9)
for method in ["quiver", "streamplot"]:
ref = field.plot(method=method)
field._update_plot(ref)
def test_vector_laplace():
"""test the laplace 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) + np.sin(y), np.sin(y) - np.cos(x)]
v = VectorField(grid, data)
vl = v.laplace("auto_periodic_neumann")
assert vl.data.shape == (2, 16, 16)
np.testing.assert_allclose(vl.data[0, ...],
-np.cos(x) - np.sin(y),
rtol=0.1,
atol=0.1)
np.testing.assert_allclose(vl.data[1, ...],
-np.sin(y) + np.cos(x),
rtol=0.1,
atol=0.1)
def test_divergence_cart(ndim):
"""test different divergence operators"""
for periodic in [True, False]:
bcs = _get_random_grid_bcs(ndim, dx="uniform", periodic=periodic, rank=1)
field = VectorField.random_uniform(bcs.grid)
res1 = field.divergence(bcs, backend="scipy").data
res2 = field.divergence(bcs, backend="numba").data
np.testing.assert_allclose(res1, res2)
def test_findiff_sph():
"""test operator for a simple spherical grid"""
grid = SphericalSymGrid(1.5, 3)
_, r1, r2 = grid.axes_coords[0]
assert grid.discretization == (0.5, )
s = ScalarField(grid, [1, 2, 4])
v = VectorField(grid, [[1, 2, 4], [0] * 3, [0] * 3])
# test gradient
grad = s.gradient(bc="value")
np.testing.assert_allclose(grad.data[0, :], [1, 3, -6])
grad = s.gradient(bc="derivative")
np.testing.assert_allclose(grad.data[0, :], [1, 3, 2])
# test divergence
div = v.divergence(bc="value", conservative=False)
np.testing.assert_allclose(div.data, [9, 3 + 4 / r1, -6 + 8 / r2])
div = v.divergence(bc="derivative", conservative=False)
np.testing.assert_allclose(div.data, [9, 3 + 4 / r1, 2 + 8 / r2])
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_pde_product_operators():
"""test inner and outer products"""
eq = PDE(
{"p": "gradient(dot(p, p) + inner(p, p)) + tensor_divergence(outer(p, p))"}
)
assert not eq.explicit_time_dependence
assert not eq.complex_valued
field = VectorField(grids.UnitGrid([4]), 1)
res = eq.solve(field, t_range=1, dt=0.1, backend="numpy", tracker=None)
np.testing.assert_allclose(res.data, field.data)
def test_findiff():
""" test operator for a simple polar grid """
grid = PolarGrid(1.5, 3)
_, _, r2 = grid.axes_coords[0]
assert grid.discretization == (0.5, )
s = ScalarField(grid, [1, 2, 4])
v = VectorField(grid, [[1, 2, 4], [0] * 3])
# test gradient
grad = s.gradient(bc="value")
np.testing.assert_allclose(grad.data[0, :], [1, 3, -6])
grad = s.gradient(bc="derivative")
np.testing.assert_allclose(grad.data[0, :], [1, 3, 2])
# test divergence
div = v.divergence(bc="value")
np.testing.assert_allclose(div.data, [5, 17 / 3, -6 + 4 / r2])
div = v.divergence(bc="derivative")
np.testing.assert_allclose(div.data, [5, 17 / 3, 2 + 4 / r2])
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_tensor_sph_symmetry():
"""test treatment of symmetric tensor field"""
grid = SphericalSymGrid(1, 16)
vf = VectorField.from_expression(grid, ["r**2", 0, 0])
vf_grad = vf.gradient({"derivative": 2})
strain = vf_grad + vf_grad.transpose()
bcs = [{"value": 0}, {"normal_derivative": [4, 0, 0]}]
strain_div = strain.divergence(bcs, safe=False)
np.testing.assert_allclose(strain_div.data[0], 8)
np.testing.assert_allclose(strain_div.data[1:], 0)
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_vector_gradient():
""" 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 = -d10.copy()
d10[:, 0] = d10[:, -1] = -7
d01[0, :] = d01[-1, :] = 7
d11 = np.cos(y)
t2 = Tensor2Field(grid, np.array([[d00, d01], [d10, d11]]))
np.testing.assert_allclose(t1.data, t2.data, rtol=0.1, atol=0.1)
v.gradient("natural", out=t1)
assert t1.data.shape == (2, 2, 16, 16)
np.testing.assert_allclose(t1.data, t2.data, rtol=0.1, atol=0.1)
def test_pde_vector():
"""test PDE with a single vector field"""
eq = PDE({"u": "vector_laplace(u) + exp(-t)"})
assert eq.explicit_time_dependence
assert not eq.complex_valued
grid = grids.UnitGrid([8, 8])
field = VectorField.random_normal(grid)
res_a = eq.solve(field, t_range=1, dt=0.01, backend="numpy", tracker=None)
res_b = eq.solve(field, t_range=1, dt=0.01, backend="numba", tracker=None)
res_a.assert_field_compatible(res_b)
np.testing.assert_allclose(res_a.data, res_b.data)
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_vector_bcs():
"""test boundary conditions on vector fields"""
grid = UnitGrid([3, 3], periodic=False)
v = VectorField.from_expression(grid, ["x", "cos(y)"])
bc_x = {"value": [0, 1]}
bc_y = {"value": [2, 3]}
bcs = [bc_x, bc_y]
s1 = v.divergence(bcs, backend="scipy").data
div = grid.make_operator("divergence", bcs)
s2 = div(v.data)
np.testing.assert_allclose(s1, s2)
def test_collections():
"""test field collections"""
grid = UnitGrid([3, 4])
sf = ScalarField.random_uniform(grid, label="sf")
vf = VectorField.random_uniform(grid, label="vf")
tf = Tensor2Field.random_uniform(grid, label="tf")
fields = FieldCollection([sf, vf, tf])
assert fields.data.shape == (7, 3, 4)
assert isinstance(str(fields), str)
fields.data[:] = 0
np.testing.assert_allclose(sf.data, 0)
np.testing.assert_allclose(vf.data, 0)
np.testing.assert_allclose(tf.data, 0)
assert fields[0] is fields["sf"]
assert fields[1] is fields["vf"]
assert fields[2] is fields["tf"]
with pytest.raises(KeyError):
fields["42"]
sf.data = 1
vf.data = 1
tf.data = 1
np.testing.assert_allclose(fields.data, 1)
assert all(np.allclose(i, 12) for i in fields.integrals)
assert all(np.allclose(i, 1) for i in fields.averages)
assert np.allclose(fields.magnitudes, np.sqrt([1, 2, 4]))
assert sf.data.shape == (3, 4)
assert vf.data.shape == (2, 3, 4)
assert tf.data.shape == (2, 2, 3, 4)
c2 = FieldBase.from_state(fields.attributes, data=fields.data)
assert c2 == fields
assert c2.grid is grid
attrs = FieldCollection.unserialize_attributes(
fields.attributes_serialized)
c2 = FieldCollection.from_state(attrs, data=fields.data)
assert c2 == fields
assert c2.grid is not grid
fields["sf"] = 2.0
np.testing.assert_allclose(sf.data, 2)
with pytest.raises(KeyError):
fields["42"] = 0
fields.plot(subplot_args=[{}, {"scale": 1}, {"colorbar": False}])