Towards AD rheology calculations

[albert-de-montserrat]

When non-linear iterations are required to during the local iterations to compute stress/strain we do require to evaluate εII and ∂εII∂τII (or same for the stress tensor). Currently we do compute the tensor and its derivative in two separate operations, while with AD we can evaluate the functions and its derivative with a single function call. The latter yields a more efficient algorithm with no loss of accuracy. Proof of concept:

function compute_τII_AD(
    v::CompositeRheology{T,N,
        0,is_parallel,
        0,is_plastic,
        Nvol,is_vol,
        false},
    εII::_T, 
    args; 
    tol=1e-6
) where {N, T, _T, is_parallel, is_plastic, Nvol, is_vol}

    # Initial guess
    τII = compute_τII_harmonic(v, εII, args)

    # Local Iterations
    iter = 0
    ϵ = 2.0 * tol
    τII_prev = τII
    while ϵ > tol
        iter += 1
        εII_n, dεII_dτII_n = compute_εII_AD_all(v, τII, args)

        τII = muladd(εII - εII_n, inv(dεII_dτII_n), τII)

        ϵ = abs(τII - τII_prev) * inv(τII)
        τII_prev = τII

    end

    return τII
end

v1 = SetDiffusionCreep("Dry Anorthite | Rybacki et al. (2006)")
v2 = SetDislocationCreep("Dry Anorthite | Rybacki et al. (2006)")
e1 = ConstantElasticity()           # elasticity
c = CompositeRheology(v1, v2, e1)   # with elasticity
args = (T=900.0, d=100e-6, τII_old=1e6, dt=1e8)
εII, τII = 2e-15, 2e6


julia> compute_τII_AD(c, εII, args) == compute_τII(c, εII, args)
true

julia> @btime compute_τII($c, $εII, $args)
  203.500 ns (0 allocations: 0 bytes)

julia> @btime compute_τII_AD($c, $εII, $args)
  162.183 ns (0 allocations: 0 bytes)

[boriskaus]

this is great - I attempted using AD to do more complicated Composite & Parallel rheologies, which was failing. Perhaps this can do the job?

[albert-de-montserrat]

I got the composites with n-Parallel bodies in series working already, nested Parallels seems a bit trickier to handle recursively - I'm not sure how of a priority that kind of rheology is anyway, the "standard" we previously implemented seems to work for those cases, but not very well anyway, I see the compiler giving up on inlining and resulting in allocations...

[boriskaus]

That's great!
The one thing that is important is that case where we use Parallel(CompositeRheology(...), LinearViscous()) where the parallel element is used as a lower viscosity cutoff and the CompositeRheology element may have a Parallel element, for example to regularise plasticity.