diff --git a/pysa/ais.py b/pysa/ais.py index f5b1bf6..2f0295f 100644 --- a/pysa/ais.py +++ b/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( diff --git a/pysa/sa.py b/pysa/sa.py index 6ee2700..036846a 100644 --- a/pysa/sa.py +++ b/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) diff --git a/pysa/simulation.py b/pysa/simulation.py index 82663b8..2917ab4 100644 --- a/pysa/simulation.py +++ b/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) diff --git a/pysa/utils.py b/pysa/utils.py index 87a24cd..6f4032f 100644 --- a/pysa/utils.py +++ b/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: diff --git a/tests/tests_ising_parallel.py b/tests/tests_ising_parallel.py index 2843e34..57a5bd6 100644 --- a/tests/tests_ising_parallel.py +++ b/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, diff --git a/tests/tests_ising_sequential.py b/tests/tests_ising_sequential.py index 398594d..1f6a0f6 100644 --- a/tests/tests_ising_sequential.py +++ b/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, diff --git a/tests/tests_qubo_parallel.py b/tests/tests_qubo_parallel.py index 373db8c..5b28b95 100644 --- a/tests/tests_qubo_parallel.py +++ b/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, diff --git a/tests/tests_qubo_sequential.py b/tests/tests_qubo_sequential.py index 62518e6..d9c6f63 100644 --- a/tests/tests_qubo_sequential.py +++ b/tests/tests_qubo_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 = 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))) @@ -237,7 +238,7 @@ def test_random_sweep_simulation_qubo(n_vars: int): # Get initial state states = np.random.randint(2, size=(1, n_vars)).astype(dtype) - + beta_idx = np.arange(1) # Compute energies energies = np.array( [qubo.get_energy(couplings, local_fields, state) for state in states]) @@ -245,7 +246,7 @@ def test_random_sweep_simulation_qubo(n_vars: int): # Simulate _, (best_state, best_energy, _, _) = simulation.simulation_sequential( 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_idx = 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_sequential( 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,