Simulator

simulator

DenseSVSimulator

DenseSVSimulator(
    n_qubits: int,
    use_gpu: bool = False,
    use_float32: bool = False,
)

Dense statevector quantum circuit simulator.

Qubit ordering: MSB-first (qubit 0 is the most significant bit). Backends: NumPy (CPU), JAX XLA JIT (CPU/GPU/TPU).

Parameters

n_qubits : number of qubits use_gpu : reserved for future CuPy/JAX GPU dispatch use_float32: use complex64 instead of complex128

Source code in dense_evolution/simulator.py

def __init__(self, n_qubits: int,
             use_gpu:     bool = False,
             use_float32: bool = False):
    if n_qubits < 1 or n_qubits > 34:
        raise ValueError(f"n_qubits must be in [1, 34], got {n_qubits}")
    self.n         = n_qubits
    self.dim       = 1 << n_qubits            # 2 ** n_qubits
    self.use_float32 = use_float32
    self.dtype     = np.complex64 if use_float32 else np.complex128
    self.xp        = jnp if HAS_JAX else np
    self._reset_sv()

set_initial_state

set_initial_state(state: Optional[ndarray] = None)

Reset the simulator.

Parameters

state : optional complex array of length 2**n. If None, resets to |0...0⟩. The array is normalised automatically.

Source code in dense_evolution/simulator.py

def set_initial_state(self, state: Optional[np.ndarray] = None):
    """
    Reset the simulator.

    Parameters
    ----------
    state : optional complex array of length 2**n.
            If None, resets to |0...0⟩.
            The array is normalised automatically.
    """
    if state is None:
        self._reset_sv()
        return
    state = np.asarray(state, dtype=self.dtype)
    if state.shape != (self.dim,):
        raise ValueError(
            f"State vector length {len(state)} != 2**{self.n} = {self.dim}")
    norm = np.linalg.norm(state)
    if norm < 1e-12:
        raise ValueError("Cannot set a zero-norm state vector")
    state = state / norm
    if HAS_JAX:
        self.sv = jnp.array(state)
    else:
        self.sv = state.copy()

apply_gate_1q

apply_gate_1q(gate: ndarray, qubit: int)

Apply a 2×2 unitary to qubit via tensor contraction.

Uses reshape + moveaxis + matmul — fully vectorised, no Python loops, compatible with both NumPy and JAX.

Source code in dense_evolution/simulator.py

def apply_gate_1q(self, gate: np.ndarray, qubit: int):
    """
    Apply a 2×2 unitary to *qubit* via tensor contraction.

    Uses reshape + moveaxis + matmul — fully vectorised,
    no Python loops, compatible with both NumPy and JAX.
    """
    if not 0 <= qubit < self.n:
        raise ValueError(f"Qubit index {qubit} out of range [0, {self.n})")
    gate = self.xp.array(gate, dtype=self.dtype)
    sv_nd      = self.sv.reshape([2] * self.n)
    sv_moved   = self.xp.moveaxis(sv_nd, qubit, -1)          # qubit axis → last
    flat_shape = (self.dim >> 1, 2)
    # matmul: (dim/2, 2) @ (2, 2).T  → (dim/2, 2)
    result     = self.xp.dot(sv_moved.reshape(flat_shape),
                             gate.T)
    self.sv    = self.xp.moveaxis(
        result.reshape([2] * self.n), -1, qubit).ravel()

apply_gate_2q

apply_gate_2q(gate: ndarray, q1: int, q2: int)

Apply a 4×4 unitary to qubits (q1, q2) via tensor contraction.

Source code in dense_evolution/simulator.py

def apply_gate_2q(self, gate: np.ndarray, q1: int, q2: int):
    """
    Apply a 4×4 unitary to qubits (q1, q2) via tensor contraction.
    """
    if q1 == q2:
        raise ValueError("Control and target qubits must differ")
    if not (0 <= q1 < self.n and 0 <= q2 < self.n):
        raise ValueError(f"Qubit indices ({q1},{q2}) out of range [0, {self.n})")
    gate = self.xp.array(gate, dtype=self.dtype)
    sv_nd      = self.sv.reshape([2] * self.n)
    sv_moved   = self.xp.moveaxis(sv_nd, (q1, q2), (-2, -1))
    flat_shape = (self.dim >> 2, 4)
    result     = self.xp.dot(sv_moved.reshape(flat_shape),
                             gate.reshape(4, 4).T)
    self.sv    = self.xp.moveaxis(
        result.reshape([2] * self.n), (-2, -1), (q1, q2)).ravel()

