association V2: use michelsen minimization + ASS, implicit AD for X i…

…f needed.

NEWS.md

@@ -8,5 +8,6 @@
 - `VT_diffusive_stability` now uses `eigmin` instead of the full eigen calculation.
 - `isstable` now works on (P,T,z) space, for the (V,T,z) space, use `VT_isstable`. there are now (P,T,z) versions of each stability function.
 - calculation of volumes,saturation pressures and critical points of CPA models now defaults to the inner cubic model when there is no association present.
-- The default association implementation now uses implicit AD to support derivatives (via michelsen's Q function), instead of propagating derivative information through the iterative procedure. 
-- the default `volume` implementation now uses implicit AD to support derivatives. instead of propagating derivative information through the iterative procedure. This allows workloads of the type: `ForwardDiff.derivative(_p -> volume(model,_p,T,z),p0)` to be efficient.
+- The default association implementation now uses a combination of accelerated successive substitution and newton optimization. While increasing allocations, the method is faster.
+- The default association implementation now uses implicit AD to support derivatives (via michelsen's Q function), instead of propagating derivative information through the iterative procedure. `Clapeyron.X` will still propagate derivative information as usual, but a new function (`Clapeyron.X_and_Δ`) just returns the primal value of the nonbonded fractions, along with the calculated association interaction energies.
+- the default `volume` implementation now uses implicit AD to support derivatives. instead of propagating derivative information through the iterative procedure. This allows workloads of the type: `ForwardDiff.derivative(_p -> property(model,_p,T,z,phase = :l,vol0 = v0),p)` to be efficiently calculated.

src/methods/property_solvers/volume.jl

@@ -18,33 +18,32 @@ function volume_compress(model,p,T,z=SA[1.0];V0=x0_volume(model,p,T,z,phase=:liq
 end
 
 function _volume_compress(model,_p,_T,_z=SA[1.0],V0=x0_volume(model,p,T,z,phase=:liquid),max_iters=100)
-    _0 = zero(_p+_T+first(_z)+oneunit(eltype(model)))
+    _0 = zero(Base.promote_eltype(model,_p,_T,_z))
     _1 = one(_0)
     isnan(V0) && return _0/_0
-    p₀ = primalval(1.0*_p)
+    p₀ = primalval(_1*_p)
+    XX = typeof(p₀)
     T = primalval(_T)
-    _nan = _0/_0
+    _nan = primalval(_0/_0)
     nan = primalval(_nan)
     logV0 = primalval(log(V0)*_1)
     z = primalval(_z)
     log_lb_v = log(primalval(lb_volume(model,z)))
-    function logstep(logVᵢ)
-        logVᵢ < log_lb_v && return zero(logVᵢ)/zero(logVᵢ)
+    function logstep(logVᵢ::TT) where TT
+        logVᵢ < log_lb_v && return TT(zero(logVᵢ)/zero(logVᵢ))
         Vᵢ = exp(logVᵢ)
         _pᵢ,_dpdVᵢ = p∂p∂V(model,Vᵢ,T,z)
         pᵢ,dpdVᵢ = primalval(_pᵢ),primalval(_dpdVᵢ) #ther could be rare cases where the model itself has derivative information.
-        dpdVᵢ > 0 && return zero(logVᵢ)/zero(logVᵢ) #inline mechanical stability.
-        abs(pᵢ-p₀) < 3eps(p₀) && return zero(Vᵢ) #this helps convergence near critical points.
+        dpdVᵢ > 0 && return TT(zero(logVᵢ)/zero(logVᵢ)) #inline mechanical stability.
+        abs(pᵢ-p₀) < 3eps(p₀) && return TT(zero(Vᵢ)) #this helps convergence near critical points.
         Δᵢ = (p₀-pᵢ)/(Vᵢ*dpdVᵢ) #(_p - pset)*κ
-        return Δᵢ
+        return TT(Δᵢ)
     end
-    function f_fixpoint(logVᵢ)
-        Δᵢ = logstep(logVᵢ)
-        logV = logVᵢ + Δᵢ
-        return logV
+    function f_fixpoint(logVᵢ::TT) where TT
+        return TT(logVᵢ + logstep(logVᵢ))
     end
 
-    logV = @nan(Solvers.fixpoint(f_fixpoint,logV0,Solvers.SSFixPoint(),rtol = 1e-12,max_iters=max_iters),nan)
+    logV = @nan(Solvers.fixpoint(f_fixpoint,logV0,Solvers.SSFixPoint(),rtol = 1e-12,max_iters=max_iters)::XX,nan)
     #netwon step to recover derivative information:
     #V = V - (p(V) - p)/(dpdV(V))
     #dVdP = -1/dpdV

src/models/SAFT/association.jl

@@ -3,8 +3,8 @@ function a_assoc(model::EoSModel, V, T, z,data=nothing)
     nn = assoc_pair_length(model)
     iszero(nn) && return _0
     isone(nn) && return a_assoc_exact_1(model,V,T,z,data)
-    X,Δ = @f(X_and_Δ,data)
-    return @f(a_assoc_impl,X,Δ)
+    _X,_Δ = @f(X_and_Δ,data)
+    return @f(a_assoc_impl,_X,_Δ)
 end
 
 """
@@ -127,8 +127,8 @@ Returns association options used in the association solver.
     return model.assoc_options
 end
 
-function Δ(model::EoSModel, V, T, z,data)
-    Δout = assoc_similar(model,typeof(V+T+first(z)))
+function Δ(model::EoSModel, V, T, z, data)
+    Δout = assoc_similar(model,@f(Base.promote_eltype))
     Δout.values .= false
     for (idx,(i,j),(a,b)) in indices(Δout)
         Δout[idx] =@f(Δ,i,j,a,b,data)
@@ -180,12 +180,12 @@ function nonzero_extrema(K)
     return _min,_max
 end
 
-function assoc_site_matrix(model,V,T,z,ad::VV = Val{false}(),data = nothing,delta = @f(__delta_assoc,data)) where VV
+function assoc_site_matrix(model,V,T,z,data = nothing,delta = @f(__delta_assoc,data))
     options = assoc_options(model)
     if !options.dense
         @warn "using sparse matrices for association is deprecated."
     end
-    return dense_assoc_site_matrix(model,V,T,z,ad,data,delta)
+    return dense_assoc_site_matrix(model,V,T,z,data,delta)
 end
 
 #this fills the zeros of the Δ vector with the corresponding mixing values
@@ -214,32 +214,19 @@ function elliott_runtime_mix!(Δ)
     return Δ
 end
 
-function maybe_ad(x::X,::Val{true}) where X
-    return x
-end
-
-function maybe_ad(x::X,::Val{false}) where X
-    return primalval(x)
-end
-
-function dense_assoc_site_matrix(model,V,T,∂z,ad::VV = Val{false}(),data=nothing,delta = @f(__delta_assoc,data)) where VV
+function dense_assoc_site_matrix(model,V,T,z,data=nothing,delta = @f(__delta_assoc,data))
     sitesparam = getsites(model)
     _sites = sitesparam.n_sites
     p = _sites.p
-    ρ = maybe_ad(N_A/V,ad)
-    #ρ = N_A/V
-    z = maybe_ad(∂z,ad)
-    #z = ∂z
+    ρ = N_A/V
     _ii::Vector{Tuple{Int,Int}} = delta.outer_indices
     _aa::Vector{Tuple{Int,Int}} = delta.inner_indices
     _idx = 1:length(_ii)
-    Δ = maybe_ad(delta.values,ad)
-    #_Δ = delta.values
+    Δ = delta.values
     TT = eltype(Δ)
     _n = sitesparam.n_sites.v
     nn = length(_n)
     K  = zeros(TT,nn,nn)
-    count = 0
     options = assoc_options(model)
     combining = options.combining
     runtime_combining = combining ∈ (:elliott_runtime,:esd_runtime)
@@ -256,13 +243,11 @@ function dense_assoc_site_matrix(model,V,T,∂z,ad::VV = Val{false}(),data=nothi
                     b = complement_index(a,ab)
                     jb = compute_index(p,j,b)
                     njb = _n[jb]
-                    count += 1
                     K[ia,jb]  = ρ*njb*z[j]*Δ[idx]
                 end
             end
         end
     end
-
     return K::Matrix{TT}
 end
 
@@ -293,32 +278,60 @@ julia> x = Clapeyron.assoc_fractions(model,2.6e-5,300.15,[0.3,0.3,0.4]) #you can
 """
 function X(model::EoSModel, V, T, z,data = nothing)
     #we return X with derivative information
-    X,Δ = X_and_Δ(model,V,T,z,data,Val{true}())
-    return X
+    nn = assoc_pair_length(model)
+    isone(nn) && return X_exact1(model,V,T,z,data)
+    X,Δ = X_and_Δ(model,V,T,z,data)
+    #bail out if there is no AD
+    if eltype(X.v) === eltype(Δ.values)
+        return X
+    end
+    X̄ = X.v
+    #K matrix with derivative information
+    K = assoc_site_matrix(model,V,T,z,data,Δ)
+    X̃ = similar(K,length(X̄))
+
+    #=
+    strategy to obtain general derivatives of nonbonded fractions with automatic differenciation:
+
+    using Implicit AD, we can update X with a "perfect newton upgrade", with the derivative information added in the last update.
+    it is equivalent to the method of Tan (2004), in the sense that we still need to solve a linear system of equations containing X.
+    but this only requires to solve one linear system, as the derivatives are carried by the number type, instead of separated.
+    
+    =#
+    mul!(X̃,K,X̄)
+    K .*= -1
+    for k in 1:size(K,1)
+        K[k,k] -= (1 + X̃[k])/X̄[k] 
+    end
+    X̃ .+= -1 ./ X̄ .+ 1
+
+    F = Solvers.unsafe_LU!(K)
+    ldiv!(F,X̃)
+    X̃ .+= X̄
+    return PackedVofV(X.p,X̃)
 end
 
-function X_and_Δ(model::EoSModel, V, T, z,data = nothing,ad::VV = Val{false}()) where VV
+function X_and_Δ(model::EoSModel, V, T, z,data = nothing)
     nn = assoc_pair_length(model)
-    #isone(nn) && return X_exact1(model,V,T,z,data)
-    options = assoc_options(model)
+    isone(nn) && return X_and_Δ_exact1(model,V,T,z,data)
+    options = assoc_options(model)::AssocOptions
     _Δ = __delta_assoc(model,V,T,z,data)
-    K = assoc_site_matrix(model,V,T,z,ad,data,_Δ)
+    K = assoc_site_matrix(model,primalval(V),T,primalval(z),data,primalval(_Δ))
     sitesparam = getsites(model)
     idxs = sitesparam.n_sites.p
     Xsol = assoc_matrix_solve(K,options)
     return PackedVofV(idxs,Xsol),_Δ
 end
 
-function assoc_matrix_solve(K,options::AssocOptions)
-    atol = options.atol
-    rtol = options.rtol
+function assoc_matrix_solve(K::AbstractMatrix{T},options::AssocOptions) where T
+    atol = T(options.atol)
+    rtol = T(options.rtol)
     max_iters = options.max_iters
-    α = options.dampingfactor
+    α = T(options.dampingfactor)
     return assoc_matrix_solve(K, α, atol ,rtol, max_iters)
 end
 
-#TODO: define implicit AD here
-function assoc_matrix_solve(K, α, atol ,rtol, max_iters)
+function assoc_matrix_solve(K::AbstractMatrix{T}, α::T, atol ,rtol, max_iters) where T
     n = LinearAlgebra.checksquare(K) #size
     #initialization procedure:
     Kmin,Kmax = nonzero_extrema(K) #look for 0 < Amin < Amax
@@ -327,41 +340,108 @@ function assoc_matrix_solve(K, α, atol ,rtol, max_iters)
     else
         f = true-Kmin
     end
-    X0 = fill(f,n) #initial point
+    f = min(f,one(f))
+    X0 = Vector{T}(undef,n)
+    Xsol = Vector{T}(undef,n)
+    X0 .= f
+    Xsol .= f
     #=
     function to solve
     find vector x that satisfies:
     (A*x .* x) + x - 1 = 0
-    Jacobian: Diagonal(Ax + 1) + Diagonal(x)*A|
-
     solved by reformulating in succesive substitution:
     x .= 1 ./ (1 .+ A*x)
+   
     #we perform a "partial multiplication". that is, we use the already calculated
     #values of the next Xi to calculate the current Xi. this seems to accelerate the convergence
-    #by around 50%
+    #by around 50% (check what the ass_matmul! function does)
+
+    note that the damping is done inside the partial multiplication. if is done outside, it causes convergence problems.
+    
+    after a number of ss iterations are done, we use newton minimization.
+    the code for the newton optimization is based on sgtpy: https://github.com/gustavochm/sgtpy/blob/336cb2a7581b22492914233e29062f5a364b47da/sgtpy/vrmie_pure/association_aux.py#L33-L57
+    
+    some notes:
+    - the linear system is solved via LU decomposition, for that, we need to allocate one (1) Matrix{T} and one (1) Vector{Int}
+    - gauss-seidel does not require an additional matrix allocation, but it is slow. (slower than SS)
+    - julia 1.10 does not have a way to make LU non-allocating, but the code is simple, so it was added as the function unsafe_LU! in the Solvers module. 
     =#
-    function fX(out,in)
-        mul!(out,K,in)
-        for i in 1:length(out)
-            #
-            Xi = in[i]
-            Kxi = out[i]
-            out[i] = 1/(1+Kxi)
-        end 
+    fx(kx,x) =  α/(1+kx) + (1-α)*x
+    function f_ss!(out,in)
+        ass_matmul!(fx,out,K,in)
         return out
     end
 
-    #successive substitution until convergence
-    return Solvers.fixpoint(fX,X0,Solvers.SSFixPoint(α),atol=atol,rtol = rtol,max_iters = max_iters)
+    #successive substitution. 50 iters
+    it_ss = (5*length(Xsol))
+    converged = false
+    for i in 1:it_ss
+        f_ss!(Xsol,X0)
+        converged,finite = Solvers.convergence(Xsol,X0,atol,rtol)
+        if converged
+            if finite
+                return Xsol
+            else
+                Xsol .= NaN
+                return Xsol 
+            end
+        end
+        X0 .= Xsol
+       # @show Xsol
+    end
+    H = Matrix{T}(undef,n,n)
+    piv = zeros(Int,n)
+    F = Solvers.unsafe_LU!(H,piv)
+    if !converged #proceed to newton minimization
+        dX = copy(Xsol)
+        KX = copy(Xsol)
+        dX_cache = copy(Xsol)
+        for i in (it_ss + 1):max_iters
+            #@show Xsol
+
+            KX = mul!(KX,K,Xsol)
+            H .= -K
+            for k in 1:size(H,1)
+                H[k,k] -= (1 + KX[k])/Xsol[k] 
+            end
+            #F already contains H and the pivots, because we refreshed H, we need to refresh
+            #the factorization too.
+            F = Solvers.unsafe_LU!(F)
+            dX .= 1 ./ Xsol .- 1 .- KX #gradient
+            ldiv!(F,dX) #we solve H/g, overwriting g
+            Xsol .-= dX
+            converged,finite = Solvers.convergence(Xsol,X0,atol,rtol,false,Inf)
+            #@show converged,finite
+            if converged
+                if !finite
+                    fill!(Xsol,NaN)
+                end
+                return Xsol
+            end
+            X0 .= Xsol
+        end
+    end
+    if !converged
+        Xsol .= NaN
+    end
+    return Xsol
 end
 
 #exact calculation of site non-bonded fraction when there is only one site
 
 function X_exact1(model,V,T,z,data = nothing)
-    xia,xjb,i,j,a,b,n,idxs = _X_exact1(model,V,T,z,data)
+    xia,xjb,i,j,a,b,n,idxs,Δijab = _X_exact1(model,V,T,z,data)
     pack_X_exact1(xia,xjb,i,j,a,b,n,idxs)
 end
 
+function X_and_Δ_exact1(model,V,T,z,data = nothing)
+    xia,xjb,i,j,a,b,n,idxs,Δijab = _X_exact1(model,V,T,z,data)
+    XX = pack_X_exact1(primalval(xia),primalval(xjb),i,j,a,b,n,idxs)
+    Δout = assoc_similar(model,@f(Base.promote_eltype))
+    Δout.values[1] = Δijab
+    return XX,Δout
+end
+
 function _X_exact1(model,V,T,z,data=nothing)
     κ = model.params.bondvol.values
     i,j = κ.outer_indices[1]
@@ -392,7 +472,7 @@ function _X_exact1(model,V,T,z,data=nothing)
     xia = -2*_c/denom
     xk_ia = kia*xia
     xjb = (1- xk_ia)/(1 - xk_ia*xk_ia)
-    return xia,xjb,i,j,a,b,n,idxs
+    return xia,xjb,i,j,a,b,n,idxs,_Δ
 end
 
 function pack_X_exact1(xia,xjb,i,j,a,b,n,idxs)

src/utils/acceleration_ss.jl

@@ -11,6 +11,24 @@ function dem!(x_dem,xₙ, xₙ₋₁, xₙ₋₂,cache = (similar(x),similar(x))
     x_dem .= xₙ .+ Δxₙ .* λ ./(1. .- λ)
 end
 
+function ass_matmul!(f::F,out,K,x) where F
+    n = length(x)
+    i_solved = 0
+    for i in 1:n
+        kxi = zero(eltype(x))
+        #strategy 3
+        Ki = @view K[i,:]
+        @inbounds for j in 1:i_solved
+            kxi += out[j]*Ki[j]
+        end
+        @inbounds for j in (i_solved+1):n
+            kxi += x[j]*Ki[j]
+        end
+        out[i] = f(kxi,x[i])
+        i_solved += 1
+    end
+end
+
 # Function to accelerate Succesive Substitution by GDEM method (2 eigenvalues)
 #=
 function gdem2(xₙ, xₙ₋₁, xₙ₋₂, xₙ₋₃)