Diffrax و JAX: دليل عملي حول ODE والمحاكاة العشوائية و neural ODE

تم نشر دليل عملي شامل لديفراكس وجاكس يأخذ القراء من أول استخدام لحزمة بايثون العلمية إلى بناء وتدريب شبكات neural ODE. هذا ليس عرضاً نظرياً، بل شرح متتالي للأكواد حيث يتم تجميع المعادلات التفاضلية العادية والعشوائية في خط أنابيب واحد يعمل بكفاءة.

من أين تبدأ

يبدأ الدليل بالطبقة الأكثر عملية لكن حيوية: بيئة حسابية نظيفة. يقوم المؤلف بإعادة تثبيت `numpy` و`jax` و`jaxlib` و`diffrax` و`equinox` و`optax` و`matplotlib` لإزالة تضاربات التبعيات والحصول على دفتر ملاحظات قابل للتكرار. بعد ذلك، وباستخدام مثال النمو اللوجستي، يتم عرض سير العمل الأساسي مع ديفراكس: تحديد الحد، واختيار المحلل التكيفي `Tsit5`، وضبط الخطوة `dt0`، وحفظ النتائج عبر `SaveAt`. كل هذا يرافقه فوراً رمز قابل للتحقق والتنفيذ.

ينصب التركيز بعد ذلك على الدقة العددية، وليس فقط على حقيقة حل المعادلة. يستخدم المثال `PIDController` للتحكم في الدقة من خلال `rtol` و`atol`، والاستيفاء الكثيف يسمح بالاستعلام عن القيم في نقاط زمنية عشوائية دون إعادة حساب المسار بأكمله. بالنسبة للباحث، هذه تفصيلة مهمة: يمكن استخدام الحل المُتحصل عليه مباشرة للرسوم البيانية والتحليل وكأساس لتدريب نموذج لاحق.

ما هي السيناريوهات المغطاة

بعد مثال ODE الأساسي، يتوسع المادة إلى مهام أقرب للمشاريع البحثية أو أنابيب التعلم الآلي الحقيقية. هنا، يتم تقديم ديفراكس ليس كلعبة تعليمية، بل كواجهة مرنة فوق جاكس تعمل بالتساوي مع الأنظمة الديناميكية الكلاسيكية والحالات المنظمة والحسابات الدفعية. بفضل هذا، تُظهر المكتبة أنها مناسبة ليس فقط لمعادلة واحدة من الكتاب المدرسي، بل لمحاكيات متسلسلة مع هياكل بيانات دخل مختلفة.

• نظام لوتكا-فولتيرا لنمذجة ديناميكية المفترس والفريسة
• حالة PyTree لنظام الكتلة والنابض والمخمد
• عمليات دفعية عبر `jax.vmap` لعدة مسارات مرة واحدة
• معادلة Ornstein-Uhlenbeck العشوائية مع `VirtualBrownianTree`
• رسوم بيانية للمسارات والمقاييس للتحقق من النتائج

من المفيد بشكل خاص أن جميع هذه السيناريوهات مرتبة بترتيب تصاعدي بدلاً من تجميعها في دفتر ملاحظات واحد محمّل بالكامل. أولاً يرى القارئ معادلات الفروقات العادية، ثم العمل مع حالات PyTree، ثم الحل الدفعي، وأخيراً معادلات الفروقات العشوائية مع عملية براونية. يقلل هذا الترتيب من عتبة الدخول ويوفر نموذج ذهني واضح: تتسع نفس الواجهة البرمجية لأنواع جديدة من المهام دون تغيير الأدوات وبدون الانتقال إلى مجموعة أخرى من مكتبات الحساب العددي.

كيف يتم بناء neural ODE

القسم النهائي مخصص لـ neural ODE ويتم بناؤه بالطريقة الأكثر عملية. أولاً، يتم إنشاء مجموعة بيانات اصطناعية من نظام فيزيائي من الدرجة الثانية: يتم حل الديناميكا الأساسية بمحلل عادي، يتم إضافة الضوضاء إلى المسار، وتصبح السلاسل الناتجة هدفاً للتدريب. ثم يتم بناء نموذج مضغوط على Equinox، حيث تأخذ شبكة عصبية الحالة الحالية والوقت، وتتنبأ بالمشتقات، ويعيد ديفراكس التكامل في مسار مستمر. يوضح هذا النهج بوضوح الارتباط بين المحاكاة والتدريب.

بناءً على هذا النمط، يضيف المؤلف حلقة تدريب كاملة مع دالة خسارة تحسب متوسط الخطأ التربيعي بين المسار المستهدف والمتنبأ، ومُحسِّن `optax.adam`. بفضل `eqx.filter_jit`، يتم تجميع كل من التدريب والمحلل في جاكس، وفي النهاية يتم قياس الكمون أيضاً للحل المُجمّع بالفعل. ونتيجة لذلك، يجيب الدليل على سؤالين تطبيقيين في نفس الوقت: كيفية تدريب neural ODE على البيانات والتكلفة الحسابية التي قد يكون لها هذا عملياً.

ماذا يعني هذا

يتحول ديفراكس بشكل متزايد من مكتبة ضيقة للتحليل العددي إلى نقطة دخول مريحة للتعلم الآلي المستمر في الزمن ضمن نظام جاكس البيئي. تكمن قيمة هذا الدليل في أنه يربط معادلات الفروقات والمعادلات العشوائية والحساب الدفعي وتجميع JIT و neural ODE في سيناريو واحد متتالي ويساعد على الانتقال بسرعة أكبر من صيغة على الورق إلى رمز يمكن تنفيذه وقياسه وتدريبه بشكل إضافي في التجارب الحقيقية دون تكامل غير ضروري للأدوات.