super_medium added to custom medium & return self for unperturbed medium

Author: weiliangjin2021

tests/test_components/test_perturbation_medium.py

@@ -362,3 +362,56 @@ def test_correct_values(dispersive):
 
     assert np.isclose(si_n + pp_large_sampled, si_index_perturb_n)
     assert np.isclose(si_k + pp_small_sampled, si_index_perturb_k)
+
+
+@pytest.mark.parametrize("unstructured", [False, True])
+def test_from_medium_field(unstructured):
+    """Test that super_medium field is properly set when calling perturbed_copy."""
+    # Setup fields to sample at
+    coords = {"x": [1, 2], "y": [3, 4], "z": [5, 6]}
+    temperature = td.SpatialDataArray(300 * np.ones((2, 2, 2)), coords=coords)
+    electron_density = td.SpatialDataArray(1e18 * np.ones((2, 2, 2)), coords=coords)
+    hole_density = td.SpatialDataArray(2e18 * np.ones((2, 2, 2)), coords=coords)
+
+    if unstructured:
+        temperature = cartesian_to_unstructured(temperature, seed=7747)
+        electron_density = cartesian_to_unstructured(electron_density, seed=7747)
+        hole_density = cartesian_to_unstructured(hole_density, seed=7747)
+
+    # Test PerturbationMedium
+    pp_real = td.ParameterPerturbation(
+        heat=td.LinearHeatPerturbation(
+            coeff=-0.01,
+            temperature_ref=300,
+            temperature_range=(200, 500),
+        ),
+    )
+
+    pmed = td.PerturbationMedium(permittivity=10, permittivity_perturbation=pp_real)
+
+    # Test without any perturbation data (returns self)
+    cmed_no_perturb = pmed.perturbed_copy()
+    assert isinstance(cmed_no_perturb, td.PerturbationMedium)
+
+    # Test with perturbation data (returns CustomMedium)
+    cmed_with_perturb = pmed.perturbed_copy(temperature, electron_density, hole_density)
+    assert isinstance(cmed_with_perturb, td.CustomMedium)
+    assert cmed_with_perturb.super_medium is pmed
+    assert hash(cmed_with_perturb.super_medium) == hash(pmed)
+
+    # Test PerturbationPoleResidue
+    pmed_pole = td.PerturbationPoleResidue(
+        eps_inf=10,
+        poles=[(1j, 3), (2j, 4)],
+        eps_inf_perturbation=pp_real,
+    )
+
+    # Test without any perturbation data (returns self)
+    cmed_pole_no_perturb = pmed_pole.perturbed_copy()
+    assert isinstance(cmed_pole_no_perturb, td.PerturbationPoleResidue)
+
+    # Test with perturbation data (returns CustomPoleResidue)
+    cmed_pole_with_perturb = pmed_pole.perturbed_copy(temperature, electron_density, hole_density)
+    assert isinstance(cmed_pole_with_perturb, td.CustomPoleResidue)
+    assert cmed_pole_with_perturb.super_medium is pmed_pole
+    assert hash(cmed_pole_with_perturb.super_medium) == hash(pmed_pole)

tidy3d/components/medium.py

@@ -91,6 +91,7 @@
     PermittivityComponent,
     PoleAndResidue,
     TensorReal,
+    annotate_type,
 )
 from .validators import _warn_potential_error, validate_name_str, validate_parameter_perturbation
 from .viz import VisualizationSpec, add_ax_if_none
@@ -875,6 +876,12 @@ class AbstractCustomMedium(AbstractMedium, ABC):
         "intersection interfaces with other structures.",
     )
 
+    super_medium: Optional[annotate_type(PerturbationMediumType)] = pd.Field(
+        None,
+        title="Super Medium",
+        description="If not ``None``, it records the super medium from which this medium was derived.",
+    )
+
     @cached_property
     @abstractmethod
     def is_isotropic(self) -> bool:
