kmerstore.h 12.4 KB
Newer Older
Mikaël Salson's avatar
Mikaël Salson committed
1 2 3 4 5 6 7 8
#ifndef KMERSTORE_H
#define KMERSTORE_H

#include <map>
#include <string>
#include <list>
#include <stdexcept>
#include <stdint.h>
9
#include <math.h>
Mikaël Salson's avatar
Mikaël Salson committed
10 11 12 13 14 15 16 17 18 19
#include "fasta.h"
#include "tools.h"

using namespace std;

class Kmer {
public:
  unsigned int count;

  Kmer();
20 21 22 23

  /**
   * This constructor is used via a IKmerStore<Kmer> index (hence the argument list)       
   */
24
  Kmer(const string &label, int strand=1);
25

Mikaël Salson's avatar
Mikaël Salson committed
26 27
  Kmer &operator+=(const Kmer &);
  static bool hasRevcompSymetry();
28 29 30 31 32

  /**
   * @return true if the element is the same as when initialised with default constructor.
   */
  bool isNull();
Mikaël Salson's avatar
Mikaël Salson committed
33 34 35 36 37 38 39 40 41 42 43 44 45
} ;
ostream &operator<<(ostream &os, const Kmer &kmer);


/* K-mer indexing */

template <class T> class IKmerStore
{
protected:
  bool revcomp_indexed;
  int k; // weight of the seed
  int s; // span of the seed (s >= k)
  string seed ;
46
  size_t nb_kmers_inserted;
Mikaël Salson's avatar
Mikaël Salson committed
47 48 49 50 51

public:

  virtual ~IKmerStore();

52 53
  list< pair <T, string> > labels;

Mikaël Salson's avatar
Mikaël Salson committed
54 55 56
  /**
   * @param input: A single FASTA file
   * @param label: label that must be associated to the given files
57
   * @post All the sequences in the FASTA files have been indexed, and the label is stored in the list of labels
Mikaël Salson's avatar
Mikaël Salson committed
58
   */
59
  void insert(Fasta& input, const string& label="", int keep_only = 0);
Mikaël Salson's avatar
Mikaël Salson committed
60 61 62 63

  /**
   * @param input: A list of FASTA files
   * @param label: label that must be associated to the given files
64
   * @post All the sequences in the FASTA files have been indexed, and the label is stored in the list of labels
Mikaël Salson's avatar
Mikaël Salson committed
65
   */
66
  void insert(list<Fasta>& input, const string& label="", int keep_only = 0);
Mikaël Salson's avatar
Mikaël Salson committed
67 68 69 70
  
  /**
   * @param input: A sequence to be cut in k-mers
   * @param label: label that must be associated to the given files
71 72 73 74
   * @param keep_only: if > 0 will keep at most the last keep_only nucleotides
   *                   of the sequence. if < 0 will keep at most the first
   *                   keep_only nucleotides of the sequence. if == 0,
   *                   will keep all the sequence.
Mikaël Salson's avatar
Mikaël Salson committed
75 76
   * @post All the k-mers in the sequence have been indexed.
   */
77
  void insert(const seqtype &sequence,
78
              const string &label,
79
              bool ignore_extended_nucleotides=true, int keep_only = 0);
Mikaël Salson's avatar
Mikaël Salson committed
80 81 82 83 84

  /**
   * @param word: a k-mer
   * @return the value for the kmer as stored in the structure (don't do revcomp)
   */ 
85
  virtual T& get(seqtype &word) = 0;
Mikaël Salson's avatar
Mikaël Salson committed
86

87 88 89 90 91
  /**
   * @return the percentage of kmers that are set in the index
   */
  float getIndexLoad() const;

92 93 94 95 96
  /**
   * @return probability that the number of kmers is 'at_least' or more in a sequence of length 'length'
   */
  double getProbabilityAtLeastOrAbove(int at_least, int length) const;

Mikaël Salson's avatar
Mikaël Salson committed
97 98 99 100 101 102 103 104 105 106
  /**
   * @return the value of k
   */
  int getK() const;

  /**
   * @return the value of s
   */
  int getS() const;

107 108 109 110 111
  /**
   * @return the seed used
   */
  string getSeed() const;

112 113 114 115 116 117
  /**
   * @param kmer: a kmer
   * @return one label associated with the kmer
   */
  string getLabel(T kmer) const;

Mikaël Salson's avatar
Mikaël Salson committed
118 119 120 121 122 123 124
  /**
   * @param seq: a sequence
   * @param no_revcomp: force not to revcomp the sequence, even if
   *                    the index was built with revcomp.
   * @return a vector of length seq.length() - getK() + 1 containing
   * for each k-mer the corresponding value in the index.
   */
125
  virtual vector<T> getResults(const seqtype &seq, bool no_revcomp=false) = 0;
Mikaël Salson's avatar
Mikaël Salson committed
126 127 128 129 130 131 132 133 134 135 136 137

  /**
   * @return true iff the revcomp is indexed
   */
  bool isRevcomp() const;

  /**
   * @param string: a k-mer
   * @return the value in the index associated to the given k-mer.
   *         isRevcomp() ==> the revcomp of the k-mer may be considered
   * @pre word.length() == this->k
   */
138
  virtual T& operator[](seqtype& word) = 0;
Mikaël Salson's avatar
Mikaël Salson committed
139 140 141 142 143 144 145
};


