def test_tensor_div_div(conservative):
"""test double divergence of a tensor field by comparison with two divergences"""
grid = SphericalSymGrid([0, 1], 64)
expr = "r * tanh((0.5 - r) * 10)"
bc = "auto_periodic_neumann"
# test radial part
tf = Tensor2Field.from_expression(grid,
[[expr, 0, 0], [0, 0, 0], [0, 0, 0]])
res = tf._apply_operator("tensor_double_divergence",
bc=bc,
conservative=conservative)
est = tf.divergence(bc).divergence(bc)
np.testing.assert_allclose(res.data[2:-2],
est.data[2:-2],
rtol=0.02,
atol=1)
# test angular part
tf = Tensor2Field.from_expression(grid,
[[0, 0, 0], [0, expr, 0], [0, 0, expr]])
res = tf._apply_operator("tensor_double_divergence",
bc=bc,
conservative=conservative)
est = tf.divergence(bc).divergence(bc)
np.testing.assert_allclose(res.data[2:-2],
est.data[2:-2],
rtol=0.02,
atol=1)
def test_tensors_basic():
"""test some tensor calculations"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [3, 4])
t1 = Tensor2Field(grid, np.full((2, 2) + grid.shape, 1))
t2 = Tensor2Field(grid, np.full((2, 2) + grid.shape, 2))
np.testing.assert_allclose(t2.average, [[2, 2], [2, 2]])
assert t1.magnitude == pytest.approx(2)
assert t1["x", "x"] == t1[0, 0]
assert t1["x", 1] == t1[0, "y"] == t1[0, 1]
t1[0, 0] = t1[0, 0]
t3 = t1 + t2
assert t3.grid == grid
np.testing.assert_allclose(t3.data, 3)
t1 += t2
np.testing.assert_allclose(t1.data, 3)
field = Tensor2Field.random_uniform(grid)
trace = field.trace()
assert isinstance(trace, ScalarField)
np.testing.assert_allclose(trace.data, field.data.trace())
t1 = Tensor2Field(grid)
t1[0, 0] = 1
t1[0, 1] = 2
t1[1, 0] = 3
t1[1, 1] = 4
for method, value in [
("min", 1),
("max", 4),
("norm", np.linalg.norm([[1, 2], [3, 4]])),
("squared_sum", 30),
("norm_squared", 30),
("trace", 5),
("invariant1", 5),
("invariant2", -1),
]:
p1 = t1.to_scalar(method)
assert p1.data.shape == grid.shape
np.testing.assert_allclose(p1.data, value)
for idx in ((1, ), (1, 2, 3), (1.5, 2), ("a", "b"), 1.0):
with pytest.raises(IndexError):
t1[idx]
t2 = FieldBase.from_state(t1.attributes, data=t1.data)
assert t1 == t2
assert t1.grid is t2.grid
attrs = Tensor2Field.unserialize_attributes(t1.attributes_serialized)
t2 = FieldBase.from_state(attrs, data=t1.data)
assert t1 == t2
assert t1.grid is not t2.grid
def test_tensor_symmetrize():
"""test advanced tensor calculations"""
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [2, 2])
t1 = Tensor2Field(grid)
t1.data[0, 0, :] = 1
t1.data[0, 1, :] = 2
t1.data[1, 0, :] = 3
t1.data[1, 1, :] = 4
# traceless = False
t2 = t1.copy()
t1.symmetrize(make_traceless=False, inplace=True)
tr = t1.trace()
assert np.all(tr.data == 5)
t1_trans = np.swapaxes(t1.data, 0, 1)
np.testing.assert_allclose(t1.data, t1_trans.data)
ts = t1.copy()
ts.symmetrize(make_traceless=False, inplace=True)
np.testing.assert_allclose(t1.data, ts.data)
# traceless = True
t2.symmetrize(make_traceless=True, inplace=True)
tr = t2.trace()
assert np.all(tr.data == 0)
t2_trans = np.swapaxes(t2.data, 0, 1)
np.testing.assert_allclose(t2.data, t2_trans.data)
ts = t2.copy()
ts.symmetrize(make_traceless=True, inplace=True)
np.testing.assert_allclose(t2.data, ts.data)
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_insert_tensor(example_grid):
"""test the `insert` method"""
f = Tensor2Field(example_grid)
a = np.random.random(f.data_shape)
c = tuple(example_grid.point_to_cell(example_grid.get_random_point()))
c_data = (Ellipsis, ) + c
p = example_grid.cell_to_point(c, cartesian=False)
f.insert(p, a)
np.testing.assert_almost_equal(f.data[c_data],
a / example_grid.cell_volumes[c])
f.insert(example_grid.get_random_point(cartesian=False), a)
np.testing.assert_almost_equal(f.integral, 2 * a)
f.data = 0 # reset
insert = example_grid.make_inserter_compiled()
c = tuple(example_grid.point_to_cell(example_grid.get_random_point()))
c_data = (Ellipsis, ) + c
p = example_grid.cell_to_point(c, cartesian=False)
insert(f.data, p, a)
np.testing.assert_almost_equal(f.data[c_data],
a / example_grid.cell_volumes[c])
insert(f.data, example_grid.get_random_point(cartesian=False), a)
np.testing.assert_almost_equal(f.integral, 2 * a)
def test_tensor_divergence_cart(ndim):
"""test different tensor divergence operators"""
bcs = _get_random_grid_bcs(ndim, dx="uniform", periodic="random", rank=2)
field = Tensor2Field.random_uniform(bcs.grid)
res1 = field.divergence(bcs, backend="scipy").data
res2 = field.divergence(bcs, backend="numba").data
assert res1.shape == (ndim,) + bcs.grid.shape
np.testing.assert_allclose(res1, res2)
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_tensors_transpose(grid):
"""test transposing tensors"""
def broadcast(arr):
return np.asarray(arr)[(..., ) + (np.newaxis, ) * grid.num_axes]
field = Tensor2Field(grid, broadcast([[0, 1], [2, 3]]))
field_T = field.transpose(label="altered")
assert field_T.label == "altered"
np.testing.assert_allclose(field_T.data, broadcast([[0, 2], [1, 3]]))
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_normal": 0}, {"value_normal": [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)
def test_tensor_invariants():
"""test the invariants"""
# dim == 1
f = Tensor2Field.random_uniform(UnitGrid([3]))
np.testing.assert_allclose(
f.to_scalar("invariant1").data,
f.to_scalar("invariant3").data)
np.testing.assert_allclose(f.to_scalar("invariant2").data, 0)
# dim == 2
f = Tensor2Field.random_uniform(UnitGrid([3, 3]))
invs = [f.to_scalar(f"invariant{i}").data for i in range(1, 4)]
np.testing.assert_allclose(2 * invs[1], invs[2])
a = np.random.uniform(0, 2 * np.pi) # pick random rotation angle
rot = Tensor2Field(f.grid)
rot.data[0, 0, ...] = np.cos(a)
rot.data[0, 1, ...] = np.sin(a)
rot.data[1, 0, ...] = -np.sin(a)
rot.data[1, 1, ...] = np.cos(a)
f_rot = rot @ f @ rot.transpose() # apply the transpose
for i, inv in enumerate(invs, 1):
np.testing.assert_allclose(
inv,
f_rot.to_scalar(f"invariant{i}").data,
err_msg=f"Mismatch in invariant {i}",
)
# dim == 3
from scipy.spatial.transform import Rotation
f = Tensor2Field.random_uniform(UnitGrid([1, 1, 1]))
rot = Tensor2Field(f.grid)
rot_mat = Rotation.from_rotvec(np.random.randn(3)).as_matrix()
rot.data = rot_mat.reshape(3, 3, 1, 1, 1)
f_rot = rot @ f @ rot.transpose() # apply the transpose
for i in range(1, 4):
np.testing.assert_allclose(
f.to_scalar(f"invariant{i}").data,
f_rot.to_scalar(f"invariant{i}").data,
err_msg=f"Mismatch in invariant {i}",
)
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_complex_tensors():
""" test some complex tensor fields """
grid = CartesianGrid([[0.1, 0.3], [-2, 3]], [3, 4])
shape = (2, 2, 2) + grid.shape
numbers = np.random.random(shape) + np.random.random(shape) * 1j
t1 = Tensor2Field(grid, numbers[0])
t2 = Tensor2Field(grid, numbers[1])
assert t1.is_complex and t2.is_complex
dot_op = t1.make_dot_operator()
# test dot product
res = dot_op(t1.data, t2.data)
for t in (t1 @ t2, t1.dot(t2)):
assert isinstance(t, Tensor2Field)
assert t.grid is grid
np.testing.assert_allclose(t.data, res)
# test without conjugate
dot_op = t1.make_dot_operator(conjugate=False)
res = t1.dot(t2, conjugate=False)
np.testing.assert_allclose(dot_op(t1.data, t2.data), res.data)
def test_from_expressions():
"""test initializing tensor fields with expressions"""
grid = UnitGrid([4, 4])
tf = Tensor2Field.from_expression(grid, [[1, 1], ["x**2", "x * y"]])
xs = grid.cell_coords[..., 0]
ys = grid.cell_coords[..., 1]
np.testing.assert_allclose(tf.data[0, 1], 1)
np.testing.assert_allclose(tf.data[0, 1], 1)
np.testing.assert_allclose(tf.data[1, 0], xs**2)
np.testing.assert_allclose(tf.data[1, 1], xs * ys)
# corner case
with pytest.raises(ValueError):
Tensor2Field.from_expression(grid, "xy")
with pytest.raises(ValueError):
Tensor2Field.from_expression(grid, ["xy"])
with pytest.raises(ValueError):
Tensor2Field.from_expression(grid, ["x"] * 3)
with pytest.raises(ValueError):
Tensor2Field.from_expression(grid, [["x"], [1, 1]])
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}])
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_smoothing_collection():
""" test smoothing of a FieldCollection """
grid = UnitGrid([3, 4], periodic=[True, False])
sf = ScalarField.random_uniform(grid)
vf = VectorField.random_uniform(grid)
tf = Tensor2Field.random_uniform(grid)
fields = FieldCollection([sf, vf, tf])
sgm = 0.5 + np.random.random()
out = fields.smooth(sigma=sgm)
for i in range(3):
np.testing.assert_allclose(out[i].data, fields[i].smooth(sgm).data)
out.data = 0
fields.smooth(sigma=sgm, out=out)
for i in range(3):
np.testing.assert_allclose(out[i].data, fields[i].smooth(sgm).data)
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)
def test_examples_tensor_sph(conservative):
"""compare derivatives of tensorial fields for spherical grids"""
# test explicit expression for which we know the results
grid = SphericalSymGrid(1, 32)
expressions = [["r**4", 0, 0], [0, "r**3", 0], [0, 0, "r**3"]]
tf = Tensor2Field.from_expression(grid, expressions)
# tensor divergence
bc = [{"derivative": 0}, {"normal_derivative": [4, 3, 3]}]
res = tf.divergence(bc, conservative=conservative)
expect = VectorField.from_expression(grid,
["2 * r**2 * (3 * r - 1)", 0, 0])
if conservative:
np.testing.assert_allclose(res.data, expect.data, rtol=0.1, atol=0.1)
else:
np.testing.assert_allclose(res.data[:, 1:-1],
expect.data[:, 1:-1],
rtol=0.1,
atol=0.1)
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)
def test_conservative_sph():
"""test whether the integral over a divergence vanishes"""
grid = SphericalSymGrid((0, 2), 50)
expr = "1 / cosh((r - 1) * 10)"
# test divergence of vector field
vf = VectorField.from_expression(grid, [expr, 0, 0])
div = vf.divergence(bc="derivative", conservative=True)
assert div.integral == pytest.approx(0, abs=1e-2)
# test laplacian of scalar field
lap = vf[0].laplace("derivative")
assert lap.integral == pytest.approx(0, abs=1e-13)
# test double divergence of tensor field
expressions = [[expr, 0, 0], [0, expr, 0], [0, 0, expr]]
tf = Tensor2Field.from_expression(grid, expressions)
res = tf._apply_operator("tensor_double_divergence",
bc="derivative",
conservative=True)
assert res.integral == pytest.approx(0, abs=1e-3)
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_insert_tensor(grid):
"""test the `insert` method"""
f = Tensor2Field(grid)
a = np.random.random(f.data_shape)
c = tuple(grid.get_random_point(coords="cell"))
c_data = (Ellipsis, ) + c
p = grid.transform(c, "cell", "grid")
f.insert(p, a)
np.testing.assert_almost_equal(f.data[c_data], a / grid.cell_volumes[c])
f.insert(grid.get_random_point(coords="grid"), a)
np.testing.assert_almost_equal(f.integral, 2 * a)
f.data = 0 # reset
insert = grid.make_inserter_compiled()
c = tuple(grid.get_random_point(coords="cell"))
c_data = (Ellipsis, ) + c
p = grid.transform(c, "cell", "grid")
insert(f.data, p, a)
np.testing.assert_almost_equal(f.data[c_data], a / grid.cell_volumes[c])
insert(f.data, grid.get_random_point(coords="grid"), a)
np.testing.assert_almost_equal(f.integral, 2 * a)
