diff --git a/config_src/external/database_comms/MOM_database_comms.F90 b/config_src/external/database_comms/MOM_database_comms.F90 new file mode 100644 index 0000000000..4c3eb38b5c --- /dev/null +++ b/config_src/external/database_comms/MOM_database_comms.F90 @@ -0,0 +1,37 @@ +!> Contains routines necessary to initialize communication with a database +module MOM_database_comms +! This file is part of MOM6. See LICENSE.md for the license. +use MOM_file_parser, only : param_file_type +use MOM_error_handler, only : MOM_error, WARNING +use database_client_interface, only : dbclient_type + +implicit none; private + +!> Control structure to store Database communication related parameters and objects +type, public :: dbcomms_CS_type + type(dbclient_type) :: client !< The Database client itself + logical :: use_dbclient !< If True, use Database within MOM6 + logical :: colocated !< If True, the orchestrator was setup in 'co-located' mode + logical :: cluster !< If True, the orchestrator has three shards or more + integer :: colocated_stride !< Sets which ranks will load the model from the file + !! e.g. mod(rank,colocated_stride) == 0 +end type dbcomms_CS_type + +public :: database_comms_init +public :: dbclient_type + +contains + +subroutine database_comms_init(param_file, CS, client_in) + type(param_file_type), intent(in ) :: param_file !< Parameter file structure + type(dbcomms_CS_type), intent(inout) :: CS !< Control structure for Database + type(dbclient_type), optional, intent(in ) :: client_in !< If present, use a previously initialized + !! Database client + + call MOM_error(WARNING,"dbcomms_init was compiled using the dummy module. If this was\n"//& + "a mistake, please follow the instructions in:\n"//& + "MOM6/config_src/external/dbclient/README.md") +end subroutine database_comms_init + +end module MOM_database_comms + diff --git a/config_src/external/database_comms/README.md b/config_src/external/database_comms/README.md new file mode 100644 index 0000000000..05f1f07259 --- /dev/null +++ b/config_src/external/database_comms/README.md @@ -0,0 +1,25 @@ +# Overview +This module is designed to be used in conjunction with the SmartSim and +SmartRedis libraries found at https://github.com/CrayLabs/. These +libraries are used to perform machine-learning inference and online +analysis using a Redis-based database. + +An earlier implementation of these routines was used in Partee et al. [2022]: +"Using Machine Learning at scale in numerical simulations with SmartSim: +An application to ocean climate modeling" (doi.org/10.1016/j.jocs.2022.101707) +to predict eddy kinetic energy for use in the MEKE module. The additional +scripts and installation instructions for compiling MOM6 for this case can +be found at: https://github.com/CrayLabs/NCAR_ML_EKE/. The substantive +code in the new implementation is part of `MOM_MEKE.F90`. + +# File description + +- `MOM_database_comms` contains just method signatures and elements of the + control structure that are imported elsewhere within the primary MOM6 + code. This includes: `dbcomms_CS_type`, `dbclient_type`, and `database_comms_init` + +- `database_client_interface.F90` contains the methods for a communication client + to transfer data and/or commands between MOM6 and a remote database. This is + roughly based on the SmartRedis library, though only the methods that are most + likely to be used with MOM6 are retained. This is to ensure that the API can be + tested without requiring MOM6 users to compile in the the full library. diff --git a/config_src/external/database_comms/database_client_interface.F90 b/config_src/external/database_comms/database_client_interface.F90 new file mode 100644 index 0000000000..9b57628921 --- /dev/null +++ b/config_src/external/database_comms/database_client_interface.F90 @@ -0,0 +1,814 @@ +module database_client_interface + +! This file is part of MOM6. See LICENSE.md for the license. + use iso_fortran_env, only : int8, int16, int32, int64, real32, real64 + + implicit none; private + + !> Dummy type for dataset + type, public :: dataset_type + private + end type dataset_type + + !> Stores all data and methods associated with the communication client that is used to communicate with the database + type, public :: dbclient_type + private + + contains + + ! Public procedures + !> Puts a tensor into the database for a variety of datatypes + generic :: put_tensor => put_tensor_float_1d, put_tensor_float_2d, put_tensor_float_3d, put_tensor_float_4d, & + put_tensor_double_1d, put_tensor_double_2d, put_tensor_double_3d, put_tensor_double_4d, & + put_tensor_int32_1d, put_tensor_int32_2d, put_tensor_int32_3d, put_tensor_int32_4d + !> Retrieve the tensor in the database into already allocated memory for a variety of datatypesm + generic :: unpack_tensor => unpack_tensor_float_1d, unpack_tensor_float_2d, & + unpack_tensor_float_3d, unpack_tensor_float_4d, & + unpack_tensor_double_1d, unpack_tensor_double_2d, & + unpack_tensor_double_3d, unpack_tensor_double_4d, & + unpack_tensor_int32_1d, unpack_tensor_int32_2d, & + unpack_tensor_int32_3d, unpack_tensor_int32_4d + + !> Decode a response code from an API function + procedure :: SR_error_parser + !> Initializes a new instance of the communication client + procedure :: initialize => initialize_client + !> Check if a communication client has been initialized + procedure :: isinitialized + !> Destructs a new instance of the communication client + procedure :: destructor + !> Rename a tensor within the database + procedure :: rename_tensor + !> Delete a tensor from the database + procedure :: delete_tensor + !> Copy a tensor within the database to a new name + procedure :: copy_tensor + !> Set a model from a file + procedure :: set_model_from_file + !> Set a model from a file on a system with multiple GPUs + procedure :: set_model_from_file_multigpu + !> Set a model from a byte string that has been loaded within the application + procedure :: set_model + !> Set a model from a byte string that has been loaded within the application on a system with multiple GPUs + procedure :: set_model_multigpu + !> Retrieve the model as a byte string + procedure :: get_model + !> Set a script from a specified file + procedure :: set_script_from_file + !> Set a script from a specified file on a system with multiple GPUS + procedure :: set_script_from_file_multigpu + !> Set a script as a byte or text string + procedure :: set_script + !> Set a script as a byte or text string on a system with multiple GPUs + procedure :: set_script_multigpu + !> Retrieve the script from the database + procedure :: get_script + !> Run a script that has already been stored in the database + procedure :: run_script + !> Run a script that has already been stored in the database with multiple GPUs + procedure :: run_script_multigpu + !> Run a model that has already been stored in the database + procedure :: run_model + !> Run a model that has already been stored in the database with multiple GPUs + procedure :: run_model_multigpu + !> Remove a script from the database + procedure :: delete_script + !> Remove a script from the database with multiple GPUs + procedure :: delete_script_multigpu + !> Remove a model from the database + procedure :: delete_model + !> Remove a model from the database with multiple GPUs + procedure :: delete_model_multigpu + !> Put a communication dataset into the database + procedure :: put_dataset + !> Retrieve a communication dataset from the database + procedure :: get_dataset + !> Rename the dataset within the database + procedure :: rename_dataset + !> Copy a dataset stored in the database into another name + procedure :: copy_dataset + !> Delete the dataset from the database + procedure :: delete_dataset + + ! Private procedures + !> Put a 1d, 32-bit real tensor into database + procedure, private :: put_tensor_float_1d + !> Put a 2d, 32-bit real tensor into database + procedure, private :: put_tensor_float_2d + !> Put a 3d, 32-bit real tensor into database + procedure, private :: put_tensor_float_3d + !> Put a 4d, 32-bit real tensor into database + procedure, private :: put_tensor_float_4d + !> Put a 1d, 64-bit real tensor into database + procedure, private :: put_tensor_double_1d + !> Put a 2d, 64-bit real tensor into database + procedure, private :: put_tensor_double_2d + !> Put a 3d, 64-bit real tensor into database + procedure, private :: put_tensor_double_3d + !> Put a 4d, 64-bit real tensor into database + procedure, private :: put_tensor_double_4d + !> Put a 1d, 32-bit integer tensor into database + procedure, private :: put_tensor_int32_1d + !> Put a 2d, 32-bit integer tensor into database + procedure, private :: put_tensor_int32_2d + !> Put a 3d, 32-bit integer tensor into database + procedure, private :: put_tensor_int32_3d + !> Put a 4d, 32-bit integer tensor into database + procedure, private :: put_tensor_int32_4d + !> Unpack a 1d, 32-bit real tensor from the database + procedure, private :: unpack_tensor_float_1d + !> Unpack a 2d, 32-bit real tensor from the database + procedure, private :: unpack_tensor_float_2d + !> Unpack a 3d, 32-bit real tensor from the database + procedure, private :: unpack_tensor_float_3d + !> Unpack a 4d, 32-bit real tensor from the database + procedure, private :: unpack_tensor_float_4d + !> Unpack a 1d, 64-bit real tensor from the database + procedure, private :: unpack_tensor_double_1d + !> Unpack a 2d, 64-bit real tensor from the database + procedure, private :: unpack_tensor_double_2d + !> Unpack a 3d, 64-bit real tensor from the database + procedure, private :: unpack_tensor_double_3d + !> Unpack a 4d, 64-bit real tensor from the database + procedure, private :: unpack_tensor_double_4d + !> Unpack a 1d, 32-bit integer tensor from the database + procedure, private :: unpack_tensor_int32_1d + !> Unpack a 2d, 32-bit integer tensor from the database + procedure, private :: unpack_tensor_int32_2d + !> Unpack a 3d, 32-bit integer tensor from the database + procedure, private :: unpack_tensor_int32_3d + !> Unpack a 4d, 32-bit integer tensor from the database + procedure, private :: unpack_tensor_int32_4d + + end type dbclient_type + + contains + + !> Decode a response code from an API function + function SR_error_parser(self, response_code) result(is_error) + class(dbclient_type), intent(in) :: self !< Receives the initialized client + integer, intent(in) :: response_code !< The response code to decode + logical :: is_error !< Indicates whether this is an error response + + is_error = .true. + end function SR_error_parser + + !> Initializes a new instance of a communication client + function initialize_client(self, cluster) + integer :: initialize_client + class(dbclient_type), intent(inout) :: self !< Receives the initialized client + logical, optional, intent(in ) :: cluster !< If true, client uses a database cluster (Default: .false.) + + initialize_client = -1 + end function initialize_client + + !> Check whether the client has been initialized + logical function isinitialized(this) + class(dbclient_type) :: this + isinitialized = .false. + end function isinitialized + + !> A destructor for the communication client + function destructor(self) + integer :: destructor + class(dbclient_type), intent(inout) :: self + + destructor = -1 + end function destructor + + !> Put a 32-bit real 1d tensor into the database + function put_tensor_float_1d(self, name, data, dims) result(code) + real(kind=real32), dimension(:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_float_1d + + !> Put a 32-bit real 2d tensor into the database + function put_tensor_float_2d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_float_2d + + !> Put a 32-bit real 3d tensor into the database + function put_tensor_float_3d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_float_3d + + !> Put a 32-bit real 4d tensor into the database + function put_tensor_float_4d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_float_4d + + !> Put a 64-bit real 1d tensor into the database + function put_tensor_double_1d(self, name, data, dims) result(code) + real(kind=real64), dimension(:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_double_1d + + !> Put a 64-bit real 2d tensor into the database + function put_tensor_double_2d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_double_2d + + !> Put a 64-bit real 3d tensor into the database + function put_tensor_double_3d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_double_3d + + !> Put a 64-bit real 4d tensor into the database + function put_tensor_double_4d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_double_4d + + !> Put a 32-bit integer 1d tensor into the database + function put_tensor_int32_1d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_int32_1d + + !> Put a 32-bit integer 2d tensor into the database + function put_tensor_int32_2d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_int32_2d + + !> Put a 32-bit integer 3d tensor into the database + function put_tensor_int32_3d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_int32_3d + + !> Put a 32-bit integer 4d tensor into the database + function put_tensor_int32_4d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:,:,:), intent(in) :: data !< Data to be sent + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function put_tensor_int32_4d + + !> Unpack a 32-bit real 1d tensor from the database + function unpack_tensor_float_1d(self, name, data, dims) result(code) + real(kind=real32), dimension(:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_float_1d + + !> Unpack a 32-bit real 2d tensor from the database + function unpack_tensor_float_2d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_float_2d + + !> Unpack a 32-bit real 3d tensor from the database + function unpack_tensor_float_3d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_float_3d + + !> Unpack a 32-bit real 4d tensor from the database + function unpack_tensor_float_4d(self, name, data, dims) result(code) + real(kind=real32), dimension(:,:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_float_4d + + !> Unpack a 64-bit real 1d tensor from the database + function unpack_tensor_double_1d(self, name, data, dims) result(code) + real(kind=real64), dimension(:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_double_1d + + !> Unpack a 64-bit real 2d tensor from the database + function unpack_tensor_double_2d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_double_2d + + !> Unpack a 64-bit real 3d tensor from the database + function unpack_tensor_double_3d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_double_3d + + !> Unpack a 64-bit real 4d tensor from the database + function unpack_tensor_double_4d(self, name, data, dims) result(code) + real(kind=real64), dimension(:,:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_double_4d + + !> Unpack a 32-bit integer 1d tensor from the database + function unpack_tensor_int32_1d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_int32_1d + + !> Unpack a 32-bit integer 2d tensor from the database + function unpack_tensor_int32_2d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_int32_2d + + !> Unpack a 32-bit integer 3d tensor from the database + function unpack_tensor_int32_3d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_int32_3d + + !> Unpack a 32-bit integer 4d tensor from the database + function unpack_tensor_int32_4d(self, name, data, dims) result(code) + integer(kind=int32), dimension(:,:,:,:), intent( out) :: data !< Data to be received + class(dbclient_type), intent(in) :: self !< Fortran communication client + character(len=*), intent(in) :: name !< The unique name used to store in the database + integer, dimension(:), intent(in) :: dims !< The length of each dimension + integer :: code + + code = -1 + end function unpack_tensor_int32_4d + + !> Move a tensor to a new name + function rename_tensor(self, old_name, new_name) result(code) + class(dbclient_type), intent(in) :: self !< The initialized Fortran communication client + character(len=*), intent(in) :: old_name !< The current name for the tensor + !! excluding null terminating character + character(len=*), intent(in) :: new_name !< The new tensor name + integer :: code + + code = -1 + end function rename_tensor + + !> Delete a tensor + function delete_tensor(self, name) result(code) + class(dbclient_type), intent(in) :: self !< The initialized Fortran communication client + character(len=*), intent(in) :: name !< The name associated with the tensor + integer :: code + + code = -1 + end function delete_tensor + + !> Copy a tensor to the destination name + function copy_tensor(self, src_name, dest_name) result(code) + class(dbclient_type), intent(in) :: self !< The initialized Fortran communication client + character(len=*), intent(in) :: src_name !< The name associated with the tensor + !! excluding null terminating character + character(len=*), intent(in) :: dest_name !< The new tensor name + integer :: code + + code = -1 + end function copy_tensor + + !> Retrieve the model from the database + function get_model(self, name, model) result(code) + class(dbclient_type), intent(in ) :: self !< An initialized communication client + character(len=*), intent(in ) :: name !< The name associated with the model + character(len=*), intent( out) :: model !< The model as a continuous buffer + integer :: code + + code = -1 + end function get_model + + !> Load the machine learning model from a file and set the configuration + function set_model_from_file(self, name, model_file, backend, device, batch_size, min_batch_size, tag, & + inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), intent(in) :: model_file !< The file storing the model + character(len=*), intent(in) :: backend !< The name of the backend + !! (TF, TFLITE, TORCH, ONNX) + character(len=*), intent(in) :: device !< The name of the device + !! (CPU, GPU, GPU:0, GPU:1...) + integer, optional, intent(in) :: batch_size !< The batch size for model execution + integer, optional, intent(in) :: min_batch_size !< The minimum batch size for model execution + character(len=*), optional, intent(in) :: tag !< A tag to attach to the model for + !! information purposes + character(len=*), dimension(:), optional, intent(in) :: inputs !< One or more names of model + !! input nodes (TF models) + character(len=*), dimension(:), optional, intent(in) :: outputs !< One or more names of model + !! output nodes (TF models) + integer :: code + + code = -1 + end function set_model_from_file + + !> Load the machine learning model from a file and set the configuration for use in multi-GPU systems + function set_model_from_file_multigpu(self, name, model_file, backend, first_gpu, num_gpus, batch_size, & + min_batch_size, tag, inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), intent(in) :: model_file !< The file storing the model + character(len=*), intent(in) :: backend !< The name of the backend + !! (TF, TFLITE, TORCH, ONNX) + integer, intent(in) :: first_gpu !< The first GPU (zero-based) + !! to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer, optional, intent(in) :: batch_size !< The batch size for model execution + integer, optional, intent(in) :: min_batch_size !< The minimum batch size for model execution + character(len=*), optional, intent(in) :: tag !< A tag to attach to the model for + !! information purposes + character(len=*), dimension(:), optional, intent(in) :: inputs !< One or more names of model + !! input nodes (TF models) + character(len=*), dimension(:), optional, intent(in) :: outputs !< One or more names of model + !! output nodes (TF models) + integer :: code + + code = -1 + end function set_model_from_file_multigpu + + !> Establish a model to run + function set_model(self, name, model, backend, device, batch_size, min_batch_size, tag, & + inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), intent(in) :: model !< The binary representation of the model + character(len=*), intent(in) :: backend !< The name of the backend (TF, TFLITE, TORCH, ONNX) + character(len=*), intent(in) :: device !< The name of the device (CPU, GPU, GPU:0, GPU:1...) + integer, intent(in) :: batch_size !< The batch size for model execution + integer, intent(in) :: min_batch_size !< The minimum batch size for model execution + character(len=*), intent(in) :: tag !< A tag to attach to the model for + !! information purposes + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of model input nodes (TF models) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of model output nodes (TF models) + integer :: code + + code = -1 + end function set_model + + !> Set a model from a byte string to run on a system with multiple GPUs + function set_model_multigpu(self, name, model, backend, first_gpu, num_gpus, batch_size, min_batch_size, tag, & + inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), intent(in) :: model !< The binary representation of the model + character(len=*), intent(in) :: backend !< The name of the backend (TF, TFLITE, TORCH, ONNX) + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer, intent(in) :: batch_size !< The batch size for model execution + integer, intent(in) :: min_batch_size !< The minimum batch size for model execution + character(len=*), intent(in) :: tag !< A tag to attach to the model for + !! information purposes + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of model input nodes (TF models) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of model output nodes (TF models) + integer :: code + + code = -1 + end function set_model_multigpu + + !> Run a model in the database using the specified input and output tensors + function run_model(self, name, inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of model input nodes (TF models) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of model output nodes (TF models) + integer :: code + + code = -1 + end function run_model + + !> Run a model in the database using the specified input and output tensors in a multi-GPU system + function run_model_multigpu(self, name, inputs, outputs, offset, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the model + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of model input nodes (TF models) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of model output nodes (TF models) + integer, intent(in) :: offset !< Index of the current image, such as a processor ID + !! or MPI rank + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function run_model_multigpu + + !> Remove a model from the database + function delete_model(self, name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to remove the model + integer :: code + + code = -1 + end function delete_model + + !> Remove a model from the database + function delete_model_multigpu(self, name, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to remove the model + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function delete_model_multigpu + + !> Retrieve the script from the database + function get_script(self, name, script) result(code) + class(dbclient_type), intent(in ) :: self !< An initialized communication client + character(len=*), intent(in ) :: name !< The name to use to place the script + character(len=*), intent( out) :: script !< The script as a continuous buffer + integer :: code + + code = -1 + end function get_script + + !> Set a script (from file) in the database for future execution + function set_script_from_file(self, name, device, script_file) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: device !< The name of the device (CPU, GPU, GPU:0, GPU:1...) + character(len=*), intent(in) :: script_file !< The file storing the script + integer :: code + + code = -1 + end function set_script_from_file + + !> Set a script (from file) in the database for future execution in a multi-GPU system + function set_script_from_file_multigpu(self, name, script_file, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: script_file !< The file storing the script + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function set_script_from_file_multigpu + + !> Set a script (from buffer) in the database for future execution + function set_script(self, name, device, script) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: device !< The name of the device (CPU, GPU, GPU:0, GPU:1...) + character(len=*), intent(in) :: script !< The file storing the script + integer :: code + + code = -1 + end function set_script + + !> Set a script (from buffer) in the database for future execution in a multi-GPU system + function set_script_multigpu(self, name, script, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: script !< The file storing the script + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function set_script_multigpu + + function run_script(self, name, func, inputs, outputs) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: func !< The name of the function in the script to call + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of script + !! input nodes (TF scripts) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of script + !! output nodes (TF scripts) + integer :: code + + code = -1 + end function run_script + + function run_script_multigpu(self, name, func, inputs, outputs, offset, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to place the script + character(len=*), intent(in) :: func !< The name of the function in the script to call + character(len=*), dimension(:), intent(in) :: inputs !< One or more names of script + !! input nodes (TF scripts) + character(len=*), dimension(:), intent(in) :: outputs !< One or more names of script + !! output nodes (TF scripts) + integer, intent(in) :: offset !< Index of the current image, such as a processor ID + !! or MPI rank + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function run_script_multigpu + + !> Remove a script from the database + function delete_script(self, name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to delete the script + integer :: code + + code = -1 + end function delete_script + + !> Remove a script_multigpu from the database + function delete_script_multigpu(self, name, first_gpu, num_gpus) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< The name to use to delete the script_multigpu + integer, intent(in) :: first_gpu !< The first GPU (zero-based) to use with the model + integer, intent(in) :: num_gpus !< The number of GPUs to use with the model + integer :: code + + code = -1 + end function delete_script_multigpu + + !> Store a dataset in the database + function put_dataset(self, dataset) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + type(dataset_type), intent(in) :: dataset !< Dataset to store in the dataset + integer :: code + + code = -1 + end function put_dataset + + !> Retrieve a dataset from the database + function get_dataset(self, name, dataset) result(code) + class(dbclient_type), intent(in ) :: self !< An initialized communication client + character(len=*), intent(in ) :: name !< Name of the dataset to get + type(dataset_type), intent( out) :: dataset !< receives the dataset + integer :: code + + code = -1 + end function get_dataset + + !> Rename a dataset stored in the database + function rename_dataset(self, name, new_name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< Original name of the dataset + character(len=*), intent(in) :: new_name !< New name of the dataset + integer :: code + + code = -1 + end function rename_dataset + + !> Copy a dataset within the database to a new name + function copy_dataset(self, name, new_name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< Source name of the dataset + character(len=*), intent(in) :: new_name !< Name of the new dataset + integer :: code + + code = -1 + end function copy_dataset + + !> Delete a dataset stored within a database + function delete_dataset(self, name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: name !< Name of the dataset to delete + integer :: code + + code = -1 + end function delete_dataset + + !> Appends a dataset to the aggregation list When appending a dataset to an aggregation list, the list will + !! automatically be created if it does not exist (i.e. this is the first entry in the list). Aggregation + !! lists work by referencing the dataset by storing its key, so appending a dataset to an aggregation list + !! does not create a copy of the dataset. Also, for this reason, the dataset must have been previously + !! placed into the database with a separate call to put_dataset(). + function append_to_list(self, list_name, dataset) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: list_name !< Name of the dataset to get + type(dataset_type), intent(in) :: dataset !< Dataset to append to the list + integer :: code + + code = -1 + end function append_to_list + + !> Delete an aggregation list + function delete_list(self, list_name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: list_name !< Name of the aggregated dataset list to delete + integer :: code + + code = -1 + end function delete_list + + !> Copy an aggregation list + function copy_list(self, src_name, dest_name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: src_name !< Name of the dataset to copy + character(len=*), intent(in) :: dest_name !< The new list name + integer :: code + + code = -1 + end function copy_list + + !> Rename an aggregation list + function rename_list(self, src_name, dest_name) result(code) + class(dbclient_type), intent(in) :: self !< An initialized communication client + character(len=*), intent(in) :: src_name !< Name of the dataset to rename + character(len=*), intent(in) :: dest_name !< The new list name + integer :: code + + code = -1 + end function rename_list + + end module database_client_interface + diff --git a/src/core/MOM.F90 b/src/core/MOM.F90 index 8fc6600b69..78170064ff 100644 --- a/src/core/MOM.F90 +++ b/src/core/MOM.F90 @@ -143,6 +143,9 @@ module MOM use MOM_porous_barriers, only : porous_widths +! Database client used for machine-learning interface +use MOM_database_comms, only : dbcomms_CS_type, database_comms_init, dbclient_type + ! ODA modules use MOM_oda_driver_mod, only : ODA_CS, oda, init_oda, oda_end use MOM_oda_driver_mod, only : set_prior_tracer, set_analysis_time, apply_oda_tracer_increments @@ -251,6 +254,8 @@ module MOM logical :: offline_tracer_mode = .false. !< If true, step_offline() is called instead of step_MOM(). !! This is intended for running MOM6 in offline tracer mode + logical :: MEKE_in_dynamics !< If .true. (default), MEKE is called in the dynamics routine otherwise + !! it is called during the tracer dynamics type(time_type), pointer :: Time !< pointer to the ocean clock real :: dt !< (baroclinic) dynamics time step [T ~> s] @@ -338,6 +343,7 @@ module MOM !! higher values use more appropriate expressions that differ at !! roundoff for non-Boussinesq cases. logical :: use_particles !< Turns on the particles package + logical :: use_dbclient !< Turns on the database client used for ML inference/analysis character(len=10) :: particle_type !< Particle types include: surface(default), profiling and sail drone. type(MOM_diag_IDs) :: IDs !< Handles used for diagnostics. @@ -404,15 +410,14 @@ module MOM type(ODA_CS), pointer :: odaCS => NULL() !< a pointer to the control structure for handling !! ensemble model state vectors and data assimilation !! increments and priors + type(dbcomms_CS_type) :: dbcomms_CS !< Control structure for database client used for online ML/AI type(porous_barrier_ptrs) :: pbv !< porous barrier fractional cell metrics - real ALLOCABLE_, dimension(NIMEMB_PTR_,NJMEM_,NKMEM_) & - :: por_face_areaU !< fractional open area of U-faces [nondim] - real ALLOCABLE_, dimension(NIMEM_,NJMEMB_PTR_,NKMEM_) & - :: por_face_areaV !< fractional open area of V-faces [nondim] - real ALLOCABLE_, dimension(NIMEMB_PTR_,NJMEM_,NK_INTERFACE_) & - :: por_layer_widthU !< fractional open width of U-faces [nondim] - real ALLOCABLE_, dimension(NIMEM_,NJMEMB_PTR_,NK_INTERFACE_) & - :: por_layer_widthV !< fractional open width of V-faces [nondim] + real ALLOCABLE_, dimension(NIMEMB_PTR_,NJMEM_,NKMEM_) :: por_face_areaU !< fractional open area of U-faces [nondim] + real ALLOCABLE_, dimension(NIMEM_,NJMEMB_PTR_,NKMEM_) :: por_face_areaV !< fractional open area of V-faces [nondim] + real ALLOCABLE_, dimension(NIMEMB_PTR_,NJMEM_,NK_INTERFACE_) :: por_layer_widthU !< fractional open width + !! of U-faces [nondim] + real ALLOCABLE_, dimension(NIMEM_,NJMEMB_PTR_,NK_INTERFACE_) :: por_layer_widthV !< fractional open width + !! of V-faces [nondim] type(particles), pointer :: particles => NULL() ! NULL() !< a pointer to the stochastics control structure end type MOM_control_struct @@ -1214,8 +1219,11 @@ subroutine step_MOM_dynamics(forces, p_surf_begin, p_surf_end, dt, dt_thermo, & ! for vertical remapping may need to be regenerated. call diag_update_remap_grids(CS%diag) - if (CS%useMEKE) call step_forward_MEKE(CS%MEKE, h, CS%VarMix%SN_u, CS%VarMix%SN_v, & - CS%visc, dt, G, GV, US, CS%MEKE_CSp, CS%uhtr, CS%vhtr) + if (CS%useMEKE .and. CS%MEKE_in_dynamics) then + call step_forward_MEKE(CS%MEKE, h, CS%VarMix%SN_u, CS%VarMix%SN_v, & + CS%visc, dt, G, GV, US, CS%MEKE_CSp, CS%uhtr, CS%vhtr, & + CS%u, CS%v, CS%tv, Time_local) + endif call disable_averaging(CS%diag) ! Advance the dynamics time by dt. @@ -1320,6 +1328,12 @@ subroutine step_MOM_tracer_dyn(CS, G, GV, US, h, Time_local) CS%t_dyn_rel_adv = 0.0 call cpu_clock_end(id_clock_tracer) ; call cpu_clock_end(id_clock_thermo) + if (CS%useMEKE .and. (.not. CS%MEKE_in_dynamics)) then + call step_forward_MEKE(CS%MEKE, h, CS%VarMix%SN_u, CS%VarMix%SN_v, & + CS%visc, CS%t_dyn_rel_adv, G, GV, US, CS%MEKE_CSp, CS%uhtr, CS%vhtr, & + CS%u, CS%v, CS%tv, Time_local) + endif + if (associated(CS%tv%T)) then call extract_diabatic_member(CS%diabatic_CSp, diabatic_halo=halo_sz) if (halo_sz > 0) then @@ -2190,6 +2204,10 @@ subroutine initialize_MOM(Time, Time_init, param_file, dirs, CS, restart_CSp, & "vertical grid files. Other values are invalid.", default=1) if (write_geom<0 .or. write_geom>2) call MOM_error(FATAL,"MOM: "//& "WRITE_GEOM must be equal to 0, 1 or 2.") + call get_param(param_file, "MOM", "USE_DBCLIENT", CS%use_dbclient, & + "If true, initialize a client to a remote database that can "//& + "be used for online analysis and machine-learning inference.",& + default=.false.) ! Check for inconsistent parameter settings. if (CS%use_ALE_algorithm .and. bulkmixedlayer) call MOM_error(FATAL, & @@ -2794,7 +2812,9 @@ subroutine initialize_MOM(Time, Time_init, param_file, dirs, CS, restart_CSp, & endif call cpu_clock_end(id_clock_MOM_init) - CS%useMEKE = MEKE_init(Time, G, US, param_file, diag, CS%MEKE_CSp, CS%MEKE, restart_CSp) + if (CS%use_dbclient) call database_comms_init(param_file, CS%dbcomms_CS) + CS%useMEKE = MEKE_init(Time, G, US, param_file, diag, CS%dbcomms_CS, CS%MEKE_CSp, CS%MEKE, & + restart_CSp, CS%MEKE_in_dynamics) call VarMix_init(Time, G, GV, US, param_file, diag, CS%VarMix) call set_visc_init(Time, G, GV, US, param_file, diag, CS%visc, CS%set_visc_CSp, restart_CSp, CS%OBC) diff --git a/src/parameterizations/lateral/MOM_MEKE.F90 b/src/parameterizations/lateral/MOM_MEKE.F90 index 2661251766..9b024e62b0 100644 --- a/src/parameterizations/lateral/MOM_MEKE.F90 +++ b/src/parameterizations/lateral/MOM_MEKE.F90 @@ -4,21 +4,33 @@ module MOM_MEKE ! This file is part of MOM6. See LICENSE.md for the license. +use iso_fortran_env, only : real32 + +use MOM_coms, only : PE_here +use MOM_database_comms, only : dbclient_type, dbcomms_CS_type +use MOM_debugging, only : hchksum, uvchksum +use MOM_cpu_clock, only : cpu_clock_id, cpu_clock_begin, cpu_clock_end, CLOCK_ROUTINE +use MOM_diag_mediator, only : post_data, register_diag_field, safe_alloc_ptr +use MOM_diag_mediator, only : diag_ctrl, time_type +use MOM_domains, only : create_group_pass, do_group_pass, group_pass_type +use MOM_domains, only : pass_vector, pass_var +use MOM_error_handler, only : MOM_error, FATAL, WARNING, NOTE, MOM_mesg, is_root_pe +use MOM_file_parser, only : read_param, get_param, log_version, param_file_type +use MOM_grid, only : ocean_grid_type +use MOM_hor_index, only : hor_index_type +use MOM_interface_heights, only : find_eta +use MOM_interpolate, only : init_external_field, time_interp_external +use MOM_interpolate, only : time_interp_external_init +use MOM_io, only : vardesc, var_desc, slasher +use MOM_isopycnal_slopes, only : calc_isoneutral_slopes +use MOM_restart, only : MOM_restart_CS, register_restart_field, query_initialized +use MOM_string_functions, only : lowercase +use MOM_time_manager, only : time_type_to_real +use MOM_unit_scaling, only : unit_scale_type +use MOM_variables, only : vertvisc_type, thermo_var_ptrs +use MOM_verticalGrid, only : verticalGrid_type +use MOM_MEKE_types, only : MEKE_type -use MOM_debugging, only : hchksum, uvchksum -use MOM_cpu_clock, only : cpu_clock_id, cpu_clock_begin, cpu_clock_end, CLOCK_ROUTINE -use MOM_diag_mediator, only : post_data, register_diag_field, safe_alloc_ptr -use MOM_diag_mediator, only : diag_ctrl, time_type -use MOM_domains, only : create_group_pass, do_group_pass, group_pass_type -use MOM_error_handler, only : MOM_error, FATAL, WARNING, NOTE, MOM_mesg -use MOM_file_parser, only : read_param, get_param, log_version, param_file_type -use MOM_grid, only : ocean_grid_type -use MOM_hor_index, only : hor_index_type -use MOM_restart, only : MOM_restart_CS, register_restart_field, query_initialized -use MOM_unit_scaling, only : unit_scale_type -use MOM_variables, only : vertvisc_type -use MOM_verticalGrid, only : verticalGrid_type -use MOM_MEKE_types, only : MEKE_type implicit none ; private @@ -26,6 +38,17 @@ module MOM_MEKE public step_forward_MEKE, MEKE_init, MEKE_alloc_register_restart, MEKE_end +! Constants for this module +integer, parameter :: NUM_FEATURES = 4 !< How many features used to predict EKE +integer, parameter :: MKE_IDX = 1 !< Index of mean kinetic energy in the feature array +integer, parameter :: SLOPE_Z_IDX = 2 !< Index of vertically averaged isopycnal slope in the feature array +integer, parameter :: RV_IDX = 3 !< Index of surface relative vorticity in the feature array +integer, parameter :: RD_DX_Z_IDX = 4 !< Index of the radius of deformation over the grid size in the feature array + +integer, parameter :: EKE_PROG = 1 !< Use prognostic equation to calcualte EKE +integer, parameter :: EKE_FILE = 2 !< Read in EKE from a file +integer, parameter :: EKE_DBCLIENT = 3 !< Infer EKE using a neural network + !> Control structure that contains MEKE parameters and diagnostics handles type, public :: MEKE_CS ; private logical :: initialized = .false. !< True if this control structure has been initialized. @@ -90,7 +113,8 @@ module MOM_MEKE logical :: kh_flux_enabled !< If true, lateral diffusive MEKE flux is enabled. logical :: initialize !< If True, invokes a steady state solver to calculate MEKE. logical :: debug !< If true, write out checksums of data for debugging - + integer :: eke_src !< Enum specifying whether EKE is stepped forward prognostically (default), + !! read in from a file, or inferred via a neural network type(diag_ctrl), pointer :: diag => NULL() !< A type that regulates diagnostics output !>@{ Diagnostic handles integer :: id_MEKE = -1, id_Ue = -1, id_Kh = -1, id_src = -1 @@ -101,19 +125,41 @@ module MOM_MEKE integer :: id_Lrhines = -1, id_Leady = -1 integer :: id_MEKE_equilibrium = -1 !>@} - + integer :: id_eke = -1 !< Handle for reading in EKE from a file ! Infrastructure integer :: id_clock_pass !< Clock for group pass calls type(group_pass_type) :: pass_MEKE !< Group halo pass handle for MEKE%MEKE and maybe MEKE%Kh_diff type(group_pass_type) :: pass_Kh !< Group halo pass handle for MEKE%Kh, MEKE%Ku, and/or MEKE%Au + + ! MEKE via Machine Learning + type(dbclient_type), pointer :: client => NULL() !< Pointer to the database client + + logical :: online_analysis !< If true, post the EKE used in MOM6 at every timestep + character(len=5) :: model_key = 'mleke' !< Key where the ML-model is stored + character(len=7) :: key_suffix !< Suffix appended to every key sent to Redis + real :: eke_max !< The maximum value of EKE considered physically reasonable + + ! Clock ids + integer :: id_client_init !< Clock id to time initialization of the client + integer :: id_put_tensor !< Clock id to time put_tensor routine + integer :: id_run_model !< Clock id to time running of the ML model + integer :: id_unpack_tensor !< Clock id to time retrieval of EKE prediction + + ! Diagnostic ids + integer :: id_mke = -1 !< Diagnostic id for surface mean kinetic energy + integer :: id_slope_z = -1 !< Diagnostic id for vertically averaged horizontal slope magnitude + integer :: id_slope_x = -1 !< Diagnostic id for isopycnal slope in the x-direction + integer :: id_slope_y = -1 !< Diagnostic id for isopycnal slope in the y-direction + integer :: id_rv = -1 !< Diagnostic id for surface relative vorticity + end type MEKE_CS contains !> Integrates forward-in-time the MEKE eddy energy equation. !! See \ref section_MEKE_equations. -subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, hv) - type(MEKE_type), intent(inout) :: MEKE !< MEKE fields +subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, hv, u, v, tv, Time) + type(MEKE_type), intent(inout) :: MEKE !< MEKE data. type(ocean_grid_type), intent(inout) :: G !< Ocean grid. type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -123,11 +169,16 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h type(vertvisc_type), intent(in) :: visc !< The vertical viscosity type. real, intent(in) :: dt !< Model(baroclinic) time-step [T ~> s]. type(MEKE_CS), intent(inout) :: CS !< MEKE control structure. - real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), intent(in) :: hu !< Accumlated zonal mass flux [H L2 ~> m3 or kg] + real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), intent(in) :: hu !< Accumlated zonal mass flux [H L2 ~> m3 or kg]. real, dimension(SZI_(G),SZJB_(G),SZK_(GV)), intent(in) :: hv !< Accumlated meridional mass flux [H L2 ~> m3 or kg] + real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), intent(inout) :: u !< Zonal velocity + real, dimension(SZI_(G),SZJB_(G),SZK_(GV)), intent(inout) :: v !< Meridional velocity + type(thermo_var_ptrs), intent(in) :: tv !< Type containing thermodynamic variables + type(time_type), intent(in) :: Time !< The time used for interpolating EKE ! Local variables real, dimension(SZI_(G),SZJ_(G)) :: & + data_eke, & ! EKE from file mass, & ! The total mass of the water column [R Z ~> kg m-2]. I_mass, & ! The inverse of mass [R-1 Z-1 ~> m2 kg-1]. depth_tot, & ! The depth of the water column [Z ~> m]. @@ -172,6 +223,7 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h real :: sdt_damp ! dt for damping [T ~> s] (sdt could be split). logical :: use_drag_rate ! Flag to indicate drag_rate is finite integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz + real(kind=real32), dimension(size(MEKE%MEKE),NUM_FEATURES) :: features_array is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB @@ -191,6 +243,8 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h return endif + select case(CS%eke_src) + case(EKE_PROG) if (CS%debug) then if (allocated(MEKE%mom_src)) & call hchksum(MEKE%mom_src, 'MEKE mom_src', G%HI, scale=US%RZ3_T3_to_W_m2*US%L_to_Z**2) @@ -246,7 +300,6 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h endif endif - ! Calculate drag_rate_visc(i,j) which accounts for the model bottom mean flow if (CS%visc_drag .and. allocated(visc%Kv_bbl_u) .and. allocated(visc%Kv_bbl_v)) then !$OMP parallel do default(shared) @@ -569,102 +622,117 @@ subroutine step_forward_MEKE(MEKE, h, SN_u, SN_v, visc, dt, G, GV, US, CS, hu, h call hchksum(MEKE%MEKE, "MEKE post-update MEKE", G%HI, haloshift=0, scale=US%L_T_to_m_s**2) endif - call cpu_clock_begin(CS%id_clock_pass) - call do_group_pass(CS%pass_MEKE, G%Domain) - call cpu_clock_end(CS%id_clock_pass) - - ! Calculate diffusivity for main model to use - if (CS%MEKE_KhCoeff>0.) then - if (.not.CS%MEKE_GEOMETRIC) then - if (CS%use_old_lscale) then - if (CS%Rd_as_max_scale) then - !$OMP parallel do default(shared) - do j=js,je ; do i=is,ie - MEKE%Kh(i,j) = (CS%MEKE_KhCoeff * & - sqrt(2.*max(0.,barotrFac2(i,j)*MEKE%MEKE(i,j))*G%areaT(i,j)) ) * & - min(MEKE%Rd_dx_h(i,j), 1.0) - enddo ; enddo - else - !$OMP parallel do default(shared) - do j=js,je ; do i=is,ie - MEKE%Kh(i,j) = CS%MEKE_KhCoeff * & - sqrt(2.*max(0., barotrFac2(i,j)*MEKE%MEKE(i,j))*G%areaT(i,j)) - enddo ; enddo - endif + case(EKE_FILE) + call time_interp_external(CS%id_eke,Time,data_eke) + do j=js,je ; do i=is,ie + MEKE%MEKE(i,j) = data_eke(i,j) * G%mask2dT(i,j) + enddo; enddo + call MEKE_lengthScales(CS, MEKE, G, GV, US, SN_u, SN_v, MEKE%MEKE, depth_tot, bottomFac2, barotrFac2, LmixScale) + case(EKE_DBCLIENT) + call pass_vector(u, v, G%Domain) + call MEKE_lengthScales(CS, MEKE, G, GV, US, SN_u, SN_v, MEKE%MEKE, depth_tot, bottomFac2, barotrFac2, LmixScale) + call ML_MEKE_calculate_features(G, GV, US, CS, MEKE%Rd_dx_h, u, v, tv, h, dt, features_array) + call predict_meke(G, CS, SIZE(h), Time, features_array, MEKE%MEKE) + case default + call MOM_error(FATAL,"Invalid method specified for calculating EKE") + end select + + call cpu_clock_begin(CS%id_clock_pass) + call do_group_pass(CS%pass_MEKE, G%Domain) + call cpu_clock_end(CS%id_clock_pass) + + ! Calculate diffusivity for main model to use + if (CS%MEKE_KhCoeff>0.) then + if (.not.CS%MEKE_GEOMETRIC) then + if (CS%use_old_lscale) then + if (CS%Rd_as_max_scale) then + !$OMP parallel do default(shared) + do j=js,je ; do i=is,ie + MEKE%Kh(i,j) = (CS%MEKE_KhCoeff * & + sqrt(2.*max(0.,barotrFac2(i,j)*MEKE%MEKE(i,j))*G%areaT(i,j)) ) * & + min(MEKE%Rd_dx_h(i,j), 1.0) + enddo ; enddo else !$OMP parallel do default(shared) do j=js,je ; do i=is,ie MEKE%Kh(i,j) = CS%MEKE_KhCoeff * & - sqrt(2.*max(0., barotrFac2(i,j)*MEKE%MEKE(i,j))) * LmixScale(i,j) + sqrt(2.*max(0., barotrFac2(i,j)*MEKE%MEKE(i,j))*G%areaT(i,j)) enddo ; enddo endif + else + !$OMP parallel do default(shared) + do j=js,je ; do i=is,ie + MEKE%Kh(i,j) = CS%MEKE_KhCoeff * & + sqrt(2.*max(0., barotrFac2(i,j)*MEKE%MEKE(i,j))) * LmixScale(i,j) + enddo ; enddo endif endif + endif - ! Calculate viscosity for the main model to use - if (CS%viscosity_coeff_Ku /=0.) then - do j=js,je ; do i=is,ie - MEKE%Ku(i,j) = CS%viscosity_coeff_Ku * sqrt(2.*max(0.,MEKE%MEKE(i,j))) * LmixScale(i,j) - enddo ; enddo - endif + ! Calculate viscosity for the main model to use + if (CS%viscosity_coeff_Ku /=0.) then + do j=js,je ; do i=is,ie + MEKE%Ku(i,j) = CS%viscosity_coeff_Ku * sqrt(2.*max(0.,MEKE%MEKE(i,j))) * LmixScale(i,j) + enddo ; enddo + endif - if (CS%viscosity_coeff_Au /=0.) then - do j=js,je ; do i=is,ie - MEKE%Au(i,j) = CS%viscosity_coeff_Au * sqrt(2.*max(0.,MEKE%MEKE(i,j))) * LmixScale(i,j)**3 - enddo ; enddo - endif + if (CS%viscosity_coeff_Au /=0.) then + do j=js,je ; do i=is,ie + MEKE%Au(i,j) = CS%viscosity_coeff_Au * sqrt(2.*max(0.,MEKE%MEKE(i,j))) * LmixScale(i,j)**3 + enddo ; enddo + endif - if (allocated(MEKE%Kh) .or. allocated(MEKE%Ku) .or. allocated(MEKE%Au)) then - call cpu_clock_begin(CS%id_clock_pass) - call do_group_pass(CS%pass_Kh, G%Domain) - call cpu_clock_end(CS%id_clock_pass) - endif + if (allocated(MEKE%Kh) .or. allocated(MEKE%Ku) .or. allocated(MEKE%Au)) then + call cpu_clock_begin(CS%id_clock_pass) + call do_group_pass(CS%pass_Kh, G%Domain) + call cpu_clock_end(CS%id_clock_pass) + endif - ! Offer fields for averaging. - if (any([CS%id_Ue, CS%id_Ub, CS%id_Ut] > 0)) & - tmp(:,:) = 0. - if (CS%id_MEKE>0) call post_data(CS%id_MEKE, MEKE%MEKE, CS%diag) - if (CS%id_Ue>0) then - do j=js,je ; do i=is,ie - tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j))) - enddo ; enddo - call post_data(CS%id_Ue, tmp, CS%diag) - endif - if (CS%id_Ub>0) then - do j=js,je ; do i=is,ie - tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j) * bottomFac2(i,j))) - enddo ; enddo - call post_data(CS%id_Ub, tmp, CS%diag) - endif - if (CS%id_Ut>0) then - do j=js,je ; do i=is,ie - tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j) * barotrFac2(i,j))) - enddo ; enddo - call post_data(CS%id_Ut, tmp, CS%diag) - endif - if (CS%id_Kh>0) call post_data(CS%id_Kh, MEKE%Kh, CS%diag) - if (CS%id_Ku>0) call post_data(CS%id_Ku, MEKE%Ku, CS%diag) - if (CS%id_Au>0) call post_data(CS%id_Au, MEKE%Au, CS%diag) - if (CS%id_KhMEKE_u>0) call post_data(CS%id_KhMEKE_u, Kh_u, CS%diag) - if (CS%id_KhMEKE_v>0) call post_data(CS%id_KhMEKE_v, Kh_v, CS%diag) - if (CS%id_src>0) call post_data(CS%id_src, src, CS%diag) - if (CS%id_decay>0) call post_data(CS%id_decay, MEKE_decay, CS%diag) - if (CS%id_GM_src>0) call post_data(CS%id_GM_src, MEKE%GM_src, CS%diag) - if (CS%id_mom_src>0) call post_data(CS%id_mom_src, MEKE%mom_src, CS%diag) - if (CS%id_GME_snk>0) call post_data(CS%id_GME_snk, MEKE%GME_snk, CS%diag) - if (CS%id_Le>0) call post_data(CS%id_Le, LmixScale, CS%diag) - if (CS%id_gamma_b>0) then - do j=js,je ; do i=is,ie - bottomFac2(i,j) = sqrt(bottomFac2(i,j)) - enddo ; enddo - call post_data(CS%id_gamma_b, bottomFac2, CS%diag) - endif - if (CS%id_gamma_t>0) then - do j=js,je ; do i=is,ie - barotrFac2(i,j) = sqrt(barotrFac2(i,j)) - enddo ; enddo - call post_data(CS%id_gamma_t, barotrFac2, CS%diag) - endif + ! Offer fields for averaging. + if (any([CS%id_Ue, CS%id_Ub, CS%id_Ut] > 0)) & + tmp(:,:) = 0. + if (CS%id_MEKE>0) call post_data(CS%id_MEKE, MEKE%MEKE, CS%diag) + if (CS%id_Ue>0) then + do j=js,je ; do i=is,ie + tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j))) + enddo ; enddo + call post_data(CS%id_Ue, tmp, CS%diag) + endif + if (CS%id_Ub>0) then + do j=js,je ; do i=is,ie + tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j) * bottomFac2(i,j))) + enddo ; enddo + call post_data(CS%id_Ub, tmp, CS%diag) + endif + if (CS%id_Ut>0) then + do j=js,je ; do i=is,ie + tmp(i,j) = sqrt(max(0., 2. * MEKE%MEKE(i,j) * barotrFac2(i,j))) + enddo ; enddo + call post_data(CS%id_Ut, tmp, CS%diag) + endif + if (CS%id_Kh>0) call post_data(CS%id_Kh, MEKE%Kh, CS%diag) + if (CS%id_Ku>0) call post_data(CS%id_Ku, MEKE%Ku, CS%diag) + if (CS%id_Au>0) call post_data(CS%id_Au, MEKE%Au, CS%diag) + if (CS%id_KhMEKE_u>0) call post_data(CS%id_KhMEKE_u, Kh_u, CS%diag) + if (CS%id_KhMEKE_v>0) call post_data(CS%id_KhMEKE_v, Kh_v, CS%diag) + if (CS%id_src>0) call post_data(CS%id_src, src, CS%diag) + if (CS%id_decay>0) call post_data(CS%id_decay, MEKE_decay, CS%diag) + if (CS%id_GM_src>0) call post_data(CS%id_GM_src, MEKE%GM_src, CS%diag) + if (CS%id_mom_src>0) call post_data(CS%id_mom_src, MEKE%mom_src, CS%diag) + if (CS%id_GME_snk>0) call post_data(CS%id_GME_snk, MEKE%GME_snk, CS%diag) + if (CS%id_Le>0) call post_data(CS%id_Le, LmixScale, CS%diag) + if (CS%id_gamma_b>0) then + do j=js,je ; do i=is,ie + bottomFac2(i,j) = sqrt(bottomFac2(i,j)) + enddo ; enddo + call post_data(CS%id_gamma_b, bottomFac2, CS%diag) + endif + if (CS%id_gamma_t>0) then + do j=js,je ; do i=is,ie + barotrFac2(i,j) = sqrt(barotrFac2(i,j)) + enddo ; enddo + call post_data(CS%id_gamma_t, barotrFac2, CS%diag) + endif end subroutine step_forward_MEKE @@ -1016,15 +1084,18 @@ end subroutine MEKE_lengthScales_0d !> Initializes the MOM_MEKE module and reads parameters. !! Returns True if module is to be used, otherwise returns False. -logical function MEKE_init(Time, G, US, param_file, diag, CS, MEKE, restart_CS) +logical function MEKE_init(Time, G, US, param_file, diag, dbcomms_CS, CS, MEKE, restart_CS, meke_in_dynamics) type(time_type), intent(in) :: Time !< The current model time. type(ocean_grid_type), intent(inout) :: G !< The ocean's grid structure. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type type(param_file_type), intent(in) :: param_file !< Parameter file parser structure. + type(dbcomms_CS_type), intent(in) :: dbcomms_CS !< Database communications control structure type(diag_ctrl), target, intent(inout) :: diag !< Diagnostics structure. type(MEKE_CS), intent(inout) :: CS !< MEKE control structure. type(MEKE_type), intent(inout) :: MEKE !< MEKE fields type(MOM_restart_CS), intent(in) :: restart_CS !< MOM restart control struct + logical, intent( out) :: meke_in_dynamics !< If true, MEKE is stepped forward in dynamics + !! otherwise in tracer dynamics ! Local variables real :: I_T_rescale ! A rescaling factor for time from the internal representation in this @@ -1033,6 +1104,8 @@ logical function MEKE_init(Time, G, US, param_file, diag, CS, MEKE, restart_CS) ! run to the representation in a restart file. real :: MEKE_restoring_timescale ! The timescale used to nudge MEKE toward its equilibrium value. real :: cdrag ! The default bottom drag coefficient [nondim]. + character(len=200) :: eke_filename, eke_varname, inputdir + character(len=16) :: eke_source_str integer :: i, j, is, ie, js, je, isd, ied, jsd, jed logical :: laplacian, biharmonic, coldStart ! This include declares and sets the variable "version". @@ -1051,75 +1124,113 @@ logical function MEKE_init(Time, G, US, param_file, diag, CS, MEKE, restart_CS) default=.false.) if (.not. MEKE_init) return CS%initialized = .true. + call get_param(param_file, mdl, "MEKE_IN_DYNAMICS", meke_in_dynamics, & + "If true, step MEKE forward with the dynamics"// & + "otherwise with the tracer timestep.", & + default=.true.) + + call get_param(param_file, mdl, "EKE_SOURCE", eke_source_str, & + "Determine the where EKE comes from:\n" // & + " 'prog': Calculated solving EKE equation\n"// & + " 'file': Read in from a file\n" // & + " 'dbclient': Retrieved from ML-database", default='prog') call MOM_mesg("MEKE_init: reading parameters ", 5) - ! Read all relevant parameters and write them to the model log. - call get_param(param_file, mdl, "MEKE_DAMPING", CS%MEKE_damping, & - "The local depth-independent MEKE dissipation rate.", & - units="s-1", default=0.0, scale=US%T_to_s) - call get_param(param_file, mdl, "MEKE_CD_SCALE", CS%MEKE_Cd_scale, & - "The ratio of the bottom eddy velocity to the column mean "//& - "eddy velocity, i.e. sqrt(2*MEKE). This should be less than 1 "//& - "to account for the surface intensification of MEKE.", & - units="nondim", default=0.) - call get_param(param_file, mdl, "MEKE_CB", CS%MEKE_Cb, & - "A coefficient in the expression for the ratio of bottom projected "//& - "eddy energy and mean column energy (see Jansen et al. 2015).",& - units="nondim", default=25.) - call get_param(param_file, mdl, "MEKE_MIN_GAMMA2", CS%MEKE_min_gamma, & - "The minimum allowed value of gamma_b^2.",& - units="nondim", default=0.0001) - call get_param(param_file, mdl, "MEKE_CT", CS%MEKE_Ct, & - "A coefficient in the expression for the ratio of barotropic "//& - "eddy energy and mean column energy (see Jansen et al. 2015).",& - units="nondim", default=50.) - call get_param(param_file, mdl, "MEKE_GMCOEFF", CS%MEKE_GMcoeff, & - "The efficiency of the conversion of potential energy "//& - "into MEKE by the thickness mixing parameterization. "//& - "If MEKE_GMCOEFF is negative, this conversion is not "//& - "used or calculated.", units="nondim", default=-1.0) - call get_param(param_file, mdl, "MEKE_GEOMETRIC", CS%MEKE_GEOMETRIC, & - "If MEKE_GEOMETRIC is true, uses the GM coefficient formulation "//& - "from the GEOMETRIC framework (Marshall et al., 2012).", default=.false.) - call get_param(param_file, mdl, "MEKE_GEOMETRIC_ALPHA", CS%MEKE_GEOMETRIC_alpha, & - "The nondimensional coefficient governing the efficiency of the GEOMETRIC \n"//& - "thickness diffusion.", units="nondim", default=0.05) - call get_param(param_file, mdl, "MEKE_EQUILIBRIUM_ALT", CS%MEKE_equilibrium_alt, & - "If true, use an alternative formula for computing the (equilibrium)"//& - "initial value of MEKE.", default=.false.) - call get_param(param_file, mdl, "MEKE_EQUILIBRIUM_RESTORING", CS%MEKE_equilibrium_restoring, & - "If true, restore MEKE back to its equilibrium value, which is calculated at "//& - "each time step.", default=.false.) - if (CS%MEKE_equilibrium_restoring) then - call get_param(param_file, mdl, "MEKE_RESTORING_TIMESCALE", MEKE_restoring_timescale, & - "The timescale used to nudge MEKE toward its equilibrium value.", units="s", & - default=1e6, scale=US%s_to_T) - CS%MEKE_restoring_rate = 1.0 / MEKE_restoring_timescale - endif + select case (lowercase(eke_source_str)) + case("file") + CS%eke_src = EKE_FILE + call time_interp_external_init + call get_param(param_file, mdl, "EKE_FILE", eke_filename, & + "A file in which to find the eddy kineteic energy variable.", & + default="eke_file.nc") + call get_param(param_file, mdl, "EKE_VARIABLE", eke_varname, & + "The name of the eddy kinetic energy variable to read from "//& + "EKE_FILE to use in MEKE.", & + default="eke") + call get_param(param_file, mdl, "INPUTDIR", inputdir, & + "The directory in which all input files are found.", & + default=".", do_not_log=.true.) + inputdir = slasher(inputdir) + + eke_filename = trim(inputdir) // trim(eke_filename) + CS%id_eke = init_external_field(eke_filename, eke_varname, domain=G%Domain%mpp_domain) + case("prog") + CS%eke_src = EKE_PROG + ! Read all relevant parameters and write them to the model log. + call get_param(param_file, mdl, "MEKE_DAMPING", CS%MEKE_damping, & + "The local depth-independent MEKE dissipation rate.", & + units="s-1", default=0.0, scale=US%T_to_s) + call get_param(param_file, mdl, "MEKE_CD_SCALE", CS%MEKE_Cd_scale, & + "The ratio of the bottom eddy velocity to the column mean "//& + "eddy velocity, i.e. sqrt(2*MEKE). This should be less than 1 "//& + "to account for the surface intensification of MEKE.", & + units="nondim", default=0.) + call get_param(param_file, mdl, "MEKE_CB", CS%MEKE_Cb, & + "A coefficient in the expression for the ratio of bottom projected "//& + "eddy energy and mean column energy (see Jansen et al. 2015).",& + units="nondim", default=25.) + call get_param(param_file, mdl, "MEKE_MIN_GAMMA2", CS%MEKE_min_gamma, & + "The minimum allowed value of gamma_b^2.",& + units="nondim", default=0.0001) + call get_param(param_file, mdl, "MEKE_CT", CS%MEKE_Ct, & + "A coefficient in the expression for the ratio of barotropic "//& + "eddy energy and mean column energy (see Jansen et al. 2015).",& + units="nondim", default=50.) + call get_param(param_file, mdl, "MEKE_GMCOEFF", CS%MEKE_GMcoeff, & + "The efficiency of the conversion of potential energy "//& + "into MEKE by the thickness mixing parameterization. "//& + "If MEKE_GMCOEFF is negative, this conversion is not "//& + "used or calculated.", units="nondim", default=-1.0) + call get_param(param_file, mdl, "MEKE_GEOMETRIC", CS%MEKE_GEOMETRIC, & + "If MEKE_GEOMETRIC is true, uses the GM coefficient formulation "//& + "from the GEOMETRIC framework (Marshall et al., 2012).", default=.false.) + call get_param(param_file, mdl, "MEKE_GEOMETRIC_ALPHA", CS%MEKE_GEOMETRIC_alpha, & + "The nondimensional coefficient governing the efficiency of the GEOMETRIC \n"//& + "thickness diffusion.", units="nondim", default=0.05) + call get_param(param_file, mdl, "MEKE_EQUILIBRIUM_ALT", CS%MEKE_equilibrium_alt, & + "If true, use an alternative formula for computing the (equilibrium)"//& + "initial value of MEKE.", default=.false.) + call get_param(param_file, mdl, "MEKE_EQUILIBRIUM_RESTORING", CS%MEKE_equilibrium_restoring, & + "If true, restore MEKE back to its equilibrium value, which is calculated at "//& + "each time step.", default=.false.) + if (CS%MEKE_equilibrium_restoring) then + call get_param(param_file, mdl, "MEKE_RESTORING_TIMESCALE", MEKE_restoring_timescale, & + "The timescale used to nudge MEKE toward its equilibrium value.", units="s", & + default=1e6, scale=US%s_to_T) + CS%MEKE_restoring_rate = 1.0 / MEKE_restoring_timescale + endif + + call get_param(param_file, mdl, "MEKE_FRCOEFF", CS%MEKE_FrCoeff, & + "The efficiency of the conversion of mean energy into "//& + "MEKE. If MEKE_FRCOEFF is negative, this conversion "//& + "is not used or calculated.", units="nondim", default=-1.0) + call get_param(param_file, mdl, "MEKE_GMECOEFF", CS%MEKE_GMECoeff, & + "The efficiency of the conversion of MEKE into mean energy "//& + "by GME. If MEKE_GMECOEFF is negative, this conversion "//& + "is not used or calculated.", units="nondim", default=-1.0) + call get_param(param_file, mdl, "MEKE_BGSRC", CS%MEKE_BGsrc, & + "A background energy source for MEKE.", units="W kg-1", & + default=0.0, scale=US%m_to_L**2*US%T_to_s**3) + call get_param(param_file, mdl, "MEKE_KH", CS%MEKE_Kh, & + "A background lateral diffusivity of MEKE. "//& + "Use a negative value to not apply lateral diffusion to MEKE.", & + units="m2 s-1", default=-1.0, scale=US%m_to_L**2*US%T_to_s) + call get_param(param_file, mdl, "MEKE_K4", CS%MEKE_K4, & + "A lateral bi-harmonic diffusivity of MEKE. "//& + "Use a negative value to not apply bi-harmonic diffusion to MEKE.", & + units="m4 s-1", default=-1.0, scale=US%m_to_L**4*US%T_to_s) + call get_param(param_file, mdl, "MEKE_DTSCALE", CS%MEKE_dtScale, & + "A scaling factor to accelerate the time evolution of MEKE.", & + units="nondim", default=1.0) + case("dbclient") + CS%eke_src = EKE_DBCLIENT + call ML_MEKE_init(diag, G, US, Time, param_file, dbcomms_CS, CS) + case default + call MOM_error(FATAL, "Invalid method selected for calculating EKE") + end select + ! GMM, make sure all params used to calculated MEKE are within the above if - call get_param(param_file, mdl, "MEKE_FRCOEFF", CS%MEKE_FrCoeff, & - "The efficiency of the conversion of mean energy into "//& - "MEKE. If MEKE_FRCOEFF is negative, this conversion "//& - "is not used or calculated.", units="nondim", default=-1.0) - call get_param(param_file, mdl, "MEKE_GMECOEFF", CS%MEKE_GMECoeff, & - "The efficiency of the conversion of MEKE into mean energy "//& - "by GME. If MEKE_GMECOEFF is negative, this conversion "//& - "is not used or calculated.", units="nondim", default=-1.0) - call get_param(param_file, mdl, "MEKE_BGSRC", CS%MEKE_BGsrc, & - "A background energy source for MEKE.", units="W kg-1", & - default=0.0, scale=US%m_to_L**2*US%T_to_s**3) - call get_param(param_file, mdl, "MEKE_KH", CS%MEKE_Kh, & - "A background lateral diffusivity of MEKE. "//& - "Use a negative value to not apply lateral diffusion to MEKE.", & - units="m2 s-1", default=-1.0, scale=US%m_to_L**2*US%T_to_s) - call get_param(param_file, mdl, "MEKE_K4", CS%MEKE_K4, & - "A lateral bi-harmonic diffusivity of MEKE. "//& - "Use a negative value to not apply bi-harmonic diffusion to MEKE.", & - units="m4 s-1", default=-1.0, scale=US%m_to_L**4*US%T_to_s) - call get_param(param_file, mdl, "MEKE_DTSCALE", CS%MEKE_dtScale, & - "A scaling factor to accelerate the time evolution of MEKE.", & - units="nondim", default=1.0) call get_param(param_file, mdl, "MEKE_KHCOEFF", CS%MEKE_KhCoeff, & "A scaling factor in the expression for eddy diffusivity "//& "which is otherwise proportional to the MEKE velocity- "//& @@ -1376,6 +1487,267 @@ logical function MEKE_init(Time, G, US, param_file, diag, CS, MEKE, restart_CS) end function MEKE_init +!> Initializer for the variant of MEKE that uses ML to predict eddy kinetic energy +subroutine ML_MEKE_init(diag, G, US, Time, param_file, dbcomms_CS, CS) + type(diag_ctrl), target, intent(inout) :: diag !< Diagnostics structure. + type(ocean_grid_type), intent(inout) :: G !< The ocean's grid structure. + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + type(time_type), intent(in) :: Time !< The current model time. + type(param_file_type), intent(in) :: param_file !< Parameter file parser structure. + type(dbcomms_CS_type), intent(in) :: dbcomms_CS !< Control structure for database communication + type(MEKE_CS), intent(inout) :: CS !< Control structure for this module + + character(len=200) :: inputdir, backend, model_filename + integer :: db_return_code, batch_size + character(len=40) :: mdl = "MOM_ML_MEKE" + + ! Store pointers in control structure + write(CS%key_suffix, '(A,I6.6)') '_', PE_here() + ! Put some basic information into the database + db_return_code = 0 + db_return_code = CS%client%put_tensor("meta"//CS%key_suffix, & + REAL([G%isd_global, G%idg_offset, G%jsd_global, G%jdg_offset]),[4]) + db_return_code + db_return_code = CS%client%put_tensor("geolat"//CS%key_suffix, G%geoLatT, shape(G%geoLatT)) + db_return_code + db_return_code = CS%client%put_tensor("geolon"//CS%key_suffix, G%geoLonT, shape(G%geoLonT)) + db_return_code + db_return_code = CS%client%put_tensor("EKE_shape"//CS%key_suffix, shape(G%geolonT), [2]) + db_return_code + + if (CS%client%SR_error_parser(db_return_code)) call MOM_error(FATAL, "Putting metadata into the database failed") + + call read_param(param_file, "INPUTDIR", inputdir) + inputdir = slasher(inputdir) + + call get_param(param_file, mdl, "BATCH_SIZE", batch_size, "Batch size to use for inference", default=1) + call get_param(param_file, mdl, "EKE_BACKEND", backend, & + "The computational backend to use for EKE inference (CPU or GPU)", default="GPU") + call get_param(param_file, mdl, "EKE_MODEL", model_filename, & + "Filename of the a saved pyTorch model to use", fail_if_missing = .true.) + call get_param(param_file, mdl, "EKE_MAX", CS%eke_max, & + "Maximum value of EKE allowed when inferring EKE", default=2., scale=US%L_T_to_m_s**2) + + ! Set the machine learning model + if (dbcomms_CS%colocated) then + if (modulo(PE_here(),dbcomms_CS%colocated_stride) == 0) then + db_return_code = CS%client%set_model_from_file(CS%model_key, trim(inputdir)//trim(model_filename), & + "TORCH", backend, batch_size=batch_size) + endif + else + if (is_root_pe()) then + db_return_code = CS%client%set_model_from_file(CS%model_key, trim(inputdir)//trim(model_filename), & + "TORCH", backend, batch_size=batch_size) + endif + endif + if (CS%client%SR_error_parser(db_return_code)) then + call MOM_error(FATAL, "MEKE: set_model failed") + endif + + call get_param(param_file, mdl, "ONLINE_ANALYSIS", CS%online_analysis, & + "If true, post EKE used in MOM6 to the database for analysis", default=.true.) + + ! Set various clock ids + CS%id_client_init = cpu_clock_id('(ML_MEKE client init)', grain=CLOCK_ROUTINE) + CS%id_put_tensor = cpu_clock_id('(ML_MEKE put tensor)', grain=CLOCK_ROUTINE) + CS%id_run_model = cpu_clock_id('(ML_MEKE run model)', grain=CLOCK_ROUTINE) + CS%id_unpack_tensor = cpu_clock_id('(ML_MEKE unpack tensor )', grain=CLOCK_ROUTINE) + + ! Diagnostics for ML_MEKE + CS%id_mke = register_diag_field('ocean_model', 'MEKE_MKE', diag%axesT1, Time, & + 'Surface mean (resolved) kinetic energy used in MEKE', 'm2 s-2', conversion=US%L_T_to_m_s**2) + CS%id_slope_z= register_diag_field('ocean_model', 'MEKE_slope_z', diag%axesT1, Time, & + 'Vertically averaged isopyncal slope magnitude used in MEKE', 'm2 s-2', conversion=US%L_T_to_m_s**2) + CS%id_slope_x= register_diag_field('ocean_model', 'MEKE_slope_x', diag%axesCui, Time, & + 'Isopycnal slope in the x-direction used in MEKE', 'm2 s-2', conversion=US%L_T_to_m_s**2) + CS%id_slope_y= register_diag_field('ocean_model', 'MEKE_slope_y', diag%axesCvi, Time, & + 'Isopycnal slope in the y-direction used in MEKE', 'm2 s-2', conversion=US%L_T_to_m_s**2) + CS%id_rv= register_diag_field('ocean_model', 'MEKE_RV', diag%axesT1, Time, & + 'Surface relative vorticity used in MEKE', 'm2 s-2', conversion=US%L_T_to_m_s**2) + +end subroutine ML_MEKE_init + +!> Calculate the various features used for the machine learning prediction +subroutine ML_MEKE_calculate_features(G, GV, US, CS, Rd_dx_h, u, v, tv, h, dt, features_array) + type(ocean_grid_type), intent(inout) :: G !< Ocean grid + type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + type(MEKE_CS), intent(in) :: CS !< Control structure for MEKE + real, dimension(SZI_(G),SZJ_(G)), intent(in ) :: Rd_dx_h !< Rossby radius of deformation over + !! the grid length scale [nondim] + real, dimension(SZIB_(G),SZJ_(G),SZK_(G)), intent(in) :: u !< Zonal velocity [L T-1 ~> m s-1] + real, dimension(SZI_(G),SZJB_(G),SZK_(G)), intent(in) :: v !< Meridional velocity [L T-1 ~> m s-1] + type(thermo_var_ptrs), intent(in) :: tv !< Type containing thermodynamic variables + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h !< Layer thickness [H ~> m or kg m-2]. + real, intent(in) :: dt !< Model(baroclinic) time-step [T ~> s]. + real(kind=real32), dimension(SIZE(h),num_features), intent( out) :: features_array + !< The array of features needed for machine + !! learning inference + + real, dimension(SZI_(G),SZJ_(G)) :: mke + real, dimension(SZI_(G),SZJ_(G)) :: slope_z + real, dimension(SZIB_(G),SZJB_(G)) :: rv_z + real, dimension(SZIB_(G),SZJB_(G)) :: rv_z_t + real, dimension(SZI_(G),SZJ_(G)) :: rd_dx_z + + real, dimension(SZIB_(G),SZJ_(G), SZK_(G)) :: h_u ! Thickness at u point + real, dimension(SZI_(G),SZJB_(G), SZK_(G)) :: h_v ! Thickness at v point + real, dimension(SZIB_(G),SZJ_(G),SZK_(G)+1) :: slope_x ! Isoneutral slope at U point + real, dimension(SZI_(G),SZJB_(G),SZK_(G)+1) :: slope_y ! Isoneutral slope at V point + real, dimension(SZIB_(G),SZJ_(G)) :: slope_x_vert_avg ! Isoneutral slope at U point + real, dimension(SZI_(G),SZJB_(G)) :: slope_y_vert_avg ! Isoneutral slope at V point + real, dimension(SZI_(G), SZJ_(G), SZK_(G)+1) :: e ! The interface heights relative to mean sea level [Z ~> m]. + real :: slope_t, u_t, v_t ! u and v interpolated to thickness point + real :: dvdx, dudy + real :: a_e, a_w, a_n, a_s, Idenom, sum_area + + integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz + + is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke + Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB + + ! Calculate various features for used to infer eddy kinetic energy + ! Linear interpolation to estimate thickness at a velocity points + do k=1,nz; do j=js-1,je+1; do i=is-1,ie+1 + h_u(I,j,k) = 0.5*(h(i,j,k)*G%mask2dT(i,j) + h(i+1,j,k)*G%mask2dT(i+1,j)) + GV%Angstrom_H + h_v(i,J,k) = 0.5*(h(i,j,k)*G%mask2dT(i,j) + h(i,j+1,k)*G%mask2dT(i,j+1)) + GV%Angstrom_H + enddo; enddo; enddo; + call find_eta(h, tv, G, GV, US, e, halo_size=2) + call calc_isoneutral_slopes(G, GV, US, h, e, tv, dt*1.e-7, .false., slope_x, slope_y) + call pass_vector(slope_x, slope_y, G%Domain) + do j=js-1,je+1; do i=is-1,ie+1 + slope_x_vert_avg(I,j) = vertical_average_interface(slope_x(i,j,:), h_u(i,j,:), GV%H_subroundoff) + slope_y_vert_avg(i,J) = vertical_average_interface(slope_y(i,j,:), h_v(i,j,:), GV%H_subroundoff) + enddo; enddo + slope_z(:,:) = 0. + + call pass_vector(slope_x_vert_avg, slope_y_vert_avg, G%Domain) + do j=js,je; do i=is,ie + ! Calculate weights for interpolation from velocity points to h points + sum_area = G%areaCu(I-1,j) + G%areaCu(I,j) + if (sum_area>0.0) then + Idenom = sqrt(0.5*G%IareaT(i,j) / sum_area) + a_w = G%areaCu(I-1,j) * Idenom + a_e = G%areaCu(I,j) * Idenom + else + a_w = 0.0 ; a_e = 0.0 + endif + + sum_area = G%areaCv(i,J-1) + G%areaCv(i,J) + if (sum_area>0.0) then + Idenom = sqrt(0.5*G%IareaT(i,j) / sum_area) + a_s = G%areaCv(i,J-1) * Idenom + a_n = G%areaCv(i,J) * Idenom + else + a_s = 0.0 ; a_n = 0.0 + endif + + ! Calculate mean kinetic energy + u_t = a_e*u(I,j,1)+a_w*u(I-1,j,1) + v_t = a_n*v(i,J,1)+a_s*v(i,J-1,1) + mke(i,j) = 0.5*( u_t*u_t + v_t*v_t ) + + ! Calculate the magnitude of the slope + slope_t = slope_x_vert_avg(I,j)*a_e+slope_x_vert_avg(I-1,j)*a_w + slope_z(i,j) = sqrt(slope_t*slope_t) + slope_t = slope_y_vert_avg(i,J)*a_n+slope_y_vert_avg(i,J-1)*a_s + slope_z(i,j) = 0.5*(slope_z(i,j) + sqrt(slope_t*slope_t))*G%mask2dT(i,j) + enddo; enddo + call pass_var(slope_z, G%Domain) + + ! Calculate relative vorticity + do J=Jsq-1,Jeq+1 ; do I=Isq-1,Ieq+1 + dvdx = (v(i+1,J,1)*G%dyCv(i+1,J) - v(i,J,1)*G%dyCv(i,J)) + dudy = (u(I,j+1,1)*G%dxCu(I,j+1) - u(I,j,1)*G%dxCu(I,j)) + ! Assumed no slip + rv_z(I,J) = (2.0-G%mask2dBu(I,J)) * (dvdx - dudy) * G%IareaBu(I,J) + enddo; enddo + ! Interpolate RV to t-point, revisit this calculation to include metrics + do j=js,je; do i=is,ie + rv_z_t(i,j) = 0.25*(rv_z(i-1,j) + rv_z(i,j) + rv_z(i-1,j-1) + rv_z(i,j-1)) + enddo; enddo + + + ! Construct the feature array + features_array(:,mke_idx) = pack(mke,.true.) + features_array(:,slope_z_idx) = pack(slope_z,.true.) + features_array(:,rd_dx_z_idx) = pack(Rd_dx_h,.true.) + features_array(:,rv_idx) = pack(rv_z_t,.true.) + + if (CS%id_rv>0) call post_data(CS%id_rv, rv_z, CS%diag) + if (CS%id_mke>0) call post_data(CS%id_mke, mke, CS%diag) + if (CS%id_slope_z>0) call post_data(CS%id_slope_z, slope_z, CS%diag) + if (CS%id_slope_x>0) call post_data(CS%id_slope_x, slope_x, CS%diag) + if (CS%id_slope_y>0) call post_data(CS%id_slope_y, slope_y, CS%diag) +end subroutine ML_MEKE_calculate_features + +!> Use the machine learning interface to predict EKE +subroutine predict_MEKE(G, CS, npts, Time, features_array, MEKE) + type(ocean_grid_type), intent(inout) :: G !< Ocean grid + type(MEKE_CS), intent(in ) :: CS !< Control structure for MEKE + integer, intent(in ) :: npts !< Number of T-grid cells on the local + !! domain + type(time_type), intent(in ) :: Time !< The current model time + real(kind=real32), dimension(npts,num_features), intent(in ) :: features_array + !< The array of features needed for machine + !! learning inference + real, dimension(SZI_(G),SZJ_(G)), intent( out) :: MEKE !< Eddy kinetic energy [L2 T-2 ~> m2 s-2] + integer :: db_return_code + character(len=255), dimension(1) :: model_out, model_in + character(len=255) :: time_suffix + real(kind=real32), dimension(SIZE(MEKE)) :: MEKE_vec + + integer :: i, j, is, ie, js, je + is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec +!> Use the database client to call a machine learning model to predict eddy kinetic energy + call cpu_clock_begin(CS%id_put_tensor) + db_return_code = CS%client%put_tensor("features"//CS%key_suffix, features_array, shape(features_array)) + call cpu_clock_end(CS%id_put_tensor) + + ! Run the ML model to predict EKE and return the result + model_out(1) = "EKE"//CS%key_suffix + model_in(1) = "features"//CS%key_suffix + call cpu_clock_begin(CS%id_run_model) + db_return_code = CS%client%run_model(CS%model_key, model_in, model_out) + call cpu_clock_end(CS%id_run_model) + if (CS%client%SR_error_parser(db_return_code)) then + call MOM_error(FATAL, "MEKE: run_model failed") + endif + call cpu_clock_begin(CS%id_unpack_tensor) + db_return_code = CS%client%unpack_tensor( model_out(1), MEKE_vec, shape(MEKE_vec) ) + call cpu_clock_end(CS%id_unpack_tensor) + + MEKE = reshape(MEKE_vec, shape(MEKE)) + do j=js,je; do i=is,ie + MEKE(i,j) = MIN(MAX(exp(MEKE(i,j)),0.),CS%eke_max) + enddo; enddo + call pass_var(MEKE,G%Domain) + + if (CS%online_analysis) then + write(time_suffix,"(F16.0)") time_type_to_real(Time) + db_return_code = CS%client%put_tensor(trim("EKE_")//trim(adjustl(time_suffix))//CS%key_suffix, MEKE, shape(MEKE)) + endif +end subroutine predict_MEKE + +!> Compute average of interface quantities weighted by the thickness of the surrounding layers +real function vertical_average_interface(h, w, h_min) + + real, dimension(:), intent(in) :: h !< Layer Thicknesses + real, dimension(:), intent(in) :: w !< Quantity to average + real, intent(in) :: h_min !< The vanishingly small layer thickness + + real :: htot, inv_htot + integer :: k, nk + + nk = size(h) + htot = h_min + do k=2,nk + htot = htot + (h(k-1)+h(k)) + enddo + inv_htot = 1./htot + + vertical_average_interface = 0. + do K=2,nk + vertical_average_interface = vertical_average_interface + (w(k)*(h(k-1)+h(k)))*inv_htot + enddo +end function vertical_average_interface + !> Allocates memory and register restart fields for the MOM_MEKE module. subroutine MEKE_alloc_register_restart(HI, US, param_file, MEKE, restart_CS) ! Arguments