[jalview.git] / src / jalview / util / HMMProbabilityDistributionAnalyser.java

package jalview.util;

import jalview.datamodel.HiddenMarkovModel;
import jalview.datamodel.SequenceI;
import jalview.io.DataSourceType;
import jalview.io.FileParse;
import jalview.io.HMMFile;
import jalview.io.StockholmFile;
import jalview.schemes.ResidueProperties;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.Vector;

/**
 * Processes probability data. The file indexes used in this program represent
 * the index of the location of a family or hmm in their respective files,
 * starting from 0.
 * 
 * @author TZVanaalten
 *
 */
public class HMMProbabilityDistributionAnalyser
{

  Vector<SequenceI> sequences;

  HiddenMarkovModel hmm;

  // contains the raw data produced
  List<ArrayList<Double>> raw = new ArrayList<>();

  // contains binned data
  Map<String, Double> binned = new HashMap<>();

  // location of the family file
  String families = "H:/Desktop/PFAM/Family/SeedFamilies.seed";

  // location of the file containing the family-clan links
  final static String FAMILIESTOCLAN = "H:/Desktop/PFAM/Family/Clanlinks.dat";

  // location of the HMM file
  String hmms = "H:/Desktop/PFAM/HMMs/Pfam-A.hmm";

  // suffix for raw file
  final static String RAW = "/Raw.csv";

  // suffix for binned file
  final static String BINNED = "/Binned.csv";

  // normalisation scale
  final static double SCALE = 1;

  // current position in file
  int currentFilePosition = 0;

  final static String NL = "\n";

  // current directory
  String currentFolder;

  boolean keepRaw = false;

  /**
   * Sets the working directory.
   * 
   * @param path
   */
  public void setFolder(String path)
  {
    currentFolder = path;
  }

  /**
   * Moves a buffered reader forward in the file by a certain amount of entries.
   * Each entry in the file is delimited by '//'.
   * 
   * @param index
   *          The index of the location in the file.
   * @param br
   * @throws IOException
   */
  public void moveLocationBy(int index, BufferedReader br)
          throws IOException
  {
    for (int i = 0; i < index; i++)
    {
      String line = br.readLine();
      while (!"//".equals(line))
      {
        line = br.readLine();
      }
    }

  }
  
  /**
   * Analyses a specified number of families and then saves the data. Before
   * analysing the data, the previous saved data will be imported and after
   * analysing this, the data is exported back into the file.
   * 
   * @param increments
   *          The number of families to read before saving.
   * @throws IOException
   */
  public void run(int increments, boolean keepRawData) throws IOException
  {
    keepRaw = keepRawData;
    try
    {
      readPreviousData(currentFolder);
      BufferedReader posReader = new BufferedReader(
              new FileReader(currentFolder + "/CurrentPosition.txt"));

      String line = posReader.readLine();
      posReader.close();
      currentFilePosition = Integer.parseInt(line);
    } catch (Exception e)
    {
      System.out.println("No previous data found");
    }



    BufferedReader inputSTO = new BufferedReader(new FileReader(families));
    BufferedReader inputHMM = new BufferedReader(new FileReader(hmms));



    moveLocationBy(currentFilePosition, inputHMM);
    moveLocationBy(currentFilePosition, inputSTO);

    int filesRead = 0;
    int i = 0;
    while (filesRead < increments)
    {

      FileParse parserSTO = new FileParse(inputSTO, "",
              DataSourceType.FILE);
      readStockholm(parserSTO);

      FileParse parserHMM = new FileParse(inputHMM, "",
              DataSourceType.FILE);
      readHMM(parserHMM);

        int count = countValidResidues();
        processData(count);
        filesRead++;

      currentFilePosition++;
      System.out.println(i);
      i++;
    }

    PrintWriter p = new PrintWriter(
            new File(currentFolder + "/CurrentPosition.txt"));
    p.print(currentFilePosition);
    p.close();
    exportData(currentFolder);
    raw.clear();
    binned.clear();

  }

