C++ 行列バランシング

行列バランシングと言えば，sinkhornですが，今回はこの論文を実装しました．

arxiv.org

（論文では自然勾配法使ってますが，ここでは，１次までしか実装してないです）

#include <iostream>
#include <Eigen/Dense>
#include <vector>
#include <cmath>
#include <time.h>
#include <cstdio>
#include <fstream>

using namespace std;
using namespace Eigen;

MatrixXd balance(MatrixXd& P, double lr){
    int N = P.rows();
    MatrixXd eta(N, N), theta(N, N);
    VectorXd eta_beta(N);

    // get eta_beta;
    for (int i = 0; i < N ; i++){
        eta_beta( i ) = ( 1.0 * N - 1.0 * i ) / (1.0 * N);
    }

    // get theta of input
    // we need all element to update p in iteration.
    double sum;
    for (int i = 0; i < N; i++){
        for(int j = 0; j < N; j++){
            sum = 0.0;
            sum = log( P(i, j) );
            if (0 < i)
                sum -= log( P( i-1, j) );
            if (0 < j)
                sum -= log( P( i, j-1 ) );
            if (0 < i and 0 < j)
                sum += log( P( i-1, j-1) );
            theta(i, j) = sum;
        }
    }

    // get eta of input
    eta( 0, N - 1 ) = P.col( N - 1 ).sum();
    eta( N - 1, 0 ) = P.row( N - 1 ).sum();
    for (int i = 2; i < N; i++){
        eta( 0, N - i ) =  eta( 0, N - i + 1) + P.col( N - i ).sum();
        eta( N - i, 0 ) =  eta( N - i + 1,0 ) + P.row( N - i ).sum();
    }

    int cnt = 0;
    while(true){
        // e-projection
        for(int i = 0; i < N ; i++){
            theta(i, 0) -= lr * ( eta(i, 0) - eta_beta( i ) );
            theta(0, i) -= lr * ( eta(0, i) - eta_beta( i ) );
        }

        // update prob by theta
        double logp_ij;
        for (int i = 0; i < N; i++){
            for (int j = 0; j < N; j++){
                logp_ij = 0.0;
                logp_ij = theta(i, j);
                if ( 0 < i )
                    logp_ij += log( P(i-1, j) );
                if ( 0 < j )
                    logp_ij += log( P(i, j-1) );
                if ( 0 < j and 0 < i )
                    logp_ij -= log( P(i-1, j-1) );

                P(i, j) = exp( logp_ij );
            }
        }

        // normalize
        P = P / P.sum();

        // update eta by prob
        eta( 0, N - 1 ) = P.col( N - 1 ).sum();
        eta( N - 1, 0 ) = P.row( N - 1 ).sum();
        for (int i = 2; i < N; i++){
            eta( 0, N - i ) =  eta( 0, N - i + 1) + P.col( N - i ).sum();
            eta( N - i, 0 ) =  eta( N - i + 1,0 ) + P.row( N - i ).sum();
        }


        cnt += 1;
        if (cnt > 2000) break;
    }

    return P;
}

int main(){
    double LO = 0.0; // minimu value of element of random matrix;
    double HI = 1.0; // maximu value of element of random matrix;
    double range = HI - LO;
    double lr = 0.1;

    int N = 25; // matrix dimenstion
    MatrixXd Q(N, N);  // output matrix

    // producting random input matrix P
    MatrixXd P = MatrixXd::Random(N, N);
    P = (P + MatrixXd::Constant(N, N, 1.0))*range/2;
    P = (P + MatrixXd::Constant(N, N, LO));

    /*
    P <<1,5,3,6,
        9,5,5,6,
        1,7,10,9,
       2,3,1,4;
   */

    P = P / P.sum();
    cout << "input matrix" << endl << P << endl;
    Q = balance( P, lr );
    cout << "output matrix" << endl << Q << endl;

    double rowsum,colsum;
    for(int i=0;i<N;i++){
        colsum = Q.row(i).sum();
        rowsum = Q.col(i).sum();
        printf("%d colsum = %f \t rowsum = %f\n",i,colsum,rowsum);
    }

}