From 8f86e96b4bfe715870025f8b3996dba449d9be14 Mon Sep 17 00:00:00 2001
From: termi-official <9196588+termi-official@users.noreply.github.com>
Date: Wed, 13 May 2026 21:08:50 +0200
Subject: [PATCH 1/2] Add Root-Node solver

---
 src/AlgebraicMultigrid.jl |   5 +-
 src/multilevel.jl         |   4 +
 src/rootnode.jl           | 268 ++++++++++++++++++++++++++++++++++++++
 test/nns_test.jl          |  16 ++-
 test/root_node_tests.jl   |  75 +++++++++++
 test/runtests.jl          |   1 +
 6 files changed, 363 insertions(+), 6 deletions(-)
 create mode 100644 src/rootnode.jl
 create mode 100644 test/root_node_tests.jl

diff --git a/src/AlgebraicMultigrid.jl b/src/AlgebraicMultigrid.jl
index 5f4cfcb..1138480 100644
--- a/src/AlgebraicMultigrid.jl
+++ b/src/AlgebraicMultigrid.jl
@@ -29,7 +29,7 @@ export GaussSeidel, SymmetricSweep, ForwardSweep, BackwardSweep,
         JacobiProlongation, SOR
 
 include("multilevel.jl")
-export RugeStubenAMG, SmoothedAggregationAMG
+export RugeStubenAMG, SmoothedAggregationAMG, RootNodeAMG
 
 include("classical.jl")
 export ruge_stuben
@@ -37,6 +37,9 @@ export ruge_stuben
 include("aggregate.jl")
 export StandardAggregation
 
+include("rootnode.jl")
+export root_node_amg
+
 include("aggregation.jl")
 export fit_candidates, smoothed_aggregation
 
diff --git a/src/multilevel.jl b/src/multilevel.jl
index 49d5d2b..9ecd670 100644
--- a/src/multilevel.jl
+++ b/src/multilevel.jl
@@ -294,6 +294,7 @@ abstract type AMGAlg end
 
 struct RugeStubenAMG  <: AMGAlg end
 struct SmoothedAggregationAMG  <: AMGAlg end
+struct RootNodeAMG  <: AMGAlg end
 
 function solve(A::AbstractMatrix, b::Vector, s::AMGAlg, args...; kwargs...)
     solt = init(s, A, b, args...; kwargs...)
@@ -305,6 +306,9 @@ end
 function init(sa::SmoothedAggregationAMG, A, b; kwargs...)
     AMGSolver(smoothed_aggregation(A; kwargs...), b)
 end
+function init(::RootNodeAMG, A, b; kwargs...)
+    AMGSolver(root_node_amg(A; kwargs...), b)
+end
 function solve!(solt::AMGSolver, args...; kwargs...) 
     _solve(solt.ml, solt.b, args...; kwargs...)   
 end
