Add support for PointSpaces with NetCDFWriter

When a point space results from computating a diagnostic,
interpolation is not run. Additionally, no space coordinate is
added to the output netcdf file.

improve tests

Fix vertical space coordinates bug

Make point space diagnostics work with dict and hdf5 writers

fix pointspace mpi

Author: imreddyTeja

NEWS.md

@@ -29,7 +29,7 @@ v0.2.12
 -------
 ## Bug fixes
 
-- `NetCDFWriter` now correctly writes purely vertical spaces.
+- `NetCDFWriter` now correctly writes purely vertical and point spaces.
 
 v0.2.11
 -------

src/clima_diagnostics.jl

@@ -1,8 +1,11 @@
 import Accessors
 import SciMLBase
+import ClimaComms
+import ClimaCore: Spaces
 
 import .Schedules: DivisorSchedule, EveryDtSchedule
-import .Writers: interpolate_field!, write_field!, sync, AbstractWriter
+import .Writers:
+    interpolate_field!, write_field!, sync, AbstractWriter, NetCDFWriter
 
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")
@@ -126,7 +129,24 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
 
         # If it is not a reduction, call the output writer as well
         if !isa_time_reduction
-            interpolate_field!(diag.output_writer, storage[i], diag, Y, p, t)
+            # no need to interpolate for point spaces
+            if axes(storage[i]) isa Spaces.PointSpace
+                # netCDFWriter expects diagnostic to be in preallocated_output_arrays
+                if diag.output_writer isa NetCDFWriter &&
+                   ClimaComms.iamroot(ClimaComms.context(storage[i]))
+                    diag.output_writer.preallocated_output_arrays[diag] =
+                        copy(parent(storage[i]))
+                end
+            else
+                interpolate_field!(
+                    diag.output_writer,
+                    storage[i],
+                    diag,
+                    Y,
+                    p,
+                    t,
+                )
+            end
             write_field!(diag.output_writer, storage[i], diag, Y, p, t)
         else
             # Add to the accumulator
@@ -224,15 +244,25 @@ function orchestrate_diagnostics(
             diagnostic_handler.storage[diag_index],
             diagnostic_handler.counters[diag_index],
         )
-
-        interpolate_field!(
-            diag.output_writer,
-            diagnostic_handler.storage[diag_index],
-            diag,
-            integrator.u,
-            integrator.p,
-            integrator.t,
-        )
+        # dont interpolate for point spaces
+        if axes(diagnostic_handler.storage[diag_index]) isa Spaces.PointSpace
+            # netCDFWriter expects diagnostic to be in preallocated_output_arrays
+            if diag.output_writer isa NetCDFWriter && ClimaComms.iamroot(
+                ClimaComms.context(diagnostic_handler.storage[diag_index]),
+            )
+                diag.output_writer.preallocated_output_arrays[diag] =
+                    copy(parent(diagnostic_handler.storage[diag_index]))
+            end
+        else
+            interpolate_field!(
+                diag.output_writer,
+                diagnostic_handler.storage[diag_index],
+                diag,
+                integrator.u,
+                integrator.p,
+                integrator.t,
+            )
+        end
     end
 
     # Save to disk

src/hdf5_writer.jl

@@ -1,5 +1,5 @@
 import ClimaComms
-import ClimaCore.InputOutput
+import ClimaCore: InputOutput, Spaces
 
 import ClimaUtilities.TimeManager: ITime
 
@@ -38,6 +38,11 @@ The name of the file is determined by the `output_short_name` of the output
 `Field`s can be read back using the `InputOutput` module in `ClimaCore`.
 """
 function write_field!(writer::HDF5Writer, field, diagnostic, u, p, t)
+    axes(field) isa Spaces.PointSpace &&
+        pkgversion(InputOutput) < v"0.14.27" &&
+        error(
+            "HDF5Writer only supports Fields with PointSpace for ClimaCore >= 0.14.27",
+        )
     var = diagnostic.variable
     time = t
 

src/netcdf_writer.jl

@@ -18,8 +18,14 @@ include("netcdf_writer_coordinates.jl")
 A struct to remap `ClimaCore` `Fields` to rectangular grids and save the output to NetCDF
 files.
 """
