add B to docstring

Author: Abdelrahman912

src/aggregation.jl

@@ -14,7 +14,7 @@ function smoothed_aggregation(A::TA, _B = nothing,
                         coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
 
     n = size(A, 1)
-    B = isnothing(_B) ? ones(T,n,1) : copy(_B)
+    B = isnothing(_B) ? ones(T,n) : copy(_B)
     @assert size(A, 1) == size(B, 1)
 
     #=max_levels, max_coarse, strength =

src/multilevel.jl

@@ -289,6 +289,38 @@ abstract type AMGAlg end
 struct RugeStubenAMG  <: AMGAlg end
 
 """
+    SmoothedAggregationAMG(B=nothing)
+
+Smoothed Aggregation AMG (Algebraic Multigrid) algorithm configuration.
+
+# Arguments
+- `B::Union{AbstractArray, Nothing}`: The **near null space** in SA-AMG represents the low energy that cannot be attenuated by relaxtion, and thus needs to be perserved across the coarse grid. 
+
+    - `B` can be:
+        - a `Vector` (e.g., for scalar PDEs),
+        - a `Matrix` (e.g., for vector PDEs or systems with multiple equations),
+        - or `nothing`.
+
+# Notes
+If `B` is set to `nothing`, it will be internally defaulted to `B = ones(T, n)`, where `T = eltype(A)` and `n = size(A, 1)`.
+
+# Examples
+
+**Poisson equation (scalar PDE):**
+```julia
+n = size(A, 1)
+B = ones(Float64, n)
+amg = SmoothedAggregationAMG(B)
+```
+
+**Linear elasticity equation in 2d (vector PDE):**
+```julia
+n = size(A, 1)          # Ndof = 2 * number of nodes
+B = zeros(Float64, n, 3) 
+B[1:2:end,   :] = [1.0, 0.0, -y]
+B[2:2:end, :] = [0.0, 1.0,  x]
+amg = SmoothedAggregationAMG(B)
+```
 """
 struct SmoothedAggregationAMG  <: AMGAlg
     B::Union{<:AbstractArray,Nothing} # `B` can be `Vector`, `Matrix`, or `nothing`

test/nns_test.jl

@@ -1,3 +1,25 @@
+## Test `B` as an argument for `SmoothedAggregationAMG`
+@testset "Different `B` as an argument for `SmoothedAggregationAMG` " begin
+    A = poisson(100); 
+    b = rand(100);
+    n = size(A,1)
+    T = eltype(A)
+
+    #1. pass `B`` as nothing (Default case)
+    x_nothing = solve(A, b, SmoothedAggregationAMG(), maxiter = 1, abstol = 1e-6)
+
+    #2. pass `B` as vector
+    B = ones(T,n)
+    x_vec = solve(A, b, SmoothedAggregationAMG(B), maxiter = 1, abstol = 1e-6)
+    @test x_vec ≈ x_nothing
+
+    #3. pass `B` as matrix
+    B = ones(T,n,1)
+    x_mat = solve(A, b, SmoothedAggregationAMG(B), maxiter = 1, abstol = 1e-6)
+    @test x_mat ≈ x_nothing
+end
+
+
 ## Test QR factorization
 @testset "fit_candidates unit test cases" begin
     cases = Vector{Tuple{SparseMatrixCSC{Float64,Int},Matrix{Float64}}}()