Python datacollection.DataCollector类代码示例

class MoneyModel(Model):
    """A model with some number of agents."""
    def __init__(self, N, width, height):
        self.num_agents = N
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(model_reporters={"Gini": compute_gini},
                agent_reporters={"Wealth": lambda a: a.wealth})
        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i)
            self.schedule.add(a)
            # Add the agent to a random grid cell
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()
    
    def run_model(self, n):
        for i in range(n):
            self.step()

class Money_Model(Model):
    def __init__(self, N, width=50, height=50, torus=True):
        self.num_agents = N
        self.schedule = RandomActivation(self)
        self.grid = MultiGrid(height, width, torus)
        self.create_agents()
        self.dc = DataCollector({"Gini": lambda m: m.compute_gini()},
                               {"Wealth": lambda a: a.wealth})
        self.running = True

    def create_agents(self):
        for i in range(self.num_agents):
            a = Money_Agent(i)
            self.schedule.add(a)
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.dc.collect(self)
        self.schedule.step()
        
    def run_model(self, steps):
        for i in range(steps):
            self.step()
    
    def compute_gini(self):
        agent_wealths = [agent.wealth for agent in self.schedule.agents]
        x = sorted(agent_wealths)
        N = self.num_agents
        B = sum( xi * (N-i) for i,xi in enumerate(x) ) / (N*sum(x))
        return (1 + (1/N) - 2*B)

class BoltzmannWealthModelNetwork(Model):
    """A model with some number of agents."""

    def __init__(self, num_agents=7, num_nodes=10):

        self.num_agents = num_agents
        self.num_nodes = num_nodes if num_nodes >= self.num_agents else self.num_agents
        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=0.5)
        self.grid = NetworkGrid(self.G)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(
            model_reporters={"Gini": compute_gini},
            agent_reporters={"Wealth": lambda _: _.wealth}
        )

        list_of_random_nodes = self.random.sample(self.G.nodes(), self.num_agents)

        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i, self)
            self.schedule.add(a)
            # Add the agent to a random node
            self.grid.place_agent(a, list_of_random_nodes[i])

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self, n):
        for i in range(n):
            self.step()

class MoneyModel(Model):
    """A simple model of an economy where agents exchange currency at random.

    All the agents begin with one unit of currency, and each time step can give
    a unit of currency to another agent. Note how, over time, this produces a
    highly skewed distribution of wealth.
    """

    def __init__(self, N, width, height):
        self.num_agents = N
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(
            model_reporters={"Gini": compute_gini},
            agent_reporters={"Wealth": lambda a: a.wealth}
        )
        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i, self)
            self.schedule.add(a)
            # Add the agent to a random grid cell
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()

    def run_model(self, n):
        for i in range(n):
            self.step()

class WorldModel(Model):
    def __init__(self, N, width, height):
        self.grid = SingleGrid(height, width, True)

        self.schedule = RandomActivation(self)
        self.num_agents = N
        self.running = True
        
        for i in range(self.num_agents):
            ethnicity = random.choice(Ethnicities)
            a = PersonAgent(unique_id=i,
                            model=self,
                            ethnicity=int(ethnicity)
                            )
            self.schedule.add(a)
            # Add the agent to a random grid cell

            self.grid.position_agent(a)
            
        self.datacollector = DataCollector(
            agent_reporters={
                "Nationalism": lambda a: a.nationalism,
                "X": lambda a: a.pos[0],
                "Y": lambda a: a.pos[1]
            }
        )

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()

class SchellingModel(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height, width, density, minority_pc, homophily):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=True)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": lambda m: m.happy}, # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.x, "y": lambda a: a.y}) 

        self.running = True

        # Set up agents
        for x in range(self.width):
            for y in range(self.height):
                if random.random() < self.density:
                    if random.random() < self.minority_pc:
                        agent_type = 1
                    else:
                        agent_type = 0

                    agent = SchellingAgent((x,y), x, y, agent_type)
                    self.grid[y][x] = agent
                    self.schedule.add(agent)

    def get_empty(self):
        '''
        Get a list of coordinate tuples of currently-empty cells.
        '''
        empty_cells = []
        for x in range(self.width):
            for y in range(self.height):
                if self.grid[y][x] is None:
                    empty_cells.append((x, y))
        return empty_cells

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0 # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False

