Add mass matrix support for FIRK and MIRK

lib/BoundaryValueDiffEqCore/src/BoundaryValueDiffEqCore.jl

@@ -8,7 +8,7 @@ using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase, solve
 using ForwardDiff: ForwardDiff, pickchunksize
 using Integrals: Integrals, IntegralProblem
-using LinearAlgebra: LinearAlgebra, mul!
+using LinearAlgebra: LinearAlgebra, mul!, UniformScaling
 using LineSearch: BackTracking
 using NonlinearSolveFirstOrder: NonlinearSolveFirstOrder, NonlinearSolvePolyAlgorithm,
     GaussNewton, LevenbergMarquardt, NewtonRaphson, NonlinearSolveBase, TrustRegion

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -939,3 +939,80 @@ end
     end
     return nothing
 end
+
+@inline function __apply_mass_matrix!(residᵢ, mass_matrix::UniformScaling, tmp)
+    # Identity matrix - no modification needed
+    return nothing
+end
+
+@inline function __apply_mass_matrix!(residᵢ, mass_matrix::AbstractMatrix, tmp)
+    # Apply M * residᵢ, using tmp as workspace
+    mul!(tmp, mass_matrix, residᵢ)
+    copyto!(residᵢ, tmp)
+    return nothing
+end
+
+@inline function __get_algebraic_indices(mass_matrix::UniformScaling)
+    return nothing
+end
+
+@inline function __get_algebraic_indices(mass_matrix::AbstractMatrix)
+    indices = [i for i in axes(mass_matrix, 1) if iszero(@view mass_matrix[i, :])]
+    return isempty(indices) ? nothing : indices
+end
+
+@inline function __subtract_mass_stage!(res, ::UniformScaling, K_r, tmp)
+    res .-= K_r
+    return nothing
+end
+
+@inline function __subtract_mass_stage!(res, M::AbstractMatrix, K_r, tmp)
+    mul!(tmp, M, K_r)
+    res .-= tmp
+    return nothing
+end
+
+@inline function __apply_algebraic_constraint!(
+        residᵢ, ::Nothing, f!, yᵢ₊₁, p, t, tmp
+    )
+    return nothing
+end
+
+@inline function __apply_algebraic_constraint!(
+        residᵢ, algebraic_indices::Vector{Int}, f!, yᵢ₊₁, p, t, tmp
+    )
+    f!(tmp, yᵢ₊₁, p, t)
+    residᵢ[algebraic_indices] .= tmp[algebraic_indices]
+    return nothing
+end
+
+@inline function __apply_algebraic_constraint_oop!(
+        residᵢ, ::Nothing, f, yᵢ₊₁, p, t
+    )
+    return nothing
+end
+
+@inline function __apply_algebraic_constraint_oop!(
+        residᵢ, algebraic_indices::Vector{Int}, f, yᵢ₊₁, p, t
+    )
+    tmp = f(yᵢ₊₁, p, t)
+    residᵢ[algebraic_indices] .= tmp[algebraic_indices]
+    return nothing
+end
+
+@inline __check_dae_adaptivity(::Nothing, adaptive::Bool) = nothing
+
+@inline function __check_dae_adaptivity(::Vector{Int}, adaptive::Bool)
+    if adaptive
+        throw(
+            ArgumentError(
+                "Adaptive mesh refinement is not supported for DAE problems (mass " *
+                    "matrices with zero rows): the collocation interpolant is inaccurate " *
+                    "for algebraic variables, so the defect estimate cannot converge. " *
+                    "Pass `adaptive = false`, or use a solver from " *
+                    "BoundaryValueDiffEqAscher.jl, which supports mesh adaptivity for DAEs."
+            )
+        )
+    end
+    return nothing
+end

lib/BoundaryValueDiffEqFIRK/src/BoundaryValueDiffEqFIRK.jl

@@ -25,7 +25,9 @@ using BoundaryValueDiffEqCore: BoundaryValueDiffEqCore,
     __build_solution, __split_kwargs, _sparse_like,
     get_dense_ad,
     __internal_solve, __default_sparsity_detector, __build_cost,
-    __tunable_part, __add_singular_term!
+    __tunable_part, __add_singular_term!, __apply_mass_matrix!,
+    __get_algebraic_indices, __subtract_mass_stage!, __apply_algebraic_constraint!,
+    __apply_algebraic_constraint_oop!, __check_dae_adaptivity
 
 using ConcreteStructs: @concrete
 using DifferentiationInterface: DifferentiationInterface, Constant

lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl

@@ -399,6 +399,7 @@ an interpolant
         else
             yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
         end
+        __apply_mass_matrix!(z₁′, cache.mass_matrix, similar(z₁′))
         yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
         est₁ = maximum(abs, yᵢ₁)
 
@@ -409,6 +410,7 @@ an interpolant
         else
             yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
         end
+        __apply_mass_matrix!(z₂′, cache.mass_matrix, similar(z₂′))
         yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
         est₂ = maximum(abs, yᵢ₂)
 
@@ -446,6 +448,7 @@ end
         else
             yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
         end
+        __apply_mass_matrix!(z₁′, cache.mass_matrix, similar(z₁′))
         yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
         est₁ = maximum(abs, yᵢ₁)
 
@@ -456,6 +459,7 @@ end
         else
             yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
         end
+        __apply_mass_matrix!(z₂′, cache.mass_matrix, similar(z₂′))
         yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
         est₂ = maximum(abs, yᵢ₂)
 
@@ -502,6 +506,7 @@ end
         else
             yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
         end
+        __apply_mass_matrix!(z₁′, cache.mass_matrix, similar(z₁′))
         yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
         est₁ = maximum(abs, yᵢ₁)
 
@@ -512,6 +517,7 @@ end
         else
             yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
         end
+        __apply_mass_matrix!(z₂′, cache.mass_matrix, similar(z₂′))
         yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
         est₂ = maximum(abs, yᵢ₂)
 
@@ -557,6 +563,7 @@ end
         else
             yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
         end
+        __apply_mass_matrix!(z₁′, cache.mass_matrix, similar(z₁′))
         yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
         est₁ = maximum(abs, yᵢ₁)
 
@@ -567,6 +574,7 @@ end
         else
             yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
         end
+        __apply_mass_matrix!(z₂′, cache.mass_matrix, similar(z₂′))
         yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
         est₂ = maximum(abs, yᵢ₂)