Optimization_Interface.cpp

Go to the documentation of this file.

/*
Created on Fri Jun 26 14:13:26 2020
Copyright 2020 Peter Rakyta, Ph.D.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

@author: Peter Rakyta, Ph.D.
*/
#include "Optimization_Interface.h"
#include "N_Qubit_Decomposition_Cost_Function.h"
#include "matrix_float.h"
#include "Adam.h"
#include "grad_descend.h"
#include "BFGS_Powell.h"
#include "Bayes_Opt.h"

#include "RL_experience.h"

#ifdef _WIN32
#include <cstdio>
#endif

#include <fstream>
#include <tbb/enumerable_thread_specific.h>

extern "C" int LAPACKE_dgesv(      int  matrix_layout, int n, int nrhs, double *a, int lda, int *ipiv, double *b, int ldb);  


Optimization_Interface::Optimization_Interface() {

    // logical value. Set true if finding the minimum number of gate layers is required (default), or false when the maximal number of two-qubit gates is used (ideal for general unitaries).
    optimize_layer_num  = false;

    // A string describing the type of the class
    type = N_QUBIT_DECOMPOSITION_CLASS_BASE;

    // The global minimum of the optimization problem
    global_target_minimum = 0;

    // logical variable indicating whether adaptive learning reate is used in the ADAM algorithm
    adaptive_eta = true;

    // parameter to contron the radius of parameter randomization around the curren tminimum
    radius = 1.0;
    randomization_rate = 0.3;

    // The chosen variant of the cost function
    cost_fnc = FROBENIUS_NORM;
    
    number_of_iters.store(0, std::memory_order_relaxed);
     
 

    // variables to calculate the cost function with first and second corrections    
    prev_cost_fnv_val = 1.0;
    correction1_scale = 1/1.7;
    correction2_scale = 1/2.0;  


    // number of utilized accelerators
    accelerator_num = 0;

    // set the trace offset
    trace_offset = 0;



    // Time spent on circuit simulation/cost function evaluation
    circuit_simulation_time = 0.0;
    // time spent on optimization
    CPU_time = 0.0;

}


Optimization_Interface::Optimization_Interface( const Optimization_Interface& other ) : Decomposition_Base(other) {

    max_inner_iterations = other.max_inner_iterations;
    random_shift_count_max = other.random_shift_count_max;
    id = other.id;

    optimize_layer_num = other.optimize_layer_num;
    identical_blocks = other.identical_blocks;
    alg = other.alg;
    cost_fnc = other.cost_fnc;
    prev_cost_fnv_val = other.prev_cost_fnv_val;
    correction1_scale = other.correction1_scale;
    correction2_scale = other.correction2_scale;
    use_cuts = other.use_cuts;
    osr_rank = other.osr_rank;
    use_softmax = other.use_softmax;
    number_of_iters.store(other.number_of_iters.load(std::memory_order_relaxed), std::memory_order_relaxed);
    adaptive_eta = other.adaptive_eta;
    radius = other.radius;
    randomization_rate = other.randomization_rate;
    accelerator_num = other.accelerator_num;
    trace_offset = other.trace_offset;
    circuit_simulation_time = other.circuit_simulation_time;
    CPU_time = other.CPU_time;

}

Optimization_Interface::Optimization_Interface( Matrix Umtx_in, int qbit_num_in, bool optimize_layer_num_in, std::map<std::string, Config_Element>& config, guess_type initial_guess_in= CLOSE_TO_ZERO, int accelerator_num_in ) : Decomposition_Base(Umtx_in, qbit_num_in, config, initial_guess_in) {

    // logical value. Set true if finding the minimum number of gate layers is required (default), or false when the maximal number of two-qubit gates is used (ideal for general unitaries).
    optimize_layer_num  = optimize_layer_num_in;

    // A string describing the type of the class
    type = N_QUBIT_DECOMPOSITION_CLASS_BASE;

    // The global minimum of the optimization problem
    global_target_minimum = 0;
    
    number_of_iters.store(0, std::memory_order_relaxed);
         

    // number of iteratrion loops in the optimization
    iteration_loops[2] = 3;

    // filling in numbers that were not given in the input
    for ( std::map<int,int>::iterator it = max_layer_num_def.begin(); it!=max_layer_num_def.end(); it++) {
        if ( max_layer_num.count( it->first ) == 0 ) {
            max_layer_num.insert( std::pair<int, int>(it->first,  it->second) );
        }
    }

    // logical variable indicating whether adaptive learning reate is used in the ADAM algorithm
    adaptive_eta = true;

    // parameter to contron the radius of parameter randomization around the curren tminimum
    radius = 1.0;
    randomization_rate = 0.3;

    // The chosen variant of the cost function
    cost_fnc = FROBENIUS_NORM;


    // variables to calculate the cost function with first and second corrections    
    prev_cost_fnv_val = 1.0;
    correction1_scale = 1/1.7;
    correction2_scale = 1/2.0; 


    // set the trace offset
    trace_offset = 0;

    // unique id indentifying the instance of the class
    std::uniform_int_distribution<> distrib_int(0, INT_MAX);  
    id = distrib_int(gen);



    // Time spent on circuit simulation/cost function evaluation
    circuit_simulation_time = 0.0;
    // time spent on optimization
    CPU_time = 0.0;

#if defined __DFE__
    // number of utilized accelerators
    accelerator_num = accelerator_num_in;
#elif defined __GROQ__
    // number of utilized accelerators
    accelerator_num = accelerator_num_in;
#else
    accelerator_num = 0;
#endif

}

Optimization_Interface::Optimization_Interface( Matrix_float Umtx_in, int qbit_num_in, bool optimize_layer_num_in, std::map<std::string, Config_Element>& config, guess_type initial_guess_in, int accelerator_num_in ) : Decomposition_Base(Umtx_in, qbit_num_in, config, initial_guess_in) {

    // logical value. Set true if finding the minimum number of gate layers is required (default), or false when the maximal number of two-qubit gates is used (ideal for general unitaries).
    optimize_layer_num  = optimize_layer_num_in;

    // A string describing the type of the class
    type = N_QUBIT_DECOMPOSITION_CLASS_BASE;

    // The global minimum of the optimization problem
    global_target_minimum = 0;

    number_of_iters.store(0, std::memory_order_relaxed);

    // number of iteratrion loops in the optimization
    iteration_loops[2] = 3;

    // filling in numbers that were not given in the input
    for ( std::map<int,int>::iterator it = max_layer_num_def.begin(); it!=max_layer_num_def.end(); it++) {
        if ( max_layer_num.count( it->first ) == 0 ) {
            max_layer_num.insert( std::pair<int, int>(it->first,  it->second) );
        }
    }

    // logical variable indicating whether adaptive learning reate is used in the ADAM algorithm
    adaptive_eta = true;

    // parameter to contron the radius of parameter randomization around the curren tminimum
    radius = 1.0;
    randomization_rate = 0.3;

    // The chosen variant of the cost function
    cost_fnc = FROBENIUS_NORM;

    // variables to calculate the cost function with first and second corrections
    prev_cost_fnv_val = 1.0;
    correction1_scale = 1/1.7;
    correction2_scale = 1/2.0;

    // set the trace offset
    trace_offset = 0;

    // unique id indentifying the instance of the class
    std::uniform_int_distribution<> distrib_int(0, INT_MAX);
    id = distrib_int(gen);

    // Time spent on circuit simulation/cost function evaluation
    circuit_simulation_time = 0.0;
    // time spent on optimization
    CPU_time = 0.0;

#if defined __DFE__
    // number of utilized accelerators
    accelerator_num = accelerator_num_in;
#elif defined __GROQ__
    // number of utilized accelerators
    accelerator_num = accelerator_num_in;
#else
    accelerator_num = 0;
#endif

}



