Add coloring of basis functions

Project.toml

@@ -15,6 +15,7 @@ ExtendableSparse = "95c220a8-a1cf-11e9-0c77-dbfce5f500b3"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GraphsColoring = "ed38bb8a-6120-4024-8d1e-8789784f251b"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
@@ -23,6 +24,8 @@ LiftedMaps = "d22a30c1-52ac-4762-a8c9-5838452405e0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 NestedUnitRanges = "032820ab-dc03-4b49-91f4-7d58d4da98b3"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SauterSchwab3D = "0a13313b-1c00-422e-8263-562364ed9544"
 SauterSchwabQuadrature = "535c7bfe-2023-5c1d-b712-654ef9d93a38"
@@ -46,12 +49,15 @@ ExtendableSparse = "1.4"
 FFTW = "0.2.3, 1"
 FastGaussQuadrature = "0.3, 0.4, 0.5, 1"
 FillArrays = "0.11, 0.12, 0.13, 1"
+GraphsColoring = "0.2.0"
 Infiltrator = "1.8.2"
 IterativeSolvers = "0.9"
 Krylov = "0.10.1"
 LiftedMaps = "0.5.1"
 LinearMaps = "3.7 - 3.9, 3.11.2"
 NestedUnitRanges = "0.2.3"
+OhMyThreads = "0.8.3"
+ProgressMeter = "1.11.0"
 Requires = "1"
 SauterSchwab3D = "0.2"
 SauterSchwabQuadrature = "2.4.0"

src/BEAST.jl

@@ -1,5 +1,8 @@
 module BEAST
 
+using GraphsColoring
+using OhMyThreads
+using ProgressMeter
 using Distributed
 using LinearAlgebra
 using SharedArrays
@@ -141,6 +144,12 @@ using SparseArrays
 
 function convolve end
 
+function progressbar(workload, verbose; color=:white, kwargs...)
+    return Progress(
+        workload; barglyphs=BarGlyphs("[=> ]"), color=color, enabled=verbose, kwargs...
+    )
+end
+
 include("utils/polynomial.jl")
 include("utils/specialfns.jl")
 include("utils/combinatorics.jl")
@@ -201,6 +210,8 @@ include("bases/tensorbasis.jl")
 include("bases/composedbasis.jl")
 include("bases/local/localcomposedbasis.jl")
 
+include("coloring.jl")
+
 include("operator.jl")
 
 include("quadrature/strategies/quadstrat.jl")

src/coloring.jl

@@ -0,0 +1,59 @@
+"""
+    function conflicts(
+        space::Space;
+        addata=assemblydata(space),
+        kwargs...,
+    )
+
+Computes conflict indices for a Space. Two elements are in conflict if they are both
+part of the support of the same basis function.
+
+# Arguments
+
+  - `space::Space`: The space.
+  - `addata=assemblydata(space)`: The assembly data for the space (default: computed using `assemblydata`).
+  - `kwargs...`: Additional keyword arguments.
+
+# Returns
+
+A tuple containing:
+
+  - `eachindex(elements)`: The indices of the elements.
+  - `ConflictFunctor(conflictindices)`: A functor that maps element indices to conflict indices.
+  - `Base.OneTo(numfunctions(space))`: The indices of the conflicts.
+"""
+function GraphsColoring.conflicts(
+    space::Space;
+    addata=assemblydata(space),
+    kwargs...,
+)
+    elements, ad, _ = addata
+
+    conflictindices = Vector{Int}[Int[] for _ in eachindex(elements)]
+
+    for elementid in eachindex(elements)
+        for i in 1:length(ad[elementid])
+            for (functionid, a) in ad[elementid, i]
+                iszero(a) && continue
+                push!(conflictindices[elementid], functionid)
+            end
+        end
+    end
+
+    for i in eachindex(conflictindices)
+        conflictindices[i] = unique(conflictindices[i])
+    end
+
+    return eachindex(elements),
+    GraphsColoring.ConflictFunctor(conflictindices),
+    Base.OneTo(numfunctions(space))
+end
+
+function color(basis, ::Any; addata=assemblydata(basis), kwargs...)
+    return GraphsColoring.color(GraphsColoring.conflictmatrix(basis; addata=addata, kwargs...), WorkstreamDSATUR).colors
+end
+
+# for serial scheduling, just return all elements in one color
+function color(basis, ::SerialScheduler; addata=assemblydata(basis), kwargs...)
+    return [collect(1:length(first(addata)))]
+end