class ForestFire(Model):
    '''
    Simple Forest Fire model.
    '''
    def __init__(self, height, width, density):
        '''
        Create a new forest fire model.

        Args:
            height, width: The size of the grid to model
            density: What fraction of grid cells have a tree in them.
        '''
        # Initialize model parameters
        self.height = height
        self.width = width
        self.density = density

        # Set up model objects
        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=False)

        self.datacollector = DataCollector(
            {"Fine": lambda m: self.count_type(m, "Fine"),
             "On Fire": lambda m: self.count_type(m, "On Fire"),
             "Burned Out": lambda m: self.count_type(m, "Burned Out")})

        # Place a tree in each cell with Prob = density
        for (contents, x, y) in self.grid.coord_iter():
            if random.random() < self.density:
                # Create a tree
                new_tree = TreeCell((x, y))
                # Set all trees in the first column on fire.
                if x == 0:
                    new_tree.condition = "On Fire"
                self.grid._place_agent((x, y), new_tree)
                self.schedule.add(new_tree)
        self.running = True

    def step(self):
        '''
        Advance the model by one step.
        '''
        self.schedule.step()
        self.datacollector.collect(self)

        # Halt if no more fire
        if self.count_type(self, "On Fire") == 0:
            self.running = False

    @staticmethod
    def count_type(model, tree_condition):
        '''
        Helper method to count trees in a given condition in a given model.
        '''
        count = 0
        for tree in model.schedule.agents:
            if tree.condition == tree_condition:
                count += 1
        return count

class Charts(Model):

    # grid height
    grid_h = 20
    # grid width
    grid_w = 20

    """init parameters "init_people", "rich_threshold", and "reserve_percent"
       are all UserSettableParameters"""
    def __init__(self, height=grid_h, width=grid_w, init_people=2, rich_threshold=10,
                 reserve_percent=50,):
        self.height = height
        self.width = width
        self.init_people = init_people
        self.schedule = RandomActivation(self)
        self.grid = MultiGrid(self.width, self.height, torus=True)
        # rich_threshold is the amount of savings a person needs to be considered "rich"
        self.rich_threshold = rich_threshold
        self.reserve_percent = reserve_percent
        # see datacollector functions above
        self.datacollector = DataCollector(model_reporters={
                                           "Rich": get_num_rich_agents,
                                           "Poor": get_num_poor_agents,
                                           "Middle Class": get_num_mid_agents,
                                           "Savings": get_total_savings,
                                           "Wallets": get_total_wallets,
                                           "Money": get_total_money,
                                           "Loans": get_total_loans},
                                           agent_reporters={
                                           "Wealth": lambda x: x.wealth})

        # create a single bank for the model
        self.bank = Bank(1, self, self.reserve_percent)

        # create people for the model according to number of people set by user
        for i in range(self.init_people):
            # set x, y coords randomly within the grid
            x = self.random.randrange(self.width)
            y = self.random.randrange(self.height)
            p = Person(i, (x, y), self, True, self.bank, self.rich_threshold)
            # place the Person object on the grid at coordinates (x, y)
            self.grid.place_agent(p, (x, y))
            # add the Person object to the model schedule
            self.schedule.add(p)

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        # tell all the agents in the model to run their step function
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self):
        for i in range(self.run_time):
            self.step()

