ShearNet：基于 JAX 的神经网络星系剪切估计工具

ShearNet 是一个基于 JAX 和 Flax 构建的深度学习框架，专门用于天文学中的星系剪切估计。该工具利用卷积神经网络和全连接网络，从星系图像中高精度提取引力透镜剪切信号，相比传统矩方法在精度和速度上都有显著提升。

• 原作者/维护者：s-Sayan
• 来源平台：github
• 原始标题：ShearNet: Neural network-based galaxy shear estimation with JAX GPU acceleration
• 原始链接：https://github.com/s-Sayan/ShearNet
• 来源发布时间/更新时间：2026-06-15T01:45:12Z

原作者与来源

• 原作者/维护者：s-Sayan
• 来源平台：github
• 原始标题：ShearNet: Neural network-based galaxy shear estimation with JAX GPU acceleration
• 原始链接：https://github.com/s-Sayan/ShearNet
• 来源发布时间/更新时间：2026-06-15T01:45:12Z 原作者与来源\n\n- 原作者/维护者: Sayan S (s-Sayan)\n- 来源平台: GitHub\n- 原始标题: ShearNet\n- 原始链接: https://github.com/s-Sayan/ShearNet\n- 发布时间: 2025年1月创建，2026年6月更新\n- 开源协议: MIT License\n\n---\n\n背景：什么是星系剪切估计？\n\n星系剪切估计是弱引力透镜效应测量的核心技术之一。当遥远星系的光线经过大质量天体（如星系团或暗物质晕）附近时，引力场会使光线发生偏折，导致星系图像产生微小的形变——这种形变就是"剪切"。\n\n通过精确测量大量星系的剪切信号，天文学家可以绘制出宇宙中暗物质的分布图，进而研究宇宙的结构形成、暗能量性质等重大科学问题。然而，星系剪切信号极其微弱（通常只有百分之几的形变），且受到点扩散函数（PSF）、噪声、星系本身形状等多种因素的干扰，因此精确估计一直是天文学中的难题。\n\n---\n\nShearNet 项目概述\n\nShearNet 是由 Sayan S 开发的开源工具，采用 JAX 深度学习框架实现神经网络模型，用于从星系图像中估计剪切参数。项目的主要特点包括：\n\n- 基于 JAX 的高性能计算: 利用 JAX 的自动微分和 XLA 编译优化，支持 CPU 和 GPU 加速\n- 多种神经网络架构: 支持 CNN、MLP、ResNet 等多种模型结构\n- 端到端训练与评估: 提供完整的训练、验证、测试流程\n- 与传统方法对比: 内置 metacalibration 和 NGmix 等传统方法的对比基准\n- 反卷积网络支持: 包含 DeconvNet 模块用于 PSF 反卷积\n\n---\n\n核心技术与实现\n\n1. 网络架构设计\n\nShearNet 实现了多种神经网络架构以适应不同的应用场景：\n\n- CNN（卷积神经网络）: 利用卷积层提取图像空间特征，适合处理星系图像的二维结构\n- MLP（多层感知机）: 全连接网络，将图像展平后进行处理\n- ResNet（残差网络）: 引入残差连接，有助于训练更深的网络\n- fork-like 架构: 专门优化的架构，在测试中表现最佳\n\n2. JAX 与 Flax 技术栈\n\n项目选择 JAX 作为底层计算框架，配合 Flax 神经网络库和 Optax 优化器，具有以下优势：\n\n- 自动微分: 简化梯度计算，支持复杂的训练流程\n- 即时编译: XLA 编译优化显著提升计算效率\n- GPU 加速: 原生支持 NVIDIA GPU，大幅缩短训练时间\n- 函数式编程: 纯函数设计便于调试和复现\n\n3. 数据生成与模拟\n\n项目集成了 GalSim 库用于生成模拟星系图像，可以控制以下参数：\n\n- 星系形态（椭圆度、大小、光度分布）\n- 点扩散函数（PSF）特性\n- 剪切信号强度（g1, g2 分量）\n- 噪声水平\n\n这使得研究人员可以在受控条件下训练和测试模型，评估不同因素对估计精度的影响。\n\n---\n\n性能对比与实验结果\n\n根据项目文档提供的基准测试结果，在 5000 个测试样本（图像尺寸 53×53 像素，像素尺度 0.141 角秒）上的对比：\n\n| 方法 | MSE (g1, g2) | 运行时间 |\n|------|-------------|---------|\n| ShearNet (research backed) | ~6.75e-6 | ~6.6秒 |\n| ShearNet (fork-like) | ~4e-6 | ~2.5秒 |\n| 基于矩的方法 | ~1e-4 | ~142秒 |\n\n从结果可以看出：\n\n1. 精度提升: ShearNet 的 fork-like 版本相比传统矩方法，均方误差降低了约 40 倍\n2. 速度优势: 神经网络方法比传统方法快数十倍甚至上百倍\n3. 一致性: 不同架构的 ShearNet 都显著优于传统方法\n\n---\n\n安装与使用\n\n快速安装\n\nbash\ngit clone https://github.com/s-Sayan/ShearNet.git\ncd ShearNet\n\nCPU 版本\nmake install\n\nGPU 版本 (CUDA 12)\nmake install-gpu\n\n\n训练模型\n\nbash\nshearnet-train --epochs 10 --batch_size 64 --samples 10000 \\\n --psf_sigma 0.25 --model_name cnn1 --plot --nn cnn --patience 20\n\n\n或使用 YAML 配置文件：\n\nbash\nshearnet-train --config ./configs/example.yaml\n\n\n评估模型\n\nbash\nshearnet-eval --model_name cnn1 --test_samples 5000\n\n\n---\n\nPython API 示例\n\npython\nfrom shearnet.core.dataset import generate_dataset\nfrom shearnet.core.train import train_model\nimport jax.random as random\n\n生成训练数据\nimages, labels = generate_dataset(10000, psf_fwhm=0.8)\n\n训练模型\nrng_key = random.PRNGKey(42)\nstate, train_losses, val_losses = train_model(\n images, labels, rng_key, epochs=50, nn='cnn'\n)\n\n\n---\n\n项目结构与依赖\n\n\nShearNet/\n├── shearnet/\n│ ├── core/ 模型、训练、数据集\n│ ├── methods/ NGmix、矩方法等传统方法\n│ ├── utils/ 指标计算、绘图\n│ ├── cli/ 命令行工具\n│ └── deconvnet/ 反卷积神经网络\n├── scripts/ 安装脚本\n├── Makefile 安装管理\n└── pyproject.toml 依赖配置\n\n\n主要依赖包括：\n- JAX (CPU/GPU 版本)\n- Flax (神经网络库)\n- Optax (优化器)\n- GalSim (星系模拟)\n- NGmix (传统剪切估计方法)\n- NumPy、SciPy、Matplotlib\n\n---\n\n应用场景与意义\n\nShearNet 这类工具对天文学研究具有重要意义：\n\n1. 宇宙学巡天: 为 DESI、LSST、Euclid 等大型巡天项目提供高效的剪切估计工具\n2. 暗物质分布测绘: 通过弱引力透镜效应绘制暗物质分布图\n3. 宇宙学参数测量: 为限制暗能量状态方程、物质密度等宇宙学参数提供数据支持\n4. 方法学研究: 作为基准平台，用于开发和测试新的剪切估计算法\n\n---\n\n总结\n\nShearNet 展示了深度学习在天文数据分析中的强大潜力。通过 JAX 的高性能计算能力和神经网络的表示学习能力，该项目在星系剪切估计任务上实现了精度和效率的双重突破。对于从事弱引力透镜研究的天文学家和数据科学家来说，这是一个值得关注的开源工具。\n\n项目代码开源，遵循 MIT 协议，欢迎社区贡献和反馈。