diff --git a/.gitignore b/.gitignore index e2f21aae..d78d6a60 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ sevenn/pretrained_potentials/SevenNet_omat sevenn/pretrained_potentials/SevenNet_omni sevenn/pretrained_potentials/SevenNet_omni_i8 sevenn/pretrained_potentials/SevenNet_omni_i12 +sevenn/pretrained_potentials/SevenNet_nano # from compile sevenn/pair_d3* diff --git a/CHANGELOG.md b/CHANGELOG.md index 14fabea4..73a24756 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ All notable changes to this project will be documented in this file. - Training features used for SevenNet-Omni: batch training, `OrderedSampler`, `grad_clip`, `onecyclelr` - Loss: MAE, L2MAE - Dataset type: `aselmdb`, `custom` +- SevenNet-Nano checkpoints for `7net-nano-4.5`, `7net-nano-5.0`, `7net-nano-5.5`, and `7net-nano-6.0` + +### Changed +- The default dtype of the rescale layer is changed to double precision. Use the `SEVENN_SHIFT_SCALE_DTYPE` environment variable for backward single-precision compatibility. +- The recommended LAMMPS ML-IAP version is bumped to `stable_22Jul2025_update4`. ### Fixed - `D3Calculator()` segfault bug when reusing the calculator within different sized `Atoms`. diff --git a/README.md b/README.md index 420f67ea..fedb87ec 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,17 @@ If you utilize the pretrained model SevenNet-Omni or multi-task training strateg } ``` +If you utilize the pretrained model SevenNet-Nano, please cite the following paper: +```bib +@article{oh_lightweight_2026, + title = {A Lightweight Universal Machine-Learning Interatomic Potential via Knowledge Distillation for Scalable Atomistic Simulations}, + doi = {10.1021/acs.jcim.6c01103}, + journal = {J. Chem. Inf. Model.}, + author = {Oh, Sangmin and You, Jinmu and Kim, Jaesun and Lee, Jiho and An, Hyungmin and Han, Seungwu and Kang, Youngho}, + year = {2026}, +} +``` + If you utilize the reEWC forgetting-aware fine-tuning strategy for continual learning of pretrained universal machine-learning interatomic potentials, please cite the following paper: ```bib @article{kim_efficient_2026, diff --git a/docs/source/cite.md b/docs/source/cite.md index 5ed6cb2b..0e84e84d 100644 --- a/docs/source/cite.md +++ b/docs/source/cite.md @@ -40,3 +40,14 @@ If you utilize the pretrained model SevenNet-Omni or multi-task training strateg year = {2026}, } ``` + +If you utilize the pretrained model SevenNet-Nano, please cite the following paper: +```bib +@article{oh_lightweight_2026, + title = {A Lightweight Universal Machine-Learning Interatomic Potential via Knowledge Distillation for Scalable Atomistic Simulations}, + doi = {10.1021/acs.jcim.6c01103}, + journal = {J. Chem. Inf. Model.}, + author = {Oh, Sangmin and You, Jinmu and Kim, Jaesun and Lee, Jiho and An, Hyungmin and Han, Seungwu and Kang, Youngho}, + year = {2026}, +} +``` diff --git a/docs/source/user_guide/ase_calculator.md b/docs/source/user_guide/ase_calculator.md index b55212cd..eea248cc 100644 --- a/docs/source/user_guide/ase_calculator.md +++ b/docs/source/user_guide/ase_calculator.md @@ -69,7 +69,7 @@ calc = SevenNetCalculator(model='7net-0', enable_cueq=True) # or enable_flash=Tr If you encounter the error `CUDA is not installed or nvcc is not available`, please ensure the `nvcc` compiler is available. Currently, CPU + D3 is not supported. -Various pretrained SevenNet models can be accessed by setting the model variable to predefined keywords like `7net-omni`, `7net-mf-ompa`, `7net-omat`, `7net-l3i5`, and `7net-0`. +Various pretrained SevenNet models can be accessed by setting the model variable to predefined keywords like `7net-omni`, `7net-nano-5.5`, `7net-mf-ompa`, `7net-omat`, `7net-l3i5`, and `7net-0`. User-trained models can be applied with the ASE calculator. In this case, the `model` parameter should be set to the checkpoint path from training. diff --git a/docs/source/user_guide/pretrained.md b/docs/source/user_guide/pretrained.md index 6d7f9d52..29818662 100644 --- a/docs/source/user_guide/pretrained.md +++ b/docs/source/user_guide/pretrained.md @@ -42,6 +42,12 @@ Multiple inference tasks are available for multi-fidelity architecture models, S - Multi-task - `mpa` (PBE+U)
`matpes_r2scan` (r²SCAN)
`omol25_low` (ωB97M-V)
and 10 more tasks +* - [SevenNet-Nano](#sevennet-nano) + - Lightweight model distilled from the SevenNet-Omni `mpa` task + - $l_{\mathrm{max}}=2$
$N_{\mathrm{layer}}=3$
parity=full
cutoff=4.5, 5.0, 5.5, or 6.0 Å + - Single-task + - N/A + * - [SevenNet-MF-ompa](#sevennet-mf-ompa) - MPtrj, sAlex, OMat24 - $l_{\mathrm{max}}=3$
$N_{\mathrm{layer}}=5$
parity=full @@ -146,6 +152,29 @@ SevenNet-Omni-i12 follows the same training strategy as [SevenNet-Omni](#sevenne |:---:|:---:|:---:|:---:| |**0.873**|0.906|0.192|0.0617| +--- +## SevenNet-Nano +> Model keywords: `7net-nano-4.5` | `7net-nano-5.0` | `7net-nano-5.5` | `7net-nano-6.0` + +SevenNet-Nano is a lightweight pretrained model released with [A Lightweight Universal Machine-Learning Interatomic Potential via Knowledge Distillation for Scalable Atomistic Simulations](https://doi.org/10.1021/acs.jcim.6c01103). It is provided as four single-task checkpoint variants with different radial cutoffs. Distilled from the SevenNet-Omni `mpa` task, SevenNet-Nano compresses broad cross-domain chemical knowledge into a 105k-parameter model, achieving over an order-of-magnitude speedup over its teacher while retaining balanced accuracy across diverse atomistic environments. + +Choose the cutoff explicitly through the model keyword: + +```python +from sevenn.calculator import SevenNetCalculator +calc = SevenNetCalculator('7net-nano-5.5') +``` + +```{table} SevenNet-Nano cutoff variants +:widths: 40 60 +| Model keyword | Cutoff | +|---------------|--------| +| `7net-nano-4.5` | 4.5 Å | +| `7net-nano-5.0` | 5.0 Å | +| `7net-nano-5.5` | 5.5 Å | +| `7net-nano-6.0` | 6.0 Å | +``` + --- ## SevenNet-MF-ompa > Model keywords: `7net-mf-ompa` | `SevenNet-mf-ompa` diff --git a/sevenn/_const.py b/sevenn/_const.py index 6c78f9c8..3d23c80b 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -66,6 +66,10 @@ SEVENNET_omni = f'{_prefix}/SevenNet_omni/checkpoint_sevennet_omni.pth' SEVENNET_omni_i8 = f'{_prefix}/SevenNet_omni_i8/checkpoint_sevennet_omni_i8.pth' SEVENNET_omni_i12 = f'{_prefix}/SevenNet_omni_i12/checkpoint_sevennet_omni_i12.pth' +SEVENNET_nano_4_5 = f'{_prefix}/SevenNet_nano/checkpoint_7net_nano_4.5.pth' +SEVENNET_nano_5_0 = f'{_prefix}/SevenNet_nano/checkpoint_7net_nano_5.0.pth' +SEVENNET_nano_5_5 = f'{_prefix}/SevenNet_nano/checkpoint_7net_nano_5.5.pth' +SEVENNET_nano_6_0 = f'{_prefix}/SevenNet_nano/checkpoint_7net_nano_6.0.pth' _git_prefix = 'https://github.com/MDIL-SNU/SevenNet/releases/download' CHECKPOINT_DOWNLOAD_LINKS = { @@ -74,6 +78,10 @@ SEVENNET_omni: f'{_git_prefix}/v0.12.0.cp/checkpoint_sevennet_omni.pth', SEVENNET_omni_i8: f'{_git_prefix}/v0.12.1.cp/checkpoint_sevennet_omni_i8.pth', SEVENNET_omni_i12: f'{_git_prefix}/v0.12.1.cp/checkpoint_sevennet_omni_i12.pth', + SEVENNET_nano_4_5: f'{_git_prefix}/v0.13.1.cp/checkpoint_7net_nano_4.5.pth', + SEVENNET_nano_5_0: f'{_git_prefix}/v0.13.1.cp/checkpoint_7net_nano_5.0.pth', + SEVENNET_nano_5_5: f'{_git_prefix}/v0.13.1.cp/checkpoint_7net_nano_5.5.pth', + SEVENNET_nano_6_0: f'{_git_prefix}/v0.13.1.cp/checkpoint_7net_nano_6.0.pth', } # to avoid torch script to compile torch_geometry.data AtomGraphDataType = Dict[str, torch.Tensor] # But it can contain 'lmp_data' diff --git a/sevenn/calculator.py b/sevenn/calculator.py index a1256b5f..a3d85da5 100644 --- a/sevenn/calculator.py +++ b/sevenn/calculator.py @@ -44,8 +44,9 @@ def __init__( Parameters ---------- model: str | Path | AtomGraphSequential, default='7net-0' - Name of pretrained models (7net-omni, 7net-mf-ompa, 7net-omat, 7net-l3i5, - 7net-0) or path to the checkpoint, deployed model or the model itself + Name of pretrained models (7net-omni, 7net-nano-5.5, 7net-mf-ompa, + 7net-omat, 7net-l3i5, 7net-0) or path to the checkpoint, deployed model + or the model itself file_type: str, default='checkpoint' one of 'checkpoint' | 'model_instance' device: str | torch.device, default='auto' @@ -255,8 +256,9 @@ def __init__( Parameters ---------- model: str | Path | AtomGraphSequential - Name of pretrained models (7net-omni, 7net-mf-ompa, 7net-omat, 7net-l3i5, - 7net-0) or path to the checkpoint, deployed model or the model itself + Name of pretrained models (7net-omni, 7net-nano-5.5, 7net-mf-ompa, + 7net-omat, 7net-l3i5, 7net-0) or path to the checkpoint, deployed model + or the model itself file_type: str, default='checkpoint' one of 'checkpoint' | 'model_instance' device: str | torch.device, default='auto' diff --git a/sevenn/main/sevenn_get_model.py b/sevenn/main/sevenn_get_model.py index 8f7bcdb6..baea7bf8 100644 --- a/sevenn/main/sevenn_get_model.py +++ b/sevenn/main/sevenn_get_model.py @@ -7,8 +7,8 @@ 'deploy LAMMPS model from the checkpoint' ) checkpoint_help = ( - 'Pretrained model name (7net-omni, 7net-omni-i8, 7net-omni-i12, etc.) ' - 'or path to checkpoint file. See documentation for all available models.' + 'Pretrained model name (7net-omni, 7net-nano-5.5, etc.) or path to ' + 'checkpoint file. See documentation for all available models.' ) output_name_help = 'filename prefix' get_parallel_help = 'deploy parallel model' diff --git a/sevenn/util.py b/sevenn/util.py index 89d54065..502b3c09 100644 --- a/sevenn/util.py +++ b/sevenn/util.py @@ -289,6 +289,14 @@ def pretrained_name_to_path(name: str) -> str: checkpoint_path = _const.SEVENNET_omni_i8 elif name in [f'{n}-omni-i12' for n in heads]: checkpoint_path = _const.SEVENNET_omni_i12 + elif name in [f'{n}-nano-4.5' for n in heads]: + checkpoint_path = _const.SEVENNET_nano_4_5 + elif name in [f'{n}-nano-5.0' for n in heads]: + checkpoint_path = _const.SEVENNET_nano_5_0 + elif name in [f'{n}-nano-5.5' for n in heads]: + checkpoint_path = _const.SEVENNET_nano_5_5 + elif name in [f'{n}-nano-6.0' for n in heads]: + checkpoint_path = _const.SEVENNET_nano_6_0 else: raise ValueError('Not a valid pretrained model name') url = _const.CHECKPOINT_DOWNLOAD_LINKS.get(checkpoint_path) @@ -331,6 +339,10 @@ def get_available_pretrained_models() -> List[str]: 'SEVENNET_omni': '7net-omni', 'SEVENNET_omni_i8': '7net-omni-i8', 'SEVENNET_omni_i12': '7net-omni-i12', + 'SEVENNET_nano_4_5': '7net-nano-4.5', + 'SEVENNET_nano_5_0': '7net-nano-5.0', + 'SEVENNET_nano_5_5': '7net-nano-5.5', + 'SEVENNET_nano_6_0': '7net-nano-6.0', } models = [] diff --git a/tests/unit_tests/test_pretrained.py b/tests/unit_tests/test_pretrained.py index 7398fe2b..e5e597a0 100644 --- a/tests/unit_tests/test_pretrained.py +++ b/tests/unit_tests/test_pretrained.py @@ -373,14 +373,14 @@ def test_7net_omni_mpa(atoms_pbc, atoms_mol): ) g1_ref_s = -1 * torch.tensor( # xx, yy, zz, xy, yz, zx - [-0.6500675, -0.0290563, -0.0290563 , 0.02576996, 0.00374571, 0.02576996] + [-0.6500675, -0.0290563, -0.0290563, 0.02576996, 0.00374571, 0.02576996] ) g2_ref_e = torch.tensor([-12.918253898620605]) g2_ref_f = torch.tensor( [ [0.0, -13.32638, 7.1434574], - [0.0, 9.442289 , -9.77207], + [0.0, 9.442289, -9.77207], [0.0, 3.8840904, 2.6286132], ] ) @@ -486,3 +486,137 @@ def test_7net_omni_i12_mpa(atoms_pbc, atoms_mol): assert acl(g2.inferred_total_energy, g2_ref_e) assert acl(g2.inferred_force, g2_ref_f) + + +@pytest.mark.parametrize( + 'name, cutoff_ref, g1_ref_e, g1_ref_f, g1_ref_s, g2_ref_e, g2_ref_f', + [ + ( + '7net-nano-4.5', + 4.5, + -3.5564712945510024, + [ + [11.843839645385742, -0.06617936491966248, -0.06617937982082367], + [-11.843839645385742, 0.06617936491966248, 0.0661793053150177], + ], + [ + 0.605344831943512, + 0.02144481986761093, + 0.021444812417030334, + -0.03438035026192665, + -0.003832500660791993, + -0.03438035771250725, + ], + -12.919640449266296, + [ + [0.0, -13.341249465942383, 7.200963020324707], + [0.0, 9.452506065368652, -9.814258575439453], + [0.0, 3.8887438774108887, 2.6132965087890625], + ], + ), + ( + '7net-nano-5.0', + 5.0, + -3.4316864360099615, + [ + [12.497757911682129, -0.009577874094247818, -0.009577933698892593], + [-12.497756958007812, 0.009577878750860691, 0.009577957913279533], + ], + [ + 0.6648563146591187, + 0.040433332324028015, + 0.04043332114815712, + -0.029892124235630035, + -0.003740792628377676, + -0.029892126098275185, + ], + -12.910758047958524, + [ + [0.0, -13.562267303466797, 7.167722702026367], + [0.0, 9.584957122802734, -9.880509376525879], + [0.0, 3.9773099422454834, 2.7127861976623535], + ], + ), + ( + '7net-nano-5.5', + 5.5, + -3.532371955806397, + [ + [12.454666137695312, -0.034893378615379333, -0.0348934531211853], + [-12.454666137695312, 0.034893378615379333, 0.03489343076944351], + ], + [ + 0.6540243029594421, + 0.021435830742120743, + 0.021435843780636787, + -0.03487098217010498, + -0.00511842779815197, + -0.03487098589539528, + ], + -12.922036098234884, + [ + [0.0, -13.198136329650879, 7.108638763427734], + [0.0, 9.389389038085938, -9.699457168579102], + [0.0, 3.8087470531463623, 2.5908186435699463], + ], + ), + ( + '7net-nano-6.0', + 6.0, + -3.484466470155766, + [ + [12.169546127319336, -0.02828906662762165, -0.028289003297686577], + [-12.169546127319336, 0.028289061039686203, 0.02828901819884777], + ], + [ + 0.6406453251838684, + 0.026448842138051987, + 0.026448845863342285, + -0.026171253994107246, + -0.003374285064637661, + -0.026171250268816948, + ], + -12.934674727250979, + [ + [0.0, -13.23235034942627, 7.185378074645996], + [0.0, 9.451717376708984, -9.761429786682129], + [0.0, 3.780632972717285, 2.576051712036133], + ], + ), + ], + ids=['7net-nano-4.5', '7net-nano-5.0', '7net-nano-5.5', '7net-nano-6.0'], +) +def test_7net_nano( + name, + cutoff_ref, + g1_ref_e, + g1_ref_f, + g1_ref_s, + g2_ref_e, + g2_ref_f, + atoms_pbc, + atoms_mol, +): + cp_path = pretrained_name_to_path(name) + cp_dict = torch.load(cp_path, map_location='cpu', weights_only=False) + assert sorted(cp_dict) == ['config', 'hash', 'model_state_dict', 'time'] + + model, config = model_from_checkpoint( + cp_path, enable_flash=False, enable_cueq=False + ) + cutoff = config['cutoff'] + assert cutoff == cutoff_ref + + g1 = AtomGraphData.from_numpy_dict(unlabeled_atoms_to_graph(atoms_pbc, cutoff)) + g2 = AtomGraphData.from_numpy_dict(unlabeled_atoms_to_graph(atoms_mol, cutoff)) + + model.set_is_batch_data(False) + g1 = model(g1) + g2 = model(g2) + + assert acl(g1.inferred_total_energy, torch.tensor(g1_ref_e)) + assert acl(g1.inferred_force, torch.tensor(g1_ref_f)) + assert acl(g1.inferred_stress, torch.tensor(g1_ref_s)) + + assert acl(g2.inferred_total_energy, torch.tensor(g2_ref_e)) + assert acl(g2.inferred_force, torch.tensor(g2_ref_f))