diff --git a/docs/source/user_guide/benchmarks/bulk_crystal.rst b/docs/source/user_guide/benchmarks/bulk_crystal.rst index 8df003d56..419ee8ca1 100644 --- a/docs/source/user_guide/benchmarks/bulk_crystal.rst +++ b/docs/source/user_guide/benchmarks/bulk_crystal.rst @@ -2,6 +2,31 @@ Bulk Crystals ============= +Materials discovery evaluation +============================== + +The materials-discovery evaluator computes Matbench Discovery classification and +regression metrics from local tables. The reference table is indexed by +``material_id`` and contains DFT hull distance, DFT formation energy, and a +unique-prototype flag. The prediction table contains ``e_form_per_atom``. + +Results are reported for the full test set, unique prototypes, and the 10,000 +unique prototypes with the lowest predicted hull distances. They include F1, +discovery acceleration factor (DAF), precision, recall, accuracy, class rates and +counts, MAE, RMSE, R², and missing-prediction counts. Predictions with +formation-energy errors above 5 eV/atom are masked before rounding to three decimal +places. Leaderboard evaluation uses the fraction of unique prototypes with an +unrounded hull distance at or below 0 eV/atom as prevalence, preventing rounding +from changing DAF. Pass this value with ``canonical=True`` and +``uniq_proto_prevalence=...``. Synthetic mode derives prevalence from the rounded +reference values. ``calc_discovery_metrics`` and ``discovery_subset_indices`` apply +the same masking and rounding as ``evaluate_discovery``. + +See ``ml_peg.analysis.bulk_crystal.materials_discovery``. WBM reference and +prediction artifacts are not included. JSON results include schema and source +framework versions. + + Lattice constants ================= diff --git a/ml_peg/analysis/bulk_crystal/materials_discovery/__init__.py b/ml_peg/analysis/bulk_crystal/materials_discovery/__init__.py new file mode 100644 index 000000000..9a6df9f60 --- /dev/null +++ b/ml_peg/analysis/bulk_crystal/materials_discovery/__init__.py @@ -0,0 +1,69 @@ +"""Materials-discovery schemas, metrics, and evaluation.""" + +from __future__ import annotations + +from ml_peg.analysis.bulk_crystal.materials_discovery.evaluation import ( + EVALUATION_DECIMALS, + MAX_E_FORM_ERROR_THRESHOLD, + MISSING_PREDICTIONS_KEY, + RESULT_SCHEMA_VERSION, + DiscoveryResults, + DiscoverySubsetResults, + SourceMetadata, + calc_discovery_metrics, + discovery_subset_indices, + evaluate_discovery, + evaluate_discovery_paths, + prepare_discovery_inputs, + write_discovery_metrics_json, +) +from ml_peg.analysis.bulk_crystal.materials_discovery.metrics import ( + MOST_STABLE_COUNT, + STABILITY_THRESHOLD, + MetricValue, + align_predictions, + classify_stable, + stable_metrics, +) +from ml_peg.analysis.bulk_crystal.materials_discovery.schema import ( + E_ABOVE_HULL, + MATERIAL_ID, + PREDICTED_FORMATION_ENERGY, + REFERENCE_COLUMNS, + REFERENCE_FORMATION_ENERGY, + UNIQUE_PROTOTYPE, + DiscoverySubset, + validate_prediction_frame, + validate_reference_frame, +) + +__all__ = [ + "EVALUATION_DECIMALS", + "MAX_E_FORM_ERROR_THRESHOLD", + "MISSING_PREDICTIONS_KEY", + "MOST_STABLE_COUNT", + "RESULT_SCHEMA_VERSION", + "STABILITY_THRESHOLD", + "DiscoveryResults", + "DiscoverySubset", + "DiscoverySubsetResults", + "MetricValue", + "SourceMetadata", + "align_predictions", + "calc_discovery_metrics", + "classify_stable", + "discovery_subset_indices", + "evaluate_discovery", + "evaluate_discovery_paths", + "prepare_discovery_inputs", + "stable_metrics", + "validate_prediction_frame", + "validate_reference_frame", + "write_discovery_metrics_json", + "E_ABOVE_HULL", + "MATERIAL_ID", + "PREDICTED_FORMATION_ENERGY", + "REFERENCE_COLUMNS", + "REFERENCE_FORMATION_ENERGY", + "UNIQUE_PROTOTYPE", +] diff --git a/ml_peg/analysis/bulk_crystal/materials_discovery/evaluation.py b/ml_peg/analysis/bulk_crystal/materials_discovery/evaluation.py new file mode 100644 index 000000000..7de2795ca --- /dev/null +++ b/ml_peg/analysis/bulk_crystal/materials_discovery/evaluation.py @@ -0,0 +1,343 @@ +"""Path- and dataframe-driven materials-discovery evaluation.""" + +from __future__ import annotations + +import json +import math +import os +from typing import Final, TypeAlias, TypedDict + +import pandas as pd + +from ml_peg.analysis.bulk_crystal.materials_discovery.metrics import ( + MetricValue, + SubsetIndices, + _align_predictions_prepared, + _calc_discovery_metrics_prepared, + _discovery_subset_indices_prepared, + _hull_distances, +) +from ml_peg.analysis.bulk_crystal.materials_discovery.schema import ( + E_ABOVE_HULL, + MATERIAL_ID, + REFERENCE_FORMATION_ENERGY, + DiscoverySubset, + _validated_reference_frame, + prediction_series, +) +from ml_peg.data.artifacts import ( + MATBENCH_DISCOVERY_ID, + MATBENCH_DISCOVERY_VERSION, + PathLike, + read_csv_artifact, +) + +RESULT_SCHEMA_VERSION: Final = 1 +MAX_E_FORM_ERROR_THRESHOLD: Final = 5.0 +EVALUATION_DECIMALS: Final = 3 +MISSING_PREDICTIONS_KEY: Final = "missing_preds" + +JsonMetricValue: TypeAlias = float | int | None +DiscoverySubsetResults: TypeAlias = dict[str, dict[str, JsonMetricValue]] + + +class SourceMetadata(TypedDict): + """Identify the upstream framework used by an evaluation result.""" + + framework: str + version: str + + +class DiscoveryResults(TypedDict): + """Serialized discovery evaluation result.""" + + schema_version: int + source: SourceMetadata + subsets: DiscoverySubsetResults + + +def prepare_discovery_inputs( + reference: pd.DataFrame, + predictions: pd.DataFrame | pd.Series, + *, + max_error_threshold: float | None = MAX_E_FORM_ERROR_THRESHOLD, + decimals: int = EVALUATION_DECIMALS, +) -> tuple[pd.DataFrame, pd.Series]: + """ + Validate, align, mask, and round discovery evaluation inputs. + + Errors strictly above the threshold become NaN before all energies are rounded + to the shared evaluation precision. + + Parameters + ---------- + reference + Discovery reference data. + predictions + Formation-energy predictions. + max_error_threshold + Maximum absolute formation-energy error to retain. + decimals + Number of decimal places used for evaluation. + + Returns + ------- + tuple[pandas.DataFrame, pandas.Series] + Prepared reference data and aligned predictions. + """ + if max_error_threshold is not None and ( + not math.isfinite(max_error_threshold) or max_error_threshold < 0 + ): + raise ValueError("max_error_threshold must be finite, non-negative, or None") + if decimals < 0: + raise ValueError("decimals must be non-negative") + + indexed_reference = _validated_reference_frame(reference) + model_predictions = ( + predictions.copy() + if isinstance(predictions, pd.Series) + else prediction_series(predictions) + ) + aligned_predictions = _align_predictions_prepared( + indexed_reference, model_predictions + ) + energy_columns = [E_ABOVE_HULL, REFERENCE_FORMATION_ENERGY] + numeric_energies = indexed_reference[energy_columns].apply(pd.to_numeric) + prepared_reference = indexed_reference.copy() + prepared_reference[energy_columns] = numeric_energies.round(decimals) + if max_error_threshold is not None: + outlier_mask = ( + aligned_predictions - numeric_energies[REFERENCE_FORMATION_ENERGY] + ).abs() > max_error_threshold + aligned_predictions = aligned_predictions.mask(outlier_mask) + + return prepared_reference, aligned_predictions.round(decimals) + + +def discovery_subset_indices( + reference: pd.DataFrame, + predictions: pd.DataFrame | pd.Series, + *, + max_error_threshold: float | None = MAX_E_FORM_ERROR_THRESHOLD, +) -> dict[DiscoverySubset, pd.Index]: + """ + Return benchmark subsets after applying artifact preprocessing. + + Parameters + ---------- + reference + Discovery reference data. + predictions + Formation-energy predictions. + max_error_threshold + Maximum absolute formation-energy error to retain. + + Returns + ------- + dict[DiscoverySubset, pandas.Index] + Material identifiers for each discovery subset. + """ + prepared_reference, prepared_predictions = prepare_discovery_inputs( + reference, + predictions, + max_error_threshold=max_error_threshold, + decimals=EVALUATION_DECIMALS, + ) + _, each_pred = _hull_distances(prepared_reference, prepared_predictions) + return _discovery_subset_indices_prepared(prepared_reference, each_pred) + + +def calc_discovery_metrics( + reference: pd.DataFrame, + predictions: pd.DataFrame | pd.Series, + *, + subset_indices: SubsetIndices | None = None, + uniq_proto_prevalence: float | None = None, + canonical: bool = False, + max_error_threshold: float | None = MAX_E_FORM_ERROR_THRESHOLD, +) -> dict[DiscoverySubset, dict[str, MetricValue]]: + """ + Calculate metrics after applying artifact masking and rounding. + + Parameters + ---------- + reference + Discovery reference data. + predictions + Formation-energy predictions. + subset_indices + Optional material identifiers for each subset. + uniq_proto_prevalence + Stable-material prevalence among unique prototypes. + canonical + Whether to require canonical leaderboard inputs. + max_error_threshold + Maximum absolute formation-energy error to retain. + + Returns + ------- + dict[DiscoverySubset, dict[str, MetricValue]] + Metrics grouped by discovery subset. + """ + prepared_reference, prepared_predictions = prepare_discovery_inputs( + reference, + predictions, + max_error_threshold=max_error_threshold, + decimals=EVALUATION_DECIMALS, + ) + metrics_by_subset, _ = _calc_discovery_metrics_prepared( + prepared_reference, + prepared_predictions, + subset_indices=subset_indices, + uniq_proto_prevalence=uniq_proto_prevalence, + canonical=canonical, + ) + return metrics_by_subset + + +def _json_safe_metric(value: MetricValue) -> JsonMetricValue: + """ + Round a metric and convert non-finite values to JSON null. + + Parameters + ---------- + value + Metric value to serialize. + + Returns + ------- + float or int or None + JSON-safe metric value. + """ + if isinstance(value, int): + return value + numeric_value = float(value) + return ( + round(numeric_value, EVALUATION_DECIMALS) + if math.isfinite(numeric_value) + else None + ) + + +def evaluate_discovery( + reference: pd.DataFrame, + predictions: pd.DataFrame | pd.Series, + *, + canonical: bool = False, + uniq_proto_prevalence: float | None = None, + max_error_threshold: float | None = MAX_E_FORM_ERROR_THRESHOLD, +) -> DiscoveryResults: + """ + Evaluate formation-energy predictions on the three discovery subsets. + + Leaderboard mode takes unrounded unique-prototype prevalence; synthetic mode + derives it from the prepared reference. + + Parameters + ---------- + reference + Discovery reference data. + predictions + Formation-energy predictions. + canonical + Whether to require canonical leaderboard inputs. + uniq_proto_prevalence + Stable-material prevalence among unique prototypes. + max_error_threshold + Maximum absolute formation-energy error to retain. + + Returns + ------- + DiscoveryResults + JSON-compatible evaluation result. + """ + prepared_reference, prepared_predictions = prepare_discovery_inputs( + reference, + predictions, + max_error_threshold=max_error_threshold, + decimals=EVALUATION_DECIMALS, + ) + raw_metrics, subset_indices = _calc_discovery_metrics_prepared( + prepared_reference, + prepared_predictions, + subset_indices=None, + uniq_proto_prevalence=uniq_proto_prevalence, + canonical=canonical, + ) + + subsets: DiscoverySubsetResults = {} + for subset in DiscoverySubset: + subset_index = subset_indices[subset] + subset_metrics = { + metric_name: _json_safe_metric(metric_value) + for metric_name, metric_value in raw_metrics[subset].items() + } + subset_metrics[MISSING_PREDICTIONS_KEY] = int( + prepared_predictions.loc[subset_index].isna().sum() + ) + subsets[str(subset)] = subset_metrics + return { + "schema_version": RESULT_SCHEMA_VERSION, + "source": { + "framework": MATBENCH_DISCOVERY_ID, + "version": MATBENCH_DISCOVERY_VERSION, + }, + "subsets": subsets, + } + + +def evaluate_discovery_paths( + reference_path: PathLike, + prediction_path: PathLike, + *, + canonical: bool = False, + uniq_proto_prevalence: float | None = None, + max_error_threshold: float | None = MAX_E_FORM_ERROR_THRESHOLD, +) -> DiscoveryResults: + """ + Load local CSV artifacts and evaluate discovery predictions without writes. + + Parameters + ---------- + reference_path + Local reference CSV path. + prediction_path + Local prediction CSV path. + canonical + Whether to require canonical leaderboard inputs. + uniq_proto_prevalence + Stable-material prevalence among unique prototypes. + max_error_threshold + Maximum absolute formation-energy error to retain. + + Returns + ------- + DiscoveryResults + JSON-compatible evaluation result. + """ + return evaluate_discovery( + read_csv_artifact(reference_path, dtype={MATERIAL_ID: str}), + read_csv_artifact(prediction_path, dtype={MATERIAL_ID: str}), + canonical=canonical, + uniq_proto_prevalence=uniq_proto_prevalence, + max_error_threshold=max_error_threshold, + ) + + +def write_discovery_metrics_json( + results: DiscoveryResults, + output_path: str | os.PathLike[str], +) -> None: + """ + Write discovery metrics as JSON. + + Parameters + ---------- + results + Evaluation results to serialize. + output_path + Destination JSON path. + """ + with open(output_path, mode="w", encoding="utf-8") as file: + json.dump(results, file, allow_nan=False, indent=2, sort_keys=True) + file.write("\n") diff --git a/ml_peg/analysis/bulk_crystal/materials_discovery/metrics.py b/ml_peg/analysis/bulk_crystal/materials_discovery/metrics.py new file mode 100644 index 000000000..88af83699 --- /dev/null +++ b/ml_peg/analysis/bulk_crystal/materials_discovery/metrics.py @@ -0,0 +1,506 @@ +"""Classification, regression, and ranking metrics for materials discovery.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Final, TypeAlias + +import numpy as np +from numpy import typing as npt +import pandas as pd +from sklearn.metrics import r2_score + +from ml_peg.analysis.bulk_crystal.materials_discovery.schema import ( + E_ABOVE_HULL, + MATERIAL_ID, + REFERENCE_FORMATION_ENERGY, + UNIQUE_PROTOTYPE, + DiscoverySubset, + _validated_reference_frame, +) + +STABILITY_THRESHOLD: Final = 0.0 +MOST_STABLE_COUNT: Final = 10_000 + +NumericValues: TypeAlias = Sequence[float | None] | pd.Series | npt.NDArray[np.generic] +MetricValue: TypeAlias = float | int +SubsetIndices: TypeAlias = Mapping[DiscoverySubset | str, pd.Index] +_ClassificationMasks: TypeAlias = tuple[pd.Series, pd.Series, pd.Series, pd.Series] + + +def _classify_stable( + each_true: NumericValues, + each_pred: NumericValues, + *, + stability_threshold: float, + fillna: bool, +) -> tuple[_ClassificationMasks, pd.Series, pd.Series]: + """ + Classify stability while retaining numeric inputs for regression. + + Parameters + ---------- + each_true + True hull distances. + each_pred + Predicted hull distances. + stability_threshold + Maximum hull distance considered stable. + fillna + Whether missing predictions count as unstable. + + Returns + ------- + tuple + Classification masks and numeric true and predicted values. + """ + if len(each_true) != len(each_pred): + raise ValueError(f"len(each_true)={len(each_true)} != {len(each_pred)=}") + + each_true_array = pd.to_numeric( + pd.Series(each_true).reset_index(drop=True), errors="coerce" + ) + each_pred_array = pd.to_numeric( + pd.Series(each_pred).reset_index(drop=True), errors="coerce" + ) + if stability_threshold is None or not np.isfinite(stability_threshold): + raise ValueError("stability_threshold must be a real number") + + actual_positive = each_true_array <= stability_threshold + actual_negative = each_true_array > stability_threshold + model_positive = each_pred_array <= stability_threshold + model_negative = each_pred_array > stability_threshold + if fillna: + missing_prediction_mask = each_pred_array.isna() + # Missing predictions count as unstable for both model class masks. + model_positive[missing_prediction_mask] = False + model_negative[missing_prediction_mask] = True + + masks = ( + actual_positive & model_positive, + actual_positive & model_negative, + actual_negative & model_positive, + actual_negative & model_negative, + ) + return masks, each_true_array, each_pred_array + + +def classify_stable( + each_true: NumericValues, + each_pred: NumericValues, + *, + stability_threshold: float = STABILITY_THRESHOLD, + fillna: bool = True, +) -> _ClassificationMasks: + """ + Return classification masks for stable and unstable materials. + + Parameters + ---------- + each_true + True hull distances. + each_pred + Predicted hull distances. + stability_threshold + Maximum hull distance considered stable. + fillna + Whether missing predictions count as unstable. + + Returns + ------- + tuple[pandas.Series, pandas.Series, pandas.Series, pandas.Series] + True-positive, false-negative, false-positive, and true-negative masks. + """ + masks, _, _ = _classify_stable( + each_true, + each_pred, + stability_threshold=stability_threshold, + fillna=fillna, + ) + return masks + + +def _safe_ratio(numerator: float | int, denominator: float | int) -> float: + """ + Divide positive-denominator values, otherwise returning NaN. + + Parameters + ---------- + numerator + Ratio numerator. + denominator + Ratio denominator. + + Returns + ------- + float + Ratio or NaN for a non-positive denominator. + """ + return numerator / denominator if denominator > 0 else float("nan") + + +def stable_metrics( + each_true: NumericValues, + each_pred: NumericValues, + *, + stability_threshold: float = STABILITY_THRESHOLD, + fillna: bool = True, +) -> dict[str, MetricValue]: + """ + Calculate classification and hull-distance regression metrics. + + Inputs should contain finite values or missing values. Artifact evaluation should + use ``evaluate_discovery``, which also masks outliers and infinities. + + Parameters + ---------- + each_true + True hull distances. + each_pred + Predicted hull distances. + stability_threshold + Maximum hull distance considered stable. + fillna + Whether missing predictions count as unstable. + + Returns + ------- + dict[str, MetricValue] + Classification and regression metrics. + """ + masks, each_true_array, each_pred_array = _classify_stable( + each_true, + each_pred, + stability_threshold=stability_threshold, + fillna=fillna, + ) + ( + true_positive_count, + false_negative_count, + false_positive_count, + true_negative_count, + ) = (int(mask.sum()) for mask in masks) + + total_positive_count = true_positive_count + false_negative_count + total_negative_count = true_negative_count + false_positive_count + classified_count = total_positive_count + total_negative_count + # Prevalence is the discovery rate from random selection over this population. + prevalence = _safe_ratio(total_positive_count, classified_count) + predicted_positive_count = true_positive_count + false_positive_count + precision = _safe_ratio(true_positive_count, predicted_positive_count) + recall = _safe_ratio(true_positive_count, total_positive_count) + true_positive_rate = recall + false_positive_rate = _safe_ratio(false_positive_count, total_negative_count) + true_negative_rate = _safe_ratio(true_negative_count, total_negative_count) + false_negative_rate = _safe_ratio(false_negative_count, total_positive_count) + + # False positives plus true negatives must account for all actual negatives. + if ( + false_positive_rate > 0 + and true_negative_rate > 0 + and not np.isclose(false_positive_rate + true_negative_rate, 1) + ): + raise ValueError( + f"FPR={false_positive_rate} and TNR={true_negative_rate} do not add up to 1" + ) + # True positives plus false negatives must account for all actual positives. + if ( + true_positive_rate > 0 + and false_negative_rate > 0 + and not np.isclose(true_positive_rate + false_negative_rate, 1) + ): + raise ValueError( + f"TPR={true_positive_rate} and FNR={false_negative_rate} do not add up to 1" + ) + + missing_pair_mask = each_true_array.isna() | each_pred_array.isna() + valid_true = each_true_array[~missing_pair_mask].to_numpy() + valid_pred = each_pred_array[~missing_pair_mask].to_numpy() + f1_score = ( + float("nan") + if precision + recall == 0 + else 2 * precision * recall / (precision + recall) + ) + if len(valid_true) == 0: + mean_absolute_error = root_mean_squared_error = float("nan") + else: + prediction_errors = valid_true - valid_pred + mean_absolute_error = float(np.abs(prediction_errors).mean()) + root_mean_squared_error = float((prediction_errors**2).mean() ** 0.5) + + return { + "F1": f1_score, + "DAF": _safe_ratio(precision, prevalence), + "Precision": precision, + "Recall": recall, + "Accuracy": _safe_ratio( + true_positive_count + true_negative_count, classified_count + ), + "TPR": true_positive_rate, + "FPR": false_positive_rate, + "TNR": true_negative_rate, + "FNR": false_negative_rate, + "TP": true_positive_count, + "FP": false_positive_count, + "TN": true_negative_count, + "FN": false_negative_count, + "MAE": mean_absolute_error, + "RMSE": root_mean_squared_error, + "R2": ( + float(r2_score(valid_true, valid_pred)) + if len(valid_true) > 1 + else float("nan") + ), + } + + +def _align_predictions_prepared( + indexed_reference: pd.DataFrame, model_predictions: pd.Series +) -> pd.Series: + """ + Align predictions to an already-validated reference index. + + Parameters + ---------- + indexed_reference + Validated reference data indexed by material ID. + model_predictions + Predictions indexed by material ID. + + Returns + ------- + pandas.Series + Numeric predictions aligned to the reference. + """ + if model_predictions.index.hasnans: + raise ValueError("discovery predictions contain missing material_id values") + model_predictions = model_predictions.copy() + model_predictions.index = model_predictions.index.astype(str) + if model_predictions.index.has_duplicates: + duplicate_ids = ( + model_predictions.index[model_predictions.index.duplicated()] + .unique() + .tolist() + ) + raise ValueError( + "discovery predictions contain duplicate material_id values: " + f"{duplicate_ids!r}" + ) + unknown_ids = model_predictions.index.difference(indexed_reference.index) + if len(unknown_ids) > 0: + rendered_ids = sorted(map(str, unknown_ids)) + raise ValueError(f"Predictions contain unknown material IDs: {rendered_ids!r}") + aligned_predictions = model_predictions.reindex(indexed_reference.index) + aligned_predictions.index.name = MATERIAL_ID + return pd.to_numeric(aligned_predictions, errors="coerce").replace( + [np.inf, -np.inf], np.nan + ) + + +def align_predictions( + reference: pd.DataFrame, model_predictions: pd.Series +) -> pd.Series: + """ + Align predictions to reference order, rejecting invalid or unknown IDs. + + Missing reference IDs remain as NaN predictions. + + Parameters + ---------- + reference + Discovery reference data. + model_predictions + Predictions indexed by material ID. + + Returns + ------- + pandas.Series + Numeric predictions aligned to the reference. + """ + return _align_predictions_prepared( + _validated_reference_frame(reference), model_predictions + ) + + +def _hull_distances( + indexed_reference: pd.DataFrame, aligned_predictions: pd.Series +) -> tuple[pd.Series, pd.Series]: + """ + Return true and predicted hull distances from prepared inputs. + + Parameters + ---------- + indexed_reference + Validated reference data indexed by material ID. + aligned_predictions + Formation-energy predictions aligned to the reference. + + Returns + ------- + tuple[pandas.Series, pandas.Series] + True and predicted hull distances. + """ + each_true = pd.to_numeric(indexed_reference[E_ABOVE_HULL], errors="coerce") + reference_formation_energy = pd.to_numeric( + indexed_reference[REFERENCE_FORMATION_ENERGY], errors="coerce" + ) + return ( + each_true, + each_true + aligned_predictions - reference_formation_energy, + ) + + +def _discovery_subset_indices_prepared( + indexed_reference: pd.DataFrame, + each_pred: pd.Series, +) -> dict[DiscoverySubset, pd.Index]: + """ + Return the three subsets from prepared reference and hull distances. + + Parameters + ---------- + indexed_reference + Validated reference data indexed by material ID. + each_pred + Predicted hull distances aligned to the reference. + + Returns + ------- + dict[DiscoverySubset, pandas.Index] + Material identifiers for each discovery subset. + """ + unique_prototype_index = indexed_reference.index[ + indexed_reference[UNIQUE_PROTOTYPE].astype(bool) + ] + most_stable_index = ( + each_pred.loc[unique_prototype_index] + .sort_values(na_position="last", kind="stable") + .head(MOST_STABLE_COUNT) + .index + ) + return { + DiscoverySubset.full_test_set: indexed_reference.index, + DiscoverySubset.unique_prototypes: unique_prototype_index, + DiscoverySubset.most_stable_10k: most_stable_index, + } + + +def _normalized_subset_indices( + subset_indices: SubsetIndices, + reference_index: pd.Index, +) -> dict[DiscoverySubset, pd.Index]: + """ + Normalize and validate given subset indices. + + Parameters + ---------- + subset_indices + Material identifiers keyed by discovery subset. + reference_index + Valid reference material identifiers. + + Returns + ------- + dict[DiscoverySubset, pandas.Index] + Validated material identifiers for every subset. + """ + normalized: dict[DiscoverySubset, pd.Index] = {} + for subset_key, identifiers in subset_indices.items(): + try: + subset = DiscoverySubset(subset_key) + except ValueError as exc: + raise ValueError(f"Unknown discovery subset {subset_key!r}") from exc + subset_index = pd.Index(identifiers, name=MATERIAL_ID) + if subset_index.has_duplicates: + raise ValueError(f"{subset} subset contains duplicate material IDs") + unknown_ids = subset_index.difference(reference_index) + if len(unknown_ids) > 0: + raise ValueError( + f"{subset} subset contains unknown material IDs: " + f"{sorted(map(str, unknown_ids))!r}" + ) + normalized[subset] = subset_index + + missing_subsets = set(DiscoverySubset) - set(normalized) + if missing_subsets: + raise ValueError( + f"Missing required discovery subsets: {sorted(map(str, missing_subsets))!r}" + ) + return normalized + + +def _calc_discovery_metrics_prepared( + indexed_reference: pd.DataFrame, + aligned_predictions: pd.Series, + *, + subset_indices: SubsetIndices | None, + uniq_proto_prevalence: float | None, + canonical: bool, +) -> tuple[ + dict[DiscoverySubset, dict[str, MetricValue]], + dict[DiscoverySubset, pd.Index], +]: + """ + Calculate metrics from validated, aligned discovery inputs. + + Parameters + ---------- + indexed_reference + Validated reference data indexed by material ID. + aligned_predictions + Formation-energy predictions aligned to the reference. + subset_indices + Optional material identifiers for each subset. + uniq_proto_prevalence + Stable-material prevalence among unique prototypes. + canonical + Whether to require canonical leaderboard inputs. + + Returns + ------- + tuple + Metrics and material identifiers grouped by discovery subset. + """ + each_true, each_pred = _hull_distances(indexed_reference, aligned_predictions) + canonical_indices = ( + _discovery_subset_indices_prepared(indexed_reference, each_pred) + if subset_indices is None + else _normalized_subset_indices(subset_indices, indexed_reference.index) + ) + metrics_by_subset = { + subset: stable_metrics( + each_true.loc[subset_index], + each_pred.loc[subset_index], + fillna=True, + ) + for subset, subset_index in canonical_indices.items() + } + + if canonical and uniq_proto_prevalence is None: + raise ValueError( + "leaderboard evaluation requires explicit unrounded " + "unique-prototype prevalence" + ) + if uniq_proto_prevalence is None: + unique_each_true = each_true.loc[ + canonical_indices[DiscoverySubset.unique_prototypes] + ] + uniq_proto_prevalence = float((unique_each_true <= STABILITY_THRESHOLD).mean()) + elif not np.isfinite(uniq_proto_prevalence) or not ( + 0 <= uniq_proto_prevalence <= 1 + ): + raise ValueError( + "uniq_proto_prevalence must be a finite fraction between 0 and 1" + ) + + daf_denominator = ( + uniq_proto_prevalence if uniq_proto_prevalence > 0 else float("nan") + ) + for subset in ( + DiscoverySubset.unique_prototypes, + DiscoverySubset.most_stable_10k, + ): + metrics_by_subset[subset]["DAF"] = ( + metrics_by_subset[subset]["Precision"] / daf_denominator + ) + return metrics_by_subset, canonical_indices diff --git a/ml_peg/analysis/bulk_crystal/materials_discovery/schema.py b/ml_peg/analysis/bulk_crystal/materials_discovery/schema.py new file mode 100644 index 000000000..97d12c2ac --- /dev/null +++ b/ml_peg/analysis/bulk_crystal/materials_discovery/schema.py @@ -0,0 +1,191 @@ +"""Schemas for materials-discovery reference and prediction artifacts.""" + +from __future__ import annotations + +from enum import Enum +from numbers import Real +from typing import Final + +import numpy as np +import pandas as pd + +from ml_peg.data.artifacts import material_id_index, validate_required_columns + +MATERIAL_ID: Final = "material_id" +E_ABOVE_HULL: Final = "e_above_hull_mp2020_corrected_ppd_mp" +REFERENCE_FORMATION_ENERGY: Final = "e_form_per_atom_mp2020_corrected" +UNIQUE_PROTOTYPE: Final = "unique_prototype" +PREDICTED_FORMATION_ENERGY: Final = "e_form_per_atom" + +REFERENCE_COLUMNS: Final[tuple[str, ...]] = ( + E_ABOVE_HULL, + REFERENCE_FORMATION_ENERGY, + UNIQUE_PROTOTYPE, +) +PREDICTION_COLUMNS: Final[tuple[str, ...]] = (PREDICTED_FORMATION_ENERGY,) + + +class DiscoverySubset(str, Enum): + """Subsets reported by materials-discovery evaluation.""" + + full_test_set = "full_test_set" + unique_prototypes = "unique_prototypes" + most_stable_10k = "most_stable_10k" + + def __str__(self) -> str: + """ + Return the serialized subset key. + + Returns + ------- + str + Serialized subset key. + """ + return self.value + + +def index_by_material_id( + dataframe: pd.DataFrame, + *, + artifact_name: str, +) -> pd.DataFrame: + """ + Return a copy indexed by validated material identifiers. + + Parameters + ---------- + dataframe + Artifact data containing material identifiers. + artifact_name + Artifact label used in error messages. + + Returns + ------- + pandas.DataFrame + Copy indexed by material ID. + """ + identifiers = material_id_index( + dataframe, id_column=MATERIAL_ID, artifact_name=artifact_name + ) + identifiers = identifiers.astype(str) + if identifiers.has_duplicates: + duplicate_ids = identifiers[identifiers.duplicated()].unique().tolist() + raise ValueError( + f"{artifact_name} contains duplicate {MATERIAL_ID!r} values after " + f"string conversion: {duplicate_ids!r}" + ) + indexed_dataframe = dataframe.drop(columns=MATERIAL_ID, errors="ignore").copy() + indexed_dataframe.index = identifiers + return indexed_dataframe + + +def _validated_frame( + dataframe: pd.DataFrame, + required_columns: tuple[str, ...], + artifact_name: str, +) -> pd.DataFrame: + """ + Validate required columns and IDs, returning an indexed copy. + + Parameters + ---------- + dataframe + Artifact data to validate. + required_columns + Column names that must be present. + artifact_name + Artifact label used in error messages. + + Returns + ------- + pandas.DataFrame + Validated copy indexed by material ID. + """ + validate_required_columns(dataframe, required_columns, artifact_name=artifact_name) + return index_by_material_id(dataframe, artifact_name=artifact_name) + + +def _validated_reference_frame(dataframe: pd.DataFrame) -> pd.DataFrame: + """ + Validate a discovery reference and return its indexed copy. + + Parameters + ---------- + dataframe + Discovery reference data. + + Returns + ------- + pandas.DataFrame + Validated copy indexed by material ID. + """ + indexed_dataframe = _validated_frame( + dataframe, REFERENCE_COLUMNS, "discovery reference" + ) + for energy_column in (E_ABOVE_HULL, REFERENCE_FORMATION_ENERGY): + numeric_values = pd.to_numeric( + indexed_dataframe[energy_column], errors="coerce" + ) + invalid_mask = ~np.isfinite(numeric_values.to_numpy(dtype=float)) + if invalid_mask.any(): + invalid_values = indexed_dataframe.loc[invalid_mask, energy_column].tolist() + raise ValueError( + f"{energy_column!r} values must be finite, got {invalid_values!r}" + ) + unique_prototype_flags = indexed_dataframe[UNIQUE_PROTOTYPE] + invalid_flags = ~unique_prototype_flags.map( + lambda value: ( + isinstance(value, (bool, np.bool_)) + or (isinstance(value, Real) and value in (0, 1)) + ) + ) + if invalid_flags.any(): + invalid_values = unique_prototype_flags[invalid_flags].tolist() + raise ValueError( + f"{UNIQUE_PROTOTYPE!r} values must be boolean, got {invalid_values!r}" + ) + return indexed_dataframe + + +def validate_reference_frame(dataframe: pd.DataFrame) -> None: + """ + Validate discovery reference columns, IDs, energies, and flags. + + Parameters + ---------- + dataframe + Discovery reference data. + """ + _validated_reference_frame(dataframe) + + +def validate_prediction_frame(dataframe: pd.DataFrame) -> None: + """ + Validate discovery prediction columns and material IDs. + + Parameters + ---------- + dataframe + Discovery prediction data. + """ + _validated_frame(dataframe, PREDICTION_COLUMNS, "discovery predictions") + + +def prediction_series(dataframe: pd.DataFrame) -> pd.Series: + """ + Return validated formation-energy predictions indexed by material ID. + + Parameters + ---------- + dataframe + Discovery prediction data. + + Returns + ------- + pandas.Series + Formation-energy predictions indexed by material ID. + """ + indexed_dataframe = _validated_frame( + dataframe, PREDICTION_COLUMNS, "discovery predictions" + ) + return indexed_dataframe[PREDICTED_FORMATION_ENERGY] diff --git a/ml_peg/data/artifacts.py b/ml_peg/data/artifacts.py new file mode 100644 index 000000000..f29cf6e8c --- /dev/null +++ b/ml_peg/data/artifacts.py @@ -0,0 +1,324 @@ +"""Typed helpers for reading and naming local benchmark artifacts.""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import date +from decimal import Decimal, InvalidOperation +from enum import Enum +import os +import re +from typing import Any, Final + +import pandas as pd + +PathLike = str | os.PathLike[str] + +MATBENCH_DISCOVERY_ID: Final = "matbench-discovery" +MATBENCH_DISCOVERY_VERSION: Final = "1.3.1" + +ISO_DATE_PATTERN: Final = re.compile(r"^\d{4}-\d{2}-\d{2}$") +MOYO_VERSION_PATTERN: Final = re.compile( + r"^[0-9]+(?:\.[0-9]+)*(?:[-+][0-9A-Za-z.-]+)?$" +) +_GEO_OPT_ANALYSIS_SUFFIX: Final = re.compile( + r"^geo-opt-symprec=([^=]+)-moyo=([^=]+)\.csv\.gz$" +) + + +class ArtifactRole(str, Enum): + """Canonical roles used in dated model artifact filenames.""" + + discovery = "discovery" + geo_opt = "geo_opt" + geo_opt_analysis = "geo_opt_analysis" + + def __str__(self) -> str: + """ + Return the role value. + + Returns + ------- + str + Serialized role value. + """ + return self.value + + +ARTIFACT_SUFFIXES: Final[dict[str, str]] = { + str(ArtifactRole.discovery): "discovery.csv.gz", + str(ArtifactRole.geo_opt): "geo-opt.jsonl.gz", +} + + +def _checked_file_path(file_path: PathLike) -> str: + """ + Return a local artifact path after confirming it names a file. + + Parameters + ---------- + file_path + Local artifact path. + + Returns + ------- + str + Validated filesystem path. + """ + normalized_path = os.fspath(file_path) + if not os.path.isfile(normalized_path): + raise FileNotFoundError(f"Artifact file not found: {normalized_path!r}") + return normalized_path + + +def read_csv_artifact(file_path: PathLike, **read_options: Any) -> pd.DataFrame: + """ + Read a CSV artifact with compression inferred from its filename. + + Parameters + ---------- + file_path + Local CSV artifact path. + **read_options + Additional options passed to :func:`pandas.read_csv`. + + Returns + ------- + pandas.DataFrame + Loaded artifact data. + """ + return pd.read_csv( + _checked_file_path(file_path), compression="infer", **read_options + ) + + +def read_jsonl_artifact(file_path: PathLike, **read_options: Any) -> pd.DataFrame: + """ + Read a line-delimited JSON artifact with transparent compression. + + Parameters + ---------- + file_path + Local JSON Lines artifact path. + **read_options + Additional options passed to :func:`pandas.read_json`. + + Returns + ------- + pandas.DataFrame + Loaded artifact data. + """ + return pd.read_json( + _checked_file_path(file_path), + lines=True, + compression="infer", + **read_options, + ) + + +def validate_required_columns( + dataframe: pd.DataFrame, + required_columns: Sequence[str], + *, + artifact_name: str = "dataframe", +) -> None: + """ + Raise if a dataframe lacks any required columns. + + Parameters + ---------- + dataframe + Dataframe to validate. + required_columns + Column names that must be present. + artifact_name + Artifact label used in error messages. + """ + missing_columns = set(required_columns) - set(dataframe.columns) + if missing_columns: + raise ValueError( + f"{artifact_name} missing required columns: {sorted(missing_columns)}" + ) + + +def material_id_index( + dataframe: pd.DataFrame, + *, + id_column: str = "material_id", + artifact_name: str = "dataframe", +) -> pd.Index: + """ + Return IDs from a column or named index after validating uniqueness. + + Parameters + ---------- + dataframe + Dataframe containing material identifiers. + id_column + Material identifier column or index name. + artifact_name + Artifact label used in error messages. + + Returns + ------- + pandas.Index + Validated material identifiers. + """ + has_id_column = id_column in dataframe.columns + has_id_index = dataframe.index.name == id_column + if not has_id_column and not has_id_index: + raise ValueError( + f"{artifact_name} must contain {id_column!r} as a column or index" + ) + + if has_id_column: + identifiers = pd.Index(dataframe[id_column], name=id_column) + if has_id_index and not identifiers.equals(dataframe.index): + raise ValueError( + f"{artifact_name} has inconsistent {id_column!r} column and index" + ) + else: + identifiers = dataframe.index.copy() + + if identifiers.hasnans: + raise ValueError(f"{artifact_name} contains missing {id_column!r} values") + if identifiers.has_duplicates: + duplicate_ids = identifiers[identifiers.duplicated()].unique().tolist() + raise ValueError( + f"{artifact_name} contains duplicate {id_column!r} values: " + f"{duplicate_ids!r}" + ) + return identifiers + + +def canonical_scientific_notation(value: float | str | Decimal) -> str: + """ + Format a positive finite number as canonical notation like ``1e-5``. + + Parameters + ---------- + value + Positive finite numeric value. + + Returns + ------- + str + Canonical scientific notation. + """ + try: + decimal_value = Decimal(str(value)) + except InvalidOperation as exc: + raise ValueError(f"Invalid numeric value {value!r}") from exc + if not decimal_value.is_finite() or decimal_value <= 0: + raise ValueError(f"Expected a positive finite number, got {value!r}") + + mantissa, _, exponent = f"{decimal_value.normalize():e}".partition("e") + return f"{mantissa.rstrip('0').rstrip('.')}e{int(exponent)}" + + +def _iso_date(value: date | str) -> str: + """ + Return a validated ``YYYY-MM-DD`` calendar date. + + Parameters + ---------- + value + Date object or ISO date string. + + Returns + ------- + str + Validated ISO date. + """ + iso_date = value.isoformat() if isinstance(value, date) else value + if not ISO_DATE_PATTERN.fullmatch(iso_date): + raise ValueError(f"Expected an ISO date, got {value!r}") + try: + date.fromisoformat(iso_date) + except ValueError as exc: + raise ValueError(f"Invalid ISO date {value!r}") from exc + return iso_date + + +def artifact_filename( + artifact_date: date | str, + role: str | ArtifactRole, + *, + symprec: float | str | Decimal | None = None, + moyo_version: str | None = None, +) -> str: + """ + Return a canonical dated basename for the requested artifact role. + + Parameters + ---------- + artifact_date + Artifact date. + role + Artifact role. + symprec + Symmetry tolerance for geometry-optimization analysis. + moyo_version + Moyo version for geometry-optimization analysis. + + Returns + ------- + str + Canonical artifact basename. + """ + iso_date = _iso_date(artifact_date) + role_value = str(role) + if role_value == ArtifactRole.geo_opt_analysis: + if symprec is None or moyo_version is None: + raise ValueError( + "symprec and moyo_version are required for geo_opt_analysis" + ) + if not MOYO_VERSION_PATTERN.fullmatch(moyo_version): + raise ValueError(f"Invalid moyo version {moyo_version!r}") + suffix = ( + f"geo-opt-symprec={canonical_scientific_notation(symprec)}" + f"-moyo={moyo_version}.csv.gz" + ) + else: + if symprec is not None or moyo_version is not None: + raise ValueError("symprec and moyo_version are only for geo_opt_analysis") + if (suffix := ARTIFACT_SUFFIXES.get(role_value)) is None: + raise ValueError(f"Unknown artifact role {role_value!r}") + return f"{iso_date}-{suffix}" + + +def parse_artifact_filename(filename: str) -> ArtifactRole: + """ + Validate a canonical artifact filename or path and return its role. + + Parameters + ---------- + filename + Artifact filename or path. + + Returns + ------- + ArtifactRole + Parsed artifact role. + """ + basename = os.path.basename(filename) + if not ISO_DATE_PATTERN.match(basename[:10]) or basename[10:11] != "-": + raise ValueError(f"Not a canonical model artifact filename: {filename!r}") + artifact_date, suffix = basename[:10], basename[11:] + _iso_date(artifact_date) + for role_value, expected_suffix in ARTIFACT_SUFFIXES.items(): + if suffix == expected_suffix: + return ArtifactRole(role_value) + if match := _GEO_OPT_ANALYSIS_SUFFIX.fullmatch(suffix): + symprec, moyo_version = match.groups() + if ( + artifact_filename( + artifact_date, + ArtifactRole.geo_opt_analysis, + symprec=symprec, + moyo_version=moyo_version, + ) + == basename + ): + return ArtifactRole.geo_opt_analysis + raise ValueError(f"Not a canonical model artifact filename: {filename!r}") diff --git a/tests/metrics/test_discovery_metrics.py b/tests/metrics/test_discovery_metrics.py new file mode 100644 index 000000000..a1b88c4eb --- /dev/null +++ b/tests/metrics/test_discovery_metrics.py @@ -0,0 +1,645 @@ +"""Tests for materials-discovery metrics and artifact handling.""" + +from __future__ import annotations + +from collections.abc import Callable +import json +import math +from pathlib import Path +import warnings + +import numpy as np +import pandas as pd +import pytest + +from ml_peg.analysis.bulk_crystal.materials_discovery import ( + E_ABOVE_HULL, + MATERIAL_ID, + PREDICTED_FORMATION_ENERGY, + REFERENCE_FORMATION_ENERGY, + UNIQUE_PROTOTYPE, + DiscoverySubset, + align_predictions, + calc_discovery_metrics, + classify_stable, + discovery_subset_indices, + evaluate_discovery, + evaluate_discovery_paths, + stable_metrics, + validate_prediction_frame, + validate_reference_frame, + write_discovery_metrics_json, +) +from ml_peg.data.artifacts import ( + ArtifactRole, + artifact_filename, + canonical_scientific_notation, + parse_artifact_filename, + read_csv_artifact, + read_jsonl_artifact, +) + +pytestmark = pytest.mark.framework("matbench-discovery") + +REPORTED_METRICS = set( + "F1 DAF Precision Recall Accuracy TPR FPR TNR FNR TP FP TN FN MAE RMSE R2 " + "missing_preds".split() +) +DAF_SUBSETS = ( + DiscoverySubset.unique_prototypes, + DiscoverySubset.most_stable_10k, +) + + +def _reference_frame( + *, + material_ids: list[str] | None = None, + each_true: list[float] | None = None, + formation_energies: list[float] | None = None, + unique_prototypes: list[bool | float | int] | None = None, +) -> pd.DataFrame: + """Build a compact valid discovery reference dataframe.""" + resolved_ids = material_ids or ["wbm-0", "wbm-1", "wbm-2", "wbm-3"] + row_count = len(resolved_ids) + return pd.DataFrame( + { + MATERIAL_ID: resolved_ids, + E_ABOVE_HULL: each_true or [-1.0, -0.5, 0.5, 1.0][:row_count], + REFERENCE_FORMATION_ENERGY: formation_energies or [0.0] * row_count, + UNIQUE_PROTOTYPE: unique_prototypes or [True] * row_count, + } + ) + + +def _prediction_frame( + material_ids: list[str], values: list[float | None] +) -> pd.DataFrame: + """Build a valid discovery prediction dataframe.""" + return pd.DataFrame({MATERIAL_ID: material_ids, PREDICTED_FORMATION_ENERGY: values}) + + +def test_stable_metrics_exact_values() -> None: + """Classification and regression metrics preserve exact source semantics.""" + metrics = stable_metrics( + [-1.0, -0.5, 0.5, 1.0], + [-0.8, 0.2, -0.1, 0.9], + ) + expected = { + "F1": 0.5, + "DAF": 1.0, + "Precision": 0.5, + "Recall": 0.5, + "Accuracy": 0.5, + "TPR": 0.5, + "FPR": 0.5, + "TNR": 0.5, + "FNR": 0.5, + "TP": 1, + "FP": 1, + "TN": 1, + "FN": 1, + "MAE": 0.4, + "RMSE": math.sqrt(0.225), + "R2": 0.64, + } + assert metrics == pytest.approx(expected) + + +def test_stable_metrics_pairs_series_by_position() -> None: + """Ignore Series labels when pairing true and predicted values.""" + each_true = pd.Series([-1.0, -0.5, 0.5, 1.0], index=[10, 11, 12, 13]) + each_pred = pd.Series([-0.8, -0.2, 0.1, 0.9], index=[13, 12, 11, 10]) + + indexed_metrics = stable_metrics(each_true, each_pred) + positional_metrics = stable_metrics(each_true.tolist(), each_pred.tolist()) + + assert indexed_metrics == pytest.approx(positional_metrics) + assert ( + indexed_metrics["TP"], + indexed_metrics["FN"], + indexed_metrics["FP"], + indexed_metrics["TN"], + ) == (2, 0, 0, 2) + + +@pytest.mark.parametrize( + ("each_true", "each_pred", "fillna", "expected_counts"), + [ + ( + [-0.1, 0.0, 0.1, np.nan, None], + [-0.1, 0.0, 0.1, 0.2, -0.2], + True, + (2, 0, 0, 1), + ), + ([0.0, -0.1], [None, 0.1], True, (0, 2, 0, 0)), + ([-0.1, 0.1, np.nan], [-0.1, np.nan, -0.2], False, (1, 0, 0, 0)), + ], + ids=["nullable-truth", "missing-predictions", "no-fill"], +) +def test_classify_stable_handles_nans( + each_true: list[float | None], + each_pred: list[float | None], + fillna: bool, + expected_counts: tuple[int, int, int, int], +) -> None: + """Nullable values retain established classification behavior.""" + masks = classify_stable(each_true, each_pred, fillna=fillna) + assert tuple(int(mask.sum()) for mask in masks) == expected_counts + + +@pytest.mark.parametrize( + ("each_true", "each_pred", "expected"), + [ + ( + [0.1, 0.2, 0.3], + [-0.1, -0.2, -0.3], + (0.0, 1.0, 0.0, np.nan, np.nan, np.nan, np.nan), + ), + ( + [-0.1, -0.2, -0.3], + [0.1, 0.2, 0.3], + (np.nan, np.nan, np.nan, 0.0, 1.0, np.nan, np.nan), + ), + ], + ids=["no-stable-class", "no-unstable-class"], +) +def test_stable_metrics_zero_classes( + each_true: list[float], + each_pred: list[float], + expected: tuple[float, ...], +) -> None: + """Absent positive or negative classes produce the expected NaN rates.""" + metrics = stable_metrics(each_true, each_pred) + metric_names = ("Precision", "FPR", "TNR", "Recall", "FNR", "DAF", "F1") + assert tuple(metrics[name] for name in metric_names) == pytest.approx( + expected, nan_ok=True + ) + + +def test_stable_metrics_regression_edge_cases() -> None: + """Regression metrics ignore missing pairs and handle fewer than two values.""" + with warnings.catch_warnings(): + warnings.simplefilter("error") + missing_metrics = stable_metrics([None, np.nan], [None, np.nan]) + assert all( + math.isnan(float(missing_metrics[name])) for name in ("MAE", "RMSE", "R2") + ) + assert math.isnan(float(stable_metrics([0.1], [0.2])["R2"])) + + filled = stable_metrics([-0.1, 0.1], [None, 0.2], fillna=True) + unfilled = stable_metrics([-0.1, 0.1], [None, 0.2], fillna=False) + assert filled["FN"] == 1 + assert unfilled["FN"] == 0 + assert tuple(filled[name] for name in ("MAE", "RMSE", "R2")) == pytest.approx( + tuple(unfilled[name] for name in ("MAE", "RMSE", "R2")), nan_ok=True + ) + + +@pytest.mark.parametrize( + ("threshold", "expected_counts"), + [(-0.1, (1, 0, 0, 2)), (0.0, (2, 0, 0, 1)), (0.1, (3, 0, 0, 0))], +) +def test_classify_stable_thresholds( + threshold: float, expected_counts: tuple[int, int, int, int] +) -> None: + """Apply the requested stability threshold to both arrays.""" + masks = classify_stable( + [-0.2, 0.0, 0.1], + [-0.2, 0.0, 0.1], + stability_threshold=threshold, + ) + assert tuple(int(mask.sum()) for mask in masks) == expected_counts + + +@pytest.mark.parametrize( + "invalid_threshold", + [None, np.nan, np.inf, -np.inf], + ids=["none", "nan", "positive-infinity", "negative-infinity"], +) +def test_classify_stable_rejects_nonfinite_thresholds( + invalid_threshold: float | None, +) -> None: + """Reject missing and non-finite stability thresholds.""" + with pytest.raises(ValueError, match="stability_threshold must be a real number"): + classify_stable( + [-0.1, 0.1], + [-0.1, 0.1], + stability_threshold=invalid_threshold, # type: ignore[arg-type] + ) + + +@pytest.mark.parametrize( + ("validator", "missing_column"), + [ + (validate_reference_frame, UNIQUE_PROTOTYPE), + (validate_prediction_frame, PREDICTED_FORMATION_ENERGY), + ], + ids=["reference", "predictions"], +) +def test_discovery_schema_validation( + validator: Callable[[pd.DataFrame], None], + missing_column: str, +) -> None: + """Schema validators accept indexed IDs and reject missing columns.""" + reference = _reference_frame().set_index(MATERIAL_ID) + predictions = _prediction_frame( + reference.index.tolist(), [0.0] * len(reference) + ).set_index(MATERIAL_ID) + dataframe = reference if validator is validate_reference_frame else predictions + validator(dataframe) + + with pytest.raises(ValueError, match="missing required columns"): + validator(dataframe.drop(columns=missing_column)) + + +@pytest.mark.parametrize( + ("column", "invalid_values", "match"), + [ + ( + MATERIAL_ID, + ["wbm-0", "wbm-0", "wbm-2", "wbm-3"], + "duplicate.*material_id", + ), + ( + MATERIAL_ID, + [1, "1", 2, 3], + "duplicate.*material_id.*string conversion", + ), + ( + UNIQUE_PROTOTYPE, + ["False", "True", "False", "True"], + "values must be boolean", + ), + (UNIQUE_PROTOTYPE, [True, None, False, True], "values must be boolean"), + (E_ABOVE_HULL, [np.inf, -0.5, 0.5, 1.0], "values must be finite"), + ( + E_ABOVE_HULL, + pd.array([-1.0, pd.NA, 0.5, 1.0], dtype="Float64"), + "values must be finite", + ), + ( + REFERENCE_FORMATION_ENERGY, + [0.0, np.nan, 0.0, 0.0], + "values must be finite", + ), + ], +) +def test_discovery_schema_rejects_invalid_values( + column: str, + invalid_values: object, + match: str, +) -> None: + """Reference fields reject duplicate IDs, non-booleans, and nonfinite energies.""" + reference = _reference_frame() + reference[column] = invalid_values + + with pytest.raises(ValueError, match=match): + validate_reference_frame(reference) + + +def test_discovery_schema_rejects_inconsistent_id_locations() -> None: + """Material IDs in a column and named index must agree.""" + reference = _reference_frame().set_index(MATERIAL_ID, drop=False) + reference.index = pd.Index( + ["other-0", "other-1", "other-2", "other-3"], name=MATERIAL_ID + ) + with pytest.raises(ValueError, match="inconsistent.*column and index"): + validate_reference_frame(reference) + + +@pytest.mark.parametrize( + "prototype_flags", + [[1, 0, 1, 0], [1.0, 0.0, 1.0, 0.0]], + ids=["integers", "floats"], +) +def test_discovery_schema_accepts_binary_prototype_flags( + prototype_flags: list[float | int], +) -> None: + """CSV schemas accept zero/one prototype flags.""" + reference = _reference_frame(unique_prototypes=prototype_flags) + + validate_reference_frame(reference) + unique_index = discovery_subset_indices( + reference, + pd.Series([0.0] * len(reference), index=reference[MATERIAL_ID]), + )[DiscoverySubset.unique_prototypes] + assert unique_index.tolist() == ["wbm-0", "wbm-2"] + + +def test_artifact_readers_support_gzip_csv_and_jsonl(tmp_path: Path) -> None: + """Artifact readers infer gzip compression for CSV and JSONL inputs.""" + dataframe = pd.DataFrame({"material_id": ["wbm-0", "wbm-1"], "value": [1, 2]}) + csv_path = tmp_path / "artifact.csv.gz" + jsonl_path = tmp_path / "artifact.jsonl.gz" + dataframe.to_csv(csv_path, index=False) + dataframe.to_json(jsonl_path, orient="records", lines=True, compression="gzip") + + pd.testing.assert_frame_equal(read_csv_artifact(csv_path), dataframe) + pd.testing.assert_frame_equal(read_jsonl_artifact(jsonl_path), dataframe) + + +@pytest.mark.parametrize( + ("value", "expected"), [(1e-5, "1e-5"), ("0.0100", "1e-2"), (2.5, "2.5e0")] +) +def test_canonical_scientific_notation(value: float | str, expected: str) -> None: + """Scientific notation is normalized for artifact names.""" + assert canonical_scientific_notation(value) == expected + + +@pytest.mark.parametrize( + ("role", "expected"), + [ + (ArtifactRole.discovery, "2026-07-18-discovery.csv.gz"), + (ArtifactRole.geo_opt, "2026-07-18-geo-opt.jsonl.gz"), + ( + ArtifactRole.geo_opt_analysis, + "2026-07-18-geo-opt-symprec=1e-5-moyo=0.12.0.csv.gz", + ), + ], +) +def test_artifact_names_round_trip(role: ArtifactRole, expected: str) -> None: + """Dated artifact names round-trip for every role.""" + filename = ( + artifact_filename("2026-07-18", role, symprec=1e-5, moyo_version="0.12.0") + if role is ArtifactRole.geo_opt_analysis + else artifact_filename("2026-07-18", role) + ) + assert filename == expected + assert parse_artifact_filename(f"/tmp/{filename}") is role + + +@pytest.mark.parametrize("invalid_date", ["2026-02-30", "2026/07/18"]) +def test_artifact_filename_rejects_invalid_dates(invalid_date: str) -> None: + """Artifact names require real ISO calendar dates.""" + with pytest.raises(ValueError, match="date"): + artifact_filename(invalid_date, ArtifactRole.discovery) + + +def test_prediction_alignment_rejects_unknown_ids_and_fills_missing() -> None: + """Alignment rejects extraneous IDs and inserts NaN for omitted references.""" + reference = _reference_frame() + reference_ids = reference[MATERIAL_ID].tolist() + aligned = align_predictions(reference, pd.Series([0.2], index=[reference_ids[1]])) + assert aligned.index.tolist() == reference_ids + assert aligned.iloc[1] == pytest.approx(0.2) + assert aligned.isna().sum() == 3 + + with pytest.raises(ValueError, match="unknown material IDs.*unknown"): + align_predictions(reference, pd.Series([0.2], index=["unknown"])) + + +def test_dataframe_evaluation_normalizes_material_ids_to_strings() -> None: + """Match numeric reference IDs with string prediction IDs.""" + reference = _reference_frame() + reference[MATERIAL_ID] = [1, 2, 3, 4] + predictions = _prediction_frame(["1", "2", "3", "4"], [-0.8, 0.2, -0.1, 0.9]) + + results = evaluate_discovery(reference, predictions) + + assert results["subsets"][str(DiscoverySubset.full_test_set)]["missing_preds"] == 0 + + +@pytest.mark.parametrize("ranking_case", ["ties", "missing"]) +def test_most_stable_10k_ranking(ranking_case: str) -> None: + """Stable sorting preserves ties and places missing predictions last.""" + material_ids = [f"wbm-{idx}" for idx in range(10_001)] + reference = _reference_frame( + material_ids=material_ids, + each_true=[0.0] * len(material_ids), + formation_energies=[0.0] * len(material_ids), + unique_prototypes=[True] * len(material_ids), + ) + if ranking_case == "ties": + prediction_values = [0.0] * len(material_ids) + else: + prediction_values = [ + *map(float, range(9_999)), + np.nan, + np.nan, + ] + predictions = pd.Series(prediction_values, index=material_ids) + + ranked_index = discovery_subset_indices(reference, predictions)[ + DiscoverySubset.most_stable_10k + ] + assert ranked_index.tolist() == material_ids[:10_000] + if ranking_case == "missing": + assert pd.isna(predictions.loc[ranked_index[-1]]) + + +def test_most_stable_ranking_uses_predicted_hull_distance() -> None: + """Ranking includes DFT hull and formation-energy offsets, not raw predictions.""" + material_ids = ["wbm-a", "wbm-b", "wbm-c", "wbm-d"] + reference = _reference_frame( + material_ids=material_ids, + each_true=[1.0, 0.0, 0.5, 0.2], + formation_energies=[0.0] * 4, + unique_prototypes=[True] * 4, + ) + predictions = pd.Series([-2.0, -0.1, -0.5, -0.2], index=material_ids) + + ranked_index = discovery_subset_indices(reference, predictions)[ + DiscoverySubset.most_stable_10k + ] + + assert ranked_index.tolist() == material_ids + assert not ranked_index.equals(predictions.sort_values().index) + + +@pytest.mark.parametrize( + ("formation_energies", "prediction_values", "expected_order"), + [ + ([10.0, 10.0, 10.0], [4.9, 5.0, 6.0], ["wbm-1", "wbm-2", "wbm-0"]), + ([0.0, 0.0, 0.0], [0.00049, 0.00041, 0.1], ["wbm-0", "wbm-1", "wbm-2"]), + ], + ids=["outlier-masked", "rounded-tie"], +) +def test_subset_ranking_applies_artifact_preprocessing( + formation_energies: list[float], + prediction_values: list[float], + expected_order: list[str], +) -> None: + """Mask outliers and round predictions before ranking subsets.""" + material_ids = ["wbm-0", "wbm-1", "wbm-2"] + reference = _reference_frame( + material_ids=material_ids, + each_true=[0.0] * 3, + formation_energies=formation_energies, + unique_prototypes=[True] * 3, + ) + predictions = pd.Series(prediction_values, index=material_ids) + + ranked_index = discovery_subset_indices( + reference=reference, predictions=predictions + )[DiscoverySubset.most_stable_10k] + + assert ranked_index.tolist() == expected_order + + +def test_subset_metrics_and_daf_override() -> None: + """Subset selection uses the given DAF prevalence.""" + reference = _reference_frame( + material_ids=[f"wbm-{idx}" for idx in range(6)], + each_true=[-0.2, -0.1, 0.1, 0.2, -0.05, 0.3], + formation_energies=[-1.0, -0.9, -0.8, -0.7, -0.6, -0.5], + unique_prototypes=[True, True, False, True, True, False], + ) + predictions = pd.Series( + [-1.1, -0.7, -0.9, -0.6, np.nan, -0.4], + index=reference[MATERIAL_ID], + ) + subset_indices = discovery_subset_indices(reference, predictions) + metrics = calc_discovery_metrics( + reference, + predictions, + subset_indices=subset_indices, + uniq_proto_prevalence=0.5, + ) + + assert set(metrics) == set(DiscoverySubset) + assert subset_indices[DiscoverySubset.most_stable_10k].tolist() == [ + "wbm-0", + "wbm-1", + "wbm-3", + "wbm-4", + ] + for subset in DAF_SUBSETS: + assert metrics[subset]["DAF"] == pytest.approx( + metrics[subset]["Precision"] / 0.5 + ) + full_metrics = metrics[DiscoverySubset.full_test_set] + full_prevalence = 3 / 6 + assert full_metrics["DAF"] == pytest.approx( + full_metrics["Precision"] / full_prevalence + ) + + +def test_calc_discovery_metrics_rejects_incomplete_subset_indices() -> None: + """Require all three subset indices when callers provide them.""" + reference = _reference_frame() + predictions = pd.Series([0.0] * len(reference), index=reference[MATERIAL_ID]) + + with pytest.raises(ValueError, match="Missing required discovery subsets"): + calc_discovery_metrics( + reference, + predictions, + subset_indices={DiscoverySubset.full_test_set: pd.Index([])}, + ) + + +def test_evaluation_masks_outliers_then_rounds_to_three_decimals() -> None: + """Evaluation masks errors above 5 eV and rounds inputs before metrics.""" + reference = _reference_frame( + material_ids=["wbm-a", "wbm-b", "wbm-c"], + each_true=[0.0004, -0.2, 0.2], + formation_energies=[-1.0, 0.0, 0.0], + unique_prototypes=[True, True, True], + ) + predictions = _prediction_frame( + ["wbm-a", "wbm-b", "wbm-c"], + [-0.9996, 5.0001, 5.0], + ) + + results = evaluate_discovery(reference, predictions) + full_metrics = results["subsets"][str(DiscoverySubset.full_test_set)] + assert set(full_metrics) == REPORTED_METRICS + assert { + metric_name: full_metrics[metric_name] + for metric_name in ("TP", "FN", "FP", "TN", "missing_preds") + } == {"TP": 1, "FN": 1, "FP": 0, "TN": 1, "missing_preds": 1} + assert {name: full_metrics[name] for name in ("MAE", "RMSE", "R2")} == { + "MAE": 2.5, + "RMSE": 3.536, + "R2": -1249.0, + } + calculated = calc_discovery_metrics(reference=reference, predictions=predictions)[ + DiscoverySubset.full_test_set + ] + assert {name: calculated[name] for name in ("TP", "FN", "FP", "TN", "MAE")} == { + "TP": 1, + "FN": 1, + "FP": 0, + "TN": 1, + "MAE": 2.5, + } + + +def test_synthetic_daf_uses_prepared_rounded_hull_labels() -> None: + """Synthetic-mode prevalence follows the documented rounded evaluation labels.""" + reference = _reference_frame( + material_ids=["wbm-a", "wbm-b"], + each_true=[0.0004, 0.1], + formation_energies=[0.0, 0.0], + unique_prototypes=[True, True], + ) + predictions = _prediction_frame(["wbm-a", "wbm-b"], [0.0004, 0.1]) + + results = evaluate_discovery(reference, predictions) + + assert results["subsets"][str(DiscoverySubset.unique_prototypes)]["DAF"] == 2.0 + + +def test_canonical_evaluation_requires_explicit_unrounded_prevalence() -> None: + """Leaderboard mode requires unrounded unique-prototype prevalence.""" + reference = _reference_frame() + predictions = _prediction_frame( + reference[MATERIAL_ID].tolist(), [-0.8, 0.2, -0.1, 0.9] + ) + with pytest.raises(ValueError, match="requires explicit unrounded"): + evaluate_discovery(reference, predictions, canonical=True) + + results = evaluate_discovery( + reference, + predictions, + canonical=True, + uniq_proto_prevalence=0.25, + ) + for subset in DAF_SUBSETS: + assert results["subsets"][str(subset)]["DAF"] == 4.0 + + +def test_path_evaluation_is_json_safe_and_writes_strict_json( + tmp_path: Path, +) -> None: + """Gzip path evaluation replaces NaNs by null and writes strict JSON.""" + reference = _reference_frame( + material_ids=["007", "008", "009", "010"], + each_true=[-0.4, -0.3, -0.2, -0.1], + ) + predictions = _prediction_frame( + reference[MATERIAL_ID].tolist(), [1.0, 1.0, 1.0, None] + ) + reference_path = tmp_path / "reference.csv.gz" + prediction_path = tmp_path / "predictions.csv.gz" + output_path = tmp_path / "metrics.json" + reference.to_csv(reference_path, index=False) + predictions.to_csv(prediction_path, index=False) + + results = evaluate_discovery_paths(reference_path, prediction_path) + assert results["schema_version"] == 1 + assert results["source"] == { + "framework": "matbench-discovery", + "version": "1.3.1", + } + full_metrics = results["subsets"][str(DiscoverySubset.full_test_set)] + assert full_metrics["Precision"] is None + assert full_metrics["missing_preds"] == 1 + json.dumps(results, allow_nan=False) + + write_discovery_metrics_json(results, output_path) + with open(output_path, encoding="utf-8") as file: + assert json.load(file) == results + + +def test_evaluators_treat_infinite_predictions_as_missing() -> None: + """Raw and JSON-safe evaluators sanitize infinite predictions.""" + reference = _reference_frame() + predictions = pd.Series([np.inf, -np.inf, 0.5, 1.0], index=reference[MATERIAL_ID]) + + results = evaluate_discovery(reference, predictions, max_error_threshold=None) + full_metrics = results["subsets"][str(DiscoverySubset.full_test_set)] + assert full_metrics["missing_preds"] == 2 + json.dumps(results, allow_nan=False) + + metrics = calc_discovery_metrics(reference, predictions) + assert metrics[DiscoverySubset.full_test_set]["FN"] == 2