Optimization_Interface::~Optimization_Interface() {


#ifdef __DFE__
    if ( Umtx.cols == Umtx.rows && qbit_num >= 2 && get_accelerator_num() > 0 ) {
        unload_dfe_lib();//releive_DFE();
    }
#endif



}


Optimization_Interface& Optimization_Interface::operator=( const Optimization_Interface& other ) {

    if ( this == &other ) {
        return *this;
    }

    Decomposition_Base::operator=(other);

    max_inner_iterations = other.max_inner_iterations;
    random_shift_count_max = other.random_shift_count_max;
    id = other.id;

    optimize_layer_num = other.optimize_layer_num;
    identical_blocks = other.identical_blocks;
    alg = other.alg;
    cost_fnc = other.cost_fnc;
    prev_cost_fnv_val = other.prev_cost_fnv_val;
    correction1_scale = other.correction1_scale;
    correction2_scale = other.correction2_scale;
    use_cuts = other.use_cuts;
    osr_rank = other.osr_rank;
    use_softmax = other.use_softmax;
    number_of_iters.store(other.number_of_iters.load(std::memory_order_relaxed), std::memory_order_relaxed);
    adaptive_eta = other.adaptive_eta;
    radius = other.radius;
    randomization_rate = other.randomization_rate;
    accelerator_num = other.accelerator_num;
    trace_offset = other.trace_offset;
    circuit_simulation_time = other.circuit_simulation_time;
    CPU_time = other.CPU_time;

    return *this;

}


void Optimization_Interface::export_current_cost_fnc(double current_minimum ){

    export_current_cost_fnc( current_minimum, optimized_parameters_mtx);
    
}



void Optimization_Interface::export_current_cost_fnc(double current_minimum, Matrix_real& parameters){

    FILE* pFile;
    std::string filename("costfuncs_and_entropy.txt");
    
    if (project_name != "") {
        filename = project_name + "_" + filename;
    }

    const char* c_filename = filename.c_str();
#ifdef _WIN32
    errno_t err = fopen_s(&pFile, c_filename, "a");
    if (err != 0) {
        pFile = NULL;
    }
#else
    pFile = fopen(c_filename, "a");
#endif

    if (pFile==NULL) {
        fputs ("File error",stderr); 
        std::string error("Cannot open file.");
        throw error;
    }
    
    Matrix input_state(Power_of_2(qbit_num),1);

    std::uniform_int_distribution<> distrib(0, qbit_num-2); 

    memset(input_state.get_data(), 0, (input_state.size()*2)*sizeof(double) ); 
    input_state[0].real = 1.0;

    matrix_base<int> qbit_sublist(1,2);
    qbit_sublist[0] = 0;//distrib(gen);
    qbit_sublist[1] = 1;//qbit_sublist[0]+1;

    double renyi_entropy = get_second_Renyi_entropy(parameters, input_state, qbit_sublist);
    
    fprintf(pFile,"%i\t%f\t%f\n", number_of_iters.load(std::memory_order_relaxed), current_minimum, renyi_entropy);
    fclose(pFile);
    
    return;
}


void Optimization_Interface::export_current_cost_fnc(double current_minimum, Matrix_real& parameters, void* void_instance) {


    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);
    
    instance->export_current_cost_fnc( current_minimum, parameters );

}


void 
Optimization_Interface::add_finalyzing_layer() {


    // creating block of gates
    Gates_block* block = new Gates_block( qbit_num );

    // adding U1 gates for final phase corrections
    for (int qbit=0; qbit<qbit_num; qbit++) {
        block->add_u1(qbit);
    }

    // adding the opeartion block to the gates
    add_gate( block );

}




void
Optimization_Interface::calc_decomposition_error(Matrix& decomposed_matrix ) {

   
     // the norm is the square root of the largest einegvalue.*/
    switch (cost_fnc) {
    case FROBENIUS_NORM:
        decomposition_error =  get_cost_function(decomposed_matrix);
        break;
    case FROBENIUS_NORM_CORRECTION1: {
        Matrix_real&& ret = get_cost_function_with_correction(decomposed_matrix, qbit_num);
        decomposition_error = ret[0] - std::sqrt(prev_cost_fnv_val)*ret[1]*correction1_scale;
        break; }
    case FROBENIUS_NORM_CORRECTION2: {
        Matrix_real&& ret = get_cost_function_with_correction2(decomposed_matrix, qbit_num);
        decomposition_error = ret[0] - std::sqrt(prev_cost_fnv_val)*(ret[1]*correction1_scale + ret[2]*correction2_scale);
        break; }
    case HILBERT_SCHMIDT_TEST:
        decomposition_error = get_hilbert_schmidt_test(decomposed_matrix);
        break;
    case HILBERT_SCHMIDT_TEST_CORRECTION1: {
        Matrix&& ret = get_trace_with_correction(decomposed_matrix, qbit_num);
        double d = 1.0/decomposed_matrix.cols;
        decomposition_error = 1 - d*d*(ret[0].real*ret[0].real+ret[0].imag*ret[0].imag+std::sqrt(prev_cost_fnv_val)*correction1_scale*(ret[1].real*ret[1].real+ret[1].imag*ret[1].imag));
        break;
    }
    case HILBERT_SCHMIDT_TEST_CORRECTION2: {
        Matrix&& ret = get_trace_with_correction2(decomposed_matrix, qbit_num);
        double d = 1.0/decomposed_matrix.cols;
        decomposition_error = 1 - d*d*(ret[0].real*ret[0].real+ret[0].imag*ret[0].imag+std::sqrt(prev_cost_fnv_val)*(correction1_scale*(ret[1].real*ret[1].real+ret[1].imag*ret[1].imag)+correction2_scale*(ret[2].real*ret[2].real+ret[2].imag*ret[2].imag)));
        break;
    }
    case SUM_OF_SQUARES:
        decomposition_error = get_cost_function_sum_of_squares(decomposed_matrix);
        break;
    case INFIDELITY:
        decomposition_error = get_infidelity(decomposed_matrix);
        break;
    case OSR_ENTANGLEMENT:
        decomposition_error = get_osr_entanglement_test(decomposed_matrix, use_cuts, osr_rank, use_softmax);
        break;
    default: {
        std::string err("Optimization_Interface::optimization_problem: Cost function variant not implmented.");
        throw err;
    } }

}



void  Optimization_Interface::final_optimization() {

     //The stringstream input to store the output messages.
     std::stringstream sstream;
     sstream << "***************************************************************" << std::endl;
     sstream << "Final fine tuning of the parameters in the " << qbit_num << "-qubit decomposition" << std::endl;
     sstream << "***************************************************************" << std::endl;
     print(sstream, 1);       

         


        //# setting the global minimum
        global_target_minimum = 0;

        if ( optimized_parameters_mtx.size() == 0 ) {
            solve_optimization_problem(NULL, 0);
        }
        else {
            current_minimum = optimization_problem(optimized_parameters_mtx.get_data());
            if ( check_optimization_solution() ) return;

            solve_optimization_problem(optimized_parameters_mtx.get_data(), parameter_num);
        }
}




