feat: Moving TYPE_TO_CLASS_MAP to Tidy3dBaseModel and allow Tidy3dBaseModel.from_file for various components

Author: momchil-flex

CHANGELOG.md

@@ -37,6 +37,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Introduced `BroadbandPulse` for exciting simulations across a wide frequency spectrum.
 - Added `interp_spec` in `ModeSpec` to allow downsampling and interpolation of waveguide modes in frequency.
 - Added warning if port mesh refinement is incompatible with the `GridSpec` in the `TerminalComponentModeler`.
+- Various types, e.g. different `Simulation` or `SimulationData` sub-classes, can be loaded from file directly with `Tidy3dBaseModel.from_file()`.
 
 ### Breaking Changes
 - Edge singularity correction at PEC and lossy metal edges defaults to `True`.

tests/test_components/test_base.py

@@ -2,8 +2,11 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 import numpy as np
 import pytest
+from pydantic.v1 import ValidationError
 
 import tidy3d as td
 from tidy3d.components.base import Tidy3dBaseModel
@@ -275,3 +278,19 @@ def test_updated_hash_and_json_with_changed_attr():
 
     assert new_hash != old_hash
     assert json_old != json_new
+
+
+def test_parse_obj_respects_subclasses():
+    class DispatchBase(Tidy3dBaseModel):
+        type: Literal["DispatchBase"] = "DispatchBase"
+        value: int
+
+    class DispatchChild(DispatchBase):
+        type: Literal["DispatchChild"] = "DispatchChild"
+
+    data = {"type": "DispatchChild", "value": 1}
+    parsed = Tidy3dBaseModel._parse_model_dict(data)
+    assert isinstance(parsed, DispatchChild)
+
+    with pytest.raises(ValidationError):
+        DispatchChild.parse_obj({"type": "DispatchBase", "value": 2})

tests/test_plugins/expressions/test_dispatch.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import pytest
+
+from tidy3d.plugins.expressions.base import Expression
+from tidy3d.plugins.expressions.variables import Constant
+
+
+def test_expression_parse_obj_round_trip():
+    expr = Constant(3.14)
+    parsed = Expression.parse_obj(expr.dict())
+    assert isinstance(parsed, Constant)
+    assert parsed.value == pytest.approx(3.14)
+
+
+def test_expression_parse_obj_rejects_unrelated_types():
+    # Simulation registers a distinct type in the global map; parsing via Expression should fail.
+    with pytest.raises(ValueError, match="Cannot parse type"):
+        Expression.parse_obj({"type": "Simulation"})

tests/test_web/test_local_cache.py

@@ -26,7 +26,11 @@
 from tidy3d.web.api import webapi as web
 from tidy3d.web.api.autograd import autograd, engine, io_utils
 from tidy3d.web.api.autograd.autograd import run as run_autograd
-from tidy3d.web.api.autograd.constants import SIM_VJP_FILE
+from tidy3d.web.api.autograd.constants import (
+    AUX_KEY_SIM_DATA_FWD,
+    AUX_KEY_SIM_DATA_ORIGINAL,
+    SIM_VJP_FILE,
+)
 from tidy3d.web.api.container import Batch, WebContainer
 from tidy3d.web.api.webapi import load_simulation_if_cached
 from tidy3d.web.cache import (
@@ -181,13 +185,37 @@ def _fake_download_file(resource_id, remote_filename, to_file=None, **kwargs):
         if str(remote_filename) == SIM_VJP_FILE:
             counters["download"] += 1
 
-    def _fake_from_file(*args, **kwargs):
-        field_map = FieldMap(tracers=())
-        return field_map
+    def _fake_postprocess_fwd(*, sim_data_combined=None, sim_original=None, aux_data=None, **_):
+        """Mimic ``autograd.postprocess_fwd`` side effects for tests."""
+        if sim_original is None:
+            sim_original = next(iter(PATH_TO_SIM.values()), None)
+        if sim_original is None:
+            sim_original = td.Simulation(
+                size=(1, 1, 1),
+                grid_spec=td.GridSpec.auto(wavelength=1.0),
+                run_time=1e-12,
+            )
+        stub_data = _FakeStubData(sim_original)
+        if aux_data is not None:
+            aux_data[AUX_KEY_SIM_DATA_ORIGINAL] = stub_data
+            aux_data[AUX_KEY_SIM_DATA_FWD] = stub_data
+        return stub_data._strip_traced_fields()
+
+    def _fake_postprocess_adj(
+        sim_data_adj=None, sim_data_orig=None, sim_data_fwd=None, sim_fields_keys=None, **_
+    ):
+        """Return zeros for every requested field key."""
+        counters["download"] += 1  # mimic VJP file download per autograd run
+        sim_fields_keys = sim_fields_keys or []
+        return dict.fromkeys(sim_fields_keys, 0.0)
+
+    def _fake_field_map_from_file(*args, **kwargs):
+        return FieldMap(tracers=())
 
     monkeypatch.setattr(io_utils, "download_file", _fake_download_file)
-    monkeypatch.setattr(autograd, "postprocess_fwd", _fake_from_file)
-    monkeypatch.setattr(FieldMap, "from_file", _fake_from_file)
+    monkeypatch.setattr(autograd, "postprocess_fwd", _fake_postprocess_fwd)
+    monkeypatch.setattr(autograd, "postprocess_adj", _fake_postprocess_adj)
+    monkeypatch.setattr(FieldMap, "from_file", _fake_field_map_from_file)
     monkeypatch.setattr(WebContainer, "_check_folder", _fake__check_folder)
     monkeypatch.setattr(web, "upload", _fake_upload)
     monkeypatch.setattr(web, "start", _fake_start)

tests/test_web/test_tidy3d_stub.py

@@ -42,6 +42,15 @@ def make_sim():
     )
 
 
