access pymatviz.io.save_fig from namespace

Author: janosh

data/mp/build_phase_diagram.py

@@ -10,12 +10,12 @@
 
 import pandas as pd
 import pymatviz
+import pymatviz as pmv
 from pymatgen.analysis.phase_diagram import PatchedPhaseDiagram
 from pymatgen.entries.compatibility import MaterialsProject2020Compatibility
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 from pymatgen.ext.matproj import MPRester
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
 from tqdm import tqdm
 
 from matbench_discovery import MP_DIR, ROOT, today
@@ -121,4 +121,4 @@
     xlabel="MP Formation Energy (eV/atom)",
     ylabel="Our Formation Energy (eV/atom)",
 )
-save_fig(ax, f"{ROOT}/tmp/{today}-our-vs-mp-formation-energies.webp", dpi=300)
+pmv.save_fig(ax, f"{ROOT}/tmp/{today}-our-vs-mp-formation-energies.webp", dpi=300)

data/mp/eda_mp_trj.py

@@ -17,8 +17,6 @@
 from matplotlib.colors import SymLogNorm
 from pymatgen.core import Composition, Element
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
-from pymatviz.utils import si_fmt
 from tqdm import tqdm
 
 from matbench_discovery import MP_DIR, PDF_FIGS, ROOT, SITE_FIGS
@@ -108,7 +106,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
     """Annotate each periodic table tile with the number of values in its histogram."""
     face_color = cmap(norm(np.sum(len(hist_vals)))) if hist_vals else "none"
     bbox = dict(facecolor=face_color, alpha=0.4, pad=2, edgecolor="none")
-    return dict(text=si_fmt(len(hist_vals), ".0f"), bbox=bbox)
+    return dict(text=pmv.si_fmt(len(hist_vals), ".0f"), bbox=bbox)
 
 
 # %% plot per-element magmom histograms
@@ -152,7 +150,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 cbar = matplotlib.colorbar.ColorbarBase(
     cbar_ax, cmap=cmap, norm=norm, orientation="horizontal"
 )
-save_fig(fig_ptable_magmoms, f"{PDF_FIGS}/mp-trj-magmoms-ptable-hists.pdf")
+pmv.save_fig(fig_ptable_magmoms, f"{PDF_FIGS}/mp-trj-magmoms-ptable-hists.pdf")
 
 
 # %% plot per-element force histograms
@@ -195,7 +193,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
     cbar_ax, cmap=cmap, norm=norm, orientation="horizontal"
 )
 
-save_fig(fig_ptable_forces, f"{PDF_FIGS}/mp-trj-forces-ptable-hists.pdf")
+pmv.save_fig(fig_ptable_forces, f"{PDF_FIGS}/mp-trj-forces-ptable-hists.pdf")
 
 
 # %% plot histogram of number of sites per element
@@ -238,7 +236,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
     cbar_title_kwds=dict(fontsize=16),
     cbar_coords=(0.18, 0.85, 0.42, 0.02),
     anno_kwds=lambda hist_vals: dict(
-        text=si_fmt(len(hist_vals), ".0f"),
+        text=pmv.si_fmt(len(hist_vals), ".0f"),
         xy=(0.8, 0.6),
         bbox=dict(pad=2, edgecolor="none", facecolor="none"),
     ),
@@ -262,7 +260,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 cbar.set_label("Number of atoms in MPtrj structures", fontsize=16)
 cbar.ax.xaxis.set_label_position("top")
 
-save_fig(fig_ptable_sites, f"{PDF_FIGS}/mp-trj-n-sites-ptable-hists.pdf")
+pmv.save_fig(fig_ptable_sites, f"{PDF_FIGS}/mp-trj-n-sites-ptable-hists.pdf")
 
 
 # %%
@@ -301,7 +299,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
     img_name += "-symlog" if isinstance(log, SymLogNorm) else "-log"
 if excl_noble:
     img_name += "-excl-noble"
-save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
+pmv.save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
 
 
 # %%
@@ -319,7 +317,7 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 img_name = "mp-trj-mp-ratio-element-counts-by-occurrence"
 if normalized:
     img_name += "-normalized"
-save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
+pmv.save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
 
 
 # %% plot formation energy per atom distribution
@@ -332,8 +330,8 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 fig.show()
 
 pdf_kwds = dict(width=500, height=300)
-# save_fig(fig, f"{PDF_FIGS}/mp-trj-e-form-hist.pdf", **pdf_kwds)
-# save_fig(fig, f"{SITE_FIGS}/mp-trj-e-form-hist.svelte")
+# pmv.save_fig(fig, f"{PDF_FIGS}/mp-trj-e-form-hist.pdf", **pdf_kwds)
+# pmv.save_fig(fig, f"{SITE_FIGS}/mp-trj-e-form-hist.svelte")
 
 
 # %% plot forces distribution
