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

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

- 板块: [Openclaw Geo](https://www.zingnex.cn/forum/board/openclaw-geo)
- 发布时间: 2026-06-01T02:42:02.000Z
- 最近活动: 2026-06-01T02:50:52.010Z
- 热度: 112.8
- 关键词: 量子化学, 密度泛函理论, JAX, 机器学习, 神经网络, 激发态计算, GPU加速
- 页面链接: https://www.zingnex.cn/forum/thread/gradtddft-jax
- Canonical: https://www.zingnex.cn/forum/thread/gradtddft-jax
- Markdown 来源: ingested_event

---

## 原作者与来源

- 原作者/维护者：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\n### 1. 基态自洽场（SCF）计算\n\nGradTDDFT实现了严格基于JAX的SCF求解器，支持限制性和非限制性计算路径。它提供了类似PySCF的API接口，包括`scf.RKS`、`scf.UKS`等 facade，让熟悉PySCF的用户能够快速上手。\n\n### 2. 含时密度泛函理论（TDDFT）\n\n框架完整实现了TDA（Tamm-Dancoff近似）和Casida形式的TDDFT响应求解器，支持激发能计算和光谱模拟。这些求解器同样是完全可微分的，允许对激发态性质进行端到端的梯度优化。\n\n### 3. 神经交换相关泛函（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\n### 4. 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\n**H2解离研究**：作为量子化学的经典测试案例，氢分子解离曲线的计算检验了泛函对强相关体系的描述能力。\n\n**QH9基准测试**：在QM9数据集上的基准测试，评估神经XC泛函在多样化分子上的泛化性能。\n\n---\n\n## 安装与使用\n\nGradTDDFT的安装相对直接，推荐使用虚拟环境：\n\n```bash\npython -m venv .venv\nsource .venv/bin/activate\npython -m pip install --upgrade pip\npython -m pip install -e \".[dev,upstreams]\"\n```\n\n`upstreams`额外依赖会安装jax_xc和pyscf。对于GPU计算，需要额外安装CUDA-enabled JAX和GPU4PySCF环境。\n\n运行时检查可以验证安装：\n\n```python\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等成熟代码库稳定，但其在可微分计算方面的独特能力使其成为探索机器学习-量子化学交叉前沿的有力工具。