fix: rebase changes for v2

Author: mahlau-flex

tests/test_package/test_log.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 import pytest
-from pydantic import ValidationError
+from pydantic import ValidationError, model_validator
 
 import tidy3d as td
 from tidy3d.exceptions import Tidy3dError
@@ -255,6 +255,35 @@ def test_logging_warning_capture():
         assert str(error_without) == str(error_with)
 
 
+def test_warning_capture_during_model_validation():
+    from tidy3d.components.base import Tidy3dBaseModel
+    from tidy3d.log import log
+
+    class _CaptureChild(Tidy3dBaseModel):
+        x: int
+
+        @model_validator(mode="after")
+        def _warn_child(self):
+            log.warning("child warning")
+            return self
+
+    class _CaptureParent(Tidy3dBaseModel):
+        child: _CaptureChild
+
+        @model_validator(mode="after")
+        def _warn_parent(self):
+            log.warning("parent warning")
+            return self
+
+    td.log.set_capture(True)
+    _CaptureParent(child={"x": 1})
+    warning_list = td.log.captured_warnings()
+    td.log.set_capture(False)
+
+    assert {"loc": [], "msg": "parent warning"} in warning_list
+    assert {"loc": ["child"], "msg": "child warning"} in warning_list
+
+
 def test_log_suppression():
     with td.log as suppressed_log:
         assert td.log._counts is not None

tidy3d/components/autograd/functions.py

@@ -262,7 +262,7 @@ def add_at(x: NDArray, indices_x: tuple, y: NDArray) -> NDArray:
 
 
 @primitive
-def _straight_through_clip(x, a_min, a_max):
+def _straight_through_clip(x: NDArray, a_min: Any, a_max: Any) -> NDArray:
     """Passthrough clip can be used to preserve gradients at the endpoints of the clip range where
     there is a discontinuity in the derivative. This is useful when values are at the endpoints but may
     have a gradient away from the boundary or in cases where numerical precision causes a function that is
@@ -271,7 +271,7 @@ def _straight_through_clip(x, a_min, a_max):
     return anp.clip(x, a_min=a_min, a_max=a_max)
 
 
-def _straight_through_clip_vjp(ans, x, a_min, a_max):
+def _straight_through_clip_vjp(ans: Any, x: NDArray, a_min: Any, a_max: Any) -> NDArray:
     """Preserve original gradient information in the backward pass up until a tolerance beyond the clip bounds."""
     tolerance = 1e-5
     mask = (x >= a_min - tolerance) & (x <= a_max + tolerance)

tidy3d/components/base.py

@@ -16,15 +16,17 @@
 from os import PathLike
 from pathlib import Path
 from types import UnionType
-from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union, get_args, get_origin, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union, get_args, get_origin
 
 import h5py
 import numpy as np
 import rich
 import xarray as xr
 import yaml
+from autograd.numpy.numpy_boxes import ArrayBox
 from autograd.tracer import isbox
-from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator
+from pydantic.functional_validators import ModelWrapValidatorHandler
 
 from tidy3d.exceptions import FileError
 from tidy3d.log import log
@@ -155,8 +157,8 @@ def _fmt_ann_literal(ann: Any) -> str:
 T = TypeVar("T", bound="Tidy3dBaseModel")
 
 
-def field_allows_scalar(field: ModelField) -> bool:
-    annotation = field.outer_type_
+def field_allows_scalar(field: FieldInfo) -> bool:
+    annotation = field.annotation
 
     def allows_scalar(a: Any) -> bool:
         origin = get_origin(a)
@@ -169,6 +171,7 @@ def allows_scalar(a: Any) -> bool:
 
     return allows_scalar(annotation)
 
+
 @total_ordering
 class Tidy3dBaseModel(BaseModel):
     """Base pydantic model that all Tidy3d components inherit from.
@@ -216,12 +219,6 @@ def _validate_name_no_special_characters(cls: type[T], name: Optional[str]) -> O
                 )
         return name
 
