Higher order implicit ad for roots

[Garren-H]

Adds IFTDuals as a dependency which handles the differentiation of root finders, with reference to #502. We can now handle mixed tags and higher order implicit AD efficiently. I took quite a lot of inspiration from the Solvers module to make the code non-allocating where possible.

Additionally, adjusted the association term to use __gradients_for_root_finders to handle AD where we do not have a closed form solution for the fraction of non-bonded sites, X.

[longemen3000]

My main question is, what is the maximum Ad level that IFTDuals supports? The maximum order of Ad ever reported was 4th order (Jacobian of crit_mix)

Nvm, impressive.

[codecov]

Codecov Report

❌ Patch coverage is 96.15385% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 77.69%. Comparing base (84f4047) to head (59a5acd).
⚠️ Report is 41 commits behind head on master.

Files with missing lines	Patch %	Lines
src/models/SAFT/association.jl	95.23%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #542      +/-   ##
==========================================
+ Coverage   77.62%   77.69%   +0.07%     
==========================================
  Files         369      370       +1     
  Lines       32831    32903      +72     
==========================================
+ Hits        25484    25564      +80     
+ Misses       7347     7339       -8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Garren-H]

Noticed a bug with IFTDuals and SVector y as input. Currently uploading a fix for this

[longemen3000]

Errors are unrelated, but i would like to see an all-green CI in this PR in particular. there are a lot of EoS that have a subproblem inside, and could benefit from good implicit (high-order) function theorem solver. for example (without checking the code):

• COSMO-SAC models
• the effective diameter on SAFT-VRQ-Mie
• the diameter in SAFT-VR-Mie is an integral problem on some parts and a solver in others.
• the MSA solver for electrolytes

I'm going to open an issue tracking those subproblems.

[Garren-H]

I have ran the CI actions on my local fork and some errors popped up with flash derivatives. Upon further inspection I have found a few issues, first being that I used inplace .*= and ldiv and/or ldiv! in cases where the jacobian returns an SMatrix. Additionally, the mixed tag seems to only work for the scalar case with a single variable being differentiated (i.e. all partial fields contains a Tuple of length 1), so have converted to a draft for now. Will fix these issues and then make the PR active again. The same tag case should work in any case, (vector with a single diff variable; scalar with a single diff variable; scalar with multiple diff variables and vector with multiple diff variables).

there are a lot of EoS that have a subproblem inside, and could benefit from good implicit (high-order) function theorem solver

Fully agree on this point and this was the intended purpose of the development of this tool, but in my case focusing mainly on SAFT

[Garren-H]

@longemen3000 I have implemented the fixes on my side and should be fine now. The CI's should be re-run, I pushed slightly early, when the new version wasn't registered yet

There is quite a lot of performance on the table for the mixed tags case when working with multiple partials and/or vectors. Don't know when I'll have time to improve the performance for that, but will look at it at some point

[longemen3000]

most test pass (including crit_mix(model,z), so it seems that the package works!) the failure in CI is for the big increase in compilation time (around 4x). is it possible to run the tests with the Vector approach to see if the full test suite runs?

That performance left in the table is in Clapeyron or IFTDuals? if there is anything i can help in that regard, let me know. In particular, regarding the linear solver part, If i remember correctly, each LU instance allocates a Matrix{T} and a Vector{Int}. i don't think you can reuse the matrix if the dual sizes are different (maybe using a view over the biggest matrix?), but that Vector{Int} can be definitely reused (that is what i did with Clapeyron.Solvers.unsafe_LU!)

In particular, the association algorithm as it is, is not the best in terms of allocations, but in terms of runtime speed. So, more allocations are not that important (if the speed is more or less the same). The SVector approach needs to be limited because it makes the function type-unstable (the result type depends on the value of the arguments.)

[Garren-H]

most test pass (including crit_mix(model,z), so it seems that the package works!) the failure in CI is for the big increase in compilation time (around 4x). is it possible to run the tests with the Vector approach to see if the full test suite runs?

I am confident that it should succeed in all cases. The tests I have for mixed tags considers all cases, all combinations of scalar and vector inputs and outputs. These tests for IFTDuals passes so the derivative logic is fine, there may be bugs with the utils however (i.e. dual promotion/demotions and obtaining the common dual type of args from generic structures) but does not seem to be the case currently.

