haplo_cluster.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 <config.h>

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

#include <libxml/tree.h>

#ifdef HAPLO_HAVE_DMALLOC
#include <dmalloc.h>
#endif

#include <jwsc/base/error.h>
#include <jwsc/base/option.h>
#include <jwsc/base/limits.h>
#include <jwsc/vector/vector.h>
#include <jwsc/matrix/matrix.h>
#include <jwsc/matrix/matrix_io.h>
#include <jwsc/matblock/matblock.h>
#include <jwsc/stat/kmeans.h>
#include <jwsc/stat/gmm.h>
#include <jwsc/stat/mmm.h>

#include "haplo_groups.h"
#include "options.h"
#include "output.h"
#include "input.h"
#include "xml.h"


#define  NUM_OPTS_NO_ARG    0  + NUM_SHARED_OPTS_NO_ARG
#define  NUM_OPTS_WITH_ARG  12 + NUM_SHARED_OPTS_WITH_ARG


#define  LABEL_COL            0
#define  NUM_CLUSTERS         4
#define  CLUSTER_TYPE         HAPLO_CLUSTER_KMEANS
#define  MEANS_OUT_FNAME      "/dev/null"
#define  WEIGHTS_OUT_FNAME    "/dev/null"
#define  RESPONSES_OUT_FNAME  "/dev/null"
#define  MEMBERS_OUT_FNAME    "/dev/stdout"


typedef enum
{
    HAPLO_CLUSTER_KMEANS,
    HAPLO_CLUSTER_GMM,
    HAPLO_CLUSTER_MMM
}
Haplo_cluster_type;


Option_no_arg opts_no_arg[NUM_OPTS_NO_ARG];

Option_with_arg opts_with_arg[NUM_OPTS_WITH_ARG];

static uint32_t num_clusters = NUM_CLUSTERS;

static Haplo_cluster_type cluster_type = CLUSTER_TYPE;

static const char* means_out_fname = MEANS_OUT_FNAME;

static const char* weights_out_fname = WEIGHTS_OUT_FNAME;

static const char* responses_out_fname = RESPONSES_OUT_FNAME;

static const char* members_out_fname = MEMBERS_OUT_FNAME;


uint32_t get_num_opts_no_arg()
{
    return NUM_OPTS_NO_ARG;
}

uint32_t get_num_opts_with_arg()
{
    return NUM_OPTS_WITH_ARG;
}

void print_usage()
{
    fprintf(stderr, "usage: haplo-cluster OPTIONS [data-fname | <stdin>]\n");
    print_options(stderr, 27, NUM_OPTS_NO_ARG, opts_no_arg, NUM_OPTS_WITH_ARG,
            opts_with_arg);
}

static Error* process_num_clusters_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'num-clusters' requires an argument");
    }
    if (sscanf(arg, "%u", &num_clusters) != 1 || num_clusters < 1)
    {
        return JWSC_EARG("Option 'num-clusters' must be > 0");
    }
    return NULL;
}

static Error* process_cluster_type_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'cluster-type' requires an argument");
    }
    if (strncmp(arg, "kmeans", 6) == 0)
    {
        cluster_type = HAPLO_CLUSTER_KMEANS;
    }
    else if (strncmp(arg, "gmm", 3) == 0)
    {
        cluster_type = HAPLO_CLUSTER_GMM;
    }
    else if (strncmp(arg, "mmm", 3) == 0)
    {
        cluster_type = HAPLO_CLUSTER_MMM;
    }
    else
    {
        return JWSC_EARG("Option 'cluster-type' must be one of {kmeans,gmm,mmm}");
    }
    return NULL;
}

Error* process_means_out_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'means-out' requires an argument");
    }
    means_out_fname = arg;
    return NULL;
}

Error* process_weights_out_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'weights-out' requires an argument");
    }
    weights_out_fname = arg;
    return NULL;
}

Error* process_responses_out_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'responses-out' requires an argument");
    }
    responses_out_fname = arg;
    return NULL;
}

Error* process_members_out_opt(Option_arg arg)
{
    if (arg == NULL)
    {
        return JWSC_EARG("Option 'members-out' requires an argument");
    }
    members_out_fname = arg;
    return NULL;
}

