skala_torch_api.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief Small CP2K wrapper around the SKALA TorchScript functional protocol.
! **************************************************************************************************
MODULE skala_torch_api
#if defined (__HAS_IEEE_EXCEPTIONS)
   USE ieee_exceptions, ONLY: ieee_all, &
                              ieee_get_halting_mode, &
                              ieee_set_halting_mode
#endif
   USE kinds, ONLY: default_string_length, &
                    dp
   USE string_utilities, ONLY: uppercase
   USE torch_api, ONLY: &
      torch_dict_type, torch_model_disable_parameter_gradients, torch_model_forward_mol_tensor, &
      torch_model_freeze_preserving_method, torch_model_load_with_metadata, torch_model_release, &
      torch_model_remap_device_constants, torch_model_promote_float32, torch_model_type, &
      torch_tensor_item_double, torch_tensor_release, torch_tensor_type, &
      torch_tensor_weighted_sum
#include "../base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_torch_api'
 
   PUBLIC :: skala_torch_model_type, skala_torch_model_load, skala_torch_model_release
   PUBLIC :: skala_torch_model_get_exc, skala_torch_model_get_exc_density
   PUBLIC :: skala_torch_model_needs_feature, skala_torch_model_protocol_version
 
   TYPE skala_torch_model_type
      PRIVATE
      INTEGER                                            :: protocol_version = -1
      CHARACTER(len=default_string_length), ALLOCATABLE, &
         DIMENSION(:)                                    :: features
      TYPE(torch_model_type)                             :: torch_model
   END TYPE skala_torch_model_type
 
CONTAINS
 
! **************************************************************************************************
!> \brief Load a SKALA TorchScript model and its feature metadata.
!> \param model ...
!> \param filename ...
!> \param fp64 Promote float32 numerical state and casts; default true for differential consistency.
! **************************************************************************************************
   SUBROUTINE skala_torch_model_load(model, filename, fp64)
      TYPE(skala_torch_model_type), INTENT(INOUT)        :: model
      CHARACTER(len=*), INTENT(IN)                       :: filename
      LOGICAL, INTENT(IN), OPTIONAL                      :: fp64
 
      CHARACTER(:), ALLOCATABLE                          :: features_json, protocol_string
      INTEGER                                            :: ios
      LOGICAL                                            :: use_fp64
 
      CALL torch_model_load_with_metadata(model%torch_model, filename, &
                                          "protocol_version", protocol_string, &
                                          "features", features_json)
      CALL torch_model_remap_device_constants(model%torch_model)
      CALL torch_model_disable_parameter_gradients(model%torch_model)
      READ (protocol_string, *, iostat=ios) model%protocol_version
      IF (ios /= 0) cpabort("Could not parse SKALA TorchScript protocol_version metadata")
      IF (model%protocol_version /= 2) THEN
         cpabort("Unsupported SKALA TorchScript protocol version")
      END IF
 
      CALL parse_feature_list(features_json, model%features)
      use_fp64 = .true.
      IF (PRESENT(fp64)) use_fp64 = fp64
      IF (use_fp64) CALL torch_model_promote_float32(model%torch_model)
      ! Preserve the exported SKALA entry point while folding constant model state.
      CALL torch_model_freeze_preserving_method(model%torch_model, "get_exc_density")
 
   END SUBROUTINE skala_torch_model_load
 
! **************************************************************************************************
!> \brief Release a loaded SKALA TorchScript model.
!> \param model ...
! **************************************************************************************************
   SUBROUTINE skala_torch_model_release(model)
      TYPE(skala_torch_model_type), INTENT(INOUT)        :: model
 
      CALL torch_model_release(model%torch_model)
      IF (ALLOCATED(model%features)) DEALLOCATE (model%features)
      model%protocol_version = -1
 
   END SUBROUTINE skala_torch_model_release
 
! **************************************************************************************************
!> \brief Check whether a loaded SKALA model requests a feature.
!> \param model ...
!> \param feature ...
!> \return ...
! **************************************************************************************************
   FUNCTION skala_torch_model_needs_feature(model, feature) RESULT(needs_feature)
      TYPE(skala_torch_model_type), INTENT(IN)           :: model
      CHARACTER(len=*), INTENT(IN)                       :: feature
      LOGICAL                                            :: needs_feature
 
      CHARACTER(len=default_string_length)               :: feature_key, model_feature
      INTEGER                                            :: i
 
      feature_key = adjustl(feature)
      CALL uppercase(feature_key)
 
      needs_feature = .false.
      IF (.NOT. ALLOCATED(model%features)) RETURN
 
      DO i = 1, SIZE(model%features)
         model_feature = adjustl(model%features(i))
         CALL uppercase(model_feature)
         IF (trim(model_feature) == trim(feature_key)) THEN
            needs_feature = .true.
            RETURN
         END IF
      END DO
 
   END FUNCTION skala_torch_model_needs_feature
 
