Line data Source code
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 218882 : PROGRAM nequip_unittest
9 :
10 2 : 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 2 : REAL(sp), DIMENSION(:, :), ALLOCATABLE :: pos, cell
33 : REAL(dp), DIMENSION(3, 3) :: hinv
34 2 : INTEGER(kind=int_8), DIMENSION(:), ALLOCATABLE :: atom_types
35 2 : INTEGER(kind=int_8), DIMENSION(:, :), ALLOCATABLE :: edge_index
36 2 : 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 2 : REAL(sp), DIMENSION(:, :), POINTER :: total_energy, atomic_energy, forces
44 2 : NULLIFY (total_energy, atomic_energy, forces)
45 :
46 : ! A box with 32 water molecules.
47 2 : ALLOCATE (pos(3, natoms))
48 : pos(:, :) = RESHAPE(REAL([ &
49 : 42.8861696_dp, -0.0556816_dp, 38.3291611_dp, &
50 : 34.2025887_dp, -0.6185484_dp, 37.3655680_dp, &
51 : 30.0803925_dp, -2.0124176_dp, 36.4807960_dp, &
52 : 28.7057911_dp, -2.6880392_dp, 36.6020983_dp, &
53 : 36.2479426_dp, -0.5163484_dp, 34.4923596_dp, &
54 : 37.6964724_dp, -0.0410872_dp, 35.0140735_dp, &
55 : 27.7606699_dp, 7.4854206_dp, 33.9276919_dp, &
56 : 28.8160999_dp, 6.4985777_dp, 34.2163608_dp, &
57 : 37.1576372_dp, 9.0188280_dp, 31.9265812_dp, &
58 : 38.6063816_dp, 9.5820079_dp, 32.3435972_dp, &
59 : 34.3031959_dp, 2.2195014_dp, 45.9880451_dp, &
60 : 33.2444139_dp, 1.3025332_dp, 46.4698427_dp, &
61 : 38.7286174_dp, -5.0541897_dp, 26.0743968_dp, &
62 : 38.3483921_dp, -6.2832846_dp, 26.9867253_dp, &
63 : 32.8642520_dp, 3.2060632_dp, 30.8971160_dp, &
64 : 31.2904088_dp, 3.0871834_dp, 30.6273977_dp, &
65 : 33.7519869_dp, -3.1383262_dp, 39.6727607_dp, &
66 : 34.6642979_dp, -3.6643859_dp, 38.6466027_dp, &
67 : 42.7173214_dp, 5.1246883_dp, 32.5883401_dp, &
68 : 41.5627455_dp, 5.5893544_dp, 33.4174902_dp, &
69 : 32.4283800_dp, 9.1182520_dp, 30.5477678_dp, &
70 : 32.6432407_dp, 10.770683_dp, 30.4842778_dp, &
71 : 31.4848670_dp, 4.6777144_dp, 37.3957194_dp, &
72 : 32.3171882_dp, -6.2287496_dp, 36.4671864_dp, &
73 : 26.6621340_dp, 3.1708123_dp, 35.6820146_dp, &
74 : 26.5271367_dp, 1.6039040_dp, 35.4883482_dp, &
75 : 32.0238236_dp, 16.918208_dp, 31.6883569_dp, &
76 : 31.4006579_dp, 7.0315610_dp, 30.2394554_dp, &
77 : 33.5264253_dp, -3.5594808_dp, 34.2636830_dp, &
78 : 34.6404855_dp, -3.2653833_dp, 35.4971482_dp, &
79 : 40.0564375_dp, -0.3054386_dp, 29.8312074_dp, &
80 : 39.4784464_dp, -1.0948314_dp, 38.3101140_dp, &
81 : 39.7040761_dp, 1.9584631_dp, 33.3902375_dp, &
82 : 38.3338570_dp, 2.6967178_dp, 42.9261945_dp, &
83 : 40.1820455_dp, -7.2199289_dp, 27.6580390_dp, &
84 : 39.3204431_dp, -8.4564252_dp, 28.1319658_dp, &
85 : 36.3876963_dp, 8.8117085_dp, 38.3545362_dp, &
86 : 36.3205637_dp, 9.0063075_dp, 36.7526001_dp, &
87 : 29.9991583_dp, -5.5637817_dp, 33.9295050_dp, &
88 : 30.7728545_dp, -5.0385870_dp, 35.1998067_dp, &
89 : 40.0592517_dp, 6.3305279_dp, 28.2579461_dp, &
90 : 40.2398360_dp, 5.1745923_dp, 29.2962956_dp, &
91 : 26.3320911_dp, 2.4393638_dp, 33.5653868_dp, &
92 : 26.9606971_dp, 1.2711078_dp, 32.5923884_dp, &
93 : 34.8372697_dp, -0.4722708_dp, 30.3824362_dp, &
94 : 35.3968813_dp, -1.9268483_dp, 30.3081837_dp, &
95 : 32.1217607_dp, -0.7333429_dp, 36.5104382_dp, &
96 : 32.2180843_dp, 7.8454304_dp, 35.6671967_dp, &
97 : 36.3780998_dp, -4.3048878_dp, 36.4539793_dp, &
98 : 35.8119275_dp, -3.0013928_dp, 27.0348937_dp, &
99 : 29.6452491_dp, 1.0652123_dp, 35.7143653_dp, &
100 : 30.3794654_dp, -0.0668146_dp, 34.9882468_dp, &
101 : 34.2149336_dp, -1.6559120_dp, 33.8876437_dp, &
102 : 34.7842435_dp, -1.0252141_dp, 32.5034832_dp, &
103 : 40.4649954_dp, 1.1467825_dp, 31.3073503_dp, &
104 : 41.3262469_dp, 0.6550803_dp, 32.4555882_dp, &
105 : 29.0210859_dp, 3.5038194_dp, 39.9087702_dp, &
106 : 29.4945426_dp, 3.7276637_dp, 41.3766138_dp, &
107 : 34.1359664_dp, -6.7533422_dp, 32.3568410_dp, &
108 : 34.9546570_dp, -5.7704242_dp, 31.4571066_dp, &
109 : 33.2532356_dp, 1.5268048_dp, 44.0562171_dp, &
110 : 33.7931669_dp, 0.5014632_dp, 43.0597590_dp, &
111 : 36.8205409_dp, 2.6214681_dp, 40.6834006_dp, &
112 : 37.5552706_dp, 1.5649832_dp, 39.7648935_dp, &
113 : 43.2099087_dp, -0.0628456_dp, 47.2593155_dp, &
114 : 29.3940583_dp, -2.3133019_dp, 37.1407883_dp, &
115 : 36.7415708_dp, -0.0838710_dp, 35.2591783_dp, &
116 : 27.9424776_dp, 6.7622961_dp, 34.5648384_dp, &
117 : 37.6812656_dp, 9.4216399_dp, 32.6478643_dp, &
118 : 33.3171290_dp, 2.0951401_dp, 45.8722265_dp, &
119 : 37.9951355_dp, 4.3611431_dp, 26.5571819_dp, &
120 : 32.1824670_dp, 2.6611503_dp, 30.4577248_dp, &
121 : 34.6538012_dp, -3.4374573_dp, 39.5889245_dp, &
122 : 42.2929833_dp, 5.9471069_dp, 32.8460995_dp, &
123 : 32.9604690_dp, 9.9050313_dp, 30.1587306_dp, &
124 : 31.4281886_dp, -5.8338304_dp, 36.6738743_dp, &
125 : 26.0563730_dp, 2.4973869_dp, 35.3486870_dp, &
126 : 32.0334927_dp, 17.3252289_dp, 30.8116013_dp, &
127 : 33.8252182_dp, -2.9520949_dp, 35.0220460_dp, &
128 : 39.4569981_dp, -0.3072759_dp, 38.9347829_dp, &
129 : 29.4846708_dp, 2.8692561_dp, 43.0061868_dp, &
130 : 39.2864184_dp, -7.6206103_dp, 27.6271147_dp, &
131 : 35.8797502_dp, 8.6515870_dp, 37.5221734_dp, &
132 : 30.3582543_dp, -4.7607656_dp, 34.3355645_dp, &
133 : 40.7098956_dp, 5.8331250_dp, 28.7558375_dp, &
134 : 26.7179083_dp, 2.2415138_dp, 32.6577297_dp, &
135 : 35.6589256_dp, -0.9968903_dp, 30.5749530_dp, &
136 : 31.5851602_dp, -1.3121804_dp, 35.9011109_dp, &
137 : 35.5489386_dp, -3.9056138_dp, 26.8214490_dp, &
138 : 29.5656616_dp, 0.4681794_dp, 34.9670711_dp, &
139 : 34.7615128_dp, -0.9569680_dp, 33.4891367_dp, &
140 : 40.4853406_dp, 0.4023620_dp, 31.9425416_dp, &
141 : 29.6728289_dp, 4.0134825_dp, 40.4505780_dp, &
142 : 34.1272286_dp, -5.8796882_dp, 31.8925999_dp, &
143 : 33.1168884_dp, 1.2338084_dp, 43.1127997_dp, &
144 770 : 37.1996993_dp, 2.5049007_dp, 39.7917126_dp], kind=sp), shape=[3, natoms])
145 :
146 2 : ALLOCATE (cell(3, 3))
147 8 : cell(1, :) = [9.85_sp, 0.0_sp, 0.0_sp]
148 8 : cell(2, :) = [0.0_sp, 9.85_sp, 0.0_sp]
149 8 : cell(3, :) = [0.0_sp, 0.0_sp, 9.85_sp]
150 :
151 26 : hinv(:, :) = inv_3x3(REAL(cell, kind=dp))
152 :
153 2 : ALLOCATE (atom_types(natoms))
154 130 : atom_types(:64) = 0 ! Hydrogen
155 66 : atom_types(65:) = 1 ! Oxygen
156 :
157 2 : WRITE (*, *) "CUDA is available: ", torch_cuda_is_available()
158 :
159 2 : filename = discover_file('NequIP/water-deployed-neq060sp.pth')
160 2 : WRITE (*, *) "Loading NequIP model from: "//TRIM(filename)
161 2 : CALL torch_model_load(model, filename)
162 2 : cutoff_str = torch_model_read_metadata(filename, "r_max")
163 2 : nequip_version = torch_model_read_metadata(filename, "nequip_version")
164 2 : READ (cutoff_str, *) cutoff
165 2 : WRITE (*, *) "Version: ", TRIM(nequip_version)
166 2 : WRITE (*, *) "Cutoff: ", cutoff
167 :
168 2 : CALL neighbor_search(nedges)
169 6 : ALLOCATE (edge_index(nedges, 2))
170 6 : ALLOCATE (edge_cell_shift(3, nedges))
171 2 : CALL neighbor_search(nedges, edge_index, edge_cell_shift)
172 2 : WRITE (*, *) "Found", nedges, "neighbor edges between", natoms, "atoms."
173 :
174 2 : CALL torch_dict_create(inputs)
175 2 : CALL torch_dict_insert(inputs, "pos", pos)
176 2 : CALL torch_dict_insert(inputs, "edge_index", edge_index)
177 2 : CALL torch_dict_insert(inputs, "edge_cell_shift", edge_cell_shift)
178 2 : CALL torch_dict_insert(inputs, "cell", cell)
179 2 : CALL torch_dict_insert(inputs, "atom_types", atom_types)
180 :
181 2 : CALL torch_dict_create(outputs)
182 2 : CALL torch_model_eval(model, inputs, outputs)
183 :
184 2 : CALL torch_dict_get(outputs, "total_energy", total_energy)
185 2 : CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
186 2 : CALL torch_dict_get(outputs, "forces", forces)
187 2 : WRITE (*, *) "Total Energy [Hartree] : ", total_energy(1, 1)/evolt
188 2 : WRITE (*, *) "FORCES: [Hartree/Bohr]: "
189 194 : DO iatom = 1, natoms
190 770 : WRITE (*, *) forces(:, iatom)*angstrom/evolt
191 : END DO
192 2 : CPASSERT(ABS(-14985.4443_dp - REAL(total_energy(1, 1), kind=dp)) < 2e-3_dp)
193 :
194 2 : CALL torch_dict_release(inputs)
195 2 : CALL torch_dict_release(outputs)
196 2 : CALL torch_model_release(model)
197 2 : DEALLOCATE (edge_index, edge_cell_shift, total_energy, atomic_energy, forces, pos, cell, atom_types)
198 :
199 8 : WRITE (*, *) "NequIP unittest was successfully :-)"
200 :
201 : CONTAINS
202 :
203 : ! **************************************************************************************************
204 : !> \brief Naive neighbor search - beware it scales O(N**2).
205 : !> \param nedges ...
206 : !> \param edge_index ...
207 : !> \param edge_cell_shift ...
208 : ! **************************************************************************************************
209 4 : SUBROUTINE neighbor_search(nedges, edge_index, edge_cell_shift)
210 : INTEGER, INTENT(OUT) :: nedges
211 : INTEGER(kind=int_8), DIMENSION(:, :), &
212 : INTENT(OUT), OPTIONAL :: edge_index
213 : REAL(sp), DIMENSION(:, :), INTENT(OUT), OPTIONAL :: edge_cell_shift
214 :
215 : INTEGER:: iatom, jatom
216 : REAL(dp), DIMENSION(3) :: s1, s2, s12, cell_shift, dx
217 :
218 4 : nedges = 0
219 388 : DO iatom = 1, natoms
220 37252 : DO jatom = 1, natoms
221 36864 : IF (iatom == jatom) CYCLE
222 583680 : s1 = MATMUL(hinv, pos(:, iatom))
223 583680 : s2 = MATMUL(hinv, pos(:, jatom))
224 145920 : s12 = s1 - s2
225 145920 : cell_shift = ANINT(s12)
226 1422720 : dx = MATMUL(cell, s12 - cell_shift)
227 146304 : IF (DOT_PRODUCT(dx, dx) <= cutoff**2) THEN
228 10736 : nedges = nedges + 1
229 10736 : IF (PRESENT(edge_index)) THEN
230 16104 : edge_index(nedges, :) = [iatom - 1, jatom - 1]
231 : END IF
232 10736 : IF (PRESENT(edge_cell_shift)) THEN
233 21472 : edge_cell_shift(:, nedges) = REAL(cell_shift, kind=sp)
234 : END IF
235 : END IF
236 : END DO
237 : END DO
238 4 : END SUBROUTINE neighbor_search
239 :
240 : END PROGRAM nequip_unittest
|
