add lazy option for loading simulation data

Author: marcorudolphflex

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - New `MediumMonitor` that returns both permittivity and permeability profiles.
 - Task names are now optional when using `run(sim)` or `Job`. When running multiple jobs (via `run_async` or `Batch`), you can also provide simulations as a list without specifying task names. The previous dictionary-based format with explicit task names is still supported.
 - Added support for `tidy3d-extras`, an optional plugin that enables more accurate local mode solving via subpixel averaging.
+- Enabled lazy loading of data via `web.load(..., lazy=True)`. When used, this returns a lightweight proxy object holding a reference to the data. On first access to any field or method, the proxy transparently loads the full object (same as with the default lazy=False).
 
 ### Changed
 - `LayerRefinementSpec` defaults to assuming structures made of different materials are interior-disjoint for more efficient mesh generation.

tests/test_web/test_tidy3d_stub.py

@@ -6,6 +6,7 @@
 import responses
 
 import tidy3d as td
+from tests.utils import AssertLogLevel
 from tidy3d.components.data.data_array import ScalarFieldDataArray
 from tidy3d.components.data.monitor_data import FieldData
 from tidy3d.components.data.sim_data import SimulationData
@@ -122,6 +123,45 @@ def test_stub_data_postprocess_logs(tmp_path):
     Tidy3dStubData.postprocess(file_path)
 
 
+@responses.activate
+def test_stub_data_lazy_loading(tmp_path):
+    """Tests the postprocess method with lazy loading of Tidy3dStubData when simulation diverged."""
+    td.log.set_capture(True)
+    sim_diverged_log = "The simulation has diverged!"
+
+    # make sim data where test diverged
+    sim_data = make_sim_data()
+    sim_data = sim_data.updated_copy(diverged=True, log=sim_diverged_log)
+    file_path = os.path.join(tmp_path, "test_diverged.hdf5")
+    sim_data.to_file(file_path)
+
+    # default case with lazy=False should output a warning
+    with AssertLogLevel("WARNING", contains_str=sim_diverged_log):
+        Tidy3dStubData.postprocess(file_path, lazy=False)
+
+    # we expect no warning in lazy mode as object should not be loaded
+    with AssertLogLevel(None):
+        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"
+
+    # 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
+
+
 def test_default_task_name():
     sim = make_sim()
     stub = Tidy3dStub(simulation=sim)

tidy3d/components/base.py

@@ -308,17 +308,30 @@ def help(self, methods: bool = False) -> None:
 
     @classmethod
     def from_file(
-        cls, fname: str, group_path: Optional[str] = None, **parse_obj_kwargs
+        cls,
+        fname: str,
+        group_path: Optional[str] = None,
+        lazy: bool = False,
+        on_load: Optional[Callable] = None,
+        **parse_obj_kwargs,
     ) -> Tidy3dBaseModel:
         """Loads a :class:`Tidy3dBaseModel` from .yaml, .json, .hdf5, or .hdf5.gz file.
 
         Parameters
         ----------
         fname : str
             Full path to the file to load the :class:`Tidy3dBaseModel` from.
-        group_path : str, optional
+        group_path : str | None = None
             Path to a group inside the file to use as the base level. Only for hdf5 files.
             Starting `/` is optional.
+        lazy : bool = False
+            Whether to load the actual data (``lazy=False``) or return a proxy that loads
+            the data when accessed (``lazy=True``).
+        on_load : Callable | None = None
+            Callback function executed once the model is fully materialized.
+            Only used if ``lazy=True``. The callback is invoked with the loaded
+            instance as its sole argument, enabling post-processing such as
+            validation, logging, or warnings checks.
         **parse_obj_kwargs
             Keyword arguments passed to either pydantic's ``parse_obj`` function when loading model.
 
@@ -331,8 +344,14 @@ def from_file(
         -------
         >>> simulation = Simulation.from_file(fname='folder/sim.json') # doctest: +SKIP
         """
