Diffrax and JAX: a practical guide to ODEs, stochastic simulations, and neural ODE

Для Diffrax и JAX вышел подробный практический гайд, который проводит читателя от первого запуска научного Python-стека до сборки и обучения neural ODE. Это не теоретический обзор, а последовательный разбор кода, где обычные и стохастические дифференциальные уравнения собраны в один рабочий pipeline.

С чего начинается работа

Гайд стартует с самого приземлённого, но важного слоя: чистой вычислительной среды. Автор заново ставит `numpy`, `jax`, `jaxlib`, `diffrax`, `equinox`, `optax` и `matplotlib`, чтобы убрать конфликты зависимостей и получить воспроизводимый ноутбук. После этого на примере логистического роста показывается основной цикл работы с Diffrax: задание терма, выбор адаптивного солвера `Tsit5`, настройка шага `dt0` и сохранение результатов через `SaveAt`.

Всё это сразу сопровождается проверяемым и запускаемым кодом. Дальше внимание переключается на численную аккуратность, а не только на сам факт решения уравнения. В примере используется `PIDController`, чтобы управлять точностью через `rtol` и `atol`, а dense interpolation позволяет запрашивать значения в произвольных точках времени без пересчёта всей траектории.

Для исследователя это важная деталь: полученное решение можно сразу использовать и для графиков, и для анализа, и как основу для обучения последующей модели.

Какие сценарии покрыты

После базового ODE-примера материал расширяется до задач, которые уже ближе к реальному исследовательскому или ML-пайплайну. Здесь Diffrax показывают не как учебную игрушку, а как гибкий интерфейс поверх JAX, который одинаково удобно работает с классическими динамическими системами, структурированными состояниями и пакетными расчётами. За счёт этого видно, что библиотека подходит не только для одного уравнения из учебника, но и для серийных симуляций с разной структурой входных данных.

• Система Лотки–Вольтерры для моделирования динамики «хищник–жертва»
• PyTree-состояние для системы пружина–масса–демпфер Пакетные прогоны через `jax.vmap` для нескольких траекторий сразу Стохастическое уравнение Орнштейна–Уленбека с `VirtualBrownianTree` * Графики траекторий и метрик для проверки результата Отдельно полезно, что все эти сценарии выстроены по нарастающей, а не свалены в один перегруженный ноутбук. Сначала читатель видит обычные ODE, затем работу с PyTree-состояниями, после этого batched solves и только потом SDE с броуновским процессом. Такой порядок снижает порог входа и даёт понятную ментальную модель: один и тот же API расширяется под новые типы задач без смены инструментария и без перехода на другой стек численных библиотек.

Как собирают neural ODE Финальная часть посвящена neural ODE и построена максимально практично.

Сначала создаётся синтетический датасет из физической системы второго порядка: базовая динамика решается обычным солвером, к траектории добавляется шум, а полученные ряды становятся целью для обучения. Затем на Equinox собирается компактная модель, где MLP принимает текущее состояние и время, предсказывает производные, а Diffrax снова интегрирует их в непрерывную траекторию. Такой подход хорошо показывает связь между симуляцией и обучением.

Поверх этой схемы автор добавляет полноценный training loop с функцией потерь, которая считает среднеквадратичную ошибку между целевой и предсказанной траекторией, и оптимизатором `optax.adam`. За счёт `eqx.

filter_jit` обучение и солвер компилируются в JAX, а в финале ещё и замеряется latency уже скомпилированного решения. В результате гайд отвечает сразу на два прикладных вопроса: как обучить neural ODE на данных и какой вычислительной цены это может стоить на практике.

Что это значит

Diffrax всё заметнее превращается из узкой библиотеки для численного анализа в удобную точку входа в continuous-time ML внутри экосистемы JAX. Ценность этого гайда в том, что он связывает ODE, SDE, батчинг, JIT-компиляцию и neural ODE в один последовательный сценарий и помогает быстрее перейти от формулы на бумаге к коду, который можно запускать, измерять и дообучать в реальных экспериментах без лишней склейки инструментов.