template <class T> 
class MapKmerStore : public IKmerStore<T>
{
public:	
146
  map<seqtype, T> store;
Mikaël Salson's avatar
Mikaël Salson committed
147 148 149 150 151 152 153

  /**
   * @param bool: if true, will also index revcomp
   * @param seed: the seed
   */
  MapKmerStore(string seed, bool=false);
  MapKmerStore(int k, bool=false);
154 155

  vector<T> getResults(const seqtype &seq, bool no_revcomp=false);
156 157
  T& get(seqtype &word);
  T& operator[](seqtype & word);
158 159 160

 private:
  void init();
Mikaël Salson's avatar
Mikaël Salson committed
161 162 163 164 165 166 167
};

template <class T> 
class ArrayKmerStore : public IKmerStore<T>
{
  T* store;
	
168
  int index(const seqtype& word) const;
169 170 171 172 173 174 175 176

  /**
   * Allocates memory for store
   * @throws bad_alloc: when there is not enough memory or when the
   *                    the total number of kmers cannot be stored in an 
   *                    unsigned int
   */
  void init();
Mikaël Salson's avatar
Mikaël Salson committed
177 178 179 180 181 182 183 184
public:
  /**
   * @param bool: if true, will also index revcomp
   * @param seed: the seed
   */
  ArrayKmerStore(string seed, bool=false);
  ArrayKmerStore(int k, bool=false);
  ~ArrayKmerStore();
185 186

  vector<T> getResults(const seqtype &seq, bool no_revcomp=false);	
187 188
  T& get(seqtype &word);
  T& operator[](seqtype & word);
189
  T& operator[](int word);
Mikaël Salson's avatar
Mikaël Salson committed
190 191 192 193 194 195 196 197 198 199
};


// IKmerStore

template<class T>
IKmerStore<T>::~IKmerStore(){}

template<class T> 
void IKmerStore<T>::insert(list<Fasta>& input,
200 201
                           const string &label,
                           int keep_only){
Mikaël Salson's avatar
Mikaël Salson committed
202
  for(list<Fasta>::iterator it = input.begin() ; it != input.end() ; it++){
203
    insert(*it, label, keep_only);
Mikaël Salson's avatar
Mikaël Salson committed
204 205 206 207 208
  }
}

