enrichment.h

/**************************************************************************
 * Copyright (c) 2004-2011 T. M. Murali                                   *
 *                                                                        *
 * This file is part of Biorithm.                                         *
 *                                                                        *
 * Biorithm is free software: you can redistribute it and/or modify       *
 * it under the terms of the GNU General Public License as published by   *
 * the Free Software Foundation, either version 3 of the License, or      *
 * (at your option) any later version.                                    *
 *                                                                        *
 * Biorithm 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 Public License for more details.                           *
 *                                                                        *
 * You should have received a copy of the GNU General Public License      *
 * along with Biorithm.  If not, see <http://www.gnu.org/licenses/>.      *
 *                                                                        *
 **************************************************************************/

#ifndef _ENRICHMENT_H_
#define _ENRICHMENT_H_

// for tolower.
#include <cctype>
#include <iostream>
#include <string>
#include <sstream>
#include <map>
#include <set>
#include <vector>
#include <cmath>
#include <algorithm>
#include <utility>
#include <cstdlib>

using namespace std;

// copying definition in murali-xmotif/global.h since histogram.* are
// now in libutil and histogram.h includes enrichment.h.
typedef double MyNT;


enum LIBENRICHMENT_TEST_TYPE {LIBENRICHMENT_NONE, LIBENRICHMENT_BONFERRONI, LIBENRICHMENT_HOLMS, LIBENRICHMENT_FALSE_DISCOVERY_RATE};

//map< LIBENRICHMENT_TEST_TYPE, string > _testTypeStrings;

inline string convertTestTypeToString(LIBENRICHMENT_TEST_TYPE test)
{
  if (LIBENRICHMENT_NONE == test)
    return("no");
  else if (LIBENRICHMENT_BONFERRONI == test)
    return("Bonferroni");
  else if (LIBENRICHMENT_HOLMS == test)
    return("Holms");
  else if (LIBENRICHMENT_FALSE_DISCOVERY_RATE == test)
    return("False discovery rate");
  else
    return("");
//  return(_testTypeStrings[test]);
}

inline string convertTestTypeToShortString(LIBENRICHMENT_TEST_TYPE test)
{
  if (LIBENRICHMENT_NONE == test)
    return("no");
  else if (LIBENRICHMENT_BONFERRONI == test)
    return("Bonferroni");
  else if (LIBENRICHMENT_HOLMS == test)
    return("Holms");
  else if (LIBENRICHMENT_FALSE_DISCOVERY_RATE == test)
    return("FDR");
  else
    return("");
//  return(_testTypeStrings[test]);
}

inline LIBENRICHMENT_TEST_TYPE  convertStringToTestType(string test)
{
//  test = tolower(test);
  if ("" == test)
    return(LIBENRICHMENT_NONE);
  else if ("Bonferroni" == test)
    return(LIBENRICHMENT_BONFERRONI);
  else if ("Holms" == test)
    return(LIBENRICHMENT_HOLMS);
  else if ("FDR" == test)
    return(LIBENRICHMENT_FALSE_DISCOVERY_RATE);
  else
    {
      cerr << "Unknown enrichment test \"" << test << "\". Will not do any correction for multiple hypothesis testing." << endl;
      return(LIBENRICHMENT_NONE);
    }
  
//  return(_testTypeStrings[test]);
}

template <class S,class T> struct EnrichmentRecord
{
  double score;
  T relation;
  int gTotal;
  int gTrue;
  int lTotal;
  int lTrue;
  set<S> annotated;

  string itoa(int in)
    {
      stringstream ss;
      ss<<in;
      string out;
      ss>>out;
      return out;
        }
  void print(ostream &fstr)
    {
      // local is the cluster. global is the universe.
      fstr << lTrue << "\t" << lTotal << "\t" << gTrue << "\t" << gTotal;
    }

  string reason()
        {
          string empty;
          return empty+itoa(lTrue)+"/"+itoa(lTotal)+" vs. "+itoa(gTrue)+"/"+itoa(gTotal);
        }
  const bool operator<(const EnrichmentRecord<S,T> &comp) const
    {
      if(comp.score!=score)
        return score<comp.score;
      return relation<comp.relation;
    }

  const bool operator==(const EnrichmentRecord<S,T> &comp) const
    {
      return (comp.score==score && relation==comp.relation);
    }
};

