dbac 0.1.0__py3-none-any.whl

dbac/__init__.py

@@ -0,0 +1,12 @@
+"""dbac: physics and numerics for double-bracket algorithmic cooling (plotting lives in `dbac.plotting`)."""
+
+import warnings
+
+import numpy as np
+
+__version__ = "0.1.0"
+
+# Silence noisy-but-harmless warnings: complex values being cast to real when
+# reading off energies/fidelities, and GCI's "edge of interval" step-size notice.
+warnings.filterwarnings("ignore", category=np.exceptions.ComplexWarning)
+warnings.filterwarnings("ignore", message=".*edge of interval reached.*")

dbac/dbqite.py

@@ -0,0 +1,192 @@
+"""
+dbqite.py
+
+One step of double-bracket quantum imaginary-time evolution (DB-QITE), the
+cooling step of DBAC, for a single qubit with Hamiltonian H = -Z.
+
+A step maps the current state psi to exp(i s rho') psi exp(-i s rho'), where
+rho' = exp(i s H) |psi><psi| exp(-i s H) is a rotated copy of psi used as the
+instruction of density matrix exponentiation (DME). All step functions share
+the signature `step(psi, s) -> density matrix`, so they can be passed to
+`dbac.gci.GCI` to run several steps.
+"""
+
+__all__ = [
+    "DME",
+    "QITE_GCI_dme",
+    "QITE_dme_select_M",
+    "initial_psi",
+    "ideal_DME",
+    "H", "H_sh", "eigs", "target_energy"
+]
+
+import numpy as np
+from scipy.linalg import expm
+from qibo import hamiltonians
+from qibo.symbols import Z
+
+# Constants and Hamiltonian setup
+L = 1
+#: The same Hamiltonian as a qibo ``SymbolicHamiltonian``, for qibo-circuit states.
+H_sh = hamiltonians.SymbolicHamiltonian(Z(0), nqubits=1)
+#: The Hamiltonian H = -Z = diag(-1, 1) whose ground state ``|0>`` DBAC cools towards.
+H = np.array([[-1.-0.j, -0.-0.j],
+              [-0.-0.j,  1.-0.j]])
+#: Eigenvalues of `H`: the ground-state energy -1 and the excited-state energy +1.
+eigs = np.array([-1., 1.])
+#: Energy of the target (ground) state ``|0>``.
+target_energy = -1
+
+def _as_density_matrix(psi: np.ndarray) -> np.ndarray:
+    """Return `psi` as a density matrix: a 2x2 input is taken to already be one
+    (e.g. the output of a previous DB-QITE step, or a mixed state), a (2, 1)
+    state vector is turned into |psi><psi|."""
+    if psi.ndim == 2 and psi.shape[0] == psi.shape[1]:
+        return psi
+    psi = psi.reshape(-1, 1)
+    return psi @ psi.conj().T
+
+def DME(instruction: np.ndarray, data: np.ndarray, t: float, M: int = 100) -> np.ndarray:
+    """
+    Apply exp(i t instruction) to `data` with Trotterized density matrix
+    exponentiation, without ever exponentiating `instruction` itself.
+
+    Each of the `M` steps is the DME channel for duration t/M with a fresh copy
+    of `instruction`, written in closed form:
+    ``cos^2 data + sin^2 instruction + i sin cos [instruction, data]``.
+    The result approaches ``exp(i t instruction) data exp(-i t instruction)``
+    with an error of order t^2 / M.
+
+    Note the argument order: the instruction comes first, unlike the functions
+    in `dbac.dme`.
+
+    Args:
+        instruction (np.ndarray): 2x2 density matrix used as the Hamiltonian
+            (the instruction state).
+        data (np.ndarray): 2x2 density matrix to transform (the data state).
+        t (float): Total DME duration.
+        M (int): Number of Trotter steps, i.e. instruction copies consumed.
+
+    Returns:
+        np.ndarray: The transformed 2x2 data density matrix.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.dbqite import DME
+        >>> zero = np.array([[1, 0], [0, 0]], dtype=complex)    # data |0><0|
+        >>> plus = np.full((2, 2), 0.5, dtype=complex)          # instruction |+><+|
+        >>> # t = pi/2 with one step is a full SWAP: the data becomes the instruction
+        >>> DME(plus, zero, t=np.pi / 2, M=1).real.round(3)
+        array([[0.5, 0.5],
+               [0.5, 0.5]])
+    """
+    _t = t / M
+    c = np.cos(_t)
+    s = np.sin(_t)
+    for _ in range(M):
+        data = c**2 * data + s**2 * instruction + 1j * s * c * (instruction @ data - data @ instruction)
+    return data
+
+def QITE_GCI_dme(psi: np.ndarray, s: float, M: int = 100) -> np.ndarray:
+    """
+    One DB-QITE step with Trotterized DME, using `M` = 100 DME steps by default.
+
+    Equivalent to ``QITE_dme_select_M(M)(psi, s)``; prefer `QITE_dme_select_M`
+    when passing a step function with a specific `M` to `dbac.gci.GCI`.
+
+    Args:
+        psi (np.ndarray): Current state, a (2, 1) state vector or 2x2 density matrix.
+        s (float): Step duration.
+        M (int): Number of DME steps.
+
+    Returns:
+        np.ndarray: The 2x2 density matrix after the step.
+    """
+    QITE_GCI_dme.label = r"DME"
+    QITE_GCI_dme.marker = '^'
+    PSI = _as_density_matrix(psi)
+    U_t = expm(1j * s * H)
+    return DME(U_t @ PSI @ U_t.conj().T, PSI, s, M=M)
+
+def QITE_dme_select_M(M: int = 100):
+    """
+    Build a DB-QITE step function that uses `M` DME (Trotter) steps.
+
+    The returned function ``step(psi, s)`` performs one DBAC step of duration
+    `s` on `psi`: it rotates a copy of the state about H to get the
+    instruction ``exp(i s H) psi exp(-i s H)``, then applies it to `psi` with
+    `DME` using `M` steps. `M` = 1 and `M` = 2 correspond to circuits A and B
+    of the paper's experiment.
+
+    Args:
+        M (int): Number of DME steps per DBAC step.
+
+    Returns:
+        Callable[[np.ndarray, float], np.ndarray]: ``step(psi, s)``, which takes
+        a (2, 1) state vector or 2x2 density matrix and returns the 2x2 density
+        matrix after the step.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.dbqite import QITE_dme_select_M, H
+        >>> from dbac.utils import rx_sweep
+        >>> psi = rx_sweep([0.5])[0]            # energy 0, halfway to |1>
+        >>> rho = QITE_dme_select_M(1)(psi, np.pi / 4)
+        >>> round(float(np.trace(rho @ H).real), 3)
+        -0.5
+    """
+    def QITE_GCI_dme(psi: np.ndarray, s: float, M: int = M) -> np.ndarray:
+        QITE_GCI_dme.label = r"DME"
+        QITE_GCI_dme.marker = '^'
+        PSI = _as_density_matrix(psi)
+        U_t = expm(1j * s * H)
+        return DME(U_t @ PSI @ U_t.conj().T, PSI, s, M=M)
+    return QITE_GCI_dme
+
+def initial_psi(psi: np.ndarray, s: float) -> np.ndarray:
+    """
+    A step function that does nothing: returns `psi` unchanged.
+
+    Useful as the "no cooling" baseline, since it has the same signature as the
+    other step functions and can be passed to `dbac.gci.GCI`.
+
+    Args:
+        psi (np.ndarray): Current state.
+        s (float): Step duration (unused).
+
+    Returns:
+        np.ndarray: The input state.
+    """
+    initial_psi.label = "Initial state"
+    initial_psi.marker = 'o'
+    return psi
+
+def ideal_DME(psi: np.ndarray, s: float) -> np.ndarray:
+    """
+    One DB-QITE step with the reflection applied exactly (the limit M -> infinity).
+
+    Same as ``QITE_dme_select_M(M)(psi, s)`` for very large `M`, but computes
+    ``exp(i s rho')`` by matrix exponential instead of DME. For this step the
+    paper gives the energy afterwards in closed form,
+    ``E0 - 2 sin^2(s) (1 - E0^2) [(1 - cos s) E0 + cos s]``.
+
+    Args:
+        psi (np.ndarray): Current state, a (2, 1) state vector or 2x2 density matrix.
+        s (float): Step duration.
+
+    Returns:
+        np.ndarray: The 2x2 density matrix after the step.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.dbqite import ideal_DME, H
+        >>> from dbac.utils import rx_sweep
+        >>> rho = ideal_DME(rx_sweep([0.5])[0], np.pi / 4)
+        >>> round(float(np.trace(rho @ H).real), 3)   # -1/sqrt(2)
+        -0.707
+    """
+    PSI = _as_density_matrix(psi)
+    U_t = expm(1j * s * H)
+    instruction = U_t @ PSI @ U_t.conj().T
+    V = expm(1j * s * instruction)
+    return V @ PSI @ V.conj().T