src/composedoperators/composedoperator.jl

@@ -47,9 +47,10 @@ integralop(a::CompSingleInt) = CompSingleKern(a.op1,a.kernel,a.op2)
 
 function assemble!(op::AbstractCompInt, tfs::Space, bfs::Space,
     store, threading = Threading{:multi};
-    quadstrat=defaultquadstrat)
+    quadstrat=defaultquadstrat,
+    kwargs...)
 
-    assemble!(integralop(op),op.tfunc(tfs),op.bfunc(bfs),store,threading;quadstrat)
+    assemble!(integralop(op),op.tfunc(tfs),op.bfunc(bfs),store,threading;quadstrat, kwargs...)
 
 end
 
@@ -138,7 +139,8 @@ defaultquadstrat(op::CompSingleInt,tfs::Space,bfs::Space) = SingleNumQStrat(6)
 ##### assembling the single integral 
 function assemble!(biop::CompSingleKern, tfs::Space, bfs::Space, store,
     threading::Type{Threading{:multi}};
-    quadstrat=defaultquadstrat(biop, tfs, bfs))
+    quadstrat=defaultquadstrat(biop, tfs, bfs),
+    kwargs...)
 
     return assemble_local_mixed!(biop, tfs, bfs, store; quadstrat)
 end

src/dyadicop.jl

@@ -35,7 +35,8 @@ end
 
 function assemble!(biop::DyadicOp, tfs::Space, bfs::Space, store,
     threading::Type{BEAST.Threading{:multi}};
-    quadstrat=defaultquadstrat(biop, tfs, bfs))
+    quadstrat=defaultquadstrat(biop, tfs, bfs),
+    kwargs...)
 
     u = assemble(biop.f, tfs)
     v = assemble(biop.g, bfs)

src/integralop.jl

@@ -172,9 +172,60 @@ function assemblechunk_body!(biop,
 end
 
 
+function assemblechunk_body_colored!(biop,
+        test_space, testad, testelementids,
+        trial_space, trialad, trialelementids,
+        qd, store, scheduler; quadstrat)
 
+    test_shapes = refspace(test_space)
+    trial_shapes = refspace(trial_space)
+
+    test_elements, test_assembly_data, test_cell_ptrs = testad
+    trial_elements, trial_assembly_data, trial_cell_ptrs = trialad
 
+    num_tshapes = numfunctions(test_shapes, domain(chart(geometry(test_space), first(geometry(test_space)))))
+    num_bshapes = numfunctions(trial_shapes, domain(chart(geometry(trial_space), first(geometry(trial_space)))))
 
+    @tasks for p in testelementids
+        @set scheduler = scheduler
+        @local zlocal = zeros(scalartype(biop, test_space, trial_space), num_tshapes, num_bshapes)
+        tcell, tptr = test_elements[p], test_cell_ptrs[p]
+
+        for q in trialelementids
+            bcell, bptr = trial_elements[q], trial_cell_ptrs[q]
+            fill!(zlocal, 0)
+            @inline qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, qd, quadstrat)
+            momintegrals!(zlocal, biop,
+                test_space,  tptr, tcell,
+                trial_space, bptr, bcell, qrule)
+            for j in 1:length(trial_assembly_data[q]), i in 1:length(test_assembly_data[p])
+                zij = zlocal[i,j]
+                for (n,b) in trial_assembly_data[q][j]
+                    iszero(b) && continue
+                    zb = zij*b
+                    for (m,a) in test_assembly_data[p][i]
+                        iszero(a) && continue
+                        store(a*zb, m, n)
+end end end end end end
+
+struct AssembleblockbodyFunctor{B,T1,T2,T3,T4,T5,T6,T7,T8,T9}
+    biop::B
+    tfs::T1
+    testelements::T2
+    testassemblydata::T3
+    bfs::T4
+    trialelements::T5
+    trialassemblydata::T6
+    quadraturedata::T7
+    zlocals::T8
+    quadstrat::T9
+end
+
+function (f::AssembleblockbodyFunctor)(testids, trialids, store)
+    assembleblock_body!(f.biop, f.tfs, testids, f.testelements, f.testassemblydata,
+        f.bfs, trialids, f.trialelements, f.trialassemblydata,
+        f.quadraturedata, f.zlocals, store; quadstrat=f.quadstrat)
+end
 
 function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space;
         quadstrat=defaultquadstrat(biop, tfs, bfs))
