Introduce z_sampling_method, remove disable_vertical_interpolation

The NetCDF writer was originally designed for ClimaAtmos. As such, it uses some
fake pressure levels for vertical sampling as default. This is generally not
applicable/useful. This commit generalized the method to sample points on the
vertical direction introducing a new type `AbstractZSamplingMethod`.

`AbstractZSamplingMethod` has to concrete subtypes:
- `LevelsMethod` (equivalent to `disable_vertical_interpolation`)
- `FakePressureLevelsMethod` (equivalent to previously default behavior)

This opens up the possibility to also define `LinearMethod`, for linear
sampling, and `LogMethod`, for log sampling.

The default was changed to `LevelsMethod` and `disable_vertical_interpolation`
was deprecated.

NEWS.md

@@ -0,0 +1,9 @@
+# NEWS
+
+v0.2.0
+-------
+
+- ![][badge-💥breaking]The `NetCDFWriter` now outputs points on the vertical levels by default.
+- ![][badge-💥breaking] `disable_vertical_interpolation` is removed in favor of `z_sampling_method`.
+
+[badge-💥breaking]: https://img.shields.io/badge/💥BREAKING-red.svg

docs/src/writers.md

@@ -7,7 +7,8 @@ Writers are needed to save the computed diagnostics.
 - `DictWriter`, to save `Field`s to dictionaries in memory;
 - `HDF5Writer`, to save `Field`s to HDF5 files.
 
-(There is an additional `DummyWriter` that does nothing. It is mostly used internally for testing and debugging.)
+(There is an additional `DummyWriter` that does nothing. It is mostly used
+internally for testing and debugging.)
 
 Users are welcome to implement their own writers. A writer has to be a subtype
 of `AbstractWriter`and has to implement the `interpolate_field!` and
@@ -31,6 +32,15 @@ and the output directory where the files should be saved. By default, the
 exist. The `NetCDFWriter` does not overwrite existing data and will error out if
 existing data is inconsistent with the new one.
 
+`NetCDFWriter`s take as one of the inputs the desired number of points along
+each of the dimensions. For the horizontal dimensions, points are sampled
+linearly. For the vertical dimension, the behavior can be customized by passing
+the `z_sampling_method` variable. When `z_sampling_method =
+ClimaDiagnostics.Writers.LevelMethod()`, points evaluated on the grid levels
+(and the provided number of points ignored), when `z_sampling_method =
+ClimaDiagnostics.Writers.FakePressureLevelMethod()`, points are sampled
+uniformly in simplified hydrostatic atmospheric model.
+
 The output in the `NetCDFWriter` roughly follows the CF conventions.
 
 Each `ScheduledDiagnostic` is output to a different file with name determined by
