From dd619adb45ae6f4278277ea93fc2db3d8b263f59 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 19 Feb 2024 15:16:28 +0900
Subject: [PATCH 01/17] Create requirements.txt

---
 pruners/requirements.txt | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 pruners/requirements.txt

diff --git a/pruners/requirements.txt b/pruners/requirements.txt
new file mode 100644
index 00000000..c3db4451
--- /dev/null
+++ b/pruners/requirements.txt
@@ -0,0 +1 @@
+numba

From 55c5739caa0c6ef12e9c78111958a7b7f51d499d Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 19 Feb 2024 15:20:47 +0900
Subject: [PATCH 02/17] Create wilcoxon_pruner_tsp_sa.py

---
 pruners/wilcoxon_pruner_tsp_sa.py | 120 ++++++++++++++++++++++++++++++
 1 file changed, 120 insertions(+)
 create mode 100644 pruners/wilcoxon_pruner_tsp_sa.py

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
new file mode 100644
index 00000000..c4d0a89e
--- /dev/null
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -0,0 +1,120 @@
+import math
+from typing import NamedTuple
+
+import numba
+import numpy as np
+import optuna
+from numpy.linalg import norm
+
+
+class SAOptions(NamedTuple):
+    max_iter: int = 1000
+    T0: float = 1.0
+    alpha: float = 1.0
+    patience: int = 300
+
+
+@numba.njit
+def simulated_annealing(vertices, initial_idxs, options: SAOptions):
+
+    def temperature(t: float):
+        # t: 0 ... 1
+        return options.T0 * (1 - t) ** options.alpha
+
+    idxs = initial_idxs.copy()
+    N = len(vertices)
+    assert len(idxs) == N
+
+    cost = sum(
+        [norm(vertices[idxs[i]] - vertices[idxs[(i + 1) % N]]) for i in range(N)]
+    )
+    best_idxs = idxs.copy()
+    best_cost = cost
+
+    remaining_patience = options.patience
+    np.random.seed(11111)
+
+    for iter in range(options.max_iter):
+        i = np.random.randint(0, N)
+        j = (i + 2 + np.random.randint(0, N - 3)) % N
+        i, j = min(i, j), max(i, j)
+        delta_cost = (
+            -norm(vertices[idxs[(i + 1) % N]] - vertices[idxs[i]])
+            - norm(vertices[idxs[j]] - vertices[idxs[(j + 1) % N]])
+            + norm(vertices[idxs[i]] - vertices[idxs[j]])
+            + norm(vertices[idxs[(i + 1) % N]] - vertices[idxs[(j + 1) % N]])
+        )
+        temp = temperature(iter / options.max_iter)
+        if np.random.rand() < math.exp(-delta_cost / temp):
+            cost += delta_cost
+            idxs[i + 1 : j + 1] = idxs[i + 1 : j + 1][::-1]
+            if cost < best_cost:
+                best_idxs[:] = idxs
+                best_cost = cost
+
+        if cost >= best_cost:
+            remaining_patience -= 1
+            if remaining_patience == 0:
+                idxs[:] = best_idxs
+                cost = best_cost
+                remaining_patience = options.patience
+
+    return best_idxs
+
+
+def make_dataset(num_vertex, num_problem, seed):
+    rng = np.random.default_rng(seed=seed)
+    dataset = []
+    for _ in range(num_problem):
+        dataset.append(
+            {
+                "vertices": rng.random((num_vertex, 2)),
+                "idxs": rng.permutation(num_vertex),
+            }
+        )
+    return dataset
+
+
+dataset = make_dataset(200, 20, seed=33333)
+rng = np.random.default_rng(seed=44444)
+
+
+def objective(trial):
+    patience = trial.suggest_int("patience", 10, 10000, log=True)
+    T0 = trial.suggest_float("T0", 0.1, 10.0, log=True)
+    alpha = trial.suggest_float("alpha", 1.1, 10.0, log=True)
+    options = SAOptions(max_iter=10000000, patience=patience, T0=T0, alpha=alpha)
+    ordering = rng.permutation(range(len(dataset)))
+    results = []
+    for i in ordering:
+        d = dataset[i]
+        result_idxs = simulated_annealing(d["vertices"], d["idxs"], options)
+        result_cost = 0.0
+        n = len(d["vertices"])
+        for j in range(n):
+            result_cost += norm(
+                d["vertices"][result_idxs[j]] - d["vertices"][result_idxs[(j + 1) % n]]
+            )
+        results.append(result_cost)
+
+        trial.report(result_cost, i)
+        if trial.should_prune():
+            # Wilcoxon pruner found that this trial was 
+            # probably worse than the current best trial.
+            # However, this trial may be in top 10% trials.
+            # So I return the current average score instead of
+            # raise optuna.TrialPruned().
+            # It provides additional information to TPESampler.
+            sum(results) / len(results)  
+
+    return sum(results) / len(results)
+
+
+if __name__ == "__main__":
+    sampler = optuna.samplers.TPESampler(seed=55555)
+    pruner = optuna.pruners.WilcoxonPruner(p_threshold=0.05)
+    study = optuna.create_study(direction="minimize", sampler=sampler, pruner=pruner)
+    study.enqueue_trial({"patience": 300, "T0": 1.0, "alpha": 1.8})  # default params
+    study.optimize(objective, n_trials=100)
+    print(f"The number of trials: {len(study.trials)}")
+    print(f"Best value: {study.best_value} (params: {study.best_params})")

