feat[frontend]: viz spec implementation and testing

Author: Gregory Roberts

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `VisualizationSpec` that allows `Medium` instances to specify color and transparency plotting attributes that override default ones.
+
 ### Changed
 - `ModeMonitor` and `ModeSolverMonitor` now use the default `td.ModeSpec()` with `num_modes=1` when `mode_spec` is not provided.
 

tests/test_components/test_viz.py

@@ -1,6 +1,7 @@
 """Tests visualization operations."""
 
 import matplotlib.pyplot as plt
+import pydantic.v1 as pd
 import pytest
 import tidy3d as td
 from tidy3d.components.viz import Polygon, set_default_labels_and_title
@@ -100,3 +101,151 @@ def test_set_default_labels_title():
         ax = set_default_labels_and_title(
             axis_labels=axis_labels, axis=2, position=0, ax=ax, plot_length_units="inches"
         )
+
+    plt.close()
+
+
+def test_make_viz_spec():
+    """
+    Tests core visualizaton spec creation.
+    """
+    viz_spec = td.VisualizationSpec(facecolor="red", edgecolor="green", alpha=0.5)
+    viz_spec = td.VisualizationSpec(facecolor="red", alpha=0.5)
+
+
+def test_unallowed_colors():
+    """
+    Tests validator for visualization spec for colors not recognized by matplotlib.
+    """
+    with pytest.raises(pd.ValidationError):
+        _ = td.VisualizationSpec(facecolor="rr", edgecolor="green", alpha=0.5)
+    with pytest.raises(pd.ValidationError):
+        _ = td.VisualizationSpec(facecolor="red", edgecolor="gg", alpha=0.5)
+
+
+def test_unallowed_alpha():
+    """
+    Tests validator for disallowed alpha values.
+    """
+    with pytest.raises(pd.ValidationError):
+        _ = td.VisualizationSpec(facecolor="red", edgecolor="green", alpha=-0.5)
+    with pytest.raises(pd.ValidationError):
+        _ = td.VisualizationSpec(facecolor="red", edgecolor="green", alpha=2.5)
+
+
+def test_plot_from_structure():
+    """
+    Tests visualization spec can be added to medium and structure plotting function can be run.
+    """
+    viz_spec = td.VisualizationSpec(facecolor="blue", edgecolor="pink", alpha=0.5)
+    medium = td.Medium(permittivity=2.25, viz_spec=viz_spec)
+    geometry = td.Box(size=(2, 0, 2))
+
+    structure = td.Structure(geometry=geometry, medium=medium)
+
+    structure.plot(z=0)
+    plt.close()
+
+
+def plot_with_viz_spec(alpha, facecolor, edgecolor=None, use_viz_spec=True):
+    """
+    Helper function for locally testing different visualization specs in structures through
+    structure plotting function.
+    """
+    if edgecolor is None:
+        viz_spec = td.VisualizationSpec(facecolor=facecolor, alpha=alpha)
+    else:
+        viz_spec = td.VisualizationSpec(facecolor=facecolor, edgecolor=edgecolor, alpha=alpha)
+
+    medium = td.Medium(permittivity=2.25)
+    if use_viz_spec:
+        medium = td.Medium(permittivity=2.25, viz_spec=viz_spec)
+
+    geometry = td.Box(size=(2, 4, 2))
+
+    structure = td.Structure(geometry=geometry, medium=medium)
+
+    structure.plot(z=1)
+    plt.show()
+
+
+def plot_with_multi_viz_spec(alphas, facecolors, edgecolors, rng, use_viz_spec=True):
+    """
+    Helper function for plotting simulations with multiple visulation specs via simluation
+    plotting function.
+    """
+    viz_specs = [
+        td.VisualizationSpec(
+            facecolor=facecolors[idx], edgecolor=edgecolors[idx], alpha=alphas[idx]
+        )
+        for idx in range(0, len(alphas))
+    ]
+    media = [td.Medium(permittivity=2.25) for idx in range(0, len(viz_specs))]
+    if use_viz_spec:
+        media = [
+            td.Medium(permittivity=2.25, viz_spec=viz_specs[idx])
+            for idx in range(0, len(viz_specs))
+        ]
+
+    structures = []
+    for idx in range(0, len(viz_specs)):
+        center = tuple(list(rng.uniform(-3, 3, 2)) + [0])
+        size = tuple(rng.uniform(1, 2, 3))
+        box = td.Box(center=center, size=size)
+
+        structures.append(td.Structure(geometry=box, medium=media[idx]))
+
+    sim = td.Simulation(
+        size=(10.0, 10.0, 10.0),
+        run_time=1e-12,
+        structures=structures,
+        grid_spec=td.GridSpec(wavelength=1.0),
+    )
+
+    sim.plot(z=0.0)
+    plt.show()
+
+
+@pytest.mark.skip(reason="Skipping test for CI, but useful for debugging locally with graphics.")
+def test_plot_from_structure_local():
+    """
+    Local test for visualizing output when specifying visualization spec.
+    """
+    plot_with_viz_spec(alpha=0.5, facecolor="red", edgecolor="blue")
+    plot_with_viz_spec(alpha=0.1, facecolor="magenta", edgecolor="cyan")
+    plot_with_viz_spec(alpha=0.9, facecolor="darkgreen", edgecolor="black")
+    plot_with_viz_spec(alpha=0.8, facecolor="brown", edgecolor="deepskyblue")
+    plot_with_viz_spec(alpha=0.2, facecolor="brown", edgecolor="deepskyblue")
+    plot_with_viz_spec(alpha=1.0, facecolor="green")
+    plot_with_viz_spec(alpha=0.75, facecolor="red", edgecolor="blue")
+    plot_with_viz_spec(alpha=0.75, facecolor="red", edgecolor="blue", use_viz_spec=False)
+
+    with pytest.raises(pd.ValidationError):
+        plot_with_viz_spec(alpha=0.5, facecolor="dark green", edgecolor="blue")
+    with pytest.raises(pd.ValidationError):
+        plot_with_viz_spec(alpha=0.5, facecolor="red", edgecolor="ble")
+    with pytest.raises(pd.ValidationError):
+        plot_with_viz_spec(alpha=1.5, facecolor="red", edgecolor="blue")
+    with pytest.raises(pd.ValidationError):
+        plot_with_viz_spec(alpha=-0.5, facecolor="red", edgecolor="blue")
+
+
+@pytest.mark.skip(reason="Skipping test for CI, but useful for debugging locally with graphics.")
+def test_plot_multi_from_structure_local(rng):
+    """
+    Local test for visualizing output when creating multiple structures with variety of
+    visualization specs.
+    """
+    plot_with_multi_viz_spec(
+        alphas=[0.5, 0.75, 0.25, 0.4],
+        facecolors=["red", "green", "blue", "orange"],
+        edgecolors=["black", "cyan", "magenta", "brown"],
+        rng=rng,
+    )
+    plot_with_multi_viz_spec(
+        alphas=[0.5, 0.75, 0.25, 0.4],
+        facecolors=["red", "green", "blue", "orange"],
+        edgecolors=["black", "cyan", "magenta", "brown"],
+        rng=rng,
+        use_viz_spec=False,
+    )