@@ -66,6 +76,15 @@ ClimaDiagnostics.Writers.sync
 Base.close
 ```
 
+Sampling methods for the vertical direction:
+```@docs
+ClimaDiagnostics.Writers.AbstractZSamplingMethod
+ClimaDiagnostics.Writers.LevelsMethod
+ClimaDiagnostics.Writers.FakePressureLevelsMethod
+Base.close
+```
+
+
 ## `DictWriter`
 
 The `DictWriter` is a in-memory writer that is particularly useful for

src/netcdf_writer.jl

@@ -16,7 +16,8 @@ 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} <: AbstractWriter
+struct NetCDFWriter{T, TS, DI, SYNC, ZSM <: AbstractZSamplingMethod} <:
+       AbstractWriter
     """The base folder where to save the files."""
     output_dir::String
 
@@ -45,10 +46,9 @@ struct NetCDFWriter{T, TS, DI, SYNC} <: AbstractWriter
     """NetCDF files that are currently open. Only the root process uses this field."""
     open_files::Dict{String, NCDatasets.NCDataset}
 
-    """ Do not interpolate on the z direction, instead evaluate on the levels. When
-    disable_vertical_interpolation is true, the num_points on the vertical direction is
-    ignored."""
-    disable_vertical_interpolation::Bool
+    """Instance of a type that determines how the points along the vertical direction are
+    sampled."""
+    z_sampling_method::ZSM
 
     """Areas of memory preallocated where the interpolation output is saved. Only the root
     process uses this."""
@@ -90,9 +90,9 @@ Keyword arguments
 - `num_points`: How many points to use along the different dimensions to interpolate the
                 fields. This is a tuple of integers, typically having meaning Long-Lat-Z,
                 or X-Y-Z (the details depend on the configuration being simulated).
-- `disable_vertical_interpolation`: Do not interpolate on the z direction, instead evaluate
-                                    at on levels. When disable_vertical_interpolation is true,
-                                    the num_points on the vertical direction is ignored.
+- `z_sampling_method`: Instance of a `AbstractZSamplingMethod` that determines how points
+                       on the vertical direction should be chosen. By default, the vertical
+                       points are sampled on the grid levels.
 - `compression_level`: How much to compress the output NetCDF file (0 is no compression, 9
                        is maximum compression).
 - `sync_schedule`: Schedule that determines when to call `NetCDF.sync` (to flush the output
@@ -106,15 +106,15 @@ function NetCDFWriter(
     space,
     output_dir;
     num_points = (180, 90, 50),
-    disable_vertical_interpolation = false,
     compression_level = 9,
     sync_schedule = ClimaComms.device(space) isa ClimaComms.CUDADevice ?
                     EveryStepSchedule() : nothing,
+    z_sampling_method = LevelsMethod(),
 )
     horizontal_space = Spaces.horizontal_space(space)
     is_horizontal_space = horizontal_space == space
 
-    if disable_vertical_interpolation
+    if z_sampling_method isa LevelsMethod
         # It is a little tricky to override the number of vertical points because we don't
         # know if the vertical direction is the 2nd (as in a plane) or 3rd index (as in a
         # box or sphere). To set this value, we check if we are on a plane or not
@@ -135,11 +135,7 @@ function NetCDFWriter(
         hpts = target_coordinates(space, num_points)
         vpts = []
     else
-        hpts, vpts = target_coordinates(
-            space,
-            num_points;
-            disable_vertical_interpolation,
-        )
+        hpts, vpts = target_coordinates(space, num_points, z_sampling_method)
     end
 
     hcoords = hcoords_from_horizontal_space(
@@ -174,14 +170,15 @@ function NetCDFWriter(
         typeof(interpolated_physical_z),
         typeof(preallocated_arrays),
         typeof(sync_schedule),
+        typeof(z_sampling_method),
     }(
         output_dir,
         Dict{String, Remapper}(),
         num_points,
         compression_level,
         interpolated_physical_z,
         Dict{String, NCDatasets.NCDataset}(),
-        disable_vertical_interpolation,
+        z_sampling_method,
         preallocated_arrays,
         sync_schedule,
         unsynced_datasets,
@@ -220,8 +217,8 @@ function interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
         else
             hpts, vpts = target_coordinates(
                 space,
-                writer.num_points;
-                writer.disable_vertical_interpolation,
+                writer.num_points,
+                writer.z_sampling_method,
             )
         end
 
@@ -233,7 +230,7 @@ function interpolate_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
 
         # When we disable vertical_interpolation, we override the vertical points with
         # the reference values for the vertical space.
-        if writer.disable_vertical_interpolation && !is_horizontal_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(
@@ -322,7 +319,7 @@ function write_field!(writer::NetCDFWriter, field, diagnostic, u, p, t)
         nc,
         space,
         writer.num_points;
-        writer.disable_vertical_interpolation,
+        writer.z_sampling_method,
         writer.interpolated_physical_z,
     )
 

src/netcdf_writer_coordinates.jl

@@ -1,3 +1,38 @@
+"""
+    AbstractZSamplingMethod
+
+The `AbstractZInterpolationMethod` defines how points along the vertical axis should be
+sampled.
+
+In other words, if a column is defined between 0 and 100 and the target number of points is
+50. How should those 50 points be chosen?
+
+Available methods are:
+- `LevelMethod`: just use the grid levels
+- `FakePressureLevelsMethod`: linearly spaced in (very) approximate atmospheric pressure levels
+"""
+abstract type AbstractZSamplingMethod end
+
+"""
+    LevelsMethod
+
+Do not perform interpolation on `z`, use directly the grid levels instead.
+"""
+struct LevelsMethod <: AbstractZSamplingMethod end
+
+"""
+    FakePressureLevelsMethod
+
+Linearly sample points from `z_min` to `z_max` in pressure levels assuming a very simplified hydrostatic balance model.
+
+Pressure is approximated with
+
+p ~ p₀ exp(-z/H)
+
+H is assumed to be 7000 m, which is a good scale height for the Earth atmosphere.
+"""
+struct FakePressureLevelsMethod <: AbstractZSamplingMethod end
+
 """
     add_dimension!(nc::NCDatasets.NCDataset,
                    name::String,
@@ -115,17 +150,12 @@ function target_coordinates(space, num_points) end
 
 function target_coordinates(
     space::S,
-    num_points;
-    disable_vertical_interpolation,
+    num_points,
+    z_sampling_method::FakePressureLevelsMethod,
 ) where {
     S <:
     Union{Spaces.CenterFiniteDifferenceSpace, Spaces.FaceFiniteDifferenceSpace},
 }
-    if disable_vertical_interpolation
-        cspace = Spaces.space(space, Grids.CellCenter())
-        return Array(parent(Fields.coordinate_field(cspace).z))[:, 1]
-    end
-
     # Exponentially spaced with base e
     #
     # We mimic something that looks like pressure levels
@@ -147,22 +177,30 @@ function target_coordinates(
     return collect(-H_EARTH * log.(range(exp_z_min, exp_z_max, num_points_z)))
 end
 
+function target_coordinates(
+    space::S,
+    num_points,
+    z_sampling_method::LevelsMethod,
+) where {
+    S <:
+    Union{Spaces.CenterFiniteDifferenceSpace, Spaces.FaceFiniteDifferenceSpace},
+}
+    cspace = Spaces.space(space, Grids.CellCenter())
+    return Array(parent(Fields.coordinate_field(cspace).z))[:, 1]
+end
+
 # Column
 function add_space_coordinates_maybe!(
     nc::NCDatasets.NCDataset,
     space::Spaces.FiniteDifferenceSpace,
     num_points_z;
-    disable_vertical_interpolation,
+    z_sampling_method,
     names = ("z",),
 )
     name, _... = names
     z_dimension_exists = dimension_exists(nc, name, (num_points_z,))
     if !z_dimension_exists
-        zpts = target_coordinates(
-            space,
-            num_points_z;
-            disable_vertical_interpolation,
-        )
+        zpts = target_coordinates(space, num_points_z, z_sampling_method)
         add_dimension!(nc, name, zpts, units = "m", axis = "Z")
     end
     return [name]
@@ -305,7 +343,7 @@ function add_space_coordinates_maybe!(
     nc::NCDatasets.NCDataset,
     space::Spaces.ExtrudedFiniteDifferenceSpace,
     num_points;
-    disable_vertical_interpolation,
+    z_sampling_method,
     interpolated_physical_z = nothing,
 )
 
@@ -330,15 +368,15 @@ function add_space_coordinates_maybe!(
             nc,
             vertical_space,
             num_points_vertic;
-            disable_vertical_interpolation,
+            z_sampling_method,
         )
     else
         vdims_names = add_space_coordinates_maybe!(
             nc,
             vertical_space,
             num_points_vertic,
             interpolated_physical_z;
-            disable_vertical_interpolation,
+            z_sampling_method,
             names = ("z_reference",),
             depending_on_dimensions = hdims_names,
         )
@@ -347,14 +385,14 @@ function add_space_coordinates_maybe!(
     return (hdims_names..., vdims_names...)
 end
 
-# Ignore the interpolated_physical_z/disable_vertical_interpolation keywords in the general
+# Ignore the interpolated_physical_z/z_sampling_method keywords in the general
 # case (we only case about the specialized one for extruded spaces)
 add_space_coordinates_maybe!(
     nc::NCDatasets.NCDataset,
     space,
     num_points;
     interpolated_physical_z = nothing,
-    disable_vertical_interpolation = false,
+    z_sampling_method = nothing,
 ) = add_space_coordinates_maybe!(nc::NCDatasets.NCDataset, space, num_points)
 
 # Elevation with topography
@@ -371,7 +409,7 @@ function add_space_coordinates_maybe!(
     num_points,
     interpolated_physical_z;
     names = ("z_reference",),
-    disable_vertical_interpolation,
+    z_sampling_method,
     depending_on_dimensions,
 )
     num_points_z = num_points
@@ -382,11 +420,8 @@ function add_space_coordinates_maybe!(
         dimension_exists(nc, name, (num_points_z,))
 
     if !z_reference_dimension_dimension_exists
-        reference_altitudes = target_coordinates(
-            space,
-            num_points_z;
-            disable_vertical_interpolation,
-        )
+        reference_altitudes =
+            target_coordinates(space, num_points_z, z_sampling_method)
         add_dimension!(nc, name, reference_altitudes; units = "m", axis = "Z")
     end
 
@@ -421,8 +456,8 @@ end
 # and combines the resulting dictionaries
 function target_coordinates(
     space::Spaces.ExtrudedFiniteDifferenceSpace,
-    num_points;
-    disable_vertical_interpolation,
+    num_points,
+    z_sampling_method,
 )
 
     hcoords = vcoords = ()
@@ -436,11 +471,8 @@ function target_coordinates(
         Spaces.vertical_topology(space),
         Spaces.staggering(space),
     )
-    vcoords = target_coordinates(
-        vertical_space,
-        num_points_vertic;
-        disable_vertical_interpolation,
-    )
+    vcoords =
+        target_coordinates(vertical_space, num_points_vertic, z_sampling_method)
 
     hcoords == vcoords == () && error("Found empty space")
 

test/integration_test.jl

@@ -113,7 +113,7 @@ end
             NCDatasets.NCDataset(joinpath(output_dir, "YO_1it_inst.nc")) do nc
                 @test nc["YO"].attrib["short_name"] == "YO"
                 @test nc["YO"].attrib["long_name"] == "YO YO, Instantaneous"
-                @test size(nc["YO"]) == (11, 10, 5, 3)
+                @test size(nc["YO"]) == (11, 10, 5, 10)
             end
 
             NCDatasets.NCDataset(
@@ -122,13 +122,13 @@ end
                 @test nc["YO"].attrib["short_name"] == "YO"
                 @test nc["YO"].attrib["long_name"] ==
                       "YO YO, average within every 2 iterations"
-                @test size(nc["YO"]) == (5, 10, 5, 3)
+                @test size(nc["YO"]) == (5, 10, 5, 10)
             end
 
             NCDatasets.NCDataset(joinpath(output_dir, "YO_3s_inst.nc")) do nc
                 @test nc["YO"].attrib["short_name"] == "YO"
                 @test nc["YO"].attrib["long_name"] == "YO YO, Instantaneous"
-                @test size(nc["YO"]) == (4, 10, 5, 3)
+                @test size(nc["YO"]) == (4, 10, 5, 10)
             end
         end
     end

test/writers.jl

@@ -41,13 +41,14 @@ end
         output_dir;
         num_points = (NUM, NUM, NUM),
         sync_schedule = ClimaDiagnostics.Schedules.DivisorSchedule(2),
+        z_sampling_method = ClimaDiagnostics.Writers.FakePressureLevelsMethod(),
     )
 
     writer_no_vert_interpolation = Writers.NetCDFWriter(
         space,
         output_dir;
         num_points = (NUM, NUM, NUM),
-        disable_vertical_interpolation = true,
+        z_sampling_method = ClimaDiagnostics.Writers.LevelsMethod(),
     )
 
     # Check Base.show