Warm start QAOA

aquapointer/digital/ansatz.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from qiskit.circuit.library import QAOAAnsatz, TwoLocal
+from qiskit import QuantumCircuit
+from scipy.optimize import minimize
+
+from aquapointer.digital.qubo_utils import get_ising_hamiltonian
+
+def quadratic_function(x, sigma, mu):
+    return x.T @ sigma @ x + mu.T @ x
+
+def QAOA_ansatz(qubo, warm_start=True, reps=1):
+    hamiltonian = get_ising_hamiltonian(qubo=qubo)
+    num_qubits = len(qubo)
+
+    if warm_start:
+        bounds = [(0, 1)] * len(qubo)
+        x0 = np.zeros(len(qubo))
+
+        result = minimize(quadratic_function, x0, args=(qubo - np.diag(np.diagonal(qubo)), np.diagonal(qubo)), bounds=bounds)
+        x_stars = result.x
+        initial_state = QuantumCircuit(num_qubits)
+        thetas = [2 * np.arcsin(np.sqrt(x_star)) for x_star in x_stars]
+
+
+        for idx, theta in enumerate(thetas):
+            initial_state.ry(theta, idx)
+
+        # QAOA ansatz circuit
+        qaoa_ansatz = QAOAAnsatz(hamiltonian, reps=reps, initial_state=initial_state)
+
+        return qaoa_ansatz
+
+    else:
+        qaoa_ansatz = QAOAAnsatz(hamiltonian, reps=reps, initial_state=initial_state)
+
+        return qaoa_ansatz
+
+
+

aquapointer/digital/loaddata.py

@@ -5,22 +5,61 @@
 
 import numpy as np
 import pickle
-
+from aquapointer.slicing import density_file_to_grid, density_slices_by_axis
 from pathlib import Path
+from aquapointer.density_canvas.DensityCanvas import DensityCanvas
 
 BASE_PATH = str(Path.cwd().parent)
 DENS_DIR = "/data/MUP1/MUP1_logfilter8_slices/"
 PP_DIR = "/data/MUP1/MUP1_logfilter8_points/"
 REG_DIR = "/registers/"
 
+RISM3D_DIR = "../data/3D-RISM_densities/"
+
 class LoadData:
 
-    def __init__(self) -> None:
+    def __init__(self, protein: str) -> None:
         self.d_list = [-1.0, -0.5, 0.0, 0.5, 1.0, 1.5]
