bmm.c

Go to the documentation of this file.

/*
 * This work is licensed under a Creative Commons 
 * Attribution-Noncommercial-Share Alike 3.0 United States License.
 * 
 *    http://creativecommons.org/licenses/by-nc-sa/3.0/us/
 * 
 * You are free:
 * 
 *    to Share - to copy, distribute, display, and perform the work
 *    to Remix - to make derivative works
 * 
 * Under the following conditions:
 * 
 *    Attribution. You must attribute the work in the manner specified by the
 *    author or licensor (but not in any way that suggests that they endorse you
 *    or your use of the work).
 * 
 *    Noncommercial. You may not use this work for commercial purposes.
 * 
 *    Share Alike. If you alter, transform, or build upon this work, you may
 *    distribute the resulting work only under the same or similar license to
 *    this one.
 * 
 * For any reuse or distribution, you must make clear to others the license
 * terms of this work. The best way to do this is by including this header.
 * 
 * Any of the above conditions can be waived if you get permission from the
 * copyright holder.
 * 
 * Apart from the remix rights granted under this license, nothing in this
 * license impairs or restricts the author's moral rights.
 */


#include <jwsc/config.h>

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <inttypes.h>

#include "jwsc/base/limits.h"
#include "jwsc/vector/vector.h"
#include "jwsc/vector/vector_math.h"
#include "jwsc/matrix/matrix.h"
#include "jwsc/matrix/matrix_math.h"
#include "jwsc/matblock/matblock.h"
#include "jwsc/prob/mv_pmf.h"
#include "jwsc/stat/stat.h"
#include "jwsc/stat/kmeans.h"
#include "jwsc/stat/bmm.h"


static void bmm_e_step_d
(
    const Matrix_d*   mu,
    const Vector_d*   pi,
    Matrix_d*         gamma,
    const Matrix_u32* x
)
{
    double   b;
    double   C_n;
    uint32_t k, K;
    uint32_t n, N;

    Vector_u32* x_n  = 0;
    Vector_d*   mu_k = 0;

    K = pi->num_elts;
    N = x->num_rows;

    create_zero_matrix_d(&gamma, N, K);

    for (n = 0; n < N; n++)
    {
        copy_vector_from_matrix_u32(&x_n, x, MATRIX_ROW_VECTOR, n);
        C_n = 0;

        for (k = 0; k < K; k++)
        {
            copy_vector_from_matrix_d(&mu_k, mu, MATRIX_ROW_VECTOR, k);

            mv_bernoulli_pmf_d(&b, mu_k, x_n);

            gamma->elts[ n ][ k ] = pi->elts[ k ] * b;

            C_n += gamma->elts[ n ][ k ];
        }

        assert(C_n > 0);

        for (k = 0; k < K; k++)
        {
            gamma->elts[ n ][ k ] /= C_n;
        }
    }

    free_vector_u32(x_n);
    free_vector_d(mu_k);
}


static void bmm_m_step_d
(
    Matrix_d*         mu,
    Vector_d*         pi,
    const Matrix_d*   gamma,
    const Matrix_u32* x
)
{
    uint32_t k, K;
    uint32_t n, N;
    uint32_t d, D;

    Vector_d* mu_k    = NULL;

    K = pi->num_elts;
    N = x->num_rows;
    D = x->num_cols;

    create_zero_vector_d(&pi, K);

    // mu
    create_matrix_d(&mu, K, D);
    for (k = 0; k < K; k++)
    {
        create_zero_vector_d(&mu_k, D);

        for (n = 0; n < N; n++)
        {
            for (d = 0; d < D; d++)
            {
                mu_k->elts[ d ] += gamma->elts[ n ][ k ]*x->elts[ n ][ d ];
            }

            pi->elts[ k ] += gamma->elts[ n ][ k ];
        }

        assert(pi->elts[ k ] > 0);

        multiply_vector_by_scalar_d(&mu_k, mu_k, 1.0 / pi->elts[ k ]);
        copy_vector_into_matrix_d(&mu, mu, mu_k, MATRIX_ROW_VECTOR, k);
    }

    // pi
    multiply_vector_by_scalar_d(&pi, pi, 1.0 / N);

    free_vector_d(mu_k);
}


