Pydantic v2

pyproject.toml

@@ -23,6 +23,7 @@ documentation = "https://docs.flexcompute.com/projects/tidy3d/en/latest/"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<3.14"
+typing-extensions = { version = "*", python = "<3.11" }
 pyroots = ">=0.5.0"
 xarray = ">=2023.08"
 importlib-metadata = ">=6.0.0"

run.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env -S poetry run python
+# ruff: noqa: F401
+
+# from tidy3d.components.base import Tidy3dBaseModel
+import numpy as np
+
+from tidy3d.components.base import Tidy3dBaseModel
+
+
+class Test(Tidy3dBaseModel):
+    a: np.ndarray
+    x: str = "test"
+
+
+m = Test(a=np.zeros(1))
+print(m.model_dump())

tidy3d/compat.py

@@ -5,4 +5,9 @@
 except ImportError:
     from xarray.core import alignment
 
-__all__ = ["alignment"]
+try:
+    from typing import Self  # Python >= 3.11
+except ImportError:  # Python <3.11
+    from typing_extensions import Self
+
+__all__ = ["alignment", "Self"]

tidy3d/components/apodization.py

@@ -1,11 +1,14 @@
 """Defines specification for apodization."""
 
+from typing import Optional
+
 import numpy as np
-import pydantic.v1 as pd
+from pydantic import Field, NonNegativeFloat, PositiveFloat, model_validator
 
+from ..compat import Self
 from ..constants import SECOND
 from ..exceptions import SetupError
-from .base import Tidy3dBaseModel, skip_if_fields_missing
+from .base import Tidy3dBaseModel
 from .types import ArrayFloat1D, Ax
 from .viz import add_ax_if_none
 
@@ -24,45 +27,40 @@ class ApodizationSpec(Tidy3dBaseModel):
 
     """
 
-    start: pd.NonNegativeFloat = pd.Field(
+    start: Optional[NonNegativeFloat] = Field(
         None,
         title="Start Interval",
         description="Defines the time at which the start apodization ends.",
         units=SECOND,
     )
 
-    end: pd.NonNegativeFloat = pd.Field(
+    end: Optional[NonNegativeFloat] = Field(
         None,
         title="End Interval",
         description="Defines the time at which the end apodization begins.",
         units=SECOND,
     )
 
-    width: pd.PositiveFloat = pd.Field(
+    width: Optional[PositiveFloat] = Field(
         None,
         title="Apodization Width",
         description="Characteristic decay length of the apodization function, i.e., the width of the ramping up of the scaling function from 0 to 1.",
         units=SECOND,
     )
 
-    @pd.validator("end", always=True, allow_reuse=True)
-    @skip_if_fields_missing(["start"])
-    def end_greater_than_start(cls, val, values):
+    @model_validator(mode="after")
+    def end_greater_than_start(self) -> Self:
         """Ensure end is greater than or equal to start."""
-        start = values.get("start")
-        if val is not None and start is not None and val < start:
+        if self.end is not None and self.start is not None and self.end < self.start:
             raise SetupError("End apodization begins before start apodization ends.")
-        return val
+        return self
 
-    @pd.validator("width", always=True, allow_reuse=True)
-    @skip_if_fields_missing(["start", "end"])
-    def width_provided(cls, val, values):
+    @model_validator(mode="after")
+    def width_provided(self) -> Self:
         """Check that width is provided if either start or end apodization is requested."""
-        start = values.get("start")
-        end = values.get("end")
-        if (start is not None or end is not None) and val is None:
+        if (self.start is not None or self.end is not None) and self.val is None:
             raise SetupError("Apodization width must be set.")
-        return val
+        return self
 
     @add_ax_if_none
     def plot(self, times: ArrayFloat1D, ax: Ax = None) -> Ax:

tidy3d/components/autograd/boxes.py

@@ -2,7 +2,7 @@
 # NOTE: we do not subclass ArrayBox since that would break autograd's internal checks
 
 import importlib
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable
 
 import autograd.numpy as anp
 from autograd.extend import VJPNode, defjvp, register_notrace
@@ -33,9 +33,9 @@ def from_arraybox(cls, box: ArrayBox) -> TidyArrayBox:
 def __array_function__(
     self: Any,
     func: Callable,
-    types: List[Any],
-    args: Tuple[Any, ...],
-    kwargs: Dict[str, Any],
+    types: list[Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
 ) -> Any:
     """
     Handle the dispatch of NumPy functions to autograd's numpy implementation.
