+ fixed alignment index calculation for negative strand
[qpalma.git] / qpalma / DatasetUtils.py
index 09d241c..b56df4a 100644 (file)
@@ -9,50 +9,38 @@
 import array
 import cPickle
 import os
+import os.path
 import pdb
 
-#from parsers import *
-from sequence_utils import SeqSpliceInfo,DataAccessWrapper,get_flat_file_size,reverse_complement,unbracket_seq,create_bracket_seq,reconstruct_dna_seq
+from sequence_utils import SeqSpliceInfo,DataAccessWrapper,get_flat_file_size,unbracket_seq,create_bracket_seq,reconstruct_dna_seq
 
+jp = os.path.join
 
-def generateDataset(global_settings):
-   """
-   This function is responsible for the dataset generation for the matches of
-   the second vmatch run and the heuristic-filtered matches of the first run.
+illumina_ga_range = (-5,40)
+#roche_454_range   =
 
-   An example in our dataset consists of:
 
-   - information on the original read:
-         * id
-         * nucleotide sequence
-         * quality vectors (at the moment: prb and chastity)
+def generateTrainingDataset(settings):
+   """
+   This function creates a training dataset.
+   """
+   dataset = []
 
-   - information on the target dna sequence:
-         * chromosome
-         * strand
-         * fragment start 
-         * fragment end positions
+   saveData('training',dataset,settings)
 
-   this information is stored in tuples which are stored then in a dictionary
-   indexed by the reads unique id.
-
-   CAUTION: The dictionary stores a list of the above defined tuples under each
-   'id'. This is because one read can be found several times on the genome.
-   """
 
+def addToDataset(map_file,dataset,settings):
    assert os.path.exists(map_file), 'Error: Can not find map file'
-   print map_2nd_file
-   assert os.path.exists(map_2nd_file), 'Error: Can not find map_2nd file'
-   self.map_file = map_file
-   self.map_2nd_file = map_2nd_file
 
-   dataset = {}
+   prb_offset = settings['prb_offset']
 
-   prb_offset = 64
+   # This tuple specifies an interval for valid Illumina Genome Analyzer quality values
+   if settings['platform'] == 'IGA':
+      quality_interval = illumina_ga_range
 
    # this means that around each vmatch hit we cut out a window of 1500 bases
    # each up-/downstream.
-   half_window_size = 1500
+   half_window_size = settings['half_window_size']
 
    instance_counter = 0
 
@@ -60,6 +48,10 @@ def generateDataset(global_settings):
    seqInfo = SeqSpliceInfo(accessWrapper,settings['allowed_fragments'])
 
    for line in open(map_file):
+      line = line.strip()
+      if line.startswith('#') or line == '':
+         continue
+
       if instance_counter > 0 and instance_counter % 5000 == 0:
          print 'processed %d examples' % instance_counter
 
@@ -70,7 +62,7 @@ def generateDataset(global_settings):
       pos      = int(slist[2])
 
       # for the time being we only consider chromosomes 1-5
-      if not chromo in range(1,6):
+      if not chromo in settings['allowed_fragments']:
          continue
 
       # convert D/P strand info to +/-
@@ -79,46 +71,36 @@ def generateDataset(global_settings):
       # QPalma uses lowercase characters
       bracket_seq = slist[4].lower()
       read_seq = unbracket_seq(bracket_seq)
-      reconstructed_dna_seq = reconstruct_dna_seq(bracket_seq)
-
-      if strand == '-' and first_round:
-         read_seq = reverse_complement(read_seq)
-         #reconstructed_dna_seq = reverse_complement(reconstructed_dna_seq)
 
+      # in prediction we do not have any bracketed reads anymore
       assert not '[' in read_seq and not ']' in read_seq
-      
+
       # we use an area of +/-  `self.half_window_size` nucleotides around the seed position
-      if pos > self.half_window_size+1:
-         us_offset = self.half_window_size
+      if pos > half_window_size+1:
+         us_offset = half_window_size
       else:
          us_offset = pos - 1 
 
