Move some backend featurese to simulation

Author: weiliangjin2021

tidy3d/components/simulation.py

@@ -2216,15 +2216,55 @@ def _modal_plane_frames(self) -> list[Structure]:
     @cached_property
     def _finalized(self) -> Simulation:
         """Return the finalized version of the simulation setup. That is, including automatic frames around mode sources and internal absorbers, and 2d strutures converted into volumetric analogues."""
+        if (
+            len(self._modal_plane_frames) == 0
+            and not self._contains_converted_volumetric_structures
+        ):
+            return self
+        return self.updated_copy(
+            grid_spec=GridSpec.from_grid(self.grid),
+            structures=self._finalized_volumetric_structures,
+        )
 
+    @cached_property
+    def _finalized_volumetric_structures(self) -> list[Structure]:
+        """Volumetric structures in the simulation, including automatic frames around mode sources and internal absorbers, and 2d strutures converted into volumetric analogues."""
         modal_frames = self._modal_plane_frames
+        if not self._contains_converted_volumetric_structures:
+            return list(self.structures) + modal_frames
+        return list(self.volumetric_structures) + modal_frames
 
-        if len(modal_frames) == 0 and not self._contains_converted_volumetric_structures:
-            return self
+    @cached_property
+    def _finalized_optical_medium_map(self) -> dict[MediumType, pydantic.NonNegativeInt]:
+        """Returns dict mapping medium to index in material in finalized simulation.
 
-        structures = list(self.volumetric_structures) + modal_frames
+        Returns
+        -------
+        Dict[:class:`.AbstractMedium`, int]
+            Mapping between distinct mediums to index in finalized simulation.
+        """
+        medium_set = {
+            structure._optical_medium for structure in self._finalized_volumetric_structures
+        }
+        medium_set.add(Structure._get_optical_medium(self.medium))
+        return {medium: index for index, medium in enumerate(medium_set)}
 
-        return self.updated_copy(grid_spec=GridSpec.from_grid(self.grid), structures=structures)
+    @cached_property
+    def _finalized_full_optical_medium_map(self) -> dict[MediumType, pydantic.NonNegativeInt]:
+        """Returns dict mapping medium to index in material in finalized simulation, including
+        parent medium.
+        """
+        full_medium_map = dict(self._finalized_optical_medium_map)
+        next_index = len(full_medium_map)
+        for medium in self._finalized_optical_medium_map.keys():
+            if (
+                isinstance(medium, AbstractCustomMedium)
+                and medium.derived_from is not None
+                and medium.derived_from not in full_medium_map
+            ):
+                full_medium_map[medium.derived_from] = next_index
+                next_index += 1
+        return full_medium_map
 
     def _validate_finalized(self) -> None:
         """Validate that after adding pec frames simulation setup is still valid."""

tidy3d/components/structure.py

@@ -26,6 +26,7 @@
 from .geometry.base import Box, Geometry
 from .geometry.utils import GeometryType, validate_no_transformed_polyslabs
 from .grid.grid import Coords
+from .material.multi_physics import MultiPhysicsMedium
 from .material.types import StructureMediumType
 from .medium import AbstractCustomMedium, CustomMedium, LossyMetalMedium, Medium, Medium2D
 from .monitor import FieldMonitor, PermittivityMonitor
@@ -257,6 +258,16 @@ def eps_diagonal(self, frequency: float, coords: Coords) -> tuple[complex, compl
             return self.medium.eps_diagonal_on_grid(frequency=frequency, coords=coords)
         return self.medium.eps_diagonal(frequency=frequency)
 
+    @staticmethod
+    def _get_optical_medium(medium):
+        """Get optical medium."""
+        return medium.optical if isinstance(medium, MultiPhysicsMedium) else medium
+
+    @property
+    def _optical_medium(self) -> StructureMediumType:
+        """Optical medium of the structure."""
+        return self._get_optical_medium(self.medium)
+
     @pydantic.validator("medium", always=True)
     @skip_if_fields_missing(["geometry"])
     def _check_2d_geometry(cls, val, values):