Support Optimisers.jl optimizers in DeepSplitting

- Add Optimisers.jl as a dependency
- Add _copy and _get_eta overloads for Optimisers.AbstractRule
- Add constructor for Optimisers.jl optimizers
- Update docs to use Flux.Optimise.Adam explicitly

Author: ChrisRackauckas

Project.toml

@@ -9,6 +9,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -31,6 +32,7 @@ DocStringExtensions = "0.9.3"
 Flux = "0.14.16, 0.15, 0.16"
 Functors = "0.4.11, 0.5"
 LinearAlgebra = "1.10"
+Optimisers = "0.3, 0.4"
 Random = "1.10"
 Reexport = "1.2.2"
 SafeTestsets = "0.1"

docs/src/getting_started.md

@@ -99,7 +99,7 @@ nn = Flux.Chain(Dense(d, hls, tanh),
     Dense(hls, hls, tanh),
     Dense(hls, 1)) # neural network used by the scheme
 
-opt = ADAM(1e-2)
+opt = Flux.Optimise.Adam(1e-2)
 
 ## Definition of the algorithm
 alg = DeepSplitting(nn,

src/DeepSplitting.jl

@@ -1,8 +1,18 @@
+using Optimisers
+
 _copy(t::Tuple) = t
 _copy(t) = t
 function _copy(opt::O) where {O <: Flux.Optimise.AbstractOptimiser}
     return O([_copy(getfield(opt, fn)) for fn in fieldnames(typeof(opt))]...)
 end
+# Support for new-style Optimisers.jl optimizers
+function _copy(opt::O) where {O <: Optimisers.AbstractRule}
+    return O([_copy(getfield(opt, fn)) for fn in fieldnames(typeof(opt))]...)
+end
+
+# Helper to get learning rate from either optimizer type
+_get_eta(opt::Flux.Optimise.AbstractOptimiser) = opt.eta
+_get_eta(opt::Optimisers.AbstractRule) = opt.eta
 
 """
     DeepSplitting(nn, K=1, opt = Flux.Optimise.Adam(0.01), λs = nothing, mc_sample =  NoSampling())
@@ -38,6 +48,7 @@ struct DeepSplitting{NN, F, O, L, MCS} <: HighDimPDEAlgorithm
     mc_sample!::MCS # Monte Carlo sample
 end
 
+# Constructor for old-style Flux.Optimise optimizers
 function DeepSplitting(
         nn;
         K = 1,
@@ -48,7 +59,22 @@ function DeepSplitting(
         O <: Flux.Optimise.AbstractOptimiser,
         L <: Union{Nothing, Vector{N}} where {N <: Number},
     }
-    isnothing(λs) ? λs = [opt.eta] : nothing
+    isnothing(λs) ? λs = [_get_eta(opt)] : nothing
+    return DeepSplitting(nn, K, opt, λs, mc_sample)
+end
+
+# Constructor for new-style Optimisers.jl optimizers
+function DeepSplitting(
+        nn;
+        K = 1,
+        opt::O,
+        λs::L = nothing,
+        mc_sample = NoSampling()
+    ) where {
+        O <: Optimisers.AbstractRule,
+        L <: Union{Nothing, Vector{N}} where {N <: Number},
+    }
+    isnothing(λs) ? λs = [_get_eta(opt)] : nothing
     return DeepSplitting(nn, K, opt, λs, mc_sample)
 end