locplot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Created by schaffer at 4/17/19

"""

import datetime
import os
from os.path import exists
from os.path import join
from shutil import move

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
import sunpy.map
from matplotlib.patches import Ellipse, Polygon
from sunpy.net.helioviewer import HelioviewerClient


def get_hmi_continuum(date):
    """
    Get HMI continuum jp2 image. Check /dat/sdc/hmi_context_cache first, if nothing is found, download from SDO via Helioviewer
    Args:
        date: datetime query

    Returns:
        path: path to the jp2 file
    """
    cache = "/dat/sdc/hmi_context_cache/"
    filename = f"HMI_continuum_{date.strftime('%Y%m%d_%H-%M-%S')}.jp2"
    if not exists(cache):
        os.makedirs(cache)
    path = join(cache, filename)
    if not exists(path):
        hv = HelioviewerClient()
        file = hv.download_jp2(date, source_id=18)
        if not file:
            raise IOError(f"Could not download file for {date}")
        move(file, path)

    return path


def make_loc_plot(coords, date, uncertainties=None):
    """
    Overplot coordionates onto an HMI image.

    Args:
        coords: helioprojective coordinates in arcsec, possible shapes are:
                (2) (x_center/y_center)          : Simple point
                (2,2) (point, x_center/y_center) : Bounding Box (upper left and lower right corner)
                (4,2) (point, x_center/y_center) : FOV (Corners of field of view)
        date (datetime.datetime):  date of observation
        uncertainties: (sX, sY): coordinate certainty (standard deviation) in arcseconds, defaults to None
    Returns:


    """

    global x_center
    corners = None
    coords = np.array(coords)

    if coords.shape == (2,):
        x_center, y_center = coords
    elif coords.shape == (2, 2):
        x_center, y_center = coords.mean(axis=0)
        xmin, ymin = coords.min(axis=1)
        xmax, ymax = coords.max(axis=1)
        corners = np.array(((xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin)))
    elif coords.shape == (4, 2):
        x_center, y_center = coords.mean(axis=0)
        corners = coords
    else:
        raise ValueError('Invalid input shape! Use (2,), or (2,2)/(4,2) (Point,X/Y)')

    # Get hmi reference image
    continuum_file = get_hmi_continuum(date)

    # Set position of the labels depending on the position
    if x_center > 0:
        xlabel_pos = x_center - 275
        ylabel_pos = y_center - 19

    if x_center < 0:
        xlabel_pos = x_center + 275
        ylabel_pos = y_center - 19

    if y_center > 0:
        ylabel_pos = y_center - 110
        xlabel_pos = x_center

    if y_center < 0:
        ylabel_pos = y_center + 80
        xlabel_pos = x_center

    if y_center == 0 and x_center == 0:
        ylabel_pos = y_center - 110
        xlabel_pos = x_center
    else:
        pass

    # Convert position from m in arcsec
    xc = [x_center] * u.arcsec
    yc = [y_center] * u.arcsec

    # Create the picture
    hmi_map = sunpy.map.Map(continuum_file)
    figure, axes = plt.subplots()
    hmi_map._default_heliographic_longitude = 0
    hmi_map.plot(axes)

    # Add labels
    title = axes.title.get_text()
    axes.annotate(title, xy=(0, 1010), ha="center", color="white", size=7)
    # Plot marker if no corners are given
    if corners is None:
        plt.plot(xc, yc, color="red", marker="x", markersize=10, mew=2)
    plt.tight_layout(pad=0)
    marker_label = f"{x_center:-7.1f}|{y_center:-7.1f}"
    axes.annotate(
        marker_label,
        xy=(xlabel_pos, ylabel_pos),
        ha="center",
        color="red",
        size=10,
    )

    # Disable automatic labels
    axes.set_title("")
    plt.xlabel("")
    plt.ylabel("")

    # Define x & y axis
    axes.set_xlim([-1100, 1100])
    axes.set_ylim([-1100, 1100])
    axes.tick_params(
        axis="y",
        direction="in",
        pad=-5,
        length=0,
        labelsize=6,
        colors="red",
        grid_linestyle="dotted",
        grid_linewidth=0.5,
        grid_color="red",
    )
    axes.tick_params(
        axis="x",
        direction="in",
        pad=-5,
        length=0,
        labelsize=6,
        colors="red",
        grid_linestyle="dotted",
        grid_linewidth=0.5,
        grid_color="red",
    )
    plt.xticks(np.arange(-1000, 1001, 200))
    plt.yticks(np.arange(-1000, 1001, 200))
    for tick in axes.yaxis.get_major_ticks():
        tick.label1.set_horizontalalignment("left")
        tick.label1.set_verticalalignment("bottom")
    for tick in axes.xaxis.get_major_ticks():
        tick.label1.set_horizontalalignment("left")
        tick.label1.set_verticalalignment("bottom")
    plt.grid(True)

    edge_length = 8

    figure.set_size_inches(edge_length, edge_length)
    axes.set_facecolor("black")
    figure.patch.set_facecolor("black")

    # Add uncertainty marker if uncertainties are given.
    if uncertainties:
        one_sigma = Ellipse(
            (x_center, y_center),
            2 * uncertainties[0],
            2 * uncertainties[1],
            edgecolor="red",
            facecolor="red",
            alpha=0.1,
            label="1 Std. Deviation",
        )
        two_sigma = Ellipse(
            (x_center, y_center),
            2 * 2 * uncertainties[0],
            2 * 2 * uncertainties[1],
            edgecolor="red",
            facecolor="red",
            alpha=0.066,
            label="2 Std. Deviations",
        )
        three_sigma = Ellipse(
            (x_center, y_center),
            3 * 2 * uncertainties[0],
            3 * 2 * uncertainties[1],
            edgecolor="red",
            facecolor="red",
            alpha=0.033,
            label="3 Std. Deviations",
        )
        axes.add_patch(one_sigma)
        axes.add_patch(two_sigma)
        axes.add_patch(three_sigma)

    # Add range box if range was specified
    if corners is not None:
        pol = Polygon(corners, closed=True, label='Field of View', fill=False, color='red')
        axes.add_patch(pol)
    return figure, axes


if __name__ == "__main__":
    fig, ax = make_loc_plot((200, -152), datetime.datetime.now())
    plt.show()