Symbolics: Keeping Array operations in derivatives

General Usage
1

Hello,
is there a way in julia symbolics keep the lazy array evaluation when computing derivatives. E.g. to get in 1 of the 2 examples below just J = 2*A as return

using Symbolics 
@variables A[1:5,1:5]
Av = Symbolics.variables(:a,1:5,1:5)
op = Symbolics.@arrayop () Av[i,j] * Av[j,i]
J1 = reshape(Symbolics.jacobian(op,Av) ,5,5)

op2 = Symbolics.@arrayop () A[i,j] * A[j,i]
J2 = reshape(Symbolics.jacobian(op2,A) ,5,5)

Instead I get:

J1 = 5×5 Matrix{Num}:
 2a₁ˏ₁  2a₂ˏ₁  2a₃ˏ₁  2a₄ˏ₁  2a₅ˏ₁
 2a₁ˏ₂  2a₂ˏ₂  2a₃ˏ₂  2a₄ˏ₂  2a₅ˏ₂
 2a₁ˏ₃  2a₂ˏ₃  2a₃ˏ₃  2a₄ˏ₃  2a₅ˏ₃
 2a₁ˏ₄  2a₂ˏ₄  2a₃ˏ₄  2a₄ˏ₄  2a₅ˏ₄
 2a₁ˏ₅  2a₂ˏ₅  2a₃ˏ₅  2a₄ˏ₅  2a₅ˏ₅
# and 

J2 = 5×5 Matrix{Num}:
 2A[1, 1]  2A[2, 1]  2A[3, 1]  2A[4, 1]  2A[5, 1]
 2A[1, 2]  2A[2, 2]  2A[3, 2]  2A[4, 2]  2A[5, 2]
 2A[1, 3]  2A[2, 3]  2A[3, 3]  2A[4, 3]  2A[5, 3]
 2A[1, 4]  2A[2, 4]  2A[3, 4]  2A[4, 4]  2A[5, 4]
 2A[1, 5]  2A[2, 5]  2A[3, 5]  2A[4, 5]  2A[5, 5]
2

Right now, no. We do not support symbolic derivatives of symbolic arrays. But it is a GSoC project we are looking to do.

3 Likes
3

Thank you for your answering.
I think this would be a really nice feature to automate numerical computations for FEM.
Many material laws are written in the form of some potential, where the gradient is then the right-hand side of the system and the Hessian the element tangent.
The performance of currently available tools is 2x-4x slower than the handwritten derivatives.

1 Like
4

Did you try automatic differentiation with e.g. Enzyme.jl for these problems?

5

It sounds like you’re looking for automatic differentiation, not symbolic differentiation?

6

I tested it sometime ago, and it was either slower or crashed completly. I will later post a minimal working example later.

1 Like
7

I considered both. I tried to

  • Compute the derivative once and compile a fast function (cse = true) with build_function
  • Mix symbolic and automatic differentiation
  • automatic differentiation only.

There is already a pretty nice software tool for this but its based on mathematica and C and its not open source https://www.wolfram.com/products/applications/acegen/.
A julia version with similar performance would be really nice.

8

FastDifferentiation.jl? Though I am not sure if that handles array ops directly either.

9

It sounds like your problem was sparse, did you use the sparse differentiation utilities of DifferentiationInterface.jl? This ecosystem has matured quite a lot recently

10

It goes in the right direction. But it has the following 2 limitations:

The two biggest limitations of FD are no support for conditionals involving FD variables and code size that scales with number of operations.

Thus, I tried do use Symbolics for this since ifelse is already supported.

11

The problem is not sparse. I built the element tangent and right-hand side by computing the Hessian. These two components are dense. Only the global tangent matrix is sparse, but this is handled by the assembler from Ferrite.jl.

1 Like
12

If you can share a complete MWE I will take a crack at differentiating it in FastDifferentiation. I added conditionals to the system a while back, although not yet the ability to differentiate through conditionals. If your code really needs this capability it would be motivation for me to add the feature.

1 Like
13

I’ve created two examples that compute the stiffness matrix for a hyperelastic material. The first benchmark successfully matches the performance of the analytical gradient approach. In this implementation, I precompute the shape function values and then calculate the Hessian.

However, I’d like to develop a more general form where I can simply input my deformation vector (called ‘x’ in my example) and compute the Hessian afterward. This is implemented in the second example, but its performance is several times slower.

For those interested: At the end of the response, I’ve included two complete files where both approaches are used to compute the full deformation of the body. Both implementations use the Ferrite.jl tutorial as their foundation.

Version 1

using Ferrite, Tensors, TimerOutputs, ProgressMeter, IterativeSolvers
using StaticArrays
using Enzyme, Chairmarks


function StaticArrays.setindex(t::AbstractTensor, v, i::Int...)
    new_data = setindex(t.data,v,i...)
    return typeof(t)(new_data)
end
function Tensors.basevec(::Type{V}) where V 
    _cache = zero(V)
    return ntuple(
            i -> setindex(_cache,one(eltype(V)),i),Val(length(_cache)))
end

function _grad!(f::F,g,x,pars) where F <:Function
    g .= Enzyme.autodiff(Reverse,f,Active(x),Const(pars))[1][1]
    nothing
end