@@ -46,11 +46,11 @@ def __array_function__(
         The instance of the class.
     func : Callable
         The NumPy function being called.
-    types : List[Any]
+    types : list[Any]
         The types of the arguments that implement __array_function__.
-    args : Tuple[Any, ...]
+    args : tuple[Any, ...]
         The positional arguments to the function.
-    kwargs : Dict[str, Any]
+    kwargs : dict[str, Any]
         The keyword arguments to the function.
 
     Returns
@@ -102,7 +102,7 @@ def __array_ufunc__(
     ufunc: Callable,
     method: str,
     *inputs: Any,
-    **kwargs: Dict[str, Any],
+    **kwargs: dict[str, Any],
 ) -> Any:
     """
     Handle the dispatch of NumPy ufuncs to autograd's numpy implementation.
@@ -117,7 +117,7 @@ def __array_ufunc__(
         The method of the ufunc being called.
     inputs : Any
         The input arguments to the ufunc.
-    kwargs : Dict[str, Any]
+    kwargs : dict[str, Any]
         The keyword arguments to the ufunc.
 
     Returns

tidy3d/components/autograd/derivative_utils.py

@@ -1,10 +1,12 @@
 # utilities for autograd derivative passing
-from __future__ import annotations
+
+from typing import Optional
 
 import numpy as np
-import pydantic.v1 as pd
 import xarray as xr
+from pydantic import Field
 
+from ...compat import Self
 from ...constants import LARGE_NUMBER
 from ..base import Tidy3dBaseModel
 from ..data.data_array import ScalarFieldDataArray, SpatialDataArray
@@ -43,32 +45,27 @@ class DerivativeSurfaceMesh(Tidy3dBaseModel):
 
     """
 
-    centers: ArrayLike = pd.Field(
-        ...,
+    centers: ArrayLike = Field(
         title="Centers",
         description="(N, 3) array storing the centers of each surface element.",
     )
 
-    areas: ArrayLike = pd.Field(
-        ...,
+    areas: ArrayLike = Field(
         title="Area Elements",
         description="(N,) array storing the first perpendicular vectors of each surface element.",
     )
 
-    normals: ArrayLike = pd.Field(
-        ...,
+    normals: ArrayLike = Field(
         title="Normals",
         description="(N, 3) array storing the normal vectors of each surface element.",
     )
 
-    perps1: ArrayLike = pd.Field(
-        ...,
+    perps1: ArrayLike = Field(
         title="Perpendiculars 1",
         description="(N, 3) array storing the first perpendicular vectors of each surface element.",
     )
 
-    perps2: ArrayLike = pd.Field(
-        ...,
+    perps2: ArrayLike = Field(
         title="Perpendiculars 1",
         description="(N, 3) array storing the first perpendicular vectors of each surface element.",
     )
@@ -77,122 +74,110 @@ class DerivativeSurfaceMesh(Tidy3dBaseModel):
 class DerivativeInfo(Tidy3dBaseModel):
     """Stores derivative information passed to the ``.compute_derivatives`` methods."""
 
-    paths: list[PathType] = pd.Field(
-        ...,
+    paths: list[PathType] = Field(
         title="Paths to Traced Fields",
         description="List of paths to the traced fields that need derivatives calculated.",
     )
 
-    E_der_map: FieldData = pd.Field(
-        ...,
+    E_der_map: FieldData = Field(
         title="Electric Field Gradient Map",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` store the '
         "multiplication of the forward and adjoint electric fields. The tangential components "
         "of this dataset is used when computing adjoint gradients for shifting boundaries. "
         "All components are used when computing volume-based gradients.",
     )
 
-    D_der_map: FieldData = pd.Field(
-        ...,
+    D_der_map: FieldData = Field(
         title="Displacement Field Gradient Map",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` store the '
         "multiplication of the forward and adjoint displacement fields. The normal component "
         "of this dataset is used when computing adjoint gradients for shifting boundaries.",
     )
 
-    E_fwd: FieldData = pd.Field(
-        ...,
+    E_fwd: FieldData = Field(
         title="Forward Electric Fields",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the '
         "forward electric fields used for computing gradients for a given structure.",
     )
 
-    E_adj: FieldData = pd.Field(
-        ...,
+    E_adj: FieldData = Field(
         title="Adjoint Electric Fields",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the '
         "adjoint electric fields used for computing gradients for a given structure.",
     )
 
-    D_fwd: FieldData = pd.Field(
-        ...,
+    D_fwd: FieldData = Field(
         title="Forward Displacement Fields",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the '
         "forward displacement fields used for computing gradients for a given structure.",
     )
 
-    D_adj: FieldData = pd.Field(
-        ...,
+    D_adj: FieldData = Field(
         title="Adjoint Displacement Fields",
         description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the '
         "adjoint displacement fields used for computing gradients for a given structure.",
     )
 
-    eps_data: PermittivityData = pd.Field(
-        ...,
+    eps_data: PermittivityData = Field(
         title="Permittivity Dataset",
         description="Dataset of relative permittivity values along all three dimensions. "
         "Used for automatically computing permittivity inside or outside of a simple geometry.",
     )
 
-    eps_in: tidycomplex = pd.Field(
+    eps_in: tidycomplex = Field(
         title="Permittivity Inside",
         description="Permittivity inside of the ``Structure``. "
         "Typically computed from ``Structure.medium.eps_model``."
         "Used when it can not be computed from ``eps_data`` or when ``eps_approx==True``.",
     )
 
-    eps_out: tidycomplex = pd.Field(
-        ...,
+    eps_out: tidycomplex = Field(
         title="Permittivity Outside",
         description="Permittivity outside of the ``Structure``. "
         "Typically computed from ``Simulation.medium.eps_model``."
         "Used when it can not be computed from ``eps_data`` or when ``eps_approx==True``.",
     )
 
-    eps_background: tidycomplex = pd.Field(
+    eps_background: tidycomplex = Field(
         None,
         title="Permittivity in Background",
         description="Permittivity outside of the ``Structure`` as manually specified by. "
         "``Structure.background_medium``. ",
     )
 
-    bounds: Bound = pd.Field(
-        ...,
+    bounds: Bound = Field(
         title="Geometry Bounds",
         description="Bounds corresponding to the structure, used in ``Medium`` calculations.",
     )
 
-    bounds_intersect: Bound = pd.Field(
-        ...,
+    bounds_intersect: Bound = Field(
         title="Geometry and Simulation Intersections Bounds",
         description="Bounds corresponding to the minimum intersection between the "
         "structure and the simulation it is contained in.",
     )
 
-    frequency: float = pd.Field(
-        ...,
+    frequency: float = Field(
         title="Frequency of adjoint simulation",
         description="Frequency at which the adjoint gradient is computed.",
     )
 
-    eps_no_structure: SpatialDataArray = pd.Field(
+    eps_no_structure: Optional[SpatialDataArray] = Field(
         None,
         title="Permittivity Without Structure",
         description="The permittivity of the original simulation without the structure that is "
         "being differentiated with respect to. Used to approximate permittivity outside of the "
         "structure for shape optimization.",
     )
 
-    eps_inf_structure: SpatialDataArray = pd.Field(
+    eps_inf_structure: Optional[SpatialDataArray] = Field(
         None,
         title="Permittivity With Infinite Structure",
         description="The permittivity of the original simulation where the structure being "
         " differentiated with respect to is inifinitely large. Used to approximate permittivity "
         "inside of the structure for shape optimization.",
     )
 
-    eps_approx: bool = pd.Field(
+    eps_approx: bool = Field(
         False,
         title="Use Permittivity Approximation",
         description="If ``True``, approximates outside permittivity using ``Simulation.medium``"
@@ -201,7 +186,7 @@ class DerivativeInfo(Tidy3dBaseModel):
         "evaluate the inside and outside relative permittivity for each geometry.",
     )
 
-    def updated_paths(self, paths: list[PathType]) -> DerivativeInfo:
+    def updated_paths(self, paths: list[PathType]) -> Self:
         """Update this ``DerivativeInfo`` with new set of paths."""
         return self.updated_copy(paths=paths)