static void init_cluster_options()
{
    uint32_t i;

    char s_name;
    const char* l_name;
    const char* desc;

    Error* (*farg)(const char*);

    init_options(opts_no_arg, opts_with_arg);

    opts.label_col = LABEL_COL;

    i = NUM_SHARED_OPTS_WITH_ARG;
    l_name = "aux-input";
    s_name = 0;
    desc   = "Auxiliary input file. Provides marker information used for imputing missing values.";
    farg   = process_aux_input_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "aux-id-cols";
    s_name = 0;
    desc   = "Auxiliary input id-cols. Comma separated ordered list of columns to use for sample identification. Prefixes the output of each sample. Count begins with 1 at the first column of the file. Set to zero to ignore the id column";
    farg   = process_aux_id_cols_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "aux-label-col";
    s_name = 0;
    desc   = "Auxiliary input label-col. Column containing the haplo group labels. Count begins with 1 at the first column of the file. Set to zero to ignore the label column.";
    farg   = process_aux_label_col_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "aux-1st-marker-col";
    s_name = 0;
    desc   = "Auxiliary input 1st-marker-col. Column containing the first marker. Use in conjunction with num-markers to specify the markers for reading. All other markers are assumed to follow this one. Count begins with 1.";
    farg   = process_aux_first_marker_col_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "aux-num-markers";
    s_name = 0;
    desc   = "Auxiliary input num-markers. Number of markers to read. Use in conjunction with 1st-marker-col to specify the markers for reading.";
    farg   = process_aux_num_markers_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "aux-marker-cols";
    s_name = 0;
    desc   = "Auxiliary input marker-cols. Comma separated ordered list of markers to use for training. Use instead of 1st-marker-col and num-markers. Count begins with 1 at the first column of the CSV file.";
    farg   = process_aux_marker_cols_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "num-clusters";
    s_name = 'k';
    desc   = "Number of clusters to use in the k-means algorithm.";
    farg   = process_num_clusters_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "cluster-type";
    s_name = 0;
    desc   = "Type of clustering to use. Must be one of {kmeans,gmm,mmm}.";
    farg   = process_cluster_type_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "means-out";
    s_name = 0;
    desc   = "Cluster means output file name. The default is no output.";
    farg   = process_means_out_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "weights-out";
    s_name = 0;
    desc   = "Cluster weights output file name. The default is no output.";
    farg   = process_weights_out_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "responses-out";
    s_name = 0;
    desc   = "Cluster responses output file name. The default is no output.";
    farg   = process_responses_out_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);

    l_name = "members-out";
    s_name = 0;
    desc   = "Cluster members output file name. The default is stdout.";
    farg   = process_members_out_opt;
    init_option_with_arg(&(opts_with_arg[i++]), l_name, s_name, desc, farg);
    assert(i == NUM_OPTS_WITH_ARG);
}

