From 478506b6fe5ffb4455b2df5737298e27364311ef Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Wed, 6 Dec 2023 08:54:31 +0100 Subject: [PATCH] _PropertyGrid: Give option to return list or mask Give the select_cells_multi_properties method an option to return a mask instead of a list of cells. list is still default. --- mesa/space.py | 39 ++++++++++++++++++++++++++------------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/mesa/space.py b/mesa/space.py index d54c718b621..8d9d2e8f587 100644 --- a/mesa/space.py +++ b/mesa/space.py @@ -690,6 +690,14 @@ class _PropertyGrid(_Grid): move_agent_to_extreme_value_cell(agent, property_name, mode, mask): Moves an agent to a cell with extreme value of a property, optionally with a mask. + Mask Usage: + Several methods in this class accept a mask as an input, which is a NumPy ndarray of boolean values. This mask + specifies the cells to be considered (True) or ignored (False) in operations. Users can create custom masks, + including neighborhood masks, to apply specific conditions or constraints. Additionally, methods that deal with + cell selection or agent movement can return either a list of cell coordinates or a mask, based on the 'return_list' + parameter. This flexibility allows for more nuanced control and customization of grid operations, catering to a wide + range of modeling requirements and scenarios. + Note: This class is not intended for direct use in user models but is currently used by the SingleGrid and MultiGrid. """ @@ -785,8 +793,8 @@ def get_neighborhood_mask( return mask def select_cells_multi_properties( - self, conditions: dict, mask: np.ndarray = None - ) -> list[Coordinate]: + self, conditions: dict, mask: np.ndarray = None, return_list: bool = True + ) -> list[Coordinate] | np.ndarray: """ Select cells based on multiple property conditions using NumPy, optionally with a mask. @@ -795,25 +803,27 @@ def select_cells_multi_properties( callables that take a single argument (the property value) and return a boolean. mask (np.ndarray, optional): A boolean mask to restrict the selection. + return_list (bool, optional): If True, return a list of coordinates, otherwise return the mask. Returns: - List[Coordinate]: Coordinates where conditions are satisfied. + Union[list[Coordinate], np.ndarray]: Coordinates where conditions are satisfied or the combined mask. """ - # If no mask is provided, use a default mask of all True values - if mask is None: - mask = np.ones((self.width, self.height), dtype=bool) - - combined_mask = mask + # Start with a mask of all True values + combined_mask = np.ones((self.width, self.height), dtype=bool) for prop_name, condition in conditions.items(): prop_layer = self.properties[prop_name].data - # Apply the condition to the property layer prop_mask = condition(prop_layer) - # Combine with the existing mask using logical AND combined_mask = np.logical_and(combined_mask, prop_mask) - selected_cells = list(zip(*np.where(combined_mask))) - return selected_cells + if mask is not None: + combined_mask = np.logical_and(combined_mask, mask) + + if return_list: + selected_cells = list(zip(*np.where(combined_mask))) + return selected_cells + else: + return combined_mask def move_agent_to_random_cell( self, agent: Agent, conditions: dict, mask: np.ndarray = None @@ -827,7 +837,10 @@ def move_agent_to_random_cell( conditions (dict): Conditions for selecting the cell. mask (np.ndarray, optional): A boolean mask to restrict the selection. """ - eligible_cells = self.select_cells_multi_properties(conditions, mask) + eligible_cells = self.select_cells_multi_properties( + conditions, mask, return_list=True + ) + if not eligible_cells: warn( f"No eligible cells found. Agent {agent.unique_id} remains in the current position.",