<commit_before>/*
Copyright 2010, HudsonAlpha Institute for Biotechnology

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.

bam2fastq.cpp

Written by Phillip Dexheimer
*/

#include "sam.h"
#include <getopt.h>
#include <map>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <algorithm>
#include <cctype>

using namespace std;

const char version[] = "1.1.0";
const char shortopts[] = "o:vhfqs";

int save_aligned = 1;
int save_unaligned = 1;
int save_filtered = 1;
int overwrite_files = 0;
int print_msgs = 1;
int strict = 0;

static struct option longopts[] = {
    { "help",         no_argument,       NULL,           'h' },
    { "version",      no_argument,       NULL,           'v' },
    { "output",       required_argument, NULL,           'o' },
    { "force",        no_argument,       NULL,           'f' },
    { "quiet",        no_argument,       NULL,           'q' },
    { "strict",       no_argument,       NULL,           's' },
    { "overwrite",    no_argument,       &overwrite_files,0  },
    { "aligned",      no_argument,       &save_aligned,   1  },
    { "no-aligned",   no_argument,       &save_aligned,   0  },
    { "unaligned",    no_argument,       &save_unaligned, 1  },
    { "no-unaligned", no_argument,       &save_unaligned, 0  },
    { "filtered",     no_argument,       &save_filtered,  1  },
    { "no-filtered",  no_argument,       &save_filtered,  0  },
    { NULL,           0,                 NULL,            0  }
};

void usage(int error=1) {
    cerr << "bam2fastq v" << version << " - extract sequences from a BAM file" << endl
         << endl
         << "Usage: bam2fastq [options] <bam file>" << endl
         << endl
         << "Options:" << endl
         << "  -o FILENAME, --output FILENAME" << endl
         << "       Specifies the name of the FASTQ file(s) that will be generated.  May" << endl
         << "       contain the special characters % (replaced with the lane number) and" << endl
         << "       # (replaced with _1 or _2 to distinguish PE reads, removed for SE reads)." << endl
         << "       [Default: s_%#_sequence.txt]" << endl
         << "  -f, --force, --overwrite" << endl
         << "       Create output files specified with --output, overwriting existing" << endl
         << "       files if necessary [Default: exit program rather than overwrite files]" << endl
         << "  --aligned" << endl
         << "  --no-aligned" << endl
         << "       Reads in the BAM that are aligned will (will not) be extracted." << endl
         << "       [Default: extract aligned reads]" << endl
         << "  --unaligned" << endl
         << "  --no-unaligned" << endl
         << "       Reads in the BAM that are not aligned will (will not) be extracted." << endl
         << "       [Default: extract unaligned reads]" << endl
         << "  --filtered" << endl
         << "  --no-filtered" << endl
         << "       Reads that are marked as failing QC checks will (will not) be extracted." << endl
         << "       [Default: extract filtered reads]" << endl
         << "  -q, --quiet" << endl
         << "       Suppress informational messages [Default: print messages]" << endl
         << "  -s, --strict" << endl
         << "       Keep bam2fastq's processing to a minimum, assuming that the BAM strictly"
         << "       meets specifications. [Default: allow some errors in the BAM]" << endl
         << endl;
    exit(error);
}

map<int, char> bases;

const string get_pair_name(const bam1_t *b) {
    return string(bam1_qname(b));
}

//Returns 0 for read 1 and 1 for read 2.
//Counterintuitive, but good for lookups
const int get_read_idx(const bam1_t *b) {
    return !(b->core.flag & BAM_FREAD1);
}

const string get_read_name(const bam1_t *b) {
    static const char suffix[2][3] = { "/1", "/2" };
    string name = get_pair_name(b);
    if (b->core.flag & BAM_FPAIRED)
        name += suffix[get_read_idx(b)];
    return name;
}

int get_lane_id(const bam1_t *b) {
    string name(get_pair_name(b));
    size_t start = name.find(":");
    if (start == string::npos)
        return 0;
    start++;
    size_t stop = name.find(":", start);
    if (stop == string::npos || stop == start)
        return 0;
    istringstream ist(string(name.begin()+start, name.begin()+stop));
    int lane;
    ist >> lane;
    return lane;
}

const string get_sequence(const bam1_t *b) {
    const uint8_t *seq = bam1_seq(b);
    size_t len = b->core.l_qseq;
    string sequence("");
    sequence.reserve(len);
    uint8_t offset = (b->core.flag & BAM_FREVERSE) ? 16 : 0;
    for (size_t i=0; i<len; i++) {
        sequence += bases[bam1_seqi(seq, i) + offset];
    }
    if (offset)
        sequence = string(sequence.rbegin(), sequence.rend());
    return sequence;
}

const string get_qualities(const bam1_t *b) {
    const uint8_t *qual = bam1_qual(b);
    size_t len = b->core.l_qseq;
    string quality("");
    quality.reserve(len);
    for (size_t i=0; i<len; i++) {
        quality += char_traits<char>::to_char_type(char_traits<char>::to_int_type(qual[i])+33);
    }
    if (b->core.flag & BAM_FREVERSE)
        quality = string(quality.rbegin(), quality.rend());
    return quality;
}

void mangle(string &name) {
    if (name.length() < 3)
        return;
    if (isdigit(name[name.length()-1]) && !isdigit(name[name.length()-2]))
        name.erase(name.length()-2);
}

vector<ofstream *> initialize_output(const string &out_template, int lane, bool isPaired) {
    vector<ofstream *> files;
    string output(out_template);

    //First, replace % in the filename with the lane number
    size_t laneMarker = output.find('%');
    if (laneMarker != string::npos) {
        if (lane == 0) {
            cerr << "The lane could not be determined from the reads.  Specify output files" << endl
                 << "(using --output) that do not include the lane number (%)" << endl;
            return files;
        }
        ostringstream laneStr;
        laneStr << lane;
        output.replace(laneMarker, 1, laneStr.str());
    }

    //Replace # with read number and open ofstreams
    size_t readMarker = output.find('#');
    if (isPaired) {
        //Replace
        if (readMarker == string::npos) {
            cerr << "The sequences in the BAM file are marked as Paired, but a single output" << endl
                 << "file is specified.  Ensure that the output filename (--output) includes" << endl
                 << "a # symbol to be replaced with the read number" << endl;
            return files;
        }
        string file1(output);
        file1.replace(readMarker, 1, "_1");
        string file2(output);
        file2.replace(readMarker, 1, "_2");
        if (print_msgs)
            cerr << "This looks like paired data from lane " << lane << "." << endl
                 << "Output will be in " << file1 << " and " << file2 << endl;

        //If we're not going to overwrite, check to see if the files exist
        if (!overwrite_files) {
            ifstream test;
            test.open(file1.c_str());
            if (test.is_open()) {
                cerr << "ERROR: " << file1 << " already exists.  Specify --force to overwrite" << endl;
                return files;
            }
            //test is not open, so don't try closing it
            test.open(file2.c_str());
            if (test.is_open()) {
                cerr << "ERROR: " << file2 << " already exists.  Specify --force to overwrite" << endl;
                return files;
            }
        }
        files.push_back(new ofstream(file1.c_str(), ios_base::out));
        files.push_back(new ofstream(file2.c_str(), ios_base::out));
    } else {
        if (readMarker != string::npos)
            output.replace(readMarker, 1, "");
        if (print_msgs)
            cerr << "This looks like single-end data from lane " << lane << "." << endl
                 << "Output will be in " << output << endl;
        if (!overwrite_files) {
            ifstream test;
            test.open(output.c_str());
            if (test.is_open()) {
                cerr << "ERROR: " << output << " already exists.  Specify --force to overwrite" << endl;
                return files;
            }
        }
        files.push_back(new ofstream(output.c_str(), ios_base::out));
    }
    return files;
}

//Effective STL, Item 7
//Except, of course, that I'm not using smart pointers
struct DeleteObject {
    template<typename T>
    void operator()(const T *ptr) const {
        delete ptr;
    }
};

void parse_bamfile(const char *bam_filename, const string &output_template) {
    samfile_t *sam = samopen(bam_filename, "rb", NULL);
    if (sam == NULL) {
        cerr << "Could not open " << bam_filename << endl;
        return;
    }
    bam1_t *read = bam_init1();
    size_t exported = 0;
    size_t all_seen = 0;
    bam_read1(sam->x.bam, read);
    bool isPaired = (read->core.flag & BAM_FPAIRED);
    int lane = get_lane_id(read);
    vector<ofstream *> output = initialize_output(output_template, lane, isPaired);
    if (output.empty())
        return;
    map<string, string> unPaired;
    map<string, string>::iterator position;
    do {
        all_seen++;
        if (!save_aligned && !(read->core.flag & BAM_FUNMAP))
            continue;
        if (!save_unaligned && (read->core.flag & BAM_FUNMAP))
            continue;
        if (!save_filtered && (read->core.flag & BAM_FQCFAIL))
            continue;
        exported++;
        ostringstream ostr;
        ostr << "@" << get_read_name(read) << endl
             << get_sequence(read) << endl
             << "+" << endl
             << get_qualities(read) << endl;
        if (isPaired) {
            //Paired-end is complicated, because both members of the pair
            //have to be output at the same position of the two files
            string pairName(get_pair_name(read));
            if (!strict)
                mangle(pairName);
            position = unPaired.find(pairName);
            if (position == unPaired.end()) {
                //I haven't seen the other member of this pair, so just save it
                unPaired[pairName] = ostr.str();
            } else {
                //Aha!  This will be the second of the two.  Dump them both,
                //then clean up
                int r_idx = get_read_idx(read);
                *output[r_idx] << ostr.str();
                r_idx = !r_idx;
                *output[r_idx] << position->second;
                unPaired.erase(position);
            }
        } else {
            //Single-end?  Just dump it
            *output[0] << ostr.str();
        }
    } while (bam_read1(sam->x.bam, read) > 0);
    //The documentation for bam_read1 says that it returns the number of
    //bytes read - which is true, unless it doesn't read any.  It returns
    //-1 for normal EOF and -2 for unexpected EOF.  So don't just wait for
    //it to return 0...
    for_each(output.begin(), output.end(), DeleteObject());
    bam_destroy1(read);
    samclose(sam);
    if (print_msgs) {
        cerr << all_seen << " sequences in the BAM file" << endl;
        cerr << exported << " sequences exported" << endl;
    }
    if (!unPaired.empty()) {
        cerr << "WARNING: " << unPaired.size() << " reads could not be matched "
             << "to a mate and were not exported" << endl;
    }
}