dbac/dme.py

@@ -0,0 +1,160 @@
+"""
+dme.py
+
+Implementations of the density matrix exponentiation (DME) channel, which
+applies ``exp(i t rho) sigma exp(-i t rho)`` to a data state `sigma` using
+copies of an instruction state `rho`, without knowing `rho`.
+
+All functions take ``(data, instruction, t, M)``: the data and instruction
+2x2 density matrices, the total duration `t`, and the number of Trotter steps
+`M` (instruction copies consumed). The approximate implementations return
+``(final_data, history)`` where `history` holds the data state after each of
+the `M` steps. They compute the same channel in different ways so they can be
+checked against each other (the paper's Fig. S2).
+"""
+
+__all__ = [
+    "dme_expm", "dme_qibo", "dme_channel", "dme_channel_equivalent", "ptrace"
+]
+
+import numpy as np
+from scipy.linalg import expm
+from qibo.quantum_info.linalg_operations import partial_trace
+from qibo.gates import SWAP
+
+def dme_expm(data, instruction, t, M=None):
+    """
+    The exact evolution ``exp(i t instruction) data exp(-i t instruction)``,
+    i.e. the M -> infinity limit that DME approximates. The reference for the
+    other implementations.
+
+    Parameters:
+        data (np.ndarray): 2x2 data density matrix.
+        instruction (np.ndarray): 2x2 instruction density matrix (the Hamiltonian).
+        t (float): Duration.
+        M: Ignored; accepted so the signature matches the other implementations.
+
+    Returns:
+        np.ndarray: The evolved 2x2 data density matrix (not a tuple).
+    """
+    V = expm(1j * t * instruction)
+    return V @ data @ V.conj().T
+
+def dme_qibo(data, instruction, t, M):
+    """
+    DME built from qibo's SWAP gate matrix and qibo's partial trace.
+
+    Same channel as `dme_channel`; useful to cross-check the numpy
+    implementation against qibo.
+
+    Parameters:
+        data (np.ndarray): 2x2 data density matrix.
+        instruction (np.ndarray): 2x2 instruction density matrix.
+        t (float): Total duration, split into `M` steps of ``t / M``.
+        M (int): Number of Trotter steps.
+
+    Returns:
+        tuple[np.ndarray, list[np.ndarray]]: The final data state, and the data
+        state after each step.
+    """
+    rho_list = []
+    _t = t / M
+    # Not qibo's `matrix_exponentiation`: in qibo 0.3.x it returns expm(phase * M)
+    # without the factor of i, i.e. a non-unitary operator.
+    u = expm(1j * _t * SWAP(0, 1).matrix())
+    dme_state = np.kron(data, instruction)
+    for _ in range(M):
+        dme_state = u @ dme_state @ u.T.conj()
+        data = partial_trace(dme_state, [1])
+        rho_list.append(data)
+        dme_state = np.kron(data, instruction)
+    return data, rho_list
+
+swap = np.array([[1, 0, 0, 0],
+                 [0, 0, 1, 0],
+                 [0, 1, 0, 0],
+                 [0, 0, 0, 1]], dtype=complex)
+
+def ptrace(matrix, index=1, n1=2, n2=2):
+    """
+    Partial trace of a bipartite ``(n1 * n2) x (n1 * n2)`` matrix.
+
+    Parameters:
+        matrix (np.ndarray): Operator on subsystems 1 (dimension `n1`) and 2
+            (dimension `n2`).
+        index (int): Which subsystem to trace **out**: 1 traces out subsystem 2
+            and keeps subsystem 1; 2 traces out subsystem 1 and keeps subsystem 2.
+        n1, n2 (int): Dimensions of the two subsystems.
+
+    Returns:
+        np.ndarray: The reduced operator.
+    """
+    current_tensor = matrix.reshape([n1, n2, n1, n2])
+    if index == 1:
+        return np.trace(current_tensor, axis1=1, axis2=3)
+    elif index == 2:
+        return np.trace(current_tensor, axis1=0, axis2=2)
+    else:
+        print('Error')
+
+def dme_channel(data, instruction, t, M):
+    """
+    DME as on hardware: apply a partial SWAP ``exp(i (t / M) SWAP)`` to
+    ``data (x) instruction``, trace out the instruction, and repeat `M` times
+    with a fresh instruction copy.
+
+    Parameters:
+        data (np.ndarray): 2x2 data density matrix.
+        instruction (np.ndarray): 2x2 instruction density matrix.
+        t (float): Total duration, split into `M` steps of ``t / M``.
+        M (int): Number of Trotter steps (instruction copies consumed).
+
+    Returns:
+        tuple[np.ndarray, list[np.ndarray]]: The final data state, and the data
+        state after each step.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.dme import dme_channel, dme_expm
+        >>> data = np.array([[1, 0], [0, 0]], dtype=complex)
+        >>> instruction = np.full((2, 2), 0.5, dtype=complex)
+        >>> exact = dme_expm(data, instruction, 1.0)
+        >>> [round(float(np.linalg.norm(exact - dme_channel(data, instruction, 1.0, M)[0])), 4)
+        ...  for M in (1, 10, 100)]   # error shrinks like 1/M
+        [0.5969, 0.0757, 0.0079]
+    """
+    rho_list = []
+    _t = t / M
+    u = expm(1j * _t * swap)
+    dme_state = np.kron(data, instruction)
+    for _ in range(M):
+        dme_state = u @ dme_state @ u.T.conj()
+        data = ptrace(dme_state)
+        rho_list.append(data)
+        dme_state = np.kron(data, instruction)
+    return data, rho_list
+
+def dme_channel_equivalent(data, instruction, t, M):
+    """
+    DME using the closed form of one channel step,
+    ``cos^2 data + sin^2 instruction + i sin cos [instruction, data]``
+    (with angle ``t / M``), which avoids building two-qubit matrices.
+
+    Parameters:
+        data (np.ndarray): 2x2 data density matrix.
+        instruction (np.ndarray): 2x2 instruction density matrix.
+        t (float): Total duration, split into `M` steps of ``t / M``.
+        M (int): Number of Trotter steps.
+
+    Returns:
+        tuple[np.ndarray, list[np.ndarray]]: The final data state, and the data
+        state after each step.
+    """
+    rho_list = []
+    _t = t / M 
+    c = np.cos(_t)
+    s = np.sin(_t)
+    for _ in range(M): 
+        data = c**2 * data + s**2 * instruction + 1j * s * c * (instruction @ data - data @ instruction)
+        rho_list.append(data)
+    return data, rho_list

