Repository: PRONTOLab/GB-25
Pull request #166: Correctness test scripts
Author: login: glwagner
Created: 2025-04-07T05:01:01Z | Merged: 2025-04-10T15:19:36Z



Diff:
diff --git a/.github/workflows/CompileOrRun.yml b/.github/workflows/CompileOrRun.yml
index 83bdc83c..9fc9e538 100644
--- a/.github/workflows/CompileOrRun.yml
+++ b/.github/workflows/CompileOrRun.yml
@@ -18,6 +18,11 @@ on:
         required: false
         default: false
         type: boolean
+      run_dir:
+        description: 'Directory the run script lives in'
+        required: false
+        default: 'simulations'
+        type: string
       grid_type:
         description: 'What type of grid to use'
         required: false
@@ -109,7 +114,7 @@ jobs:
         env:
           grid_type: ${{ inputs.grid_type }}
           XLA_FLAGS: ${{ inputs.sharded && '--xla_force_host_platform_device_count=4 --xla_dump_to=xla_dump' || '--xla_dump_to=xla_dump' }}
-          RUNDIR: ${{ inputs.sharded && 'sharding' || 'simulations' }}
+          RUNDIR: ${{ inputs.run_dir }}
         run: |
           earlyoom -m ${{ inputs.earlyoom_threshold }} -s 100 -r 120 --prefer 'julia' &
           julia --color=yes --project -O${{ inputs.julia_optlevel }} $RUNDIR/${{ inputs.sim_type }}_simulation_${{ inputs.compile_or_run }}.jl
diff --git a/.github/workflows/Correctness.yml b/.github/workflows/Correctness.yml
new file mode 100644
index 00000000..cb60de24
--- /dev/null
+++ b/.github/workflows/Correctness.yml
@@ -0,0 +1,51 @@
+name: "Correctness"
+
+on:
+  pull_request:
+    paths:
+      - '.github/workflows/Correctness.yml'
+      - '.github/workflows/CompileOrRun.yml'
+      - 'correctness/correctness_baroclinic_instability_simulation_run.jl'
+      - 'Project.toml'
+      - 'src/**'
+  push:
+    branches:
+      - main
+      - release-*
+    tags: '*'
+    paths:
+      - '.github/workflows/Correctness.yml'
+      - '.github/workflows/CompileOrRun.yml'
+      - 'correctness/correctness_baroclinic_instability_simulation_run.jl'
+      - 'Project.toml'
+      - 'src/**'
+  workflow_dispatch:
+  schedule:
+    - cron: '12 07,19 * * *'
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  run_correctness:
+    name: Correctness - Julia ${{ matrix.julia_version }} - ${{ matrix.grid_type }} - ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+          os: ['ubuntu-24.04', 'ubuntu-22.04-arm']
+          grid_type: ['simple_lat_lon']
+    uses: ./.github/workflows/CompileOrRun.yml
+    with:
+      sim_type: 'correctness_baroclinic_instability'
+      grid_type: ${{ matrix.grid_type }}
+      run_dir: 'correctness'
+      julia_version: '1.11'
+      os: ${{ matrix.os }}
+      xla_runtime: 'IFRT'
+      compile_or_run: 'run'
+      earlyoom_threshold: 4
+      julia_optlevel: 0
+
diff --git a/.github/workflows/Run.yml b/.github/workflows/Run.yml
index ff7ac6f2..911aeb60 100644
--- a/.github/workflows/Run.yml
+++ b/.github/workflows/Run.yml
@@ -81,10 +81,10 @@ jobs:
       sim_type: 'sharded_baroclinic_instability'
       grid_type: ${{ matrix.grid_type }}
       sharded: true
+      run_dir: 'sharding'
       julia_version: '1.11'
       os: ${{ matrix.os }}
       xla_runtime: 'IFRT'
       compile_or_run: 'run'
       earlyoom_threshold: 4
       julia_optlevel: 0
-
diff --git a/Project.toml b/Project.toml
index 065a1a71..bda426a4 100644
--- a/Project.toml
+++ b/Project.toml
@@ -41,11 +41,11 @@ Dates = "1"
 HDF5_jll = "=1.14.3"
 KernelAbstractions = "0.9.34"
 MPI = "0.20.22"
-Oceananigans = "0.96.15"
+Oceananigans = "0.96.19"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"
-Reactant = "0.2.66"
+Reactant = "0.2.68"
 SeawaterPolynomials = "0.3.9"
 Serialization = "1"
 
