comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-29 06:35:28 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add MACE support to TorchScript machine learning interface

Browse source
This commit is contained in:
xysun25 2026-07-10 18:30:32 +02:00
parent d224794830
commit d987b39481
1 changed files with 27 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

32
src/manybody_potential.F
View file

@ -54,7 +54,7 @@ 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, &
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
@ -105,11 +105,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_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 +132,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 +178,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)
@ -204,6 +206,15 @@ CONTAINS
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 nequip_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,7 +589,7 @@ 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_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)
@ -599,6 +610,7 @@ CONTAINS
any_tersoff = .FALSE.
any_allegro = .FALSE.
any_nequip = .FALSE.
any_mace = .FALSE.
any_siepmann = .FALSE.
any_ace = .FALSE.
any_deepmd = .FALSE.
@ -673,6 +685,16 @@ CONTAINS
CALL nequip_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 nequip_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
END IF
! starting the force loop
DO ilist = 1, nonbonded%nlists
neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)