NLSolversBase.jl/src/objective_types/twicedifferentiable.jl at ce51aa93493dcaceac8eee1a8c7c24a84a1165e0 · devmotion/NLSolversBase.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# Used for objectives and solvers where the gradient and Hessian is available/exists
mutable struct TwiceDifferentiable{T<:Real,TDF<:AbstractArray,TH<:AbstractMatrix,TX<:AbstractArray} <: AbstractObjective
    f
    df
    fdf
    dfh
    fdfh
    h
    F::T
    DF::TDF
    H::TH
    x_f::TX
    x_df::TX
    x_h::TX
    f_calls::Int
    df_calls::Int
    h_calls::Int
end
# compatibility with old constructor
function TwiceDifferentiable(f, g, fg, h, x::TX, F::T = real(zero(eltype(x))), G::TG = alloc_DF(x, F), H::TH = alloc_H(x, F); inplace::Bool = true) where {T<:Real, TG<:AbstractArray, TH<:AbstractMatrix, TX<:AbstractArray}
    x_f, x_df, x_h = x_of_nans(x), x_of_nans(x), x_of_nans(x)

    g! = df!_from_df(g, F, inplace)
    fg! = fdf!_from_fdf(fg, F, inplace)
    h! = h!_from_h(h, F, inplace)

    TwiceDifferentiable{T,TG,TH,TX}(f, g!, fg!, nothing, nothing, h!,
                                        copy(F), copy(G), copy(H),
                                        x_f, x_df, x_h,
                                        0, 0, 0)
end

function TwiceDifferentiable(f, g, h,
                             x::AbstractArray,
                             F::Real = real(zero(eltype(x))),
                             G::AbstractArray = alloc_DF(x, F),
                             H::AbstractMatrix = alloc_H(x, F); inplace = true)
    g! = df!_from_df(g, F, inplace)
    h! = h!_from_h(h, F, inplace)

    fg! = make_fdf(x, F, f, g!)
    x_f, x_df, x_h = x_of_nans(x), x_of_nans(x), x_of_nans(x)

    return TwiceDifferentiable(f, g!, fg!, nothing, nothing, h!, F, G, H, x_f, x_df, x_h, 0, 0, 0)
end


function TwiceDifferentiable(f, g,
                             x_seed::AbstractArray,
                             F::Real = real(zero(eltype(x_seed))); autodiff::Union{AbstractADType,Symbol,Bool} = :finite, inplace::Bool = true)
    g! = df!_from_df(g, F, inplace)
    fg! = make_fdf(x_seed, F, f, g!)

    backend = get_adtype(autodiff)
    hess_prep = DI.prepare_hessian(f, backend, x_seed)
    function h!(_h, _x)
        DI.hessian!(f, _h, hess_prep, backend, _x)
        return _h
    end
    TwiceDifferentiable(f, g!, fg!, h!, x_seed, F)
end

TwiceDifferentiable(d::NonDifferentiable, x_seed::AbstractArray = d.x_f, F::Real = real(zero(eltype(x_seed))); autodiff::Union{AbstractADType,Symbol,Bool} = :finite) =
    TwiceDifferentiable(d.f, x_seed, F; autodiff = autodiff)

function TwiceDifferentiable(d::OnceDifferentiable, x_seed::AbstractArray = d.x_f,
                             F::Real = real(zero(eltype(x_seed))); autodiff::Union{AbstractADType,Symbol,Bool} = :finite)
    backend = get_adtype(autodiff)
    hess_prep = DI.prepare_hessian(d.f, backend, x_seed)
    function h!(_h, _x)
        DI.hessian!(d.f, _h, hess_prep, backend, _x)
        return _h
    end
    return TwiceDifferentiable(d.f, d.df, d.fdf, h!, x_seed, F, gradient(d))
end

function TwiceDifferentiable(f, x::AbstractArray, F::Real = real(zero(eltype(x)));
                             autodiff::Union{AbstractADType,Symbol,Bool} = :finite, inplace::Bool = true)
    backend = get_adtype(autodiff)
    grad_prep = DI.prepare_gradient(f, backend, x)
    hess_prep = DI.prepare_hessian(f, backend, x)
    function g!(_g, _x)
        DI.gradient!(f, _g, grad_prep, backend, _x)
        return nothing
    end
    function fg!(_g, _x)
        y, _ = DI.value_and_gradient!(f, _g, grad_prep, backend, _x)
        return y
    end
    function h!(_h, _x)
        DI.hessian!(f, _h, hess_prep, backend, _x)
        return _h
    end
    TwiceDifferentiable(f, g!, fg!, h!, x, F)
end

function hv_product!(obj::TwiceDifferentiable, x, v)
    H = hessian!(obj, x)
    return H*v
end