Added OTFlow Layers

src/DiffEqFlux.jl

@@ -32,13 +32,14 @@
 
 include("ffjord.jl")
 include("neural_de.jl")
-
+include("otflow.jl")
 include("collocation.jl")
 include("multiple_shooting.jl")
 
 export NeuralODE, NeuralDSDE, NeuralSDE, NeuralCDDE, NeuralDAE, AugmentedNDELayer,
        NeuralODEMM
 export FFJORD, FFJORDDistribution
+export OTFlow, OTFlowDistribution
 export DimMover
 
 export EpanechnikovKernel, UniformKernel, TriangularKernel, QuarticKernel, TriweightKernel,

src/otflow.jl

@@ -0,0 +1,98 @@
+struct OTFlow <: AbstractLuxLayer
+    d::Int          # Input dimension
+    m::Int          # Hidden dimension
+    r::Int          # Rank for low-rank approximation
+end
+
+# Constructor with default rank
+OTFlow(d::Int, m::Int; r::Int=min(10,d)) = OTFlow(d, m, r)
+
+# Initialize parameters and states
+function Lux.initialparameters(rng::AbstractRNG, l::OTFlow)
+    w = randn(rng, Float32, l.m) .* 0.01
+    A = randn(rng, Float32, l.r, l.d + 1) .* 0.01
+    b = zeros(Float32, l.d + 1)
+    c = zero(Float32)
+    K0 = randn(rng, Float32, l.m, l.d + 1) .* 0.01
+    K1 = randn(rng, Float32, l.m, l.m) .* 0.01
+    b0 = zeros(Float32, l.m)
+    b1 = zeros(Float32, l.m)
+
+    ps = (w=w, A=A, b=b, c=c, K0=K0, K1=K1, b0=b0, b1=b1)
+    st = NamedTuple()
+    return ps, st
+end
+
+σ(x) = log(exp(x) + exp(-x))
+σ′(x) = tanh(x)
+σ′′(x) = 1 - tanh(x)^2
+
+function resnet_forward(x::AbstractVector, t::Real, ps)
+    s = vcat(x, t)
+    u0 = σ.(ps.K0 * s .+ ps.b0)
+    u1 = u0 .+ σ.(ps.K1 * u0 .+ ps.b1)  # h=1 as in paper
+    return u1
+end
+
+function potential(x::AbstractVector, t::Real, ps)
+    s = vcat(x, t)
+    N = resnet_forward(x, t, ps)
+    quadratic_term = 0.5 * s' * (ps.A' * ps.A) * s
+    linear_term = ps.b' * s
+    return ps.w' * N + quadratic_term + linear_term + ps.c
+end
+
+function gradient(x::AbstractVector, t::Real, ps, d::Int)
+    s = vcat(x, t)
+    u0 = σ.(ps.K0 * s .+ ps.b0)
+    z1 = ps.w .+ ps.K1' * (σ′.(ps.K1 * u0 .+ ps.b1) .* ps.w)
+    z0 = ps.K0' * (σ′.(ps.K0 * s .+ ps.b0) .* z1)
+
+    grad = z0 + (ps.A' * ps.A) * s + ps.b
+    return grad[1:d] 
+end
+
+function trace(x::AbstractVector, t::Real, ps, d::Int)
+    s = vcat(x, t)
+    u0 = σ.(ps.K0 * s .+ ps.b0)
+    z1 = ps.w .+ ps.K1' * (σ′.(ps.K1 * u0 .+ ps.b1) .* ps.w)
+
+    K0_E = ps.K0[:, 1:d]
+    A_E = ps.A[:, 1:d]
+
+    t0 = sum(σ′′.(ps.K0 * s .+ ps.b0) .* z1 .* (K0_E .^ 2))
+    J = Diagonal(σ′.(ps.K0 * s .+ ps.b0)) * K0_E
+    t1 = sum(σ′′.(ps.K1 * u0 .+ ps.b1) .* ps.w .* (ps.K1 * J) .^ 2)
+    trace_A = tr(A_E' * A_E)
+
+    return t0 + t1 + trace_A
+end
+
+function (l::OTFlow)(xt::Tuple{AbstractVector, Real}, ps, st)
+    x, t = xt
+    v = -gradient(x, t, ps, l.d)  # v = -∇Φ
+    tr = -trace(x, t, ps, l.d)   # tr(∇v) = -tr(∇²Φ)
+    return (v, tr), st
+end
+
+function simple_loss(x::AbstractVector, t::Real, l::OTFlow, ps)
+    (v, tr), _ = l((x, t), ps, NamedTuple())
+    return sum(v.^2) / 2 - tr
+end
+
+function manual_gradient(x::AbstractVector, t::Real, l::OTFlow, ps)
+    s = vcat(x, t)
+    u0 = σ.(ps.K0 * s .+ ps.b0)
+    u1 = u0 .+ σ.(ps.K1 * u0 .+ ps.b1)
+
+    v = -gradient(x, t, ps, l.d)
+    tr = -trace(x, t, ps, l.d)
+
+    # Simplified gradients (not full implementation)
+    grad_w = u1
+    grad_A = (ps.A * s) * s'
+
+    return (w=grad_w, A=grad_A, b=similar(ps.b), c=0.0, 
+            K0=zeros(l.m, l.d+1), K1=zeros(l.m, l.m), 
+            b0=zeros(l.m), b1=zeros(l.m))
+end

test/cnf_tests.jl

@@ -34,7 +34,7 @@ end
         return -mean(logpx)
     end
 
-    @testset "ADType: $(adtype)" for adtype in (Optimization.AutoForwardDiff(),
+    @testitem "ADType: $(adtype)" for adtype in (Optimization.AutoForwardDiff(),
         Optimization.AutoReverseDiff(), Optimization.AutoTracker(),
         Optimization.AutoZygote(), Optimization.AutoFiniteDiff())
         @testset "regularize = $(regularize) & monte_carlo = $(monte_carlo)" for regularize in (

test/collocation_tests.jl

@@ -53,4 +53,4 @@
             @test_throws ErrorException collocate_data(data, t, kernel, 0.001)
         end
     end
-end
+end

test/neural_dae_tests.jl

@@ -69,4 +69,4 @@
         optprob = Optimization.OptimizationProblem(optfunc, p)
         res = Optimization.solve(optprob, BFGS(; initial_stepnorm = 0.0001))
     end
-end
+end

test/neural_de_tests.jl

@@ -240,7 +240,7 @@ end
     @test first(layer(r, ps, st))[:, :, :, 1] == r[:, :, 1, :]
 end
 
-@testitem "Neural DE CUDA" tags=[:cuda] skip=:(using LuxCUDA; !LuxCUDA.functional()) begin
+@testitem "Neural DE CUDA" tags=[:cuda] begin
     using LuxCUDA, Zygote, OrdinaryDiffEq, StochasticDiffEq, Test, Random, ComponentArrays
     import Flux
 
@@ -323,4 +323,4 @@ end
             end
         end
     end
-end
+end

test/otflow_tests.jl

@@ -0,0 +1,65 @@
+@testitem "Tests for OTFlow Layer Functionality" begin
+    using Lux, LuxCore, Random, LinearAlgebra, Test, ComponentArrays, Flux, DiffEqFlux
+    rng = Xoshiro(0)
+    d = 2 
+    m = 4 
+    r = 2 
+    otflow = OTFlow(d, m; r=r)
+    ps, st = Lux.setup(rng, otflow)
+    ps = ComponentArray(ps)
+
+    x = Float32[1.0, 2.0]
+    t = 0.5f0
+
+    @testitem "Forward Pass" begin
+        (v, tr), st_new = otflow((x, t), ps, st)
+        @test length(v) == d 
+        @test isa(tr, Float32)
+        @test st_new == st 
+    end
+
+    @testitem "Potential Function" begin
+        phi = potential(x, t, ps)
+        @test isa(phi, Float32)
+    end
+
+    @testitem "Gradient Consistency" begin
+        grad = gradient(x, t, ps, d)
+        (v, _), _ = otflow((x, t), ps, st)
+        @test length(grad) == d
+        @test grad ≈ -v atol=1e-5  # v = -∇Φ
+    end
+
+    @testitem "Trace Consistency" begin
+        tr_manual = trace(x, t, ps, d)
+        (_, tr_forward), _ = otflow((x, t), ps, st)
+        @test tr_manual ≈ -tr_forward atol=1e-5
+    end
+
+    @testitem "ODE Integration" begin
+        x0 = Float32[1.0, 1.0]
+        tspan = (0.0f0, 1.0f0)
+        x_traj, t_vec = simple_ode_solve(otflow, x0, tspan, ps, st; dt=0.01f0)
+        @test size(x_traj) == (d, length(t_vec))
+        @test all(isfinite, x_traj)
+        @test x_traj[:, end] != x0
+    end
+
+    @testitem "Loss Function" begin
+        loss_val = simple_loss(x, t, otflow, ps)
+        @test isa(loss_val, Float32)
+        @test isfinite(loss_val)
+    end
+
+    @testitem "Manual Gradient" begin
+        grads = manual_gradient(x, t, otflow, ps)
+        @test haskey(grads, :w) && length(grads.w) == m
+        @test haskey(grads, :A) && size(grads.A) == (r, d+1)
+        @test haskey(grads, :b) && length(grads.b) == d+1
+        @test haskey(grads, :c) && isa(grads.c, Float32)
+        @test haskey(grads, :K0) && size(grads.K0) == (m, d+1)
+        @test haskey(grads, :K1) && size(grads.K1) == (m, m)
+        @test haskey(grads, :b0) && length(grads.b0) == m
+        @test haskey(grads, :b1) && length(grads.b1) == m
+    end
+end

test/stiff_nested_ad_tests.jl

@@ -27,7 +27,7 @@
         end
     end
 
-    @testset "Solver: $(nameof(typeof(solver)))" for solver in (
+    @testitem "Solver: $(nameof(typeof(solver)))" for solver in (
         KenCarp4(), Rodas5(), RadauIIA5())
         neuralde = NeuralODE(model, tspan, solver; saveat = t, reltol = 1e-7, abstol = 1e-9)
         ps, st = Lux.setup(Xoshiro(0), neuralde)