  /**
   * Analyses all families and then saves the data. Before analysing the data,
   * the previous saved data will be imported and after analysing this, the data
   * is exported back into the file.
   * 
   * @param increments
   *          The number of families to read before saving.
   * @throws IOException
   */
  public void runToEnd(boolean keepRawData) throws IOException
  {
    keepRaw = keepRawData;
    BufferedReader inputSTO = null;
    BufferedReader inputHMM = null;
    int size = 0;
    try
    {
      readPreviousData(currentFolder);
      BufferedReader posReader = new BufferedReader(
              new FileReader(currentFolder + "/CurrentPosition.txt"));

      String line = posReader.readLine();
      posReader.close();
      currentFilePosition = Integer.parseInt(line);
      readPreviousData(currentFolder);
      
      inputSTO = new BufferedReader(new FileReader(families));
      inputHMM = new BufferedReader(new FileReader(hmms));
    } catch (Exception e)
    {
      System.out.println("No or incomplete previous data found");
    }

    

    moveLocationBy(currentFilePosition, inputHMM);
    moveLocationBy(currentFilePosition, inputSTO);

    int filesRead = 0;
    int i = 0;
    inputSTO.mark(20);
    String check = inputSTO.readLine();
    inputSTO.reset();
    while (!"".equals(check) && !" ".equals(check) && check != null)
      {
      inputSTO.mark(20);
      String line = inputSTO.readLine();
      inputSTO.reset();

        FileParse parserSTO = new FileParse(inputSTO, "",
                DataSourceType.FILE);
        readStockholm(parserSTO);

        FileParse parserHMM = new FileParse(inputHMM, "",
                DataSourceType.FILE);
        readHMM(parserHMM);

        int count = countValidResidues();
        processData(count);
        filesRead++;

        currentFilePosition++;
        System.out.println(i);
        i++;
      inputSTO.mark(20);
      check = inputSTO.readLine();
      inputSTO.reset();
      }


    PrintWriter p = new PrintWriter(
            new File(currentFolder + "/CurrentPosition.txt"));
    p.print(currentFilePosition);
    p.close();
    exportData(currentFolder);
    raw.clear();
    binned.clear();

  }

  /**
   * Reads the previous data from both files
   * 
   * @param source
   * @throws IOException
   */
  public void readPreviousData(String source) throws IOException
  {
    readBinned(source);
    if (keepRaw)
    {
      readRaw(source);
    }
  }

  /**
   * Reads the previous data from the binned file.
   * 
   * @param source
   * @throws IOException
   */
  public void readBinned(String source) throws IOException
  {
    BufferedReader input = new BufferedReader(
            new FileReader(source + BINNED));
    String line = input.readLine();
    binned = new HashMap<>();
    while (!("".equals(line) || line == null))
    {
      Scanner scanner = new Scanner(line);
      scanner.useDelimiter(",");
      String key = scanner.next();
      String value = scanner.next();
      binned.put(key, Double.valueOf(value));
      scanner.close();
      line = input.readLine();
    }

    input.close();
  }

  /**
   * Reads the previous data from the raw file.
   * 
   * @param source
   * @throws IOException
   */
  public void readRaw(String source) throws IOException
  {
    BufferedReader input = new BufferedReader(new FileReader(source + RAW));
    String line = input.readLine();
    if (line == null)
    {
      input.close();
      return;
    }
    Scanner numberScanner = new Scanner(line);
    numberScanner.useDelimiter(",");
    raw = new ArrayList<>();
    while (numberScanner.hasNext())
    {
      numberScanner.next();
      raw.add(new ArrayList<Double>());
    }
    numberScanner.close();

    line = input.readLine();
    while (!("".equals(line) || line == null))
    {
      Scanner scanner = new Scanner(line);
      scanner.useDelimiter(",");

      int i = 0;
      while (scanner.hasNext())
      {
        String value;
        value = scanner.next();
        if (!value.equals("EMPTY"))
        {
          raw.get(i).add(Double.parseDouble(value));
        }
        else
        {
          raw.get(i).add(null);
        }

        i++;
      }
      scanner.close();
      line = input.readLine();
    }

    input.close();
  }

  /**
   * Counts the number of valid residues in the sequence.
   * 
   * @return
   */
  public int countValidResidues()
  {
    int count = 0;

    for (int width = 0; width < sequences.size(); width++)
    {
      for (int length = 1; length < hmm.getLength() + 1; length++)
      {
        char symbol;
        int alignPos;
        alignPos = hmm.getNodeAlignmentColumn(length);

        symbol = sequences.get(width).getCharAt(alignPos);
        if (ResidueProperties.aminoBackgroundFrequencies
                .containsKey(symbol))
        {
          count++;
        }
      }
    }

    return count;
  }