template<class T> 
void IKmerStore<T>::insert(Fasta& input,
209 210
                           const string &label,
                           int keep_only){
Mikaël Salson's avatar
Mikaël Salson committed
211
  for (int r = 0; r < input.size(); r++) {
212
    insert(input.sequence(r), label, true, keep_only);
Mikaël Salson's avatar
Mikaël Salson committed
213
  }
214

215
  labels.push_back(make_pair(T(label, 1), label)) ;
Mikaël Salson's avatar
Mikaël Salson committed
216 217 218
}

template<class T> 
219
void IKmerStore<T>::insert(const seqtype &sequence,
220
                           const string &label,
221 222 223 224 225 226 227 228 229 230
                           bool ignore_extended_nucleotides,
                           int keep_only){
  size_t start_indexing = 0;
  size_t end_indexing = sequence.length();
  if (keep_only > 0 && sequence.length() > (size_t)keep_only) {
    start_indexing = sequence.length() - keep_only;
  } else if (keep_only < 0 && sequence.length() > (size_t) -keep_only) {
    end_indexing = -keep_only;
  }
  for(size_t i = start_indexing ; i + s < end_indexing + 1 ; i++) {
231 232
    seqtype substr = sequence.substr(i, s);
    seqtype kmer = spaced(substr, seed);
233

234
    if (ignore_extended_nucleotides && has_extended_nucleotides(kmer))
235 236
      continue;

Mikaël Salson's avatar
Mikaël Salson committed
237 238
    int strand = 1;
    if (revcomp_indexed && T::hasRevcompSymetry()) {
239
      seqtype rc_kmer = revcomp(kmer);
Mikaël Salson's avatar
Mikaël Salson committed
240 241 242 243 244
      if (rc_kmer.compare(kmer) < 0) {
        strand = -1;
        kmer = rc_kmer;
      }
    }
245 246 247 248 249
    T &this_kmer = this->get(kmer);
    if (this_kmer.isNull()) {
      nb_kmers_inserted++;
    }
    this_kmer += T(label, strand);
Mikaël Salson's avatar
Mikaël Salson committed
250
    if (revcomp_indexed && ! T::hasRevcompSymetry()) {
251
      seqtype rc_kmer = spaced(revcomp(substr), seed);
252 253 254 255
      T &this_rc_kmer = this->get(rc_kmer);
      if (this_rc_kmer.isNull())
        nb_kmers_inserted++;
      this_rc_kmer += T(label, -1);
Mikaël Salson's avatar
Mikaël Salson committed
256 257 258
    }
  }
}
259 260 261 262 263

template<class T>
float IKmerStore<T>::getIndexLoad() const {
  return nb_kmers_inserted*1. / (1 << (2 * k));
}
Mikaël Salson's avatar
Mikaël Salson committed
264

265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286

template<class T>
double IKmerStore<T>::getProbabilityAtLeastOrAbove(int at_least, int length) const {

  // n: number of kmers in the sequence
  int n = length - getS() + 1;
  float index_load = getIndexLoad() ;

  double proba = 0;
  double probability_having_system = pow(index_load, at_least);
  double probability_not_having_system = pow(1 - index_load, n - at_least);
  for (int i=at_least; i<=n; i++) {
    proba += nChoosek(n, i) * probability_having_system * probability_not_having_system;
    probability_having_system *= index_load;
    probability_not_having_system /= (1 - index_load);
  }

  return proba;
}



Mikaël Salson's avatar
Mikaël Salson committed
287 288 289 290 291 292 293 294 295 296
template<class T>
int IKmerStore<T>::getK() const {
  return k;
}

template<class T>
int IKmerStore<T>::getS() const {
  return s;
}

297 298 299 300 301
template<class T>
string IKmerStore<T>::getSeed() const {
  return seed;
}

302 303 304 305 306 307 308
template<class T>
string IKmerStore<T>::getLabel(T kmer) const {
  for (typename list< pair<T, string> >::const_iterator it = labels.begin(); it != labels.end(); ++it)
    if (it->first == kmer)
      return it->second ;
  return "" ;
}
309 310

