C++ Kmer类代码示例

TEST(Kmer, canonicalize)
{
	Kmer::setLength(4);
	Kmer canonical("ATGC");
	Kmer nonCanonical("GCAT");
	Kmer palindrome("ACGT");

	Kmer kmer = canonical;
	kmer.canonicalize();
	EXPECT_EQ(canonical, kmer);

	kmer = nonCanonical;
	kmer.canonicalize();
	EXPECT_EQ(canonical, kmer);

	kmer = palindrome;
	kmer.canonicalize();
	EXPECT_EQ(palindrome, kmer);

	Kmer::setLength(5);
	Kmer oddLength("GCTCG");
	Kmer oddLengthCanonical("CGAGC");

	kmer = oddLength;
	kmer.canonicalize();
	EXPECT_EQ(oddLengthCanonical, kmer);
}

void NetworkSequenceCollection::processSequenceExtension(
		uint64_t groupID, uint64_t branchID, const Kmer& seq,
		const ExtensionRecord& extRec, int multiplicity)
{
	switch(m_state)
	{
		case NAS_TRIM:
			return processLinearSequenceExtension(groupID, branchID,
					seq, extRec, multiplicity, m_trimStep);
		case NAS_ASSEMBLE:
		case NAS_COVERAGE:
			return processLinearSequenceExtension(groupID, branchID,
					seq, extRec, multiplicity, UINT_MAX);
		case NAS_DISCOVER_BUBBLES:
			return processSequenceExtensionPop(groupID, branchID,
					seq, extRec, multiplicity,
					opt::bubbleLen - opt::kmerSize + 1);
		case NAS_WAITING:
			if (m_finishedGroups.count(groupID) == 0) {
				logger(0) << "error: unexpected seqext message: "
					"state: " << m_state << " "
					"gid: " << groupID << " bid: " << branchID << " "
					"seq: " << seq.str() << '\n';
				assert(false);
			}
			break;
		default:
			logger(0) << "error: unexpected seqext message: "
				"state: " << m_state << " "
				"gid: " << groupID << " bid: " << branchID << " "
				"seq: " << seq.str() << '\n';
			assert(false);
			break;
	}
}

void GraphPath::reverseContent(GraphPath & newPath) const {

#ifdef CONFIG_ASSERT
	assert(newPath.size() == 0);
#endif

	newPath.setKmerLength(getKmerLength());

	for(int i = size() - 1 ; i >= 0 ; --i) {
		Kmer element;
		at(i, &element);

		// the false here drops support for colored data (SOLiD)
		// anyway who cares

		Kmer newElement = element.complementVertex(getKmerLength(), false);
		newPath.push_back(&newElement);
	}

#ifdef CONFIG_ASSERT
	assert(size() == newPath.size());
	assert(getKmerLength() == newPath.getKmerLength());
#endif

}

int GraphPath::getRequiredNumberOfBytes() const {

	int position = 0;

	uint32_t elements = size();

	int operationSize = sizeof(uint32_t);
	//memcpy(buffer + position, &elements, operationSize);
	position += operationSize;


#ifdef CONFIG_ASSERT
	uint32_t kmerLength = getKmerLength();
	assert(kmerLength > 0);
#endif

	//memcpy(buffer + position, &kmerLength, operationSize);
	position += operationSize;

	//cout << "[DEBUG] GraphPath::dump kmerLength " << kmerLength << endl;

	for(int i = 0 ; i < (int)elements ; i ++) {
		Kmer value;
		at(i, &value);
		//position += value.dump(buffer + position);
		position += value.getRequiredNumberOfBytes();
	}

	return position;


}