+def is_lazy_object(data):
+    assert set(data.__dict__.keys()) == {
+        "_lazy_fname",
+        "_lazy_group_path",
+        "_lazy_parse_obj_kwargs",
+    }
+    return True
+
+
 def make_sim_data(file_size_gb=0.001):
     """Makes a simulation data."""
     N = int(2.528e8 / 4 * file_size_gb)
@@ -146,22 +155,16 @@ def test_stub_data_lazy_loading(tmp_path):
         sim_data = Tidy3dStubData.postprocess(file_path, lazy=True)
 
     sim_data_copy = sim_data.copy()
-    assert type(sim_data).__name__ == "SimulationDataProxy"
-    assert type(sim_data_copy).__name__ == "SimulationDataProxy"
+
+    # variable dict should only contain metadata to load the data, not the data itself
+    assert is_lazy_object(sim_data)
 
     # the type should be still SimulationData despite being lazy
     assert isinstance(sim_data, SimulationData)
 
-    # variable dict should only contain metadata to load the data, not the data itself
-    assert set(sim_data.__dict__.keys()) == {
-        "_lazy_fname",
-        "_lazy_group_path",
-        "_lazy_parse_obj_kwargs",
-    }
-
     # we expect a warning from the lazy object if some field is accessed
     with AssertLogLevel("WARNING", contains_str=sim_diverged_log):
-        _ = sim_data.monitor_data
+        _ = sim_data_copy.monitor_data
 
 
 @pytest.mark.parametrize(
@@ -185,6 +188,7 @@ def test_stub_pathlike_roundtrip(tmp_path, path_builder):
 
     # Simulation data stub roundtrip
     sim_data = make_sim_data()
+    sim_data = sim_data.updated_copy(log="log")
     stub_data = Tidy3dStubData(data=sim_data)
     data_path = path_builder(tmp_path, "pathlike_data.hdf5")
     stub_data.to_file(data_path)

tests/test_web/test_webapi.py

@@ -15,6 +15,7 @@
 from responses import matchers
 
 import tidy3d as td
+from tests.test_web.test_tidy3d_stub import is_lazy_object
 from tidy3d import Simulation
 from tidy3d.__main__ import main
 from tidy3d.components.data.data_array import ScalarFieldDataArray
@@ -958,21 +959,18 @@ def test_run_with_flexible_containers_offline_lazy(monkeypatch, tmp_path):
     apply_common_patches(monkeypatch, tmp_path, taskid_to_sim=taskid_to_sim)
 
     data = run(sim_container, task_name=task_name, folder_name="PROJECT", path=str(out_dir))
-
+    assert is_lazy_object(data[0])
     assert isinstance(data, list) and len(data) == 3
 
-    assert isinstance(data[0], SimulationData)
-    assert data[0].__class__.__name__ == "SimulationDataProxy"
-
     assert isinstance(data[1], dict)
     assert "sim2" in data[1]
+    assert is_lazy_object(data[1]["sim2"])
     assert isinstance(data[1]["sim2"], SimulationData)
-    assert data[1]["sim2"].__class__.__name__ == "SimulationDataProxy"
 
+    assert is_lazy_object(data[2][0])
     assert isinstance(data[2], tuple)
-    assert data[2][0].__class__.__name__ == "SimulationDataProxy"
+    assert is_lazy_object(data[2][1][0])
     assert isinstance(data[2][1], list)
-    assert data[2][1][0].__class__.__name__ == "SimulationDataProxy"
 
     assert data[0].simulation == sim1
     assert data[1]["sim2"].simulation == sim2