Generalize getting raw matrix from dict, comment out for scattering

gplugins/common/base_models/simulation.py

@@ -9,7 +9,32 @@
 from numpy.typing import NDArray
 from pydantic import BaseModel, ConfigDict, computed_field
 
-from ..types import CapacitanceDict
+from ..types import CapacitanceDict, ScatteringDict
+
+
+def _raw_matrix_from_dict(dict_matrix: CapacitanceDict | ScatteringDict) -> NDArray:
+    """Converts dictionary formatted matrix results to a NumPy array.
+
+    Args:
+        dict_matrix: Dictionary with matrix results in ``(port_i, port_j): result`` configuration.
+
+    Returns:
+        ndarray: A matrix representation of the connections.
+    """
+    n = int(np.sqrt(len(dict_matrix)))
+    matrix = np.zeros((n, n))
+
+    port_to_index_map = {}
+    for iname, jname in dict_matrix.keys():
+        if iname not in port_to_index_map:
+            port_to_index_map[iname] = len(port_to_index_map) + 1
+        if jname not in port_to_index_map:
+            port_to_index_map[jname] = len(port_to_index_map) + 1
+
+    for (iname, jname), c in dict_matrix.items():
+        matrix[port_to_index_map[iname], port_to_index_map[jname]] = c
+
+    return matrix
 
 
 class ElectrostaticResults(BaseModel):
@@ -24,25 +49,19 @@ class ElectrostaticResults(BaseModel):
     @computed_field
     @cached_property
     def raw_capacitance_matrix(self) -> NDArray:
-        n = int(np.sqrt(len(self.capacitance_matrix)))
-        matrix = np.zeros((n, n))
-
-        port_to_index_map = {}
-        for iname, jname in self.capacitance_matrix.keys():
-            if iname not in port_to_index_map:
-                port_to_index_map[iname] = len(port_to_index_map) + 1
-            if jname not in port_to_index_map:
-                port_to_index_map[jname] = len(port_to_index_map) + 1
-
-        for (iname, jname), c in self.capacitance_matrix.items():
-            matrix[port_to_index_map[iname], port_to_index_map[jname]] = c
-
-        return matrix
+        """Capacitance matrix as a NumPy array."""
+        return _raw_matrix_from_dict(self.capacitance_matrix)
 
 
 class DrivenFullWaveResults(BaseModel):
     """Results class for driven full-wave simulations."""
 
-    scattering_matrix: Any  # TODO convert to SDict or similar
+    scattering_matrix: Any  # TODO convert dataframe to ScatteringDict
     mesh_location: Path | None = None
     field_file_locations: Sequence[Path] | None = None
+
+    # @computed_field
+    # @cached_property
+    # def raw_scattering_matrix(self) -> NDArray:
+    #     """Scattering matrix as a NumPy array."""
+    #     return _raw_matrix_from_dict(self.scattering_matrix)

gplugins/common/types.py

@@ -6,6 +6,7 @@
 
 RFMaterialSpec = dict[str, dict[str, float | int]]
 CapacitanceDict = dict[tuple[str, str], float]
+ScatteringDict = dict[tuple[str, str], float]
 AnyShapelyPolygon = Annotated[
     GeometryCollection | MultiPolygon | Polygon,
     PlainSerializer(lambda x: x.wkb_hex, when_used="json"),