scientific-programming-exercises/exam/ex13/main.py

import matplotlib.pyplot as plt
import numpy as np

from model import prepare_model
from backend import total_living_cells
from executor import execute_tick

def cells_alive_at_ticks(model, ticks):
	for i in range(ticks):
		cell_count = total_living_cells(model)
		yield [i, cell_count]
		execute_tick(model)

plot_data = {
	0.3 : ("r", "p=0.3")
	, 0.5: ("g", "p=0.5")
	, 0.7: ("b", "p=0.7")
}

handles = []
for p, (color, label) in plot_data.items():
	model = prepare_model(100, p)
	living_cells = np.array(list(cells_alive_at_ticks(model, 10)))
	plot, = plt.plot(living_cells[:, 0], living_cells[:, 1], color, label=label)
	handles.append(plot)

plt.legend(handles=handles)
plt.show()