Add IS_IPYTHON global (#198)

janosh · web-flow · commit 73bf7f870b1a · 2024-09-05T19:31:47.000Z
* add IS_IPYTHON global

* fix outdated function names in readme

* breaking: remove structure_2d()'s deprecated colors kwarg, now only elem_colors works

* fix test_structure_2d_color_warning
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -8,7 +8,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.3
+    rev: v0.6.4
     hooks:
       - id: ruff
         args: [--fix]
@@ -88,6 +88,6 @@ repos:
           - typescript-eslint
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.378
+    rev: v1.1.379
     hooks:
       - id: pyright
diff --git a/assets/histogram-ecdf.svg b/assets/histogram-ecdf.svg
diff --git a/examples/make_assets/histogram.py b/examples/make_assets/histogram.py
@@ -54,4 +54,4 @@
     pmv.powerups.add_ecdf_line(fig, trace_idx=idx)
 fig.show()
 
-pmv.io.save_and_compress_svg(fig, "plot-histogram-ecdf")
+pmv.io.save_and_compress_svg(fig, "histogram-ecdf")
diff --git a/pymatviz/__init__.py b/pymatviz/__init__.py
@@ -94,6 +94,8 @@
     pass  # package not installed
 
 
+IS_IPYTHON = hasattr(__builtins__, "__IPYTHON__")
+
 # define a sensible order for crystal systems across plots
 crystal_sys_order = (
     "cubic hexagonal trigonal tetragonal orthorhombic monoclinic triclinic".split()
@@ -104,7 +106,6 @@
 }
 spg_num_to_symbol = dict(sorted(spg_num_to_symbol.items()))  # sort
 
-
 px.defaults.labels |= {
     "gap expt": "Experimental band gap (eV)",
 } | Key.val_label_dict()
diff --git a/pymatviz/structure_viz.py b/pymatviz/structure_viz.py
@@ -1,6 +1,6 @@
 """2D plots of pymatgen structures with matplotlib.
 
-plot_structure_2d() and its helpers get_rot_matrix() and unit_cell_to_lines() were
+structure_2d() and its helpers get_rot_matrix() and unit_cell_to_lines() were
 inspired by ASE https://wiki.fysik.dtu.dk/ase/ase/visualize/visualize.html#matplotlib.
 """
 
@@ -138,7 +138,6 @@ def structure_2d(
     n_cols: int = 4,
     subplot_kwargs: dict[str, Any] | None = None,
     subplot_title: Callable[[Structure, str | int], str] | None = None,
-    **kwargs: Any,
 ) -> plt.Axes | tuple[plt.Figure, np.ndarray[plt.Axes]]:
     """Plot pymatgen structures in 2D with matplotlib.
 
@@ -161,22 +160,22 @@ def structure_2d(
     plot_atoms(AseAtomsAdaptor().get_atoms(mp_19017), rotation="10x,10y,0z", radii=0.5)
 
     # pymatviz
-    from pymatviz import plot_structure_2d
+    from pymatviz import structure_2d
 
-    plot_structure_2d(mp_19017)
+    structure_2d(mp_19017)
     ```
 
     Multiple structures in single figure example:
 
     ```py
     from pymatgen.ext.matproj import MPRester
-    from pymatviz import plot_structure_2d
+    from pymatviz import structure_2d
 
     structures = {
         (mp_id := f"mp-{idx}"): MPRester().get_structure_by_material_id(mp_id)
         for idx in range(1, 5)
     }
-    plot_structure_2d(structures)
+    structure_2d(structures)
     ```
 
     Args:
@@ -228,22 +227,13 @@ class used to plot chemical bonds. Allowed are edgecolor, facecolor, color,
             subplot titles. Receives the structure and its key or index when passed as
             a dict or pandas.Series. Defaults to None in which case the title is the
             structure's material id if available, otherwise its formula and space group.
-        **kwargs: Unused.
 
     Raises:
         ValueError: On invalid site_labels.
 
     Returns:
         plt.Axes: matplotlib Axes instance with plotted structure.
     """
-    if "colors" in kwargs:
-        warnings.warn(
-            "plot_structure_2d's colors keyword was deprecated on 2024-07-06 and will "
-            "be removed soon. Use 'elem_colors' instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        elem_colors = kwargs.pop("colors")
     if isinstance(struct, Structure):
         ax = ax or plt.gca()
 
@@ -382,11 +372,13 @@ class used to plot chemical bonds. Allowed are edgecolor, facecolor, color,
                     fallback_color = "gray"
                     face_color = elem_colors.get(elem_symbol, fallback_color)
                     if elem_symbol not in elem_colors:
-                        warnings.warn(
-                            f"{elem_symbol=} not in colors, using {fallback_color}",
-                            UserWarning,
-                            stacklevel=2,
+                        elem_color_symbols = ", ".join(elem_colors)
+                        warn_msg = (
+                            f"{elem_symbol=} not in elem_colors, using "
+                            f"{fallback_color=}\nelement color palette specifies the "
+                            f"following elements: {elem_color_symbols}"
                         )
+                        warnings.warn(warn_msg, UserWarning, stacklevel=2)
                     wedge = Wedge(
                         xy,
                         radius,
@@ -456,7 +448,7 @@ class used to plot chemical bonds. Allowed are edgecolor, facecolor, color,
 
         if show_bonds:
             warnings.warn(
-                "Warning: the show_bonds feature of plot_structure_2d() is "
+                "Warning: the show_bonds feature of structure_2d() is "
                 "experimental. Issues and PRs with improvements welcome.",
                 category=ExperimentalWarning,
                 stacklevel=2,
@@ -472,8 +464,8 @@ class used to plot chemical bonds. Allowed are edgecolor, facecolor, color,
 
             # If structure doesn't have any oxidation states yet, guess them from
             # chemical composition. Use CrystalNN and other strategies to better
-            # estimate bond connectivity. Use getattr on site.specie since it's often a
-            # pymatgen Element which has no oxi_state
+            # estimate bond connectivity. Use hasattr("oxi_state") on site.specie since
+            # it's often a pymatgen Element which has no oxi_state
             if not any(
                 hasattr(getattr(site, "specie", None), "oxi_state") for site in struct
             ):
diff --git a/readme.md b/readme.md
@@ -77,11 +77,11 @@ See [`pymatviz/ptable/ptable_matplotlib.py`](pymatviz/ptable/ptable_matplotlib.p
 
 See [`examples/mlff_phonons.ipynb`](https://github.com/janosh/pymatviz/blob/main/examples/mlff_phonons.ipynb) for usage example.
 
-|           [`plot_phonon_bands(bands_dict)`](pymatviz/phonons.py)           |             [`plot_phonon_dos(doses_dict)`](pymatviz/phonons.py)             |
-| :------------------------------------------------------------------------: | :--------------------------------------------------------------------------: |
-|                              ![phonon-bands]                               |                                ![phonon-dos]                                 |
-| [`plot_phonon_bands_and_dos(bands_dict, doses_dict)`](pymatviz/phonons.py) | [`plot_phonon_bands_and_dos(single_bands, single_dos)`](pymatviz/phonons.py) |
-|                      ![phonon-bands-and-dos-mp-2758]                       |                       ![phonon-bands-and-dos-mp-23907]                       |
+|           [`phonon_bands(bands_dict)`](pymatviz/phonons.py)           |             [`phonon_dos(doses_dict)`](pymatviz/phonons.py)             |
+| :-------------------------------------------------------------------: | :---------------------------------------------------------------------: |
+|                            ![phonon-bands]                            |                              ![phonon-dos]                              |
+| [`phonon_bands_and_dos(bands_dict, doses_dict)`](pymatviz/phonons.py) | [`phonon_bands_and_dos(single_bands, single_dos)`](pymatviz/phonons.py) |
+|                    ![phonon-bands-and-dos-mp-2758]                    |                    ![phonon-bands-and-dos-mp-23907]                     |
 
 [phonon-bands]: https://github.com/janosh/pymatviz/raw/main/assets/phonon-bands-mp-2758.svg
 [phonon-dos]: https://github.com/janosh/pymatviz/raw/main/assets/phonon-dos-mp-2758.svg
@@ -114,9 +114,9 @@ See [`pymatviz/sankey.py`](pymatviz/sankey.py).
 
 See [`pymatviz/structure_viz.py`](pymatviz/structure_viz.py). Currently structure plotting is only supported with `matplotlib` in 2d. 3d interactive plots (probably with `plotly`) are on the road map.
 
-| [`plot_structure_2d(mp_19017)`](pymatviz/structure_viz.py) | [`plot_structure_2d(mp_12712)`](pymatviz/structure_viz.py) |
-| :--------------------------------------------------------: | :--------------------------------------------------------: |
-|  ![struct-2d-mp-19017-Li4Mn0.8Fe1.6P4C1.6O16-disordered]   |        ![struct-2d-mp-12712-Hf9Zr9Pd24-disordered]         |
+|  [`structure_2d(mp_19017)`](pymatviz/structure_viz.py)  | [`structure_2d(mp_12712)`](pymatviz/structure_viz.py) |
+| :-----------------------------------------------------: | :---------------------------------------------------: |
+| ![struct-2d-mp-19017-Li4Mn0.8Fe1.6P4C1.6O16-disordered] |      ![struct-2d-mp-12712-Hf9Zr9Pd24-disordered]      |
 
 ![matbench-phonons-structures-2d]
 
@@ -129,14 +129,14 @@ See [`pymatviz/histogram.py`](pymatviz/histogram.py).
 |                              ![spg-num-hist-matplotlib]                              |                                  ![spg-symbol-hist-matplotlib]                                   |
 |   [`spacegroup_bar([65, 134, 225, ...], backend="plotly")`](pymatviz/histogram.py)   |   [`spacegroup_bar(["C2/m", "P-43m", "Fm-3m", ...], backend="plotly")`](pymatviz/histogram.py)   |
 |                                ![spg-num-hist-plotly]                                |                                    ![spg-symbol-hist-plotly]                                     |
-| [`elements_hist(compositions, log=True, bar_values='count')`](pymatviz/histogram.py) |          [`plot_histogram({'key1': values1, 'key2': values2})`](pymatviz/histogram.py)           |
-|                                   ![elements-hist]                                   |                                      ![plot-histogram-ecdf]                                      |
+| [`elements_hist(compositions, log=True, bar_values='count')`](pymatviz/histogram.py) |             [`histogram({'key1': values1, 'key2': values2})`](pymatviz/histogram.py)             |
+|                                   ![elements-hist]                                   |                                        ![histogram-ecdf]                                         |
 
 [spg-symbol-hist-plotly]: https://github.com/janosh/pymatviz/raw/main/assets/spg-symbol-hist-plotly.svg
 [spg-num-hist-plotly]: https://github.com/janosh/pymatviz/raw/main/assets/spg-num-hist-plotly.svg
 [spg-num-hist-matplotlib]: https://github.com/janosh/pymatviz/raw/main/assets/spg-num-hist-matplotlib.svg
 [spg-symbol-hist-matplotlib]: https://github.com/janosh/pymatviz/raw/main/assets/spg-symbol-hist-matplotlib.svg
-[plot-histogram-ecdf]: https://github.com/janosh/pymatviz/raw/main/assets/plot-histogram-ecdf.svg
+[histogram-ecdf]: https://github.com/janosh/pymatviz/raw/main/assets/histogram-ecdf.svg
 
 ## Scatter Plots
 
@@ -161,9 +161,9 @@ See [`pymatviz/scatter.py`](pymatviz/scatter.py).
 
 See [`pymatviz/xrd.py`](pymatviz/xrd.py).
 
-| [`plot_xrd_pattern(pattern)`](pymatviz/xrd.py) | [`plot_xrd_pattern({key1: patt1, key2: patt2})`](pymatviz/xrd.py) |
-| :--------------------------------------------: | :---------------------------------------------------------------: |
-|                 ![xrd-pattern]                 |                      ![xrd-pattern-multiple]                      |
+| [`xrd_pattern(pattern)`](pymatviz/xrd.py) | [`xrd_pattern({key1: patt1, key2: patt2})`](pymatviz/xrd.py) |
+| :---------------------------------------: | :----------------------------------------------------------: |
+|              ![xrd-pattern]               |                   ![xrd-pattern-multiple]                    |
 
 [xrd-pattern]: https://github.com/janosh/pymatviz/raw/main/assets/xrd-pattern.svg
 [xrd-pattern-multiple]: https://github.com/janosh/pymatviz/raw/main/assets/xrd-pattern-multiple.svg
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1 +1 @@
-"""pymatviz tests."""
+"""pymatviz test suite."""
diff --git a/tests/test_pkg.py b/tests/test_pkg.py
@@ -7,7 +7,7 @@
 from importlib.metadata import version
 from types import ModuleType
 
-from pymatviz import PKG_DIR, PKG_NAME, __version__
+from pymatviz import IS_IPYTHON, PKG_DIR, PKG_NAME, __version__
 
 
 def test_pkg_metadata() -> None:
@@ -73,3 +73,7 @@ def test_all_modules_reexported() -> None:
 
     finally:
         sys.modules[PKG_NAME] = pymatviz
+
+
+def test_is_ipython() -> None:
+    assert not IS_IPYTHON
diff --git a/tests/test_structure_viz.py b/tests/test_structure_viz.py
@@ -10,6 +10,7 @@
 from pymatgen.core import Structure
 
 import pymatviz as pmv
+from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
 from pymatviz.enums import ElemColorScheme, Key
 
 
@@ -165,20 +166,35 @@ def subplot_title(struct: Structure, key: str | int) -> str:
 
 def test_structure_2d_color_warning() -> None:
     # Copernicium is not in the default color scheme
-    copernicium = "Cn"
-    struct = Structure(np.eye(3) * 5, [copernicium] * 2, coords=coords)
+    elem_symbol = "Cn"
+    struct = Structure(np.eye(3) * 5, [elem_symbol] * 2, coords=coords)
+    fallback_color = "gray"
+
+    for elem_colors in ElemColorScheme:
+        if elem_colors == ElemColorScheme.jmol:
+            elem_color_symbols = ", ".join(ELEM_COLORS_JMOL)
+        elif elem_colors == ElemColorScheme.vesta:
+            elem_color_symbols = ", ".join(ELEM_COLORS_VESTA)
+        else:
+            raise ValueError(f"Unexpected {elem_colors=}")
 
-    for colors in ElemColorScheme:
         with pytest.warns(
-            UserWarning, match=f"{copernicium!r} not in colors, using gray"
+            UserWarning,
+            match=f"{elem_symbol=} not in elem_colors, using "
+            f"{fallback_color=}\nelement color palette specifies the "
+            f"following elements: {elem_color_symbols}",
         ):
-            pmv.structure_2d(struct, elem_colors=colors)
+            pmv.structure_2d(struct, elem_colors=elem_colors)
 
     # create custom color scheme missing an element
     custom_colors = {"Fe": "red"}
     struct = Structure(np.eye(3) * 5, ["Fe", "O"], coords=coords)
 
-    with pytest.warns(UserWarning, match="'O' not in colors, using gray"):
+    with pytest.warns(
+        UserWarning,
+        match="elem_symbol='O' not in elem_colors, using fallback_color='gray'\nelement"
+        " color palette specifies the following elements: Fe",
+    ):
         pmv.structure_2d(struct, elem_colors=custom_colors)
 
 

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-"""pymatviz tests."""`
	`1`	`+"""pymatviz test suite."""`