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_freeze_preserving_method, torch_model_load_with_metadata, torch_model_release, &
23 : torch_model_remap_device_constants, torch_model_promote_float32, 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 : !> \param fp64 Promote float32 numerical state and casts; default true for differential consistency.
53 : ! **************************************************************************************************
54 17 : SUBROUTINE skala_torch_model_load(model, filename, fp64)
55 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
56 : CHARACTER(len=*), INTENT(IN) :: filename
57 : LOGICAL, INTENT(IN), OPTIONAL :: fp64
58 :
59 17 : CHARACTER(:), ALLOCATABLE :: features_json, protocol_string
60 : INTEGER :: ios
61 : LOGICAL :: use_fp64
62 :
63 : CALL torch_model_load_with_metadata(model%torch_model, filename, &
64 : "protocol_version", protocol_string, &
65 17 : "features", features_json)
66 17 : CALL torch_model_remap_device_constants(model%torch_model)
67 17 : CALL torch_model_disable_parameter_gradients(model%torch_model)
68 17 : READ (protocol_string, *, IOSTAT=ios) model%protocol_version
69 17 : IF (ios /= 0) CPABORT("Could not parse SKALA TorchScript protocol_version metadata")
70 17 : IF (model%protocol_version /= 2) THEN
71 0 : CPABORT("Unsupported SKALA TorchScript protocol version")
72 : END IF
73 :
74 17 : CALL parse_feature_list(features_json, model%features)
75 17 : use_fp64 = .TRUE.
76 17 : IF (PRESENT(fp64)) use_fp64 = fp64
77 17 : IF (use_fp64) CALL torch_model_promote_float32(model%torch_model)
78 : ! Preserve the exported SKALA entry point while folding constant model state.
79 17 : CALL torch_model_freeze_preserving_method(model%torch_model, "get_exc_density")
80 :
81 17 : END SUBROUTINE skala_torch_model_load
82 :
83 : ! **************************************************************************************************
84 : !> \brief Release a loaded SKALA TorchScript model.
85 : !> \param model ...
86 : ! **************************************************************************************************
87 0 : SUBROUTINE skala_torch_model_release(model)
88 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
89 :
90 0 : CALL torch_model_release(model%torch_model)
91 0 : IF (ALLOCATED(model%features)) DEALLOCATE (model%features)
92 0 : model%protocol_version = -1
93 :
94 0 : END SUBROUTINE skala_torch_model_release
95 :
96 : ! **************************************************************************************************
97 : !> \brief Check whether a loaded SKALA model requests a feature.
98 : !> \param model ...
99 : !> \param feature ...
100 : !> \return ...
101 : ! **************************************************************************************************
102 0 : FUNCTION skala_torch_model_needs_feature(model, feature) RESULT(needs_feature)
103 : TYPE(skala_torch_model_type), INTENT(IN) :: model
104 : CHARACTER(len=*), INTENT(IN) :: feature
105 : LOGICAL :: needs_feature
106 :
107 : CHARACTER(len=default_string_length) :: feature_key, model_feature
108 : INTEGER :: i
109 :
110 0 : feature_key = ADJUSTL(feature)
111 0 : CALL uppercase(feature_key)
112 :
113 0 : needs_feature = .FALSE.
114 0 : IF (.NOT. ALLOCATED(model%features)) RETURN
115 :
116 0 : DO i = 1, SIZE(model%features)
117 0 : model_feature = ADJUSTL(model%features(i))
118 0 : CALL uppercase(model_feature)
119 0 : IF (TRIM(model_feature) == TRIM(feature_key)) THEN
120 0 : needs_feature = .TRUE.
121 : RETURN
122 : END IF
123 : END DO
124 :
125 0 : END FUNCTION skala_torch_model_needs_feature
126 :
127 : ! **************************************************************************************************
128 : !> \brief Return the loaded SKALA TorchScript protocol version.
129 : !> \param model ...
130 : !> \return ...
131 : ! **************************************************************************************************
132 0 : FUNCTION skala_torch_model_protocol_version(model) RESULT(protocol_version)
133 : TYPE(skala_torch_model_type), INTENT(IN) :: model
134 : INTEGER :: protocol_version
135 :
136 0 : protocol_version = model%protocol_version
137 :
138 0 : END FUNCTION skala_torch_model_protocol_version
139 :
140 : ! **************************************************************************************************
141 : !> \brief Evaluate the SKALA exchange-correlation energy density.
142 : !> \param model ...
143 : !> \param inputs ...
144 : !> \param exc_density ...
145 : ! **************************************************************************************************
146 0 : SUBROUTINE skala_torch_model_get_exc_density(model, inputs, exc_density)
147 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
148 : TYPE(torch_dict_type), INTENT(IN) :: inputs
149 : TYPE(torch_tensor_type), INTENT(INOUT) :: exc_density
150 :
151 : #if defined (__HAS_IEEE_EXCEPTIONS)
152 : LOGICAL, DIMENSION(5) :: ieee_halt
153 :
154 : CALL ieee_get_halting_mode(IEEE_ALL, ieee_halt)
155 : CALL ieee_set_halting_mode(IEEE_ALL, .FALSE.)
156 : #endif
157 0 : CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
158 : #if defined (__HAS_IEEE_EXCEPTIONS)
159 : CALL ieee_set_halting_mode(IEEE_ALL, ieee_halt)
160 : #endif
161 :
162 0 : END SUBROUTINE skala_torch_model_get_exc_density
163 :
164 : ! **************************************************************************************************
165 : !> \brief Evaluate the weighted SKALA exchange-correlation energy.
166 : !> \param model ...
167 : !> \param inputs ...
168 : !> \param grid_weights ...
169 : !> \param exc_tensor ...
170 : !> \param exc ...
171 : ! **************************************************************************************************
172 284 : SUBROUTINE skala_torch_model_get_exc(model, inputs, grid_weights, exc_tensor, exc)
173 : TYPE(skala_torch_model_type), INTENT(INOUT) :: model
174 : TYPE(torch_dict_type), INTENT(IN) :: inputs
175 : TYPE(torch_tensor_type), INTENT(IN) :: grid_weights
176 : TYPE(torch_tensor_type), INTENT(INOUT) :: exc_tensor
177 : REAL(KIND=dp), INTENT(OUT) :: exc
178 :
179 : TYPE(torch_tensor_type) :: exc_density
180 :
181 : #if defined (__HAS_IEEE_EXCEPTIONS)
182 : LOGICAL, DIMENSION(5) :: ieee_halt
183 :
184 : CALL ieee_get_halting_mode(IEEE_ALL, ieee_halt)
185 : CALL ieee_set_halting_mode(IEEE_ALL, .FALSE.)
186 : #endif
187 284 : CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
188 284 : CALL torch_tensor_weighted_sum(exc_density, grid_weights, exc_tensor)
189 284 : CALL torch_tensor_release(exc_density)
190 284 : exc = torch_tensor_item_double(exc_tensor)
191 : #if defined (__HAS_IEEE_EXCEPTIONS)
192 : CALL ieee_set_halting_mode(IEEE_ALL, ieee_halt)
193 : #endif
194 :
195 284 : END SUBROUTINE skala_torch_model_get_exc
196 :
197 : ! **************************************************************************************************
198 : !> \brief Parse a TorchScript extra_files JSON list of feature names.
199 : !> \param features_json ...
200 : !> \param features ...
201 : ! **************************************************************************************************
202 17 : SUBROUTINE parse_feature_list(features_json, features)
203 : CHARACTER(len=*), INTENT(IN) :: features_json
204 : CHARACTER(len=default_string_length), &
205 : ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: features
206 :
207 : INTEGER :: end_pos, feature_count, i, pos, quote1, &
208 : quote2, start_pos
209 :
210 17 : feature_count = 0
211 17 : pos = 1
212 153 : DO
213 170 : quote1 = INDEX(features_json(pos:), '"')
214 170 : IF (quote1 == 0) EXIT
215 153 : start_pos = pos + quote1
216 153 : quote2 = INDEX(features_json(start_pos:), '"')
217 153 : IF (quote2 == 0) EXIT
218 153 : feature_count = feature_count + 1
219 153 : pos = start_pos + quote2
220 : END DO
221 :
222 17 : IF (feature_count == 0) CPABORT("SKALA TorchScript model does not list any features")
223 51 : ALLOCATE (features(feature_count))
224 170 : features = ""
225 :
226 : pos = 1
227 170 : DO i = 1, feature_count
228 153 : quote1 = INDEX(features_json(pos:), '"')
229 153 : start_pos = pos + quote1
230 153 : quote2 = INDEX(features_json(start_pos:), '"')
231 153 : end_pos = start_pos + quote2 - 2
232 153 : features(i) = features_json(start_pos:end_pos)
233 170 : pos = start_pos + quote2
234 : END DO
235 :
236 17 : END SUBROUTINE parse_feature_list
237 :
238 0 : END MODULE skala_torch_api
|