@@ -343,8 +341,8 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 fig.update_yaxes(type="log")
 fig.show()
 
-# save_fig(fig, f"{PDF_FIGS}/mp-trj-forces-hist.pdf", **pdf_kwds)
-# save_fig(fig, f"{SITE_FIGS}/mp-trj-forces-hist.svelte")
+# pmv.save_fig(fig, f"{PDF_FIGS}/mp-trj-forces-hist.pdf", **pdf_kwds)
+# pmv.save_fig(fig, f"{SITE_FIGS}/mp-trj-forces-hist.svelte")
 
 
 # %% plot hydrostatic stress distribution
@@ -354,8 +352,8 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 fig.update_yaxes(type="log")
 fig.show()
 
-# save_fig(fig, f"{PDF_FIGS}/mp-trj-stresses-hist.pdf", **pdf_kwds)
-# save_fig(fig, f"{SITE_FIGS}/mp-trj-stresses-hist.svelte")
+# pmv.save_fig(fig, f"{PDF_FIGS}/mp-trj-stresses-hist.pdf", **pdf_kwds)
+# pmv.save_fig(fig, f"{SITE_FIGS}/mp-trj-stresses-hist.svelte")
 
 
 # %% plot magmoms distribution
@@ -365,8 +363,8 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 fig.update_yaxes(type="log")
 fig.show()
 
-# save_fig(fig, f"{PDF_FIGS}/mp-trj-magmoms-hist.pdf", **pdf_kwds)
-# save_fig(fig, f"{SITE_FIGS}/mp-trj-magmoms-hist.svelte")
+# pmv.save_fig(fig, f"{PDF_FIGS}/mp-trj-magmoms-hist.pdf", **pdf_kwds)
+# pmv.save_fig(fig, f"{SITE_FIGS}/mp-trj-magmoms-hist.svelte")
 
 
 # %%
@@ -393,8 +391,8 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 
 fig.show()
 img_name = "mp-vs-mp-trj-vs-wbm-arity-hist"
-save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
-save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=280)
+pmv.save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
+pmv.save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=280)
 
 
 # %% calc n_sites from per-site atomic numbers
@@ -472,5 +470,5 @@ def tile_count_anno(hist_vals: list[Any]) -> dict[str, Any]:
 img_name = "mp-trj-n-sites-hist"
 if log_y:
     img_name += "-log"
-save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
-# save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=300)
+pmv.save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
+# pmv.save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=300)

data/mp/get_mp_energies.py

@@ -8,13 +8,11 @@
 import os
 
 import pandas as pd
+import pymatviz as pmv
 from aviary.wren.utils import get_protostructure_label_from_spglib
 from mp_api.client import MPRester
 from pymatgen.core import Structure
 from pymatviz.enums import Key
-from pymatviz.io import save_fig  # noqa: F401
-from pymatviz.powerups import annotate_metrics
-from pymatviz.utils import PLOTLY
 from tqdm import tqdm
 
 from matbench_discovery import STABILITY_THRESHOLD, today
@@ -55,7 +53,7 @@
 df_spg = pd.json_normalize(df_mp.pop("symmetry"))[["number", "symbol"]]
 df_mp["spacegroup_symbol"] = df_spg.symbol.to_numpy()
 
-df_mp.energy_type.value_counts().plot.pie(backend=PLOTLY, autopct="%1.1f%%")
+df_mp.energy_type.value_counts().plot.pie(backend=pmv.utils.PLOTLY, autopct="%1.1f%%")
 # GGA: 72.2%, GGA+U: 27.8%
 
 
@@ -96,10 +94,12 @@
     title=f"{today} - {len(df_mp):,} MP entries",
 )
 
-annotate_metrics(df_mp.formation_energy_per_atom, df_mp.decomposition_enthalpy)
+pmv.powerups.annotate_metrics(
+    df_mp.formation_energy_per_atom, df_mp.decomposition_enthalpy
+)
 # result on 2023-01-10: plots match. no correlation between formation energy and
 # decomposition enthalpy. R^2 = -1.571, MAE = 1.604
-# save_fig(ax, f"{module_dir}/mp-decomp-enth-vs-e-form.webp", dpi=300)
+# pmv.save_fig(ax, f"{module_dir}/mp-decomp-enth-vs-e-form.webp", dpi=300)
 
 
 # %% scatter plot energy above convex hull vs decomposition enthalpy
