From c30ae232bd10b0549a29f6723eb2bd3e407ed295 Mon Sep 17 00:00:00 2001
From: Christopher Lorton <christopher.lorton@gatesfoundation.org>
Date: Tue, 17 Dec 2024 21:17:16 +0000
Subject: [PATCH] Implement add_array_property

---
 src/laser_core/laserframe.py | 30 ++++++++++++++++++++++++++++--
 tests/test_laserframe.py     | 20 ++++++++++++++++++++
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/laser_core/laserframe.py b/src/laser_core/laserframe.py
index 4bf56da..7b86690 100644
--- a/src/laser_core/laserframe.py
+++ b/src/laser_core/laserframe.py
@@ -60,7 +60,7 @@ def add_scalar_property(self, name: str, dtype=np.uint32, default=0) -> None:
         Add a scalar property to the class.
 
         This method initializes a new scalar property for the class instance. The property is
-        stored as a NumPy array with a specified data type and default value.
+        stored as a 1-D NumPy array (scalar / entry) with a specified data type and default value.
 
         Parameters:
 
@@ -81,7 +81,7 @@ def add_vector_property(self, name: str, length: int, dtype=np.uint32, default=0
         """
         Adds a vector property to the object.
 
-        This method initializes a new property with the given name as a NumPy array.
+        This method initializes a new property with the given name as a 2-D NumPy array (vector per entry).
 
         The array will have a shape of (length, self._capacity) and will be filled
         with the specified default value. The data type of the array elements is
@@ -103,6 +103,32 @@ def add_vector_property(self, name: str, length: int, dtype=np.uint32, default=0
         setattr(self, name, np.full((length, self._capacity), default, dtype=dtype))
         return
 
+    def add_array_property(self, name: str, shape: tuple, dtype=np.uint32, default=0) -> None:
+        """
+        Adds an array property to the object.
+
+        This method initializes a new property with the given name as a multi-dimensional NumPy array.
+
+        The array will have the given shape (note that there is no implied dimension of size self._capacity),
+        datatype (default is np.uint32), and default value (default is 0).
+
+        Parameters:
+
+            name (str): The name of the property to be added.
+            shape (tuple): The shape of the array.
+            dtype (data-type, optional): The desired data-type for the array, default is np.uint32.
+            default (scalar, optional): The default value to fill the array with, default is 0.
+
+        Returns:
+
+            None
+
+        """
+
+        # initialize the property to a NumPy array with given shape, dtype, and default value
+        setattr(self, name, np.full(shape, default, dtype=dtype))
+        return
+
     @property
     def count(self) -> int:
         """
diff --git a/tests/test_laserframe.py b/tests/test_laserframe.py
index 8291736..bc820b4 100644
--- a/tests/test_laserframe.py
+++ b/tests/test_laserframe.py
@@ -84,6 +84,26 @@ def test_add_vectory_property_with_value(self):
         assert np.all(pop.events == 1)
         assert pop.events.shape == (365, 1024)
 
+    def test_add_array_property(self):
+        pop = LaserFrame(1024)
+        pop.add_array_property("events", (365, 1024))
+        assert np.all(pop.events == 0)
+        assert pop.events.shape == (365, 1024)
+
+    def test_add_array_property_with_value(self):
+        pop = LaserFrame(1024)
+        pop.add_array_property("events", (365, 1024), default=42)
+        assert np.all(pop.events == 42)
+        assert pop.events.shape == (365, 1024)
+
+    def test_add_array_property_with_dtype(self):
+        pop = LaserFrame(1024)
+        default = np.float32(-3.14159265)
+        pop.add_array_property("events", (365, 1024), dtype=np.float32, default=default)
+        assert np.all(pop.events == default)
+        assert pop.events.shape == (365, 1024)
+        assert pop.events.dtype == np.float32
+
     def test_add_agents(self):
         pop = LaserFrame(1024)
         istart, iend = pop.add(100)