Add NetCDFWriter support for long-lat boxes

This commit makes identifying coordinates for resampling more
coordinate agnostic.

This commit introduces a workaround for
CliMA/ClimaCore.jl#1936

More fundamentally, CliMA/ClimaCore.jl#1936 is
probably a reflection of the fact that boxes with longlat points are not
supported extremely well.

Here, we make the decision to work with LatLong boxes (instead of the
more natural LongLat ones, more natural because Long is the X axis).
LongLat boxes will fail

NEWS.md

@@ -1,5 +1,17 @@
 # NEWS
 
+v0.2.5
+-------
+
+## Features
+
+### Add support for box spaces with LatLong points in `NetCDFWriter`.
+
+The `NetCDFWriter` can now work with regional boxes with `LatLong` points. Due
+to incompatibility in `ClimaCore`, only `LatLong` points are supported (and not
+`LongLat` points). This means that the box has to be created with latitude on
+the x axis and longitude on the y axis.
+
 v0.2.4
 -------
 

src/netcdf_writer_coordinates.jl

@@ -231,10 +231,10 @@ function target_coordinates(
 )
     num_points_x, num_points_y = num_points
     FT = Spaces.undertype(space)
-    xmin = FT(domain.interval1.coord_min.x)
-    xmax = FT(domain.interval1.coord_max.x)
-    ymin = FT(domain.interval2.coord_min.y)
-    ymax = FT(domain.interval2.coord_max.y)
+    xmin = Geometry.tofloat(domain.interval1.coord_min)
+    xmax = Geometry.tofloat(domain.interval1.coord_max)
+    ymin = Geometry.tofloat(domain.interval2.coord_min)
+    ymax = Geometry.tofloat(domain.interval2.coord_max)
     xpts = collect(range(xmin, xmax, num_points_x))
     ypts = collect(range(ymin, ymax, num_points_y))
     return (xpts, ypts)
@@ -248,8 +248,8 @@ function target_coordinates(
 )
     num_points_x, _... = num_points
     FT = Spaces.undertype(space)
-    xmin = FT(domain.coord_min.x)
-    xmax = FT(domain.coord_max.x)
+    xmin = Geometry.tofloat(domain.coord_min)
+    xmax = Geometry.tofloat(domain.coord_max)
     xpts = collect(range(xmin, xmax, num_points_x))
     return (xpts)
 end
@@ -269,24 +269,39 @@ function target_coordinates(
 end
 
 # Box
+function _islatlonbox(domain)
+    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,
     num_points,
-    ::Domains.RectangleDomain;
-    names = ("x", "y"),
+    domain::Domains.RectangleDomain;
+    names = _islatlonbox(domain) ? ("lon", "lat") : ("x", "y"),
 )
-    xname, yname = names
-    num_points_x, num_points_y = num_points
-    x_dimension_exists = dimension_exists(nc, xname, (num_points_x,))
-    y_dimension_exists = dimension_exists(nc, yname, (num_points_y,))
+    name1, name2 = names
+    num_points1, num_points2 = num_points
 
-    if !x_dimension_exists && !y_dimension_exists
-        xpts, ypts = target_coordinates(space, num_points)
-        add_dimension!(nc, xname, xpts; units = "m", axis = "X")
-        add_dimension!(nc, yname, ypts; units = "m", axis = "Y")
+    if _islatlonbox(domain)
+        units, axis = ("degrees_east", "degrees_north"), ("X", "Y")
+        # ClimaCore assumes LatLon, but we really want LongLat, so we need to flip
+        num_points1, num_points2 = num_points2, num_points1
+    else
+        units, axis = ("m", "m"), ("X", "Y")
     end
-    return [xname, yname]
+
+    dim1_exists = dimension_exists(nc, name1, (num_points1,))
+    dim2_exists = dimension_exists(nc, name2, (num_points2,))
+
+    if !dim1_exists && !dim2_exists
+        pts1, pts2 = target_coordinates(space, (num_points1, num_points2))
+        add_dimension!(nc, name1, pts1; units = units[1], axis = axis[1])
+        add_dimension!(nc, name2, pts2; units = units[2], axis = axis[2])
+    end
+
+    return [name1, name2]
 end
 
 # Plane