// comparator for EnrichmentRecords.
// struct compareEnrichmentRecords :
//   public binary_function< EnrichmentRecord<class S,class T>,
//                           EnrichmentRecord<class S,class T>, bool> 
// {
//   bool operator()(const EnrichmentRecord<class S,class T> &__x,
//                   const EnrichmentRecord<class S,class T> & __y) const
//     {
//       return __x.score < __y.score; }
// };



template <class S,class T> class Enrichment {
public:
  Enrichment();

  void clear();
  
  double computeHyperGeometricProbability(int globTotal, int globTrue, int locTotal, int locTrue);

  void addUnannotated(int numS,int numT);

  bool check(double value, LIBENRICHMENT_TEST_TYPE test, double alpha, unsigned int numTests, unsigned int rank);

  double correct(double value, LIBENRICHMENT_TEST_TYPE test, double alpha, unsigned int numTests, unsigned int rank);

  void check(vector<EnrichmentRecord<S,T> > &in, vector<EnrichmentRecord<S,T> > &out,
             LIBENRICHMENT_TEST_TYPE test, double alpha);
                

  void getEnrichments(set<S> &in, vector<EnrichmentRecord<S,T> > &out);

  void getEnrichmentsGSEA(set<S> &in, vector<EnrichmentRecord<S,T> > &out);

  bool isPair(const S &s, const T &t) const;

  void loadPairs(const vector<pair<S,T> > &in);

  void loadPairs(const map< S, set<T> > &in);

private:
  double HyperGeoHelper(int globTotal, int globTrue, int locTotal, int locTrue);
  double factorial(int);
  double factorialLog(int);
  double factorialEstimate(int);
  double factorialEstimateLog(int);
  map<S,set<T> > sXt;
  map<T,set<S> > tXs;
  int unS,unT;
  string itoa(int in)
    {
      stringstream ss;
      ss<<in;
      string out;
      ss>>out;
      return out;
    }
  vector<double> logFactCache;
};

#ifndef PI_DEFINED
#define PI_DEFINED 1
const double PI = 4*atan2(1.0, 1.0);
//#define PI 3.14159265
#endif //

template <class S,class T>
Enrichment<S,T>::Enrichment()
    : sXt(), tXs(), unS(0), unT(0)
{
        
}

template <class S,class T>
void Enrichment<S,T>::clear()
{
  sXt.clear();
  tXs.clear();
  unS=unT=0;
}

template <class S,class T>
bool Enrichment<S,T>::isPair(const S &s, const T &t) const
{
  typename map<S,set<T> >::const_iterator sitr;
  typename set< T >::const_iterator titr;
  sitr = sXt.find(s);
  if (sXt.end() == sitr)
    return(false);
  titr = sitr->second.find(t);
  if (sitr->second.end() == titr)
    return(false);
  return(true);
}


template <class S,class T>
void Enrichment<S,T>::loadPairs(const vector<pair<S,T> > &in)
{
        for(int i=0;i<in.size();i++)
        {
                sXt[in[i].first].insert(in[i].second);
                tXs[in[i].second].insert(in[i].first);
        }
}

template <class S,class T>
void Enrichment<S,T>::loadPairs(const map < S, set<T> > &in)
{
  sXt = in;
  for (typename map< S, set<T> >::const_iterator itr = in.begin(); itr != in.end(); itr++)
    {
      for (typename set<T>::const_iterator titr = itr->second.begin(); titr != itr->second.end(); titr++)
        tXs[*titr].insert(itr->first);
    }
}

template <class S,class T>
void Enrichment<S,T>::addUnannotated(int a,int b)
{
        unS=a;
        unT=b;
}