From cd37c3c8152e0527f6d920136a94ac7ed2afd0c9 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 19 Feb 2024 15:28:37 +0900
Subject: [PATCH 03/17] fix lint

---
 pruners/wilcoxon_pruner_tsp_sa.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index c4d0a89e..ddce9bae 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -99,13 +99,13 @@ def objective(trial):
 
         trial.report(result_cost, i)
         if trial.should_prune():
-            # Wilcoxon pruner found that this trial was 
+            # Wilcoxon pruner found that this trial was
             # probably worse than the current best trial.
             # However, this trial may be in top 10% trials.
             # So I return the current average score instead of
             # raise optuna.TrialPruned().
             # It provides additional information to TPESampler.
-            sum(results) / len(results)  
+            sum(results) / len(results)
 
     return sum(results) / len(results)
 

From 84c967fdd0f4db0a750ad4c37c77331c464ab51d Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 19 Feb 2024 15:29:55 +0900
Subject: [PATCH 04/17] fix lint

---
 pruners/wilcoxon_pruner_tsp_sa.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index ddce9bae..92579656 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -25,9 +25,7 @@ def temperature(t: float):
     N = len(vertices)
     assert len(idxs) == N
 
-    cost = sum(
-        [norm(vertices[idxs[i]] - vertices[idxs[(i + 1) % N]]) for i in range(N)]
-    )
+    cost = sum([norm(vertices[idxs[i]] - vertices[idxs[(i + 1) % N]]) for i in range(N)])
     best_idxs = idxs.copy()
     best_cost = cost
 

From 51f6ce9e5f91fdcde3a5641c1dbb18bf35951b0e Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 19 Feb 2024 15:33:42 +0900
Subject: [PATCH 05/17] fix isort

---
 pruners/wilcoxon_pruner_tsp_sa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index 92579656..0c729537 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -3,8 +3,8 @@
 
 import numba
 import numpy as np
-import optuna
 from numpy.linalg import norm
+import optuna
 
 
 class SAOptions(NamedTuple):

From 853a319a39a6b40962f887b35ceb589d96062a48 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:02:29 +0900
Subject: [PATCH 06/17] remove numba from requirements.txt

---
 pruners/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pruners/requirements.txt b/pruners/requirements.txt
index c3db4451..8b137891 100644
--- a/pruners/requirements.txt
+++ b/pruners/requirements.txt
@@ -1 +1 @@
-numba
+

From 9a7b013d9362611137c1296373c3547941d66fe3 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:15:10 +0900
Subject: [PATCH 07/17] Update wilcoxon_pruner_tsp_sa.py

---
 pruners/wilcoxon_pruner_tsp_sa.py | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index 0c729537..de3fa5a9 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -1,7 +1,6 @@
 import math
 from typing import NamedTuple
 
-import numba
 import numpy as np
 from numpy.linalg import norm
 import optuna
@@ -14,7 +13,6 @@ class SAOptions(NamedTuple):
     patience: int = 300
 
 
-@numba.njit
 def simulated_annealing(vertices, initial_idxs, options: SAOptions):
 
     def temperature(t: float):
@@ -73,18 +71,21 @@ def make_dataset(num_vertex, num_problem, seed):
     return dataset
 
 
