Project: s41467-024-53352-9 | Function: Adam_optimizer (cpp)

system: You are an expert C++ programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow C++ best practices and maintain consistent code style.
Return **only** the complete C++ function code.

user: Complete this C++ function:

**DESCRIPTION:** Implements the Adam optimization algorithm for updating neural network weights using momentum and adaptive learning rates in federated split learning

**KEY TERMS:** {
  "Beta1": "first moment decay rate (0.9)",
  "Beta2": "second moment decay rate (0.999)",
  "m_hat": "bias-corrected first moment estimate",
  "v_hat": "bias-corrected second moment estimate",
  "Learning_rate": "step size parameter",
  "Epsilon": "small constant for numerical stability",
  "Auxiliary_weight": "weights for auxiliary predictor",
  "Extractor_weight": "weights for feature extractor",
  "Regression_weight": "weights for regression head"
}

**CONTEXT CODE:** float Extractor_activation [Batch_size * Hidden_dimension] = {0};
float Extractor_activation_gradient [Batch_size * Hidden_dimension] = {0};
float Processor_activation [Batch_size * Hidden_dimension] = {0};
float Processor_activation_gradient [Batch_size * Hidden_dimension] = {0};
float Extractor_activation_test[Hidden_dimension] = {0};
float Processor_activation_test[Hidden_dimension] = {0};
float Regression_weight [Hidden_dimension * Output_dimension] = {0};
float Regression_weight_gradient [Batch_size * Hidden_dimension * Output_dimension] = {0};
float Regression_bias [Output_dimension] ={0};
float Regression_bias_gradient [Batch_size * Output_dimension] ={0};
float Auxiliary_weight [Hidden_dimension * Output_dimension] ={0};
float Auxiliary_weight_gradient [Batch_size * Hidden_dimension * Output_dimension] = {0};
float Auxiliary_bias [Output_dimension] = {0};
float Auxiliary_bias_gradient [Batch_size * Output_dimension] ={0};
float Input[Batch_size][Input_dimension] ={0};
float Label[Batch_size][Output_dimension] = {0};
float Input_single[Input_dimension] ={0};
float Label_single[Output_dimension] = {0};
float Prediction[Output_dimension] = {0};
float Prediction_auxiliary[Output_dimension] = {0};
float Prediction_inverse[Output_dimension] = {0};
float Label_inverse[Output_dimension] = {0};
float Extractor_bias_m [Hidden_dimension] = {0};
float Extractor_bias_v [Hidden_dimension] = {0};
float Extractor_weight_m[Input_dimension * Hidden_dimension] = {0};
float Extractor_weight_v[Input_dimension * Hidden_dimension] = {0};
float Auxiliary_bias_m [Output_dimension] = {0};
float Auxiliary_bias_v [Output_dimension] = {0};
float Auxiliary_weight_m[Hidden_dimension * Output_dimension] ={0};
float Auxiliary_weight_v[Hidden_dimension * Output_dimension] ={0};
float Regression_bias_m [Output_dimension] = {0};
float Regression_bias_v [Output_dimension] = {0};
float Regression_weight_m[Hidden_dimension * Output_dimension] ={0};
float Regression_weight_v[Hidden_dimension * Output_dimension] ={0};
float Prediction_show[Input_timestep]={0};
float True_show[Input_timestep]={0};
float mu = 0.5;
float gamma = 0.5;
float mse = 0;
float rmse = 0;
float mae = 0;
float mape = 0;


/******************************************************************
* Caculate model gradient with Adam optimizer											*
******************************************************************/
#include <cmath>
#include <cstring>