+        if lazy:
+            Proxy = _make_lazy_proxy(cls, on_load=on_load)  # staticmethod usage
+            return Proxy(fname, group_path, parse_obj_kwargs)
         model_dict = cls.dict_from_file(fname=fname, group_path=group_path)
-        return cls.parse_obj(model_dict, **parse_obj_kwargs)
+        obj = cls.parse_obj(model_dict, **parse_obj_kwargs)
+        if not lazy and on_load is not None:
+            on_load(obj)
+        return obj
 
     @classmethod
     def dict_from_file(cls, fname: str, group_path: Optional[str] = None) -> dict:
@@ -1201,3 +1220,77 @@ def to_sci(value: float, exponent: int, precision: int) -> str:
         sci_max = to_sci(max_val, common_exponent, precision)
 
         return sci_min, sci_max
+
+
+def _make_lazy_proxy(
+    target_cls: type,
+    on_load: Optional[Callable[[Any], None]] = None,
+) -> type:
+    """
+    Return a lazy-loading proxy subclass of ``target_cls``.
+
+    Parameters
+    ----------
+    target_cls : type
+        Must implement ``dict_from_file`` and ``parse_obj``.
+    on_load : Callable[[Any], None] | None = None
+        A function to call with the fully loaded instance once loaded.
+
+    Returns
+    -------
+    type
+        A class named ``<TargetClsName>Proxy`` with init args:
+        ``(fname, group_path, parse_obj_kwargs)``.
+    """
+    proxy_name = f"{target_cls.__name__}Proxy"
+
+    class _LazyProxy(target_cls):
+        def __init__(
+            self, fname: str, group_path: Optional[str], parse_obj_kwargs: Optional[dict[str, Any]]
+        ):
+            object.__setattr__(self, "_lazy_fname", fname)
+            object.__setattr__(self, "_lazy_group_path", group_path)
+            object.__setattr__(self, "_lazy_parse_obj_kwargs", dict(parse_obj_kwargs or {}))
+
+        def copy(self, **kwargs):
+            """Return another lazy proxy instead of materializing."""
+            return _LazyProxy(
+                self._lazy_fname,
+                self._lazy_group_path,
+                {**self._lazy_parse_obj_kwargs, **kwargs},
+            )
+
+        def __getattribute__(self, name: str):
+            if name in (
+                "__class__",
+                "__dict__",
+                "__weakref__",
+                "__post_root_validators__",
+                "copy",  # <-- avoid materializing just for copy
+            ) or name.startswith("_lazy_"):
+                return object.__getattribute__(self, name)
+
+            d = object.__getattribute__(self, "__dict__")
+            if "_lazy_fname" in d:  # sentinel: not loaded yet
+                fname = d["_lazy_fname"]
+                group_path = d["_lazy_group_path"]
+                kwargs = d["_lazy_parse_obj_kwargs"]
+
+                model_dict = target_cls.dict_from_file(fname=fname, group_path=group_path)
+                target = target_cls.parse_obj(model_dict, **kwargs)
+
+                d.clear()
+                d.update(target.__dict__)
+                object.__setattr__(self, "__class__", target_cls)
+                object.__setattr__(self, "__fields_set__", set(target.__fields_set__))
+                private_attrs = getattr(target, "__private_attributes__", {}) or {}
+                for attr_name in private_attrs:
+                    object.__setattr__(self, attr_name, getattr(target, attr_name))
+
+                if on_load is not None:
+                    on_load(self)
+
+            return object.__getattribute__(self, name)
+
+    _LazyProxy.__name__ = proxy_name
+    return _LazyProxy

tidy3d/components/data/sim_data.py

@@ -903,6 +903,13 @@ class SimulationData(AbstractYeeGridSimulationData):
         sim_data.to_file(fname='path/to/file.hdf5') # Save a SimulationData object to a HDF5 file
         sim_data = SimulationData.from_file(fname='path/to/file.hdf5') # Load a SimulationData object from a HDF5 file.
 