diff --git a/correctness/correctness_baroclinic_instability_simulation_run.jl b/correctness/correctness_baroclinic_instability_simulation_run.jl
new file mode 100644
index 00000000..9bbc140a
--- /dev/null
+++ b/correctness/correctness_baroclinic_instability_simulation_run.jl
@@ -0,0 +1,79 @@
+using GordonBell25
+using Oceananigans
+using Reactant
+
+throw_error = true
+include_halos = true
+rtol = 0
+atol = 1e-7
+
+model_kw = (
+    halo = (6, 6, 6),
+    Δt = 1e-9,
+)
+
+Nx = Ny = Nz = 16
+rarch = Oceananigans.Architectures.ReactantState()
+varch = CPU()
+rmodel = GordonBell25.baroclinic_instability_model(rarch, Nx, Ny, Nz; model_kw...)
+vmodel = GordonBell25.baroclinic_instability_model(varch, Nx, Ny, Nz; model_kw...)
+@show vmodel
+@show rmodel
+
+ui = 1e-3 .* rand(size(vmodel.velocities.u)...)
+vi = 1e-3 .* rand(size(vmodel.velocities.v)...)
+set!(vmodel, u=ui, v=vi)
+GordonBell25.sync_states!(rmodel, vmodel)
+
+@jit Oceananigans.initialize!(rmodel)
+Oceananigans.initialize!(vmodel)
+
+@jit Oceananigans.TimeSteppers.update_state!(rmodel)
+Oceananigans.TimeSteppers.update_state!(vmodel)
+
+@info "After initialization and update state:"
+GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
+
+GordonBell25.sync_states!(rmodel, vmodel)
+rfirst! = @compile sync=true raise=true GordonBell25.first_time_step!(rmodel)
+@time rfirst!(rmodel)
+@time GordonBell25.first_time_step!(vmodel)
+
+@info "After first time step:"
+GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
+
+rstep! = @compile sync=true raise=true GordonBell25.time_step!(rmodel)
+
+@info "Warm up:"
+@time rstep!(rmodel)
+@time rstep!(rmodel)
+@time GordonBell25.time_step!(vmodel)
+@time GordonBell25.time_step!(vmodel)
+
+@info "Time step with Reactant:"
+for _ in 1:10
+    @time rstep!(rmodel)
+end
+
+@info "Time step vanilla:"
+for _ in 1:10
+    @time GordonBell25.time_step!(vmodel)
+end
+
+@info "After 10 steps:"
+GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
+
+GordonBell25.sync_states!(rmodel, vmodel)
+@jit Oceananigans.TimeSteppers.update_state!(rmodel)
+
+@info "After syncing and updating state again:"
+GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
+
+Nt = 100
+rNt = ConcreteRNumber(Nt)
+rloop! = @compile sync=true raise=true GordonBell25.loop!(rmodel, rNt)
+@time rloop!(rmodel, rNt)
+@time GordonBell25.loop!(vmodel, Nt)
+
+@info "After a loop of 100 steps:"
+GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
diff --git a/src/GordonBell25.jl b/src/GordonBell25.jl
index 2faa7aeb..bc0be3ed 100644
--- a/src/GordonBell25.jl
+++ b/src/GordonBell25.jl
@@ -7,6 +7,7 @@ include("timestepping_utils.jl")
 include("data_free_ocean_climate_model.jl")
 include("baroclinic_instability_model.jl")
 include("sharding_utils.jl")
+include("correctness.jl")
 include("precompile.jl")
 
 end # module