tidy3d/__init__.py

@@ -284,6 +284,7 @@
     SpaceTimeModulation,
 )
 from .components.transformation import RotationAroundAxis
+from .components.viz import VisualizationSpec
 
 # config
 from .config import config
@@ -534,6 +535,7 @@ def set_logging_level(level: str) -> None:
     "HeuristicPECStaircasing",
     "PECConformal",
     "SurfaceImpedance",
+    "VisualizationSpec",
     "EMESimulation",
     "EMESimulationData",
     "EMEMonitor",

tidy3d/components/geometry/base.py

@@ -47,6 +47,7 @@
     ARROW_LENGTH,
     PLOT_BUFFER,
     PlotParams,
+    VisualizationSpec,
     add_ax_if_none,
     arrow_style,
     equal_aspect,
@@ -446,6 +447,7 @@ def plot(
         z: float = None,
         ax: Ax = None,
         plot_length_units: LengthUnit = None,
+        viz_spec: VisualizationSpec = None,
         **patch_kwargs,
     ) -> Ax:
         """Plot geometry cross section at single (x,y,z) coordinate.
@@ -462,6 +464,8 @@ def plot(
             Matplotlib axes to plot on, if not specified, one is created.
         plot_length_units : LengthUnit = None
             Specify units to use for axis labels, tick labels, and the title.
+        viz_spec : VisualizationSpec = None
+            Plotting parameters associated with a medium to use instead of defaults.
         **patch_kwargs
             Optional keyword arguments passed to the matplotlib patch plotting of structure.
             For details on accepted values, refer to
@@ -477,7 +481,10 @@ def plot(
         axis, position = self.parse_xyz_kwargs(x=x, y=y, z=z)
         shapes_intersect = self.intersections_plane(x=x, y=y, z=z)
 
-        plot_params = self.plot_params.include_kwargs(**patch_kwargs)
+        plot_params = self.plot_params
+        if viz_spec is not None:
+            plot_params = plot_params.override_with_viz_spec(viz_spec)
+        plot_params = plot_params.include_kwargs(**patch_kwargs)
 
         # for each intersection, plot the shape
         for shape in shapes_intersect:

tidy3d/components/medium.py

@@ -82,7 +82,7 @@
     TensorReal,
 )
 from .validators import _warn_potential_error, validate_name_str, validate_parameter_perturbation
-from .viz import add_ax_if_none
+from .viz import VisualizationSpec, add_ax_if_none
 
 # evaluate frequency as this number (Hz) if inf
 FREQ_EVAL_INF = 1e50
@@ -646,6 +646,12 @@ class AbstractMedium(ABC, Tidy3dBaseModel):
         description="Modulation spec applied on top of the base medium properties.",
     )
 
+    viz_spec: Optional[VisualizationSpec] = pd.Field(
+        None,
+        title="Visualization Specification",
+        description="Plotting specification for visualizing medium.",
+    )
+
     @cached_property
     def _nonlinear_models(self) -> NonlinearSpec:
         """The nonlinear models in the nonlinear_spec."""

tidy3d/components/scene.py

@@ -466,6 +466,8 @@ def _get_structure_plot_params(self, mat_index: int, medium: Medium) -> PlotPara
             # regular medium
             facecolor = MEDIUM_CMAP[(mat_index - 1) % len(MEDIUM_CMAP)]
             plot_params = plot_params.copy(update={"facecolor": facecolor})
+            if medium.viz_spec is not None:
+                plot_params = plot_params.override_with_viz_spec(medium.viz_spec)
 
         return plot_params
 
@@ -1068,6 +1070,8 @@ def _get_structure_eps_plot_params(
         """Constructs the plot parameters for a given medium in scene.plot_eps()."""
 
         plot_params = plot_params_structure.copy(update={"linewidth": 0})
+        if medium.viz_spec is not None:
+            plot_params = plot_params.override_with_viz_spec(medium.viz_spec)
         if alpha is not None:
             plot_params = plot_params.copy(update={"alpha": alpha})
 
@@ -1390,6 +1394,8 @@ def _get_structure_heat_charge_property_plot_params(
         """
 
         plot_params = plot_params_structure.copy(update={"linewidth": 0})