int GraphPath::load(const char * buffer) {
	int position = 0;

	uint32_t elements = 0;
	int operationSize = sizeof(uint32_t);

	memcpy(&elements, buffer + position, operationSize);
	position += operationSize;

	uint32_t kmerLength = 0;
	memcpy(&kmerLength, buffer + position, operationSize);
	position += operationSize;

	setKmerLength(kmerLength);
	//cout << "[DEBUG] GraphPath::load kmerLength " << kmerLength << endl;

	for(int i = 0 ; i < (int)elements ; i ++) {
		Kmer value;
		position += value.load(buffer + position);
		push_back(&value);
	}

	//cout << "DEBUG] loaded " << size() << " items for GraphPath" << endl;

	return position;
}

int64 HashGraph::Trim(int minLength)
{
    vector<Contig> contigs;
    Assemble(contigs);

    int total = 0;
#pragma omp parallel for
    for (int i = 0; i < (int)contigs.size(); ++i)
    {
        if (contigs[i].IsTangle() && contigs[i].Size() < kmerLength + minLength - 1)
        {
            Kmer kmer;
            for (int j = 0; j+1 < kmerLength; ++j)
                kmer.AddRight(contigs[i][j]);
            for (int j = kmerLength-1; j < contigs[i].Size(); ++j)
            {
                kmer.AddRight(contigs[i][j]);
                KmerNode *node = GetNode(kmer);
                if (node != NULL)
                    node->SetDeadFlag();
            }

#pragma omp atomic
            ++total;
        }
    }

    Refresh();

    LogMessage("trim %lld dead ends\n", total);

    return total;
}

void f2() {
    srand(time(NULL));
    int size=1536;
    map<int,int> counts;

    uint64_t samples=100000000;
    uint64_t base=rand();
    int wordSize=63;
    Kmer kmer;
    kmer.setU64(0,base);

    int average=samples/size;
    while(samples--) {
        uint64_t second=rand();
        kmer.setU64(1,second);
        int rank=vertexRank(&kmer,size,wordSize,false);
        counts[rank]++;
    }
    vector<int> data;
    for(int i=0; i<size; i++) {
        data.push_back(counts[i]);
        //cout<<i<<" "<<counts[i]<<endl;
    }
    int deviation=average/10;
    int min=average-deviation;
    int max=average+deviation;
    for(int i=0; i<size; i++) {
        if(counts[i]>=max) {
            cout<<counts[i]<<" and Max="<<max<<endl;
        }
        assert(counts[i]<max);
        assert(counts[i]>min);
    }
}

void ContigGraph::BuildEdgeCountTable()
{
    edge_count_table_.clear();
    edge_count_table_.set_kmer_size(kmer_size_+1);
#pragma omp parallel for schedule(static, 1)
    for (int64_t i = 0; i < (int64_t)vertices_.size(); ++i)
    {
        for (int strand = 0; strand < 2; ++strand)
        {
            ContigGraphVertexAdaptor current(&vertices_[i], strand);

            Kmer kmer = current.end_kmer(kmer_size_);
            kmer.resize(kmer_size_+1);
            for (int x = 0; x < 4; ++x)
            {
                if (current.out_edges()[x])
                {
                    kmer.set_base(kmer_size_, x);
                    edge_count_table_.InsertVertex(kmer);
                }
            }
        }
    }
    edge_count_table_.ClearCount();
}