diff --git a/src/correctness.jl b/src/correctness.jl
new file mode 100644
index 00000000..52140e2e
--- /dev/null
+++ b/src/correctness.jl
@@ -0,0 +1,123 @@
+using Reactant
+using Printf
+
+function compare_parent(name, ψ1, ψ2; rtol=1e-8, atol=sqrt(eps(eltype(ψ1))))
+    ψ1 = Array(parent(ψ1))
+    ψ2 = Array(parent(ψ2))
+    Nx, Ny, Nz = size(ψ1)
+    ψ2 = view(ψ2, 1:Nx, 1:Ny, 1:Nz) # assuming that ψ1 is smaller than ψ2
+    δ = ψ1 .- ψ2
+    idxs = findmax(abs, δ)[2]
+    approx_equal = isapprox(ψ1, ψ2; rtol, atol)
+    @printf("(%4s) ψ₁ ≈ ψ₂: %-5s, max|ψ₁|, max|ψ₂|: %.15e, %.15e, max|δ|: %.15e at %d %d %d \n",
+            name, approx_equal, maximum(abs, ψ1), maximum(abs, ψ2), maximum(abs, δ), idxs.I...)
+    return approx_equal
+end
+
+function compare_interior(name, ψ1, ψ2; rtol=1e-8, atol=sqrt(eps(eltype(ψ1))))
+    ψ1 = Array(interior(ψ1))
+    ψ2 = Array(interior(ψ2))
+    δ = ψ1 .- ψ2
+    idxs = findmax(abs, δ)[2]
+    approx_equal = isapprox(ψ1, ψ2; rtol, atol)
+    @printf("(%4s) ψ₁ ≈ ψ₂: %-5s, max|ψ₁|, max|ψ₂|: %.15e, %.15e, max|δ|: %.15e at %d %d %d \n",
+            name, approx_equal, maximum(abs, ψ1), maximum(abs, ψ2), maximum(abs, δ), idxs.I...)
+    return approx_equal
+end
+
+function compare_states(m1, m2; rtol=1e-8, atol=sqrt(eps(eltype(m1.grid))),
+                        include_halos=false, throw_error=false)
+
+    approx_equal = true
+    compare_fields = include_halos ? compare_parent : compare_interior
+
+    Ψ1 = Oceananigans.fields(m1)
+    Ψ2 = Oceananigans.fields(m2)
+
+    for name in keys(Ψ1)
+        approx_equal *= compare_fields(name, Ψ1[name], Ψ2[name]; rtol, atol)
+
+        if !(name ∈ (:w, :η))
+            Gⁿ1 = m1.timestepper.Gⁿ
+            Gⁿ2 = m2.timestepper.Gⁿ
+            approx_equal *= compare_fields("Gⁿ.$name", Gⁿ1[name], Gⁿ2[name]; rtol, atol)
+
+            G⁻1 = m1.timestepper.G⁻
+            G⁻2 = m2.timestepper.G⁻
+            approx_equal *= compare_fields("G⁻.$name", G⁻1[name], G⁻2[name]; rtol, atol)
+        end
+    end
+
+    if m1.free_surface isa Oceananigans.SplitExplicitFreeSurface
+        names = (:U, :V, :η)
+        Φ1 = NamedTuple(name => getproperty(m1.free_surface.filtered_state, name) for name in names)
+        Φ2 = NamedTuple(name => getproperty(m2.free_surface.filtered_state, name) for name in names)
+        for name in keys(Φ1)
+            approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
+        end
+    end
+
+    if m1.closure isa Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivity
+        names = (:κu, :κc, :κe, :Le, :Jᵇ)
+        Φ1 = NamedTuple(name => getproperty(m1.diffusivity_fields, name) for name in names)
+        Φ2 = NamedTuple(name => getproperty(m2.diffusivity_fields, name) for name in names)
+        for name in keys(Φ1)
+            approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
+        end
+    end
+
+    if m1.closure isa Oceananigans.TurbulenceClosures.TKEDissipationVerticalDiffusivity
+        names = (:κu, :κc, :κe, :κϵ, :Le, :Lϵ)
+        Φ1 = NamedTuple(name => getproperty(m1.diffusivity_fields, name) for name in names)
+        Φ2 = NamedTuple(name => getproperty(m2.diffusivity_fields, name) for name in names)
+        for name in keys(Φ1)
+            approx_equal *= compare_fields(name, Φ1[name], Φ2[name]; rtol, atol)
+        end
+    end
+
+    if throw_error && !approx_equal
+        error("There is a discrepancy between the models!")
+    end
+
+    return nothing
+end
+  
+function sync_states!(m1, m2)
+    Ψ1 = Oceananigans.fields(m1)
+    Ψ2 = Oceananigans.fields(m2)
+    for name in keys(Ψ1)
+        ψ1 = Ψ1[name]
+        ψ2 = Ψ2[name]
+        loc = Oceananigans.Fields.location(ψ1)
+        ψ2r = Reactant.to_rarray(interior(ψ2))
+        @jit copy_interior!(ψ1, ψ2r)
+    end
+    return nothing
+end
+
+function sync_parent_states!(m1, m2)
+    Ψ1 = Oceananigans.fields(m1)
+    Ψ2 = Oceananigans.fields(m2)
+    for name in keys(Ψ1)
+        ψ1p = parent(Ψ1[name])
+        ψ2p = parent(Ψ2[name])
+        copyto!(ψ1p, Reactant.to_rarray(ψ2p))	
+    end
+    
+    return nothing
+end
+  
+using KernelAbstractions
+
+function copy_interior!(rf, vf)
+    grid = rf.grid
+    arch = Oceananigans.Architectures.architecture(grid)
+    Oceananigans.Utils.launch!(arch, grid, size(rf), _copy_interior!, interior(rf), vf)
+    return nothing
+end
+
+@kernel function _copy_interior!(rf, vf)
+    i, j, k = @index(Global, NTuple)
+    @inbounds rf[i, j, k] = vf[i, j, k]
+end
+