template <class S,class T>
void Enrichment<S,T>::getEnrichments(set<S> &in, vector<EnrichmentRecord<S,T> > &out)
{
  map<T,int> counts;
  map<T,set<S> > annot;
//  vector<EnrichmentRecord<S,T> > out;
  int locTotal=0;
  for(typename set<S>::iterator i=in.begin();i!=in.end();i++)
    // loop over every element (gene) in in.
    {
      if(sXt.find(*i)!=sXt.end())
        // *i has an annotation of type T.
        {
//           typename set<T>::iterator begin=sXt[*i].begin();
//           typename set<T>::iterator end=sXt[*i].end();
          for(typename set<T>::iterator j = sXt[*i].begin(); j != sXt[*i].end(); j++)
            {
              // increment count for type T (function) element *j.
              counts[*j]++;
              // keep in mind that *j annotates *i.
              annot[*j].insert(*i);
            }
          // increment #elements (genes) in in that have an annotation of type T.
          locTotal++;
        }
    }
  for(typename map<T,int>::iterator i=counts.begin();i!=counts.end();i++)
    {
      int locTrue=i->second;
      int globTrue=tXs[i->first].size();
      int globTotal=sXt.size()+unS;
      double value= computeHyperGeometricProbability(globTotal, globTrue, locTotal, locTrue);
      if(value<=1)
        {
          EnrichmentRecord<S,T> e;
          e.score=value;
          e.gTotal=globTotal;
          e.gTrue=globTrue;
          e.lTotal=locTotal;
          e.lTrue=locTrue;
          e.relation=i->first;
          e.annotated=annot[i->first];
          out.push_back(e);
        }
    }
  sort(out.begin(),out.end());//, compareEnrichmentRecords);
//  return out;         
}


template <class S,class T>
void Enrichment<S,T>::getEnrichmentsGSEA(set<S> &in,
                                         // sorted order 

                                         vector<EnrichmentRecord<S,T> > &out)
{
  map<T,int> counts;
  map<T,set<S> > annot;
//  vector<EnrichmentRecord<S,T> > out;
  int locTotal=0;
  for(typename set<S>::iterator i=in.begin();i!=in.end();i++)
    // loop over every element (gene) in in.
    {
      if(sXt.find(*i)!=sXt.end())
        // *i has an annotation of type T.
        {
//           typename set<T>::iterator begin=sXt[*i].begin();
//           typename set<T>::iterator end=sXt[*i].end();
          for(typename set<T>::iterator j = sXt[*i].begin(); j != sXt[*i].end(); j++)
            {
              // increment count for type T (function) element *j.
              counts[*j]++;
              // keep in mind that *j annotates *i.
              annot[*j].insert(*i);
            }
          // increment #elements (genes) in in that have an annotation of type T.
          locTotal++;
        }
    }
  for(typename map<T,int>::iterator i=counts.begin();i!=counts.end();i++)
    {
      int locTrue=i->second;
      int globTrue=tXs[i->first].size();
      int globTotal=sXt.size()+unS;
      double value= computeHyperGeometricProbability(globTotal, globTrue, locTotal, locTrue);
      if(value<=1)
        {
          EnrichmentRecord<S,T> e;
          e.score=value;
          e.gTotal=globTotal;
          e.gTrue=globTrue;
          e.lTotal=locTotal;
          e.lTrue=locTrue;
          e.relation=i->first;
          e.annotated=annot[i->first];
          out.push_back(e);
        }
    }
  sort(out.begin(),out.end());//, compareEnrichmentRecords);
//  return out;         
}


template <class S,class T>
bool
Enrichment<S,T>::check(double value, LIBENRICHMENT_TEST_TYPE test, double alpha, unsigned int numTests,
                         unsigned int rank)
{
  double correctedValue = correct(value, test, alpha, numTests, rank);
  return(correctedValue <= alpha);
}

template <class S,class T>
double
Enrichment<S,T>::correct(double value, LIBENRICHMENT_TEST_TYPE test, double alpha, unsigned int numTests,
                         unsigned int rank)
{
  double correctedValue = value;
  switch(test)
    {
      // don't change alpha! keep alpha the same but change the
      // score/pvalue.
      case LIBENRICHMENT_NONE:
//            myAlpha=alpha;
        break;
      case LIBENRICHMENT_BONFERRONI:
//            myAlpha=alpha/n;
        correctedValue = numTests*correctedValue;
            break;
          case LIBENRICHMENT_HOLMS:
//            myAlpha=alpha/(n-i2+1);
            correctedValue = (numTests - rank + 1)*correctedValue;
            break;
          case LIBENRICHMENT_FALSE_DISCOVERY_RATE:
//            myAlpha=(alpha*i2)/n;
            correctedValue = (numTests*1.0/rank)*correctedValue;
            break;
        }
  return(correctedValue);
}


