diff --git a/tests/test_components/autograd/test_autograd.py b/tests/test_components/autograd/test_autograd.py index cfd55079b0..64fb7e5445 100644 --- a/tests/test_components/autograd/test_autograd.py +++ b/tests/test_components/autograd/test_autograd.py @@ -25,12 +25,12 @@ MINIMUM_SPACING_FRACTION, ) from tidy3d.components.autograd.derivative_utils import DerivativeInfo +from tidy3d.components.autograd.field_map import FieldMap from tidy3d.components.autograd.utils import is_tidy_box from tidy3d.components.data.data_array import DataArray from tidy3d.exceptions import AdjointError from tidy3d.plugins.polyslab import ComplexPolySlab from tidy3d.web import run, run_async -from tidy3d.web.api.autograd.utils import FieldMap from ...utils import SIM_FULL, AssertLogLevel, run_emulated, tracer_arr @@ -1174,6 +1174,21 @@ def objective(*params): ag.grad(objective)(params0) +def test_sim_hash_changes_with_traced_keys(): + """Ensure the model hash accounts for autograd traced paths.""" + + sim_traced = SIM_FULL.copy() + original_field_map = sim_traced._strip_traced_fields() + + structures = list(sim_traced.structures) + structures[0] = structures[0].to_static() + sim_modified = sim_traced.updated_copy(structures=tuple(structures)) + + modified_field_map = sim_modified._strip_traced_fields() + assert original_field_map != modified_field_map + assert sim_traced._hash_self() != sim_modified._hash_self() + + def test_sim_traced_override_structures(): """Make sure that sims with traced override structures are handled properly.""" diff --git a/tidy3d/components/autograd/field_map.py b/tidy3d/components/autograd/field_map.py new file mode 100644 index 0000000000..101e0b56bd --- /dev/null +++ b/tidy3d/components/autograd/field_map.py @@ -0,0 +1,76 @@ +"""Typed containers for autograd traced field metadata.""" + +from __future__ import annotations + +import json +from typing import Any, Callable + +import pydantic.v1 as pydantic + +from tidy3d.components.autograd.types import AutogradFieldMap, dict_ag +from tidy3d.components.base import Tidy3dBaseModel +from tidy3d.components.types import ArrayLike, tidycomplex + + +class Tracer(Tidy3dBaseModel): + """Representation of a single traced element within a model.""" + + path: tuple[Any, ...] = pydantic.Field( + ..., + title="Path to the traced object in the model dictionary.", + ) + data: float | tidycomplex | ArrayLike = pydantic.Field(..., title="Tracing data") + + +class FieldMap(Tidy3dBaseModel): + """Collection of traced elements.""" + + tracers: tuple[Tracer, ...] = pydantic.Field( + ..., + title="Collection of Tracers.", + ) + + @property + def to_autograd_field_map(self) -> AutogradFieldMap: + """Convert to ``AutogradFieldMap`` autograd dictionary.""" + return dict_ag({tracer.path: tracer.data for tracer in self.tracers}) + + @classmethod + def from_autograd_field_map(cls, autograd_field_map: AutogradFieldMap) -> FieldMap: + """Initialize from an ``AutogradFieldMap`` autograd dictionary.""" + tracers = [] + for path, data in autograd_field_map.items(): + tracers.append(Tracer(path=path, data=data)) + return cls(tracers=tuple(tracers)) + + +def _encoded_path(path: tuple[Any, ...]) -> str: + """Return a stable JSON representation for a traced path.""" + return json.dumps(list(path), separators=(",", ":"), ensure_ascii=True) + + +class TracerKeys(Tidy3dBaseModel): + """Collection of traced field paths.""" + + keys: tuple[tuple[Any, ...], ...] = pydantic.Field( + ..., + title="Collection of tracer keys.", + ) + + def encoded_keys(self) -> list[str]: + """Return the JSON-encoded representation of keys.""" + return [_encoded_path(path) for path in self.keys] + + @classmethod + def from_field_mapping( + cls, + field_mapping: AutogradFieldMap, + *, + sort_key: Callable[[tuple[Any, ...]], str] | None = None, + ) -> TracerKeys: + """Construct keys from an autograd field mapping.""" + if sort_key is None: + sort_key = _encoded_path + + sorted_paths = tuple(sorted(field_mapping.keys(), key=sort_key)) + return cls(keys=sorted_paths) diff --git a/tidy3d/components/base.py b/tidy3d/components/base.py index de0a1df789..990859d21f 100644 --- a/tidy3d/components/base.py +++ b/tidy3d/components/base.py @@ -38,6 +38,7 @@ # If json string is larger than ``MAX_STRING_LENGTH``, split the string when storing in hdf5 MAX_STRING_LENGTH = 1_000_000_000 FORBID_SPECIAL_CHARACTERS = ["/"] +TRACED_FIELD_KEYS_ATTR = "__tidy3d_traced_field_keys__" def cache(prop): @@ -767,6 +768,9 @@ def add_data_to_file(data_dict: dict, group_path: str = "") -> None: add_data_to_file(data_dict=value, group_path=subpath) add_data_to_file(data_dict=self.dict()) + traced_keys_payload = self._serialized_traced_field_keys() + if traced_keys_payload: + f_handle.attrs[TRACED_FIELD_KEYS_ATTR] = traced_keys_payload @classmethod def dict_from_hdf5_gz( @@ -1054,6 +1058,19 @@ def insert_value(x, path: tuple[str, ...], sub_dict: dict): return self.parse_obj(self_dict) + def _serialized_traced_field_keys(self) -> Optional[str]: + """Return a serialized, order-independent representation of traced field paths.""" + + field_mapping = self._strip_traced_fields() + if not field_mapping: + return None + + # TODO: remove this deferred import once TracerKeys is decoupled from Tidy3dBaseModel. + from tidy3d.components.autograd.field_map import TracerKeys + + tracer_keys = TracerKeys.from_field_mapping(field_mapping) + return tracer_keys.json(separators=(",", ":"), ensure_ascii=True) + def to_static(self) -> Tidy3dBaseModel: """Version of object with all autograd-traced fields removed.""" diff --git a/tidy3d/web/api/autograd/io_utils.py b/tidy3d/web/api/autograd/io_utils.py index 775cb69fc3..afd9f01410 100644 --- a/tidy3d/web/api/autograd/io_utils.py +++ b/tidy3d/web/api/autograd/io_utils.py @@ -4,10 +4,10 @@ import tempfile import tidy3d as td +from tidy3d.components.autograd.field_map import FieldMap, TracerKeys from tidy3d.web.core.s3utils import download_file, upload_file # type: ignore from .constants import SIM_FIELDS_KEYS_FILE, SIM_VJP_FILE -from .utils import FieldMap, TracerKeys def upload_sim_fields_keys(sim_fields_keys: list[tuple], task_id: str, verbose: bool = False): diff --git a/tidy3d/web/api/autograd/utils.py b/tidy3d/web/api/autograd/utils.py index 5d86eed822..f58f2f3c34 100644 --- a/tidy3d/web/api/autograd/utils.py +++ b/tidy3d/web/api/autograd/utils.py @@ -4,12 +4,8 @@ import typing import numpy as np -import pydantic as pd import tidy3d as td -from tidy3d.components.autograd.types import AutogradFieldMap, dict_ag -from tidy3d.components.base import Tidy3dBaseModel -from tidy3d.components.types import ArrayLike, tidycomplex """ E and D field gradient map calculation helpers. """ @@ -79,46 +75,3 @@ def get_field_key(dim: str, fld_data: typing.Union[td.FieldData, td.Permittivity mult = cmp_1 * cmp_2 field_components[key_1] = mult return fld_1.updated_copy(**field_components) - - -class Tracer(Tidy3dBaseModel): - """Class to store a single traced field.""" - - path: tuple[typing.Any, ...] = pd.Field( - ..., - title="Path to the traced object in the model dictionary.", - ) - - data: typing.Union[float, tidycomplex, ArrayLike] = pd.Field(..., title="Tracing data") - - -class FieldMap(Tidy3dBaseModel): - """Class to store a collection of traced fields.""" - - tracers: tuple[Tracer, ...] = pd.Field( - ..., - title="Collection of Tracers.", - ) - - @property - def to_autograd_field_map(self) -> AutogradFieldMap: - """Convert to ``AutogradFieldMap`` autograd dictionary.""" - return dict_ag({tracer.path: tracer.data for tracer in self.tracers}) - - @classmethod - def from_autograd_field_map(cls, autograd_field_map) -> FieldMap: - """Initialize from an ``AutogradFieldMap`` autograd dictionary.""" - tracers = [] - for path, data in autograd_field_map.items(): - tracers.append(Tracer(path=path, data=data)) - - return cls(tracers=tuple(tracers)) - - -class TracerKeys(Tidy3dBaseModel): - """Class to store a collection of tracer keys.""" - - keys: tuple[tuple[typing.Any, ...], ...] = pd.Field( - ..., - title="Collection of tracer keys.", - )