void Optimization_Interface::solve_layer_optimization_problem( int num_of_parameters, Matrix_real solution_guess) {


    switch ( alg ) {
        case ADAM:
            solve_layer_optimization_problem_ADAM( num_of_parameters, solution_guess);
            break;
        case ADAM_BATCHED:
            solve_layer_optimization_problem_ADAM_BATCHED( num_of_parameters, solution_guess);
            break;
        case GRAD_DESCEND:
            solve_layer_optimization_problem_GRAD_DESCEND( num_of_parameters, solution_guess);
            break;
        case AGENTS:
            solve_layer_optimization_problem_AGENTS( num_of_parameters, solution_guess);
            break;
        case COSINE:
            solve_layer_optimization_problem_COSINE( num_of_parameters, solution_guess);
            break;
        case GRAD_DESCEND_PARAMETER_SHIFT_RULE:
            solve_layer_optimization_problem_GRAD_DESCEND_PARAMETER_SHIFT_RULE( num_of_parameters, solution_guess);
            break;
        case AGENTS_COMBINED:
            solve_layer_optimization_problem_AGENTS_COMBINED( num_of_parameters, solution_guess);
            break;
        case BFGS:
            solve_layer_optimization_problem_BFGS( num_of_parameters, solution_guess);
            break;
        case BAYES_OPT:
            solve_layer_optimization_problem_BAYES_OPT( num_of_parameters, solution_guess);
            break;
        case BAYES_AGENTS:
            solve_layer_optimization_problem_BAYES_AGENTS( num_of_parameters, solution_guess);
            break;
        case BFGS2:
            solve_layer_optimization_problem_BFGS2( num_of_parameters, solution_guess);
            break;
        default:
            std::string error("Optimization_Interface::solve_layer_optimization_problem: unimplemented optimization algorithm");
            throw error;
    }

    if ( use_float ) {
        optimized_parameters_mtx_float = Matrix_real_float(1, optimized_parameters_mtx.size());
        for (int idx=0; idx<optimized_parameters_mtx.size(); idx++) {
            optimized_parameters_mtx_float[idx] = static_cast<float>(optimized_parameters_mtx[idx]);
        }
    }

}

void Optimization_Interface::randomize_parameters( Matrix_real& input, Matrix_real& output, const double f0  ) {

    // random generator of real numbers   
    std::uniform_real_distribution<> distrib_prob(0.0, 1.0);
    std::uniform_real_distribution<> distrib_real(-2*M_PI, 2*M_PI);


    double radius_loc;
    if ( config.count("Randomized_Radius") > 0 ) {
        config["Randomized_Radius"].get_property( radius_loc );  
    }
    else {
        radius_loc = radius;
    }

    const int num_of_parameters = input.size();

    int changed_parameters = 0;
    for ( int jdx=0; jdx<num_of_parameters; jdx++) {
        if ( distrib_prob(gen) <= randomization_rate ) {
            output[jdx] = (cost_fnc!=VQE) ? input[jdx] + distrib_real(gen)*std::sqrt(f0)*radius_loc : input[jdx] + distrib_real(gen)*radius_loc;
            changed_parameters++;
        }
        else {
            output[jdx] = input[jdx];
        }
    }

#ifdef __MPI__  
        //MPI_Bcast( (void*)output.get_data(), num_of_parameters, MPI_DOUBLE, 0, MPI_COMM_WORLD);
#endif


}


double Optimization_Interface::optimization_problem( double* parameters ) {

    // get the transformed matrix with the gates in the list
    Matrix_real parameters_mtx(parameters, 1, parameter_num );
      
    
    return optimization_problem( parameters_mtx );


}


double Optimization_Interface::optimization_problem( Matrix_real& parameters ) {

    // get the transformed matrix with the gates in the list
    if ( parameters.size() != parameter_num ) {
        std::stringstream sstream;
     sstream << "Optimization_Interface::optimization_problem: Number of free paramaters should be " << parameter_num << ", but got " << parameters.size() << std::endl;
        print(sstream, 0);
        std::string err("Optimization_Interface::optimization_problem: Wrong number of parameters.");
        throw err;
    }  
    
    // Float32 circuit application is useful for experimental hot paths, but
    // optimizer cost values must stay in double precision.  The objective and
    // gradients guide convergence and fidelity; evaluating them through
    // Matrix_float introduces enough drift to steer the solver to bad minima.
    const bool use_float_cost_path = false;
    if ( use_float_cost_path && use_float ) {
        static tbb::enumerable_thread_specific<Matrix_real_float> parameters_float_tls;
        static tbb::enumerable_thread_specific<Matrix_float> matrix_new_tls;
        Matrix_real_float& parameters_float = parameters_float_tls.local();
        Matrix_float& matrix_new = matrix_new_tls.local();
        parameters.copy_to(parameters_float);
        Umtx_float.copy_to(matrix_new);
        Gates_block::apply_to( parameters_float, matrix_new );
        return calculate_cost_function(matrix_new, NULL);
    }

    static tbb::enumerable_thread_specific<Matrix> matrix_new_tls;
    Matrix& matrix_new = matrix_new_tls.local();
    Umtx.copy_to(matrix_new);
    Gates_block::apply_to( parameters, matrix_new );

    return calculate_cost_function(matrix_new, NULL);

}


double Optimization_Interface::calculate_cost_function( Matrix& matrix_new, Matrix* ret_temp ) {

    switch (cost_fnc) {
    case FROBENIUS_NORM:
        return get_cost_function(matrix_new, trace_offset);
    case FROBENIUS_NORM_CORRECTION1: {
        Matrix_real&& ret = get_cost_function_with_correction(matrix_new, qbit_num, trace_offset);
        return ret[0] - std::sqrt(prev_cost_fnv_val)*ret[1]*correction1_scale; }
    case FROBENIUS_NORM_CORRECTION2: {
        Matrix_real&& ret = get_cost_function_with_correction2(matrix_new, qbit_num, trace_offset);
        return ret[0] - std::sqrt(prev_cost_fnv_val)*(ret[1]*correction1_scale + ret[2]*correction2_scale); }
    case HILBERT_SCHMIDT_TEST:
        if ( ret_temp != NULL ) {
            QGD_Complex16 trace_temp = get_trace(matrix_new);
            (*ret_temp)[0].real = trace_temp.real;
            (*ret_temp)[0].imag = trace_temp.imag;
            double d = 1.0/matrix_new.cols;
            return 1 - d*d*(trace_temp.real*trace_temp.real + trace_temp.imag*trace_temp.imag);
        }
        return get_hilbert_schmidt_test(matrix_new);
    case HILBERT_SCHMIDT_TEST_CORRECTION1: {
        Matrix&& ret = get_trace_with_correction(matrix_new, qbit_num);
        double d = 1.0/matrix_new.cols;
        if ( ret_temp != NULL ) {
            for (int idx=0; idx<3; idx++) {
                (*ret_temp)[idx].real = ret[idx].real;
                (*ret_temp)[idx].imag = ret[idx].imag;
            }
        }
        return 1 - d*d*(ret[0].real*ret[0].real+ret[0].imag*ret[0].imag+std::sqrt(prev_cost_fnv_val)*correction1_scale*(ret[1].real*ret[1].real+ret[1].imag*ret[1].imag)); }
    case HILBERT_SCHMIDT_TEST_CORRECTION2: {
        Matrix&& ret = get_trace_with_correction2(matrix_new, qbit_num);
        double d = 1.0/matrix_new.cols;
        if ( ret_temp != NULL ) {
            for (int idx=0; idx<4; idx++) {
                (*ret_temp)[idx].real = ret[idx].real;
                (*ret_temp)[idx].imag = ret[idx].imag;
            }
        }
        return 1 - d*d*(ret[0].real*ret[0].real+ret[0].imag*ret[0].imag+std::sqrt(prev_cost_fnv_val)*(correction1_scale*(ret[1].real*ret[1].real+ret[1].imag*ret[1].imag)+correction2_scale*(ret[2].real*ret[2].real+ret[2].imag*ret[2].imag))); }
    case SUM_OF_SQUARES:
        return get_cost_function_sum_of_squares(matrix_new);
    case INFIDELITY:
        if ( ret_temp != NULL ) {
            QGD_Complex16 trace_temp = get_trace(matrix_new);
            (*ret_temp)[0].real = trace_temp.real;
            (*ret_temp)[0].imag = trace_temp.imag;
            double d = matrix_new.cols;
            return 1.0-((trace_temp.real*trace_temp.real+trace_temp.imag*trace_temp.imag)/d+1)/(d+1);
        }
        return get_infidelity(matrix_new);
    case OSR_ENTANGLEMENT:
        return get_osr_entanglement_test(matrix_new, use_cuts, osr_rank, use_softmax);
    default: {
        std::string err("Optimization_Interface::optimization_problem: Cost function variant not implmented.");
        throw err;
    } }

}

