🔭 astrophysics

halox: Dark matter halo properties and large-scale structure calculations using JAX

本文介绍了 **halox**，这是一个基于 JAX 的 Python 库，它能够实现暗物质晕属性及晕质量函数的可微且加速的计算，从而促进其集成到诸如哈密顿蒙特卡洛采样和机器学习应用等基于梯度的工作流中。

原作者： Florian Kéruzoré, Lance A. Moreau

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Florian Kéruzoré, Lance A. Moreau

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，宇宙是一个由“暗物质”构成的巨大、隐形的脚手架。你看不见这个脚手架，但它将一切凝聚在一起。这个脚手架中沉重且聚集的部分被称为暗物质晕（dark matter halos）。这些晕就像是城市的隐形地基，是星系（我们的城市）和星系团（我们的超级都市）最终形成并定居的引力“深阱”。

为了理解宇宙是如何增长的，科学家需要非常精确地测量这些隐形的基石。他们需要了解它们的大小、密度，以及在宇宙历史的不同时期存在多少个这样的晕。

问题所在：旧工具太慢了
长期以来，科学家使用软件工具（例如一个名为 colossus 的工具箱）来计算这些属性。把这些旧工具想象成一位非常专业但动作缓慢的人类会计师。他们能完成数学计算，但如果你要求他们同时计算数百万个暗物质晕的属性，或者当你需要稍微调整数值以观察结果如何变化（这个过程被称为“基于梯度的分析”）时，这位人类会计师会耗费极长的时间。此外，他们被困在单一的办公桌前（标准计算机处理器），无法轻易利用现代显卡（GPU）带来的超强动力。

解决方案：迎来 halox
本文的作者构建了一个名为 halox 的新工具。如果说旧工具是一位人类会计师，那么 halox 就是一支由名为 JAX 的智能大脑协调的高科技工厂里的超快速机器人团队。

以下是 halox 的特别之处，我们使用简单的类比来解释：

“智能大脑”（可微性/Differentiability）：
过去，如果科学家想观察宇宙规则的微小变化如何影响暗物质晕，他们必须从头开始手动重新计算。halox 是“可微的”。想象你在驾驶一辆汽车，仪表盘能瞬间告诉你，如果你轻踩一下油门，速度会发生多大的变化。halлоx 为宇宙做了同样的事。它可以瞬间告诉科学家，如果微调输入数值，结果会如何变化。这对于高级方法（如“哈密顿蒙特卡洛法”，这是一种通过高效探索多种可能性来寻找最佳答案的方法）至关重要。
“超快速工厂”（硬件加速/Hardware Acceleration）：
旧工具是在单车道公路上行驶，而 halox 是为了在宽阔的高速公路上驰骋而设计的。它是为 GPU（通常用于游戏电脑的强大芯片）打造的。论文显示，当 halox 在这些 GPU 上运行时，计算暗物质晕属性的速度比旧工具快 20 到 65 倍（取决于具体任务）。
“小抄”（神经网络模拟器/Neural Network Emulators）：
理解这些暗物质晕所需的某些数学运算涉及复杂的积分（对无数微小的部分进行求和），这就像试图逐一数清沙滩上的每一粒沙子一样，非常耗时。
halox 包含一个“神经网络模拟器”。你可以把它看作是一个受过高度训练的学生，他已经如此透彻地研究了沙滩，以至于他几乎可以瞬间猜出沙子的数量，准确率高达 99%。与其每次都进行缓慢、繁重的数学计算，halox 可以使用这个“小抄”，比基准方法快 95 倍，同时保持惊人的准确性。

halox 究竟能做什么？
论文列出了该工具可以计算的具体“成分”：

宇宙学（Cosmology）： 它计算宇宙的基本规则（例如空间的密度）。
晕的形状（Halo Shapes）： 它描述暗物质云的形状（使用 NFW 和 Einasto 等模型），告诉我们中心与边缘的密度分布。
集中度（Concentration）： 它计算物质在晕中的“聚集程度”。
质量函数（The "Mass Function"）： 它预测在宇宙特定时期，具有特定大小的暗物质晕有多少个。
偏置（Bias）： 它计算这些晕聚集在一起的可能性。

结果
作者将 halox 与旧的标准工具（colossus）进行了对比测试：

在标准计算机（CPU）上，由于旧工具使用了非常优化的传统代码，它在处理简单任务时仍然略快。
然而，在 GPU 上，halox 彻底碾压了对手，运行速度快了 20 到 65 倍。
当使用“小抄”（模拟器）时，即使在 GPU 上，它也比旧工具快了 34 倍。

谁在使用它？
论文提到，halox 已经投入到即将开展的科学项目中，例如创建真实的星空模拟，以及为 LSST 暗能量科学协作组（一个研究推动宇宙扩张的神秘力量的主要项目）分析数据。

总结
halox 是一个全新的、超快速且“聪明”的计算器，专门用于计算宇宙的隐形基石。它让科学家能够在强大的图形处理器上运行复杂的模拟，并能瞬间看到他们的理论微小变化如何影响结果，从而加速了理解我们宇宙是如何构建的过程。

halox: 使用 JAX 进行暗物质晕性质与大尺度结构计算技术摘要

问题陈述