  /**
   * Processes data, and stores it in both a raw and binned format.
   * 
   * @param count
   */
  public void processData(int count)
  {
    int rawPos = 0;
    if (keepRaw)
    {
      raw.add(new ArrayList<Double>());
      rawPos = raw.size() - 1;
    }

    for (int width = 0; width < sequences.size(); width++)
    {
      for (int length = 1; length < hmm.getLength() + 1; length++)
      {
        char symbol;
        int alignPos;
        alignPos = hmm.getNodeAlignmentColumn(length);
        
        symbol = sequences.get(width).getCharAt(alignPos);
        if (ResidueProperties.aminoBackgroundFrequencies
                .containsKey(symbol))
        {
          Double prob;
          Float bfreq;
          Double llr;
          prob = hmm.getMatchEmissionProbability(alignPos, symbol);
          bfreq = ResidueProperties.aminoBackgroundFrequencies.get(symbol);
          llr = Math.log(prob / bfreq);
          if (keepRaw)
          {
            raw.get(rawPos).add(llr);
          }

          String output;
          output = String.format("%.1f", llr);
          if ("-0.0".equals(output))
          {
            output = "0.0";
          }
          if (binned.containsKey(output))
          {
            double prev = binned.get(output);
            prev += (SCALE / count);
            binned.put(output, prev);

          }
          else
          {
            binned.put(output, SCALE / count);
          }
        }
      }
    }
  }


  /**
   * Reads in the sequence data from a Stockholm file.
   * 
   * @param source
   * @throws IOException
   */
  public void readStockholm(FileParse source) throws IOException
  {
    StockholmFile file = new StockholmFile(source);
    sequences = file.getSeqs();
  }

  /**
   * Reads in the HMM data from a HMMer file.
   * 
   * @param source
   * @throws IOException
   */
  public void readHMM(FileParse source) throws IOException
  {

    HMMFile file = new HMMFile(source);
    file.parse();
    hmm = file.getHMM();

  }

  /**
   * Exports both the binned and raw data into separate files.
   * 
   * @param location
   * @throws FileNotFoundException
   */
  public void exportData(String location) throws FileNotFoundException
  {
    PrintWriter writerBin = new PrintWriter(new File(location + BINNED));
    for (Map.Entry<String, Double> entry : binned.entrySet())
    {
      writerBin.println(entry.getKey() + "," + entry.getValue());
    }
    writerBin.close();
    if (keepRaw)
    {

    PrintWriter writerRaw = new PrintWriter(new File(location + RAW));
    
    StringBuilder identifier = new StringBuilder();
    
    for (int i = 1; i < raw.size() + 1; i++)
    {
      identifier.append("Fam " + i + ",");
    }
    
    writerRaw.println(identifier);
    
    boolean rowIsEmpty = false;
    int row = 0;
    while (!rowIsEmpty)
    {
      rowIsEmpty = true;
      StringBuilder string = new StringBuilder();
      for (int column = 0; column < raw.size(); column++)
      {
        if (raw.get(column).size() <= row)
        {
          string.append("EMPTY,");
        }
        else
        {
          string.append(raw.get(column).get(row) + ",");
          rowIsEmpty = false;
        }
      }
      row++;
      writerRaw.println(string);
    }
    writerRaw.close();

    }

  }

  /**
   * Prints the specified family on the console.
   * 
   * @param index
   * @throws IOException
   */
  public void printFam(int index) throws IOException
  {
    BufferedReader br = new BufferedReader(new FileReader(families));

    moveLocationBy(index, br);

    String line = br.readLine();

    while (!"//".equals(line))
    {
      System.out.println(line);
      line = br.readLine();
    }
    System.out.println(line);
    br.close();

  }

  /**
   * Prints the specified HMM on the console.
   * 
   * @param index
   * @throws IOException
   */
  public void printHMM(int index) throws IOException
  {
    BufferedReader br = new BufferedReader(new FileReader(hmms));

    moveLocationBy(index, br);

    String line = br.readLine();

    while (!"//".equals(line))
    {
      System.out.println(line);
      line = br.readLine();
    }
    System.out.println(line);
    br.close();

  }

  /**
   * Prints the specified family to a .sto file.
   * 
   * @param index
   * @throws IOException
   */
  public void exportFam(int index, String location) throws IOException
  {
    BufferedReader br = new BufferedReader(new FileReader(families));

    moveLocationBy(index, br);

    String line = br.readLine();
    PrintWriter writer = new PrintWriter(
            new FileOutputStream(new File(location), true));
    while (!"//".equals(line))
    {
      writer.println(line);
      line = br.readLine();
    }
    writer.println(line);
    writer.close();
    br.close();

  }