-        self.densities = self.load_density_slices(path=BASE_PATH + DENS_DIR)
-        self.plane_points = self.load_plane_points(path=BASE_PATH + PP_DIR)
-        self.register_positions = self.load_register_positions(path=BASE_PATH + REG_DIR)
-        self.rescaled_register_positions = self.load_rescaled_register_positions(path=BASE_PATH + REG_DIR)
+
+        if protein == 'MUP1':
+            self.densities = self.load_density_slices(path=BASE_PATH + DENS_DIR)
+            self.plane_points = self.load_plane_points(path=BASE_PATH + PP_DIR)
+
+            self.register_positions = self.load_register_positions(path=BASE_PATH + REG_DIR)
+            self.rescaled_register_positions = self.load_rescaled_register_positions(path=BASE_PATH + REG_DIR)
+
+        elif protein in ["1NNC", "bromoD", "dehydratase", "HIV1", "test_from_Watsite"]:
+            grid = density_file_to_grid("../data/3D-RISM_densities/test_from_Watsite/prot_3drism.O.1.dx")
+            self.plane_points, self.densities = density_slices_by_axis(grid, axis=np.array([0, 0, 1]), distances=np.array([10, 20, 30]))
+            self.rescaled_register_positions = self.get_rescaled_register_positions()
+            # with open(RISM3D_DIR + protein + '/reg_rescaled_positions.pkl', 'rb') as handle:
+            #     self.rescaled_register_positions = pickle.load(handle)
+            # with open(RISM3D_DIR + protein + '/slices.pkl', 'rb') as handle:
+            #     self.densities = pickle.load(handle)
+
+        else:
+            print(f"there is no 3D RISM data for {protein}")
+
+
+    def get_rescaled_register_positions(self):
+        origin = (-20, -20)
+        length = 40
+        npoints = 80
+        canvas = DensityCanvas(
+            origin=origin,
+            length_x=length,
+            length_y=length,
+            npoints_x=npoints,
+            npoints_y=npoints,
+        )
+        rescaled_positions = []
+        for density in self.densities:    
+            canvas.set_density_from_slice(density)
+            canvas.set_poisson_disk_lattice(spacing=(2,10))
+            rescaled_positions.append(canvas._lattice._coords)
+
+        return rescaled_positions
+
 
     def load_density_slices(self, path: str) -> list[np.ndarray]:
         r"""The 3D-RISM density slices are saved as pickled files in the folder MUP1.
@@ -42,6 +81,7 @@ def load_density_slices(self, path: str) -> list[np.ndarray]:
 
         return densities
 
+
     def load_plane_points(self, path: str) -> list[np.ndarray]:
         r"""Load slice coordinates (these are 3D coordinates in
         angstroms, they are needed at the very end to map

aquapointer/digital/qubo.py

@@ -5,17 +5,60 @@
 
 import numpy as np
 from qiskit.quantum_info import SparsePauliOp
-from aquapointer.digital.loaddata import LoadData
 from aquapointer.digital.qubo_utils import gaussian, gaussian_mixture, gamma, Vij
+from aquapointer.density_canvas.DensityCanvas import DensityCanvas
 
 class Qubo:
 
-    def __init__(self, loaddata: LoadData) -> None:
-        self.ld = loaddata
-        self.qubo_matrices = self.get_qubo_matrices(densities=self.ld.densities, rescaled_positions=self.ld.rescaled_register_positions)
+    def __init__(self, densities, rescaled_register_positions) -> None:
+
+        self.qubo_matrices = self.get_qubo_matrices(densities=densities, rescaled_positions=rescaled_register_positions)
+        # self.qubo_matrices = self.get_qubo_matrices_canvas(densities=densities)
         hamiltonians = [self.get_ising_hamiltonian(qubo=qubo) for qubo in self.qubo_matrices]
         self.qubo_hamiltonian_pairs = list(zip(self.qubo_matrices, hamiltonians))
 
+    def get_qubo_matrices_canvas(self, densities: list[np.ndarray]) -> list[np.ndarray]:
+        estimated_variance = 50
+        estimated_amplitude = 6
+
+        origin = (-20, -20)
+        length = 40
+        npoints = densities[0].shape[0]
+        canvas = DensityCanvas(
+            origin=origin,
+            length_x=length,
+            length_y=length,
+            npoints_x=npoints,
+            npoints_y=npoints,
+        )
+        qubo_matrices = []
+        for density in densities:    
+            canvas.set_density_from_slice(density)
+            # canvas.set_poisson_disk_lattice(spacing=(2,10))
+            canvas.set_rectangular_lattice(num_x=8, num_y=8, spacing=4)
+            canvas.calculate_pubo_coefficients(
+                p = 2, #order of the PUBO, p=2 effectively creates a QUBO
+                params = [estimated_amplitude, estimated_variance]
+            )
+
+            coefficients = canvas._pubo["coeffs"]
+            linear = coefficients[1]
+            quadratic = coefficients[2]
+
+            qubo = np.zeros((len(linear), len(linear)))
+
+            for i, key in enumerate(linear.keys()):
+                qubo[i][i] = linear[key]
+
+            for key in quadratic.keys():
+                qubo[key] = quadratic[key]
+                qubo[key[::-1]] = quadratic[key]
+
+            qubo_matrices.append(qubo)
+
+        return qubo_matrices
+
+
     def get_qubo_matrices(self, densities: list[np.ndarray], rescaled_positions: list[np.ndarray]) -> list[np.ndarray]:
         r""" Given the density slices and rescaled positions of the registers,
         one can compute the corresponding QUBO matrices.