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

1import matplotlib.pyplot as plt 

2from .util import parse_data 

3 

4def 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