Add the ability to use NetCDF

.github/workflows/ci.yml

@@ -124,6 +124,31 @@ jobs:
         env:
           TEST_GROUP: "timestepping"
 
+  netcdf:
+    name: NetCDF Writer - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.12'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "netcdf"
+
   enthalpy:
     name: Enthalpy Model - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaSeaIce"
 uuid = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.4.5"
+version = "0.4.6"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -12,6 +12,12 @@ RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 
+[weakdeps]
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+
+[extensions]
+ClimaSeaIceNCDatasetsExt = "NCDatasets"
+
 [compat]
 Adapt = "3, 4"
 JLD2 = "0.6.2"
@@ -27,4 +33,4 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["MPI", "Test"]
+test = ["MPI", "Test", "NCDatasets"]

ext/ClimaSeaIceNCDatasetsExt.jl

@@ -0,0 +1,44 @@
+module ClimaSeaIceNCDatasetsExt
+
+using NCDatasets
+using Oceananigans
+using ClimaSeaIce
+
+#####
+##### Variable attributes
+#####
+
+default_horizontal_velocity_attributes(::RectilinearGrid) = Dict(
+    "u" => Dict("long_name" => "Velocity in the +x-direction.", "units" => "m/s"),
+    "v" => Dict("long_name" => "Velocity in the +y-direction.", "units" => "m/s"))
+
+default_horizontal_velocity_attributes(::LatitudeLongitudeGrid) = Dict(
+    "u" => Dict("long_name" => "Velocity in the zonal direction (+ = east).", "units" => "m/s"),
+    "v" => Dict("long_name" => "Velocity in the meridional direction (+ = north).", "units" => "m/s"))
+
+default_horizontal_velocity_attributes(::OrthogonalSphericalShellGrid) = Dict(
+    "u" => Dict("long_name" => "Velocity in the i-direction (+ = increasing i).", "units" => "m/s"),
+    "v" => Dict("long_name" => "Velocity in the j-direction (+ = increasing j).", "units" => "m/s")
+)
+
+default_horizontal_velocity_attributes(ibg::ImmersedBoundaryGrid) = default_horizontal_velocity_attributes(ibg.underlying_grid)
+
+default_sea_ice_attributes() = Dict(
+    "h" => Dict("long_name" => "Sea ice thickness.", "units" => "m"),
+    "ℵ" => Dict("long_name" => "Sea ice concentration.", "units" => "-")
+)
+
+function __init__()
+    OCNE = Base.get_extension(Oceananigans, :OceananigansNCDatasetsExt)
+    if !isnothing(OCNE)
+        @eval begin
+            function $OCNE.default_output_attributes(model::SeaIceModel)
+                velocity_attrs = default_horizontal_velocity_attributes(model.grid)
+                tracer_attrs = default_sea_ice_attributes()
+                return merge(velocity_attrs, tracer_attrs)
+            end
+        end
+    end
+end
+
+end

test/runtests.jl

@@ -39,6 +39,10 @@ if TEST_GROUP == "all" || TEST_GROUP == "checkpointing"
     include("test_checkpointing.jl")
 end
 
+if TEST_GROUP == "all" || TEST_GROUP == "netcdf"
+    include("test_netcdf_writer.jl")
+end
+
 if TEST_GROUP == "all" || TEST_GROUP == "distributed"
     include("test_distributed_sea_ice.jl")
 end

test/test_netcdf_writer.jl

@@ -0,0 +1,57 @@
+using ClimaSeaIce
+using ClimaSeaIce.SeaIceDynamics
+using ClimaSeaIce.SeaIceThermodynamics
+using Test
+
+using NCDatasets
+using Oceananigans
+using Oceananigans.OutputWriters: NetCDFWriter
+
+function test_netcdf_writer(arch)
+    Nx, Ny = 16, 16
+    Lx, Ly = 100, 100
+    Δt = 1
+
+    grid = RectilinearGrid(arch, size=(Nx, Ny), x=(0, Lx), y=(0, Ly), topology=(Bounded, Bounded, Flat))
+
+    for dynamics in (nothing, SeaIceMomentumEquation(grid))
+        ice_thermodynamics = SlabSeaIceThermodynamics(grid)
+        model = SeaIceModel(grid; dynamics, ice_thermodynamics)
+
+        set!(model, h=0.5, ℵ=1)
+
+        simulation = Simulation(model, Δt=Δt, stop_iteration=5)
+
+        filepath = "test_sea_ice_netcdf.nc"
+
+        simulation.output_writers[:nc] = NetCDFWriter(model, Oceananigans.prognostic_fields(model);
+                                                      filename = filepath,
+                                                      schedule = IterationInterval(1),
+                                                      overwrite_existing = true)
+
+        run!(simulation)
+
+        ds = NCDataset(filepath)
+
+        # Check that the time dimension exists
+        @test haskey(ds.dim, "time")
+        @test length(ds["time"]) == 6  # iteration 0 through 5
+
+        # Check that sea ice fields are present
+        @test haskey(ds, "h")
+        @test haskey(ds, "ℵ")
+
+        # Check that velocity fields are present when dynamics are enabled
+        if !isnothing(dynamics)
+            @test haskey(ds, "u")
+            @test haskey(ds, "v")
+        end
+
+        close(ds)
+        rm(filepath, force=true)
+    end
+end
+
+@testset "NetCDFWriter Tests" begin
+    test_netcdf_writer(CPU())
+end