From d987b39481e158343fbb9845e6e92681df31f7ee Mon Sep 17 00:00:00 2001 From: xysun25 Date: Fri, 10 Jul 2026 18:30:32 +0200 Subject: [PATCH] Add MACE support to TorchScript machine learning interface --- src/manybody_potential.F | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/src/manybody_potential.F b/src/manybody_potential.F index 9de5d3949d..b03c1b0550 100644 --- a/src/manybody_potential.F +++ b/src/manybody_potential.F @@ -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)