eqc-bio 0.1.2__py3-none-any.whl

eqc_bio/protein_design/__init__.py

@@ -0,0 +1,30 @@
+# (C) Quantum Computing Inc., 2024.
+from .fixed_backbone import ProteinDesignFixed
+from .viewer import (
+    build_designed_pdb,
+    get_backbone_dihedrals,
+    get_rotamer_angles,
+    view_pdb,
+)
+from .rotamers import (
+    download_rotamer_library,
+    load_rotamer_library,
+    parse_dunbrack_library,
+    resolve_rotamer_library,
+    check_rotamer_counts,
+    get_chi_angles,
+)
+
+__all__ = [
+    "ProteinDesignFixed",
+    "build_designed_pdb",
+    "get_backbone_dihedrals",
+    "get_rotamer_angles",
+    "view_pdb",
+    "download_rotamer_library",
+    "load_rotamer_library",
+    "parse_dunbrack_library",
+    "resolve_rotamer_library",
+    "check_rotamer_counts",
+    "get_chi_angles",
+]

eqc_bio/protein_design/base.py

@@ -0,0 +1,168 @@
+import os
+import sys
+import warnings
+import logging
+import numpy as np
+import pandas as pd
+
+from eqc_models.base import ModelSolver, SolutionResults
+from eqc_models.solvers.qciclient import Dirac3CloudSolver
+from eqc_models import QuadraticModel
+from eqc_bio.utils import timer
+from eqc_bio.utils import metis_partition
+
+warnings.filterwarnings("ignore")
+log = logging.getLogger(name=__name__)
+
+# Seconds per unit for SolutionResults.time_units
+TIME_UNIT_SECONDS = {"ns": 1e-9, "s": 1.0}
+
+
+class ProteinDesign:
+    def __init__(self, solver: ModelSolver = None):
+        # Defaults to a cloud solver configured from QCI_API_URL/QCI_TOKEN
+        self.solver = solver
+
+    def get_dirac3_time(self, results: SolutionResults):
+        """Return the device time of a solve in seconds."""
+        device_time = results.device_time
+        if device_time is None:
+            log.warning("Device timing not available for this solve")
+            return 0.0
+
+        return np.sum(device_time) * TIME_UNIT_SECONDS[results.time_units]
+
+    def run_dirac3(self, J, C, sum_constraint, n_vars, num_samples):
+        model = QuadraticModel(C, J)
+        model.upper_bound = np.array([1] * n_vars)
+
+        solver = self.solver
+        if solver is None:
+            solver = Dirac3CloudSolver()
+
+        results = solver.solve(
+            model,
+            sum_constraint=sum_constraint,
+            relaxation_schedule=self.relaxation_schedule,
+            num_samples=num_samples,
+            mean_photon_number=self.mean_photon_number,
+            quantum_fluctuation_coefficient=self.quantum_fluctuation_coefficient,
+        )
+        log.info(results)
+
+        dirac3_time = self.get_dirac3_time(results)
+
+        return results, dirac3_time
+
+    def run_dirac3_iterative(
+        self,
+        J,
+        C,
+        sum_constraint,
+        n_vars,
+        positions,
+        position_graph,
+        position_index_hash,
+    ):
+        batch_count = int(round(n_vars / self.batch_size))
+
+        X = np.random.rand(n_vars)
+
+        if sum_constraint is not None:
+            X = X * (sum_constraint / np.sum(X))
+
+        tmp_group_list = metis_partition(position_graph, k=batch_count)
+        tmp_hash = dict(enumerate(positions))
+        position_groups = []
+        for tmp_group in tmp_group_list:
+            position_groups.append([tmp_hash[item] for item in tmp_group])
+
+        groups = []
+        for position_group in position_groups:
+            tmp_list = []
+            for position in position_group:
+                tmp_list += position_index_hash[position]
+            groups.append(tmp_list)
+
+        sols = []
+        energies = []
+        dirac3_time = 0.0
+        for epoch in range(self.n_epochs):
+            order = np.random.permutation(batch_count)
+            for i in order:
+                idx_p = groups[i]
+                idx_q = np.setdiff1d(np.arange(n_vars), idx_p)
+                Cp = C[idx_p]
+                Cq = C[idx_q]
+                Xq = X[idx_q]
+                Jpp = J[np.ix_(idx_p, idx_p)]
+                Jpq = J[np.ix_(idx_p, idx_q)]
+                Jqp = J[np.ix_(idx_q, idx_p)]
+                Jqq = J[np.ix_(idx_q, idx_q)]
+
+                tmp_v = 2 * Jpq @ Xq
+                tmp_v = tmp_v.reshape((tmp_v.shape[0], 1))
+                tmp_C = Cp + tmp_v
+                tmp_J = Jpp
+
+                if sum_constraint is not None:
+                    tmp_sum_constraint = sum_constraint - np.sum(Xq)
+                else:
+                    tmp_sum_constraint = None
+
+                sub_response, sub_dirac3_time = self.run_dirac3(
+                    tmp_J,
+                    tmp_C,
+                    max(1.0, tmp_sum_constraint),
+                    tmp_C.shape[0],
+                    self.n_samples,
+                )
+                dirac3_time += sub_dirac3_time
+
+                sub_sols = sub_response.solutions
+
+                glob_loc_hash = {}
+                for loc_k, glob_k in enumerate(idx_p):
+                    glob_loc_hash[glob_k] = loc_k
+
+                for sub_sol in sub_sols:
+                    sub_sol = np.array(sub_sol)
+                    feasible_flag = True
+                    for position in position_groups[i]:
+                        glob_indexes = position_index_hash[position]
+                        loc_indexes = [
+                            glob_loc_hash[item] for item in glob_indexes
+                        ]
+                        if np.sum(sub_sol[loc_indexes]) == 0:
+                            feasible_flag = False
+                            break
+
+                    if feasible_flag:
+                        break
+
+                Xp = sub_sol
+                fct = np.sum(Xp)
+                if fct != 0:
+                    fct = tmp_sum_constraint / fct
+                Xp = fct * np.array(Xp)
+                X[idx_p] = Xp
+
+                sol = np.array(X)
+                energy = sol.T @ J @ sol + sol.T @ C
+                energy = energy[0]
+
+                sols.append(sol)
+                energies.append(energy)
+
+        # Device time is summed over the sub-solves and returned separately
+        results = SolutionResults(
+            solutions=np.array(sols),
+            energies=np.array(energies),
+            counts=np.ones(len(sols)),
+            objectives=None,
+            run_time=None,
+            preprocessing_time=None,
+            postprocessing_time=None,
+        )
+
+        return results, dirac3_time