feat(viz): Tidy3D Matplotlib Styles on Import (#2428)

* add automatic tidy3d matplotlib styles

* refactor utils to remove get_hex_formatted_string function that is used only once

* add unit test for setting the matplotlib styles on import

* add changelog entry and export for the the reset_previous_style function

* add docstring to rcParams test

* add test for validating the material library rich representation

* cleanup unused imports

* fix formatting issues

* move viz.py into its own module to prepare for refactoring

* add the flexcompute tint palette

* cleanup viz structure

* refactor: split up large viz.py and add documentation

* refactor: re-add fix to plot_sim_3d

* fix: rename reset_previous_style to restore_matplotlib_rcparams

* fix: typo

* fix: move tidy3d style into mplstyle file

* revert typo

* fix: resolve merge conflict

* fix: use exact module

* fix: lint issues

Author: frederikschubertflex

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `eps_lim` keyword argument to `Simulation.plot_eps()` for manual control over the permittivity color limits.
 - Added `thickness` parameter to `LossyMetalMedium` for computing surface impedance of a thin conductor.
 - `priority` field in `Structure` and `MeshOverrideStructure` for setting the behavior in structure overlapping region. When its value is `None`, the priority is automatically determined based on the material property and simulation's `structure_priority_mode`.
+- Automatically apply `matplotlib` styles when importing `tidy3d` which can be reverted via the `td.restore_matplotlib_rcparams()` function.
 
 ### Fixed
 - Fixed bug in broadband adjoint source creation when forward simulation had a pulse amplitude greater than 1 or a nonzero pulse phase.

docs/api/index.rst

@@ -33,6 +33,7 @@ API |:computer:|
     constants
     abstract_base
     abstract_models
+    viz
 
 .. include:: /api/simulation.rst
 .. include:: /api/boundary_conditions.rst
@@ -58,4 +59,5 @@ API |:computer:|
 .. include:: /api/plugins/index.rst
 .. include:: /api/constants.rst
 .. include:: /api/abstract_base.rst
-.. include:: /api/abstract_models.rst
+.. include:: /api/abstract_models.rst
+.. include:: /api/viz.rst

docs/api/viz.rst

@@ -0,0 +1,11 @@
+.. currentmodule:: tidy3d
+
+Visualization and Plotting
+==========================
+
+.. autosummary::
+    :toctree: _autosummary/
+    :template: module.rst
+
+    tidy3d.restore_matplotlib_rcparams
+    tidy3d.VisualizationSpec

pyproject.toml

@@ -17,6 +17,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 documentation = "https://docs.flexcompute.com/projects/tidy3d/en/latest/"
+include = [{ path = "tidy3d/style.mplstyle", format = ["sdist", "wheel"] }]
 
 [tool.poetry.urls]
 "Bug Tracker" = "https://github.com/flexcompute/tidy3d/issues"
@@ -82,8 +83,8 @@ pyswarms = { version = "*", optional = true }
 
 # pytorch
 torch = [
-    { version = "^2.2.0", source = "PyPI", platform = "darwin", optional = true },
-    { version = "^2.2.0", source = "torch-cpu", platform = "!=darwin", optional = true },
+  { version = "^2.2.0", source = "PyPI", platform = "darwin", optional = true },
+  { version = "^2.2.0", source = "torch-cpu", platform = "!=darwin", optional = true },
 ]
 
 # scikit-rf
@@ -275,7 +276,7 @@ ignore = [
 # TODO: remove --assert=plain when https://github.com/scipy/scipy/issues/22236 is resolved
 addopts = "--cov=tidy3d --doctest-modules -n auto --dist worksteal --assert=plain -m 'not numerical' "
 markers = [
-    "numerical: marks numerical tests for adjoint gradients that require running simulations (deselect with '-m \"not numerical\"')",
+  "numerical: marks numerical tests for adjoint gradients that require running simulations (deselect with '-m \"not numerical\"')",
 ]
 env = ["MPLBACKEND=Agg", "OMP_NUM_THREADS=1"]
 doctest_optionflags = "NORMALIZE_WHITESPACE ELLIPSIS"

tests/test_components/test_viz.py

@@ -6,7 +6,7 @@
 import pytest
 import tidy3d as td
 from tidy3d import Box, Medium, Simulation, Structure
-from tidy3d.components.viz import Polygon, set_default_labels_and_title
+from tidy3d.components.viz import Polygon, restore_matplotlib_rcparams, set_default_labels_and_title
 from tidy3d.constants import inf
 from tidy3d.exceptions import Tidy3dKeyError
 
@@ -222,7 +222,7 @@ def plot_with_multi_viz_spec(alphas, facecolors, edgecolors, rng, use_viz_spec=T
 
 def test_no_matlab_install(monkeypatch):
     """Test that the `VisualizationSpec` only throws a warning on validation if matplotlib is not installed."""
-    monkeypatch.setattr("tidy3d.components.viz.MATPLOTLIB_IMPORTED", False)
+    monkeypatch.setattr("tidy3d.components.viz.visualization_spec.MATPLOTLIB_IMPORTED", False)
 
     EXPECTED_WARNING_MSG_PIECE = (
         "matplotlib was not successfully imported, but is required to validate colors"
@@ -330,3 +330,10 @@ def test_sim_plot_structures_fill():
     for patch in structure_patches[:1]:
         assert patch.get_fill(), "Should be filled when True"
         assert patch.get_facecolor() != "none", "Face color should be set"
+
+
+def test_tidy3d_matplotlib_style_application_on_import():
+    """Test restore_matplotlib_rcparams() to reset the automatically applied matplotlib.rcParams"""
+    assert mpl.rcParams.get("axes.edgecolor") == "#ECEBEA"
+    restore_matplotlib_rcparams()
+    assert mpl.rcParams.get("axes.edgecolor") == mpl.rcParamsDefault.get("axes.edgecolor")

tests/test_material_library/test_material_library.py

@@ -1,6 +1,7 @@
 """Tests material library functions and pretty printing"""
 
 import tidy3d as td
+from rich.console import Console
 from tidy3d.material_library.material_library import MaterialItemUniaxial
 
 
@@ -9,6 +10,12 @@ def test_material_library_summary():
     print(td.material_library)
 
 
+def test_material_library_rich_console():
+    """Test the rich representation of the material library which validates its styles etc."""
+    console = Console()
+    console.print(td.material_library)
+
+
 def test_material_summary():
     """Test the string method for each material in the material library."""
 

tidy3d/__init__.py

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