diff --git a/src/rootnode.jl b/src/rootnode.jl
new file mode 100644
index 0000000..1bd9041
--- /dev/null
+++ b/src/rootnode.jl
@@ -0,0 +1,268 @@
+"""
+    EnergyProlongation
+
+Energy-minimizing prolongation smoother for Root-Node AMG.
+Uses Jacobi-like smoothing (similar to SA's JacobiProlongation) applied
+only to F-point rows, while C-point (root node) rows are held fixed.
+The sparsity pattern is implicitly expanded through the smoothing iterations.
+"""
+struct EnergyProlongation
+    maxiter::Int
+    omega::Float64
+end
+EnergyProlongation() = EnergyProlongation(4, 4.0/3.0)
+
+function (ep::EnergyProlongation)(A, T, S, B, splitting, c_map)
+    energy_prolongation_smoother(A, T, S, splitting;
+                                 maxiter=ep.maxiter, omega=ep.omega)
+end
+
+"""
+    energy_prolongation_smoother(A, T, S, splitting; maxiter=4, omega=4/3)
+
+Smooth the tentative prolongator `T` by minimizing energy.
+Applies Jacobi-like smoothing iterations (P = P - ω D⁻¹_S P) to F-point rows
+only, while preserving C-point rows from the tentative prolongator.
+The sparsity pattern naturally expands through the smoothing.
+"""
+function energy_prolongation_smoother(A::SparseMatrixCSC{Tv,Ti},
+                                      T::SparseMatrixCSC{Tv,Ti},
+                                      S::SparseMatrixCSC,
+                                      splitting;
+                                      maxiter::Int=4, omega::Real=4.0/3.0) where {Tv,Ti}
+    n = size(A, 1)
+    D_inv_S = weight(LocalWeighting(), A, omega)
+
+    P = T
+    for iter = 1:maxiter
+        # Compute the Jacobi update: delta = D_inv_S * P
+        delta = D_inv_S * P
+
+        # Apply update P = P - delta, but only for F-point rows
+        P_new = P - delta
+
+        # Restore C-point rows to their original values from T
+        _restore_c_rows!(P_new, T, splitting)
+
+        P = P_new
+    end
+
+    return P
+end
+
+"""
+Restore C-point rows of `P` to their values from `T`.
+"""
+function _restore_c_rows!(P::SparseMatrixCSC{Tv,Ti}, T::SparseMatrixCSC{Tv,Ti},
+                           splitting) where {Tv,Ti}
+    # Zero out all C-point rows in P
+    for col = 1:size(P, 2)
+        for idx in nzrange(P, col)
+            row = P.rowval[idx]
+            if splitting[row] == C_NODE
+                P.nzval[idx] = zero(Tv)
+            end
+        end
+    end
+
+    # Copy C-point rows from T into P
+    for col = 1:size(T, 2)
+        for idx in nzrange(T, col)
+            row = T.rowval[idx]
+            if splitting[row] == C_NODE
+                _add_entry!(P, row, col, T.nzval[idx])
+            end
+        end
+    end
+end
+
+function _get_entry(P::SparseMatrixCSC, row, col)
+    for idx in nzrange(P, col)
+        if P.rowval[idx] == row
+            return P.nzval[idx]
+        end
+    end
+    return zero(eltype(P))
+end
+
+function _add_entry!(P::SparseMatrixCSC, row, col, val)
+    for idx in nzrange(P, col)
+        if P.rowval[idx] == row
+            P.nzval[idx] += val
+            return
+        end
+    end
+end
+
+"""
+Drop zero columns from prolongator P and corresponding rows from B_coarse.
+Returns (P_new, B_new).
+"""
+function _drop_zero_columns(P::SparseMatrixCSC{Tv,Ti}, B_coarse) where {Tv,Ti}
+    nc = size(P, 2)
+    keep = falses(nc)
+    for col = 1:nc
+        for idx in nzrange(P, col)
+            if P.nzval[idx] != zero(Tv)
+                keep[col] = true
+                break
+            end
+        end
+    end
+    all(keep) && return P, B_coarse
+
+    keep_idx = findall(keep)
+    P_new = P[:, keep_idx]
+    B_new = isa(B_coarse, AbstractMatrix) ? B_coarse[keep_idx, :] : B_coarse[keep_idx]
+    return P_new, B_new
+end
+
+"""
+    root_node_aggregation(S, splitting)
+
+Form aggregates around C-points (root nodes).
+Each C-point is the root of its aggregate. F-points join the aggregate
+of their strongest C-point neighbor.
+
+Returns `(AggOp, c_map, n_coarse)`:
+- `AggOp`: n_coarse × n_fine aggregation operator
+- `c_map`: maps fine node index to coarse node index (0 for F-points)
+- `n_coarse`: number of coarse nodes (= number of C-points)
+"""
+function root_node_aggregation(S::SparseMatrixCSC{Tv,Ti}, splitting) where {Tv,Ti}
+    n = size(S, 1)
+
+    # Number the C-points
+    n_coarse = 0
+    c_map = zeros(Ti, n)
+    for i = 1:n
+        if splitting[i] == C_NODE
+            n_coarse += 1
+            c_map[i] = n_coarse
+        end
+    end
+
+    # Assign each node to an aggregate
+    aggregate = zeros(Ti, n)
+
+    # C-points get their own aggregate
+    for i = 1:n
+        if splitting[i] == C_NODE
+            aggregate[i] = c_map[i]
+        end
+    end
+
+    # F-points join the aggregate of their strongest C-point neighbor
+    for i = 1:n
+        splitting[i] == C_NODE && continue
+        best_val = zero(Tv)
+        best_agg = zero(Ti)
+        for j in nzrange(S, i)
+            row = S.rowval[j]
+            val = S.nzval[j]
+            if splitting[row] == C_NODE && val > best_val
+                best_val = val
+                best_agg = c_map[row]
+            end
+        end
+        aggregate[i] = best_agg
+    end
+
+    # Build sparse aggregation operator (n_coarse × n_fine)
+    I = Ti[]
+    J = Ti[]
+    V = Tv[]
+    for i = 1:n
+        if aggregate[i] > 0
+            push!(I, aggregate[i])
+            push!(J, Ti(i))
+            push!(V, one(Tv))
+        end
+    end
+
+    AggOp = sparse(I, J, V, n_coarse, n)
+    return AggOp, c_map, n_coarse
+end
+
+function root_node_amg(A::TA,
+                    ::Type{Val{bs}}=Val{1};
+                    B = nothing,
+                    symmetry = HermitianSymmetry(),
+                    strength = SymmetricStrength(),
+                    CF = RS(),
+                    smooth = EnergyProlongation(),
+                    presmoother = GaussSeidel(),
+                    postsmoother = GaussSeidel(),
+                    improve_candidates = GaussSeidel(iter=4),
+                    max_levels = 10,
+                    max_coarse = 10,
+                    keep = false,
+                    verbose = false,
+                    coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
+
+    n = size(A, 1)
+    B = isnothing(B) ? ones(T, n) : copy(B)
+    @assert size(A, 1) == size(B, 1)
+
+    levels = Vector{Level{TA, TA, Adjoint{T, TA}}}()
+    bsr_flag = false
+    w = MultiLevelWorkspace(Val{bs}, eltype(A))
+    residual!(w, size(A, 1))
+
+    while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
+        A, B, bsr_flag = extend_hierarchy_rn!(levels, strength, CF, smooth,
+                                improve_candidates, keep, A, B,
+                                symmetry, bsr_flag, verbose)
+        size(A, 1) == 0 && break
+        coarse_x!(w, size(A, 1))
+        coarse_b!(w, size(A, 1))
+        residual!(w, size(A, 1))
+    end
+
+    cs = coarse_solver(A)
+    ml = MultiLevel(levels, A, cs, presmoother, postsmoother, w)
+
+    if verbose
+        @info ml
+    end
+
+    return ml
+end
+
+function extend_hierarchy_rn!(levels, strength, CF, smooth,
+                              improve_candidates, keep, A, B,
+                              symmetry, bsr_flag, verbose = false)
+
+    # Strength of connection
+    if symmetry isa HermitianSymmetry
+        S, _T = strength(A, bsr_flag)
+    else
+        S, _T = strength(adjoint(A), bsr_flag)
+    end
+
+    # C/F splitting
+    S_copy = copy(S)
+    remove_diag!(S_copy)
+    splitting = RS_CF_splitting(S_copy, adjoint(S_copy))
+
+    # Root-node aggregation
+    AggOp, c_map, n_coarse = root_node_aggregation(S, splitting)
+
+    # Improve candidates
+    b = zeros(size(A, 1), size(B, 2))
+    improve_candidates(A, B, b)
+    T_tent, B_coarse = fit_candidates(AggOp, B)
+
+    # Energy-minimizing prolongation
+    P = smooth(A, T_tent, S, B_coarse, splitting, c_map)
+    P, B_coarse = _drop_zero_columns(P, B_coarse)
+    R = construct_R(symmetry, P)
+
+    RAP = R * A * P
+
+    push!(levels, Level(A, P, R))
+
+    bsr_flag = true
+
+    RAP, B_coarse, bsr_flag
+end
diff --git a/test/nns_test.jl b/test/nns_test.jl
index 3e86d2d..b489c59 100644
--- a/test/nns_test.jl
+++ b/test/nns_test.jl
@@ -194,12 +194,17 @@ end
         @load "lin_elastic_2d.jld2" A b B
         A = SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, A.nzval)
 