static void write_means(const Matrix_d* means)
{
    FILE*       fp;
    uint32_t    i, j;
    xmlDoc*     xml_doc   = NULL;
    xmlNode*    xml_root  = NULL;
    xmlNode*    xml_node  = NULL;
    xmlNode*    xml_child = NULL;
    char        xml_buf[256] = {0};
    Error*      err;

    if ((err = open_output(&fp, &xml_doc, "haplo-cluster-means-out",
                    "haplo-cluster-means-out.dtd", means_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if (opts.output_format == HAPLO_OUTPUT_XML)
    {
        xml_root = xmlDocGetRootElement(xml_doc);
    }

    for (i = 0; i < means->num_rows; i++)
    {
        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-4d", i+1);
                for (j = 0; j < means->num_cols; j++)
                {
                    fprintf(fp, "  %6.3f", means->elts[ i ][ j ]);
                }
                fprintf(fp, "\n");
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%d", i+1);
                for (j = 0; j < means->num_cols; j++)
                {
                    fprintf(fp, ",%.3f", means->elts[ i ][ j ]);
                }
                fprintf(fp, "\n");
                break;
            case HAPLO_OUTPUT_XML:
                xml_node = XMLNewChild(xml_root, "mean", NULL);
                snprintf(xml_buf, 256, "%d", i+1);
                XMLNewProp(xml_node, "number", xml_buf);
                for (j = 0; j < means->num_cols; j++)
                {
                    snprintf(xml_buf, 256, "%.3f", means->elts[ i ][ j ]);
                    xml_child = XMLNewChild(xml_node, "marker", xml_buf);
                    snprintf(xml_buf, 256, "%d", j+1);
                    XMLNewProp(xml_child, "number", xml_buf);
                }
                break;
        }
    }

    if ((err = close_output(fp, xml_doc, means_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }
}

static void write_mmm_means(const Matblock_d* means)
{
    FILE*       fp;
    uint32_t    i, j, k;
    uint32_t    n;
    xmlDoc*     xml_doc   = NULL;
    xmlNode*    xml_root  = NULL;
    xmlNode*    xml_node  = NULL;
    xmlNode*    xml_child = NULL;
    char        xml_buf[256] = {0};
    Error*      err;

    if ((err = open_output(&fp, &xml_doc, "haplo-cluster-means-out",
                    "haplo-cluster-means-out.dtd", means_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if (opts.output_format == HAPLO_OUTPUT_XML)
    {
        xml_root = xmlDocGetRootElement(xml_doc);
    }

    for (i = 0; i < means->num_mats; i++)
    {
        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-4d\n", i+1);
                for (j = 0; j < means->num_rows; j++)
                {
                    fprintf(fp, "%-4d", j+1);
                    for (k = 0; k < means->num_cols; k++)
                    {
                        fprintf(fp, "  %6.3f", means->elts[ i ][ j ][ k ]);
                    }
                    fprintf(fp, "\n");
                }
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%d\n", i+1);
                for (j = 0; j < means->num_rows; j++)
                {
                    fprintf(fp, "%d", j+1);
                    for (k = 0; k < means->num_cols; k++)
                    {
                        fprintf(fp, ",%.3f", means->elts[ i ][ j ][ k ]);
                    }
                    fprintf(fp, "\n");
                }
                break;
            case HAPLO_OUTPUT_XML:
                xml_node = XMLNewChild(xml_root, "mean", NULL);
                snprintf(xml_buf, 256, "%d", i+1);
                XMLNewProp(xml_node, "number", xml_buf);
                for (j = 0; j < means->num_rows; j++)
                {
                    n = 0;
                    for (k = 0; k < means->num_cols; k++)
                    {
                        n += snprintf(xml_buf+n, 256, "%.3f ", 
                                means->elts[ i ][ j ][ k ]);
                        assert(n < 256);
                    }
                    xml_child = XMLNewChild(xml_node, "marker", xml_buf);
                    snprintf(xml_buf, 256, "%d", j+1);
                    XMLNewProp(xml_child, "number", xml_buf);
                }
                break;
        }
    }

    if ((err = close_output(fp, xml_doc, means_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }
}

static void write_members_header
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    FILE*              fp
)
{
    uint32_t id;

    if (!opts.header_out || opts.output_format == HAPLO_OUTPUT_XML)
        return;

    if (ids)
    {
        for (id = 0; id < ids->num_rows; id++)
        {
            switch (opts.output_format)
            {
                case HAPLO_OUTPUT_TXT:
                    fprintf(fp, "ID %-7d  ", id+1);
                    break;
                case HAPLO_OUTPUT_CSV:
                    fprintf(fp, "ID %d,", id+1);
                    break;
                case HAPLO_OUTPUT_XML:
                    break;
            }
        }
    }

    if (labels)
    {
        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-10s  ", "Label");
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%s,", "Label");
                break;
            case HAPLO_OUTPUT_XML:
                break;
        }
    }

    switch (opts.output_format)
    {
        case HAPLO_OUTPUT_TXT:
            fprintf(fp, "%6s\n", "Member");
            break;
        case HAPLO_OUTPUT_CSV:
            fprintf(fp, "%s\n", "Member");
            break;
        case HAPLO_OUTPUT_XML:
            break;
    }
}

static void write_members
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Vector_i32*  members
)
{
    FILE*       fp;
    uint32_t    i;
    xmlDoc*     xml_doc   = NULL;
    xmlNode*    xml_root  = NULL;
    xmlNode*    xml_node  = NULL;
    char*       xml_buf   = NULL;
    uint32_t    xml_len;
    Error*      err;

    if ((err = open_output(&fp, &xml_doc, "haplo-cluster-members-out",
                    "haplo-cluster-members-out.dtd", members_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    write_members_header(ids, labels, fp);

    if (opts.output_format == HAPLO_OUTPUT_XML)
    {
        xml_root = xmlDocGetRootElement(xml_doc);
    }

    for (i = 0; i < members->num_elts; i++)
    {
        if (opts.output_format == HAPLO_OUTPUT_XML)
        {
            xml_node = XMLNewChild(xml_root, "sample", NULL);
            xml_len = 256;
            xml_buf = malloc(xml_len*sizeof(char));
            snprintf(xml_buf, xml_len, "%d", i+1);
            XMLNewProp(xml_node, "number", xml_buf);
            free(xml_buf);
        }

        write_ids(ids, i, HAPLO_SEP_SUFFIX, fp, xml_node);
        write_label(labels, i, HAPLO_SEP_SUFFIX, fp, xml_node);

        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-6d\n", members->elts[ i ]+1);
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%d\n", members->elts[ i ]+1);
                break;
            case HAPLO_OUTPUT_XML:
                xml_len = 256;
                xml_buf = malloc(xml_len*sizeof(char));
                snprintf(xml_buf, xml_len, "%d", members->elts[ i ]+1);
                XMLNewChild(xml_node, "member", xml_buf);
                free(xml_buf);
                break;
        }
    }

    if ((err = close_output(fp, xml_doc, members_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }
}

static void write_weights(const Vector_d* weights)
{
    FILE*       fp;
    uint32_t    i;
    xmlDoc*     xml_doc   = NULL;
    xmlNode*    xml_root  = NULL;
    xmlNode*    xml_node  = NULL;
    char        xml_buf[256] = {0};
    Error*      err;

    if ((err = open_output(&fp, &xml_doc, "haplo-cluster-weights-out",
                    "haplo-cluster-weights-out.dtd", weights_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if (opts.output_format == HAPLO_OUTPUT_XML)
    {
        xml_root = xmlDocGetRootElement(xml_doc);
    }

    for (i = 0; i < weights->num_elts; i++)
    {
        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-4d  %5.3f\n", i+1, weights->elts[ i ]);
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%d,%.3f\n", i+1, weights->elts[ i ]);
                break;
            case HAPLO_OUTPUT_XML:
                snprintf(xml_buf, 256, "%.3f", weights->elts[ i ]);
                xml_node = XMLNewChild(xml_root, "weight", xml_buf);
                snprintf(xml_buf, 256, "%d", i+1);
                XMLNewProp(xml_node, "number", xml_buf);
                break;
        }
    }

    if ((err = close_output(fp, xml_doc, weights_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }
}

static void write_responses_header
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Matrix_d*    responses,
    FILE*              fp
)
{
    uint32_t i, j;

    if (!opts.header_out || opts.output_format == HAPLO_OUTPUT_XML)
        return;

    if (ids)
    {
        for (i = 0; i < ids->num_rows; i++)
        {
            switch (opts.output_format)
            {
                case HAPLO_OUTPUT_TXT:
                    fprintf(fp, "ID %-7d  ", i+1);
                    break;
                case HAPLO_OUTPUT_CSV:
                    fprintf(fp, "ID %d,", i+1);
                    break;
                case HAPLO_OUTPUT_XML:
                    break;
            }
        }
    }

    if (labels)
    {
        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-10s  ", "Label");
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%s,", "Label");
                break;
            case HAPLO_OUTPUT_XML:
                break;
        }
    }

    switch (opts.output_format)
    {
        case HAPLO_OUTPUT_TXT:
            fprintf(fp, "%8s %2d", "Response", 1);
            for (j = 1; j < responses->num_cols; j++)
            {
                fprintf(fp, "  %8s %2d", "Response", j+1);
            }
            break;
        case HAPLO_OUTPUT_CSV:
            fprintf(fp, "%s %d", "Response", 1);
            for (j = 1; j < responses->num_cols; j++)
            {
                fprintf(fp, ",%s %d", "Response", j+1);
            }
            break;
        case HAPLO_OUTPUT_XML:
            break;
    }
    fprintf(fp, "\n");
}

static void write_responses
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Matrix_d*    responses
)
{
    FILE*       fp;
    uint32_t    i, j;
    xmlDoc*     xml_doc   = NULL;
    xmlNode*    xml_root  = NULL;
    xmlNode*    xml_node  = NULL;
    xmlNode*    xml_child = NULL;
    char*       xml_buf   = NULL;
    uint32_t    xml_len;
    Error*      err;

    if ((err = open_output(&fp, &xml_doc, "haplo-cluster-responses-out",
                    "haplo-cluster-responses-out.dtd", responses_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    write_responses_header(ids, labels, responses, fp);

    if (opts.output_format == HAPLO_OUTPUT_XML)
    {
        xml_root = xmlDocGetRootElement(xml_doc);
    }

    for (i = 0; i < responses->num_rows; i++)
    {
        if (opts.output_format == HAPLO_OUTPUT_XML)
        {
            xml_node = XMLNewChild(xml_root, "sample", NULL);
            xml_len = 256;
            xml_buf = malloc(xml_len*sizeof(char));
            snprintf(xml_buf, xml_len, "%d", i+1);
            XMLNewProp(xml_node, "number", xml_buf);
            free(xml_buf);
        }

        write_ids(ids, i, HAPLO_SEP_SUFFIX, fp, xml_node);
        write_label(labels, i, HAPLO_SEP_SUFFIX, fp, xml_node);

        switch (opts.output_format)
        {
            case HAPLO_OUTPUT_TXT:
                fprintf(fp, "%-11.4f", responses->elts[ i ][ 0 ]);
                for (j = 1; j < responses->num_cols; j++)
                {
                    fprintf(fp, "  %-11.4f", responses->elts[ i ][ j ]);
                }
                fprintf(fp, "\n");
                break;
            case HAPLO_OUTPUT_CSV:
                fprintf(fp, "%.3f", responses->elts[ i ][ 0 ]);
                for (j = 1; j < responses->num_cols; j++)
                {
                    fprintf(fp, ",%.4f", responses->elts[ i ][ j ]);
                }
                fprintf(fp, "\n");
                break;
            case HAPLO_OUTPUT_XML:
                xml_len = 256;
                xml_buf = malloc(xml_len*sizeof(char));
                for (j = 0; j < responses->num_cols; j++)
                {
                    snprintf(xml_buf, xml_len, "%.4f", responses->elts[i][j]);
                    xml_child = XMLNewChild(xml_node, "response", xml_buf);
                    snprintf(xml_buf, xml_len, "%d", j+1);
                    XMLNewProp(xml_child, "number", xml_buf);
                }
                free(xml_buf);
                break;
        }
    }

    if ((err = close_output(fp, xml_doc, responses_out_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }
}

static void write_kmeans_results
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Matrix_d*    means, 
    const Vector_i32*  members
)
{
    write_means(means);
    write_members(ids, labels, members);
}

static void assign_gmm_members
(
    Vector_i32**    members_out,
    const Matrix_d* responses
)
{
    uint32_t i, j, k;

    create_vector_i32(members_out, responses->num_rows);

    for (i = 0; i < responses->num_rows; i++)
    {
        k = 0;
        for (j = 1; j < responses->num_cols; j++)
        {
            if (responses->elts[ i ][ j ] > responses->elts[ i ][ k ])
            {
                k = j;
            }
        }

        (*members_out)->elts[ i ] = k;
    }
}

static void assign_mmm_members
(
    Vector_i32**    members_out,
    const Matrix_d* responses
)
{
    uint32_t i, j, k;

    create_vector_i32(members_out, responses->num_rows);

    for (i = 0; i < responses->num_rows; i++)
    {
        k = 0;
        for (j = 1; j < responses->num_cols; j++)
        {
            if (responses->elts[ i ][ j ] > responses->elts[ i ][ k ])
            {
                k = j;
            }
        }

        (*members_out)->elts[ i ] = k;
    }
}

static void assign_mmm_aux_members
(
    Vector_i32**        members_out,
    const Matblock_d*   means, 
    const Vector_d*     weights,
    const Matblock_u32* markers
)
{
    uint32_t k, K;
    uint32_t n, N;
    uint32_t d, D;
    uint32_t m, M;
    double   p_min;
    double   p, q;

    Matrix_u32* x_n  = NULL;
    Matrix_d*   mu_k = NULL;

    K = weights->num_elts;
    N = markers->num_mats;
    D = markers->num_rows;
    M = markers->num_cols;

    create_vector_i32(members_out, N);

    for (n = 0; n < N; n++)
    {
        p_min = 0;
        copy_matrix_from_matblock_u32(&x_n, markers, MATBLOCK_MAT_MATRIX, n);

        for (k = 0; k < K; k++)
        {
            copy_matrix_from_matblock_d(&mu_k, means, MATBLOCK_MAT_MATRIX, k);

            q = 0.0;
            for (d = 0; d < D; d++)
            {
                for (m = 0; m < M; m++)
                {
                    if (x_n->elts[ d ][ m ])
                    {
                        q += log((mu_k->elts[ d ][ m ] < JWSC_MIN_LOG_ARG) ? JWSC_MIN_LOG_ARG : mu_k->elts[ d ][ m ]);
                    }
                }
            }

            p = log((weights->elts[ k ] < JWSC_MIN_LOG_ARG) ? JWSC_MIN_LOG_ARG : weights->elts[ k ]) + q;
#if defined(JWSC_HAVE_ISNAN) && defined(JWSC_HAVE_ISINF)
            assert(!isnan(p) && !isinf(p));
#endif

            if (p < p_min)
            {
                p_min = p;
                (*members_out)->elts[ n ] = k;
            }
        }
    }

    free_matrix_u32(x_n);
    free_matrix_d(mu_k);
}

static void write_gmm_results
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Matrix_d*    means, 
    const Matblock_d*  covars, 
    const Vector_d*    weights,
    const Matrix_d*    responses,
    const Vector_i32*  members
)
{
    write_means(means);
    write_weights(weights);
    write_responses(ids, labels, responses);
    write_members(ids, labels, members);
}

static void write_mmm_results
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Matblock_d*  means, 
    const Vector_d*    weights,
    const Matrix_d*    responses,
    const Vector_i32*  members
)
{
    write_mmm_means(means);
    write_weights(weights);
    write_responses(ids, labels, responses);
    write_members(ids, labels, members);
}

static void write_mmm_aux_results
(
    const Matblock_u8* ids,
    const Vector_u32*  labels,
    const Vector_i32*  members
)
{
    write_members(ids, labels, members);
}

static void cluster_kmeans
(
    const Matblock_u8* ids, 
    const Vector_u32*  labels,
    const Matrix_i32*  markers
)
{
    uint32_t num_rows, num_cols;
    uint32_t row, col;

    Matrix_d*    markers_d = NULL;
    Matrix_d*    means     = NULL;
    Vector_i32*  members   = NULL;

    num_rows = markers->num_rows;
    num_cols = markers->num_cols;
    create_matrix_d(&markers_d, num_rows, num_cols);

    for (row = 0; row < num_rows; row++)
    {
        for (col = 0; col < num_cols; col++)
        {
            markers_d->elts[ row ][ col ] = (double)markers->elts[ row ][ col ];
        }
    }

    kmeans_d(&means, &members, markers_d, num_clusters, 1);
    write_kmeans_results(ids, labels, means, members);

    free_matrix_d(markers_d);
    free_matrix_d(means);
    free_vector_i32(members);
}

static void cluster_gmm
(
    const Matblock_u8* ids, 
    const Vector_u32*  labels,
    const Matrix_i32*  markers
)
{
    uint32_t num_rows, num_cols;
    uint32_t row, col;

    Matrix_d*    markers_d = NULL;
    Matrix_d*    means     = NULL;
    Matblock_d*  covars    = NULL;
    Vector_d*    weights   = NULL;
    Matrix_d*    responses = NULL;
    Vector_i32*  members   = NULL;

    num_rows = markers->num_rows;
    num_cols = markers->num_cols;
    create_matrix_d(&markers_d, num_rows, num_cols);

    for (row = 0; row < num_rows; row++)
    {
        for (col = 0; col < num_cols; col++)
        {
            markers_d->elts[ row ][ col ] = (double)markers->elts[ row ][ col ];
        }
    }

    train_gmm_d(&means, &covars, &weights, &responses, markers_d, num_clusters);
    assign_gmm_members(&members, responses);
    write_gmm_results(ids, labels, means, covars, weights, responses, members);

    free_matrix_d(markers_d);
    free_matrix_d(means);
    free_matblock_d(covars);
    free_vector_d(weights);
    free_matrix_d(responses);
    free_vector_i32(members);
}

static void cluster_mmm
(
    const Matblock_u8* ids, 
    const Vector_u32*  labels,
    const Matrix_i32*  markers,
    const Matblock_u8* aux_ids, 
    const Vector_u32*  aux_labels,
    const Matrix_i32*  aux_markers
)
{
    uint32_t s, num_samples;
    uint32_t m, num_markers;
    uint32_t num_marker_vals;

    Matblock_u32* markers_d = NULL;
    Matblock_d*   means     = NULL;
    Vector_d*     weights   = NULL;
    Matrix_d*     responses = NULL;
    Vector_i32*   members   = NULL;

    num_marker_vals = 0; 
    num_samples = markers->num_rows;
    num_markers = markers->num_cols;

    for (s = 0; s < num_samples; s++)
    {
        for (m = 0; m < num_markers; m++)
        {
            if (markers->elts[ s ][ m ] >= num_marker_vals)
            {
                num_marker_vals = markers->elts[ s ][ m ]+1;
            }
        }
    }
    assert(num_marker_vals);

    create_zero_matblock_u32(&markers_d, num_samples, num_markers,
            num_marker_vals);
    for (s = 0; s < num_samples; s++)
    {
        for (m = 0; m < num_markers; m++)
        {
            markers_d->elts[ s ][ m ][ markers->elts[s][m] ] = 1;
        }
    }

    // DEBUG
    for (s = 0; s < num_samples; s++)
    {
        for (m = 0; m < num_markers; m++)
        {
            uint32_t v, sum = 0;
            for (v = 0; v < num_marker_vals; v++)
            {
                if (markers_d->elts[ s ][ m ][ v ])
                {
                    sum++;
                }
            }
            assert(sum == 1);
        }
    }

    train_mmm_d(&means, &weights, &responses, markers_d, num_clusters);
    assign_mmm_members(&members, responses);
    write_mmm_results(ids, labels, means, weights, responses, members);

    if (aux_markers)
    {
        num_samples = aux_markers->num_rows;
        assert(num_markers == aux_markers->num_cols);

        create_zero_matblock_u32(&markers_d, num_samples, num_markers,
                num_marker_vals);
        for (s = 0; s < num_samples; s++)
        {
            for (m = 0; m < num_markers; m++)
            {
                markers_d->elts[ s ][ m ][ aux_markers->elts[s][m] ] = 1;
            }
        }

        assign_mmm_aux_members(&members, means, weights, markers_d);
        write_mmm_aux_results(aux_ids, aux_labels, members);
    }

    free_matblock_u32(markers_d);
    free_matblock_d(means);
    free_vector_d(weights);
    free_matrix_d(responses);
    free_vector_i32(members);
}

int main(int argc, const char** argv)
{
    int         argi;
    const char* data_fname = "/dev/stdin";
    Error*      err;

    Matblock_u8* ids       = NULL;
    Vector_u32*  labels    = NULL;
    Matrix_i32*  markers   = NULL;

    Matblock_u8* aux_ids       = NULL;
    Vector_u32*  aux_labels    = NULL;
    Matrix_i32*  aux_markers   = NULL;

    init_cluster_options();

    if ((err = process_options(argc, argv, &argi, NUM_OPTS_NO_ARG, opts_no_arg,
                    NUM_OPTS_WITH_ARG, opts_with_arg)) != NULL)
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if ((argc - argi) == 1)
    {
        data_fname = argv[ argi ];
    }

    if ((err = read_haplo_groups(opts.labels_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if ((err = read_input(&ids, &labels, &markers, data_fname)))
    {
        print_error_msg_exit("haplo-cluster", err->msg);
    }

    if (opts.aux_input_fname)
    {
        if ((err = read_aux_input(&aux_ids, &aux_labels, &aux_markers,
                        opts.aux_input_fname)))
        {
            print_error_msg_exit("haplo-impute", err->msg);
        }
    }

    switch (cluster_type)
    {
        case HAPLO_CLUSTER_KMEANS:
            cluster_kmeans(ids, labels, markers);
            break;
        case HAPLO_CLUSTER_GMM:
            cluster_gmm(ids, labels, markers);
            break;
        case HAPLO_CLUSTER_MMM:
            cluster_mmm(ids, labels, markers, aux_ids, aux_labels, aux_markers);
            break;
    }

    free_matblock_u8(ids);
    free_vector_u32(labels);
    free_matrix_i32(markers);

    if (get_num_unhandled_errors() > 0)
    {
        print_error_msg_exit("haplo-cluster", "Unhandled errors");
    }

    return EXIT_SUCCESS;
}