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)