ALIGNN test on WBM data

[pbenner]

Added ALIGNN test on WBM data, achieving a MAE of 0.0916. The code needs some double checking to make sure that matbench-discovery is correctly used. coGN would be the next step.

[janosh]

This is great, thank you @pbenner! 🎉

Sorry for the slow reply here. Was about to hit send this morning when the day suddenly turned super packed. Needless to say, I'm excited about this 1st external model submission to Matbench Discovery!

I'd thought about running ALIGNN myself but expected it to be similar to MEGNet and so in the end decided against it. But looks like I was wrong. Performance is quite a bit better than expected!

I'll refactor the code here a bit and then start updating all the site metrics. You didn't by any chance log the training and test runs for ALIGNN with wandb?

[pbenner]

Sorry no, had some minor trouble with wandb and excluded it, but should be easy to add to the existing code.

[janosh]

@pbenner Just out of curiosity, did you try the pre-trained mp_e_form_alignnn?

from alignn.pretrained import get_figshare_model

model = get_figshare_model("mp_e_form_alignnn")

If not, I'll run it. Would be good to know how it compares to the ALIGNN specifically trained for this benchmark.

Also, looks like we can avoid writing all structures to disk as POSCAR and then reading them in again as Jarvis Atoms. We can convert them directly using JarvisAtomsAdaptor (see fffd6d7).

from pymatgen.io.jarvis import JarvisAtomsAdaptor

[janosh]

@pbenner Could you briefly explain the need for alignn-2023.01.10.patch (formerly train_alignn_patch.txt)?

[pbenner]

I didn't know about the JarvisAtomsAdaptor, that's a good hint!

I tried the pretrained model and got a MAE of 0.1095, so slightly worse. However, I'm not sure if I have to apply some corrections to the energies, so training the model myself seemed safer. In general, predictive performance is quite sensitive to the number of training samples, so it seems best to me to always use the exact same training data instead of pretrained models (at least if the goal is to compare model architectures).

The ALIGNN patch has multiple purposes. First, it fixes some issues when there is no test set (see ALIGNN Issue #104). In this case, we actually forked a test set, but it might be better to use the entire data, as mentioned above, especially if the test set by chance contains some important outliers. Second, the Checkpoint() handler in ALIGNN does not define a score name (see ALIGNN train.py), so it will just save the last two models during training. With this patch, also the best model in terms of accuracy on the validation set is saved, which is the one I use for predictions. This is important, because I used a relatively large n_early_stopping in case the validation accuracy shows a double descent (see Figure 10). Maybe one has to consider here also how individual models are trained, since longer training might be beneficial.

[janosh]

I also ran the pre-trained ALIGNN before merging this PR just to check how it compares. The MAE doesn't tell the whole story. As a classifier, it significantly underperforms the ALIGNN trained from scratch for MBD.

Raw pre-trained ALIGNN (mp_e_form_alignn)

After subtracting the legacy corrections ALIGNN was trained on and adding the new MP 2020 corrections:

Slightly better classification performance but still bad.

---

Btw, you can find both the old and new MP corrections in matbench_discovery.data.df_wbm so updating model preds is just two lines: subtract e_correction_per_atom_mp_legacy and add e_correction_per_atom_mp2020.

Speaking of checkpoints, could you send me the checkpoint that was used to make the predictions in this PR? Would be good to provide that for reproducibility.

[pbenner]

Sure, the model is relatively small, so I just added it to the repo: Alignn best model

[CompRhys]

@pbenner do you know whether it's possible to run your alignn model in a uip-gnn setup where it first relaxes itself according to it's own predictions? In some of the metrics we see what we've been calling a characteristic fire-hook curve which we believe is a function of the task mismatch seen in the IS2RE task that is mostly resolved with uip-relaxed-IS2RE.

[janosh]

@CompRhys The original ALIGNN doesn't predict forces. They released a force field version of ALIGNN a year later in 2022 which looks similar to M3GNet. You could use the original ALIGNN to do repeated energy preds to get finite-difference forces but every time I've seen that it worked poorly. The ALIGNN FF would be interesting to add though.

[pbenner]

Just a small update. I am currently testing the ALIGNN-FF model for computing relaxed structures:
https://github.com/pbenner/matbench-discovery/blob/main/models/alignn/alignn_relax.py
It's quite slow, so I have to see how to get results within some reasonable time frame.

[janosh]

Nice! 🎉

For CHGnet and M3GNet, I used np.array_split() to split the unrelaxed structures into 100 subsets and then launched slurm jobs to relax each subset in parallel. Takes ~2 hours to get a complete set of predictions on A100s.

matbench-discovery/models/chgnet/test_chgnet.py

Lines 67 to 69 in 4a90dee

	df_in: pd.DataFrame = np.array_split(
	pd.read_json(data_path).set_index("material_id"), slurm_array_task_count
	)[slurm_array_task_id - 1]

[pbenner]

Thanks! It takes around 10 days on 4 A100 GPUs. Results should be done next week.

[janosh]

Sweet, looking forward to it!

[pbenner]

The relaxed predictions are finally done! You can find the results here:
https://github.com/pbenner/matbench-discovery/tree/alignn/models/alignn

I'm not filing a pull request, because my repository contains the full set of relaxed structures, which you might want to upload to figshare. I had a lot of fun computing the relaxed structures. Some of them converged very poorly, which is why this computation took so long. Also it seems that the prediction error went up. Unfortunately, there is also a bug in ALIGNN or ASE:

Traceback (most recent call last):
  File "alignn_relax.py", line 96, in <module>
  File "alignn_relax.py", line 88, in alignn_relax
  File "../alignn/ff/ff.py", line 310, in optimize_atoms
  File "../alignn/lib/python3.9/site-packages/ase/optimize/optimize.py", line 269, in run
  File "../alignn/lib/python3.9/site-packages/ase/optimize/optimize.py", line 156, in run
  File "../alignn/lib/python3.9/site-packages/ase/optimize/optimize.py", line 129, in irun
  File "../alignn/lib/python3.9/site-packages/ase/optimize/optimize.py", line 108, in call_observers
  File "../alignn/lib/python3.9/site-packages/ase/io/trajectory.py", line 132, in write
  File "../alignn/lib/python3.9/site-packages/ase/io/trajectory.py", line 156, in _write_atoms
  File "../alignn/lib/python3.9/site-packages/ase/io/trajectory.py", line 381, in write_atoms
  File "../alignn/lib/python3.9/site-packages/ase/io/ulm.py", line 400, in write
  File "../alignn/lib/python3.9/site-packages/ase/io/ulm.py", line 325, in fill
OSError: [Errno 24] Too many open files

which forced me to turn off any logging/trajectories. To debug we would have to recompute some trajectories, but the ones that converge poorly take several days to finish.

[janosh]

The relaxed predictions are finally done! You can find the results here:

Excellent. That's great news!

I'm not filing a pull request, because my repository contains the full set of relaxed structures, which you might want to upload to figshare.

Sounds good. I'll start prepping a PR with your code and results then?

the ones that converge poorly take several days to finish.

What?! That sounds problematic. You're saying some of the 100 test set batches (~2500 structures) took several days? If so, suggests lack of robustness to data variety in the model.