Reviews:
user: login: glwagner
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T05:03:32Z
user: login: glwagner
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T05:04:02Z
user: login: glwagner
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T05:04:16Z
user: login: giordano
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:37:12Z
user: login: giordano
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T19:37:24Z
user: login: glwagner
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:54:43Z
user: login: giordano
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T13:55:54Z
user: login: glwagner
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T19:39:14Z

Review comments:
diff_hunk: @@ -8,6 +8,7 @@
     return N² * z + Δb * μ + 1e-2 * Δb * randn()
 end
 
+@inline random_velocity(λ, φ, z) = 1e-3 * randn()
path: src/baroclinic_instability_model.jl
user: login: glwagner
body: ```suggestion
```
created_at: 2025-04-07T05:03:31Z
author_association: COLLABORATOR
diff_hunk: @@ -51,10 +52,28 @@ function sync_states!(m1, m2)
     Ψ1 = Oceananigans.fields(m1)
     Ψ2 = Oceananigans.fields(m2)
     for name in keys(Ψ1)
-        ψ1p = parent(Ψ1[name])
-        ψ2p = parent(Ψ2[name])
-        copyto!(ψ1p, Reactant.to_rarray(ψ2p))
+        ψ1 = Ψ1[name]
+        ψ2 = Ψ2[name]
+        loc = Oceananigans.Fields.location(ψ1)
+        #if loc[3] != Nothing
path: src/correctness.jl
user: login: glwagner
body: ```suggestion
```
created_at: 2025-04-07T05:04:02Z
author_association: COLLABORATOR
diff_hunk: @@ -51,10 +52,27 @@ function sync_states!(m1, m2)
     Ψ1 = Oceananigans.fields(m1)
     Ψ2 = Oceananigans.fields(m2)
     for name in keys(Ψ1)
-        ψ1p = parent(Ψ1[name])
-        ψ2p = parent(Ψ2[name])
-        copyto!(ψ1p, Reactant.to_rarray(ψ2p))
+        ψ1 = Ψ1[name]
+        ψ2 = Ψ2[name]
+        loc = Oceananigans.Fields.location(ψ1)
+        ψ2r = Reactant.to_rarray(interior(Ψ2[name]))
+        #@jit donated_args=:none copy_interior!(interior(ψ1), ψ2r)
path: src/correctness.jl
user: login: glwagner
body: ```suggestion
```
created_at: 2025-04-07T05:04:16Z
author_association: COLLABORATOR
diff_hunk: @@ -45,7 +45,8 @@ Oceananigans = "0.96.18"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"
-Reactant = "0.2.68"
+Reactant = "0.2.67"
+Reactant_jll = "0.0.132"
path: Project.toml
user: login: giordano
body: ```suggestion
```
created_at: 2025-04-08T19:37:12Z
author_association: COLLABORATOR
diff_hunk: @@ -45,7 +45,8 @@ Oceananigans = "0.96.18"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"
-Reactant = "0.2.68"
+Reactant = "0.2.67"
path: Project.toml
user: login: giordano
body: ```suggestion
Reactant = "0.2.68"
```
created_at: 2025-04-08T19:37:24Z
author_association: COLLABORATOR
diff_hunk: @@ -51,6 +51,8 @@ else
     rank = Reactant.Distributed.local_rank()
 end
 
+NxNy = 6^2 * 384 * 192
+
path: sharding/sharded_baroclinic_instability_simulation_run.jl
user: login: glwagner
body: ```suggestion
```
created_at: 2025-04-09T13:54:43Z
author_association: COLLABORATOR
diff_hunk: @@ -81,6 +81,26 @@ jobs:
       sim_type: 'sharded_baroclinic_instability'
       grid_type: ${{ matrix.grid_type }}
       sharded: true
