Add support for purely vertical spaces with NetCDF

Leveraging the new feature introduced in ClimaCore, this commit adds
support to interpolating purely vertical fields and saving them as a
NetCDF file.

NEWS.md

@@ -1,5 +1,11 @@
 # NEWS
 
+v0.2.12
+-------
+## Bug fixes
+
+- `NetCDFWriter` now correctly writes purely vertical spaces.
+
 v0.2.11
 -------
 ## Bug fixes

src/netcdf_writer.jl

@@ -109,7 +109,7 @@ Keyword arguments
 - `start_date`: Date of the beginning of the simulation.
 """
 function NetCDFWriter(
-    space,
+    space::Spaces.AbstractSpace,
     output_dir;
     num_points = (180, 90, 50),
     compression_level = 9,
@@ -118,6 +118,10 @@ function NetCDFWriter(
     z_sampling_method = LevelsMethod(),
     start_date = nothing,
 )
+    has_horizontal_space =
+        space isa Spaces.ExtrudedFiniteDifferenceSpace ||
+        space isa Spaces.AbstractSpectralElementSpace
+
     horizontal_space = Spaces.horizontal_space(space)
     is_horizontal_space = horizontal_space == space
 
@@ -194,6 +198,63 @@ function NetCDFWriter(
     )
 end
 
+function NetCDFWriter(
+    space::Spaces.Spaces.FiniteDifferenceSpace,
+    output_dir;
+    num_points = (180, 90, 50),
+    compression_level = 9,
+    sync_schedule = ClimaComms.device(space) isa ClimaComms.CUDADevice ?
+                    EveryStepSchedule() : nothing,
+    z_sampling_method = LevelsMethod(),
+    start_date = nothing,
+)
+    if z_sampling_method isa LevelsMethod
+        num_vpts = Meshes.nelements(Grids.vertical_topology(space).mesh)
+        @warn "Disabling vertical interpolation, the provided number of points is ignored (using $num_vpts)"
+        num_points = (num_vpts,)
+    end
+    vpts = target_coordinates(space, num_points, z_sampling_method)
+    target_zcoords = Geometry.ZPoint.(vpts)
+    remapper = Remapper(space; target_zcoords)
+
+    comms_ctx = ClimaComms.context(space)
+
+    coords_z = Fields.coordinate_field(space).z
+    maybe_move_to_cpu =
+        ClimaComms.device(coords_z) isa ClimaComms.CUDADevice &&
+        ClimaComms.iamroot(comms_ctx) ? Array : identity
+
+    interpolated_physical_z = maybe_move_to_cpu(interpolate(remapper, coords_z))
+
+    preallocated_arrays =
+        ClimaComms.iamroot(comms_ctx) ?
+        Dict{ScheduledDiagnostic, ClimaComms.array_type(space)}() :
+        Dict{ScheduledDiagnostic, Nothing}()
+
+    unsynced_datasets = Set{NCDatasets.NCDataset}()
+
+    return NetCDFWriter{
+        typeof(num_points),
+        typeof(interpolated_physical_z),
+        typeof(preallocated_arrays),
+        typeof(sync_schedule),
+        typeof(z_sampling_method),
+        typeof(start_date),
+    }(
+        output_dir,
+        Dict{String, Remapper}(),
+        num_points,
+        compression_level,
+        interpolated_physical_z,
+        Dict{String, NCDatasets.NCDataset}(),
+        z_sampling_method,
+        preallocated_arrays,
+        sync_schedule,
+        unsynced_datasets,
+        start_date,
+    )
+end
+
 """
     interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
 
@@ -205,61 +266,62 @@ function interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
 
     space = axes(field)
 
-    horizontal_space = Spaces.horizontal_space(space)
+    has_horizontal_space = !(space isa Spaces.FiniteDifferenceSpace)
 
-    # We have to deal with to cases: when we have an horizontal slice (e.g., the
-    # surface), and when we have a full space. We distinguish these cases by checking if
-    # the given space has the horizontal_space attribute. If not, it is going to be a
-    # SpectralElementSpace2D and we don't have to deal with the z coordinates.
-    is_horizontal_space = horizontal_space == space
+    if has_horizontal_space
+        horizontal_space = Spaces.horizontal_space(space)
+
+        # We have to deal with to cases: when we have an horizontal slice (e.g., the
+        # surface), and when we have a full space. We distinguish these cases by checking if
+        # the given space has the horizontal_space attribute. If not, it is going to be a
+        # SpectralElementSpace2D and we don't have to deal with the z coordinates.
+        is_horizontal_space = horizontal_space == space
+    end
 
     # Prepare the remapper if we don't have one for the given variable. We need one remapper
     # per variable (not one per diagnostic since all the time reductions return the same
     # type of space).
 
