Prettify legend

zonnewaarneming.py

@@ -12,6 +12,7 @@ import time
 import matplotlib.pyplot as plt
 import pandas as pd
 import zmq
+import re
 import struct
 import numpy as np
 from telescope import Telescope
@@ -58,12 +59,19 @@ cas_a = SkyCoord.from_name("Cassiopeia A")
 dwl = EarthLocation(lon=Angle("6:23:46.21 degrees"),
                     lat=Angle("52:48:43.27 degrees"))
 
-sourcenames = ["Cassiopeia A", "Cygnus A", "Perseus A", "Taurus A", "Zon"]
+sourcenames = ["Cassiopeia A\n(supernovarestant)", "Cygnus A\n(sterrenstelsel)", "Perseus A\n(sterrenstelsel)", "Taurus A (Krabnevel)", "Andromeda (sterrenstelsel)", "Zon"]
+
+pattern = re.compile(r'[ \n]+\(.*\)')
+def clean_name(prettyname):
+    """Strip things between brackets"""
+    #return re.sub(r'[ \n]+\(.*\)', '', prettyname)
+    return pattern.sub('', prettyname)
+
 sources = {}
 for sourcename in sourcenames:
     if sourcename == "Zon":
         continue
-    sources[sourcename] = SkyCoord.from_name(sourcename)
+    sources[sourcename] = SkyCoord.from_name(clean_name(sourcename))
 
 # create plot
 plt.ion()
@@ -75,14 +83,31 @@ ax.axvline(x=0, linestyle='--', color='grey', linewidth=.5)
 
 scatterplots = {}
 for sourcename in sourcenames:
-    scatterplots[sourcename] = ax.scatter([], [], s=3, label=sourcename)
+    scatterplots[sourcename] = ax.scatter([], [], s=3, label=None)
 
 #ax.set_ylim((0, None))
 ax.set_xlabel("Afstand tot middelpunt (graden)")
 ax.set_ylabel("Signaalsterkte (ongekalibreerd)")
 fig.canvas.set_window_title('Live waarneming Dwingeloo radiotelescoop op 1330 MHz')
 ax.set_title('Live waarneming Dwingeloo radiotelescoop op 1330 MHz')
-ax.legend()
+
+def update_legend(ax, sourcename):
+    somethingchanged = False
+    scatterplot = scatterplots[sourcename]
+    if scatterplot.get_label()[0] == "_":
+        num_points = scatterplot.get_offsets().data.shape[0]
+        if num_points > 5:
+            scatterplot.set_label(sourcename)
+            somethingchanged = True
+            print("Adding", sourcename.replace("\n", " "), "to legend")
+    if somethingchanged:
+        legend = ax.get_legend()
+        if legend is not None:
+            legend.remove()
+        ax.legend()
+
+plt.show()
+plt.pause(0.001)
 
 prev_points = None
 separation = None
@@ -116,10 +141,12 @@ while not stoprequest:
             separation = source_separation
 
     if prev_points is not None:
-        prev_points.remove()
+        try:
+            prev_points.remove()
+        except:
+            pass
 
     if separation > 6 * u.deg:
-        print("Far from everything, plotting nothing")
         continue
 
     source_altaz = sources[closest_source].transform_to(dwl_altaz)
@@ -134,7 +161,7 @@ while not stoprequest:
 
     scatterplot = scatterplots[closest_source]
 
-    print(Time.now(), separation, az.to(u.deg).value, alt.to(u.deg).value, signal, closest_source, file=datafile, flush=True, sep=',')
+    print(Time.now(), separation, az.to(u.deg).value, alt.to(u.deg).value, signal, clean_name(closest_source), file=datafile, flush=True, sep=',')
 
     if signal > maxsignal:
         maxsignal = signal
@@ -146,6 +173,9 @@ while not stoprequest:
     prev_points = ax.scatter([separation.deg], [signal], s=50, color='red')
     ax.set_ylim((2800, 1.1 * maxsignal))
 
+    update_legend(ax, closest_source)
+
     # show the plot
     plt.show()
+
     plt.pause(0.0001)