1 from collections import defaultdict

2 from . import BaseAgent

3

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