correcting plot functions for ModeSolver

tidy3d/plugins/mode/mode_solver.py

@@ -1217,9 +1217,9 @@ def plot(
 
         """
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, _ = self._center_and_lims()
+        a_center, h_lim, v_lim, _, plot_sim_size = self._center_and_lims()
 
-        return self.simulation.plot(
+        return self.simulation.updated_copy(size=plot_sim_size).plot(
             x=a_center[0],
             y=a_center[1],
             z=a_center[2],
@@ -1265,12 +1265,12 @@ def plot_eps(
         """
 
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, _ = self._center_and_lims()
+        a_center, h_lim, v_lim, _, plot_sim_size = 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(
+        return self.simulation.updated_copy(size=plot_sim_size).plot_eps(
             x=a_center[0],
             y=a_center[1],
             z=a_center[2],
@@ -1324,12 +1324,12 @@ def plot_structures_eps(
         """
 
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, _ = self._center_and_lims()
+        a_center, h_lim, v_lim, _, plot_sim_size = 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(
+        return self.simulation.updated_copy(size=plot_sim_size).plot_structures_eps(
             x=a_center[0],
             y=a_center[1],
             z=a_center[2],
@@ -1365,9 +1365,9 @@ def plot_grid(
         """
 
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, _ = self._center_and_lims()
+        a_center, h_lim, v_lim, _, plot_sim_size = self._center_and_lims()
 
-        return self.simulation.plot_grid(
+        return self.simulation.updated_copy(size=plot_sim_size).plot_grid(
             x=a_center[0], y=a_center[1], z=a_center[2], hlim=h_lim, vlim=v_lim, ax=ax, **kwargs
         )
 
@@ -1389,11 +1389,11 @@ def plot_pml(
         """
 
         # Get the mode plane normal axis, center, and limits.
-        a_center, h_lim, v_lim, t_axes = self._center_and_lims()
+        a_center, h_lim, v_lim, t_axes, plot_sim_size = self._center_and_lims()
 
         # Plot the mode plane is ax=None.
         if not ax:
-            ax = self.simulation.plot(
+            ax = self.simulation.updated_copy(size=plot_sim_size).plot(
                 x=a_center[0],
                 y=a_center[1],
                 z=a_center[2],
@@ -1472,21 +1472,24 @@ def _center_and_lims(self) -> Tuple[List, List, List, List]:
         _, (h_min_s, v_min_s) = Box.pop_axis(self.simulation.bounds[0], axis=n_axis)
         _, (h_max_s, v_max_s) = Box.pop_axis(self.simulation.bounds[1], axis=n_axis)
 
-        h_min = a_center[n_axis] - self.plane.size[t_axes[0]] / 2
-        h_max = a_center[n_axis] + self.plane.size[t_axes[0]] / 2
-        v_min = a_center[n_axis] - self.plane.size[t_axes[1]] / 2
-        v_max = a_center[n_axis] + self.plane.size[t_axes[1]] / 2
+        h_min = self.plane.center[t_axes[0]] - self.plane.size[t_axes[0]] / 2
+        h_max = self.plane.center[t_axes[0]] + self.plane.size[t_axes[0]] / 2
+        v_min = self.plane.center[t_axes[1]] - self.plane.size[t_axes[1]] / 2
+        v_max = self.plane.center[t_axes[1]] + self.plane.size[t_axes[1]] / 2
 
-        h_lim = [
-            h_min if abs(h_min) < abs(h_min_s) else h_min_s,
-            h_max if abs(h_max) < abs(h_max_s) else h_max_s,
-        ]
-        v_lim = [
-            v_min if abs(v_min) < abs(v_min_s) else v_min_s,
-            v_max if abs(v_max) < abs(v_max_s) else v_max_s,
-        ]
+        h_lim = [h_min, h_max]
+        v_lim = [v_min, v_max]
+
+        # Define the size of the simulation to be plotted to ensure that it contains the ModeSolver plane
+        h_plane_size = h_lim[1] - h_lim[0]
+        v_plane_size = v_lim[1] - v_lim[0]
+
+        plane_size = [h_plane_size, v_plane_size]
+        plane_size.insert(a_center.index(0.0), 1)
+
+        plot_sim_size = tuple(plane_size[i] + self.simulation.size[i] for i in range(3))
 
-        return a_center, h_lim, v_lim, t_axes
+        return a_center, h_lim, v_lim, t_axes, plot_sim_size
 
     def _validate_modes_size(self):
         """Make sure that the total size of the modes fields is not too large."""