add planar and radial flow; updating docs

Author: zuhengxu

docs/src/api.md

@@ -57,8 +57,8 @@ NormalizingFlows.optimize
 
 ## Available Flows
 
-`NormalizingFlows.jl` provides two commonly used normalizing flows: `RealNVP` and 
-`Neural Spline Flow (NSF)`.
+`NormalizingFlows.jl` provides two commonly used normalizing flows---`RealNVP` and
+`Neural Spline Flow (NSF)`---and two simple flows---`Planar Flow` and `Radial Flow`.
 
 ### RealNVP (Affine Coupling Flow)
 
@@ -78,7 +78,14 @@ NormalizingFlows.NSF_layer
 NormalizingFlows.NeuralSplineCoupling
 ```
 
-##  Utility Functions
+#### Planar and Radial Flows
+
+```@docs
+NormalizingFlows.planarflow
+NormalizingFlows.radialflow
+```
+
+## Utility Functions
 
 ```@docs
 NormalizingFlows.create_flow

example/demo_RealNVP.jl

@@ -48,7 +48,7 @@ adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
 checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
     rng, 
-    elbo,        # using elbo_batch instead of elbo achieves 4-5 times speedup 
+    elbo_batch,        # using elbo_batch instead of elbo achieves 4-5 times speedup 
     flow,
     logp,
     sample_per_iter;

example/demo_planar_flow.jl

@@ -20,16 +20,9 @@ logp = Base.Fix1(logpdf, target)
 ######################################
 # setup planar flow
 ######################################
-function create_planar_flow(n_layers::Int, q₀)
-    d = length(q₀)
-    Ls = [PlanarLayer(d) for _ in 1:n_layers]
-    ts = reduce(∘, Ls)
-    return transformed(q₀, ts)
-end
-
 @leaf MvNormal
 q0 = MvNormal(zeros(T, 2), ones(T, 2))
-flow = create_planar_flow(10, q0)
+flow = planarflow(q0, 10; paramtype=T)
 flow_untrained = deepcopy(flow)
 
 ######################################

example/demo_radial_flow.jl

@@ -19,17 +19,10 @@ logp = Base.Fix1(logpdf, target)
 ######################################
 # setup radial flow 
 ######################################
-function create_radial_flow(n_layers::Int, q₀)
-    d = length(q₀)
-    Ls = [RadialLayer(d) for _ in 1:n_layers]
-    ts = reduce(∘, Ls)
-    return transformed(q₀, ts)
-end
-
 # create a 10-layer radial flow
 @leaf MvNormal
 q0 = MvNormal(zeros(T, 2), ones(T, 2))
-flow = create_radial_flow(10, q0)
+flow = radialflow(q0, 10; paramtype=T)
 
 flow_untrained = deepcopy(flow)
 

src/NormalizingFlows.jl

@@ -129,12 +129,14 @@ end
 
 # interface of contructing common flow layers
 include("flows/utils.jl")
+include("flows/planar_radial.jl")
 include("flows/realnvp.jl")
 
 using MonotonicSplines
 include("flows/neuralspline.jl")
 
 export create_flow
+export planarflow, radialflow
 export AffineCoupling, RealNVP_layer, realnvp
 export NeuralSplineCoupling, NSF_layer, nsf
 

src/flows/planar_radial.jl

@@ -0,0 +1,60 @@
+"""
+    planarflow(q0, nlayers; paramtype = Float64)
+
+Construct a Planar Flow by stacking `nlayers` `Bijectors.PlanarLayer` blocks
+on top of a base distribution `q0`.
+
+Arguments
+- `q0::Distribution{Multivariate,Continuous}`: base distribution (e.g., `MvNormal(zeros(d), I)`).
+- `nlayers::Int`: number of planar layers to compose.
+
+Keyword Arguments
+- `paramtype::Type{T} = Float64`: parameter element type (use `Float32` for GPU friendliness).
+
+Returns
+- `Bijectors.TransformedDistribution` representing the planar flow.
+
+Example
+- `q0 = MvNormal(zeros(2), I); flow = planarflow(q0, 10)`
+- `x = rand(flow, 128); lp = logpdf(flow, x)`
+"""
+function planarflow(
+    q0::Distribution{Multivariate,Continuous},  
+    nlayers::Int;                   
+    paramtype::Type{T} = Float64,   
+) where {T<:AbstractFloat}
+    dim = length(q0)
+    Ls = [Flux._paramtype(paramtype, Bijectors.PlanarLayer(dim)) for _ in 1:nlayers]
+    return create_flow(Ls, q0)
+end
+
+
+"""
+    radialflow(q0, nlayers; paramtype = Float64)
+
+Construct a Radial Flow by stacking `nlayers` `Bijectors.RadialLayer` blocks
+on top of a base distribution `q0`.
+
+Arguments
+- `q0::Distribution{Multivariate,Continuous}`: base distribution (e.g., `MvNormal(zeros(d), I)`).
+- `nlayers::Int`: number of radial layers to compose.
+
+Keyword Arguments
+- `paramtype::Type{T} = Float64`: parameter element type (use `Float32` for GPU friendliness).
+
+Returns
+- `Bijectors.TransformedDistribution` representing the radial flow.
+
+Example
+- `q0 = MvNormal(zeros(2), I); flow = radialflow(q0, 6)`
+- `x = rand(flow); lp = logpdf(flow, x)`
+"""
+function radialflow(
+    q0::Distribution{Multivariate,Continuous},  
+    nlayers::Int;                   
+    paramtype::Type{T} = Float64,   
+) where {T<:AbstractFloat}
+    dim = length(q0)
+    Ls = [Flux._paramtype(paramtype, Bijectors.RadialLayer(dim)) for _ in 1:nlayers]
+    return create_flow(Ls, q0)
+end

