Coverage for agent_model/agents/atmosphere_equalizer.py: 100%

37 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-05-04 13:14 +0700

1from collections import defaultdict 

2from . import BaseAgent 

3 

4class AtmosphereEqualizerAgent(BaseAgent): 

5 def __init__(self, *args, **kwargs): 

6 # -- NON_SERIALIZED 

7 self.atms = {} 

8 super().__init__(*args, **kwargs) 

9 

10 def register(self, record_initial_state=True): 

11 self.atms = {a: self.model.agents[a] for a in self.flows['in']['atmosphere']['connections']} 

12 for agent_id in self.atms.keys(): 

13 for direction in ('in', 'out'): 

14 conns = self.flows[direction]['atmosphere']['connections'] 

15 if agent_id not in conns: 

16 conns.append(agent_id) 

17 super().register(record_initial_state) 

18 

19 def step(self, dT=1): 

20 if not self.registered: 

21 self.register() 

22 volumes = {} # agent_type: m3 

23 current = {} # agent_type: {atmo_currency: kg} 

24 total_atm = defaultdict(float) # atmo_currency: kg 

25 for agent_id, agent in self.atms.items(): 

26 volumes[agent_id] = agent.properties['volume']['value'] * agent.amount 

27 current[agent_id] = agent.view(view='atmosphere') 

28 for currency, amount in current[agent_id].items(): 

29 total_atm[currency] += amount 

30 

31 total_volume = sum(volumes.values()) 

32 for agent_id, agent in self.atms.items(): 

33 atm_ratio = volumes[agent_id] / total_volume 

34 targets = {k: v * atm_ratio for k, v in total_atm.items()} 

35 deltas = {k: v - current[agent_id][k] for k, v in targets.items()} 

36 for currency, delta in deltas.items(): 

37 if delta != 0: 

38 # TODO: Add these to flows records 

39 agent.increment(currency, delta) 

40 inflow = abs(max(0, delta)) 

41 outflow = abs(min(0, delta)) 

42 self.records['flows']['in'][currency][agent_id].append(inflow) 

43 self.records['flows']['out'][currency][agent_id].append(outflow) 

44