From 5f34480f6f69fbcdc3e2a74f39fa6ffbbb952f0d Mon Sep 17 00:00:00 2001
From: Ewout ter Hoeven <E.M.terHoeven@student.tudelft.nl>
Date: Mon, 4 Dec 2023 22:42:18 +0100
Subject: [PATCH] PropertyLayer: Fix handling of 'condition' callable in
 set_cells

Resolved an inconsistency in the set_cells method of our grid class. The method's documentation stated that the 'condition' argument should be a callable (such as a lambda function or a NumPy ufunc), but the implementation incorrectly treated 'condition' as a NumPy array.

This update rectifies the implementation to align with the documented behavior. Now, the 'condition' argument is correctly called with the grid data as input, and its output (expected to be a boolean array) is used for conditional in-place updates of the grid. This change ensures that the function operates correctly when provided with callable conditions, fulfilling the intended functionality.

Includes:
- Calling 'condition' with self.data and using its output for conditional updates.
- Adjusted error handling to check the output of the callable, ensuring it's a NumPy array with the correct shape.
- Updated comments within the method for clarity.
---
 mesa/space.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/mesa/space.py b/mesa/space.py
index 6f55c05ce33..a921f2de202 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -571,20 +571,23 @@ def set_cells(self, value, condition=None):
 
         Args:
             value: The value to be used for the update.
-            condition: (Optional) A callable that returns a boolean array when applied to the data.
+            condition: (Optional) A callable (like a lambda function or a NumPy ufunc)
+                       that returns a boolean array when applied to the data.
         """
         if condition is None:
             np.copyto(self.data, value)  # In-place update
         else:
-            # Ensure condition is a boolean array of the same shape as self.data
+            # Call the condition and check if the result is a boolean array
+            condition_result = condition(self.data)
             if (
-                not isinstance(condition, np.ndarray)
-                or condition.shape != self.data.shape
+                not isinstance(condition_result, np.ndarray)
+                or condition_result.shape != self.data.shape
             ):
                 raise ValueError(
-                    "Condition must be a NumPy array with the same shape as the grid."
+                    "Result of condition must be a NumPy array with the same shape as the grid."
                 )
-            np.copyto(self.data, value, where=condition)  # Conditional in-place update
+            # Conditional in-place update
+            np.copyto(self.data, value, where=condition_result)
 
     def modify_cell(self, position: Coordinate, operation, value=None):
         """