Q: multithreading - threads

[Audrius-St]

Should have posted this in Symbolics.jl issues, which I've now done. Sorry.

Hello,

I would like to use multithreading in Symbolics.jl to accelerate a symbolic computation. In an application, expand_derivatives(), becomes a bottleneck for large symbolic expressions with terms with higher order partial derivatives, so I would like to break up the long expression and perform the calculation in parts in a multithreaded for loop.

The MNWE, (a continuation from Issue #687 | "Interfacing - building an expression tree to use custom rules") is

# test_symbolics_main.jl

using Symbolics, SymbolicUtils, SymbolicUtils.Rewriters

begin

    @variables x, t
    @variables ϕ(x, t)
    @variables a, d

    Dx = Differential(x)
    Dt = Differential(t)

    # A rewriter to everse the partial derivatives order w.r.t. t and x
    # such that the partial derivative Dt is on the right:
    # Dt(Dx) +> Dx(Dt)
    rewriter = Prewalk(Chain([@rule Dt(Dx(~f)) => Dx(Dt(~f))]))

    # Unwrap the expression from Num to expose the SymbolicUtils expression tree for rewriting
    # and rewrap to Num
    reverse_DtDx(x) = Symbolics.wrap(rewriter(Symbolics.unwrap(x)))

    # Test PDF: scalar reaction diffusion equation
    reaction_diffusion = Dict(Dt(ϕ) => d*(Dx^2)(ϕ) - a*ϕ^2) # Eq 1

    out_file = "test_symbolics_main.txt"
    io = open(out_file, "w")

    # 0th order expansion
    γ₀ = ϕ 
    write(io, string("γ₀ = ", γ₀, "\n\n"))
    println(string("γ₀ = ", γ₀, "\n\n"))

    # 1st order expansion
    γ₁ = Dt(γ₀) 

    # Substitute the partial derivative Dt() with the reaction-diffusion pde
    γ₁ = substitute(γ₁, reaction_diffusion)
    write(io, string("γ₁ = ", γ₁, "\n\n"))
    println(string("γ₁ = ", γ₁, "\n\n"))

    # 2nd order expansion
    γ₂ = Dt(γ₁)

    # Reverse the partial derivatives order w.r.t. t and x
    # such that the partial derivative Dt is on the right:
    # Dt(Dx) => Dx(Dt)
    γ₂ = substitute(γ₂, reaction_diffusion) 
    γ₂ = expand_derivatives(γ₂) 
    γ₂ = reverse_DtDx(γ₂) 
    write(io, string("γ₂ = ", γ₂, "\n"))
    println(string("γ₂ = ", γ₂, "\n"))

    # Substitute the partial derivative Dt() with the reaction-diffusion equation
    γ₂ = substitute(γ₂, reaction_diffusion) 
    γ₂ = expand_derivatives(γ₂)
    γ₂ = expand(γ₂)
    write(io, string("γ₂ = ", γ₂, "\n\n"))
    println(string("γ₂ = ", γ₂, "\n\n"))

    # 3rd order expansion
    γ₃ = Dt(γ₂) 

    # Substitute the partial derivative Dt() with the reaction-diffusion equation
    γ₃ = substitute(γ₃, reaction_diffusion) 
    γ₃ = expand_derivatives(γ₃)

    # Reverse the partial derivatives order w.r.t. t and x
    γ₃ = reverse_DtDx(γ₃)
    write(io, string("γ₃ = ", γ₃, "\n"))
    println(string("γ₃ = ", γ₃, "\n"))

    # Substitute the partial derivative Dt() with the reaction-diffusion equation
    γ₃ = substitute(γ₃, reaction_diffusion)
    γ₃ = expand_derivatives(γ₃) 
    γ₃ = expand(γ₃)
    write(io, string("γ₃ = ", γ₃, "\n\n"))
    println(string("γ₃ = ", γ₃, "\n\n"))

    # Expose the γ₃ tree
    γ₃_tree = Symbolics.value(γ₃)

    # Extract the γ₃ arguments - terms and rewrap with Num
    γ₃_args = Symbolics.wrap(arguments(γ₃_tree))

    # Issue: "TypeError: in Atomic, in T, expected 
    # T<:Union{
    # Bool, Float16, Float32, Float64, Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}, 
    # got Type{Num}"
    γ₄ = Threads.Atomic{Num}(0)

    # 4th order expansion
    # Multi-threaded for performance
    Threads.@threads for i in eachindex(γ₃_args) 
        γ₄ᵢ = Dt(γ₃_args[i])

        # Substitute the partial derivative Dt() with the reaction-diffusion equation
        γ₄ᵢ = substitute(γ₄ᵢ, reaction_diffusion) 
        γ₄ᵢ = expand_derivatives(γ₄ᵢ) 

        # Reverse the partial derivatives order w.r.t. t and x
        γ₄ᵢ = reverse_DtDx(γ₄ᵢ)

        # Substitute the partial derivative Dt() with the reaction-diffusion equation
        γ₄ᵢ = substitute(γ₄ᵢ, reaction_diffusion) 
        γ₄ᵢ = expand_derivatives(γ₄ᵢ)
        γ₄ᵢ = expand(γ₄ᵢ) 

        Threads.atomic_add!(γ₄, γ₄ᵢ)

    end

    write(io, string("γ₄ = ", γ₄, "\n\n"))
    println(string("γ₄ = ", γ₄, "\n\n"))

    # Higher order expansions . . .

    close(io)

end

which returns

TypeError: in Atomic, in T, expected T<:Union{Bool, Float16, Float32, Float64, Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}, got Type{Num}

for the line γ₄ = Threads.Atomic{Num}(0)