That performance left in the table is in Clapeyron or IFTDuals? if there is anything i can help in that regard, let me know.

Performance with IFTDuals. During the recursive process we constructs intermediate buffer arrays for the storage of duals. These buffers are short lived as we only use them to loop through the partials, solve each partial and then stack the result into a dual. Lets say we have want to solve z=g(x,y,w) where each input is a diff variable. We denote the sizes here by Nz, Nx, .... We assume that x is the inner most Dual, y is the 2nd layer of nesting and w is the outer most nesting. So a Tag like Tag{:w, Dual{Tag{:y, Dual{Tag{:x}}}}}. If we want the 3rd dual, we already have the second order dual (value field if the 3rd order dual) solved. We then solve for dz/dw first (the partial.value.value field). This is a Nz x Nw matrix, which is reconstructed into the Vector of Duals. Next we solve for d2z/dwdx (the partials.value.partials) field. How we solve this is by considering a single partial direction for x, i.e. solving d2z/dwdx_i at a time, maintaining the Nz x Nw size matrix in solves. This is repeated for each index in x. We then do the same for y in the partials.partials.value field. And finally we solve for the partials.partials.partials field. Here we solve for each d3z/dw dx_i dy_j where we loop over i and j. So we need to loop, store the derivatives in a buffer, and only afterwards can we reconstruct the partials field, and then combine with the value field. The process of constructing these intermediate buffers allocates quite substantially, which slows down the code. So we either need a way to make the buffer non-allocating. I tried SArrays (locally), but this seems to allocate even more. So a potential may be to rewrite this as a generated or equivalent where we circumvent the intermediate storage buffers. Haven't come up with a solution yet and there are other things which require my attention, so unsure when I'll get to looking at making improvements again. Some source code is

https://github.com/Garren-H/IFTDuals.jl/blob/70440b567552c148e49ed889dbff02aee5fa218f/src/derivatives.jl#L143C1-L151C4

function store_ift_cache(y::Y, BNi::B, neg_A, PT=Partials{N,V}) where {T,V<:Real,N,DT<:Dual{T,V,N},Y<:AbstractVector{V},B<:AbstractVector{DT}} # offload logic to be more efficient with storage types
    dual_cache = Matrix{V}(undef, length(y), N)
    for i in 1:N
        BNi_i = extract_partials_(BNi, i) # get nested Duals
        dy = extract_partials_(y, i) # get nested Duals
        dual_cache[:, i] .= ift_(dy, BNi_i, neg_A)
    end
    return make_dual(y, DT, PT, dual_cache) # construct Dual numbers
end

Here make_dual takes the value field y, Dual type, Partials type and partials field dual_cache and constructs the vector of duals. This is in the same tag case, the mixed tag case allocates either Vectors, Matrices or 3rd order Tensor/array as storage buffers. It allocates both interediate storage buffer to solve the partials.value field and the partials.partials fields. Source can be found here: https://github.com/Garren-H/IFTDuals.jl/blob/70440b567552c148e49ed889dbff02aee5fa218f/src/derivatives.jl#L203C1-L221C4, https://github.com/Garren-H/IFTDuals.jl/blob/70440b567552c148e49ed889dbff02aee5fa218f/src/derivatives.jl#L243C1-L252C1, https://github.com/Garren-H/IFTDuals.jl/blob/70440b567552c148e49ed889dbff02aee5fa218f/src/derivatives.jl#L256C1-L265C4. These are the main sources where performance can be improved.

In particular, regarding the linear solver part, If i remember correctly, each LU instance allocates a Matrix{T} and a Vector{Int}. i don't think you can reuse the matrix if the dual sizes are different (maybe using a view over the biggest matrix?), but that Vector{Int} can be definitely reused (that is what i did with Clapeyron.Solvers.unsafe_LU!)

The LU factorization is always reused. That is for a vector z we will always have an LU decomposition and this is used to solve for all partials (see logic above, where we maintain matrix algebra even for nested duals)

The SVector approach needs to be limited because it makes the function type-unstable (the result type depends on the value of the arguments.)

Just for my sanity, I thought type instability occurs when the compiler cannot infer the types (either input or return) of variables. Since these are known (as part of the signature) I was under the impression that the SVector construction would be faster.

[longemen3000]

My decision for now is to merge this, and stabilize performance along the way. Even in the worst case, using IFTDuals is a clear improvement over the the existent IFT solver