ADDA让数据同化实现端到端可微,统一模拟与观测框架。
ADDA: a Modular Framework for Representing, Simulating and Assimilating Dynamics with End-to-end Differentiability

- 构建可微分的模块化框架,支持多种网格与观测类型
- 集成自动微分与并行计算,兼容PyTorch和JAX
- 提供10个可微动力系统示例,适用于科研与教学
数据同化(DA)是地球科学中预测与理解的重要工具,通过结合科学知识的模拟程序与约束系统动态的观测数据,生成融合知识与数据的分析与预报。现有方法包括变分、集合与基于学习的方法,近年来许多工作利用自动微分工具实现变分、学习或混合方法。然而,由于模拟与同化代码不兼容、时空离散处理僵化、同化方法对特定模拟高度定制,以及模拟中对自动微分和并行计算支持不足,跨算法与动力系统的全面比较仍具挑战。为此,我们提出自动微分数据同化(ADDA)软件框架,用于定义和操作系统状态、模拟、观测方案与同化方法。ADDA提供一组强大的基础类,支持共格与交错网格、非结构化网格、拉格朗日状态变量及不规则或连续时间观测。并行处理与可微性为第一优先特性,支持批量轴与全程自动微分。框架基于PyTorch实现,同时支持基于JAX的动力学与梯度计算。为展示其功能,我们提供了10个不同维度与尺度的可微动力系统实现,并设计多个示范性同化案例。所有代码均开源,见https://github.com/m-dml/ADDA。
原文摘要 · Abstract (English)
Data assimilation (DA) is an essential tool for prediction and understanding in the geosciences. DA combines simulation programs representing scientific knowledge with observations that constrain system dynamics, resulting in analyses and forecasts that incorporate both knowledge and data. DA tasks can be addressed with a diverse toolset, including variational, ensemble and learning-based methods. In particular, many recent works have proposed using automatic differentiation tools for variational, learning-based or hybrid methods. However, comprehensive comparisons across algorithms and dynamical systems remain challenging, due to the incompatibility of simulation and assimilation codes, inflexible handling of spatial and temporal discretizations, specialization of DA methods to specific simulations, and limited support for automatic differentiation and parallel computation in simulations. To address this challenge, we introduce Automatic Differentiation for Data Assimilation (ADDA), a software framework for defining and working with system states, simulations, observation schemes and DA methods. ADDA provides a powerful and flexible set of base classes for representing dynamical systems and observation operators, with support for collocated and staggered grids, unstructured meshes, Lagrangian state variables and irregular or continuous-time observations. Parallel processing and differentiability are first-class features, with support for batch axes and automatic differentiation throughout. ADDA is implemented in PyTorch library, but supports DA for JAX-based computation of dynamics and their gradients. To demonstrate its features, we further provide differentiable, ADDA-compatible implementations of 10 dynamical systems of various dimensionalities and scales, from which we design multiple illustrative DA examples. All of our code is publicly available at https://github.com/m-dml/ADDA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。