fix(autograd): include traced keys in HDF5 hash input

Author: yaugenst-flex

tests/test_components/autograd/test_autograd.py

@@ -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."""
 

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)

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."""
 

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):

tidy3d/web/api/autograd/utils.py

@@ -4,12 +4,9 @@
 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
+from tidy3d.components.autograd.field_map import FieldMap, Tracer, TracerKeys
 
 """ E and D field gradient map calculation helpers. """
 
@@ -81,44 +78,14 @@ def get_field_key(dim: str, fld_data: typing.Union[td.FieldData, td.Permittivity
     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.",
-    )
+__all__ = [
+    "E_to_D",
+    "FieldMap",
+    "Tracer",
+    "TracerKeys",
+    "derivative_map_D",
+    "derivative_map_E",
+    "derivative_map_H",
+    "get_derivative_maps",
+    "multiply_field_data",
+]