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

Author: yaugenst-flex

tests/test_components/autograd/test_autograd.py

@@ -10,6 +10,7 @@
 
 import autograd as ag
 import autograd.numpy as anp
+import h5py
 import matplotlib.pylab as plt
 import numpy as np
 import numpy.testing as npt
@@ -25,12 +26,14 @@
     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.base import TRACED_FIELD_KEYS_ATTR
 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 tidy3d.web.api.autograd import autograd as autograd_module
 
 from ...utils import SIM_FULL, AssertLogLevel, run_emulated, tracer_arr
 
@@ -1174,6 +1177,124 @@ 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_hdf5_records_traced_keys(tmp_path):
+    """HDF5 exports should include traced-key metadata for caching."""
+
+    sim_traced = SIM_FULL.copy()
+    expected_payload = sim_traced._serialized_traced_field_keys()
+    assert expected_payload, "simulation fixture must yield traced keys"
+
+    sim_traced.attrs.pop(TRACED_FIELD_KEYS_ATTR, None)
+
+    export_path = tmp_path / "sim_traced.hdf5"
+    sim_traced.to_hdf5(str(export_path))
+
+    with h5py.File(export_path, "r") as handle:
+        assert TRACED_FIELD_KEYS_ATTR in handle.attrs
+        assert handle.attrs[TRACED_FIELD_KEYS_ATTR] == expected_payload
+
+    static_export = tmp_path / "sim_traced_static.hdf5"
+    sim_traced.attrs[TRACED_FIELD_KEYS_ATTR] = expected_payload
+    sim_static = sim_traced.to_static()
+    sim_static.to_hdf5(str(static_export))
+
+    with h5py.File(static_export, "r") as handle:
+        assert TRACED_FIELD_KEYS_ATTR in handle.attrs
+        assert handle.attrs[TRACED_FIELD_KEYS_ATTR] == expected_payload
+
+
+def test_web_run_duplicate_simulations(monkeypatch):
+    """Repeated simulation objects should reuse cached data without hash mismatches."""
+
+    sim = SIM_FULL.copy()
+    sim.attrs.pop(TRACED_FIELD_KEYS_ATTR, None)
+
+    copy_calls = {"count": 0}
+
+    class DummyData:
+        def __init__(self, label: str):
+            self.label = label
+
+        def copy(self):
+            copy_calls["count"] += 1
+            return DummyData(f"{self.label}_copy{copy_calls['count']}")
+
+    dummy = DummyData("root")
+
+    def fake_run_autograd(*args, **kwargs):
+        return dummy
+
+    monkeypatch.setattr("tidy3d.web.api.run.run_autograd", fake_run_autograd)
+
+    results = web.run([sim, sim])
+
+    assert isinstance(results, list)
+    assert len(results) == 2
+    assert results[0] is dummy
+    assert results[1] is not dummy
+    assert copy_calls["count"] == 1
+
+
+def test_autograd_run_does_not_mutate_input_attrs(monkeypatch):
+    """Autograd run should attach traced metadata only to the exported static copy."""
+
+    sim = SIM_FULL.copy()
+    sim.attrs.pop(TRACED_FIELD_KEYS_ATTR, None)
+    payload = sim._serialized_traced_field_keys()
+    assert payload
+
+    captured: dict[str, typing.Any] = {}
+
+    def fake_run_primitive(
+        sim_fields,
+        sim_original,
+        task_name,
+        aux_data,
+        local_gradient,
+        max_num_adjoint_per_fwd,
+        **run_kwargs,
+    ):
+        captured["sim_original"] = sim_original
+        captured["payload"] = sim_original.attrs.get(TRACED_FIELD_KEYS_ATTR)
+        captured["sim_fields"] = sim_fields
+        captured["aux_data"] = aux_data
+        return sim_fields
+
+    def fake_postprocess_run(traced_fields_data, aux_data):
+        captured["postprocess_data"] = traced_fields_data
+        captured["postprocess_aux"] = aux_data
+        return "sentinel"
+
+    monkeypatch.setattr(autograd_module, "_run_primitive", fake_run_primitive)
+    monkeypatch.setattr(autograd_module, "postprocess_run", fake_postprocess_run)
+
+    result = autograd_module._run(simulation=sim, task_name="dummy")
+
+    assert result == "sentinel"
+    assert sim.attrs.get(TRACED_FIELD_KEYS_ATTR) is None
+    assert captured["payload"] == payload
+    assert captured["sim_original"] is not sim
+    assert captured["sim_original"].attrs.get(TRACED_FIELD_KEYS_ATTR) == payload
+    assert captured["postprocess_data"] == captured["sim_fields"]
+    assert captured["postprocess_aux"] is captured["aux_data"]
+    assert captured["postprocess_aux"] == {}
+
+
 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

@@ -39,6 +39,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):
@@ -524,7 +525,8 @@ def to_json(self, fname: str) -> None:
         -------
         >>> simulation.to_json(fname='folder/sim.json') # doctest: +SKIP
         """
-        json_string = self._json(indent=INDENT_JSON_FILE)
+        export_model = self.to_static()
+        json_string = export_model._json(indent=INDENT_JSON_FILE)
         self._warn_if_contains_data(json_string)
         with open(fname, "w", encoding="utf-8") as file_handle:
             file_handle.write(json_string)
@@ -586,7 +588,8 @@ def to_yaml(self, fname: str) -> None:
         -------
         >>> simulation.to_yaml(fname='folder/sim.yaml') # doctest: +SKIP
         """
