Update_BEAST_operators.pyΒΆ

#! /usr/bin/env python

desc = """
Update the BEAST operators using the values suggested in the output log.

BEAST versions supported: 2.3.0 only

<logfile> should be the name of the log file output by Beast that has the new
operator tuning values
<target_xml> is the Beast XML file where you want the new operator tuning values

The script writes a new XML file named <target_xml>-updated

Version 0.1.2
Cymon J. Cox Jun 2015
"""
import sys
import os
import argparse
import textwrap
import shutil
import subprocess
import signal
import xml.etree.ElementTree as ET

OK_VERSIONS = ["v2.3.0"]

def check_version():

    p = subprocess.Popen("beast blah45", stdout=subprocess.PIPE,
            stderr=subprocess.PIPE, shell=True)
    stdout, stderr = p.communicate()
    version = stdout.split("\n")[1].strip().split(" ")[1].rstrip(",")
    if not version in OK_VERSIONS:
        print "Beast version %s is not supported" % version
        sys.exit()
    return version

def collect_operator_values(filename):

    f = open(filename, "r")
    lines = f.readlines()
    f.close()

    all_count = 0
    updated_count = 0
    scaleFactors = {}
    deltas = {}

    ovs = [{},{}]

    ilines = iter(lines)

    while True:

        line = ilines.next()

        if line.startswith("Appending"):
            print "It looks like you have a resumed run and therefore maybe" + \
                  " two or more acceptance blocks: refusing to continue."
            sys.exit()

        if line.startswith("Operator"):
            #Loop through until next blank line:                
            while True:
                line = ilines.next()
                if line.strip() == "":
                    break
                else:
                    #collect operators
                    if "Try" in line:
                        bits = line.split()
                        value = bits[-1]
                        operstr = bits[0]
                        oper, name = bits[0].split("(")
                        name = name[:-1]
                        if oper == "SubtreeSlide":
                            ovs[0][name] = value
                        else:
                            ovs[1][name] = value
            break

    return ovs


def update_XML_runfile(ovs, target):

    count_sts = 0
    count_so  = 0

    tree = ET.parse(target)
    root = tree.getroot()
    for oper in root.iter("operator"):
        #print oper.attrib
        if oper.attrib["spec"] == "SubtreeSlide":
            if oper.attrib["id"] in ovs[0].keys():
                attr = oper.attrib["id"]
                value = ovs[0][attr]
                #print "found SubtreesSlide %s" % attr
                if "size" in oper.attrib.keys():
                    print "Updating SubtreeSlide 'size' attribute of %s from %s to %s " % \
                            (attr, oper.get("size"), value)
                else:
                    print "Adding SubtreeSlide 'size' attribute to %s with value %s" % \
                            (attr, value)
                oper.set("size", value)
                count_sts += 1
        else:
            if oper.attrib["id"] in ovs[1].keys():
                attr = oper.attrib["id"]
                value = ovs[1][attr]
                #print "found ScaleOperator %s" % oper.attrib["id"]
                print "Updating ScaleOperator of %s to %s" % (attr, value)
                oper.set("scaleFactor", value)
                count_so += 1

    if count_so != len(ovs[1].keys()):
        print "Failed to find all operators:"
        for oper in ovs[1]:
            print oper
            print "Exiting."
            sys.exit()
    if count_sts != len(ovs[0].keys()):
        print "Failed to find all operators:"
        for oper in ovs[0]:
            print oper
            print "Exiting."
            sys.exit()
    else:
        print "-="*40
        print "Updated %s subtreeSlides" % len(ovs[0])
        print "Updated %s scaleFactors" % len(ovs[1])
        tree.write(target + "-updated", xml_declaration=True, method="xml",
                encoding="UTF-8")
        print "Written new file: %s" % target + "-updated"
        print "Done."
    return

def main(filename, target):

    version = check_version()

    if not filename:
        print "Error: No filename specified."
        print "Aborting."
        sys.exit(1)
    else:
        if not os.path.exists(filename):
            print "Error: cannot find file %s" % filename
            print "Aborting."
            sys.exit(1)

    # Find the new values from log file
    ovs = collect_operator_values(filename)
    # Update the XML
    update_XML_runfile(ovs, target)

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=textwrap.dedent(desc),
            )
    parser.add_argument("logfile",
                        help="BEAST log filename",
                        )
    parser.add_argument("target_xml",
                        help="BEAST XML filename",
                        )
    args = parser.parse_args()
    main(args.logfile, args.target_xml)