feature[frontend]: improve printing of materials and material library

Author: Gregory Roberts

CHANGELOG.md

@@ -17,10 +17,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Classmethod `from_frequency_range` in `GaussianPulse` for generating a pulse whose amplitude in the frequency_range [fmin, fmax] is maximized, which is particularly useful for running broadband simulations.
 - Differentiable function `td.plugins.autograd.interpolate_spline` for 1D linear, quadratic, and cubic spline interpolation, supporting differentiation with respect to the interpolated values (`y_points`) and optional endpoint derivative constraints.
 - `SteadyEnergyBandMonitor` in the Charge solver.
+- Pretty printing enabled with `rich.print` for the material library, materials, and their variants. In notebooks, this can be accessed using `rich.print` or `display`, or by evaluating the material library, a material, or a variant in a cell.
 
 ### Changed
 - Performance enhancement for adjoint gradient calculations by optimizing field interpolation.
 - Auto grid in EME simulations with multiple `freqs` provided uses the largest instead of raising an error.
+- Named mediums now display by name for brevity; materials/variants print concise summaries including references.
 
 ### Fixed
 - Fixed `reverse` property of `td.Scene.plot_structures_property()` to also reverse the colorbar.

tests/test_material_library/__init__.py

@@ -0,0 +1,132 @@
+"""Tests material library functions and pretty printing"""
+
+import tidy3d as td
+from tidy3d.material_library.material_library import MaterialItemUniaxial
+
+
+def test_material_library_summary():
+    """Test to make sure we can print the material library without error."""
+    print(td.material_library)
+
+
+def test_material_summary():
+    """Test the string method for each material in the material library."""
+
+    for _, material in td.material_library.items():
+        print(material)
+
+
+def test_variant_summary():
+    """Test the string method for each variant in the material library."""
+
+    for _, material in td.material_library.items():
+        # graphene in the material library is run differently than the other materials and
+        # doesn't have the variant structure so we exclude any materials that are in this
+        # format
+        if hasattr(material, "variants"):
+            for variant in material.variants:
+                print(variant)
+
+
+def test_material_library_medium_repr():
+    """Test the new repr method does not error for material library variants."""
+
+    for material_key, material in td.material_library.items():
+        if hasattr(material, "variants"):
+            for variant_key in material.variants:
+                mat = td.material_library[material_key]
+                med = td.material_library[material_key][variant_key]
+                print(med)
+                print(repr(med))
+
+                if (type(med) is not td.MultiPhysicsMedium) and (
+                    type(mat) is not MaterialItemUniaxial
+                ):
+                    assert med.__repr__() == med.name, "Expected repr to return just the name"
+
+
+def test_medium_repr():
+    """Test the new repr method does not error for regular media with and without names and
+    that names are returned correctly by repr when they exist."""
+
+    material_name = "material"
+    test_media = [
+        td.Medium(permittivity=1.5**2),
+        td.Medium(permittivity=1.5**2, name=material_name),
+        td.material_library["SiO2"]["Horiba"].updated_copy(name=None),
+    ]
+    noname_medium_in_dict = {"medium": test_media[0]}
+
+    str_noname_medium = str(test_media[0])
+    repr_noname_medium = test_media[0].__repr__()
+    str_noname_medium_dict = str(noname_medium_in_dict)
+
+    assert (
+        "type='Medium' permittivity=2.25 conductivity=0.0" in str_noname_medium
+    ), "Expected medium information in string"
+    assert (
+        "Medium(attrs={}, name=None, frequency_range=None" in repr_noname_medium
+    ), "Expcted medium information in repr"
+    assert repr_noname_medium in str_noname_medium_dict, "Expected repr in dictionary string"
+
+    for medium in test_media:
+        repr_str = medium.__repr__()
+
+    assert (
+        test_media[1].__repr__() == material_name
+    ), "Expected repr to return just the material name."
+
+
+def test_variant_str():
+    """Test one of the materials for some expected output in variant printing."""
+
+    printed_SiO2 = str(td.material_library["SiO2"].variants["Horiba"])
+
+    assert "eps_inf: 1.0" in printed_SiO2, "Expected eps_inf in SiO2 printed string"
+    assert "poles: 1" in printed_SiO2, "Expected 1 pole in SiO2 printed string"
+
+    printed_SiO2_Palik_lossless = str(td.material_library["SiO2"].variants["Palik_Lossless"])
+
+    assert (
+        "eps_inf: 1.5385442336875639" in printed_SiO2_Palik_lossless
+    ), "Expected eps_inf in SiO2 printed string"
+    assert "poles: 2" in printed_SiO2_Palik_lossless, "Expected 1 pole in SiO2 printed string"
+
+    printed_SiO2_Palik_lossy = str(td.material_library["SiO2"].variants["Palik_Lossy"])
+
+    assert (
+        "eps_inf: 2.1560362571240765" in printed_SiO2_Palik_lossy
+    ), "Expected eps_inf in SiO2 printed string"
+    assert "poles: 5" in printed_SiO2_Palik_lossy, "Expected 1 pole in SiO2 printed string"
+
+
+def test_material_str():
+    """Test one of the materials for some expected output in variant printing."""
+
+    printed_Ag = str(td.material_library["Ag"])
+
+    assert (
+        "Default Variant: Rakic1998BB" in printed_Ag
+    ), "Expected default variant in printed string"
+    assert "RakicLorentzDrude1998" in printed_Ag, "Expected variant in printed string"
+
+    printed_Au = str(td.material_library["Au"])
+
+    assert "Olmon2012evaporated" in printed_Au, "Expected default variant in printed string"
+    assert "Olmon2012evaporated" in printed_Au, "Expected variant in printed string"
+
+
+def test_material_library_str():
+    """Test the material library string method for expected output."""
+
+    printed_library = str(td.material_library)
+
+    assert (
+        "Key: Polycarbonate, Name: Polycarbonate, Default Variant: Sultanova2009, # Variants: 2"
+        in printed_library
+    ), "Expected information in material library"
+    assert (
+        "- Key: WS2, Name: Tungsten Disulfide, Default Variant: Li2014, # Variants: 1"
+        in printed_library
+    ), "Expected information in material library"
+    assert "- Key: graphene" in printed_library, "Expected information in material library"

tests/test_material_library/test_material_library.py

@@ -1418,6 +1418,13 @@ def derivative_eps_complex_volume(
 
         return vjp_value.sum("f")
 
+    def __repr__(self):
+        """If the medium has a name, use it as the representation. Otherwise, use the default representation."""
+        if self.name:
+            return self.name
+        else:
+            return super().__repr__()
+
 
 class AbstractCustomMedium(AbstractMedium, ABC):
     """A spatially varying medium."""
@@ -7244,7 +7251,7 @@ def from_dispersive_medium(cls, medium: DispersiveMedium, thickness: float) -> M
             The 2D equivalent of the given 3D medium.
         """
         med = cls._weighted_avg([medium], [thickness])
-        return Medium2D(ss=med, tt=med, frequency_range=medium.frequency_range)
+        return Medium2D(ss=med, tt=med, frequency_range=medium.frequency_range, name=medium.name)
 
     @classmethod
     def from_anisotropic_medium(

tidy3d/components/medium.py

@@ -40,6 +40,14 @@
 # Arrow length in inches
 ARROW_LENGTH = 0.3
 
+FLEXCOMPUTE_COLORS = {
+    "brand_green": 0x00643C,
+    "brand_tan": 0xB8A18B,
+    "brand_blue": 0x6DB5DD,
+    "brand_purple": 0x8851AD,
+    "brand_black": 0x000000,
+}
+
 """ Decorators """