resolve circular dependencies

Author: frederikschubertflex

tidy3d/components/material/solver_types.py

@@ -10,7 +10,7 @@
     SemiconductorMedium,
 )
 from tidy3d.components.material.tcad.heat import ThermalSpecType
-from tidy3d.components.medium.types import MediumType, MediumType3D
+from tidy3d.components.medium import MediumType, MediumType3D
 
 OpticalMediumType = MediumType
 ElectricalMediumType = MediumType

tidy3d/components/material/tcad/charge.py

@@ -7,7 +7,7 @@
 import pydantic.v1 as pd
 
 from tidy3d.components.data.data_array import SpatialDataArray
-from tidy3d.components.medium.base import AbstractMedium
+from tidy3d.components.medium import AbstractMedium
 from tidy3d.components.tcad.doping import DopingBoxType
 from tidy3d.components.tcad.types import (
     BandGapNarrowingModelType,

tidy3d/components/medium/__init__.py

@@ -3,17 +3,15 @@
 
 from __future__ import annotations
 
-from typing import Union
-
 from .nonlinear import (
     KerrNonlinearity,
     NonlinearModel,
     NonlinearSpec,
     NonlinearSusceptibility,
     TwoPhotonAbsorption,
 )
-from .base import AbstractMedium
-from .dispersionless import Medium, PECMedium, CustomMedium, CustomIsotropicMedium
+from .base import AbstractMedium, AbstractCustomMedium
+from .dispersionless import Medium, PECMedium, CustomMedium, PEC
 from .anisotropic import (
     AnisotropicMedium,
     CustomAnisotropicMedium,
@@ -39,54 +37,17 @@
     LossyMetalMedium,
     SurfaceImpedanceFitterParam,
 )
-from .medium_2d import Medium2D
-from .perturbation import PerturbationMedium, PerturbationPoleResidue
-
-
-NonlinearModelType = Union[NonlinearSusceptibility, TwoPhotonAbsorption, KerrNonlinearity]
-IsotropicUniformMediumType = Union[
-    Medium, LossyMetalMedium, PoleResidue, Sellmeier, Lorentz, Debye, Drude, PECMedium
-]
-IsotropicCustomMediumType = Union[
-    CustomPoleResidue,
-    CustomSellmeier,
-    CustomLorentz,
-    CustomDebye,
-    CustomDrude,
-]
-IsotropicCustomMediumInternalType = Union[IsotropicCustomMediumType, CustomIsotropicMedium]
-IsotropicMediumType = Union[IsotropicCustomMediumType, IsotropicUniformMediumType]
-MediumType3D = Union[
-    Medium,
-    AnisotropicMedium,
-    PECMedium,
-    PoleResidue,
-    Sellmeier,
-    Lorentz,
-    Debye,
-    Drude,
-    FullyAnisotropicMedium,
-    CustomMedium,
-    CustomPoleResidue,
-    CustomSellmeier,
-    CustomLorentz,
-    CustomDebye,
-    CustomDrude,
-    CustomAnisotropicMedium,
-    PerturbationMedium,
-    PerturbationPoleResidue,
-    LossyMetalMedium,
-]
-MediumType = Union[MediumType3D, Medium2D, AnisotropicMediumFromMedium2D]
-
-PEC = PECMedium(name="PEC")
-PEC2D = Medium2D(ss=PEC, tt=PEC)
+from .medium_2d import Medium2D, MediumType, MediumType3D, PEC2D
+from .perturbation import PerturbationMedium, PerturbationPoleResidue, AbstractPerturbationMedium
 
+AnisotropicMedium.update_forward_refs()
 
 __all__ = [
     "PEC",
     "PEC2D",
     "AbstractMedium",
+    "AbstractCustomMedium",
+    "AbstractPerturbationMedium",
     "AnisotropicMedium",
     "AnisotropicMediumFromMedium2D",
     "CustomAnisotropicMedium",
@@ -106,6 +67,8 @@
     "LossyMetalMedium",
     "Medium",
     "Medium2D",
+    "MediumType",
+    "MediumType3D",
     "NonlinearModel",
     "NonlinearSpec",
     "NonlinearSusceptibility",

tidy3d/components/medium/anisotropic.py

@@ -24,6 +24,11 @@
 
 from .base import AbstractCustomMedium, AbstractMedium
 from .dispersionless import CustomMedium, Medium, PECMedium
+from .types import (
+    IsotropicCustomMediumInternalType,
+    IsotropicCustomMediumType,
+    IsotropicUniformMediumType,
+)
 from .utils import ensure_freq_in_range
 
 

tidy3d/components/medium/base.py

@@ -38,6 +38,7 @@
 from tidy3d.exceptions import ValidationError
 from tidy3d.log import log
 
+from .nonlinear import NonlinearModel, NonlinearSpec, NonlinearSusceptibility
 from .utils import ensure_freq_in_range
 
 
@@ -223,6 +224,8 @@ def is_custom(self) -> bool:
     @cached_property
     def is_fully_anisotropic(self) -> bool:
         """Whether the medium is fully anisotropic."""
+        from .anisotropic import FullyAnisotropicMedium
+
         return isinstance(self, FullyAnisotropicMedium)
 
     @cached_property

tidy3d/components/medium/dispersionless.py

@@ -752,6 +752,8 @@ def _medium(self):
         """Internal representation in the form of
         either `CustomIsotropicMedium` or `CustomAnisotropicMedium`.
         """
+        from .anisotropic import CustomAnisotropicMediumInternal
+
         self_dict = self.dict(exclude={"type", "eps_dataset"})
         # isotropic
         if self.eps_dataset is None:
@@ -1235,3 +1237,7 @@ def _derivative_field_cmp(
         vjp_array = vjp_array.reshape(eps_data.shape)
 
         return vjp_array
+
+
+PECMedium.update_forward_refs()
+PEC = PECMedium(name="PEC")

tidy3d/components/medium/medium_2d.py

@@ -12,13 +12,54 @@
 from tidy3d.exceptions import ValidationError
 from tidy3d.log import log
 
-from .anisotropic import AnisotropicMedium, AnisotropicMediumFromMedium2D
+from .anisotropic import (
+    AnisotropicMedium,
+    AnisotropicMediumFromMedium2D,
+    CustomAnisotropicMedium,
+    FullyAnisotropicMedium,
+)
 from .base import AbstractMedium
-from .dispersionless import Medium, PECMedium
-from .dispersive import PoleResidue
+from .dispersionless import PEC, CustomMedium, Medium, PECMedium
+from .dispersive import (
+    CustomDebye,
+    CustomDrude,
+    CustomLorentz,
+    CustomPoleResidue,
+    CustomSellmeier,
+    Debye,
+    Drude,
+    Lorentz,
+    PoleResidue,
+    Sellmeier,
+)
 from .dispersive_abc import DispersiveMedium
+from .lossy_metal import LossyMetalMedium
+from .perturbation import PerturbationMedium, PerturbationPoleResidue
+from .types import IsotropicUniformMediumType
 from .utils import ensure_freq_in_range
 
+MediumType3D = Union[
+    Medium,
+    AnisotropicMedium,
+    PECMedium,
+    PoleResidue,
+    Sellmeier,
+    Lorentz,
+    Debye,
+    Drude,
+    FullyAnisotropicMedium,
+    CustomMedium,
+    CustomPoleResidue,
+    CustomSellmeier,
+    CustomLorentz,
+    CustomDebye,
+    CustomDrude,
+    CustomAnisotropicMedium,
+    PerturbationMedium,
+    PerturbationPoleResidue,
+    LossyMetalMedium,
+]
+
 
 class Medium2D(AbstractMedium):
     """2D diagonally anisotropic medium.
@@ -421,3 +462,7 @@ def is_comp_pec_2d(self, comp: Axis, axis: Axis):
             ax_coord=Medium(), plane_coords=self.elements.values(), axis=axis
         )
         return isinstance(elements_3d[comp], PECMedium)
+
+
+MediumType = Union[MediumType3D, Medium2D, AnisotropicMediumFromMedium2D]
+PEC2D = Medium2D(ss=PEC, tt=PEC)

tidy3d/components/medium/nonlinear.py

@@ -13,11 +13,8 @@
 from tidy3d.log import log
 
 from .constants import NONLINEAR_DEFAULT_NUM_ITERS, NONLINEAR_MAX_NUM_ITERS
-from .dispersionless import Medium
-from .dispersive import DispersiveMedium
 
 if TYPE_CHECKING:
-    from . import NonlinearModelType
     from .base import AbstractMedium
 
 
@@ -27,6 +24,8 @@ class NonlinearModel(ABC, Tidy3dBaseModel):
 
     def _validate_medium_type(self, medium: AbstractMedium):
         from .base import AbstractCustomMedium
+        from .dispersionless import Medium
+        from .dispersive import DispersiveMedium
 
         """Check that the model is compatible with the medium."""
         if isinstance(medium, AbstractCustomMedium):
@@ -516,6 +515,9 @@ def complex_fields(self) -> bool:
         return self.use_complex_fields
 
 
+NonlinearModelType = Union[NonlinearSusceptibility, TwoPhotonAbsorption, KerrNonlinearity]
+
+
 class NonlinearSpec(ABC, Tidy3dBaseModel):
     """Abstract specification for adding nonlinearities to a medium.
 

tidy3d/components/medium/perturbation.py

@@ -20,8 +20,9 @@
 from tidy3d.exceptions import SetupError
 
 from .base import AbstractCustomMedium, AbstractMedium
-from .dispersionless import Medium
-from .dispersive import PoleResidue
+from .dispersionless import CustomMedium, Medium
+from .dispersive import CustomPoleResidue, PoleResidue
+from .dispersive_abc import DispersiveMedium
 
 
 class AbstractPerturbationMedium(ABC, Tidy3dBaseModel):

tidy3d/components/medium/types.py

@@ -1,6 +1,20 @@
 from typing import Union
 
-NonlinearModelType = Union[NonlinearSusceptibility, TwoPhotonAbsorption, KerrNonlinearity]
+from .dispersionless import CustomIsotropicMedium, Medium, PECMedium
+from .dispersive import (
+    CustomDebye,
+    CustomDrude,
+    CustomLorentz,
+    CustomPoleResidue,
+    CustomSellmeier,
+    Debye,
+    Drude,
+    Lorentz,
+    PoleResidue,
+    Sellmeier,
+)
+from .lossy_metal import LossyMetalMedium
+
 IsotropicUniformMediumType = Union[
     Medium, LossyMetalMedium, PoleResidue, Sellmeier, Lorentz, Debye, Drude, PECMedium
 ]
@@ -13,25 +27,3 @@
 ]
 IsotropicCustomMediumInternalType = Union[IsotropicCustomMediumType, CustomIsotropicMedium]
 IsotropicMediumType = Union[IsotropicCustomMediumType, IsotropicUniformMediumType]
-MediumType3D = Union[
-    Medium,
-    AnisotropicMedium,
-    PECMedium,
-    PoleResidue,
-    Sellmeier,
-    Lorentz,
-    Debye,
-    Drude,
-    FullyAnisotropicMedium,
-    CustomMedium,
-    CustomPoleResidue,
-    CustomSellmeier,
-    CustomLorentz,
-    CustomDebye,
-    CustomDrude,
-    CustomAnisotropicMedium,
-    PerturbationMedium,
-    PerturbationPoleResidue,
-    LossyMetalMedium,
-]
-MediumType = Union[MediumType3D, Medium2D, AnisotropicMediumFromMedium2D]