refactor imports and fix light mode inputs in MatterViz widgets in dark VSCode Jupyter notebooks

- fix widget asset paths in `matterviz.py`

Author: janosh

assets/scripts/cluster/composition/cluster_compositions_matbench.py

@@ -24,11 +24,7 @@
 from pymatgen.util.string import htmlify
 
 import pymatviz as pmv
-from pymatviz.cluster.composition import (
-    EmbeddingMethod,
-    matminer_featurize,
-    one_hot_encode,
-)
+import pymatviz.cluster.composition as pcc
 from pymatviz.enums import Key
 
 
@@ -64,7 +60,7 @@ def process_dataset(
     target_key: str,
     target_label: str,
     target_symbol: str,
-    embed_method: EmbeddingMethod,
+    embed_method: pcc.EmbeddingMethod,
     projection: ProjectionMethod,
     n_components: int,
     **kwargs: Any,
@@ -113,9 +109,9 @@ def process_dataset(
     if embeddings_dict is None:
         # Create embeddings
         if embed_method == "one-hot":
-            embeddings = one_hot_encode(compositions)
+            embeddings = pcc.one_hot_encode(compositions)
         elif embed_method in ["magpie", "matscholar_el"]:
-            embeddings = matminer_featurize(compositions, preset=embed_method)
+            embeddings = pcc.matminer_featurize(compositions, preset=embed_method)
         else:
             raise ValueError(f"Unknown {embed_method=}")
 
@@ -226,7 +222,9 @@ def annotate_top_points(row: pd.Series) -> dict[str, Any] | None:
     "K<sub>VRH</sub>",
 )
 plot_combinations: list[
-    tuple[str, str, str, str, EmbeddingMethod, ProjectionMethod, int, dict[str, Any]]
+    tuple[
+        str, str, str, str, pcc.EmbeddingMethod, ProjectionMethod, int, dict[str, Any]
+    ]
 ] = [
     # 1. Steels with PCA (2D) - shows clear linear trends
     (*mb_steels, "magpie", "pca", 2, dict(x=0.01, xanchor="left")),

examples/matbench/dielectric/explore_dielectric.py

@@ -14,10 +14,7 @@
 
 # %%
 import plotly.express as px
-from matbench_discovery.structure.prototype import (
-    count_wyckoff_positions,
-    get_protostructure_label,
-)
+from matbench_discovery.structure import prototype
 from matminer.datasets import load_dataset
 from tqdm import tqdm
 
@@ -36,10 +33,12 @@
 proto_label_key = f"{Key.protostructure}_moyo"
 n_wyckoff_pos_key = f"{Key.n_wyckoff_pos}_moyo"
 df_diel[proto_label_key] = [
-    get_protostructure_label(struct)
+    prototype.get_protostructure_label(struct)
     for struct in tqdm(df_diel[Key.structure], desc="Getting Wyckoff strings")
 ]
-df_diel[n_wyckoff_pos_key] = df_diel[proto_label_key].map(count_wyckoff_positions)
+df_diel[n_wyckoff_pos_key] = df_diel[proto_label_key].map(
+    prototype.count_wyckoff_positions
+)
 
 df_diel[Key.crystal_system] = df_diel[Key.spg_num].map(pmv.utils.spg_to_crystal_sys)
 

examples/matbench/log_g+kvrh/explore_log_g+krvh.py

@@ -15,10 +15,7 @@
 import numpy as np
 import pandas as pd
 import plotly.express as px
-from matbench_discovery.structure.prototype import (
-    count_wyckoff_positions,
-    get_protostructure_label,
-)
+from matbench_discovery.structure import prototype
 from matminer.datasets import load_dataset
 from pymatgen.core import Structure
 from tqdm import tqdm
@@ -37,7 +34,7 @@
 )
 df_sym[Key.crystal_system] = df_sym["number"].map(pmv.utils.spg_to_crystal_sys)
 df_grvh[Key.protostructure] = [
-    get_protostructure_label(struct)
+    prototype.get_protostructure_label(struct)
     for struct in tqdm(df_grvh[Key.structure], desc="matbench_log_gvrh Wyckoff strings")
 ]
 df_kvrh[Key.protostructure] = df_grvh[Key.protostructure]
@@ -46,7 +43,9 @@
     df[[Key.spg_num, Key.wyckoff_symbols]] = df_sym[["number", "wyckoffs"]]
     df[Key.crystal_system] = df_sym[Key.crystal_system]
 
-    df[Key.n_wyckoff_pos] = df[Key.protostructure].map(count_wyckoff_positions)
+    df[Key.n_wyckoff_pos] = df[Key.protostructure].map(
+        prototype.count_wyckoff_positions
+    )
     df[Key.formula] = [x.formula for x in df[Key.structure]]
 
 

pymatviz/structure/helpers.py

@@ -16,12 +16,7 @@
 from pymatgen.core.periodic_table import Element
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
-from pymatviz.colors import (
-    ELEM_COLORS_ALLOY,
-    ELEM_COLORS_JMOL,
-    ELEM_COLORS_PASTEL,
-    ELEM_COLORS_VESTA,
-)
+from pymatviz import colors
 from pymatviz.enums import ElemColorScheme, Key, SiteCoords
 from pymatviz.typing import Xyz
 from pymatviz.utils import df_ptable, pick_max_contrast_color
@@ -103,7 +98,7 @@ def get_struct_prop(
 )
 
 
-def _get_site_symbol(site: PeriodicSite) -> str:
+def get_site_symbol(site: PeriodicSite) -> str:
     """Get a single element symbol for a site.
 
     Handles disordered sites by picking the element with the highest fraction.
@@ -268,14 +263,8 @@ def get_elem_colors(
     """Get element colors based on the provided scheme or custom dictionary."""
     if isinstance(elem_colors, dict):
         return elem_colors
-    if str(elem_colors) == str(ElemColorScheme.jmol):
-        return ELEM_COLORS_JMOL  # type: ignore[return-value]
-    if str(elem_colors) == str(ElemColorScheme.vesta):
-        return ELEM_COLORS_VESTA  # type: ignore[return-value]
-    if str(elem_colors) == str(ElemColorScheme.alloy):
-        return ELEM_COLORS_ALLOY  # type: ignore[return-value]
-    if str(elem_colors) == str(ElemColorScheme.pastel):
-        return ELEM_COLORS_PASTEL  # type: ignore[return-value]
+    if color_dict := getattr(colors, f"ELEM_COLORS_{str(elem_colors).upper()}", None):
+        return color_dict
     raise ValueError(
         f"colors must be a dict or one of ('{', '.join(ElemColorScheme)}')"
     )
@@ -313,12 +302,12 @@ def generate_site_label(
         return None
 
     if site_labels == "symbol":
-        return _get_site_symbol(site)
+        return get_site_symbol(site)
     if site_labels == "species":
         return site.species_string  # Use full species string for disordered
 
     label_text = ""
-    symbol = _get_site_symbol(site)  # Majority element symbol of site
+    symbol = get_site_symbol(site)  # Majority element symbol of site
 
     if isinstance(site_labels, dict):
         # Use provided label for symbol, else symbol itself, or empty if not found &
@@ -423,7 +412,7 @@ def format_coord(coord_val: float) -> str:
     return out_text
 
 
-def _process_element_color(raw_color_from_map: ColorType) -> str:
+def normalize_elem_color(raw_color_from_map: ColorType) -> str:
     """Process a color from the element color map into a consistent RGB string format.
 
     Args:
@@ -538,7 +527,7 @@ def draw_site(
     raw_color_from_map = elem_colors.get(majority_species.symbol, "gray")
 
     # Process the color from the map into a string format
-    atom_color = _process_element_color(raw_color_from_map)
+    atom_color = normalize_elem_color(raw_color_from_map)
 
     site_hover_text = get_site_hover_text(site, hover_text, majority_species, float_fmt)
 
@@ -582,7 +571,7 @@ def draw_site(
         fig.add_scatter(**scatter_kwargs, row=row, col=col)
 
 
-def _create_disordered_site_legend_name(
+def get_disordered_site_legend_name(
     sorted_species: list[tuple[Species, float]], *, is_image: bool = False
 ) -> str:
     """Create a legend name for a disordered site showing all elements with occupancies.
@@ -699,7 +688,7 @@ def draw_disordered_site(
     sorted_species = sorted(species.items(), key=lambda x: x[1], reverse=True)
 
     # Create a combined legend name showing all elements with occupancies
-    legend_name = _create_disordered_site_legend_name(sorted_species, is_image=is_image)
+    legend_name = get_disordered_site_legend_name(sorted_species, is_image=is_image)
 
     # Set up legendgroup - use site_idx if not provided to group all parts together
     if legendgroup is None:
@@ -724,7 +713,7 @@ def draw_disordered_site(
             elem_symbol = element_species.symbol
             raw_color_from_map = elem_colors.get(elem_symbol, "gray")
 
-            atom_color = _process_element_color(raw_color_from_map)
+            atom_color = normalize_elem_color(raw_color_from_map)
 
             # Calculate the angular span for this species based on occupancy
             angle_span = 2 * np.pi * occupancy
@@ -735,7 +724,7 @@ def draw_disordered_site(
 
             # Generate the spherical wedge mesh
             x_coords, y_coords, z_coords, i_indices, j_indices, k_indices = (
-                _generate_spherical_wedge_mesh(
+                get_spherical_wedge_mesh(
                     center=coords,
                     radius=wedge_radius,
                     start_angle=current_angle,
@@ -847,7 +836,7 @@ def draw_disordered_site(
         for species_idx, (element_species, occupancy) in enumerate(sorted_species):
             elem_symbol = element_species.symbol
             raw_color_from_map = elem_colors.get(elem_symbol, "gray")
-            atom_color = _process_element_color(raw_color_from_map)
+            atom_color = normalize_elem_color(raw_color_from_map)
 
             # Calculate angular width for this species based on occupancy
             angular_width = 2 * math.pi * occupancy
@@ -984,7 +973,7 @@ def draw_disordered_site(
 MAX_PIE_SLICE_POINTS = 20  # Base points for 2D pie slices
 
 
-def _generate_spherical_wedge_mesh(
+def get_spherical_wedge_mesh(
     center: np.ndarray,
     radius: float,
     start_angle: float,
@@ -1496,8 +1485,8 @@ def parse_color(color_val: Any) -> str:
 
             if current_bond_color_setting is True:
                 # Default gradient: use element colors of bonded sites
-                elem1_symbol = _get_site_symbol(site1)
-                elem2_symbol = _get_site_symbol(site2)
+                elem1_symbol = get_site_symbol(site1)
+                elem2_symbol = get_site_symbol(site2)
                 color1_rgb_str = parse_color(_elem_colors.get(elem1_symbol, "gray"))
                 color2_rgb_str = parse_color(_elem_colors.get(elem2_symbol, "gray"))
                 color_for_segment_calc = (color1_rgb_str, color2_rgb_str)

pymatviz/structure/plotly.py

@@ -242,7 +242,7 @@ def structure_2d(
                 zip(struct_i, rotated_coords_all_sites, strict=False)
             ):
                 # Determine legend parameters for primary sites
-                symbol = helpers._get_site_symbol(site)
+                symbol = helpers.get_site_symbol(site)
                 legendgroup = None
                 showlegend = False
                 if site_labels == "legend":
@@ -621,7 +621,7 @@ def structure_3d(
                     else site_idx_loop % len(struct_i)
                 )
 
-                symbol = helpers._get_site_symbol(site)
+                symbol = helpers.get_site_symbol(site)
 
                 # Determine legend parameters for primary sites only
                 legendgroup = None

pymatviz/widgets/matterviz.py

@@ -34,7 +34,7 @@ def fetch_widget_asset(filename: str, version_override: str | None = None) -> st
     asset_version = version_override or PKG_VERSION
 
     # Paths
-    local_path = f"{os.path.dirname(__file__)}/build/{filename}"
+    local_path = f"{os.path.dirname(__file__)}/web/build/{filename}"
     cache_dir = f"{os.path.expanduser('~/.cache/pymatviz/build')}/{asset_version}"
     os.makedirs(cache_dir, exist_ok=True)
     cache_path = f"{cache_dir}/{filename}"

pymatviz/widgets/mime.py

@@ -89,13 +89,9 @@ def _register_renderers() -> None:
 
 def register_matterviz_widgets() -> None:
     """Register widgets for automatic display."""
-    from pymatviz.process_data import (
-        is_composition_like,
-        is_structure_like,
-        is_trajectory_like,
-    )
-
-    WIDGET_REGISTRY[is_trajectory_like] = "trajectory"
-    WIDGET_REGISTRY[is_composition_like] = "composition"
-    WIDGET_REGISTRY[is_structure_like] = "structure"
+    from pymatviz import process_data as pd
+
+    WIDGET_REGISTRY[pd.is_trajectory_like] = StructureType
+    WIDGET_REGISTRY[pd.is_composition_like] = CompositionType
+    WIDGET_REGISTRY[pd.is_structure_like] = StructureType
     _register_renderers()

pymatviz/widgets/web/anywidget.ts

@@ -32,6 +32,23 @@ function inject_app_css(theme_type?: ThemeType): void {
       background-color: transparent !important;
     }
 
+    /* Dark mode input styling for Jupyter notebooks and interactive windows in VSCode  */
+    :is(.vscode-dark, .dark-theme, [data-jp-theme-light="false"]) :is(input, textarea, select) {
+      background-color: #2d2d2d;
+      color: #ffffff;
+      border: 1px solid #555555;
+      border-radius: 4px;
+      padding: 6px 8px;
+    }
+    :is(.vscode-dark, .dark-theme, [data-jp-theme-light="false"]) :is(input, textarea, select):focus {
+      outline: none;
+      border-color: #007acc;
+      box-shadow: 0 0 0 2px rgba(0, 122, 204, 0.2);
+    }
+    :is(.vscode-dark, .dark-theme, [data-jp-theme-light="false"]) :is(input, textarea)::placeholder {
+      color: #888888;
+    }
+
     /* scope global styles to matterviz widgets to prevent site styles leaking into notebook styles */
     /* this is brittle, will break should component CSS classes in matterviz change, try to find better solution */
     div:is(.structure, .trajectory, .composition) {

pymatviz/widgets/web/deno.jsonc

@@ -1,8 +1,7 @@
 {
   "tasks": {
     "dev": "deno run -A --node-modules-dir npm:vite dev",
-    "build": "rm -rf build && deno run -A --node-modules-dir npm:vite build",
-    "ext-install": "deno task build && find . -name 'matterviz-v*.vsix' -type f | head -1 | xargs -I {} code --install-extension {}"
+    "build": "rm -rf build && deno run -A --node-modules-dir npm:vite build"
   },
   "lock": false,
   "lint": {

tests/files/.pytest-split-durations

@@ -489,7 +489,7 @@
     "tests/structure/test_structure_helpers.py::test_angles_to_rotation_matrix[30x,45y,60z-expected_shape1]": 0.00016524898819625378,
     "tests/structure/test_structure_helpers.py::test_angles_to_rotation_matrix_invalid_input": 0.00013866700464859605,
     "tests/structure/test_structure_helpers.py::test_constants": 9.791701450012624e-05,
-    "tests/structure/test_structure_helpers.py::test_create_disordered_site_legend_name": 0.00012099897139705718,
+    "tests/structure/test_structure_helpers.py::test_get_disordered_site_legend_name": 0.00012099897139705718,
     "tests/structure/test_structure_helpers.py::test_draw_bonds[False-bond_kwargs2-plotted_sites_coords_param2-0]": 0.0062469589756801724,
     "tests/structure/test_structure_helpers.py::test_draw_bonds[False-bond_kwargs4-None-16]": 0.008371708041522652,
     "tests/structure/test_structure_helpers.py::test_draw_bonds[True-None-None-160]": 0.026459541986696422,