Merge branch 'main' into xk/ows-papa

xkykai · web-flow · commit ff9fdc0aa254 · 2026-04-03T16:17:02.000-04:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -31,6 +31,7 @@ env:
   CDSAPI_URL: "https://cds.climate.copernicus.eu/api"
   CDSAPI_KEY: ${{ secrets.CDSAPI_KEY }}
   DATADEPS_ALWAYS_ACCEPT: true
+  JULIA_PKG_SERVER_REGISTRY_PREFERENCE: 'eager'
 
 #####
 ##### CPU tests (GitHub-hosted runners)
@@ -98,6 +99,9 @@ jobs:
     container:
       image: ghcr.io/numericalearth/numerical-earth-docker-images:test-julia_1.12.5
       options: --gpus=all
+      volumes:
+        # Mount host `/usr/local` so that we can delete some stuff afterwards
+        - /usr/local:/host-usr-local
     timeout-minutes: 120
     strategy:
       fail-fast: false
@@ -117,6 +121,23 @@ jobs:
         run: git config --global --add safe.directory ${PWD}
       - name: Copy LocalPreferences
         run: cp -v /usr/local/share/julia/environments/numericalearth/LocalPreferences.toml test/.
+      - name: df before cleanup
+        run: |
+          df -hT
+      - name: Clean up old CUDA toolkits
+        # Save storage by deleting old CUDA toolkits, we'll use v13
+        run: |
+          rm -rf /host-usr-local/cuda-12.*
+      - name: df after cleanup
+        run: |
+          df -hT
+      - name: Update registry
+        shell: julia --color=yes {0}
+        run: |
+          using Pkg
+          Pkg.Registry.rm("General")
+          Pkg.Registry.add("General")
+          Pkg.Registry.update()
       - name: Run tests
         run: |
           earlyoom -m 3 -s 100 -r 300 --prefer 'julia' &
diff --git a/Project.toml b/Project.toml
@@ -52,7 +52,7 @@ NumericalEarthWOAExt = "WorldOceanAtlasTools"
 
 [compat]
 Adapt = "4"
-Breeze = "0.4"
+Breeze = "0.4.4"
 CDSAPI = "2.2.1"
 CFTime = "0.1, 0.2"
 ClimaSeaIce = "0.4.4, 0.5"
@@ -69,7 +69,7 @@ JLD2 = "0.4, 0.5, 0.6"
 KernelAbstractions = "0.9"
 MeshArrays = "0.3, 0.4, 0.5"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.104.2, 0.105, 0.106"
+Oceananigans = "0.106.3"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
 PythonCall = "0.9.28"
diff --git a/ext/NumericalEarthReactantExt.jl b/ext/NumericalEarthReactantExt.jl
@@ -7,7 +7,7 @@ using Oceananigans.DistributedComputations: Distributed
 using NumericalEarth: EarthSystemModel
 
 import Oceananigans
-import Oceananigans.Models: initialization_update_state!
+import Oceananigans.TimeSteppers: reconcile_state!
 
 const OceananigansReactantExt = Base.get_extension(
      Oceananigans, :OceananigansReactantExt
@@ -18,6 +18,6 @@ const ReactantOSIM{I, A, O, F, C} = Union{
     EarthSystemModel{I, A, O, F, C, <:Distributed{ReactantState}},
 }
 
-initialization_update_state!(model::ReactantOSIM) = nothing
+reconcile_state!(model::ReactantOSIM) = nothing
 
 end # module NumericalEarthReactantExt
diff --git a/src/Bathymetry/orca_grid.jl b/src/Bathymetry/orca_grid.jl
diff --git a/src/EarthSystemModels/EarthSystemModels.jl b/src/EarthSystemModels/EarthSystemModels.jl
@@ -47,10 +47,10 @@ import Thermodynamics as AtmosphericThermodynamics
 import Oceananigans: fields, prognostic_fields, prognostic_state, restore_prognostic_state!
 import Oceananigans.Architectures: architecture
 import Oceananigans.Fields: set!
-import Oceananigans.Models: NaNChecker, default_nan_checker, initialization_update_state!
+import Oceananigans.Models: NaNChecker, default_nan_checker
 import Oceananigans.OutputWriters: default_included_properties
 import Oceananigans.Simulations: timestepper, reset!, initialize!, iteration
-import Oceananigans.TimeSteppers: time_step!, update_state!, time
+import Oceananigans.TimeSteppers: time_step!, update_state!, time, reconcile_state!
 import Oceananigans.Utils: prettytime
 
 include("components.jl")
diff --git a/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl b/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_fluxes.jl
@@ -202,9 +202,8 @@ end
                                                      liquidus, ocean_properties, ℰ, u★)
 
     # Store interface values and heat flux
