Repository: dwavesystems/dimod

Issue:
as_vartype function
This would fill a similar role to `dimod.decorators.vartype_argument` in that it would cast `'SPIN'`, `{-1, 1}`, etc. to the appropriate vartype.


Discussion:

Patch:
diff --git a/dimod/__init__.py b/dimod/__init__.py
index 6f3e3e87..114c9886 100644
--- a/dimod/__init__.py
+++ b/dimod/__init__.py
@@ -58,5 +58,4 @@ import dimod.response
 from dimod.utilities import *
 import dimod.utilities
 
-from dimod.vartypes import *
-import dimod.vartypes
+from dimod.vartypes import as_vartype, Vartype, SPIN, BINARY
diff --git a/dimod/binary_quadratic_model.py b/dimod/binary_quadratic_model.py
index a8895ba2..8fde1d9c 100644
--- a/dimod/binary_quadratic_model.py
+++ b/dimod/binary_quadratic_model.py
@@ -57,7 +57,7 @@ from numbers import Number
 
 import numpy as np
 
-from six import itervalues, iteritems, iterkeys, PY2
+from six import itervalues, iteritems
 
 from dimod.decorators import vartype_argument
 from dimod.serialization.utils import array2bytes, bytes2array
@@ -151,6 +151,10 @@ class BinaryQuadraticModel(abc.Sized, abc.Container, abc.Iterable):
         vartype (:class:`.Vartype`):
             The model's type. One of :class:`.Vartype.SPIN` or :class:`.Vartype.BINARY`.
 
+        variables (keysview):
+            The variables in the binary quadratic model as a dictionary keys
+            view object.
+
         adj (dict):
             The model's interactions as nested dicts.
             In graphic representation, where variables are nodes and interactions
@@ -279,10 +283,7 @@ class BinaryQuadraticModel(abc.Sized, abc.Container, abc.Iterable):
     @property
     def variables(self):
         """Return binary quadratic model's variables as a dictionary view object."""
-        if PY2:
-            return set(self.linear)
-        else:
-            return self.linear.keys()
+        return abc.KeysView(self.linear)
 
 ##################################################################################################
 # vartype properties
diff --git a/dimod/decorators.py b/dimod/decorators.py
index b47d3895..7a37fbde 100644
--- a/dimod/decorators.py
+++ b/dimod/decorators.py
@@ -36,7 +36,7 @@ from six import iteritems, integer_types
 from dimod.compatibility23 import getargspec
 from dimod.core.structured import Structured
 from dimod.exceptions import BinaryQuadraticModelStructureError
-from dimod.vartypes import Vartype
+from dimod.vartypes import as_vartype
 
 
 def bqm_index_labels(f):
@@ -215,6 +215,9 @@ def vartype_argument(*arg_names):
         constrained by :func:`vartype_argument`, or it can use a keyword
         arguments `dict`.
 
+    See also:
+        :func:`~dimod.as_vartype`
+
     """
     # by default, constrain only one argument, the 'vartype`
     if not arg_names:
@@ -229,21 +232,7 @@ def vartype_argument(*arg_names):
             except KeyError:
                 raise TypeError('vartype argument missing')
 
-            if isinstance(vartype, Vartype):
-                return
-
-            try:
-                if isinstance(vartype, str):
-                    vartype = Vartype[vartype]
-                else:
-                    vartype = Vartype(vartype)
-
-            except (ValueError, KeyError):
-                raise TypeError(("expected input vartype to be one of: "
-                                 "Vartype.SPIN, 'SPIN', {-1, 1}, "
-                                 "Vartype.BINARY, 'BINARY', or {0, 1}."))
-
-            kwargs[name] = vartype
+            kwargs[name] = as_vartype(vartype)
 
         @wraps(f)
         def new_f(*args, **kwargs):
diff --git a/dimod/sampleset.py b/dimod/sampleset.py
index 48f0ba39..29e3ebba 100644
--- a/dimod/sampleset.py
+++ b/dimod/sampleset.py
@@ -1154,8 +1154,8 @@ class SampleSet(abc.Iterable, abc.Sized):
 
             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.
+                truncating. Note that this sort order is maintained in the
+                returned SampleSet.
 
         Returns:
             :obj:`.SampleSet`
@@ -1179,14 +1179,95 @@ class SampleSet(abc.Iterable, abc.Sized):
             1 +1 +1 +1 +1 +1      5       1
             ['SPIN', 2 rows, 2 samples, 5 variables]
 
+        See:
+            :meth:`SampleSet.slice`
+
         """
