How to add branch labels from node data

Help and Getting Started
1

User question, originally submitted via email

I’ve imported a beast tree using augur import beast following this tutorial and would like to some of the metadata as branch labels. I followed the example workflow and successfully created an auspice.json tree using augur export v2, and if I click on the tips, I see metadata I imported from the beast tree - but I would like the metadata to show as branch labels, like this:

image
image1490×664 40.8 KB

How can I achieve this? I know that the information is in the beast_data.json that I pass to augur export v2, but I couldn’t find an option how to specify what to add as branch labels.

1 Like
2

As far as I’m aware, augur export doesn’t yet have the ability to easily add branch labels.

But with a small script, this is quite easy.

Here’s how to add a branch label based on the information in the beast_data.json that contains metadata and is produced by augur import beast:

python add-branch-labels.py \
  --input-tree auspice.json \
  --input-node-data beast_data.json \
  --attribute posterior \
  --output auspice_with_extra_branch_label.json

The script is pasted at the bottom.

For input tree use the tree as it comes out of augur export v2, the beast_data.json is the metadata file that comes out of augur import beast. attribute is what you want to turn from metadata into a branch label. The final tree will be at what ever you pass to --output.

If you want to make the branch labels automatically selected - rather than None being selected as branch label by adding or changing the “display_defaults”:

"display_defaults": {
   "branch_label": "posterior"  # Add this line to set default branch label
},

Here’s the script:

"""
add-branch-labels.py
"""

import argparse
import json


def extract_attribute(node_data, attribute: str) -> dict[str, str]:
    data = {}
    for name, node in node_data["nodes"].items():
        if attribute in node:
            data[name] = node[attribute]
    return data


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Add extra branch labels",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--input-tree",
        type=str,
        metavar="JSON",
        required=True,
        help="input Auspice JSON",
    )
    parser.add_argument(
        "--input-node-data", type=str, required=True, help="node data file"
    )
    parser.add_argument(
        "--attribute",
        type=str,
        required=True,
        help="Attribute from node data to add as branch label",
    )
    parser.add_argument(
        "--output",
        type=str,
        metavar="JSON",
        required=True,
        help="output Auspice JSON",
    )
    args = parser.parse_args()

    with open(args.input_tree, "r") as f:
        auspice_json = json.load(f)

    with open(args.input_node_data, "r") as f:
        node_data = extract_attribute(json.load(f), args.attribute)

    def attach_labels(n: dict[str,dict]):  # closure
        if n["name"] in node_data:
            if "branch_attrs" not in n:
                n["branch_attrs"] = {}
            if "labels" not in n["branch_attrs"]:
                n["branch_attrs"]["labels"] = {}
            n["branch_attrs"]["labels"][
                args.attribute
            ] = node_data[n["name"]]

        if "children" in n:
            for c in n["children"]:
                attach_labels(c)

    attach_labels(auspice_json["tree"])

    with open(args.output, "w") as f:
        json.dump(auspice_json, f, indent=2)