double bmm_log_likelihood_d
(
    const Matrix_d*   mu,
    const Vector_d*   pi,
    const Matrix_u32* x
)
{
    uint32_t n, N;
    uint32_t k, K;
    double   p, q;
    double   ll;

    Vector_u32* x_n  = NULL;
    Vector_d*   mu_k = NULL;

    N = x->num_rows;
    K = pi->num_elts;

    ll = 0;

    for (n = 0; n < N; n++)
    {
        p = 0;
        copy_vector_from_matrix_u32(&x_n, x, MATRIX_ROW_VECTOR, n);

        for (k = 0; k < K; k++)
        {
            copy_vector_from_matrix_d(&mu_k, mu, MATRIX_ROW_VECTOR, k);

            mv_bernoulli_pmf_d(&q, mu_k, x_n);

            p += pi->elts[ k ] * q;

#if defined(JWSC_HAVE_ISNAN) && defined(JWSC_HAVE_ISINF)
            assert(!isnan(p) && !isinf(p));
#endif
        }

        ll += log((p < JWSC_MIN_LOG_ARG) ? JWSC_MIN_LOG_ARG : p);
    }

    free_vector_u32(x_n);
    free_vector_d(mu_k);

    return ll;
}


void bmm_kmeans_init_d
(
    Matrix_d**        mu_out,
    Vector_d**        pi_out,
    const Matrix_u32* x,
    uint32_t          K
)
{
    uint32_t d, D;
    uint32_t n, nn, N;
    uint32_t k;

    Vector_i32* members = NULL;
    Vector_u32* counts  = NULL;
    Matrix_d*   x_d     = NULL;
    Matrix_d*   x_k     = NULL;

    D = x->num_cols;
    N = x->num_rows;

    create_matrix_d(&x_d, N, D);
    for (n = 0; n < N; n++)
    {
        for (d = 0; d < D; d++)
        {
            x_d->elts[ n ][ d ] = (double)x->elts[ n ][ d ];
        }
    }

    init_kmeans_randomly_d(mu_out, x_d, K);
    kmeans_d(mu_out, &members, x_d, K, 2);

    create_zero_vector_u32(&counts, K);

    for (n = 0; n < N; n++)
    {
        k = members->elts[ n ];
        counts->elts[ k ]++;
    }
    for (k = 0; k < K; k++)
    {
        create_matrix_d(&x_k, counts->elts[ k ], D);

        nn = 0;
        for (n = 0; n < N; n++)
        {
            if (k == members->elts[ n ])
            {
                for (d = 0; d < D; d++)
                {
                    x_k->elts[ nn ][ d ] = x_d->elts[ n ][ d ];
                }
                nn++;
            }
        }
    }

    create_vector_d(pi_out, K);
    for (k = 0; k < K; k++)
    {
        (*pi_out)->elts[ k ] = counts->elts[ k ] / (double)N;
    }

    free_vector_i32(members);
    free_vector_u32(counts);
    free_matrix_d(x_k);
    free_matrix_d(x_d);
}


void train_bmm_d
(
    Matrix_d**        mu_out,
    Vector_d**        pi_out,
    const Matrix_u32* x,
    uint32_t          K
)
{
    double   ll, ll_prev;

    Matrix_d*   mu;
    Vector_d*   pi;
    Matrix_d*   gamma = 0;

    bmm_kmeans_init_d(mu_out, pi_out, x, K);
    mu    = *mu_out;
    pi    = *pi_out;

    create_matrix_d(&gamma, x->num_rows, pi->num_elts);

    ll_prev = 0;
    ll      = log(0);

    while (fabs(ll - ll_prev) > 1.0e-10)
    {
        bmm_e_step_d(mu, pi, gamma, x);
        bmm_m_step_d(mu, pi, gamma, x);

        ll_prev = ll;
        ll = bmm_log_likelihood_d(mu, pi, x);
    }

    free_matrix_d(gamma);
}