mirror of
https://github.com/cp2k/cp2k.git
synced 2026-07-21 06:25:15 -04:00
Add MACE support interface (#5580)
Co-authored-by: Ole Schütt <ole@schuett.name>
This commit is contained in:
xysun
GitHub
parent
290501bb7f
commit
a0df33dd6a
22 changed files with 568 additions and 77 deletions
BIN
data/MACE/MACE_scratch_run-3.model-cp2k.pth
Normal file
BIN
data/MACE/MACE_scratch_run-3.model-cp2k.pth
Normal file
Binary file not shown.
|
|
@ -6,6 +6,7 @@ titlesonly:
|
|||
maxdepth: 2
|
||||
---
|
||||
nequip
|
||||
mace
|
||||
nnp
|
||||
pao-ml
|
||||
deepmd
|
||||
|
|
|
|||
60
docs/methods/machine_learning/mace.md
Normal file
60
docs/methods/machine_learning/mace.md
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
# MACE
|
||||
|
||||
MACE is a framework for building interatomic potentials using higher-order equivariant
|
||||
message-passing neural networks. The methodology is described in detail in the literature by Batatia
|
||||
et al. (2022).
|
||||
|
||||
**Note:** Running MACE requires a CP2K build with LibTorch support
|
||||
|
||||
Like the [NequIP/Allegro](https://manual.cp2k.org/trunk/methods/machine_learning/nequip.html)
|
||||
interface, CP2K runs MACE models through the generic LibTorch interface: a trained MACE model is
|
||||
exported once to a self-contained TorchScript file (`.pth`), which CP2K then loads and evaluates at
|
||||
runtime. No Python interpreter is involved during the simulation.
|
||||
|
||||
## Exporting a MACE model for CP2K
|
||||
|
||||
Wraps a trained MACE model (`.model`) and compiles it to CP2K-loadable TorchScrip file with the
|
||||
helper script `cp2k/tools/mace/create_cp2k_model.py`:
|
||||
|
||||
```shell
|
||||
python create_cp2k_model.py my_mace.model --dtype float64
|
||||
# -> writes my_mace.model-cp2k.pth
|
||||
```
|
||||
|
||||
This conversion only needs to be performed once on a machine with both `torch` and `mace` installed.
|
||||
The resulting `*.pth` file uses the same tensor and metadata format as the NequIP interface (inputs
|
||||
`pos`, `edge_index`, `edge_cell_shift`, `cell`, `atom_types`; outputs `atomic_energy`, `forces`,
|
||||
`virial`) and embeds the metadata (`num_types`, `r_max`, `type_names`, `model_dtype`) that CP2K
|
||||
reads to build the neighbour graph. The `torch` version used for export must be compatible with the
|
||||
LibTorch version linked into CP2K.
|
||||
|
||||
## Input Section
|
||||
|
||||
Inference is configured through the [MACE](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.MACE)
|
||||
section within the `&NONBONDED` forcefield parameters:
|
||||
|
||||
```text
|
||||
&FORCEFIELD
|
||||
&NONBONDED
|
||||
&MACE
|
||||
ATOMS Cu
|
||||
POT_FILE_NAME MACE/my_mace.model-cp2k.pth
|
||||
&END MACE
|
||||
&END NONBONDED
|
||||
&END FORCEFIELD
|
||||
```
|
||||
|
||||
- [ATOMS](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.MACE.ATOMS): a list of elements/kinds; the
|
||||
mapping to the model type list must be consistent with the coordinates in `&COORDS`/`&TOPOLOGY`.
|
||||
- [POT_FILE_NAME](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.MACE.POT_FILE_NAME): path to the
|
||||
exported MACE model.
|
||||
|
||||
MACE is a message-passing model with a non-local receptive field. As with NequIP, the interface
|
||||
evaluates the full system on every MPI rank and divides the energy, forces, and virial by the number
|
||||
of ranks.
|
||||
|
||||
## Further Resources
|
||||
|
||||
- **MACE:** Paper [](#Batatia2022) and source code at
|
||||
[github.com/ACEsuit/mace](https://github.com/ACEsuit/mace).
|
||||
- **e3nn:** For an introduction to Euclidean neural networks, visit [e3nn.org](https://e3nn.org).
|
||||
|
|
@ -220,8 +220,8 @@ of each atom.
|
|||
|
||||
## Torch (PyTorch C++ library)
|
||||
|
||||
LibTorch is the C++ distribution of PyTorch. CP2K uses it for the NequIP interface and for GauXC
|
||||
Skala models.
|
||||
LibTorch is the C++ distribution of PyTorch. CP2K uses it for the NequIP and MACE interfaces and for
|
||||
GauXC Skala models.
|
||||
|
||||
- LibTorch can be downloaded from the
|
||||
[PyTorch installation page](https://pytorch.org/get-started/locally/).
|
||||
|
|
|
|||
|
|
@ -355,7 +355,7 @@ list(
|
|||
manybody_deepmd.F
|
||||
manybody_gal21.F
|
||||
manybody_gal.F
|
||||
manybody_nequip.F
|
||||
manybody_e3nn.F
|
||||
manybody_potential.F
|
||||
manybody_siepmann.F
|
||||
manybody_tersoff.F
|
||||
|
|
|
|||
|
|
@ -66,7 +66,7 @@ MODULE bibliography
|
|||
Kantorovich2008a, Wellendorff2012, Niklasson2014, Borstnik2014, &
|
||||
Rayson2009, Grimme2011, Fattebert2002, Andreussi2012, &
|
||||
Khaliullin2007, Khaliullin2008, Merlot2014, Lin2009, Lin2013, Lin2016ACE, &
|
||||
Batzner2022, DelBen2015, Souza2002, Umari2002, Stengel2009, &
|
||||
Batzner2022, Batatia2022, DelBen2015, Souza2002, Umari2002, Stengel2009, &
|
||||
Luber2014, Berghold2011, DelBen2015b, Campana2009, &
|
||||
Schiffmann2015, Bruck2014, Rappe1992, Ceriotti2012, &
|
||||
Ceriotti2010, Walewski2014, Monkhorst1976, MacDonald1978, Worlton1972, &
|
||||
|
|
@ -164,6 +164,13 @@ CONTAINS
|
|||
source="Nat. Commun.", volume="13", pages="2453", &
|
||||
year=2022, doi="10.1038/s41467-022-29939-5")
|
||||
|
||||
CALL add_reference(key=Batatia2022, &
|
||||
authors=s2a("I. Batatia", "D. P. Kovacs", "G. N. C. Simm", "C. Ortner", "G. Csanyi"), &
|
||||
title="MACE: Higher order equivariant message passing neural networks "// &
|
||||
"for fast and accurate force fields", &
|
||||
source="Adv. Neural Inf. Process. Syst.", volume="35", pages="11423-11436", &
|
||||
year=2022, doi="10.48550/arXiv.2206.07697")
|
||||
|
||||
CALL add_reference(key=VandenCic2006, &
|
||||
authors=s2a("E. Vanden-Eijnden", "G. Ciccotti"), &
|
||||
title="Second-order integrators for Langevin equations with holonomic constraints", &
|
||||
|
|
|
|||
|
|
@ -39,14 +39,9 @@ MODULE fist_neighbor_list_control
|
|||
section_vals_val_get
|
||||
USE kinds, ONLY: dp
|
||||
USE message_passing, ONLY: mp_para_env_type
|
||||
USE pair_potential_types, ONLY: ace_type,&
|
||||
allegro_type,&
|
||||
gal21_type,&
|
||||
gal_type,&
|
||||
nequip_type,&
|
||||
pair_potential_pp_type,&
|
||||
siepmann_type,&
|
||||
tersoff_type
|
||||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, gal21_type, gal_type, mace_type, nequip_type, &
|
||||
pair_potential_pp_type, siepmann_type, tersoff_type
|
||||
USE particle_types, ONLY: particle_type
|
||||
#include "./base/base_uses.f90"
|
||||
|
||||
|
|
@ -236,6 +231,9 @@ CONTAINS
|
|||
IF (ANY(potparm%pot(ikind, jkind)%pot%type == nequip_type)) THEN
|
||||
full_nl(ikind, jkind) = .TRUE.
|
||||
END IF
|
||||
IF (ANY(potparm%pot(ikind, jkind)%pot%type == mace_type)) THEN
|
||||
full_nl(ikind, jkind) = .TRUE.
|
||||
END IF
|
||||
IF (ANY(potparm%pot(ikind, jkind)%pot%type == ace_type)) THEN
|
||||
full_nl(ikind, jkind) = .TRUE.
|
||||
END IF
|
||||
|
|
|
|||
|
|
@ -23,8 +23,8 @@ MODULE fist_nonbond_env_types
|
|||
USE kinds, ONLY: default_string_length,&
|
||||
dp
|
||||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, gal21_type, gal_type, nequip_type, pair_potential_pp_release, &
|
||||
pair_potential_pp_type, siepmann_type, tersoff_type
|
||||
ace_type, allegro_type, gal21_type, gal_type, mace_type, nequip_type, &
|
||||
pair_potential_pp_release, pair_potential_pp_type, siepmann_type, tersoff_type
|
||||
USE torch_api, ONLY: torch_model_release,&
|
||||
torch_model_type
|
||||
#include "./base/base_uses.f90"
|
||||
|
|
@ -497,6 +497,10 @@ CONTAINS
|
|||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
END IF
|
||||
IF (ANY(potparm%pot(idim, jdim)%pot%type == mace_type)) THEN
|
||||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
END IF
|
||||
IF (ANY(potparm%pot(idim, jdim)%pot%type == allegro_type)) THEN
|
||||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
fist_nonbond_env%ij_kind_full_fac(idim, jdim) = 0.5_dp
|
||||
|
|
|
|||
|
|
@ -38,9 +38,9 @@ MODULE fist_nonbond_force
|
|||
USE message_passing, ONLY: mp_comm_type
|
||||
USE pair_potential_coulomb, ONLY: potential_coulomb
|
||||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, deepmd_type, gal21_type, gal_type, nequip_type, nosh_nosh, &
|
||||
nosh_sh, pair_potential_pp_type, pair_potential_single_type, sh_sh, siepmann_type, &
|
||||
tersoff_type
|
||||
ace_type, allegro_type, deepmd_type, gal21_type, gal_type, mace_type, nequip_type, &
|
||||
nosh_nosh, nosh_sh, pair_potential_pp_type, pair_potential_single_type, sh_sh, &
|
||||
siepmann_type, tersoff_type
|
||||
USE particle_types, ONLY: particle_type
|
||||
USE shell_potential_types, ONLY: get_shell,&
|
||||
shell_kind_type
|
||||
|
|
@ -231,6 +231,7 @@ CONTAINS
|
|||
full_nl = ANY(pot%type == tersoff_type) .OR. ANY(pot%type == siepmann_type) &
|
||||
.OR. ANY(pot%type == gal_type) .OR. ANY(pot%type == gal21_type) &
|
||||
.OR. ANY(pot%type == nequip_type) .OR. ANY(pot%type == allegro_type) &
|
||||
.OR. ANY(pot%type == mace_type) &
|
||||
.OR. ANY(pot%type == ace_type) .OR. ANY(pot%type == deepmd_type)
|
||||
IF ((.NOT. full_nl) .AND. (atom_a == atom_b)) THEN
|
||||
fac_ei = 0.5_dp*fac_ei
|
||||
|
|
@ -722,6 +723,7 @@ CONTAINS
|
|||
full_nl = ANY(pot%type == tersoff_type) .OR. ANY(pot%type == siepmann_type) &
|
||||
.OR. ANY(pot%type == gal_type) .OR. ANY(pot%type == gal21_type) &
|
||||
.OR. ANY(pot%type == nequip_type) .OR. ANY(pot%type == allegro_type) &
|
||||
.OR. ANY(pot%type == mace_type) &
|
||||
.OR. ANY(pot%type == ace_type) .OR. ANY(pot%type == deepmd_type)
|
||||
IF ((.NOT. full_nl) .AND. (atom_a == atom_b)) THEN
|
||||
fac_ei = fac_ei*0.5_dp
|
||||
|
|
|
|||
|
|
@ -64,10 +64,11 @@ MODULE force_fields_all
|
|||
spline_nonbond_control
|
||||
USE pair_potential_coulomb, ONLY: potential_coulomb
|
||||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, deepmd_type, ea_type, lj_charmm_type, lj_type, nequip_type, &
|
||||
nn_type, nosh_nosh, nosh_sh, pair_potential_lj_create, pair_potential_pp_create, &
|
||||
pair_potential_pp_type, pair_potential_single_add, pair_potential_single_clean, &
|
||||
pair_potential_single_copy, pair_potential_single_type, sh_sh, siepmann_type, tersoff_type
|
||||
ace_type, allegro_type, deepmd_type, ea_type, lj_charmm_type, lj_type, mace_type, &
|
||||
nequip_type, nn_type, nosh_nosh, nosh_sh, pair_potential_lj_create, &
|
||||
pair_potential_pp_create, pair_potential_pp_type, pair_potential_single_add, &
|
||||
pair_potential_single_clean, pair_potential_single_copy, pair_potential_single_type, &
|
||||
sh_sh, siepmann_type, tersoff_type
|
||||
USE particle_types, ONLY: allocate_particle_set,&
|
||||
particle_type
|
||||
USE physcon, ONLY: bohr
|
||||
|
|
@ -2194,7 +2195,7 @@ CONTAINS
|
|||
atmname == inp_info%nonbonded%pot(j)%pot%at2) THEN
|
||||
SELECT CASE (inp_info%nonbonded%pot(j)%pot%type(1))
|
||||
CASE (ea_type, tersoff_type, siepmann_type, nequip_type, &
|
||||
allegro_type, deepmd_type, ace_type)
|
||||
allegro_type, deepmd_type, ace_type, mace_type)
|
||||
! Charge is zero for EAM, TERSOFF and SIEPMANN type potential
|
||||
! Do nothing..
|
||||
CASE DEFAULT
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ MODULE force_fields_input
|
|||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, b4_type, bm_type, deepmd_type, do_potential_single_allocation, &
|
||||
ea_type, eam_pot_type, ft_pot_type, ft_type, ftd_type, gal21_type, gal_type, gp_type, &
|
||||
gw_type, ip_type, ipbv_pot_type, lj_charmm_type, nequip_pot_type, nequip_type, &
|
||||
gw_type, ip_type, ipbv_pot_type, lj_charmm_type, mace_type, nequip_pot_type, nequip_type, &
|
||||
no_potential_single_allocation, pair_potential_p_type, pair_potential_reallocate, &
|
||||
potential_single_allocation, siepmann_type, tab_pot_type, tab_type, tersoff_type, wl_type
|
||||
USE shell_potential_types, ONLY: shell_p_create,&
|
||||
|
|
@ -109,8 +109,8 @@ CONTAINS
|
|||
CHARACTER(LEN=default_string_length), &
|
||||
DIMENSION(:), POINTER :: atm_names
|
||||
INTEGER :: nace, nb4, nbends, nbm, nbmhft, nbmhftd, nbonds, nchg, ndeepmd, neam, ngal, &
|
||||
ngal21, ngd, ngp, nimpr, nipbv, nlj, nnequip, nopbend, nshell, nsiepmann, ntab, ntersoff, &
|
||||
ntors, ntot, nubs, nwl
|
||||
ngal21, ngd, ngp, nimpr, nipbv, nlj, nmace, nnequip, nopbend, nshell, nsiepmann, ntab, &
|
||||
ntersoff, ntors, ntot, nubs, nwl
|
||||
LOGICAL :: explicit, unique_spline
|
||||
REAL(KIND=dp) :: min_eps_spline_allowed
|
||||
TYPE(input_info_type), POINTER :: inp_info
|
||||
|
|
@ -135,9 +135,11 @@ CONTAINS
|
|||
SELECT CASE (ff_type%ff_type)
|
||||
CASE (do_ff_charmm, do_ff_amber, do_ff_g96, do_ff_g87)
|
||||
CALL section_vals_val_get(ff_section, "PARM_FILE_NAME", c_val=ff_type%ff_file_name)
|
||||
|
||||
IF (TRIM(ff_type%ff_file_name) == "") THEN
|
||||
CPABORT("Force Field Parameter's filename is empty! Please check your input file.")
|
||||
END IF
|
||||
|
||||
CASE (do_ff_undef)
|
||||
! Do Nothing
|
||||
CASE DEFAULT
|
||||
|
|
@ -326,6 +328,19 @@ CONTAINS
|
|||
CALL read_ace_section(inp_info%nonbonded, tmp_section2, ntot)
|
||||
END IF
|
||||
|
||||
tmp_section2 => section_vals_get_subs_vals(tmp_section, "MACE")
|
||||
CALL section_vals_get(tmp_section2, explicit=explicit, n_repetition=nmace)
|
||||
ntot = nlj + nwl + neam + ngd + nipbv + nbmhft + nbmhftd + nb4 + nbm + ngp + ntersoff + &
|
||||
ngal + ngal21 + nsiepmann + nnequip + ntab + ndeepmd + nace
|
||||
IF (explicit) THEN
|
||||
! avoid repeating the mace section for each pair
|
||||
CALL section_vals_val_get(tmp_section2, "ATOMS", c_vals=atm_names)
|
||||
nmace = nmace - 1 + SIZE(atm_names) + (SIZE(atm_names)*SIZE(atm_names) - SIZE(atm_names))/2
|
||||
! MACE reuses the nequip_pot_type storage (set%nequip), hence nequip=.TRUE. here
|
||||
CALL pair_potential_reallocate(inp_info%nonbonded, 1, ntot + nmace, nequip=.TRUE.)
|
||||
CALL read_mace_section(inp_info%nonbonded, tmp_section2, ntot)
|
||||
END IF
|
||||
|
||||
END IF
|
||||
|
||||
tmp_section => section_vals_get_subs_vals(ff_section, "NONBONDED14")
|
||||
|
|
@ -541,6 +556,7 @@ CONTAINS
|
|||
ipbv%a(15) = 12917180227.21_dp
|
||||
ELSE IF (((at1(1:1) == 'O') .AND. (at2(1:1) == 'H')) .OR. &
|
||||
((at1(1:1) == 'H') .AND. (at2(1:1) == 'O'))) THEN
|
||||
|
||||
ipbv%rcore = 2.95_dp ! a.u.
|
||||
|
||||
ipbv%m = -0.004025691139759147_dp ! Hartree/a.u.
|
||||
|
|
@ -602,6 +618,7 @@ CONTAINS
|
|||
ft%d = cp_unit_to_cp2k(0.499_dp, "eV*angstrom^8")
|
||||
ELSE IF (((at1(1:2) == 'NA') .AND. (at2(1:2) == 'CL')) .OR. &
|
||||
((at1(1:2) == 'CL') .AND. (at2(1:2) == 'NA'))) THEN
|
||||
|
||||
ft%a = cp_unit_to_cp2k(1256.31_dp, "eV")
|
||||
ft%c = cp_unit_to_cp2k(7.00_dp, "eV*angstrom^6")
|
||||
ft%d = cp_unit_to_cp2k(8.676_dp, "eV*angstrom^8")
|
||||
|
|
@ -826,6 +843,54 @@ CONTAINS
|
|||
|
||||
END SUBROUTINE read_nequip_section
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief Reads the MACE section
|
||||
!> \param nonbonded ...
|
||||
!> \param section ...
|
||||
!> \param start ...
|
||||
!> \author Xinyue Sun
|
||||
! **************************************************************************************************
|
||||
SUBROUTINE read_mace_section(nonbonded, section, start)
|
||||
TYPE(pair_potential_p_type), POINTER :: nonbonded
|
||||
TYPE(section_vals_type), POINTER :: section
|
||||
INTEGER, INTENT(IN) :: start
|
||||
|
||||
CHARACTER(LEN=default_string_length) :: pot_file_name
|
||||
CHARACTER(LEN=default_string_length), &
|
||||
DIMENSION(:), POINTER :: atm_names
|
||||
INTEGER :: isec, jsec, n_items
|
||||
TYPE(nequip_pot_type) :: mace
|
||||
|
||||
n_items = 1
|
||||
isec = 1
|
||||
n_items = isec*n_items
|
||||
CALL section_vals_val_get(section, "ATOMS", c_vals=atm_names)
|
||||
CALL section_vals_val_get(section, "POT_FILE_NAME", c_val=pot_file_name)
|
||||
|
||||
mace%pot_file_name = discover_file(pot_file_name)
|
||||
! MACE models use standardized units: Angstrom, eV and eV/Angstrom
|
||||
mace%unit_length = "angstrom"
|
||||
mace%unit_energy = "eV"
|
||||
mace%unit_forces = "eV/Angstrom"
|
||||
! MACE models are exported to speak the same metadata/tensor dialect as NequIP
|
||||
CALL read_nequip_data(mace)
|
||||
CALL check_cp2k_atom_names_in_torch(atm_names, mace%type_names_torch)
|
||||
|
||||
DO isec = 1, SIZE(atm_names)
|
||||
DO jsec = isec, SIZE(atm_names)
|
||||
nonbonded%pot(start + n_items)%pot%type = mace_type
|
||||
nonbonded%pot(start + n_items)%pot%at1 = atm_names(isec)
|
||||
nonbonded%pot(start + n_items)%pot%at2 = atm_names(jsec)
|
||||
CALL uppercase(nonbonded%pot(start + n_items)%pot%at1)
|
||||
CALL uppercase(nonbonded%pot(start + n_items)%pot%at2)
|
||||
nonbonded%pot(start + n_items)%pot%set(1)%nequip = mace
|
||||
nonbonded%pot(start + n_items)%pot%rcutsq = mace%rcutsq
|
||||
n_items = n_items + 1
|
||||
END DO
|
||||
END DO
|
||||
|
||||
END SUBROUTINE read_mace_section
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief Reads the LJ section
|
||||
!> \param nonbonded ...
|
||||
|
|
@ -1272,11 +1337,13 @@ CONTAINS
|
|||
! Calculate p_inv the inverse of the matrix p
|
||||
p_inv(:, :) = 0.0_dp
|
||||
CALL invert_matrix(p, p_inv, eval_error)
|
||||
|
||||
IF (eval_error >= 1.0E-8_dp) THEN
|
||||
CALL cp_warn(__LOCATION__, &
|
||||
"The polynomial fit for the BUCK4RANGES potential is only accurate to "// &
|
||||
TRIM(cp_to_string(eval_error)))
|
||||
END IF
|
||||
|
||||
! Get the 6 coefficients of the 5th-order polynomial -> x(1:6)
|
||||
! and the 4 coefficients of the 3rd-order polynomial -> x(7:10)
|
||||
x(:) = MATMUL(p_inv(:, :), v(:))
|
||||
|
|
|
|||
|
|
@ -15,9 +15,9 @@
|
|||
! **************************************************************************************************
|
||||
MODULE input_cp2k_mm
|
||||
USE bibliography, ONLY: &
|
||||
Batzner2022, Bochkarev2024, Clabaut2020, Clabaut2021, Devynck2012, Dick1958, Drautz2019, &
|
||||
Foiles1986, Lysogorskiy2021, Mitchell1993, Musaelian2023, Siepmann1995, Tan2025, &
|
||||
Tersoff1988, Tosi1964a, Tosi1964b, Wang2018, Yamada2000, Zeng2023
|
||||
Batatia2022, Batzner2022, Bochkarev2024, Clabaut2020, Clabaut2021, Devynck2012, Dick1958, &
|
||||
Drautz2019, Foiles1986, Lysogorskiy2021, Mitchell1993, Musaelian2023, Siepmann1995, &
|
||||
Tan2025, Tersoff1988, Tosi1964a, Tosi1964b, Wang2018, Yamada2000, Zeng2023
|
||||
USE cp_output_handling, ONLY: cp_print_key_section_create,&
|
||||
debug_print_level,&
|
||||
high_print_level,&
|
||||
|
|
@ -1185,6 +1185,10 @@ CONTAINS
|
|||
CALL section_add_subsection(section, subsection)
|
||||
CALL section_release(subsection)
|
||||
|
||||
CALL create_MACE_section(subsection)
|
||||
CALL section_add_subsection(section, subsection)
|
||||
CALL section_release(subsection)
|
||||
|
||||
CALL create_DEEPMD_section(subsection)
|
||||
CALL section_add_subsection(section, subsection)
|
||||
CALL section_release(subsection)
|
||||
|
|
@ -1483,6 +1487,48 @@ CONTAINS
|
|||
|
||||
END SUBROUTINE create_NEQUIP_section
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief This section specifies the input parameters for MACE potential type
|
||||
!> \param section the section to create
|
||||
!> \author Xinyue Sun
|
||||
! **************************************************************************************************
|
||||
SUBROUTINE create_MACE_section(section)
|
||||
TYPE(section_type), POINTER :: section
|
||||
|
||||
TYPE(keyword_type), POINTER :: keyword
|
||||
|
||||
CPASSERT(.NOT. ASSOCIATED(section))
|
||||
CALL section_create(section, __LOCATION__, name="MACE", &
|
||||
description="This section specifies the input parameters for MACE potential type, "// &
|
||||
"a higher-order equivariant message-passing neural network. "// &
|
||||
"The MACE model must be exported to a TorchScript file that takes a single "// &
|
||||
"dictionary argument (see the create_cp2k_model.py helper). "// &
|
||||
"Requires linking with libtorch library from <https://pytorch.org/cppdocs/installing.html>.", &
|
||||
citations=[Batatia2022], n_keywords=1, n_subsections=0, repeats=.FALSE.)
|
||||
|
||||
NULLIFY (keyword)
|
||||
|
||||
CALL keyword_create(keyword, __LOCATION__, name="ATOMS", &
|
||||
description="Defines the atomic kinds involved in the MACE potential. "// &
|
||||
"Provide a list of each element, making sure that the mapping from the ATOMS list "// &
|
||||
"to MACE atom types is correct. This mapping should also be consistent for the "// &
|
||||
"atomic coordinates as specified in the sections COORDS or TOPOLOGY.", &
|
||||
usage="ATOMS {KIND 1} {KIND 2} .. {KIND N}", type_of_var=char_t, &
|
||||
n_var=-1)
|
||||
CALL section_add_keyword(section, keyword)
|
||||
CALL keyword_release(keyword)
|
||||
|
||||
CALL keyword_create(keyword, __LOCATION__, name="POT_FILE_NAME", &
|
||||
variants=["MODEL_FILE_NAME"], &
|
||||
description="Specifies the filename that contains the exported MACE model. "// &
|
||||
"MACE models use standardized units (Angstrom for length, eV for energy, "// &
|
||||
"eV/Angstrom for forces), so no unit keywords are required.", &
|
||||
usage="POT_FILE_NAME {FILENAME}", default_lc_val=" ")
|
||||
CALL section_add_keyword(section, keyword)
|
||||
CALL keyword_release(keyword)
|
||||
|
||||
END SUBROUTINE create_MACE_section
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief This section specifies the input parameters for ACE potential type
|
||||
!> \param section the section to create
|
||||
|
|
|
|||
|
|
@ -6,13 +6,16 @@
|
|||
!--------------------------------------------------------------------------------------------------!
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief Shared TorchScript evaluation path for e3nn-based equivariant message-passing
|
||||
!> potentials (NequIP, Allegro and MACE).
|
||||
!> \par History
|
||||
!> Implementation of NequIP and Allegro potentials - [gtocci] 2022
|
||||
!> Index mapping of atoms from .xyz to Allegro config.yaml file - [mbilichenko] 2024
|
||||
!> Refactoring and update to NequIP version >= v0.7.0 - [gtocci] 2026
|
||||
!> Renamed manybody_nequip -> manybody_e3nn as it now also serves MACE - [xysun] 2026
|
||||
!> \author Gabriele Tocci
|
||||
! **************************************************************************************************
|
||||
MODULE manybody_nequip
|
||||
MODULE manybody_e3nn
|
||||
|
||||
USE atomic_kind_types, ONLY: atomic_kind_type
|
||||
USE cell_types, ONLY: cell_type
|
||||
|
|
@ -28,7 +31,8 @@ MODULE manybody_nequip
|
|||
dp,&
|
||||
int_8
|
||||
USE message_passing, ONLY: mp_para_env_type
|
||||
USE pair_potential_types, ONLY: nequip_pot_type,&
|
||||
USE pair_potential_types, ONLY: mace_type,&
|
||||
nequip_pot_type,&
|
||||
nequip_type,&
|
||||
pair_potential_pp_type,&
|
||||
pair_potential_single_type
|
||||
|
|
@ -43,10 +47,10 @@ MODULE manybody_nequip
|
|||
IMPLICIT NONE
|
||||
|
||||
PRIVATE
|
||||
PUBLIC :: nequip_energy_store_force_virial, &
|
||||
nequip_add_force_virial
|
||||
PUBLIC :: e3nn_energy_store_force_virial, &
|
||||
e3nn_add_force_virial
|
||||
|
||||
CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'manybody_nequip'
|
||||
CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'manybody_e3nn'
|
||||
|
||||
TYPE, PRIVATE :: nequip_work_type
|
||||
INTEGER :: target_pot_type
|
||||
|
|
@ -89,7 +93,7 @@ CONTAINS
|
|||
!> Refactoring and unifying NequIP and Allegro - [gtocci] 2026
|
||||
!> \author Gabriele Tocci - University of Zurich
|
||||
! **************************************************************************************************
|
||||
SUBROUTINE nequip_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
SUBROUTINE e3nn_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
atomic_kind_set, potparm, r_last_update_pbc, &
|
||||
pot_total, fist_nonbond_env, para_env, use_virial, &
|
||||
target_pot_type)
|
||||
|
|
@ -107,7 +111,7 @@ CONTAINS
|
|||
LOGICAL, INTENT(IN) :: use_virial
|
||||
INTEGER, INTENT(IN) :: target_pot_type
|
||||
|
||||
CHARACTER(LEN=*), PARAMETER :: routineN = 'nequip_energy_store_force_virial'
|
||||
CHARACTER(LEN=*), PARAMETER :: routineN = 'e3nn_energy_store_force_virial'
|
||||
|
||||
INTEGER :: handle
|
||||
TYPE(nequip_data_type), POINTER :: neq_data
|
||||
|
|
@ -132,7 +136,8 @@ CONTAINS
|
|||
|
||||
CALL setup_neq_data(fist_nonbond_env, neq_data, neq_pot, nequip_work)
|
||||
|
||||
IF (nequip_work%target_pot_type == nequip_type) THEN
|
||||
IF (nequip_work%target_pot_type == nequip_type .OR. &
|
||||
nequip_work%target_pot_type == mace_type) THEN
|
||||
CALL prepare_edges_shifts_nequip(nequip_work)
|
||||
ELSE
|
||||
CALL prepare_edges_shifts_allegro(nequip_work)
|
||||
|
|
@ -146,7 +151,7 @@ CONTAINS
|
|||
CALL release_nequip_work(nequip_work)
|
||||
|
||||
CALL timestop(handle)
|
||||
END SUBROUTINE nequip_energy_store_force_virial
|
||||
END SUBROUTINE e3nn_energy_store_force_virial
|
||||
|
||||
! **************************************************************************************************
|
||||
!> \brief ...
|
||||
|
|
@ -564,7 +569,8 @@ CONTAINS
|
|||
n_atoms = SIZE(nequip_work%particle_set)
|
||||
|
||||
! for allegro ensure ghost atoms are included in the evaluation
|
||||
IF (nequip_work%target_pot_type /= nequip_type) THEN
|
||||
IF (nequip_work%target_pot_type /= nequip_type .AND. &
|
||||
nequip_work%target_pot_type /= mace_type) THEN
|
||||
! label atoms in the local edges
|
||||
DO i = 1, SIZE(nequip_work%local_edges, 2)
|
||||
atom_a = INT(nequip_work%local_edges(1, i))
|
||||
|
|
@ -708,7 +714,8 @@ CONTAINS
|
|||
DO iat_use = 1, SIZE(neq_data%use_indices)
|
||||
iat = neq_data%use_indices(iat_use)
|
||||
! Only apply the local mask for Allegro models
|
||||
IF (nequip_work%target_pot_type /= nequip_type) THEN
|
||||
IF (nequip_work%target_pot_type /= nequip_type .AND. &
|
||||
nequip_work%target_pot_type /= mace_type) THEN
|
||||
IF (.NOT. nequip_work%sum_energy(iat)) CYCLE
|
||||
END IF
|
||||
|
||||
|
|
@ -717,7 +724,8 @@ CONTAINS
|
|||
CALL torch_tensor_release(t_energy)
|
||||
pot_total = pot_total*pot%unit_energy_val
|
||||
|
||||
IF (nequip_work%target_pot_type == nequip_type) THEN
|
||||
IF (nequip_work%target_pot_type == nequip_type .OR. &
|
||||
nequip_work%target_pot_type == mace_type) THEN
|
||||
neq_data%force = neq_data%force/REAL(nequip_work%para_env%num_pe, dp)
|
||||
pot_total = pot_total/REAL(nequip_work%para_env%num_pe, dp)
|
||||
END IF
|
||||
|
|
@ -728,7 +736,8 @@ CONTAINS
|
|||
|
||||
neq_data%virial(:, :) = RESHAPE(v_ptr, [3, 3])*pot%unit_energy_val
|
||||
CALL torch_tensor_release(t_virial)
|
||||
IF (nequip_work%target_pot_type == nequip_type) THEN
|
||||
IF (nequip_work%target_pot_type == nequip_type .OR. &
|
||||
nequip_work%target_pot_type == mace_type) THEN
|
||||
neq_data%virial = neq_data%virial/REAL(nequip_work%para_env%num_pe, dp)
|
||||
END IF
|
||||
END IF
|
||||
|
|
@ -745,7 +754,7 @@ CONTAINS
|
|||
!> Sum forces, virial to nonbond - [gtocci] 2026
|
||||
!> \author Gabriele Tocci - University of Zurich
|
||||
! **************************************************************************************************
|
||||
SUBROUTINE nequip_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
SUBROUTINE e3nn_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
|
||||
REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: f_nonbond, pv_nonbond
|
||||
LOGICAL, INTENT(IN) :: use_virial
|
||||
|
|
@ -764,6 +773,6 @@ CONTAINS
|
|||
f_nonbond(1:3, iat) = f_nonbond(1:3, iat) + neq_data%force(1:3, iat_use)
|
||||
END DO
|
||||
|
||||
END SUBROUTINE nequip_add_force_virial
|
||||
END SUBROUTINE e3nn_add_force_virial
|
||||
|
||||
END MODULE manybody_nequip
|
||||
END MODULE manybody_e3nn
|
||||
|
|
@ -30,6 +30,8 @@ MODULE manybody_potential
|
|||
ace_energy_store_force_virial
|
||||
USE manybody_deepmd, ONLY: deepmd_add_force_virial,&
|
||||
deepmd_energy_store_force_virial
|
||||
USE manybody_e3nn, ONLY: e3nn_add_force_virial,&
|
||||
e3nn_energy_store_force_virial
|
||||
USE manybody_eam, ONLY: get_force_eam
|
||||
USE manybody_gal, ONLY: destroy_gal_arrays,&
|
||||
gal_energy,&
|
||||
|
|
@ -39,8 +41,6 @@ MODULE manybody_potential
|
|||
gal21_energy,&
|
||||
gal21_forces,&
|
||||
setup_gal21_arrays
|
||||
USE manybody_nequip, ONLY: nequip_add_force_virial,&
|
||||
nequip_energy_store_force_virial
|
||||
USE manybody_siepmann, ONLY: destroy_siepmann_arrays,&
|
||||
print_nr_ions_siepmann,&
|
||||
setup_siepmann_arrays,&
|
||||
|
|
@ -54,8 +54,9 @@ MODULE manybody_potential
|
|||
USE message_passing, ONLY: mp_para_env_type
|
||||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, deepmd_type, ea_type, eam_pot_type, gal21_pot_type, gal21_type, &
|
||||
gal_pot_type, gal_type, nequip_type, pair_potential_pp_type, pair_potential_single_type, &
|
||||
siepmann_pot_type, siepmann_type, tersoff_pot_type, tersoff_type
|
||||
gal_pot_type, gal_type, mace_type, nequip_type, pair_potential_pp_type, &
|
||||
pair_potential_single_type, siepmann_pot_type, siepmann_type, tersoff_pot_type, &
|
||||
tersoff_type
|
||||
USE particle_types, ONLY: particle_type
|
||||
USE util, ONLY: sort
|
||||
#include "./base/base_uses.f90"
|
||||
|
|
@ -105,11 +106,11 @@ CONTAINS
|
|||
INTEGER, DIMENSION(:), POINTER :: glob_loc_list_a, work_list
|
||||
INTEGER, DIMENSION(:, :), POINTER :: glob_loc_list, list, sort_list
|
||||
LOGICAL :: any_ace, any_allegro, any_deepmd, &
|
||||
any_gal, any_gal21, any_nequip, &
|
||||
any_siepmann, any_tersoff
|
||||
any_gal, any_gal21, any_mace, &
|
||||
any_nequip, any_siepmann, any_tersoff
|
||||
REAL(KIND=dp) :: drij, embed, pot_ace, pot_allegro, &
|
||||
pot_deepmd, pot_loc, pot_nequip, qr, &
|
||||
rab2_max, rij(3)
|
||||
pot_deepmd, pot_loc, pot_mace, &
|
||||
pot_nequip, qr, rab2_max, rij(3)
|
||||
REAL(KIND=dp), DIMENSION(3) :: cell_v, cvi
|
||||
REAL(KIND=dp), DIMENSION(:, :), POINTER :: glob_cell_v
|
||||
REAL(KIND=dp), POINTER :: fembed(:)
|
||||
|
|
@ -132,6 +133,7 @@ CONTAINS
|
|||
any_gal21 = .FALSE.
|
||||
any_allegro = .FALSE.
|
||||
any_nequip = .FALSE.
|
||||
any_mace = .FALSE.
|
||||
any_ace = .FALSE.
|
||||
any_deepmd = .FALSE.
|
||||
CALL timeset(routineN, handle)
|
||||
|
|
@ -177,6 +179,7 @@ CONTAINS
|
|||
pot => potparm%pot(ikind, jkind)%pot
|
||||
any_tersoff = any_tersoff .OR. ANY(pot%type == tersoff_type)
|
||||
any_nequip = any_nequip .OR. ANY(pot%type == nequip_type)
|
||||
any_mace = any_mace .OR. ANY(pot%type == mace_type)
|
||||
any_ace = any_ace .OR. ANY(pot%type == ace_type)
|
||||
any_allegro = any_allegro .OR. ANY(pot%type == allegro_type)
|
||||
any_deepmd = any_deepmd .OR. ANY(pot%type == deepmd_type)
|
||||
|
|
@ -189,7 +192,7 @@ CONTAINS
|
|||
! NEQUIP
|
||||
IF (any_nequip) THEN
|
||||
NULLIFY (glob_loc_list, glob_cell_v, glob_loc_list_a)
|
||||
CALL nequip_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
CALL e3nn_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
atomic_kind_set, potparm, r_last_update_pbc, &
|
||||
pot_nequip, fist_nonbond_env, &
|
||||
para_env, use_virial, nequip_type)
|
||||
|
|
@ -198,12 +201,21 @@ CONTAINS
|
|||
! ALLEGRO
|
||||
IF (any_allegro) THEN
|
||||
NULLIFY (glob_loc_list, glob_cell_v, glob_loc_list_a)
|
||||
CALL nequip_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
CALL e3nn_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
atomic_kind_set, potparm, r_last_update_pbc, &
|
||||
pot_allegro, fist_nonbond_env, &
|
||||
para_env, use_virial, allegro_type)
|
||||
pot_manybody = pot_manybody + pot_allegro
|
||||
END IF
|
||||
! MACE (reuses the NequIP message-passing evaluation path)
|
||||
IF (any_mace) THEN
|
||||
NULLIFY (glob_loc_list, glob_cell_v, glob_loc_list_a)
|
||||
CALL e3nn_energy_store_force_virial(nonbonded, particle_set, local_particles, cell, &
|
||||
atomic_kind_set, potparm, r_last_update_pbc, &
|
||||
pot_mace, fist_nonbond_env, &
|
||||
para_env, use_virial, mace_type)
|
||||
pot_manybody = pot_manybody + pot_mace
|
||||
END IF
|
||||
! ACE
|
||||
IF (any_ace) THEN
|
||||
CALL ace_energy_store_force_virial(particle_set, cell, atomic_kind_set, potparm, &
|
||||
|
|
@ -578,8 +590,8 @@ CONTAINS
|
|||
INTEGER, DIMENSION(:), POINTER :: glob_loc_list_a, work_list
|
||||
INTEGER, DIMENSION(:, :), POINTER :: glob_loc_list, list, sort_list
|
||||
LOGICAL :: any_ace, any_allegro, any_deepmd, &
|
||||
any_gal, any_gal21, any_nequip, &
|
||||
any_siepmann, any_tersoff
|
||||
any_gal, any_gal21, any_mace, &
|
||||
any_nequip, any_siepmann, any_tersoff
|
||||
REAL(KIND=dp) :: f_eam, fac, fr(3), ptens11, ptens12, ptens13, ptens21, ptens22, ptens23, &
|
||||
ptens31, ptens32, ptens33, rab(3), rab2, rab2_max, rtmp(3)
|
||||
REAL(KIND=dp), DIMENSION(3) :: cell_v, cvi
|
||||
|
|
@ -599,6 +611,7 @@ CONTAINS
|
|||
any_tersoff = .FALSE.
|
||||
any_allegro = .FALSE.
|
||||
any_nequip = .FALSE.
|
||||
any_mace = .FALSE.
|
||||
any_siepmann = .FALSE.
|
||||
any_ace = .FALSE.
|
||||
any_deepmd = .FALSE.
|
||||
|
|
@ -660,7 +673,7 @@ CONTAINS
|
|||
END DO
|
||||
END DO
|
||||
IF (any_nequip) THEN
|
||||
CALL nequip_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
CALL e3nn_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
END IF
|
||||
|
||||
! ALLEGRO
|
||||
|
|
@ -670,7 +683,17 @@ CONTAINS
|
|||
END DO
|
||||
END DO
|
||||
IF (any_allegro) THEN
|
||||
CALL nequip_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
CALL e3nn_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
END IF
|
||||
|
||||
! MACE (reuses the NequIP force/virial accumulation)
|
||||
DO ikind = 1, nkinds
|
||||
DO jkind = ikind, nkinds
|
||||
any_mace = any_mace .OR. ANY(potparm%pot(ikind, jkind)%pot%type == mace_type)
|
||||
END DO
|
||||
END DO
|
||||
IF (any_mace) THEN
|
||||
CALL e3nn_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
|
||||
END IF
|
||||
|
||||
! starting the force loop
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ MODULE pair_potential
|
|||
USE pair_potential_types, ONLY: &
|
||||
ace_type, allegro_type, b4_type, bm_type, compare_pot, deepmd_type, ea_type, ft_type, &
|
||||
ftd_type, gal21_type, gal_type, gp_type, gw_type, ip_type, list_pot, lj_charmm_type, &
|
||||
lj_type, multi_type, nequip_type, nn_type, pair_potential_pp_type, &
|
||||
lj_type, mace_type, multi_type, nequip_type, nn_type, pair_potential_pp_type, &
|
||||
pair_potential_single_type, potential_single_allocation, siepmann_type, tab_type, &
|
||||
tersoff_type, wl_type
|
||||
USE pair_potential_util, ONLY: ener_pot,&
|
||||
|
|
@ -188,7 +188,8 @@ CONTAINS
|
|||
DO k = 1, SIZE(pot%type)
|
||||
SELECT CASE (pot%type(k))
|
||||
CASE (lj_type, lj_charmm_type, wl_type, gw_type, ft_type, ftd_type, ip_type, &
|
||||
b4_type, bm_type, gp_type, ea_type, allegro_type, nequip_type, tab_type, deepmd_type, ace_type)
|
||||
b4_type, bm_type, gp_type, ea_type, allegro_type, nequip_type, mace_type, tab_type, &
|
||||
deepmd_type, ace_type)
|
||||
pot%no_pp = .FALSE.
|
||||
CASE (tersoff_type)
|
||||
pot%no_mb = .FALSE.
|
||||
|
|
@ -206,7 +207,7 @@ CONTAINS
|
|||
END SELECT
|
||||
! Special case for EAM
|
||||
SELECT CASE (pot%type(k))
|
||||
CASE (ea_type, nequip_type, allegro_type, deepmd_type, ace_type)
|
||||
CASE (ea_type, nequip_type, allegro_type, mace_type, deepmd_type, ace_type)
|
||||
pot%no_mb = .FALSE.
|
||||
END SELECT
|
||||
END DO
|
||||
|
|
@ -605,7 +606,7 @@ CONTAINS
|
|||
nvar = 5 + nvar
|
||||
CASE (ea_type)
|
||||
nvar = 4 + nvar
|
||||
CASE (nequip_type)
|
||||
CASE (nequip_type, mace_type)
|
||||
nvar = 1 + nvar
|
||||
CASE (allegro_type)
|
||||
nvar = 1 + nvar
|
||||
|
|
@ -688,7 +689,7 @@ CONTAINS
|
|||
pot_par(nk, 2) = potparm%pot(i, j)%pot%set(1)%eam%drhoar
|
||||
pot_par(nk, 3) = potparm%pot(i, j)%pot%set(1)%eam%acutal
|
||||
pot_par(nk, 4) = potparm%pot(i, j)%pot%set(1)%eam%npoints
|
||||
CASE (nequip_type, allegro_type)
|
||||
CASE (nequip_type, allegro_type, mace_type)
|
||||
pot_par(nk, 1) = str2id( &
|
||||
TRIM(potparm%pot(i, j)%pot%set(1)%nequip%pot_file_name))
|
||||
CASE (ace_type)
|
||||
|
|
|
|||
|
|
@ -54,9 +54,10 @@ MODULE pair_potential_types
|
|||
gal21_type = 18, &
|
||||
tab_type = 19, &
|
||||
deepmd_type = 20, &
|
||||
ace_type = 21
|
||||
ace_type = 21, &
|
||||
mace_type = 22
|
||||
|
||||
INTEGER, PUBLIC, PARAMETER, DIMENSION(21) :: list_pot = [nn_type, &
|
||||
INTEGER, PUBLIC, PARAMETER, DIMENSION(22) :: list_pot = [nn_type, &
|
||||
lj_type, &
|
||||
lj_charmm_type, &
|
||||
ft_type, &
|
||||
|
|
@ -76,7 +77,8 @@ MODULE pair_potential_types
|
|||
gal21_type, &
|
||||
tab_type, &
|
||||
deepmd_type, &
|
||||
ace_type]
|
||||
ace_type, &
|
||||
mace_type]
|
||||
|
||||
! Shell model
|
||||
INTEGER, PUBLIC, PARAMETER :: nosh_nosh = 0, &
|
||||
|
|
@ -468,7 +470,7 @@ CONTAINS
|
|||
CASE (deepmd_type)
|
||||
IF ((pot1%set(i)%deepmd%deepmd_file_name == pot2%set(i)%deepmd%deepmd_file_name) .AND. &
|
||||
(pot1%set(i)%deepmd%atom_deepmd_type == pot2%set(i)%deepmd%atom_deepmd_type)) mycompare = .TRUE.
|
||||
CASE (nequip_type, allegro_type)
|
||||
CASE (nequip_type, allegro_type, mace_type)
|
||||
IF ((pot1%set(i)%nequip%pot_file_name == pot2%set(i)%nequip%pot_file_name) .AND. &
|
||||
(pot1%set(i)%nequip%unit_length == pot2%set(i)%nequip%unit_length) .AND. &
|
||||
(pot1%set(i)%nequip%unit_forces == pot2%set(i)%nequip%unit_forces) .AND. &
|
||||
|
|
@ -1289,7 +1291,7 @@ CONTAINS
|
|||
CALL pair_potential_goodwin_create(p%pot(i)%pot%set(std_dim)%goodwin)
|
||||
CASE (ea_type)
|
||||
CALL pair_potential_eam_create(p%pot(i)%pot%set(std_dim)%eam)
|
||||
CASE (nequip_type, allegro_type)
|
||||
CASE (nequip_type, allegro_type, mace_type)
|
||||
CALL pair_potential_nequip_create(p%pot(i)%pot%set(std_dim)%nequip)
|
||||
CASE (ace_type)
|
||||
CALL pair_potential_ace_create(p%pot(i)%pot%set(std_dim)%ace)
|
||||
|
|
|
|||
3
tests/Fist/regtest-mace/TEST_FILES.toml
Normal file
3
tests/Fist/regtest-mace/TEST_FILES.toml
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
# Test of MACE using libtorch https://pytorch.org/cppdocs/installing.html
|
||||
"mace_test.inp" = [{matcher="M011", tol=1.0E-9, ref=-1532.363300232407482}]
|
||||
#EOF
|
||||
34
tests/Fist/regtest-mace/mace-test.xyz
Normal file
34
tests/Fist/regtest-mace/mace-test.xyz
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
32
|
||||
Lattice="7.23 0.0 0.0 0.0 7.23 0.0 0.0 0.0 7.23" Properties=species:S:1:pos:R:3 pbc="T T T"
|
||||
Cu 7.19929385 0.07184150 7.15208419
|
||||
Cu 7.14664045 2.10236709 2.01651087
|
||||
Cu 1.82143618 0.04226192 1.92285339
|
||||
Cu 1.99446521 1.95857840 7.03556683
|
||||
Cu 0.04124875 0.03433693 3.81793753
|
||||
Cu 0.13296440 1.50725440 5.36672362
|
||||
Cu 2.05785380 7.16421454 5.34153878
|
||||
Cu 1.87904775 2.29484159 3.46181587
|
||||
Cu 7.14343690 3.63361819 0.04539203
|
||||
Cu 0.07856581 5.42264104 2.00907147
|
||||
Cu 1.70046840 3.49032697 1.45196525
|
||||
Cu 1.52838588 5.29338639 0.08402179
|
||||
Cu 7.04010983 3.63297407 3.45547313
|
||||
Cu 0.04993241 5.06858718 5.39256856
|
||||
Cu 1.57620067 3.46938961 5.22644546
|
||||
Cu 1.85045246 5.47919762 3.50191702
|
||||
Cu 3.66469285 0.20246133 0.01048150
|
||||
Cu 3.65200112 1.80572076 1.95822174
|
||||
Cu 5.62157919 7.09211077 1.57513403
|
||||
Cu 5.42582769 1.92125447 7.13092135
|
||||
Cu 3.74438701 7.22849521 3.62250140
|
||||
Cu 3.71553234 1.93544475 5.27911967
|
||||
Cu 5.41897600 6.88436492 5.32462967
|
||||
Cu 5.23975470 1.60760854 3.77619340
|
||||
Cu 3.72354623 3.71850028 0.15023152
|
||||
Cu 3.53953689 5.32915887 1.66932471
|
||||
Cu 5.31356798 3.64843433 1.81519742
|
||||
Cu 5.24884208 5.54500604 0.06504144
|
||||
Cu 3.76661054 3.88873128 3.46537226
|
||||
Cu 3.74258866 5.40276336 5.55936212
|
||||
Cu 5.45073160 3.94041922 5.40526077
|
||||
Cu 5.72305460 5.42694152 3.73428797
|
||||
41
tests/Fist/regtest-mace/mace_test.inp
Normal file
41
tests/Fist/regtest-mace/mace_test.inp
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
&GLOBAL
|
||||
PRINT_LEVEL LOW
|
||||
PROJECT mace_test
|
||||
RUN_TYPE ENERGY_FORCE
|
||||
&END GLOBAL
|
||||
|
||||
&FORCE_EVAL
|
||||
METHOD FIST
|
||||
#STRESS_TENSOR ANALYTICAL
|
||||
&MM
|
||||
&FORCEFIELD
|
||||
&NONBONDED
|
||||
&MACE
|
||||
ATOMS Cu
|
||||
POT_FILE_NAME MACE/MACE_scratch_run-3.model-cp2k.pth
|
||||
&END MACE
|
||||
&END NONBONDED
|
||||
&END FORCEFIELD
|
||||
&POISSON
|
||||
&EWALD
|
||||
EWALD_TYPE none
|
||||
&END EWALD
|
||||
&END POISSON
|
||||
&END MM
|
||||
&PRINT
|
||||
&FORCES ON
|
||||
&END FORCES
|
||||
#&STRESS_TENSOR ON
|
||||
#&END STRESS_TENSOR
|
||||
&END PRINT
|
||||
&SUBSYS
|
||||
&CELL
|
||||
CELL_FILE_FORMAT XYZ
|
||||
CELL_FILE_NAME ./mace-test.xyz
|
||||
&END CELL
|
||||
&TOPOLOGY
|
||||
COORD_FILE_FORMAT XYZ
|
||||
COORD_FILE_NAME ./mace-test.xyz
|
||||
&END TOPOLOGY
|
||||
&END SUBSYS
|
||||
&END FORCE_EVAL
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
# Directories have been reordered according the execution time needed for a gfortran pdbg run using 2 MPI tasks
|
||||
# in case a new directory is added just add it at the top of the list..
|
||||
# the order will be regularly checked and modified...
|
||||
Fist/regtest-mace libtorch
|
||||
QS/regtest-floquet
|
||||
QS/regtest-dft-d4-auto-ref libdftd4 libint !ifx
|
||||
QS/regtest-dft-d2-auto-ref libint !ifx
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ run_test ./tools/pao-ml/pao-validate.py --threshold=1e-5 --model="tests/QS/regte
|
|||
run_test ./tools/vibronic_spec/main.py ./tools/vibronic_spec/example/example_config.toml
|
||||
|
||||
run_test mypy --strict ./tools/pao-ml/
|
||||
run_test mypy --strict ./tools/mace/create_cp2k_model.py
|
||||
run_test mypy --strict ./tools/minimax_tools/minimax_to_fortran_source.py
|
||||
run_test mypy --strict ./tools/dashboard/generate_dashboard.py
|
||||
run_test mypy --strict ./tools/regtesting/optimize_test_dirs.py
|
||||
|
|
|
|||
190
tools/mace/create_cp2k_model.py
Normal file
190
tools/mace/create_cp2k_model.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
# pylint: disable=wrong-import-position
|
||||
"""
|
||||
Export a trained MACE model to a TorchScript file that can be loaded by CP2K.
|
||||
|
||||
The export script must be executed in an environment where MACE is installed.
|
||||
|
||||
Required packages:
|
||||
* torch: required for loading and converting the model.
|
||||
* e3nn: required by MACE for JIT compilation.
|
||||
* mace-torch: required because the trained MACE model is loaded using
|
||||
torch and its original model classes must be available.
|
||||
|
||||
The MACE version used for export should be compatible with the version
|
||||
used for training the model.
|
||||
|
||||
After exporting, the CP2K-compatible `.pth` file can be loaded by CP2K
|
||||
through the LibTorch interface. No MACE, e3nn, or Python installation is
|
||||
required during CP2K simulations.
|
||||
|
||||
Usage:
|
||||
python create_cp2k_model.py my_mace.model --dtype float64 --head default
|
||||
-> writes my_mace.model-cp2k.pth
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
|
||||
|
||||
import torch
|
||||
from e3nn.util import jit # type: ignore
|
||||
from e3nn.util.jit import compile_mode # type: ignore
|
||||
|
||||
# Z -> chemical symbol table (index == atomic number). Covers the dummy X plus
|
||||
# all 118 known elements, matching CP2K's src/common/periodic_table.F.
|
||||
chemical_symbols = ["X"] + """
|
||||
H He Li Be B C N O F Ne Na Mg Al Si P S Cl Ar K Ca Sc Ti V Cr Mn Fe Co
|
||||
Ni Cu Zn Ga Ge As Se Br Kr Rb Sr Y Zr Nb Mo Tc Ru Rh Pd Ag Cd In Sn Sb
|
||||
Te I Xe Cs Ba La Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Lu Hf Ta W Re Os
|
||||
Ir Pt Au Hg Tl Pb Bi Po At Rn Fr Ra Ac Th Pa U Np Pu Am Cm Bk Cf Es Fm
|
||||
Md No Lr Rf Db Sg Bh Hs Mt Ds Rg Cn Nh Fl Mc Lv Ts Og
|
||||
""".split()
|
||||
|
||||
|
||||
@compile_mode("script")
|
||||
class CP2K_MACE(torch.nn.Module):
|
||||
"""MACE model wrapped for CP2K's single-dictionary libtorch interface."""
|
||||
|
||||
head: torch.Tensor
|
||||
|
||||
def __init__(self, model: Any, head: Optional[str] = None) -> None:
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.register_buffer("atomic_numbers", model.atomic_numbers)
|
||||
self.register_buffer("r_max", model.r_max)
|
||||
self.register_buffer("num_interactions", model.num_interactions)
|
||||
self.num_types: int = int(model.atomic_numbers.shape[0])
|
||||
|
||||
if not hasattr(model, "heads"):
|
||||
model.heads = [None]
|
||||
if head is not None and head in model.heads:
|
||||
head_idx = model.heads.index(head)
|
||||
else:
|
||||
head_idx = len(model.heads) - 1
|
||||
self.register_buffer(
|
||||
"head", torch.tensor(head_idx, dtype=torch.long).unsqueeze(0)
|
||||
)
|
||||
|
||||
for param in self.model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(
|
||||
self, data: Dict[str, torch.Tensor]
|
||||
) -> Dict[str, Optional[torch.Tensor]]:
|
||||
pos = data["pos"]
|
||||
edge_index = data["edge_index"]
|
||||
unit_shifts = data["edge_cell_shift"]
|
||||
cell = data["cell"].view(3, 3)
|
||||
atom_types = data["atom_types"].to(torch.long)
|
||||
n_atoms = pos.shape[0]
|
||||
|
||||
# one-hot node attributes expected by MACE
|
||||
node_attrs = torch.nn.functional.one_hot(
|
||||
atom_types, num_classes=self.num_types
|
||||
).to(pos.dtype)
|
||||
|
||||
# real-space edge shift vectors: unit_shifts @ cell (Angstrom)
|
||||
shifts = torch.matmul(unit_shifts, cell)
|
||||
|
||||
batch = torch.zeros(n_atoms, dtype=torch.long, device=pos.device)
|
||||
ptr = torch.tensor([0, n_atoms], dtype=torch.long, device=pos.device)
|
||||
|
||||
mace_in: Dict[str, torch.Tensor] = {
|
||||
"positions": pos,
|
||||
"node_attrs": node_attrs,
|
||||
"edge_index": edge_index,
|
||||
"shifts": shifts,
|
||||
"unit_shifts": unit_shifts,
|
||||
"cell": cell,
|
||||
"batch": batch,
|
||||
"ptr": ptr,
|
||||
"head": self.head,
|
||||
}
|
||||
|
||||
out = self.model(
|
||||
mace_in,
|
||||
training=False,
|
||||
compute_force=True,
|
||||
compute_virials=True,
|
||||
compute_stress=False,
|
||||
compute_displacement=True,
|
||||
)
|
||||
|
||||
node_energy = out["node_energy"]
|
||||
forces = out["forces"]
|
||||
virials = out["virials"]
|
||||
|
||||
atomic_energy: Optional[torch.Tensor] = None
|
||||
if node_energy is not None:
|
||||
atomic_energy = node_energy.unsqueeze(-1)
|
||||
|
||||
virial: Optional[torch.Tensor] = None
|
||||
if virials is not None:
|
||||
virial = virials.view(3, 3)
|
||||
|
||||
return {
|
||||
"atomic_energy": atomic_energy,
|
||||
"forces": forces,
|
||||
"virial": virial,
|
||||
}
|
||||
|
||||
|
||||
def build_metadata(model: Any, dtype: str) -> Dict[str, str]:
|
||||
z = model.atomic_numbers.to(torch.long).tolist()
|
||||
type_names = " ".join(chemical_symbols[int(zi)] for zi in z)
|
||||
return {
|
||||
"num_types": str(len(z)),
|
||||
"r_max": repr(float(model.r_max)),
|
||||
"type_names": type_names,
|
||||
"model_dtype": dtype,
|
||||
"allow_tf32": "0",
|
||||
# MACE uses a single global cutoff; leave per-edge-type cutoffs empty so
|
||||
# CP2K's reader falls back to filling the cutoff matrix with r_max**2.
|
||||
"per_edge_type_cutoff": "",
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("model_path", help="Path to the trained MACE .model file")
|
||||
p.add_argument(
|
||||
"--head", default=None, help="Model head to export (multi-head models)"
|
||||
)
|
||||
p.add_argument("--dtype", choices=["float64", "float32"], default="float64")
|
||||
p.add_argument(
|
||||
"--output", default=None, help="Output path (default: <model>-cp2k.pth)"
|
||||
)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
model = torch.load(args.model_path, map_location="cpu")
|
||||
if args.dtype == "float64":
|
||||
model = model.double()
|
||||
else:
|
||||
print("Converting model to float32, this may cause loss of precision.")
|
||||
model = model.float()
|
||||
model = model.to("cpu")
|
||||
|
||||
wrapper = CP2K_MACE(model, head=args.head)
|
||||
wrapper.eval()
|
||||
|
||||
scripted = jit.compile(wrapper)
|
||||
|
||||
metadata = build_metadata(model, args.dtype)
|
||||
extra_files = {k: v.encode("utf-8") for k, v in metadata.items()}
|
||||
|
||||
out_path = args.output or (args.model_path + "-cp2k.pth")
|
||||
scripted.save(out_path, _extra_files=extra_files)
|
||||
|
||||
print(f"Wrote CP2K MACE model to: {out_path}")
|
||||
print("Metadata:")
|
||||
for k, v in metadata.items():
|
||||
print(f" {k}: {v}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue