Viscosity jacobian

[albert-de-montserrat]

Small prototype

using GeoParams, ForwardDiff, StaticArrays
import GeoParams.MaterialParameters.ConstitutiveRelationships.lambda

function petsc_strain_rate!(εij, ∂u∂xi) 
    # indices may need to be changed accordingly, im being lazy here
    idx = LinearIndices(@SMatrix rand(2,2)) 
    for j in axes(idx, 2)
        for i in axes(idx, 1)
            εij[idx[i,j]] = 0.5 * (∂u∂xi[idx[i,j]] + ∂u∂xi[idx[j,i]])
        end
    end
end

function foo(τij, εij, c)
    # indices to be changed here too
    εII  = second_invariant(εij[1], εij[2], εij[3])
    τII  = second_invariant(τij[1], τij[2], τij[3])
    P    = 1e6
    Pf   = 0.5e6
    K    = 1.0
    args = (K=K, P=P, τII=τII, T=900.0, d=100e-6, τII_old=1e6, dt=1e8)
    
    η      = compute_viscosity_εII(c, εII, args)
    F      = compute_yieldfunction(c, args)
    λ      = lambda(F, c, η, args.K, args.dt)
    τII_pl = 2 * η * (εII - λ * ∂Q∂τII(c, τII, args)) 
    η_vep  = τII_pl * 0.5 * inv.(εij)
    @. τij = 2 * η_vep * εij
end

# Define a range of rheological components
v = LinearViscous(; η=1e22)
el = SetConstantElasticity(; G=5e10, Kb=1e11)
pl = DruckerPrager_regularised(; C=1e6, ϕ = 30)  
c  = CompositeRheology(v, el, pl)
τij = @MVector [1e9, -1e9, 0.0, 0.0] # dont know petsc order here
εij = similar(τij)
∂u∂xi = (@MVector rand(4)) .* 1e-11
petsc_strain_rate!(εij, ∂u∂xi)
using GeoParams, ForwardDiff, StaticArrays
import GeoParams.MaterialParameters.ConstitutiveRelationships.lambda

function petsc_strain_rate!(εij, ∂u∂xi) 
    # indices may need to be changed accordingly, im being lazy here
    idx = LinearIndices(@SMatrix rand(2,2)) 
    for j in axes(idx, 2)
        for i in axes(idx, 1)
            εij[idx[i,j]] = 0.5 * (∂u∂xi[idx[i,j]] + ∂u∂xi[idx[j,i]])
        end
    end
end

function foo(τij, εij, c)
    # indices to be changed here too
    εII  = second_invariant(εij[1], εij[2], εij[3])
    τII  = second_invariant(τij[1], τij[2], τij[3])
    P    = 1e6
    Pf   = 0.5e6
    K    = 1.0
    args = (K=K, P=P, τII=τII, T=900.0, d=100e-6, τII_old=1e6, dt=1e8)
    
    η      = compute_viscosity_εII(c, εII, args)
    F      = compute_yieldfunction(c, args)
    λ      = lambda(F, c, η, args.K, args.dt)
    τII_pl = 2 * η * (εII - λ * ∂Q∂τII(c, τII, args)) 
    η_vep  = τII_pl * 0.5 * inv.(εij)
    @. τij = 2 * η_vep * εij
end

# Define a range of rheological components
v = LinearViscous(; η=1e22)
el = SetConstantElasticity(; G=5e10, Kb=1e11)
pl = DruckerPrager_regularised(; C=1e6, ϕ = 30)  
c  = CompositeRheology(v, el, pl)
τij = @MVector [1e9, -1e9, 0.0, 0.0] # dont know petsc order here
εij = similar(τij)
∂u∂xi = (@MVector rand(4)) .* 1e-11
petsc_strain_rate!(εij, ∂u∂xi)
Jacobian = ForwardDiff.jacobian((τij, εij) -> foo(τij, εij, c), τij, εij)

[boriskaus]

lgtm but tests need fixing

[boriskaus]

also note that the petsc strainrate vector involves 4 components in 2D (3D=9)

[albert-de-montserrat]

The length of the strain rate is actually 4, but the GeoParams function to compute the invariant assumes tensor symmetry, thats why I take only 3 components.

[albert-de-montserrat]

julia> @btime ForwardDiff.jacobian(S_closed, τij, ∂u∂xi);
  445.500 ns (6 allocations: 1.02 KiB)

julia> @btime jacobian(Forward, ∂u∂xi -> foo(τij, ∂u∂xi, c), ∂u∂xi);
  3.500 μs (55 allocations: 3.00 KiB) # enzyme

FD seems the best option here.

[boriskaus]

So this kind of works for me with the PETSc code; yet, please note that the compute_viscosity_εII routine does not seem to take viscoelasticity into account properly.

A workaround is this:

 τII =  compute_τII(c, εII + eps(), args);
 η_vep = τII/(εII + eps())/2.0
 @. τij = 2 * η_vep * εij

I guess it is better, though, to update compute_viscosity_εII to take elasticity into account.

[albert-de-montserrat]

My bad, that function already exists as in here

[boriskaus]

Don't think that works as it should. Also can't handle this case:

v = LinearViscous(; η=η)
vlow = LinearViscous(; η=1e-3)      # lower cutoff
el = SetConstantElasticity(; G=G, ν=0.5)
pl = DruckerPrager_regularised(; C=0.15, ϕ = 0, η_vp = 1e-3)

comp = Parallel(CompositeRheology(v,el, pl),vlow)

[vchuravy]

FD seems the best option here.

What is a full reproducer for this? Taking the code in the example I get:

julia> Jacobian = ForwardDiff.jacobian((τij, εij) -> foo(τij, ∂u∂xi, c), τij, ∂u∂xi)
ERROR: UndefVarError: `lambda` not defined

[albert-de-montserrat]

I updated the MWE above, it required an import. It should run now. Thanks for taking a look!

[albert-de-montserrat]

Don't think that works as it should. Also can't handle this case:

v = LinearViscous(; η=η)
vlow = LinearViscous(; η=1e-3)      # lower cutoff
el = SetConstantElasticity(; G=G, ν=0.5)
pl = DruckerPrager_regularised(; C=0.15, ϕ = 0, η_vp = 1e-3)

comp = Parallel(CompositeRheology(v,el, pl),vlow)

Both compute_elastoviscosity_εII and compute_viscosity_εII are meant to work only for the common serial case (i.e. rheology = CompositeRheology(v,el, pl)), and fully ignores plasticity. There was a bug in compute_elastoviscosity_εII for composite rheology, I pushed the patch.

[boriskaus]

Don't think that works as it should. Also can't handle this case:

v = LinearViscous(; η=η)
vlow = LinearViscous(; η=1e-3)      # lower cutoff
el = SetConstantElasticity(; G=G, ν=0.5)
pl = DruckerPrager_regularised(; C=0.15, ϕ = 0, η_vp = 1e-3)

comp = Parallel(CompositeRheology(v,el, pl),vlow)

Both compute_elastoviscosity_εII and compute_viscosity_εII are meant to work only for the common serial case only (i.e. rheology = CompositeRheology(v,el, pl)), and fully ignores plasticity. There was a bug in compute_elastoviscosity_εII for composite rheology, I pushed the patch.

can this be extended? this is the way to add a lower viscosity cutoff in a differential manner