From 60401decea17c35e328f40932ca11d45f5e561eb Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 4 Jun 2024 23:20:53 -0700 Subject: [PATCH] make spaces consistent --- ext/TempestRegridderExt.jl | 97 +++++++++++++++++++++----------------- test/data_handling.jl | 13 ++++- 2 files changed, 65 insertions(+), 45 deletions(-) diff --git a/ext/TempestRegridderExt.jl b/ext/TempestRegridderExt.jl index c4f7ad61..6e361091 100644 --- a/ext/TempestRegridderExt.jl +++ b/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, diff --git a/test/data_handling.jl b/test/data_handling.jl index bb651f99..fd0f0f64 100644 --- a/test/data_handling.jl +++ b/test/data_handling.jl @@ -26,7 +26,7 @@ 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 @@ -34,7 +34,16 @@ ClimaComms.init(context) 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)