void
verify_node_orig(kg_node_t * node, unsigned kmer_length) {
    assert( false && "TODO FIX! REVERSED KMER ENDIANNESS" );
    int double_kmer_length = kmer_length << 1;
#ifdef LARGE_KMERS
    Kmer mask;
    mask.createMask(double_kmer_length);
#else
    Kmer mask = (Kmer(1) << double_kmer_length) - 1;
#endif
    Kmer kmer = node->kmer;
    Kmer rc_kmer = reverseComplement(kmer, kmer_length);
    char leftmost_base = (kmer >> (double_kmer_length - 2)) & 0x3;
    char rightmost_base = kmer & 0x3;

    for (int i = 0 ; i < 4 ; ++ i) {
        // check on the left side
        kg_node_t * node2 = node->left[i];
        int count = node->left_count[i];

        if (node2) {
            assert (count != 0);
            if (count > 0) {
                Kmer kmer2 = KMER_PREPEND(kmer, i, double_kmer_length, mask);
                assert(kmer2 == node2->kmer);
                assert(node2->right[(int)rightmost_base] == node);
                assert(node2->right_count[(int)rightmost_base] == count);
            } else {
                Kmer kmer2 = KMER_APPEND(rc_kmer, i ^ 0x3, double_kmer_length, mask);
                assert(kmer2 == node2->kmer);
                assert(node2->left[rightmost_base ^ 0x3] == node);
                assert(node2->left_count[rightmost_base ^ 0x3] == count);
            }
        } else {
            assert (count == 0);
        }


        // check on the right side
        node2 = node->right[i];
        count = node->right_count[i];

        if (node2) {
            assert (count != 0);
            if (count > 0) {
                Kmer kmer2 = KMER_APPEND(kmer, i, double_kmer_length, mask);
                assert(kmer2 == node2->kmer);
                assert(node2->left[(int)leftmost_base] == node);
                assert(node2->left_count[(int)leftmost_base] == count);
            } else {
                Kmer kmer2 = KMER_PREPEND(rc_kmer, i ^ 0x3, double_kmer_length, mask);
                assert(kmer2 == node2->kmer);
                assert(node2->right[leftmost_base ^ 0x3] == node);
                assert(node2->right_count[leftmost_base ^ 0x3] == count);
            }
        } else {
            assert (count == 0);
        }
    }
}

bool HashGraph::AddEdgesFromSequence(const Sequence &seq)
{
    if (seq.Size() < kmerLength)
        return false;

    bool flag = false;
    Kmer kmer;
    for (int i = 0; i < kmerLength-1; ++i)
        kmer.AddRight(seq[i]);
    for (int i = kmerLength-1; i < seq.Size(); ++i)
    {
        kmer.AddRight(seq[i]);

        KmerNodeAdapter adp = GetNodeAdapter(kmer);
        if (!adp.IsNull())
        {
            flag = true;
            adp.Increase();
            if (i >= (int)kmerLength)
            {
                adp.AddInEdge(3 - seq[i-kmerLength]);
            }

            if (i+1 < seq.Size())
            {
                adp.AddOutEdge(seq[i+1]);
            }
        }
    }

    return flag;
}

void SpuriousSeedAnnihilator::writeCheckpointForSeeds(){

	/* write the Seeds checkpoint */
	if(m_parameters->writeCheckpoints() && !m_parameters->hasCheckpoint("Seeds")){

		ofstream f(m_parameters->getCheckpointFile("Seeds").c_str());
		cout<<"Rank "<<m_parameters->getRank()<<" is writing checkpoint Seeds"<<endl;
		int count=(*m_seeds).size();

		f.write((char*)&count,sizeof(int));

		for(int i=0;i<(int)(*m_seeds).size();i++){
			int length=(*m_seeds)[i].size();
			f.write((char*)&length,sizeof(int));

			for(int j=0;j<(int)(*m_seeds)[i].size();j++){
				Kmer theKmer;
				(*m_seeds)[i].at(j,&theKmer);
				theKmer.write(&f);

				CoverageDepth coverageValue=0;
				coverageValue=(*m_seeds)[i].getCoverageAt(j);
				f.write((char*)&coverageValue,sizeof(CoverageDepth));
			}
		}
		f.close();
	}
}

