rework model card from foundation.md (model tables, downloads, fine-tuning) + ASL license

README.md

@@ -11,11 +11,14 @@ tags:
   - materials-science
 ---
 
-# GRACE foundation models
+# Pretrained GRACE Foundation Models
 
-This repository distributes **GRACE** (Graph Atomic Cluster Expansion) machine-learning
-interatomic potentials — foundation models fitted with
-[GRACEmaker / tensorpotential](https://github.com/ICAMS/grace-tensorpotential).
+This repository distributes pretrained **GRACE** (Graph Atomic Cluster Expansion) machine-learning
+interatomic potentials, fitted with [GRACEmaker / tensorpotential](https://github.com/ICAMS/grace-tensorpotential).
+The model tables and usage notes below are adapted from the
+[GRACEmaker documentation](https://github.com/ICAMS/grace-tensorpotential/blob/master/docs/gracemaker/foundation.md).
+
+> **Use the "Full Name" column** to refer to a model in LAMMPS and ASE.
 
 ## License
 
@@ -23,21 +26,117 @@ interatomic potentials — foundation models fitted with
 License (ASL).** By downloading or using any model from this repository you agree to the terms
 of the ASL: <https://github.com/ICAMS/grace-tensorpotential/blob/master/LICENSE.md>
 
-## Contents
+---
+
+## SMAX models
+
+Reference: [arXiv](https://arxiv.org/abs/2602.23489)
+
+The **SMAX** (**Maximum Entropy**) models are trained on a chemistry-agnostic dataset generated via a multicomponent
+maximum information entropy structure generation protocol.
+
+Unlike traditional datasets that focus on low-energy equilibrium structures, SMAX is constructed to deliberately sample
+broad and diverse regions of configurational space. This provides a robust physical prior for atomic interactions across the entire periodic table, enabling accurate modeling of large-strain phase transformations, defects in complex alloys, and reaction barriers in catalytic systems.
+
+**Custom Cutoffs & Interaction Ranges:** SMAX models utilize a **custom element-dependent cutoff radius** ranging from
+**5.0 Å to 7.5 Å**.
+
+**Recommendation:** For most general-purpose applications, we recommend using the **SMAX-OMAT** models.
+They offer the best balance of structural robustness (from SMAX) and high-precision energy/force accuracy (from OMat24).
+
+#### Single-layer, local models
+
+| Model Name | Full Name | Size | κ_SRME | Description |
+| :---- | :---- | :---- | :---- | :---- |
+| GRACE-1L-SMAX-L | GRACE-1L-SMAX-large | large | 0.696 | Single-layer local (SMAX) |
+| **GRACE-1L-SMAX-OMAT-L** | **GRACE-1L-SMAX-OMAT-large** | **large** | 0.338 | **Single-layer local (SMAX + OMat24)** |
+
+#### Two-layer, semilocal models
+
+| Model Name | Full Name | Size | κ_SRME | Description |
+| :---- | :---- | :---- | :---- | :---- |
+| GRACE-2L-SMAX-M | GRACE-2L-SMAX-medium | medium | 0.469 | Two-layer semi-local (SMAX) |
+| GRACE-2L-SMAX-L | GRACE-2L-SMAX-large | large | 0.444 | Two-layer semi-local (SMAX) |
+| **GRACE-2L-SMAX-OMAT-M** | **GRACE-2L-SMAX-OMAT-medium** | **medium** | 0.197 | **Two-layer semi-local (SMAX + OMat24)** |
+| **GRACE-2L-SMAX-OMAT-L** | **GRACE-2L-SMAX-OMAT-large** | **large** | 0.191 | **Two-layer semi-local (SMAX + OMat24)** |
+
+## OMAT models
+
+Reference: [npj Comp. Mat.](https://www.nature.com/articles/s41524-026-01979-1), [arXiv](https://arxiv.org/abs/2508.17936)
+
+The base models (**-OMAT**) are trained on the [OMat24](https://huggingface.co/datasets/fairchem/OMAT24#omat24-dataset) dataset.
+The fine-tuned versions (**-OMAT-ft-E**) are derived from these base models by fine-tuning with more emphasis on energies.
+All models listed use a fixed **6 Å cutoff**.
+
+#### Single-layer, local models
+
+| Model Name | Full Name | Size | κ_SRME | Description |
+|:---|:---|:---|:---|:---|
+| **GRACE-1L-OMAT** | GRACE-1L-OMAT | small | 0.398 | Single-layer local |
+| GRACE-1L-OMAT-M-base | GRACE-1L-OMAT-medium-base | medium | 0.380 | Single-layer local (base) |
+| **GRACE-1L-OMAT-M** | GRACE-1L-OMAT-medium-ft-E | medium | 0.417 | Single-layer local (finetuned on energy) |
+| GRACE-1L-OMAT-L-base | GRACE-1L-OMAT-large-base | large | **0.354** | Single-layer local (base) |
+| **GRACE-1L-OMAT-L** | GRACE-1L-OMAT-large-ft-E | large | 0.383 | Single-layer local (finetuned on energy) |
+
+#### Two-layer, semilocal models
+
+| Model Name | Full Name | Size | κ_SRME | Description |
+|:---|:---|:---|:---|:---|
+| **GRACE-2L-OMAT** | GRACE-2L-OMAT | small | 0.288 | Two-layer semi-local |
+| GRACE-2L-OMAT-M-base | GRACE-2L-OMAT-medium-base | medium | 0.212 | Two-layer semi-local (base) |
+| **GRACE-2L-OMAT-M** | GRACE-2L-OMAT-medium-ft-E | medium | 0.217 | Two-layer semi-local (finetuned on energy) |
+| GRACE-2L-OMAT-L-base | GRACE-2L-OMAT-large-base | large | **0.165** | Two-layer semi-local (base) |
+| **GRACE-2L-OMAT-L** | GRACE-2L-OMAT-large-ft-E | large | 0.186 | Two-layer semi-local (finetuned on energy) |
+
+## OAM models
+
+Reference: [npj Comp. Mat.](https://www.nature.com/articles/s41524-026-01979-1), [arXiv](https://arxiv.org/abs/2508.17936)
+
+These models are first pre-trained on **OMat24** and then fine-tuned on a combination of the [sAlex](https://huggingface.co/datasets/fairchem/OMAT24#salex-dataset) dataset (10.4M structures) and the [MPtraj](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842?file=41619375) dataset (1.58M structures).
+
+#### Single-layer, local models
+
+| Model Name | Full Name | Size | F1 | κ_SRME | Description |
+| :--- |:---| :--- |:---|:---| :--- |
+| GRACE-1L-OAM | GRACE-1L-OAM | small | 0.824 | 0.516 | Single-layer local |
+| GRACE-1L-OAM-M | GRACE-1L-OMAT-medium-ft-AM | medium | 0.800 | 0.411 | Single-layer local |
+| GRACE-1L-OAM-L | GRACE-1L-OMAT-large-ft-AM | large | 0.815 | **0.377** | Single-layer local |
+
+#### Two-layer, semilocal models
+
+| Model Name | Full Name | Size | F1 | κ_SRME | Description |
+| :--- |:--- | :--- |:---| :--- | :--- |
+| GRACE-2L-OAM | GRACE-2L-OAM | small | 0.880 | 0.294 | Two-layer semi-local |
+| GRACE-2L-OAM-M | GRACE-2L-OMAT-medium-ft-AM | medium | 0.881 | 0.200 | Two-layer semi-local |
+| GRACE-2L-OAM-L | GRACE-2L-OMAT-large-ft-AM | large | 0.889 | **0.168** | Two-layer semi-local |
+
+---
+
+## Repository layout
 
 Each model is provided as two `.tar.gz` archives:
 
-- `models/<model>-model.tar.gz` — the deployable TensorFlow `SavedModel` directory
+- `models/<full-name>-model.tar.gz` — the deployable TensorFlow `SavedModel` directory
   (used by the ASE `TPCalculator`). UQ-enabled models embed an uncertainty head, and many
   archives also bundle `kokkos.npz` for LAMMPS-Kokkos deployment.
-- `checkpoints/<model>-checkpoint.tar.gz` — the training checkpoint
-  (`model.yaml`, `checkpoint.index`, `checkpoint.data-*`, and `gmm_artifacts.npz` for UQ models).
+- `checkpoints/<full-name>-checkpoint.tar.gz` — the training checkpoint required for
+  fine-tuning (`model.yaml`, `checkpoint.index`, `checkpoint.data-*`, plus `gmm_artifacts.npz`
+  for UQ models).
+
+## Downloading Foundation Models
+
+The recommended way is the `grace_models` utility shipped with `tensorpotential`:
+
+- `grace_models list` — view the list of available models.
+- `grace_models download <model_name>` — download a specific model (SavedModel).
+- `grace_models checkpoint <model_name>` — download a model's training checkpoint.
 
-## Usage
+By default all foundation models are stored in `$HOME/.cache/grace`; override with the
+`GRACE_CACHE` environment variable. The downloaded models can be used for simulations in
+**ASE** and **LAMMPS** — see the
+[GRACEmaker quickstart](https://github.com/ICAMS/grace-tensorpotential/blob/master/docs/quickstart.md).
 
-These models are normally fetched automatically by the `tensorpotential` / GRACEmaker
-tooling (`grace_models`, `TPCalculator`, foundation-model presets). They can also be
-downloaded directly, e.g.:
+You can also fetch archives directly from this repo:
 
 ```python
 from huggingface_hub import hf_hub_download
@@ -48,7 +147,7 @@ path = hf_hub_download(
 )
 ```
 
-## Versioning
+### Versioning
 
 Revisions are pinned via immutable git tags so downloads are reproducible:
 
@@ -56,3 +155,67 @@ Revisions are pinned via immutable git tags so downloads are reproducible:
 - **`kk`** — model archives that additionally bundle the LAMMPS-Kokkos export (`kokkos.npz`).
 
 `main` always points at the latest revision.
+
+---
+
+## Fine-tuning foundation models
+
+Fine-tuning is performed from **checkpoints** (not SavedModels). Run `grace_models list`
+to see which models expose a `CHECKPOINT:` field; download with
+`grace_models checkpoint <MODEL-NAME>` (or it is fetched automatically when needed).
+
+Generate a fine-tuning `input.yaml` interactively with `gracemaker -t`, or add to the
+`potential` section manually:
+
+```yaml
+potential:
+  finetune_foundation_model: GRACE-1L-OAM
+  shift: auto  # automatically align FM energies with your DFT reference data
+# reduce_elements: True # if True - select from original models only elements present in the CURRENT dataset
+```
+
+### Automatic energy shift correction (`shift: auto`)
+
+When `shift: auto` is set alongside `finetune_foundation_model`, GRACEmaker computes optimal
+per-element energy shifts that minimize the difference between the foundation model's
+predictions and your reference DFT data, injecting them via `ConstantScaleShiftTarget` before
+training. This is useful when your DFT settings differ from the foundation model's training
+data (systematic energy offset) and leads to faster convergence.
+
+`shift: auto` is independent of `data::reference_energy`. Typical usage:
+
+* `reference_energy: 0` (or dict) — set your DFT reference frame
+* `shift: auto` — align the FM output to your DFT reference frame
+
+### Learning rate & initial metrics
+
+Use a small `learning_rate` and evaluate initial metrics:
+
+```yaml
+fit:
+  opt_params: {learning_rate: 0.001, ... }
+  eval_init_stats: True
+```
+
+### Frozen-weights fine-tuning
+
+To mitigate catastrophic forgetting, restrict training to selected variables via
+`trainable_variable_names`:
+
+```yaml
+# GRACE-2L models
+fit:
+  trainable_variable_names: ["I2/reducing_", "rho/reducing_", "I1/reducing_"]
+```
+
+```yaml
+# GRACE-1L models
+fit:
+  trainable_variable_names: ["rho/reducing_"]
+```
+
+Discover variable names with:
+
+```bash
+grace_utils -p ~/.cache/grace/checkpoints/SOME_FOUNDATIONAL_MODEL_NAME/model.yaml summary -v 1
+```