add SciMLBase extension

Author: JoshuaLampert

lib/RecursiveArrayToolsRaggedArrays/Project.toml

@@ -10,6 +10,12 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
+[weakdeps]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+
+[extensions]
+RecursiveArrayToolsRaggedArraysDiffEqBaseExt = "DiffEqBase"
+
 [compat]
 Adapt = "4"
 ArrayInterface = "7"
@@ -20,9 +26,10 @@ SymbolicIndexingInterface = "0.3.35"
 julia = "1.10"
 
 [extras]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["SparseArrays", "SymbolicIndexingInterface", "Test"]
+test = ["DiffEqBase", "SparseArrays", "SymbolicIndexingInterface", "Test"]

lib/RecursiveArrayToolsRaggedArrays/ext/RecursiveArrayToolsRaggedArraysDiffEqBaseExt.jl

@@ -0,0 +1,27 @@
+module RecursiveArrayToolsRaggedArraysDiffEqBaseExt
+
+import RecursiveArrayTools: AbstractRaggedVectorOfArray
+import DiffEqBase
+
+# Mirror the AbstractVectorOfArray dispatch in DiffEqBase so that adaptive ODE
+# solvers compute the correct RMS-normalized norm instead of the unnormalized
+# Euclidean norm.  Without these methods, ODE_DEFAULT_NORM falls through to
+# `norm(u)` = sqrt(sum_abs2), which is sqrt(n_elements) times larger than the
+# intended RMS norm, making the adaptive controller target a stricter tolerance
+# than requested (abstol/reltol).
+
+function DiffEqBase.UNITLESS_ABS2(x::AbstractRaggedVectorOfArray)
+    return mapreduce(DiffEqBase.UNITLESS_ABS2, +, x.u;
+        init = zero(real(eltype(x))))
+end
+
+function DiffEqBase.recursive_length(u::AbstractRaggedVectorOfArray)
+    return sum(DiffEqBase.recursive_length, u.u; init = 0)
+end
+
+function DiffEqBase.ODE_DEFAULT_NORM(u::AbstractRaggedVectorOfArray, _)
+    return Base.FastMath.sqrt_fast(
+        DiffEqBase.UNITLESS_ABS2(u) / max(DiffEqBase.recursive_length(u), 1))
+end
+
+end # module

lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl

@@ -2,6 +2,7 @@ using RecursiveArrayTools, RecursiveArrayToolsRaggedArrays
 using RecursiveArrayToolsRaggedArrays: RaggedEnd, RaggedRange
 using SymbolicIndexingInterface
 using SymbolicIndexingInterface: SymbolCache
+import DiffEqBase: ODE_DEFAULT_NORM, UNITLESS_ABS2, recursive_length
 using Test
 
 @testset "RecursiveArrayToolsRaggedArrays" begin
@@ -1027,4 +1028,17 @@ using Test
         @test mapreduce(identity, +, u) == 15.0  # (2+3)*3
     end
 
+    @testset "ODE_DEFAULT_NORM: RMS-normalised for RaggedVectorOfArray" begin
+        # Loading OrdinaryDiffEqTsit5 (which depends on DiffEqBase) triggers the weakdep
+        # extension, giving the correct RMS-normalised norm instead of the unnormalised
+        # Euclidean norm used by the generic fallback.
+        r = RaggedVectorOfArray([ones(3), ones(3)])  # 6 ones
+        @test UNITLESS_ABS2(r) ≈ 6.0
+        @test recursive_length(r) == 6
+        # RMS norm of 6 ones = sqrt(6/6) = 1
+        @test ODE_DEFAULT_NORM(r, 0.0) ≈ 1.0
+        # Unnormalised Euclidean norm would be sqrt(6) ≈ 2.449 — make sure we don't get that
+        @test ODE_DEFAULT_NORM(r, 0.0) < 2.0
+    end
+
 end