Merge pull request #154 from jackraymond/feature/specify_SRT

Add srts option to SpinReversalComposite

Author: randomir

dwave/preprocessing/composites/spin_reversal_transform.py

@@ -137,63 +137,100 @@ def _reorder_variables(sampleset: dimod.SampleSet,
 
     @dimod.decorators.nonblocking_sample_method
     def sample(self, bqm: dimod.BinaryQuadraticModel, *,
-               num_spin_reversal_transforms: int = 1,
+               srts: typing.Optional[np.ndarray] = None,
+               num_spin_reversal_transforms: typing.Optional[int] = None,
                **kwargs,
                ):
         """Sample from the binary quadratic model.
 
         Args:
             bqm: Binary quadratic model to be sampled from.
 
+            srts: A boolean NumPy array with shape
+                ``(num_spin_reversal_transforms, bqm.num_variables)``.
+                True indicates a flip and False indicates no flip; applied to
+                in the order given by ``bqm.variables``.
+                If this is not specified as an input values are generated uniformly
+                at random from the class pseudo-random number generator.
+
             num_spin_reversal_transforms:
                 Number of spin reversal transform runs.
                 A value of ``0`` will not transform the problem.
                 If you specify a nonzero value, each spin reversal transform
                 will result in an independent run of the child sampler.
+                If ``srts`` is set then ``num_spin_reversal_transforms``
+                is inferred by the shape, otherwise the default is 1.
 
         Returns:
             A sample set. Note that for a sampler that returns ``num_reads`` samples,
             the sample set will contain ``num_reads*num_spin_reversal_transforms`` samples.
 
+        Raises:
+            ValueError: If ``srts`` is inconsistent with
+            ``num_spin_reversal_transforms`` or the binary quadratic model.
+
         Examples:
-            This example runs 100 spin reversals applied to one variable of a QUBO problem.
+            This example runs 10 spin reversals applied to an unfrustrated chain
+            of length 6.
+
+            Using the lowest energy (ground) state returned, you can define a
+            special SRT that transforms all programmed couplers to be ferromagnetic 
+            (ground state to all 1).
 
             >>> from dimod import ExactSolver
+            >>> import numpy as np
             >>> from dwave.preprocessing.composites import SpinReversalTransformComposite
             >>> base_sampler = ExactSolver()
             >>> composed_sampler = SpinReversalTransformComposite(base_sampler)
             ...
-            >>> Q = {('a', 'a'): -1, ('b', 'b'): -1, ('a', 'b'): 2}
-            >>> response = composed_sampler.sample_qubo(Q,
-            ...               num_spin_reversal_transforms=100)
-            >>> len(response)
-            400
+            >>> num_var = 6
+            >>> num_spin_reversal_transforms = 10
+            >>> J = {(i, i+1): np.random.choice([-1,1]) for i in range(num_var-1)}
+            >>> h = {i: 0 for i in range(num_var)}
+            >>> response = composed_sampler.sample_ising(h, J,
+            ...               num_spin_reversal_transforms=num_spin_reversal_transforms)
+            >>> len(response) == 2**num_var * num_spin_reversal_transforms
+            True
+            >>> srts = np.array([[response.first.sample[i] != 1 for i in range(num_var)]])
+            >>> response = composed_sampler.sample_ising(h, J,
+            ...               srts=srts, num_reads=1)
+            >>> sum(response.record.num_occurrences) == 2**num_var
+            True
         """
         sampler = self._child
 
+        if srts is not None:
+            nsrt, num_bqm_var = srts.shape
+            if num_bqm_var != bqm.num_variables:
+                raise ValueError('srt shape is inconsistent with the bqm')
+            if num_spin_reversal_transforms is not None:
+                if num_spin_reversal_transforms != nsrt:
+                    raise ValueError('srt shape is inconsistent with num_spin_reversal_transforms')
+            else:
+                num_spin_reversal_transforms = nsrt
+        elif num_spin_reversal_transforms is None:
+            num_spin_reversal_transforms = 1
+
         # No SRTs, so just pass the problem through
-        if not num_spin_reversal_transforms or not bqm.num_variables:
+        if num_spin_reversal_transforms == 0 or not bqm.num_variables:
             sampleset = sampler.sample(bqm, **kwargs)
             # yield twice because we're using the @nonblocking_sample_method
             yield sampleset  # this one signals done()-ness
             yield sampleset  # this is the one actually used by the user
             return
 
+        if srts is None:
+            srts = self.rng.random((num_spin_reversal_transforms, bqm.num_variables)) > .5
+
         # we'll be modifying the BQM, so make a copy
         bqm = bqm.copy()
 
-        # We maintain the Leap behavior that num_spin_reversal_transforms == 1
-        # corresponds to a single problem with randomly flipped variables.
-
-        # Get the SRT matrix
-        SRT = self.rng.random((num_spin_reversal_transforms, bqm.num_variables)) > .5
-
         # Submit the problems
         samplesets: typing.List[dimod.SampleSet] = []
         flipped = np.zeros(bqm.num_variables, dtype=bool)  # what variables are currently flipped
         for i in range(num_spin_reversal_transforms):
             # determine what needs to be flipped
-            transform = flipped != SRT[i, :]
+            transform = flipped != srts[i, :]
 
             # apply the transform
             for v, flip in zip(bqm.variables, transform):
@@ -213,10 +250,10 @@ def sample(self, bqm: dimod.BinaryQuadraticModel, *,
         # Undo the SRTs according to vartype
         if bqm.vartype is Vartype.BINARY:
             for i, sampleset in enumerate(samplesets):
-                sampleset.record.sample[:, SRT[i, :]] = 1 - sampleset.record.sample[:, SRT[i, :]]
+                sampleset.record.sample[:, srts[i, :]] = 1 - sampleset.record.sample[:, srts[i, :]]
         elif bqm.vartype is Vartype.SPIN:
             for i, sampleset in enumerate(samplesets):
-                sampleset.record.sample[:, SRT[i, :]] *= -1
+                sampleset.record.sample[:, srts[i, :]] *= -1
         else:
             raise RuntimeError("unexpected vartype")
         if num_spin_reversal_transforms == 1:

releasenotes/notes/srts-for-SpinReversalComposite-70d17f0cc8192ca7.yaml

@@ -0,0 +1,6 @@
+---
+prelude: >
+    Add srts option to SpinReversalTransformComposite
+features:
+  - |
+    Allow non-random specification of spin reversal transforms

tests/test_spin_reversal_transform.py

@@ -19,7 +19,6 @@
 
 from dwave.preprocessing.composites import SpinReversalTransformComposite
 
-
 @dimod.testing.load_sampler_bqm_tests(SpinReversalTransformComposite(dimod.ExactSolver()))
 class TestSpinTransformComposite(unittest.TestCase):
     def test_instantiation(self):
@@ -45,7 +44,7 @@ def test_empty_bqm_composition(self):
 
         sampler = SpinReversalTransformComposite(dimod.RandomSampler())
         bqm = dimod.BinaryQuadraticModel('SPIN')
-        sampleset = sampler.sample(bqm, num_spin_reversals=1)
+        sampleset = sampler.sample(bqm, num_spin_reversal_transforms=1)
         self.assertEqual(len(sampleset.variables), 0)
 
     def test_concatenation_stripping(self):
@@ -73,12 +72,13 @@ def test_sampleset_size(self):
         sampler = SpinReversalTransformComposite(dimod.RandomSampler())
         for num_spin_reversal_transforms in [1, 2]:
             for num_reads in [1, 3]:
-                sampleset = sampler.sample_ising(
-                    {0: 1}, {},
-                    num_spin_reversal_transforms=num_spin_reversal_transforms,
+                with self.subTest(f'{num_reads} {num_spin_reversal_transforms}'):
+                    sampleset = sampler.sample_ising(
+                        {0: 1}, {},
+                        num_spin_reversal_transforms=num_spin_reversal_transforms,
                     num_reads=num_reads)
-                self.assertTrue(sum(sampleset.record.num_occurrences) ==
-                                num_reads*num_spin_reversal_transforms)
+                    self.assertTrue(sum(sampleset.record.num_occurrences) ==
+                                    num_reads*num_spin_reversal_transforms)
 
     def test_empty(self):
         # Check that empty BQMs are handled
@@ -151,7 +151,6 @@ def sample(self, bqm):
         ss1 = SpinReversalTransformComposite(Sampler(), seed=42).sample(bqm)
         ss2 = SpinReversalTransformComposite(Sampler(), seed=42).sample(bqm)
         ss3 = SpinReversalTransformComposite(Sampler(), seed=35).sample(bqm)
-
         self.assertTrue((ss1.record == ss2.record).all())
         self.assertFalse((ss1.record == ss3.record).all())
 
@@ -230,3 +229,45 @@ def sample(bqm):
         self.assertTrue(hasattr(sampleset,'info'))
         self.assertEqual(sampleset.info, {'has_some': True})
 
+    def test_srts_argument(self):
+        # All 1 ground state
+        class Sampler:
+            def sample(self, bqm):
+                return dimod.SampleSet.from_samples_bqm([-1] * bqm.num_variables, bqm)
+        num_var = 10
+        bqm = dimod.BinaryQuadraticModel(
+            {i: -1 for i in range(num_var)}, {}, 0, 'SPIN')  
+        
+        sampler = Sampler()
+        ss = sampler.sample(bqm)
+        samples = ss.record.sample
+        sampler = SpinReversalTransformComposite(sampler)
+        srts = np.zeros(shape=(1, num_var), dtype=bool)  # 
+        ss = sampler.sample(bqm, srts=srts)
+        self.assertTrue(np.all(ss.record.sample == samples),
+                        "Neutral srts leaves result unpermuted.")
+
+        ss = sampler.sample(bqm, srts=np.empty(shape=(0, num_var)))
+        self.assertTrue(np.all(ss.record.sample == samples),
+                        "Empty srts also allows pass through "
+                        "(just like num_spin_reversals=0).")
+        
+        srts = np.ones(shape=(1, num_var), dtype=bool)
+        ss = sampler.sample(bqm, srts=srts)
+        self.assertTrue(np.all(ss.record.sample == -samples),
+                        "Flip-all srts inverts the order")
+        
+        with self.subTest('srt shape'):
+            num_spin_reversal_transforms = 3
+            srts = np.unique(np.random.random(size=(num_spin_reversal_transforms, num_var)) > 0.5, axis=0)
+            ss = sampler.sample(bqm, srts=srts)
+            self.assertEqual(np.sum(ss.record.num_occurrences), srts.shape[0],
+                             "Apply 3 srtss")
+            self.assertTrue(np.all(srts == (ss.record.sample==1)))
+            
+            with self.assertRaises(ValueError):
+                # Inconsistent arguments
+                ss = sampler.sample(bqm, srts=srts, num_spin_reversal_transforms=num_spin_reversal_transforms+1)
+            
+
+