Optimization.jl/lib/OptimizationSpeedMapping/src/OptimizationSpeedMapping.jl at 65e95eff40a0562368e01e49b1caa0984d52e42a · SciML/Optimization.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
102
103
104
105
106
107
108
109
110
111
112
module OptimizationSpeedMapping

using Reexport
@reexport using OptimizationBase
using SpeedMapping, SciMLBase

export SpeedMappingOpt

struct SpeedMappingOpt end

SciMLBase.allowsbounds(::SpeedMappingOpt) = true
SciMLBase.allowscallback(::SpeedMappingOpt) = false
SciMLBase.has_init(opt::SpeedMappingOpt) = true
SciMLBase.requiresgradient(opt::SpeedMappingOpt) = true

function _retcode(opt_res)
    status = opt_res.status
    status === :first_order && return SciMLBase.ReturnCode.Success
    status === :max_time && return SciMLBase.ReturnCode.MaxTime
    status in (:max_iter, :max_eval) && return SciMLBase.ReturnCode.MaxIters
    return SciMLBase.ReturnCode.Failure
end

function __map_optimizer_args(
        cache::OptimizationBase.OptimizationCache,
        opt::SpeedMappingOpt;
        callback = nothing,
        maxiters::Union{Number, Nothing} = nothing,
        maxtime::Union{Number, Nothing} = nothing,
        abstol::Union{Number, Nothing} = nothing,
        reltol::Union{Number, Nothing} = nothing,
    )

    # add optimiser options from kwargs
    mapped_args = (;)

    if !(isnothing(maxiters))
        @info "maxiters defines maximum gradient calls for $(opt)"
        mapped_args = (; mapped_args..., maps_limit = maxiters)
    end

    if !(isnothing(maxtime))
        mapped_args = (; mapped_args..., time_limit = maxtime)
    end

    if !isnothing(abstol)
        @SciMLMessage(
            lazy"common abstol is currently not used by $(opt)",
            cache.verbose,
            :unsupported_kwargs
        )
    end

    if !isnothing(reltol)
        @SciMLMessage(
            lazy"common reltol is currently not used by $(opt)",
            cache.verbose,
            :unsupported_kwargs
        )
    end

    return mapped_args
end

function SciMLBase.__solve(cache::OptimizationCache{O}) where {O <: SpeedMappingOpt}
    local x

    _loss = function (θ)
        x = cache.f.f(θ, cache.p)
        return first(x)
    end

    if isnothing(cache.f.grad)
        @info "SpeedMapping's ForwardDiff AD backend is used to calculate the gradient information."
    end

    maxiters = OptimizationBase._check_and_convert_maxiters(cache.solver_args.maxiters)
    maxtime = OptimizationBase._check_and_convert_maxtime(cache.solver_args.maxtime)
    opt_args = __map_optimizer_args(
        cache,
        cache.opt,
        maxiters = maxiters,
        maxtime = maxtime,
        abstol = cache.solver_args.abstol,
        reltol = cache.solver_args.reltol;
        cache.solver_args...,
    )

    t0 = time()
    opt_res = SpeedMapping.speedmapping(
        cache.u0;
        f = _loss,
        (g!) = cache.f.grad,
        lower = cache.lb,
        upper = cache.ub,
        opt_args...,
    )
    t1 = time()
    opt_ret = _retcode(opt_res)
    stats = OptimizationBase.OptimizationStats(; time = t1 - t0)
    return SciMLBase.build_solution(
        cache,
        cache.opt,
        opt_res.minimizer,
        _loss(opt_res.minimizer);
        original = opt_res,
        retcode = opt_ret,
        stats = stats,
    )
end

end