add plotting to path integrals

tidy3d/components/viz.py

@@ -4,7 +4,7 @@
 
 from functools import wraps
 from html import escape
-from typing import Any
+from typing import Any, Self
 
 import matplotlib.pyplot as plt
 import pydantic.v1 as pd
@@ -73,18 +73,15 @@ def _plot(*args, **kwargs) -> Ax:
 """ plot parameters """
 
 
-class PlotParams(Tidy3dBaseModel):
-    """Stores plotting parameters / specifications for a given model."""
+class AbstractPlotParams(Tidy3dBaseModel):
+    """Abstract class for storing plotting parameters.
+    Corresponds with select properties of ``matplotlib.artist.Artist``.
+    """
 
     alpha: Any = pd.Field(1.0, title="Opacity")
-    edgecolor: Any = pd.Field(None, title="Edge Color", alias="ec")
-    facecolor: Any = pd.Field(None, title="Face Color", alias="fc")
-    fill: bool = pd.Field(True, title="Is Filled")
-    hatch: str = pd.Field(None, title="Hatch Style")
     zorder: float = pd.Field(None, title="Display Order")
-    linewidth: pd.NonNegativeFloat = pd.Field(1, title="Line Width", alias="lw")
 
-    def include_kwargs(self, **kwargs) -> PlotParams:
+    def include_kwargs(self, **kwargs) -> Self:
         """Update the plot params with supplied kwargs."""
         update_dict = {
             key: value
@@ -101,6 +98,32 @@ def to_kwargs(self) -> dict:
         return kwarg_dict
 
 
+class PathPlotParams(AbstractPlotParams):
+    """Stores plotting parameters / specifications for a path.
+    Corresponds with select properties of ``matplotlib.lines.Line2D``.
+    """
+
+    color: Any = pd.Field(None, title="Color", alias="c")
+    linewidth: pd.NonNegativeFloat = pd.Field(2, title="Line Width", alias="lw")
+    linestyle: str = pd.Field("--", title="Line Style", alias="ls")
+    marker: Any = pd.Field("o", title="Marker Style")
+    markeredgecolor: Any = pd.Field(None, title="Marker Edge Color", alias="mec")
+    markerfacecolor: Any = pd.Field(None, title="Marker Face Color", alias="mfc")
+    markersize: pd.NonNegativeFloat = pd.Field(10, title="Marker Size", alias="ms")
+
+
+class PlotParams(AbstractPlotParams):
+    """Stores plotting parameters / specifications for a given model.
+    Corresponds with select properties of ``matplotlib.patches.Patch``.
+    """
+
+    edgecolor: Any = pd.Field(None, title="Edge Color", alias="ec")
+    facecolor: Any = pd.Field(None, title="Face Color", alias="fc")
+    fill: bool = pd.Field(True, title="Is Filled")
+    hatch: str = pd.Field(None, title="Hatch Style")
+    linewidth: pd.NonNegativeFloat = pd.Field(1, title="Line Width", alias="lw")
+
+
 # defaults for different tidy3d objects
 plot_params_geometry = PlotParams()
 plot_params_structure = PlotParams()

tidy3d/plugins/microwave/custom_path_integrals.py

@@ -12,7 +12,8 @@
 from ...components.data.data_array import FreqDataArray, FreqModeDataArray, TimeDataArray
 from ...components.data.monitor_data import FieldData, FieldTimeData, ModeSolverData
 from ...components.geometry.base import Geometry
-from ...components.types import ArrayFloat2D, Axis, Bound, Coordinate
+from ...components.types import ArrayFloat2D, Ax, Axis, Bound, Coordinate
+from ...components.viz import add_ax_if_none
 from ...constants import MICROMETER, fp_eps
 from ...exceptions import DataError, SetupError
 from .path_integrals import (
@@ -22,6 +23,13 @@
     MonitorDataTypes,
     VoltageIntegralAxisAligned,
 )
+from .viz import (
+    ARROW_CURRENT,
+    plot_params_current_path,
+    plot_params_voltage_minus,
+    plot_params_voltage_path,
+    plot_params_voltage_plus,
+)
 
 FieldParameter = Literal["E", "H"]
 
@@ -172,8 +180,8 @@ def from_circular_path(
             A path integral defined on a circular path.
         """
 
-        # Helper for generating x,y vertices around a circle in the local coordinate frame
-        def generate_circle_coordinates(radius, num_points, clockwise):
+        def generate_circle_coordinates(radius: float, num_points: int, clockwise: bool):
+            """Helper for generating x,y vertices around a circle in the local coordinate frame."""
             sign = 1.0
             if clockwise:
                 sign = -1.0
@@ -199,6 +207,7 @@ def generate_circle_coordinates(radius, num_points, clockwise):
 
     @cached_property
     def is_closed_contour(self) -> bool:
+        """Returns ``true`` when the first vertex equals the last vertex."""
         return np.isclose(
             self.vertices[0, :],
             self.vertices[-1, :],
@@ -261,6 +270,55 @@ def compute_voltage(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
         voltage = VoltageIntegralAxisAligned._set_data_array_attributes(voltage)
         return voltage
 
+    @add_ax_if_none
+    def plot(
+        self,
+        x: float = None,
+        y: float = None,
+        z: float = None,
+        ax: Ax = None,
+        **path_kwargs,
+    ) -> Ax:
+        """Plot path integral at single (x,y,z) coordinate.
+
+        Parameters
+        ----------
+        x : float = None
+            Position of plane in x direction, only one of x,y,z can be specified to define plane.
+        y : float = None
+            Position of plane in y direction, only one of x,y,z can be specified to define plane.
+        z : float = None
+            Position of plane in z direction, only one of x,y,z can be specified to define plane.
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+        **path_kwargs
+            Optional keyword arguments passed to the matplotlib plotting of the line.
+            For details on accepted values, refer to
+            `Matplotlib's documentation <https://tinyurl.com/36marrat>`_.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+        """
+        axis, position = Geometry.parse_xyz_kwargs(x=x, y=y, z=z)
+        if axis != self.main_axis or not np.isclose(position, self.position, rtol=fp_eps):
+            return ax
+
+        plot_params = plot_params_voltage_path.include_kwargs(**path_kwargs)
+        plot_kwargs = plot_params.to_kwargs()
+        xs = self.vertices[:, 0]
+        ys = self.vertices[:, 1]
+        ax.plot(xs, ys, markevery=[0, -1], **plot_kwargs)
+
+        # Plot special end points
+        end_kwargs = plot_params_voltage_plus.include_kwargs(**path_kwargs).to_kwargs()
+        start_kwargs = plot_params_voltage_minus.include_kwargs(**path_kwargs).to_kwargs()
+        ax.plot(xs[0], ys[0], **start_kwargs)
+        ax.plot(xs[-1], ys[-1], **end_kwargs)
+
+        return ax
+
 
 class CustomCurrentIntegral2D(CustomPathIntegral2D):
     """Class for computing conduction current via Ampère's circuital law on a custom path.
@@ -283,3 +341,53 @@ def compute_current(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
         current = self.compute_integral(field="H", em_field=em_field)
         current = CurrentIntegralAxisAligned._set_data_array_attributes(current)
         return current
+
+    @add_ax_if_none
+    def plot(
+        self,
+        x: float = None,
+        y: float = None,
+        z: float = None,
+        ax: Ax = None,
+        **path_kwargs,
+    ) -> Ax:
+        """Plot path integral at single (x,y,z) coordinate.
+
+        Parameters
+        ----------
+        x : float = None
+            Position of plane in x direction, only one of x,y,z can be specified to define plane.
+        y : float = None
+            Position of plane in y direction, only one of x,y,z can be specified to define plane.
+        z : float = None
+            Position of plane in z direction, only one of x,y,z can be specified to define plane.
+        ax : matplotlib.axes._subplots.Axes = None
+            Matplotlib axes to plot on, if not specified, one is created.
+        **path_kwargs
+            Optional keyword arguments passed to the matplotlib plotting of the line.
+            For details on accepted values, refer to
+            `Matplotlib's documentation <https://tinyurl.com/36marrat>`_.
+
+        Returns
+        -------
+        matplotlib.axes._subplots.Axes
+            The supplied or created matplotlib axes.
+        """
+        axis, position = Geometry.parse_xyz_kwargs(x=x, y=y, z=z)
+        if axis != self.main_axis or not np.isclose(position, self.position, rtol=fp_eps):
+            return ax
+
+        plot_params = plot_params_current_path.include_kwargs(**path_kwargs)
+        plot_kwargs = plot_params.to_kwargs()
+        xs = self.vertices[:, 0]
+        ys = self.vertices[:, 1]
+        ax.plot(xs, ys, **plot_kwargs)
+
+        # Add arrow at start of contour
+        ax.annotate(
+            "",
+            xytext=(xs[0], ys[0]),
+            xy=(xs[1], ys[1]),
+            arrowprops=ARROW_CURRENT,
+        )
+        return ax