Repository: dwavesystems/dimod

Issue:
data_vectors should have either numpy array values, or a list values, but not both.
https://github.com/dwavesystems/dimod/blob/7b75e47ce4fec541e432f84367ba58393934b941/dimod/response.py#L40

Making it be more than one thing requires the parser of the response to inspect the object before using it. If we want to have the benefits of a numpy array for some of the data_vectors, I think it's worth it to make everything a numpy array

If we don't need it to be a numpy array, might as well make them all lists?

Discussion:

Patch:
diff --git a/dimod/response.py b/dimod/response.py
index d3b2c38d..59d96329 100644
--- a/dimod/response.py
+++ b/dimod/response.py
@@ -87,21 +87,24 @@ class Response(Iterable, Sized):
         self._samples_matrix = samples_matrix
         num_samples, num_variables = samples_matrix.shape
 
-        if not isinstance(data_vectors, dict):
+        if not isinstance(data_vectors, Mapping):
             raise TypeError("expected 'data_vectors' to be a dict")
         if 'energy' not in data_vectors:
             raise ValueError("energy must be provided")
         else:
-            data_vectors = data_vectors.copy()  # shallow copy
-            data_vectors['energy'] = np.asarray(data_vectors['energy'])
-        for vector in data_vectors.values():
-            # todo - check that is a vector and that has the right length
-            if isinstance(vector, (np.ndarray, list)):
-                if len(vector) != num_samples:
-                    raise ValueError(("expected data vector {} (length {}) to be a vector of length {}"
-                                      "").format(vector, len(vector), num_samples))
-            else:
-                raise TypeError("expected data vector {} to be a list of NumPy array".format(vector))
+            data_vectors = dict(data_vectors)  # shallow copy
+
+        for key, vector in iteritems(data_vectors):
+            try:
+                data_vectors[key] = vector = np.asarray(vector)
+            except (ValueError, TypeError):
+                raise TypeError("expected data vector {} to be array-like".format(key))
+
+            shape = vector.shape
+            if not shape or shape[0] != num_samples:
+                raise ValueError(("expected data vector {} (shape {}) to have {} rows"
+                                  "").format(key, vector.shape, num_samples))
+
         self._data_vectors = data_vectors
 
         # vartype is checked by the decorator
@@ -824,10 +827,13 @@ class Response(Iterable, Sized):
     # Viewing a Response
     ###############################################################################################
 
-    def samples(self, sorted_by='energy'):
+    def samples(self, n=None, sorted_by='energy'):
         """Iterate over the samples in the response.
 
         Args:
+            n (int, optional, default=None):
+                The maximum number of samples to provide. If None, all are provided.
+
             sorted_by (str/None, optional, default='energy'):
                 Selects the `data_vector` used to sort the samples. If None, the samples are yielded in
                 the order given by the samples matrix.
@@ -861,13 +867,21 @@ class Response(Iterable, Sized):
             {'a': -1, 'b': 1}
 
         """
+        num_samples = len(self)
+
+        if n is not None:
+            for sample in itertools.islice(self.samples(n=None, sorted_by=sorted_by), n):
+                yield sample
+            return
+
         if sorted_by is None:
-            order = np.arange(len(self))
+            order = np.arange(num_samples)
         else:
             order = np.argsort(self.data_vectors[sorted_by])
 
         samples = self.samples_matrix
         label_mapping = self.label_to_idx
+
         for idx in order:
             yield SampleView(idx, self)
 


Test patch:
diff --git a/tests/test_response.py b/tests/test_response.py
index 73c1091e..6c5f827e 100644
--- a/tests/test_response.py
+++ b/tests/test_response.py
@@ -54,7 +54,7 @@ class TestResponse(unittest.TestCase):
         npt.assert_equal(samples_matrix, response.samples_matrix)
         npt.assert_allclose(energies, response.data_vectors['energy'])
 
-    def test_data_vector_copy(self):
+    def test_data_vectors_copy(self):
         samples_matrix = np.matrix([[0, 1, 0, 1],
                                     [1, 0, 1, 0],
                                     [0, 0, 0, 0],
@@ -66,6 +66,80 @@ class TestResponse(unittest.TestCase):
 
         self.assertIsNot(response.data_vectors, data_vectors)
 
+    def test_data_vectors_are_arrays(self):
+        samples_matrix = np.matrix([[0, 1, 0, 1],
+                                    [1, 0, 1, 0],
+                                    [0, 0, 0, 0],
+                                    [1, 1, 1, 1]])
+        energies = [2, 2, 0, 4]
+        num_occurrences = [1, 1, 2, 1]
+        objects = [object() for __ in range(4)]
+
+        data_vectors = {'energy': energies, 'occurences': num_occurrences, 'objects': objects}
+
+        response = dimod.Response(samples_matrix, data_vectors, dimod.BINARY)
+
+        self.assertEqual(len(response.data_vectors), 3)
+
+        for key in data_vectors:
+            self.assertIn(key, response.data_vectors)
+
+            vector = response.data_vectors[key]
+
+            self.assertIsInstance(vector, np.ndarray)
+
+            self.assertEqual(vector.shape, (4,))
+
+    def test_data_vectors_wrong_length(self):
+        samples_matrix = np.matrix([[0, 1, 0, 1],
+                                    [1, 0, 1, 0],
+                                    [0, 0, 0, 0],
+                                    [1, 1, 1, 1]])
+        energies = [2, 2, 0, 4]
+        num_occurrences = [1, 1, 2, 1, 1]
+        objects = [object() for __ in range(4)]
+
+        data_vectors = {'energy': energies, 'occurences': num_occurrences, 'objects': objects}
+
+        with self.assertRaises(ValueError):
+            response = dimod.Response(samples_matrix, data_vectors, dimod.BINARY)
+
+    def test_data_vectors_not_array_like(self):
+        samples_matrix = np.matrix([[0, 1, 0, 1],
+                                    [1, 0, 1, 0],
+                                    [0, 0, 0, 0],
+                                    [1, 1, 1, 1]])
+        energies = [2, 2, 0, 4]
+        num_occurrences = 'hi there'
+        objects = [object() for __ in range(4)]
+
+        data_vectors = {'energy': energies, 'occurences': num_occurrences, 'objects': objects}
+
+        with self.assertRaises(ValueError):
+            response = dimod.Response(samples_matrix, data_vectors, dimod.BINARY)
+
+    def test_samples_num_limited(self):
+        samples_matrix = np.matrix([[0, 1, 0, 1],
+                                    [1, 0, 1, 0],
+                                    [0, 0, 0, 0],
+                                    [1, 1, 1, 1]])
+        energies = [2, 2, 0, 4]
+        num_occurrences = [1, 1, 2, 1]
+        objects = [object() for __ in range(4)]
+
+        data_vectors = {'energy': energies, 'occurences': num_occurrences, 'objects': objects}
+
+        response = dimod.Response(samples_matrix, data_vectors, dimod.BINARY)
+
+        samples_list = list(response.samples())
+
+        self.assertEqual(len(samples_list), 4)
+
+        shortened_samples_list = list(response.samples(3))
+
+        self.assertEqual(len(shortened_samples_list), 3)
+        self.assertEqual(shortened_samples_list, samples_list[0:3])
+
     def test_instantiation_without_energy(self):
         samples_matrix = np.matrix([[0, 1, 0, 1],
                                     [1, 0, 1, 0],