Better dist (#81)

* use better dist * update Dense api usage * add I to sigma
impICNF · Jul 16, 2022 · 55be9f0 · 55be9f0
1 parent 98c51ef
commit 55be9f0
Show file tree

Hide file tree

Showing 9 changed files with 15 additions and 12 deletions.
diff --git a/Project.toml b/Project.toml
@@ -11,6 +11,7 @@ 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"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -40,6 +41,7 @@ DataFrames = "1.2"
 DifferentialEquations = "7.0"
 Distributions = "0.25"
 DistributionsAD = "0.6"
+FillArrays = "0.13"
 Flux = "0.13"
 IterTools = "1.3"
 MLJBase = "0.20"

diff --git a/src/ICNF.jl b/src/ICNF.jl
@@ -8,6 +8,7 @@ module ICNF
         DifferentialEquations,
         Distributions,
         DistributionsAD,
+        FillArrays,
         Flux,
         IterTools,
         MLJBase,

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

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

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

diff --git a/src/ffjord.jl b/src/ffjord.jl
@@ -32,7 +32,7 @@ function FFJORD{T}(
         nn,
         nvars::Integer,
         ;
-        basedist::Distribution=MvNormal(zeros(T, nvars), Diagonal(ones(T, nvars))),
+        basedist::Distribution=MvNormal(Zeros{T}(nvars), one(T)*I),
         tspan::Tuple{T, T}=convert(Tuple{T, T}, (0, 1)),
 
         solvealg_test::SciMLBase.AbstractODEAlgorithm=default_solvealg,

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

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

diff --git a/test/smoke_tests.jl b/test/smoke_tests.jl
@@ -43,7 +43,7 @@
             nn = PlanarNN(nvars, tanh)
         else
             nn = Chain(
-                Dense(nvars, nvars, tanh),
+                Dense(nvars => nvars, tanh),
             )
         end
         icnf = mt{tp}(nn, nvars; acceleration=cr)
@@ -128,7 +128,7 @@
                 nn = PlanarNN(nvars, tanh)
             else
                 nn = Chain(
-                    Dense(nvars, nvars, tanh),
+                    Dense(nvars => nvars, tanh),
                 )
             end
             icnf = mt{tp}(nn, nvars; acceleration=cr)
@@ -147,7 +147,7 @@
                 nn = PlanarNN(nvars, tanh)
             else
                 nn = Chain(
-                    Dense(nvars, nvars, tanh),
+                    Dense(nvars => nvars, tanh),
                 )
             end
             icnf = mt{tp}(nn, nvars; acceleration=cr)
@@ -174,7 +174,7 @@
             nn = PlanarNN(nvars, tanh; cond=true)
         else
             nn = Chain(
-                Dense(nvars*2, nvars, tanh),
+                Dense(2*nvars => nvars, tanh),
             )
         end
         icnf = mt{tp}(nn, nvars; acceleration=cr)
@@ -262,7 +262,7 @@
                 nn = PlanarNN(nvars, tanh; cond=true)
             else
                 nn = Chain(
-                    Dense(nvars*2, nvars, tanh),
+                    Dense(2*nvars => nvars, tanh),
                 )
             end
             icnf = mt{tp}(nn, nvars; acceleration=cr)
@@ -281,7 +281,7 @@
                 nn = PlanarNN(nvars, tanh; cond=true)
             else
                 nn = Chain(
-                    Dense(nvars*2, nvars, tanh),
+                    Dense(2*nvars => nvars, tanh),
                 )
             end
             icnf = mt{tp}(nn, nvars; acceleration=cr)