double Optimization_Interface::calculate_cost_function( Matrix_float& matrix_new, Matrix_float* ret_temp ) {

    // Experimental helper only.  Optimizer objective/gradient callbacks must
    // use calculate_cost_function(Matrix&, ...) so the solver is guided by
    // double precision cost values.
    switch (cost_fnc) {
    case FROBENIUS_NORM:
        return get_cost_function(matrix_new, trace_offset);
    case FROBENIUS_NORM_CORRECTION1: {
        Matrix_real_float ret = get_cost_function_with_correction(matrix_new, qbit_num, trace_offset);
        return ret[0] - std::sqrt(prev_cost_fnv_val)*ret[1]*correction1_scale; }
    case FROBENIUS_NORM_CORRECTION2: {
        Matrix_real_float ret = get_cost_function_with_correction2(matrix_new, qbit_num, trace_offset);
        return ret[0] - std::sqrt(prev_cost_fnv_val)*(ret[1]*correction1_scale + ret[2]*correction2_scale); }
    case HILBERT_SCHMIDT_TEST: {
        QGD_Complex16 trace_temp = get_trace(matrix_new);
        if ( ret_temp != NULL ) {
            (*ret_temp)[0].real = static_cast<float>(trace_temp.real);
            (*ret_temp)[0].imag = static_cast<float>(trace_temp.imag);
        }
        double d = 1.0/matrix_new.cols;
        return 1 - d*d*(trace_temp.real*trace_temp.real + trace_temp.imag*trace_temp.imag); }
    case INFIDELITY: {
        QGD_Complex16 trace_temp = get_trace(matrix_new);
        if ( ret_temp != NULL ) {
            (*ret_temp)[0].real = static_cast<float>(trace_temp.real);
            (*ret_temp)[0].imag = static_cast<float>(trace_temp.imag);
        }
        double d = matrix_new.cols;
        return 1.0-((trace_temp.real*trace_temp.real+trace_temp.imag*trace_temp.imag)/d+1)/(d+1); }
    case HILBERT_SCHMIDT_TEST_CORRECTION1:
    case HILBERT_SCHMIDT_TEST_CORRECTION2:
    case SUM_OF_SQUARES:
    case OSR_ENTANGLEMENT: {
        if (cost_fnc == OSR_ENTANGLEMENT) {
            return get_osr_entanglement_test(matrix_new, use_cuts, osr_rank, use_softmax);
        }
        Matrix matrix_new64 = matrix_new.to_float64();
        return calculate_cost_function(matrix_new64, NULL); }
    default: {
        std::string err("Optimization_Interface::calculate_cost_function(Matrix_float&): Cost function variant not implmented.");
        throw err;
    } }

}

#ifdef __DFE__

Matrix_real 
Optimization_Interface::optimization_problem_batched_DFE( std::vector<Matrix_real>& parameters_vec) {


    Matrix_real cost_fnc_mtx(parameters_vec.size(), 1);

    int gatesNum, gateSetNum, redundantGateSets;
    DFEgate_kernel_type* DFEgates = convert_to_batched_DFE_gates( parameters_vec, gatesNum, gateSetNum, redundantGateSets );                        
            
    Matrix_real trace_DFE_mtx(gateSetNum, 3);
        
        
        
    increment_num_iters(static_cast<int>(parameters_vec.size()));
       
        
   
#ifdef __MPI__
    // the number of decomposing layers are divided between the MPI processes

    int mpi_gateSetNum = gateSetNum / world_size;
    int mpi_starting_gateSetIdx = gateSetNum/world_size * current_rank;

    Matrix_real mpi_trace_DFE_mtx(mpi_gateSetNum, 3);
        

    increment_num_iters(mpi_gateSetNum);
        

    {
        DFE_Lib_Read_Lock dfe_lock;
        calcqgdKernelDFE( Umtx.rows, Umtx.cols, DFEgates+mpi_starting_gateSetIdx*gatesNum, gatesNum, mpi_gateSetNum, trace_offset, mpi_trace_DFE_mtx.get_data() );
    }

    int bytes = mpi_trace_DFE_mtx.size()*sizeof(double);
    MPI_Allgather(mpi_trace_DFE_mtx.get_data(), bytes, MPI_BYTE, trace_DFE_mtx.get_data(), bytes, MPI_BYTE, MPI_COMM_WORLD);

#else
    {
        DFE_Lib_Read_Lock dfe_lock;
        calcqgdKernelDFE( Umtx.rows, Umtx.cols, DFEgates, gatesNum, gateSetNum, trace_offset, trace_DFE_mtx.get_data() );
    }
                                                                      
#endif  // __MPI__
      

    // calculate the cost function
    if ( cost_fnc == FROBENIUS_NORM ) {
        for ( int idx=0; idx<parameters_vec.size(); idx++ ) {
            cost_fnc_mtx[idx] = 1-trace_DFE_mtx[idx*3]/Umtx.cols;
        }
    }
    else if ( cost_fnc == FROBENIUS_NORM_CORRECTION1 ) {
        for ( int idx=0; idx<parameters_vec.size(); idx++ ) {
            cost_fnc_mtx[idx] = 1-(trace_DFE_mtx[idx*3] + std::sqrt(prev_cost_fnv_val)*trace_DFE_mtx[idx*3+1]*correction1_scale)/Umtx.cols;
        }
    }
    else if ( cost_fnc == FROBENIUS_NORM_CORRECTION2 ) {
        for ( int idx=0; idx<parameters_vec.size(); idx++ ) {
            cost_fnc_mtx[idx] = 1-(trace_DFE_mtx[idx*3] + std::sqrt(prev_cost_fnv_val)*(trace_DFE_mtx[idx*3+1]*correction1_scale + trace_DFE_mtx[idx*3+2]*correction2_scale))/Umtx.cols;
        }
    }
    else {
        std::string err("Optimization_Interface::optimization_problem_batched: Cost function variant not implmented for DFE.");
        throw err;
    }


       


    delete[] DFEgates;
    
    return cost_fnc_mtx;

}
#endif



#ifdef __GROQ__


double 
Optimization_Interface::optimization_problem_Groq( Matrix_real& parameters, int chosen_device) {

    throw std::string("Optimization_Interface::optimization_problem_Groq should be implemented in derrived classes.");
    
    return 0.0;

}



Matrix_real 
Optimization_Interface::optimization_problem_batched_Groq( std::vector<Matrix_real>& parameters_vec) {

    int task_num = parameters_vec.size();
    
    Matrix_real cost_fnc_mtx(task_num, 1);
    
    if ( get_initialize_id() != id ) {     
        init_groq_sv_lib(accelerator_num, id);
    }
    
    if ( accelerator_num == 1 ) {
    
        int chosen_device = 0;

        for( int idx=0; idx<task_num; idx++ ) {    
            cost_fnc_mtx[idx] = optimization_problem_Groq( parameters_vec[idx], chosen_device ); 
        }
            
    }
    else if ( accelerator_num == 2 ) {
    
        for( int idx=0; idx<task_num; idx=idx+2 ) {    
            tbb::parallel_invoke(
                [&]() { cost_fnc_mtx[idx] = optimization_problem_Groq( parameters_vec[idx], 0 );  },
                [&]() { if ( (idx+1) < task_num )  cost_fnc_mtx[idx+1] = optimization_problem_Groq( parameters_vec[idx+1], 1 );  }     
            );            
        }
        
    }
    else {
        throw std::string("Unsupported number of Groq accelerators.");
    }
    
    return cost_fnc_mtx;


}
#endif

