class TestSingleNetworkGrid(unittest.TestCase): GRAPH_SIZE = 10 def setUp(self): ''' Create a test network grid and populate with Mock Agents. ''' G = nx.complete_graph(TestSingleNetworkGrid.GRAPH_SIZE) self.space = NetworkGrid(G) self.agents = [] for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE): a = MockAgent(i, None) self.agents.append(a) self.space.place_agent(a, pos) def test_agent_positions(self): ''' Ensure that the agents are all placed properly. ''' for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE): a = self.agents[i] assert a.pos == pos def test_get_neighbors(self): assert len(self.space.get_neighbors( 0, include_center=True)) == TestSingleNetworkGrid.GRAPH_SIZE assert len(self.space.get_neighbors( 0, include_center=False)) == TestSingleNetworkGrid.GRAPH_SIZE - 1 def test_move_agent(self): initial_pos = 1 agent_number = 1 final_pos = TestSingleNetworkGrid.GRAPH_SIZE - 1 _agent = self.agents[agent_number] assert _agent.pos == initial_pos assert _agent in self.space.G.node[initial_pos]['agent'] assert _agent not in self.space.G.node[final_pos]['agent'] self.space.move_agent(_agent, final_pos) assert _agent.pos == final_pos assert _agent not in self.space.G.node[initial_pos]['agent'] assert _agent in self.space.G.node[final_pos]['agent'] def test_is_cell_empty(self): assert not self.space.is_cell_empty(0) assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1) def test_get_cell_list_contents(self): assert self.space.get_cell_list_contents([0]) == [self.agents[0]] assert self.space.get_cell_list_contents( list(range(TestSingleNetworkGrid.GRAPH_SIZE))) == [ self.agents[0], self.agents[1], self.agents[2] ] def test_get_all_cell_contents(self): assert self.space.get_all_cell_contents() == [ self.agents[0], self.agents[1], self.agents[2] ]
class TestSingleNetworkGrid(unittest.TestCase): GRAPH_SIZE = 10 def setUp(self): ''' Create a test network grid and populate with Mock Agents. ''' G = nx.complete_graph(TestSingleNetworkGrid.GRAPH_SIZE) self.space = NetworkGrid(G) self.agents = [] for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE): a = MockAgent(i, None) self.agents.append(a) self.space.place_agent(a, pos) def test_agent_positions(self): ''' Ensure that the agents are all placed properly. ''' for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE): a = self.agents[i] assert a.pos == pos def test_get_neighbors(self): assert len(self.space.get_neighbors(0, include_center=True)) == TestSingleNetworkGrid.GRAPH_SIZE assert len(self.space.get_neighbors(0, include_center=False)) == TestSingleNetworkGrid.GRAPH_SIZE - 1 def test_move_agent(self): initial_pos = 1 agent_number = 1 final_pos = TestSingleNetworkGrid.GRAPH_SIZE - 1 _agent = self.agents[agent_number] assert _agent.pos == initial_pos assert _agent in self.space.G.node[initial_pos]['agent'] assert _agent not in self.space.G.node[final_pos]['agent'] self.space.move_agent(_agent, final_pos) assert _agent.pos == final_pos assert _agent not in self.space.G.node[initial_pos]['agent'] assert _agent in self.space.G.node[final_pos]['agent'] def test_is_cell_empty(self): assert not self.space.is_cell_empty(0) assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1) def test_get_cell_list_contents(self): assert self.space.get_cell_list_contents([0]) == [self.agents[0]] assert self.space.get_cell_list_contents(list(range(TestSingleNetworkGrid.GRAPH_SIZE))) == [self.agents[0], self.agents[1], self.agents[2]] def test_get_all_cell_contents(self): assert self.space.get_all_cell_contents() == [self.agents[0], self.agents[1], self.agents[2]]
class VirusModel(Model): """A virus model with some number of agents""" def __init__(self): # self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob) # self.G = nx.erdos_renyi_graph(n=3, p=0.5) self.G = nx.Graph() self.G.add_node(0) self.G.add_node(1) self.G.add_node(2) self.G.add_node(3) self.G.add_node(4) self.G.add_node(4) self.G.add_edge(0, 1) self.G.add_edge(0, 2) self.G.add_edge(0, 3) self.G.add_edge(0, 4) self.G.add_edge(0, 5) self.G.add_edge(1, 4) self.G.add_edge(4, 5) self.grid = NetworkGrid(self.G) self.rooms = {} self.rooms[0] = {"name": "Wejście", "rates": {}} self.rooms[1] = {"name": "Czytelnia", "rates": {"Nauka": 2}} self.rooms[2] = {"name": "Chillout", "rates": {"Relaks": 10}} self.rooms[3] = {"name": "Biuro", "rates": {"Praca": 1.5}} self.rooms[4] = {"name": "Toaleta", "rates": {"Toaleta": 30}} self.rooms[5] = { "name": "Kawiarnia", "rates": { "Jedzenie": 12, "Kultura": 0.5 } } collector_dict = {} for i, room in enumerate(self.rooms): collector_dict[self.rooms[i]["name"]] = lambda model, i=i: len( model.grid.get_cell_list_contents([i])) - 1 self.datacollector = DataCollector(collector_dict) self.schedule = RandomActivation(self) # Create agents for i, node in enumerate(self.G.nodes()): r = RoomAgent(i, self, self.rooms[i]["name"], self.rooms[i]["rates"]) self.schedule.add(r) # Add the agent to the node self.grid.place_agent(r, node) self.prob_needs = { "Jedzenie": [4, 0.6], "Toaleta": [2, 0.6], "Relaks": [5, 1] } self.prob_studs = { "Nauka": [2, 1.5], "Praca": [0, 0.5], "Kultura": [0, 1.0] } self.prob_works = { "Nauka": [0, 0.3], "Praca": [6, 1.0], "Kultura": [0, 0.2] } self.prob_tours = { "Nauka": [0, 0.3], "Praca": [0, 0.5], "Kultura": [1, 1.0] } self.prob_local = { "Nauka": [1, 0.7], "Praca": [2, 0.9], "Kultura": [1, 1.0] } # godziny 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 self.rate_studs = [ 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0 ] self.rate_works = [ 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0 ] self.rate_tours = [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 3, 3, 3, 3, 4, 4, 4, 6, 6, 4, 3, 2, 0, 0 ] self.rate_local = [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 4, 4, 2, 2, 4, 5, 6, 6, 6, 3, 0, 0, 0 ] self.running = True self.datacollector.collect(self) self.tm = 0 * 60 self.count = 0 def get_sample(self, probs): ret = {} for k, [m, s] in probs.items(): tm = int(np.clip(np.random.normal(m, s) * 60, 15, 600)) ret[k] = tm return ret def step(self): # prepare list for the satisfied agents self.satisfied = [] # add new agents hour = int(self.tm / 60) if (hour > 23): hour = 0 for i in range(self.rate_studs[hour]): a = HumanAgent(100 + self.count, self, self.get_sample(self.prob_needs), self.get_sample(self.prob_studs)) self.schedule.add(a) self.grid.place_agent(a, 0) self.count += 1 for i in range(self.rate_works[hour]): a = HumanAgent(100 + self.count, self, self.get_sample(self.prob_needs), self.get_sample(self.prob_works)) self.schedule.add(a) self.grid.place_agent(a, 0) self.count += 1 # update system self.schedule.step() # collect data self.datacollector.collect(self) # make time step self.tm = self.tm + 1 if (self.tm > 24 * 60): self.datacollector.get_model_vars_dataframe().to_csv("one_day.csv") self.tm = 0 # remove satisfied agents from the system for a in self.satisfied: print(a.unique_id, a.goals, "is satisfied") self.grid.move_agent(a, 0) self.grid._remove_agent(a, 0) self.schedule.remove(a) def run_model(self, n): for i in range(n): self.step() def find_best_room(self, goal): #print("Looking for room for", goal) for i, room in enumerate(self.rooms): #print("Room", room, self.rooms[room]["rates"]) if goal in self.rooms[room]["rates"]: return room return -1