Use Adapt to keep type on GPU

Fix data creation

Author: SCiarella

Project.toml

@@ -4,6 +4,7 @@ authors = ["Simone Ciarella <[email protected]>, Luisa Orozco <l.oroz
 version = "0.0.3"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
@@ -34,7 +35,6 @@ YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AttentionLayer = "3ee63b08-73c5-50c8-acc9-f395aa68c39a"
 CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 ConvolutionalNeuralOperators = "d769ba41-1544-53e8-a779-241a28c31cef"
@@ -55,8 +55,8 @@ NeuralClosure = {rev = "main", url = "https://github.com/DEEPDIP-project/NeuralC
 AttentionCNN = ["AttentionLayer"]
 CNO = ["ConvolutionalNeuralOperators"]
 CoupledNODECUDA_ext = ["cuDNN", "CUDSS"]
-NavierStokes = ["IncompressibleNavierStokes", "NeuralClosure", "Adapt"]
-fno = ["NeuralOperators", "Adapt"]
+NavierStokes = ["IncompressibleNavierStokes", "NeuralClosure"]
+fno = ["NeuralOperators"]
 
 [compat]
 Adapt = "4"
@@ -104,4 +104,4 @@ NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "IncompressibleNavierStokes", "NeuralClosure", "Adapt", "cuDNN", "NeuralOperators"]
+test = ["Test", "IncompressibleNavierStokes", "NeuralClosure", "cuDNN", "NeuralOperators"]

ext/NavierStokes/create_data.jl

@@ -1,4 +1,5 @@
 using DifferentialEquations
+using JLD2
 using IncompressibleNavierStokes: right_hand_side!, apply_bc_u!, momentum!, project!
 
 function create_les_data_projected(;
@@ -75,9 +76,9 @@ function create_les_data_projected(;
         t in tdatapoint && return true
         return false
     end
-    all_ules = Array{T}(undef, (nles[1] + 2, nles[1]+2, D, length(tdatapoint)-1))
-    all_c = Array{T}(undef, (nles[1]+2, nles[1]+2, D, length(tdatapoint)-1))
-    all_t = Array{T}(undef, (length(tdatapoint)-1))
+    all_ules = Array{T}(undef, (nles[1] + 2, nles[1]+2, D, length(tdatapoint)))
+    all_c = Array{T}(undef, (nles[1]+2, nles[1]+2, D, length(tdatapoint)))
+    all_t = Array{T}(undef, (length(tdatapoint)))
     idx = Ref(1)
     Fdns = INS.create_right_hand_side(dns, psolver)
     p = scalarfield(les)
@@ -118,6 +119,9 @@ function create_les_data_projected(;
     u_current = u  # Initial condition
     prob = ODEProblem(rhs!, u_current, nothing, nothing)
 
+    # Store the data at t=0
+    filter_callback((; u = u, t = T(0)))
+
     any(u -> any(isnan, u), u_current) &&
         @warn "Solution contains NaNs. Probably dt is too large."
 
@@ -134,7 +138,7 @@ function create_les_data_projected(;
 
         sol = solve(
             prob, sciml_solver; u0 = u_current, p = nothing,
-            adaptive = false, dt = Δt, save_end = true, callback = cb,
+            adaptive = true, save_end = true, callback = cb,
             tspan = tspan_chunk, tstops = tdatapoint
         )
 

src/models/cnn.jl

@@ -24,7 +24,7 @@ A tuple `(chain, params, state)` where
 - `state`: The state of the model.
 """
 function cnn(;
-        T = Float32,
+        T,
         D,
         data_ch,
         radii,
@@ -37,12 +37,13 @@ function cnn(;
     r, c, σ, b = radii, channels, activations, use_bias
 
     if use_cuda
-        dev = Lux.gpu_device()
+        dev = x -> adapt(CuArray, x)
     else
         dev = Lux.cpu_device()
     end
 
-    @warn "*** CNN is using the following device: $(dev) "
+    T = eltype(T(0.0))
+    @warn "*** CNN is using the following device: $(dev) and type $(T)"
 
     # Weight initializer
     glorot_uniform_T(rng::Random.AbstractRNG, dims...) = glorot_uniform(rng, T, dims...)
@@ -71,6 +72,7 @@ function cnn(;
         )...,
         decollocate
     )
+
     chain = Chain(layers...)
     params, state = Lux.setup(rng, chain)
     state = state |> dev

src/train.jl

@@ -4,6 +4,7 @@ using Zygote: Zygote
 using Optimization: Optimization
 using OptimizationOptimisers: OptimizationOptimisers
 using ChainRulesCore: ignore_derivatives
+using Adapt: adapt
 
 function train(model, ps, st, train_dataloader, loss_function;
         nepochs = 50,
@@ -12,7 +13,7 @@ function train(model, ps, st, train_dataloader, loss_function;
         cpu::Bool = false,
         λ = nothing,
         kwargs...)
-    dev = cpu ? identity : Lux.gpu_device()
+    dev = cpu ? identity : x -> adapt(CuArray, x)
     if !cpu
         ps, st = (ps, st) .|> dev
     end