structures.py

import matplotlib.pyplot as plt
import numpy as np
from mesa import Model, Agent
from mesa.time import RandomActivation
import random
from mesa.space import SingleGrid
from matplotlib.animation import FuncAnimation
from matplotlib.colors import ListedColormap



class SchellingAgent(Agent):
    def __init__(self, unique_id, model, group, tolerance):
        super().__init__(unique_id, model)
        self.group = group
        self.tolerance = tolerance
        self.position = None
        self.satisfied = None

    def set_position(self, position):
        self.position = position

    def fraction_of_similar_neighbors(self):
        """Count the number of similar agents in the neighborhood."""
        count = 0
        neighbors = self.model.grid.get_neighbors(self.position, moore = True)
        total = len(list(neighbors))

        if total ==0:
            return 1
        for neighbor in neighbors:
            if neighbor.group == self.group:
                count += 1
        return count/total
    
    def step(self):
        s_i = self.fraction_of_similar_neighbors()
        if s_i >= self.tolerance: 
            self.satisfied = 1
        else:
            self.satisfied = 0 

        if s_i < 1.0: # J'ai au moins un étranger dans mes voisins
            if not self.satisfied:
                find = False
                last_position  = self.position

                empty_positions = [(i, j) for i in range(self.model.grid.width) for j in range(self.model.grid.height) if self.model.grid.is_cell_empty((i, j))]

                for new_position in random.sample(empty_positions, len(empty_positions)):
                    
                    self.set_position(new_position)  
                    
                    if self.fraction_of_similar_neighbors() >= self.tolerance:
                        find = True
                        self.model.grid.remove_agent(self)
                        self.model.grid.place_agent(self, new_position)
                        break  
                if not find:
                    self.set_position(last_position)
            else:
                self.tolerance = min(self.model.df+self.tolerance, self.model.fmax)
                prob = random.uniform(0, 1)
                if prob < self.model.p_move:
                    empty_positions = [(i, j) for i in range(self.model.grid.width) for j in range(self.model.grid.height) if self.model.grid.is_cell_empty((i, j))]

                    for new_position in random.sample(empty_positions, len(empty_positions)):
                        self.model.grid.remove_agent(self)
                        self.model.grid.place_agent(self, new_position)
                        self.set_position(new_position)  
        else:
            self.tolerance = max(self.tolerance-self.model.df, self.model.fmin)


        



class SchellingModel(Model):
    def __init__(self, width, height, N, initial_tolerance, df, p_move, fmin, fmax):
        self.fmax = fmax
        self.fmin = fmin
        self.num_agents = N  # Total number of agents
        self.df = df # Increase tolerance step 
        self.p_move = p_move # Probability of randomly moving 
        self.grid = SingleGrid(width, height, torus=True)  # Create a grid for agent movement
        self.schedule = RandomActivation(self)  # Schedule for agent activation


      # Create agents and place them on the grid
        for i in range(self.num_agents):
            group = 'native' if i < N // 2 else 'immigrant'  # Example: half natives, half immigrants
            agent = SchellingAgent(i, self, group, initial_tolerance)
            self.schedule.add(agent)

            # Place agent on a random empty cell in the grid
            placed = False
            while not placed:
                x = self.random.randrange(self.grid.width)
                y = self.random.randrange(self.grid.height)
                if self.grid.is_cell_empty((x, y)):  # Check if the cell is empty
                    self.grid.place_agent(agent, (x, y))
                    agent.set_position((x, y))
                    placed = True  # Mark as placed

    def step(self):
        """Advance the model by one step."""
        self.schedule.step()  # Call the step method for each agent

    def get_agent_count(self, group):
        """Get the count of agents in a specific group."""
        count = sum(1 for agent in self.schedule.agents if agent.group == group)
        return count


    def display_grid(self):
        """Display the grid with agents."""
        grid_array = np.zeros((self.grid.width, self.grid.height))  

        for agent in self.schedule.agents:
            x, y = agent.position
            if agent.group == 'native':
                grid_array[x, y] = 1  
            else:
                grid_array[x, y] = 2  

        cmap = ListedColormap(['white', 'blue', 'red'])

        plt.imshow(grid_array, cmap=cmap, origin='upper')
        plt.colorbar(ticks=[0, 1, 2], label='Agent Group')
        plt.title("Initial State of the Schelling Model")

        plt.scatter([], [], c='blue', label='Native') 
        plt.scatter([], [], c='red', label='Immigrant')  
        plt.legend(title="Agent Types")
        
        plt.show()
    
    def run_model(self, n):
        """Run the model for a specified number of steps."""
        for _ in range(n):
            self.step()