fix ptable_heatmap return type plt.axes->plt.Axes

add more Key enum attrs
don't rely on convenience function re-exports in tests

pymatviz/enums.py

@@ -177,8 +177,9 @@ class Key(LabelEnum):
     protostructure = "protostructure", "Protostructure Label"
     # Deprecated name for the protostructure
     wyckoff = "wyckoff", "AFLOW-style Label with Chemical System"
-    # AFLOW-style prototype label
-    prototype = "prototype", "AFLOW-style Prototype Label"
+    wyckoff_spglib = "wyckoff_spglib", "Wyckoff Label (spglib)"
+    prototype = "prototype", "Prototype Label"
+    aflow_prototype = "aflow_prototype", "AFLOW-style Prototype Label"
     # AFLOW-style prototype label that has been canonicalized
     canonical_proto = "canonical_prototype", "Canonical AFLOW-style Prototype"
     # Deprecated name for the canonical_proto
@@ -306,6 +307,10 @@ class Key(LabelEnum):
     poisson_ratio = "poisson_ratio", "Poisson's Ratio"
     hardness = "hardness", "Hardness (Mohs scale)"
     elastic_tensor = "elastic_tensor", "Elastic Tensor"
+    elastic_tensor_voigt = "elastic_tensor_voigt", "Voigt Elastic Tensor"
+    elastic_tensor_reuss = "elastic_tensor_reuss", "Reuss Elastic Tensor"
+    elastic_tensor_vrh = "elastic_tensor_vrh", "Voigt-Reuss-Hill Elastic Tensor"
+    toughness = "toughness", "Toughness (MPa)"
     yield_strength = "yield_strength", "Yield Strength (MPa)"
     tensile_strength = "tensile_strength", "Tensile Strength (MPa)"
     ductility = "ductility", "Ductility (%)"
@@ -455,10 +460,21 @@ class Key(LabelEnum):
     tn = "TN", "True Negatives"
     tpr = "TPR", "True Positive Rate"
     fpr = "FPR", "False Positive Rate"
+    tnr = "TNR", "True Negative Rate"
+    fnr = "FNR", "False Negative Rate"
     mae = "MAE", "Mean Absolute Error"
     r2 = "R²", "R² Score"
+    pearson = "Pearson", "Pearson Correlation"
+    spearman = "Spearman", "Spearman Correlation"
+    kendall = "Kendall", "Kendall Correlation"
     rmse = "RMSE", "Root Mean Squared Error"
     mape = "MAPE", "Mean Absolute Percentage Error"
+    variance = "variance", "Variance"
+    std_dev = "std_dev", "Standard Deviation"
+    iqr = "IQR", "Interquartile Range"
+    outlier = "outlier", "Outlier"
+    error = "error", "Error"
+    residuals = "residuals", "Residuals"
     prc = "PRC", "Precision-Recall Curve"
     prc_curve = "prc_curve", "PRC Curve"
     precision = "precision", "Precision"
@@ -524,10 +540,29 @@ class Key(LabelEnum):
     # Performance Indicators
     fom = "figure_of_merit", "Figure of Merit"  # codespell:ignore
     power_factor = "power_factor", "Power Factor"
+    zt = "ZT", "ZT"
+    efficiency = "efficiency", "Efficiency"
+    capacity = "capacity", "Capacity"
+    rate = "rate", "Rate"
+    lifetime = "lifetime", "Lifetime"
+    stability = "stability", "Stability"
+    selectivity = "selectivity", "Selectivity"
+    yield_ = "yield", "Yield"  # underscore because reserved keyword
+    activity = "activity", "Activity"
+    performance = "performance", "Performance"
+    gain = "gain", "Gain"
+    power = "power", "Power"
+    current = "current", "Current"
+    voltage = "voltage", "Voltage"
+    resistance = "resistance", "Resistance"
+    impedance = "impedance", "Impedance"
+    capacitance = "capacitance", "Capacitance"
 
     # Environmental Indicators
     toxicity = "toxicity", "Toxicity Index"
     recyclability = "recyclability", "Recyclability Score"
