コード例 #1
0
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
コード例 #2
0
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()
コード例 #3
0
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)
コード例 #4
0
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)
コード例 #5
0
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
コード例 #6
0
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)