@@ -488,20 +503,43 @@ function hcoords_from_horizontal_space(
     return [Geometry.LatLongPoint(hc2, hc1) for hc1 in hpts[1], hc2 in hpts[2]]
 end
 
+
+# Workaround for https://github.com/CliMA/ClimaCore.jl/issues/1936
+#
+# Note that here we are assuming that lat is first
+Geometry._coordinate(pt::Geometry.LatLongPoint, ::Val{1}) =
+    Geometry.LatPoint(pt.lat)
+Geometry._coordinate(pt::Geometry.LatLongPoint, ::Val{2}) =
+    Geometry.LongPoint(pt.long)
+
 function hcoords_from_horizontal_space(
     space::Spaces.SpectralElementSpace2D,
     domain::Domains.RectangleDomain,
     hpts,
 )
-    return [Geometry.XYPoint(hc1, hc2) for hc1 in hpts[1], hc2 in hpts[2]]
+    XYPointType = typeof(
+        Geometry.product_coordinates(
+            domain.interval1.coord_max,
+            domain.interval2.coord_max,
+        ),
+    )
+
+    # ClimaCore works with LatLon boxes, but we really care about LongLat boxes
+    if _islatlonbox(domain)
+        hpts1, hpts2 = hpts[2], hpts[1]
+    else
+        hpts1, hpts2 = hpts[1], hpts[2]
+    end
+
+    return [XYPointType(hc1, hc2) for hc1 in hpts1, hc2 in hpts2]
 end
 
 function hcoords_from_horizontal_space(
     space::Spaces.SpectralElementSpace1D,
-    domain::Domains.IntervalDomain,
+    domain::Domains.IntervalDomain{PointType},
     hpts,
-)
-    return [Geometry.XPoint(hc1) for hc1 in hpts]
+) where {PointType}
+    return [PointType(hc1) for hc1 in hpts]
 end
 
 """

test/TestTools.jl

@@ -62,6 +62,66 @@ function SphericalShellSpace(;
     )
 end
 
+function BoxSpace(;
+    FT = Float64,
+    xlim = (-FT(1), FT(1)),
+    ylim = (-FT(1), FT(1)),
+    height = 10.0,
+    lonlat = false,
+    nelements = (10, 10),
+    zelem = 10,
+    npolynomial = 4,
+    context = ClimaComms.context(),
+)
+    vertdomain = ClimaCore.Domains.IntervalDomain(
+        ClimaCore.Geometry.ZPoint(FT(0)),
+        ClimaCore.Geometry.ZPoint(FT(height));
+        boundary_names = (:bottom, :top),
+    )
+    vertmesh = ClimaCore.Meshes.IntervalMesh(vertdomain; nelems = zelem)
+    if pkgversion(ClimaCore) >= v"0.14.10"
+        vert_center_space = ClimaCore.Spaces.CenterFiniteDifferenceSpace(
+            ClimaComms.device(context),
+            vertmesh,
+        )
+    else
+        vert_center_space =
+            ClimaCore.Spaces.CenterFiniteDifferenceSpace(vertmesh)
+    end
+
+    # NOTE: here we assume LatPoint is first
+    XPoint = lonlat ? ClimaCore.Geometry.LatPoint : ClimaCore.Geometry.XPoint
+    YPoint = lonlat ? ClimaCore.Geometry.LongPoint : ClimaCore.Geometry.YPoint
+    periodic = !lonlat
+    boundary_names = (:first, :second)
+
+    domain_x = ClimaCore.Domains.IntervalDomain(
+        XPoint(xlim[1]),
+        XPoint(xlim[2]);
+        periodic,
+        boundary_names,
+    )
+
+    domain_y = ClimaCore.Domains.IntervalDomain(
+        YPoint(ylim[1]),
+        YPoint(ylim[2]);
+        periodic,
+        boundary_names,
+    )
+
+    horzdomain = ClimaCore.Domains.RectangleDomain(domain_x, domain_y)
+    horzmesh =
+        ClimaCore.Meshes.RectilinearMesh(horzdomain, nelements[1], nelements[2])
+    horztopology = ClimaCore.Topologies.Topology2D(context, horzmesh)
+    quad = ClimaCore.Spaces.Quadratures.GLL{npolynomial + 1}()
+    horzspace = ClimaCore.Spaces.SpectralElementSpace2D(horztopology, quad)
+
+    return ClimaCore.Spaces.ExtrudedFiniteDifferenceSpace(
+        horzspace,
+        vert_center_space,
+    )
+end
+
 """
     create_problem()
 

test/writers.jl

@@ -133,6 +133,79 @@ end
         t,
     )
 
+    # Check boxes
+    xyboxspace = BoxSpace()
+    xyboxfield = Fields.coordinate_field(xyboxspace).z
+
+    xyboxwriter = Writers.NetCDFWriter(
+        xyboxspace,
+        output_dir;
+        num_points = (NUM, NUM, NUM),
+    )
+    xyboxdiagnostic = ClimaDiagnostics.ScheduledDiagnostic(;
+        variable = ClimaDiagnostics.DiagnosticVariable(;
+            compute!,
+            short_name = "ABC",
+        ),
+        output_short_name = "my_short_name_xybox",
+        output_long_name = "My Long Name",
+        output_writer = xyboxwriter,
+    )
+    xyboxu = (; xyboxfield)
+    Writers.interpolate_field!(
+        xyboxwriter,
+        xyboxfield,
+        xyboxdiagnostic,
+        xyboxu,
+        p,
+        t,
+    )
+    Writers.write_field!(xyboxwriter, xyboxfield, xyboxdiagnostic, xyboxu, p, t)
+
+    longlatboxspace = BoxSpace(; lonlat = true)
+    longlatboxfield = Fields.coordinate_field(longlatboxspace).z
+
+    longlatboxwriter = Writers.NetCDFWriter(
+        longlatboxspace,
+        output_dir;
+        num_points = (NUM, NUM, NUM),
+    )
+    longlatboxdiagnostic = ClimaDiagnostics.ScheduledDiagnostic(;
+        variable = ClimaDiagnostics.DiagnosticVariable(;
+            compute!,
+            short_name = "ABC",
+        ),
+        output_short_name = "my_short_name_longlatbox",
+        output_long_name = "My Long Name",
+        output_writer = longlatboxwriter,
+    )
+    longlatboxu = (; longlatboxfield)
+    Writers.interpolate_field!(
+        longlatboxwriter,
+        longlatboxfield,
+        longlatboxdiagnostic,
+        longlatboxu,
+        p,
+        t,
+    )
+    Writers.write_field!(
+        longlatboxwriter,
+        longlatboxfield,
+        longlatboxdiagnostic,
+        longlatboxu,
+        p,
+        t,
+    )
+    # Write a second time, to check consistency
+    Writers.write_field!(
+        longlatboxwriter,
+        longlatboxfield,
+        longlatboxdiagnostic,
+        longlatboxu,
+        p,
+        t,
+    )
+
     ###############
     # Performance #
     ###############