-dataset = make_dataset(200, 20, seed=33333)
+N_DATAPOINTS, N_TRIALS = 20, 100
+dataset = make_dataset(200, N_DATAPOINTS, seed=33333)
 rng = np.random.default_rng(seed=44444)
+count = 0
 
 
 def objective(trial):
-    patience = trial.suggest_int("patience", 10, 10000, log=True)
+    patience = trial.suggest_int("patience", 10, 1000, log=True)
     T0 = trial.suggest_float("T0", 0.1, 10.0, log=True)
     alpha = trial.suggest_float("alpha", 1.1, 10.0, log=True)
-    options = SAOptions(max_iter=10000000, patience=patience, T0=T0, alpha=alpha)
+    options = SAOptions(max_iter=10000, patience=patience, T0=T0, alpha=alpha)
     ordering = rng.permutation(range(len(dataset)))
     results = []
     for i in ordering:
+        count += 1
         d = dataset[i]
         result_idxs = simulated_annealing(d["vertices"], d["idxs"], options)
         result_cost = 0.0
@@ -97,13 +98,8 @@ def objective(trial):
 
         trial.report(result_cost, i)
         if trial.should_prune():
-            # Wilcoxon pruner found that this trial was
-            # probably worse than the current best trial.
-            # However, this trial may be in top 10% trials.
-            # So I return the current average score instead of
-            # raise optuna.TrialPruned().
-            # It provides additional information to TPESampler.
-            sum(results) / len(results)
+            sum(results) / len(results)  # An advanced technique
+            # raise optuna.TrialPruned()
 
     return sum(results) / len(results)
 
@@ -113,6 +109,7 @@ def objective(trial):
     pruner = optuna.pruners.WilcoxonPruner(p_threshold=0.05)
     study = optuna.create_study(direction="minimize", sampler=sampler, pruner=pruner)
     study.enqueue_trial({"patience": 300, "T0": 1.0, "alpha": 1.8})  # default params
-    study.optimize(objective, n_trials=100)
+    study.optimize(objective, n_trials=N_TRIALS)
     print(f"The number of trials: {len(study.trials)}")
     print(f"Best value: {study.best_value} (params: {study.best_params})")
+    print(f"Number of evaluations: {count} / {N_DATAPOINTS * N_TRIALS}")

From 451eef827d28342054f7fe6190b8ba268914915c Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:19:57 +0900
Subject: [PATCH 08/17] Create wilcoxon_pruner.yml

---
 .github/workflows/wilcoxon_pruner.yml | 37 +++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 .github/workflows/wilcoxon_pruner.yml

diff --git a/.github/workflows/wilcoxon_pruner.yml b/.github/workflows/wilcoxon_pruner.yml
new file mode 100644
index 00000000..9ac608a6
--- /dev/null
+++ b/.github/workflows/wilcoxon_pruner.yml
@@ -0,0 +1,37 @@
+name: dask_ml
+
+on:
+  schedule:
+    - cron: '0 15 * * *'
+  pull_request:
+    paths:
+      - 'pruners/**'
+      - '.github/workflows/wilcoxon_pruner.yml'
+
+jobs:
+  examples:
+    if: (github.event_name == 'schedule' && github.repository == 'optuna/optuna-examples') || (github.event_name != 'schedule')
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: setup-python${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install (Python)
+      run: |
+        python -m pip install --upgrade pip
+        pip install --progress-bar off -U setuptools
+        pip install git+https://github.com/optuna/optuna.git
+        python -c 'import optuna'
+
+        pip install -r pruners/requirements.txt
+    - name: Run examples
+      run: |
+        python pruners/wilcoxon_pruner_tsp_sa.py
+      env:
+        OMP_NUM_THREADS: 1

From c4d001606df5aa4e9a05fbfeba1de58da3d1af72 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:26:13 +0900
Subject: [PATCH 09/17] fix flake8

---
 pruners/wilcoxon_pruner_tsp_sa.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index de3fa5a9..ab602740 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -78,6 +78,7 @@ def make_dataset(num_vertex, num_problem, seed):
 
 
 def objective(trial):
+    global count
     patience = trial.suggest_int("patience", 10, 1000, log=True)
     T0 = trial.suggest_float("T0", 0.1, 10.0, log=True)
     alpha = trial.suggest_float("alpha", 1.1, 10.0, log=True)

From 878636c84e26de92915d7e6924724b4eccb59940 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:28:23 +0900
Subject: [PATCH 10/17] avoid math.exp error

---
 pruners/wilcoxon_pruner_tsp_sa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index ab602740..77b7926c 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -41,7 +41,7 @@ def temperature(t: float):
             + norm(vertices[idxs[(i + 1) % N]] - vertices[idxs[(j + 1) % N]])
         )
         temp = temperature(iter / options.max_iter)
