Repository: dwavesystems/dimod

Issue:
SampleSet.trim or similar
Something like
```
samples.trim(n, 'energy')
```
which would return a new sampleset containing the first n rows of the sampleset ordered by energy.

We could then do things like
```
samples.aggregate().trim(n)
```

Discussion:

Patch:
diff --git a/dimod/sampleset.py b/dimod/sampleset.py
index f93633af..87179821 100644
--- a/dimod/sampleset.py
+++ b/dimod/sampleset.py
@@ -912,7 +912,51 @@ class SampleSet(abc.Iterable, abc.Sized):
         # dev note: we don't check the energies as they should be the same
         # for individual samples
 
-        return self.__class__(record, self.variables, self.info, self.vartype)
+        return type(self)(record, self.variables, self.info, self.vartype)
+
+    def truncate(self, n, sorted_by='energy'):
+        """Create a new SampleSet with rows truncated after n.
+
+        Args:
+            n (int):
+                Maximum number of rows in the returned sample set.
+
+            sorted_by (str/None, optional, default='energy'):
+                Selects the record field used to sort the samples before
+                truncating. Note that sample order is maintained in the
+                underlying array.
+
+        Returns:
+            :obj:`.SampleSet`
+
+        Examples:
+            >>> sampleset = dimod.SampleSet.from_samples(np.ones((5, 5)), dimod.SPIN, energy=5)
+            >>> print(sampleset)
+                0   1   2   3   4  energy  num_occ.
+            0  +1  +1  +1  +1  +1       5         1
+            1  +1  +1  +1  +1  +1       5         1
+            2  +1  +1  +1  +1  +1       5         1
+            3  +1  +1  +1  +1  +1       5         1
+            4  +1  +1  +1  +1  +1       5         1
+
+            [ 5 rows, 5 variables ]
+            >>> print(sampleset.truncate(2))
+                0   1   2   3   4  energy  num_occ.
+            0  +1  +1  +1  +1  +1       5         1
+            1  +1  +1  +1  +1  +1       5         1
+
+            [ 2 rows, 5 variables ]
+
+        """
+        record = self.record
+
+        if sorted_by is None:
+            record = record[:n]
+        else:
+            sort_indices = np.argsort(record[sorted_by])
+            record = record[sort_indices[:n]]
+
+        return type(self)(record, self.variables, self.info, self.vartype)
 
     ###############################################################################################
     # Serialization
diff --git a/docs/reference/sampleset.rst b/docs/reference/sampleset.rst
index d9705d80..96267201 100644
--- a/docs/reference/sampleset.rst
+++ b/docs/reference/sampleset.rst
@@ -14,6 +14,14 @@ sample_like Objects
 
    as_samples
 
+Utility Functions
+=================
+
+.. autosummary::
+   :toctree: generated/
+
+   concatenate
+
 SampleSet
 =========
 
@@ -28,72 +36,30 @@ Properties
 .. autosummary::
    :toctree: generated/
 
+   SampleSet.first
+   SampleSet.info
    SampleSet.record
    SampleSet.variables
-   SampleSet.info
    SampleSet.vartype
-   SampleSet.first
 
 
 Methods
 -------
 
-Viewing a Response
-~~~~~~~~~~~~~~~~~~
-
 .. autosummary::
    :toctree: generated/
 
-   SampleSet.samples
+   SampleSet.aggregate
+   SampleSet.change_vartype
+   SampleSet.copy
    SampleSet.data
-
-Constructing a Response
-~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autosummary::
-   :toctree: generated/
-
+   SampleSet.done
+   SampleSet.from_future
    SampleSet.from_samples
    SampleSet.from_samples_bqm
-   SampleSet.from_future
-
-Transformations
-~~~~~~~~~~~~~~~
-
-.. autosummary::
-   :toctree: generated/
-
-   SampleSet.change_vartype
+   SampleSet.from_serializable
    SampleSet.relabel_variables
-
-Copy
-~~~~
-
-.. autosummary::
-   :toctree: generated/
-
-   SampleSet.copy
-
-Aggregate
-~~~~~~~~~
-
-.. autosummary::
-   :toctree: generated/
-
-   SampleSet.aggregate
-
-Done
-~~~~
-
-.. autosummary::
-   :toctree: generated/
-
-   SampleSet.done
-
-Utility Functions
-=================
-
-.. autosummary::
-   :toctree: generated/
-
-   concatenate
+   SampleSet.samples
+   SampleSet.to_pandas_dataframe
+   SampleSet.to_serializable
+   SampleSet.truncate


Test patch:
diff --git a/tests/test_sampleset.py b/tests/test_sampleset.py
index 9eb9a300..a36a6b2c 100644
--- a/tests/test_sampleset.py
+++ b/tests/test_sampleset.py
@@ -241,6 +241,34 @@ class TestSampleSet(unittest.TestCase):
         self.assertEqual(samples,
                          dimod.SampleSet.from_samples(([-1, 1], 'ab'), dimod.SPIN, energy=1))
 
+    def test_truncate(self):
+        bqm = dimod.BinaryQuadraticModel.from_ising({v: -1 for v in range(100)}, {})
+        samples = dimod.SampleSet.from_samples_bqm(np.tril(np.ones(100)), bqm.binary)
+
+        # by default should be in reverse order
+        new = samples.truncate(10)
+        self.assertEqual(len(new), 10)
+        for n, sample in enumerate(new.samples()):
+            for v, val in sample.items():
+                if v > 100 - n - 1:
+                    self.assertEqual(val, 0)
+                else:
+                    self.assertEqual(val, 1)
+
+    def test_truncate_unordered(self):
+        bqm = dimod.BinaryQuadraticModel.from_ising({v: -1 for v in range(100)}, {})
+        samples = dimod.SampleSet.from_samples_bqm(np.triu(np.ones(100)), bqm.binary)
+
+        # now undordered
+        new = samples.truncate(10, sorted_by=None)
+        self.assertEqual(len(new), 10)
+        for n, sample in enumerate(new.samples()):
+            for v, val in sample.items():
+                if v < n:
+                    self.assertEqual(val, 0)
+                else:
+                    self.assertEqual(val, 1)
+
 
 class TestSampleSetIteration(unittest.TestCase):
     def test_data_reverse(self):