everything importable for doctests

Author: yaugenst-flex

tidy3d/plugins/adjoint/components/__init__.py

@@ -4,8 +4,7 @@
 from .data.data_array import JaxDataArray
 from .data.dataset import JaxPermittivityDataset
 from .data.monitor_data import JaxModeData
-from .data.sim_data import JaxSimulationData
-from .geometry import JaxBox, JaxComplexPolySlab, JaxPolySlab
+from .geometry import JaxBox, JaxComplexPolySlab, JaxGeometryGroup, JaxPolySlab
 from .medium import JaxAnisotropicMedium, JaxCustomMedium, JaxMedium
 from .simulation import JaxSimulation
 from .structure import JaxStructure, JaxStructureStaticGeometry, JaxStructureStaticMedium
@@ -22,7 +21,6 @@
     "JaxStructureStaticMedium",
     "JaxStructureStaticGeometry",
     "JaxSimulation",
-    "JaxSimulationData",
     "JaxModeData",
     "JaxPermittivityDataset",
     "JaxDataArray",

tidy3d/plugins/adjoint/components/geometry.py

@@ -24,7 +24,7 @@
     PolySlab,
 )
 from tidy3d.components.monitor import FieldMonitor, PermittivityMonitor
-from tidy3d.components.types import ArrayFloat2D, Bound, Coordinate2D  # , annotate_type
+from tidy3d.components.types import ArrayFloat2D, Bound, Coordinate2D
 from tidy3d.constants import MICROMETER, fp_eps
 from tidy3d.exceptions import AdjointError
 from tidy3d.log import log

tidy3d/plugins/adjoint/components/medium.py

@@ -290,7 +290,7 @@ class JaxCustomMedium(CustomMedium, AbstractJaxMedium):
         jax_field=True,
     )
 
-    @model_validator(pre=True)
+    @model_validator(mode="before")
     def _pre_deprecation_dataset(data):
         """Don't allow permittivity as a field until we support it."""
         if data.get("permittivity") or data.get("conductivity"):

tidy3d/plugins/adjoint/components/simulation.py

@@ -31,7 +31,7 @@
 from tidy3d.components.simulation import Simulation
 from tidy3d.components.structure import Structure
 from tidy3d.components.subpixel_spec import Staircasing, SubpixelSpec
-from tidy3d.components.types import Ax, annotate_type
+from tidy3d.components.types import Ax, discriminated_union
 from tidy3d.constants import HERTZ, SECOND
 from tidy3d.exceptions import AdjointError
 from tidy3d.log import log
@@ -69,7 +69,7 @@
 )
 
 OutputMonitorTypes = (DiffractionMonitor, FieldMonitor, ModeMonitor)
-OutputMonitorType = tuple[annotate_type(Union[OutputMonitorTypes]), ...]
+OutputMonitorType = tuple[discriminated_union(Union[OutputMonitorTypes]), ...]
 
 
 class JaxInfo(Tidy3dBaseModel):
@@ -122,7 +122,7 @@ class JaxInfo(Tidy3dBaseModel):
 class JaxSimulation(Simulation, JaxObject):
     """A :class:`.Simulation` registered with jax."""
 
-    input_structures: tuple[annotate_type(JaxStructureType), ...] = Field(
+    input_structures: tuple[discriminated_union(JaxStructureType), ...] = Field(
         (),
         title="Input Structures",
         description="Tuple of jax-compatible structures"
@@ -164,7 +164,7 @@ class JaxSimulation(Simulation, JaxObject):
         units=SECOND,
     )
 
-    @field_validator("output_monitors", always=True)
+    @field_validator("output_monitors")
     def _output_monitors_colocate_false(val):
         """Make sure server-side colocation is off."""
         new_vals = []

tidy3d/plugins/adjoint/components/types.py

@@ -3,6 +3,7 @@
 from typing import Any, Union
 
 import numpy as np
+from pydantic_core import core_schema
 
 from tidy3d.components.type_util import _add_schema
 
@@ -29,16 +30,16 @@ class NumpyArrayType(np.ndarray):
     """Subclass of ``np.ndarray`` with a schema defined for pydantic."""
 
     @classmethod
-    # TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.
-    # Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
-    def __modify_schema__(cls, field_schema):
-        """Sets the schema of np.ndarray object."""
-
-        schema = dict(
-            title="npdarray",
-            type="numpy.ndarray",
-        )
-        field_schema.update(schema)
+    def __get_pydantic_core_schema__(cls, source, handler):
+        return core_schema.no_info_plain_validator_function(lambda v, _: np.asarray(v))
+
+    @classmethod
+    def __get_pydantic_json_schema__(cls, core_schema, handler):
+        return {
+            "title": "npdarray",
+            "type": "numpy.ndarray",
+            "items": {},
+        }
 
 
 _add_schema(JaxArrayType, title="JaxArray", field_type_str="jax.numpy.ndarray")

tidy3d/plugins/design/parameter.py

@@ -63,7 +63,7 @@ class ParameterNumeric(Parameter, ABC):
         description="(min, max) range within which are allowed values for the variable. Is inclusive of max value.",
     )
 
-    @field_validator("span", always=True)
+    @field_validator("span")
     def _span_valid(val):
         """Span min <= span max."""
         span_min, span_max = val

tidy3d/plugins/expressions/__init__.py

@@ -1,22 +1,44 @@
 from .base import Expression
 from .functions import Cos, Exp, Log, Log10, Sin, Sqrt, Tan
 from .metrics import ModeAmp, ModePower, generate_validation_data
+from .operators import (
+    Abs,
+    Add,
+    Divide,
+    FloorDivide,
+    MatMul,
+    Modulus,
+    Multiply,
+    Negate,
+    Power,
+    Subtract,
+)
 from .variables import Constant, Variable
 
 __all__ = [
-    "Expression",
+    "Abs",
+    "Add",
     "Constant",
-    "Variable",
-    "ModeAmp",
-    "ModePower",
-    "generate_validation_data",
-    "Sin",
     "Cos",
-    "Tan",
+    "Divide",
     "Exp",
+    "Expression",
+    "FloorDivide",
     "Log",
     "Log10",
+    "MatMul",
+    "ModeAmp",
+    "ModePower",
+    "Modulus",
+    "Multiply",
+    "Negate",
+    "Power",
+    "Sin",
     "Sqrt",
+    "Subtract",
+    "Tan",
+    "Variable",
+    "generate_validation_data",
 ]
 
 # The following code dynamically collects all classes that are subclasses of Expression
@@ -41,4 +63,4 @@
             _local_vars[name] = obj
 
 for cls in _model_classes:
-    cls.update_forward_refs(**_local_vars)
+    cls.model_rebuild(force=True)

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -219,7 +219,7 @@ def _validate_3d_simulation(cls, val):
             )
         return val
 
-    @model_validator("radiation_monitors")
+    @model_validator(mode="after")
     def _validate_radiation_monitors(self):
         freqs = set(self.freqs)
         for rad_mon in self.radiation_monitors:

tidy3d/plugins/smatrix/ports/coaxial_lumped.py

@@ -88,7 +88,7 @@ def _center_not_inf(val):
             raise ValidationError("'center' can not contain 'td.inf' terms.")
         return val
 
-    @model_validator("inner_diameter", always=True)
+    @model_validator(mode="after")
     def _ensure_inner_diameter_is_smaller(self):
         """Ensures that the inner diameter is smaller than the outer diameter, so that the final
         shape is an annulus."""