-        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10;B=B)
         x_wonns, residuals_wonns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10)
+        println("SA without NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
+        @test !(A * x_wonns ≈ b)
 
-        println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
-        println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
+        x_nns, residuals_nns = solve(A, b, SmoothedAggregationAMG(), log=true, reltol=1e-10;B=B)
+        println("SA with NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
+        @test A * x_nns ≈ b
 
+        x_nns, residuals_nns = solve(A, b, RootNodeAMG(), log=true, reltol=1e-10, B=B)
+        println("Root-Node with NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
+        @test A * x_nns ≈ b
 
         #test QR factorization on linear elasticity
         aggregate = StandardAggregation()
@@ -210,8 +215,6 @@ end
         @test B ≈ Q * (Q' * B)
 
         # Check convergence
-        @test !(A * x_wonns ≈ b)
-        @test A * x_nns ≈ b
 
     end
 
@@ -236,6 +239,9 @@ end
         println("No NNS: final residual at iteration ", length(residuals_wonns), ": ", residuals_wonns[end])
         println("With NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
 
+        x_nns, residuals_nns = solve(A, b, RootNodeAMG(); aggregate=StandardAggregation(), log=true, reltol=1e-10, B=B, max_levels=2)
+        println("Root-Node with NNS: final residual at iteration ", length(residuals_nns), ": ", residuals_nns[end])
+        @test A * x_nns ≈ b
 
         # test QR factorization on bending beam
         # Aggregation
diff --git a/test/root_node_tests.jl b/test/root_node_tests.jl
new file mode 100644
index 0000000..f461fd0
--- /dev/null
+++ b/test/root_node_tests.jl
@@ -0,0 +1,75 @@
+using AlgebraicMultigrid
+import AlgebraicMultigrid as AMG
+
+## Energy Prolongation Tests
+@testset "Energy Prolongation" begin
+    @testset "Energy is reduced" begin
+        A = float.(poisson(50))
+        n = size(A, 1)
+        B = ones(n)
+
+        # Build strength, splitting, aggregation, tentative P
+        S, _ = AMG.SymmetricStrength()(A)
+        splitting = AMG.RS()(copy(S))
+        AggOp, c_map, nc = AMG.root_node_aggregation(S, splitting)
+        T, Bc = fit_candidates(AggOp, B)
+
+        # Compute energy before smoothing: trace(T' * A * T)
+        energy_before = tr(Matrix(T' * A * T))
+
+        # Apply energy prolongation
+        P = AMG.energy_prolongation_smoother(A, T, S, splitting)
+        energy_after = tr(Matrix(P' * A * P))
+
+        @test energy_after < energy_before
+    end
+
+    @testset "Null space is approximately preserved" begin
+        A = float.(poisson(50))
+        n = size(A, 1)
+        B = ones(n)
+
+        S, _ = AMG.SymmetricStrength()(A)
+        splitting = AMG.RS()(copy(S))
+        AggOp, c_map, nc = AMG.root_node_aggregation(S, splitting)
+        T, Bc = fit_candidates(AggOp, B)
+
+        P = AMG.energy_prolongation_smoother(A, T, S, splitting)
+
+        # C-point rows exactly preserve the null space
+        # F-point rows are smoothed, so overall null space is approximately preserved
+        @test norm(P * Bc - B) / norm(B) < 0.1
+    end
+
+    @testset "C-point rows are unchanged" begin
+        A = float.(poisson(50))
+        n = size(A, 1)
+        B = ones(n)
+
+        S, _ = AMG.SymmetricStrength()(A)
+        splitting = AMG.RS()(copy(S))
+        AggOp, c_map, nc = AMG.root_node_aggregation(S, splitting)
+        T, Bc = fit_candidates(AggOp, B)
+
+        P = AMG.energy_prolongation_smoother(A, T, S, splitting)
+
+        # C-point rows should be identical in T and P
+        for i = 1:n
+            if splitting[i] == AMG.C_NODE
+                for col = 1:nc
+                    @test AMG._get_entry(P, i, col) ≈ AMG._get_entry(T, i, col) atol=1e-14
+                end
+            end
+        end
+    end
+end
+
+## Root-Node AMG Solver Tests
+@testset "Root-Node AMG" begin
+    @testset "Scalar-valued problem" begin
+        A = float.(poisson(100))
+        b = A * ones(100)
+        x = solve(A, b, RootNodeAMG(), reltol=1e-10)
+        @test norm(A * x - b) / norm(b) < 1e-8
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 64fda95..391d185 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,6 +8,7 @@ using FileIO
 using Random: seed!
 
 include("sa_tests.jl")
+include("root_node_tests.jl")
 include("cycle_tests.jl")
 include("nns_test.jl")
 

From 2daee183e8dcfd793971e7d205748f343c07a346 Mon Sep 17 00:00:00 2001
From: termi-official <9196588+termi-official@users.noreply.github.com>
Date: Wed, 13 May 2026 21:16:11 +0200
Subject: [PATCH 2/2] Add precon builder

---
 src/AlgebraicMultigrid.jl |  2 +-
 src/precs.jl              | 20 ++++++++++++++++++++
 test/runtests.jl          |  3 +++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/AlgebraicMultigrid.jl b/src/AlgebraicMultigrid.jl
index 1138480..f46633e 100644
--- a/src/AlgebraicMultigrid.jl
+++ b/src/AlgebraicMultigrid.jl
@@ -47,6 +47,6 @@ include("preconditioner.jl")
 export aspreconditioner
 
 include("precs.jl")
-export SmoothedAggregationPreconBuilder, RugeStubenPreconBuilder
+export SmoothedAggregationPreconBuilder, RugeStubenPreconBuilder, RootNodePreconBuilder
 
 end # module
diff --git a/src/precs.jl b/src/precs.jl
index 59b683f..8526ac0 100644
--- a/src/precs.jl
+++ b/src/precs.jl
@@ -36,3 +36,23 @@ end
 function (b::RugeStubenPreconBuilder)(A::AbstractSparseMatrixCSC, p)
     return (aspreconditioner(ruge_stuben(SparseMatrixCSC(A), Val{b.blocksize}; b.kwargs...)), I)
 end
+
+
+"""
+   RootNodePreconBuilder(;blocksize=1, kwargs...)
+
+Return callable object constructing a  left algebraic multigrid preconditioner after Ruge & Stüben
+to be used with the `precs` API of LinearSolve.
+"""
+struct RootNodePreconBuilder{Tk}
+    blocksize::Int
+    kwargs::Tk
+end
+
+function RootNodePreconBuilder(; blocksize = 1, kwargs...)
+    return RootNodePreconBuilder(blocksize, kwargs)
+end
+
+function (b::RootNodePreconBuilder)(A::AbstractSparseMatrixCSC, p)
+    return (aspreconditioner(root_node_amg(SparseMatrixCSC(A), Val{b.blocksize}; b.kwargs...)), I)
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 391d185..46bd3f1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -348,6 +348,9 @@ for sz in [ (10,10), (20,20), (50,50)]
 
     strategy = KrylovJL_CG(precs = SmoothedAggregationPreconBuilder())
     @test solve(prob, strategy, atol=1.0e-14) ≈ u0 rtol = 1.0e-8
+
+    strategy = KrylovJL_CG(precs = RootNodePreconBuilder())
+    @test solve(prob, strategy, atol=1.0e-14) ≈ u0 rtol = 1.0e-7
 end
 
 end