-        json_string = self._json_string
+        export_model = self.to_static()
+        json_string = export_model._json()
         self._warn_if_contains_data(json_string)
         model_dict = json.loads(json_string)
         with open(fname, "w+", encoding="utf-8") as file_handle:
@@ -792,8 +795,15 @@ def to_hdf5(
         >>> simulation.to_hdf5(fname='folder/sim.hdf5') # doctest: +SKIP
         """
 
+        export_model = self.to_static()
+        traced_keys_payload = export_model.attrs.get(TRACED_FIELD_KEYS_ATTR)
+
+        if traced_keys_payload is None:
+            traced_keys_payload = self.attrs.get(TRACED_FIELD_KEYS_ATTR)
+        if traced_keys_payload is None:
+            traced_keys_payload = self._serialized_traced_field_keys()
         with h5py.File(fname, "w") as f_handle:
-            json_str = self._json_string
+            json_str = export_model._json()
             for ind in range(ceil(len(json_str) / MAX_STRING_LENGTH)):
                 ind_start = int(ind * MAX_STRING_LENGTH)
                 ind_stop = min(int(ind + 1) * MAX_STRING_LENGTH, len(json_str))
@@ -816,14 +826,16 @@ def add_data_to_file(data_dict: dict, group_path: str = "") -> None:
 
                     # if a tuple, assign each element a unique key
                     if isinstance(value, (list, tuple)):
-                        value_dict = self.tuple_to_dict(tuple_values=value)
+                        value_dict = export_model.tuple_to_dict(tuple_values=value)
                         add_data_to_file(data_dict=value_dict, group_path=subpath)
 
                     # if a dict, recurse
                     elif isinstance(value, dict):
                         add_data_to_file(data_dict=value, group_path=subpath)
 
-            add_data_to_file(data_dict=self.dict())
+            add_data_to_file(data_dict=export_model.dict())
+            if traced_keys_payload:
+                f_handle.attrs[TRACED_FIELD_KEYS_ATTR] = traced_keys_payload
 
     @classmethod
     def dict_from_hdf5_gz(
@@ -1101,6 +1113,22 @@ def insert_value(x, path: tuple[str, ...], sub_dict: dict):
 
         return self.parse_obj(self_dict)
 
+    def _serialized_traced_field_keys(
+        self, field_mapping: AutogradFieldMap | None = None
+    ) -> Optional[str]:
+        """Return a serialized, order-independent representation of traced field paths."""
+
+        if field_mapping is None:
+            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/autograd.py

@@ -14,6 +14,7 @@
     MAX_NUM_ADJOINT_PER_FWD,
     MAX_NUM_TRACED_STRUCTURES,
 )
+from tidy3d.components.base import TRACED_FIELD_KEYS_ATTR
 from tidy3d.components.types.workflow import WorkflowDataType, WorkflowType
 from tidy3d.exceptions import AdjointError
 from tidy3d.web.api.asynchronous import DEFAULT_DATA_DIR
@@ -409,9 +410,14 @@ def _run(
     aux_data = {}
 
     # run our custom @primitive, passing the traced fields first to register with autograd
+    sim_static = simulation.to_static()
+    traced_keys_payload = simulation._serialized_traced_field_keys()
+    if traced_keys_payload:
+        sim_static.attrs[TRACED_FIELD_KEYS_ATTR] = traced_keys_payload
+
     traced_fields_data = _run_primitive(
         traced_fields_sim,  # if you pass as a kwarg it will not trace :/
-        sim_original=simulation.to_static(),
+        sim_original=sim_static,
         task_name=task_name,
         aux_data=aux_data,
         local_gradient=local_gradient,
@@ -433,16 +439,22 @@ def _run_async(
     task_names = simulations.keys()
 
     traced_fields_sim_dict = {}
+    sims_original = {}
     for task_name in task_names:
-        traced_fields_sim_dict[task_name] = setup_run(simulation=simulations[task_name])
+        simulation = simulations[task_name]
+        traced_fields = setup_run(simulation=simulation)
+        traced_fields_sim_dict[task_name] = traced_fields
+        sim_static = simulation.to_static()
+        if traced_fields:
+            traced_keys_payload = simulation._serialized_traced_field_keys()
+            if traced_keys_payload:
+                sim_static.attrs[TRACED_FIELD_KEYS_ATTR] = traced_keys_payload
+        sims_original[task_name] = sim_static
     traced_fields_sim_dict = dict_ag(traced_fields_sim_dict)
 
     # TODO: shortcut primitive running for any items with no tracers?
 
     aux_data_dict = {task_name: {} for task_name in task_names}
-    sims_original = {
-        task_name: simulation.to_static() for task_name, simulation in simulations.items()
-    }
     traced_fields_data_dict = _run_async_primitive(
         traced_fields_sim_dict,  # if you pass as a kwarg it will not trace :/
         sims_original=sims_original,

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