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
« prev ^ index » next coverage.py v7.2.3, created at 2023-05-04 13:14 +0700
1from collections import defaultdict
2from . import BaseAgent
4class AtmosphereEqualizerAgent(BaseAgent):
5 def __init__(self, *args, **kwargs):
6 # -- NON_SERIALIZED
7 self.atms = {}
8 super().__init__(*args, **kwargs)
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)
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
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)