-    # TODO: Expand this once we support spatial reductions
+    # TODO: Expand this once we support spatial reductions.
+    # TODO: More generally, this can be clean up to have less conditionals
+    # depending on the type of space and use dispatch instead
     if !haskey(writer.remappers, var.short_name)
 
         # hpts, vpts are ranges of numbers
-        # hcoords, zcoords are ranges of Geometry.Points
-
-        zcoords = []
-
-        if is_horizontal_space
-            hpts = target_coordinates(space, writer.num_points)
-            vpts = []
+        # target_hcoords, target_zcoords are ranges of Geometry.Points
+
+        target_zcoords = nothing
+        target_hcoords = nothing
+
+        if has_horizontal_space
+            if is_horizontal_space
+                hpts = target_coordinates(space, writer.num_points)
+                vpts = []
+            else
+                hpts, vpts = target_coordinates(
+                    space,
+                    writer.num_points,
+                    writer.z_sampling_method,
+                )
+            end
+
+            target_hcoords = hcoords_from_horizontal_space(
+                horizontal_space,
+                Meshes.domain(Spaces.topology(horizontal_space)),
+                hpts,
+            )
         else
-            hpts, vpts = target_coordinates(
+            vpts = target_coordinates(
                 space,
                 writer.num_points,
                 writer.z_sampling_method,
             )
         end
 
-        hcoords = hcoords_from_horizontal_space(
-            horizontal_space,
-            Meshes.domain(Spaces.topology(horizontal_space)),
-            hpts,
-        )
-
-        # When we disable vertical_interpolation, we override the vertical points with
-        # the reference values for the vertical space.
-        if writer.z_sampling_method isa LevelsMethod && !is_horizontal_space
-            # We need Array(parent()) because we want an array of values, not a DataLayout
-            # of Points
-            vpts = Array(
-                parent(
-                    space.grid.vertical_grid.center_local_geometry.coordinates,
-                ),
-            )[
-                :,
-                1,
-            ]
-        end
+        target_zcoords = Geometry.ZPoint.(vpts)
 
-        zcoords = [Geometry.ZPoint(p) for p in vpts]
-
-        writer.remappers[var.short_name] = Remapper(space, hcoords, zcoords)
+        writer.remappers[var.short_name] =
+            Remapper(space, target_hcoords, target_zcoords)
     end
 
     remapper = writer.remappers[var.short_name]
@@ -314,9 +376,7 @@ function write_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
     interpolated_field =
         maybe_move_to_cpu(writer.preallocated_output_arrays[diagnostic])
 
-    if islatlonbox(
-        Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))),
-    )
+    if islatlonbox(space)
         # ClimaCore works with LatLong points, but we want to have longitude
         # first in the output, so we have to flip things
         perm = collect(1:length(size(interpolated_field)))
@@ -401,14 +461,8 @@ function write_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
         nc["date"][time_index] = string(start_date + Dates.Millisecond(1000t))
     end
 
-    # TODO: It would be nice to find a cleaner way to do this
-    if length(dim_names) == 3
-        v[time_index, :, :, :] = interpolated_field
-    elseif length(dim_names) == 2
-        v[time_index, :, :] = interpolated_field
-    elseif length(dim_names) == 1
-        v[time_index, :] = interpolated_field
-    end
+    colons = ntuple(_ -> Colon(), length(dim_names))
+    v[time_index, colons...] = interpolated_field
 
     # Add file to list of files that might need manual sync
     push!(writer.unsynced_datasets, nc)

src/netcdf_writer_coordinates.jl

@@ -164,7 +164,7 @@ function target_coordinates(
     # We assume H to be 7000, which is a good scale height for the Earth atmosphere
     H_EARTH = 7000
 
-    num_points_z = num_points[]
+    num_points_z = last(num_points)
     FT = Spaces.undertype(space)
     topology = Spaces.topology(space)
     vert_domain = topology.mesh.domain
@@ -195,9 +195,12 @@ function add_space_coordinates_maybe!(
     num_points_z;
     z_sampling_method,
     names = ("z",),
+    interpolated_physical_z = nothing, # Not needed here, but needed for consistency of
+    # interface and dispatch
 )
     name, _... = names
-    z_dimension_exists = dimension_exists(nc, name, (num_points_z,))
+    z_dimension_exists = dimension_exists(nc, name, num_points_z)
+
     if !z_dimension_exists
         zpts = target_coordinates(space, num_points_z, z_sampling_method)
         add_dimension!(nc, name, zpts, units = "m", axis = "Z")
@@ -273,14 +276,21 @@ function target_coordinates(
     return (longpts, latpts)
 end
 
-islatlonbox(domain) = false
+islatlonbox(space::Spaces.FiniteDifferenceSpace) = false
+islatlonbox(space::Domains.AbstractDomain) = false
+function islatlonbox(space::Spaces.AbstractSpace)
+    return islatlonbox(
+        Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))),
+    )
+end
 
 # Box
 function islatlonbox(domain::Domains.RectangleDomain)
     return domain.interval1.coord_max isa Geometry.LatPoint &&
            domain.interval2.coord_max isa Geometry.LongPoint
 end
 
