Optimization improvements

[mhauru]

Implementation of the work started in #2214. Work in progress. Needs

• More and better tests.
• A general clean-up.
• Docstrings all over.
• Taking care of the TODOs littered around.

@torfjelde

Closes #2227

[mhauru]

This is at a stage where getting a review would be good. Please be pedantic and show no mercy, I need to learn the ways of the project. Some questions I've been wondering about myself:

1. Functions names, especially exported ones. estimate_mode and link were two that I hesitated on.
2. What's the convention for when to use implicit and explicit imports? I see a lot of source files mixing the two.
3. Do we want to keep using ModeResult? That's the type that OptimInterface.jl returns its results in. I made Optimization.jl return the same type for compatibility, but would probably have designed the return type differently if given free rein.
4. I've copied over tests from OptimInterface.jl's tests with minimal adaptation, resulting in some overlap. If the whole Optim.jl interface is on its way out though, might be worth leaving the overlap there, so that one can delete the OptimInterface.jl tests and not lose anything significant.
5. Depending on what we do with ModeResult and OptimInterface.jl tests, I might split off some of the tests to a new file of tests just for OptimizationCore.jl.

[torfjelde]

Having a look now:)

[torfjelde]

Alrighty, here comes comments! 😬

Do let me know if anything is unclear, as there are quite a few of them.

I think going forward I'll try to keep the reviews a bit "themed" so that there isn't too much feedback at once, but we instead can do several iterations. But given that I'm not feeling 100%, I figured I'd do as much reviewing as possible when I'm feeling okay-ish.

Overall, I also have a few comments:

1. See https://docs.julialang.org/en/v1/manual/documentation/#Writing-Documentation for official docs on docstrings 👍
2. Though I can see that the ModeEstimationProblem is in a way a nice abstraction, now that we're going to depend fully on Optimization.jl, it seems a bit overkill in terms of what is needed. IIUC it's only really used in estimate_mode, and so can't we just put all this functionailty in estimate_mode directly? IMO this will make the code a more readable and easier to follow, as it just becomes "1) set up SciMLBase.OptimizationProblem, 2) solve it, and 3) convert into human readable format".
3. I'd recommend merging the files OptimisationCore.jl and Optimisation.jl into a single one. Moreover, in particular for the sake of reviewing, try to avoid copy-pasting code form existing files to other files before the majority of the reviewing has been done, as it makes it very difficult to see exactly what is new code and what is simply moved code.

Hope this is helpful!

[mhauru]

Thanks! I'll just leave comments now an get to fixing things later. Might take until Monday, tomorrow is quite full of other stuff.

Though I can see that the ModeEstimationProblem is in a way a nice abstraction, now that we're going to depend fully on Optimization.jl, it seems a bit overkill in terms of what is needed. IIUC it's only really used in estimate_mode, and so can't we just put all this functionailty in estimate_mode directly? IMO this will make the code a more readable and easier to follow, as it just becomes "1) set up SciMLBase.OptimizationProblem, 2) solve it, and 3) convert into human readable format".

We can put all of it in estimate_mode, yes, though I would at least keep some subfunctions to separate concerns. I do find the abstraction clarifies it for me. The main reason for the struct was that whether the objective function and the initial value (and correspondingly the returned solution) have been transformed to unconstrained space or not need to go hand in hand. The initial value is a bit nasty in that it's just a vector of numbers, that carries no information about the space it lives in. By putting it all in one struct you have to try harder to accidentally end up in a situation where you run an unconstrained space optimisation with the constrained space initial values, or some other ill combination of things.

I'd recommend merging the files OptimisationCore.jl and Optimisation.jl into a single one. Moreover, in particular for the sake of reviewing, try to avoid copy-pasting code form existing files to other files before the majority of the reviewing has been done, as it makes it very difficult to see exactly what is new code and what is simply moved code.

Fair point about copy-paste, sorry about that. I would usually make moving and editing separate commits, but this time the code went through a liquefy-and-reassemble stage in the beginning, and the first few commits do a bit of everything.

[torfjelde]

We can put all of it in estimate_mode, yes, though I would at least keep some subfunctions to separate concerns. I do find the abstraction clarifies it for me

I can see that; as I said, I think they make sense. But it does add quite a bit of an overhead for a uninitiated reader who just wants to see what's up.

he main reason for the struct was that whether the objective function and the initial value (and correspondingly the returned solution) have been transformed to unconstrained space or not need to go hand in hand.

Makes a lot of sense! But I think after one of the changes suggested, this shouldn't really be a worry anymore.

Specifically, you can just do the following once:

    ld = p.log_density
    varinfo = DynamicPPL.link(
        DynamicPPL.unflatten(ld.varinfo, copy(p.init_value))
        ld.model
    )
    ld = Accessors.@set ld.varinfo = varinfo

if we need to work in unconstrained space, and then use

    varinfo_new = DynamicPPL.unflatten(varinfo, solution.u)
    # `getparams` performs the invlinking, if needed, etc.
    vns_vals_iter = getparams(prob.model, varinfo_new)
    syms = map(Symbol ∘ first, vns_vals_iter)
    vals = map(last, vns_vals_iter)
    return ModeResult(
        NamedArray(vals, syms),
        solution,
        -solution.objective,
        prob.log_density
    )

(copy-paste from inline comments above) once you need to convert to a ModeResult. With this, there shouldn't be any other handling of unconstrained, etc.

Does that make sense?

Fair point about copy-paste, sorry about that. I would usually make moving and editing separate commits, but this time the code went through a liquefy-and-reassemble stage in the beginning, and the first few commits do a bit of everything.

No worries:) I also should have mentioned this before, and I every much know what you mean:)

[coveralls]

Pull Request Test Coverage Report for Build 9221611696

Details

• 0 of 120 (0.0%) changed or added relevant lines in 2 files are covered.
• 5 unchanged lines in 2 files lost coverage.
• Overall coverage remained the same at 0.0%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
ext/TuringOptimExt.jl	0	20	0.0%
src/optimisation/Optimisation.jl	0	100	0.0%

Files with Coverage Reduction	New Missed Lines	%
ext/TuringOptimExt.jl	2	0.0%
src/optimisation/Optimisation.jl	3	0.0%

Totals
Change from base Build 9141965383:	0.0%
Covered Lines:	0
Relevant Lines:	1501

---

💛 - Coveralls

[torfjelde]

miiiinor final comments, but looking really good now:)

[yebai]

Looks good — feel free to merge!

[torfjelde]

Fantastic stuff @mhauru 👏 I'm very happy about this change and you've done really great here:)

And sorry about the delayed reply; had a busy week with presentations, as you know.

Feel free to merge (and we can bump the Project.toml afterwards to avoid having to re-run tests)