@@ -117,4 +117,4 @@
     title=f"{n_above_line:,} / {len(df_mp):,} = {n_above_line / len(df_mp):.1%} "
     "MP materials with\nenergy_above_hull - decomposition_enthalpy.clip(0) > 0.1"
 )
-# save_fig(ax, f"{module_dir}/mp-e-above-hull-vs-decomp-enth.webp", dpi=300)
+# pmv.save_fig(ax, f"{module_dir}/mp-e-above-hull-vs-decomp-enth.webp", dpi=300)

data/wbm/compare_cse_vs_ce_mp_2020_corrections.py

@@ -8,13 +8,13 @@
 import json
 
 import pandas as pd
+import pymatviz as pmv
 from pymatgen.entries.compatibility import (
     MaterialsProject2020Compatibility,
     MaterialsProjectCompatibility,
 )
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
 from tqdm import tqdm
 
 from matbench_discovery import ROOT, today
@@ -96,7 +96,7 @@
 
 ax.axline((0, 0), slope=1, color="gray", linestyle="dashed", zorder=-1)
 
-save_fig(ax, f"{ROOT}/tmp/{today}-ce-vs-cse-corrections-outliers.pdf")
+pmv.save_fig(ax, f"{ROOT}/tmp/{today}-ce-vs-cse-corrections-outliers.pdf")
 
 
 # %%
@@ -117,7 +117,7 @@
 # insight: all materials for which ComputedEntry and ComputedStructureEntry give
 # different formation energies are oxides or sulfides for which MP 2020 compat takes
 # into account structural information to make more accurate corrections.
-save_fig(ax, f"{ROOT}/tmp/{today}-ce-vs-cse-e-form-outliers.pdf")
+pmv.save_fig(ax, f"{ROOT}/tmp/{today}-ce-vs-cse-e-form-outliers.pdf")
 
 
 # %% below code resulted in

data/wbm/compile_wbm_test_set.py

@@ -22,7 +22,6 @@
 )
 from pymatgen.entries.computed_entries import ComputedStructureEntry
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
 from tqdm import tqdm
 
 from matbench_discovery import PDF_FIGS, SITE_FIGS, WBM_DIR, today
@@ -478,13 +477,13 @@ def fix_bad_struct_index_mismatch(material_id: str) -> str:
 
 # %%
 img_name = "hist-wbm-e-form-per-atom"
-save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
+pmv.save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
 # recommended to upload SVG to vecta.io/nano for compression
-# save_fig(fig, f"{img_name}.svg", width=800, height=300)
+# pmv.save_fig(fig, f"{img_name}.svg", width=800, height=300)
 
 # make full data range visible in PDF
 # fig.layout.xaxis.range = [-12, 82]
-save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf")
+pmv.save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf")
 
 
 # %%

data/wbm/eda_wbm.py

@@ -12,7 +12,6 @@
 from matplotlib.colors import SymLogNorm
 from pymatgen.core import Composition, Structure
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
 from pymatviz.utils import PLOTLY, si_fmt, si_fmt_int
 
 from matbench_discovery import PDF_FIGS, ROOT, SITE_FIGS, STABILITY_THRESHOLD
@@ -91,7 +90,7 @@
     )
     if log:
         filename += "-symlog" if isinstance(log, SymLogNorm) else "-log"
-    save_fig(ax_mp_cnt, f"{PDF_FIGS}/{filename}.pdf")
+    pmv.save_fig(ax_mp_cnt, f"{PDF_FIGS}/{filename}.pdf")
 
 
 # %% ratio of WBM to MP counts
@@ -105,7 +104,7 @@
 img_name = "wbm-mp-ratio-element-counts-by-occurrence"
 if normalized:
     img_name += "-normalized"
-save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
+pmv.save_fig(ax_ptable, f"{PDF_FIGS}/{img_name}.pdf")
 
 
 # %% export element counts by WBM step to JSON
@@ -139,8 +138,8 @@
     fig.layout.margin = dict(l=0, r=0, b=0, t=0)
     fig.show()
     svg_path = f"{module_dir}/figs/wbm-elements.svg"
-    # save_fig(fig, svg_path, width=1000, height=500)
-    save_fig(fig, f"{PDF_FIGS}/{dataset}-element-{count_mode}-counts.pdf")
+    # pmv.save_fig(fig, svg_path, width=1000, height=500)
+    pmv.save_fig(fig, f"{PDF_FIGS}/{dataset}-element-{count_mode}-counts.pdf")
 
 
 # %% histogram of energy distance to MP convex hull for WBM
