Partially reverting #2534 and reorganizing mode validation

fix: Python 3.9 compatibility for isinstance with Union types

Use get_args() to extract types from IsotropicUniformMediumType Union
for isinstance check. Fixes test_validation_from_simulation failure.

Related to PR #2534

Author: momchil-flex

tests/test_components/test_mode.py

@@ -89,6 +89,11 @@ def test_validation_from_simulation():
         monitors=[],
     )
 
+    reg_geometry = td.Structure(
+        geometry=td.Box.from_bounds((-100, -1, -100), (100, 1, 0)),
+        medium=td.Medium(permittivity=4.0, conductivity=1e-4),
+    )
+
     inf_geometry = td.Structure(
         geometry=td.Box.from_bounds((-td.inf, -1, -100), (td.inf, 1, 0)),
         medium=td.Medium(permittivity=4.0, conductivity=1e-4),
@@ -117,32 +122,23 @@ def test_validation_from_simulation():
         direction="+",
     )
 
+    # First test that a mode object can be added if there's no problem with the geometries
+    _ = sim.updated_copy(structures=[reg_geometry], monitors=[rot_monitor])
+
     # Test that transforming a geometry with an infinite extent raises an error
     with pytest.raises(SetupError):
-        sim.updated_copy(
-            structures=[inf_geometry],
-            monitors=[rot_monitor],
-        )
+        sim.updated_copy(structures=[inf_geometry], monitors=[rot_monitor])
 
     # Test that transforming an anisotropic medium raises an error
     with pytest.raises(SetupError):
-        sim.updated_copy(
-            structures=[anisotropic_geometry],
-            monitors=[rot_monitor],
-        )
+        sim.updated_copy(structures=[anisotropic_geometry], monitors=[rot_monitor])
 
     # Same thing with a ModeSource
     with pytest.raises(SetupError):
-        sim.updated_copy(
-            structures=[inf_geometry],
-            sources=[rot_source],
-        )
+        sim.updated_copy(structures=[inf_geometry], sources=[rot_source])
 
     with pytest.raises(SetupError):
-        sim.updated_copy(
-            structures=[anisotropic_geometry],
-            monitors=[rot_monitor],
-        )
+        sim.updated_copy(structures=[anisotropic_geometry], sources=[rot_source])
 
     # Same thing with ModeSimulation
     with pytest.raises(SetupError):

tidy3d/components/mode/mode_solver.py

@@ -6,7 +6,7 @@
 
 from functools import wraps
 from math import isclose
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, Union, get_args
 
 import numpy as np
 import pydantic.v1 as pydantic
@@ -218,14 +218,17 @@ def _post_init_validators(self) -> None:
             mode_spec=self.mode_spec,
             plane=self.plane,
             sim_geom=self.simulation.geometry,