-        if np.random.rand() < math.exp(-delta_cost / temp):
+        if delta_cost <= 0.0 or np.random.rand() < math.exp(-delta_cost / temp):
             cost += delta_cost
             idxs[i + 1 : j + 1] = idxs[i + 1 : j + 1][::-1]
             if cost < best_cost:

From 83658b2f12a1ade575a1934fafce857e2263017a Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 21 Feb 2024 11:29:48 +0900
Subject: [PATCH 11/17] add scipy to requirements.txt

---
 pruners/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pruners/requirements.txt b/pruners/requirements.txt
index 8b137891..9a635b91 100644
--- a/pruners/requirements.txt
+++ b/pruners/requirements.txt
@@ -1 +1 @@
-
+scipy

From cdb7de32afacb5a48096592d5e5fb62e51244d75 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 26 Feb 2024 15:33:16 +0900
Subject: [PATCH 12/17] modify name

---
 .github/workflows/wilcoxon_pruner.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/wilcoxon_pruner.yml b/.github/workflows/wilcoxon_pruner.yml
index 9ac608a6..63568018 100644
--- a/.github/workflows/wilcoxon_pruner.yml
+++ b/.github/workflows/wilcoxon_pruner.yml
@@ -1,4 +1,4 @@
-name: dask_ml
+name: pruners
 
 on:
   schedule:

From ccba822fd89c98918546ab63ea0cc334014e469f Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Mon, 26 Feb 2024 16:02:21 +0900
Subject: [PATCH 13/17] Rename wilcoxon_pruner.yml to pruners.yml

---
 .github/workflows/{wilcoxon_pruner.yml => pruners.yml} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename .github/workflows/{wilcoxon_pruner.yml => pruners.yml} (100%)

diff --git a/.github/workflows/wilcoxon_pruner.yml b/.github/workflows/pruners.yml
similarity index 100%
rename from .github/workflows/wilcoxon_pruner.yml
rename to .github/workflows/pruners.yml

From 49dab7554209aae69978704fc4d52e7400d0b61c Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 28 Feb 2024 19:39:38 +0900
Subject: [PATCH 14/17] fix a bug

---
 pruners/wilcoxon_pruner_tsp_sa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index 77b7926c..7ddb658e 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -99,7 +99,7 @@ def objective(trial):
 
         trial.report(result_cost, i)
         if trial.should_prune():
-            sum(results) / len(results)  # An advanced technique
+            return sum(results) / len(results)  # An advanced workaround
             # raise optuna.TrialPruned()
 
     return sum(results) / len(results)

From df980046d927a162b14e5d7ddd9394caf2c64749 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 28 Feb 2024 20:56:53 +0900
Subject: [PATCH 15/17] Update wilcoxon_pruner_tsp_sa.py

---
 pruners/wilcoxon_pruner_tsp_sa.py | 123 +++++++++++++++++++-----------
 1 file changed, 77 insertions(+), 46 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index 7ddb658e..2491ab4a 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -1,39 +1,57 @@
 import math
-from typing import NamedTuple
+import sys
+from dataclasses import dataclass
 
 import numpy as np
-from numpy.linalg import norm
 import optuna
+from numpy.linalg import norm
 
 
-class SAOptions(NamedTuple):
-    max_iter: int = 1000
+@dataclass
+class SAOptions:
+    max_iter: int = 10000
     T0: float = 1.0
-    alpha: float = 1.0
-    patience: int = 300
+    alpha: float = 2.0
+    patience: int = 50
+
+
+def tsp_cost(vertices: np.ndarray, idxs: np.ndarray) -> float:
+    return norm(vertices[idxs] - vertices[np.roll(idxs, 1)], axis=-1).sum()
+
 
