rearrange, make sure that dualcache works

Author: jClugstor

ext/LinearSolveForwardDiffExt.jl

@@ -34,17 +34,15 @@ const DualAbstractLinearProblem = Union{
     DualLinearProblem, DualALinearProblem, DualBLinearProblem}
 
 LinearSolve.@concrete mutable struct DualLinearCache
-    cache
+    linear_cache
     prob
     alg
-    A
-    b
     partials_A
     partials_b
 end
 
 function linearsolve_forwarddiff_solve(cache::DualLinearCache, alg, args...; kwargs...)
-    sol = solve!(cache.cache, alg, args...; kwargs...)
+    sol = solve!(cache.linear_cache, alg, args...; kwargs...)
     uu = sol.u
 
     # Solves Dual partials separately 
@@ -53,7 +51,7 @@ function linearsolve_forwarddiff_solve(cache::DualLinearCache, alg, args...; kwa
 
     rhs_list = xp_linsolve_rhs(uu, ∂_A, ∂_b)
 
-    new_A = nodual_value(cache.prob.A)
+    new_A = nodual_value(cache.A)
     partial_prob = LinearProblem(new_A, rhs_list[1])
     partial_cache = init(partial_prob, alg, args...; kwargs...)
 
@@ -67,44 +65,6 @@ function linearsolve_forwarddiff_solve(cache::DualLinearCache, alg, args...; kwa
     sol, partial_sols
 end
 
-function SciMLBase.solve(prob::DualAbstractLinearProblem, args...; kwargs...)
-    return solve(prob, nothing, args...; kwargs...)
-end
-
-function SciMLBase.solve(prob::DualAbstractLinearProblem, ::Nothing, args...;
-        assump = OperatorAssumptions(issquare(prob.A)), kwargs...)
-    return solve(prob, LinearSolve.defaultalg(prob.A, prob.b, assump), args...; kwargs...)
-end
-
-function SciMLBase.solve(prob::DualAbstractLinearProblem,
-        alg::LinearSolve.SciMLLinearSolveAlgorithm, args...; kwargs...)
-    solve!(init(prob, alg, args...; kwargs...))
-end
-
-function linearsolve_dual_solution(
-        u::Number, partials, dual_type)
-    return dual_type(u, partials)
-end
-
-function linearsolve_dual_solution(
-        u::AbstractArray, partials, dual_type)
-    partials_list = RecursiveArrayTools.VectorOfArray(partials)
-    return map(((uᵢ, pᵢ),) -> dual_type(uᵢ, Partials(Tuple(pᵢ))),
-        zip(u, partials_list[i, :] for i in 1:length(partials_list[1])))
-end
-
-get_dual_type(x::Dual) = typeof(x)
-get_dual_type(x::AbstractArray{<:Dual}) = eltype(x)
-get_dual_type(x) = nothing
-
-partial_vals(x::Dual) = ForwardDiff.partials(x)
-partial_vals(x::AbstractArray{<:Dual}) = map(ForwardDiff.partials, x)
-partial_vals(x) = nothing
-
-nodual_value(x) = x
-nodual_value(x::Dual) = ForwardDiff.value(x)
-nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
-
 function xp_linsolve_rhs(uu, ∂_A::Union{<:Partials, <:AbstractArray{<:Partials}},
         ∂_b::Union{<:Partials, <:AbstractArray{<:Partials}})
     A_list = partials_to_list(∂_A)
@@ -130,25 +90,30 @@ function xp_linsolve_rhs(
     b_list
 end
 
-function partials_to_list(partial_matrix::Vector)
-    p = eachindex(first(partial_matrix))
-    [[partial[i] for partial in partial_matrix] for i in p]
+function SciMLBase.solve(prob::DualAbstractLinearProblem, args...; kwargs...)
+    return solve(prob, nothing, args...; kwargs...)
 end
 
-function partials_to_list(partial_matrix)
-    p = length(first(partial_matrix))
-    m, n = size(partial_matrix)
-    res_list = fill(zeros(m, n), p)
-    for k in 1:p
-        res = zeros(m, n)
-        for i in 1:m
-            for j in 1:n
-                res[i, j] = partial_matrix[i, j][k]
-            end
-        end
-        res_list[k] = res
-    end
-    return res_list
+function SciMLBase.solve(prob::DualAbstractLinearProblem, ::Nothing, args...;
+        assump = OperatorAssumptions(issquare(prob.A)), kwargs...)
+    return solve(prob, LinearSolve.defaultalg(prob.A, prob.b, assump), args...; kwargs...)
+end
+
+function SciMLBase.solve(prob::DualAbstractLinearProblem,
+        alg::LinearSolve.SciMLLinearSolveAlgorithm, args...; kwargs...)
+    solve!(init(prob, alg, args...; kwargs...))
+end
+
+function linearsolve_dual_solution(
+        u::Number, partials, dual_type)
+    return dual_type(u, partials)
+end
+
+function linearsolve_dual_solution(
+        u::AbstractArray, partials, dual_type)
+    partials_list = RecursiveArrayTools.VectorOfArray(partials)
+    return map(((uᵢ, pᵢ),) -> dual_type(uᵢ, Partials(Tuple(pᵢ))),
+        zip(u, partials_list[i, :] for i in 1:length(partials_list[1])))
 end
 
 function SciMLBase.init(
@@ -164,6 +129,7 @@ function SciMLBase.init(
         assumptions = OperatorAssumptions(issquare(prob.A)),
         sensealg = LinearSolveAdjoint(),
         kwargs...)
+
     new_A = nodual_value(prob.A)
     new_b = nodual_value(prob.b)
 
@@ -177,7 +143,7 @@ function SciMLBase.init(
         maxiters = maxiters, verbose = verbose, Pl = Pl, Pr = Pr, assumptions = assumptions,
         sensealg = sensealg, kwargs...)
 
-    return DualLinearCache(non_partial_cache, prob, alg, new_A, new_b, ∂_A, ∂_b)
+    return DualLinearCache(non_partial_cache, prob, alg, ∂_A, ∂_b)
 end
 
 function SciMLBase.solve!(cache::DualLinearCache, args...; kwargs...)
@@ -198,4 +164,100 @@ function SciMLBase.solve!(cache::DualLinearCache, args...; kwargs...)
     )
 end
 
+# If setting A or b for DualLinearCache, also set it for the underlying LinearCache
+function Base.setproperty!(dc::DualLinearCache, sym::Symbol, val)
+    # If the property is A or b, also update it in the LinearCache
+    if sym === :A || sym === :b
+        if hasproperty(dc, :linear_cache)
+            setproperty!(dc.linear_cache, sym, nodual_value(val))
+        end
+    end
+
+    # Update the partials if setting A or b
+    if sym === :A
+        setfield!(dc, :partials_A, partial_vals(val))
+    elseif  sym === :b
+        setfield!(dc, :partials_b, partial_vals(val))
+    end
+
+    return val
+end
+
+function Base.getproperty(dc::DualLinearCache, sym::Symbol)
+    if sym === :A
+        return dc.linear_cache.A
+    elseif sym === :b
+        return dc.linear_cache.b
+    else
+        getfield(dc,sym)
+    end
+end
+
+function SciMLBase.reinit!(cache::DualLinearCache;
+        A = nothing,
+        b = cache.b,
+        u = cache.u,
+        p = nothing,
+        reuse_precs = false)
+    (; alg, cacheval, abstol, reltol, maxiters, verbose, assumptions, sensealg) = cache
+
+    isfresh = !isnothing(A)
+    precsisfresh = !reuse_precs && (isfresh || !isnothing(p))
+    isfresh |= cache.isfresh
+    precsisfresh |= cache.precsisfresh
+
+    A = isnothing(A) ? cache.A : A
+    b = isnothing(b) ? cache.b : b
+    u = isnothing(u) ? cache.u : u
+    p = isnothing(p) ? cache.p : p
+    Pl = cache.Pl
+    Pr = cache.Pr
+
+    cache.A = A
+    cache.b = b
+    cache.u = u
+    cache.p = p
+    cache.Pl = Pl
+    cache.Pr = Pr
+    cache.isfresh = true
+    cache.precsisfresh = precsisfresh
+    nothing
+end
+
+# Helper functions for Dual numbers
+get_dual_type(x::Dual) = typeof(x)
+get_dual_type(x::AbstractArray{<:Dual}) = eltype(x)
+get_dual_type(x) = nothing
+
+partial_vals(x::Dual) = ForwardDiff.partials(x)
+partial_vals(x::AbstractArray{<:Dual}) = map(ForwardDiff.partials, x)
+partial_vals(x) = nothing
+
+nodual_value(x) = x
+nodual_value(x::Dual) = ForwardDiff.value(x)
+nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)
+
+
+function partials_to_list(partial_matrix::Vector)
+    p = eachindex(first(partial_matrix))
+    [[partial[i] for partial in partial_matrix] for i in p]
+end
+
+function partials_to_list(partial_matrix)
+    p = length(first(partial_matrix))
+    m, n = size(partial_matrix)
+    res_list = fill(zeros(m, n), p)
+    for k in 1:p
+        res = zeros(m, n)
+        for i in 1:m
+            for j in 1:n
+                res[i, j] = partial_matrix[i, j][k]
+            end
+        end
+        res_list[k] = res
+    end
+    return res_list
+end
+
+
 end