开源代码：对老潘论文复现，格式自己调吧

import torch import torch.nn as nn import torch.optim as optim import numpy as np import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D import plotly.graph_objects as go from plotly.subplots import make_subplots from tqdm import tqdm import warnings warnings.filterwarnings('ignore') # ----------------------------------------------------------------------------- # 1. 论文核心物理参数定义（严格复刻实验数据） # ----------------------------------------------------------------------------- class QuantumExperimentParams: def __init__(self): # 原子与光子基本参数 self.atom_mass = 85.4678e-3 / 6.022e23 # 87Rb原子质量 (kg) self.hbar = 1.0546e-34 # 约化普朗克常数 (J·s) self.photon_wavelength = 780e-9 # 散射光子波长 (m) self.photon_momentum = self.hbar * 2 * np.pi / self.photon_wavelength # ħk (kg·m/s) # 光学镊子参数 self.trap_depths = np.linspace(0.60e-3, 10.49e-3, 20) # 陷阱深度范围 (K) self.axial_trap_freq = 40e3 # 轴向陷阱频率 (Hz) self.radial_trap_freq = 300e3 # 径向陷阱频率 (Hz) # 量子态参数 self.residual_phonon_num = np.linspace(0.08, 0.37, 20) # 残余声子数 n̄ self.eta_base = self.photon_momentum / (2 * self.hbar * np.sqrt(np.pi * self.axial_trap_freq / self.atom_mass)) # 实验几何参数（单位：μm，便于可视化） self.atom_position = np.array([0.0, 0.0, 0.0]) # 原子初始位置 (μm) self.laser_beam_radius = 0.01 # 激光束半径 (μm) self.tweezer_length = 2.0 # 光学镊子长度 (μm) # ----------------------------------------------------------------------------- # 2. 2025前沿深度学习模型：量子增强NeRF (Q-NeRF) + 物理先验Transformer # ----------------------------------------------------------------------------- class PhysicConstraintTransformer(nn.Module): """物理约束Transformer：嵌入海森堡原理和干涉可见度公式""" def __init__(self, d_model=128, nhead=8): super().__init__() self.transformer = nn.Transformer( d_model=d_model, nhead=nhead, num_encoder_layers=3, num_decoder_layers=3, dim_feedforward=512, batch_first=True ) self.uncertainty_embedding = nn.Linear(2, d_model) # 位置/动量不确定性嵌入 self.visibility_head = nn.Sequential( nn.Linear(d_model, 64), nn.ReLU(), nn.Linear(64, 1), nn.Sigmoid() # 可见度输出范围 [0,1] ) self.hbar_embed = nn.Parameter(torch.tensor([1.0546e-34], dtype=torch.float32)) # 物理常数嵌入 self.photon_momentum = None # 后续绑定论文中的光子动量 def forward(self, pos_uncertainty, mom_uncertainty, n_bar): # 海森堡约束：Δx·Δp ≥ ħ/2（硬约束损失） heisenberg_constraint = pos_uncertainty * mom_uncertainty - self.hbar_embed / 2 heisenberg_loss = torch.relu(-heisenberg_constraint).mean() # 特征嵌入与Transformer编码（扩展维度以匹配Transformer输入要求） pos_uncertainty = pos_uncertainty.unsqueeze(1) # (1,) → (1,1) mom_uncertainty = mom_uncertainty.unsqueeze(1) # (1,) → (1,1) x = torch.cat([pos_uncertainty, mom_uncertainty], dim=-1) # (1,2) x = self.uncertainty_embedding(x) # (1, d_model) x = self.transformer(x, x) # (1, d_model) # 可见度预测（融合论文公式先验） V_pred = self.visibility_head(x).squeeze(-1) # (1,) eta = self.photon_momentum / (2 * mom_uncertainty) # (1,1) eta_eff = eta * torch.sqrt(2 * n_bar + 1) # (1,) V_physics = torch.exp(-2 * eta_eff ** 2) # (1,) # 物理一致性损失 physics_consistency_loss = torch.mean((V_pred - V_physics) ** 2) return V_pred, heisenberg_loss + physics_consistency_loss class QNeRF(nn.Module): """量子增强NeRF：渲染原子三维动量波函数（复数形式）""" def __init__(self, hidden_dim=256): super().__init__() self.mlp = nn.Sequential( nn.Linear(3 + 1, hidden_dim), # 3D位置 + 陷阱深度 → 输入特征 (N,4) nn.LayerNorm(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 2) # 输出：波函数实部 + 虚部 (N,2) ) def forward(self, xyz, trap_depth): """ 输入： xyz: 三维空间坐标 (N, 3)，N为网格点数 trap_depth: 陷阱深度 (1,) → 需扩展为 (N,1) 输出：复数波函数模 (N,) """ n_points = xyz.shape[0] # 获取网格点数N # 关键修复：将trap_depth从(1,)扩展为(N,1)，匹配xyz维度 trap_depth_expanded = trap_depth.repeat(n_points, 1) # (1,) → (N,1) input_feat = torch.cat([xyz, trap_depth_expanded], dim=-1) # (N,3)+(N,1)=(N,4) psi = self.mlp(input_feat) # (N, 2)：实部、虚部 psi_magnitude = torch.sqrt(psi[:, 0] ** 2 + psi[:, 1] ** 2) # 波函数模（概率密度） return psi, psi_magnitude # ----------------------------------------------------------------------------- # 3. 数据集生成：基于论文物理模型合成训练数据 # ----------------------------------------------------------------------------- def generate_quantum_dataset(params, num_samples=5000): dataset = [] for trap_depth in params.trap_depths: for n_bar in params.residual_phonon_num: # 计算动量不确定性Δp（论文：Δp ∝ (trap_depth)^(1/4)） mom_uncertainty = 0.78e-27 + (1.60e-27 - 0.78e-27) * (trap_depth / 10.49e-3) ** 0.25 # 海森堡原理：Δx = ħ/(2Δp)（转换为μm便于可视化） pos_uncertainty = (params.hbar / (2 * mom_uncertainty)) * 1e6 # 理论可见度 eta = params.photon_momentum / (2 * mom_uncertainty) eta_eff = eta * np.sqrt(2 * n_bar + 1) visibility = np.exp(-2 * eta_eff ** 2) # 生成原子周围三维网格点（μm单位） x = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 30) y = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 30) z = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 30) xyz = np.array(np.meshgrid(x, y, z)).reshape(3, -1).T # (27000, 3) dataset.append({ "xyz": xyz, "trap_depth": trap_depth, "mom_uncertainty": mom_uncertainty, "pos_uncertainty": pos_uncertainty, "visibility": visibility, "n_bar": n_bar }) return dataset # ----------------------------------------------------------------------------- # 4. 模型训练：Q-NeRF + 物理约束Transformer联合训练 # ----------------------------------------------------------------------------- def train_models(params, dataset, epochs=50): # 自动适配CPU/GPU device = "cuda" if torch.cuda.is_available() else "cpu" qnerf = QNeRF().to(device) phys_transformer = PhysicConstraintTransformer().to(device) # 绑定论文中的光子动量到Transformer（确保维度正确） phys_transformer.photon_momentum = torch.tensor([params.photon_momentum], dtype=torch.float32).to(device) optimizer = optim.AdamW( list(qnerf.parameters()) + list(phys_transformer.parameters()), lr=1e-4, weight_decay=1e-5 ) criterion = nn.MSELoss() for epoch in tqdm(range(epochs), desc="Training Models"): total_loss = 0.0 for batch in dataset: # 数据预处理（确保所有张量维度正确） xyz = torch.tensor(batch["xyz"], dtype=torch.float32).to(device) # (27000, 3) trap_depth = torch.tensor([batch["trap_depth"]], dtype=torch.float32).to(device) # (1,) mom_uncertainty = torch.tensor([batch["mom_uncertainty"]], dtype=torch.float32).to(device) # (1,) pos_uncertainty = torch.tensor([batch["pos_uncertainty"]], dtype=torch.float32).to(device) # (1,) visibility = torch.tensor([batch["visibility"]], dtype=torch.float32).to(device) # (1,) n_bar = torch.tensor([batch["n_bar"]], dtype=torch.float32).to(device) # (1,) # Q-NeRF预测波函数（修复后维度匹配） psi_pred, psi_magnitude_pred = qnerf(xyz, trap_depth) # 物理约束Transformer预测可见度 V_pred, physics_loss = phys_transformer(pos_uncertainty, mom_uncertainty, n_bar) # 波函数损失（高斯波包理论，论文中原子基态为高斯波函数） psi_theory = torch.exp(-torch.norm(xyz, dim=1) ** 2 / (2 * pos_uncertainty ** 2)).to(device) # (27000,) wavefunction_loss = criterion(psi_magnitude_pred, psi_theory) # 总损失：波函数损失 + 可见度损失 + 物理约束损失 total_batch_loss = wavefunction_loss + criterion(V_pred, visibility) + physics_loss total_loss += total_batch_loss.item() # 反向传播与参数更新 optimizer.zero_grad() total_batch_loss.backward() optimizer.step() # 每10个epoch打印训练日志 if (epoch + 1) % 10 == 0: avg_loss = total_loss / len(dataset) print(f"Epoch [{epoch+1}/{epochs}], Avg Loss: {avg_loss:.6f}") # 保存训练好的模型 torch.save(qnerf.state_dict(), "qnerf_2025.pth") torch.save(phys_transformer.state_dict(), "phys_transformer_2025.pth") return qnerf, phys_transformer # ----------------------------------------------------------------------------- # 5. 三维可视化：静态图(Matplotlib) + 交互式图(Plotly) # ----------------------------------------------------------------------------- def visualize_3d_experiment(params, qnerf, trap_depth=10.49e-3): """生成两种三维效果图： 1. Matplotlib静态图：原子波函数+实验装置+干涉条纹（保存为PNG） 2. Plotly交互式图：支持旋转/缩放（浏览器打开） """ device = "cuda" if torch.cuda.is_available() else "cpu" qnerf.eval() # 模型切换到评估模式 with torch.no_grad(): # 禁用梯度计算，加快推理 # 计算关键物理参数（μm单位，便于可视化） mom_uncertainty = 0.78e-27 + (1.60e-27 - 0.78e-27) * (trap_depth / 10.49e-3) ** 0.25 pos_uncertainty = (params.hbar / (2 * mom_uncertainty)) * 1e6 eta = params.photon_momentum / (2 * mom_uncertainty) visibility = np.exp(-2 * eta ** 2) # 生成三维网格并预测原子动量波函数 x = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 50) y = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 50) z = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 50) x_grid, y_grid, z_grid = np.meshgrid(x, y, z) xyz = np.array([x_grid.ravel(), y_grid.ravel(), z_grid.ravel()]).T # (125000, 3) xyz_tensor = torch.tensor(xyz, dtype=torch.float32).to(device) trap_depth_tensor = torch.tensor([trap_depth], dtype=torch.float32).to(device) # (1,) _, psi_magnitude = qnerf(xyz_tensor, trap_depth_tensor) # 调用修复后的forward psi_magnitude_grid = psi_magnitude.cpu().numpy().reshape(x_grid.shape) # 生成单光子干涉条纹（z=1μm平面，论文中干涉条纹为正弦分布） z_plane = 1.0 # 干涉平面z坐标 (μm) x_interf = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 200) y_interf = np.linspace(-3 * pos_uncertainty, 3 * pos_uncertainty, 200) x_interf_grid, y_interf_grid = np.meshgrid(x_interf, y_interf) # 论文干涉强度公式：I = 0.5*(1 + V*cos(2πx/(λ/2)))，λ为光子波长 interf_intensity = 0.5 * (1 + visibility * np.cos(2 * np.pi * x_interf_grid / (params.photon_wavelength * 1e6 / 2))) # -------------------------- # 1. Matplotlib静态三维图（保存为PNG，便于分享） # -------------------------- fig = plt.figure(figsize=(15, 5)) # 子图1：原子三维动量波函数（高斯等值面，论文核心量子态） ax1 = fig.add_subplot(131, projection='3d') # 取中间5个等值面，清晰展示波函数分布 levels = np.linspace(psi_magnitude_grid.min(), psi_magnitude_grid.max(), 6)[1:-1] for level in levels: ax1.contour3D(x_grid, y_grid, z_grid, psi_magnitude_grid, levels=[level], cmap='viridis', alpha=0.6) ax1.set_title(f'Atom Momentum Wavefunction\n(Trap Depth: {trap_depth*1e3:.2f} mK)', fontsize=10) ax1.set_xlabel('X (μm)') ax1.set_ylabel('Y (μm)') ax1.set_zlabel('Z (μm)') # 子图2：实验装置（光学镊子+87Rb原子，复刻论文图2a） ax2 = fig.add_subplot(132, projection='3d') # 光学镊子（十字激光束，论文中852nm激光形成陷阱） z_tweezer = np.linspace(-params.tweezer_length/2, params.tweezer_length/2, 100) ax2.plot(np.zeros_like(z_tweezer), np.zeros_like(z_tweezer), z_tweezer, 'b-', linewidth=5, label='Tweezer Laser (Z-axis)') x_tweezer = np.linspace(-params.tweezer_length/2, params.tweezer_length/2, 100) ax2.plot(x_tweezer, np.zeros_like(x_tweezer), np.zeros_like(x_tweezer), 'b-', linewidth=5, label='Tweezer Laser (X-axis)') # 87Rb原子（红色球，位于陷阱中心） ax2.scatter([0], [0], [0], color='red', s=200, label='87Rb Atom') ax2.set_title('Experimental Setup', fontsize=10) ax2.set_xlabel('X (μm)') ax2.set_ylabel('Y (μm)') ax2.set_zlabel('Z (μm)') ax2.legend() # 子图3：单光子干涉条纹（复刻论文图3b） ax3 = fig.add_subplot(133, projection='3d') surf = ax3.plot_surface(x_interf_grid, y_interf_grid, np.full_like(x_interf_grid, z_plane), facecolors=plt.cm.hot(interf_intensity), alpha=0.8, shade=False) ax3.set_title(f'Single-Photon Interference Fringes\n(Visibility: {visibility:.2%})', fontsize=10) ax3.set_xlabel('X (μm)') ax3.set_ylabel('Y (μm)') ax3.set_zlabel('Z (μm)') # 添加颜色条表示干涉强度 cbar = fig.colorbar(surf, ax=ax3, shrink=0.5, aspect=10) cbar.set_label('Intensity', fontsize=8) plt.tight_layout() plt.savefig('quantum_experiment_3d_static.png', dpi=300, bbox_inches='tight') plt.show() # -------------------------- # 2. Plotly交互式三维图（浏览器打开，支持旋转/缩放） # -------------------------- fig = make_subplots(rows=1, cols=2, specs=[[{'type': 'scene'}, {'type': 'scene'}]]) # 子图1：原子动量波函数（散点图+颜色编码概率密度） # 采样5000个点避免图面拥挤 sample_idx = np.random.choice(xyz.shape[0], 5000, replace=False) x_sample = xyz[sample_idx, 0] y_sample = xyz[sample_idx, 1] z_sample = xyz[sample_idx, 2] psi_sample = psi_magnitude.cpu().numpy()[sample_idx] fig.add_trace( go.Scatter3d( x=x_sample, y=y_sample, z=z_sample, mode='markers', marker=dict(size=2, color=psi_sample, colorscale='viridis', opacity=0.7), name='Wavefunction Probability' ), row=1, col=1 ) # 子图2：干涉条纹+原子位置 fig.add_trace( go.Surface( x=x_interf_grid, y=y_interf_grid, z=np.full_like(x_interf_grid, z_plane), surfacecolor=interf_intensity, colorscale='hot', opacity=0.8, name='Interference Fringes' ), row=1, col=2 ) fig.add_trace( go.Scatter3d( x=[0], y=[0], z=[0], mode='markers', marker=dict(size=10, color='red'), name='87Rb Atom' ), row=1, col=2 ) # 布局优化 fig.update_layout( title=f'Einstein-Bohr Gedankenexperiment (Trap Depth: {trap_depth*1e3:.2f} mK)', width=1200, height=600, scene1=dict(title='Atom Momentum Wavefunction', xaxis_title='X (μm)', yaxis_title='Y (μm)', zaxis_title='Z (μm)'), scene2=dict(title='Interference Fringes', xaxis_title='X (μm)', yaxis_title='Y (μm)', zaxis_title='Z (μm)') ) # 保存为HTML文件 fig.write_html('quantum_experiment_3d_interactive.html') print("✅ 交互式三维图已保存为 'quantum_experiment_3d_interactive.html'，请用浏览器打开") # ----------------------------------------------------------------------------- # 6. 主函数：端到端运行（一键训练+可视化） # ----------------------------------------------------------------------------- if __name__ == "__main__": # 初始化实验参数（严格复刻论文数据） params = QuantumExperimentParams() # 生成训练数据集（5000个样本，平衡训练速度与精度） print("📊 Generating quantum dataset...") dataset = generate_quantum_dataset(params, num_samples=5000) # 训练深度学习模型（50个epoch，CPU/GPU自动适配） print(":rocket: Training Q-NeRF + Physics-Constrained Transformer...") qnerf, phys_transformer = train_models(params, dataset, epochs=50) # 生成三维可视化结果（最大陷阱深度10.49 mK，对应最高干涉可见度） print("🎨 Generating 3D visualizations...") visualize_3d_experiment(params, qnerf, trap_depth=10.49e-3)