apply_cx

apply_cx(ctrl: int, tgt: int)

CX (CNOT) gate.

JAX path: matrix contraction via apply_gate_2q. NumPy path: fully vectorised index swap — no Python loops.

Source code in dense_evolution/simulator.py

def apply_cx(self, ctrl: int, tgt: int):
    """
    CX (CNOT) gate.

    JAX path: matrix contraction via apply_gate_2q.
    NumPy path: fully vectorised index swap — no Python loops.
    """
    if ctrl == tgt:
        raise ValueError("Control and target qubits must differ")
    if not (0 <= ctrl < self.n and 0 <= tgt < self.n):
        raise ValueError(f"Qubit indices ({ctrl},{tgt}) out of range [0, {self.n})")
    if HAS_JAX:
        cx_mat = jnp.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [0, 0, 0, 1],
            [0, 0, 1, 0],
        ], dtype=self.dtype)
        self.apply_gate_2q(cx_mat, ctrl, tgt)
    else:
        self.sv = _cx_numpy(np.array(self.sv), self.n, ctrl, tgt)

apply_cz

apply_cz(ctrl: int, tgt: int)

CZ gate.

JAX path: matrix contraction via apply_gate_2q. NumPy path: fully vectorised sign flip — no Python loops.

Source code in dense_evolution/simulator.py

def apply_cz(self, ctrl: int, tgt: int):
    """
    CZ gate.

    JAX path: matrix contraction via apply_gate_2q.
    NumPy path: fully vectorised sign flip — no Python loops.
    """
    if ctrl == tgt:
        raise ValueError("Control and target qubits must differ")
    if not (0 <= ctrl < self.n and 0 <= tgt < self.n):
        raise ValueError(f"Qubit indices ({ctrl},{tgt}) out of range [0, {self.n})")
    if HAS_JAX:
        cz_mat = jnp.array([
            [1, 0, 0,  0],
            [0, 1, 0,  0],
            [0, 0, 1,  0],
            [0, 0, 0, -1],
        ], dtype=self.dtype)
        self.apply_gate_2q(cz_mat, ctrl, tgt)
    else:
        self.sv = _cz_numpy(np.array(self.sv), self.n, ctrl, tgt)

apply_rx

apply_rx(qubit: int, theta: float)

Apply a parameterized RX gate using the active backend (NumPy/JAX).

Source code in dense_evolution/simulator.py

def apply_rx(self, qubit: int, theta: float):
    """Apply a parameterized RX gate using the active backend (NumPy/JAX)."""
    cos, sin = self.xp.cos(theta / 2), self.xp.sin(theta / 2)
    mat = self.xp.array([[cos, -1j * sin], [-1j * sin, cos]], dtype=self.dtype)
    self.apply_gate_1q(mat, qubit)

apply_ry

apply_ry(qubit: int, theta: float)

Apply a parameterized RY gate using the active backend (NumPy/JAX).

Source code in dense_evolution/simulator.py

def apply_ry(self, qubit: int, theta: float):
    """Apply a parameterized RY gate using the active backend (NumPy/JAX)."""
    cos, sin = self.xp.cos(theta / 2), self.xp.sin(theta / 2)
    mat = self.xp.array([[cos, -sin], [sin, cos]], dtype=self.dtype)
    self.apply_gate_1q(mat, qubit)

apply_rz

apply_rz(qubit: int, theta: float)

Apply a parameterized RZ gate using the active backend (NumPy/JAX).

Source code in dense_evolution/simulator.py

def apply_rz(self, qubit: int, theta: float):
    """Apply a parameterized RZ gate using the active backend (NumPy/JAX)."""
    exp_neg = self.xp.exp(-1j * theta / 2)
    exp_pos = self.xp.exp(1j * theta / 2)
    mat = self.xp.array([[exp_neg, 0.0], [0.0, exp_pos]], dtype=self.dtype)
    self.apply_gate_1q(mat, qubit)

measure

measure(qubit_idx: int) -> int

Projective measurement on qubit_idx.

Returns 0 or 1 and collapses the statevector. Uses MSB-first physical bit index: phys = n - 1 - qubit_idx.

BUG FIX (original): the original NumPy collapse wrote sv_reshaped[:, 1 if result == 0 else 0, :] = 0.0 which zeroed the wrong basis state (0 when result=1, 1 when result=0) and never normalised the JAX path.

