Add NequIP & Allegro to matbench leaderboard

[kavanase]

Description

This PR adds initial foundation potentials from the NequIP GNN and Allegro architectures, including both compliant and non-compliant (OMat) models.

Checklist

Please check the following items before submitting your PR:

• I created a new folder and YAML metadata file models/<arch_name>/<model_variant>.yml for my submission. arch_name is the name of the architecture and model_variant.yml includes things like author details, training set names and important hyperparameters.
• I added the new model as a new attribute Model.<arch_name> enum on the Model enum in enums.py.
• I uploaded the energy/force/stress model prediction file for the WBM test set to Figshare or another cloud storage service (<yyyy-mm-dd>-<model_variant>-preds.csv.gz).
• I uploaded the model-relaxed structures file to Figshare or another cloud storage service in JSON lines format (<yyyy-mm-dd>-wbm-IS2RE-FIRE.jsonl.gz). JSON Lines allows fast loading of small numbers of structures with pandas.read_json(lines=True, nrows=100) for inspection.
• I uploaded the phonon predictions to Figshare or another cloud storage service (<yyyy-mm-dd>-kappa-103-FIRE-<values-of-dist|fmax|symprec>.gz).
• I included the urls to the Figshare files in the YAML metadata file (models/<arch_name>/<model_variant>.yml). If not using Figshare I have included the urls to the cloud storage service in the description of the PR.
• I included the test script (test_<arch_name>_<task>.py for task in discovery, kappa, diatomics) that generated the prediction files.

Additional Information (Optional)

• I included a training script (train_<arch_name>.py) if I trained a model specifically for this benchmark.
• I included a readme.md with additional details about my model.

@coderabbitai ignore

[kavanase]

@janosh I think this is nearly ready to go. The remaining issues that seem to be causing the pre-commit test failures are the k_SRME values in the configs. When I try running scripts/evals/kappa_srme.py, I get this same error for all models, so not sure what's causing this? I'm using the main branch, with it's requirements (i.e. pymatviz 0.16.0)

❯ python scripts/evals/kappa_srme.py
Evaluating kappa metrics for 38 models...
Skipping AlchemBERT: no kappa_103_path found

Processing AlphaNet-v1-OMA...
	✗ Error processing AlphaNet-v1-OMA: "None of [<Key.mat_id: 'material_id'>] are in the columns"
Skipping ALIGNN: no kappa_103_path found

Processing Allegro-OAM-L...
	✗ Error processing Allegro-OAM-L: "None of [<Key.mat_id: 'material_id'>] are in the columns"

Processing Allegro-MP-L...
	✗ Error processing Allegro-MP-L: "None of [<Key.mat_id: 'material_id'>] are in the columns"
Skipping BOWSR: no kappa_103_path found

Processing CHGNet...
	✗ Error processing CHGNet: "None of [<Key.mat_id: 'material_id'>] are in the columns"
Skipping CGCNN: no kappa_103_path found
Skipping CGCNN+P: no kappa_103_path found

Then also for the geo-opt analysis script:

