revert to composition and allow a single impedance spec to be used for all modes

Author: dmarek-flex

tests/test_components/test_microwave.py

@@ -200,10 +200,13 @@ def make_mw_sim(
     size_port = [0, sim_width, size_sim[2]]
     center_port = [0, 0, center_sim[2]]
     impedance_specs = (td.AutoImpedanceSpec(),) * 4
-    mode_spec = td.MicrowaveModeSpec(
+    microwave_spec = td.MicrowaveModeSpec(
+        impedance_specs=impedance_specs,
+    )
+    mode_spec = td.ModeSpec(
         num_modes=4,
         target_neff=1.8,
-        impedance_specs=impedance_specs,
+        microwave_spec=microwave_spec,
     )
 
     mode_monitor = td.ModeMonitor(
@@ -484,7 +487,12 @@ def test_impedance_spec_validation():
     with pytest.raises(pd.ValidationError):
         _ = td.CustomImpedanceSpec(voltage_spec=None, current_spec=None)
 
-    _ = td.MicrowaveModeSpec(num_modes=4, impedance_specs=(both, voltage_only, current_only, None))
+    _ = td.ModeSpec(
+        num_modes=4,
+        microwave_spec=td.MicrowaveModeSpec(
+            impedance_specs=(both, voltage_only, current_only, None)
+        ),
+    )
 
 
 def test_path_integral_factory_voltage_validation():
@@ -570,21 +578,19 @@ def test_make_path_integrals_validation():
         voltage_spec=v_spec,
         current_spec=i_spec,
     )
-    microwave_mode_spec = td.MicrowaveModeSpec(impedance_specs=(impedance_spec,))
+    microwave_spec = td.MicrowaveModeSpec(impedance_specs=(impedance_spec,))
 
     # Test successful creation
-    voltage_integrals, current_integrals = make_path_integrals(microwave_mode_spec, mode_monitor)
+    voltage_integrals, current_integrals = make_path_integrals(microwave_spec, mode_monitor)
     assert len(voltage_integrals) == 1
     assert len(current_integrals) == 1
     assert voltage_integrals[0] is not None
     assert current_integrals[0] is not None
 
     # Test with None specs - when both are None, use_automatic_setup is True
     # This means current integrals will be auto-generated, not None
-    microwave_mode_spec_none = td.MicrowaveModeSpec(impedance_specs=(None,))
-    voltage_integrals, current_integrals = make_path_integrals(
-        microwave_mode_spec_none, mode_monitor
-    )
+    microwave_spec_none = td.MicrowaveModeSpec(impedance_specs=(None,))
+    voltage_integrals, current_integrals = make_path_integrals(microwave_spec_none, mode_monitor)
     assert len(voltage_integrals) == mode_monitor.mode_spec.num_modes
     assert len(current_integrals) == mode_monitor.mode_spec.num_modes
     assert all(vi is None for vi in voltage_integrals)
@@ -601,7 +607,7 @@ def test_make_path_integrals_construction_errors(monkeypatch):
     v_spec = td.AxisAlignedVoltageIntegralSpec(center=(1, 2, 3), size=(0, 0, 1), sign="-")
 
     impedance_spec = td.CustomImpedanceSpec(voltage_spec=v_spec, current_spec=None)
-    microwave_mode_spec = td.MicrowaveModeSpec(impedance_specs=(impedance_spec,))
+    microwave_spec = td.MicrowaveModeSpec(impedance_specs=(impedance_spec,))
 
     # Mock make_voltage_integral to raise an exception
     def mock_make_voltage_integral(path_spec):
@@ -614,7 +620,7 @@ def mock_make_voltage_integral(path_spec):
 
     # This should raise a SetupError due to construction failure
     with pytest.raises(SetupError, match="Failed to construct path integrals"):
-        make_path_integrals(microwave_mode_spec, mode_monitor)
+        make_path_integrals(microwave_spec, mode_monitor)
 
 
 def test_path_integral_factory_composite_current():
@@ -662,14 +668,14 @@ def test_path_integral_factory_mixed_specs():
     # Test with mixed specs - some None, some specified
     impedance_spec1 = td.CustomImpedanceSpec(voltage_spec=v_spec, current_spec=None)
     impedance_spec2 = td.CustomImpedanceSpec(voltage_spec=None, current_spec=i_spec)
-    microwave_mode_spec = td.MicrowaveModeSpec(
-        num_modes=2,
+    microwave_spec = td.MicrowaveModeSpec(
         impedance_specs=(
             impedance_spec1,
             impedance_spec2,
         ),
     )
-    voltage_integrals, current_integrals = make_path_integrals(microwave_mode_spec, mode_monitor)
+
+    voltage_integrals, current_integrals = make_path_integrals(microwave_spec, mode_monitor)
 
     assert len(voltage_integrals) == 2
     assert len(current_integrals) == 2
@@ -697,11 +703,15 @@ def test_mode_solver_with_microwave_mode_spec():
 
     plane = td.Box(center=(0, 0, 0), size=(0, 10 * width, 2 * height + metal_thickness))
     num_modes = 3
-    impedance_specs = (td.AutoImpedanceSpec(), None, None)
-    mode_spec = td.MicrowaveModeSpec(
+    # impedance_specs = (td.AutoImpedanceSpec(), None, None)
+    impedance_specs = td.AutoImpedanceSpec()
+    microwave_spec = td.MicrowaveModeSpec(
+        impedance_specs=impedance_specs,
+    )
+    mode_spec = td.ModeSpec(
         num_modes=num_modes,
         target_neff=2.2,
-        impedance_specs=impedance_specs,
+        microwave_spec=microwave_spec,
     )
     mms = ModeSolver(
         simulation=stripline_sim,
@@ -729,7 +739,8 @@ def test_mode_solver_with_microwave_mode_spec():
         ),
     )
     impedance_specs = (custom_spec, None, None)
-    mms = mms.updated_copy(path="mode_spec/", impedance_specs=impedance_specs)
+    microwave_spec_custom = td.MicrowaveModeSpec(impedance_specs=impedance_specs)
+    mms = mms.updated_copy(path="mode_spec/", microwave_spec=microwave_spec_custom)
     mms_data: ModeSolverData = mms.data
 
     # _, ax = plt.subplots(1, 1, tight_layout=True, figsize=(15, 15))
@@ -738,4 +749,10 @@ def test_mode_solver_with_microwave_mode_spec():
     # plt.show()
     mms_data.to_dataframe()
 
-    assert np.all(np.isclose(mms_data.microwave_data.Z0.real, 28.6, 0.2))
+    assert np.all(np.isclose(mms_data.microwave_data.Z0.real.sel(mode_index=0), 28.6, 0.2))
+
+    # Make sure a single spec can be used
+    microwave_spec_custom = td.MicrowaveModeSpec(impedance_specs=custom_spec)
+    mms = mms.updated_copy(path="mode_spec/", microwave_spec=microwave_spec_custom)
+    mms_data: ModeSolverData = mms.data
+    assert np.all(np.isclose(mms_data.microwave_data.Z0.real.sel(mode_index=0), 28.6, 0.2))

tests/test_components/test_mode.py

@@ -394,10 +394,12 @@ def test_mode_spec_with_microwave_mode_spec():
     impedance_specs = (td.AutoImpedanceSpec(),)
 
     # Should work when impedance_specs matches num_modes
-    mw_mode_spec = td.MicrowaveModeSpec(num_modes=1, impedance_specs=impedance_specs)
-    assert mw_mode_spec.num_modes == 1
-    assert len(mw_mode_spec.impedance_specs) == 1
+    mw_spec = td.MicrowaveModeSpec(impedance_specs=impedance_specs)
+    mode_spec = td.ModeSpec(num_modes=1, microwave_spec=mw_spec)
+    assert mode_spec.num_modes == 1
+    assert len(mode_spec.microwave_spec.impedance_specs) == 1
 
     # Should fail when impedance_specs doesn't match num_modes
     with pytest.raises(pydantic.ValidationError):
-        td.MicrowaveModeSpec(num_modes=2, impedance_specs=impedance_specs)
+        mw_spec_2 = td.MicrowaveModeSpec(impedance_specs=impedance_specs)
+        td.ModeSpec(num_modes=2, microwave_spec=mw_spec_2)

tidy3d/components/microwave/microwave_mode_spec.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Union
 
 import pydantic.v1 as pd
 
@@ -12,12 +12,10 @@
     AutoImpedanceSpec,
     ImpedanceSpecTypes,
 )
-from tidy3d.components.mode_spec import ModeSpec
 from tidy3d.components.types import annotate_type
-from tidy3d.exceptions import SetupError
 
 
-class MicrowaveModeSpec(MicrowaveBaseModel, ModeSpec):
+class MicrowaveModeSpec(MicrowaveBaseModel):
     """
     The :class:`.MicrowaveModeSpec` class specifies how quantities related to transmission line
     modes and microwave waveguides are computed. For example, it defines the paths for line integrals, which are used to
@@ -26,13 +24,12 @@ class MicrowaveModeSpec(MicrowaveBaseModel, ModeSpec):
     Example
     -------
     >>> import tidy3d as td
-    >>> # Using automatic impedance calculation
-    >>> mode_spec_auto = td.MicrowaveModeSpec(
+    >>> # Using automatic impedance calculation (single spec, will be duplicated for all modes)
+    >>> mode_spec_auto = td.ModeSpec(
     ...     num_modes=2,
-    ...     target_neff=1.5,
-    ...     impedance_specs=(td.AutoImpedanceSpec(), td.AutoImpedanceSpec())
+    ...     microwave_spec=td.MicrowaveModeSpec(impedance_specs=td.AutoImpedanceSpec())
     ... )
-    >>> # Using custom impedance specification
+    >>> # Using custom impedance specification for multiple modes
     >>> voltage_spec = td.AxisAlignedVoltageIntegralSpec(
     ...     center=(0, 0, 0), size=(0, 0, 1), sign="+"
     ... )
@@ -42,47 +39,36 @@ class MicrowaveModeSpec(MicrowaveBaseModel, ModeSpec):
     >>> custom_impedance = td.CustomImpedanceSpec(
     ...     voltage_spec=voltage_spec, current_spec=current_spec
     ... )
-    >>> mode_spec_custom = td.MicrowaveModeSpec(
+    >>> mode_spec_custom = td.ModeSpec(
     ...     num_modes=1,
-    ...     impedance_specs=(custom_impedance,)
+    ...     microwave_spec=td.MicrowaveModeSpec(impedance_specs=custom_impedance)
     ... )
     """
 
-    impedance_specs: tuple[Optional[annotate_type(ImpedanceSpecTypes)], ...] = pd.Field(
+    impedance_specs: Union[
+        annotate_type(ImpedanceSpecTypes),
+        tuple[Optional[annotate_type(ImpedanceSpecTypes)], ...],
+    ] = pd.Field(
         ...,
         title="Impedance Specifications",
         description="Field controls how the impedance is calculated for each mode calculated by the mode solver. "
-        "The number of impedance specifications should match the number of modes field. "
+        "Can be a single impedance specification (which will be applied to all modes) or a tuple of specifications "
+        "(one per mode). The number of impedance specifications should match the number of modes field. "
         "When an impedance specification of ``None`` is used, the impedance calculation will be "
         "ignored for the associated mode.",
     )
 
+    @cached_property
+    def _impedance_specs_as_tuple(self) -> tuple[Optional[ImpedanceSpecTypes]]:
+        """Gets the impedance_specs field converted to a tuple."""
+        if isinstance(self.impedance_specs, Union[tuple, list]):
+            return tuple(self.impedance_specs)
+        return (self.impedance_specs,)
+
     @cached_property
     def _using_auto_current_spec(self) -> bool:
         """Checks whether at least one of the modes will require an auto setup of the current path specification."""
         return any(
-            isinstance(impedance_spec, AutoImpedanceSpec) for impedance_spec in self.impedance_specs
+            isinstance(impedance_spec, AutoImpedanceSpec)
+            for impedance_spec in self._impedance_specs_as_tuple
         )
-
-    @cached_property
-    def num_impedance_specs(self) -> int:
-        """The number of impedance specifications to be used."""
-        return len(self.impedance_specs)
-
-    @pd.validator("impedance_specs", always=True)
-    def check_impedance_specs_consistent_with_num_modes(cls, val, values):
-        """Check that the number of impedance specifications is equal to the number of modes."""
-        if val is None:
-            return val
-
-        num_modes = values.get("num_modes", 1)
-        valid_number_impedance_specs = len(val) == num_modes
-
-        if not valid_number_impedance_specs:
-            raise SetupError(
-                f"Given {len(val)} impedance specifications in the 'MicrowaveModeSpec', "
-                f"but the number of modes requested is {num_modes}. Please ensure that the "
-                "number of impedance specifications is equal to the number of modes."
-            )
-
-        return val

tidy3d/components/microwave/path_integrals/path_integral_factory.py

@@ -86,7 +86,7 @@ def make_path_integrals(
     monitor: Union[ModeMonitor, ModeSolverMonitor],
 ) -> tuple[tuple[Optional[VoltageIntegralTypes]], tuple[Optional[CurrentIntegralTypes]]]:
     """
-    Given an impedance specification and monitor, create the voltage and
+    Given a microwave mode specification and monitor, create the voltage and
     current path integrals used for the impedance computation.
 
     Parameters
@@ -112,7 +112,11 @@ def make_path_integrals(
 
     v_integrals = []
     i_integrals = []
-    for idx, impedance_spec in enumerate(microwave_mode_spec.impedance_specs):
+
+    # Handle case where impedance spec is a single ImpedanceSpecTypes
+    impedance_specs = microwave_mode_spec._impedance_specs_as_tuple
+
+    for idx, impedance_spec in enumerate(impedance_specs):
         if impedance_spec is None:
             # Do not calculate impedance for this mode
             v_integrals.append(None)

tidy3d/components/mode/mode_solver.py

@@ -35,10 +35,9 @@
     LossyMetalMedium,
 )
 from tidy3d.components.microwave.data.dataset import MicrowaveModeDataset
-from tidy3d.components.microwave.microwave_mode_spec import MicrowaveModeSpec
 from tidy3d.components.microwave.path_integrals.path_integral_factory import make_path_integrals
 from tidy3d.components.mode_spec import ModeSpec
-from tidy3d.components.monitor import ModeMonitor, ModeSolverMonitor, ModeSpecTypes
+from tidy3d.components.monitor import ModeMonitor, ModeSolverMonitor
 from tidy3d.components.scene import Scene
 from tidy3d.components.simulation import Simulation
 from tidy3d.components.source.field import ModeSource
@@ -157,7 +156,7 @@ class ModeSolver(Tidy3dBaseModel):
         discriminator=TYPE_TAG_STR,
     )
 
-    mode_spec: ModeSpecTypes = pydantic.Field(
+    mode_spec: ModeSpec = pydantic.Field(
         ...,
         title="Mode specification",
         description="Container with specifications about the modes to be solved for.",
@@ -439,9 +438,9 @@ def _num_cells_freqs_modes(self) -> tuple[int, int, int]:
         return num_cells, num_freqs, num_modes
 
     @property
-    def _is_microwave_mode_spec(self) -> bool:
-        """Check if the mode_spec is a MicrowaveModeSpec."""
-        return isinstance(self.mode_spec, MicrowaveModeSpec)
+    def _has_microwave_mode_spec(self) -> bool:
+        """Check if the mode_spec is using a MicrowaveModeSpec."""
+        return self.mode_spec.microwave_spec is not None
 
     def solve(self) -> ModeSolverData:
         """:class:`.ModeSolverData` containing the field and effective index data.
@@ -550,7 +549,7 @@ def data_raw(self) -> ModeSolverData:
 
         mode_solver_data = self._filter_components(mode_solver_data)
         # Calculate and add the characteristic impedance
-        if self._is_microwave_mode_spec:
+        if self._has_microwave_mode_spec:
             mode_solver_data = self._add_microwave_data(mode_solver_data)
         return mode_solver_data
 
@@ -1377,10 +1376,12 @@ def _make_path_integrals(
     ) -> tuple[tuple[Optional[VoltageIntegralTypes]], tuple[Optional[CurrentIntegralTypes]]]:
         """Wrapper for making path integrals from the MicrowaveModeSpec. Note: overriden in the backend to support
         auto creation of path integrals."""
-        if not self._is_microwave_mode_spec:
-            raise ValueError("Cannot make path integrals for non-MicrowaveModeSpec")
+        if not self._has_microwave_mode_spec:
+            raise ValueError(
+                "Cannot make path integrals for when 'mode_spec' is not a 'MicrowaveModeSpec'."
+            )
         return make_path_integrals(
-            self.mode_spec,
+            self.mode_spec.microwave_spec,
             self.to_monitor(name=MODE_MONITOR_NAME),
         )
 
@@ -1392,6 +1393,9 @@ def _add_microwave_data(self, mode_solver_data: ModeSolverData) -> ModeSolverDat
         Z0_list = []
         V_list = []
         I_list = []
+        if len(voltage_integrals) == 1 and self.mode_spec.num_modes > 1:
+            voltage_integrals = voltage_integrals * self.mode_spec.num_modes
+            current_integrals = current_integrals * self.mode_spec.num_modes
         for mode_index in range(self.mode_spec.num_modes):
             vi = voltage_integrals[mode_index]
             ci = current_integrals[mode_index]