Make ArrayStore use int32 indices

HISTORY.md

@@ -10,7 +10,7 @@
   ({pr}`397`)
 - **Backwards-incompatible:** Rename `measure_*` columns to `measures_*` in
   `as_pandas` ({pr}`396`)
-- Add ArrayStore data structure ({pr}`395`, {pr}`398`)
+- Add ArrayStore data structure ({pr}`395`, {pr}`398`, {pr}`400`)
 - Add GradientOperatorEmitter to support OMG-MEGA and OG-MAP-Elites ({pr}`348`)
 
 #### Improvements

ribs/archives/_array_store.py

@@ -113,7 +113,7 @@ def __init__(self, field_desc, capacity):
             "capacity": capacity,
             "occupied": np.zeros(capacity, dtype=bool),
             "n_occupied": 0,
-            "occupied_list": np.empty(capacity, dtype=int),
+            "occupied_list": np.empty(capacity, dtype=np.int32),
             "updates": np.array([0, 0]),
         }
 
@@ -168,7 +168,7 @@ def occupied(self):
 
     @property
     def occupied_list(self):
-        """numpy.ndarray: Integer array listing all occupied indices in the
+        """numpy.ndarray: int32 array listing all occupied indices in the
         store."""
         return readonly(
             self._props["occupied_list"][:self._props["n_occupied"]])
@@ -207,7 +207,7 @@ def retrieve(self, indices, fields=None):
         Raises:
             ValueError: Invalid field name provided.
         """
-        indices = np.asarray(indices)
+        indices = np.asarray(indices, dtype=np.int32)
         occupied = readonly(self._props["occupied"][indices])
 
         data = {}
@@ -363,7 +363,7 @@ def resize(self, capacity):
         self._props["occupied"][:cur_capacity] = cur_occupied
 
         cur_occupied_list = self._props["occupied_list"]
-        self._props["occupied_list"] = np.empty(capacity, dtype=int)
+        self._props["occupied_list"] = np.empty(capacity, dtype=np.int32)
         self._props["occupied_list"][:cur_capacity] = cur_occupied_list
 
         for name, cur_arr in self._fields.items():

tests/archives/array_store_test.py

@@ -138,6 +138,28 @@ def test_add_duplicate_indices(store):
     assert np.all(store.occupied_list == [3])
 
 
+def test_dtypes(store):
+    store.add(
+        [3, 5],
+        {
+            "objective": [1.0, 2.0],
+            "measures": [[1.0, 2.0], [3.0, 4.0]],
+            "solution": [np.zeros(10), np.ones(10)],
+        },
+        {},  # Empty extra_args.
+        [],  # Empty transforms.
+    )
+
+    _, data = store.retrieve([5, 3])
+
+    # Index is always int32, and other fields were defined as float32 in the
+    # `store` fixture.
+    assert data["index"].dtype == np.int32
+    assert data["objective"].dtype == np.float32
+    assert data["measures"].dtype == np.float32
+    assert data["solution"].dtype == np.float32
+
+
 def test_retrieve_duplicate_indices(store):
     store.add(
         [3],
@@ -400,7 +422,7 @@ def test_as_pandas(store):
         "solution_8",
         "solution_9",
     ]).all()
-    assert (df.dtypes == [int] + [np.float32] * 13).all()
+    assert (df.dtypes == [np.int32] + [np.float32] * 13).all()
     assert len(df) == 2
 
     row0 = np.concatenate(([3, 1.0, 1.0, 2.0], np.zeros(10)))