test/flow.jl

@@ -128,3 +128,137 @@ end
         end
     end
 end
+
+
+
+@testset "Planar flow" begin
+    Random.seed!(123)
+    
+    dim = 5
+    nlayers = 10
+    for T in [Float32, Float64]
+        # Create a nsf 
+        q₀ = MvNormal(zeros(T, dim), I)
+        @leaf MvNormal
+
+        flow = NormalizingFlows.planarflow(q₀, nlayers; paramtype=T)
+
+        @testset "Sampling and density estimation for type: $T" begin
+            ys = rand(flow, 100) 
+            ℓs = logpdf(flow, ys)
+
+            @test size(ys) == (dim, 100)
+            @test length(ℓs) == 100            
+
+            @test eltype(ys) == T
+            @test eltype(ℓs) == T
+        end
+            
+
+        @testset "Inverse compatibility for type: $T" begin
+            x = rand(q₀)
+            y, lj_fwd = Bijectors.with_logabsdet_jacobian(flow.transform, x)
+            x_reconstructed, lj_bwd = Bijectors.with_logabsdet_jacobian(inverse(flow.transform), y)
+
+            @test x ≈ x_reconstructed rtol=1e-4
+            @test lj_fwd ≈ -lj_bwd rtol=1e-4
+
+            x_batch = rand(q₀, 10)
+            y_batch, ljs_fwd = Bijectors.with_logabsdet_jacobian(flow.transform, x_batch)
+            x_batch_reconstructed, ljs_bwd = Bijectors.with_logabsdet_jacobian(inverse(flow.transform), y_batch)
+
+            @test x_batch ≈ x_batch_reconstructed rtol=1e-4
+            @test ljs_fwd ≈ -ljs_bwd rtol=1e-4
+        end
+
+
+        @testset "ELBO test for type: $T" begin
+            μ = randn(T, dim)
+            Σ = Diagonal(rand(T, dim) .+ T(1e-3))
+            target = MvNormal(μ, Σ)
+            logp(z) = logpdf(target, z)
+
+            # Compute ELBO
+            batchsize = 64
+            elbo_value = elbo(Random.default_rng(), flow, logp, batchsize)
+            elbo_batch_value = elbo_batch(Random.default_rng(), flow, logp, batchsize)
+
+            # test when batchsize == 1
+            batchsize_single = 1
+            elbo_value_single = elbo(Random.default_rng(), flow, logp, batchsize_single)
+
+            # test elbo_value is not NaN and not Inf
+            @test isfinite(elbo_value)
+            @test isfinite(elbo_batch_value)
+            @test isfinite(elbo_value_single)
+        end
+    end
+end
+
+
+
+@testset "Radial flow" begin
+    Random.seed!(123)
+    
+    dim = 5
+    nlayers = 10
+    for T in [Float32, Float64]
+        # Create a nsf 
+        q₀ = MvNormal(zeros(T, dim), I)
+        @leaf MvNormal
+
+        flow = NormalizingFlows.radialflow(q₀, nlayers; paramtype=T)
+
+        @testset "Sampling and density estimation for type: $T" begin
+            ys = rand(flow, 100) 
+            ℓs = logpdf(flow, ys)
+
+            @test size(ys) == (dim, 100)
+            @test length(ℓs) == 100            
+
+            @test eltype(ys) == T
+            @test eltype(ℓs) == T
+        end
+            
+
+        @testset "Inverse compatibility for type: $T" begin
+            x = rand(q₀)
+            y, lj_fwd = Bijectors.with_logabsdet_jacobian(flow.transform, x)
+            x_reconstructed, lj_bwd = Bijectors.with_logabsdet_jacobian(inverse(flow.transform), y)
+
+            @test x ≈ x_reconstructed rtol=1e-4
+            @test lj_fwd ≈ -lj_bwd rtol=1e-4
+
+            x_batch = rand(q₀, 10)
+            y_batch, ljs_fwd = Bijectors.with_logabsdet_jacobian(flow.transform, x_batch)
+            x_batch_reconstructed, ljs_bwd = Bijectors.with_logabsdet_jacobian(inverse(flow.transform), y_batch)
+
+            @test x_batch ≈ x_batch_reconstructed rtol=1e-4
+            @test ljs_fwd ≈ -ljs_bwd rtol=1e-4
+        end
+
+
+        @testset "ELBO test for type: $T" begin
+            μ = randn(T, dim)
+            Σ = Diagonal(rand(T, dim) .+ T(1e-3))
+            target = MvNormal(μ, Σ)
+            logp(z) = logpdf(target, z)
+
+            # Compute ELBO
+            batchsize = 64
+            elbo_value = elbo(Random.default_rng(), flow, logp, batchsize)
+            elbo_batch_value = elbo_batch(Random.default_rng(), flow, logp, batchsize)
+
+            # test when batchsize == 1
+            batchsize_single = 1
+            elbo_value_single = elbo(Random.default_rng(), flow, logp, batchsize_single)
+
+            # test elbo_value is not NaN and not Inf
+            @test isfinite(elbo_value)
+            @test isfinite(elbo_batch_value)
+            @test isfinite(elbo_value_single)
+        end
+    end
+end
+
+