// .getResults()
Mikaël Salson's avatar
Mikaël Salson committed
311
template<class T>
312 313 314 315
vector<T> MapKmerStore<T>::getResults(const seqtype &seq, bool no_revcomp) {

  int s = IKmerStore<T>::getS();

Mikaël Salson's avatar
Mikaël Salson committed
316 317 318 319 320
  if ((int)seq.length() < s - 1) {
    return vector<T>(0);
  }
  vector<T> result(seq.length() - s + 1);
  for (size_t i=0; i + s < seq.length() + 1; i++) {
321
    seqtype kmer = spaced(seq.substr(i, s), IKmerStore<T>::seed);
322
    //    seqtype kmer = seq.substr(i, s);
Mikaël Salson's avatar
Mikaël Salson committed
323
    // cout << kmer << endl << kmer0 << endl << endl ;
324
    if (IKmerStore<T>::revcomp_indexed && no_revcomp) {
Mikaël Salson's avatar
Mikaël Salson committed
325 326 327 328 329 330 331 332
      result[i] = get(kmer);
    } else {
      result[i] = (*this)[kmer];
    }
  }
  return result;
}

333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348
template<class T>
vector<T> ArrayKmerStore<T>::getResults(const seqtype &seq, bool no_revcomp) {
  
  int s = IKmerStore<T>::getS();

  int N = (int)seq.length();

  if (N < s - 1) {
    return vector<T>(0);
  }
  vector<T> result(N - s + 1);

  /* Read once the sequence, convert it to int* */
  int* intseq = new int[N];
  for (int i=0; i<N; i++)
    {
349
      intseq[i] = nuc_to_int(seq[i]);
350 351 352 353 354 355 356 357 358
    }

  /* Compute results */
  for (size_t i=0; (int) i+s < N+1; i++) {
    int kmer = spaced_int(intseq + i, IKmerStore<T>::seed);
    if (IKmerStore<T>::revcomp_indexed && no_revcomp) {
      result[i] = store[kmer]; // getfromint(kmer); // store[kmer];
      // cout << i << "/" << N << "  " << kmer << result[i] << endl ;
    } else {
359
      result[i] = (*this)[kmer]; // Deals with revcomp
360 361 362 363 364 365 366
    }
  }

  delete[] intseq ;
  return result;
}

Mikaël Salson's avatar
Mikaël Salson committed
367 368 369 370 371 372 373 374 375 376 377 378 379 380 381

template<class T>
bool IKmerStore<T>::isRevcomp() const {
  return revcomp_indexed;
}

// MapKmerStore

template <class T>
MapKmerStore<T>::MapKmerStore(string seed, bool revcomp){
  this->seed = seed;   
  int k = seed_weight(seed);
  this->k = k;
  this->s = seed.size();
  this->revcomp_indexed = revcomp;
382
  init();
Mikaël Salson's avatar
Mikaël Salson committed
383 384 385 386 387 388 389 390
}

template <class T>
MapKmerStore<T>::MapKmerStore(int k, bool revcomp){
  this->seed = seed_contiguous(k);
  this->k = k;
  this->s = k;
  this->revcomp_indexed = revcomp;
391 392 393 394 395 396
  init();
}

template <class T>
void MapKmerStore<T>::init() {
  this->nb_kmers_inserted = 0;
Mikaël Salson's avatar
Mikaël Salson committed
397 398 399
}

template <class T> 
400
T& MapKmerStore<T>::get(seqtype& word){
Mikaël Salson's avatar
Mikaël Salson committed
401 402 403 404
  return store[word];
}

template <class T> 
405
T& MapKmerStore<T>::operator[](seqtype& word){
Mikaël Salson's avatar
Mikaël Salson committed
406
  if (this->isRevcomp() && T::hasRevcompSymetry()) {
407
    seqtype rc_kmer = revcomp(word);
Mikaël Salson's avatar
Mikaël Salson committed
408 409 410 411 412 413 414 415 416
    if (rc_kmer.compare(word) < 0)
      word = rc_kmer;
  }
  return store[word];
}

