Update to [email protected] (#43)

Author: SCiarella

Project.toml

@@ -33,20 +33,20 @@ CairoMakie = "0.12, 0.13"
 ChainRulesCore = "1.25.1"
 ChainRulesTestUtils = "1.13.0"
 ComponentArrays = "0.15"
-DifferentialEquations = "7.16.0"
+DifferentialEquations = "7"
 FFTW = "1"
-Images = "0.26.2"
+Images = "0.26"
 JuliaFormatter = "2"
-KernelAbstractions = "0.9.34"
+KernelAbstractions = "0.9"
 Lux = "1"
-LuxCUDA = "0.3.3"
+LuxCUDA = "0.3"
 LuxCore = "1"
 NNlib = "0.9"
 OpenSSL_jll = "3.0.13"
-Optimization = "4.1.1"
-OptimizationOptimisers = "0.3.7"
-TestImages = "1.9.0"
-Zygote = "0.6.76"
+Optimization = "4"
+OptimizationOptimisers = "0.3"
+TestImages = "1.9"
+Zygote = "0.7"
 julia = "1.11"
 
 [extras]

src/convolution.jl

@@ -66,7 +66,8 @@ function ChainRulesCore.rrule(::typeof(convolve), x, k)
     fft_k = fft(k, (2, 3))
 
     function convolve_pb(y_bar)
-        ffty_bar = fft(y_bar, (1, 2))
+        yb = unthunk(y_bar)
+        ffty_bar = fft(yb, (1, 2))
 
         if CUDA.functional() && k isa CuArray
             x_bar_re = CUDA.zeros(Float32, size(x))
@@ -150,15 +151,15 @@ end
 
 function apply_masked_convolution(y, k, mask)
     # to get the correct k i have to reshape+mask+trim
-    # TODO: i don't like this...
-    # ! Zygote does not like that you reuse variable names so, this makes it even uglier with the definition of k2 and k3
-    # ! also Zygote wants the mask to be explicitely defined as a vector so i have to pull it out from the tuple via mask=masks[i]
+    # ! Zygote does not like that you reuse variable names so k2 and k3 needs to be defined
+    # ! also Zygote wants the mask to be explicitely defined as a vector so mask_kernel is needed
 
     # Apply the mask to the kernel
     k2 = mask_kernel(k, mask)
 
-    # Adjust the kernel size to match the input dimensions
+    ## Adjust the kernel size to match the input dimensions
     k3 = trim_kernel(k2, size(y))
+    #k3 = k2
 
     # Apply the convolution
     y = convolve(y, k3)
@@ -178,20 +179,20 @@ end
 
 function ChainRulesCore.rrule(::typeof(trim_kernel), k, sizex)
     y = trim_kernel(k, sizex)
-    if k isa CuArray
-        k_bar = CUDA.zeros(Float32, size(k))
-    else
-        k_bar = zeros(Float32, size(k))
-    end
+    k_bar = similar(k, Float32)
 
     function trim_kernel_pullback(y_bar)
-        k_bar[:, 1:size(y_bar)[2], 1:size(y_bar)[3]] .= y_bar
+        yb = unthunk(y_bar)
+        sz2, sz3 = size(yb, 2), size(yb, 3)
+        k_bar .= 0  # clear first to be safe
+        k_bar[:, 1:sz2, 1:sz3] .= yb
         return NoTangent(), k_bar, NoTangent()
     end
     return y, trim_kernel_pullback
 end
 
 
+
 function mask_kernel(k, mask)
     permutedims(permutedims(k, [2, 3, 1]) .* mask, [3, 1, 2])
 end

src/downsample.jl

@@ -56,7 +56,7 @@ function ChainRulesCore.rrule(
 
         dk_pb!(backend, workgroupsize)(
             x_filter_bar,
-            result_bar,
+            unthunk(result_bar),
             down_factor;
             ndrange = downsampled_size,
         )

src/models.jl

@@ -383,7 +383,7 @@ function ((;)::CNO)(x, params, state)
         )
         # concatenate with the corresponding bottleneck
         y = cat(y, bottlenecks_out[i], dims = D + 1)
-        # apply the last bottleneck
+        # apply the last bottleneck that combines the two branches
         # ! do not forget to reverse the bottleneck ranges
         y = apply_masked_convolution(
             y,

test/data_train.jld2

@@ -8,6 +8,12 @@ Don't add your tests to runtests.jl. Instead, create files named
 
 The file will be automatically included inside a `@testset` with title "Title For My Test".
 =#
+
+# Helper function to check if a variable is on the GPU
+function is_on_gpu(x)
+    return x isa CuArray || (x isa SubArray && is_on_gpu(x.parent))
+end
+
 for (root, dirs, files) in walkdir(@__DIR__)
     for file in files
         if isnothing(match(r"^test-.*\.jl$", file))

test/runtests.jl

@@ -82,42 +82,41 @@ batch = 4
         pairs = @. Symbol(k) => v
         (; pairs...)
     end
-    data_train = []
-    data_i = namedtupleload("data_train.jld2")
-    push!(data_train, hcat(data_i))
+    data_train = load("data_train.jld2", "data_train")
 
     # Create the io array
     NS = Base.get_extension(CoupledNODE, :NavierStokes)
-    io_train = NS.create_io_arrays_posteriori(data_train, setup)
+    io_train = NS.create_io_arrays_posteriori(data_train, setup[1])
 
     # Create the dataloader
     θ = device(copy(θ_start))
     nunroll = 2
     nunroll_valid = 2
     dataloader_post = NS.create_dataloader_posteriori(
-        io_train[1];
+        io_train;
         nunroll = nunroll,
         rng = Random.Xoshiro(24),
-        device = device,
     )
+    u, t = dataloader_post()
 
     # Create the right hand side and the loss
     dudt_nn = NS.create_right_hand_side_with_closure(setup[1], psolver, closure, st)
+    griddims = ((:) for _ = 1:D)
     loss = CoupledNODE.create_loss_post_lux(
-        dudt_nn;
+        dudt_nn,
+        griddims,
+        griddims;
         sciml_solver = Tsit5(),
-        dt = T(conf["params"]["Δt"]),
-        use_cuda = false,
+        force_cpu = true,
     )
     callbackstate = trainstate = nothing
 
 
     # For testing reason, explicitely set up the probelm
     # Notice that this is automatically done in CoupledNODE
     u, t = dataloader_post()
-    griddims = ((:) for _ = 1:(ndims(u)-2))
-    x = u[griddims..., :, 1]
-    y = u[griddims..., :, 2:end] # remember to discard sol at the initial time step
+    x = u[griddims..., :, 1, 1]
+    y = u[griddims..., :, 1, 2:end] # remember to discard sol at the initial time step
     tspan, dt, prob, pred = nothing, nothing, nothing, nothing # initialize variable outside allowscalar do.
     dt = @views t[2:2] .- t[1:1]
     dt = only(Array(dt))
@@ -126,9 +125,7 @@ batch = 4
     end
     tspan = get_tspan(t)
     prob = ODEProblem(dudt_nn, x, tspan, θ)
-    pred = Array(
-        solve(prob, Tsit5(); u0 = x, p = θ, adaptive = false, saveat = Array(t), dt = dt),
-    )
+    pred = Array(solve(prob, Tsit5(); u0 = x, p = θ, adaptive = true, saveat = Array(t)))
 
     # Test the forward pass
     @test size(pred[:, :, :, 2:end]) == size(y)
@@ -226,54 +223,50 @@ end
         pairs = @. Symbol(k) => v
         (; pairs...)
     end
-    data_train = []
-    data_i = namedtupleload("data_train.jld2")
-    push!(data_train, hcat(data_i))
+    data_train = load("data_train.jld2", "data_train")
 
     # Create the io array
     NS = Base.get_extension(CoupledNODE, :NavierStokes)
-    io_train = NS.create_io_arrays_posteriori(data_train, setup)
+    io_train = NS.create_io_arrays_posteriori(data_train, setup[1], device)
 
     # Create the dataloader
     θ = device(copy(θ_start))
     nunroll = 2
     nunroll_valid = 2
     dataloader_post = NS.create_dataloader_posteriori(
-        io_train[1];
+        io_train;
         nunroll = nunroll,
         rng = Random.Xoshiro(24),
         device = device,
     )
+    u, t = dataloader_post()
 
     # Create the right hand side and the loss
     dudt_nn = NS.create_right_hand_side_with_closure(setup[1], psolver, closure, st)
+    griddims = ((:) for _ = 1:D)
     loss = CoupledNODE.create_loss_post_lux(
-        dudt_nn;
+        dudt_nn,
+        griddims,
+        griddims;
         sciml_solver = Tsit5(),
-        dt = T(conf["params"]["Δt"]),
-        use_cuda = true,
     )
     callbackstate = trainstate = nothing
 
 
     # For testing reason, explicitely set up the probelm
     # Notice that this is automatically done in CoupledNODE
-    u, t = dataloader_post()
-    griddims = ((:) for _ = 1:(ndims(u)-2))
-    x = u[griddims..., :, 1]
-    y = u[griddims..., :, 2:end] # remember to discard sol at the initial time step
-    tspan, dt, prob, pred = nothing, nothing, nothing, nothing # initialize variable outside allowscalar do.
-    dt = CUDA.allowscalar() do
-        t[2] .- t[1]
+    x, y = nothing, nothing
+    CUDA.allowscalar() do
+        x = u[griddims..., :, 1, 1]
+        y = u[griddims..., :, 1, 2:end] # remember to discard sol at the initial time step
     end
+    tspan, dt, prob, pred = nothing, nothing, nothing, nothing # initialize variable outside allowscalar do.
     function get_tspan(t)
         return (Array(t)[1], Array(t)[end])
     end
     tspan = get_tspan(t)
     prob = ODEProblem(dudt_nn, x, tspan, θ)
-    pred = Array(
-        solve(prob, Tsit5(); u0 = x, p = θ, adaptive = false, saveat = Array(t), dt = dt),
-    )
+    pred = Array(solve(prob, Tsit5(); u0 = x, p = θ, adaptive = true, saveat = Array(t)))
 
     # Test the forward pass
     @test size(pred[:, :, :, 2:end]) == size(y)

test/test-couplednode_posterior.jl

@@ -83,18 +83,16 @@ batch = 4
         pairs = @. Symbol(k) => v
         (; pairs...)
     end
-    data_train = []
-    data_i = namedtupleload("data_train.jld2")
-    push!(data_train, hcat(data_i))
+    data_train = load("data_train.jld2", "data_train")
 
     # Create the io array
     NS = Base.get_extension(CoupledNODE, :NavierStokes)
-    io_train = NS.create_io_arrays_priori(data_train, setup)
+    io_train = NS.create_io_arrays_priori(data_train, setup[1])
 
     # Create the dataloader
     θ = device(copy(θ_start))
     dataloader_prior = NS.create_dataloader_prior(
-        io_train[1];
+        io_train;
         batchsize = 4,
         rng = Random.Xoshiro(24),
         device = device,
@@ -186,25 +184,23 @@ end
         pairs = @. Symbol(k) => v
         (; pairs...)
     end
-    data_train = []
-    data_i = namedtupleload("data_train.jld2")
-    push!(data_train, hcat(data_i))
+    data_train = load("data_train.jld2", "data_train")
 
     # Create the io array
     NS = Base.get_extension(CoupledNODE, :NavierStokes)
-    io_train = NS.create_io_arrays_priori(data_train, setup)
+    io_train = NS.create_io_arrays_priori(data_train, setup[1], device)
 
     # Create the dataloader
     θ = device(copy(θ_start))
     dataloader_prior = NS.create_dataloader_prior(
-        io_train[1];
+        io_train;
         batchsize = 4,
         rng = Random.Xoshiro(24),
         device = device,
     )
     train_data_priori = dataloader_prior()
-    @test isa(train_data_priori[1], CuArray)
-    @test isa(train_data_priori[2], CuArray)
+    @test is_on_gpu(train_data_priori[1])
+    @test is_on_gpu(train_data_priori[2])
 
     l0 = CoupledNODE.loss_priori_lux(closure, θ, st, train_data_priori)[1]
     @test isnan(l0) == false

test/test-couplednode_prior.jl

@@ -112,9 +112,8 @@ end
         @test isa(y, CuArray)
 
 
-        return
-        u_in = rand(T, size(u))
-        tgt = rand(T, size(u))
+        u_in = CUDA.rand(T, size(u))
+        tgt = CUDA.rand(T, size(u))
         function loss(θ, batch = 16)
             yout = model(u_in, θ, st)[1]
             return sum(abs2, (yout .- tgt))

test/test-fullmodel.jl

@@ -8,7 +8,7 @@ using Lux: Lux
 using CUDA
 using LuxCUDA
 using ConvolutionalNeuralOperators:
-    convolve, apply_masked_convolution, trim_kernel, get_kernel
+    convolve, apply_masked_convolution, trim_kernel, get_kernel, mask_kernel
 using Zygote: Zygote
 using Test  # Importing the Test module for @test statements
 using AbstractFFTs: fft, ifft