make spaces consistent

ext/TempestRegridderExt.jl

@@ -170,25 +170,6 @@ function reshape_cgll_sparse_to_field!(
     ClimaCore.Topologies.dss!(target, topology)
 end
 
-"""
-    swap_space(field, new_space)
-
-Update the space of a ClimaCore.Fields.Field object. Returns a new Field
-object with the same values as the original field, but on the new space.
-This is needed to correctly read in regridded files that may be reused
-between simulations.
-
-# Arguments
-- `field`: The ClimaCore.Fields.Field object to swap the space of.
-- `new_space`: The new ClimaCore.Spaces.AbstractSpace to assign to the field.
-"""
-function swap_space(field, new_space)
-    return ClimaCore.Fields.Field(
-        ClimaCore.Fields.field_values(field),
-        new_space,
-    )
-end
-
 """
     read_from_hdf5(REGIRD_DIR, hd_outfile_root, time, varname,
         space)
@@ -225,9 +206,10 @@ function read_from_hdf5(REGRID_DIR, hd_outfile_root, time, varname, space)
     field = ClimaCore.InputOutput.read_field(hdfreader, varname)
     Base.close(hdfreader)
 
-    # Ensure the field is on the correct space when reusing regridded files
-    #  between simulations
-    return swap_space(field, space)
+    # Ensure the field is on the correct space
+    @assert axes(field) == space
+
+    return field
 end
 
 
@@ -283,6 +265,44 @@ function write_to_hdf5(
     Base.close(hdfwriter)
 end
 
+"""
+    construct_singleton_space(space)
+
+Constructs a singleton space from a given space. This is not ideal, but
+necessary to use TempestRemap regridding, which is not MPI or GPU compatible.
+"""
+function construct_singleton_space(space)
+    # If doesn't make sense to regrid with GPUs/MPI processes
+    cpu_context =
+        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
+
+    # Check if input space was constructed using `spacefillingcurve`
+    use_spacefillingcurve =
+        ClimaCore.Spaces.topology(space).elemorder isa CartesianIndices ?
+        false : true
+
+    mesh = ClimaCore.Spaces.topology(space).mesh
+    topology = nothing
+    if use_spacefillingcurve
+        topology = ClimaCore.Topologies.Topology2D(
+            cpu_context,
+            mesh,
+            ClimaCore.Topologies.spacefillingcurve(mesh),
+        )
+    else
+        topology = ClimaCore.Topologies.Topology2D(cpu_context, mesh)
+    end
+    Nq =
+        ClimaCore.Spaces.Quadratures.polynomial_degree(
+            ClimaCore.Spaces.quadrature_style(space),
+        ) + 1
+    space_singleton = ClimaCore.Spaces.SpectralElementSpace2D(
+        topology,
+        ClimaCore.Spaces.Quadratures.GLL{Nq}(),
+    )
+
+    return space_singleton
+end
 
 """
     function hdwrite_regridfile_rll_to_cgll(
@@ -334,23 +354,9 @@ function hdwrite_regridfile_rll_to_cgll(
     weightfile = joinpath(REGRID_DIR, outfile_root * "_remap_weights.nc")
 
     # If doesn't make sense to regrid with GPUs/MPI processes
-    cpu_context =
-        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
-
-    # Note: this topology gives us `space == space_undistributed` in the undistributed
-    # case (as desired), which wouldn't hold if we used `spacefillingcurve` here.
-    topology = ClimaCore.Topologies.Topology2D(
-        cpu_context,
-        ClimaCore.Spaces.topology(space).mesh,
-    )
-    Nq =
-        ClimaCore.Spaces.Quadratures.polynomial_degree(
-            ClimaCore.Spaces.quadrature_style(space),
-        ) + 1
-    space_undistributed = ClimaCore.Spaces.SpectralElementSpace2D(
-        topology,
-        ClimaCore.Spaces.Quadratures.GLL{Nq}(),
-    )
+    space_singleton = construct_singleton_space(space)
+    topology_singleton = ClimaCore.Spaces.topology(space_singleton)
+    cpu_context = ClimaCore.Spaces.topology(space_singleton).context
 
     if isfile(datafile_cgll) == false
         isdir(REGRID_DIR) ? nothing : mkpath(REGRID_DIR)
@@ -362,7 +368,7 @@ function hdwrite_regridfile_rll_to_cgll(
         ClimaCoreTempestRemap.rll_mesh(meshfile_rll; nlat = nlat, nlon = nlon)
 
         # write cgll mesh, overlap mesh and weight file
-        ClimaCoreTempestRemap.write_exodus(meshfile_cgll, topology)
+        ClimaCoreTempestRemap.write_exodus(meshfile_cgll, topology_singleton)
         ClimaCoreTempestRemap.overlap_mesh(
             meshfile_overlap,
             meshfile_rll,
@@ -372,6 +378,10 @@ function hdwrite_regridfile_rll_to_cgll(
         # 'in_np = 1' and 'mono = true' arguments ensure mapping is conservative and monotone
         # Note: for a kwarg not followed by a value, set it to true here (i.e. pass 'mono = true' to produce '--mono')
         # Note: out_np = degrees of freedom = polynomial degree + 1
+        Nq =
+            ClimaCore.Spaces.Quadratures.polynomial_degree(
+                ClimaCore.Spaces.quadrature_style(space),
+            ) + 1
         kwargs = (; out_type = out_type, out_np = Nq)
         kwargs = mono ? (; (kwargs)..., in_np = 1, mono = mono) : kwargs
         ClimaCoreTempestRemap.remap_weights(
@@ -400,8 +410,8 @@ function hdwrite_regridfile_rll_to_cgll(
 
     target_unique_idxs =
         out_type == "cgll" ?
-        collect(ClimaCore.Spaces.unique_nodes(space_undistributed)) :
-        collect(ClimaCore.Spaces.all_nodes(space_undistributed))
+        collect(ClimaCore.Spaces.unique_nodes(space_singleton)) :
+        collect(ClimaCore.Spaces.all_nodes(space_singleton))
     target_unique_idxs_i =
         map(row -> target_unique_idxs[row][1][1], row_indices)
     target_unique_idxs_j =
@@ -412,7 +422,7 @@ function hdwrite_regridfile_rll_to_cgll(
 
     R = (; target_idxs = target_unique_idxs, row_indices = row_indices)
 
-    offline_field = ClimaCore.Fields.zeros(FT, space_undistributed)
+    offline_field = ClimaCore.Fields.zeros(FT, space_singleton)
 
     times = [DateTime(0)]
     # Save regridded HDF5 file for each variable in `varnames`
@@ -452,6 +462,7 @@ function hdwrite_regridfile_rll_to_cgll(
             length(times),
         )
 
+        # Save regridded HDF5 file for each time slice
         map(
             x -> write_to_hdf5(
                 REGRID_DIR,

test/data_handling.jl

@@ -26,15 +26,24 @@ ClimaComms.init(context)
         if regridder_type == :TempestRegridder && Sys.iswindows()
             continue
         end
-        for FT in (Float32, Float64)
+        for FT in (Float32, Float64), use_spacefillingcurve in (true, false)
             radius = FT(6731e3)
             helem = 40
             Nq = 4
 
             horzdomain = ClimaCore.Domains.SphereDomain(radius)
             horzmesh =
                 ClimaCore.Meshes.EquiangularCubedSphere(horzdomain, helem)
-            horztopology = ClimaCore.Topologies.Topology2D(context, horzmesh)
+            if use_spacefillingcurve
+                horztopology = ClimaCore.Topologies.Topology2D(
+                    context,
+                    horzmesh,
+                    ClimaCore.Topologies.spacefillingcurve(horzmesh),
+                )
+            else
+                horztopology =
+                    ClimaCore.Topologies.Topology2D(context, horzmesh)
+            end
             quad = ClimaCore.Spaces.Quadratures.GLL{Nq}()
             target_space =
                 ClimaCore.Spaces.SpectralElementSpace2D(horztopology, quad)