2022-06-6--sunpy-help.py

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
import sunpy.map
from astropy.coordinates import SkyCoord
from matplotlib.animation import FuncAnimation
from sunpy.coordinates import RotatedSunFrame
from sunpy.coordinates.utils import GreatArc
from sunpy.net import Fido
from sunpy.net import attrs as a


################################################################################
# These two functions are going to be added to sunpy in version 4.1
# They are copied here so this script runs on 4.0
def _bresenham(x1, y1, x2, y2):
    """
    Returns an array of all pixel coordinates which the line defined by `x1, y1` and
    `x2, y2` crosses. Uses Bresenham's line algorithm to enumerate the pixels along
    a line. This was adapted from ginga.
    Parameters
    ----------
    x1, y1, x2, y2 :`int`
    References
    ----------
    * https://github.com/ejeschke/ginga/blob/c8ceaf8e559acc547bf25661842a53ed44a7b36f/ginga/BaseImage.py#L503
    * http://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm
    """
    for x in [x1, y1, x2, y2]:
        if type(x) not in (int, np.int64):
            raise TypeError('All pixel coordinates must be of type int')
    dx = abs(x2 - x1)
    dy = abs(y2 - y1)
    sx = 1 if x1 < x2 else -1
    sy = 1 if y1 < y2 else -1
    err = dx - dy
    res = []
    x, y = x1, y1
    while True:
        res.append((x, y))
        if (x == x2) and (y == y2):
            break
        e2 = 2 * err
        if e2 > -dy:
            err = err - dy
            x += sx
        if e2 < dx:
            err = err + dx
            y += sy
    return np.array(res)


def extract_along_coord(smap, coord):
    """
    Return the value of the image array at every point along the coordinate.
    For a given coordinate ``coord``, find all the pixels that cross the coordinate
    and extract the values of the image array in ``smap`` at these points. This is done by applying
    `Bresenham's line algorithm <http://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm>`_
    between the consecutive pairs of points in the coordinate and then indexing the data
    array of ``smap`` at those points.
    Parameters
    ----------
    smap : `~sunpy.map.GenericMap`
    coord : `~astropy.coordinates.SkyCoord`
        Coordinate along which to extract intensity
    Returns
    -------
    intensity : `~astropy.units.Quantity`
    loop_coord : `~astropy.coordinates.SkyCoord`
    """
    if not len(coord.shape) or coord.shape[0] < 2:
        raise ValueError('At least two points are required for extracting intensity along a '
                         'line. To extract points at single coordinates, use '
                         'sunpy.map.maputils.sample_at_coords.')
    if not all(sunpy.map.contains_coordinate(smap, coord)):
        raise ValueError('At least one coordinate is not within the bounds of the map.'
                         'To extract the intensity along a coordinate, all points must fall within '
                         'the bounds of the map.')
    # Find pixels between each loop segment
    px, py = smap.wcs.world_to_array_index(coord)
    pix = []
    for i in range(len(px)-1):
        b = _bresenham(px[i], py[i], px[i+1], py[i+1])
        # Pop the last one, unless this is the final entry because the first point
        # of the next section will be the same
        if i < (len(px) - 2):
            b = b[:-1]
        pix.append(b)
    pix = np.vstack(pix)

    intensity = u.Quantity(smap.data[pix[:, 0], pix[:, 1]], smap.unit)
    coord_new = smap.pixel_to_world(pix[:, 1]*u.pix, pix[:, 0]*u.pix)

    return intensity, coord_new
################################################################################

# Fetch data
res = Fido.search(a.Time("2022-5-10T00:00:00", "2022-5-10T23:00:00"),
                  a.Instrument("aia"), a.Wavelength(171 * u.angstrom),
                  a.Sample(1 * u.hour))
res1 = Fido.fetch(res)

# Make a mapsequence
seq = sunpy.map.Map(res1, sequence=True)

# Now we want to create a stack of co-rotating submaps
# To do this in a nice way we are going to define a centre point of interest in
# the first frame
centre_point = SkyCoord(-50, -425, unit=u.arcsec, frame=seq[0].coordinate_frame)

