text: import numpy as np
  import pandas as pd
  from numpy.random import permutation
  import time
  import logging
  
  try:
      import matplotlib.pyplot as plt
      import matplotlib
  except ImportError:
      print('Imported epitope_hotspot without matplotlib.')
  
  
  import statsmodels.api as sm
  from HLAPredCache import grabKmer, convertHLAAsterisk
  from objhist import objhist
  
  __all__ = ['getBAMat',
             'get2DBAMat',
             'populationBindingMap',
             'plotBindingGrid',
             'BAMat2map',
             'computeHLAHotspots',
             'plotKmerHotpots',
             'plotSiteHotpots',
             'plotAlignmentBindingGrid']
  
  _MAX_IC50 = 15
  
  """Set up logger for printing progress of long-running simulations and analysis"""
  logger = logging.getLogger('epitope_hotspot')
  logger.setLevel(logging.INFO)
  formatter = logging.Formatter('%(levelname)s:%(asctime)s:%(message)s')
  
  """Create console handler and set level to INFO"""
  ch = logging.StreamHandler()
  ch.setLevel(logging.INFO)
  ch.setFormatter(formatter)
  logger.addHandler(ch)
  
  def get2DBAMat(peptides, hlas, ba):
      """Generate a pd.DataFrame of log(IC50) binding affinities
      with index of peptides and columns of HLA alleles.
  
      If HLA alleles are keys in dict then they are sorted() to generate the matrix.
  
      Some predictions may be nan if they are not found in ba.
      
      Parameters
      ----------
      peptides : list/iterator
          Peptides whose HLA binding affinities should be in ba.
      hlas : dictionary or list of HLA alleles
          HLA alleles as keys in dict or elements in list
      ba : hlaPredCache
          Should contain log(IC50) for all HLA allele and peptide combinations.
  
      Returns
      -------
      bamat : pd.DatFrame
          Containing all HLA:peptide pairs of log(IC50) predictions with shape [index: hlaList x columns: peptides]
      """
      try:
          hlaList = sorted(hlas.keys())
      except AttributeError:
          hlaList = hlas
  
      hlaList = convertHLAAsterisk(hlaList)
  
      bamat = _MAX_IC50 * np.ones((len(hlaList), len(peptides)))
      for hlai, h in enumerate(hlaList):
          for pepi, pep in enumerate(peptides):
              bamat[hlai, pepi] = ba[(h, pep)]
      baDf = pd.DataFrame(bamat, index = hlaList, columns = peptides)
      return baDf
  
  def getBAMat(seqs, hlas, ba, k=9, gapless=False):
      """Generate a 3D np.ndarray of log(IC50) binding affinities
      with shape [nKmers x nSeqs x nHLAs]
  
      If HLA alleles are keys in dict then they are sorted() to generate the matrix.
  
      Some predictions may be nan if they are not found in ba.
  
      By default all kmers containing a gap have Nan affinity.
      
      Parameters
      ----------
      seqs : list/iterator
          Aligned AA sequences.
      hlas : dictionary or list of HLA alleles
          HLA alleles as keys in dict or elements in list
      ba : hlaPredCache
          Should contain log(IC50) for all HLA allele and kmer combinations.
      k : int
          Width of the peptides to use for HLA binding prediction.
      gapless : bool
          If True then skip gap characters when grabbing kmers from the alignment.
          Kmers starting with a gap are returned as nan.
          NOTE: This option will "break" the alignment by skipping gapped sites.
                Do not use if you intend to make an integrated site-based map.
  
      Returns
      -------
      bamat : np.ndarray
          Array of shape [nKmers x nSeqs x nHLAs]
      """
  
      """Do all seqs have the same length"""
      if len(np.unique(list(map(len, seqs)))) > 1:
          raise ValueError('Sequences must all have the same length (%s)' % np.unique(list(map(len, seqs))))
  
      nHLAs = len(hlas)
      nSeqs = len(seqs)
      nSites = int(np.median(list(map(len, seqs))))
      nKmers = nSites - k + 1
  
      try:
          hlaList = sorted(hlas.keys())
      except AttributeError:
          hlaList = hlas
  
      baMat = np.nan * np.ones((nKmers, nSeqs, nHLAs), dtype = np.float64)
  
      """Initialize the 3 dimensional binding affinity matrix"""
      for seqi, s in enumerate(seqs):
          for meri in range(nKmers):
              gappedMer, gaplessMer = grabKmer(s, meri, k=k)
              if gapless:
                  """Use the gapped 9mer to preserve the alignment precisely"""
                  mer = gaplessMer
              else:
                  mer = gappedMer
              
              if mer is None:
                  baMat[meri, seqi,:] = np.nan
              else:
                  for hlai, h in enumerate(hlaList):
                      baMat[meri, seqi, hlai] = ba[(h, mer)]
      return baMat
  
  def BAMat2map(mat, hlas, siteMap=True, k=9, phenotypicFrequencies=False):
      """Take a matrix containing HLA binding data organized by [kmers, seqs, hlas]
      and return a vector of length nKmers or nSites that is:
  
      (1) Weighted by and aggregated across HLA allele frequencies (optionally phenotypicFrequencies)
      (2) Summed over seqs
      (3) Optionall integrated within kmers (i.e. translated to sites)
  
      To work properly this should be matrix of 1/0's indicating binding for each HLA:kmer pair
      over the alignment. But if phenotypicFrequencies is False then its simply
      a weighted sum over HLAs, summed over seqs and optionally integrated ove kmers.
  
      If hlas is a list of alleles then it assumes uniform weights.
  
      Parameters
      ----------
      mat : ndarray
          Matrix of binding affinities (or values that have been operated on)
          with shape [nKmers x nSeqs x nHLAs]
      hlas : dict or list
          Allele frequencies as key/value pairs.
          A uniform distribution is assumed over all alleles if a list is given.
      siteMap : bool
          Return a site-wise map if True, otherwise return a kmer-based map.
          A site-wise map is integrated over kmers.
      k : int
          Width of the peptide for HLA binding affinity prediction
      phenotypicFrequencies : bool
          Convert HLA allele frequencies to phenotypic frequencies
          (sums homozygous and heterozygous probabilities,
          assuming no linkage disequilibrium with other allelels)
  
      Returns
      -------
      vec : ndarray shape [1,]
          Vector/map of HLA binding by kmer or optionally by site.
      """
  
      try:
          hlaList = sorted(hlas.keys())
      except AttributeError:
          logger.info('No HLA frequencies given, assuming uniform distribution.')
          hlaList = sorted(hlas)
          hlas = {h:1./len(hlaList) for h in hlaList}
      
      freqVec = np.array([hlas[h] for h in hlaList])
  
      nKmers, nSeqs, nHLAs = mat.shape
  
      """Multiply element-wise by hla allele frequencies before summing"""
      weightedMat = mat * np.tile(freqVec.reshape((1, 1, nHLAs)), (nKmers, nSeqs, 1))
  
      if phenotypicFrequencies:
          """Convert total allele frequency into phenotypic frequency"""
          locusFreq = np.zeros((2, nKmers, nSeqs))
          for loci, locus in enumerate('AB'):
              locusInd = [i for i, h in enumerate(hlaList) if h[0]==locus]
              tmpFreq = weightedMat[:,:, locusInd].sum(axis=2)
              locusFreq[loci,:,:] = tmpFreq**2 + 2*tmpFreq*(1-tmpFreq)
          hlaBindingFreq = locusFreq[0,:,:]*locusFreq[1,:,:] + locusFreq[0,:,:]*(1-locusFreq[1,:,:]) + locusFreq[1,:,:]*(1-locusFreq[0,:,:])
      else:
          """Compute the total frequency of the alleles that will bind."""
          hlaBindingFreq = weightedMat.sum(axis=2)
  
      """Sum across sequences"""
      vec = hlaBindingFreq.sum(axis=1)
  
      if siteMap:
          """Sum across sites for each kmer, towards a Pr(epitope) by site"""
          return np.convolve(vec, np.ones(k))
      else:
          return vec
  
  def populationBindingMap(align, hlas, ba, bindingThreshold=5, k=9, atLeastFracSeqs=0):
      """Compute the fraction of the population that will bind the kmer
      at each position, averaged across seqs.
  
      Parameters
      ----------
      align : list or pd.Series
          Set of amino acid sequences.
      hlas : dict
          HLA allele frequencies as key/value pairs (e.g. (A*0201:0.54))
          Allele frequencies should sum to 1 across all HLA-A or HLA-B alleles (separately).
      ba : hlaPredCache
          Dictionary of log binding affinity for all kmers in align and all HLAs in hlas
      bindingThreshold : float
          Units of logIC50 (or whatever is in ba)
      k : int
          Width of the mer for computing HLA binding predictions
      atLeastFracSeqs : float fraction
          Fraction of sequences required for each HLA to be considered a "binder" there.
  
      Returns
      -------
      vec : ndarray of shape[nKmers, 1]
          Vector with element for each start position
          indicating fraction of population binding the kmer, on average
      """
      """baMat is shape [nKmers x nSeqs x nHLAs]"""
      baMat = (getBAMat(align, hlas, ba) < bindingThreshold).astype(np.float64)
      
      if atLeastFracSeqs>0:
          """Squashes baMat into [nKmers x 1 x nHLAs] indicating
          if each HLA binds at least X pct of sequences"""
          baMat = (baMat.mean(axis=1)[:, None,:] > atLeastFracSeqs).astype(np.float64)
  
      """If baMat is not squashed then vec is an average over sequences"""
      emap = BAMat2map(baMat, hlas, siteMap=False, k=9, phenotypicFrequencies=True)/baMat.shape[1]
      return emap
  
  def computeHLAHotspots(seqs, hlas, ba, hlaMask=None, bindingThreshold=5, nPerms=100, k=9, returnDf=True):
      """Creates a map of epitope hotspots in the aligned sequences in seqs,
          based on the alleles and their frequencies in hlas e.g. {'A*0201':0.48, 'B*2701':0.35}
  
      Returns a pdf (sums to 1 across kmers) with the relative Pr that a kmer is an epitope,
          given the hlas, their frequencies and the sequences.
      Frequencies can be relative, but PDFs will be normalized as though all alleles and whole pop is represented.
  
      Computes Pr for each site by integrating across all overlapping kmers. 
  
      Also returns raw p-values for each kmer/site, based on the null hypothesis that
      the HLA predictions are random. Achieved by permuting the kmer dimension of the binding matrix.
      This means that under the null two related alleles will still have related binding affinities.
  
      Does not test against the null that HLA binding is uniform (although that also might be interesting).
  
      Parameters
      ----------
      seqs : pd.Series or any collection
          Aligned amino acid sequences.
      hlas : dict or list
          Dict of HLA alleles and population frequencies.
          If list of alleles then it assumes uniform frequencies.
      ba : hlaPredCache (dict-like)
          Dictionary of predicted binding affinities (e.g. ba[(h,pep)] = prediction)
      hlaMask : ndarray [seqs x alleles]
          Mask of alleles associated with each sequence (1 = expressing, 0 = not expressing)
      bindingThreshold : float
          Threshold for what is considered a HLA binder.
      nPerms : int
          Number of permutations for creating the null distribution.
      k : int
          Length of the kmer peptide.
  
      Returns
      -------
      kmerPdf : ndarray [kmers]
          Probability that the kmer is a predicted epitope restricted by any HLA allele.
      kmerPvalues : ndarray [kmers]
          One-sided p-value that the observed Pr that a kmer is an epitope is greater than under the null.
      sitePdf : ndarray [kmers]
          Probability that the site is in a predicted epitope restricted by any HLA allele.
      sitePvalues : ndarray [kmers]
          One-sided p-value that the observed Pr that a site is in an epitope is greater than under the null.
      kmerNull : ndarray [kmers x perms]
          Null distribution of kmer PDFs that was used to compute the p-values.
      siteNull : ndarray [kmers x perms]
          Null distribution of site PDFs that was used to compute the p-values.
  
      Example
      -------
      >> kmerPdf,kmerPvalues,sitePdf,sitePvalues,kmerNull,siteNull = computeHLAHotspots(seqDf.seq, hlaFreq, ba)"""
  
      nPerms = int(nPerms)
      nHLAs = len(hlas)
      nSeqs = len(seqs)
      nSites = int(np.median(list(map(len, seqs))))
      nKmers = nSites-k+1
  
      if hlaMask is None:
          hlaMask = np.ones((nSeqs, nHLAs))
  
      hlaMaskMat = np.tile(hlaMask, (nKmers, 1, 1)).astype(np.float64)
  
      try:
          hlaList = sorted(hlas.keys())
      except AttributeError:
          hlaList = sorted(hlas)
          hlas = {h:1./len(hlaList) for h in hlaList}
      freqVec = np.array([hlas[h] for h in hlaList])
      
      baMat = getBAMat(seqs, hlaList, ba)
      
      """Identify binders with a 1"""
      binders = (baMat < bindingThreshold).astype(np.float64)
  
      """Multiply element-wise by hla frequencies before summing (to help with permutations later)"""
      weightedBinders = binders * np.tile(freqVec.reshape((1, 1, nHLAs)), (nKmers, nSeqs, 1))
  
      """Multiply by the HLA mask"""
      weightedBinders = weightedBinders * hlaMaskMat
      
      startT = time.time()
      logger.info('Starting computation of observed PDF')
      obsPdf, obsSitePdf = _computePDF(weightedBinders, k)
      
      startT = time.time()
      logger.info('Starting permutation test')
      shuffledPr = np.zeros((nKmers, nPerms))
      
      """Only save the sites that were in k kmers (so dim=0 is nKmers not nSites)"""
      shuffledSitePr = np.zeros((nKmers, nPerms))
      for permi in range(nPerms):
          if nPerms >= 100 and (permi % (nPerms//10)) == 0:
              logger.info('Completed %1.0f%% of permutations', 100 * (permi + 1.) / nPerms)
          """Shuffle the kmer dimension for each HLA, but shuffle the same for all sequences
              (hotspots induced by a dominant HLA allele will still look like hotspots in the null
              so we can est. the p-value of the hotspot correctly)"""
          shuffWB = weightedBinders.copy()
          for hlai in range(nHLAs):
              shuffWB[:,:, hlai] = shuffWB[permutation(nKmers),:, hlai]
          pdf, sitePdf = _computePDF(shuffWB, k)
          shuffledPr[:, permi] = pdf
          
          """Don't include the sites at the begining which are in < k kmers"""
          shuffledSitePr[:, permi] = sitePdf[(k-1):]
      logger.info('Completed hotspot computation in %1.1f seconds', time.time()-startT)
  
      """All kmers should be the same under the null so I can combine permutations across kmers"""
      shuffledPr = shuffledPr.flatten()
      """Most sites are the same under the null so I combine. The significance of a hotspot in the first few
         sites will be underestimated since they will be artificially low from being in <k kmers"""
      shuffledSitePr = shuffledSitePr.flatten()
  
      """One-sided pvalue that the observed Pr that a kmer is an epitope is greater than under the null"""
      pvalues = (np.array([(pr < shuffledPr).sum() for pr in obsPdf], dtype=np.float64) + 1) / (len(shuffledPr) + 1)
      
      sitePvalues = (np.array([(pr < shuffledSitePr).sum() for pr in obsSitePdf], dtype=np.float64) + 1) / (len(shuffledSitePr) + 1)
  
      kmerNull = np.reshape(shuffledPr, (nKmers, nPerms))
      siteNull = np.reshape(shuffledSitePr, (nKmers, nPerms))
      if returnDf:
          kmerDf = pd.DataFrame({'KmerPDF':obsPdf,
                                 'P-value':pvalues})
          siteDf = pd.DataFrame({'SitePDF':obsSitePdf,
                                 'P-value':sitePvalues})
          return kmerDf, siteDf
      else:
          return obsPdf, pvalues, obsSitePdf, sitePvalues, kmerNull, siteNull
  
  def _computePDF(wb, k):
      """Used by computeHLAHotspots in the permutation test.
      Parameter wb is weighted binders [kmers x seqs x hlas]
      Returns pdf, sitePdf"""
  
      """Sum across sequences and then across HLAs"""
      pdf = wb.sum(axis=1).sum(axis=1)
  
      """Sum across sites for each kmer, towards a Pr(epitope) by site"""
      sitePdf = np.convolve(pdf, np.ones(k))
  
      """Normalize pdf to sum to 1"""
      return pdf/pdf.sum(), sitePdf/sitePdf.sum()
  
  
  def plotKmerHotpots(pdf, pvalues=None, cutoff=0.025, pvalueCut=True, seqTicks=''):
      xvec = np.arange(len(pdf))
  
      if pvalueCut and not pvalues is None:
          """Interpret cutoff as p-value alpha threshold"""
          sigInd = pvalues < cutoff
      elif not pvalues is None:
          """Interpret cutoff as a percentile"""
          tmp = sorted(pvalues)[int(np.ceil(cutoff * len(pvalues)))]
          sigInd = pvalues < tmp
      else:
          sigInd = np.zeros(pdf.shape, dtype=bool)
      
      plt.clf()
      axh = plt.gca()
      if sigInd.sum() > 0:
          mn = min(pdf[sigInd])
          """Plot cutoff line"""
          axh.plot([-1, xvec[-1]+1], [mn, mn], '--', color='red', lw=2, alpha=0.5)
  
      for x, y in zip(xvec, pdf):
          axh.plot([x, x], [0, y], '-', color='gray', alpha=0.7)
      axh.plot(xvec[sigInd], pdf[sigInd], 'o', color='red')
      axh.plot(xvec[~sigInd], pdf[~sigInd], 'o', color='black')
      plt.xlabel('9mer start position')
      
      plt.xlim((-1, xvec[-1]+1))
      axh.spines['right'].set_visible(False)
      axh.spines['top'].set_visible(False)
      axh.xaxis.set_ticks_position('none')
      axh.yaxis.set_ticks_position('left')
      if len(seqTicks)>0:
          axh.xticks(np.arange(len(seqTicks)), [s for s in seqTicks])
  
  def plotSiteHotpots(pdf, pvalues, cutoff = 0.025, pvalueCut = True):
      xvec = np.arange(len(pdf))
      if pvalueCut:
          sigInd = pvalues < cutoff
      else:
          tmp = sorted(pvalues)[int(ceil(cutoff * len(pvalues)))]
          sigInd = pvalues < tmp
      
      plt.clf()
      axh = plt.gca()
      if sigInd.sum() > 0:
          mn = min(pdf[sigInd])
          """Plot cutoff line"""
          axh.plot([-1, xvec[-1]+1], [mn, mn], '--', color='red', lw=2, alpha=0.5)
      
      axh.plot(xvec, pdf, '-', color='gray', alpha=0.7)
      axh.plot(xvec[~sigInd], pdf[~sigInd], 'o', color='gray', alpha=0.7)
      axh.plot(xvec[sigInd], pdf[sigInd], 'o', color='red')
      plt.xlabel('AA site number')
      plt.ylabel('PDF of T-cell pressure')
      plt.xlim((-1, xvec[-1]+1))
   
  def plotBindingGrid(hlas, peptides, ba, annotateIC50=True, limitToBinders=False):
      """Heatmap of IC50 for a set of HLA alleles and peptides"""
      hlas = sorted(hlas, reverse =True)
      """Make a pd.DataFrame [hla x peptides] full of binding affinity IC50"""
      bamat = _MAX_IC50 * np.ones((len(hlas), len(peptides)))
      for hlai, h in enumerate(hlas):
          for pepi, pep in enumerate(peptides):
              bamat[hlai, pepi] = ba[(h, pep)]
      baDf = pd.DataFrame(bamat, index=hlas, columns=peptides)
      if limitToBinders:
          ind = (bamat < np.log(1000)).any(axis=1)
          baDf = baDf.iloc[ind]
      _plotBindingGridCommon(baDf.index, peptides, baDf, annotateIC50=annotateIC50)
  
  def plotAlignmentBindingGrid(hlas, align, ba, topN=5, k=9, annotateIC50=True, limitToBinders=False, returnDf=False):
      """Heatmap of IC50 for a set of HLA alleles and all kmers in an alignment.
      The heatmap value is the minimum IC50 of the topN most common variants of that peptide."""
      hlas = sorted(hlas, reverse = True)
  
      nHLAs = len(hlas)
      nSeqs = len(align)
      nSites = int(np.median(list(map(len, align))))
      nKmers = nSites - k + 1
  
      """[hla,kmer]"""
      bamat = _MAX_IC50 * np.ones((nHLAs, nKmers))
  
      peptides=[]
      for hlai, h in enumerate(hlas):
          for meri in np.arange(bamat.shape[1]):
              seqs = _sliceAlign(align, (meri, meri+k))
              """Remove any sequences with a gap at any position"""
              seqs = [s for s in seqs if s.find('-')==-1]
  
              topSeqs = [s for s in objhist(seqs).sortedKeys(reverse=True)][:topN]
              tmp = np.array([ba[(h, pep)] for pep in topSeqs])
              bamat[hlai, meri] = np.nanmin(tmp)
  
              if hlai == 0:
                  peptides.append(topSeqs[0])
  
      baDf = pd.DataFrame(bamat, index = hlas, columns = peptides)
  
      if limitToBinders:
          ind = (bamat < np.log(1000)).any(axis=1)
          baDf = baDf.iloc[ind]
      if returnDf:
          return baDf
      _plotBindingGridCommon(baDf.index, peptides, baDf, annotateIC50 = annotateIC50)
  
  def _plotBindingGridCommon(hlas, peptides, baDf, annotateIC50):
      """Heatmap of IC50 for a set of HLA alleles and peptides"""
      qmat = baDf.values.copy()
      quantization = {4:0.0, 5:0.2, 6:0.4 , 7:0.6 , 16:0.9}
  
      for bound in sorted(list(quantization.keys()), reverse = True):
          qmat[baDf.values < bound] = quantization[bound]
      qmat[np.isnan(qmat)] = 0.9
  
      """Heatmap of covariation"""
      cdict = {'green'  :  ((0, 0, 0), (1, 0, 0)),
               'red':  ((0, 1, 1), (1, 0, 0)),
               'blue' :  ((0, 0, 0), (1, 0, 0))}
  
      """Generate the colormap with 1024 interpolated values"""
      heatCmap = matplotlib.colors.LinearSegmentedColormap('my_colormap', cdict, 1024)
      #heatCmap = get_cmap('hot')
  
      plt.clf()
      axh = plt.subplot2grid((10, 1), (1, 0), rowspan=9)
      labelTxtProp = dict(family='monospace', size='medium', weight='bold', color='white', ha='center', va='center')
  
      axh.pcolor(qmat, cmap=heatCmap, vmin=0, vmax=1)
      if annotateIC50:
          for hlai, h in enumerate(hlas):
              for pepi, pep in enumerate(peptides):
                  if baDf.values[hlai, pepi] <= np.log(500):
                      txtArgs =  (np.round(np.exp(baDf.values[hlai, pepi]), -1))
                      plt.text(pepi+0.5,
                               hlai+0.5,
                               '%1.0f' % txtArgs,
                               **labelTxtProp)
  
      axh.colorbar(fraction=0.05,
                   values=[0, 0.2, 0.4, 0.6, 1],
                   boundaries=[0, 50, 150, 400, 1000],
                   label='IC50 (nM)')
  
      axh.grid(color='white')
      plt.xlim((0, len(peptides)))
      plt.ylim((0, len(hlas)))
      plt.yticks(np.arange(len(hlas)) + 0.5, hlas, fontname='Consolas')
      if len(peptides) < 50:
          axh.xaxis.tick_top()
          plt.xticks(np.arange(len(peptides)) + 0.5,
                     peptides,
                     rotation=90,
                     fontname='Consolas')
      else:
          plt.xticks(np.arange(0, len(peptides), 10)+0.5, np.arange(0, len(peptides), 10))
  
      plt.xlabel('Peptides')
      plt.ylabel('HLA alleles')
  
  def _sliceAlign(align, region, sites=False):
      """Return a region of the alignment where region is (start, end)
      OR if sites is True then include all sites in region (not range)"""
      if region is None:
          return align
      elif sites:
          return align.map(lambda seq: ''.join([seq[r] for r in region]))
      else:
          return align.map(lambda seq: seq[region[0]:region[1]])