@@ -216,9 +215,9 @@
     MbdKey.e_form_raw: "e-form-uncorrected",
 }[e_col]
 img_name = f"hist-wbm-{suffix}"
-save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
-# save_fig(fig, f"./figs/{img_name}.svg", width=800, height=500)
-save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=600, height=300)
+pmv.save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
+# pmv.save_fig(fig, f"./figs/{img_name}.svg", width=800, height=500)
+pmv.save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=600, height=300)
 
 
 # %%
@@ -252,8 +251,8 @@
 
 fig.show()
 
-save_fig(fig, f"{SITE_FIGS}/mp-elemental-ref-energies.svelte")
-save_fig(fig, f"{PDF_FIGS}/mp-elemental-ref-energies.pdf")
+pmv.save_fig(fig, f"{SITE_FIGS}/mp-elemental-ref-energies.svelte")
+pmv.save_fig(fig, f"{PDF_FIGS}/mp-elemental-ref-energies.pdf")
 
 
 # %% plot 2d and 3d t-SNE projections of one-hot encoded element vectors summed by
@@ -326,8 +325,8 @@
 fig.layout.title.update(text="WBM Spacegroup Sunburst", x=0.5, font_size=14)
 fig.layout.margin = dict(l=0, r=0, t=30, b=0)
 fig.show()
-save_fig(fig, f"{SITE_FIGS}/spacegroup-sunburst-wbm.svelte")
-save_fig(fig, f"{PDF_FIGS}/spacegroup-sunburst-wbm.pdf")
+pmv.save_fig(fig, f"{SITE_FIGS}/spacegroup-sunburst-wbm.svelte")
+pmv.save_fig(fig, f"{PDF_FIGS}/spacegroup-sunburst-wbm.pdf")
 
 
 # %%
@@ -337,8 +336,8 @@
 fig.layout.title.update(text="MP Spacegroup Sunburst", x=0.5, font_size=14)
 fig.layout.margin = dict(l=0, r=0, t=30, b=0)
 fig.show()
-save_fig(fig, f"{SITE_FIGS}/spacegroup-sunburst-mp.svelte")
-save_fig(fig, f"{PDF_FIGS}/spacegroup-sunburst-mp.pdf")
+pmv.save_fig(fig, f"{SITE_FIGS}/spacegroup-sunburst-mp.svelte")
+pmv.save_fig(fig, f"{PDF_FIGS}/spacegroup-sunburst-mp.pdf")
 # would be good to have consistent order of crystal systems between sunbursts but not
 # controllable yet
 # https://github.com/plotly/plotly.py/issues/4115
@@ -366,8 +365,8 @@
 
 fig.show()
 img_name = "mp-vs-wbm-arity-hist"
-save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
-save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=280)
+pmv.save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
+pmv.save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=280)
 
 
 # %% find large structures that changed symmetry during relaxation

matbench_discovery/enums.py

@@ -3,7 +3,7 @@
 from enum import StrEnum, unique
 from typing import Self
 
-from pymatviz.utils import html_tag
+import pymatviz as pmv
 
 
 class LabelEnum(StrEnum):
@@ -136,7 +136,7 @@ class TestSubset(LabelEnum):
     full = "full", "Full Test Set"
 
 
-eV_per_atom = html_tag(  # noqa: N816
+eV_per_atom = pmv.html_tag(  # noqa: N816
     "(eV/atom)", tag="span", style="font-size: 0.8em; font-weight: lighter;"
 )
 

models/chgnet/analyze_chgnet.py

@@ -8,8 +8,6 @@
 import pymatviz as pmv
 from pymatgen.core import Structure
 from pymatviz.enums import Key
-from pymatviz.io import save_fig
-from pymatviz.utils import PLOTLY
 
 from matbench_discovery import PDF_FIGS
 from matbench_discovery import plots as plots
@@ -54,7 +52,7 @@
     y=e_form_2000,
     hover_name=Key.mat_id,
     hover_data=[Key.formula],
-    backend=PLOTLY,
+    backend=pmv.utils.PLOTLY,
     title=f"{len(df_diff)} structures have > {min_e_diff} eV/atom energy diff after "
     "longer relaxation",
 )
@@ -88,7 +86,7 @@
     formula = struct.composition.reduced_formula
     ax.set_title(f"{idx}. {formula} (spg={spg_num})\n{row.Index}", fontweight="bold")
 
-save_fig(fig, f"{PDF_FIGS}/chgnet-bad-relax-structures.pdf")
+pmv.save_fig(fig, f"{PDF_FIGS}/chgnet-bad-relax-structures.pdf")
 
 
 # %% ensure all CHGNet static predictions (direct energy without any structure