int main (int argc, char *argv[]) {
    bases[1] = 'A';
    bases[2] = 'C';
    bases[4] = 'G';
    bases[8] = 'T';
    bases[15] = 'N';
    //Complements (original index + 16)
    bases[17] = 'T';
    bases[18] = 'G';
    bases[20] = 'C';
    bases[24] = 'A';
    bases[31] = 'N';
    string output_template("s_%#_sequence.txt");
    int ch;
    while ((ch = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1)
        switch (ch) {
            case 'v':
                cerr << "bam2fastq v" << version << endl;
                exit(0);
            case 'h' :
                usage(0);
            case 'o' :
                output_template = string(optarg);
                break;
            case 'f' :
                overwrite_files = 1;
                break;
            case 'q' :
                print_msgs = 0;
                break;
            case 's' :
                strict = 1;
                break;
            case '?' : //Unrecognized option
                usage(2);
            //The remaining options will set the appropriate variable themselves
        };
    argc -= optind;
    argv += optind;
    if (argc == 0)
        usage(1);
    parse_bamfile(argv[0], output_template);
    return 0;
}<commit_msg>write unpaired reads to a separate file<commit_after>/*
Copyright 2010, HudsonAlpha Institute for Biotechnology

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.

bam2fastq.cpp

Written by Phillip Dexheimer
*/

#include "sam.h"
#include <getopt.h>
#include <map>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <algorithm>
#include <cctype>

using namespace std;

const char version[] = "1.1.0";
const char shortopts[] = "o:vhfqs";

int save_aligned = 1;
int save_unaligned = 1;
int save_filtered = 1;
int overwrite_files = 0;
int print_msgs = 1;
int strict = 0;

static struct option longopts[] = {
    { "help",         no_argument,       NULL,           'h' },
    { "version",      no_argument,       NULL,           'v' },
    { "output",       required_argument, NULL,           'o' },
    { "force",        no_argument,       NULL,           'f' },
    { "quiet",        no_argument,       NULL,           'q' },
    { "strict",       no_argument,       NULL,           's' },
    { "overwrite",    no_argument,       &overwrite_files,0  },
    { "aligned",      no_argument,       &save_aligned,   1  },
    { "no-aligned",   no_argument,       &save_aligned,   0  },
    { "unaligned",    no_argument,       &save_unaligned, 1  },
    { "no-unaligned", no_argument,       &save_unaligned, 0  },
    { "filtered",     no_argument,       &save_filtered,  1  },
    { "no-filtered",  no_argument,       &save_filtered,  0  },
    { NULL,           0,                 NULL,            0  }
};

void usage(int error=1) {
    cerr << "bam2fastq v" << version << " - extract sequences from a BAM file" << endl
         << endl
         << "Usage: bam2fastq [options] <bam file>" << endl
         << endl
         << "Options:" << endl
         << "  -o FILENAME, --output FILENAME" << endl
         << "       Specifies the name of the FASTQ file(s) that will be generated.  May" << endl
         << "       contain the special characters % (replaced with the lane number) and" << endl
         << "       # (replaced with _1 or _2 to distinguish PE reads, removed for SE reads)." << endl
         << "       [Default: s_%#_sequence.txt]" << endl
         << "  -f, --force, --overwrite" << endl
         << "       Create output files specified with --output, overwriting existing" << endl
         << "       files if necessary [Default: exit program rather than overwrite files]" << endl
         << "  --aligned" << endl
         << "  --no-aligned" << endl
         << "       Reads in the BAM that are aligned will (will not) be extracted." << endl
         << "       [Default: extract aligned reads]" << endl
         << "  --unaligned" << endl
         << "  --no-unaligned" << endl
         << "       Reads in the BAM that are not aligned will (will not) be extracted." << endl
         << "       [Default: extract unaligned reads]" << endl
         << "  --filtered" << endl
         << "  --no-filtered" << endl
         << "       Reads that are marked as failing QC checks will (will not) be extracted." << endl
         << "       [Default: extract filtered reads]" << endl
         << "  -q, --quiet" << endl
         << "       Suppress informational messages [Default: print messages]" << endl
         << "  -s, --strict" << endl
         << "       Keep bam2fastq's processing to a minimum, assuming that the BAM strictly"
         << "       meets specifications. [Default: allow some errors in the BAM]" << endl
         << endl;
    exit(error);
}

map<int, char> bases;

const string get_pair_name(const bam1_t *b) {
    return string(bam1_qname(b));
}

//Returns 0 for read 1 and 1 for read 2.
//Counterintuitive, but good for lookups
const int get_read_idx(const bam1_t *b) {
    return !(b->core.flag & BAM_FREAD1);
}

const string get_read_name(const bam1_t *b) {
    static const char suffix[2][3] = { "/1", "/2" };
    string name = get_pair_name(b);
    if (b->core.flag & BAM_FPAIRED)
        name += suffix[get_read_idx(b)];
    return name;
}

int get_lane_id(const bam1_t *b) {
    string name(get_pair_name(b));
    size_t start = name.find(":");
    if (start == string::npos)
        return 0;
    start++;
    size_t stop = name.find(":", start);
    if (stop == string::npos || stop == start)
        return 0;
    istringstream ist(string(name.begin()+start, name.begin()+stop));
    int lane;
    ist >> lane;
    return lane;
}

const string get_sequence(const bam1_t *b) {
    const uint8_t *seq = bam1_seq(b);
    size_t len = b->core.l_qseq;
    string sequence("");
    sequence.reserve(len);
    uint8_t offset = (b->core.flag & BAM_FREVERSE) ? 16 : 0;
    for (size_t i=0; i<len; i++) {
        sequence += bases[bam1_seqi(seq, i) + offset];
    }
    if (offset)
        sequence = string(sequence.rbegin(), sequence.rend());
    return sequence;
}

const string get_qualities(const bam1_t *b) {
    const uint8_t *qual = bam1_qual(b);
    size_t len = b->core.l_qseq;
    string quality("");
    quality.reserve(len);
    for (size_t i=0; i<len; i++) {
        quality += char_traits<char>::to_char_type(char_traits<char>::to_int_type(qual[i])+33);
    }
    if (b->core.flag & BAM_FREVERSE)
        quality = string(quality.rbegin(), quality.rend());
    return quality;
}

void mangle(string &name) {
    if (name.length() < 3)
        return;
    if (isdigit(name[name.length()-1]) && !isdigit(name[name.length()-2]))
        name.erase(name.length()-2);
}

vector<ofstream *> initialize_output(const string &out_template, int lane) {
    vector<ofstream *> files;
    string output(out_template);

    //First, replace % in the filename with the lane number
    size_t laneMarker = output.find('%');
    if (laneMarker != string::npos) {
        if (lane == 0) {
            cerr << "The lane could not be determined from the reads.  Specify output files" << endl
                 << "(using --output) that do not include the lane number (%)" << endl;
            return files;
        }
        ostringstream laneStr;
        laneStr << lane;
        output.replace(laneMarker, 1, laneStr.str());
    }

    //Replace # with read number and open ofstreams
    size_t readMarker = output.find('#');

    //Replace
    if (readMarker == string::npos) {
        cerr << "The sequences in the BAM file are marked as Paired, but a single output" << endl
             << "file is specified.  Ensure that the output filename (--output) includes" << endl
             << "a # symbol to be replaced with the read number" << endl;
        return files;
    }
    string file1(output);
    file1.replace(readMarker, 1, "_1");
    string file2(output);
    file2.replace(readMarker, 1, "_2");
    string file3(output);
    file3.replace(readMarker, 1, "_M");
    if (print_msgs)
        cerr << "This looks like paired data from lane " << lane << "." << endl
             << "Output will be in " << file1 << " and " << file2 << endl
             << "Single-end reads will be in " << file3 << endl;

    //If we're not going to overwrite, check to see if the files exist
    if (!overwrite_files) {
        ifstream test;
        test.open(file1.c_str());
        if (test.is_open()) {
            cerr << "ERROR: " << file1 << " already exists.  Specify --force to overwrite" << endl;
            return files;
        }
        //test is not open, so don't try closing it
        test.open(file2.c_str());
        if (test.is_open()) {
            cerr << "ERROR: " << file2 << " already exists.  Specify --force to overwrite" << endl;
            return files;
        }
    }
    files.push_back(new ofstream(file1.c_str(), ios_base::out));
    files.push_back(new ofstream(file2.c_str(), ios_base::out));
    files.push_back(new ofstream(file3.c_str(), ios_base::out));

    return files;
}

//Effective STL, Item 7
//Except, of course, that I'm not using smart pointers
struct DeleteObject {
    template<typename T>
    void operator()(const T *ptr) const {
        delete ptr;
    }
};

void parse_bamfile(const char *bam_filename, const string &output_template) {
    samfile_t *sam = samopen(bam_filename, "rb", NULL);
    if (sam == NULL) {
        cerr << "Could not open " << bam_filename << endl;
        return;
    }
    bam1_t *read = bam_init1();
    size_t exported = 0;
    size_t all_seen = 0;

    bam_read1(sam->x.bam, read);
    int lane = get_lane_id(read);
    vector<ofstream *> output = initialize_output(output_template, lane);
    if (output.empty())
        return;

    map<string, string> unPaired;
    map<string, string>::iterator position;
    do {
        all_seen++;
        if (!save_aligned && !(read->core.flag & BAM_FUNMAP))
            continue;
        if (!save_unaligned && (read->core.flag & BAM_FUNMAP))
            continue;
        if (!save_filtered && (read->core.flag & BAM_FQCFAIL))
            continue;

        exported++;
        ostringstream ostr;
        ostr << "@" << get_read_name(read) << endl
             << get_sequence(read) << endl
             << "+" << endl
             << get_qualities(read) << endl;

        //Paired-end is complicated, because both members of the pair
        //have to be output at the same position of the two files

        // Is this an unpaired read in a BAM with pairs? write to the _M file
        if( !(read->core.flag & BAM_FPAIRED) ) {
            *output[2] << ostr.str();
        } else {

            // Search for the pair in the map
            string pairName(get_pair_name(read));
            if (!strict)
                mangle(pairName);
            position = unPaired.find(pairName);
            if (position == unPaired.end()) {
                //I haven't seen the other member of this pair, so just save it
                unPaired[pairName] = ostr.str();
            } else {
                //Aha!  This will be the second of the two.  Dump them both,
                //then clean up
                int r_idx = get_read_idx(read);
                *output[r_idx] << ostr.str();
                r_idx = !r_idx;
                *output[r_idx] << position->second;
                unPaired.erase(position);
            }
        }

    } while (bam_read1(sam->x.bam, read) > 0);
    //The documentation for bam_read1 says that it returns the number of
    //bytes read - which is true, unless it doesn't read any.  It returns
    //-1 for normal EOF and -2 for unexpected EOF.  So don't just wait for
    //it to return 0...
    for_each(output.begin(), output.end(), DeleteObject());
    bam_destroy1(read);
    samclose(sam);
    if (print_msgs) {
        cerr << all_seen << " sequences in the BAM file" << endl;
        cerr << exported << " sequences exported" << endl;
    }
    if (!unPaired.empty()) {
        cerr << "WARNING: " << unPaired.size() << " reads could not be matched "
             << "to a mate and were not exported" << endl;
    }
}

int main (int argc, char *argv[]) {
    bases[1] = 'A';
    bases[2] = 'C';
    bases[4] = 'G';
    bases[8] = 'T';
    bases[15] = 'N';
    //Complements (original index + 16)
    bases[17] = 'T';
    bases[18] = 'G';
    bases[20] = 'C';
    bases[24] = 'A';
    bases[31] = 'N';
    string output_template("s_%#_sequence.txt");
    int ch;
    while ((ch = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1)
        switch (ch) {
            case 'v':
                cerr << "bam2fastq v" << version << endl;
                exit(0);
            case 'h' :
                usage(0);
            case 'o' :
                output_template = string(optarg);
                break;
            case 'f' :
                overwrite_files = 1;
                break;
            case 'q' :
                print_msgs = 0;
                break;
            case 's' :
                strict = 1;
                break;
            case '?' : //Unrecognized option
                usage(2);
            //The remaining options will set the appropriate variable themselves
        };
    argc -= optind;
    argv += optind;
    if (argc == 0)
        usage(1);
    parse_bamfile(argv[0], output_template);
    return 0;
}
<|endoftext|>


<commit_before>//@author A0094446X
#include "stdafx.h"
#include <QApplication>
#include <QList>
#include "task_panel_manager.h"
namespace You {
namespace GUI {

using Date = boost::gregorian::date;
const QString MainWindow::TaskPanelManager::TASK_COLUMN_1 = "Hidden ID Column";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_2 = "Index";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_3 = "Description";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_4 = "Deadline";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_5 = "Priority";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_6 = "Dependencies";

MainWindow::TaskPanelManager::TaskPanelManager(MainWindow* const parentGUI)
: BaseManager(parentGUI) {
}

MainWindow::TaskPanelManager::~TaskPanelManager() {
}

void MainWindow::TaskPanelManager::setup() {
	QStringList columnHeaders({
		TASK_COLUMN_1,
		TASK_COLUMN_2,
		TASK_COLUMN_3,
		TASK_COLUMN_4,
		TASK_COLUMN_5,
		TASK_COLUMN_6
	});
	QTreeWidget* taskTreePanel = parentGUI->ui.taskTreePanel;
	connect(taskTreePanel, SIGNAL(itemSelectionChanged()),
		parentGUI, SLOT(taskSelected()));
	taskTreePanel->setColumnCount(columnHeaders.size());
	taskTreePanel->setHeaderItem(createItem(columnHeaders).release());
	// TODO(angathorion): remove magic constants.
	QHeaderView* header = taskTreePanel->header();
	header->setStretchLastSection(true);

	for (int i = 1; i < columnHeaders.size(); ++i) {
		if (i == 2) {
			continue;
		}
		header->resizeSection(i, header->defaultSectionSize());
	}

	taskTreePanel->header()->setMinimumSectionSize(75);
	taskTreePanel->setColumnHidden(0, true);
	taskTreePanel->setColumnHidden(5, true);
}

void MainWindow::TaskPanelManager::addTask(const Task& task) {
	std::unique_ptr<QTreeWidgetItem> item(createItem(task));
	parentGUI->ui.taskTreePanel->addTopLevelItem(item.release());
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::addSubtask(QTreeWidgetItem* parent,
	const QStringList& rowStrings) {
	std::unique_ptr<QTreeWidgetItem> item(createItem(rowStrings));
	parent->addChild(item.release());
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::editTask(const Task& task) {
	QList<QTreeWidgetItem*> items = findItems(task.getID());
	assert(items.length() == 1);
	QTreeWidgetItem item = *items.at(0);
	QStringList wstr = taskToStrVec(task);
	*items.at(0) = *createItem(wstr);
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::deleteTask(Task::ID taskID) {
	QList<QTreeWidgetItem*> items = findItems(taskID);
	assert(items.length() == 1);
	deleteTask(items.at(0));
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::deleteTask(QTreeWidgetItem* task) {
	delete task;
}

std::unique_ptr<QTreeWidgetItem> MainWindow::TaskPanelManager::createItem(
	const Task& task) {
	return createItem(taskToStrVec(task));
}

std::unique_ptr<QTreeWidgetItem> MainWindow::TaskPanelManager::createItem(
	const QStringList& rowStrings) {
	return std::make_unique<QTreeWidgetItem>(rowStrings);
}

QList<QTreeWidgetItem*> MainWindow::TaskPanelManager::findItems(
	You::Controller::Task::ID taskID) const {
	return parentGUI->ui.taskTreePanel->findItems(
		boost::lexical_cast<QString>(taskID), 0);
}

QStringList MainWindow::TaskPanelManager::taskToStrVec(
	const You::Controller::Task& task) {
	QStringList result;

	// Insert id
	result.push_back(boost::lexical_cast<QString>(task.getID()));

	// Insert dummy count
	result.push_back("0");

	// Insert description
	result.push_back(QString::fromStdWString(task.getDescription()));

	// Insert deadline
	if (task.getDeadline() == Task::NEVER) {
		result.push_back(QString("Never"));
	} else {
		result.push_back(boost::lexical_cast<QString>(task.getDeadline()));
		task.getDeadline();
	}

	// Insert priority
	QString priority[] { "High", "Normal" };

	switch (task.getPriority()) {
	case Task::Priority::HIGH:
		result.push_back(priority[0]);
	case Task::Priority::NORMAL:
		result.push_back(priority[1]);
	}

	// Insert dependencies
	std::wstringstream ss;
	if (task.getDependencies().size() != 0) {
		ss << task.getDependencies().at(0);
		for (int i = 1; i < task.getDependencies().size(); i++) {
			ss << ", " << task.getDependencies().at(i);
		}
		result.push_back(QString::fromStdWString(ss.str()));
	} else {
		result.push_back("None");
	}

	return result;
}

void MainWindow::TaskPanelManager::updateRowNumbers() {
	int rowNum = 0;
	for (QTreeWidgetItemIterator it(parentGUI->ui.taskTreePanel); *it; ++it) {
		(*it)->setData(1, Qt::DisplayRole, rowNum++);
	}
}

bool MainWindow::TaskPanelManager::isPastDue(Task::Time deadline) {
	return true;
}

bool MainWindow::TaskPanelManager::isDueAfter(
		Task::Time deadline, int daysLeft) {
	Date by = Date(deadline.date());
	Date today = boost::gregorian::day_clock::universal_day();
	if (by.day_of_year() - today.day_of_year() == daysLeft) {
		return true;
	} else {
		return false;
	}
}

}  // namespace GUI
}  // namespace You
<commit_msg>Switched to match to local_time. Previously did not match.<commit_after>//@author A0094446X
#include "stdafx.h"
#include <QApplication>
#include <QList>
#include "task_panel_manager.h"
namespace You {
namespace GUI {

using Date = boost::gregorian::date;
const QString MainWindow::TaskPanelManager::TASK_COLUMN_1 = "Hidden ID Column";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_2 = "Index";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_3 = "Description";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_4 = "Deadline";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_5 = "Priority";
const QString MainWindow::TaskPanelManager::TASK_COLUMN_6 = "Dependencies";

MainWindow::TaskPanelManager::TaskPanelManager(MainWindow* const parentGUI)
: BaseManager(parentGUI) {
}

MainWindow::TaskPanelManager::~TaskPanelManager() {
}

void MainWindow::TaskPanelManager::setup() {
	QStringList columnHeaders({
		TASK_COLUMN_1,
		TASK_COLUMN_2,
		TASK_COLUMN_3,
		TASK_COLUMN_4,
		TASK_COLUMN_5,
		TASK_COLUMN_6
	});
	QTreeWidget* taskTreePanel = parentGUI->ui.taskTreePanel;
	connect(taskTreePanel, SIGNAL(itemSelectionChanged()),
		parentGUI, SLOT(taskSelected()));
	taskTreePanel->setColumnCount(columnHeaders.size());
	taskTreePanel->setHeaderItem(createItem(columnHeaders).release());
	// TODO(angathorion): remove magic constants.
	QHeaderView* header = taskTreePanel->header();
	header->setStretchLastSection(true);

	for (int i = 1; i < columnHeaders.size(); ++i) {
		if (i == 2) {
			continue;
		}
		header->resizeSection(i, header->defaultSectionSize());
	}

	taskTreePanel->header()->setMinimumSectionSize(75);
	taskTreePanel->setColumnHidden(0, true);
	taskTreePanel->setColumnHidden(5, true);
}

void MainWindow::TaskPanelManager::addTask(const Task& task) {
	std::unique_ptr<QTreeWidgetItem> item(createItem(task));
	parentGUI->ui.taskTreePanel->addTopLevelItem(item.release());
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::addSubtask(QTreeWidgetItem* parent,
	const QStringList& rowStrings) {
	std::unique_ptr<QTreeWidgetItem> item(createItem(rowStrings));
	parent->addChild(item.release());
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::editTask(const Task& task) {
	QList<QTreeWidgetItem*> items = findItems(task.getID());
	assert(items.length() == 1);
	QTreeWidgetItem item = *items.at(0);
	QStringList wstr = taskToStrVec(task);
	*items.at(0) = *createItem(wstr);
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::deleteTask(Task::ID taskID) {
	QList<QTreeWidgetItem*> items = findItems(taskID);
	assert(items.length() == 1);
	deleteTask(items.at(0));
	updateRowNumbers();
}

void MainWindow::TaskPanelManager::deleteTask(QTreeWidgetItem* task) {
	delete task;
}

std::unique_ptr<QTreeWidgetItem> MainWindow::TaskPanelManager::createItem(
	const Task& task) {
	return createItem(taskToStrVec(task));
}

std::unique_ptr<QTreeWidgetItem> MainWindow::TaskPanelManager::createItem(
	const QStringList& rowStrings) {
	return std::make_unique<QTreeWidgetItem>(rowStrings);
}

QList<QTreeWidgetItem*> MainWindow::TaskPanelManager::findItems(
	You::Controller::Task::ID taskID) const {
	return parentGUI->ui.taskTreePanel->findItems(
		boost::lexical_cast<QString>(taskID), 0);
}

QStringList MainWindow::TaskPanelManager::taskToStrVec(
	const You::Controller::Task& task) {
	QStringList result;

	// Insert id
	result.push_back(boost::lexical_cast<QString>(task.getID()));

	// Insert dummy count
	result.push_back("0");

	// Insert description
	result.push_back(QString::fromStdWString(task.getDescription()));

	// Insert deadline
	if (task.getDeadline() == Task::NEVER) {
		result.push_back(QString("Never"));
	} else {
		result.push_back(boost::lexical_cast<QString>(task.getDeadline()));
		task.getDeadline();
	}

	// Insert priority
	QString priority[] { "High", "Normal" };

	switch (task.getPriority()) {
	case Task::Priority::HIGH:
		result.push_back(priority[0]);
	case Task::Priority::NORMAL:
		result.push_back(priority[1]);
	}

	// Insert dependencies
	std::wstringstream ss;
	if (task.getDependencies().size() != 0) {
		ss << task.getDependencies().at(0);
		for (int i = 1; i < task.getDependencies().size(); i++) {
			ss << ", " << task.getDependencies().at(i);
		}
		result.push_back(QString::fromStdWString(ss.str()));
	} else {
		result.push_back("None");
	}

	return result;
}

void MainWindow::TaskPanelManager::updateRowNumbers() {
	int rowNum = 0;
	for (QTreeWidgetItemIterator it(parentGUI->ui.taskTreePanel); *it; ++it) {
		(*it)->setData(1, Qt::DisplayRole, rowNum++);
	}
}

bool MainWindow::TaskPanelManager::isPastDue(Task::Time deadline) {
	return true;
}

bool MainWindow::TaskPanelManager::isDueAfter(
		Task::Time deadline, int daysLeft) {
	Date by = Date(deadline.date());
	Date today = boost::gregorian::day_clock::local_day();
	if (by.day_of_year() - today.day_of_year() == daysLeft) {
		return true;
	} else {
		return false;
	}
}

}  // namespace GUI
}  // namespace You
<|endoftext|>


<commit_before>//
// Copyright (c) 2015 CNRS
// Copyright (c) 2015 Wandercraft, 86 rue de Paris 91400 Orsay, France.
//
// This file is part of Pinocchio
// Pinocchio is free software: you can redistribute it
// and/or modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation, either version
// 3 of the License, or (at your option) any later version.
//
// Pinocchio is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// General Lesser Public License for more details. You should have
// received a copy of the GNU Lesser General Public License along with
// Pinocchio If not, see
// <http://www.gnu.org/licenses/>.

#ifndef __se3_urdf_hpp__
#define __se3_urdf_hpp__

#include <urdf_model/model.h>
#include <urdf_parser/urdf_parser.h>

#include <iostream>
#include <boost/foreach.hpp>
#include "pinocchio/multibody/model.hpp"

#include <exception>

namespace urdf
{
  typedef boost::shared_ptr<ModelInterface> ModelInterfacePtr;
  typedef boost::shared_ptr<const Joint> JointConstPtr;
  typedef boost::shared_ptr<const Link> LinkConstPtr;
  typedef boost::shared_ptr<Link> LinkPtr;
  typedef boost::shared_ptr<const Inertial> InertialConstPtr;
}

namespace se3
{
  namespace urdf
  {
    inline Inertia convertFromUrdf( const ::urdf::Inertial& Y )
    {
      const ::urdf::Vector3 & p = Y.origin.position;
      const ::urdf::Rotation & q = Y.origin.rotation;

      const Eigen::Vector3d com(p.x,p.y,p.z);
      const Eigen::Matrix3d & R = Eigen::Quaterniond(q.w,q.x,q.y,q.z).matrix();

      Eigen::Matrix3d I; I << Y.ixx,Y.ixy,Y.ixz
			   ,  Y.ixy,Y.iyy,Y.iyz
			   ,  Y.ixz,Y.iyz,Y.izz;
      return Inertia(Y.mass,com,R*I*R.transpose());
    }

    inline SE3 convertFromUrdf( const ::urdf::Pose & M )
    {
      const ::urdf::Vector3 & p = M.position;
      const ::urdf::Rotation & q = M.rotation;
      return SE3( Eigen::Quaterniond(q.w,q.x,q.y,q.z).matrix(), Eigen::Vector3d(p.x,p.y,p.z));
    }

    enum AxisCartesian { AXIS_X, AXIS_Y, AXIS_Z, AXIS_UNALIGNED };
    inline AxisCartesian extractCartesianAxis( const ::urdf::Vector3 & axis )
    {
      if( (axis.x==1.0)&&(axis.y==0.0)&&(axis.z==0.0) )
	return AXIS_X;
      else if( (axis.x==0.0)&&(axis.y==1.0)&&(axis.z==0.0) )
	return AXIS_Y;
      else if( (axis.x==0.0)&&(axis.y==0.0)&&(axis.z==1.0) )
	return AXIS_Z;
      else
	return AXIS_UNALIGNED;
    }

inline void parseTree( ::urdf::LinkConstPtr link, Model & model, const SE3 & placementOffset = SE3::Identity()) throw (std::invalid_argument)
{

  ::urdf::JointConstPtr joint = link->parent_joint;
  SE3 nextPlacementOffset = SE3::Identity(); // OffSet of the next link. In case we encounter a fixed joint, we need to propagate the length of its attached body to next joint.

  // std::cout << " *** " << link->name << "    < attached by joint ";
  // if(joint)
  //   std::cout << "#" << link->parent_joint->name << std::endl;
  // else std::cout << "###ROOT" << std::endl;

  // std::cout << "placementOffset: " << placementOffset << std::endl;

  if(joint!=NULL)
  {
    assert(link->getParent()!=NULL);

    if (!link->inertial && joint->type != ::urdf::Joint::FIXED)
    {
      const std::string exception_message (link->name + " - spatial inertia information missing.");
      throw std::invalid_argument(exception_message);
    }

    Model::Index parent = (link->getParent()->parent_joint==NULL) ? (model.existJointName("root_joint") ? model.getJointId("root_joint") : 0) :
                                                                    model.getJointId( link->getParent()->parent_joint->name );
    //std::cout << joint->name << " === " << parent << std::endl;

    const SE3 & jointPlacement = placementOffset*convertFromUrdf(joint->parent_to_joint_origin_transform);

    const Inertia & Y = (link->inertial) ?  convertFromUrdf(*link->inertial) :
    Inertia::Identity();

    bool visual = (link->visual) ? true : false;

    //std::cout << "Parent = " << parent << std::endl;
    //std::cout << "Placement = " << (Matrix4)jointPlacement << std::endl;

    switch(joint->type)
    {
      case ::urdf::Joint::REVOLUTE:
      case ::urdf::Joint::CONTINUOUS: // Revolute with no joint limits
      {

        Eigen::VectorXd maxEffort;
        Eigen::VectorXd velocity;
        Eigen::VectorXd lowerPosition;
        Eigen::VectorXd upperPosition;

        if (joint->limits)
        {
          maxEffort.resize(1);      maxEffort     << joint->limits->effort;
          velocity.resize(1);       velocity      << joint->limits->velocity;
          lowerPosition.resize(1);  lowerPosition << joint->limits->lower;
          upperPosition.resize(1);  upperPosition << joint->limits->upper;
        }

        Eigen::Vector3d jointAxis(Eigen::Vector3d::Zero());
        AxisCartesian axis = extractCartesianAxis(joint->axis);
        switch(axis)
        {
          case AXIS_X:
            model.addBody(  parent, JointModelRX(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Y:
            model.addBody(  parent, JointModelRY(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Z:
            model.addBody(  parent, JointModelRZ(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_UNALIGNED:
            jointAxis= Eigen::Vector3d( joint->axis.x,joint->axis.y,joint->axis.z );
            jointAxis.normalize();
            model.addBody(  parent, JointModelRevoluteUnaligned(jointAxis),
                            jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          default:
            assert( false && "Fatal Error while extracting revolute joint axis");
            break;
        }
        break;
      }
      case ::urdf::Joint::PRISMATIC:
      {
        Eigen::VectorXd maxEffort = Eigen::VectorXd(0.);
        Eigen::VectorXd velocity = Eigen::VectorXd(0.);
        Eigen::VectorXd lowerPosition = Eigen::VectorXd(0.);
        Eigen::VectorXd upperPosition = Eigen::VectorXd(0.);

        if (joint->limits)
        {
          maxEffort.resize(1);      maxEffort     << joint->limits->effort;
          velocity.resize(1);       velocity      << joint->limits->velocity;
          lowerPosition.resize(1);  lowerPosition << joint->limits->lower;
          upperPosition.resize(1);  upperPosition << joint->limits->upper;
        }
        AxisCartesian axis = extractCartesianAxis(joint->axis);
        switch(axis)
        {
          case AXIS_X:
            model.addBody(  parent, JointModelPX(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Y:
            model.addBody(  parent, JointModelPY(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Z:
            model.addBody(  parent, JointModelPZ(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_UNALIGNED:
            std::cerr << "Bad axis = (" << joint->axis.x <<"," << joint->axis.y << "," << joint->axis.z << ")" << std::endl;
            assert(false && "Only X, Y or Z axis are accepted." );
            break;
          default:
            assert( false && "Fatal Error while extracting prismatic joint axis");
            break;
        }
        break;
      }
      case ::urdf::Joint::FIXED:
      {
        // In case of fixed joint, if link has inertial tag:
        //    -add the inertia of the link to his parent in the model
        // Otherwise do nothing.
        // In all cases:
        //    -let all the children become children of parent
        //    -inform the parser of the offset to apply
        //    -add fixed body in model to display it in gepetto-viewer
        if (link->inertial)
        {
          model.mergeFixedBody(parent, jointPlacement, Y); //Modify the parent inertia in the model
        }

        SE3 ptjot_se3 = convertFromUrdf(link->parent_joint->parent_to_joint_origin_transform);

        //transformation of the current placement offset
        nextPlacementOffset = placementOffset*ptjot_se3;

        //add the fixed Body in the model for the viewer
        model.addFixedBody(parent,nextPlacementOffset,link->name,visual);

        BOOST_FOREACH(::urdf::LinkPtr child_link,link->child_links)
        {
          child_link->setParent(link->getParent() );  //skip the fixed generation
        }
        break;
      }
      default:
      {
        std::cerr << "The joint type " << joint->type << " is not supported." << std::endl;
        assert(false && "Only revolute, prismatic and fixed joints are accepted." );
        break;
      }
    }

  }
          else if (link->getParent() != NULL)
          {
            const std::string exception_message (link->name + " - joint information missing.");
            throw std::invalid_argument(exception_message);
          }

  BOOST_FOREACH(::urdf::LinkConstPtr child,link->child_links)
  {
    parseTree(child, model, nextPlacementOffset);
  }
}

    template <typename D>
    void parseTree( ::urdf::LinkConstPtr link, Model & model, const SE3 & placementOffset , const JointModelBase<D> &  root_joint  )
    {
      const Inertia & Y = (link->inertial) ?
      convertFromUrdf(*link->inertial)
      : Inertia::Identity();
      model.addBody( 0, root_joint, placementOffset, Y , "root_joint", link->name, true );
      BOOST_FOREACH(::urdf::LinkConstPtr child,link->child_links)
      {
        parseTree(child, model, SE3::Identity());
      }
    }

    template <typename D>
    Model buildModel( const std::string & filename, const JointModelBase<D> &  root_joint )
    {
      Model model;

      ::urdf::ModelInterfacePtr urdfTree = ::urdf::parseURDFFile (filename);
      if (urdfTree)
        parseTree(urdfTree->getRoot(), model, SE3::Identity(), root_joint);
      else
      {
        std::cerr << "The URDF tree seems to be empty" << std::endl; // In future, raise exception instead
      }

      return model;
    }

    inline Model buildModel( const std::string & filename)
    {
      Model model;

      ::urdf::ModelInterfacePtr urdfTree = ::urdf::parseURDFFile (filename);
      if (urdfTree)
        parseTree(urdfTree->getRoot(), model, SE3::Identity());
      else
      {
        std::cerr << "The URDF tree seems to be empty" << std::endl; // In future, raise exception instead
      }

      return model;
    }

  } // namespace urdf
} // namespace se3

#endif // ifndef __se3_urdf_hpp__

<commit_msg>[C++] Raise an exception when the urdf tree is badly formated<commit_after>//
// Copyright (c) 2015 CNRS
// Copyright (c) 2015 Wandercraft, 86 rue de Paris 91400 Orsay, France.
//
// This file is part of Pinocchio
// Pinocchio is free software: you can redistribute it
// and/or modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation, either version
// 3 of the License, or (at your option) any later version.
//
// Pinocchio is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// General Lesser Public License for more details. You should have
// received a copy of the GNU Lesser General Public License along with
// Pinocchio If not, see
// <http://www.gnu.org/licenses/>.

#ifndef __se3_urdf_hpp__
#define __se3_urdf_hpp__

#include <urdf_model/model.h>
#include <urdf_parser/urdf_parser.h>

#include <iostream>
#include <boost/foreach.hpp>
#include "pinocchio/multibody/model.hpp"

#include <exception>

namespace urdf
{
  typedef boost::shared_ptr<ModelInterface> ModelInterfacePtr;
  typedef boost::shared_ptr<const Joint> JointConstPtr;
  typedef boost::shared_ptr<const Link> LinkConstPtr;
  typedef boost::shared_ptr<Link> LinkPtr;
  typedef boost::shared_ptr<const Inertial> InertialConstPtr;
}

namespace se3
{
  namespace urdf
  {
    inline Inertia convertFromUrdf( const ::urdf::Inertial& Y )
    {
      const ::urdf::Vector3 & p = Y.origin.position;
      const ::urdf::Rotation & q = Y.origin.rotation;

      const Eigen::Vector3d com(p.x,p.y,p.z);
      const Eigen::Matrix3d & R = Eigen::Quaterniond(q.w,q.x,q.y,q.z).matrix();

      Eigen::Matrix3d I; I << Y.ixx,Y.ixy,Y.ixz
			   ,  Y.ixy,Y.iyy,Y.iyz
			   ,  Y.ixz,Y.iyz,Y.izz;
      return Inertia(Y.mass,com,R*I*R.transpose());
    }

    inline SE3 convertFromUrdf( const ::urdf::Pose & M )
    {
      const ::urdf::Vector3 & p = M.position;
      const ::urdf::Rotation & q = M.rotation;
      return SE3( Eigen::Quaterniond(q.w,q.x,q.y,q.z).matrix(), Eigen::Vector3d(p.x,p.y,p.z));
    }

    enum AxisCartesian { AXIS_X, AXIS_Y, AXIS_Z, AXIS_UNALIGNED };
    inline AxisCartesian extractCartesianAxis( const ::urdf::Vector3 & axis )
    {
      if( (axis.x==1.0)&&(axis.y==0.0)&&(axis.z==0.0) )
	return AXIS_X;
      else if( (axis.x==0.0)&&(axis.y==1.0)&&(axis.z==0.0) )
	return AXIS_Y;
      else if( (axis.x==0.0)&&(axis.y==0.0)&&(axis.z==1.0) )
	return AXIS_Z;
      else
	return AXIS_UNALIGNED;
    }

inline void parseTree( ::urdf::LinkConstPtr link, Model & model, const SE3 & placementOffset = SE3::Identity()) throw (std::invalid_argument)
{

  ::urdf::JointConstPtr joint = link->parent_joint;
  SE3 nextPlacementOffset = SE3::Identity(); // OffSet of the next link. In case we encounter a fixed joint, we need to propagate the length of its attached body to next joint.

  // std::cout << " *** " << link->name << "    < attached by joint ";
  // if(joint)
  //   std::cout << "#" << link->parent_joint->name << std::endl;
  // else std::cout << "###ROOT" << std::endl;

  // std::cout << "placementOffset: " << placementOffset << std::endl;

  if(joint!=NULL)
  {
    assert(link->getParent()!=NULL);

    if (!link->inertial && joint->type != ::urdf::Joint::FIXED)
    {
      const std::string exception_message (link->name + " - spatial inertia information missing.");
      throw std::invalid_argument(exception_message);
    }

    Model::Index parent = (link->getParent()->parent_joint==NULL) ? (model.existJointName("root_joint") ? model.getJointId("root_joint") : 0) :
                                                                    model.getJointId( link->getParent()->parent_joint->name );
    //std::cout << joint->name << " === " << parent << std::endl;

    const SE3 & jointPlacement = placementOffset*convertFromUrdf(joint->parent_to_joint_origin_transform);

    const Inertia & Y = (link->inertial) ?  convertFromUrdf(*link->inertial) :
    Inertia::Identity();

    bool visual = (link->visual) ? true : false;

    //std::cout << "Parent = " << parent << std::endl;
    //std::cout << "Placement = " << (Matrix4)jointPlacement << std::endl;

    switch(joint->type)
    {
      case ::urdf::Joint::REVOLUTE:
      case ::urdf::Joint::CONTINUOUS: // Revolute with no joint limits
      {

        Eigen::VectorXd maxEffort;
        Eigen::VectorXd velocity;
        Eigen::VectorXd lowerPosition;
        Eigen::VectorXd upperPosition;

        if (joint->limits)
        {
          maxEffort.resize(1);      maxEffort     << joint->limits->effort;
          velocity.resize(1);       velocity      << joint->limits->velocity;
          lowerPosition.resize(1);  lowerPosition << joint->limits->lower;
          upperPosition.resize(1);  upperPosition << joint->limits->upper;
        }

        Eigen::Vector3d jointAxis(Eigen::Vector3d::Zero());
        AxisCartesian axis = extractCartesianAxis(joint->axis);
        switch(axis)
        {
          case AXIS_X:
            model.addBody(  parent, JointModelRX(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Y:
            model.addBody(  parent, JointModelRY(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Z:
            model.addBody(  parent, JointModelRZ(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_UNALIGNED:
            jointAxis= Eigen::Vector3d( joint->axis.x,joint->axis.y,joint->axis.z );
            jointAxis.normalize();
            model.addBody(  parent, JointModelRevoluteUnaligned(jointAxis),
                            jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          default:
            assert( false && "Fatal Error while extracting revolute joint axis");
            break;
        }
        break;
      }
      case ::urdf::Joint::PRISMATIC:
      {
        Eigen::VectorXd maxEffort = Eigen::VectorXd(0.);
        Eigen::VectorXd velocity = Eigen::VectorXd(0.);
        Eigen::VectorXd lowerPosition = Eigen::VectorXd(0.);
        Eigen::VectorXd upperPosition = Eigen::VectorXd(0.);

        if (joint->limits)
        {
          maxEffort.resize(1);      maxEffort     << joint->limits->effort;
          velocity.resize(1);       velocity      << joint->limits->velocity;
          lowerPosition.resize(1);  lowerPosition << joint->limits->lower;
          upperPosition.resize(1);  upperPosition << joint->limits->upper;
        }
        AxisCartesian axis = extractCartesianAxis(joint->axis);
        switch(axis)
        {
          case AXIS_X:
            model.addBody(  parent, JointModelPX(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Y:
            model.addBody(  parent, JointModelPY(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_Z:
            model.addBody(  parent, JointModelPZ(), jointPlacement, Y,
                            maxEffort, velocity, lowerPosition, upperPosition,
                            joint->name,link->name, visual );
            break;
          case AXIS_UNALIGNED:
            std::cerr << "Bad axis = (" << joint->axis.x <<"," << joint->axis.y << "," << joint->axis.z << ")" << std::endl;
            assert(false && "Only X, Y or Z axis are accepted." );
            break;
          default:
            assert( false && "Fatal Error while extracting prismatic joint axis");
            break;
        }
        break;
      }
      case ::urdf::Joint::FIXED:
      {
        // In case of fixed joint, if link has inertial tag:
        //    -add the inertia of the link to his parent in the model
        // Otherwise do nothing.
        // In all cases:
        //    -let all the children become children of parent
        //    -inform the parser of the offset to apply
        //    -add fixed body in model to display it in gepetto-viewer
        if (link->inertial)
        {
          model.mergeFixedBody(parent, jointPlacement, Y); //Modify the parent inertia in the model
        }

        SE3 ptjot_se3 = convertFromUrdf(link->parent_joint->parent_to_joint_origin_transform);

        //transformation of the current placement offset
        nextPlacementOffset = placementOffset*ptjot_se3;

        //add the fixed Body in the model for the viewer
        model.addFixedBody(parent,nextPlacementOffset,link->name,visual);

        BOOST_FOREACH(::urdf::LinkPtr child_link,link->child_links)
        {
          child_link->setParent(link->getParent() );  //skip the fixed generation
        }
        break;
      }
      default:
      {
        std::cerr << "The joint type " << joint->type << " is not supported." << std::endl;
        assert(false && "Only revolute, prismatic and fixed joints are accepted." );
        break;
      }
    }

  }
          else if (link->getParent() != NULL)
          {
            const std::string exception_message (link->name + " - joint information missing.");
            throw std::invalid_argument(exception_message);
          }

  BOOST_FOREACH(::urdf::LinkConstPtr child,link->child_links)
  {
    parseTree(child, model, nextPlacementOffset);
  }
}

    template <typename D>
    void parseTree( ::urdf::LinkConstPtr link, Model & model, const SE3 & placementOffset , const JointModelBase<D> &  root_joint  )
    {
      const Inertia & Y = (link->inertial) ?
      convertFromUrdf(*link->inertial)
      : Inertia::Identity();
      model.addBody( 0, root_joint, placementOffset, Y , "root_joint", link->name, true );
      BOOST_FOREACH(::urdf::LinkConstPtr child,link->child_links)
      {
        parseTree(child, model, SE3::Identity());
      }
    }

    template <typename D>
    Model buildModel( const std::string & filename, const JointModelBase<D> &  root_joint )
    {
      Model model;

      ::urdf::ModelInterfacePtr urdfTree = ::urdf::parseURDFFile (filename);
      if (urdfTree)
        parseTree(urdfTree->getRoot(), model, SE3::Identity(), root_joint);
      else
      {
        const std::string exception_message ("The file " + filename + " does not contain a valid URDF model.");
        throw std::invalid_argument(exception_message);
      }

      return model;
    }

    inline Model buildModel( const std::string & filename)
    {
      Model model;

      ::urdf::ModelInterfacePtr urdfTree = ::urdf::parseURDFFile (filename);
      if (urdfTree)
        parseTree(urdfTree->getRoot(), model, SE3::Identity());
      else
      {
        const std::string exception_message ("The file " + filename + " does not contain a valid URDF model.");
        throw std::invalid_argument(exception_message);
      }

      return model;
    }

  } // namespace urdf
} // namespace se3

#endif // ifndef __se3_urdf_hpp__

<|endoftext|>


<commit_before>//---------------------------------------------------------
// Copyright 2016 Ontario Institute for Cancer Research
// Written by Jared Simpson (jared.simpson@oicr.on.ca)
//---------------------------------------------------------
//
// nanopolish_trainmodel - train a new pore model from
// the FAST5 output of a basecaller
//
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <vector>
#include <inttypes.h>
#include <assert.h>
#include <math.h>
#include <sys/time.h>
#include <algorithm>
#include <sstream>
#include <set>
#include <omp.h>
#include <getopt.h>
#include "htslib/faidx.h"
#include "nanopolish_poremodel.h"
#include "nanopolish_squiggle_read.h"
#include "nanopolish_methyltrain.h"
#include "training_core.hpp"
#include "profiler.h"

//
// Getopt
//
#define SUBPROGRAM "trainmodel"

static const char *TRAINMODEL_VERSION_MESSAGE =
SUBPROGRAM " Version " PACKAGE_VERSION "\n"
"Written by Jared Simpson.\n"
"\n"
"Copyright 2016 Ontario Institute for Cancer Research\n";

static const char *TRAINMODEL_USAGE_MESSAGE =
"Usage: " PACKAGE_NAME " " SUBPROGRAM " [OPTIONS] input.fofn\n"
"Train a new pore model using the basecalled reads in input.fofn\n"
"\n"
"  -v, --verbose                        display verbose output\n"
"      --version                        display version\n"
"      --help                           display this help and exit\n"
"\nReport bugs to " PACKAGE_BUGREPORT "\n\n";

namespace opt
{
    static unsigned int verbose;
    static std::string fofn_file;
}

static const char* shortopts = "v";

enum { OPT_HELP = 1, OPT_VERSION };

static const struct option longopts[] = {
    { "verbose",     no_argument,       NULL, 'v' },
    { "help",        no_argument,       NULL, OPT_HELP },
    { "version",     no_argument,       NULL, OPT_VERSION },
    { NULL, 0, NULL, 0 }
};

void parse_trainmodel_options(int argc, char** argv)
{
    bool die = false;
    for (char c; (c = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1;) {
        std::istringstream arg(optarg != NULL ? optarg : "");
        switch (c) {
            case '?': die = true; break;
            case 'v': opt::verbose++; break;
            case OPT_HELP:
                std::cout << TRAINMODEL_USAGE_MESSAGE;
                exit(EXIT_SUCCESS);
            case OPT_VERSION:
                std::cout << TRAINMODEL_VERSION_MESSAGE;
                exit(EXIT_SUCCESS);
        }
    }

    if (argc - optind < 1) {
        std::cerr << SUBPROGRAM ": not enough arguments\n";
        die = true;
    }

    if (argc - optind > 1) {
        std::cerr << SUBPROGRAM ": too many arguments\n";
        die = true;
    }

    if (die)
    {
        std::cout << "\n" << TRAINMODEL_USAGE_MESSAGE;
        exit(EXIT_FAILURE);
    }

    opt::fofn_file = argv[optind++];
}

int trainmodel_main(int argc, char** argv)
{
    parse_trainmodel_options(argc, argv);

    std::ifstream fofn_reader(opt::fofn_file);
    std::string fast5_name;

    // Read input
    std::vector<SquiggleRead*> reads;
    while(getline(fofn_reader, fast5_name)) {
        fprintf(stderr, "Loading %s\n", fast5_name.c_str());
        reads.push_back(new SquiggleRead(fast5_name, fast5_name));
    }
    fprintf(stderr, "Loaded %zu reads\n", reads.size());

    //
    unsigned int basecalled_k = 5; // TODO: infer this
    size_t num_kmers = gDNAAlphabet.get_num_strings(basecalled_k);
    unsigned int training_strand = T_IDX; // template training for now

    typedef std::vector<StateTrainingData> TrainingDataVector;
    typedef std::vector<TrainingDataVector> KmerTrainingData;

    // This vector is indexed by read, then kmer, then event
    std::vector<KmerTrainingData> read_training_data;

    size_t read_idx = 0;
    for(auto* read : reads) {

        // Initialize a vector-of-vectors to hold training data for this read
        read_training_data.push_back(KmerTrainingData());
        KmerTrainingData& kmer_training_data = read_training_data.back();

        // Initialize the events-by-kmer vector
        kmer_training_data.resize(num_kmers);
        printf("KTD: %zu\n", kmer_training_data.size());

        const std::string& read_sequence = read->read_sequence;
        size_t n_kmers = read_sequence.size() - basecalled_k + 1;
        for(size_t ki = 0; ki < n_kmers; ++ki) {

            IndexPair event_range_for_kmer = read->base_to_event_map[ki].indices[training_strand];

            // skip kmers without events and with multiple events
            if(event_range_for_kmer.start == -1 ||
               event_range_for_kmer.start != event_range_for_kmer.stop) {
                continue;
            }

            std::string kmer = read_sequence.substr(ki, basecalled_k);
            size_t kmer_rank = gDNAAlphabet.kmer_rank(kmer.c_str(), basecalled_k);
            assert(kmer_rank < num_kmers);

            size_t event_idx = event_range_for_kmer.start;
            double level = read->events[training_strand][event_idx].mean;
            double stdv = read->events[training_strand][event_idx].stdv;

            printf("read: %zu kmer: %s ki: %zu lvl: %.2lf dur: %.5lf\n", read_idx, kmer.c_str(), kmer_rank, level, read->events[training_strand][event_idx].duration);
            StateTrainingData std(level, stdv, read->pore_model[training_strand].var);
            kmer_training_data[kmer_rank].push_back(std);
        }

        read_idx++;
        /*
        std::string kmer(basecalled_k, 'A');
        for(size_t ki = 0; ki < num_kmers; ki++) {
            size_t num_events = kmer_training_data[ki].size();
            printf("%s:", kmer.c_str());
            for(size_t ei = 0; ei < num_events; ++ei) {
                printf("%.2lf\t", kmer_training_data[ki][ei].level_mean);
            }
            printf("\n");

            gDNAAlphabet.lexicographic_next(kmer);
        }
        */
    }

    // Select the read with the most events as the "baseline" read for generating the model
    size_t max_events = 0;
    size_t max_events_index = 0;
    for(size_t rti = 0; rti < read_training_data.size(); ++rti) {
        auto& kmer_training_data = read_training_data[rti];
        size_t total_events = 0;

        for(size_t ki = 0; ki < kmer_training_data.size(); ++ki) {
            total_events += kmer_training_data[ki].size();
        }
        printf("read %zu has %zu events (max: %zu, %zu)\n", rti, total_events, max_events, max_events_index);

        if(total_events > max_events) {
            max_events = total_events;
            max_events_index = rti;
        }
    }

    // Set the initial pore model based on the read with the most events
    PoreModel pore_model(basecalled_k);
    pore_model.states.resize(num_kmers);
    pore_model.scaled_states.resize(num_kmers);
    pore_model.scaled_params.resize(num_kmers);

    pore_model.shift = 0.0;
    pore_model.scale = 1.0;
    pore_model.drift = 0.0;
    pore_model.var = 1.0;
    pore_model.scale_sd = 1.0;
    pore_model.var_sd = 1.0;

    auto& kmer_training_data_for_selected = read_training_data[max_events_index];

    std::vector<bool> use_kmer(num_kmers, false);
    for(size_t ki = 0; ki < kmer_training_data_for_selected.size(); ++ki) {
        std::vector<double> values;
        std::stringstream ss;
        for(size_t ei = 0; ei < kmer_training_data_for_selected[ki].size(); ++ei) {
            values.push_back(kmer_training_data_for_selected[ki][ei].level_mean);
            ss << values.back() << " ";
        }

        // Set the kmer's mean parameter to be the median of the recorded values
        std::sort(values.begin(), values.end());

        size_t n = values.size();
        double median;
        if(n == 0) {
            median = 0.0f;
        } else {
            if(n % 2 == 0) {
                median = (values[n / 2 - 1] + values[n/2]) / 2.0f;
            } else {
                median = values[n/2];
            }

            // mark this kmer as valid
            use_kmer[ki] = true;
            pore_model.states[ki].level_mean = median;
            pore_model.states[ki].level_stdv = 1.0;
            printf("k: %zu median: %.2lf values: %s\n", ki, median, ss.str().c_str());
        }
    }
    pore_model.bake_gaussian_parameters();

    // Apply model to each read
    for(auto* read: reads) {
        read->pore_model[training_strand] = pore_model;
    }

    // Recalibrate read
    for(auto* read: reads) {

        // We generate a vector of event-to-kmer mapping to use the recalibration linear solver
        std::vector<EventAlignment> alignment;
        const std::string& read_sequence = read->read_sequence;
        size_t n_kmers = read_sequence.size() - basecalled_k + 1;
        for(size_t ki = 0; ki < n_kmers; ++ki) {

            IndexPair event_range_for_kmer = read->base_to_event_map[ki].indices[training_strand];

            // skip kmers without events and with multiple events
            if(event_range_for_kmer.start == -1 ||
               event_range_for_kmer.start != event_range_for_kmer.stop) {
                continue;
            }

            std::string kmer = read_sequence.substr(ki, basecalled_k);
            size_t kmer_rank = gDNAAlphabet.kmer_rank(kmer.c_str(), basecalled_k);
            assert(kmer_rank < num_kmers);

            size_t event_idx = event_range_for_kmer.start;

            // Only use this kmer if it is part of the initial model
            if(use_kmer[kmer_rank]) {
                EventAlignment ea;
                // ref data
                ea.ref_name = ""; // not needed
                ea.ref_kmer = kmer;
                ea.ref_position = ki;
                ea.read_idx = -1; // not needed
                ea.strand_idx = training_strand;
                ea.event_idx = event_idx;
                ea.rc = false;
                ea.model_kmer = kmer;
                ea.hmm_state = 'M'; // recalibration code only uses "M" alignments
                alignment.push_back(ea);
            }
        }

        recalibrate_model(*read,
                          training_strand,
                          alignment,
                          &gDNAAlphabet,
                          false);

        const PoreModel& read_model = read->pore_model[training_strand];
        printf("[recalibration] events: %zu alignment: %zu shift: %.2lf scale: %.2lf drift: %.4lf var: %.2lf\n",
            read->events[training_strand].size(),
            alignment.size(),
            read_model.shift,
            read_model.scale,
            read_model.drift,
            read_model.var);

    }

    // Deallocate input reads
    for(auto* read : reads) {
        delete read;
    }

    return 0;
}
<commit_msg>refactoring<commit_after>//---------------------------------------------------------
// Copyright 2016 Ontario Institute for Cancer Research
// Written by Jared Simpson (jared.simpson@oicr.on.ca)
//---------------------------------------------------------
//
// nanopolish_trainmodel - train a new pore model from
// the FAST5 output of a basecaller
//
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <vector>
#include <inttypes.h>
#include <assert.h>
#include <math.h>
#include <sys/time.h>
#include <algorithm>
#include <sstream>
#include <set>
#include <omp.h>
#include <getopt.h>
#include "htslib/faidx.h"
#include "nanopolish_poremodel.h"
#include "nanopolish_squiggle_read.h"
#include "nanopolish_methyltrain.h"
#include "training_core.hpp"
#include "profiler.h"

//
// Typedefs
//
typedef std::vector<StateTrainingData> TrainingDataVector;
typedef std::vector<TrainingDataVector> KmerTrainingData;

//
// Getopt
//
#define SUBPROGRAM "trainmodel"

static const char *TRAINMODEL_VERSION_MESSAGE =
SUBPROGRAM " Version " PACKAGE_VERSION "\n"
"Written by Jared Simpson.\n"
"\n"
"Copyright 2016 Ontario Institute for Cancer Research\n";

static const char *TRAINMODEL_USAGE_MESSAGE =
"Usage: " PACKAGE_NAME " " SUBPROGRAM " [OPTIONS] input.fofn\n"
"Train a new pore model using the basecalled reads in input.fofn\n"
"\n"
"  -v, --verbose                        display verbose output\n"
"      --version                        display version\n"
"      --help                           display this help and exit\n"
"\nReport bugs to " PACKAGE_BUGREPORT "\n\n";

namespace opt
{
    static unsigned int verbose;
    static std::string fofn_file;
}

static const char* shortopts = "v";

enum { OPT_HELP = 1, OPT_VERSION };

static const struct option longopts[] = {
    { "verbose",     no_argument,       NULL, 'v' },
    { "help",        no_argument,       NULL, OPT_HELP },
    { "version",     no_argument,       NULL, OPT_VERSION },
    { NULL, 0, NULL, 0 }
};

void parse_trainmodel_options(int argc, char** argv)
{
    bool die = false;
    for (char c; (c = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1;) {
        std::istringstream arg(optarg != NULL ? optarg : "");
        switch (c) {
            case '?': die = true; break;
            case 'v': opt::verbose++; break;
            case OPT_HELP:
                std::cout << TRAINMODEL_USAGE_MESSAGE;
                exit(EXIT_SUCCESS);
            case OPT_VERSION:
                std::cout << TRAINMODEL_VERSION_MESSAGE;
                exit(EXIT_SUCCESS);
        }
    }

    if (argc - optind < 1) {
        std::cerr << SUBPROGRAM ": not enough arguments\n";
        die = true;
    }

    if (argc - optind > 1) {
        std::cerr << SUBPROGRAM ": too many arguments\n";
        die = true;
    }

    if (die)
    {
        std::cout << "\n" << TRAINMODEL_USAGE_MESSAGE;
        exit(EXIT_FAILURE);
    }

    opt::fofn_file = argv[optind++];
}

std::vector<EventAlignment> generate_alignment_to_basecalls(const SquiggleRead* read,
                                                            const size_t k,
                                                            const size_t strand_idx,
                                                            const std::vector<bool>* use_kmer = NULL)
{
    size_t num_kmers_in_alphabet = gDNAAlphabet.get_num_strings(k);
    std::vector<EventAlignment> alignment;
    const std::string& read_sequence = read->read_sequence;
    size_t n_kmers = read_sequence.size() - k + 1;
    for(size_t ki = 0; ki < n_kmers; ++ki) {

        IndexPair event_range_for_kmer = read->base_to_event_map[ki].indices[strand_idx];

        // skip kmers without events and with multiple events
        if(event_range_for_kmer.start == -1 ||
           event_range_for_kmer.start != event_range_for_kmer.stop) {
            continue;
        }

        std::string kmer = read_sequence.substr(ki, k);
        size_t kmer_rank = gDNAAlphabet.kmer_rank(kmer.c_str(), k);
        assert(kmer_rank < num_kmers_in_alphabet);

        size_t event_idx = event_range_for_kmer.start;

        // Check if this kmer is marked as being useful
        if(use_kmer == NULL || use_kmer->at(kmer_rank)) {
            EventAlignment ea;
            // ref data
            ea.ref_name = ""; // not needed
            ea.ref_kmer = kmer;
            ea.ref_position = ki;
            ea.read_idx = -1; // not needed
            ea.strand_idx = strand_idx;
            ea.event_idx = event_idx;
            ea.rc = false;
            ea.model_kmer = kmer;
            ea.hmm_state = 'M'; // recalibration code only uses "M" alignments
            alignment.push_back(ea);
        }
    }

    return alignment;
}

KmerTrainingData alignment_to_training_data(const SquiggleRead* read,
                                            const std::vector<EventAlignment>& alignment,
                                            const size_t k,
                                            size_t read_idx)
{
    size_t num_kmers_in_alphabet = gDNAAlphabet.get_num_strings(k);
    KmerTrainingData kmer_training_data(num_kmers_in_alphabet);

    for(auto const& a : alignment) {
        size_t kmer_rank = gDNAAlphabet.kmer_rank(a.model_kmer.c_str(), k);
        assert(kmer_rank < num_kmers_in_alphabet);

        double level = read->events[a.strand_idx][a.event_idx].mean;
        double stdv = read->events[a.strand_idx][a.event_idx].stdv;

        StateTrainingData std(level, stdv, read->pore_model[a.strand_idx].var);
        kmer_training_data[kmer_rank].push_back(std);

        //fprintf(tsv_writer, "%zu\t%s\t%.2lf\t%.5lf\n", read_idx, kmer.c_str(), level, read->events[training_strand][event_idx].duration);
    }
    return kmer_training_data;
}

int trainmodel_main(int argc, char** argv)
{
    parse_trainmodel_options(argc, argv);

    std::ifstream fofn_reader(opt::fofn_file);
    std::string fast5_name;

    // Read input
    std::vector<SquiggleRead*> reads;
    while(getline(fofn_reader, fast5_name)) {
        fprintf(stderr, "Loading %s\n", fast5_name.c_str());
        reads.push_back(new SquiggleRead(fast5_name, fast5_name));
    }
    fprintf(stderr, "Loaded %zu reads\n", reads.size());

    //
    unsigned int basecalled_k = 5; // TODO: infer this
    size_t num_kmers = gDNAAlphabet.get_num_strings(basecalled_k);
    unsigned int training_strand = T_IDX; // template training for now

    // This vector is indexed by read, then kmer, then event
    std::vector<KmerTrainingData> read_training_data;

    FILE* tsv_writer = fopen("trainmodel.tsv", "w");
    fprintf(tsv_writer, "read_idx\tkmer\tlevel_mean\tduration\n");

    size_t read_idx = 0;
    for(auto* read : reads) {

        // extract alignment of events to k-mers
        std::vector<EventAlignment> alignment =
            generate_alignment_to_basecalls(read,
                                            basecalled_k,
                                            training_strand,
                                            NULL);

        // convert the alignment into model training data for this read
        KmerTrainingData training_data =
            alignment_to_training_data(read,
                                       alignment,
                                       basecalled_k,
                                       read_idx);

        read_training_data.push_back(training_data);
        read_idx++;
    }

    // Select the read with the most events as the "baseline" read for generating the model
    size_t max_events = 0;
    size_t max_events_index = 0;
    for(size_t rti = 0; rti < read_training_data.size(); ++rti) {
        auto& kmer_training_data = read_training_data[rti];
        size_t total_events = 0;

        for(size_t ki = 0; ki < kmer_training_data.size(); ++ki) {
            total_events += kmer_training_data[ki].size();
        }
        printf("read %zu has %zu events (max: %zu, %zu)\n", rti, total_events, max_events, max_events_index);

        if(total_events > max_events) {
            max_events = total_events;
            max_events_index = rti;
        }
    }

    // Set the initial pore model based on the read with the most events
    PoreModel pore_model(basecalled_k);
    pore_model.states.resize(num_kmers);
    pore_model.scaled_states.resize(num_kmers);
    pore_model.scaled_params.resize(num_kmers);

    pore_model.shift = 0.0;
    pore_model.scale = 1.0;
    pore_model.drift = 0.0;
    pore_model.var = 1.0;
    pore_model.scale_sd = 1.0;
    pore_model.var_sd = 1.0;

    auto& kmer_training_data_for_selected = read_training_data[max_events_index];

    std::vector<bool> use_kmer(num_kmers, false);
    for(size_t ki = 0; ki < kmer_training_data_for_selected.size(); ++ki) {
        std::vector<double> values;
        std::stringstream ss;
        for(size_t ei = 0; ei < kmer_training_data_for_selected[ki].size(); ++ei) {
            values.push_back(kmer_training_data_for_selected[ki][ei].level_mean);
            ss << values.back() << " ";
        }

        // Set the kmer's mean parameter to be the median of the recorded values
        std::sort(values.begin(), values.end());

        size_t n = values.size();
        double median;
        if(n == 0) {
            median = 0.0f;
        } else {
            if(n % 2 == 0) {
                median = (values[n / 2 - 1] + values[n/2]) / 2.0f;
            } else {
                median = values[n/2];
            }

            // mark this kmer as valid
            use_kmer[ki] = true;
            pore_model.states[ki].level_mean = median;
            pore_model.states[ki].level_stdv = 1.0;
            printf("k: %zu median: %.2lf values: %s\n", ki, median, ss.str().c_str());
        }
    }
    pore_model.bake_gaussian_parameters();

    // Apply model to each read
    for(auto* read: reads) {
        read->pore_model[training_strand] = pore_model;
    }

    // Recalibrate read
    for(auto* read: reads) {

        // generate an alignment between the RNN output and the basecalled read
        std::vector<EventAlignment> alignment =
            generate_alignment_to_basecalls(read,
                                            basecalled_k,
                                            training_strand,
                                            &use_kmer);

        // recalibrate shift/scale/etc
        recalibrate_model(*read,
                          training_strand,
                          alignment,
                          &gDNAAlphabet,
                          false);

        const PoreModel& read_model = read->pore_model[training_strand];
        printf("[recalibration] events: %zu alignment: %zu shift: %.2lf scale: %.2lf drift: %.4lf var: %.2lf\n",
            read->events[training_strand].size(),
            alignment.size(),
            read_model.shift,
            read_model.scale,
            read_model.drift,
            read_model.var);

    }

    // Deallocate input reads
    for(auto* read : reads) {
        delete read;
    }

    return 0;
}
<|endoftext|>


<commit_before>/* ---------------------------------------------------------------------
 * Numenta Platform for Intelligent Computing (NuPIC)
 * Copyright (C) 2016, Numenta, Inc.  Unless you have an agreement
 * with Numenta, Inc., for a separate license for this software code, the
 * following terms and conditions apply:
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU Affero Public License for more details.
 *
 * You should have received a copy of the GNU Affero Public License
 * along with this program.  If not, see http://www.gnu.org/licenses.
 *
 * http://numenta.org/licenses/
 * ---------------------------------------------------------------------
 */

/** @file
 * Define the ScalarEncoder and PeriodicScalarEncoder
 */

#ifndef NTA_ENCODERS_SCALAR
#define NTA_ENCODERS_SCALAR

#include <nupic/encoders/Base.hpp>

namespace nupic
{
  /** Encodes a floating point number as a contiguous block of 1s.
   *
   * @b Description
   * A ScalarEncoder encodes a numeric (floating point) value into an array
   * of bits. The output is 0's except for a contiguous block of 1's. The
   * location of this contiguous block varies continuously with the input value.
   *
   * Conceptually, the set of possible outputs is a set of "buckets". If there
   * are m buckets, the ScalarEncoder distributes m points along the domain
   * [minValue, maxValue], including the endpoints. To figure out the bucket
   * index of an input, it rounds the input to the nearest of these points.
   *
   * This approach is different from the PeriodicScalarEncoder because two
   * buckets, the first and last, are half as wide as the rest, since fewer
   * numbers in the input domain will round to these endpoints. This behavior
   * makes sense because, for example, with the input space [1, 10] and 10
   * buckets, 1.49 is in the first bucket and 1.51 is in the second.
   */
  class ScalarEncoder : public Encoder
  {
  public:
    /**
     * Constructs a ScalarEncoder
     *
     * @param w The number of bits that are set to encode a single value -- the
     *   "width" of the output signal
     * @param minValue The minimum value of the input signal, inclusive.
     * @param maxValue The maximum value of the input signal, inclusive.
     * @param clipInput Whether to allow input values outside the [minValue, maxValue]
     *   range.  If true, the input will be clipped to minValue or maxValue.
     *
     * There are three mutually exclusive parameters that determine the overall
     * size of of the output. Only one of these should be nonzero:
     *
     * @param n The number of bits in the output. Must be greater than or equal to w.
     * @param radius Two inputs separated by more than the radius have
     *   non-overlapping representations. Two inputs separated by less than the
     *   radius will in general overlap in at least some of their bits. You can
     *   think of this as the radius of the input.
     * @param resolution Two inputs separated by greater than, or equal to the
     *   resolution are guaranteed to have different representations.
     */
    ScalarEncoder(int w, double minValue, double maxValue, int n, double radius,
                  double resolution, bool clipInput);
    ~ScalarEncoder() override;

    virtual void encodeIntoArray(const ArrayBase & input, UInt output[],
                                 bool learn) override;
    virtual int getWidth() const override { return n_; }

  private:
    int w_;
    int n_;
    double minValue_;
    double maxValue_;
    double resolution_;
    bool clipInput_;
  }; // end class ScalarEncoder

  /** Encodes a floating point number as a block of 1s that might wrap around.
   *
   * @b Description
   * A PeriodicScalarEncoder encodes a numeric (floating point) value into an
   * array of bits. The output is 0's except for a contiguous block of 1's that
   * may wrap around the edge. The location of this contiguous block varies
   * continuously with the input value.
   *
   * Conceptually, the set of possible outputs is a set of "buckets". If there
   * are m buckets, the PeriodicScalarEncoder plots m equal-width bands along
   * the domain [minValue, maxValue]. The bucket index of an input is simply its
   * band index.
   *
   * Because of the equal-width buckets, the rounding differs from the
   * ScalarEncoder. In cases where the ScalarEncoder would put 1.49 in the first
   * bucket and 1.51 in the second, the PeriodicScalarEncoder will put 1.99 in
   * the first bucket and 2.0 in the second.
   */
  class PeriodicScalarEncoder : public Encoder
  {
  public:
    /**
     * Constructs a PeriodicScalarEncoder
     *
     * @param w The number of bits that are set to encode a single value -- the
     *   "width" of the output signal
     * @param minValue The minimum value of the input signal, inclusive.
     * @param maxValue The maximum value of the input signal, exclusive. All
     *   inputs will be strictly less than this value.
     *
     * There are three mutually exclusive parameters that determine the overall
     * size of the output. Only one of these should be nonzero:
     *
     * @param n The number of bits in the output. Must be greater than or equal
     *   to w.
     * @param radius Two inputs separated by more than the radius have
     *   non-overlapping representations. Two inputs separated by less than the
     *   radius will in general overlap in at least some of their bits. You can
     *   think of this as the radius of the input.
     * @param resolution Two inputs separated by greater than, or equal to the
     *   resolution are guaranteed to have different representations.
     */
    PeriodicScalarEncoder(int w, double minValue, double maxValue, int n,
                          double radius, double resolution);
    ~PeriodicScalarEncoder() override;

    virtual void encodeIntoArray(const ArrayBase & input, UInt output[],
                                 bool learn) override;
    virtual int getWidth() const override { return n_; }

  private:
    int w_;
    int n_;
    double minValue_;
    double maxValue_;
    double resolution_;
  }; // end class PeriodicScalarEncoder
} // end namespace nupic

#endif // NTA_ENCODERS_SCALAR
<commit_msg>Explicitly mark all derived virtual methods as virtual.<commit_after>/* ---------------------------------------------------------------------
 * Numenta Platform for Intelligent Computing (NuPIC)
 * Copyright (C) 2016, Numenta, Inc.  Unless you have an agreement
 * with Numenta, Inc., for a separate license for this software code, the
 * following terms and conditions apply:
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU Affero Public License for more details.
 *
 * You should have received a copy of the GNU Affero Public License
 * along with this program.  If not, see http://www.gnu.org/licenses.
 *
 * http://numenta.org/licenses/
 * ---------------------------------------------------------------------
 */

/** @file
 * Define the ScalarEncoder and PeriodicScalarEncoder
 */

#ifndef NTA_ENCODERS_SCALAR
#define NTA_ENCODERS_SCALAR

#include <nupic/encoders/Base.hpp>

namespace nupic
{
  /** Encodes a floating point number as a contiguous block of 1s.
   *
   * @b Description
   * A ScalarEncoder encodes a numeric (floating point) value into an array
   * of bits. The output is 0's except for a contiguous block of 1's. The
   * location of this contiguous block varies continuously with the input value.
   *
   * Conceptually, the set of possible outputs is a set of "buckets". If there
   * are m buckets, the ScalarEncoder distributes m points along the domain
   * [minValue, maxValue], including the endpoints. To figure out the bucket
   * index of an input, it rounds the input to the nearest of these points.
   *
   * This approach is different from the PeriodicScalarEncoder because two
   * buckets, the first and last, are half as wide as the rest, since fewer
   * numbers in the input domain will round to these endpoints. This behavior
   * makes sense because, for example, with the input space [1, 10] and 10
   * buckets, 1.49 is in the first bucket and 1.51 is in the second.
   */
  class ScalarEncoder : public Encoder
  {
  public:
    /**
     * Constructs a ScalarEncoder
     *
     * @param w The number of bits that are set to encode a single value -- the
     *   "width" of the output signal
     * @param minValue The minimum value of the input signal, inclusive.
     * @param maxValue The maximum value of the input signal, inclusive.
     * @param clipInput Whether to allow input values outside the [minValue, maxValue]
     *   range.  If true, the input will be clipped to minValue or maxValue.
     *
     * There are three mutually exclusive parameters that determine the overall
     * size of of the output. Only one of these should be nonzero:
     *
     * @param n The number of bits in the output. Must be greater than or equal to w.
     * @param radius Two inputs separated by more than the radius have
     *   non-overlapping representations. Two inputs separated by less than the
     *   radius will in general overlap in at least some of their bits. You can
     *   think of this as the radius of the input.
     * @param resolution Two inputs separated by greater than, or equal to the
     *   resolution are guaranteed to have different representations.
     */
    ScalarEncoder(int w, double minValue, double maxValue, int n, double radius,
                  double resolution, bool clipInput);
    ~ScalarEncoder() override;

    virtual void encodeIntoArray(const ArrayBase & input, UInt output[],
                                 bool learn) override;
    virtual int getWidth() const override { return n_; }

  private:
    int w_;
    int n_;
    double minValue_;
    double maxValue_;
    double resolution_;
    bool clipInput_;
  }; // end class ScalarEncoder

  /** Encodes a floating point number as a block of 1s that might wrap around.
   *
   * @b Description
   * A PeriodicScalarEncoder encodes a numeric (floating point) value into an
   * array of bits. The output is 0's except for a contiguous block of 1's that
   * may wrap around the edge. The location of this contiguous block varies
   * continuously with the input value.
   *
   * Conceptually, the set of possible outputs is a set of "buckets". If there
   * are m buckets, the PeriodicScalarEncoder plots m equal-width bands along
   * the domain [minValue, maxValue]. The bucket index of an input is simply its
   * band index.
   *
   * Because of the equal-width buckets, the rounding differs from the
   * ScalarEncoder. In cases where the ScalarEncoder would put 1.49 in the first
   * bucket and 1.51 in the second, the PeriodicScalarEncoder will put 1.99 in
   * the first bucket and 2.0 in the second.
   */
  class PeriodicScalarEncoder : public Encoder
  {
  public:
    /**
     * Constructs a PeriodicScalarEncoder
     *
     * @param w The number of bits that are set to encode a single value -- the
     *   "width" of the output signal
     * @param minValue The minimum value of the input signal, inclusive.
     * @param maxValue The maximum value of the input signal, exclusive. All
     *   inputs will be strictly less than this value.
     *
     * There are three mutually exclusive parameters that determine the overall
     * size of the output. Only one of these should be nonzero:
     *
     * @param n The number of bits in the output. Must be greater than or equal
     *   to w.
     * @param radius Two inputs separated by more than the radius have
     *   non-overlapping representations. Two inputs separated by less than the
     *   radius will in general overlap in at least some of their bits. You can
     *   think of this as the radius of the input.
     * @param resolution Two inputs separated by greater than, or equal to the
     *   resolution are guaranteed to have different representations.
     */
    PeriodicScalarEncoder(int w, double minValue, double maxValue, int n,
                          double radius, double resolution);
    virtual ~PeriodicScalarEncoder() override;

    virtual void encodeIntoArray(const ArrayBase & input, UInt output[],
                                 bool learn) override;
    virtual int getWidth() const override { return n_; }

  private:
    int w_;
    int n_;
    double minValue_;
    double maxValue_;
    double resolution_;
  }; // end class PeriodicScalarEncoder
} // end namespace nupic

#endif // NTA_ENCODERS_SCALAR
<|endoftext|>