diff --git a/mesa/examples/advanced/wolf_sheep/model.py b/mesa/examples/advanced/wolf_sheep/model.py index c92214f6cc9..0d06a707cf3 100644 --- a/mesa/examples/advanced/wolf_sheep/model.py +++ b/mesa/examples/advanced/wolf_sheep/model.py @@ -22,37 +22,6 @@ from mesa.experimental.devs import ABMSimulator -def is_trapped_in_wall( - cell, wall_coord, width, height -): # true if cell is trapped of walls - north = (cell.coordinate[0] - 1, cell.coordinate[1]) - south = (cell.coordinate[0] + 1, cell.coordinate[1]) - east = (cell.coordinate[0], cell.coordinate[1] + 1) - west = (cell.coordinate[0], cell.coordinate[1] - 1) - - coord = (cell.coordinate[0], cell.coordinate[1]) - - # 'corner' cases (pun intended) - if coord == (0, 0): # top left corner - return {east, south}.issubset(wall_coord) - if coord == (height - 1, 0): # bottom left corner - return {north, east}.issubset(wall_coord) - if coord == (0, width - 1): # top right corner - return {west, south}.issubset(wall_coord) - if coord == (height - 1, width - 1): # bottom right corner - return {north, west}.issubset(wall_coord) - if coord[0] == 0: # for cells at top row - return {south, west, east}.issubset(wall_coord) - if coord[1] == 0: # for cells at leftmost column - return {north, south, east}.issubset(wall_coord) - if coord[0] == height - 1: # for cells at the bottom row - return {north, east, west}.issubset(wall_coord) - if coord[1] == width - 1: # for cells at rightmost column - return {north, south, west}.issubset(wall_coord) - - return {north, south, west, east}.issubset(wall_coord) - - class WolfSheep(Model): """Wolf-Sheep Predation Model. @@ -140,17 +109,25 @@ def __init__( self.grid.add_property_layer(PropertyLayer.from_data("wall", wall_arr)) - possible_cells = [] - for cell in self.grid.all_cells.cells: - if ( - ( - cell.coordinate[0], - cell.coordinate[1], - ) - not in wall_coord - and not is_trapped_in_wall(cell, wall_coord, width, height) - ): # so we don't create an animal at wall cells. and make sure the animal is not trapped in walls - possible_cells.append(cell) + def is_wall(row, col): + return ( + True + if row < 0 or col < 0 or row >= height or col >= width # corner case + else wall_arr[row][col] + ) + + def is_trapped_in_walls(row, col): + return ( + is_wall(row + 1, col) + and is_wall(row - 1, col) + and is_wall(row, col + 1) + and is_wall(row, col - 1) + ) + + possible_cells = self.grid.all_cells.select( + lambda cell: not wall_arr[cell.coordinate[0]][cell.coordinate[1]] + and not is_trapped_in_walls(cell.coordinate[0], cell.coordinate[1]) + ).cells # so we don't create an animal at wall cells. and make sure the animal is not trapped in walls # Create sheep: Sheep.create_agents(