Matrix_real 
Optimization_Interface::optimization_problem_batched( std::vector<Matrix_real>& parameters_vec) {




             
#if defined __DFE__
    if ( Umtx.cols == Umtx.rows && get_accelerator_num() > 0 ) {
        return optimization_problem_batched_DFE( parameters_vec );
    }
#elif defined __GROQ__
    if ( Umtx.cols == 1 && get_accelerator_num() > 0 ) {
        return optimization_problem_batched_Groq( parameters_vec );
    }
#endif 


    Matrix_real cost_fnc_mtx(static_cast<int>(parameters_vec.size()), 1);
    int parallel = get_parallel_configuration();
    
#ifdef __MPI__


    // the number of decomposing layers are divided between the MPI processes

    int batch_element_num           = parameters_vec.size();
    int mpi_batch_element_num       = batch_element_num / world_size;
    int mpi_batch_element_remainder = batch_element_num % world_size;

    if ( mpi_batch_element_remainder > 0 ) {
        std::string err("Optimization_Interface::optimization_problem_batched: The size of the batch should be divisible with the number of processes.");
        throw err;
    }

    int mpi_starting_batchIdx = mpi_batch_element_num * current_rank;


    Matrix_real cost_fnc_mtx_loc(mpi_batch_element_num, 1);

    auto calculate_local_batch_element = [&](int idx) {
        cost_fnc_mtx_loc[idx] = optimization_problem( parameters_vec[idx + mpi_starting_batchIdx] );
    };

    if ( parallel == 0 ) {
        for (int idx=0; idx<mpi_batch_element_num; ++idx) {
            calculate_local_batch_element(idx);
        }
    }
    else {
        int work_batch = 1;
        tbb::parallel_for( tbb::blocked_range<int>(0, (int)mpi_batch_element_num, work_batch), [&](tbb::blocked_range<int> r) {
            for (int idx=r.begin(); idx<r.end(); ++idx) { 
                calculate_local_batch_element(idx);
            }
        });
    }

    //number_of_iters = number_of_iters + mpi_batch_element_num; 



    int bytes = cost_fnc_mtx_loc.size()*sizeof(double);
    MPI_Allgather(cost_fnc_mtx_loc.get_data(), bytes, MPI_BYTE, cost_fnc_mtx.get_data(), bytes, MPI_BYTE, MPI_COMM_WORLD);


#else

    auto calculate_batch_element = [&](int idx) {
        cost_fnc_mtx[idx] = optimization_problem( parameters_vec[idx] );
    };

    if ( parallel == 0 ) {
        for (int idx=0; idx<(int)parameters_vec.size(); ++idx) {
            calculate_batch_element(idx);
        }
    }
    else {
        int work_batch = 1;
        tbb::parallel_for( tbb::blocked_range<int>(0, (int)parameters_vec.size(), work_batch), [&](tbb::blocked_range<int> r) {
            for (int idx=r.begin(); idx<r.end(); ++idx) { 
                calculate_batch_element(idx);
            }
        });
    }


#endif  // __MPI__
       
     
    return cost_fnc_mtx;
        
}





double Optimization_Interface::optimization_problem( Matrix_real parameters, void* void_instance, Matrix ret_temp) {

    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);
    instance->increment_num_iters();

    // Keep the Matrix_float path compiled for experiments, but never use it
    // for optimizer objective evaluations.  Cost/trace calculations need
    // double precision even when float32 circuit application is enabled
    // elsewhere.
    const bool use_float_cost_path = false;
    if ( use_float_cost_path && instance->get_use_float() ) {
        static tbb::enumerable_thread_specific<Matrix_real_float> parameters_float_tls;
        static tbb::enumerable_thread_specific<Matrix_float> matrix_new_tls;
        static tbb::enumerable_thread_specific<Matrix_float> ret_temp_float_tls;
        Matrix_real_float& parameters_float = parameters_float_tls.local();
        Matrix_float& matrix_new = matrix_new_tls.local();
        Matrix_float& ret_temp_float = ret_temp_float_tls.local();
        parameters.copy_to(parameters_float);
        instance->Umtx_float.copy_to(matrix_new);
        instance->Gates_block::apply_to( parameters_float, matrix_new );
        if (ret_temp_float.rows != ret_temp.rows || ret_temp_float.cols != ret_temp.cols || ret_temp_float.stride != ret_temp.stride) {
            ret_temp_float = Matrix_float(ret_temp.rows, ret_temp.cols, ret_temp.stride);
        }
        return instance->calculate_cost_function(matrix_new, &ret_temp_float);
    }

    static tbb::enumerable_thread_specific<Matrix> matrix_new_tls;
    Matrix& matrix_new = matrix_new_tls.local();
    instance->Umtx.copy_to(matrix_new);
    instance->Gates_block::apply_to( parameters, matrix_new );

    return instance->calculate_cost_function(matrix_new, &ret_temp);


}



double Optimization_Interface::optimization_problem_non_static( Matrix_real parameters, void* void_instance) {

    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);
    Matrix ret(1,3);
    double cost_func = instance->optimization_problem(parameters, void_instance, ret);
    return cost_func;
}




double Optimization_Interface::optimization_problem( Matrix_real parameters, void* void_instance){
    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);
    return instance->optimization_problem_non_static(parameters, void_instance);
}







void Optimization_Interface::optimization_problem_grad( Matrix_real parameters, void* void_instance, Matrix_real& grad ) {

    // The function value at x0
    double f0;

    // calculate the approximate gradient
    optimization_problem_combined( parameters, void_instance, &f0, grad);

}





