1 import matplotlib.pyplot as plt

2 from .util import parse_data

3

4 def plot_agent(data, agent, category, exclude=[], include=[], i=None, j=None, ax=None):

5 """Helper function for plotting model data

6

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