(git:b279b6b)
nequip_unittest.F
Go to the documentation of this file.
1 !--------------------------------------------------------------------------------------------------!
2 ! CP2K: A general program to perform molecular dynamics simulations !
3 ! Copyright 2000-2024 CP2K developers group <https://cp2k.org> !
4 ! !
5 ! SPDX-License-Identifier: GPL-2.0-or-later !
6 !--------------------------------------------------------------------------------------------------!
7 
8 PROGRAM nequip_unittest
9 
10  USE cp_files, ONLY: discover_file
11  USE kinds, ONLY: default_path_length,&
12  dp,&
13  int_8,&
14  sp
15  USE mathlib, ONLY: inv_3x3
16  USE physcon, ONLY: angstrom,&
17  evolt
18  USE torch_api, ONLY: &
19  torch_cuda_is_available, torch_dict_create, torch_dict_get, torch_dict_insert, &
20  torch_dict_release, torch_dict_type, torch_model_eval, torch_model_load, &
21  torch_model_read_metadata, torch_model_release, torch_model_type
22 #include "./base/base_uses.f90"
23 
24  IMPLICIT NONE
25 
26  CHARACTER(LEN=default_path_length) :: filename, cutoff_str, nequip_version
27  REAL(dp) :: cutoff
28 
29  ! Inputs.
30  INTEGER, PARAMETER :: natoms = 96
31  INTEGER :: iatom, nedges
32  REAL(sp), DIMENSION(3, natoms) :: pos
33  REAL(dp), DIMENSION(3, 3):: cell, hinv
34  INTEGER(kind=int_8), DIMENSION(natoms):: atom_types
35  INTEGER(kind=int_8), DIMENSION(:, :), ALLOCATABLE:: edge_index
36  REAL(sp), DIMENSION(:, :), ALLOCATABLE:: edge_cell_shift
37 
38  ! Torch objects.
39  TYPE(torch_model_type) :: model
40  TYPE(torch_dict_type) :: inputs, outputs
41 
42  ! Outputs.
43  REAL(sp), DIMENSION(:, :), POINTER :: total_energy, atomic_energy, forces
44  NULLIFY (total_energy, atomic_energy, forces)
45 
46  ! A box with 32 water molecules.
47  pos = reshape(real([ &
48  42.8861696_dp, -0.0556816_dp, 38.3291611_dp, &
49  34.2025887_dp, -0.6185484_dp, 37.3655680_dp, &
50  30.0803925_dp, -2.0124176_dp, 36.4807960_dp, &
51  28.7057911_dp, -2.6880392_dp, 36.6020983_dp, &
52  36.2479426_dp, -0.5163484_dp, 34.4923596_dp, &
53  37.6964724_dp, -0.0410872_dp, 35.0140735_dp, &
54  27.7606699_dp, 7.4854206_dp, 33.9276919_dp, &
55  28.8160999_dp, 6.4985777_dp, 34.2163608_dp, &
56  37.1576372_dp, 9.0188280_dp, 31.9265812_dp, &
57  38.6063816_dp, 9.5820079_dp, 32.3435972_dp, &
58  34.3031959_dp, 2.2195014_dp, 45.9880451_dp, &
59  33.2444139_dp, 1.3025332_dp, 46.4698427_dp, &
60  38.7286174_dp, -5.0541897_dp, 26.0743968_dp, &
61  38.3483921_dp, -6.2832846_dp, 26.9867253_dp, &
62  32.8642520_dp, 3.2060632_dp, 30.8971160_dp, &
63  31.2904088_dp, 3.0871834_dp, 30.6273977_dp, &
64  33.7519869_dp, -3.1383262_dp, 39.6727607_dp, &
65  34.6642979_dp, -3.6643859_dp, 38.6466027_dp, &
66  42.7173214_dp, 5.1246883_dp, 32.5883401_dp, &
67  41.5627455_dp, 5.5893544_dp, 33.4174902_dp, &
68  32.4283800_dp, 9.1182520_dp, 30.5477678_dp, &
69  32.6432407_dp, 10.770683_dp, 30.4842778_dp, &
70  31.4848670_dp, 4.6777144_dp, 37.3957194_dp, &
71  32.3171882_dp, -6.2287496_dp, 36.4671864_dp, &
72  26.6621340_dp, 3.1708123_dp, 35.6820146_dp, &
73  26.5271367_dp, 1.6039040_dp, 35.4883482_dp, &
74  32.0238236_dp, 16.918208_dp, 31.6883569_dp, &
75  31.4006579_dp, 7.0315610_dp, 30.2394554_dp, &
76  33.5264253_dp, -3.5594808_dp, 34.2636830_dp, &
77  34.6404855_dp, -3.2653833_dp, 35.4971482_dp, &
78  40.0564375_dp, -0.3054386_dp, 29.8312074_dp, &
79  39.4784464_dp, -1.0948314_dp, 38.3101140_dp, &
80  39.7040761_dp, 1.9584631_dp, 33.3902375_dp, &
81  38.3338570_dp, 2.6967178_dp, 42.9261945_dp, &
82  40.1820455_dp, -7.2199289_dp, 27.6580390_dp, &
83  39.3204431_dp, -8.4564252_dp, 28.1319658_dp, &
84  36.3876963_dp, 8.8117085_dp, 38.3545362_dp, &
85  36.3205637_dp, 9.0063075_dp, 36.7526001_dp, &
86  29.9991583_dp, -5.5637817_dp, 33.9295050_dp, &
87  30.7728545_dp, -5.0385870_dp, 35.1998067_dp, &
88  40.0592517_dp, 6.3305279_dp, 28.2579461_dp, &
89  40.2398360_dp, 5.1745923_dp, 29.2962956_dp, &
90  26.3320911_dp, 2.4393638_dp, 33.5653868_dp, &
91  26.9606971_dp, 1.2711078_dp, 32.5923884_dp, &
92  34.8372697_dp, -0.4722708_dp, 30.3824362_dp, &
93  35.3968813_dp, -1.9268483_dp, 30.3081837_dp, &
94  32.1217607_dp, -0.7333429_dp, 36.5104382_dp, &
95  32.2180843_dp, 7.8454304_dp, 35.6671967_dp, &
96  36.3780998_dp, -4.3048878_dp, 36.4539793_dp, &
97  35.8119275_dp, -3.0013928_dp, 27.0348937_dp, &
98  29.6452491_dp, 1.0652123_dp, 35.7143653_dp, &
99  30.3794654_dp, -0.0668146_dp, 34.9882468_dp, &
100  34.2149336_dp, -1.6559120_dp, 33.8876437_dp, &
101  34.7842435_dp, -1.0252141_dp, 32.5034832_dp, &
102  40.4649954_dp, 1.1467825_dp, 31.3073503_dp, &
103  41.3262469_dp, 0.6550803_dp, 32.4555882_dp, &
104  29.0210859_dp, 3.5038194_dp, 39.9087702_dp, &
105  29.4945426_dp, 3.7276637_dp, 41.3766138_dp, &
106  34.1359664_dp, -6.7533422_dp, 32.3568410_dp, &
107  34.9546570_dp, -5.7704242_dp, 31.4571066_dp, &
108  33.2532356_dp, 1.5268048_dp, 44.0562171_dp, &
109  33.7931669_dp, 0.5014632_dp, 43.0597590_dp, &
110  36.8205409_dp, 2.6214681_dp, 40.6834006_dp, &
111  37.5552706_dp, 1.5649832_dp, 39.7648935_dp, &
112  43.2099087_dp, -0.0628456_dp, 47.2593155_dp, &
113  29.3940583_dp, -2.3133019_dp, 37.1407883_dp, &
114  36.7415708_dp, -0.0838710_dp, 35.2591783_dp, &
115  27.9424776_dp, 6.7622961_dp, 34.5648384_dp, &
116  37.6812656_dp, 9.4216399_dp, 32.6478643_dp, &
117  33.3171290_dp, 2.0951401_dp, 45.8722265_dp, &
118  37.9951355_dp, 4.3611431_dp, 26.5571819_dp, &
119  32.1824670_dp, 2.6611503_dp, 30.4577248_dp, &
120  34.6538012_dp, -3.4374573_dp, 39.5889245_dp, &
121  42.2929833_dp, 5.9471069_dp, 32.8460995_dp, &
122  32.9604690_dp, 9.9050313_dp, 30.1587306_dp, &
123  31.4281886_dp, -5.8338304_dp, 36.6738743_dp, &
124  26.0563730_dp, 2.4973869_dp, 35.3486870_dp, &
125  32.0334927_dp, 17.3252289_dp, 30.8116013_dp, &
126  33.8252182_dp, -2.9520949_dp, 35.0220460_dp, &
127  39.4569981_dp, -0.3072759_dp, 38.9347829_dp, &
128  29.4846708_dp, 2.8692561_dp, 43.0061868_dp, &
129  39.2864184_dp, -7.6206103_dp, 27.6271147_dp, &
130  35.8797502_dp, 8.6515870_dp, 37.5221734_dp, &
131  30.3582543_dp, -4.7607656_dp, 34.3355645_dp, &
132  40.7098956_dp, 5.8331250_dp, 28.7558375_dp, &
133  26.7179083_dp, 2.2415138_dp, 32.6577297_dp, &
134  35.6589256_dp, -0.9968903_dp, 30.5749530_dp, &
135  31.5851602_dp, -1.3121804_dp, 35.9011109_dp, &
136  35.5489386_dp, -3.9056138_dp, 26.8214490_dp, &
137  29.5656616_dp, 0.4681794_dp, 34.9670711_dp, &
138  34.7615128_dp, -0.9569680_dp, 33.4891367_dp, &
139  40.4853406_dp, 0.4023620_dp, 31.9425416_dp, &
140  29.6728289_dp, 4.0134825_dp, 40.4505780_dp, &
141  34.1272286_dp, -5.8796882_dp, 31.8925999_dp, &
142  33.1168884_dp, 1.2338084_dp, 43.1127997_dp, &
143  37.1996993_dp, 2.5049007_dp, 39.7917126_dp], kind=sp), shape=[3, natoms])
144 
145  cell(1, :) = [9.85_dp, 0.0_dp, 0.0_dp]
146  cell(2, :) = [0.0_dp, 9.85_dp, 0.0_dp]
147  cell(3, :) = [0.0_dp, 0.0_dp, 9.85_dp]
148 
149  hinv = inv_3x3(cell)
150 
151  atom_types(:64) = 0 ! Hydrogen
152  atom_types(65:) = 1 ! Oxygen
153 
154  WRITE (*, *) "CUDA is available: ", torch_cuda_is_available()
155 
156  filename = discover_file('NequIP/water.pth')
157  WRITE (*, *) "Loading NequIP model from: "//trim(filename)
158  CALL torch_model_load(model, filename)
159  cutoff_str = torch_model_read_metadata(filename, "r_max")
160  nequip_version = torch_model_read_metadata(filename, "nequip_version")
161  READ (cutoff_str, *) cutoff
162  WRITE (*, *) "Version: ", trim(nequip_version)
163  WRITE (*, *) "Cutoff: ", cutoff
164 
165  CALL neighbor_search(nedges)
166  ALLOCATE (edge_index(nedges, 2))
167  ALLOCATE (edge_cell_shift(3, nedges))
168  CALL neighbor_search(nedges, edge_index, edge_cell_shift)
169  WRITE (*, *) "Found", nedges, "neighbor edges between", natoms, "atoms."
170 
171  CALL torch_dict_create(inputs)
172  CALL torch_dict_insert(inputs, "pos", pos)
173  CALL torch_dict_insert(inputs, "edge_index", edge_index)
174  CALL torch_dict_insert(inputs, "edge_cell_shift", edge_cell_shift)
175  CALL torch_dict_insert(inputs, "cell", real(cell, kind=sp))
176  CALL torch_dict_insert(inputs, "atom_types", atom_types)
177 
178  CALL torch_dict_create(outputs)
179  CALL torch_model_eval(model, inputs, outputs)
180 
181  CALL torch_dict_get(outputs, "total_energy", total_energy)
182  CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
183  CALL torch_dict_get(outputs, "forces", forces)
184  WRITE (*, *) "Total Energy [Hartree] : ", total_energy(1, 1)/evolt
185  WRITE (*, *) "FORCES: [Hartree/Bohr]: "
186  DO iatom = 1, natoms
187  WRITE (*, *) forces(:, iatom)*angstrom/evolt
188  END DO
189  cpassert(abs(-14985.615_dp - real(total_energy(1, 1), kind=dp)) < 2e-3_dp)
190 
191  CALL torch_dict_release(inputs)
192  CALL torch_dict_release(outputs)
193  CALL torch_model_release(model)
194  DEALLOCATE (edge_index, edge_cell_shift, total_energy, atomic_energy, forces)
195 
196  WRITE (*, *) "NequIP unittest was successfully :-)"
197 
198 CONTAINS
199 
200 ! **************************************************************************************************
201 !> \brief Naive neighbor search - beware it scales O(N**2).
202 !> \param nedges ...
203 !> \param edge_index ...
204 !> \param edge_cell_shift ...
205 ! **************************************************************************************************
206  SUBROUTINE neighbor_search(nedges, edge_index, edge_cell_shift)
207  INTEGER, INTENT(OUT) :: nedges
208  INTEGER(kind=int_8), DIMENSION(:, :), &
209  INTENT(OUT), OPTIONAL :: edge_index
210  REAL(sp), DIMENSION(:, :), INTENT(OUT), OPTIONAL :: edge_cell_shift
211 
212  INTEGER:: iatom, jatom
213  REAL(dp), DIMENSION(3) :: s1, s2, s12, cell_shift, dx
214 
215  nedges = 0
216  DO iatom = 1, natoms
217  DO jatom = 1, natoms
218  IF (iatom == jatom) cycle
219  s1 = matmul(hinv, pos(:, iatom))
220  s2 = matmul(hinv, pos(:, jatom))
221  s12 = s1 - s2
222  cell_shift = anint(s12)
223  dx = matmul(cell, s12 - cell_shift)
224  IF (dot_product(dx, dx) <= cutoff**2) THEN
225  nedges = nedges + 1
226  IF (PRESENT(edge_index)) THEN
227  edge_index(nedges, :) = [iatom - 1, jatom - 1]
228  END IF
229  IF (PRESENT(edge_cell_shift)) THEN
230  edge_cell_shift(:, nedges) = real(cell_shift, kind=sp)
231  END IF
232  END IF
233  END DO
234  END DO
235  END SUBROUTINE neighbor_search
236 
237 END PROGRAM nequip_unittest
atom_types
Define the atom type and its sub types.
Definition: atom_types.F:15
cp_files
Utility routines to open and close files. Tracking of preconnections.
Definition: cp_files.F:16
cp_files::discover_file
character(len=default_path_length) function, public discover_file(file_name)
Checks various locations for a file name.
Definition: cp_files.F:510
kinds
Defines the basic variable types.
Definition: kinds.F:23
kinds::int_8
integer, parameter, public int_8
Definition: kinds.F:54
kinds::dp
integer, parameter, public dp
Definition: kinds.F:34
kinds::default_path_length
integer, parameter, public default_path_length
Definition: kinds.F:58
kinds::sp
integer, parameter, public sp
Definition: kinds.F:33
mathlib
Collection of simple mathematical functions and subroutines.
Definition: mathlib.F:15
mathlib::inv_3x3
pure real(kind=dp) function, dimension(3, 3), public inv_3x3(a)
Returns the inverse of the 3 x 3 matrix a.
Definition: mathlib.F:516
physcon
Definition of physical constants:
Definition: physcon.F:68
physcon::evolt
real(kind=dp), parameter, public evolt
Definition: physcon.F:183
physcon::angstrom
real(kind=dp), parameter, public angstrom
Definition: physcon.F:144
torch_api
Definition: torch_api.F:7
torch_api::torch_dict_release
subroutine, public torch_dict_release(dict)
Releases a Torch dictionary and all its ressources.
Definition: torch_api.F:920
torch_api::torch_model_load
subroutine, public torch_model_load(model, filename)
Loads a Torch model from given "*.pth" file. (In Torch lingo models are called modules)
Definition: torch_api.F:944
torch_api::torch_dict_create
subroutine, public torch_dict_create(dict)
Creates an empty Torch dictionary.
Definition: torch_api.F:896
torch_api::torch_model_release
subroutine, public torch_model_release(model)
Releases a Torch model and all its ressources.
Definition: torch_api.F:1004
torch_api::torch_model_eval
subroutine, public torch_model_eval(model, inputs, outputs)
Evaluates the given Torch model. (In Torch lingo this operation is called forward())
Definition: torch_api.F:971
torch_api::torch_model_read_metadata
character(:) function, allocatable, public torch_model_read_metadata(filename, key)
Reads metadata entry from given "*.pth" file. (In Torch lingo they are called extra files)
Definition: torch_api.F:1028
torch_api::torch_cuda_is_available
logical function, public torch_cuda_is_available()
Returns true iff the Torch CUDA backend is available.
Definition: torch_api.F:1080
nequip_unittest
program nequip_unittest
Definition: nequip_unittest.F:8
neighbor_search
subroutine neighbor_search(nedges, edge_index, edge_cell_shift)
Naive neighbor search - beware it scales O(N**2).
Definition: nequip_unittest.F:207