Source code in dense_evolution/simulator.py

def measure(self, qubit_idx: int) -> int:
    """
    Projective measurement on *qubit_idx*.

    Returns 0 or 1 and collapses the statevector.
    Uses MSB-first physical bit index: phys = n - 1 - qubit_idx.

    BUG FIX (original): the original NumPy collapse wrote
        sv_reshaped[:, 1 if result == 0 else 0, :] = 0.0
    which zeroed the *wrong* basis state (0 when result=1, 1 when result=0)
    and never normalised the JAX path.
    """
    if not 0 <= qubit_idx < self.n:
        raise ValueError(
            f"Qubit {qubit_idx} out of range [0, {self.n})")

    phys   = self.n - 1 - qubit_idx
    stride = 1 << phys

    # ── compute marginal probabilities ──────────────────────────
    if HAS_JAX:
        probs   = jnp.abs(self.sv) ** 2
        sv_nd   = probs.reshape([2] * self.n)
        mv      = jnp.moveaxis(sv_nd, phys, 0)
        prob_0  = float(jnp.sum(mv[0]))
        prob_1  = float(jnp.sum(mv[1]))
    else:
        sv_res  = self.sv.reshape(-1, 2, stride)
        prob_0  = float(np.sum(np.abs(sv_res[:, 0, :]) ** 2))
        prob_1  = float(np.sum(np.abs(sv_res[:, 1, :]) ** 2))

    total = prob_0 + prob_1
    if total < 1e-12:
        raise RuntimeError("Statevector norm is zero — cannot measure")
    prob_0 /= total
    prob_1 /= total

    result = int(np.random.choice([0, 1], p=[prob_0, prob_1]))

    # ── collapse ────────────────────────────────────────────────
    # Zero out the amplitudes corresponding to the *opposite* outcome.
    zero_slot = 1 - result     # if result=0, zero slot 1; if result=1, zero slot 0

    if HAS_JAX:
        sv_nd  = self.sv.reshape([2] * self.n)
        mv     = jnp.moveaxis(sv_nd, phys, 0)
        mv     = mv.at[zero_slot].set(0.0 + 0j)
        self.sv = jnp.moveaxis(mv, 0, phys).ravel()
    else:
        sv_res = self.sv.reshape(-1, 2, stride)
        sv_res[:, zero_slot, :] = 0.0
        self.sv = sv_res.ravel()

    self.normalize()
    return result

run_circuit_with_chunking

run_circuit_with_chunking(
    circuit: List, chunk_size: int = 500
)

Execute a circuit in chunks to avoid JIT recompilation on large variable-length circuits.

Each chunk is a separate _compile_and_run_circuit_jit call with a fixed-size ops array, allowing XLA to cache each size.

Source code in dense_evolution/simulator.py

def run_circuit_with_chunking(self, circuit: List, chunk_size: int = 500):
    """
    Execute a circuit in chunks to avoid JIT recompilation on
    large variable-length circuits.

    Each chunk is a separate _compile_and_run_circuit_jit call
    with a fixed-size ops array, allowing XLA to cache each size.
    """
    target = QuantumTranspiler.transpile(circuit)
    for i in range(0, len(target), chunk_size):
        self.run_circuit_jit_beast_mode(target[i: i + chunk_size])

get_probabilities

get_probabilities() -> np.ndarray

Return measurement probability distribution as a NumPy float64 array.

Source code in dense_evolution/simulator.py

def get_probabilities(self) -> np.ndarray:
    """Return measurement probability distribution as a NumPy float64 array."""
    probs = np.array(self.xp.abs(self.sv) ** 2, dtype=np.float64)
    # guard against floating-point leakage outside [0, 1]
    probs = np.clip(probs, 0.0, 1.0)
    total = probs.sum()
    if total > 1e-12:
        probs /= total
    return probs

get_statevector

get_statevector() -> np.ndarray

Return the current statevector as a NumPy complex array.

Source code in dense_evolution/simulator.py

def get_statevector(self) -> np.ndarray:
    """Return the current statevector as a NumPy complex array."""
    return np.array(self.sv, dtype=self.dtype)

memory_mb

memory_mb() -> float

Statevector memory footprint in megabytes.

Source code in dense_evolution/simulator.py

def memory_mb(self) -> float:
    """Statevector memory footprint in megabytes."""
    bytes_per_element = 8 if self.use_float32 else 16   # complex64=8, complex128=16
    return self.dim * bytes_per_element / 1_000_000