-struct NetCDFWriter{T, TS, DI, SYNC, ZSM <: AbstractZSamplingMethod, DATE} <:
-       AbstractWriter
+struct NetCDFWriter{
+    T,
+    TS,
+    DI,
+    SYNC,
+    ZSM <: Union{AbstractZSamplingMethod, Nothing},
+    DATE,
+} <: AbstractWriter
     """The base folder where to save the files."""
     output_dir::String
 
@@ -262,6 +268,42 @@ function NetCDFWriter(
     )
 end
 
+function NetCDFWriter(
+    space::Spaces.Spaces.PointSpace,
+    output_dir;
+    compression_level = 9,
+    sync_schedule = ClimaComms.device(space) isa ClimaComms.CUDADevice ?
+                    EveryStepSchedule() : nothing,
+    start_date = nothing,
+    kwargs...,
+)
+    comms_ctx = ClimaComms.context(space)
+    preallocated_arrays =
+        ClimaComms.iamroot(comms_ctx) ?
+        Dict{ScheduledDiagnostic, ClimaComms.array_type(space)}() :
+        Dict{ScheduledDiagnostic, Nothing}()
+    unsynced_datasets = Set{NCDatasets.NCDataset}()
+    return NetCDFWriter{
+        Nothing,
+        Nothing,
+        typeof(preallocated_arrays),
+        typeof(sync_schedule),
+        Nothing,
+        typeof(start_date),
+    }(
+        output_dir,
+        Dict{String, Remapper}(),
+        nothing,
+        compression_level,
+        nothing,
+        Dict{String, NCDatasets.NCDataset}(),
+        nothing,
+        preallocated_arrays,
+        sync_schedule,
+        unsynced_datasets,
+        start_date,
+    )
+end
 """
     interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
 
@@ -278,7 +320,7 @@ function interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
     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
+        # We have to deal with two 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.
@@ -391,6 +433,11 @@ function write_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
         interpolated_field = permutedims(interpolated_field, perm)
     end
 
+    if space isa Spaces.PointSpace
+        # If the space is a point space, we have to remove the singleton dimension
+        interpolated_field = interpolated_field[]
+    end
+
     FT = Spaces.undertype(space)
 
     output_path =
@@ -459,7 +506,8 @@ function write_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
         # We already have something in the file
         v = nc["$(var.short_name)"]
         temporal_size, spatial_size... = size(v)
-        spatial_size == size(interpolated_field) ||
+        interpolated_size = size(interpolated_field)
+        spatial_size == interpolated_size ||
             error("incompatible dimensions for $(var.short_name)")
     else
         v = NCDatasets.defVar(

src/netcdf_writer_coordinates.jl

@@ -270,6 +270,19 @@ function add_space_coordinates_maybe!(
     return [name]
 end
 
+# PointSpace
+function add_space_coordinates_maybe!(
+    nc::NCDatasets.NCDataset,
+    space::Spaces.PointSpace,
+    num_points_z;
+    z_sampling_method,
+    names = (),
+    interpolated_physical_z = nothing, # Not needed here, but needed for consistency of
+    # interface and dispatch
+)
+    return []
+end
+
 add_space_coordinates_maybe!(
     nc::NCDatasets.NCDataset,
     space::Spaces.AbstractSpectralElementSpace,
@@ -338,6 +351,7 @@ function target_coordinates(
     return (longpts, latpts)
 end
 
+islatlonbox(space::Spaces.PointSpace) = false
 islatlonbox(space::Spaces.FiniteDifferenceSpace) = false
 islatlonbox(space::Domains.AbstractDomain) = false
 function islatlonbox(space::Spaces.AbstractSpace)
@@ -533,16 +547,15 @@ function add_space_coordinates_maybe!(
     z_sampling_method,
     depending_on_dimensions,
 )
-    num_points_z = num_points
     name, _... = names
 
     # Add z_reference
     z_reference_dimension_dimension_exists =
-        dimension_exists(nc, name, (num_points_z,))
+        dimension_exists(nc, name, num_points)
 
     if !z_reference_dimension_dimension_exists
         reference_altitudes =
-            target_coordinates(space, num_points_z, z_sampling_method)
+            target_coordinates(space, num_points, z_sampling_method)
         add_dimension!(nc, name, reference_altitudes; units = "m", axis = "Z")
     end
 

test/writers.jl

@@ -8,6 +8,9 @@ import ProfileCanvas
 import NCDatasets
 import ClimaCore
 import ClimaCore.Fields
+import ClimaCore.Spaces
+import ClimaCore.Geometry
+import ClimaComms
 
 import ClimaDiagnostics
 import ClimaDiagnostics.Writers
@@ -261,6 +264,106 @@ end
         end
     end
 
+    ###############
+    # Point Space #
+    ###############
+    point_val = 3.14
+    point_space =
+        Spaces.PointSpace(ClimaComms.context(), Geometry.ZPoint(point_val))
+    point_field = Fields.coordinate_field(point_space)
+    point_writer = Writers.NetCDFWriter(point_space, output_dir)
+
+    point_u = (; field = point_field)
+
+    point_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(;
+        variable = ClimaDiagnostics.DiagnosticVariable(;
+            compute!,
+            short_name = "ABC",
+        ),
+        output_short_name = "my_short_name_point",
+        output_long_name = "My Long Name Point",
+        output_writer = point_writer,
+    )
+    point_writer.preallocated_output_arrays[point_diagnostic] = [point_val]
+    # No interpolation needed for point space
+    Writers.write_field!(
+        point_writer,
+        point_field,
+        point_diagnostic,
+        point_u,
+        p,
+        t,
+    )
+    # Write a second time
+    Writers.write_field!(
+        point_writer,
+        point_field,
+        point_diagnostic,
+        point_u,
+        p,
+        t,
+    )
+    close(point_writer)
+
+    NCDatasets.NCDataset(joinpath(output_dir, "my_short_name_point.nc")) do nc
+        @test nc["ABC"][:] == [point_val, point_val]
+    end
+
+    ###################
+    # Horizontal Space#
+    ###################
+
+    horizontal_space = ClimaCore.Spaces.level(space, 1)
+    horizontal_field = Fields.coordinate_field(horizontal_space).z
+    horizontal_writer = Writers.NetCDFWriter(
+        horizontal_space,
+        output_dir;
+        num_points = (NUM, 2NUM),
+    )
+    horizontal_u = (; field = horizontal_field)
+
+    horizontal_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(;
+        variable = ClimaDiagnostics.DiagnosticVariable(;
+            compute!,
+            short_name = "ABC",
+        ),
+        output_short_name = "my_short_name_horizontal",
+        output_long_name = "My Long Name Point Horizontal",
+        output_writer = horizontal_writer,
+    )
+
+    Writers.interpolate_field!(
+        horizontal_writer,
+        horizontal_field,
+        horizontal_diagnostic,
+        horizontal_u,
+        p,
+        t,
+    )
+    Writers.write_field!(
+        horizontal_writer,
+        horizontal_field,
+        horizontal_diagnostic,
+        horizontal_u,
+        p,
+        t,
+    )
+    # Write a second time
+    Writers.write_field!(
+        horizontal_writer,
+        horizontal_field,
+        horizontal_diagnostic,
+        horizontal_u,
+        p,
+        t,
+    )
+    close(horizontal_writer)
+    NCDatasets.NCDataset(
+        joinpath(output_dir, "my_short_name_horizontal.nc"),
+    ) do nc
+        @test size(nc["ABC"]) == (2, NUM, 2NUM)
+    end
+
     ###############
     # Performance #
     ###############