fix beta indexing for PT

pysa/ais.py

@@ -77,7 +77,7 @@ def combine_logZf(solution: pd.DataFrame):
 
 
 @numba.njit(fastmath=True, nogil=True, parallel=False)
-def get_log_omega(betas: Vector, energies: Vector):
+def get_log_omega(betas: Vector, beta_idx: List[int], energies: Vector):
     '''Takes in the temperatures and energies of an annealing
     procedure and outputs what log_omega the AIS weight used for
     calculating the log of the partition function'''
@@ -89,7 +89,7 @@ def get_log_omega(betas: Vector, energies: Vector):
 
     # Sort temperatures and states
     cur_betas = betas[_sort]
-    cur_energies = energies[_sort]
+    cur_energies = energies[beta_idx[_sort]]
 
     if cur_betas[0] != 0:
         raise ValueError(

pysa/sa.py

@@ -319,15 +319,16 @@ def _simulate_core():
                 # Start timer
                 t_ini = time()
 
-                # Get states and energy
+                # Get states and energy of replicas
                 w = _init_strategy()
-
+                # Assign initial temperatures sequentially to the replicas
+                beta_idx = np.arange(num_replicas)
                 # End timer
                 t_end = time()
 
                 return simulation(self._module.update_spin,
                                   pysa.simulation.random_sweep, couplings,
-                                  local_fields, *w, betas, num_sweeps,
+                                  local_fields, *w, beta_idx, betas, num_sweeps,
                                   get_part_fun, use_pt), t_end - t_ini
 
         elif update_strategy == 'sequential':
@@ -339,13 +340,15 @@ def _simulate_core():
 
                 # Get states and energy
                 w = _init_strategy()
+                # Assign initial temperatures sequentially to the replicas
+                beta_idx = np.arange(num_replicas)
 
                 # End timer
                 t_end = time()
 
                 return simulation(self._module.update_spin,
                                   pysa.simulation.sequential_sweep, couplings,
-                                  local_fields, *w, betas, num_sweeps,
+                                  local_fields, *w, beta_idx, betas, num_sweeps,
                                   get_part_fun, use_pt), t_end - t_ini
 
         else:
@@ -358,7 +361,7 @@ def _simulate():
             t_ini = time()
 
             # Simulate
-            ((out_states, out_energies, out_betas, out_log_omegas),
+            ((out_states, out_energies, out_beta_idx, out_log_omegas),
              (best_state, best_energy, best_sweeps,
               ns)), init_time = _simulate_core()
 
@@ -376,13 +379,12 @@ def _simulate():
             # Sort output temperatures (if required)
             if sort_output_temps:
 
-                # Argsort temperatures
-                _sort = np.argsort(out_betas)[::-1]
-
                 # Sort temperatures and states
-                out_betas = out_betas[_sort]
-                out_states = out_states[_sort]
-                out_energies = out_energies[_sort]
+                out_betas = betas
+                out_states = out_states[out_beta_idx]
+                out_energies = out_energies[out_beta_idx]
+            else:
+                out_betas = betas[out_beta_idx]
 
             with np.errstate(divide='ignore'):
                 out_temps = np.divide(1, out_betas)

pysa/simulation.py

@@ -127,6 +127,7 @@ def simulation_parallel(update_spin: UpdateSpinFunction,
                         local_fields: Vector,
                         states: List[State],
                         energies: List[float],
+                        beta_idx: List[int],
                         betas: List[float],
                         n_sweeps: int,
                         get_part_fun: bool = False,
@@ -142,18 +143,19 @@ def simulation_parallel(update_spin: UpdateSpinFunction,
     _best_energy = np.copy(energies)
     _best_state = np.copy(states)
     _best_sweeps = np.zeros(n_replicas, dtype=np.int32)
-
+    betas_sorted = np.empty_like(betas)
     log_omegas = np.zeros(n_sweeps)
 
     # For each run ...
     for s in range(n_sweeps):
-
+        for k in range(n_replicas):
+            betas_sorted[beta_idx[k]] = betas[k]
         # ... apply sweep for each replica ...
         for k in numba.prange(n_replicas):
 
             # Apply sweep
             energies[k] += sweep(update_spin, couplings, local_fields,
-                                 states[k], betas[k])
+                                 states[k], betas_sorted[k])
 
             # Store best state
             if energies[k] < _best_energy[k]:
@@ -163,10 +165,10 @@ def simulation_parallel(update_spin: UpdateSpinFunction,
 
         # ... and pt move.
         if use_pt:
-            utils.pt(states, energies, betas)
+            utils.pt(states, energies, beta_idx, betas)
         # Calculate the weights for the partition function
         if get_part_fun:
-            log_omegas[s] = get_log_omega(betas, energies)
+            log_omegas[s] = get_log_omega(betas, beta_idx, energies)
 
     # Get lowest energy
     best_pos = np.argmin(_best_energy)
@@ -175,8 +177,8 @@ def simulation_parallel(update_spin: UpdateSpinFunction,
     best_sweeps = _best_sweeps[best_pos]
 
     # Return states and energies
-    return ((states, energies, betas, log_omegas), (best_state, best_energy,
-                                                    best_sweeps, s + 1))
+    return ((states, energies, beta_idx, log_omegas), (best_state, best_energy,
+                                                       best_sweeps, s + 1))
 
 
 @numba.njit(fastmath=True, nogil=True, parallel=False)
@@ -186,6 +188,7 @@ def simulation_sequential(update_spin: UpdateSpinFunction,
                           local_fields: Vector,
                           states: List[State],
                           energies: List[float],
+                          beta_idx: List[int],
                           betas: List[float],
                           n_sweeps: int,
                           get_part_fun: bool = False,
@@ -201,18 +204,19 @@ def simulation_sequential(update_spin: UpdateSpinFunction,
     best_state = states[0]
     best_energy = energies[0]
     best_sweeps = 0
-
+    betas_sorted = np.empty_like(betas)
     log_omegas = np.zeros(n_sweeps)
 
     # For each run ...
     for s in range(n_sweeps):
-
+        for k in range(n_replicas):
+            betas_sorted[beta_idx[k]] = betas[k]
         # ... apply sweep for each replica ...
         for k in range(n_replicas):
 
             # Apply sweep
             energies[k] += sweep(update_spin, couplings, local_fields,
-                                 states[k], betas[k])
+                                 states[k], betas_sorted[k])
 
             # Store best state
             if energies[k] < best_energy:
@@ -222,14 +226,14 @@ def simulation_sequential(update_spin: UpdateSpinFunction,
 
         # ... and pt move.
         if use_pt:
-            utils.pt(states, energies, betas)
+            utils.pt(states, energies, beta_idx, betas)
         # Calculate the weights for the partition function
         if get_part_fun:
-            log_omegas[s] = get_log_omega(betas, energies)
+            log_omegas[s] = get_log_omega(betas, beta_idx, energies)
 
     # Return states and energies
-    return ((states, energies, betas, log_omegas), (best_state, best_energy,
-                                                    best_sweeps, s + 1))
+    return ((states, energies, beta_idx, log_omegas), (best_state, best_energy,
+                                                       best_sweeps, s + 1))
 
 
 @numba.njit(fastmath=True, nogil=True, parallel=False)

pysa/utils.py

@@ -24,10 +24,17 @@
 
 
 @numba.njit(cache=True, fastmath=True, nogil=True, parallel=False)
-def pt(states: List[State], energies: List[float],
+def pt(states: List[State], energies: List[float], beta_idx: List[int],
        betas: List[float]) -> NoReturn:
     """
   Parallel tempering move.
+    states: [n_replicas, ...]  Array of replicas
+    energies: [n_replicas] Array of energies of each replica
+    beta_idx: [n_replicas] The replica index currently assigned to each beta,
+        i.e. inverse temperature K is currently used for simulating replica beta_idx[K]
+    betas: [n_replicas] Sequential array of inverse temperatures.
+
+    This function only modifies the order of beta_idx.
   """
 
     # Get number of replicas
@@ -43,11 +50,12 @@ def pt(states: List[State], energies: List[float],
         k2 = n_replicas - k - 2
 
         # Compute delta energy
-        de = (energies[k1] - energies[k2]) * (betas[k1] - betas[k2])
+        de = (energies[beta_idx[k1]] - energies[beta_idx[k2]]) * (betas[k1] -
+                                                                  betas[k2])
 
         # Accept/reject following Metropolis
         if de >= 0 or np.random.random() < np.exp(de):
-            betas[k1], betas[k2] = betas[k2], betas[k1]
+            beta_idx[k1], beta_idx[k2] = beta_idx[k2], beta_idx[k1]
 
 
 def get_problem_matrix(instance: List[Tuple[int, int, float]]) -> Matrix:

tests/tests_ising_parallel.py

@@ -190,6 +190,7 @@ def test_pt(n_vars: int, n_replicas: int):
     # Get random temperatures
     betas = np.random.random(n_replicas) * 20
     betas_copy = np.copy(betas)
+    beta_idx = np.arange(n_replicas)
 
     # Get initial state
     states = 2 * np.random.randint(2,
@@ -202,7 +203,7 @@ def test_pt(n_vars: int, n_replicas: int):
     energies_copy = np.copy(energies)
 
     # Apply PT
-    utils.pt(states, energies, betas)
+    utils.pt(states, energies, beta_idx, betas)
 
     # Check that betas are just shuffled
     assert (np.allclose(sorted(betas), sorted(betas_copy)))
@@ -235,7 +236,7 @@ def test_random_sweep_simulation_ising(n_vars: int):
 
     # Fix temperature
     betas = np.array([1], dtype=dtype)
-
+    beta_idx = np.arange(1)
     # Get initial state
     states = 2 * np.random.randint(2, size=(1, n_vars)).astype(dtype) - 1
 
@@ -246,7 +247,7 @@ def test_random_sweep_simulation_ising(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_parallel(
         ising.update_spin, simulation.random_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,
@@ -268,7 +269,7 @@ def test_sequential_sweep_simulation_ising(n_vars: int):
 
     # Fix temperature
     betas = np.array([1], dtype=dtype)
-
+    beta_idx = np.arange(1)
     # Get initial state
     states = 2 * np.random.randint(2, size=(1, n_vars)).astype(dtype) - 1
 
@@ -279,7 +280,7 @@ def test_sequential_sweep_simulation_ising(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_parallel(
         ising.update_spin, simulation.sequential_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,

tests/tests_ising_sequential.py

@@ -190,6 +190,7 @@ def test_pt(n_vars: int, n_replicas: int):
     # Get random temperatures
     betas = np.random.random(n_replicas) * 20
     betas_copy = np.copy(betas)
+    beta_idx = np.arange(n_replicas)
 
     # Get initial state
     states = 2 * np.random.randint(2,
@@ -202,7 +203,7 @@ def test_pt(n_vars: int, n_replicas: int):
     energies_copy = np.copy(energies)
 
     # Apply PT
-    utils.pt(states, energies, betas)
+    utils.pt(states, energies, beta_idx, betas)
 
     # Check that betas are just shuffled
     assert (np.allclose(sorted(betas), sorted(betas_copy)))
@@ -235,7 +236,7 @@ def test_random_sweep_simulation_ising(n_vars: int):
 
     # Fix temperature
     betas = np.array([1], dtype=dtype)
-
+    beta_idx = np.arange(1)
     # Get initial state
     states = 2 * np.random.randint(2, size=(1, n_vars)).astype(dtype) - 1
 
@@ -246,7 +247,7 @@ def test_random_sweep_simulation_ising(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_sequential(
         ising.update_spin, simulation.random_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,
@@ -268,7 +269,7 @@ def test_sequential_sweep_simulation_ising(n_vars: int):
 
     # Fix temperature
     betas = np.array([1], dtype=dtype)
-
+    beta_idx = np.arange(1)
     # Get initial state
     states = 2 * np.random.randint(2, size=(1, n_vars)).astype(dtype) - 1
 
@@ -279,7 +280,7 @@ def test_sequential_sweep_simulation_ising(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_sequential(
         ising.update_spin, simulation.sequential_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,

tests/tests_qubo_parallel.py

@@ -190,6 +190,7 @@ def test_pt(n_vars: int, n_replicas: int):
     # Get random temperatures
     betas = np.random.random(n_replicas) * 20
     betas_copy = np.copy(betas)
+    beta_idx = np.arange(n_replicas)
 
     # Get initial state
     states = np.random.randint(2, size=(n_replicas, n_vars)).astype(dtype)
@@ -201,7 +202,7 @@ def test_pt(n_vars: int, n_replicas: int):
     energies_copy = np.copy(energies)
 
     # Apply PT
-    utils.pt(states, energies, betas)
+    utils.pt(states, energies, beta_idx, betas)
 
     # Check that betas are just shuffled
     assert (np.allclose(sorted(betas), sorted(betas_copy)))
@@ -234,7 +235,7 @@ def test_random_sweep_simulation_qubo(n_vars: int):
 
     # Fix temperature
     betas = np.array([10], dtype=dtype)
-
+    beta_idx = np.arange(1)
     # Get initial state
     states = np.random.randint(2, size=(1, n_vars)).astype(dtype)
 
@@ -245,7 +246,7 @@ def test_random_sweep_simulation_qubo(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_parallel(
         qubo.update_spin, simulation.random_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,
@@ -267,7 +268,7 @@ def test_sequential_sweep_simulation_qubo(n_vars: int):
 
     # Fix temperature
     betas = np.array([10], dtype=dtype)
-
+    beta = np.arange(1)
     # Get initial state
     states = np.random.randint(2, size=(1, n_vars)).astype(dtype)
 
@@ -278,7 +279,7 @@ def test_sequential_sweep_simulation_qubo(n_vars: int):
     # Simulate
     _, (best_state, best_energy, _, _) = simulation.simulation_parallel(
         qubo.update_spin, simulation.sequential_sweep, couplings, local_fields,
-        states, energies, betas, 10000)
+        states, energies, beta_idx, betas, 10000)
 
     # Check that best energy is correct
     assert (np.isclose(best_energy,