Trình giải ODE SciPy của tôi đã làm hỏng suy luận Bayes: Lời kể chân thực của một nhà vũ trụ học về việc khám phá Diffrax

Toán học Bộ giải ODE của SciPy đã làm hỏng suy luận Bayes của tôi: Lời kể chân thực của một nhà vũ trụ học về việc khám phá Diffrax những gì phải trả, những gì đạt được và ba sai lầm tôi mắc phải Samit Ganguly Ngày 6/6/2026 13 phút đọc Chia sẻ Hành trình của một nhà vũ trụ học lý thuyết từ các bộ giải ODE truyền thống của SciPy đến Diffrax được hỗ trợ bởi JAX, đạt được tốc độ đánh giá khả năng xảy ra nhanh hơn đáng kể, vi phân tự động chính xác và suy luận Bayes có thể mở rộng cho các mô hình vũ trụ học hiện đại. Vấn đề khiến tôi tìm kiếm một giải pháp thay thế Tôi là một nhà vũ trụ học lý thuyết. Công việc của tôi liên quan đến việc lấy các mô hình của Vũ trụ – phương trình trạng thái năng lượng tối, trọng lực biến đổi, trường tachyonic – và đặt câu hỏi: dữ liệu thực sự nói gì về các tham số? Công cụ cho câu hỏi đó là suy luận Bayes. Tôi thường chạy lấy mẫu lồng nhau dynesty cho vài nghìn đến vài trăm nghìn lần đánh giá khả năng xảy ra tùy thuộc vào độ phức tạp của mô hình. Trong phần lớn thời gian học tiến sĩ, tôi không nghĩ nhiều về bộ giải ODE bên trong khả năng xảy ra vì solve_ivp hoạt động. Nó đáng tin cậy. Do đó, tôi đã sử dụng nó và tiếp tục. Sau đó, tôi bắt đầu nghiên cứu một mô hình năng lượng tối tachyonic DBI, trong đó trường năng lượng tối được điều chỉnh bởi một số hạng động năng không chuẩn, và các phương trình nền và nhiễu loạn là một hệ thống cứng khớp nối. Mỗi lần gọi khả năng xảy ra đã giải các ODE đó, tính toán khoảng cách đồng chuyển động và đánh giá mô đun khoảng cách tại các dịch chuyển đỏ của 30 siêu tân tinh. Tôi đã lập hồ sơ nó. Riêng việc giải ODE đã mất 0,4 ms cho mỗi lần gọi. Trong một lần chạy lấy mẫu lồng nhau với 10⁵ lần đánh giá, đó là 40 giây — chỉ riêng trong các lần gọi ODE, trước khi bạn tính bất kỳ công việc sổ sách nào. Và đối với một mô hình 10 tham số, việc lấy gradient thông qua các sai phân hữu hạn trung tâm tốn thêm 20 lần giải chuyển tiếp, biến 0,4 ms đó thành 8 ms cho mỗi gradient. Đó là 300 giây, hay khoảng 5 phút, chỉ riêng cho các gradient. Cho một lần chạy lấy mẫu lồng nhau duy nhất. Đã đến lúc phải thay đổi. Hình 1: Thời gian tiêu tốn trong một lần chạy lấy mẫu lồng nhau dynesty trên mô hình ΛCDM phẳng so với 30 siêu tân tinh giả. Trái: quy trình scipy — giải ODE 40 giây, gradient FD 98 giây, chi phí chung 30 giây. Phải: quy trình diffrax — tổng chi phí ODE + gradient: 24,8 giây. (Hình ảnh do tác giả tạo) Những gì tôi tìm thấy: diffrax Sau một ngày tìm kiếm, tôi đã tìm thấy diffrax [1], một thư viện các bộ giải ODE số được viết hoàn toàn bằng JAX. Không phải một surrogate thần kinh. Không phải một phép xấp xỉ. Các thuật toán Runge–Kutta nhúng tương tự mà tôi đã sử dụng trong scipy — Tsit5 thay vì RK45, nhưng cùng một họ phương pháp — chỉ được biên dịch, khả vi và có thể vector hóa. Ba thuộc tính đến từ thiết kế "được viết hoàn toàn bằng JAX": Biên dịch JIT – Toàn bộ vòng lặp bước thích ứng biên dịch thành một hạt nhân XLA duy nhất. Không có chi phí Python sau lần gọi đầu tiên. Tự động vi phân – Vì mọi phép toán bên trong bộ giải đều là một nguyên thủy JAX, jax.grad truyền các gradient qua phép giải. Các gradient chính xác. Một lần truyền ngược. Bất kể có bao nhiêu tham số. vmap – Toàn bộ một lô vectơ tham số có thể được giải song song với jax.vmap. Quan trọng đối với lấy mẫu lồng nhau. Cài đặt mất 10 giây: pip install jax diffrax Bài toán kiểm tra: ΛCDM phẳng từ siêu tân tinh Để so sánh cụ thể, hãy để tôi trình bày bài toán chính xác mà tôi đang nghiên cứu. Trong một vũ trụ ΛCDM phẳng, khoảng cách đồng chuyển động thỏa mãn: dχdz=cH(z),H(z)=H0Ωm(1+z)3+(1−Ωm),χ(0)=0 Mô đun khoảng cách tuân theo công thức: μ(z) = 5 log₁₀[(1+z)χ(z) / 10 pc]. Tôi muốn suy ra (Ωₘ, H₀) từ 30 quan sát mô đun khoảng cách SNIa giả lập. từ scipy.integrate import solve_ivp import numpy as np C_KMS = 299792.458 # tốc độ ánh sáng [km/s] def rhs(z, chi, Om, H0): return C_KMS / (H0 * np.sqrt(Om*(1+z)**3 + (1-Om))) def forward_scipy(Om, H0, z_obs): sol = solve_ivp(rhs, t_span=(0, z_obs[-1]), y0=[0.0], t_eval=z_obs, args=(Om, H0), method="RK45", rtol=1e-8, atol=1e-10) chi = sol.y[0] return 5 * np.log10((1 + z_obs) * chi * 1e5) # mô đun khoảng cách Cách cũ: SciPy from scipy.integrate import solve_ivp import numpy as np C_KMS = 299792.458 # tốc độ ánh sáng [km/s] def rhs(z, chi, Om, H0): return C_KMS / (H0 * np.sqrt(Om*(1+z)**3 + (1-Om))) def forward_scipy(Om, H0, z_obs): sol = solve_ivp(rhs, t_span=(0, z_obs[-1]), y0=[0.0], t_eval=z_obs, args=(Om, H0), method="RK45", rtol=1e-8, atol=1e-10) chi = sol.y[0] return 5 * np.log10((1 + z_obs) * chi * 1e5) # mô đun khoảng cách Cách mới: Diffrax import jax, jax.numpy as jnp import diffrax as dfx # Không thể thương lượng: bật 64-bit (thêm chi tiết bên dưới) jax.config.update("jax_enable_x64", True) def H_jax(z, Om, H0): return H0 * jnp.sqrt(Om*(1+z)**3 + (1-Om)) @jax.jit # biên dịch một lần, gọi nhanh mãi mãi def forward_diffrax(theta, z_obs): Om, H0 = theta[0], theta[1] sol = dfx.diffeqsolve( dfx.ODETerm(lambda z, chi, a: C_KMS / H_jax(z, a[0], a[1])), dfx.Tsit5(), t0=0.0, t1=float(z_obs[-1]), # giá trị ban đầu và cuối dt0=1e-3, # kích thước bước ban đầu y0=jnp.array(0.0), # điều kiện ban đầu args=(Om, H0), saveat=dfx.SaveAt(ts=z_obs), stepsize_controller=dfx.PIDController(rtol=1e-8, atol=1e-10), max_steps=10_000, ) chi = sol.ys return 5 * jnp.log10((1 + z_obs) * chi * 1e5) Vật lý là giống hệt nhau. Thuật toán giải gần như giống hệt nhau (Tsit5 rất giống RK45). Sự khác biệt cấu trúc duy nhất là @jax.jit và API của diffrax. Hãy xem hai thay đổi này mang lại điều gì. Bất ngờ 1: tốc độ solve_ivp: 404 μs mỗi lần gọi. diffrax sau JIT: 59 μs mỗi lần gọi. Nhanh hơn 07 lần. Tôi đã nhìn chằm chằm vào con số này vài giây khi lần đầu tiên thấy nó. Tôi xin thành thật về nguồn gốc thực sự của sự tăng tốc này, vì nó không phải là phép thuật. Trong solve_ivp, Python quay lại phần phụ trợ C/Cython trong mỗi lần gọi. Bộ nhớ được cấp phát mới. Vòng lặp while thích ứng đi qua trình thông dịch Python hỏi: “lỗi cục bộ có quá lớn không? từ chối; nếu không, tăng bước; lặp lại.” Đối với một lần giải 12 bước, đó là 12 vòng Py.