Generate iTOL "binary data" annotation file from pandas dataframe

itol_binary.py

import numpy as np
import pandas as pd
from ete3 import PhyloTree

def generate_binary(df: pd.DataFrame, save_file: str) -> None:
    # df columns must be labels, values, values, ..., values.
    num_cols = df.shape[1]-1
    shapes = ",".join(["1"]*num_cols)
    colors = ",".join(["#ff0000"]*num_cols)
    labels = ",".join([f"{i+1}" for i in range(num_cols)])
    output = f"""DATASET_BINARY\n
SEPARATOR COMMA\n
DATASET_LABEL,label_binary\n
COLOR,{colors}\n
FIELD_SHAPES,{shapes}\n
FIELD_LABELS,{labels}\n
DATA\n""" + df.to_csv(header=False, index=False)


    with open(save_file, "w") as f:
        f.write(output)
# tree = PhyloTree("tree.nwk")
# leaf_labels = [leaf.name for leaf in tree.get_leaves()]
leaf_labels = ["A", "B", "D", "E"]
matrix = np.random.randint(0, 2, size=(len(leaf_labels),100))

generate_binary(pd.concat([pd.Series(leaf_labels), pd.DataFrame(matrix)], axis=1), "binary.txt")