void Optimization_Interface::optimization_problem_combined_non_static( Matrix_real parameters, void* void_instance, double* f0, Matrix_real& grad ) {

    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);

    int parallel = instance->get_parallel_configuration();

    // the number of free parameters
    int parameter_num_loc = instance->get_parameter_num();

    // the variant of the cost function
    cost_function_type cost_fnc = instance->get_cost_function_variant();

    // value of the cost function from the previous iteration to weigth the correction to the trace
    double prev_cost_fnv_val = instance->get_previous_cost_function_value();
    double correction1_scale    = instance->get_correction1_scale();
    double correction2_scale    = instance->get_correction2_scale();    

    int qbit_num = instance->get_qbit_num();
    int trace_offset_loc = instance->get_trace_offset();

    // Gradient-driven optimization is sensitive to low-precision objective
    // values.  Leave this float32 combined path in place as experimental code,
    // but force production optimizer cost/gradient evaluation through the
    // double precision branch below.
    const bool use_float_combined_cost_path = false;
    if ( use_float_combined_cost_path && instance->get_use_float() ) {
        static tbb::enumerable_thread_specific<Matrix_real_float> parameters_float_tls;
        Matrix_real_float& parameters_float = parameters_float_tls.local();
        parameters.copy_to(parameters_float);
        Matrix_float Umtx_loc = instance->get_Umtx_float();
        static tbb::enumerable_thread_specific<std::vector<Matrix_float>> combined_result_tls;
        std::vector<Matrix_float>& combined_result = combined_result_tls.local();
        instance->Gates_block::apply_to_combined( parameters_float, Umtx_loc, parallel, combined_result );
        Matrix_float& matrix_new = combined_result[0];

        Matrix_float trace_tmp(1,3);
        *f0 = instance->calculate_cost_function(matrix_new, &trace_tmp);

        Matrix Upartial;
        Matrix_float Upartial_float;
        Matrix matrix_new64;
        if (cost_fnc == SUM_OF_SQUARES) {
            matrix_new64 = matrix_new.to_float64();
            Upartial = get_deriv_sum_of_squares(matrix_new64);
        }
        else if (cost_fnc == OSR_ENTANGLEMENT) {
            Upartial_float = get_deriv_osr_entanglement(matrix_new, use_cuts, osr_rank, use_softmax);
        }

        auto calculate_gradient_component = [&](int idx) {
            double grad_comp;
            Matrix_float& deriv_mtx = combined_result[static_cast<size_t>(idx) + 1];
            switch (cost_fnc) {
            case FROBENIUS_NORM:
                grad_comp = (get_cost_function(deriv_mtx, trace_offset_loc) - 1.0);
                break;
            case FROBENIUS_NORM_CORRECTION1: {
                Matrix_real_float deriv_tmp = get_cost_function_with_correction(deriv_mtx, qbit_num, trace_offset_loc);
                grad_comp = (deriv_tmp[0] - std::sqrt(prev_cost_fnv_val)*deriv_tmp[1]*correction1_scale - 1.0);
                break;
            }
            case FROBENIUS_NORM_CORRECTION2: {
                Matrix_real_float deriv_tmp = get_cost_function_with_correction2(deriv_mtx, qbit_num, trace_offset_loc);
                grad_comp = (deriv_tmp[0] - std::sqrt(prev_cost_fnv_val)*(deriv_tmp[1]*correction1_scale + deriv_tmp[2]*correction2_scale) - 1.0);
                break;
            }
            case HILBERT_SCHMIDT_TEST: {
                double d = 1.0/deriv_mtx.cols;
                QGD_Complex16 deriv_tmp = get_trace(deriv_mtx);
                grad_comp = -2.0*d*d*trace_tmp[0].real*deriv_tmp.real-2.0*d*d*trace_tmp[0].imag*deriv_tmp.imag;
                break;
            }
            case INFIDELITY: {
                double d = deriv_mtx.cols;
                QGD_Complex16 deriv_tmp = get_trace(deriv_mtx);
                grad_comp = -2.0/d/(d+1)*trace_tmp[0].real*deriv_tmp.real-2.0/d/(d+1)*trace_tmp[0].imag*deriv_tmp.imag;
                break;
            }
            case SUM_OF_SQUARES: {
                Matrix deriv64 = deriv_mtx.to_float64();
                grad_comp = real_trace_conj_dot(Upartial, deriv64);
                break;
            }
            case OSR_ENTANGLEMENT:
                grad_comp = real_trace_conj_dot(Upartial_float, deriv_mtx);
                break;
            case HILBERT_SCHMIDT_TEST_CORRECTION1:
            case HILBERT_SCHMIDT_TEST_CORRECTION2: {
                Matrix deriv64 = deriv_mtx.to_float64();
                Matrix matrix_new_for_trace = matrix_new.to_float64();
                Matrix trace_tmp64(1,3);
                instance->calculate_cost_function(matrix_new_for_trace, &trace_tmp64);
                if (cost_fnc == HILBERT_SCHMIDT_TEST_CORRECTION1) {
                    Matrix&& deriv_tmp = get_trace_with_correction(deriv64, qbit_num);
                    double d = 1.0/deriv64.cols;
                    grad_comp = -2.0*d*d*(trace_tmp64[0].real*deriv_tmp[0].real+trace_tmp64[0].imag*deriv_tmp[0].imag+std::sqrt(prev_cost_fnv_val)*correction1_scale*(trace_tmp64[1].real*deriv_tmp[1].real+trace_tmp64[1].imag*deriv_tmp[1].imag));
                }
                else {
                    Matrix&& deriv_tmp = get_trace_with_correction2(deriv64, qbit_num);
                    double d = 1.0/deriv64.cols;
                    grad_comp = -2.0*d*d*(trace_tmp64[0].real*deriv_tmp[0].real+trace_tmp64[0].imag*deriv_tmp[0].imag+std::sqrt(prev_cost_fnv_val)*(correction1_scale*(trace_tmp64[1].real*deriv_tmp[1].real+trace_tmp64[1].imag*deriv_tmp[1].imag) + correction2_scale*(trace_tmp64[2].real*deriv_tmp[2].real+trace_tmp64[2].imag*deriv_tmp[2].imag)));
                }
                break;
            }
            default: {
                std::string err("Optimization_Interface::optimization_problem_combined: Cost function variant not implmented.");
                throw err;
            } }
            grad[idx] = grad_comp;
        };

        if ( parallel == 0 ) {
            for (int idx=0; idx<parameter_num_loc; ++idx) {
                calculate_gradient_component(idx);
            }
        }
        else {
            int work_batch = 10;
            tbb::parallel_for( tbb::blocked_range<int>(0,parameter_num_loc,work_batch), [&](tbb::blocked_range<int> r) {
                for (int idx=r.begin(); idx<r.end(); ++idx) {
                    calculate_gradient_component(idx);
                }
            });
        }

        instance->increment_num_iters(parameter_num_loc + 1);
        std::stringstream sstream;
        sstream << *f0 << std::endl;
        instance->print(sstream, 5);
        return;
    }

#ifdef __DFE__