-        record = self.record
+        return self.slice(n, sorted_by=sorted_by)
+
+    def slice(self, *slice_args, **kwargs):
+        """Create a new SampleSet with rows sliced according to standard Python
+        slicing syntax.
+
+        Args:
+            start (int, optional, default=None):
+                Start index for `slice`.
+
+            stop (int):
+                Stop index for `slice`.
+
+            step (int, optional, default=None):
+                Step value for `slice`.
+
+            sorted_by (str/None, optional, default='energy'):
+                Selects the record field used to sort the samples before
+                slicing. Note that `sorted_by` determines the sample order in
+                the returned SampleSet.
+
+        Returns:
+            :obj:`.SampleSet`
+
+        Examples:
+
+            >>> import numpy as np
+            ...
+            >>> sampleset = dimod.SampleSet.from_samples(np.diag(range(1, 11)), dimod.BINARY, energy=range(10))
+            >>> print(sampleset)
+               0  1  2  3  4  5  6  7  8  9 energy num_oc.
+            0  1  0  0  0  0  0  0  0  0  0      0       1
+            1  0  1  0  0  0  0  0  0  0  0      1       1
+            2  0  0  1  0  0  0  0  0  0  0      2       1
+            3  0  0  0  1  0  0  0  0  0  0      3       1
+            4  0  0  0  0  1  0  0  0  0  0      4       1
+            5  0  0  0  0  0  1  0  0  0  0      5       1
+            6  0  0  0  0  0  0  1  0  0  0      6       1
+            7  0  0  0  0  0  0  0  1  0  0      7       1
+            8  0  0  0  0  0  0  0  0  1  0      8       1
+            9  0  0  0  0  0  0  0  0  0  1      9       1
+            ['BINARY', 10 rows, 10 samples, 10 variables]
+
+            >>> # the first 3 samples by energy == truncate(3)
+            >>> print(sampleset.slice(3))
+               0  1  2  3  4  5  6  7  8  9 energy num_oc.
+            0  1  0  0  0  0  0  0  0  0  0      0       1
+            1  0  1  0  0  0  0  0  0  0  0      1       1
+            2  0  0  1  0  0  0  0  0  0  0      2       1
+            ['BINARY', 3 rows, 3 samples, 10 variables]
+
+            >>> # the last 3 samples by energy
+            >>> print(sampleset.slice(-3, None))
+               0  1  2  3  4  5  6  7  8  9 energy num_oc.
+            0  0  0  0  0  0  0  0  1  0  0      7       1
+            1  0  0  0  0  0  0  0  0  1  0      8       1
+            2  0  0  0  0  0  0  0  0  0  1      9       1
+            ['BINARY', 3 rows, 3 samples, 10 variables]
+
+            >>> # every second sample in between (skip the top and the bottom 3)
+            >>> print(sampleset.slice(3, -3, 2))
+               0  1  2  3  4  5  6  7  8  9 energy num_oc.
+            0  0  0  0  1  0  0  0  0  0  0      3       1
+            1  0  0  0  0  0  1  0  0  0  0      5       1
+            ['BINARY', 2 rows, 2 samples, 10 variables]
+
+        """
+        # handle `sorted_by` kwarg with a default value in a python2-compatible way
+        sorted_by = kwargs.pop('sorted_by', 'energy')
+        if kwargs:
+            # be strict about allowed kwargs: throw the same error as python3 would
+            raise TypeError('slice got an unexpected '
+                            'keyword argument {!r}'.format(kwargs.popitem()[0]))
+
+        # follow Python's slice syntax
+        if slice_args:
+            selector = slice(*slice_args)
+        else:
+            selector = slice(None)
 
         if sorted_by is None:
-            record = record[:n]
+            record = self.record[selector]
         else:
-            sort_indices = np.argsort(record[sorted_by])
-            record = record[sort_indices[:n]]
+            sort_indices = np.argsort(self.record[sorted_by])
+            record = self.record[sort_indices[selector]]
 
         return type(self)(record, self.variables, copy.deepcopy(self.info),
                           self.vartype)
diff --git a/dimod/variables.py b/dimod/variables.py
index 11a65ecf..810b09a4 100644
--- a/dimod/variables.py
+++ b/dimod/variables.py
@@ -133,6 +133,9 @@ class Variables(abc.Sequence, abc.Set):
         index = self.index
 
         for old, new in mapping.items():
