Coverage for agent_model/viz.py: 0%
26 statements
« prev ^ index » next coverage.py v7.2.3, created at 2023-05-04 13:20 +0700
« prev ^ index » next coverage.py v7.2.3, created at 2023-05-04 13:20 +0700
1import matplotlib.pyplot as plt
2from .util import parse_data
4def plot_agent(data, agent, category, exclude=[], include=[], i=None, j=None, ax=None):
5 """Helper function for plotting model data
7 Plotting function which takes model-exported data, agent name,
8 one of (flows, growth, storage, deprive), exclude, and i:j
9 """
10 i = i if i is not None else 0
11 j = j if j is not None else data['step_num'][-1]
12 ax = ax if ax is not None else plt
13 if category == 'flows':
14 path = [agent, 'flows', '*', '*', 'SUM', f'{i}:{j}']
15 flows = parse_data(data['agents'], path)
16 for direction in ('in', 'out'):
17 if direction not in flows:
18 continue
19 for currency, values in flows[direction].items():
20 label = f'{direction}_{currency}'
21 if currency in exclude or label in exclude:
22 continue
23 ax.plot(range(i, j), values, label=label)
24 elif category in {'storage', 'attributes'}:
25 path = [agent, category, '*', f'{i}:{j}']
26 parsed = parse_data(data['agents'], path)
27 for field, values in parsed.items():
28 if field in exclude or (include and field not in include):
29 continue
30 ax.plot(range(i, j), values, label=field)
31 ax.legend()
32 return ax