BatchNLPKernels.jl/src/api/hess.jl at 6e8cd0341e5da7138ffba36781f9b3ce4d413f0f · LearningToOptimize/BatchNLPKernels.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
"""
    lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0)

Evaluate Hessian coordinates for a batch of points.
"""
function lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Θ::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0)
    H_view = _maybe_view(bm, :hprod_work, X)
    lagrangian_hessian!(bm, X, Θ, Y, H_view; obj_weight=obj_weight)
    return H_view
end

"""
    lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0)

Evaluate Hessian coordinates for a batch of points.
"""
function lagrangian_hessian!(bm::BatchModel, X::AbstractMatrix, Y::AbstractMatrix; obj_weight=1.0)
    Θ = _repeat_params(bm, X)
    lagrangian_hessian!(bm, X, Θ, Y; obj_weight=obj_weight)
end

function lagrangian_hessian!(
    bm::BatchModel,
    X::AbstractMatrix,
    Θ::AbstractMatrix,
    Y::AbstractMatrix,
    H::AbstractMatrix;
    obj_weight=1.0,
)
    batch_size = size(X, 2)
    @lencheck batch_size eachcol(X) eachcol(Θ) eachcol(Y) eachcol(H)
    @lencheck bm.model.meta.nvar eachrow(X)
    @lencheck length(bm.model.θ) eachrow(Θ)
    @lencheck bm.model.meta.ncon eachrow(Y)
    @lencheck bm.model.meta.nnzh eachrow(H)
    _assert_batch_size(batch_size, bm.batch_size)
    backend = _get_backend(bm.model)

    fill!(H, zero(eltype(H)))
    _obj_lagrangian_hessian!(backend, H, bm.model.objs, X, Θ, obj_weight)
    _con_lagrangian_hessian!(backend, H, bm.model.cons, X, Θ, Y)
    return H
end

function _obj_lagrangian_hessian!(backend, H, objs, X, Θ, obj_weight)
    shessian_batch!(backend, H, nothing, objs, X, Θ, obj_weight, zero(eltype(H)))
    _obj_lagrangian_hessian!(backend, H, objs.inner, X, Θ, obj_weight)
    synchronize(backend)
end
function _obj_lagrangian_hessian!(backend, H, objs::ExaModels.ObjectiveNull, X, Θ, obj_weight) end

function _con_lagrangian_hessian!(backend, H, cons, X, Θ, Y)
    shessian_batch!(backend, H, nothing, cons, X, Θ, Y, zero(eltype(H)))
    _con_lagrangian_hessian!(backend, H, cons.inner, X, Θ, Y)
    synchronize(backend)
end
function _con_lagrangian_hessian!(backend, H, cons::ExaModels.ConstraintNull, X, Θ, Y) end

function shessian_batch!(
    backend::B,
    y1,
    y2,
    f,
    X,
    Θ,
    adj,
    adj2,
) where {B<:KernelAbstractions.Backend}
    if !isempty(f.itr)
        batch_size = size(X, 2)
        kerh_batch(backend)(y1, y2, f.f, f.itr, X, Θ, adj, adj2; ndrange = (length(f.itr), batch_size))
    end
end

function shessian_batch!(
    backend::B,
    y1,
    y2,
    f,
    X,
    Θ,
    adj::AbstractMatrix,
    adj2,
) where {B<:KernelAbstractions.Backend}
    if !isempty(f.itr)
        batch_size = size(X, 2)
        kerh2_batch(backend)(y1, y2, f.f, f.itr, X, Θ, adj, adj2; ndrange = (length(f.itr), batch_size))
    end
end