+            if old not in self:
+                continue
+
             label[index[old]] = new
             index[new] = index[old]
             del index[old]
diff --git a/dimod/vartypes.py b/dimod/vartypes.py
index af482d97..164642db 100644
--- a/dimod/vartypes.py
+++ b/dimod/vartypes.py
@@ -12,48 +12,55 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 #
-# ================================================================================================
+# =============================================================================
 """
 Enumeration of valid variable types for binary quadratic models.
 
 Examples:
-    This example shows easy access to different Vartypes, which are in the main
-    namespace.
+
+    :class:`.Vartype` is an :py:class:`~enum.Enum`. Each vartype has a value and
+    a name.
 
     >>> vartype = dimod.SPIN
-    >>> print(vartype)
-    Vartype.SPIN
+    >>> vartype.name
+    'SPIN'
+    >>> vartype.value == {-1, +1}
+    True
+
     >>> vartype = dimod.BINARY
-    >>> print(vartype)
-    Vartype.BINARY
-    >>> vartype = dimod.Vartype.SPIN
-    >>> print(vartype)
-    Vartype.SPIN
-    >>> isinstance(vartype, dimod.Vartype)
+    >>> vartype.name
+    'BINARY'
+    >>> vartype.value == {0, 1}
     True
 
-    This example shows access by value or name.
+    The :func:`.as_vartype` function allows the user to provide several
+    convenient forms.
 
-    >>> print(dimod.Vartype({0, 1}))
-    Vartype.BINARY
-    >>> print(dimod.Vartype['SPIN'])
-    Vartype.SPIN
+    >>> from dimod import as_vartype
 
-    This example uses the `.value` parameter to validate.
+    >>> as_vartype(dimod.SPIN) is dimod.SPIN
+    True
+    >>> as_vartype('SPIN') is dimod.SPIN
+    True
+    >>> as_vartype({-1, 1}) is dimod.SPIN
+    True
 
-    >>> sample = {'u': -1, 'v': 1}
-    >>> vartype = dimod.Vartype.SPIN
-    >>> all(val in vartype.value for val in sample.values())
+    >>> as_vartype(dimod.BINARY) is dimod.BINARY
+    True
+    >>> as_vartype('BINARY') is dimod.BINARY
+    True
+    >>> as_vartype({0, 1}) is dimod.BINARY
     True
 
 """
 import enum
 
-__all__ = ['Vartype', 'SPIN', 'BINARY']
+__all__ = ['as_vartype', 'Vartype', 'SPIN', 'BINARY']
 
 
 class Vartype(enum.Enum):
-    """An :py:class:`~enum.Enum` over the types of variables for the binary quadratic model.
+    """An :py:class:`~enum.Enum` over the types of variables for the binary
+    quadratic model.
 
     Attributes:
         SPIN (:class:`.Vartype`): Vartype for spin-valued models; variables of
@@ -68,3 +75,40 @@ class Vartype(enum.Enum):
 
 SPIN = Vartype.SPIN
 BINARY = Vartype.BINARY
+
+
+def as_vartype(vartype):
+    """Cast various inputs to a valid vartype object.
+
+    Args:
+        vartype (:class:`.Vartype`/str/set):
+            Variable type. Accepted input values:
+
+            * :class:`.Vartype.SPIN`, ``'SPIN'``, ``{-1, 1}``
+            * :class:`.Vartype.BINARY`, ``'BINARY'``, ``{0, 1}``
+
+    Returns:
+        :class:`.Vartype`: Either :class:`.Vartype.SPIN` or
+        :class:`.Vartype.BINARY`.
+
+    See also:
+        :func:`~dimod.decorators.vartype_argument`
+
+    """
+    if isinstance(vartype, Vartype):
+        return vartype
+
+    try:
+        if isinstance(vartype, str):
+            vartype = Vartype[vartype]
+        elif isinstance(vartype, frozenset):
+            vartype = Vartype(vartype)
+        else:
+            vartype = Vartype(frozenset(vartype))
+
+    except (ValueError, KeyError):
+        raise TypeError(("expected input vartype to be one of: "
+                         "Vartype.SPIN, 'SPIN', {-1, 1}, "
+                         "Vartype.BINARY, 'BINARY', or {0, 1}."))
+
+    return vartype
diff --git a/docs/reference/vartypes.rst b/docs/reference/vartypes.rst
index fe664899..0ef79df1 100644
--- a/docs/reference/vartypes.rst
+++ b/docs/reference/vartypes.rst
@@ -8,3 +8,4 @@ Vartype
 
 .. currentmodule:: dimod
 .. autoclass:: Vartype