-    def __init__(self, **kwargs: Any) -> None:
-        """Init method, includes post-init validators."""
-        log.begin_capture()
-        super().__init__(**kwargs)
-        log.end_capture(self)
-
     def __init_subclass__(cls: type[T], **kwargs: Any) -> None:
         """Injects a constant discriminator field before Pydantic builds the model.
 
@@ -239,6 +236,16 @@ def __init_subclass__(cls: type[T], **kwargs: Any) -> None:
         setattr(cls, TYPE_TAG_STR, tag)
         TYPE_TO_CLASS_MAP[tag] = cls
 
+        if "__tidy3d_end_capture__" not in cls.__dict__:
+
+            @model_validator(mode="after")
+            def __tidy3d_end_capture__(self: T) -> T:
+                if log._capture:
+                    log.end_capture(self)
+                return self
+
+            cls.__tidy3d_end_capture__ = __tidy3d_end_capture__
+
         super().__init_subclass__(**kwargs)
 
     @classmethod
@@ -248,33 +255,95 @@ def __pydantic_init_subclass__(cls: type[T], **kwargs: Any) -> None:
         # add docstring once pydantic is done constructing the class
         cls.__doc__ = cls.generate_docstring()
 
+    @model_validator(mode="wrap")
+    @classmethod
+    def _capture_validation_warnings(
+        cls: type[T],
+        data: Any,
+        handler: ModelWrapValidatorHandler[T],
+    ) -> T:
+        if not log._capture:
+            return handler(data)
+
+        log.begin_capture()
+        try:
+            return handler(data)
+        except Exception:
+            log.abort_capture()
+            raise
+
     def __hash__(self) -> int:
         """Hash method."""
-        try:
-            return super().__hash__(self)
-        except TypeError:
-            return hash(self.model_dump_json())
+        return self._recursive_hash(self)
+
+    @staticmethod
+    def _recursive_hash(value: Any) -> int:
+        # Handle Autograd ArrayBoxes
+        if isinstance(value, ArrayBox):
+            # Unwrap the underlying numpy array and recurse
+            return Tidy3dBaseModel._recursive_hash(value._value)
+        if isinstance(value, np.ndarray):
+            # numpy arrays are not hashable by default, use byte representation
+            v_hash = hashlib.md5(value.tobytes()).hexdigest()
+            return hash(v_hash)
+        if isinstance(value, (xr.DataArray, xr.Dataset)):
+            # we choose to not hash data arrays as this would require a lot of careful handling of units, metadata.
+            # technically this is incorrect, but should never lead to bugs in current implementation
+            return hash(str(value.__class__.__name__))
+        if isinstance(value, str):
+            # this if-case is necessary because length-1 string would lead to infinite recursion in sequence case below
+            return hash(value)
+        if isinstance(value, Sequence):
+            # this assumes all objects in lists are hashable by default and do not require special handling
+            v_hash = tuple([Tidy3dBaseModel._recursive_hash(vi) for vi in value])
+            return hash(v_hash)
+        if isinstance(value, dict):
+            to_hash_list = []
+            for k, v in value.items():
+                v_hash = Tidy3dBaseModel._recursive_hash(v)
+                to_hash_list.append((k, v_hash))
+            return hash(tuple(to_hash_list))
+        if isinstance(value, Tidy3dBaseModel):
+            # This function needs to take special care because of mutable attributes inside of frozen pydantic models
+            to_hash_list = []
+            for k, v in dict(value).items():
+                if k == "attrs":
+                    continue
+                v_hash = Tidy3dBaseModel._recursive_hash(v)
+                to_hash_list.append((k, v_hash))
+            # attrs is mutable, use serialized output as safe hashing option
+            if value.attrs:
+                attrs_str = value._attrs_digest()
+                attrs_hash = hash(attrs_str)
+                to_hash_list.append(("attrs", attrs_hash))
+            return hash(tuple(to_hash_list))
+        return hash(value)
 
     def _hash_self(self) -> str:
         """Hash this component with ``hashlib`` in a way that is the same every session."""
         bf = io.BytesIO()
         self.to_hdf5(bf)
         return hashlib.md5(bf.getvalue()).hexdigest()
 
-
-    @pydantic.validator("*", pre=True, allow_reuse=True)
-    def coerce_numpy_scalars_for_model(cls, v: Any, field: ModelField) -> Any:
+    @model_validator(mode="before")
+    @classmethod
+    def coerce_numpy_scalars_for_model(cls, data: Any) -> Any:
         """