  public void exportFile(BufferedReader br, String location)
          throws IOException
  {
    String line = br.readLine();
    PrintWriter writer = new PrintWriter(
            new FileOutputStream(new File(location), true));
    while (!"//".equals(line))
    {
      writer.println(line);
      line = br.readLine();
    }
    writer.println(line);
    writer.close();


  }

  public String getHMMName(int index) throws IOException
  {
    String name;

    BufferedReader nameFinder = new BufferedReader(new FileReader(hmms));

    moveLocationBy(index, nameFinder);

    nameFinder.readLine();

    Scanner scanner = new Scanner(nameFinder.readLine());
    name = scanner.next();
    name = scanner.next();
    scanner.close();
    return name;
  }

  public String getFamilyName(int index) throws IOException
  {
    String name;

    BufferedReader nameFinder = new BufferedReader(
            new FileReader(families));

    moveLocationBy(index, nameFinder);

    nameFinder.readLine();

    Scanner scanner = new Scanner(nameFinder.readLine());
    name = scanner.next();
    name = scanner.next();
    name = scanner.next();
    scanner.close();
    return name;
  }

  /**
   * Prints the specified family to a .hmm file in the current directory.
   * 
   * @param index
   * @throws IOException
   */
  public void exportHMM(int index, String location) throws IOException
  {


    BufferedReader br = new BufferedReader(new FileReader(hmms));

    moveLocationBy(index, br);

    String line = br.readLine();

    PrintWriter writer = new PrintWriter(
            new FileOutputStream(new File(location), true));
    while (!"//".equals(line))
    {
      writer.println(line);
      line = br.readLine();
    }
    writer.println(line);
    writer.close();
    br.close();

  }
  
  /**
   * Clears all raw, binned and current position data in the current directory.
   * 
   * @throws FileNotFoundException
   */
  public void clear() throws FileNotFoundException
  {
    PrintWriter pos = new PrintWriter(
            currentFolder + "/CurrentPosition.txt");
    pos.println("0");
    
    PrintWriter raw = new PrintWriter(currentFolder + RAW);
    
    PrintWriter bin = new PrintWriter(currentFolder + BINNED);
    
    pos.close();
    bin.close();
    raw.close();
  }

  public void sortIntoClans(String directory) throws IOException
  {
    BufferedReader clanFinder = new BufferedReader(new FileReader(FAMILIESTOCLAN));
    BufferedReader familyReader = new BufferedReader(
            new FileReader(families));
    BufferedReader hmmReader = new BufferedReader(new FileReader(hmms));
    HashMap<String, Integer> clanIndexes = new HashMap<>();
    int filePos = 0; 
    int clanCount = 0;
    String line;
    line = clanFinder.readLine();
    
    while (!"".equals(line) && !" ".equals(line) && line != null)
    {
     String clanName;
      boolean inClan = false;
     while (!(line.indexOf("//") > -1))
     {
       
      if (line.indexOf("#=GF CL") > -1)
      {
          inClan = true;
        Scanner scanner = new Scanner(line);
        scanner.next();
        scanner.next();
        clanName = scanner.next();
          scanner.close();
        
        if (!clanIndexes.containsKey(clanName))
        {
          clanIndexes.put(clanName, clanCount);
            clanCount++;
        }


          Integer clanI = clanIndexes.get(clanName);
          String clanPath = directory + "/Clan" + clanI.toString();
          File clanFolder = new File(clanPath);
          String famPath = clanPath + "/Families.sto";
          String hmmPath = clanPath + "/HMMs.hmm";
          if (!clanFolder.exists())
        {
            clanFolder.mkdir();
        }
          exportFile(familyReader, famPath);
          exportFile(hmmReader, hmmPath);

      }
        line = clanFinder.readLine();
      }
      if (!inClan)
      {
        moveLocationBy(1, familyReader);
        moveLocationBy(1, hmmReader);
      }
      filePos++;
      System.out.println(filePos + " files read.");
      line = clanFinder.readLine();

     }
    clanFinder.close();
      
    }

  public String getFamilies()
  {
    return families;
  }

  public void setFamilies(String families)
  {
    this.families = currentFolder + families;
  }

  public String getHmms()
  {
    return hmms;
  }

  public void setHmms(String hmms)
  {
    this.hmms = currentFolder + hmms;
  }
    


  }