❯ python scripts/evals/geo_opt.py
Traceback (most recent call last):
  File "/Users/kavanase/Packages/matbench-discovery/scripts/evals/geo_opt.py", line 40, in <module>
    df_dft_analysis = pd.read_csv(dft_analysis_file, index_col=0)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pandas/io/parsers/readers.py", line 1026, in read_csv
    return _read(filepath_or_buffer, kwds)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pandas/io/parsers/readers.py", line 620, in _read
    parser = TextFileReader(filepath_or_buffer, **kwds)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pandas/io/parsers/readers.py", line 1620, in __init__
    self._engine = self._make_engine(f, self.engine)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pandas/io/parsers/readers.py", line 1880, in _make_engine
    self.handles = get_handle(
                   ^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pandas/io/common.py", line 765, in get_handle
    handle = gzip.GzipFile(  # type: ignore[assignment]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kavanase/miniconda3/envs/py312/lib/python3.12/gzip.py", line 192, in __init__
    fileobj = self.myfileobj = builtins.open(filename, mode or 'rb')
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: '/Users/kavanase/Packages/matbench-discovery/data/wbm/dft-geo-opt-symprec=1e-5-moyo=0.4.4.csv.gz'

The discovery analysis script works fine though.

[kavanase]

Trying to use the compile-WBM script, I get:

❯ python data/wbm/compile_wbm_test_set.py
/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pymatgen/core/composition.py:1365: UserWarning: No Pauling electronegativity for Ne. Setting to NaN. This has no physical meaning, and is mainly done to avoid errors caused by the code expecting a float.
  syms: list[str] = sorted(sym_amt, key=lambda x: [get_el_sp(x).X, x])
/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pymatgen/core/composition.py:1365: UserWarning: No Pauling electronegativity for He. Setting to NaN. This has no physical meaning, and is mainly done to avoid errors caused by the code expecting a float.
  syms: list[str] = sorted(sym_amt, key=lambda x: [get_el_sp(x).X, x])
/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pymatgen/core/composition.py:1365: UserWarning: No Pauling electronegativity for Ar. Setting to NaN. This has no physical meaning, and is mainly done to avoid errors caused by the code expecting a float.
  syms: list[str] = sorted(sym_amt, key=lambda x: [get_el_sp(x).X, x])
Traceback (most recent call last):
  File "/Users/kavanase/Packages/matbench-discovery/data/wbm/compile_wbm_test_set.py", line 33, in <module>
    from matbench_discovery.structure import prototype
  File "/Users/kavanase/Packages/matbench-discovery/matbench_discovery/structure/prototype.py", line 16, in <module>
    from pymatviz.typing import AnyStructure
ImportError: cannot import name 'AnyStructure' from 'pymatviz.typing' (/Users/kavanase/miniconda3/envs/py312/lib/python3.12/site-packages/pymatviz/typing.py)

(with pymatviz==0.16.0)

[janosh]

very nice work! 👍 thanks a lot for this submission @kavanase! 🙏 i'm very excited to add NequiP and Allegro to the benchmark!

thanks for reporting the data/wbm/compile_wbm_test_set.py import error. should be fixed in f3419a2

re the failing scripts/evals/kappa_srme.py, I haven't taken a closer look yet but at first glance seems like the column names are unexpected (e.g. mp_id instead of material_id). i'll take a closer look. no need for you to debug

[janosh]

also, apologies about all the coderabbit spam. looks like it went absolutely crazy on this PR. i wonder if there's a gh CLI command to delete all bot comments on a PR

[janosh]

this script seems to mostly bring the kappa pred files into the expected format (except there's still an extra nesting level that doesn't play well with pd.read_json). need to change trains now but will have another go later.

standardize_kappa_preds.py

"""Convert kappa predictions to Matbench format."""

import json
import os
import shutil

import pandas as pd
from pymatviz.enums import Key

from matbench_discovery import today
from matbench_discovery.enums import MbdKey, Model

DRY_RUN = False
col_map = {
    "mp_id": Key.mat_id,
    "name": Key.formula,
    "max_stress": Key.max_stress,
    "reached_max_steps": "reached_max_steps",
    "imaginary_freqs": Key.has_imag_ph_modes,
    "frequencies": Key.ph_freqs,
    "kappa_TOT_RTA": MbdKey.kappa_tot_rta,
    "kappa_P_RTA": MbdKey.kappa_p_rta,
    "kappa_C": MbdKey.kappa_c,
    "weights": Key.mode_weights,
    "qpoints": Key.q_points,
    "mode_kappa_TOT": MbdKey.mode_kappa_tot_rta,
    "kappa_TOT_ave": MbdKey.kappa_tot_avg,
    "initial_space_group_number": Key.spg_num,
    "mode_kappa_TOT_ave": MbdKey.mode_kappa_tot_avg,
}
cols_to_drop = [
    "DFT_kappa_TOT_ave",
    "SRD",
    "SRE",
    "SRME",
    "errors",
    "error_traceback",
    "symprec_tests",
    "redirected_to_symm",
]

models_updated: dict[str, dict[str, str | list[str]]] = {}


for model in Model:
    if model.kappa_103_path is None or not os.path.isfile(model.kappa_103_path):
        print(f"{model.label=} kappa_103_path={model.kappa_103_path!r} does not exist")
        continue

    df_ml = pd.read_json(model.kappa_103_path)
    if MbdKey.mode_kappa_tot_avg in df_ml:
        print(f"{model.label=} already processed")
        continue  # skip if already processed

    old_cols = df_ml.columns.copy()
    df_ml = df_ml.rename(columns=col_map)

    renamed_cols = {
        str(old_col): str(new_col)
        for old_col, new_col in zip(old_cols, df_ml)
        if old_col != new_col
    }
    old_cols = df_ml.columns.copy()
    df_ml = df_ml.drop(columns=cols_to_drop, errors="ignore")
    dropped_cols = [col for col in old_cols if col not in df_ml]

    # if dataframe was modified, backup the original file
    if renamed_cols or dropped_cols:
        models_updated[model.key] = {
            "renamed_cols": renamed_cols,
            "dropped_cols": dropped_cols,
        }
        if not DRY_RUN:
            shutil.copy(model.kappa_103_path, f"{model.kappa_103_path}.bak")
            df_ml.to_json(model.kappa_103_path)
            print(f"{model.label=} new column names: {df_ml.columns}")
    else:
        print(f"{model.label=} no changes")


print(f"\n\n{len(models_updated)=}")

for model_key, actions in models_updated.items():
    print(f"{model_key=} {actions=}")

# write updated models to file
with open(f"{today}-kappa-model-files-updated.json", mode="w") as file:
    json.dump(models_updated, file, indent=2)

[kavanase]

Thanks for this @janosh! 🙏

[samblau]

Excited to see this merged soon!

[kavanase]

@kavanase very nearly ready to merge. there are some TODOs left

matbench-discovery/models/allegro/allegro-MP-L-0.1.yml

Line 56 in 9aea794

checkpoint_url: https://url.of/model-checkpoint-license # TODO: Zenodo? / Nequip.net

can we address those now that the checkpoints are up https://www.nequip.net/models/mir-group/Allegro-MP-L:0.1?

Fixed now! I missed this entry for the Allegro models on the first draft 🤦

[kavanase]

@janosh thanks very much for pushing through with this!
I see the Tests / scripts (scripts/evals/discovery.py) (pull_request) appears as passing, but looking at the outputs it's showing download errors for the updated NequIP/Allegro figshare files?

[janosh]

thanks again for this submission @kavanase! outstanding work 👍 🚀

been hoping for the day when Allegro and Nequip are added to the leaderboard for almost 2 years now! so many people kept asking me why those we're missing and all I could say is lack of time...
very glad that's now over 🙏

but looking at the outputs it's showing download errors for the updated NequIP/Allegro figshare files?

i think i just forgot to re-publish the Figshare article after uploading these files. should be good now

[kavanase]

Great to see this merged, thanks very much @janosh !!