GradTDDFT：基于JAX的可微分含时密度泛函理论研究工具包

GradTDDFT是一个基于JAX的量子化学计算框架，专注于可微分的含时密度泛函理论(TDDFT)计算，支持神经网络交换相关泛函训练、激发态计算和GPU加速，为机器学习与量子化学的交叉研究提供了强大工具。

• 原作者/维护者：STOKES-DOT
• 来源平台：github
• 原始标题：GradTDDFT: JAX research toolkit for differentiable TDDFT
• 原始链接：https://github.com/STOKES-DOT/GradTDDFT
• 来源发布时间/更新时间：2026-06-01T02:42:02Z

原作者与来源

• 原作者/维护者：STOKES-DOT
• 来源平台：github
• 原始标题：GradTDDFT: JAX research toolkit for differentiable TDDFT
• 原始链接：https://github.com/STOKES-DOT/GradTDDFT
• 来源发布时间/更新时间：2026-06-01T02:42:02Z 原作者与来源\n\n- 原作者/维护者： STOKES-DOT\n- 来源平台： GitHub\n- 原始标题： GradTDDFT: JAX research toolkit for differentiable TDDFT\n- 原始链接： https://github.com/STOKES-DOT/GradTDDFT\n- 发布时间： 2026年6月\n\n---\n\n背景：量子化学与机器学习的交汇点\n\n密度泛函理论（DFT）是现代量子化学计算的基石，它通过电子密度而非复杂的波函数来描述多电子系统，在计算效率和精度之间取得了出色的平衡。含时密度泛函理论（TDDFT）则进一步扩展了这一框架，使其能够处理分子的激发态性质，在光化学、材料科学和药物设计等领域有着广泛应用。\n\n然而，传统DFT/TDDFT面临一个根本性挑战：交换相关（XC）泛函的选择。虽然Kohn-Sham方程在理论上是精确的，但实际计算中必须使用近似的XC泛函。数十年来，化学家开发了LDA、GGA、杂化泛函等一系列近似方法，但每种方法都有其适用范围和局限性。\n\n近年来，机器学习为解决这一难题提供了新思路：能否让神经网络学习一个更精确的XC泛函？这需要一个能够自动微分的计算框架，以便通过梯度下降优化神经网络的参数。GradTDDFT正是为此而生。\n\n---\n\n核心功能：可微分量子化学计算\n\nGradTDDFT是一个基于JAX的研究级代码库，提供了端到端的可微分DFT/TDDFT计算能力。JAX作为Google开发的机器学习框架，不仅支持GPU加速，更重要的是提供了强大的自动微分功能，这正是训练神经XC泛函的关键。\n\n该工具包的主要功能模块包括：\n\n1. 基态自洽场（SCF）计算\n\nGradTDDFT实现了严格基于JAX的SCF求解器，支持限制性和非限制性计算路径。它提供了类似PySCF的API接口，包括scf.RKS、scf.UKS等 facade，让熟悉PySCF的用户能够快速上手。\n\n2. 含时密度泛函理论（TDDFT）\n\n框架完整实现了TDA（Tamm-Dancoff近似）和Casida形式的TDDFT响应求解器，支持激发能计算和光谱模拟。这些求解器同样是完全可微分的，允许对激发态性质进行端到端的梯度优化。\n\n3. 神经交换相关泛函（Neural XC）\n\n这是GradTDDFT最具创新性的功能。用户可以定义由神经网络参数化的XC泛函，形式为：\n\n\ne_xc^NN(r) = Σ_k c_k^θ(r) · e_k^semilocal(r) + c_pt2^θ(r) · e_pt2(r) + c_hf^θ(r) · e_hf(r)\n\n\n其中半局域基组可以包含LDA、GGA甚至MGGA组件，神经网络的输出是空间变化的组合系数。这种设计既保留了物理直觉（基于传统泛函的线性组合），又赋予了模型足够的灵活性来学习复杂的电子相关效应。\n\n4. GPU加速与大规模计算\n\n通过与GPU4PySCF的集成，GradTDDFT能够利用GPU的并行计算能力加速SCF和TDDFT计算，这对于训练数据生成和大型分子系统的研究尤为重要。\n\n---\n\n技术架构：模块化设计\n\nGradTDDFT采用清晰的三层架构分离不同关注点：\n\n分子与计算基础设施层（gto、scf、dft、tddft模块）：负责分子结构表示、积分计算、网格生成、SCF迭代和TDDFT响应矩阵构建。这些模块提供了类似PySCF的API，降低了学习成本。\n\n传统XC后端层（xc_backend模块）：处理传统密度泛函的解析和计算。GradTDDFT通过jax_xc库支持广泛的LDA、GGA和MGGA泛函，包括PBE、B3LYP、HSE等常用泛函。\n\n神经XC构建层（neural_xc模块）：提供神经网络XC泛函的构建、训练和响应计算功能。这是整个框架的核心创新，支持多种架构（如残差网络、简单MLP）和训练策略。\n\n---\n\n支持的泛函与组件\n\nGradTDDFT提供了丰富的XC组件支持：\n\n基础组件：lda_x、lda_c_pw、lda_c_vwn、lda_c_vwn_rpa等LDA相关泛函；gga_x_b88、gga_x_pbe、gga_c_lyp、gga_c_pbe等GGA泛函。\n\n复合泛函：支持PySCF风格的复合字符串解析，如svwn、pbe、pbe0、b3lyp、bhandhlyp、hse03、hse06等。这些复合泛函会被解析为显式的本地子通道组合。\n\n元GGA（MGGA）：支持scan_x、scan_c、r2scan_x、tpss_c等元GGA泛函，提供更精确的动能密度依赖描述。\n\n用户还可以启用实验性组件，框架会动态发现已安装的jax_xc泛函，尽管这些组件可能需要额外的验证。\n\n---\n\n训练工作流与应用场景\n\nGradTDDFT提供了多种研究脚本和训练工作流：\n\n基态训练：针对分子基态能量和性质的神经网络XC泛函训练，可用于构建数据驱动的精确泛函。\n\n激发态训练：直接针对激发态能量的训练，让神经网络学习捕捉电子激发的关键物理。\n\nH2解离研究：作为量子化学的经典测试案例，氢分子解离曲线的计算检验了泛函对强相关体系的描述能力。\n\nQH9基准测试：在QM9数据集上的基准测试，评估神经XC泛函在多样化分子上的泛化性能。\n\n---\n\n安装与使用\n\nGradTDDFT的安装相对直接，推荐使用虚拟环境：\n\nbash\npython -m venv .venv\nsource .venv/bin/activate\npython -m pip install --upgrade pip\npython -m pip install -e \".[dev,upstreams]\"\n\n\nupstreams额外依赖会安装jax_xc和pyscf。对于GPU计算，需要额外安装CUDA-enabled JAX和GPU4PySCF环境。\n\n运行时检查可以验证安装：\n\npython\nimport jax\nimport td_graddft\nfrom td_graddft.xc_backend import jax_xc_backend_info\n\nprint(jax.devices()) 显示可用计算设备\nprint(jax_xc_backend_info()) 显示XC后端信息\n\n\n---\n\n实际意义与研究价值\n\nGradTDDFT代表了量子化学计算方法学的一个重要发展方向：将现代机器学习工具与传统第一性原理计算深度融合。通过自动微分，研究者可以系统性地优化XC泛函，甚至让数据本身"发现"新的密度泛函形式。\n\n对于材料科学和药物设计领域，更精确的激发态计算意味着更可靠的光谱预测和光化学性质评估。对于基础物理研究，神经XC泛函可能揭示传统近似中缺失的电子相关物理图像。\n\n作为一个活跃的研究原型，GradTDDFT的API设计兼顾了实验灵活性和实用性。虽然它可能不如PySCF等成熟代码库稳定，但其在可微分计算方面的独特能力使其成为探索机器学习-量子化学交叉前沿的有力工具。