+# Greedy solution for initial guess.
+def tsp_greedy(vertices: np.ndarray) -> np.ndarray:
+    idxs = [0]
+    for _ in range(len(vertices) - 1):
+        dists_from_last = norm(vertices[idxs[-1], None] - vertices, axis=-1)
+        dists_from_last[idxs] = np.inf
+        idxs.append(np.argmin(dists_from_last))
+    return np.array(idxs)
 
-def simulated_annealing(vertices, initial_idxs, options: SAOptions):
+
+# A minimal implementation of TSP solver using simulated annealing on 2-opt neighbors.
+def tsp_simulated_annealing(vertices: np.ndarray, options: SAOptions) -> np.ndarray:
 
     def temperature(t: float):
         # t: 0 ... 1
         return options.T0 * (1 - t) ** options.alpha
 
-    idxs = initial_idxs.copy()
     N = len(vertices)
-    assert len(idxs) == N
 
-    cost = sum([norm(vertices[idxs[i]] - vertices[idxs[(i + 1) % N]]) for i in range(N)])
+    idxs = tsp_greedy(vertices)
+    cost = tsp_cost(vertices, idxs)
     best_idxs = idxs.copy()
     best_cost = cost
-
     remaining_patience = options.patience
-    np.random.seed(11111)
 
     for iter in range(options.max_iter):
+
         i = np.random.randint(0, N)
         j = (i + 2 + np.random.randint(0, N - 3)) % N
         i, j = min(i, j), max(i, j)
+        # Reverse the order of vertices between range [i+1, j].
+
+        # cost difference by 2-opt reversal
         delta_cost = (
             -norm(vertices[idxs[(i + 1) % N]] - vertices[idxs[i]])
             - norm(vertices[idxs[j]] - vertices[idxs[(j + 1) % N]])
@@ -41,14 +59,18 @@ def temperature(t: float):
             + norm(vertices[idxs[(i + 1) % N]] - vertices[idxs[(j + 1) % N]])
         )
         temp = temperature(iter / options.max_iter)
-        if delta_cost <= 0.0 or np.random.rand() < math.exp(-delta_cost / temp):
+        if delta_cost <= 0.0 or np.random.random() < math.exp(-delta_cost / temp):
+            # accept the 2-opt reversal
             cost += delta_cost
             idxs[i + 1 : j + 1] = idxs[i + 1 : j + 1][::-1]
             if cost < best_cost:
                 best_idxs[:] = idxs
                 best_cost = cost
+                remaining_patience = options.patience
 
-        if cost >= best_cost:
+        if cost > best_cost:
+            # If the best solution is not updated for "patience" iteratoins,
+            # restart from the best solution.
             remaining_patience -= 1
             if remaining_patience == 0:
                 idxs[:] = best_idxs
@@ -58,59 +80,68 @@ def temperature(t: float):
     return best_idxs
 
 
-def make_dataset(num_vertex, num_problem, seed):
+def make_dataset(num_vertex: int, num_problem: int, seed: int = 0) -> np.ndarray:
     rng = np.random.default_rng(seed=seed)
-    dataset = []
-    for _ in range(num_problem):
-        dataset.append(
-            {
-                "vertices": rng.random((num_vertex, 2)),
-                "idxs": rng.permutation(num_vertex),
-            }
-        )
-    return dataset
+    return rng.random((num_problem, num_vertex, 2))
+
+
+dataset = make_dataset(
+    num_vertex=100,
+    num_problem=50,
+)
 
+N_TRIALS = 50
 
-N_DATAPOINTS, N_TRIALS = 20, 100
-dataset = make_dataset(200, N_DATAPOINTS, seed=33333)
-rng = np.random.default_rng(seed=44444)
+# We set a very small number of SA iterations for demonstration purpose.
+# In practice, you should set a larger number of iterations.
+N_SA_ITER = 10000
 count = 0
 
 
-def objective(trial):
+def objective(trial: optuna.Trial) -> float:
     global count
-    patience = trial.suggest_int("patience", 10, 1000, log=True)
-    T0 = trial.suggest_float("T0", 0.1, 10.0, log=True)
-    alpha = trial.suggest_float("alpha", 1.1, 10.0, log=True)
-    options = SAOptions(max_iter=10000, patience=patience, T0=T0, alpha=alpha)
-    ordering = rng.permutation(range(len(dataset)))
+    options = SAOptions(
+        max_iter=N_SA_ITER,
+        T0=trial.suggest_float("T0", 0.01, 10.0, log=True),
+        alpha=trial.suggest_float("alpha", 1.0, 10.0, log=True),
+        patience=trial.suggest_int("patience", 10, 1000, log=True),
+    )
     results = []
