add terminal modifications for autograd smatrix

Author: Gregory Roberts

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -4,7 +4,7 @@
 
 from typing import Optional, Union
 
-import numpy as np
+import autograd.numpy as np
 import pydantic.v1 as pd
 
 from tidy3d.components.base import cached_property
@@ -202,7 +202,11 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> TerminalPortData
 
         port_names = [port.name for port in self.ports]
 
-        values = np.zeros(
+        a_values = np.zeros(
+            (len(self.freqs), len(port_names), len(port_names)),
+            dtype=complex,
+        )
+        b_values = np.zeros(
             (len(self.freqs), len(port_names), len(port_names)),
             dtype=complex,
         )
@@ -211,21 +215,24 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> TerminalPortData
             "port_out": port_names,
             "port_in": port_names,
         }
-        a_matrix = TerminalPortDataArray(values, coords=coords)
-        b_matrix = a_matrix.copy(deep=True)
+        a_matrix = TerminalPortDataArray(a_values, coords=coords)
+        b_matrix = TerminalPortDataArray(b_values, coords=coords)
 
         # Tabulate the reference impedances at each port and frequency
         port_impedances = self._port_reference_impedances(batch_data=batch_data)
 
-        # loop through source ports
-        for port_in in self.ports:
+        for _port_in_idx, port_in in enumerate(self.ports):
             sim_data = batch_data[self._task_name(port=port_in)]
             a, b = self.compute_power_wave_amplitudes_at_each_port(port_impedances, sim_data)
             indexer = {"f": a.f, "port_in": port_in.name, "port_out": a.port}
-            a_matrix.loc[indexer] = a
-            b_matrix.loc[indexer] = b
+
+            a_data = np.expand_dims(a.data, axis=1)
+            b_data = np.expand_dims(b.data, axis=1)
+            a_matrix = a_matrix.with_updated_data(data=a_data, coords=indexer)
+            b_matrix = b_matrix.with_updated_data(data=b_data, coords=indexer)
 
         s_matrix = self.ab_to_s(a_matrix, b_matrix)
+
         return s_matrix
 
     @pd.validator("simulation")
@@ -315,9 +322,12 @@ def compute_power_wave_amplitudes_at_each_port(
 
         for port_out in self.ports:
             V_out, I_out = self.compute_port_VI(port_out, sim_data)
-            indexer = {"port": port_out.name}
-            V_matrix.loc[indexer] = V_out
-            I_matrix.loc[indexer] = I_out
+            indexer = {"f": V_out.f, "port": port_out.name}
+
+            V_out_data = np.expand_dims(V_out.data, axis=1)
+            I_out_data = np.expand_dims(I_out.data, axis=1)
+            V_matrix = V_matrix.with_updated_data(data=V_out_data, coords=indexer)
+            I_matrix = V_matrix.with_updated_data(data=I_out_data, coords=indexer)
 
         V_numpy = V_matrix.values
         I_numpy = I_matrix.values