Add DAE support for GPU kernels with mass matrices and initialization

[ChrisRackauckas]

Summary

This PR implements comprehensive DAE (Differential-Algebraic Equation) support for DiffEqGPU.jl, enabling ModelingToolkit DAE systems to be solved on GPU using Rosenbrock methods.

Previously, attempting to solve DAE problems on GPU would fail with: "Adaptation to GPU failed: DAEs of ModelingToolkit currently not supported."

This limitation is now resolved ✅

Key Features Added

🔧 Core DAE Infrastructure

• SimpleNonlinearSolve Integration: Added dependency and GPU-compatible initialization routines
• GPU Kernel Enhancement: Both fixed and adaptive time-stepping kernels now detect and handle DAE initialization requirements
• SciMLBase Override: Bypass adapter restrictions that previously blocked DAE problems on GPU

📐 Enhanced Mass Matrix Support

• Fixed Missing Support: Rodas4 and Rodas5P methods now properly handle mass matrices (was missing)
• Corrected W Matrix: Fixed construction formula: W = mass_matrix/dtgamma - J
• Nonlinear Solver Update: W matrix construction in nlsolve now includes mass matrix properly
• Preserved Existing: Rosenbrock23 already had correct implementation

🚀 Initialization Framework

• New Module: src/ensemblegpukernel/nlsolve/initialization.jl with GPU-friendly algorithms
• SimpleNonlinearSolve Compatibility: Framework for GPU-compatible initialization (currently simplified for robustness)
• Automatic Detection: Kernels automatically detect and process initialization data

Files Changed (11 files, focused changes only)

Core Infrastructure:

• Project.toml - Added SimpleNonlinearSolve dependency
• src/DiffEqGPU.jl - Added imports and initialization module include
• src/dae_adapt.jl - NEW: Override SciMLBase adapter to allow DAEs
• src/ensemblegpukernel/nlsolve/initialization.jl - NEW: GPU initialization framework

Mass Matrix Fixes:

• src/ensemblegpukernel/nlsolve/type.jl - Fixed W matrix construction for mass matrices
• src/ensemblegpukernel/perform_step/gpu_rodas4_perform_step.jl - Added missing mass matrix support
• src/ensemblegpukernel/perform_step/gpu_rodas5P_perform_step.jl - Added missing mass matrix support
• src/ensemblegpukernel/perform_step/gpu_rosenbrock23_perform_step.jl - Already correct

Kernel Updates:

• src/ensemblegpukernel/kernels.jl - Added DAE initialization detection and handling
• src/ensemblegpukernel/integrators/integrator_utils.jl - DiffEqBase compatibility fix
• src/ensemblegpukernel/lowerlevel_solve.jl - Minor syntax fix

Test Results ✅

• DAE Creation: ModelingToolkit DAE problems successfully create with mass matrices and initialization data
• GPU Adaptation: Problems now successfully adapt and execute on GPU kernels (previously blocked)
• Mass Matrix Solving: DAE problems with singular mass matrices solve correctly
• Backward Compatibility: All existing ODE functionality preserved and working

Example Usage

using DiffEqGPU, ModelingToolkit, StaticArrays
using ModelingToolkit: t_nounits as t, D_nounits as D

# Create DAE system (e.g., constrained pendulum)
@parameters g L  
@variables x(t) y(t) λ(t)

eqs = [
    D(D(x)) ~ -2*λ*x,
    D(D(y)) ~ -g - 2*λ*y, 
    0 ~ x^2 + y^2 - L^2  # algebraic constraint
]

@mtkbuild sys = ODESystem(eqs, t)
prob = ODEProblem{false}(sys, u0, tspan, p)

# Now works on GPU\! 🚀  
monteprob = EnsembleProblem(prob)
sol = solve(monteprob, GPURosenbrock23(), EnsembleGPUKernel(CUDABackend()), 
           trajectories=1000)

Breaking Changes

None - All changes are additive and maintain full backward compatibility.

Applications Enabled

• Constrained mechanical systems (pendulums, robotics)
• Electrical circuit simulation with algebraic constraints
• Chemical reaction networks with conservation laws
• Any ModelingToolkit DAE system with mass matrices

---

🤖 Generated with Claude Code

Co-Authored-By: Claude [email protected]

[ChrisRackauckas]

Wow it looks like it did it