GSOC Goal: Objective function and Jacobian for exact solver

[msricher]

Implement the objective function and Jacobian for the exact solver.

[sparsh-karna]

Hi,

I was working on the new objective function and jacobian implementation

function objective_function!(f::Vector{Float64}, rxn_system::ReactionSystem, K_eqs::AbstractVector{Float64}, xi::Vector{Float64})
    activities = rxn_system.concs_init .+ rxn_system.stoich * xi

    if any(x -> x <= 0, activities)
        f .= Inf
        return nothing
    end

    for r in 1:rxn_system.n_reaction
        term1 = K_eqs[r]
        term2 = 1.0

        for i in 1:rxn_system.n_species
            nu_ir = rxn_system.stoich[i, r]
            if nu_ir < 0
                term1 *= activities[i]^(-nu_ir)
            elseif nu_ir > 0
                term2 *= activities[i]^(nu_ir)
            end
        end
        f[r] = term1 - term2
    end
    return nothing
end

function jacobian_function!(J::Matrix{Float64}, rxn_system::ReactionSystem, K_eqs::AbstractVector{Float64}, xi::Vector{Float64})
    R = rxn_system.n_reaction
    N = rxn_system.n_species

    activities = rxn_system.concs_init .+ rxn_system.stoich * xi

    if any(x -> x <= 0, activities)
        J .= Inf
        return nothing
    end

    P_reactants = ones(R)
    P_products = ones(R)
    for r in 1:R
        for i in 1:N
            nu_ir = rxn_system.stoich[i, r]
            if nu_ir < 0
                P_reactants[r] *= activities[i]^(-nu_ir)
            elseif nu_ir > 0
                P_products[r] *= activities[i]^(nu_ir)
            end
        end
    end

    for r in 1:R, s in 1:R
        sum_deriv1 = 0.0
        for i in 1:N
            if rxn_system.stoich[i, r] < 0
                sum_deriv1 += (-rxn_system.stoich[i, r] / activities[i]) * rxn_system.stoich[i, s]
            end
        end
        term1_deriv = K_eqs[r] * P_reactants[r] * sum_deriv1

        sum_deriv2 = 0.0
        for i in 1:N
            if rxn_system.stoich[i, r] > 0
                sum_deriv2 += (rxn_system.stoich[i, r] / activities[i]) * rxn_system.stoich[i, s]
            end
        end
        term2_deriv = P_products[r] * sum_deriv2

        J[r, s] = term1_deriv - term2_deriv
    end
    return nothing
end

The objective function is more or less the same, but i changed the jacobian to use logarithmic differentiation, early the looping was very inefficient and was yielding a time complexity O(n^2), now i think it will be faster was only single loop is running.

As discussed earlier in the meeting i incorporated it in the solve function and tried to run it.

function solve(
    rxn_system::ReactionSystem,
    K_eqs::AbstractVector{Float64};
    maxiters::Integer=1000,
    abstol::Real=1.0e-9,
    reltol::Real=0.0,
)
    # Define wrappers that unpack the parameter tuple p = (rxn_system, K_eqs)
    f_obj = (du, u, p) -> objective_function!(du, p[1], p[2], u)
    f_jac = (J, u, p) -> jacobian_function!(J, p[1], p[2], u)

    nls_func = NonlinearFunction(f_obj; jac=f_jac)
    u0 = zeros(Float64, rxn_system.n_reaction)
    params = (rxn_system, K_eqs)
    problem = NonlinearProblem(nls_func, u0, params)
    solution = NonlinearSolve.solve(problem, TrustRegion(); maxiters=maxiters, abstol=abstol, reltol=reltol)

    if SciMLBase.successful_retcode(solution)
        rxn_system.concs .= rxn_system.concs_init .+ rxn_system.stoich * solution.u
    end

    return solution
end

...check the full code in commit f4334f50.

     Testing Running tests...
Test Summary:                | Pass  Total  Time
Biological Oscillator System |    6      6  3.2s
Test Summary: | Pass  Total     Time
System 1      |    3      3  1m40.3s
Test Summary: | Pass  Total     Time
System 2      |    3      3  3m01.7s
Test Summary:                  | Pass  Total  Time
Constructor and Property Tests |    7      7  0.1s
Test Summary:          | Pass  Total  Time
Rate Calculation Tests |    2      2  0.0s
Test Summary:           | Pass  Total  Time
Mass Conservation Tests |    8      8  0.0s
     Testing NICE tests passed 

[sparsh-karna]

As suggested by @msricher , I used the ForwardDiff.jl to compute the correctness of the jacobian implementation.

@testset "Jacobian Verification" begin
    # Use System 1 for Jacobian testing since it's well-behaved
    stoich = Float64[
        -0.5 -0.5;
        1.0 -1.0;
        0.0 1.0
    ]
    keq_vals = [1.0, 0.1]
    concs = [1.0, 0.2, 0.4]
    rxn_system = ReactionSystem(stoich, concs, keq_vals; φ=1.0)

    # Wrapper function for ForwardDiff - needs to be out-of-place
    objective_wrapper = (xi) -> begin
        f_out = similar(xi, rxn_system.n_reaction)
        NICE.objective_function!(f_out, rxn_system, keq_vals, xi)
        return f_out
    end

    # Test Jacobian accuracy at multiple random points
    for _ in 1:20
        # Small random perturbations to avoid negative concentrations
        xi = 0.01 * rand(rxn_system.n_reaction)

        # Skip if this would give negative concentrations
        if any(rxn_system.concs_init .+ rxn_system.stoich * xi .<= 0)
            continue
        end

        # Compare analytical Jacobian to automatic differentiation
        J_analytical = zeros(rxn_system.n_reaction, rxn_system.n_reaction)
        NICE.jacobian_function!(J_analytical, rxn_system, keq_vals, xi)

        J_ad = ForwardDiff.jacobian(objective_wrapper, xi)

        @test isapprox(J_analytical, J_ad; atol=1e-8, rtol=1e-8)
    end
end

i ran and got positive results.

Jacobian Verification |   20     20  2.3s
     Testing NICE tests passed