Size mismatch in dual number handling

[ChrisRackauckas]

MWE:

using NaNMath
using LinearAlgebra: Diagonal
using ForwardDiff
using OrdinaryDiffEqRosenbrock
function mwe_dae!(du, u, params, t)

    (; k_f, C, N, S, Rs, ) = params
    Ap = 3.14e-4
    Av = 3.8e-4
    Tsh = t->233.15 + 1800*t 
    Kshf = 17.344
    k_ice = 2.4
    hf0 = .01315
    R = 8.3145
    ϕ = 0.96
    pch = 13.33
    ΔHsub = 2.838e6
    Mw = 18
    
    hf, Tf, μ0, μ2 = u
    hd = hf0 - hf
    μ1 = hd * Ap
    Qshf = Av*Kshf*(Tsh(t) - Tf)
    Tsub = Tf - Qshf/k_ice/Ap*hf
    k = k_f(Tsub)

    # exit hatch for bad parameter guesses or solver iterations
    if Tf < 0 # || Tsub < 0u"K"
        du .= NaN
        return
    end

    porefac = μ1^2*μ0/μ2^3 
    Rp = C*NaNMath.sqrt(R*Tsub/Mw)/ϕ*hd * sqrt(sqrt(NaNMath.sqrt(porefac))) 
    if !isfinite(Rp)
        Rp = 0.0 # No pores to resist flow
    end

    Rsp = Rs + Rp
    psub = 3.597e12 * exp(-6149.1/Tsub)
    dmdt = - Ap*(psub-pch)/Rsp
    Qsub = dmdt*ΔHsub

    dhf_dt = min(0.0, dmdt/ϕ/Ap) # Cap dhf_dt at 0: no desublimation
    dμ0_dt = -k/2*μ0^2 - N*Ap*dhf_dt
    dμ2_dt = k*μ1^2 - S*Ap*dhf_dt # dhf_dt < 0 for drying proceeding

    du[1] = dhf_dt
    du[2] = Qsub + Qshf
    du[3] = dμ0_dt
    du[4] = dμ2_dt

    return 
end

mwe_daef = ODEFunction(mwe_dae!, mass_matrix=Diagonal([1.0, 0.0, 1.0, 1.0]))
u0 = [0.01315, 233.15, 0, 0]
tspan = (1439.28, 7200)

function mwe_loss(x) 
    badps = (N = x*2.5055e14, S = 7.344991e-14, C = 0.02147, Rs = 1.9017e-5, k_f = x->8.38e-16)
    solve(ODEProblem(mwe_daef, u0, tspan, badps), Rodas4())
end
mwe_loss(1.0)
ForwardDiff.derivative(mwe_loss, 1.0)

[jClugstor]

Issue was that some parts of the dual cache were being set up with the size of b instead of u. This was causing issues because the A and b values in the LinearProblem were:

infil> A
2×1 Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}, Float64, 1}}:
 Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}}(-1.3310918624566178e13,-0.0)
 Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}}(1.0e-8,0.0)

infil> b
2-element Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}, Float64, 1}}:
 Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}}(-1.720721691078986e20,-0.0)
 Dual{ForwardDiff.Tag{typeof(mwe_loss), Float64}}(233.15,0.0)

Isn't this an overdetermined system anyways? Do we need to give the derivative of the least squares answer in this case or something?

This should allow it to run: #842

But the results do not seem correct:

julia> A = reshape([ForwardDiff.Dual(2.0, 1.0), ForwardDiff.Dual(3.0, 1.0)], 2, 1)  # 2×1 matrix
2×1 Matrix{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(2.0,1.0)
 Dual{Nothing}(3.0,1.0)

julia> b = [ForwardDiff.Dual(5.0, 1.0), ForwardDiff.Dual(8.0, 9.0)]
2-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(5.0,1.0)
 Dual{Nothing}(8.0,9.0)

julia> backslash_u = A\b
1-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(2.615384615384616,1.2189349112426036)

julia> prob = LinearProblem(A,b)
LinearProblem. In-place: true
b: 2-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(5.0,1.0)
 Dual{Nothing}(8.0,9.0)

julia> solve(prob)
retcode: Success
u: 1-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(2.615384615384615,1.2248520710059172)

julia> sol = solve(prob)
retcode: Success
u: 1-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(2.615384615384615,1.2248520710059172)

julia> A*backslash_u
2-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(5.230769230769232,5.053254437869823)
 Dual{Nothing}(7.846153846153848,6.272189349112427)

julia> A*sol.u
2-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(5.23076923076923,5.065088757396449)
 Dual{Nothing}(7.846153846153845,6.289940828402367)

julia> b
2-element Vector{ForwardDiff.Dual{Nothing, Float64, 1}}:
 Dual{Nothing}(5.0,1.0)
 Dual{Nothing}(8.0,9.0)

it also seems that backslash gives incorrect answers.

[ChrisRackauckas]

backslash is fine here, it's the QR result.

[jClugstor]

My fix doesn't quite match the backslash result in the Duals. I'll need to fix that up and add tests.