@@ -194,11 +245,18 @@ function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space;
         trial_elements, trial_assembly_data,
         quadrature_data, zlocals = assembleblock_primer(biop, tfs, bfs; quadstrat=qs)
 
-    return (test_ids, trial_ids, store) ->
-        assembleblock_body!(biop,
-            tfs, test_ids,   test_elements,  test_assembly_data,
-            bfs, trial_ids, trial_elements, trial_assembly_data,
-            quadrature_data, zlocals, store; quadstrat=qs)
+    return AssembleblockbodyFunctor(
+        biop,
+        tfs,
+        test_elements,
+        test_assembly_data,
+        bfs,
+        trial_elements,
+        trial_assembly_data,
+        quadrature_data,
+        zlocals,
+        qs,
+    )
 
     # if CompScienceMeshes.refines(tgeo, bgeo)
     #     return (test_ids, trial_ids, store) -> begin
@@ -267,10 +325,10 @@ function assembleblock_primer(biop, tfs, bfs;
 
     qd = quaddata(biop, tshapes, bshapes, test_elements, bsis_elements, quadstrat)
 
-    zlocals = Matrix{scalartype(biop, tfs, bfs)}[]
+    zlocals = Channel{Matrix{scalartype(biop, tfs, bfs)}}(2*Threads.nthreads())
 
-    for i in 1:Threads.nthreads()
-        push!(zlocals, zeros(scalartype(biop, tfs, bfs), num_tshapes, num_bshapes))
+    for _ in 1:2*Threads.nthreads()
+        put!(zlocals, zeros(scalartype(biop, tfs, bfs), num_tshapes, num_bshapes))
     end
 
     return test_elements, tad, bsis_elements, bad, qd, zlocals
@@ -306,27 +364,30 @@ function assembleblock_body!(biop::IntegralOperator,
     @assert maximum(active_test_el_ids) <= length(test_elements) "$(maximum(active_test_el_ids)), $(length(test_elements))"
     @assert maximum(active_trial_el_ids) <= length(bsis_elements) "$(maximum(active_trial_el_ids)), $(length(bsis_elements))"
 
+    zlocal = take!(zlocals)
     for p in active_test_el_ids
         tcell = test_elements[p]
         for q in active_trial_el_ids
             bcell = bsis_elements[q]
 
-            fill!(zlocals[Threads.threadid()], 0)
+            fill!(zlocal, 0)
             qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, quadrature_data, quadstrat)
-            momintegrals!(zlocals[Threads.threadid()], biop,
+            momintegrals!(zlocal, biop,
                 tfs, p, tcell,
                 bfs, q, bcell, qrule)
 
-            for j in 1 : size(zlocals[Threads.threadid()],2)
-                for i in 1 : size(zlocals[Threads.threadid()],1)
+            for j in 1 : size(zlocal,2)
+                for i in 1 : size(zlocal,1)
                     for (n,b) in trial_assembly_data[q,j]
                         n′ = get(trial_id_in_blk, n, 0)
                         n′ == 0 && continue
                         for (m,a) in test_assembly_data[p,i]
                             m′ = get(test_id_in_blk, m, 0)
                             m′ == 0 && continue
-                            store(a*zlocals[Threads.threadid()][i,j]*b, m′, n′)
-end end end end end end end
+                            store(a*zlocal[i,j]*b, m′, n′)
+    end end end end end end 
+    put!(zlocals, zlocal)
+end
 
 # function assembleblock_body_trial_refines_test!(biop::IntegralOperator,
 #         tfs, test_ids, test_elements, test_assembly_data,

src/localop.jl

@@ -59,7 +59,8 @@ end
 
 function assemble!(biop::LocalOperator, tfs::Space, bfs::Space, store,
         threading::Type{Threading{:multi}};
-        quadstrat=defaultquadstrat)
+        quadstrat=defaultquadstrat,
+        kwargs...)
 
         quadstrat = quadstrat(biop, tfs, bfs)
 
@@ -79,7 +80,8 @@ end
 
 function assemble!(biop::LocalOperator, tfs::Space, bfs::Space, store,
     threading::Type{Threading{:single}};
-    quadstrat=defaultquadstrat)
+    quadstrat=defaultquadstrat,
+    kwargs...)
 
     quadstrat = quadstrat(biop, tfs, bfs)