class Foraging(Model):
    
    number_of_bean = 0
    number_of_corn = 0
    number_of_soy = 0
    
    def __init__(self, width=50, height=50, torus=True, num_bug=50, seed=42, strategy=None):
        super().__init__(seed=seed)
        self.number_of_bug = num_bug
        if not(strategy in ["stick", "switch"]):
            raise TypeError("'strategy' must be one of {stick, switch}")
        self.strategy = strategy
        
        self.grid = SingleGrid(width, height, torus)
        self.schedule = RandomActivation(self)
        data = {"Bean": lambda m: m.number_of_bean,
                "Corn": lambda m: m.number_of_corn,
                "Soy": lambda m: m.number_of_soy,
                "Bug": lambda m: m.number_of_bug,
                }
        self.datacollector = DataCollector(data)
        
        # create foods
        self._populate(Bean)
        self._populate(Corn)
        self._populate(Soy)
        
        # create bugs
        for i in range(self.number_of_bug):
            pos = self.grid.find_empty()
            bug = Bug(i, self)
            bug.strategy = self.strategy
            self.grid.place_agent(bug, pos)
            self.schedule.add(bug)
    
    def step(self):
        self.schedule.step()
        self.datacollector.collect(self)
        
        if not(self.grid.exists_empty_cells()):
            self.running = False
    
    def _populate(self, food_type):
        prefix = "number_of_{}"
        
        counter = 0
        while counter < food_type.density * (self.grid.width * self.grid.height):
            pos = self.grid.find_empty()
            food = food_type(counter, self)
            self.grid.place_agent(food, pos)
            self.schedule.add(food)
            food_name = food_type.__name__.lower()
            attr_name = prefix.format(food_name)
            val = getattr(self, attr_name)
            val += 1
            setattr(self, attr_name, val)
            counter += 1

class Schelling(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height=20, width=20, density=0.8, minority_pc=0.2, homophily=3):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=True)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": "happy"},  # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]})

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)
        for cell in self.grid.coord_iter():
            x = cell[1]
            y = cell[2]
            if self.random.random() < self.density:
                if self.random.random() < self.minority_pc:
                    agent_type = 1
                else:
                    agent_type = 0

                agent = SchellingAgent((x, y), self, agent_type)
                self.grid.position_agent(agent, (x, y))
                self.schedule.add(agent)

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False

class SchellingModel(Model):
    """
    Model class for the Schelling segregation model.
    """

    def __init__(self, height, width, density, minority_pc, homophily):
        """
        """

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=True)

        self.happy = 0
        self.total_agents = 0
        self.datacollector = DataCollector(
            {"unhappy": lambda m: m.total_agents - m.happy},
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[X], "y": lambda a: a.pos[Y]},
        )

        self.running = True

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)
        for cell, x, y in self.grid.coord_iter():
            if random.random() < self.density:
                if random.random() < self.minority_pc:
                    agent_type = 1
                else:
                    agent_type = 0

                agent = SchellingAgent(self.total_agents, agent_type)
                self.grid.position_agent(agent, x, y)
                self.schedule.add(agent)
                self.total_agents += 1

    def step(self):
        """
        Run one step of the model. If All agents are happy, halt the model.
        """
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.total_agents:
            self.running = False

class PD_Model(Model):
    '''
    Model class for iterated, spatial prisoner's dilemma model.
    '''

    schedule_types = {"Sequential": BaseScheduler,
                      "Random": RandomActivation,
                      "Simultaneous": SimultaneousActivation}

    # This dictionary holds the payoff for this agent,
    # keyed on: (my_move, other_move)

    payoff = {("C", "C"): 1,
              ("C", "D"): 0,
              ("D", "C"): 1.6,
              ("D", "D"): 0}

    def __init__(self, height, width, schedule_type, payoffs=None):
        '''
        Create a new Spatial Prisoners' Dilemma Model.

        Args:
            height, width: Grid size. There will be one agent per grid cell.
            schedule_type: Can be "Sequential", "Random", or "Simultaneous".
                           Determines the agent activation regime.
            payoffs: (optional) Dictionary of (move, neighbor_move) payoffs.
        '''
        self.running = True
        self.grid = SingleGrid(height, width, torus=True)
        self.schedule_type = schedule_type
        self.schedule = self.schedule_types[self.schedule_type](self)

        # Create agents
        for x in range(width):
            for y in range(height):
                agent = PD_Agent((x, y), self)
                self.grid.place_agent(agent, (x, y))
                self.schedule.add(agent)

        self.datacollector = DataCollector({
            "Cooperating_Agents":
            lambda m: len([a for a in m.schedule.agents if a.move == "C"])
        })

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()

    def run(self, n):
        '''
        Run the model for a certain number of steps.
        '''
        for _ in range(n):
            self.step()

