From 17ff71a8c8887c57b5741cc340875a297488f85b Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Wed, 28 Feb 2024 20:12:44 -0800
Subject: [PATCH 01/23] feat: implement plot_component (E/H x/y/z)

---
 femwell/maxwell/waveguide.py | 81 ++++++++++++++++++++++++++++++++++--
 1 file changed, 78 insertions(+), 3 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index d9cdd1cb..d677e4b9 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -1,7 +1,7 @@
 """Waveguide analysis based on https://doi.org/10.1080/02726340290084012."""
 from dataclasses import dataclass
 from functools import cached_property
-from typing import List, Tuple
+from typing import List, Literal, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -280,8 +280,83 @@ def plot(self, field, plot_vectors=False, colorbar=True, direction="y", title="E
             direction=direction,
         )
 
-    def show(self, field, **kwargs):
-        self.plot(field=field, **kwargs)
+    def plot_component(
+            self,
+            field_name: Literal["E", "H"],
+            component: Literal["x", "y", "z"],
+            part: Literal["real", "imag", "abs"]="real",
+            boundaries: bool=True,
+            colorbar: bool=False,
+            ax: Axes=None):
+        from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+        if part == "real":
+            conv_func = np.real
+        elif part == "imag":
+            conv_func = np.imag
+        elif part == "abs":
+            conv_func = np.abs
+        else:
+            raise ValueError("A valid part is 'real', 'imag' or 'abs'.")
+
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = plt.gcf()
+
+        if field_name == "E":
+            mfield = self.E
+        elif field_name == "H":
+            mfield = self.H
+        else:
+            raise ValueError("A valid field_name is 'E' or 'H'.")
+
+        (mfield_t, mfield_t_basis), (mfield_n, mfield_n_basis) = self.basis.split(mfield)
+
+        if component == "x" or component == "y":
+            plot_basis = mfield_t_basis.with_element(ElementVector(ElementDG(ElementTriP1())))
+            mfield_xy = plot_basis.project(mfield_t_basis.interpolate(conv_func(mfield_t)))
+            (mfield_x, mfield_x_basis), (mfield_y, mfield_y_basis) = plot_basis.split(mfield_xy)
+            if component == "x":
+                mfield_x_basis.plot(mfield_x, ax=ax, shading="gouraud")
+            else:
+                mfield_y_basis.plot(mfield_y, ax=ax, shading="gouraud")
+        elif component == "z":
+            # plot_basis = mfield_n_basis.with_element(ElementVector(ElementDG(ElementTriP1())))
+            plot_basis = mfield_n_basis
+            mfield_z, mfield_z_basis = mfield_n, mfield_n_basis
+            mfield_z_basis.plot(conv_func(mfield_z), ax=ax, shading="gouraud")
+        else:
+            raise ValueError("A valid component is 'x', 'y' or 'z'.")
+
+        if boundaries:
+            # self.basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
+            # for subdomain in self.basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
+            #     self.basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="w")
+            #plot_basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
+            for subdomain in plot_basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
+                plot_basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="k")
+        if colorbar:
+            divider = make_axes_locatable(ax)
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+            plt.colorbar(ax.collections[-1], cax=cax)
+
+        ax.set_title(f"{field_name}{component} ({part}. part)")
+
+        return fig, ax
+
+    def show(
+        self,
+        field_name: Literal["E", "H"],
+        part: Literal["real", "imag", "abs"]="real",
+        boundaries: bool=True,
+        colorbar: bool=False,
+    ):
+        fig, axs = plt.subplots(1,3, subplot_kw=dict(aspect=1))
+
+        for id_ax, comp in enumerate("xyz"):
+            self.plot_component(field_name, comp, part, boundaries, colorbar, axs[id_ax])
+        plt.tight_layout()
         plt.show()
 
     def plot_intensity(

From 3020694c84135e5525e16bb439149a0b62716f1e Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Wed, 28 Feb 2024 20:25:18 -0800
Subject: [PATCH 02/23] doc: update examples

---
 docs/photonics/examples/effective_area.py  | 2 +-
 docs/photonics/examples/selecting_modes.py | 6 +++---
 docs/photonics/examples/waveguide_modes.py | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/docs/photonics/examples/effective_area.py b/docs/photonics/examples/effective_area.py
index 0b36083a..dbf3f6c2 100644
--- a/docs/photonics/examples/effective_area.py
+++ b/docs/photonics/examples/effective_area.py
@@ -92,7 +92,7 @@
 
         for mode in modes:
             if mode.tm_fraction > 0.5:
-                # mode.show(np.real(mode.E))
+                # mode.show("E", part="real")
                 print(f"Effective refractive index: {mode.n_eff:.4f}")
                 print(f"Effective mode area: {mode.calculate_effective_area(field='y'):.4f}")
                 print(f"Mode transversality: {mode.transversality}")
diff --git a/docs/photonics/examples/selecting_modes.py b/docs/photonics/examples/selecting_modes.py
index dd57eb32..635722fd 100644
--- a/docs/photonics/examples/selecting_modes.py
+++ b/docs/photonics/examples/selecting_modes.py
@@ -124,7 +124,7 @@
 modes = compute_modes(basis0, epsilon, wavelength=wavelength, num_modes=4)
 
 for mode in modes:
-    mode.show(mode.E.real, direction="x")
+    mode.show("E", part="real")
 
 print(f"The effective index of the SiN mode is {np.real(modes[2].n_eff)}")
 
@@ -158,7 +158,7 @@
 modes = compute_modes(basis0, epsilon, wavelength=wavelength, num_modes=2)
 
 for mode in modes:
-    mode.show(mode.E.real, direction="x")
+    mode.show("E", part="real")
 
 print(f"The effective index of the SiN mode is {np.real(modes[0].n_eff)}")
 
@@ -180,7 +180,7 @@
 modes = compute_modes(basis0, epsilon, wavelength=wavelength, num_modes=2, n_guess=1.62)
 
 for mode in modes:
-    mode.show(mode.E.real, direction="x")
+    mode.show("E", part="real")
 
 print(f"The effective index of the SiN mode is {np.real(modes[1].n_eff)}")
 
diff --git a/docs/photonics/examples/waveguide_modes.py b/docs/photonics/examples/waveguide_modes.py
index 72a5062e..fd44d5e3 100644
--- a/docs/photonics/examples/waveguide_modes.py
+++ b/docs/photonics/examples/waveguide_modes.py
@@ -83,8 +83,8 @@
 modes = compute_modes(basis0, epsilon, wavelength=wavelength, num_modes=2, order=2)
 for mode in modes:
     print(f"Effective refractive index: {mode.n_eff:.4f}")
-    mode.show(mode.E.real, colorbar=True, direction="x")
-    mode.show(mode.E.imag, colorbar=True, direction="x")
+    mode.show("E", part="real", colorbar=True)
+    mode.show("E", part="imag", colorbar=True) 
 
 
 # %% [markdown]

From 78e3fa06610257885b83defba50c60927b7e3657 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Wed, 28 Feb 2024 20:35:20 -0800
Subject: [PATCH 03/23] refactor: black compliant

---
 docs/photonics/examples/waveguide_modes.py |  2 +-
 femwell/maxwell/waveguide.py               | 26 ++++++++++++----------
 femwell/pn_analytical.py                   |  4 +---
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/docs/photonics/examples/waveguide_modes.py b/docs/photonics/examples/waveguide_modes.py
index fd44d5e3..e44fe938 100644
--- a/docs/photonics/examples/waveguide_modes.py
+++ b/docs/photonics/examples/waveguide_modes.py
@@ -84,7 +84,7 @@
 for mode in modes:
     print(f"Effective refractive index: {mode.n_eff:.4f}")
     mode.show("E", part="real", colorbar=True)
-    mode.show("E", part="imag", colorbar=True) 
+    mode.show("E", part="imag", colorbar=True)
 
 
 # %% [markdown]
diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index d677e4b9..18e31f79 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -1,4 +1,5 @@
 """Waveguide analysis based on https://doi.org/10.1080/02726340290084012."""
+
 from dataclasses import dataclass
 from functools import cached_property
 from typing import List, Literal, Tuple
@@ -281,13 +282,14 @@ def plot(self, field, plot_vectors=False, colorbar=True, direction="y", title="E
         )
 
     def plot_component(
-            self,
-            field_name: Literal["E", "H"],
-            component: Literal["x", "y", "z"],
-            part: Literal["real", "imag", "abs"]="real",
-            boundaries: bool=True,
-            colorbar: bool=False,
-            ax: Axes=None):
+        self,
+        field_name: Literal["E", "H"],
+        component: Literal["x", "y", "z"],
+        part: Literal["real", "imag", "abs"] = "real",
+        boundaries: bool = True,
+        colorbar: bool = False,
+        ax: Axes = None,
+    ):
         from mpl_toolkits.axes_grid1 import make_axes_locatable
 
         if part == "real":