+
+    # For best results, shuffle the evaluation order in each trial.
+    ordering = np.random.permutation(len(dataset))
     for i in ordering:
         count += 1
-        d = dataset[i]
-        result_idxs = simulated_annealing(d["vertices"], d["idxs"], options)
-        result_cost = 0.0
-        n = len(d["vertices"])
-        for j in range(n):
-            result_cost += norm(
-                d["vertices"][result_idxs[j]] - d["vertices"][result_idxs[(j + 1) % n]]
-            )
+        result_idxs = tsp_simulated_annealing(vertices=dataset[i], options=options)
+        result_cost = tsp_cost(dataset[i], result_idxs)
         results.append(result_cost)
 
         trial.report(result_cost, i)
         if trial.should_prune():
-            return sum(results) / len(results)  # An advanced workaround
+            print(
+                f"[{trial.number}] Pruned at {len(results)}/{len(dataset)}",
+                file=sys.stderr,
+            )
             # raise optuna.TrialPruned()
 
+            # Return the current predicted value when pruned.
+            # This is a workaround for the problem that
+            # current TPE sampler cannot utilize pruned trials effectively.
+            return sum(results) / len(results)
+
+    print(
+        f"[{trial.number}] Not pruned ({len(results)}/{len(dataset)})", file=sys.stderr
+    )
     return sum(results) / len(results)
 
 
 if __name__ == "__main__":
-    sampler = optuna.samplers.TPESampler(seed=55555)
-    pruner = optuna.pruners.WilcoxonPruner(p_threshold=0.05)
+    np.random.seed(0)
+    sampler = optuna.samplers.TPESampler(seed=1)
+    pruner = optuna.pruners.WilcoxonPruner(p_threshold=0.1)
     study = optuna.create_study(direction="minimize", sampler=sampler, pruner=pruner)
-    study.enqueue_trial({"patience": 300, "T0": 1.0, "alpha": 1.8})  # default params
+    study.enqueue_trial({"T0": 1.0, "alpha": 2.0, "patience": 50})  # default params
     study.optimize(objective, n_trials=N_TRIALS)
     print(f"The number of trials: {len(study.trials)}")
     print(f"Best value: {study.best_value} (params: {study.best_params})")
-    print(f"Number of evaluations: {count} / {N_DATAPOINTS * N_TRIALS}")
+    print(f"Number of evaluations: {count} / {N_TRIALS * len(dataset)}")

From fcdfcaa540eb308e7c0840ad6b2476574b364512 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 28 Feb 2024 20:58:44 +0900
Subject: [PATCH 16/17] fix black

---
 pruners/wilcoxon_pruner_tsp_sa.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index 2491ab4a..f9b642a9 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -129,9 +129,7 @@ def objective(trial: optuna.Trial) -> float:
             # current TPE sampler cannot utilize pruned trials effectively.
             return sum(results) / len(results)
 
-    print(
-        f"[{trial.number}] Not pruned ({len(results)}/{len(dataset)})", file=sys.stderr
-    )
+    print(f"[{trial.number}] Not pruned ({len(results)}/{len(dataset)})", file=sys.stderr)
     return sum(results) / len(results)
 
 

From a5b1442387fe973a317c1e3f501c6c139ef42d80 Mon Sep 17 00:00:00 2001
From: Hiroki Takizawa <contact@hiroki-takizawa.name>
Date: Wed, 28 Feb 2024 21:00:49 +0900
Subject: [PATCH 17/17] fix isort

---
 pruners/wilcoxon_pruner_tsp_sa.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pruners/wilcoxon_pruner_tsp_sa.py b/pruners/wilcoxon_pruner_tsp_sa.py
index f9b642a9..2ee3c75a 100644
--- a/pruners/wilcoxon_pruner_tsp_sa.py
+++ b/pruners/wilcoxon_pruner_tsp_sa.py
@@ -1,10 +1,10 @@
+from dataclasses import dataclass
 import math
 import sys
-from dataclasses import dataclass
 
 import numpy as np
-import optuna
 from numpy.linalg import norm
+import optuna
 
 
 @dataclass