feat(medium): track parent medium from which the medium is derived

Author: weiliangjin2021

tests/test_components/test_perturbation_medium.py

@@ -362,3 +362,60 @@ 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 Medium)
+    cmed_no_perturb = pmed.perturbed_copy()
+    assert isinstance(cmed_no_perturb, td.Medium)
+    assert cmed_no_perturb.super_medium is pmed
+    assert hash(cmed_no_perturb.super_medium) == hash(pmed)
+
+    # 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 PoleResidue)
+    cmed_pole_no_perturb = pmed_pole.perturbed_copy()
+    assert isinstance(cmed_pole_no_perturb, td.PoleResidue)
+    assert cmed_pole_no_perturb.super_medium is pmed_pole
+    assert hash(cmed_pole_no_perturb.super_medium) == hash(pmed_pole)
+
+    # 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

@@ -197,6 +197,12 @@ class AbstractMedium(ABC, Tidy3dBaseModel):
         description="Plotting specification for visualizing medium.",
     )
 
+    super_medium: Optional["MediumType"] = pd.Field(  # noqa: UP037
+        None,
+        title="Super Medium",
+        description="If not ``None``, it records the super medium from which this medium was derived.",
+    )
+
     @cached_property
     def _nonlinear_models(self) -> list:
         """The nonlinear models in the nonlinear_spec."""
@@ -1156,6 +1162,10 @@ def _derivative_field_cmp(
 
 
 # PEC keyword
+# Resolve forward references early to allow instantiation of builtins like PEC.
+AbstractMedium.update_forward_refs(MediumType=Any)
+
+
 class PECMedium(AbstractMedium):
     """Perfect electrical conductor class.
 
@@ -6795,6 +6805,7 @@ def perturbed_copy(
 
         if all(x is None for x in [temperature, electron_density, hole_density]):
             new_dict.pop("subpixel")
+            new_dict["super_medium"] = self
             return Medium.parse_obj(new_dict)
 
         permittivity_field = self.permittivity + ParameterPerturbation._zeros_like(
@@ -6836,6 +6847,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)
 
@@ -7001,6 +7013,7 @@ def perturbed_copy(
 
         if all(x is None for x in [temperature, electron_density, hole_density]):
             new_dict.pop("subpixel")
+            new_dict["super_medium"] = self
             return PoleResidue.parse_obj(new_dict)
 
         zeros = ParameterPerturbation._zeros_like(temperature, electron_density, hole_density)
@@ -7050,6 +7063,7 @@ 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)