Added: mode plane plotting functions (#1397).

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Support for differentiation with respect to `GeometryGroup.geometries` elements.
 - Users can now export `SimulationData` to MATLAB `.mat` files with the `to_mat_file` method.
+- `ModeSolver` methods to plot the mode plane simulation components, including `.plot()`, `.plot_eps()`, `.plot_structures_eps()`, `.plot_grid()`, and `.plot_pml()`.
 
 ### Changed
 

tests/test_plugins/test_mode_solver.py

@@ -939,3 +939,40 @@ def test_mode_solver_relative():
     new_freqs = np.array(freqs) * 1.01
     ms = ms.updated_copy(freqs=new_freqs)
     _ = ms._data_on_yee_grid_relative(basis=basis)
+
+
+def test_mode_solver_plot():
+    """Test mode plane plotting functions"""
+
+    simulation = td.Simulation(
+        size=SIM_SIZE,
+        grid_spec=td.GridSpec(wavelength=1.0),
+        structures=[WAVEGUIDE],
+        run_time=1e-12,
+        symmetry=(0, 0, 1),
+        boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()),
+        sources=[SRC],
+    )
+    mode_spec = td.ModeSpec(
+        num_modes=3,
+        target_neff=2.0,
+        num_pml=[8, 4],
+    )
+    freqs = [td.C_0 / 0.9, td.C_0 / 1.0, td.C_0 / 1.1]
+    ms = ModeSolver(
+        simulation=simulation,
+        plane=PLANE,
+        mode_spec=mode_spec,
+        freqs=freqs,
+        direction="-",
+        colocate=False,
+    )
+    _, ax = plt.subplots(2, 2, figsize=(12, 8), tight_layout=True)
+    ms.plot(ax=ax[0, 0])
+    ms.plot_eps(freq=200e14, alpha=0.7, ax=ax[0, 1])
+    ms.plot_structures_eps(freq=200e14, alpha=0.8, cbar=True, reverse=False, ax=ax[1, 0])
+    ms.plot_grid(linewidth=0.3, ax=ax[1, 0])
+    ms.plot(ax=ax[1, 1])
+    ms.plot_pml(ax=ax[1, 1])
+    ms.plot_grid(linewidth=0.3, ax=ax[1, 1])
+    plt.close()

tidy3d/plugins/mode/mode_solver.py

@@ -10,6 +10,8 @@
 import numpy as np
 import pydantic.v1 as pydantic
 import xarray as xr
+from matplotlib.collections import PatchCollection
+from matplotlib.patches import Rectangle
 
 from ...components.base import Tidy3dBaseModel, cached_property, skip_if_fields_missing
 from ...components.boundary import PML, Absorber, Boundary, BoundarySpec, PECBoundary, StablePML
@@ -41,6 +43,7 @@
     Symmetry,
 )
 from ...components.validators import validate_freqs_min, validate_freqs_not_empty
+from ...components.viz import plot_params_pml
 from ...constants import C_0
 from ...exceptions import SetupError, ValidationError
 from ...log import log
@@ -1141,6 +1144,288 @@ def plot_field(
             **sel_kwargs,
         )
 