-        Wildcard field validator: coerce numpy scalars / size-1 arrays to native Python
+        coerce numpy scalars / size-1 arrays to native Python
         scalars, but only for fields whose annotations allow scalars.
         """
-        if not field_allows_scalar(field):
-            return v
+        if not isinstance(data, dict):
+            return data
+
+        for name, field in cls.model_fields.items():
+            if name not in data or not field_allows_scalar(field):
+                continue
 
-        if isinstance(v, np.generic) or (isinstance(v, np.ndarray) and v.size == 1):
-            return v.item()
+            v = data[name]
+            if isinstance(v, np.generic) or (isinstance(v, np.ndarray) and v.size == 1):
+                data[name] = v.item()
 
-        return v
+        return data
 
     @classmethod
     def _get_type_value(cls, obj: dict[str, Any]) -> str:

tidy3d/components/medium.py

@@ -11,6 +11,15 @@
 import autograd.numpy as np
 import numpy as npo
 from autograd.differential_operators import tensor_jacobian_product
+from numpy.typing import NDArray
+from pydantic import (
+    Field,
+    NonNegativeFloat,
+    PositiveFloat,
+    PositiveInt,
+    field_validator,
+    model_validator,
+)
 
 from tidy3d.components.autograd.utils import pack_complex_vec
 from tidy3d.constants import (
@@ -2419,11 +2428,11 @@ def _derivative_field_cmp_custom(
         E_der_map: ElectromagneticFieldDataset,
         spatial_data: SpatialDataArray,
         dim: str,
-        freqs: NDArray,
+        freqs: ArrayFloat,
         bounds: Optional[Bound] = None,
         component: str = "real",
         interp_method: Optional[InterpMethod] = None,
-    ) -> NDArray:
+    ) -> ArrayGeneric:
         """Compute the derivative with respect to a material property component."""
         param_coords = {axis: np.asarray(spatial_data.coords[axis]) for axis in "xyz"}
         eps_shape = [len(param_coords[axis]) for axis in "xyz"]
@@ -2609,7 +2618,18 @@ def _interp_axis(
         elif component == "real":
             values = values.real
 
-        return values.sum(axis=-1).reshape(eps_shape)
+        vjp_array = values.sum(axis=-1).reshape(eps_shape)
+
+        # match derivative dtype to the underlying dataset
+        target_array = getattr(spatial_data, "values", None)
+        if target_array is None and hasattr(spatial_data, "data"):
+            target_array = spatial_data.data
+        if target_array is not None:
+            target_dtype = np.asarray(target_array).dtype
+            if not np.issubdtype(target_dtype, np.complexfloating):
+                vjp_array = np.real(vjp_array).astype(target_dtype, copy=False)
+
+        return vjp_array
 
 
 """ Dispersive Media """

tidy3d/components/simulation.py

@@ -4258,19 +4258,6 @@ def _validate_low_freq_smoothing(self) -> Self:
                 )
         return self
 
-    @model_validator(mode="after")
-    def _validate_scene(self) -> Self:
-        _ = self.scene
-        self._validate_no_structures_pml()
-        self._validate_tfsf_nonuniform_grid()
-        self._validate_tfsf_aux_sources()
-        self._validate_nonlinear_specs()
-        self._validate_custom_source_time()
-        self._validate_mode_objects()
-        self._warn_rf_license()
-        self._validate_internal_abc_no_fully_anisotropic()
-        return self
-
     @model_validator(mode="after")
     def _warn_rf_license(self) -> Self:
         """
