Add index to get_random_elite() and elite_with_behavior() in archives (…

…icaros-usc#129)

- Add index to get_random_elite() and elite_with_behavior()
- Fix calls to these methods in tests
- Fix calls to these methods in examples / tutorials
- Force CVTArchive to return int instead of sometimes returning np.integer

examples/tutorials/arm_repertoire.ipynb

@@ -429,8 +429,8 @@
     "fig, ax = plt.subplots(2, 2, figsize=(8, 8))\n",
     "ax = ax.ravel()\n",
     "for i in range(len(ax)):\n",
-    "    solution, objective, behavior, _ = archive.get_random_elite()\n",
-    "    visualize(solution, link_lengths, objective, ax[i])"
+    "    sol, obj, *_ = archive.get_random_elite()\n",
+    "    visualize(sol, link_lengths, obj, ax[i])"
    ]
   },
   {
@@ -459,7 +459,7 @@
     }
    ],
    "source": [
-    "sol, obj, beh, _ = archive.elite_with_behavior([0, 0])\n",
+    "sol, obj, *_ = archive.elite_with_behavior([0, 0])\n",
     "_, ax = plt.subplots()\n",
     "if sol is not None:\n",
     "    visualize(sol, link_lengths, obj, ax)"

examples/tutorials/lunar_lander.ipynb

@@ -496,7 +496,7 @@
     }
    ],
    "source": [
-    "sol, obj, beh, _ = archive.elite_with_behavior([-0.4, -0.10])\n",
+    "sol, *_ = archive.elite_with_behavior([-0.4, -0.10])\n",
     "if sol is not None:\n",
     "    display_video(sol)"
    ]
@@ -530,7 +530,7 @@
     }
    ],
    "source": [
-    "sol, obj, beh, _ = archive.elite_with_behavior([0.4, -0.10])\n",
+    "sol, *_ = archive.elite_with_behavior([0.4, -0.10])\n",
     "if sol is not None:\n",
     "    display_video(sol)"
    ]
@@ -564,7 +564,7 @@
     }
    ],
    "source": [
-    "sol, obj, beh, _ = archive.elite_with_behavior([0.0, -0.10])\n",
+    "sol, *_ = archive.elite_with_behavior([0.0, -0.10])\n",
     "if sol is not None:\n",
     "    display_video(sol)"
    ]

ribs/archives/_archive_base.py

@@ -374,10 +374,15 @@ def add(self, solution, objective_value, behavior_values, metadata=None):
     def elite_with_behavior(self, behavior_values):
         """Gets the elite with behavior vals in the same bin as those specified.
 
+        .. note:: Values like ``index`` and ``metadata`` are often not used.
+            In such cases, consider ignoring them::
+
+                sol, obj, beh, *_ = archive.elite_with_behavior(...)
+
         Args:
             behavior_values (array-like): Coordinates in behavior space.
         Returns:
-            tuple: 4-element tuple for the elite if it is found:
+            tuple: 5-element tuple for the elite if it is found:
 
                 **solution** (:class:`numpy.ndarray`): Parameters for the
                 solution.
@@ -389,25 +394,38 @@ def elite_with_behavior(self, behavior_values):
                 space coordinates of the elite (may not be exactly the same as
                 those specified).
 
+                **index** (int or tuple of int): Index of the entry in the
+                archive. See :attr:`get_index` for more info.
+
                 **metadata** (object): Metadata for the solution.
 
-            If there is no elite in the bin, a tuple of (None, None, None, None)
-            is returned. Thus, something like
-            ``sol, obj, beh, meta = archive.elite_with_behavior(...)`` still
-            works).
+            If there is no elite in the bin, each of the above values is None.
+            Thus, something like
+            ``sol, obj, beh, idx, meta = archive.elite_with_behavior(...)``
+            still works).
         """
         index = self.get_index(np.asarray(behavior_values))
         if self._occupied[index]:
-            return (self._solutions[index], self._objective_values[index],
-                    self._behavior_values[index], self._metadata[index])
-        return (None, None, None, None)
+            return (
+                self._solutions[index],
+                self._objective_values[index],
+                self._behavior_values[index],
+                index,
+                self._metadata[index],
+            )
+        return (None, None, None, None, None)
 
     @require_init
     def get_random_elite(self):
         """Selects an elite uniformly at random from one of the archive's bins.
 
+        .. note:: Values like ``index`` and ``metadata`` are often not used.
+            In such cases, consider ignoring them::
+
+                sol, obj, beh, *_ = archive.get_random_elite()
+
         Returns:
-            tuple: 4-element tuple containing:
+            tuple: 5-element tuple containing:
 
                 **solution** (:class:`numpy.ndarray`): Parameters for the
                 solution.
@@ -418,6 +436,9 @@ def get_random_elite(self):
                 **behavior_values** (:class:`numpy.ndarray`): Behavior space
                 coordinates.
 
+                **index** (int or tuple of int): Index of the entry in the
+                archive. See :attr:`get_index` for more info.
+
                 **metadata** (object): Metadata for the solution.
         Raises:
             IndexError: The archive is empty.
@@ -428,8 +449,13 @@ def get_random_elite(self):
         random_idx = self._rand_buf.get(len(self._occupied_indices))
         index = self._occupied_indices[random_idx]
 
-        return (self._solutions[index], self._objective_values[index],
-                self._behavior_values[index], self._metadata[index])
+        return (
+            self._solutions[index],
+            self._objective_values[index],
+            self._behavior_values[index],
+            index,
+            self._metadata[index],
+        )
 
     def data(self):
         """Returns columns containing all data in the archive.
@@ -479,11 +505,13 @@ def data(self):
                 **all_metadata** (:class:`numpy.ndarray` -- shape (n_entries,)):
                 Object array with metadata of all entries.
         """
-        return (self._solutions[self._occupied_indices_cols],
-                self._objective_values[self._occupied_indices_cols],
-                self._behavior_values[self._occupied_indices_cols],
-                self._occupied_indices,
-                self._metadata[self._occupied_indices_cols])
+        return (
+            self._solutions[self._occupied_indices_cols],
+            self._objective_values[self._occupied_indices_cols],
+            self._behavior_values[self._occupied_indices_cols],
+            self._occupied_indices,
+            self._metadata[self._occupied_indices_cols],
+        )
 
     def as_pandas(self, include_solutions=True, include_metadata=False):
         """Converts the archive into a Pandas dataframe.

ribs/archives/_cvt_archive.py

@@ -253,7 +253,8 @@ def get_index(self, behavior_values):
         Returns:
             int: Centroid index.
         """
+        # We use int() here since these methods may return a numpy integer.
         if self._use_kd_tree:
-            return self._centroid_kd_tree.query(behavior_values)[1]
+            return int(self._centroid_kd_tree.query(behavior_values)[1])
 
-        return self._brute_force_nn_numba(behavior_values, self._centroids)
+        return int(self._brute_force_nn_numba(behavior_values, self._centroids))

tests/archives/archive_base_test.py

@@ -109,27 +109,31 @@ def test_solution_dim_correct(data):
 
 
 def test_elite_with_behavior_gets_correct_elite(data):
-    sol, obj, beh, meta = data.archive_with_entry.elite_with_behavior(
-        data.behavior_values)
+    (sol, obj, beh, idx,
+     meta) = data.archive_with_entry.elite_with_behavior(data.behavior_values)
     assert (sol == data.solution).all()
     assert obj == data.objective_value
     assert (beh == data.behavior_values).all()
+    assert isinstance(idx, (int, tuple))  # Exact value depends on archive.
     assert meta == data.metadata
 
 
 def test_elite_with_behavior_returns_none(data):
-    sol, obj, beh, meta = data.archive.elite_with_behavior(data.behavior_values)
+    (sol, obj, beh, idx,
+     meta) = data.archive.elite_with_behavior(data.behavior_values)
     assert sol is None
     assert obj is None
     assert beh is None
+    assert idx is None
     assert meta is None
 
 
 def test_random_elite_gets_single_elite(data):
-    sol, obj, beh, meta = data.archive_with_entry.get_random_elite()
+    sol, obj, beh, idx, meta = data.archive_with_entry.get_random_elite()
     assert np.all(sol == data.solution)
     assert obj == data.objective_value
     assert np.all(beh == data.behavior_values)
+    assert isinstance(idx, (int, tuple))  # Exact value depends on archive.
     assert meta == data.metadata