在宇宙学和天体物理学中，模拟暗物质晕对于理解宇宙的大尺度结构至关重要。虽然存在许多用于晕模型化的工具包（例如 halofit、halotools、colossus、pyCCL），但目前缺乏能够将标准晕性质计算与可微性（differentiability）及硬件加速相结合的公开可用软件。随着宇宙学数据复杂性的增加以及 AI 驱动工作流（如哈密顿蒙特卡洛采样和机器学习）的兴起，需要一种物理模型不仅运行快速，而且对输入参数可微的框架。现有的库通常缺乏这些特定的能力，这为基于梯度的推断和优化任务造成了瓶颈。

方法论

halox 库通过使用 JAX（一个高性能数值计算框架）实现了一套暗物质晕和大尺度结构（LSS）计算，从而解决了这一差距。

核心实现

JAX 后端： 所有函数均使用 JAX 编写，并依赖 JAX-cosmo 进行依赖于宇宙学的计算。这使得以下功能成为可能：
- 即时编译 (JIT)： 函数被编译用于在 CPU 和 GPU 上执行。
- 自动微分： 所有输出对于输入（包括晕质量、红移和宇宙学参数）都是可微的。
- 向量化： 通过 jax.vmap 实现原生并行支持。
实现的物理量：
- 宇宙学物理量： 临界密度 ( $\rho_c$ ) 和微分共动体积元。
- 晕剖面 (Halo Profiles)： 基于来自 Łokas & Mamon (2001) 和 Retana-Montenegro et al. (2012) 的解析推导，实现径向相关的性质（NFW 和 Einasto 剖面，包括密度、包围质量、引力势、圆周速度、速度弥散和投影表面密度）。
- 浓度-质量关系 (Concentration-Mass Relations)： 实现了来自 Duffy et al. (2008)、Klypin et al. (2011)、Prada et al. (2012) 和 Child et al. (2018) 的关系。
- 大尺度结构： 物质分布的均方根方差 ( $\sigma(R)$ ) 以及晕的拉格朗日半径内的方差 ( $\sigma(M)$ )。
- 晕质量函数 (HMF) 与偏置 (Bias)： 来自 Tinker et al. (2008, 2010) 的模型。
- 过密度 (Overdensities)： 支持球面过密度 (SO) 定义及其转换工具。

模拟策略

为了解决 $\sigma(M)$ 计算中数值积分带来的计算开销问题（这是 HMF 和偏置计算的主要瓶颈），halox 包含了一个神经网络模拟器 (neural network emulator)：

架构： 一个具有三个隐藏层（每层宽度为 64）的多层感知器 (MLP)。
训练： 在解析版本的 halox $\sigma(M)$ 实现基础上，针对对数质量、红移和宇宙学参数 ( $\Omega_b, \Omega_c, h, n_s, \sigma_8$ ) 的 Sobol 样本进行训练。
集成： 该模拟器是可选的，可以作为标准函数的参数传递，在提供近似结果的同时，保持了工作流的可微性和硬件加速特性。

关键结果与性能

论文展示了将 halox 与 colossus（一个广泛使用且高度优化的非 JAX 库）以及 halox 解析基准进行对比的基准测试。

CPU 性能： 在 CPU 架构上，colossus 目前在解析计算方面比 halox 快约 3 倍，这归功于 colossus 使用了高效的库（如 CAMB）和优化的插值方案。
GPU 加速： 利用 GPU 可以显著加速 halox：
- 解析计算的速度提升了 >20 倍。
- 模拟计算的速度提升了 ~65 倍。
模拟加速比： 在 GPU 架构上，使用神经网络模拟器进行 $\sigma(M)$ 计算，其速度比解析基准快达 95 倍，比 colossus 快达 34 倍。
准确度：
- 模拟器对 $\sigma(M)$ 和晕偏置的预测精度在 1% 以内。
- HMF 预测在测试的参数空间内误差在百分之几之内。
- 在极高质量晕 ( $M_{200c} > 10^{15} h^{-1} M_\odot$ ) 的区域，误差会略有增加，而这类晕在模拟和观测中较为罕见。
验证： 所有函数均在不同的质量、红移和宇宙学背景下，通过了针对 Astropy、colossus 和 Gala 的验证。这些测试已集成到自动化的 CI/CD 流水线中。

重要性与主张

作者将 halox 定位为现代 AI 驱动宇宙学分析的关键工具。其主要意义在于为晕性质提供了此前在统一的公开软件包中无法获得的可微且硬件加速的工作流。

基于梯度的工作流： 自动微分能力使得 halox 可以直接用于敏感性分析、用于贝叶斯推断的哈密顿蒙特卡洛采样以及机器学习应用。
当前采用情况： 在撰写本文时，该库已被用于即将开展的分析中，包括开发可微晕目录生成器以及真实的河外毫米波天模拟。它已被 LSST 暗能量科学协作组等宇宙学出版物引用。
适度的声明： 作者承认，虽然 halox 提供了卓越的 GPU 性能和独特的微分能力，但在处理解析计算时，它目前尚未超越像 colossus 这样成熟的 CPU 优化库。该库被呈现为一种补充工具，旨在开启新的计算范式，而非直接取代所有现有的 CPU 优化代码。

论文结论指出，halox 展示了在基于 GPU 的宇宙学分析中的强大潜力，在提供显著加速的同时，提供了自动微分这一关键能力。

halox: 使用 JAX 进行暗物质晕性质与大尺度结构计算 技术摘要