+.. autofunction:: as_vartype


Test patch:
diff --git a/tests/test_sampleset.py b/tests/test_sampleset.py
index 176627ee..d4379a4b 100644
--- a/tests/test_sampleset.py
+++ b/tests/test_sampleset.py
@@ -412,6 +412,65 @@ class TestTruncate(unittest.TestCase):
                     self.assertEqual(val, 1)
 
 
+class TestSlice(unittest.TestCase):
+    def test_typical(self):
+        bqm = dimod.BinaryQuadraticModel.from_ising({v: -1 for v in range(100)}, {})
+        sampleset = dimod.SampleSet.from_samples_bqm(np.tril(np.ones(100)), bqm.binary)
+
+        # `:10` is equal to `truncate(10)`
+        self.assertEqual(sampleset.slice(10), sampleset.truncate(10))
+
+    def test_unordered(self):
+        bqm = dimod.BinaryQuadraticModel.from_ising({v: -1 for v in range(100)}, {})
+        sampleset = dimod.SampleSet.from_samples_bqm(np.triu(np.ones(100)), bqm.binary)
+
+        # `:10` but for the unordered case
+        self.assertEqual(sampleset.slice(10, sorted_by=None), sampleset.truncate(10, sorted_by=None))
+
+    def test_null_slice(self):
+        energies = list(range(10))
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=energies)
+
+        self.assertTrue((sampleset.slice().record.energy == energies).all())
+
+    def test_slice_stop_only(self):
+        energies = list(range(10))
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=energies)
+
+        self.assertTrue((sampleset.slice(3).record.energy == energies[:3]).all())
+        self.assertTrue((sampleset.slice(-3).record.energy == energies[:-3]).all())
+
+    def test_slice_range(self):
+        energies = list(range(10))
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=energies)
+
+        self.assertTrue((sampleset.slice(3, 5).record.energy == energies[3:5]).all())
+        self.assertTrue((sampleset.slice(3, -3).record.energy == energies[3:-3]).all())
+        self.assertTrue((sampleset.slice(-3, None).record.energy == energies[-3:]).all())
+
+    def test_slice_stride(self):
+        energies = list(range(10))
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=energies)
+
+        self.assertTrue((sampleset.slice(3, -3, 2).record.energy == energies[3:-3:2]).all())
+        self.assertTrue((sampleset.slice(None, None, 2).record.energy == energies[::2]).all())
+        self.assertTrue((sampleset.slice(None, None, -1).record.energy == energies[::-1]).all())
+
+    def test_custom_ordering(self):
+        custom = list(range(10))
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=None, custom=custom)
+
+        self.assertTrue((sampleset.slice(3, sorted_by='custom').record.custom == custom[:3]).all())
+        self.assertTrue((sampleset.slice(3, -3, sorted_by='custom').record.custom == custom[3:-3]).all())
+        self.assertTrue((sampleset.slice(None, None, -1, sorted_by='custom').record.custom == custom[::-1]).all())
+
+    def test_kwargs(self):
+        sampleset = dimod.SampleSet.from_samples(np.ones((10, 1)), dimod.SPIN, energy=None)
+
+        with self.assertRaises(TypeError):
+            sampleset.slice(1, sortedby='invalid-kwarg')
+
+
 class TestIteration(unittest.TestCase):
     def test_data_reverse(self):
         bqm = dimod.BinaryQuadraticModel.from_ising({}, {'ab': -1})
diff --git a/tests/test_variables.py b/tests/test_variables.py
index 532c99ed..6e5f3370 100644
--- a/tests/test_variables.py
+++ b/tests/test_variables.py
@@ -43,6 +43,16 @@ class TestDuplicates(unittest.TestCase):
         self.assertEqual(len(variables), 1)
 
 
+class TestRelabel(unittest.TestCase):
+    def test_permissive(self):
+        variables = Variables([0, 1])
+
+        # relabels a non-existant variable 2
+        variables.relabel({0: 'a', 1: 'b', 2: 'c'})
+
+        self.assertEqual(variables, Variables('ab'))
+
+
 class TestList(unittest.TestCase):
     iterable = list(range(5))