+
 function add_space_coordinates_maybe!(
     nc::NCDatasets.NCDataset,
     space::Spaces.SpectralElementSpace2D,
@@ -416,14 +426,14 @@ function add_space_coordinates_maybe!(
         vdims_names = add_space_coordinates_maybe!(
             nc,
             vertical_space,
-            num_points_vertic;
+            (num_points_vertic,);
             z_sampling_method,
         )
     else
         vdims_names = add_space_coordinates_maybe!(
             nc,
             vertical_space,
-            num_points_vertic,
+            (num_points_vertic,),
             interpolated_physical_z;
             z_sampling_method,
             names = ("z_reference",),

test/TestTools.jl

@@ -14,6 +14,28 @@ function ColumnCenterFiniteDifferenceSpace(
     context = ClimaComms.SingletonCommsContext();
     FT = Float64,
 )
+    return _column(
+        zelem,
+        ClimaCore.Spaces.CenterFiniteDifferenceSpace,
+        context,
+        FT,
+    )
+end
+
+function ColumnFaceFiniteDifferenceSpace(
+    zelem = 10,
+    context = ClimaComms.SingletonCommsContext();
+    FT = Float64,
+)
+    return _column(
+        zelem,
+        ClimaCore.Spaces.FaceFiniteDifferenceSpace,
+        context,
+        FT,
+    )
+end
+
+function _column(zelem, constructor, context, FT)
     zlim = (FT(0.0), FT(1.0))
     domain = ClimaCore.Domains.IntervalDomain(
         ClimaCore.Geometry.ZPoint(zlim[1]),
@@ -22,9 +44,10 @@ function ColumnCenterFiniteDifferenceSpace(
     )
     mesh = ClimaCore.Meshes.IntervalMesh(domain, nelems = zelem)
     topology = ClimaCore.Topologies.IntervalTopology(context, mesh)
-    return ClimaCore.Spaces.CenterFiniteDifferenceSpace(topology)
+    return constructor(topology)
 end
 
+
 function SphericalShellSpace(;
     radius = 6371.0,
     height = 10.0,

test/writers.jl

@@ -6,6 +6,7 @@ using Profile
 using BenchmarkTools
 import ProfileCanvas
 import NCDatasets
+import ClimaCore
 import ClimaCore.Fields
 
 import ClimaDiagnostics
@@ -15,7 +16,7 @@ include("TestTools.jl")
 
 # The temporary directory where we write the file cannot be in /tmp, it has
 # to be on disk
-output_dir = mktempdir(".")
+output_dir = mktempdir(pwd())
 
 @testset "DictWriter" begin
     writer = Writers.DictWriter()
@@ -217,6 +218,41 @@ end
         t,
     )
 
+    # Check columns
+    if pkgversion(ClimaCore) >= v"0.14.23"
+        # Center space
+        for (i, colspace) in enumerate((
+            ColumnCenterFiniteDifferenceSpace(),
+            ColumnFaceFiniteDifferenceSpace(),
+        ))
+            colfield = Fields.coordinate_field(colspace).z
+
+            colwriter =
+                Writers.NetCDFWriter(colspace, output_dir; num_points = (NUM,))
+            coldiagnostic = ClimaDiagnostics.ScheduledDiagnostic(;
+                variable = ClimaDiagnostics.DiagnosticVariable(;
+                    compute!,
+                    short_name = "ABC",
+                ),
+                output_short_name = "my_short_name_c$(i)",
+                output_long_name = "My Long Name",
+                output_writer = colwriter,
+            )
+            colu = (; colfield)
+            Writers.interpolate_field!(
+                colwriter,
+                colfield,
+                coldiagnostic,
+                colu,
+                p,
+                t,
+            )
+            Writers.write_field!(colwriter, colfield, coldiagnostic, colu, p, t)
+            # Write a second time, to check consistency
+            Writers.write_field!(colwriter, colfield, coldiagnostic, colu, p, t)
+        end
+    end
+
     ###############
     # Performance #
     ###############