-      if pos+self.half_window_size < seqInfo.chromo_sizes[chromo]:
-         ds_offset = self.half_window_size
+      if pos+half_window_size < seqInfo.getFragmentSize(chromo):
+         ds_offset = half_window_size
       else:
-         ds_offset = seqInfo.chromo_sizes[chromo]-pos-1
+         ds_offset = seqInfo.getFragmentSize(chromo)-pos-1
          
       genomicSeq_start = pos - us_offset
       genomicSeq_stop  = pos + ds_offset
 
       dna_seq = seqInfo.get_seq_and_scores(chromo,strand,genomicSeq_start,genomicSeq_stop,True)
       
-      #if strand == '-':
-      #   dna_part = dna_seq[us_offset-len(read_seq):us_offset]
-      #else:
-      #   dna_part = dna_seq[us_offset:us_offset+len(read_seq)]
-
-      ## check whether everything is ok
-      #   try:
-      #      currentDNASeq, currentAcc, currentDon = seqInfo.get_seq_and_scores(chromo,strand,genomicSeq_start,genomicSeq_stop)
-      #   except:
-      #      problem_ctr += 1
-      #      continue
-
       # In order to save some space we use a signed char to store the
       # qualities. Each quality element can range as follows: -128 <= elem <= 127
       
-      prb = array.array('b',map(lambda x: ord(x)-self.prb_offset,slist[5]))
+      q_values = map(lambda x: ord(x)-prb_offset,slist[5])
+
+      if settings['perform_checks']:
+         for entry in q_values:
+            assert quality_interval[0] <= entry <= quality_interval[1], 'Error: Your read has invalid quality values: %d' % entry
+
+      prb = array.array('b',q_values)
 
       currentSeqInfo = (id,chromo,strand,genomicSeq_start,genomicSeq_stop)
       currentQualities = [prb]
@@ -131,12 +113,60 @@ def generateDataset(global_settings):
 
    print 'Total examples processed: %d' % instance_counter
 
-   ddir = global_settings['dataset_dir']
-   dataset_fn        = jp(ddir,'data.pickle')
-   dataset_keys_fn   = jp(ddir,'data.keys.pickle')
+   return dataset
+
+
+def generatePredictionDataset(settings):
+   """
+   This function is responsible for the prediction dataset generation for the matches of
+   the second vmatch run and the heuristic-filtered matches of the first run.
+
+   An example in our dataset consists of:
+
+   - information on the original read:
+         * id
+         * nucleotide sequence
+         * quality vectors (at the moment: prb and chastity)
+
+   - information on the target dna sequence:
+         * chromosome
+         * strand
+         * fragment start 
+         * fragment end positions
+
+   this information is stored in tuples which are stored then in a dictionary
+   indexed by the reads unique id.
+
+   CAUTION: The dictionary stores a list of the above defined tuples under each
+   'id'. This is because one read can be found several times on the genome.
+   """
+
+   dataset = {}
+
+   map_file = settings['spliced_reads_fn']
+   dataset = addToDataset(map_file,dataset,settings)
+
+   map_file = jp(settings['approximation_dir'],'map.vm.spliced')
+   dataset = addToDataset(map_file,dataset,settings)
+
+   saveData('prediction',dataset,settings)
+
+
+def saveData(prefix,dataset,settings):
+   """
+   """
 
-   assert not os.path.exists(dataset_fn), 'The data_file already exists!'
-   assert not os.path.exists(dataset_keys_fn), 'The data_keys file already exists!'
+   if prefix == 'prediction':
+      dataset_fn        = settings['prediction_dataset_fn']
+      dataset_keys_fn   = settings['prediction_dataset_keys_fn']
+   elif prefix == 'training':
+      dataset_fn        = settings['training_dataset_fn']
+      dataset_keys_fn   = settings['training_dataset_keys_fn']
+   else:
+      assert False
+
+   #assert not os.path.exists(dataset_fn), 'The data file already exists!'
+   #assert not os.path.exists(dataset_keys_fn), 'The data keys file already exists!'
 
    # saving new dataset and single keys as well
    cPickle.dump(dataset,open(dataset_fn,'w+'),protocol=2)