diff --git a/data/MACE/MACE_scratch_run-3.model-cp2k.pth b/data/MACE/MACE_scratch_run-3.model-cp2k.pth new file mode 100644 index 0000000000..f0ff697cb9 Binary files /dev/null and b/data/MACE/MACE_scratch_run-3.model-cp2k.pth differ diff --git a/docs/methods/machine_learning/index.md b/docs/methods/machine_learning/index.md index 004f3c7a6e..f3ae029b1e 100644 --- a/docs/methods/machine_learning/index.md +++ b/docs/methods/machine_learning/index.md @@ -6,6 +6,7 @@ titlesonly: maxdepth: 2 --- nequip +mace nnp pao-ml deepmd diff --git a/docs/methods/machine_learning/mace.md b/docs/methods/machine_learning/mace.md new file mode 100644 index 0000000000..7e3bf66da5 --- /dev/null +++ b/docs/methods/machine_learning/mace.md @@ -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). diff --git a/docs/technologies/libraries.md b/docs/technologies/libraries.md index 2e5bf4ebbf..86a2cb3451 100644 --- a/docs/technologies/libraries.md +++ b/docs/technologies/libraries.md @@ -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/). diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 89417a3cf5..a016d826a6 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -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 diff --git a/src/common/bibliography.F b/src/common/bibliography.F index 621576c26d..38ca082835 100644 --- a/src/common/bibliography.F +++ b/src/common/bibliography.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", & diff --git a/src/fist_neighbor_list_control.F b/src/fist_neighbor_list_control.F index f16cc70787..0e23752054 100644 --- a/src/fist_neighbor_list_control.F +++ b/src/fist_neighbor_list_control.F @@ -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 diff --git a/src/fist_nonbond_env_types.F b/src/fist_nonbond_env_types.F index 26360ee060..045df06365 100644 --- a/src/fist_nonbond_env_types.F +++ b/src/fist_nonbond_env_types.F @@ -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 diff --git a/src/fist_nonbond_force.F b/src/fist_nonbond_force.F index 6de5a4894d..0d6680b31c 100644 --- a/src/fist_nonbond_force.F +++ b/src/fist_nonbond_force.F @@ -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 diff --git a/src/force_fields_all.F b/src/force_fields_all.F index 6f9106752b..798fe280b8 100644 --- a/src/force_fields_all.F +++ b/src/force_fields_all.F @@ -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 diff --git a/src/force_fields_input.F b/src/force_fields_input.F index 89bcf55132..1df71ce40d 100644 --- a/src/force_fields_input.F +++ b/src/force_fields_input.F @@ -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(:)) diff --git a/src/input_cp2k_mm.F b/src/input_cp2k_mm.F index aeb0be36a1..4938f9cfea 100644 --- a/src/input_cp2k_mm.F +++ b/src/input_cp2k_mm.F @@ -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 .", & + 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 diff --git a/src/manybody_nequip.F b/src/manybody_e3nn.F similarity index 94% rename from src/manybody_nequip.F rename to src/manybody_e3nn.F index 9938d0b82b..58e5dc64c6 100644 --- a/src/manybody_nequip.F +++ b/src/manybody_e3nn.F @@ -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,10 +93,10 @@ 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, & - atomic_kind_set, potparm, r_last_update_pbc, & - pot_total, fist_nonbond_env, para_env, use_virial, & - target_pot_type) + 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) TYPE(fist_neighbor_type), POINTER :: nonbonded TYPE(particle_type), POINTER :: particle_set(:) @@ -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 diff --git a/src/manybody_potential.F b/src/manybody_potential.F index 9de5d3949d..a3a706a7f1 100644 --- a/src/manybody_potential.F +++ b/src/manybody_potential.F @@ -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,21 +192,30 @@ 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, & - atomic_kind_set, potparm, r_last_update_pbc, & - pot_nequip, fist_nonbond_env, & - para_env, use_virial, nequip_type) + 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) pot_manybody = pot_manybody + pot_nequip END IF ! 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, & - atomic_kind_set, potparm, r_last_update_pbc, & - pot_allegro, fist_nonbond_env, & - para_env, use_virial, allegro_type) + 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 diff --git a/src/pair_potential.F b/src/pair_potential.F index 6e2d4af4f5..3f56b1de2c 100644 --- a/src/pair_potential.F +++ b/src/pair_potential.F @@ -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) diff --git a/src/pair_potential_types.F b/src/pair_potential_types.F index 413632847a..23aacc5225 100644 --- a/src/pair_potential_types.F +++ b/src/pair_potential_types.F @@ -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) diff --git a/tests/Fist/regtest-mace/TEST_FILES.toml b/tests/Fist/regtest-mace/TEST_FILES.toml new file mode 100644 index 0000000000..4e86987486 --- /dev/null +++ b/tests/Fist/regtest-mace/TEST_FILES.toml @@ -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 \ No newline at end of file diff --git a/tests/Fist/regtest-mace/mace-test.xyz b/tests/Fist/regtest-mace/mace-test.xyz new file mode 100644 index 0000000000..de5d503ae3 --- /dev/null +++ b/tests/Fist/regtest-mace/mace-test.xyz @@ -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 diff --git a/tests/Fist/regtest-mace/mace_test.inp b/tests/Fist/regtest-mace/mace_test.inp new file mode 100644 index 0000000000..feabc680c5 --- /dev/null +++ b/tests/Fist/regtest-mace/mace_test.inp @@ -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 diff --git a/tests/TEST_DIRS b/tests/TEST_DIRS index b20afb6390..4bcedaf889 100644 --- a/tests/TEST_DIRS +++ b/tests/TEST_DIRS @@ -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 diff --git a/tools/docker/scripts/test_misc.sh b/tools/docker/scripts/test_misc.sh index fe39415cc9..d3fc3dc645 100755 --- a/tools/docker/scripts/test_misc.sh +++ b/tools/docker/scripts/test_misc.sh @@ -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 diff --git a/tools/mace/create_cp2k_model.py b/tools/mace/create_cp2k_model.py new file mode 100644 index 0000000000..16b097b0ce --- /dev/null +++ b/tools/mace/create_cp2k_model.py @@ -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: -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()