Add metadata to Tidy3D components

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Support for `.gz` files in `Simulation` version updater.
 - Warning if a nonuniform custom medium is intersecting `PlaneWave`, `GaussianBeam`, `AstigmaticGaussianBeam`, `FieldProjectionCartesianMonitor`, `FieldProjectionAngleMonitor`, `FieldProjectionKSpaceMonitor`, and `DiffractionMonitor`.
+- Tidy3D objects may store arbitrary metadata in an `.attrs` dictionary.
 
 ### Changed
 

tests/test_components/test_base.py

@@ -178,3 +178,51 @@ def test_special_characters_in_name():
     """Test error if special characters are in a component's name."""
     with pytest.raises(ValueError):
         mnt = td.FluxMonitor(size=(1, 1, 0), freqs=np.array([1, 2, 3]) * 1e12, name="mnt/flux")
+
+
+def test_attrs(tmp_path):
+    """Test the ``.attrs`` metadata feature."""
+
+    # attrs initialize to empty dict
+    obj = td.Medium()
+    assert obj.attrs == {}
+
+    # or they can be initialized directly
+    obj = td.Medium(attrs={"foo": "attr"})
+    assert obj.attrs == {"foo": "attr"}
+
+    # this is still not allowed though
+    with pytest.raises(TypeError):
+        obj.attrs = {}
+
+    # attrs can be modified
+    obj.attrs["foo"] = "bar"
+    assert obj.attrs == {"foo": "bar"}
+
+    # attrs persist with regular copies
+    obj2 = obj.copy()
+    assert obj2.attrs == obj.attrs
+
+    # attrs persist with updated copies
+    obj3 = obj2.updated_copy(permittivity=2.0)
+    assert obj3.attrs == obj2.attrs
+
+    # attrs are in the json strings
+    obj_json = obj3.json()
+    assert '{"foo": "bar"}' in obj_json
+
+    # attrs are in the dict()
+    obj_dict = obj3.dict()
+    assert obj_dict["attrs"] == {"foo": "bar"}
+
+    # objects saved and loaded from file still have attrs
+    for extension in ("hdf5", "json"):
+        path = str(tmp_path / ("obj." + extension))
+        obj.to_file(path)
+        obj4 = obj.from_file(path)
+        assert obj4.attrs == obj.attrs
+
+    # test attrs that can't be serialized
+    obj.attrs["not_serializable"] = type
+    with pytest.raises(TypeError):
+        obj.json()

tests/test_components/test_grid.py

@@ -4,7 +4,6 @@
 
 import tidy3d as td
 from tidy3d.components.grid.grid import Coords, FieldGrid, Grid
-from tidy3d.components.types import TYPE_TAG_STR
 from tidy3d.exceptions import SetupError
 
 
@@ -42,7 +41,8 @@ def test_grid():
     assert np.all(g.centers.y == np.array([-1.5, -0.5, 0.5, 1.5]))
     assert np.all(g.centers.z == np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5]))
 
-    for s in g.sizes.dict(exclude={TYPE_TAG_STR}).values():
+    for dim in "xyz":
+        s = g.sizes.dict()[dim]
         assert np.all(np.array(s) == 1.0)
 
     assert np.all(g.yee.E.x.x == np.array([-0.5, 0.5]))
@@ -211,9 +211,12 @@ def test_sim_grid():
         boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()),
     )
 
-    for c in sim.grid.centers.dict(exclude={TYPE_TAG_STR}).values():
+    for dim in "xyz":
+        c = sim.grid.centers.dict()[dim]
         assert np.all(c == np.array([-1.5, -0.5, 0.5, 1.5]))
-    for b in sim.grid.boundaries.dict(exclude={TYPE_TAG_STR}).values():
+
+    for dim in "xyz":
+        b = sim.grid.boundaries.dict()[dim]
         assert np.all(b == np.array([-2, -1, 0, 1, 2]))
 
 
@@ -261,9 +264,12 @@ def test_sim_pml_grid():
         run_time=1e-12,
     )
 
-    for c in sim.grid.centers.dict(exclude={TYPE_TAG_STR}).values():
+    for dim in "xyz":
+        c = sim.grid.centers.dict()[dim]
         assert np.all(c == np.array([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]))
-    for b in sim.grid.boundaries.dict(exclude={TYPE_TAG_STR}).values():
+
+    for dim in "xyz":
+        b = sim.grid.boundaries.dict()[dim]
         assert np.all(b == np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4]))
 
 
@@ -279,10 +285,12 @@ def test_sim_discretize_vol():
 
     subgrid = sim.discretize(vol)
 
-    for b in subgrid.boundaries.dict(exclude={TYPE_TAG_STR}).values():
+    for dim in "xyz":
+        b = subgrid.boundaries.dict()[dim]
         assert np.all(b == np.array([-1, 0, 1]))
 
-    for c in subgrid.centers.dict(exclude={TYPE_TAG_STR}).values():
+    for dim in "xyz":
+        c = subgrid.centers.dict()[dim]
         assert np.all(c == np.array([-0.5, 0.5]))
 
     _ = td.Box(size=(6, 6, 0))

tidy3d/components/base.py

@@ -192,6 +192,19 @@ def _special_characters_not_in_name(cls, values):
                     )
         return values
 
+    attrs: dict = pydantic.Field(
+        {},
+        title="Attributes",
+        description="Dictionary storing arbitrary metadata for a Tidy3D object. "
+        "This dictionary can be freely used by the user for storing data without affecting the "
+        "operation of Tidy3D as it is not used internally. "
+        "Note that, unlike regular Tidy3D fields, ``attrs`` are mutable. "
+        "For example, the following is allowed for setting an ``attr`` ``obj.attrs['foo'] = bar``. "
+        "Also note that `Tidy3D`` will raise a ``TypeError`` if ``attrs`` contain objects "
+        "that can not be serialized. One can check if ``attrs`` are serializable "
+        "by calling ``obj.json()``.",
+    )
+
     def copy(self, **kwargs) -> Tidy3dBaseModel:
         """Copy a Tidy3dBaseModel.  With ``deep=True`` as default."""
         if "deep" in kwargs and kwargs["deep"] is False:

tidy3d/components/grid/grid.py

@@ -8,7 +8,7 @@
 from ..base import Tidy3dBaseModel
 from ..data.data_array import DataArray, SpatialDataArray, ScalarFieldDataArray
 from ..data.dataset import UnstructuredGridDatasetType, UnstructuredGridDataset
-from ..types import ArrayFloat1D, Axis, TYPE_TAG_STR, InterpMethod, Literal
+from ..types import ArrayFloat1D, Axis, InterpMethod, Literal
 from ..geometry.base import Box
 
 from ...exceptions import SetupError
@@ -43,7 +43,7 @@ class Coords(Tidy3dBaseModel):
     @property
     def to_dict(self):
         """Return a dict of the three Coord1D objects as numpy arrays."""
-        return {key: np.array(value) for key, value in self.dict(exclude={TYPE_TAG_STR}).items()}
+        return {key: self.dict()[key] for key in "xyz"}
 
     @property
     def to_list(self):
@@ -386,9 +386,7 @@ def num_cells(self) -> Tuple[int, int, int]:
         >>> grid = Grid(boundaries=coords)
         >>> Nx, Ny, Nz = grid.num_cells
         """
-        return [
-            len(coords1d) - 1 for coords1d in self.boundaries.dict(exclude={TYPE_TAG_STR}).values()
-        ]
+        return [len(self.boundaries.dict()[dim]) - 1 for dim in "xyz"]
 
     @property
     def _primal_steps(self) -> Coords:
@@ -412,7 +410,7 @@ def _dual_steps(self) -> Coords:
             applied.
         """
 
-        primal_steps = self._primal_steps.dict(exclude={TYPE_TAG_STR})
+        primal_steps = {dim: self._primal_steps.dict()[dim] for dim in "xyz"}
         dsteps = {key: (psteps + np.roll(psteps, 1)) / 2 for (key, psteps) in primal_steps.items()}
 
         return Coords(**dsteps)

tidy3d/components/viz.py

@@ -95,7 +95,7 @@ def include_kwargs(self, **kwargs) -> PlotParams:
     def to_kwargs(self) -> dict:
         """Export the plot parameters as kwargs dict that can be supplied to plot function."""
         kwarg_dict = self.dict()
-        for ignore_key in ("type",):
+        for ignore_key in ("type", "attrs"):
             kwarg_dict.pop(ignore_key)
         return kwarg_dict