From 261a3f5505291bac5f53f490c3a41f0e63228a87 Mon Sep 17 00:00:00 2001
From: Yaraslau <yaraslau.tamashevich@gmail.com>
Date: Thu, 12 Oct 2023 02:21:36 +0300
Subject: [PATCH] Add label for geometrical object (#2631)

* Add label for geometrical object

* Fix formatting

* Change default behaviour and check label for None

* fix bug
---
 python/geom.py          |  5 ++++-
 python/simulation.py    |  4 ++++
 python/visualization.py | 46 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/python/geom.py b/python/geom.py
index b4756ab21..6642bf31b 100755
--- a/python/geom.py
+++ b/python/geom.py
@@ -1067,7 +1067,9 @@ class GeometricObject:
     ```
     """
 
-    def __init__(self, material=Medium(), center=Vector3(), epsilon_func=None):
+    def __init__(
+        self, material=Medium(), center=Vector3(), epsilon_func=None, label=None
+    ):
         """
         Construct a `GeometricObject`.
 
@@ -1095,6 +1097,7 @@ def __init__(self, material=Medium(), center=Vector3(), epsilon_func=None):
             epsilon_func.eps = True
             material = epsilon_func
 
+        self.label = label
         self.material = material
         self.center = Vector3(*center)
 
diff --git a/python/simulation.py b/python/simulation.py
index 390dc6121..08eaa61a5 100644
--- a/python/simulation.py
+++ b/python/simulation.py
@@ -4745,6 +4745,7 @@ def plot2D(
         output_plane: Optional[Volume] = None,
         fields: Optional = None,
         labels: bool = False,
+        label_geometry: bool = True,
         eps_parameters: Optional[dict] = None,
         boundary_parameters: Optional[dict] = None,
         source_parameters: Optional[dict] = None,
@@ -4797,6 +4798,8 @@ def plot2D(
           no fields are superimposed.
         * `labels`: if `True`, then labels will appear over each of the simulation
           elements. Defaults to `False`.
+        * `label_geometry`: if `True`, then labels will appear over each of the geometry
+          elements. Defaults to `True`.
         * `eps_parameters`: a `dict` of optional plotting parameters that override the
           default parameters for the geometry.
             - `interpolation='spline36'`: interpolation algorithm used to upsample the pixels.
@@ -4871,6 +4874,7 @@ def plot2D(
             output_plane=output_plane,
             fields=fields,
             labels=labels,
+            label_geometry=label_geometry,
             eps_parameters=eps_parameters,
             boundary_parameters=boundary_parameters,
             source_parameters=source_parameters,
diff --git a/python/visualization.py b/python/visualization.py
index a0545ec14..ff1589a73 100644
--- a/python/visualization.py
+++ b/python/visualization.py
@@ -92,6 +92,13 @@
 
 default_label_parameters = {"label_color": "r", "offset": 20, "label_alpha": 0.3}
 
+default_label_geometry_parameters = {
+    "label_color": "w",
+    "arrow_color": "r",
+    "offset": 20,
+    "label_alpha": 0.8,
+}
+
 # Used to remove the elements of a dictionary (dict_to_filter) that
 # don't correspond to the keyword arguments of a particular
 # function (func_with_kwargs.)
@@ -138,6 +145,10 @@ def place_label(
     offset = label_parameters["offset"]
     alpha = label_parameters["label_alpha"]
     color = label_parameters["label_color"]
+    if "arrow_color" in label_parameters:
+        arrow_color = label_parameters["arrow_color"]
+    else:
+        arrow_color = color
 
     if x > centerx:
         xtext = -offset
@@ -156,7 +167,9 @@ def place_label(
         ha="center",
         va="bottom",
         bbox=dict(boxstyle="round,pad=0.2", fc=color, alpha=alpha),
-        arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=0.5", color=color),
+        arrowprops=dict(
+            arrowstyle="->", connectionstyle="arc3,rad=0.5", color=arrow_color
+        ),
     )
     return ax
 
@@ -507,6 +520,7 @@ def plot_eps(
     eps_parameters: Optional[dict] = None,
     colorbar_parameters: Optional[dict] = None,
     frequency: Optional[float] = None,
+    label_geometry: bool = False,
 ) -> Union[Axes, Any]:
     # consolidate plotting parameters
     if eps_parameters is None:
@@ -615,6 +629,34 @@ def plot_eps(
                 colorbar_parameters=colorbar_parameters,
             )
 
+        if label_geometry:
+            for el in sim.geometry:
+                if sim_size.x == 0:
+                    center_first = el.center.y
+                    center_second = el.center.z
+                    sim_first = sim_center.y
+                    sim_second = sim_center.z
+                elif sim_size.y == 0:
+                    center_first = el.center.x
+                    center_second = el.center.z
+                    sim_first = sim_center.x
+                    sim_second = sim_center.z
+                elif sim_size.z == 0:
+                    center_first = el.center.x
+                    center_second = el.center.y
+                    sim_first = sim_center.x
+                    sim_second = sim_center.y
+                if el.label is not None:
+                    ax = place_label(
+                        ax,
+                        el.label,
+                        center_first,
+                        center_second,
+                        sim_first,
+                        sim_second,
+                        label_parameters=default_label_geometry_parameters,
+                    )
+
         ax.set_xlabel(xlabel)
         ax.set_ylabel(ylabel)
         return ax
@@ -903,6 +945,7 @@ def plot2D(
     output_plane: Optional[Volume] = None,
     fields: Optional = None,
     labels: bool = False,
+    label_geometry: bool = False,
     eps_parameters: Optional[dict] = None,
     boundary_parameters: Optional[dict] = None,
     source_parameters: Optional[dict] = None,
@@ -945,6 +988,7 @@ def plot2D(
             eps_parameters=eps_parameters,
             colorbar_parameters=colorbar_parameters,
             frequency=frequency,
+            label_geometry=label_geometry,
         )
 
     # Plot boundaries