phylomatic

#!/usr/bin/gawk -f

# Standalone phylomatic 

# Copyright (c) 2012-2021, Campbell Webb <cw@camwebb.info>

# Global variables: Parent[], Bl[], Taxon[], Note[], DaughtersN[],
#                   HasBL, Warning, RootNode
# Global constants: NOCLEAN, FYIN, FYOUT, TAXAFILE, PHYLOFILE

BEGIN{

  read_args()
  # read input newick
  new2fyt()

  # convert and exit
  if (FYIN) {
    printfy()
    exit 0
  }
  
  phylomatic()

  # output as fy or newick
  if (FYOUT) {
    printfy()
    exit 0
  }
  else {
    fy2new()
    exit 0
  }
}

## ----- NEW2FYT ---------------------------------------------------------

function new2fyt(          newick, i, j , k, atn, iname, taxa, b) {

  RS="\x04"
  while ((getline < PHYLOFILE) > 0)
    newick = $0
  if (!newick)
    error("Reading phylo file failed")

  # no returns allowed
  gsub(/[\n\r]/, "", newick)
  # make sure names adhere to parentheses or commas (see tests below)
  newick = gensub(/([(),]) +/,"\\1","G",newick)
  
  # checks
  if (gsub(/\(/,"(",newick) != gsub(/\)/,")",newick))
    error("Newick, mismatched parentheses")
  if (!index(newick, ";") || !index(newick, "(") || !index(newick, ")") || \
      !index(newick, ","))
    error("Newick, missing ( , ) ; characters")
  if (newick ~ /:[0-9.-]+[eE][0-9-]+/)
    error("Newick, branch lengths contain sci notation")
      
  j = k = -1
  RootNode = 0

  # Move through the newick format tree character by character
  i = 1
  while (i < length(newick)) {
    # ignore spaces not in names or notes
    if (substr(newick,i,1) == " ")
      i++

    # the end (if not the last char)?
    else if (substr(newick,i,1) == ";")
      break
    
    # descend a branch, create Parent
    else if (substr(newick,i,1) == "(") {
      j++
      Parent[j] = k
      Taxon[j] = ""
      DaughtersN[k]++
      k = j
      i++
    }
      
    # sibling taxa
    else if (substr(newick, i, 1) == ",")
      i++

    # back up a node to len, keep track of locn with atn
    else if (substr(newick, i, 1 ) == ")") {
      k = Parent[k]
      atn = Parent[atn]
      i++
    }

    # Interior name
    else if ((substr(newick, i, 1 ) ~ /[A-Za-z\-_\']/) &&   \
             (substr(newick, i-1, 1 ) == ")")) {
      iname = ""
          
      while(substr(newick, i, 1) !~ /[:,\)\[;]/) {
        iname = iname substr(newick, i, 1)
        i++
      }
      # replaces spaces with underscores
      gsub(/ /,"_",iname)
      Taxon[atn] = iname
    }

    # Note
    else if (substr(newick, i, 1 ) == "[") {
      while(substr(newick, i+1, 1) != "]") {
        i++
        Note[atn] = Note[atn] substr(newick, i, 1 )
      }
      i += 2  # from ...not<e>]... to ...note]<.>..
    }
    
    # branch length coming up
    else if (substr(newick, i, 1 ) == ":") {
      b = ""
      i++
      while(substr(newick, i, 1) ~ /[ 0-9.-]/) {
        b = b substr(newick, i, 1)
        i++
      }
      # clean any spaces after colon
      gsub(/ /,"",b)
      if (b != "") {HasBL = 1}
      Bl[atn] = b
    }

    # default - it's a new Taxon name
    # TODO - fix this. anytime you find a non ,:;()[] character, check to 
    # see if previous VALID character is , or ). If it is, or if character
    # is ', begin a new Taxon name. If it's not, ignore. This should handle
    # whitespace, newlines and other garbage (but non-fatal) characters 
    # that currently crash the input.
    else if ((substr(newick, i, 1 ) ~ /[A-Za-z\-_]\'/) &&    \
             (substr(newick, i-1, 1 ) ~ /[(,]/)) {
      taxa = ""
      taxa = taxa substr(newick, i, 1)
      i++
    
      # Keep adding more chars
      while ((substr(newick, i, 1) !~ /[,\):\[]/) &&    \
             (i < length(newick))) {
        taxa = taxa substr(newick, i, 1)
        i++
      }
      gsub(/ +$/,"",taxa)
      gsub(/ /,"_",taxa)
      
      # A new name means a new node
      j++
      atn = j
      Parent[j] = k
      Taxon[j] = taxa
      DaughtersN[k]++
    }
    else i++
  }
  
  # Set the RootNode Parent
  Parent[0] = "-1"

}

## ----- PHYLOMATIC ---------------------------------------------------------

function phylomatic(       divbl, endbit, ftaxa, i, j, keep, kpd, matched, \
                           nodetmp, newnode, nkpd, nmatched, node, ntaxatrees,\
                           p, shiftedbl, taxa, thisnode, tobecleaned, x, xbl, \
                           z1, z2, possNewRoot, wascleaned, k) {
  if (HasBL)
    ageToRoot()

  # Read taxa file
  RS="\x04"
  while ((getline < TAXAFILE) > 0)
    ftaxa = $0
  if (!ftaxa)
    error("Reading taxa file failed")
  
  # clean taxa file
  gsub(/\r\n/,"\n",ftaxa);  # fix windows
  gsub(/\r/,"\n",ftaxa);    # fix mac
  gsub(/\n\n*$/,"",ftaxa);  # clean empty newlines at end
  gsub(/[\t]/,"",ftaxa);   # important - tabs were causing hangs

  # split lines
  ntaxatrees = split(ftaxa, taxa, "\n")
  if (ntaxatrees > 10000)
    error("More than 10,000 taxa. May take too long.\n"     \
          "  Comment out the relevant Warning in source"\
          " if you really want to do this.")

  # ** Graft **

  # counter for additional nodes
  newnode = 1

  for (i = 1; i <= ntaxatrees; i++) {
    delete node
    # parse slashes
    nnodes = split(taxa[i], nodetmp, "/")
    # reverse the indices, so 1 is the terminal one
    for (j = 1; j <= nnodes; j++) {
      # clean in the same way names in the phylo were cleaned
      gsub(/ +$/,"",nodetmp[j])
      gsub(/^ +/,"",nodetmp[j])
      gsub(/ /,"_",nodetmp[j])
      node[nnodes - j+1] =  nodetmp[j]
    }
    #  climb through the levels of nesting from terminal to basal
    for (j = 1; j <= nnodes; j++) {
      if (matched[i])
        break
      
      # looping through all nodes in megatree that have taxa names
      # NOTE: should be quicker to reverse this and loop once through
      #   mega tree names, and multiple times through taxa file names.
      #   will require a rewrite
      for (k in Parent) {
        if (!Taxon[k])
          continue

        # compare with the Taxon name at every node in the megatree
        # case insensitive
        if (tolower(node[j]) == tolower(Taxon[k])) {
          if (DEBUG)
            print "found " Taxon[k] " for t " i " j " j "/" nnodes \
              > "/dev/stderr"

          # if matches, add to megatree:
          
          # terminal to terminal match:
          if ((DaughtersN[k] == 0) && (j==1)) {
            # flag the existing node and break (no graft, just prune)
            keep[k] =1
            matched[i]=1
            nmatched++
            break
          }
          # grafting needed
          else {
            # deal with branch length adjustments
            if (HasBL) {
              if (DaughtersN[k] == 0) {
                # need to shift the BL of tip
                ageNode(Parent[k])
                divbl = Age[Parent[k]] / j
                if (!shiftedbl[k])    {
                  Bl[k] = divbl
                  shiftedbl[k] = 1
                }
              }
              # internal match
              else {
                ageNode(k)
                divbl = Age[k] / (j-1)
              }
            }

            x = k

            # starting with the node distal to the matched one
            for (p = j-1; p >= 1; p--) {
              thisnode = "pm" newnode
              if (HasBL)
                Bl[thisnode] = divbl
              Parent[thisnode] = x
              DaughtersN[x]++
              Taxon[thisnode] = node[p]
              if (p == 1)
                keep[thisnode] = 1
              x = thisnode
              newnode++
            }
            matched[i]=1
            nmatched++
            break
          }
        }
      } 
    }
  }

  if(!nmatched) error("No taxa in common")

  # ** Prune **
  
  for (i in Parent) {
    # terminals only with matching taxa
    if ((DaughtersN[i] == 0) && keep[i]) {
      j = i
      # move up through tree
      # ...was: while ((Parent[j] != "NULL") || (!Parent[j])) {
      #   because the RootNode's Parent was sometimes ""; safer to use -1
      while (Parent[j] != "-1") {
        # for possible cleaning. kpd = kept parent / daughter. Later,
        # if a node has only been visited via it's single root from
        # retained terminal, and not via two or more it can be cleaned
        # out
        kpd[Parent[j]][j]=1

        # keep the parent in the pruned tree (before cleaning)
        j = Parent[j]
        keep[j]=1
      }
    }
    # and keep the root, never reached in previous loop
    if (Parent[j] == "-1")
      keep[j]=1
  }

  # now remove the elements not kept:
  for (i in Parent)
    if (!keep[i])
      delete Parent[i]  # Careful to only loop through Parent[] after this!
  
  # Finally, if phy is to be cleaned:
  if (!NOCLEAN) {

    # nkpd is number of times the Parent in an edge is passed through
    # if only one, the node can be cleaned out
    for (i in kpd)
      for (j in kpd[i])
        nkpd[i]++

    for (i in Parent) {
      # visit paths to root from all terminal nodes
      if (!DaughtersN[i]) {

        if (DEBUG)
          print "node " i ", Parent " Parent[i] ", tax " Taxon[i] \
            > "/dev/stderr"
        
        j = i
        # move up through tree, not considering the root, and skipping
        #   paths through tree already visited
        while ((Parent[j] != "-1") && !wascleaned[j]) {
          if (DEBUG)
            print "  considering Parent of " j " -> " Parent[j] " Taxon " \
              Taxon[Parent[j]]  > "/dev/stderr"
          x = Parent[j]
          xbl = 0
          # if the Parent node is only passed through once, i.e. not a node
          #   in cleaned phylogeny
          while ((nkpd[x] == 1) && (Parent[x] != "-1")) {
            tobecleaned = x
            xbl += Bl[tobecleaned]
            x = Parent[x]

            if (DEBUG)
              print "    deleted " tobecleaned " Taxon " Taxon[tobecleaned] \
                > "/dev/stderr"
            delete Parent[tobecleaned]
          }
          # connect j to x
          Parent[j] = x
          Bl[j] += xbl
          # record that this node has been visited during cleaning
          wascleaned[j] = 1
          # now try the next one...
          j = Parent[j]
          # record a node (may be several) that is a node away from
          #   the original root
          if (Parent[Parent[j]] == "-1")
            possNewRoot = j
        }
        # set a counter to find dangling
        
        # set new RootNode
        # if (Parent[j] == "-1") RootNode = j
      }
    }

    ## if the original root is now a tail
    if (nkpd[RootNode] == 1) {
      delete Parent["-1"]
      RootNode = possNewRoot
    }
      
    # NULL was added to the index of Parent, clear it:
    # delete Parent["-1"]
  }


  # Missing taxa
  if (nmatched < ntaxatrees) {
    Warning = "Warning: " ntaxatrees - nmatched " taxa not matched: "
    for (i = 1; i <= ntaxatrees; i++)
      if (!matched[i])
        Warning = Warning "'" taxa[i] "', "
    gsub(/, $/,"",Warning)
  }

  close("/dev/stderr")

}

function ageToRoot(     i, j) {
  TotalBL = 0
  # assumes an ultrametric megatree
  for (i in Parent) {
    # find the first term
    if (!DaughtersN[i]) {
      j = i
      # move up through tree
      # WAS: while ((Parent[j] != "NULL") || (!Parent[j])) {
      while (Parent[j] != "-1") {
        TotalBL = TotalBL + Bl[j]
        j = Parent[j]
      }
      break
    }
  }
}

function ageNode(innode,      j, sumBL) {
  
  if (Age[innode])
    return
  else {
    j = innode
    # move up through tree
    while (Parent[j] != "-1") {
      sumBL = sumBL + Bl[j]
      j = Parent[j]
    }
  }
  Age[innode] = TotalBL - sumBL
  if (DEBUG)
    print "Age of " innode " is " Age[innode] > "/dev/stderr"
  close("/dev/stderr")

}

function fy2new(                 x , y, n, first , mark, tmp)  {
  # (c) 2011 Cam Webb

  # fy2new (recursive version)
  # converts a tabular fy format phylogeny to Newick, parenthetical format

  # fy-format = tab-delimited string fields:
  #   1. nodeID
  #   2. Parent node nodeID (must be either "-1", "", or "NULL" for
  #      root node)
  #   3. branch length to Parent node (an integer or float; missing for no BLs)
  #   4. node name (terminal or interior node)
  #   5. (optional) notes attached node

  # check for RootNode:
  if (RootNode == "")
    error("No RootNode found")

  # lDaughter, rSister needed for down-pass recursivity. Create:

  # initialize
  for (x in Parent) {
    lDaughter[x] = "NULL"
    rSister[x] = "NULL"
    first[x] = 1
    n++
  }

  # special case of a single node:
  if (n == 1) {
    print Taxon[x] ";"
    return
  }

  # create lDaughter, rSister
  for (x in Parent) {
    # starting at terms
    if (!DaughtersN[x]) {
      y = x
      # while not yet at the RootNode (allow diff coding schemes)
      while (y != RootNode) {
        # is this the first daughter?
        if (lDaughter[Parent[y]] == "NULL")
          lDaughter[Parent[y]] = y
        # if not, find the dangling sister
        else {
          # start at lDaughter
          mark = lDaughter[Parent[y]]
          # move to node with an empty rSister
          while (rSister[mark] != "NULL")
            mark = rSister[mark]
          rSister[mark] = y
        }
        
        # test for refollowing old routes
        if (first[Parent[y]] == 1) {
          y = Parent[y]
          first[y] = 0
        }
        else
          break
      }
    }
  }

  # Now recurse through levels, creating out string in tmp

  tmp = ""
  
  print downPar(RootNode, tmp) ";"
  
  if (Warning)
    print Warning > "/dev/stderr"
}

function downPar(atn, tmp,           x, which, tmpnext ) {
  which = 0

  # if terminal, go no further
  if (!DaughtersN[atn]) {
    tmp = gensub(/ /,"_","G",Taxon[atn])
    if ((HasBL) && ( Bl[atn])) 
      tmp = tmp ":" Bl[atn]
    if (Note[atn])
      # TODO: escape any [, ] in note 
      tmp = tmp "[" Note[atn] "]" 
  }
  else {
    
    x = lDaughter[atn]
    tmp =  "("
    tmp = tmp downPar(x, tmpnext[which])
    
    x = rSister[x]; which++
    
    while (x != "NULL") {
      tmp = tmp  ","
      tmp = tmp downPar(x, tmpnext[which])
      x = rSister[x]; which++
    }
    tmp = tmp  ")"
    tmp = tmp  gensub(/ /,"_","G",Taxon[atn])
    
    if ((HasBL) && ( Bl[atn])) 
      tmp = tmp  ":" Bl[atn]
    if (Note[atn])
      tmp = tmp "[" Note[atn] "]" 
  }
  
  return tmp
}

function printfy(     i) {

  # n = asorti(Parent, sortidx)   # old method
  # for (i = 1; i <= n; i++) {
  # print sortidx[i], Parent[sortidx[i]],         \
  #   Bl[sortidx[i]], Taxon[sortidx[i]] }

  PROCINFO["sorted_in"] = "@ind_num_asc"
  OFS="|"
  for (i in Parent)
    print i, Parent[i], Bl[i], Taxon[i], Note[i]  # DaughtersN[i],  
  
  if (Warning)
    print Warning > "/dev/stderr"
}

function error( msg ) {
  print "Error: " msg 
  exit 1
}

function read_args(    arg) {
  # Read arguments, check, give usage. Long args to avoid Gawk clash
  NOCLEAN = FYIN = FYOUT = TAXAFILE = PHYLOFILE = 0
  for (arg = 1; arg < ARGC; arg++) {
    if      (ARGV[arg] == "--noclean" ) NOCLEAN     = 1
    else if (ARGV[arg] == "--new2fy") FYIN      = 1 # print fy of phylo, exit
    else if (ARGV[arg] == "--fyout" ) FYOUT     = 1 # output fy not newick
    else if (ARGV[arg] == "--taxa"  ) TAXAFILE  = ARGV[arg+1]
    else if (ARGV[arg] == "--newick") PHYLOFILE = ARGV[arg+1]
  }
  
  if ((!TAXAFILE && !FYIN) || !PHYLOFILE) {
    print "Usages: phylomatic [ --noclean --fyout  ] "                 \
      "--newick <phylo> --taxa <taxa>\n"                             \
      "        phylomatic             --new2fy   --newick <phylo>"     \
      > "/dev/stderr"
    close("/dev/stderr")
    exit 1
  }
}