BacterAI/make_figure_growth_front.py at master · jensenlab/BacterAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import argparse
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import utils

plt.rcParams["font.family"] = "sans-serif"
plt.rcParams["font.sans-serif"] = "Arial"
plt.rcParams["svg.fonttype"] = "none"

def main(
    bacterai_data, random_data, experiment_folder, n_bins
):
    fig, axs = plt.subplots(nrows=1, ncols=1, figsize=(10, 6))

    width = 0.5

    bins = np.arange(0, 1.01, 1 / n_bins)
    rand, _ = np.histogram(random_data["fitness"], bins)
    bact, _ = np.histogram(bacterai_data["fitness"], bins)

    rand = rand / len(random_data)
    bact = bact / len(bacterai_data)

    x = np.arange(n_bins)
    r1 = axs.bar(
        x + width / 2,
        rand,
        width,
        color="dodgerblue",
    )
    r2 = axs.bar(
        x + 1.5 * width,
        bact,
        width,
        color="k",
    )
    axs.set_ylabel("Count Density")
    axs.set_xlabel("Fitness")
    axs.bar_label(r1, padding=2, fmt="%.2f", fontsize=7.5, rotation=90)
    axs.bar_label(r2, padding=2, fmt="%.2f", fontsize=7.5, rotation=90)

    bin_labels = [f"{x:.2f}" for x in np.arange(0, 1.01, 1 / n_bins)]
    print(bin_labels)
    axs.set_xticks(np.arange(0, n_bins + 1, 1))
    axs.set_xticklabels(bin_labels)
    plt.xticks(rotation=45)
    axs.set_yscale('log')
    for tick in axs.xaxis.get_majorticklabels():
        tick.set_horizontalalignment("right")

    plt.legend(["Random", "BacterAI"])
    fig.tight_layout()

    fig_path = os.path.join(experiment_folder, "growth_distribution_hist")
    for file_ext in (".png", ".svg"):
        plt.savefig(f"{fig_path}{file_ext}", dpi=400)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="BacterAI Growth Front Figure")

    parser.add_argument(
        "path",
        type=str,
        help="The path to the experiment folder",
    )

    parser.add_argument(
        "-r",
        "--rounds",
        type=int,
        required=True,
        help="The number of rounds to plot out starting at 1",
    )

    parser.add_argument(
        "-num",
        "--num_ingredients",
        type=int,
        required=False,
        default=20,
        help="The number of experiment ingredients",
    )

    args = parser.parse_args()

    # Second plot
    data = utils.combined_round_data(args.path, max_n=args.rounds)

    # originally Randoms (1) SGO CH1 17f3 mapped_data
    random_path = "published_data/SGO random/experiment_data.csv"
    rand_data = pd.read_csv(random_path, index_col=None)

    if "is_redo" in rand_data.columns:
        rand_data = rand_data[~rand_data["is_redo"]]

    main(data, rand_data, args.path, args.num_ingredients)