+    Optionally, the simulation data can be loaded in a lazy mode, which only holds a reference until a field is accessed
+    or a method is applied. This is useful to save I/O operations and memory.
+
+    .. code-block:: python
+
+        sim_data = SimulationData.from_file(fname='path/to/file.hdf5', lazy=True) # Does not contain data until accessed.
+
     See Also
     --------
 

tidy3d/web/api/tidy3d_stub.py

@@ -82,24 +82,28 @@ def from_file(cls, file_path: str) -> WorkflowType:
 
         data = json.loads(json_str)
         type_ = data["type"]
-        if type_ == "Simulation":
-            sim = Simulation.from_file(file_path)
-        elif type_ == "ModeSolver":
-            sim = ModeSolver.from_file(file_path)
-        elif type_ == "HeatSimulation":
-            sim = HeatSimulation.from_file(file_path)
-        elif type_ == "HeatChargeSimulation":
-            sim = HeatChargeSimulation.from_file(file_path)
-        elif type_ == "EMESimulation":
-            sim = EMESimulation.from_file(file_path)
-        elif type_ == "ModeSimulation":
-            sim = ModeSimulation.from_file(file_path)
-        elif type_ == "VolumeMesher":
-            sim = VolumeMesher.from_file(file_path)
-        elif type_ == "ModalComponentModeler":
-            sim = ModalComponentModeler.from_file(file_path)
-        elif type_ == "TerminalComponentModeler":
-            sim = TerminalComponentModeler.from_file(file_path)
+
+        supported_classes = [
+            Simulation,
+            ModeSolver,
+            HeatSimulation,
+            HeatChargeSimulation,
+            EMESimulation,
+            ModeSimulation,
+            VolumeMesher,
+            ModalComponentModeler,
+            TerminalComponentModeler,
+        ]
+
+        class_map = {cls.__name__: cls for cls in supported_classes}
+
+        if type_ not in class_map:
+            raise ValueError(
+                f"Unsupported type '{type_}'. Supported types: {list(class_map.keys())}"
+            )
+
+        sim_class = class_map[type_]
+        sim = sim_class.from_file(file_path)
 
         return sim
 
@@ -202,7 +206,9 @@ class Tidy3dStubData(BaseModel, TaskStubData):
     data: WorkflowDataType
 
     @classmethod
-    def from_file(cls, file_path: str) -> WorkflowDataType:
+    def from_file(
+        cls, file_path: str, lazy: bool = False, on_load: Optional[Callable] = None
+    ) -> WorkflowDataType:
         """Loads a Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
         from .yaml, .json, or .hdf5 file.
 
@@ -211,6 +217,14 @@ def from_file(cls, file_path: str) -> WorkflowDataType:
         file_path : str
             Full path to the .yaml or .json or .hdf5 file to load the
             Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`] from.
+        lazy : bool = False
+            Whether to load the actual data (``lazy=False``) or return a proxy that loads
+            the data when accessed (``lazy=True``).
+        on_load : Callable | None = None
+            Callback function executed once the model is fully materialized.
+            Only used if ``lazy=True``. The callback is invoked with the loaded
+            instance as its sole argument, enabling post-processing such as
+            validation, logging, or warnings checks.
 
         Returns
         -------
@@ -227,24 +241,28 @@ def from_file(cls, file_path: str) -> WorkflowDataType:
 
         data = json.loads(json_str)
         type_ = data["type"]
-        if type_ == "SimulationData":
-            sim_data = SimulationData.from_file(file_path)
-        elif type_ == "ModeSolverData":
-            sim_data = ModeSolverData.from_file(file_path)
-        elif type_ == "HeatSimulationData":
-            sim_data = HeatSimulationData.from_file(file_path)
-        elif type_ == "HeatChargeSimulationData":
-            sim_data = HeatChargeSimulationData.from_file(file_path)
-        elif type_ == "EMESimulationData":
-            sim_data = EMESimulationData.from_file(file_path)
-        elif type_ == "ModeSimulationData":
-            sim_data = ModeSimulationData.from_file(file_path)
-        elif type_ == "VolumeMesherData":
-            sim_data = VolumeMesherData.from_file(file_path)
-        elif type_ == "ModalComponentModelerData":
-            sim_data = ModalComponentModelerData.from_file(file_path)
-        elif type_ == "TerminalComponentModelerData":
-            sim_data = TerminalComponentModelerData.from_file(file_path)
+
+        supported_data_classes = [
+            SimulationData,
+            ModeSolverData,
+            HeatSimulationData,
+            HeatChargeSimulationData,
+            EMESimulationData,
+            ModeSimulationData,
+            VolumeMesherData,
+            ModalComponentModelerData,
+            TerminalComponentModelerData,
+        ]
+
+        data_class_map = {cls.__name__: cls for cls in supported_data_classes}
+
+        if type_ not in data_class_map:
+            raise ValueError(
+                f"Unsupported data type '{type_}'. Supported types: {list(data_class_map.keys())}"
+            )
+
+        data_class = data_class_map[type_]
+        sim_data = data_class.from_file(file_path, lazy=lazy, on_load=on_load)
 
         return sim_data
 
@@ -265,7 +283,7 @@ def to_file(self, file_path: str):
         self.data.to_file(file_path)
 
     @classmethod
-    def postprocess(cls, file_path: str) -> WorkflowDataType:
+    def postprocess(cls, file_path: str, lazy: bool = True) -> WorkflowDataType:
         """Load .yaml, .json, or .hdf5 file to
         Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`] instance.
 
@@ -274,16 +292,28 @@ def postprocess(cls, file_path: str) -> WorkflowDataType:
         file_path : str
             Full path to the .yaml or .json or .hdf5 file to save the
             Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`] to.