@@ -6748,7 +6755,7 @@ def perturbed_copy(
         electron_density: CustomSpatialDataType = None,
         hole_density: CustomSpatialDataType = None,
         interp_method: InterpMethod = "linear",
-    ) -> Union[Medium, CustomMedium]:
+    ) -> Union[PerturbationMedium, CustomMedium]:
         """Sample perturbations on provided heat and/or charge data and return 'CustomMedium'.
         Any of temperature, electron_density, and hole_density can be 'None'. If all passed
         arguments are 'None' then a 'Medium' object is returned. All provided fields must have
@@ -6780,10 +6787,14 @@ def perturbed_copy(
 
         Returns
         -------
-        Union[Medium, CustomMedium]
+        Union[PerturbationMedium, CustomMedium]
             Medium specification after application of heat and/or charge data.
         """
 
+        # in the absence of perturbation
+        if all(x is None for x in [temperature, electron_density, hole_density]):
+            return self
+
         new_dict = self.dict(
             exclude={
                 "permittivity_perturbation",
@@ -6793,10 +6804,6 @@ def perturbed_copy(
             }
         )
 
-        if all(x is None for x in [temperature, electron_density, hole_density]):
-            new_dict.pop("subpixel")
-            return Medium.parse_obj(new_dict)
-
         permittivity_field = self.permittivity + ParameterPerturbation._zeros_like(
             temperature, electron_density, hole_density
         )
@@ -6836,6 +6843,7 @@ def perturbed_copy(
         new_dict["permittivity"] = permittivity_field
         new_dict["conductivity"] = conductivity_field
         new_dict["interp_method"] = interp_method
+        new_dict["super_medium"] = self
 
         return CustomMedium.parse_obj(new_dict)
 
@@ -6959,7 +6967,7 @@ def perturbed_copy(
         electron_density: CustomSpatialDataType = None,
         hole_density: CustomSpatialDataType = None,
         interp_method: InterpMethod = "linear",
-    ) -> Union[PoleResidue, CustomPoleResidue]:
+    ) -> Union[PerturbationPoleResidue, CustomPoleResidue]:
         """Sample perturbations on provided heat and/or charge data and return 'CustomPoleResidue'.
         Any of temperature, electron_density, and hole_density can be 'None'. If all passed
         arguments are 'None' then a 'PoleResidue' object is returned. All provided fields must have
@@ -6991,18 +6999,18 @@ def perturbed_copy(
 
         Returns
         -------
-        Union[PoleResidue, CustomPoleResidue]
+        Union[PerturbationPoleResidue, CustomPoleResidue]
             Medium specification after application of heat and/or charge data.
         """
 
+        # in the absence of perturbation
+        if all(x is None for x in [temperature, electron_density, hole_density]):
+            return self
+
         new_dict = self.dict(
             exclude={"eps_inf_perturbation", "poles_perturbation", "perturbation_spec", "type"}
         )
 
-        if all(x is None for x in [temperature, electron_density, hole_density]):
-            new_dict.pop("subpixel")
-            return PoleResidue.parse_obj(new_dict)
-
         zeros = ParameterPerturbation._zeros_like(temperature, electron_density, hole_density)
 
         eps_inf_field = self.eps_inf + zeros
@@ -7050,12 +7058,28 @@ def perturbed_copy(
         new_dict["eps_inf"] = eps_inf_field
         new_dict["poles"] = poles_field
         new_dict["interp_method"] = interp_method
+        new_dict["super_medium"] = self
 
         return CustomPoleResidue.parse_obj(new_dict)
 
 
 # types of mediums that can be used in Simulation and Structures
 
+PerturbationMediumType = Union[PerturbationMedium, PerturbationPoleResidue]
+
+
+# Update forward references for all Custom medium classes that inherit from AbstractCustomMedium
+def _get_all_subclasses(cls):
+    """Recursively get all subclasses of a class."""
+    all_subclasses = []
+    for subclass in cls.__subclasses__():
+        all_subclasses.append(subclass)
+        all_subclasses.extend(_get_all_subclasses(subclass))
+    return all_subclasses
+
+
+for _custom_medium_cls in _get_all_subclasses(AbstractCustomMedium):
+    _custom_medium_cls.update_forward_refs()
 
 MediumType3D = Union[
     Medium,