class VirusOnNetwork(Model):
    """A virus model with some number of agents"""

    def __init__(self, num_nodes=10, avg_node_degree=3, initial_outbreak_size=1, virus_spread_chance=0.4,
                virus_check_frequency=0.4, recovery_chance=0.3, gain_resistance_chance=0.5):

        self.num_nodes = num_nodes
        prob = avg_node_degree / self.num_nodes
        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
        self.grid = NetworkGrid(self.G)
        self.schedule = RandomActivation(self)
        self.initial_outbreak_size = initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
        self.virus_spread_chance = virus_spread_chance
        self.virus_check_frequency = virus_check_frequency
        self.recovery_chance = recovery_chance
        self.gain_resistance_chance = gain_resistance_chance

        self.datacollector = DataCollector({"Infected": number_infected,
                                            "Susceptible": number_susceptible,
                                            "Resistant": number_resistant})

        # Create agents
        for i, node in enumerate(self.G.nodes()):
            a = VirusAgent(i, self, State.SUSCEPTIBLE, self.virus_spread_chance, self.virus_check_frequency,
                           self.recovery_chance, self.gain_resistance_chance)
            self.schedule.add(a)
            # Add the agent to the node
            self.grid.place_agent(a, node)

        # Infect some nodes
        infected_nodes = self.random.sample(self.G.nodes(), self.initial_outbreak_size)
        for a in self.grid.get_cell_list_contents(infected_nodes):
            a.state = State.INFECTED

        self.running = True
        self.datacollector.collect(self)

    def resistant_susceptible_ratio(self):
        try:
            return number_state(self, State.RESISTANT) / number_state(self, State.SUSCEPTIBLE)
        except ZeroDivisionError:
            return math.inf

    def step(self):
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self, n):
        for i in range(n):
            self.step()

class InspectionModel(Model):
    '''
    Simple Restaurant Inspection model.
    '''
    def __init__(self, height, width, density):
        '''
        Create a new restaurant inspection model.

        Args:
            height, width: The size of the grid to model
            density: What fraction of grid cells have a restaurant in them.
        '''
        # Initialize model parameters
        self.height = height
        self.width = width
        self.density = density

        # Set up model objects
        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=False)

        self.datacollector = DataCollector(
            {"Good": lambda m: self.count_type(m, "Good"),
             "Bad": lambda m: self.count_type(m, "Bad")})

        # Place a restaurant in each cell with Prob = density
        for (contents, x, y) in self.grid.coord_iter():
            if random.random() < self.density:
                # Create a restaurant
                new_restaurant = RestaurantCell((x, y))
                self.grid._place_agent((x, y), new_restaurant)
                self.schedule.add(new_restaurant)
        self.running = True

    def step(self):
        '''
        Advance the model by one step.
        '''
        self.schedule.step()
        self.datacollector.collect(self)

    @staticmethod
    def count_type(model, restaurant_hygiene):
        '''
        Helper method to count restaurants in a given condition in a given model.
        '''
        count = 0
        for restaurant in model.schedule.agents:
            if restaurant.hygiene == restaurant_hygiene and restaurant.rating != 'Closed':
                count += 1
        return count

class SchellingModel(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height, width, density, type_pcs=[.2, .2, .2, .2, .2]):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.type_pcs = type_pcs

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=False)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": lambda m: m.happy},  # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]})

        self.running = True

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)

        total_agents = self.height * self.width * self.density
        agents_by_type = [total_agents*val for val in self.type_pcs]

        for loc, types in enumerate(agents_by_type):
            for i in range(int(types)):
                pos = self.grid.find_empty()
                agent = SchellingAgent(pos, self, loc)
                self.grid.position_agent(agent, pos)
                self.schedule.add(agent)

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False