Premchan369/Q-TensorFormer

1"""2Quantum Feature Encoding Layers.3 4PennyLane-based quantum circuits wrapped as PyTorch nn.Module layers.5 6Components:7  - QuantumAngleEmbedding: Classical data → rotation angles on qubits8  - QuantumAmplitudeEmbedding: Encodes data as quantum amplitudes9  - EntanglementMonitor: Estimates entanglement via attention patterns10  - ClassicalQuantumFallback: MLP-based fallback when PennyLane unavailable11"""12 13import torch14import torch.nn as nn15import torch.nn.functional as F16import math17from typing import Optional, Tuple, List18 19try:20    import pennylane as qml21    HAS_PENNYLANE = True22except ImportError:23    HAS_PENNYLANE = False24 25 26class QuantumAngleEmbedding(nn.Module):27    """28    Encodes classical features into quantum states via angle encoding.29 30    Circuit: RX(input) → [RY(θ) → CNOT ladder] × n_layers → ⟨Z_i⟩31 32    Parameters33    ----------34    n_qubits : int35        Number of qubits (4-8 for NISQ compatibility).36    n_layers : int37        Number of variational circuit layers.38    n_outputs : int or None39        Number of expectation values to measure. Default: n_qubits.40    diff_method : str41        Differentiation method. 'backprop' for batched inputs,42        'parameter-shift' for hardware compatibility.43    """44 45    def __init__(self, n_qubits: int = 4, n_layers: int = 2,46                 n_outputs: int = None, diff_method: str = "backprop"):47        super().__init__()48        if not HAS_PENNYLANE:49            raise ImportError(50                "PennyLane is required for quantum layers. "51                "Install with: pip install pennylane"52            )53 54        self.n_qubits = n_qubits55        self.n_layers = n_layers56        self.n_outputs = n_outputs or n_qubits57 58        dev = qml.device("default.qubit", wires=n_qubits)59 60        @qml.qnode(dev, interface="torch", diff_method=diff_method)61        def circuit(inputs, weights):62            # Angle encoding63            for i in range(n_qubits):64                qml.RX(inputs[..., i], wires=i)65 66            # Variational layers with entanglement67            for layer in range(n_layers):68                for i in range(n_qubits):69                    qml.RY(weights[layer, i], wires=i)70                # Nearest-neighbor CNOT ladder71                for i in range(n_qubits - 1):72                    qml.CNOT(wires=[i, i + 1])73                # Cyclic entanglement for >2 qubits74                if n_qubits > 2:75                    qml.CNOT(wires=[n_qubits - 1, 0])76 77            # Measure PauliZ expectation values78            return [qml.expval(qml.PauliZ(i)) for i in range(self.n_outputs)]79 80        weight_shapes = {"weights": (n_layers, n_qubits)}81        self.qlayer = qml.qnn.TorchLayer(circuit, weight_shapes)82 83    def forward(self, x: torch.Tensor) -> torch.Tensor:84        """85        Args:86            x: (*batch, n_qubits) — classical inputs mapped to rotation angles87        Returns:88            (*batch, n_outputs) — PauliZ expectation values in [-1, 1]89        """90        return self.qlayer(x)91 92 93class EntanglementMonitor(nn.Module):94    """95    Estimates entanglement entropy from attention patterns.96 97    Uses attention distribution entropy as a classical proxy98    for quantum entanglement entropy. Avoids expensive quantum99    state tomography during training.100 101    Parameters102    ----------103    n_qubits : int104        Number of qubits in the simulated quantum system.105    subsystem_a : list of ints or None106        Qubit indices for subsystem A (bipartition).107    """108 109    def __init__(self, n_qubits: int = 4,110                 subsystem_a: Optional[List[int]] = None):111        super().__init__()112        self.n_qubits = n_qubits113        if subsystem_a is None:114            subsystem_a = list(range(n_qubits // 2))115        self.subsystem_a = subsystem_a116 117    def forward(self, attention_weights: torch.Tensor) -> torch.Tensor:118        """119        Estimate entanglement from attention distributions.120 121        Args:122            attention_weights: (batch, heads, seq_len, seq_len)123                Softmax-normalized attention weights.124 125        Returns:126            (batch, heads) — estimated entanglement entropy per head127        """128        eps = 1e-8129        entropy = -torch.sum(130            attention_weights * torch.log(attention_weights + eps),131            dim=-1132        )  # (batch, heads, seq_len)133        return entropy.mean(dim=-1)  # (batch, heads)134 135 136class ClassicalQuantumFallback(nn.Module):137    """138    Classical MLP fallback when PennyLane is unavailable.139 140    Uses sinusoidal activations to mimic quantum rotation gate behavior.141    """142 143    def __init__(self, n_qubits: int = 4, n_layers: int = 2,144                 n_outputs: int = None):145        super().__init__()146        n_outputs = n_outputs or n_qubits147        layers = []148        in_dim = n_qubits149        for _ in range(n_layers):150            layers.extend([151                nn.Linear(in_dim, n_qubits * 2),152                nn.SiLU(),  # Smooth activation like quantum gates153            ])154            in_dim = n_qubits * 2155        layers.append(nn.Linear(in_dim, n_outputs))156        layers.append(nn.Tanh())  # Bound output to [-1, 1] like expectation values157        self.net = nn.Sequential(*layers)158 159    def forward(self, x: torch.Tensor) -> torch.Tensor:160        return self.net(x)161 162 163def create_quantum_embedding(input_dim: int, n_qubits: int = 4,164                              n_layers: int = 2, output_dim: int = None,165                              embedding_type: str = "angle") -> nn.Module:166    """167    Factory for quantum embedding layers.168 169    Args:170        input_dim: Input feature dimension.171        n_qubits: Number of qubits.172        n_layers: Circuit depth.173        output_dim: Output dimension.174        embedding_type: 'angle' or 'amplitude'.175 176    Returns:177        Quantum embedding nn.Module (or classical fallback if no PennyLane).178    """179    output_dim = output_dim or n_qubits180 181    if not HAS_PENNYLANE:182        print("[WARN] PennyLane not installed. Using classical fallback.")183        return nn.Sequential(184            nn.Linear(input_dim, n_qubits),185            ClassicalQuantumFallback(n_qubits, n_layers, output_dim),186            nn.Linear(output_dim, output_dim),187        )188 189    if embedding_type == "angle":190        return nn.Sequential(191            nn.Linear(input_dim, n_qubits),192            QuantumAngleEmbedding(n_qubits, n_layers, output_dim),193        )194    elif embedding_type == "amplitude":195        return nn.Sequential(196            nn.Linear(input_dim, 2 ** n_qubits),197            nn.Softmax(dim=-1),198            # Amplitude embedding would go here199            nn.Linear(2 ** n_qubits, output_dim),200        )201    else:202        raise ValueError(f"Unknown embedding type: {embedding_type}")203