+    biodegradability = "biodegradability", "Biodegradability Score"
+    sustainability = "sustainability", "Sustainability Index"
 
     # Economic Factors
     raw_material_cost = "raw_material_cost", "Raw Material Cost ($/kg)"
@@ -542,6 +577,7 @@ class Key(LabelEnum):
     count = "count", "Count"  # type: ignore[assignment]
     heat_val = "heat_val", "Heatmap Value"
     piezoelectric_tensor = "piezoelectric_tensor", "Piezoelectric Tensor"
+    dielectric_tensor = "dielectric_tensor", "Dielectric Tensor"
 
 
 class Task(LabelEnum):

pymatviz/ptable/ptable_matplotlib.py

@@ -80,7 +80,7 @@ def ptable_heatmap(
     fmt: str | None = None,
     cbar_fmt: str | None = None,
     show_scale: bool | None = None,
-) -> plt.axes:
+) -> plt.Axes:
     """Plot a heatmap across the periodic table.
 
     Args:

tests/ptable/test_ptable_plotly.py

@@ -10,8 +10,9 @@
 from plotly.exceptions import PlotlyError
 from plotly.graph_objs import Figure
 
-from pymatviz import df_ptable, ptable_heatmap_plotly
 from pymatviz.enums import ElemCountMode, Key
+from pymatviz.ptable import ptable_heatmap_plotly
+from pymatviz.utils import df_ptable
 
 
 if TYPE_CHECKING:

tests/test_bar.py

@@ -7,7 +7,7 @@
 import plotly.graph_objects as go
 import pytest
 
-from pymatviz import spacegroup_bar
+from pymatviz.bar import spacegroup_bar
 from pymatviz.utils import BACKENDS, MATPLOTLIB, PLOTLY
 from tests.conftest import y_pred, y_true
 

tests/test_correlation.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from pymatviz import marchenko_pastur
+from pymatviz.correlation import marchenko_pastur
 from tests.conftest import np_rng
 
 

tests/test_cumulative.py

@@ -3,7 +3,7 @@
 import matplotlib.pyplot as plt
 import pytest
 
-from pymatviz import cumulative_error, cumulative_residual
+from pymatviz.cumulative import cumulative_error, cumulative_residual
 from tests.conftest import y_pred, y_true
 
 

tests/test_histogram.py

@@ -7,7 +7,7 @@
 import plotly.graph_objects as go
 import pytest
 
-from pymatviz import elements_hist, histogram
+from pymatviz.histogram import elements_hist, histogram
 from pymatviz.utils import BACKENDS, MATPLOTLIB
 from tests.conftest import df_regr, y_true
 

tests/test_relevance.py

@@ -5,7 +5,7 @@
 import matplotlib.pyplot as plt
 import pytest
 
-from pymatviz import precision_recall_curve, roc_curve
+from pymatviz.relevance import precision_recall_curve, roc_curve
 from tests.conftest import df_x_y_clf
 
 

tests/test_scatter.py

@@ -8,7 +8,7 @@
 import plotly.graph_objects as go
 import pytest
 
-from pymatviz import (
+from pymatviz.scatter import (
     density_hexbin,
     density_hexbin_with_hist,
     density_scatter,

tests/test_sunburst.py

@@ -5,7 +5,7 @@
 import plotly.graph_objects as go
 import pytest
 
-from pymatviz import spacegroup_sunburst
+from pymatviz.sunburst import spacegroup_sunburst
 
 
 if TYPE_CHECKING:

tests/test_uncertainty.py

@@ -6,7 +6,7 @@
 import pandas as pd
 import pytest
 
-from pymatviz import error_decay_with_uncert, qq_gaussian
+from pymatviz.uncertainty import error_decay_with_uncert, qq_gaussian
 from tests.conftest import DfOrArrays, df_regr, xs, y_pred, y_true