@@ -4545,16 +4532,18 @@ def aux_fields(self) -> list[str]:
                 fields += medium.nonlinear_spec.aux_fields
         return fields
 
-    def _validate_internal_abc_no_fully_anisotropic(self) -> None:
+    @model_validator(mode="after")
+    def _validate_internal_abc_no_fully_anisotropic(self) -> Self:
         """Error if internal absorber intersect fully anisotropic mediums."""
 
         total_structures = [self.scene.background_structure, *list(self.structures)]
 
         for abc in self._shifted_internal_absorbers:
-            mediums = Scene.intersecting_media(abc, total_structures)
+            mediums = Scene.intersecting_media(abc, tuple(total_structures))
 
             if any(isinstance(med, FullyAnisotropicMedium) for med in mediums):
                 raise SetupError("A 'InternalAbsorber' cannot cross a 'FullyAnisotropicMedium'.")
+        return self
 
     """ Pre submit validation (before web.upload()) """
 

tidy3d/log.py

@@ -122,8 +122,9 @@ class Logger:
     highest level of the captures messages.
 
     Messages can also be captured for post-processing. That can be enabled through 'set_capture' to
-    record all warnings emitted during model validation. A structured copy of all validation
-    messages can then be recovered through 'captured_warnings'.
+    record warnings emitted during model validation (and other explicit begin/end capture regions,
+    e.g. validation routines like ``validate_pre_upload``). A structured copy of captured warnings
+    can then be recovered through 'captured_warnings'.
     """
 
     _static_cache = set()
@@ -178,8 +179,8 @@ def __exit__(
     def begin_capture(self) -> None:
         """Start capturing log stack for consolidated validation log.
 
-        This method is used before any model validation starts and is included in the initialization
-        of 'BaseModel'. It must be followed by a corresponding 'end_capture'.
+        This method should be called before a validation routine starts. It must be followed by a
+        corresponding 'end_capture'.
         """
         if not self._capture:
             return
@@ -190,11 +191,23 @@ def begin_capture(self) -> None:
         else:
             self._stack = [stack_item]
 
+    def abort_capture(self) -> None:
+        """Undo the last ``begin_capture()`` call.
+
+        This is used when validation fails before reaching the corresponding ``end_capture()``.
+        """
+        if not self._stack:
+            return
+
+        self._stack.pop()
+        if len(self._stack) == 0:
+            self._stack = None
+
     def end_capture(self, model: BaseModel) -> None:
         """End capturing log stack for consolidated validation log.
 
-        This method is used after all model validations and is included in the initialization of
-        'BaseModel'. It must follow a corresponding 'begin_capture'.
+        This method should be called after a validation routine ends. It must follow a
+        corresponding 'begin_capture'.
         """
         if not self._stack:
             return

tidy3d/plugins/expressions/metrics.py

@@ -32,7 +32,7 @@ def generate_validation_data(expr: Expression) -> dict[str, xr.Dataset]:
     dict[str, xr.Dataset]
         The combined validation data.
     """
-    metrics = set(expr.filter(target_type=Metric))
+    metrics = set(expr.filter(target_type=Metric))  # type: ignore[type-abstract]
     combined_data = {k: v for metric in metrics for k, v in metric._validation_data.items()}
     return combined_data
 

tidy3d/plugins/invdes/region.py

@@ -181,7 +181,7 @@ class TopologyDesignRegion(DesignRegion):
         "Supplying ``False`` will completely leave out the override structure.",
     )
 
-    priority: int = pd.Field(
+    priority: Optional[int] = Field(
         None,
         title="Priority",
         description="Priority of the structure applied in structure overlapping region. "