@@ -333,7 +335,7 @@ def plot_component(
             # self.basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
             # for subdomain in self.basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
             #     self.basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="w")
-            #plot_basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
+            # plot_basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
             for subdomain in plot_basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
                 plot_basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="k")
         if colorbar:
@@ -348,11 +350,11 @@ def plot_component(
     def show(
         self,
         field_name: Literal["E", "H"],
-        part: Literal["real", "imag", "abs"]="real",
-        boundaries: bool=True,
-        colorbar: bool=False,
+        part: Literal["real", "imag", "abs"] = "real",
+        boundaries: bool = True,
+        colorbar: bool = False,
     ):
-        fig, axs = plt.subplots(1,3, subplot_kw=dict(aspect=1))
+        fig, axs = plt.subplots(1, 3, subplot_kw=dict(aspect=1))
 
         for id_ax, comp in enumerate("xyz"):
             self.plot_component(field_name, comp, part, boundaries, colorbar, axs[id_ax])
diff --git a/femwell/pn_analytical.py b/femwell/pn_analytical.py
index 4be4381f..79d8b8bb 100644
--- a/femwell/pn_analytical.py
+++ b/femwell/pn_analytical.py
@@ -33,9 +33,7 @@ def dn_carriers(wavelength: float, dN: float, dP: float) -> float:
         return -2.98 * 1e-22 * np.power(dN, 1.016) - 1.25 * 1e-18 * np.power(dP, 0.835)
     else:
         wavelength *= 1e-6
-        return -3.64 * 1e-10 * wavelength**2 * dN - 3.51 * 1e-6 * wavelength**2 * np.power(
-            dP, 0.8
-        )
+        return -3.64 * 1e-10 * wavelength**2 * dN - 3.51 * 1e-6 * wavelength**2 * np.power(dP, 0.8)
 
 
 def dalpha_carriers(wavelength: float, dN: float, dP: float) -> float:

From 9ca4582532e4cf1376668df5496f7b700e5685d5 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Wed, 28 Feb 2024 21:22:16 -0800
Subject: [PATCH 04/23] fix: black issue with pn_analytical

---
 femwell/pn_analytical.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/femwell/pn_analytical.py b/femwell/pn_analytical.py
index 79d8b8bb..4be4381f 100644
--- a/femwell/pn_analytical.py
+++ b/femwell/pn_analytical.py
@@ -33,7 +33,9 @@ def dn_carriers(wavelength: float, dN: float, dP: float) -> float:
         return -2.98 * 1e-22 * np.power(dN, 1.016) - 1.25 * 1e-18 * np.power(dP, 0.835)
     else:
         wavelength *= 1e-6
-        return -3.64 * 1e-10 * wavelength**2 * dN - 3.51 * 1e-6 * wavelength**2 * np.power(dP, 0.8)
+        return -3.64 * 1e-10 * wavelength**2 * dN - 3.51 * 1e-6 * wavelength**2 * np.power(
+            dP, 0.8
+        )
 
 
 def dalpha_carriers(wavelength: float, dN: float, dP: float) -> float:

From cdb32f582bc73e55e4f8100230224b19ef49747d Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 20:36:08 -0800
Subject: [PATCH 05/23] naming: use field instead of field_name in waveguide.py

---
 femwell/maxwell/waveguide.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 18e31f79..6e8f2b4e 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -283,7 +283,7 @@ def plot(self, field, plot_vectors=False, colorbar=True, direction="y", title="E
 
     def plot_component(
         self,
-        field_name: Literal["E", "H"],
+        field: Literal["E", "H"],
         component: Literal["x", "y", "z"],
         part: Literal["real", "imag", "abs"] = "real",
         boundaries: bool = True,
@@ -306,12 +306,12 @@ def plot_component(
         else:
             fig = plt.gcf()
 
-        if field_name == "E":
+        if field == "E":
             mfield = self.E
-        elif field_name == "H":
+        elif field == "H":
             mfield = self.H
         else:
-            raise ValueError("A valid field_name is 'E' or 'H'.")
+            raise ValueError("A valid field is 'E' or 'H'.")
 
         (mfield_t, mfield_t_basis), (mfield_n, mfield_n_basis) = self.basis.split(mfield)
 
@@ -343,7 +343,7 @@ def plot_component(
             cax = divider.append_axes("right", size="5%", pad=0.05)
             plt.colorbar(ax.collections[-1], cax=cax)
 
-        ax.set_title(f"{field_name}{component} ({part}. part)")
+        ax.set_title(f"{field}{component} ({part}. part)")
 
         return fig, ax
 

From d4ae6c08a68093c0974cf10829ba3f0d70bcde00 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 20:54:39 -0800
Subject: [PATCH 06/23] clean: remove unused code

---
 femwell/maxwell/waveguide.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 6e8f2b4e..baab6a3f 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -332,10 +332,6 @@ def plot_component(
             raise ValueError("A valid component is 'x', 'y' or 'z'.")
 
         if boundaries:
-            # self.basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
-            # for subdomain in self.basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
-            #     self.basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="w")
-            # plot_basis.mesh.draw(ax=ax, boundaries=True, boundaries_only=True)
             for subdomain in plot_basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
                 plot_basis.mesh.restrict(subdomain).draw(ax=ax, boundaries_only=True, color="k")
         if colorbar:

From bd592abe93e69dac27a64d55e80fa650249d9e11 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 20:55:51 -0800
Subject: [PATCH 07/23] feat: mode.show() is now backward compatible!

---
 femwell/maxwell/waveguide.py | 68 +++++++++++++++++++++++++++++++-----
 1 file changed, 59 insertions(+), 9 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index baab6a3f..1e8b2b24 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from functools import cached_property
 from typing import List, Literal, Tuple
+from warnings import warn
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -343,19 +344,68 @@ def plot_component(
 
         return fig, ax
 
+    # def show(
+    #     self,
+    #     field: Literal["E", "H"] ,
+    #     part: Literal["real", "imag", "abs"] = "real",
+    #     boundaries: bool = True,
+    #     colorbar: bool = False,
+    # ):
+    #     fig, axs = plt.subplots(1, 3, subplot_kw=dict(aspect=1))
+
+    #     for id_ax, comp in enumerate("xyz"):
+    #         self.plot_component(field, comp, part, boundaries, colorbar, axs[id_ax])
+    #     plt.tight_layout()
+    #     plt.show()
+
     def show(
         self,
-        field_name: Literal["E", "H"],
-        part: Literal["real", "imag", "abs"] = "real",
-        boundaries: bool = True,
-        colorbar: bool = False,
+        field: Literal["E", "H"] | NDArray,
+        **kwargs
     ):
-        fig, axs = plt.subplots(1, 3, subplot_kw=dict(aspect=1))
+        if type(field) is np.ndarray:
+            warn("The behavior of passing an array directly to `show` "
+                 + "is deprecated and will be removed in the future. "
+                 + "Use `plot` instead.", DeprecationWarning, stacklevel=2)
+            self.plot(field=field, **kwargs)
+            plt.show()
+        else:
+            from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+            part = kwargs.get("part", "real")
+            boundaries = kwargs.get("boundaries", True)
+            colorbar = kwargs.get("colorbar", False)
+            direction = kwargs.get("direction", "y")
+            plot_vectors = kwargs.get("plot_vectors", False)
+            title = kwargs.get("title", "E")
+
+            if plot_vectors is True:
+                if field == "E":
+                    mfield = self.E
+                elif field == "H":
+                    mfield = self.H
+
+                if part == "real":
+                    conv_func = np.real
+                elif part == "imag":
+                    conv_func = np.imag
+                elif part == "abs":
+                    conv_func = np.abs
+                else:
+                    raise ValueError("A valid part is 'real', 'imag' or 'abs'.")
+
+                plot_mode(
+                    self.basis, conv_func(mfield), plot_vectors=True,
+                    colorbar=colorbar, direction=direction, title=title
+                )
+            else:
+                rc = (3, 1) if direction != "x" else (1, 3)
+                fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))
 
-        for id_ax, comp in enumerate("xyz"):
-            self.plot_component(field_name, comp, part, boundaries, colorbar, axs[id_ax])
-        plt.tight_layout()
-        plt.show()
+                for id_ax, comp in enumerate("xyz"):
+                    self.plot_component(field, comp, part, boundaries, colorbar, axs[id_ax])
+            plt.tight_layout()
+            plt.show()
 
     def plot_intensity(
         self,

From f77b35692acf9fb90cf3ac12c0f666c5851e087c Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 20:56:56 -0800
Subject: [PATCH 08/23] doc: update coplanar_waveguide plot func

---
 femwell/examples/coplanar_waveguide.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/femwell/examples/coplanar_waveguide.py b/femwell/examples/coplanar_waveguide.py
index 47dae27b..a481677e 100644
--- a/femwell/examples/coplanar_waveguide.py
+++ b/femwell/examples/coplanar_waveguide.py
@@ -127,7 +127,7 @@ def mesh_coax(filename, radius_inner, radius_outer):
     )
     print("propagation constants", 1 / modes.n_effs)
 
-    modes[0].show(modes[0].E.real, plot_vectors=True)
+    modes[0].show("E", part="real", plot_vectors=True)
 
     from skfem import *
     from skfem.helpers import *
@@ -139,7 +139,7 @@ def current_form(w):
     currents = np.zeros((len(conductors), len(modes)))
 
     for mode_i, mode in enumerate(modes):
-        mode.show(mode.H.real, plot_vectors=True)
+        modes[0].show("H", part="real", plot_vectors=True)
 
         (ht, ht_basis), (hz, hz_basis) = mode.basis.split(mode.H)
         for conductors_i, conductor in enumerate(conductors):

From 8d26699c35094a7ef17fc9281f0c77ca323a5f29 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 21:07:41 -0800
Subject: [PATCH 09/23] doc: add missing colorbar for coplanar_waveguide

---
 femwell/examples/coplanar_waveguide.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/femwell/examples/coplanar_waveguide.py b/femwell/examples/coplanar_waveguide.py
index a481677e..a1fb39cf 100644
--- a/femwell/examples/coplanar_waveguide.py
+++ b/femwell/examples/coplanar_waveguide.py
@@ -127,7 +127,7 @@ def mesh_coax(filename, radius_inner, radius_outer):
     )
     print("propagation constants", 1 / modes.n_effs)
 
-    modes[0].show("E", part="real", plot_vectors=True)
+    modes[0].show("E", part="real", plot_vectors=True, colorbar=True)
 
     from skfem import *
     from skfem.helpers import *
@@ -139,7 +139,7 @@ def current_form(w):
     currents = np.zeros((len(conductors), len(modes)))
 
     for mode_i, mode in enumerate(modes):
-        modes[0].show("H", part="real", plot_vectors=True)
+        modes[0].show("H", part="real", plot_vectors=True, colorbar=True)
 
         (ht, ht_basis), (hz, hz_basis) = mode.basis.split(mode.H)
         for conductors_i, conductor in enumerate(conductors):

From 17e12d1dcf7e88a576af679c154d91428c0b185d Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 21:09:04 -0800
Subject: [PATCH 10/23] feat: waveguide.show() allow to plot normal and tangent
 fields

---
 femwell/maxwell/waveguide.py | 33 +++++++++++----------------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 1e8b2b24..a3e11f4f 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -285,7 +285,7 @@ def plot(self, field, plot_vectors=False, colorbar=True, direction="y", title="E
     def plot_component(
         self,
         field: Literal["E", "H"],
-        component: Literal["x", "y", "z"],
+        component: Literal["x", "y", "z", "n", "t"],
         part: Literal["real", "imag", "abs"] = "real",
         boundaries: bool = True,
         colorbar: bool = False,
@@ -324,13 +324,15 @@ def plot_component(
                 mfield_x_basis.plot(mfield_x, ax=ax, shading="gouraud")
             else:
                 mfield_y_basis.plot(mfield_y, ax=ax, shading="gouraud")
-        elif component == "z":
-            # plot_basis = mfield_n_basis.with_element(ElementVector(ElementDG(ElementTriP1())))
+        elif component == "t":
+            plot_basis = mfield_t_basis
+            mfield_t_basis.plot(conv_func(mfield_t), ax=ax, shading="gouraud")
+        elif component == "z" or component == "n":
             plot_basis = mfield_n_basis
             mfield_z, mfield_z_basis = mfield_n, mfield_n_basis
             mfield_z_basis.plot(conv_func(mfield_z), ax=ax, shading="gouraud")
         else:
-            raise ValueError("A valid component is 'x', 'y' or 'z'.")
+            raise ValueError("A valid component is 'x', 'y', 'z', 'n' or 't'.")
 
         if boundaries:
             for subdomain in plot_basis.mesh.subdomains.keys() - {"gmsh:bounding_entities"}:
@@ -380,24 +382,11 @@ def show(
             title = kwargs.get("title", "E")
 
             if plot_vectors is True:
-                if field == "E":
-                    mfield = self.E
-                elif field == "H":
-                    mfield = self.H
-
-                if part == "real":
-                    conv_func = np.real
-                elif part == "imag":
-                    conv_func = np.imag
-                elif part == "abs":
-                    conv_func = np.abs
-                else:
-                    raise ValueError("A valid part is 'real', 'imag' or 'abs'.")
-
-                plot_mode(
-                    self.basis, conv_func(mfield), plot_vectors=True,
-                    colorbar=colorbar, direction=direction, title=title
-                )
+                rc = (2, 1) if direction != "x" else (1, 3)
+                fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))
+
+                self.plot_component(field, "t", part, boundaries, colorbar, axs[0])
+                self.plot_component(field, "n", part, boundaries, colorbar, axs[1])
             else:
                 rc = (3, 1) if direction != "x" else (1, 3)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))

From b551588d2a90dc30cbc49f962789838b273eefa3 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 21:13:09 -0800
Subject: [PATCH 11/23] refactor: black compliant

---
 femwell/maxwell/waveguide.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index a3e11f4f..517a72cc 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -366,9 +366,13 @@ def show(
         **kwargs
     ):
         if type(field) is np.ndarray:
-            warn("The behavior of passing an array directly to `show` "
-                 + "is deprecated and will be removed in the future. "
-                 + "Use `plot` instead.", DeprecationWarning, stacklevel=2)
+            warn(
+                "The behavior of passing an array directly to `show` "
+                + "is deprecated and will be removed in the future. "
+                + "Use `plot` instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             self.plot(field=field, **kwargs)
             plt.show()
         else:

From 6c0bd6cd6ba2e04b29a59e07803f3ea96e2ff456 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 21:16:27 -0800
Subject: [PATCH 12/23] refactor: black compliant2

---
 femwell/maxwell/waveguide.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 517a72cc..338aa4fd 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -360,11 +360,7 @@ def plot_component(
     #     plt.tight_layout()
     #     plt.show()
 
-    def show(
-        self,
-        field: Literal["E", "H"] | NDArray,
-        **kwargs
-    ):
++    def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
         if type(field) is np.ndarray:
             warn(
                 "The behavior of passing an array directly to `show` "

From 37efdda3b110211fed7174e3186f4185f614ab61 Mon Sep 17 00:00:00 2001
From: Dj1312 <18169686+Dj1312@users.noreply.github.com>
Date: Sun, 3 Mar 2024 21:23:08 -0800
Subject: [PATCH 13/23] fix: missing space after refactor

---
 femwell/maxwell/waveguide.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 338aa4fd..708b0a3e 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -360,7 +360,7 @@ def plot_component(
     #     plt.tight_layout()
     #     plt.show()
 
-+    def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
+    def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
         if type(field) is np.ndarray:
             warn(
                 "The behavior of passing an array directly to `show` "

From fc3868079fca749516869358a5e7bf440241d414 Mon Sep 17 00:00:00 2001
From: Lucas Grosjean <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 12:55:26 -0800
Subject: [PATCH 14/23] fix: wrong dimensions in subplots

Co-authored-by: Helge Gehring <42973196+HelgeGehring@users.noreply.github.com>
---
 femwell/maxwell/waveguide.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 708b0a3e..f0862486 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -382,7 +382,7 @@ def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
             title = kwargs.get("title", "E")
 
             if plot_vectors is True:
-                rc = (2, 1) if direction != "x" else (1, 3)
+                rc = (2, 1) if direction != "x" else (1, 2)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))
 
                 self.plot_component(field, "t", part, boundaries, colorbar, axs[0])

From d3932b16bb5e0a1a23c377e1a2d80e8810a11bf9 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 18:25:08 -0800
Subject: [PATCH 15/23] refactor: discard use of **kwargs in Mode.show()

---
 femwell/maxwell/waveguide.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index f0862486..d3a5c9f3 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -360,7 +360,16 @@ def plot_component(
     #     plt.tight_layout()
     #     plt.show()
 
-    def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
+    def show(
+        self,
+        field: Literal["E", "H"] | NDArray,
+        part: Literal["real", "imag", "abs"] = "real",
+        plot_vectors: bool = False,
+        boundaries: bool = True,
+        colorbar: bool = False,
+        direction: Literal["x", "y"] = "x",
+        title: str = "E",
+    ):
         if type(field) is np.ndarray:
             warn(
                 "The behavior of passing an array directly to `show` "
@@ -369,18 +378,11 @@ def show(self, field: Literal["E", "H"] | NDArray, **kwargs):
                 DeprecationWarning,
                 stacklevel=2,
             )
-            self.plot(field=field, **kwargs)
+            self.plot(field, plot_vectors, colorbar, direction, title)
             plt.show()
         else:
             from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-            part = kwargs.get("part", "real")
-            boundaries = kwargs.get("boundaries", True)
-            colorbar = kwargs.get("colorbar", False)
-            direction = kwargs.get("direction", "y")
-            plot_vectors = kwargs.get("plot_vectors", False)
-            title = kwargs.get("title", "E")
-
             if plot_vectors is True:
                 rc = (2, 1) if direction != "x" else (1, 2)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))

From 98783406d963b8ddc63a9f34d35884ccab7bbd6c Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 18:33:21 -0800
Subject: [PATCH 16/23] fix: remove plt.gcf() from Mode.plot_component()

---
 femwell/maxwell/waveguide.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index d3a5c9f3..c5fc12f8 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -303,9 +303,8 @@ def plot_component(
             raise ValueError("A valid part is 'real', 'imag' or 'abs'.")
 
         if ax is None:
-            fig, ax = plt.subplots()
-        else:
-            fig = plt.gcf()
+            fig = plt.figure()
+            ax = fig.add_subplot(111)
 
         if field == "E":
             mfield = self.E
@@ -344,7 +343,7 @@ def plot_component(
 
         ax.set_title(f"{field}{component} ({part}. part)")
 
-        return fig, ax
+        return ax
 
     # def show(
     #     self,

From c941eceb9131eb30e404f368b2ae1453818f978c Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 18:34:54 -0800
Subject: [PATCH 17/23] feat: allow Mode.show() to plot the intensity

---
 femwell/maxwell/waveguide.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index c5fc12f8..83fe133c 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -361,7 +361,7 @@ def plot_component(
 
     def show(
         self,
-        field: Literal["E", "H"] | NDArray,
+        field: Literal["E", "H", "I"] | NDArray,
         part: Literal["real", "imag", "abs"] = "real",
         plot_vectors: bool = False,
         boundaries: bool = True,
@@ -383,11 +383,18 @@ def show(
             from mpl_toolkits.axes_grid1 import make_axes_locatable
 
             if plot_vectors is True:
+                if field == "I":
+                    return ValueError(
+                        "'plot_vectors' is used to plot the tangential components "
+                        + "of a field. Thus it can be used only with 'E' or 'H'."
+                    )
                 rc = (2, 1) if direction != "x" else (1, 2)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))
 
                 self.plot_component(field, "t", part, boundaries, colorbar, axs[0])
                 self.plot_component(field, "n", part, boundaries, colorbar, axs[1])
+            elif field == "I":
+                fig, ax = self.plot_intensity(ax=None, colorbar=colorbar)
             else:
                 rc = (3, 1) if direction != "x" else (1, 3)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))

From 3523bf40f64382fab218775d27985790356770a2 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 18:43:09 -0800
Subject: [PATCH 18/23] fix: add missing title in Mode.show("I")

---
 femwell/maxwell/waveguide.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 83fe133c..568ac6ea 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -395,6 +395,7 @@ def show(
                 self.plot_component(field, "n", part, boundaries, colorbar, axs[1])
             elif field == "I":
                 fig, ax = self.plot_intensity(ax=None, colorbar=colorbar)
+                ax.set_title(f"Intensity")
             else:
                 rc = (3, 1) if direction != "x" else (1, 3)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))

From c8037af1e7b686e1679b55f03d6257e0bc14f16f Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Mon, 4 Mar 2024 18:44:45 -0800
Subject: [PATCH 19/23] doc: update ex. plot_intensity() to .show("I")

---
 docs/photonics/examples/waveguide_modes.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/docs/photonics/examples/waveguide_modes.py b/docs/photonics/examples/waveguide_modes.py
index e44fe938..49c0249c 100644
--- a/docs/photonics/examples/waveguide_modes.py
+++ b/docs/photonics/examples/waveguide_modes.py
@@ -92,11 +92,7 @@
 # +
 
 # %%
-fig, ax = plt.subplots()
-modes[0].plot_intensity(ax=ax)
-plt.title("Normalized Intensity")
-plt.tight_layout()
-plt.show()
+modes[0].show("I", colorbar=True)
 
 # %% [markdown]
 # -

From c0f62d6857bb2d9ef98af1598e0aef0ef9051e89 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Tue, 5 Mar 2024 01:12:15 -0800
Subject: [PATCH 20/23] feat: improve show() and plot_component()

---
 femwell/maxwell/waveguide.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 568ac6ea..5ee748c7 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import cached_property
-from typing import List, Literal, Tuple
+from typing import Callable, List, Literal, Optional, Tuple
 from warnings import warn
 
 import matplotlib.pyplot as plt
@@ -286,7 +286,7 @@ def plot_component(
         self,
         field: Literal["E", "H"],
         component: Literal["x", "y", "z", "n", "t"],
-        part: Literal["real", "imag", "abs"] = "real",
+        part: Literal["real", "imag", "abs"] | Callable = "real",
         boundaries: bool = True,
         colorbar: bool = False,
         ax: Axes = None,
@@ -299,6 +299,8 @@ def plot_component(
             conv_func = np.imag
         elif part == "abs":
             conv_func = np.abs
+        elif isinstance(part, Callable):
+            conv_func = part
         else:
             raise ValueError("A valid part is 'real', 'imag' or 'abs'.")
 
@@ -362,12 +364,12 @@ def plot_component(
     def show(
         self,
         field: Literal["E", "H", "I"] | NDArray,
-        part: Literal["real", "imag", "abs"] = "real",
+        part: Literal["real", "imag", "abs"] | Callable = "real",
         plot_vectors: bool = False,
         boundaries: bool = True,
         colorbar: bool = False,
         direction: Literal["x", "y"] = "x",
-        title: str = "E",
+        title: Optional[str] = None,
     ):
         if type(field) is np.ndarray:
             warn(

From a691a80709b5b04992c59e3620c09e7696630be2 Mon Sep 17 00:00:00 2001
From: lucasgrjn <18169686+lucasgrjn@users.noreply.github.com>
Date: Wed, 6 Mar 2024 17:24:59 -0800
Subject: [PATCH 21/23] doc: add docstring for Mode.show()

---
 femwell/maxwell/waveguide.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 5ee748c7..b6f93160 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -343,7 +343,7 @@ def plot_component(
             cax = divider.append_axes("right", size="5%", pad=0.05)
             plt.colorbar(ax.collections[-1], cax=cax)
 
-        ax.set_title(f"{field}{component} ({part}. part)")
+        ax.set_title(f"{field}{component} ({conv_func.__name__}. part)")
 
         return ax
 
@@ -371,6 +371,18 @@ def show(
         direction: Literal["x", "y"] = "x",
         title: Optional[str] = None,
     ):
+        """Plots the different quantities associated with a field.
+
+        Args:
+            field ("E", "H", "I"): Field of interest, can be the electric field, the magnetic field or the intensity of the mode.
+            part ("real", "imag", "abs", callable): Function to use to preprocess the field to be plotted. Defaults to "real".
+            plot_vectors (bool): If set to True, plot the normal and tangential component
+            boundaries (bool): Superimpose the mesh boundaries on the plot. Defaults to True.
+            colorbar (bool): Adds a colorbar to the plot. Defaults to False.
+            direction ("x", "y"): Orientation of the plots ("x" for horizontal and "y" for vertical) Defaults to "x".
+        Returns:
+            Tuple[Figure, Axes]: Figure and axes of the plot.
+        """
         if type(field) is np.ndarray:
             warn(
                 "The behavior of passing an array directly to `show` "

From 5496c79048f4b383a1ab3e437a5215b34bc0f2ce Mon Sep 17 00:00:00 2001
From: Helge Gehring <42973196+HelgeGehring@users.noreply.github.com>
Date: Wed, 6 Mar 2024 17:53:23 -0800
Subject: [PATCH 22/23] Add suptitle to show

---
 femwell/maxwell/waveguide.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index b6f93160..83774733 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -416,6 +416,7 @@ def show(
 
                 for id_ax, comp in enumerate("xyz"):
                     self.plot_component(field, comp, part, boundaries, colorbar, axs[id_ax])
+            fig.suptitle(title if title else field)
             plt.tight_layout()
             plt.show()
 

From a573be70b2b7e05d91c3819abe92af3859dc6c4b Mon Sep 17 00:00:00 2001
From: Helge Gehring <42973196+HelgeGehring@users.noreply.github.com>
Date: Wed, 6 Mar 2024 17:53:35 -0800
Subject: [PATCH 23/23] remove title for Intensity

---
 femwell/maxwell/waveguide.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/femwell/maxwell/waveguide.py b/femwell/maxwell/waveguide.py
index 83774733..0b4b124b 100644
--- a/femwell/maxwell/waveguide.py
+++ b/femwell/maxwell/waveguide.py
@@ -409,7 +409,6 @@ def show(
                 self.plot_component(field, "n", part, boundaries, colorbar, axs[1])
             elif field == "I":
                 fig, ax = self.plot_intensity(ax=None, colorbar=colorbar)
-                ax.set_title(f"Intensity")
             else:
                 rc = (3, 1) if direction != "x" else (1, 3)
                 fig, axs = plt.subplots(*rc, subplot_kw=dict(aspect=1))