template <class S,class T>
void Enrichment<S,T>::check(vector<EnrichmentRecord<S,T> > &in,
                            vector<EnrichmentRecord<S,T> > &out,
                            LIBENRICHMENT_TEST_TYPE test, double alpha)
{
//  vector<EnrichmentRecord<S,T> > out;
  unsigned int n=in.size();
  for(unsigned int i=0;i<n;i++)
    {
      double myAlpha=alpha;
      unsigned int i2= i + 1;
      in[i].score = correct(in[i].score, test, alpha, n, i2);
      if(in[i].score>myAlpha)
        break;
      out.push_back(in[i]);
    }
//  return out;
}       

template <class S,class T>
double Enrichment<S,T>::factorial( int n ){
  double result = 1;
  for(int i=n; i > 1; i-- )
    result*=i;
  return( result );
}

template <class S,class T>
double Enrichment<S,T>::factorialLog(int n){
  logFactCache.resize(n + 1);
  // while loop is just resizing!
//   while(logFactCache.size()<n+1)
//     logFactCache.push_back(0);
  if(logFactCache[n]!=0)
    return logFactCache[n];
  
  double result=0;
  for(int i=1;i<=n;i++)
    result+=log(double(i));
  logFactCache[n]=result;
  return result;
}

template <class S,class T>
double Enrichment<S,T>::factorialEstimate( int n ){
  //Note: This method computes an approximation of the factorial using
  //Stirling's formula.
  if( n <= 0 )
    return( 1 );
  // be exact if n is small, since the Stirling approximation is
  // not good for small n.
  if (n <= 20){ return (factorial(n)); }
  double result = sqrt(2*PI*n)*pow(n/exp((double)1),n);
  return( result );
}

template <class S,class T>
double Enrichment<S,T>::factorialEstimateLog( int n ){
        if( n <= 0 ){ return 0; }
        // be exact if n is small, since the Stirling approximation is
        // not good for small n.
        if (n <= 20){ return (factorialLog(n)); }
        // exp(x) is e^x, so the argument to the second log is (n/e).
        double result = 0.5*log(2*PI*n) + n*log(n/exp((double)1));
        return( result );
}
template <class S,class T>
double Enrichment<S,T>::HyperGeoHelper( int globTotal, int globTrue, int locTotal, int locTrue ){
        // double result = 
        //      factorialEstimateLog(globTrue) + \
        //      factorialEstimateLog(globTotal - globTrue) + \
        //      factorialEstimateLog(locTotal) + \
        //      factorialEstimateLog(globTotal-locTotal) - \
        //      factorialEstimateLog(locTrue) - \
        //      factorialEstimateLog(globTrue-locTrue) - \
        //      factorialEstimateLog(globTotal-globTrue+locTrue-locTotal) - \
        //      factorialEstimateLog(locTotal - locTrue) - \
        //      factorialEstimateLog(globTotal);

        // I can be sure that locTrue <= globTrue and locTrue <=
        // locTotal, if I called this method from Enrichment::computeHyperGeometricProbability().
        double result = 
          // start with choose(locTotal, locTrue)
          factorialEstimateLog(locTotal) 
          - factorialEstimateLog(locTrue) 
          - factorialEstimateLog(locTotal - locTrue) 
          // multiply by choose(globTotal-locTotal, globTrue-locTrue)
          + factorialEstimateLog(globTotal-locTotal) 
          - factorialEstimateLog(globTrue-locTrue) 
          - factorialEstimateLog(globTotal-globTrue+locTrue-locTotal) 
          // divide by choose(globTotal globTrue)
          + factorialEstimateLog(globTrue)
          + factorialEstimateLog(globTotal - globTrue) 
          - factorialEstimateLog(globTotal)
;

        return exp(result);
}
template <class S,class T> double Enrichment<S,T>::computeHyperGeometricProbability(int globTotal, int globTrue,
                                                                                    int locTotal, int locTrue){
        double sum=0;
        for(int i=locTrue; i<=locTotal && i<=globTrue; i++ )
                sum+=HyperGeoHelper(globTotal,globTrue,locTotal,i);
        return sum;
}

#endif