From 7fd1b13651da70ca3731d9c2e8e7a84ed4eeb88a Mon Sep 17 00:00:00 2001
From: Andrej Dolenc <andr3.dolenc@gmail.com>
Date: Sun, 7 Apr 2024 23:37:52 +0000
Subject: [PATCH] Allow exiting the anneal early via .min_energy

---
 README.md           |  2 ++
 examples/nqueens.py | 36 ++++++++++++++++++++++++++++++++++++
 simanneal/anneal.py |  5 +++++
 3 files changed, 43 insertions(+)
 create mode 100644 examples/nqueens.py

diff --git a/README.md b/README.md
index 26a647e..413ff28 100644
--- a/README.md
+++ b/README.md
@@ -121,6 +121,8 @@ tsp.set_schedule(auto_schedule)
 itinerary, miles = tsp.anneal()
 ```
 
+Sometimes you have a problem for which an optimal solution exists and you can calculate its energy. One such example is the n-queens problem, where 0 queens attacking each-other is optimal. In such cases you may specify that you want the algorithm to terminate after finding the solution by setting `min_energy` member to the energy of the optimal solution.
+
 ## Extra data dependencies
 
 You might have noticed that the `energy` function above requires a `cities` dict 
diff --git a/examples/nqueens.py b/examples/nqueens.py
new file mode 100644
index 0000000..8a3320f
--- /dev/null
+++ b/examples/nqueens.py
@@ -0,0 +1,36 @@
+import random
+from simanneal import Annealer
+
+
+class NQueensProblem(Annealer):
+    def move(self):
+        """Move a queen on a random column to some different row."""
+        column = random.randrange(len(self.state))
+        new_row = random.randrange(len(self.state))
+        self.state[column] = new_row
+
+    def energy(self):
+        """Calculates the number of attacks among the queens."""
+        e = 0
+        for i in range(len(self.state)):
+            for j in range(i + 1, len(self.state)):
+                e += self.state[i] == self.state[j]
+                e += abs(i - j) == abs(self.state[i] - self.state[j])
+        return e
+
+
+if __name__ == '__main__':
+    init_state = list(range(10))
+    random.shuffle(init_state)
+    print(init_state)
+
+    nqueens = NQueensProblem(init_state)
+    nqueens.set_schedule(nqueens.auto(minutes=0.2))
+    # stop as soon as we hit 0 attacks
+    nqueens.min_energy = 0
+    nqueens.copy_strategy = "slice"
+    state, e = nqueens.anneal()
+
+    print()
+    print("number of attacks: %i" % e)
+    print(state)
diff --git a/simanneal/anneal.py b/simanneal/anneal.py
index 3b70d81..cfb5c3f 100644
--- a/simanneal/anneal.py
+++ b/simanneal/anneal.py
@@ -40,6 +40,7 @@ class Annealer(object):
     Tmin = 2.5
     steps = 50000
     updates = 100
+    min_energy = None
     copy_strategy = 'deepcopy'
     user_exit = False
     save_state_on_exit = False
@@ -192,6 +193,8 @@ def anneal(self):
         prevEnergy = E
         self.best_state = self.copy_state(self.state)
         self.best_energy = E
+        if self.min_energy is not None and self.best_energy <= self.min_energy:
+            return self.best_state, self.best_energy
         trials = accepts = improves = 0
         if self.updates > 0:
             updateWavelength = self.steps / self.updates
@@ -222,6 +225,8 @@ def anneal(self):
                 if E < self.best_energy:
                     self.best_state = self.copy_state(self.state)
                     self.best_energy = E
+                    if self.min_energy is not None and self.best_energy <= self.min_energy:
+                        return self.best_state, self.best_energy
             if self.updates > 1:
                 if (step // updateWavelength) > ((step - 1) // updateWavelength):
                     self.update(