def test_direction_change():
# Test custom (non-uniform) grid coordinates
# A uniform field pointing in the x direction
v = np.zeros((4, 4, 4, 3))
# Make first layer of vectors point in x-direction
v[0:2, :, :, 0] = 1
# After that layer, make vectors point in y-direction
v[2:4, :, :, 1] = 1
grid = VectorGrid(v, grid_spacing=[1, 1, 1])
seed = np.array([0, 0, 0])
step_size = 0.1
tracer = StreamTracer(100, step_size)
tracer.trace(seed, grid)
sline = tracer.xs[0]
# Check that initial steps are in x-direction
# Check diff in y/z directions is zero
assert np.allclose(np.diff(sline[0:10, 1:2]), 0)
# Check there's a diff in x direction
assert np.allclose(np.diff(sline[0:10, 0]), 0.1)
# Check that field line changes direction
assert sline[0, 1] == 0
# Check that fline leaves box in y-direction
assert sline[-1, 1] > 3 - step_size
assert 0 < sline[-1, 0] < 3
def test_cyclic_field():
# A uniform field pointing in the x direction
v = np.zeros((100, 100, 100, 3))
# Make all vectors point in the x-direction
v[:, :, :, 0] = 1
# Mismatch vectors on the x-faces
v[0, :, :, 0] = -1
spacing = [1, 1, 1]
cyclic = [True, False, False]
with pytest.raises(AssertionError, match='Arrays are not equal'):
VectorGrid(v, spacing, cyclic=cyclic)
# Check that with cyclic off no error is thrown
cyclic = [False, False, False]
VectorGrid(v, spacing, cyclic=cyclic)
def uniform_x_field():
# A uniform field pointing in the x direction
v = np.zeros((101, 101, 101, 3))
# Make all vectors point in the x-direction
v[:, :, :, 0] = 1
spacing = [1, 1, 1]
return VectorGrid(v, spacing)
def vector_grid(output):
"""
Create a `streamtracer.VectorGrid` object from an `~pfsspy.Output`.
"""
from streamtracer import VectorGrid
# The indexing order on the last index is (phi, s, r)
vectors = output.bg.copy()
# Correct s direction for coordinate system distortion
sqrtsg = output.grid._sqrtsg_correction
# phi correction
with np.errstate(divide='ignore', invalid='ignore'):
vectors[..., 0] /= sqrtsg
# Technically where s=0 Bphi is now infinite, but because this is
# singular and Bphi doesn't matter, just set it to a large number
vectors[~np.isfinite(vectors[..., 0]), 0] = 0.8e+308
# s correction
vectors[..., 1] *= -sqrtsg
grid_spacing = output.grid._grid_spacing
# Cyclic only in the phi direction
# (theta direction becomes singular at the poles so it is not cyclic)
cyclic = [True, False, False]
origin_coord = [0, -1, 0]
vector_grid = VectorGrid(vectors,
grid_spacing,
cyclic=cyclic,
origin_coord=origin_coord)
return vector_grid
def test_bounds(dir):
v = np.zeros((3, 3, 3, 3))
# Make all vectors point along the specified dimension
v[:, :, :, dir] = 1
spacing = [1, 1, 1]
grid = VectorGrid(v, spacing)
seed = np.array([[0.5, 0.5, 0.5]])
tracer = StreamTracer(max_steps=10, step_size=1.0)
tracer.trace(seed, grid)
expected = np.roll(np.array([1.5, 0.5, 0.5]), dir)
assert (tracer.xs[0][-1, :] == expected).all()
expected = np.array([0.5, 0.5, 0.5])
assert (tracer.xs[0][0, :] == expected).all()
def test_bad_input(tracer, uniform_x_field):
# Check input validation
seed = np.array([0, 0, 0])
with pytest.raises(ValueError, match='seeds must be a 2D array'):
tracer.trace(np.array([[[1], [1]], [[1], [1]]]), uniform_x_field)
with pytest.raises(ValueError, match='seeds must have shape'):
tracer.trace(np.array([1, 1]), uniform_x_field)
with pytest.raises(ValueError,
match='grid must be an instance of StreamTracer'):
tracer.trace(seed, 1)
with pytest.raises(ValueError, match='Direction must be -1, 1 or 0'):
tracer.trace(seed, uniform_x_field, direction=2)
with pytest.raises(ValueError, match='vectors must be a 4D array'):
VectorGrid(np.array([1]), [1, 1, 1])
with pytest.raises(ValueError, match='vectors must have shape'):
VectorGrid(np.zeros((1, 1, 1, 2)), [1, 1, 1])
with pytest.raises(ValueError, match='grid spacing must have shape'):
VectorGrid(np.zeros((1, 1, 1, 3)), [1, 1])
def test_grid_points():
# A uniform field pointing in the x direction
v = np.zeros((100, 100, 100, 3))
# Make all vectors point in the x-direction
spacing = [2, 3, 4]
origin_coord = [4, 9, 16]
grid = VectorGrid(v, spacing, origin_coord=origin_coord)
assert grid.xcoords[0] == 4
assert grid.xcoords[1] == 6
assert grid.xcoords[-1] == 4 + 99 * 2
assert grid.ycoords[0] == 9
assert grid.ycoords[1] == 12
assert grid.ycoords[-1] == 9 + 99 * 3
assert grid.zcoords[0] == 16
assert grid.zcoords[1] == 20
assert grid.zcoords[-1] == 16 + 99 * 4
def test_origin(tracer, origin_coord, ds):
# A uniform field pointing in the x direction
v = np.zeros((4, 4, 4, 3))
# Make all vectors point in the x-direction
v[:, :, :, 0] = 1
spacing = [1, 1, 1]
grid = VectorGrid(v, spacing, origin_coord=origin_coord)
seed = np.array([2., 2., 2.])
tracer = StreamTracer(100, ds)
tracer.trace(seed, grid, direction=1)
fline = tracer.xs[0]
# Check first field line coordinate is the seed
np.testing.assert_equal(fline[0, :], seed)
# Should stop before the edge of the box
end_coord = seed
end_coord[0] = grid.xcoords[-1] - ds
np.testing.assert_almost_equal(fline[-1, :], end_coord)
def test_coords():
# Test custom (non-uniform) grid coordinates
# A uniform field pointing in the x direction
v = np.zeros((4, 4, 4, 3))
# Make all vectors point diagonally from one corner to the other
v[:, :, :, :] = 1
xcoords = [0, 1, 2, 10]
ycoords = [0, 3, 6, 10]
zcoords = [0, 8, 9, 10]
grid = VectorGrid(v, grid_coords=[xcoords, ycoords, zcoords])
seed = np.array([0, 0, 0])
tracer = StreamTracer(100, 1)
tracer.trace(seed, grid)
# Check that step sizes are all 1
sline = tracer.xs[0]
ds = np.diff(np.linalg.norm(sline, axis=1))
np.testing.assert_equal(ds[0], 1)
np.testing.assert_almost_equal(ds[1:], 1)
# Check that first/last steps are outside box
assert np.all(sline[0] < 0 + tracer.ds)
assert np.all(sline[-1] > 10 - tracer.ds)