added

val option in SourceTime plotting

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 - Bug in remote file transfer when client environment has no correct certificate authority pem file install locally. 
 
+- `SourceTime` plotting methods `.plot()` and `.plot_spectrum()` accept a `val` kwarg, which selects which part of the data (`'real'`, `'imag'`, or `'abs'`) to plot, rather than plotting all at once.
 
 ## [1.9.1] - 2023-3-06
 

tests/test_components/test_source.py

@@ -14,8 +14,16 @@
 
 
 def test_plot_source_time():
-    ST.plot(times=[1e-15, 2e-15, 3e-15], ax=AX)
-    ST.plot_spectrum(times=[1e-15, 2e-15, 3e-15], num_freqs=4, ax=AX)
+
+    for val in ("real", "imag", "abs"):
+        ST.plot(times=[1e-15, 2e-15, 3e-15], val=val, ax=AX)
+        ST.plot_spectrum(times=[1e-15, 2e-15, 3e-15], num_freqs=4, val=val, ax=AX)
+
+    with pytest.raises(ValueError):
+        ST.plot(times=[1e-15, 2e-15, 3e-15], val="blah", ax=AX)
+
+    with pytest.raises(ValueError):
+        ST.plot_spectrum(times=[1e-15, 2e-15, 3e-15], num_freqs=4, val="blah", ax=AX)
 
     # uneven spacing in times
     with pytest.raises(SetupError):

tidy3d/components/data/sim_data.py

@@ -10,7 +10,7 @@
 from ..base import Tidy3dBaseModel
 from ..simulation import Simulation
 from ..boundary import BlochBoundary
-from ..types import Ax, Axis, annotate_type, Literal
+from ..types import Ax, Axis, annotate_type, Literal, PlotVal
 from ..viz import equal_aspect, add_ax_if_none
 from ...log import DataError, log, Tidy3dKeyError, ValidationError
 
@@ -245,7 +245,7 @@ def plot_field(  # pylint:disable=too-many-arguments, too-many-locals, too-many-
         self,
         field_monitor_name: str,
         field_name: str,
-        val: Literal["real", "imag", "abs"] = "real",
+        val: PlotVal = "real",
         eps_alpha: float = 0.2,
         robust: bool = True,
         vmin: float = None,

tidy3d/components/source.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from .base import Tidy3dBaseModel, cached_property, DATA_ARRAY_MAP
-from .types import Direction, Polarization, Ax, FreqBound, ArrayLike, Axis
+from .types import Direction, Polarization, Ax, FreqBound, ArrayLike, Axis, PlotVal
 from .validators import assert_plane, validate_name_str, get_value
 from .data.dataset import FieldDataset
 from .geometry import Box
@@ -99,7 +99,7 @@ def spectrum(
         return dt * dft_matrix @ time_amps
 
     @add_ax_if_none
-    def plot(self, times: ArrayLike[float, 1], ax: Ax = None) -> Ax:
+    def plot(self, times: ArrayLike[float, 1], val: PlotVal = "real", ax: Ax = None) -> Ax:
         """Plot the complex-valued amplitude of the source time-dependence.
 
         Parameters
@@ -108,6 +108,8 @@ def plot(self, times: ArrayLike[float, 1], ax: Ax = None) -> Ax:
             Array of times (seconds) to plot source at.
             To see source time amplitude for a specific :class:`Simulation`,
             pass ``simulation.tmesh``.
+        val : Literal['real', 'imag', 'abs'] = 'real'
+            Which part of the spectrum to plot.
         ax : matplotlib.axes._subplots.Axes = None
             Matplotlib axes to plot on, if not specified, one is created.
 
@@ -119,20 +121,27 @@ def plot(self, times: ArrayLike[float, 1], ax: Ax = None) -> Ax:
         times = np.array(times)
         amp_complex = self.amp_time(times)
 
-        ax.plot(times, amp_complex.real, color="blueviolet", label="real")
-        ax.plot(times, amp_complex.imag, color="crimson", label="imag")
-        ax.plot(times, np.abs(amp_complex), color="k", label="abs")
+        if val == "real":
+            ax.plot(times, amp_complex.real, color="blueviolet", label="real")
+        elif val == "imag":
+            ax.plot(times, amp_complex.imag, color="crimson", label="imag")
+        elif val == "abs":
+            ax.plot(times, np.abs(amp_complex), color="k", label="abs")
+        else:
+            raise ValueError(f"Plot 'val' option of '{val}' not recognized.")
         ax.set_xlabel("time (s)")
         ax.set_title("source amplitude")
         ax.legend()
         ax.set_aspect("auto")
         return ax
 
+    # pylint:disable=too-many-arguments
     @add_ax_if_none
     def plot_spectrum(
         self,
         times: ArrayLike[float, 1],
         num_freqs: int = 101,
+        val: PlotVal = "real",
         ax: Ax = None,
         complex_fields: bool = False,
     ) -> Ax:
@@ -170,9 +179,14 @@ def plot_spectrum(
 
         spectrum = self.spectrum(times=times, dt=dt, freqs=freqs, complex_fields=complex_fields)
 
-        ax.plot(freqs, spectrum.real, color="blueviolet", label="real")
-        ax.plot(freqs, spectrum.imag, color="crimson", label="imag")
-        ax.plot(freqs, np.abs(spectrum), color="k", label="abs")
+        if val == "real":
+            ax.plot(freqs, spectrum.real, color="blueviolet", label="real")
+        elif val == "imag":
+            ax.plot(freqs, spectrum.imag, color="crimson", label="imag")
+        elif val == "abs":
+            ax.plot(freqs, np.abs(spectrum), color="k", label="abs")
+        else:
+            raise ValueError(f"Plot 'val' option of '{val}' not recognized.")
         ax.set_xlabel("frequency (Hz)")
         ax.set_title("source spectrum")
         ax.legend()

tidy3d/components/types.py

@@ -179,6 +179,7 @@ def __modify_schema__(cls, field_schema):
 """ plotting """
 
 Ax = Axes
+PlotVal = Literal["real", "imag", "abs"]
 
 """ mode tracking """