Adding ModeSimulation.plot and more arguments to mode plotting

Author: momchil-flex

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `EMEFieldMonitor` now supports `interval_space`.
 - `Simulation.precision` option allows to select `"double"` precision for very high-accuracy results. Note that this is very rarely needed, and doubles the simulation computational weight and correpsondingly FlexCredit cost.
 - Added material type `PMCMedium` for perfect magnetic conductor.
+- `ModeSimulation.plot()` method that plots the mode simulation plane by default, or the containing FDTD simulation if any of ``x``, ``y``, or ``z`` is passed. 
 
 ### Changed
 - Switched to an analytical gradient calculation for spatially-varying pole-residue models (`CustomPoleResidue`).

tests/test_components/test_mode.py

@@ -176,6 +176,8 @@ def get_mode_sim():
 def test_mode_sim():
     with AssertLogLevel(None):
         sim = get_mode_sim()
+        _ = sim.plot(ax=AX)
+        _ = sim.plot(ax=AX, fill_structures=False, hlim=(-1, 1), vlim=(-1, 1))
         _ = sim.plot(y=0, ax=AX)
         _ = sim.plot_mode_plane(ax=AX)
         _ = sim.plot_eps_mode_plane(ax=AX)

tidy3d/components/mode/mode_solver.py

@@ -2117,6 +2117,9 @@ def plot_field(
     def plot(
         self,
         ax: Ax = None,
+        hlim: Optional[tuple[float, float]] = None,
+        vlim: Optional[tuple[float, float]] = None,
+        fill_structures: bool = True,
         **patch_kwargs,
     ) -> Ax:
         """Plot the mode plane simulation's components.
@@ -2125,6 +2128,12 @@ def plot(
         ----------
         ax : matplotlib.axes._subplots.Axes = None
             Matplotlib axes to plot on, if not specified, one is created.
+        hlim : Tuple[float, float] = None
+            The x range if plotting on xy or xz planes, y range if plotting on yz plane.
+        vlim : Tuple[float, float] = None
+            The z range if plotting on xz or yz planes, y plane if plotting on xy plane.
+        fill_structures : bool = True
+            Whether to fill structures with color or just draw outlines.
 
         Returns
         -------
@@ -2139,20 +2148,26 @@ def plot(
 
         """
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, _ = self._center_and_lims(
+        a_center, hlim_plane, vlim_plane, _ = self._center_and_lims(
             simulation=self.simulation, plane=self.plane
         )
 
+        if hlim is None:
+            hlim = hlim_plane
+        if vlim is None:
+            vlim = vlim_plane
+
         ax = self.simulation.plot(
             x=a_center[0],
             y=a_center[1],
             z=a_center[2],
-            hlim=h_lim,
-            vlim=v_lim,
+            hlim=hlim,
+            vlim=vlim,
             source_alpha=0,
             monitor_alpha=0,
             lumped_element_alpha=0,
             ax=ax,
+            fill_structures=fill_structures,
             **patch_kwargs,
         )
 

tidy3d/components/mode/simulation.py

@@ -361,6 +361,75 @@ def from_mode_solver(
         )
         return mode_sim
 
+    def plot(
+        self,
+        x: Optional[float] = None,
+        y: Optional[float] = None,
+        z: Optional[float] = None,
+        ax: Ax = None,
+        source_alpha: Optional[float] = 0,
+        monitor_alpha: Optional[float] = 0,
+        lumped_element_alpha: Optional[float] = 0,
+        hlim: Optional[tuple[float, float]] = None,
+        vlim: Optional[tuple[float, float]] = None,
+        fill_structures: bool = True,
+        **patch_kwargs,
+    ) -> Ax:
+        """Plot the mode simulation. If any of ``x``, ``y``, or ``z`` is provided, the potentially
+        larger FDTD simulation containing the mode plane is plotted at the desired location.
+        Otherwise, the mode plane is plotted by default.
+
+        Parameters
+        ----------
+        fill_structures : bool = True
+            Whether to fill structures with color or just draw outlines.
+        x : float = None
+            position of plane in x direction, only one of x, y, z must be specified to define plane.
+        y : float = None
+            position of plane in y direction, only one of x, y, z must be specified to define plane.
+        z : float = None
+            position of plane in z direction, only one of x, y, z must be specified to define plane.
+        source_alpha : float = 0
+            Opacity of the sources. If ``None``, uses Tidy3d default.
+        monitor_alpha : float = 0
+            Opacity of the monitors. If ``None``, uses Tidy3d default.
+        lumped_element_alpha : float = 0
+            Opacity of the lumped elements. If ``None``, uses Tidy3d default.
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+        hlim : Tuple[float, float] = None
+            The x range if plotting on xy or xz planes, y range if plotting on yz plane.
+        vlim : Tuple[float, float] = None
+            The z range if plotting on xz or yz planes, y plane if plotting on xy plane.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+        """
+
+        if x is not None or y is not None or z is not None:
+            return super().plot(
+                x=x,
+                y=y,
+                z=z,
+                ax=ax,
+                source_alpha=source_alpha,
+                monitor_alpha=monitor_alpha,
+                lumped_element_alpha=lumped_element_alpha,
+                hlim=hlim,
+                vlim=vlim,
+                fill_structures=fill_structures,
+                **patch_kwargs,
+            )
+        return self._mode_solver.plot(
+            ax=ax,
+            hlim=hlim,
+            vlim=vlim,
+            fill_structures=fill_structures,
+            **patch_kwargs,
+        )
+
     def plot_mode_plane(
         self,
         ax: Ax = None,