//std::cout << "number of qubits: " << instance->qbit_num << std::endl;
//tbb::tick_count t0_DFE = tbb::tick_count::now();/////////////////////////////////    
if ( Umtx.cols == Umtx.rows && instance->qbit_num >= 5 && instance->get_accelerator_num() > 0 ) {

    int gatesNum, redundantGateSets, gateSetNum;
    DFEgate_kernel_type* DFEgates = instance->convert_to_DFE_gates_with_derivates( parameters, gatesNum, gateSetNum, redundantGateSets );

    Matrix&& Umtx_loc = instance->get_Umtx();   
    Matrix_real trace_DFE_mtx(gateSetNum, 3);


#ifdef __MPI__
    // the number of decomposing layers are divided between the MPI processes

    int mpi_gateSetNum = gateSetNum / instance->world_size;
    int mpi_starting_gateSetIdx = gateSetNum/instance->world_size * instance->current_rank;

    Matrix_real mpi_trace_DFE_mtx(mpi_gateSetNum, 3);
    
    instance->increment_num_iters(mpi_gateSetNum);

    {
        DFE_Lib_Read_Lock dfe_lock;
        calcqgdKernelDFE( Umtx_loc.rows, Umtx_loc.cols, DFEgates+mpi_starting_gateSetIdx*gatesNum, gatesNum, mpi_gateSetNum, trace_offset_loc, mpi_trace_DFE_mtx.get_data() );
    }

    int bytes = mpi_trace_DFE_mtx.size()*sizeof(double);
    MPI_Allgather(mpi_trace_DFE_mtx.get_data(), bytes, MPI_BYTE, trace_DFE_mtx.get_data(), bytes, MPI_BYTE, MPI_COMM_WORLD);

#else

    instance->increment_num_iters(gateSetNum);

    {
        DFE_Lib_Read_Lock dfe_lock;
        calcqgdKernelDFE( Umtx_loc.rows, Umtx_loc.cols, DFEgates, gatesNum, gateSetNum, trace_offset_loc, trace_DFE_mtx.get_data() );
    }

#endif  

    std::stringstream sstream;
    sstream << *f0 << " " << 1.0 - trace_DFE_mtx[0]/Umtx_loc.cols << " " << trace_DFE_mtx[1]/Umtx_loc.cols << " " << trace_DFE_mtx[2]/Umtx_loc.cols << std::endl;
    instance->print(sstream, 5);   
    
  
    if ( cost_fnc == FROBENIUS_NORM ) {
        *f0 = 1-trace_DFE_mtx[0]/Umtx_loc.cols;
    }
    else if ( cost_fnc == FROBENIUS_NORM_CORRECTION1 ) {
        *f0 = 1 - (trace_DFE_mtx[0] + std::sqrt(prev_cost_fnv_val)*trace_DFE_mtx[1]*correction1_scale)/Umtx_loc.cols;
    }
    else if ( cost_fnc == FROBENIUS_NORM_CORRECTION2 ) {
        *f0 = 1 - (trace_DFE_mtx[0] + std::sqrt(prev_cost_fnv_val)*(trace_DFE_mtx[1]*correction1_scale + trace_DFE_mtx[2]*correction2_scale))/Umtx_loc.cols;
    }
    else {
        std::string err("Optimization_Interface::optimization_problem_combined: Cost function variant not implmented.");
        throw err;
    }

    //double f0_DFE = *f0;

    //Matrix_real grad_components_DFE_mtx(1, parameter_num_loc);
    for (int idx=0; idx<parameter_num_loc; idx++) {

        if ( cost_fnc == FROBENIUS_NORM ) {
            grad[idx] = -trace_DFE_mtx[3*(idx+1)]/Umtx_loc.cols;
        }
        else if ( cost_fnc == FROBENIUS_NORM_CORRECTION1 ) {
            grad[idx] = -(trace_DFE_mtx[3*(idx+1)] + std::sqrt(prev_cost_fnv_val)*trace_DFE_mtx[3*(idx+1)+1]*correction1_scale)/Umtx_loc.cols;
        }
        else if ( cost_fnc == FROBENIUS_NORM_CORRECTION2 ) {
            grad[idx] = -(trace_DFE_mtx[3*(idx+1)] + std::sqrt(prev_cost_fnv_val)*(trace_DFE_mtx[3*(idx+1)+1]*correction1_scale + trace_DFE_mtx[3*(idx+1)+2]*correction2_scale))/Umtx_loc.cols;
        }
        else {
            std::string err("Optimization_Interface::optimization_problem_combined: Cost function variant not implmented.");
            throw err;
        }

        //grad_components_DFE_mtx[idx] = gsl_vector_get( grad, idx );


    }

    delete[] DFEgates;

//tbb::tick_count t1_DFE = tbb::tick_count::now();/////////////////////////////////
//std::cout << "uploaded data to DFE: " << (int)(gatesNum*gateSetNum*sizeof(DFEgate_kernel_type)) << " bytes" << std::endl;
//std::cout << "time elapsed DFE: " << (t1_DFE-t0_DFE).seconds() << ", expected time: " << (((double)Umtx_loc.rows*(double)Umtx_loc.cols*gatesNum*gateSetNum/get_chained_gates_num()/4 + 4578*3*get_chained_gates_num()))/350000000 + 0.001<< std::endl;

}
else {

#endif

#ifdef __DFE__
tbb::tick_count t0_CPU = tbb::tick_count::now();
#endif

    // vector containing gradients of the transformed matrix
    Matrix trace_tmp(1,3);
    Matrix Umtx_loc = instance->get_Umtx();
    static tbb::enumerable_thread_specific<std::vector<Matrix>> combined_result_tls;
    std::vector<Matrix>& combined_result = combined_result_tls.local();
    instance->apply_to_combined( parameters, Umtx_loc, parallel, combined_result );
    Matrix& matrix_new = combined_result[0];

    *f0 = instance->calculate_cost_function(matrix_new, &trace_tmp);

    Matrix Upartial;
    if (cost_fnc == SUM_OF_SQUARES) {
        Upartial = get_deriv_sum_of_squares(matrix_new);
    } else if (cost_fnc == OSR_ENTANGLEMENT) {
        Upartial = get_deriv_osr_entanglement(matrix_new, use_cuts, osr_rank, use_softmax);
    }


    auto calculate_gradient_component = [&](int idx) {
        double grad_comp;
        Matrix& deriv_mtx = combined_result[static_cast<size_t>(idx) + 1];
        switch (cost_fnc) {
        case FROBENIUS_NORM:
            grad_comp = (get_cost_function(deriv_mtx, trace_offset_loc) - 1.0);
            break;
        case FROBENIUS_NORM_CORRECTION1: {
            Matrix_real deriv_tmp = get_cost_function_with_correction( deriv_mtx, qbit_num, trace_offset_loc );
            grad_comp = (deriv_tmp[0] - std::sqrt(prev_cost_fnv_val)*deriv_tmp[1]*correction1_scale - 1.0);
            break;
        }
        case FROBENIUS_NORM_CORRECTION2: {
            Matrix_real deriv_tmp = get_cost_function_with_correction2( deriv_mtx, qbit_num, trace_offset_loc );
            grad_comp = (deriv_tmp[0] - std::sqrt(prev_cost_fnv_val)*(deriv_tmp[1]*correction1_scale + deriv_tmp[2]*correction2_scale) - 1.0);
            break;
        }
        case HILBERT_SCHMIDT_TEST: {
            double d = 1.0/deriv_mtx.cols;
            QGD_Complex16 deriv_tmp = get_trace(deriv_mtx);
            grad_comp = -2.0*d*d*trace_tmp[0].real*deriv_tmp.real-2.0*d*d*trace_tmp[0].imag*deriv_tmp.imag;
            break;
        }
        case HILBERT_SCHMIDT_TEST_CORRECTION1: {
            Matrix&&  deriv_tmp = get_trace_with_correction( deriv_mtx, qbit_num);
            double d = 1.0/deriv_mtx.cols;
            grad_comp = -2.0*d*d* (trace_tmp[0].real*deriv_tmp[0].real+trace_tmp[0].imag*deriv_tmp[0].imag+std::sqrt(prev_cost_fnv_val)*correction1_scale*(trace_tmp[1].real*deriv_tmp[1].real+trace_tmp[1].imag*deriv_tmp[1].imag));
            break;
        }
        case HILBERT_SCHMIDT_TEST_CORRECTION2: {
            Matrix&& deriv_tmp = get_trace_with_correction2( deriv_mtx, qbit_num);
            double d = 1.0/deriv_mtx.cols;
            grad_comp = -2.0*d*d* (trace_tmp[0].real*deriv_tmp[0].real+trace_tmp[0].imag*deriv_tmp[0].imag+std::sqrt(prev_cost_fnv_val)*(correction1_scale*(trace_tmp[1].real*deriv_tmp[1].real+trace_tmp[1].imag*deriv_tmp[1].imag) + correction2_scale*(trace_tmp[2].real*deriv_tmp[2].real+trace_tmp[2].imag*deriv_tmp[2].imag)));
            break;
        }
        case SUM_OF_SQUARES:
        case OSR_ENTANGLEMENT:
            grad_comp = real_trace_conj_dot(Upartial, deriv_mtx);
            /*{
                Matrix matrix_new = instance->get_Umtx().copy();
                auto paramcopy = parameters.copy();
                paramcopy[idx] += 1e-10;
                instance->apply_to( paramcopy, matrix_new );
                double f1 = instance->get_cost_function_variant() == SUM_OF_SQUARES ? get_cost_function_sum_of_squares(matrix_new) : get_osr_entanglement_test(matrix_new, use_cuts, osr_rank, use_softmax);
                double check = (f1 - *f0) / 1e-10;
                //printf("%d: %g %g\n", idx, grad_comp, check);
                grad_comp = check;
            }*/
            break;
        case INFIDELITY: {
            double d = deriv_mtx.cols;
            QGD_Complex16 deriv_tmp = get_trace(deriv_mtx);
            grad_comp = -2.0/d/(d+1)*trace_tmp[0].real*deriv_tmp.real-2.0/d/(d+1)*trace_tmp[0].imag*deriv_tmp.imag;
            break;
        }
        default: {
            std::string err("Optimization_Interface::optimization_problem_combined: Cost function variant not implmented.");
            throw err;
        } }
        
        grad[idx] = grad_comp;
    };

    if ( parallel == 0 ) {
        for (int idx=0; idx<parameter_num_loc; ++idx) {
            calculate_gradient_component(idx);
        }
    }
    else {
        int work_batch = 10;
        tbb::parallel_for( tbb::blocked_range<int>(0,parameter_num_loc,work_batch), [&](tbb::blocked_range<int> r) {
            for (int idx=r.begin(); idx<r.end(); ++idx) {
                calculate_gradient_component(idx);
            }
        });
    }

    instance->increment_num_iters(parameter_num_loc + 1);

    std::stringstream sstream;
    sstream << *f0 << std::endl;
    instance->print(sstream, 5);   

#ifdef __DFE__
}
#endif


}



