update defaults (#87)

impICNF · Aug 14, 2022 · 4ef6577 · 4ef6577
1 parent 65c5758
commit 4ef6577
Show file tree

Hide file tree

Showing 10 changed files with 20 additions and 13 deletions.
diff --git a/Project.toml b/Project.toml
@@ -7,12 +7,12 @@ version = "0.1.0"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
@@ -36,12 +36,12 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Adapt = "3.3"
 CUDA = "3.7"
 DataFrames = "1.2"
-DifferentialEquations = "7.0"
 Distributions = "0.25"
 DistributionsAD = "0.6"
 FillArrays = "0.13"
 Flux = "0.13"
 IterTools = "1.3"
+LineSearches = "7.1"
 MLJBase = "0.20"
 MLJModelInterface = "1.3"
 MLUtils = "0.2"

diff --git a/src/ICNF.jl b/src/ICNF.jl
@@ -4,12 +4,12 @@ module ICNF
         Adapt,
         CUDA,
         DataFrames,
-        DifferentialEquations,
         Distributions,
         DistributionsAD,
         FillArrays,
         Flux,
         IterTools,
+        LineSearches,
         MLJBase,
         MLJModelInterface,
         MLUtils,

diff --git a/src/cond_ffjord.jl b/src/cond_ffjord.jl
@@ -28,7 +28,7 @@ function CondFFJORD{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/src/cond_planar.jl b/src/cond_planar.jl
@@ -28,7 +28,7 @@ function CondPlanar{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/src/cond_rnode.jl b/src/cond_rnode.jl
@@ -28,7 +28,7 @@ function CondRNODE{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/src/defaults.jl b/src/defaults.jl
@@ -1,15 +1,22 @@
+default_tspan = (0, 1)
 default_solvealg = Tsit5(
     ;
     thread=OrdinaryDiffEq.True(),
 )
 default_sensealg = InterpolatingAdjoint(
     ;
     autodiff=true,
+    checkpointing=false,
+    noisemixing=false,
     chunk_size=0,
     autojacvec=ZygoteVJP(),
 )
 default_optimizer = Dict(
-    FluxOptApp => Flux.AMSGrad(),
-    OptimOptApp => BFGS(),
-    SciMLOptApp => Optimisers.AMSGrad(),
+    FluxOptApp => Flux.AMSGrad(0.001, (0.9, 0.999), eps()),
+    OptimOptApp => BFGS(
+        alphaguess=InitialHagerZhang(),
+        linesearch=HagerZhang(),
+        manifold=Flat(),
+    ),
+    SciMLOptApp => Optimisers.AMSGrad(0.001, (0.9, 0.999), eps()),
 )
diff --git a/src/ffjord.jl b/src/ffjord.jl
@@ -30,7 +30,7 @@ function FFJORD{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/src/planar.jl b/src/planar.jl
@@ -53,7 +53,7 @@ function Planar{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/src/rnode.jl b/src/rnode.jl
@@ -30,7 +30,7 @@ function RNODE{T, AT}(
         nvars::Integer,
         ;
         basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
-        tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
+        tspan::Tuple{T, T}=convert(Tuple{T, T}, default_tspan),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,
         solvealg_train::SciMLBase.AbstractODEAlgorithm=default_solvealg,

diff --git a/test/smoke_tests.jl b/test/smoke_tests.jl
@@ -23,7 +23,7 @@
         Optimization.AutoTracker(),
         Optimization.AutoFiniteDiff(),
     ]
-    go_mds = Any[ICNF.default_optimizer[FluxOptApp], ICNF.default_optimizer[OptimOptApp]]
+    go_mds = Any[ICNF.default_optimizer[FluxOptApp], ICNF.default_optimizer[OptimOptApp], ICNF.default_optimizer[SciMLOptApp]]
     nvars_ = (1:2)
     n_epochs = 2
     batch_size = 8