+    def plot(
+        self,
+        ax: Ax = None,
+        **patch_kwargs,
+    ) -> Ax:
+        """Plot the mode plane simulation's components.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+
+        See Also
+        ---------
+
+        **Notebooks**
+            * `Visualizing geometries in Tidy3D: Plotting Materials <../../notebooks/VizSimulation.html#Plotting-Materials>`_
+
+        """
+        # Get the mode plane normal axis, center, and limits.
+        a_center, h_lim, v_lim, _ = self._center_and_lims()
+
+        return self.simulation.plot(
+            x=a_center[0],
+            y=a_center[1],
+            z=a_center[2],
+            hlim=h_lim,
+            vlim=v_lim,
+            source_alpha=0,
+            monitor_alpha=0,
+            lumped_element_alpha=0,
+            ax=ax,
+            **patch_kwargs,
+        )
+
+    def plot_eps(
+        self,
+        freq: float = None,
+        alpha: float = None,
+        ax: Ax = None,
+    ) -> Ax:
+        """Plot the mode plane simulation's components.
+        The permittivity is plotted in grayscale based on its value at the specified frequency.
+
+        Parameters
+        ----------
+        freq : float = None
+            Frequency to evaluate the relative permittivity of all mediums.
+            If not specified, evaluates at infinite frequency.
+        alpha : float = None
+            Opacity of the structures being plotted.
+            Defaults to the structure default alpha.
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+
+        See Also
+        ---------
+
+        **Notebooks**
+            * `Visualizing geometries in Tidy3D: Plotting Permittivity <../../notebooks/VizSimulation.html#Plotting-Permittivity>`_
+        """
+
+        # Get the mode plane normal axis, center, and limits.
+        a_center, h_lim, v_lim, _ = self._center_and_lims()
+
+        # Plot at central mode frequency if freq is not provided.
+        f = freq if freq is not None else self.freqs[len(self.freqs) // 2]
+
+        return self.simulation.plot_eps(
+            x=a_center[0],
+            y=a_center[1],
+            z=a_center[2],
+            freq=f,
+            alpha=alpha,
+            hlim=h_lim,
+            vlim=v_lim,
+            source_alpha=0,
+            monitor_alpha=0,
+            lumped_element_alpha=0,
+            ax=ax,
+        )
+
+    def plot_structures_eps(
+        self,
+        freq: float = None,
+        alpha: float = None,
+        cbar: bool = True,
+        reverse: bool = False,
+        ax: Ax = None,
+    ) -> Ax:
+        """Plot the mode plane simulation's components.
+        The permittivity is plotted in grayscale based on its value at the specified frequency.
+
+        Parameters
+        ----------
+        freq : float = None
+            Frequency to evaluate the relative permittivity of all mediums.
+            If not specified, evaluates at infinite frequency.
+        alpha : float = None
+            Opacity of the structures being plotted.
+            Defaults to the structure default alpha.
+        cbar : bool = True
+            Whether to plot a colorbar for the relative permittivity.
+        reverse : bool = False
+            If ``False``, the highest permittivity is plotted in black.
+            If ``True``, it is plotteed in white (suitable for black backgrounds).
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+
+        See Also
+        ---------
+
+        **Notebooks**
+            * `Visualizing geometries in Tidy3D: Plotting Permittivity <../../notebooks/VizSimulation.html#Plotting-Permittivity>`_
+        """
+
+        # Get the mode plane normal axis, center, and limits.
+        a_center, h_lim, v_lim, _ = self._center_and_lims()
+
+        # Plot at central mode frequency if freq is not provided.
+        f = freq if freq is not None else self.freqs[len(self.freqs) // 2]
+
+        return self.simulation.plot_structures_eps(
+            x=a_center[0],
+            y=a_center[1],
+            z=a_center[2],
+            freq=f,
+            alpha=alpha,
+            cbar=cbar,
+            reverse=reverse,
+            hlim=h_lim,
+            vlim=v_lim,
+            ax=ax,
+        )
+
+    def plot_grid(
+        self,
+        ax: Ax = None,
+        **kwargs,
+    ) -> Ax:
+        """Plot the mode plane cell boundaries as lines.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+        **kwargs
+            Optional keyword arguments passed to the matplotlib ``LineCollection``.
+            For details on accepted values, refer to
+            `Matplotlib's documentation <https://tinyurl.com/2p97z4cn>`_.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+        """
+
+        # Get the mode plane normal axis, center, and limits.
+        a_center, h_lim, v_lim, _ = self._center_and_lims()
+
+        return self.simulation.plot_grid(
+            x=a_center[0], y=a_center[1], z=a_center[2], hlim=h_lim, vlim=v_lim, ax=ax, **kwargs
+        )
+
+    def plot_pml(
+        self,
+        ax: Ax = None,
+    ) -> Ax:
+        """Plot the mode plane absorbing boundaries.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+        """
+
+        # Get the mode plane normal axis, center, and limits.
+        a_center, h_lim, v_lim, t_axes = self._center_and_lims()
+
+        # Plot the mode plane is ax=None.
+        if not ax:
+            ax = self.simulation.plot(
+                x=a_center[0],
+                y=a_center[1],
+                z=a_center[2],
+                hlim=h_lim,
+                vlim=v_lim,
+                source_alpha=0,
+                monitor_alpha=0,
+                ax=ax,
+            )
+
+        # Mode plane grid.
+        plane_grid = self.grid_snapped.centers.to_list
+        coord_0 = plane_grid[t_axes[0]][1:-1]
+        coord_1 = plane_grid[t_axes[1]][1:-1]
+
+        # Number of PML layers in ModeSpec.
+        num_pml_0 = self.mode_spec.num_pml[0]
+        num_pml_1 = self.mode_spec.num_pml[1]
+
+        # Calculate PML thickness.
+        pml_thick_0_plus = 0
+        pml_thick_0_minus = 0
+        if num_pml_0 > 0:
+            pml_thick_0_plus = coord_0[-1] - coord_0[-num_pml_0 - 1]
+            pml_thick_0_minus = coord_0[num_pml_0] - coord_0[0]
+
+        pml_thick_1_plus = 0
+        pml_thick_1_minus = 0
+        if num_pml_1 > 0:
+            pml_thick_1_plus = coord_1[-1] - coord_1[-num_pml_1 - 1]
+            pml_thick_1_minus = coord_1[num_pml_1] - coord_1[0]
+
+        # Mode Plane width and height
+        mp_w = coord_0[-1] - coord_0[0]
+        mp_h = coord_1[-1] - coord_1[0]
+
+        # Plot the absorbing layers.
+        if num_pml_0 > 0 or num_pml_1 > 0:
+            pml_rect = []
+            if pml_thick_0_minus > 0:
+                pml_rect.append(Rectangle((coord_0[0], coord_1[0]), pml_thick_0_minus, mp_h))
+            if pml_thick_0_plus > 0:
+                pml_rect.append(
+                    Rectangle((coord_0[-num_pml_0 - 1], coord_1[0]), pml_thick_0_plus, mp_h)
+                )
+            if pml_thick_1_minus > 0:
+                pml_rect.append(Rectangle((coord_0[0], coord_1[0]), mp_w, pml_thick_1_minus))
+            if pml_thick_1_plus > 0:
+                pml_rect.append(
+                    Rectangle((coord_0[0], coord_1[-num_pml_1 - 1]), mp_w, pml_thick_1_plus)
+                )
+
+            pc = PatchCollection(
+                pml_rect,
+                alpha=plot_params_pml.alpha,
+                facecolor=plot_params_pml.facecolor,
+                edgecolor=plot_params_pml.edgecolor,
+                hatch=plot_params_pml.hatch,
+                zorder=plot_params_pml.zorder,
+            )
+            ax.add_collection(pc)
+
+        return ax
+
+    def _center_and_lims(self) -> Tuple[List, List, List, List]:
+        """Get the mode plane center and limits."""
+
+        n_axis, t_axes = self.plane.pop_axis([0, 1, 2], self.normal_axis)
+        a_center = [None, None, None]
+        a_center[n_axis] = self.plane.center[n_axis]
+        h_lim = [
+            a_center[n_axis] - self.plane.size[t_axes[0]] / 2,
+            a_center[n_axis] + self.plane.size[t_axes[0]] / 2,
+        ]
+        v_lim = [
+            a_center[n_axis] - self.plane.size[t_axes[1]] / 2,
+            a_center[n_axis] + self.plane.size[t_axes[1]] / 2,
+        ]
+        return a_center, h_lim, v_lim, t_axes
+
     def _validate_modes_size(self):
         """Make sure that the total size of the modes fields is not too large."""
         monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME)