Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Small CP2K wrapper around the SKALA TorchScript functional protocol.
10 : ! **************************************************************************************************
11 : MODULE skala_torch_api
12 : #if defined (__HAS_IEEE_EXCEPTIONS)
13 : USE ieee_exceptions, ONLY: ieee_all, &
14 : ieee_get_halting_mode, &
15 : ieee_set_halting_mode
16 : #endif
17 : USE kinds, ONLY: default_string_length, &
18 : dp
19 : USE string_utilities, ONLY: uppercase
20 : USE torch_api, ONLY: &
21 : torch_dict_type, torch_model_disable_parameter_gradients, torch_model_forward_mol_tensor, &
22 : torch_model_load_with_metadata, torch_model_release, torch_model_remap_device_constants, &
23 : torch_model_type, &
24 : torch_tensor_item_double, torch_tensor_release, torch_tensor_type, &
25 : torch_tensor_weighted_sum
26 : #include "./base/base_uses.f90"
27 :
28 : IMPLICIT NONE
29 :
30 : PRIVATE
31 :
32 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_torch_api'
33 :
34 : PUBLIC :: skala_torch_model_type, skala_torch_model_load, skala_torch_model_release
35 : PUBLIC :: skala_torch_model_get_exc, skala_torch_model_get_exc_density
36 : PUBLIC :: skala_torch_model_needs_feature, skala_torch_model_protocol_version
37 :
38 : TYPE skala_torch_model_type
39 : PRIVATE
40 : INTEGER :: protocol_version = -1
41 : CHARACTER(len=default_string_length), ALLOCATABLE, &
42 : DIMENSION(:) :: features
43 : TYPE(torch_model_type) :: torch_model
44 : END TYPE skala_torch_model_type
45 :
46 : CONTAINS
47 :
48 : ! **************************************************************************************************
49 : !> \brief Load a SKALA TorchScript model and its feature metadata.
50 : !> \param model ...
51 : !> \param filename ...
52 : ! **************************************************************************************************
53 97 : SUBROUTINE skala_torch_model_load(model, filename)
54 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
55 : CHARACTER(len=*), INTENT(IN) :: filename
56 :
57 97 : CHARACTER(:), ALLOCATABLE :: features_json, protocol_string
58 : INTEGER :: ios
59 :
60 : CALL torch_model_load_with_metadata(model%torch_model, filename, &
61 : "protocol_version", protocol_string, &
62 97 : "features", features_json)
63 97 : CALL torch_model_remap_device_constants(model%torch_model)
64 97 : CALL torch_model_disable_parameter_gradients(model%torch_model)
65 97 : READ (protocol_string, *, IOSTAT=ios) model%protocol_version
66 97 : IF (ios /= 0) CPABORT("Could not parse SKALA TorchScript protocol_version metadata")
67 97 : IF (model%protocol_version /= 2) THEN
68 0 : CPABORT("Unsupported SKALA TorchScript protocol version")
69 : END IF
70 :
71 97 : CALL parse_feature_list(features_json, model%features)
72 :
73 97 : END SUBROUTINE skala_torch_model_load
74 :
75 : ! **************************************************************************************************
76 : !> \brief Release a loaded SKALA TorchScript model.
77 : !> \param model ...
78 : ! **************************************************************************************************
79 0 : SUBROUTINE skala_torch_model_release(model)
80 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
81 :
82 0 : CALL torch_model_release(model%torch_model)
83 0 : IF (ALLOCATED(model%features)) DEALLOCATE (model%features)
84 0 : model%protocol_version = -1
85 :
86 0 : END SUBROUTINE skala_torch_model_release
87 :
88 : ! **************************************************************************************************
89 : !> \brief Check whether a loaded SKALA model requests a feature.
90 : !> \param model ...
91 : !> \param feature ...
92 : !> \return ...
93 : ! **************************************************************************************************
94 0 : FUNCTION skala_torch_model_needs_feature(model, feature) RESULT(needs_feature)
95 : TYPE(skala_torch_model_type), INTENT(IN) :: model
96 : CHARACTER(len=*), INTENT(IN) :: feature
97 : LOGICAL :: needs_feature
98 :
99 : CHARACTER(len=default_string_length) :: feature_key, model_feature
100 : INTEGER :: i
101 :
102 0 : feature_key = ADJUSTL(feature)
103 0 : CALL uppercase(feature_key)
104 :
105 0 : needs_feature = .FALSE.
106 0 : IF (.NOT. ALLOCATED(model%features)) RETURN
107 :
108 0 : DO i = 1, SIZE(model%features)
109 0 : model_feature = ADJUSTL(model%features(i))
110 0 : CALL uppercase(model_feature)
111 0 : IF (TRIM(model_feature) == TRIM(feature_key)) THEN
112 0 : needs_feature = .TRUE.
113 : RETURN
114 : END IF
115 : END DO
116 :
117 0 : END FUNCTION skala_torch_model_needs_feature
118 :
119 : ! **************************************************************************************************
120 : !> \brief Return the loaded SKALA TorchScript protocol version.
121 : !> \param model ...
122 : !> \return ...
123 : ! **************************************************************************************************
124 0 : FUNCTION skala_torch_model_protocol_version(model) RESULT(protocol_version)
125 : TYPE(skala_torch_model_type), INTENT(IN) :: model
126 : INTEGER :: protocol_version
127 :
128 0 : protocol_version = model%protocol_version
129 :
130 0 : END FUNCTION skala_torch_model_protocol_version
131 :
132 : ! **************************************************************************************************
133 : !> \brief Evaluate the SKALA exchange-correlation energy density.
134 : !> \param model ...
135 : !> \param inputs ...
136 : !> \param exc_density ...
137 : ! **************************************************************************************************
138 0 : SUBROUTINE skala_torch_model_get_exc_density(model, inputs, exc_density)
139 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
140 : TYPE(torch_dict_type), INTENT(IN) :: inputs
141 : TYPE(torch_tensor_type), INTENT(INOUT) :: exc_density
142 :
143 : #if defined (__HAS_IEEE_EXCEPTIONS)
144 : LOGICAL, DIMENSION(5) :: ieee_halt
145 :
146 : CALL ieee_get_halting_mode(IEEE_ALL, ieee_halt)
147 : CALL ieee_set_halting_mode(IEEE_ALL, .FALSE.)
148 : #endif
149 0 : CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
150 : #if defined (__HAS_IEEE_EXCEPTIONS)
151 : CALL ieee_set_halting_mode(IEEE_ALL, ieee_halt)
152 : #endif
153 :
154 0 : END SUBROUTINE skala_torch_model_get_exc_density
155 :
156 : ! **************************************************************************************************
157 : !> \brief Evaluate the weighted SKALA exchange-correlation energy.
158 : !> \param model ...
159 : !> \param inputs ...
160 : !> \param grid_weights ...
161 : !> \param exc_tensor ...
162 : !> \param exc ...
163 : ! **************************************************************************************************
164 386 : SUBROUTINE skala_torch_model_get_exc(model, inputs, grid_weights, exc_tensor, exc)
165 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
166 : TYPE(torch_dict_type), INTENT(IN) :: inputs
167 : TYPE(torch_tensor_type), INTENT(IN) :: grid_weights
168 : TYPE(torch_tensor_type), INTENT(INOUT) :: exc_tensor
169 : REAL(KIND=dp), INTENT(OUT) :: exc
170 :
171 : TYPE(torch_tensor_type) :: exc_density
172 :
173 : #if defined (__HAS_IEEE_EXCEPTIONS)
174 : LOGICAL, DIMENSION(5) :: ieee_halt
175 :
176 : CALL ieee_get_halting_mode(IEEE_ALL, ieee_halt)
177 : CALL ieee_set_halting_mode(IEEE_ALL, .FALSE.)
178 : #endif
179 386 : CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
180 386 : CALL torch_tensor_weighted_sum(exc_density, grid_weights, exc_tensor)
181 386 : CALL torch_tensor_release(exc_density)
182 386 : exc = torch_tensor_item_double(exc_tensor)
183 : #if defined (__HAS_IEEE_EXCEPTIONS)
184 : CALL ieee_set_halting_mode(IEEE_ALL, ieee_halt)
185 : #endif
186 :
187 386 : END SUBROUTINE skala_torch_model_get_exc
188 :
189 : ! **************************************************************************************************
190 : !> \brief Parse a TorchScript extra_files JSON list of feature names.
191 : !> \param features_json ...
192 : !> \param features ...
193 : ! **************************************************************************************************
194 97 : SUBROUTINE parse_feature_list(features_json, features)
195 : CHARACTER(len=*), INTENT(IN) :: features_json
196 : CHARACTER(len=default_string_length), &
197 : ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: features
198 :
199 : INTEGER :: end_pos, feature_count, i, pos, quote1, &
200 : quote2, start_pos
201 :
202 97 : feature_count = 0
203 97 : pos = 1
204 873 : DO
205 970 : quote1 = INDEX(features_json(pos:), '"')
206 970 : IF (quote1 == 0) EXIT
207 873 : start_pos = pos + quote1
208 873 : quote2 = INDEX(features_json(start_pos:), '"')
209 873 : IF (quote2 == 0) EXIT
210 873 : feature_count = feature_count + 1
211 873 : pos = start_pos + quote2
212 : END DO
213 :
214 97 : IF (feature_count == 0) CPABORT("SKALA TorchScript model does not list any features")
215 291 : ALLOCATE (features(feature_count))
216 970 : features = ""
217 :
218 : pos = 1
219 970 : DO i = 1, feature_count
220 873 : quote1 = INDEX(features_json(pos:), '"')
221 873 : start_pos = pos + quote1
222 873 : quote2 = INDEX(features_json(start_pos:), '"')
223 873 : end_pos = start_pos + quote2 - 2
224 873 : features(i) = features_json(start_pos:end_pos)
225 970 : pos = start_pos + quote2
226 : END DO
227 :
228 97 : END SUBROUTINE parse_feature_list
229 :
230 0 : END MODULE skala_torch_api
|
