kestrelquantum / QuantumCollocation.jl

Quantum Optimal Control with Direct Collocation
MIT License
29 stars 7 forks source link

[Feature]: Automated gradient testing #144

Open andgoldschmidt opened 1 month ago

andgoldschmidt commented 1 month ago

Feature Description

Background

A big part of the code is writing gradients for integrators, objectives, and constraints for custom objects.

ForwardDiff.jl allows us to calculate analytic gradients, so we should have a function that automates testing the accuracy of our gradients and Jacobians. The function should dispatch over generic Integrators, Objectives, or Constraints, separately. The function should make informative error reporting about where and how the gradients or Jacobians differ.

A useful starting point is the undocumented function show_diffs in test/test_utils.jl, which starts this process of informative error reporting. An updated version of that function is included here:

function show_diffs(A::AbstractMatrix, B::AbstractMatrix; return_indices=false, atol=0, components=nothing)
    @assert size(A) == size(B)
    matrix_is_square = size(A, 1) == size(A, 2)
    indices = []
    for (i, (a, b)) in enumerate(zip(A, B))
        inds = Tuple(CartesianIndices(A)[i])
        if matrix_is_square
            if !isapprox(a, b; atol=atol) && inds[1] ≤ inds[2]
                println((a, b), " @ ", inds)
            end
        else
            if !isapprox(a, b; atol=atol)
                println((a, b), " @ ", inds)
                if return_indices
                    push!(indices, inds)
                end
            end
        end
    end
    return return_indices ? indices : nothing
end

These functions can be added to the _testutils.jl file, and should use a few kinds of NamedTrajectory's to be robust to different cases when checking derivatives (e.g. named_trajectory_type_1(...) in _testutils.jl).

Suggested checklist

Importance

1 (lowest)

What does this feature affect?

Other information

No response