13 USE omp_lib,
ONLY: omp_init_lock,&
47#include "./base/base_uses.f90"
68 INTEGER,
INTENT(IN) :: ikind
69 CHARACTER(LEN=default_path_length),
INTENT(IN) :: pao_model_file
72 CHARACTER(len=*),
PARAMETER :: routinen =
'pao_model_load'
74 CHARACTER(LEN=default_string_length) :: kind_name
75 CHARACTER(LEN=default_string_length), &
76 ALLOCATABLE,
DIMENSION(:) :: feature_kind_names
77 INTEGER :: handle, jkind, kkind, pao_basis_size, z
82 CALL timeset(routinen, handle)
83 CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set, atomic_kind_set=atomic_kind_set)
85 IF (pao%iw > 0)
WRITE (pao%iw,
'(A)')
" PAO| Loading PyTorch model from: "//trim(pao_model_file)
105 ALLOCATE (model%feature_kinds(
SIZE(feature_kind_names)))
106 model%feature_kinds(:) = -1
107 DO jkind = 1,
SIZE(feature_kind_names)
108 DO kkind = 1,
SIZE(atomic_kind_set)
109 IF (trim(atomic_kind_set(kkind)%name) == trim(feature_kind_names(jkind)))
THEN
110 model%feature_kinds(jkind) = kkind
113 IF (model%feature_kinds(jkind) < 0)
THEN
115 WRITE (pao%iw,
'(A)')
" PAO| ML-model supports feature kind '"// &
116 trim(feature_kind_names(jkind))//
"' that is not present in subsys."
121 DO jkind = 1,
SIZE(atomic_kind_set)
122 IF (all(model%feature_kinds /= atomic_kind_set(jkind)%kind_number))
THEN
124 WRITE (pao%iw,
'(A)')
" PAO| ML-Model lacks feature kind '"// &
125 trim(atomic_kind_set(jkind)%name)//
"' that is present in subsys."
130 CALL get_qs_kind(qs_kind_set(ikind), basis_set=basis_set, pao_basis_size=pao_basis_size)
132 IF (model%version /= 1) &
133 cpabort(
"Model version not supported.")
134 IF (trim(model%kind_name) .NE. trim(kind_name)) &
135 cpabort(
"Kind name does not match.")
136 IF (model%atomic_number /= z) &
137 cpabort(
"Atomic number does not match.")
138 IF (trim(model%prim_basis_name) .NE. trim(basis_set%name)) &
139 cpabort(
"Primary basis set name does not match.")
140 IF (model%prim_basis_size /= basis_set%nsgf) &
141 cpabort(
"Primary basis set size does not match.")
142 IF (model%pao_basis_size /= pao_basis_size) &
143 cpabort(
"PAO basis size does not match.")
145 CALL omp_init_lock(model%lock)
146 CALL timestop(handle)
159 CHARACTER(len=*),
PARAMETER :: routinen =
'pao_model_predict'
161 INTEGER :: acol, arow, handle, iatom
162 REAL(
dp),
DIMENSION(:, :),
POINTER :: block_x
165 CALL timeset(routinen, handle)
171 IF (
SIZE(block_x) == 0) cycle
172 iatom = arow; cpassert(arow == acol)
173 CALL predict_single_atom(pao, qs_env, iatom, block_x=block_x)
178 CALL timestop(handle)
193 REAL(
dp),
DIMENSION(:, :),
INTENT(INOUT) :: forces
195 CHARACTER(len=*),
PARAMETER :: routinen =
'pao_model_forces'
197 INTEGER :: acol, arow, handle, iatom
198 REAL(
dp),
DIMENSION(:, :),
POINTER :: block_g
201 CALL timeset(routinen, handle)
207 iatom = arow; cpassert(arow == acol)
208 IF (
SIZE(block_g) == 0) cycle
209 CALL predict_single_atom(pao, qs_env, iatom, block_g=block_g, forces=forces)
214 CALL timestop(handle)
227 SUBROUTINE predict_single_atom(pao, qs_env, iatom, block_X, block_G, forces)
230 INTEGER,
INTENT(IN) :: iatom
231 REAL(
dp),
DIMENSION(:, :),
OPTIONAL :: block_x, block_g, forces
233 INTEGER :: ikind, jatom, jkind, jneighbor, m, n, &
235 INTEGER,
ALLOCATABLE,
DIMENSION(:) :: neighbors_index
236 INTEGER,
DIMENSION(:),
POINTER :: blk_sizes_pao, blk_sizes_pri
237 REAL(
dp),
DIMENSION(3) :: ri, rij, rj
238 REAL(kind=
dp),
ALLOCATABLE,
DIMENSION(:) :: neighbors_distance
239 REAL(
sp),
ALLOCATABLE,
DIMENSION(:, :) :: features, outer_grad, relpos
240 REAL(
sp),
DIMENSION(:, :),
POINTER :: predicted_xblock, relpos_grad
246 TYPE(
qs_kind_type),
DIMENSION(:),
POINTER :: qs_kind_set
249 predicted_xblock_tensor, &
250 relpos_grad_tensor, relpos_tensor
252 CALL dbcsr_get_info(pao%matrix_Y, row_blk_size=blk_sizes_pri, col_blk_size=blk_sizes_pao)
253 n = blk_sizes_pri(iatom)
254 m = blk_sizes_pao(iatom)
259 particle_set=particle_set, &
260 atomic_kind_set=atomic_kind_set, &
261 qs_kind_set=qs_kind_set, &
264 CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
265 model => pao%models(ikind)
266 cpassert(model%version > 0)
267 CALL omp_set_lock(model%lock)
271 ALLOCATE (neighbors_distance(natoms), neighbors_index(natoms))
272 ri = particle_set(iatom)%r
274 rj = particle_set(jatom)%r
275 rij =
pbc(ri, rj, cell)
276 neighbors_distance(jatom) = dot_product(rij, rij)
278 CALL sort(neighbors_distance, natoms, neighbors_index)
279 cpassert(neighbors_index(1) == iatom)
282 ALLOCATE (relpos(3, model%num_neighbors))
283 relpos(:, :) = 0.0_sp
284 DO jneighbor = 1, min(model%num_neighbors, natoms - 1)
285 jatom = neighbors_index(jneighbor + 1)
286 rj = particle_set(jatom)%r
287 rij =
pbc(ri, rj, cell)
288 relpos(:, jneighbor) = real(
angstrom*rij, kind=
sp)
292 ALLOCATE (features(
SIZE(model%feature_kinds), model%num_neighbors))
293 features(:, :) = 0.0_sp
294 DO jneighbor = 1, min(model%num_neighbors, natoms - 1)
295 jatom = neighbors_index(jneighbor + 1)
296 jkind = particle_set(jatom)%atomic_kind%kind_number
297 WHERE (model%feature_kinds == jkind) features(:, jneighbor) = 1.0_sp
312 NULLIFY (predicted_xblock)
313 CALL torch_dict_get(model_outputs,
"xblock", predicted_xblock_tensor)
315 cpassert(
SIZE(predicted_xblock, 1) == n .AND.
SIZE(predicted_xblock, 2) == m)
316 IF (
PRESENT(block_x))
THEN
317 block_x = reshape(predicted_xblock, [n*m, 1])
321 IF (
PRESENT(block_g))
THEN
322 ALLOCATE (outer_grad(n, m))
323 outer_grad(:, :) = real(reshape(block_g, [n, m]), kind=
sp)
327 NULLIFY (relpos_grad)
329 cpassert(
SIZE(relpos_grad, 1) == 3 .AND.
SIZE(relpos_grad, 2) == model%num_neighbors)
330 DO jneighbor = 1, min(model%num_neighbors, natoms - 1)
331 jatom = neighbors_index(jneighbor + 1)
332 forces(iatom, :) = forces(iatom, :) + relpos_grad(:, jneighbor)*
angstrom
333 forces(jatom, :) = forces(jatom, :) - relpos_grad(:, jneighbor)*
angstrom
345 DEALLOCATE (neighbors_distance, neighbors_index, relpos, features)
346 CALL omp_unset_lock(model%lock)
348 END SUBROUTINE predict_single_atom
Define the atomic kind types and their sub types.
subroutine, public get_atomic_kind(atomic_kind, fist_potential, element_symbol, name, mass, kind_number, natom, atom_list, rcov, rvdw, z, qeff, apol, cpol, mm_radius, shell, shell_active, damping)
Get attributes of an atomic kind.
Handles all functions related to the CELL.
subroutine, public dbcsr_iterator_next_block(iterator, row, column, block, block_number_argument_has_been_removed, row_size, col_size, row_offset, col_offset)
...
logical function, public dbcsr_iterator_blocks_left(iterator)
...
subroutine, public dbcsr_iterator_stop(iterator)
...
subroutine, public dbcsr_get_info(matrix, nblkrows_total, nblkcols_total, nfullrows_total, nfullcols_total, nblkrows_local, nblkcols_local, nfullrows_local, nfullcols_local, my_prow, my_pcol, local_rows, local_cols, proc_row_dist, proc_col_dist, row_blk_size, col_blk_size, row_blk_offset, col_blk_offset, distribution, name, matrix_type, group)
...
subroutine, public dbcsr_iterator_start(iterator, matrix, shared, dynamic, dynamic_byrows)
...
Defines the basic variable types.
integer, parameter, public dp
integer, parameter, public default_string_length
integer, parameter, public default_path_length
integer, parameter, public sp
Interface to the message passing library MPI.
Module for equivariant PAO-ML based on PyTorch.
subroutine, public pao_model_predict(pao, qs_env)
Fills paomatrix_X based on machine learning predictions.
subroutine, public pao_model_forces(pao, qs_env, matrix_g, forces)
Calculate forces contributed by machine learning.
subroutine, public pao_model_load(pao, qs_env, ikind, pao_model_file, model)
Loads a PAO-ML model.
Types used by the PAO machinery.
Define the data structure for the particle information.
Definition of physical constants:
real(kind=dp), parameter, public angstrom
subroutine, public get_qs_env(qs_env, atomic_kind_set, qs_kind_set, cell, super_cell, cell_ref, use_ref_cell, kpoints, dft_control, mos, sab_orb, sab_all, qmmm, qmmm_periodic, sac_ae, sac_ppl, sac_lri, sap_ppnl, sab_vdw, sab_scp, sap_oce, sab_lrc, sab_se, sab_xtbe, sab_tbe, sab_core, sab_xb, sab_xtb_pp, sab_xtb_nonbond, sab_almo, sab_kp, sab_kp_nosym, particle_set, energy, force, matrix_h, matrix_h_im, matrix_ks, matrix_ks_im, matrix_vxc, run_rtp, rtp, matrix_h_kp, matrix_h_im_kp, matrix_ks_kp, matrix_ks_im_kp, matrix_vxc_kp, kinetic_kp, matrix_s_kp, matrix_w_kp, matrix_s_ri_aux_kp, matrix_s, matrix_s_ri_aux, matrix_w, matrix_p_mp2, matrix_p_mp2_admm, rho, rho_xc, pw_env, ewald_env, ewald_pw, active_space, mpools, input, para_env, blacs_env, scf_control, rel_control, kinetic, qs_charges, vppl, rho_core, rho_nlcc, rho_nlcc_g, ks_env, ks_qmmm_env, wf_history, scf_env, local_particles, local_molecules, distribution_2d, dbcsr_dist, molecule_kind_set, molecule_set, subsys, cp_subsys, oce, local_rho_set, rho_atom_set, task_list, task_list_soft, rho0_atom_set, rho0_mpole, rhoz_set, ecoul_1c, rho0_s_rs, rho0_s_gs, do_kpoints, has_unit_metric, requires_mo_derivs, mo_derivs, mo_loc_history, nkind, natom, nelectron_total, nelectron_spin, efield, neighbor_list_id, linres_control, xas_env, virial, cp_ddapc_env, cp_ddapc_ewald, outer_scf_history, outer_scf_ihistory, x_data, et_coupling, dftb_potential, results, se_taper, se_store_int_env, se_nddo_mpole, se_nonbond_env, admm_env, lri_env, lri_density, exstate_env, ec_env, harris_env, dispersion_env, gcp_env, vee, rho_external, external_vxc, mask, mp2_env, bs_env, kg_env, wanniercentres, atprop, ls_scf_env, do_transport, transport_env, v_hartree_rspace, s_mstruct_changed, rho_changed, potential_changed, forces_up_to_date, mscfg_env, almo_scf_env, gradient_history, variable_history, embed_pot, spin_embed_pot, polar_env, mos_last_converged, eeq, rhs)
Get the QUICKSTEP environment.
Define the quickstep kind type and their sub types.
subroutine, public get_qs_kind(qs_kind, basis_set, basis_type, ncgf, nsgf, all_potential, tnadd_potential, gth_potential, sgp_potential, upf_potential, se_parameter, dftb_parameter, xtb_parameter, dftb3_param, zatom, zeff, elec_conf, mao, lmax_dftb, alpha_core_charge, ccore_charge, core_charge, core_charge_radius, paw_proj_set, paw_atom, hard_radius, hard0_radius, max_rad_local, covalent_radius, vdw_radius, gpw_type_forced, harmonics, max_iso_not0, max_s_harm, grid_atom, ngrid_ang, ngrid_rad, lmax_rho0, dft_plus_u_atom, l_of_dft_plus_u, n_of_dft_plus_u, u_minus_j, u_of_dft_plus_u, j_of_dft_plus_u, alpha_of_dft_plus_u, beta_of_dft_plus_u, j0_of_dft_plus_u, occupation_of_dft_plus_u, dispersion, bs_occupation, magnetization, no_optimize, addel, laddel, naddel, orbitals, max_scf, eps_scf, smear, u_ramping, u_minus_j_target, eps_u_ramping, init_u_ramping_each_scf, reltmat, ghost, floating, name, element_symbol, pao_basis_size, pao_model_file, pao_potentials, pao_descriptors, nelec)
Get attributes of an atomic kind.
subroutine, public torch_dict_release(dict)
Releases a Torch dictionary and all its ressources.
subroutine, public torch_tensor_backward(tensor, outer_grad)
Runs autograd on a Torch tensor.
subroutine, public torch_dict_get(dict, key, tensor)
Retrieves a Torch tensor from a Torch dictionary.
subroutine, public torch_model_load(model, filename)
Loads a Torch model from given "*.pth" file. (In Torch lingo models are called modules)
subroutine, public torch_dict_create(dict)
Creates an empty Torch dictionary.
subroutine, public torch_tensor_grad(tensor, grad)
Returns the gradient of a Torch tensor which was computed by autograd.
subroutine, public torch_dict_insert(dict, key, tensor)
Inserts a Torch tensor into a Torch dictionary.
subroutine, public torch_tensor_release(tensor)
Releases a Torch tensor and all its ressources.
subroutine, public torch_model_forward(model, inputs, outputs)
Evaluates the given Torch model.
All kind of helpful little routines.
Provides all information about an atomic kind.
Type defining parameters related to the simulation cell.
stores all the informations relevant to an mpi environment
PAO-ML model for a single atomic kind.
Provides all information about a quickstep kind.