@generated function _grad_and_hessian(f,
            x::Tensor{2,N,T,L},pars) where {N,L,T}
    xbb = basevec(Tensor{2,N,Float64})
    g = zero(MMatrix{N,N})
    hh = ntuple(_ -> zero(MMatrix{N,N}),Val(L))

    return quote 
        fill!.($hh,0.0)
        Enzyme.autodiff(Forward,_grad!,Const(f),
        BatchDuplicated($g,$hh),BatchDuplicatedNoNeed(x,$xbb),Const(pars))
        
        return $g,$hh
    end
end
function grad_and_hessian(f,x::Tensor{2,N,T,L},
            pars = nothing) where {N,L,T}
    g,hh = _grad_and_hessian(f,x,pars)
    H = Tensor{4,N}(((hh...)...,))
    G = Tensor{2,N}((g...,))
    return G,H
end

struct NeoHooke
    μ::Float64
    λ::Float64
end

function Ψ(∇ux,mp::NeoHooke)
    μ = mp.μ
    λ = mp.λ
    F = one(∇ux) + ∇ux
    C = tdot(F)
    Ic = tr(C)
    J = sqrt(det(C))
    return μ / 2 * (Ic - 3 - 2log(J)) + λ/2 * (J - 1)^2
end

x = rand(Tensor{2,3,Float64})
G,H = grad_and_hessian(Ψ,x,NeoHooke(1.0,1.0))
mp = NeoHooke(1.0,1.0)
@b grad_and_hessian($Ψ,$x,$mp)

Version 2

using Enzyme
using StaticArrays

using Enzyme
using Chairmarks
using Ferrite
using Ferrite.Tensors: Vec, Tensor, eᵢ, tdot
using StaticArrays: SVector, setindex, SMatrix
using Bumper

struct NeoHooke
    μ::Float64
    λ::Float64
end

function Ψ(∇ux,mp::NeoHooke)
    μ = mp.μ
    λ = mp.λ
    F = one(∇ux) + ∇ux
    C = tdot(F)
    Ic = tr(C)
    J = sqrt(det(C))
    return μ / 2 * (Ic - 3 - 2log(J)) + λ/2 * (J - 1)^2
end

function ham_fun(ue::AbstractVector{T},pars::P) where {T,P}
    _,cv,mp = pars
    
    μ = mp.μ
    λ = mp.λ
    int = zero(T)
    for qp in 1:getnquadpoints(cv)
        dΩ = getdetJdV(cv, qp)
        ∇ux = function_gradient(cv, qp, ue)
        F = one(∇ux) + ∇ux
        C = tdot(F)
        Ic = tr(C)
        J = sqrt(det(C))
        int += (μ / 2 * (Ic - 3 - 2log(J)) + λ/2 * (J - 1)^2)*dΩ 
    end
    int
end

function full_grad!(f::F,g,x,pars) where {F <:Function}
    Enzyme.autodiff(Reverse,f,Duplicated(x,g),Const(pars))
end




function _grad_and_hessian!(grad,hess,
    f::F,_x::AbstractVector{T},pars::P,vl::Val) where {F,T,P}

    #INFO: Bumper Version
    @no_escape begin
        L = length(_x)
        x = @alloc(T, L)
        copyto!(x,_x)
        

        xbb = ntuple(i -> begin 
                v = @alloc(T,L)
                v .= zero(T)
                v[i] = one(eltype(x))
                v 
            end,vl)
  
        g = @alloc(T,L)
        g .= zero(T)
        hh = ntuple(_ -> begin
                v = @alloc(T,L)
                v .= zero(T)
                v 
        end,vl)

        Enzyme.autodiff(Forward,full_grad!,Const(f),
        BatchDuplicated(g,hh),BatchDuplicatedNoNeed(x,xbb),Const(pars))
        
        copyto!(grad,g)
        for i in axes(hess,1), j in axes(hess,2)
            hess[i,j] = hh[j][i]
        end
        nothing
    end

    # #INFO: Default Version (Slower)
    # x = _x   

    # xbb = ntuple(i -> begin 
    #         v = zero(x)
    #         v[i] = one(eltype(x))
    #         v 
    #     end,vl)

    # g = grad
    # hh = ntuple(i -> zero(x),vl)

    # Enzyme.autodiff(Forward,full_grad!,Const(f),
    #     BatchDuplicated(g,hh),BatchDuplicatedNoNeed(x,xbb),Const(pars))
    
    # for i in axes(hess,1), j in axes(hess,2)
    #     hess[i,j] = hh[j][i]
    # end

end

function grad_and_hessian(f,x,pars) 
    n = length(x)
    grad = zeros(n)
    hess = zeros(n,n)
    _grad_and_hessian!(grad,hess,f,x,pars,Val(n))
    return grad,hess
end


x = rand(12)

# Generate a grid
N = 20
L = 1.0
left = zero(Vec{3})
right = L * ones(Vec{3})
grid = generate_grid(Tetrahedron, (N, N, N), left, right)

# Material parameters
E = 10.0
ν = 0.3
μ = E / (2(1 + ν))
λ = (E * ν) / ((1 + ν) * (1 - 2ν))
mp = NeoHooke(μ, λ)

# Finite element base
ip = Lagrange{RefTetrahedron, 1}()^3
qr = QuadratureRule{RefTetrahedron}(1)
qr_facet = FacetQuadratureRule{RefTetrahedron}(1)
cv = CellValues(qr, ip)
cell = first(CellIterator(grid))
reinit!(cv, cell)
p = (cell,cv,mp)
grad_and_hessian(ham_fun,x,p)
@b grad_and_hessian($ham_fun,$x,$p)

ADfull_Full.jl (5.2 KB)
ADhalf_Full.jl (6.8 KB)