// ArrayKmerStore

template <class T> 
417
ArrayKmerStore<T>::ArrayKmerStore(int k, bool revcomp) {
Mikaël Salson's avatar
Mikaël Salson committed
418 419 420 421
  this->seed = seed_contiguous(k);
  this->k = k;
  this->s = k;
  this->revcomp_indexed = revcomp;
422
  init();
Mikaël Salson's avatar
Mikaël Salson committed
423 424 425 426 427 428 429 430 431 432
}


template <class T> 
ArrayKmerStore<T>::ArrayKmerStore(string seed, bool revcomp){
  this->seed = seed; 
  int k = seed_weight(seed);
  this->k = k;
  this->s = seed.size();
  this->revcomp_indexed = revcomp;
433 434 435 436 437
  init();
}

template <class T>
void ArrayKmerStore<T>::init() {
438
  this->nb_kmers_inserted = 0;
439
  if ((size_t)(this->k << 1) >= sizeof(int) * 8)
Mikaël Salson's avatar
Mikaël Salson committed
440
    throw std::bad_alloc();
441
  store = new T[(unsigned int)1 << (this->k << 1)];
Mikaël Salson's avatar
Mikaël Salson committed
442 443 444 445 446 447 448 449 450 451 452 453 454 455
  if (! store)
    throw std::bad_alloc();
}

template <class T> 
ArrayKmerStore<T>::~ArrayKmerStore(){
  delete [] store;
}

/**	Returns index of word in an array of size 4^k
	Considering B(A) = 0, B(C) = 1, B(T) = 2, B(G) = 3
	index_word = \sum_{i=0}^{k-1}B(word[i])*4^(k-i-1)
**/
template <class T> 
456
int ArrayKmerStore<T>::index(const seqtype& word) const{
Mikaël Salson's avatar
Mikaël Salson committed
457
  return dna_to_int(word, this->k);
Mikaël Salson's avatar
Mikaël Salson committed
458 459 460
}

template <class T> 
461
T& ArrayKmerStore<T>::get(seqtype& word){
Mikaël Salson's avatar
Mikaël Salson committed
462 463 464 465
  return store[index(word)];
}

template <class T> 
466
T& ArrayKmerStore<T>::operator[](seqtype& word){
467 468 469 470 471
  return (*this)[index(word)];
}

template <class T> 
T& ArrayKmerStore<T>::operator[](int word){
Mikaël Salson's avatar
Mikaël Salson committed
472
  if (this->isRevcomp() && T::hasRevcompSymetry()) {
473 474
    int rc_kmer = revcomp_int(word, IKmerStore<T>::k);
    if (rc_kmer < word)
Mikaël Salson's avatar
Mikaël Salson committed
475 476
      word = rc_kmer;
  }
477
  return store[word];
Mikaël Salson's avatar
Mikaël Salson committed
478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498
}


/**
 * KmerStoreFactory is a factory that allows to create an index that best fits
 * your needs!
 */
class KmerStoreFactory {
 public:
  template<class T>
    static IKmerStore<T> *createIndex(string seed, bool revcomp=false);
  template<class T>
    static IKmerStore<T> *createIndex(int k, bool revcomp=false);
};

template<class T>
IKmerStore<T> *KmerStoreFactory::createIndex(string seed, bool revcomp) {
  IKmerStore<T> *index;
  try{
    index = new ArrayKmerStore<T>(seed, revcomp);
  }catch(exception e){
499
    cout << "  (using a MapKmer to fit into memory)" << endl;
Mikaël Salson's avatar
Mikaël Salson committed
500 501 502 503 504 505 506 507 508 509 510
    index = new MapKmerStore<T>(seed, revcomp);
  }
  return index;
}

template<class T>
IKmerStore<T> *KmerStoreFactory::createIndex(int k, bool revcomp) {
  return createIndex<T>(seed_contiguous(k), revcomp);
}

#endif