+      run_dir: 'sharding'
+      julia_version: '1.11'
+      os: ${{ matrix.os }}
+      xla_runtime: 'IFRT'
+      compile_or_run: 'run'
+      earlyoom_threshold: 4
+      julia_optlevel: 0
+
+  run_correctness:
+    name: Correctness - Julia ${{ matrix.julia_version }} - ${{ matrix.grid_type }} - ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+          os: ['ubuntu-24.04', 'ubuntu-22.04-arm']
+          grid_type: ['simple_lat_lon']
+    uses: ./.github/workflows/CompileOrRun.yml
+    with:
+      sim_type: 'correctness_baroclinic_instability'
+      grid_type: ${{ matrix.grid_type }}
+      run_dir: 'correctness'
path: .github/workflows/Run.yml
user: login: giordano
body: Maybe this could just be moved to a separate file (with different paths triggers)?
created_at: 2025-04-09T13:55:54Z
author_association: COLLABORATOR
diff_hunk: @@ -81,6 +81,26 @@ jobs:
       sim_type: 'sharded_baroclinic_instability'
       grid_type: ${{ matrix.grid_type }}
       sharded: true
+      run_dir: 'sharding'
+      julia_version: '1.11'
+      os: ${{ matrix.os }}
+      xla_runtime: 'IFRT'
+      compile_or_run: 'run'
+      earlyoom_threshold: 4
+      julia_optlevel: 0
+
+  run_correctness:
+    name: Correctness - Julia ${{ matrix.julia_version }} - ${{ matrix.grid_type }} - ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+          os: ['ubuntu-24.04', 'ubuntu-22.04-arm']
+          grid_type: ['simple_lat_lon']
+    uses: ./.github/workflows/CompileOrRun.yml
+    with:
+      sim_type: 'correctness_baroclinic_instability'
+      grid_type: ${{ matrix.grid_type }}
+      run_dir: 'correctness'
path: .github/workflows/Run.yml
user: login: glwagner
body: I'm ok with that 
created_at: 2025-04-09T19:39:14Z
author_association: COLLABORATOR

PR comments:
user: login: giordano
created_at: 2025-04-07T13:22:23Z
body: @simone-silvestri note that there are merge conflicts.
author_association: COLLABORATOR
user: login: wsmoses
created_at: 2025-04-08T17:39:13Z
body: @glwagner so CI is now green, does that means correctness passes?
author_association: MEMBER
user: login: giordano
created_at: 2025-04-08T17:52:31Z
body: That's not actually tested in CI yet.
author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2025-04-08T18:11:01Z
body: We could test it in CI.
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-04-08T18:46:29Z
body: I'm getting this going locally, then I will add something that would cause CI to fail if correctness fails
author_association: COLLABORATOR
user: login: giordano
created_at: 2025-04-09T14:16:50Z
body: https://github.com/PRONTOLab/GB-25/actions/runs/14359045228/job/40255468090?pr=166#step:10:102
```
ERROR: LoadError: There is a discrepancy between the models!
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] 
```
:smile: 

Could the error message show where the discrepancy is coming from?
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-04-09T19:39:38Z
body: > https://github.com/PRONTOLab/GB-25/actions/runs/14359045228/job/40255468090?pr=166#step:10:102
> 
> ```
> ERROR: LoadError: There is a discrepancy between the models!
> Stacktrace:
>  [1] error(s::String)
>    @ Base ./error.jl:35
>  [2] 
> ```
> 
> 😄
> 
> Could the error message show where the discrepancy is coming from?

Heh, yes
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-04-09T19:41:05Z
body: It can be seen actually by inspecting the log, eg:

https://github.com/PRONTOLab/GB-25/actions/runs/14359045228/job/40255468090?pr=166#step:10:56

What are you thinking of to make this clearer? If we remove the two `show`s that occur before the error, maybe it will be easier to read?
author_association: COLLABORATOR
user: login: giordano
created_at: 2025-04-10T14:25:56Z
body: All green now :rocket:  can we merge it?

But [I still think](https://github.com/PRONTOLab/GB-25/pull/166#discussion_r2035435360) it'd be nicer to split the "correctness" jobs to their own workflow file, they don't share much with the "run" jobs.

***Edit***: I did that
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-04-10T14:57:29Z
body: great! merge away!
author_association: COLLABORATOR