! **************************************************************************************************
!> \brief Return the loaded SKALA TorchScript protocol version.
!> \param model ...
!> \return ...
! **************************************************************************************************
   FUNCTION skala_torch_model_protocol_version(model) RESULT(protocol_version)
      TYPE(skala_torch_model_type), INTENT(IN)           :: model
      INTEGER                                            :: protocol_version
 
      protocol_version = model%protocol_version
 
   END FUNCTION skala_torch_model_protocol_version
 
! **************************************************************************************************
!> \brief Evaluate the SKALA exchange-correlation energy density.
!> \param model ...
!> \param inputs ...
!> \param exc_density ...
! **************************************************************************************************
   SUBROUTINE skala_torch_model_get_exc_density(model, inputs, exc_density)
      TYPE(skala_torch_model_type), INTENT(INOUT)        :: model
      TYPE(torch_dict_type), INTENT(IN)                  :: inputs
      TYPE(torch_tensor_type), INTENT(INOUT)             :: exc_density
 
#if defined (__HAS_IEEE_EXCEPTIONS)
      LOGICAL, DIMENSION(5)                              :: ieee_halt
 
      CALL ieee_get_halting_mode(ieee_all, ieee_halt)
      CALL ieee_set_halting_mode(ieee_all, .false.)
#endif
      CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
#if defined (__HAS_IEEE_EXCEPTIONS)
      CALL ieee_set_halting_mode(ieee_all, ieee_halt)
#endif
 
   END SUBROUTINE skala_torch_model_get_exc_density
 
! **************************************************************************************************
!> \brief Evaluate the weighted SKALA exchange-correlation energy.
!> \param model ...
!> \param inputs ...
!> \param grid_weights ...
!> \param exc_tensor ...
!> \param exc ...
! **************************************************************************************************
   SUBROUTINE skala_torch_model_get_exc(model, inputs, grid_weights, exc_tensor, exc)
      TYPE(skala_torch_model_type), INTENT(INOUT)        :: model
      TYPE(torch_dict_type), INTENT(IN)                  :: inputs
      TYPE(torch_tensor_type), INTENT(IN)                :: grid_weights
      TYPE(torch_tensor_type), INTENT(INOUT)             :: exc_tensor
      REAL(kind=dp), INTENT(OUT)                         :: exc
 
      TYPE(torch_tensor_type)                            :: exc_density
 
#if defined (__HAS_IEEE_EXCEPTIONS)
      LOGICAL, DIMENSION(5)                              :: ieee_halt
 
      CALL ieee_get_halting_mode(ieee_all, ieee_halt)
      CALL ieee_set_halting_mode(ieee_all, .false.)
#endif
      CALL torch_model_forward_mol_tensor(model%torch_model, "get_exc_density", inputs, exc_density)
      CALL torch_tensor_weighted_sum(exc_density, grid_weights, exc_tensor)
      CALL torch_tensor_release(exc_density)
      exc = torch_tensor_item_double(exc_tensor)
#if defined (__HAS_IEEE_EXCEPTIONS)
      CALL ieee_set_halting_mode(ieee_all, ieee_halt)
#endif
 
   END SUBROUTINE skala_torch_model_get_exc
 
! **************************************************************************************************
!> \brief Parse a TorchScript extra_files JSON list of feature names.
!> \param features_json ...
!> \param features ...
! **************************************************************************************************
   SUBROUTINE parse_feature_list(features_json, features)
      CHARACTER(len=*), INTENT(IN)                       :: features_json
      CHARACTER(len=default_string_length), &
         ALLOCATABLE, DIMENSION(:), INTENT(OUT)          :: features
 
      INTEGER                                            :: end_pos, feature_count, i, pos, quote1, &
                                                            quote2, start_pos
 
      feature_count = 0
      pos = 1
      DO
         quote1 = index(features_json(pos:), '"')
         IF (quote1 == 0) EXIT
         start_pos = pos + quote1
         quote2 = index(features_json(start_pos:), '"')
         IF (quote2 == 0) EXIT
         feature_count = feature_count + 1
         pos = start_pos + quote2
      END DO
 
      IF (feature_count == 0) cpabort("SKALA TorchScript model does not list any features")
      ALLOCATE (features(feature_count))
      features = ""
 
      pos = 1
      DO i = 1, feature_count
         quote1 = index(features_json(pos:), '"')
         start_pos = pos + quote1
         quote2 = index(features_json(start_pos:), '"')
         end_pos = start_pos + quote2 - 2
         features(i) = features_json(start_pos:end_pos)
         pos = start_pos + quote2
      END DO
 
   END SUBROUTINE parse_feature_list
 
END MODULE skala_torch_api