+        lazy : bool = False
+            Whether to load the actual data (``lazy=False``) or return a proxy that loads
+            the data when accessed (``lazy=True``).
 
         Returns
         -------
         Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
             An instance of the component class calling ``load``.
         """
-        stub_data = Tidy3dStubData.from_file(file_path)
+        stub_data = Tidy3dStubData.from_file(
+            file_path, lazy=lazy, on_load=cls._check_convergence_and_warnings
+        )
+        if not lazy:
+            cls._check_convergence_and_warnings(stub_data)
+        return stub_data
 
-        check_log_msg = "For more information, check 'SimulationData.log' or use "
-        check_log_msg += "'web.download_log(task_id)'."
+    @staticmethod
+    def _check_convergence_and_warnings(stub_data: WorkflowDataType) -> None:
+        """Check convergence, divergence, and warnings in the solver log and emit log messages."""
+        check_log_msg = (
+            "For more information, check 'SimulationData.log' or use 'web.download_log(task_id)'."
+        )
         warned_about_warnings = False
 
         if isinstance(stub_data, SimulationData):
@@ -313,5 +343,3 @@ def postprocess(cls, file_path: str) -> WorkflowDataType:
             and not warned_about_warnings
         ):
             log.warning("Warning messages were found in the solver log. " + check_log_msg)
-
-        return stub_data

tidy3d/web/api/webapi.py

@@ -969,6 +969,7 @@ def load(
     replace_existing: bool = True,
     verbose: bool = True,
     progress_callback: Optional[Callable[[float], None]] = None,
+    lazy: bool = False,
 ) -> WorkflowDataType:
     """
     Download and Load simulation results into :class:`.SimulationData` object.
@@ -998,6 +999,9 @@ def load(
         If ``True``, will print progressbars and status, otherwise, will run silently.
     progress_callback : Callable[[float], None] = None
         Optional callback function called when downloading file with ``bytes_in_chunk`` as argument.
+    lazy : bool = False
+        Whether to load the actual data (``lazy=False``) or return a proxy that loads
+        the data when accessed (``lazy=True``).
 
     Returns
     -------
@@ -1019,7 +1023,7 @@ def load(
         else:
             console.log(f"loading simulation from {path}")
 
-    stub_data = Tidy3dStubData.postprocess(path)
+    stub_data = Tidy3dStubData.postprocess(path, lazy=lazy)
     return stub_data