From 2c51f37f8f7f99a485dff9b24989ba3ba47985c1 Mon Sep 17 00:00:00 2001
From: tylerflex <tyler@flexcompute.com>
Date: Fri, 12 Nov 2021 09:12:54 -0800
Subject: [PATCH] fixed plot_field and added source spectrum calcultor

---
 tidy3d/components/data.py   | 25 ++++++++++++++++---------
 tidy3d/components/source.py | 24 ++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/tidy3d/components/data.py b/tidy3d/components/data.py
index b5a304f2f7..f427294dfb 100644
--- a/tidy3d/components/data.py
+++ b/tidy3d/components/data.py
@@ -535,7 +535,6 @@ def plot_field(
         val: Literal["real", "imag", "abs"] = "real",
         freq: float = None,
         time: float = None,
-        cbar: bool = None,
         eps_alpha: float = 0.2,
         ax: Ax = None,
         **kwargs,
@@ -558,10 +557,10 @@ def plot_field(
             What part of the field to plot (in )
         freq: float = None
             If monitor is a :class:`FieldMonitor`, specifies the frequency (Hz) to plot the field.
+            Also sets the frequency at which the permittivity is evaluated at (if dispersive).
+            By default, chooses permittivity as frequency goes to infinity.
         time: float = None
             if monitor is a :class:`FieldTimeMonitor`, specifies the time (sec) to plot the field.
-        cbar: bool = True
-            if True (default), will include colorbar
         eps_alpha : float = 0.2
             Opacity of the structure permittivity.
             Must be between 0 and 1 (inclusive).
@@ -589,11 +588,11 @@ def plot_field(
         xr_data = monitor_data.data_dict.get(field_name).data
 
         # select the frequency or time value
-        if "f" in monitor_data.coords:
+        if "f" in xr_data.coords:
             if freq is None:
                 raise DataError("'freq' must be supplied to plot a FieldMonitor.")
             field_data = xr_data.interp(f=freq)
-        elif "t" in monitor_data.coords:
+        elif "t" in xr_data.coords:
             if time is None:
                 raise DataError("'time' must be supplied to plot a FieldMonitor.")
             field_data = xr_data.interp(t=time)
@@ -620,14 +619,22 @@ def plot_field(
             field_data = abs(field_data)
 
         # plot the field
-        xy_coords = list("xyz")
-        xy_coords.pop(axis)
-        field_data.plot(ax=ax, x=xy_coords[0], y=xy_coords[1])
+        xy_coord_labels = list("xyz")
+        xy_coord_labels.pop(axis)
+        x_coord_label, y_coord_label = xy_coord_labels
+        field_data.plot(ax=ax, x=x_coord_label, y=y_coord_label)
 
         # plot the simulation epsilon
         ax = self.simulation.plot_structures_eps(
-            freq=freq, cbar=cbar, x=x, y=y, z=z, alpha=eps_alpha, ax=ax
+            freq=freq, cbar=False, x=x, y=y, z=z, alpha=eps_alpha, ax=ax
         )
+
+        # set the limits based on the xarray coordinates min and max
+        x_coord_values = field_data.coords[x_coord_label]
+        y_coord_values = field_data.coords[y_coord_label]
+        ax.set_xlim(min(x_coord_values), max(x_coord_values))
+        ax.set_ylim(min(y_coord_values), max(y_coord_values))
+
         return ax
 
     def export(self, fname: str) -> None:
diff --git a/tidy3d/components/source.py b/tidy3d/components/source.py
index ff8ab27e98..79651af489 100644
--- a/tidy3d/components/source.py
+++ b/tidy3d/components/source.py
@@ -42,6 +42,30 @@ def amp_time(self, time: float) -> complex:
             Complex-valued source amplitude at that time..
         """
 
+    def spectrum(self, times: Array[float], freqs: Array[float], dt: float) -> complex:
+        """Complex-valued source spectrum as a function of frequency
+
+        Parameters
+        ----------
+        times : np.ndarray
+            Times to use to evaluate spectrum Fourier transform.
+            (Typically the simulation time mesh).
+        freqs : np.ndarray
+            Frequencies in Hz to evaluate spectrum at.
+        dt : float or np.ndarray
+            Time step to weight FT integral with.
+            If array, use to weigh each of the time intervals in ``times``.
+
+        Returns
+        -------
+        np.ndarray
+            Complex-valued array (of len(freqs)) containing spectrum at those frequencies.
+        """
+
+        # (Nf, Nt) matrix that gives DFT when matrix multiplied with signal
+        dft_matrix = np.exp(2j * np.pi * freqs[:, None] * times) / (2 * np.pi)
+        return dt * dft_matrix @ self.amp_time(times)
+
     @add_ax_if_none
     def plot(self, times: Array[float], ax: Ax = None) -> Ax:
         """Plot the complex-valued amplitude of the source time-dependence.