void Optimization_Interface::optimization_problem_combined( Matrix_real parameters, void* void_instance, double* f0, Matrix_real& grad ){
    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);
    instance->optimization_problem_combined_non_static(parameters, void_instance, f0, grad );
    return;
}


void Optimization_Interface::optimization_problem_combined( Matrix_real parameters, double* f0, Matrix_real grad ) {

    optimization_problem_combined( parameters, this, f0, grad );
    return;
}



void Optimization_Interface::optimization_problem_combined_unitary( Matrix_real parameters, void* void_instance, Matrix& Umtx, std::vector<Matrix>& Umtx_deriv ) {
    // vector containing gradients of the transformed matrix
    Optimization_Interface* instance = reinterpret_cast<Optimization_Interface*>(void_instance);

    int parallel = instance->get_parallel_configuration();

    Matrix Umtx_loc = instance->get_Umtx();
    static tbb::enumerable_thread_specific<std::vector<Matrix>> combined_result_tls;
    std::vector<Matrix>& combined_result = combined_result_tls.local();
    instance->apply_to_combined( parameters, Umtx_loc, parallel, combined_result );
    combined_result[0].copy_to(Umtx);
    Umtx_deriv.resize(combined_result.size() - 1);
    for (size_t idx = 1; idx < combined_result.size(); ++idx) {
        combined_result[idx].copy_to(Umtx_deriv[idx - 1]);
    }



  



}


void Optimization_Interface::optimization_problem_combined_unitary( Matrix_real parameters, Matrix& Umtx, std::vector<Matrix>& Umtx_deriv ) {

    optimization_problem_combined_unitary( parameters, this, Umtx, Umtx_deriv);
    return;
    
}


cost_function_type 
Optimization_Interface::get_cost_function_variant() {

    return cost_fnc;

}


void 
Optimization_Interface::set_cost_function_variant( cost_function_type variant  ) {

    cost_fnc = variant;

    std::stringstream sstream;
    sstream << "Optimization_Interface::set_cost_function_variant: Cost function variant set to " << cost_fnc << std::endl;
    print(sstream, 2);   


}



void Optimization_Interface::set_max_inner_iterations( int max_inner_iterations_in  ) {

    max_inner_iterations = max_inner_iterations_in;
    
}



void Optimization_Interface::set_random_shift_count_max( int random_shift_count_max_in  ) {

    random_shift_count_max = random_shift_count_max_in;

}


void Optimization_Interface::set_optimizer( optimization_aglorithms alg_in ) {

    alg = alg_in;

    switch ( alg ) {
        case ADAM:
            max_inner_iterations = 100000; 
            random_shift_count_max = 100;
            max_outer_iterations = 1;
            return;

        case ADAM_BATCHED:
            max_inner_iterations = 2500;
            random_shift_count_max = 3;
            max_outer_iterations = 1;
            return;

        case GRAD_DESCEND:
            max_inner_iterations = 10000;
            random_shift_count_max = 1;  
            max_outer_iterations = 100000000; 
            return;

        case COSINE:
            max_inner_iterations = 2500;
            random_shift_count_max = 3;
            max_outer_iterations = 1;
            return;
            
        case GRAD_DESCEND_PARAMETER_SHIFT_RULE:
            max_inner_iterations = 2500;
            random_shift_count_max = 3;
            max_outer_iterations = 1;
            return;                       

        case AGENTS:
            max_inner_iterations = 2500;
            random_shift_count_max = 3;
            max_outer_iterations = 1;
            return;

        case AGENTS_COMBINED:
            max_inner_iterations = 2500;
            random_shift_count_max = 3;
            max_outer_iterations = 1;
            return;

        case BFGS:
            max_inner_iterations = 10000;
            random_shift_count_max = 1;  
            max_outer_iterations = 100000000; 
            return;

        case BFGS2:
            max_inner_iterations = 100000;
            random_shift_count_max = 100;
            max_outer_iterations = 1;
            return;
        
        case BAYES_OPT:
            max_inner_iterations = 100;
            random_shift_count_max = 100;
            max_outer_iterations = 1;
            return;
        case BAYES_AGENTS:
            max_inner_iterations = 100;
            random_shift_count_max = 100;
            max_outer_iterations = 1;
            return;

        default:
            std::string error("Optimization_Interface::set_optimizer: unimplemented optimization algorithm");
            throw error;
    }



}




double 
Optimization_Interface::get_previous_cost_function_value() {

    return prev_cost_fnv_val;

}



double 
Optimization_Interface::get_correction1_scale() {

    return correction1_scale;

}



double 
Optimization_Interface::get_correction2_scale() {

    return correction2_scale;

}






int 
Optimization_Interface::get_num_iters() {

    return number_of_iters.load(std::memory_order_relaxed);

}


void Optimization_Interface::increment_num_iters( int delta ) {

    number_of_iters.fetch_add(delta, std::memory_order_relaxed);

}


void 
Optimization_Interface::set_custom_gate_structure( Gates_block* gate_structure_in ) {

    release_gates();

    set_qbit_num( gate_structure_in->get_qbit_num() );

    combine( gate_structure_in );

}


int 
Optimization_Interface::get_trace_offset() {

    return trace_offset;

}


void 
Optimization_Interface::set_trace_offset(int trace_offset_in) {


    if ( (trace_offset_in + Umtx.cols) > Umtx.rows ) {
        std::string error("Optimization_Interface::set_trace_offset: trace offset must be smaller or equal to the difference of the rows and columns in the input unitary.");
        throw error;

    }

    
    trace_offset = trace_offset_in;


    std::stringstream sstream;
    sstream << "Optimization_Interface::set_trace_offset: trace offset set to " << trace_offset << std::endl;
    print(sstream, 2);   

}


#ifdef __DFE__

void 
Optimization_Interface::upload_Umtx_to_DFE() {
    if (Umtx.cols == Umtx.rows) {
        init_dfe_lib_and_upload(accelerator_num, qbit_num, id, Umtx);
    }

}


#endif

int 
Optimization_Interface::get_accelerator_num() {

    return accelerator_num;

}


void Optimization_Interface::set_osr_params( std::vector<std::vector<int>> use_cuts_in, int osr_rank_in, bool use_softmax_in )
{
    use_cuts = use_cuts_in;
    osr_rank = osr_rank_in;
    use_softmax = use_softmax_in;
    //std::stringstream sstream;
    //sstream << "Optimization_Interface::set_osr_params: OSR entanglement test parameters set. osr_rank: " << osr_rank << ", use_softmax: " << use_softmax << std::endl;
    //print(sstream, 2);
}