-            msg_prefix="Mode solver",
         )
         self._warn_thick_pml(simulation=self.simulation, plane=self.plane, mode_spec=self.mode_spec)
         self._validate_rotate_structures()
 
     @classmethod
     def _warn_thick_pml(
-        cls, simulation: Simulation, plane: Box, mode_spec: ModeSpec, warn_str: str = "'ModeSolver'"
+        cls,
+        simulation: Simulation,
+        plane: Box,
+        mode_spec: ModeSpec,
+        msg_prefix: str = "'ModeSolver'",
     ):
         """Warn if the pml covers a significant portion of the mode plane."""
         coord_0, coord_1 = cls._plane_grid(
@@ -239,7 +242,7 @@ def _warn_thick_pml(
         for i in (0, 1):
             if 2 * effective_num_pml[i] > (WARN_THICK_PML_PERCENT / 100) * num_cells[i]:
                 log.warning(
-                    f"{warn_str}:  "
+                    f"{msg_prefix}: "
                     f"The mode solver pml in tangential axis '{i}' "
                     f"covers more than '{WARN_THICK_PML_PERCENT}%' of the "
                     "mode plane cells. Consider using a larger mode plane "
@@ -254,9 +257,7 @@ def _mode_plane(plane: Box, sim_geom: Box) -> Box:
         return Box.from_bounds(*mode_plane_bnds)
 
     @classmethod
-    def _validate_mode_plane_radius(
-        cls, mode_spec: ModeSpec, plane: Box, sim_geom: Box, msg_prefix: str = ""
-    ):
+    def _validate_mode_plane_radius(cls, mode_spec: ModeSpec, plane: Box, sim_geom: Box):
         """Validate that the radius of a mode spec with a bend is not smaller than half the size of
         the plane along the radial direction."""
 
@@ -271,21 +272,14 @@ def _validate_mode_plane_radius(
 
         if np.abs(mode_spec.bend_radius) < mode_plane.size[radial_ax] / 2:
             raise ValueError(
-                f"{msg_prefix} bend radius is smaller than half the mode plane size "
+                "Mode solver bend radius is smaller than half the mode plane size "
                 "along the radial axis, which can produce wrong results."
             )
 
     def _validate_rotate_structures(self) -> None:
         """Validate that structures can be rotated if angle_rotation is True."""
-        if not self.mode_spec.angle_rotation:
-            return
-        try:
+        if np.abs(self.mode_spec.angle_theta) > 0 and self.mode_spec.angle_rotation:
             _ = self._rotate_structures
-        except Exception as e:
-            raise SetupError(
-                "Mode object defined with 'angle_rotation=True' but failed "
-                f"to create rotated structures: {e!s}"
-            ) from e
 
     @cached_property
     def normal_axis(self) -> Axis:
@@ -628,28 +622,41 @@ def _rotate_structures(self) -> list[Structure]:
         translate_coords = [0, 0, 0]
         translate_coords[idx_u] = mnt_center[idx_u]
         translate_coords[idx_v] = mnt_center[idx_v]
+        translate_kwargs = dict(zip("xyz", translate_coords))
+        # Rotation arguments
+        rotate_kwargs = {"angle": theta, "axis": self.bend_axis_3d}
 
-        rotated_structures = []
-        for structure in Scene.intersecting_structures(self.plane, self.simulation.structures):
-            if not isinstance(structure.medium, IsotropicUniformMediumType):
-                raise NotImplementedError(
-                    "Mode solver plane intersects an unsupported "
-                    "medium. Only uniform isotropic media are supported for the plane rotation. "
-                )
+        structs_in = Scene.intersecting_structures(self.plane, self.simulation.structures)
+        return self._make_rotated_structures(structs_in, translate_kwargs, rotate_kwargs)
+
+    @staticmethod
+    def _make_rotated_structures(
+        structures: list[Structure], translate_kwargs: dict, rotate_kwargs: dict
+    ):
+        try:
+            rotated_structures = []
+            for structure in structures:
+                if not isinstance(structure.medium, get_args(IsotropicUniformMediumType)):
+                    raise NotImplementedError(
+                        "Mode solver plane intersects an unsupported medium. "
+                        "Only uniform isotropic media are supported for the plane rotation."
+                    )
 
-            # Rotate and apply translations
-            geometry = structure.geometry
-            geometry = (
-                geometry.translated(
-                    x=-translate_coords[0], y=-translate_coords[1], z=-translate_coords[2]
+                # Rotate and apply translations
+                geometry = structure.geometry
+                geometry = (
+                    geometry.translated(**{key: -val for key, val in translate_kwargs.items()})
+                    .rotated(**rotate_kwargs)
+                    .translated(**translate_kwargs)
                 )
-                .rotated(theta, axis=self.bend_axis_3d)
-                .translated(x=translate_coords[0], y=translate_coords[1], z=translate_coords[2])
-            )
 
-            rotated_structures.append(structure.updated_copy(geometry=geometry))
+                rotated_structures.append(structure.updated_copy(geometry=geometry))
 
-        return rotated_structures
+            return rotated_structures
+        except Exception as e:
+            raise SetupError(
+                f"'angle_rotation' set to True but could not rotate structures: {e!s}"
+            ) from e
 
     @cached_property
     def rotated_bend_center(self) -> list:

tidy3d/components/mode/simulation.py

@@ -215,12 +215,6 @@ def plane_in_sim_bounds(cls, val, values):
 
     def _post_init_validators(self) -> None:
         """Call validators taking `self` that get run after init."""
-        ModeSolver._validate_mode_plane_radius(
-            mode_spec=self.mode_spec,
-            plane=self.plane,
-            sim_geom=self.geometry,
-            msg_prefix="'ModeSimulation'",
-        )
         _ = self._mode_solver
         _ = self.grid
 

tidy3d/components/simulation.py

@@ -3710,15 +3710,34 @@ def _validate_mode_objects(self) -> None:
         """Create a ModeSolver for each mode object in order to validate."""
         from .mode.mode_solver import ModeSolver
 
+        def validate_mode_object(mode_obj: Union[ModeSource, AbstractModeMonitor], msg_prefix: str):
+            # Warn if pml is too thick
+            ModeSolver._warn_thick_pml(
+                simulation=self,
+                plane=mode_obj.geometry,
+                mode_spec=mode_obj.mode_spec,
+                msg_prefix=msg_prefix,
+            )
+            # Error if mode plane radius is too small
+            ModeSolver._validate_mode_plane_radius(
+                mode_spec=mode_obj.mode_spec,
+                plane=mode_obj.geometry,
+                sim_geom=self.geometry,
+            )
+            # Test if structures can be rotated if ``angle_rotation=True``
+            theta = mode_obj.mode_spec.angle_theta
+            if np.abs(theta) > 0 and mode_obj.mode_spec.angle_rotation:
+                structs_in = Scene.intersecting_structures(mode_obj.geometry, self.structures)
+                translate_kwargs = dict(zip("xyz", mode_obj.center))
+                _, axes = mode_obj.pop_axis([0, 1, 2], mode_obj.size.index(0.0))
+                # we just pick one of the in-plane axes to test the roation
+                rotate_kwargs = {"angle": theta, "axis": axes[1]}
+                ModeSolver._make_rotated_structures(structs_in, translate_kwargs, rotate_kwargs)
+
         for imnt, monitor in enumerate(self.monitors):
             if isinstance(monitor, AbstractModeMonitor):
                 try:
-                    _ = ModeSolver(
-                        mode_spec=monitor.mode_spec,
-                        plane=monitor.geometry,
-                        simulation=self,
-                        freqs=monitor.freqs,
-                    )
+                    validate_mode_object(mode_obj=monitor, msg_prefix=f"'monitors[{imnt}]'")
                 except Exception as e:
                     raise SetupError(
                         f"Monitor at 'monitors[{imnt}]' failed validation: {e!s}"
@@ -3727,12 +3746,7 @@ def _validate_mode_objects(self) -> None:
         for isrc, source in enumerate(self.sources):
             if isinstance(source, ModeSource):
                 try:
-                    _ = ModeSolver(
-                        mode_spec=source.mode_spec,
-                        plane=source.geometry,
-                        simulation=self,
-                        freqs=source.source_time.freq0,
-                    )
+                    validate_mode_object(mode_obj=source, msg_prefix=f"'sources[{isrc}]'")
                 except Exception as e:
                     raise SetupError(f"Source at 'sources[{isrc}]' failed validation: {e!s}") from e