Simu: add function to quickly plot geometry

simulations/lib/location/example.py

@@ -3,6 +3,7 @@
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import axes3d
 
+import location as loc
 from antenna import Receiver, Emitter
 
 # 2D showcase
@@ -16,16 +17,7 @@ antennae = [
 ]
 
 fig, ax = plt.subplots()
-
+loc.plot_geometry(ax, [source], antennae)
-ax.set_title("Geometry of Emitter(s) and Antennae")
-ax.set_ylabel("y")
-ax.set_xlabel("x")
-ax.plot(*source.x, '*', label="Emitter")
-
-for j, ant in enumerate(antennae):
-    ax.plot(*ant.x, '+', label="Antenna {}".format(j))
-
-ax.legend()
 fig.show()
 
 # 3D showcase

simulations/lib/location/location.py

@@ -7,6 +7,56 @@ def distance(x1, x2):
     """
     return np.sqrt( np.sum( (x1 - x2)**2, axis=0) )
 
+def plot_geometry(ax, emitters=[], antennae=[], unit='m'):
+    """
+    Show the geometry of emitters and antennae in a square plot.
+
+    Parameters
+    ----------
+    ax - matplotlib.Axes
+        The axis object to plot the geometry on.
+    emitters - list of Locations
+        The Emitter objects to plot.
+    antennae - list of Locations
+        The Receiver objects to plot.
+
+    Returns
+    -------
+    ax - matplotlib.Axes
+        The axis object containing the plotted geometry.
+    annots - dict of list of matplotlib.text.Annotation
+        The dictionary is split up into a list of annotations
+        belonging to the emitters, and one for the antennae.
+    """
+
+    ax.grid()
+    ax.set_title("Geometry of Emitter(s) and Antennae")
+    ax.set_ylabel("y ({})".format(unit))
+    ax.set_xlabel("x ({})".format(unit))
+    ax.margins(0.3)
+    ax.set_aspect('equal', 'datalim') # make it a square plot
+
+    annots = {}
+    for k, locs in {"E": emitters, "A": antennae}.items():
+        if k == "E":
+            marker='*'
+            prefix = k
+        elif k == "A":
+            marker="o"
+            prefix = k
+
+        # create the list of annotations
+        if k not in annots:
+            annots[k] = []
+
+        # plot marker and create annotation
+        for j, loc in enumerate(locs):
+            label = "{}{}".format(prefix, j)
+            ax.plot(*loc.x, marker=marker, label=label)
+            annots[k].append(ax.annotate(label, loc.x))
+
+    return ax, annots
+
 class Location:
     """
     A location is a point designated by a spatial coordinate x.