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_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_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)
def test_cyclic(uniform_x_field):
# Check the cyclic option
maxsteps = 4
tracer = StreamTracer(maxsteps, 0.1)
seed = np.array([99.95, 50, 50])
uniform_x_field.cyclic = [True, False, False]
tracer.trace(seed, uniform_x_field, direction=1)
fline = tracer.xs[0]
# Check that the tracer uses all the steps available
assert len(fline) == maxsteps
assert tracer.max_steps == maxsteps
# Check that the cyclic boundary works properly
np.testing.assert_equal(fline[0, :], seed)
np.testing.assert_almost_equal(fline[1, :], np.array([0.05, 50, 50]))
np.testing.assert_almost_equal(fline[2, :], np.array([0.15, 50, 50]))
# Check that going the other way across the boundary (through zero) works
seed = np.array([0.1, 50, 50])
tracer.trace(seed, uniform_x_field, direction=-1)
fline = tracer.xs[0]
np.testing.assert_equal(fline[0, :], seed)
np.testing.assert_almost_equal(fline[1, :], np.array([0, 50, 50]))
np.testing.assert_almost_equal(fline[2, :], np.array([99.9, 50, 50]))
# Check that turning cyclic off interrupts the tracing
uniform_x_field.cyclic = [False, False, False]
# Going forwards
seed = np.array([99.9, 50, 50])
tracer.trace(seed, uniform_x_field, direction=1)
assert len(tracer.xs[0]) == 2
# Going backwards
seed = np.array([0.1, 50, 50])
tracer.trace(seed, uniform_x_field, direction=-1)
assert len(tracer.xs[0]) == 2
class FortranTracer(Tracer):
r"""
Tracer using Fortran code.
Parameters
----------
max_steps: int
Maximum number of steps each streamline can take before stopping.
step_size : float
Step size as a fraction of cell size at the equator.
Notes
-----
Because the stream tracing is done in spherical coordinates, there is a
singularity at the poles (ie. :math:`s = \pm 1`), which means seeds placed
directly on the poles will not go anywhere.
"""
def __init__(self, max_steps=1000, step_size=0.01):
try:
from streamtracer import StreamTracer
except ModuleNotFoundError as e:
raise RuntimeError(
'Using FortranTracer requires the streamtracer module, '
'but streamtracer could not be loaded') from e
self.max_steps = max_steps
self.step_size = step_size
self.tracer = StreamTracer(max_steps, step_size)
@staticmethod
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 trace(self, seeds, output):
self.validate_seeds(seeds)
x, y, z = self.coords_to_xyz(seeds, output)
r, lat, phi = astrocoords.cartesian_to_spherical(x, y, z)
# Force 360deg wrapping
phi = astrocoords.Longitude(phi).to_value(u.rad)
s = np.sin(lat).to_value(u.dimensionless_unscaled)
rho = np.log(r)
seeds = np.atleast_2d(np.stack((phi, s, rho), axis=-1))
# Get a grid
vector_grid = self.vector_grid(output)
# Do the tracing
self.tracer.trace(seeds, vector_grid)
xs = self.tracer.xs
rots = self.tracer.ROT
if np.any(rots == 1):
warnings.warn(
'At least one field line ran out of steps during tracing.\n'
'You should probably increase max_steps '
f'(currently set to {self.max_steps}) and try again.')
xs = [
np.stack(pfsspy.coords.strum2cart(x[:, 2], x[:, 1], x[:, 0]),
axis=-1) for x in xs
]
flines = [
fieldline.FieldLine(x[:, 0], x[:, 1], x[:, 2], output) for x in xs
]
return fieldline.FieldLines(flines)