-    @inbounds T★[i, j, 1] = Tᵦ
-    @inbounds S★[i, j, 1] = Sᵦ
     @inbounds 𝒬ⁱⁿᵗ[i, j, 1] = 𝒬ⁱᵒ
+    store_interface_state!(flux_formulation, T★, S★, i, j, Tᵦ, Sᵦ)
 
     # =============================================
     # Part 4: Salt flux
diff --git a/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_heat_flux_formulations.jl b/src/EarthSystemModels/InterfaceComputations/sea_ice_ocean_heat_flux_formulations.jl
@@ -207,6 +207,14 @@ const ConductiveFluxTEF{FT} = ThreeEquationHeatFlux{<:ConductiveFlux, <:Abstract
 @inline extract_internal_temperature(::IceBathHeatFlux{FT},   i, j) where FT = zero(FT)
 @inline extract_internal_temperature(flux::ConductiveFluxTEF, i, j) = @inbounds flux.internal_temperature[i, j, 1]
 
+# For IceBathHeatFlux, T★ and S★ are views into ocean surface fields so we skip writing.
+# For ThreeEquationHeatFlux, T★ and S★ are dedicated interface fields.
+@inline store_interface_state!(::IceBathHeatFlux, T★, S★, i, j, Tᵦ, Sᵦ) = nothing
+@inline function store_interface_state!(::ThreeEquationHeatFlux, T★, S★, i, j, Tᵦ, Sᵦ)
+    @inbounds T★[i, j, 1] = Tᵦ
+    @inbounds S★[i, j, 1] = Sᵦ
+end
+
 """
     compute_interface_heat_flux(flux::ThreeEquationHeatFlux, ocean_state, ice_state, liquidus, ocean_properties, ℰ, u★)
 
diff --git a/src/EarthSystemModels/earth_system_model.jl b/src/EarthSystemModels/earth_system_model.jl
@@ -61,15 +61,14 @@ function reset!(model::ESM)
 end
 
 # Make sure to initialize the exchanger here
-function initialization_update_state!(model::ESM)
+function initialize!(model::ESM)
     initialize!(model.interfaces.exchanger, model)
-    update_state!(model)
     return nothing
 end
 
-function initialize!(model::ESM)
-    # initialize!(model.ocean)
+function reconcile_state!(model::ESM)
     initialize!(model.interfaces.exchanger, model)
+    update_state!(model)
     return nothing
 end
 
@@ -209,7 +208,7 @@ function EarthSystemModel(atmosphere, ocean, sea_ice;
     # Make sure the initial temperature of the ocean
     # is not below freezing and above melting near the surface
     above_freezing_ocean_temperature!(ocean, interfaces.exchanger.grid, sea_ice)
-    initialization_update_state!(earth_system_model)
+    reconcile_state!(earth_system_model)
 
     return earth_system_model
 end
diff --git a/src/EarthSystemModels/time_step_earth_system_model.jl b/src/EarthSystemModels/time_step_earth_system_model.jl
@@ -2,11 +2,14 @@ using .InterfaceComputations:
     compute_atmosphere_ocean_fluxes!,
     compute_sea_ice_ocean_fluxes!
 
+using Oceananigans.TimeSteppers: maybe_prepare_first_time_step!
 using ClimaSeaIce: SeaIceModel, SeaIceThermodynamics
 using Oceananigans.Grids: φnode
 using Printf
 
 function time_step!(coupled_model::EarthSystemModel, Δt; callbacks=[])
+    maybe_prepare_first_time_step!(coupled_model, callbacks)
+    
     ocean = coupled_model.ocean
     sea_ice = coupled_model.sea_ice
     atmosphere = coupled_model.atmosphere
@@ -30,7 +33,7 @@ function time_step!(coupled_model::EarthSystemModel, Δt; callbacks=[])
     return nothing
 end
 
-function update_state!(coupled_model::EarthSystemModel)
+function update_state!(coupled_model::EarthSystemModel, callbacks=[])
 
     # The three components
     ocean      = coupled_model.ocean
diff --git a/test/Project.toml b/test/Project.toml
@@ -45,7 +45,7 @@ JLD2 = "0.4, 0.5, 0.6"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.104.2, 0.105"
+Oceananigans = "0.106"
 PythonCall = "0.9.28"
 Reactant = "0.2.235"
 Scratch = "1"
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -87,7 +87,7 @@ function __init__()
     # Download few datasets for tests
     for dataset in test_datasets
         time_resolution = dataset isa ECCO2Daily ? Day(1) : Month(1)
-        end_date = start_date + 2 * time_resolution
+        end_date = start_date + 1 * time_resolution
         dates = start_date:time_resolution:end_date
 
         temperature_metadata = Metadata(:temperature; dataset, dates)
diff --git a/test/test_orca_grid.jl b/test/test_orca_grid.jl
@@ -23,7 +23,7 @@ end
 @testset "ORCAGrid with ORCA1 dataset on $(arch)" for arch in test_architectures
     south_rows_to_remove = 43
     grid = ORCAGrid(arch; dataset=ORCA1(), Nz=5, z=(-5000, 0), halo=(4, 4, 4), south_rows_to_remove)
-    @test grid.underlying_grid.Ny == 333 - south_rows_to_remove
+    @test grid.underlying_grid.Ny == 332 - south_rows_to_remove
 
     grid = ORCAGrid(arch; dataset=ORCA1(), Nz=5, z=(-5000, 0), halo=(4, 4, 4), south_rows_to_remove=0)
 
@@ -33,7 +33,7 @@ end
     @test underlying isa Oceananigans.Grids.OrthogonalSphericalShellGrid
     @test underlying isa TripolarGrid
     @test underlying.Nx == 362
-    @test underlying.Ny == 333
+    @test underlying.Ny == 332
     @test underlying.Nz == 5
 
     # Coordinates span near-global domain
@@ -51,7 +51,7 @@ end
     @test grid isa TripolarGrid
     @test !(grid isa ImmersedBoundaryGrid)
     @test grid.Nx == 362
-    @test grid.Ny == 333 - default_south_rows_to_remove(ORCA1())
+    @test grid.Ny == 332 - default_south_rows_to_remove(ORCA1())
     @test grid.Nz == 5
 end
 
@@ -63,7 +63,7 @@ end
     @test grid isa ImmersedBoundaryGrid
     underlying = grid.underlying_grid
     @test underlying.Nx == 362
-    @test underlying.Ny == 333 - Nremove
+    @test underlying.Ny == 332 - Nremove
     @test underlying.Nz == 5
 end
 
@@ -115,8 +115,8 @@ end
     # At interior points, Face[j] should be < Center[j] in latitude.
     imid = Nx ÷ 2
     φF   = grid.φᶜᶠᵃ[imid, 1:Ny]
-    φC   = grid.φᶜᶜᵃ[imid, 1:Ny-1]  # Center has Ny-1 interior points
-    nsouth = count(j -> φF[j] < φC[j], 1:length(φC))
+    φC   = grid.φᶜᶜᵃ[imid, 1:Ny]
+    nsouth = count(j -> φF[j] < φC[j], 1:Ny)
     @test nsouth / length(φC) > 0.95
 
     # Periodic overlap: first and last unique columns should be consistent
diff --git a/test/test_reactant.jl b/test/test_reactant.jl
@@ -1,6 +1,5 @@
 using Test
 using Reactant
-using Oceananigans.Models: initialization_update_state!
 using Oceananigans: Oceananigans
 using Oceananigans.Architectures: ReactantState
 using Oceananigans.Grids: Bounded, Flat, LatitudeLongitudeGrid, Periodic