+        if medium.viz_spec is not None:
+            plot_params = plot_params.override_with_viz_spec(medium.viz_spec)
         if alpha is not None:
             plot_params = plot_params.copy(update={"alpha": alpha})
 

tidy3d/components/structure.py

@@ -136,7 +136,9 @@ def plot(
         matplotlib.axes._subplots.Axes
             The supplied or created matplotlib axes.
         """
-        return self.geometry.plot(x=x, y=y, z=z, ax=ax, **patch_kwargs)
+        return self.geometry.plot(
+            x=x, y=y, z=z, ax=ax, viz_spec=self.medium.viz_spec, **patch_kwargs
+        )
 
 
 class Structure(AbstractStructure):

tidy3d/components/viz.py

@@ -4,11 +4,12 @@
 
 from functools import wraps
 from html import escape
-from typing import Any
+from typing import Any, Dict, Optional
 
 import matplotlib.pyplot as plt
 import matplotlib.ticker as ticker
 import pydantic.v1 as pd
+from matplotlib.colors import is_color_like
 from matplotlib.patches import ArrowStyle, PathPatch
 from matplotlib.path import Path
 from numpy import array, concatenate, inf, ones
@@ -92,6 +93,10 @@ def include_kwargs(self, **kwargs) -> AbstractPlotParams:
         }
         return self.copy(update=update_dict)
 
+    def override_with_viz_spec(self, viz_spec) -> AbstractPlotParams:
+        """Override plot params with supplied VisualizationSpec."""
+        return self.include_kwargs(**dict(viz_spec))
+
     def to_kwargs(self) -> dict:
         """Export the plot parameters as kwargs dict that can be supplied to plot function."""
         kwarg_dict = self.dict()
@@ -165,6 +170,46 @@ class PlotParams(AbstractPlotParams):
 arrow_style = ArrowStyle.Simple(head_length=12, head_width=9, tail_width=4)
 
 
+def is_valid_color(value: str) -> str:
+    if not is_color_like(value):
+        raise pd.ValidationError(f"{value} is not a valid plotting color")
+
+    return value
+
+
+class VisualizationSpec(Tidy3dBaseModel):
+    """Defines specification for visualization when used with plotting functions."""
+
+    facecolor: str = pd.Field(
+        "",
+        title="Face color",
+        description="Color applied to the faces in visualization.",
+    )
+
+    edgecolor: Optional[str] = pd.Field(
+        "",
+        title="Edge color",
+        description="Color applied to the edges in visualization.",
+    )
+
+    alpha: Optional[pd.confloat(ge=0.0, le=1.0)] = pd.Field(
+        1.0,
+        title="Opacity",
+        description="Opacity/alpha value in plotting between 0 and 1.",
+    )
+
+    @pd.validator("facecolor", always=True)
+    def validate_color(value: str) -> str:
+        return is_valid_color(value)
+
+    @pd.validator("edgecolor", always=True)
+    def validate_and_copy_color(value: str, values: Dict[str, Any]) -> str:
+        if (value == "") and "facecolor" in values:
+            return is_valid_color(values["facecolor"])
+
+        return is_valid_color(value)
+
+
 """=================================================================================================
 Descartes modified from https://pypi.org/project/descartes/ for Shapely >= 1.8.0