void SeedingData::writeCheckpoints(){

	/* write the Seeds checkpoint */
	if(m_parameters->writeCheckpoints() && !m_parameters->hasCheckpoint("SimpleSeeds")){

		ofstream f(m_parameters->getCheckpointFile("SimpleSeeds").c_str());
		ostringstream buffer;

		cout<<"Rank "<<m_parameters->getRank()<<" is writing checkpoint SimpleSeeds"<<endl;

		vector<GraphPath> * seeds = & m_SEEDING_seeds;

		int count=(*seeds).size();

		buffer.write((char*)&count, sizeof(int));

		for(int i=0;i<(int)(*seeds).size();i++){
			int length=(*seeds)[i].size();
			buffer.write((char*)&length, sizeof(int));

			for(int j=0;j<(int)(*seeds)[i].size();j++){
				Kmer theKmer;
				(*seeds)[i].at(j,&theKmer);
				theKmer.write(&buffer);

				CoverageDepth coverageValue=0;
				coverageValue=(*seeds)[i].getCoverageAt(j);
				buffer.write((char*)&coverageValue, sizeof(CoverageDepth));
				flushFileOperationBuffer(false, &buffer, &f, CONFIG_FILE_IO_BUFFER_SIZE);
			}
		}
                flushFileOperationBuffer(true, &buffer, &f, CONFIG_FILE_IO_BUFFER_SIZE);
		f.close();
	}
}

void GraphPath::push_back(const Kmer*a){

#ifdef ASSERT
	assert(m_kmerLength!=0);
#endif

	if(!canBeAdded(a)){
		if(!m_errorRaised){
			cout<<"Error: can not add "<<a->idToWord(m_kmerLength,false)<<endl;
			cout<<"last objects:"<<endl;
			int count=16;
			int iterator=size()-count;
			while(iterator<size()){
				Kmer theObject;
				at(iterator,&theObject);

				cout<<" ["<<iterator<<"] ------> "<<theObject.idToWord(m_kmerLength,false)<<endl;

				iterator++;
			}

			m_errorRaised=true;
		}

		return;
	}

#ifdef CONFIG_PATH_STORAGE_DEFAULT
	m_vertices.push_back(*a);
#elif defined(CONFIG_PATH_STORAGE_BLOCK)

	writeObjectInBlock(a);
#endif
}

void HashGraph::InsertSequence(const Sequence &seq, uint64 prefix, uint64 mask)
{
    if (seq.Size() < kmerLength)
        return;

    Kmer kmer;
    for (int i = 0; i < kmerLength-1; ++i)
        kmer.AddRight(seq[i]);
    for (int i = kmerLength-1; i < seq.Size(); ++i)
    {
        kmer.AddRight(seq[i]);
        Kmer key = kmer;
        Kmer rev_comp = kmer;
        rev_comp.ReverseComplement();
        if (rev_comp < kmer)
            key = rev_comp;

        if ((key.Hash() & mask) == prefix)
        {
            KmerNodeAdapter adp(InsertKmer(kmer), kmer);
            if (i >= (int)kmerLength)
            {
                adp.AddInEdge(3 - seq[i-kmerLength]);
            }

            if (i+1 < seq.Size())
            {
                adp.AddOutEdge(seq[i+1]);
            }
        }
    }
}

void SeedingData::loadCheckpoint(){
	cout<<"Rank "<<m_parameters->getRank()<<" is reading checkpoint Seeds"<<endl;

	ifstream f(m_parameters->getCheckpointFile("Seeds").c_str());
	int n=0;
	f.read((char*)&n,sizeof(int));
	for(int i=0;i<n;i++){
		GraphPath seed;
		seed.setKmerLength(m_parameters->getWordSize());
		int vertices=0;
		f.read((char*)&vertices,sizeof(int));
		for(int j=0;j<vertices;j++){
			Kmer kmer;
			kmer.read(&f);
			seed.push_back(&kmer);

			CoverageDepth coverageValue=0;

			f.read((char*)&coverageValue,sizeof(CoverageDepth));
			seed.addCoverageValue(coverageValue);
		}

		seed.computePeakCoverage();

		m_SEEDING_seeds.push_back(seed);
	}
	cout<<"Rank "<<m_parameters->getRank()<<" loaded "<<n<<" seeds from checkpoint Seeds"<<endl;
	f.close();
}