FIST: Fix CMake build and conventions for NequIP

cp2k · Jan 29, 2023 · 8f18a4e · 8f18a4e
1 parent 1ed42ca
commit 8f18a4e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -296,6 +296,7 @@ list(
   manybody_eam.F
   manybody_gal21.F
   manybody_gal.F
+  manybody_nequip.F
   manybody_potential.F
   manybody_quip.F
   manybody_siepmann.F
@@ -750,6 +751,7 @@ list(
   topology_util.F
   topology_xtl.F
   topology_xyz.F
+  torch_api.F
   transport_env_types.F
   transport.F
   uff_vdw_radii_table.F
@@ -1512,7 +1514,8 @@ list(
   "graph"
   "dumpdcd"
   "xyz2dcd"
-  "libcp2k_unittest")
+  "libcp2k_unittest"
+  "nequip_unittest")
 
 add_executable(cp2k-bin start/cp2k.F)
 set_target_properties(
@@ -1526,6 +1529,7 @@ add_executable(graph metadyn_tools/graph.F)
 add_executable(dumpdcd motion/dumpdcd.F)
 add_executable(xyz2dcd motion/xyz2dcd.F)
 add_executable(libcp2k_unittest start/libcp2k_unittest.c)
+add_executable(nequip_unittest nequip_unittest.F)
 
 set_target_properties(cp2k-bin PROPERTIES CXX_STANDARD 11)
 set_target_properties(cp2k-bin PROPERTIES IMPORTED_LINK_INTERFACE_LANGUAGES

diff --git a/src/force_fields_input.F b/src/force_fields_input.F
@@ -2317,6 +2317,9 @@ SUBROUTINE read_nequip_data(nequip)
       CHARACTER(len=*), PARAMETER                        :: routineN = 'read_nequip_data'
 
       CHARACTER(LEN=default_path_length)                 :: cutoff_str
+      INTEGER                                            :: handle
+
+      CALL timeset(routineN, handle)
 
       nequip%nequip_version = torch_model_read_metadata(nequip%nequip_file_name, "nequip_version")
       cutoff_str = torch_model_read_metadata(nequip%nequip_file_name, "r_max")
@@ -2332,6 +2335,7 @@ SUBROUTINE read_nequip_data(nequip)
       nequip%unit_cell_val = 1.0_dp
       nequip%unit_cell_val = cp_unit_to_cp2k(nequip%unit_cell_val, nequip%unit_cell)
 
+      CALL timestop(handle)
    END SUBROUTINE read_nequip_data
 
 ! **************************************************************************************************

diff --git a/src/manybody_nequip.F b/src/manybody_nequip.F
@@ -218,8 +218,8 @@ SUBROUTINE nequip_energy_store_force_virial(nonbonded, particle_set, cell, atomi
 
       CHARACTER(LEN=*), PARAMETER :: routineN = 'nequip_energy_store_force_virial'
 
-      INTEGER :: atom_a, atom_b, i, iat, iat_use, iend, ifirst, igrp, ikind, ilast, ilist, ipair, &
-         istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, nedges, nedges_tot, &
+      INTEGER :: atom_a, atom_b, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, ilast, ilist, &
+         ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, nedges, nedges_tot, &
          nloc_size, npairs, nunique
       INTEGER(kind=int_8), ALLOCATABLE                   :: atom_types(:)
       INTEGER(kind=int_8), ALLOCATABLE, DIMENSION(:, :)  :: edge_index, t_edge_index, temp_edge_index
@@ -238,6 +238,8 @@ SUBROUTINE nequip_energy_store_force_virial(nonbonded, particle_set, cell, atomi
       TYPE(pair_potential_single_type), POINTER          :: pot
       TYPE(torch_dict_type)                              :: inputs, outputs
 
+      CALL timeset(routineN, handle)
+
       NULLIFY (total_energy, atomic_energy, forces)
       n_atoms = SIZE(particle_set)
       ALLOCATE (use_atom(n_atoms))
@@ -438,6 +440,7 @@ SUBROUTINE nequip_energy_store_force_virial(nonbonded, particle_set, cell, atomi
 
       DEALLOCATE (forces)
 
+      CALL timestop(handle)
    END SUBROUTINE nequip_energy_store_force_virial
 
 ! **************************************************************************************************
@@ -453,8 +456,6 @@ SUBROUTINE nequip_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_
       REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT)      :: f_nonbond, pv_nonbond
       LOGICAL, INTENT(IN)                                :: use_virial
 
-      CHARACTER(LEN=*), PARAMETER :: routineN = 'nequip_add_force_virial'
-
       INTEGER                                            :: iat, iat_use
       REAL(KIND=dp), DIMENSION(3, 3)                     :: virial
       TYPE(nequip_data_type), POINTER                    :: nequip_data