# Define the height and width of the cutout
cutout_width = 700 * u.arcsec
cutout_height = 650 * u.arcsec

aia_seq = []
arc_coords = []
intensities = []
for img in seq:
    # Now we define a coordinate frame which rotates with the sun.
    # The "metaframe" RotatedSunFrame takes the point we want to rotate and an
    # amount of solar rotation to apply (the difference in time between the
    # first image and this image.)
    rotating_frame = RotatedSunFrame(base=centre_point, duration=(img.date - seq[0].date).to(u.s))
    # We then transform the rotated point to the coordinate system of this image
    centre = rotating_frame.transform_to(img.coordinate_frame)
    # and use it to define the lower and upper corners of our box
    bl = SkyCoord(centre.Tx - cutout_width/2 , centre.Ty - cutout_height/2, frame=img.coordinate_frame)
    aia_seq.append(img.submap(bottom_left=bl, height=cutout_height, width=cutout_width))


    # We use the centre coordinate to define the slit so it rotates at the same
    # rate as the submap
    start = SkyCoord(centre.Tx - 50*u.arcsec, -280 * u.arcsec, frame=img.coordinate_frame)
    end = SkyCoord(start.Tx, start.Ty + 110*u.arcsec, frame=img.coordinate_frame)
    great_arc = GreatArc(start, end)

    intensity, coords = extract_along_coord(img, great_arc.coordinates())
    arc_coords.append(coords)
    intensities.append(intensity)

seq = sunpy.map.Map(aia_seq, sequence=True)


# Plot the maps with sunpy
# def no_grid(fig, axes, smap):
#     ind = seq.maps.index(smap)
#     lon, lat = axes.coords
#     lon.grid(draw_grid=False)
#     lat.grid(draw_grid=False)

#     coords = arc_coords[ind]
#     line, = axes.plot_coord(coords, color="blue")
#     # This function must return a list of items to remove in the next frame
#     return [line]

# # Interactive plot
# # seq.peek(plot_function=no_grid)

# # Save movie plot
# ani = seq.plot(plot_function=no_grid)
# ani.save("test_sunpy_track.mp4")


# Animate the map and the line next to each other
def animate(frame, ax1, im, intensity_line, remove_items):
    ind = frame % len(aia_seq)

    # Remove any items from the last iteration we don't want to still be in the
    # frame
    while remove_items:
        remove_items.pop().remove()

    # We have to reset the WCS to set it to the one for the next frame
    ax1.reset_wcs(aia_seq[ind].wcs)
    # Set the data for the map plot
    im.set_data(aia_seq[ind].data)

    # Reset the labels
    ax1.set_title(aia_seq[ind].latex_name)
    ax1.set_xlabel("Helioprojective Longitude [arcsec]")
    ax1.set_ylabel("Helioprojective Latitude [arcsec]")

    # Plot the line for the slit
    coord_line, = ax1.plot_coord(arc_coords[ind], color="blue")
    # Remove this line next time around
    remove_items.append(coord_line)

    # Update the line plots data
    intensity_line.set_data(np.arange(len(intensities[ind])), intensities[ind])

# We are going to use this list as a way of passing variables between calls to
# animate
remove_items = []

fig = plt.figure(figsize=(10, 5))
ax1 = plt.subplot(1, 2, 1, projection=aia_seq[0])
im = aia_seq[0].plot(axes=ax1)

ax2 = plt.subplot(1, 2, 2)
intensity_line, = ax2.plot(intensities[0])
max_y = np.max([i.value.max() for i in intensities])
ax2.set_ylim((0, max_y))
ax2.set_xlabel("Pixels along slit")
ax2.set_ylabel(f"Intensity along slit [{intensities[0].unit}]")

ani = FuncAnimation(fig, animate, fargs=(ax1, im, intensity_line, remove_items))
ani.save("sunpy_test_slit.mp4")

sunpy_test_slit.mp4