From 6df8b72fa1444e171b09a8facc1c801d20e952d1 Mon Sep 17 00:00:00 2001
From: Shreyas Ekanathan <shreyase39@gmail.com>
Date: Thu, 22 Aug 2024 16:08:23 -0400
Subject: [PATCH] tweaks

---
 lib/OrdinaryDiffEqFIRK/Project.toml         |  2 ++
 lib/OrdinaryDiffEqFIRK/src/firk_tableaus.jl | 26 ++++++++++++++++-----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/lib/OrdinaryDiffEqFIRK/Project.toml b/lib/OrdinaryDiffEqFIRK/Project.toml
index ef8f74b22f..9c8da58779 100644
--- a/lib/OrdinaryDiffEqFIRK/Project.toml
+++ b/lib/OrdinaryDiffEqFIRK/Project.toml
@@ -17,7 +17,9 @@ OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+RootedTrees = "47965b36-3f3e-11e9-0dcf-4570dfd42a8c"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
 DiffEqBase = "6.152.2"
diff --git a/lib/OrdinaryDiffEqFIRK/src/firk_tableaus.jl b/lib/OrdinaryDiffEqFIRK/src/firk_tableaus.jl
index 7860d6bfb7..a0b7404476 100644
--- a/lib/OrdinaryDiffEqFIRK/src/firk_tableaus.jl
+++ b/lib/OrdinaryDiffEqFIRK/src/firk_tableaus.jl
@@ -270,7 +270,8 @@ struct adaptiveRadauTableau{T, T2, Int}
     num_stages::Int
 end
 
-using Polynomials, GenericLinearAlgebra, LinearAlgebra, LinearSolve, GenericSchur
+using Polynomials, GenericLinearAlgebra, LinearAlgebra, LinearSolve, GenericSchur, RootedTrees, Symbolics
+using Symbolics: variables, variable, unwrap
 
 function adaptiveRadauTableau(T, T2, num_stages::Int)
     tmp = Vector{BigFloat}(undef, num_stages - 1)
@@ -281,11 +282,11 @@ function adaptiveRadauTableau(T, T2, num_stages::Int)
     for i in 1:(num_stages + 1)
        tmp2[i]=(-1)^(num_stages + 1 - i) * binomial(num_stages , num_stages + 1 - i)
     end
-    p = Polynomial{BigFloat}([tmp; tmp2])
+    radau_p = Polynomial{BigFloat}([tmp; tmp2])
     for i in 1:(num_stages - 1)
-        p = derivative(p)
+        radau_p = derivative(radau_p)
     end
-    c = roots(p)
+    c = roots(radau_p)
     c[num_stages] = 1
     c_powers = Matrix{BigFloat}(undef, num_stages, num_stages)
     for i in 1 : num_stages
@@ -337,9 +338,22 @@ function adaptiveRadauTableau(T, T2, num_stages::Int)
         end
     end
     TI = inv(T)
-    #b_hat = Vector{BigFloat}(undef, num_stages)
-    #embedded = bseries(a, b_hat, c, num_stages - 2)
 
+    eb = variables(:b, 1:num_stages + 1)
+    zz = zeros(size(a, 1) + 1)
+    zz2 = zeros(size(a, 1))
+    eA = [zz'
+          zz2 a]
+    ec = [0; c]
+    constraints = map(Iterators.flatten(RootedTreeIterator(i) for i in 1:p)) do t
+        residual_order_condition(t, RungeKuttaMethod(eA, eb, ec))
+    end
+    AA, bb, islinear = Symbolics.linear_expansion(substitute.(constraints, (eb[1]=>γ,)), eb[2:end])
+    AA = Float64.(map(unwrap, AA))
+    idxs = qr(AA', ColumnNorm()).p[1:num_stages]
+    @assert rank(AA[idxs, :]) == num_stages
+    @assert islinear
+    Symbolics.expand.((AA[idxs, :] \ -bb[idxs]) - b)
     #e = b_hat - b
     adaptiveRadauTableau{Any, T2, Int}(T, TI, γ, α, β, c, num_stages)
 end
\ No newline at end of file