Use optimised sa of neal

pyqubo/utils/solver.py

@@ -12,37 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import dimod
+import neal
 import numpy as np
 
 
-def solve_qubo(qubo, num_reads=10, num_sweeps=1000, beta_range=(1.0, 50.0)):
-    """Solve QUBO with Simulated Annealing (SA) provided by dimod.
+def solve_qubo(qubo, num_reads=10, sweeps=1000, beta_range=(1.0, 50.0)):
+    """Solve QUBO with Simulated Annealing (SA) provided by neal.
     
     Args:
         qubo (dict[(label, label), float]): The QUBO to be solved.
         
         num_reads (int, default=10): Number of run repetitions of SA.
         
-        num_sweeps (int, default=1000): Number of iterations in each run of SA.
+        sweeps (int, default=1000): Number of iterations in each run of SA.
          
         beta_range (tuple(float, float), default=(1.0, 50.0)): Tuple of start beta and end beta.
          
     Returns:
          dict[label, bit]: The solution of SA.
+    
+    >>> from pyqubo import Spin, solve_qubo
+    >>> s1, s2, s3 = Spin("s1"), Spin("s2"), Spin("s3")
+    >>> H = (2*s1 + 4*s2 + 6*s3)**2
+    >>> model = H.compile()
+    >>> qubo, offset = model.to_qubo()
+    >>> solution = solve_qubo(qubo)
     """
     max_abs_value = float(max(abs(v) for v in qubo.values()))
     scale_qubo = {k: float(v) / max_abs_value for k, v in qubo.items()}
-    sa = dimod.reference.SimulatedAnnealingSampler()
+    sa = neal.SimulatedAnnealingSampler()
     sa_computation = sa.sample_qubo(scale_qubo, num_reads=num_reads,
-                                    num_sweeps=num_sweeps, beta_range=beta_range)
+                                    sweeps=sweeps, beta_range=beta_range)
     best = np.argmin(sa_computation.record.energy)
     best_solution = list(sa_computation.record.sample[best])
     return dict(zip(sa_computation.variable_labels, best_solution))
 
 
-def solve_ising(linear, quad, num_reads=10, num_sweeps=1000, beta_range=(1.0, 50.0)):
-    """Solve Ising model with Simulated Annealing (SA) provided by dimod.
+def solve_ising(linear, quad, num_reads=10, sweeps=1000, beta_range=(1.0, 50.0)):
+    """Solve Ising model with Simulated Annealing (SA) provided by neal.
 
     Args:
         linear (dict[label, float]): The linear parameter of the Ising model.
@@ -51,19 +58,26 @@ def solve_ising(linear, quad, num_reads=10, num_sweeps=1000, beta_range=(1.0, 50
 
         num_reads (int, default=10): Number of run repetitions of SA.
 
-        num_sweeps (int, default=1000): Number of iterations in each run of SA.
+        sweeps (int, default=1000): Number of iterations in each run of SA.
 
         beta_range (tuple(float, float), default=(1.0, 50.0)): Tuple of start beta and end beta.
 
     Returns:
          dict[label, bit]: The solution of SA.
+    
+    >>> from pyqubo import Spin, solve_ising
+    >>> s1, s2, s3 = Spin("s1"), Spin("s2"), Spin("s3")
+    >>> H = (2*s1 + 4*s2 + 6*s3)**2
+    >>> model = H.compile()
+    >>> linear, quad, offset = model.to_ising()
+    >>> solution = solve_ising(linear, quad)
     """
     max_abs_value = float(max(abs(v) for v in (list(quad.values()) + list(linear.values()))))
     scale_linear = {k: float(v) / max_abs_value for k, v in linear.items()}
     scale_quad = {k: float(v) / max_abs_value for k, v in quad.items()}
-    sa = dimod.reference.SimulatedAnnealingSampler()
+    sa = neal.SimulatedAnnealingSampler()
     sa_computation = sa.sample_ising(scale_linear, scale_quad, num_reads=num_reads,
-                                     num_sweeps=num_sweeps, beta_range=beta_range)
+                                     sweeps=sweeps, beta_range=beta_range)
     best = np.argmin(sa_computation.record.energy)
     best_solution = list(sa_computation.record.sample[best])
     return dict(zip(sa_computation.variable_labels, best_solution))

requirements.txt

@@ -8,3 +8,4 @@ ipykernel==4.9.0
 runipy==0.1.5
 matplotlib==2.2.3
 networkx==2.0
+dwave-neal==0.4.2

setup.py

@@ -17,7 +17,8 @@ def __getattribute__(self, name):
 
 install_requires = ['dimod>=0.7.4',
                     'numpy>=1.14.0,<2.0.0',
-                    'six>=1.10.0,<2.0.0']
+                    'six>=1.10.0,<2.0.0',
+                    'dwave-neal>=0.4.2']
 
 packages = ['pyqubo', 'pyqubo.core', 'pyqubo.utils']
 

test/test_solver.py

@@ -28,11 +28,11 @@ def create_number_partition_model():
     def test_solve_qubo(self):
         model = TestSolver.create_number_partition_model()
         qubo, offset = model.to_qubo()
-        solution = solve_qubo(qubo, num_reads=1, num_sweeps=10)
+        solution = solve_qubo(qubo, num_reads=1, sweeps=10)
         self.assertTrue(solution == {'s1': 0, 's2': 0, 's3': 1} or {'s1': 1, 's2': 1, 's3': 0})
 
     def test_solve_ising(self):
         model = TestSolver.create_number_partition_model()
         linear, quad, offset = model.to_ising()
-        solution = solve_ising(linear, quad, num_reads=1, num_sweeps=10)
+        solution = solve_ising(linear, quad, num_reads=1, sweeps=10)
         self.assertTrue(solution == {'s1': -1, 's2': -1, 's3': 1} or {'s1': 1, 's2': 1, 's3': -1})