fix(tidy3d): FXC-4541-avoid-mpl-import-for-wasm

Author: marcorudolphflex

tests/test_components/test_scene.py

@@ -9,6 +9,7 @@
 import pytest
 
 import tidy3d as td
+import tidy3d.components.scene as scene_mod
 from tidy3d.components.scene import MAX_NUM_MEDIUMS
 from tidy3d.components.viz import STRUCTURE_EPS_CMAP, STRUCTURE_EPS_CMAP_R
 from tidy3d.exceptions import SetupError
@@ -200,6 +201,45 @@ def test_structure_eps_color_mapping():
     assert np.allclose(pp_max_reverse.facecolor, expected_max_reverse)
 
 
+@pytest.mark.parametrize(
+    "medium, eps_min, eps_max, reverse, expected",
+    [
+        pytest.param(
+            td.Medium(permittivity=1.0), 1.0, 5.0, False, (1.0, 1.0, 1.0, 1.0), id="min-forward"
+        ),
+        pytest.param(
+            td.Medium(permittivity=5.0), 1.0, 5.0, False, (0.0, 0.0, 0.0, 1.0), id="max-forward"
+        ),
+        pytest.param(
+            td.Medium(permittivity=1.0), 1.0, 5.0, True, (0.0, 0.0, 0.0, 1.0), id="min-reverse"
+        ),
+        pytest.param(
+            td.Medium(permittivity=5.0), 1.0, 5.0, True, (1.0, 1.0, 1.0, 1.0), id="max-reverse"
+        ),
+        pytest.param(
+            td.Medium(permittivity=2.0), 2.0, 2.0, False, (0.5, 0.5, 0.5, 1.0), id="equal-forward"
+        ),
+        pytest.param(
+            td.Medium(permittivity=2.0), 2.0, 2.0, True, (0.5, 0.5, 0.5, 1.0), id="equal-reverse"
+        ),
+    ],
+)
+def test_structure_eps_color_mapping_no_matplotlib(
+    monkeypatch, medium, eps_min, eps_max, reverse, expected
+):
+    monkeypatch.setattr(scene_mod, "mpl", None)
+
+    params = SCENE_FULL._get_structure_eps_plot_params(
+        medium=medium,
+        freq=1,
+        eps_min=eps_min,
+        eps_max=eps_max,
+        reverse=reverse,
+    )
+
+    assert np.allclose(params.facecolor, expected)
+
+
 def test_num_mediums():
     """Make sure we error if too many mediums supplied."""
 

tidy3d/components/scene.py

@@ -11,7 +11,7 @@
     import matplotlib.pylab as plt
     from mpl_toolkits.axes_grid1 import make_axes_locatable
 except ImportError:
-    pass
+    mpl = None
 import pydantic.v1 as pd
 
 from tidy3d.components.material.tcad.charge import (
@@ -1446,14 +1446,24 @@ def _get_structure_eps_plot_params(
             plot_params = plot_params.copy(update={"edgecolor": "k", "linewidth": 1})
         else:
             eps_medium = medium._eps_plot(frequency=freq, eps_component=eps_component)
-            active_norm = (
-                norm if norm is not None else mpl.colors.Normalize(vmin=eps_min, vmax=eps_max)
-            )
-            color_value = float(active_norm(eps_medium))
+            if norm is not None:
+                color_value = float(norm(eps_medium))
+            elif mpl is not None:
+                active_norm = mpl.colors.Normalize(vmin=eps_min, vmax=eps_max)
+                color_value = float(active_norm(eps_medium))
+            else:
+                if eps_max == eps_min:
+                    color_value = 0.5
+                else:
+                    color_value = (eps_medium - eps_min) / (eps_max - eps_min)
             color_value = min(1.0, max(0.0, color_value))
-            cmap_name = _get_colormap(reverse=reverse)
-            cmap = mpl.cm.get_cmap(cmap_name)
-            rgba = tuple(float(component) for component in cmap(color_value))
+            if mpl is not None:
+                cmap_name = _get_colormap(reverse=reverse)
+                cmap = mpl.cm.get_cmap(cmap_name)
+                rgba = tuple(float(component) for component in cmap(color_value))
+            else:
+                gray_value = color_value if reverse else 1.0 - color_value
+                rgba = (gray_value, gray_value, gray_value, 1.0)
             plot_params = plot_params.copy(update={"facecolor": rgba})
 
         return plot_params