dbac/gci.py

@@ -0,0 +1,226 @@
+"""
+GCI: Group Commutator Iteration, the driver that runs the DBAC recursion.
+
+`GCI` applies a step function (from `dbac.dbqite`) repeatedly, feeding each
+output state back in as the next input, and records energy, fidelity and the
+states along the way. Plotting of the results produced by
+`GCI.execute_gci` lives in `dbac.plotting.gci.plot_gci_results`, so this module has
+no matplotlib dependency.
+"""
+
+__all__ = [
+    "GCI",
+    "max_excitation"
+]
+
+import warnings
+
+from qibo import hamiltonians, Circuit
+from qibo.symbols import Z
+import numpy as np
+
+class GCI:
+    """
+    Run the DBAC recursion: apply a step function `k` times, each time to the
+    output of the previous step, and track energy and fidelity.
+
+    The step durations are chosen in one of two modes, selected by `nmb_steps`:
+
+    - **fixed schedule** (``nmb_steps=None``): `times_dbr` lists the duration
+      of each step, ``[s_1, ..., s_k]``.
+    - **greedy line search** (``nmb_steps=k``): `times_dbr` lists candidate
+      durations; every step uses the candidate that gives the lowest energy.
+      A warning is emitted if the best candidate is at the edge of the list.
+
+    Nothing is computed until `execute_gci` is called.
+
+    Parameters:
+        gci_step (Callable): Step function ``step(psi, s) -> state``, e.g.
+            ``dbac.dbqite.QITE_dme_select_M(1)`` or ``dbac.dbqite.ideal_DME``.
+        state_of_recursion (np.ndarray): Initial state, a (2, 1) state vector
+            or 2x2 density matrix (or a qibo ``Circuit``).
+        nmb_steps (int or None): Number of steps for the line-search mode, or
+            None for the fixed-schedule mode.
+        times_dbr (sequence[float]): The schedule (fixed mode) or the
+            candidate durations (line-search mode).
+        please_be_verbose (bool): Unused; kept for backwards compatibility.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.gci import GCI
+        >>> from dbac.dbqite import QITE_dme_select_M
+        >>> from dbac.utils import rx_sweep
+        >>> psi = rx_sweep([0.5])[0]
+        >>> results = GCI(QITE_dme_select_M(1), psi, None, [np.pi / 4, np.pi / 4]).execute_gci()
+        >>> [round(float(F), 3) for F in results[-1]]   # fidelity with |0> after 0, 1, 2 steps
+        [0.5, 0.75, 0.875]
+    """
+    def __init__(self, gci_step, state_of_recursion, nmb_steps, times_dbr, please_be_verbose=True):
+        self.gci_step = gci_step
+        self.state_of_recursion = state_of_recursion
+        self.psi0 = state_of_recursion
+        self.nmb_steps = nmb_steps
+        self.times_dbr = times_dbr
+        self.please_be_verbose = please_be_verbose
+        self.states = [state_of_recursion]
+
+        # Fixed problem setup for 1-qubit Z Hamiltonian
+        self.L = 1
+        self.H_sh = hamiltonians.SymbolicHamiltonian(Z(0), nqubits=1)
+        self.H = np.array([[-1.-0.j, -0.-0.j],
+                           [-0.-0.j,  1.-0.j]])
+        self.eigs = np.array([-1., 1.])
+        self.target_energy = -1
+
+        # Initial tracking
+        self.results = []
+        self.fidelities = []
+        self.s_mins = [0]
+        self.state_0 = np.array([1, 0], dtype=complex)
+        self.fidelities.append(self.fidelity_with_0(state_of_recursion))
+        self.results.append(self.compute_moments(state_of_recursion))
+
+        if isinstance(state_of_recursion, Circuit):
+            self.circuit_depths = [0]
+            self.circuits = [state_of_recursion]
+
+    def fidelity_with_0(self, state):
+        """Fidelity of a state vector or density matrix with ``|0>``."""
+        shape = np.shape(state)
+        if shape[0] == shape[1]:  # density matrix
+            return np.real(np.conj(self.state_0).T @ state @ self.state_0)
+        else:  # state vector
+            dm = state @ state.T.conj()
+            return np.real(np.conj(self.state_0).T @ dm @ self.state_0)
+
+    def compute_moments(self, psi):
+        """
+        Energy moments of a state with respect to H = -Z.
+
+        Returns:
+            tuple: ``(E, E2, variance)`` with ``E = <H>``, ``E2 = <H^2>`` and
+            ``variance = E2 - E^2``. For a density matrix the values are
+            complex-typed (take ``np.real``); for a qibo circuit only ``E`` is
+            computed.
+        """
+        if isinstance(psi, Circuit):
+            return self.H_sh.expectation(psi(self.psi0.T[0]).state()), None, None
+        shape = np.shape(psi)
+        if shape[0] == shape[1]:  # density matrix
+            E = np.trace(psi @ self.H)
+            S = np.trace(psi @ self.H @ self.H)
+        else:  # state vector
+            E = (psi.conj().T @ self.H @ psi)[0, 0].real
+            S = (psi.conj().T @ self.H @ self.H @ psi)[0, 0].real
+        return E, S, S - E**2
+
+    def execute_gci(self):
+        """
+        Run the recursion and return everything it recorded.
+
+        Every list in the result has one entry for the initial state followed
+        by one entry per step.
+
+        Returns:
+            tuple: An 11-tuple, best read from both ends (``results[0]`` for
+            energies, ``results[-1]`` for fidelities, ``results[-2]`` for the
+            final state):
+
+            0. energies: ``Tr(rho H)`` after each step (take ``np.real``).
+            1. step sizes: the durations used, starting with 0.
+            2. variances: energy variance after each step.
+            3. durations: cumulative evolution time, the sum of ``s_j ** 2``.
+            4. label: display label of the step function.
+            5. marker: matplotlib marker of the step function.
+            6. circuit depths: for qibo-circuit states only, else None.
+            7. circuits: for qibo-circuit states only, else None.
+            8. states: the state after each step.
+            9. final state: the last entry of `states`.
+            10. fidelities: fidelity with ``|0>`` after each step.
+        """
+        if self.nmb_steps is None:
+            for s in self.times_dbr:
+                self.state_of_recursion = self.gci_step(self.state_of_recursion, s)
+                self.results.append(self.compute_moments(self.state_of_recursion))
+                self.fidelities.append(self.fidelity_with_0(self.state_of_recursion))
+                self.states.append(self.state_of_recursion)
+            self.s_mins = [0] + list(self.times_dbr)
+        else:
+            for k in range(self.nmb_steps):
+                test_es = [self.compute_moments(self.gci_step(self.state_of_recursion, dt))[0] for dt in self.times_dbr]
+                ind = np.argmin(test_es)
+                if ind == 0 or ind == len(self.times_dbr) - 1:
+                    warnings.warn(f"{k} {ind} {self.times_dbr[ind]} --- Warning edge of interval reached")
+                s_min = self.times_dbr[ind]
+                self.state_of_recursion = self.gci_step(self.state_of_recursion, s_min)
+                self.results.append(self.compute_moments(self.state_of_recursion))
+                self.fidelities.append(self.fidelity_with_0(self.state_of_recursion))
+                self.s_mins.append(s_min)
+                self.states.append(self.state_of_recursion)
+                if isinstance(self.state_of_recursion, Circuit):
+                    self.circuit_depths.append(len(self.state_of_recursion.queue) * 3)
+                    self.circuits.append(self.state_of_recursion)
+
+        energies = [r[0] for r in self.results]
+        variances = [r[2] for r in self.results]
+
+        # Add default display labels if not already set
+        if not hasattr(self.gci_step, 'exponent'):
+            self.gci_step.exponent = 2
+        if not hasattr(self.gci_step, 'label'):
+            self.gci_step.label = ''
+        if not hasattr(self.gci_step, 'marker'):
+            self.gci_step.marker = '-'
+
+        durations = (
+            [sum(s**self.gci_step.exponent for s in self.s_mins[1:x]) for x in range(1, len(self.s_mins) + 1)]
+            if self.gci_step.exponent > 0 else self.s_mins
+        )
+
+        return (
+            energies, self.s_mins, variances,
+            durations,
+            self.gci_step.label, self.gci_step.marker,
+            self.circuit_depths if isinstance(self.state_of_recursion, Circuit) else None,
+            self.circuits if isinstance(self.state_of_recursion, Circuit) else None,
+            self.states,
+            self.state_of_recursion,
+            self.fidelities
+        )
+
+def max_excitation(thetas, fidelity, target_fidelity, initial_tolerance=1e-5, max_tolerance=0.2, step=1e-3):
+    """
+    Find where a cooling curve first reaches a target fidelity.
+
+    Scans ``(thetas[i], fidelity[i])`` in order and returns the first
+    ``thetas[i]`` whose fidelity is within a tolerance of `target_fidelity`.
+    The tolerance starts at `initial_tolerance` and grows by `step` until a
+    match is found or it exceeds `max_tolerance`. With a sweep ordered from
+    hot to cold, the result is the hottest initial state that is still cooled
+    to the target: everything colder ends up at least as cold.
+
+    Parameters:
+        thetas (sequence[float]): x-values of the curve, e.g. the initial
+            fidelities or angles of a sweep.
+        fidelity (sequence[float]): Final fidelity for each x-value.
+        target_fidelity (float): Fidelity to look for.
+        initial_tolerance (float): Starting tolerance.
+        max_tolerance (float): Largest tolerance tried before giving up.
+        step (float): Increment of the tolerance.
+
+    Returns:
+        float or None: The matching x-value, or None if no point is within
+        `max_tolerance` of the target.
+
+    Example:
+        >>> from dbac.gci import max_excitation
+        >>> max_excitation([0.1, 0.2, 0.3, 0.4], [0.5, 0.7, 0.9, 0.95], target_fidelity=0.9)
+        0.3
+    """
+    tolerance = initial_tolerance
+    while tolerance <= max_tolerance:
+        for xi, yi in zip(thetas, fidelity):
+            if abs(yi - target_fidelity) <= tolerance:
+                return xi
+        tolerance += step
+    return None