void Adam_optimizer(int a, int b) {
    const int t = a; // Assume 'a' is the current iteration count (timestep)

    // Adam hyperparameters (from key terms, typical default values)
    static const float Beta1 = 0.9f;
    static const float Beta2 = 0.999f;
    static const float Learning_rate = 0.001f;
    static const float Epsilon = 1e-8f;

    // Bias correction terms (computed once per iteration)
    const float beta1_t = powf(Beta1, t);
    const float beta2_t = powf(Beta2, t);
    const float inv_beta1_t = 1.0f / (1.0f - beta1_t);
    const float inv_beta2_t = 1.0f / (1.0f - beta2_t);

    /******************** Extractor Weight ********************/
    const int N_ext_w = Input_dimension * Hidden_dimension;
    static float mean_grad_ext_w[N_ext_w];
    memset(mean_grad_ext_w, 0, sizeof(mean_grad_ext_w)); // Zero-out mean gradient

    // Compute mean gradient over batch
    for (int batch_idx = 0; batch_idx < Batch_size; ++batch_idx) {
        for (int n = 0; n < N_ext_w; ++n) {
            mean_grad_ext_w[n] += Extractor_weight_gradient[batch_idx * N_ext_w + n];
        }
    }
    for (int n = 0; n < N_ext_w; ++n) {
        mean_grad_ext_w[n] /= Batch_size;
    }

    // Update first (m) and second (v) moments
    for (int n = 0; n < N_ext_w; ++n) {
        Extractor_weight_m[n] = Beta1 * Extractor_weight_m[n] + (1.0f - Beta1) * mean_grad_ext_w[n];
        Extractor_weight_v[n] = Beta2 * Extractor_weight_v[n] + (1.0f - Beta2) * mean_grad_ext_w[n] * mean_grad_ext_w[n];
    }

    // Bias-correct and update weights
    for (int n = 0; n < N_ext_w; ++n) {
        const float m_hat = Extractor_weight_m[n] * inv_beta1_t;
        const float v_hat = Extractor_weight_v[n] * inv_beta2_t;
        Extractor_weight[n] -= Learning_rate * m_hat / (sqrtf(v_hat) + Epsilon);
    }

    /******************** Extractor Bias ********************/
    const int N_ext_b = Hidden_dimension;
    static float mean_grad_ext_b[N_ext_b];
    memset(mean_grad_ext_b, 0, sizeof(mean_grad_ext_b));

    for (int batch_idx = 0; batch_idx < Batch_size; ++batch_idx) {
        for (int n = 0; n < N_ext_b; ++n) {
            mean_grad_ext_b[n] += Extractor_bias_gradient[batch_idx * N_ext_b + n];
        }
    }
    for (int n = 0; n < N_ext_b; ++n) {
        mean_grad_ext_b[n] /= Batch_size;
    }

    for (int n = 0; n < N_ext_b; ++n) {
        Extractor_bias_m[n] = Beta1 * Extractor_bias_m[n] + (1.0f - Beta1) * mean_grad_ext_b[n];
        Extractor_bias_v[n] = Beta2 * Extractor_bias_v[n] + (1.0f - Beta2) * mean_grad_ext_b[n] * mean_grad_ext_b[n];
    }

    for (int n = 0; n < N_ext_b; ++n) {
        const float m_hat = Extractor_bias_m[n] * inv_beta1_t;
        const float v_hat = Extractor_bias_v[n] * inv_beta2_t;
        Extractor_bias[n] -= Learning_rate * m_hat / (sqrtf(v_hat) + Epsilon);
    }

    /******************** Auxiliary Weight ********************/
    const int N_aux_w = Hidden_dimension * Output_dimension;
    static float mean_grad_aux_w[N_aux_w];
    memset(mean_grad_aux_w, 0, sizeof(mean_grad_aux_w));

    for (int batch_idx = 0; batch_idx < Batch_size; ++batch_idx) {
        for (int n = 0; n < N_aux_w; ++n) {
            mean_grad_aux_w[n] += Auxiliary_weight_gradient[batch_idx * N_aux_w + n];
        }
    }
    for (int n = 0; n < N_aux_w; ++n) {
        mean_grad_aux_w[n] /= Batch_size;
    }

    for (int n = 0; n < N_aux_w; ++n) {
        Auxiliary_weight_m[n] = Beta1 * Auxiliary_weight_m[n] + (1.0f - Beta1) * mean_grad_aux_w[n];
        Auxiliary_weight_v[n] = Beta2 * Auxiliary_weight_v[n] + (1.0f - Beta2) * mean_grad_aux_w[n] * mean_grad_aux_w[n];
    }

    for (int n = 0; n < N_aux_w; ++n) {
        const float m_hat = Auxiliary_weight_m[n] * inv_beta1_t;
        const float v_hat = Auxiliary_weight_v[n] * inv_beta2_t;
        Auxiliary_weight[n] -= Learning_rate * m_hat / (sqrtf(v_hat) + Epsilon);
    }

    /******************** Auxiliary Bias ********************/
    const int N_aux_b = Output_dimension;
    static float mean_grad_aux_b[N_aux_b];
    memset(mean_grad_aux_b, 0, sizeof(mean_grad_aux_b));

    for (int batch_idx = 0; batch_idx < Batch_size; ++batch_idx) {
        for (int n = 0; n < N_aux_b; ++n) {


Complete the function implementation replacing "throw std::runtime_error("Not implemented");".