dbac/hbac.py

@@ -0,0 +1,100 @@
+"""
+hbac.py
+
+Heat-bath algorithmic cooling (HBAC) baseline used for comparison against
+DBAC. HBAC repeatedly SWAPs entropy onto a fresh, highly-polarized "reset"
+qubit and discards it, compressing the joint state towards its ground state
+one qubit at a time.
+"""
+
+__all__ = ["refresh_reset_qubit", "compress_density_matrix", "hbac"]
+
+import numpy as np
+
+from dbac.dme import ptrace
+
+
+def refresh_reset_qubit(rho, epsilon_b):
+    """
+    Reset the last qubit of a joint density matrix to a thermal state with
+    polarization `epsilon_b`, by rescaling each 2x2 block of `rho` that acts
+    on that qubit.
+
+    Parameters:
+        rho (ndarray): 2^n x 2^n density matrix of n qubits.
+        epsilon_b (float): Polarization of the reset qubit.
+
+    Returns:
+        ndarray: New density matrix after resetting the last qubit.
+    """
+    d = rho.shape[0]
+    assert rho.shape == (d, d), "Input must be a square matrix"
+    assert (d & (d - 1)) == 0, "Dimension must be a power of 2"
+
+    num_blocks = d // 2  # since each M_ij is 2x2
+    new_rho = np.zeros_like(rho, dtype=complex)
+
+    thermal_matrix = np.array([
+        [(1 + epsilon_b) / 2, 0],
+        [0, (1 - epsilon_b) / 2],
+    ])
+
+    for i in range(num_blocks):
+        for j in range(num_blocks):
+            row_start, row_end = i * 2, i * 2 + 2
+            col_start, col_end = j * 2, j * 2 + 2
+            m_ij = rho[row_start:row_end, col_start:col_end]
+            new_rho[row_start:row_end, col_start:col_end] = np.trace(m_ij) * thermal_matrix
+
+    return new_rho
+
+
+def compress_density_matrix(rho):
+    """
+    Sort-and-compress step of HBAC: reorder the populations of `rho` in
+    descending order, discarding coherences, to approximate the state
+    achievable by the optimal (sorting) unitary.
+
+    Parameters:
+        rho (ndarray): Density matrix (Hermitian, positive semidefinite, trace 1).
+
+    Returns:
+        ndarray: Diagonal density matrix with sorted probabilities.
+    """
+    assert rho.shape[0] == rho.shape[1], "Matrix must be square"
+    assert np.allclose(rho, rho.conj().T), "Matrix must be Hermitian"
+
+    probs = np.real(np.diag(rho))
+    sorted_probs = np.sort(probs)[::-1]
+    return np.diag(sorted_probs)
+
+
+def hbac(rho, epsilon_b=0.8, step=1):
+    """
+    Run `step` rounds of heat-bath algorithmic cooling (reset + compress) on
+    `rho`, then trace out every qubit except the first (the target qubit,
+    which the sorting step cools).
+
+    Parameters:
+        rho (ndarray): 2^n x 2^n density matrix, last qubit is the reset qubit.
+        epsilon_b (float): Polarization of the fresh reset qubit each round,
+            i.e. populations ``(1 +- epsilon_b) / 2``.
+        step (int): Number of reset-and-compress rounds.
+
+    Returns:
+        ndarray: 2x2 reduced density matrix of the first qubit.
+
+    Example:
+        >>> import numpy as np
+        >>> from dbac.hbac import hbac
+        >>> mixed = np.eye(2) / 2                      # maximally mixed qubit
+        >>> rho = np.kron(np.kron(mixed, mixed), mixed)
+        >>> hbac(rho, epsilon_b=0.8, step=6).real.round(3)
+        array([[0.985, 0.   ],
+               [0.   , 0.015]])
+    """
+    for _ in range(step):
+        rho = refresh_reset_qubit(rho, epsilon_b)
+        rho = compress_density_matrix(rho)
+    d = rho.shape[0]
+    return ptrace(rho, index=1, n1=2, n2=int(d / 2))