diff --git a/docs/src/postprocessor.md b/docs/src/postprocessor.md index 76569afb74..0c88e728c1 100644 --- a/docs/src/postprocessor.md +++ b/docs/src/postprocessor.md @@ -1,8 +1,8 @@ # PostProcessor -This module contains functions for postprocessing model data that was saved during the simulation -by `ClimaCoupler.Diagnostics`. This module is used for offline regridding, slicing and spatial -averages. It can also handle data from other sources (e.g., NCEP reanalysis). +This module contains functions for postprocessing model data that was saved during the simulation +by `ClimaCoupler.Diagnostics`. This module is used for offline regridding, slicing and spatial +averages. It can also handle data from other sources (e.g., NCEP reanalysis). ## Diagnostics API @@ -18,8 +18,3 @@ ClimaCoupler.PostProcessor.DataPackage ``` -## Diagnostics Internal Functions - -```@docs -ClimaCoupler.PostProcessor.read_remapped_field -``` diff --git a/docs/src/regridder.md b/docs/src/regridder.md index 24facc7961..5b823c2a09 100644 --- a/docs/src/regridder.md +++ b/docs/src/regridder.md @@ -18,6 +18,7 @@ ClimaCoupler.Regridder.remap_field_cgll_to_rll ClimaCoupler.Regridder.land_fraction ClimaCoupler.Regridder.update_surface_fractions! ClimaCoupler.Regridder.combine_surfaces! +ClimaCoupler.Regridder.rcgll2latlonz ``` @@ -28,4 +29,7 @@ ClimaCoupler.Regridder.reshape_cgll_sparse_to_field! ClimaCoupler.Regridder.hdwrite_regridfile_rll_to_cgll ClimaCoupler.Regridder.write_datafile_cc ClimaCoupler.Regridder.binary_mask +ClimaCoupler.Regridder.read_remapped_field +ClimaCoupler.Regridder.get_coords +ClimaCoupler.Regridder.get_time ``` \ No newline at end of file diff --git a/src/PostProcessor.jl b/src/PostProcessor.jl index e76f8bbac4..309dba4fe9 100644 --- a/src/PostProcessor.jl +++ b/src/PostProcessor.jl @@ -10,7 +10,7 @@ export PostProcessedData, ZLatLonData, ZLatData, LatLonData, LatData, RawData, D using Statistics using NCDatasets: NCDataset -using ClimaCoupler.Regridder: remap_field_cgll_to_rll +using ClimaCoupler: Regridder using ClimaCore: Fields # data types for postprocessing @@ -93,8 +93,8 @@ function postprocess( isdir(DIR) ? nothing : mkpath(DIR) datafile_latlon = (datafile_latlon == nothing) ? datafile_latlon = DIR * "/remapped_" * string(name) * ".nc" : datafile_latlon - remap_field_cgll_to_rll(name, raw_data, DIR, datafile_latlon, nlat = nlat, nlon = nlon) - new_data, coords = read_remapped_field(name, datafile_latlon) + Regridder.remap_field_cgll_to_rll(name, raw_data, DIR, datafile_latlon, nlat = nlat, nlon = nlon) + new_data, coords = Regridder.read_remapped_field(name, datafile_latlon) raw_tag = length(size(new_data)) == 3 ? ZLatLonData() : LatLonData() package = DataPackage(raw_tag, name, new_data, coords = coords) else @@ -135,23 +135,4 @@ function postprocess( end -""" - read_remapped_field(name::Symbol, datafile_latlon::String, lev_name = "z") - -Extract data and coordinates from `datafile_latlon`. -""" -function read_remapped_field(name::Symbol, datafile_latlon::String, lev_name = "z") - out = NCDataset(datafile_latlon, "r") do nc - lon = nc["lon"][:] - lat = nc["lat"][:] - lev = lev_name in keys(nc) ? nc[lev_name][:] : Float64(-999) - var = nc[name][:] - coords = (; lon = lon, lat = lat, lev = lev) - - (var, coords) - end - - return out -end - end # module diff --git a/src/Regridder.jl b/src/Regridder.jl index 0eca829bd7..f83ff96f37 100644 --- a/src/Regridder.jl +++ b/src/Regridder.jl @@ -26,14 +26,15 @@ export write_to_hdf5, combine_surfaces!, combine_surfaces_from_sol!, binary_mask, - nans_to_zero + nans_to_zero, + cgll2latlonz #= Converts NaNs to zeros of the same type. =# nans_to_zero(v) = isnan(v) ? typeof(v)(0) : v """ - reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R) + reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R, ::Spaces.SpectralElementSpace2D) Reshapes a sparse vector array `in_array` (CGLL, raw output of the TempestRemap), and uses its data to populate the input Field object `field`. @@ -43,15 +44,17 @@ Redundant nodes are populated using `dss` operations. - `field`: [Fields.Field] object populated with the input array. - `in_array`: [Array] input used to fill `field`. - `R`: [NamedTuple] containing `target_idxs` and `row_indices` used for indexing. +- `space`: [Spaces.SpectralElementSpace2D] 2d space to which we are mapping. """ -function reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R) +function reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::SubArray, R, ::Spaces.SpectralElementSpace2D) field_array = parent(field) fill!(field_array, zero(eltype(field_array))) Nf = size(field_array, 3) + # populate the field by iterating over the sparse vector per face for (n, row) in enumerate(R.row_indices) - it, jt, et = (view(R.target_idxs[1], n), view(R.target_idxs[2], n), view(R.target_idxs[3], n)) + it, jt, et = (view(R.target_idxs[1], n), view(R.target_idxs[2], n), view(R.target_idxs[3], n)) # cgll_x, cgll_y, elem for f in 1:Nf field_array[it, jt, f, et] .= in_array[row] end @@ -64,6 +67,47 @@ function reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R) Spaces.dss!(Fields.field_values(field), topology, hspace.quadrature_style) end +""" + reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R, ::Spaces.ExtrudedFiniteDifferenceSpace) + +Reshapes a sparse vector array `in_array` (CGLL, raw output of the TempestRemap), +and uses its data to populate the input Field object `field`. +Redundant nodes are populated using `dss` operations. + +# Arguments +- `field`: [Fields.Field] object populated with the input array. +- `in_array`: [Array] input used to fill `field`. +- `R`: [NamedTuple] containing `target_idxs` and `row_indices` used for indexing. +- `space`: [Spaces.ExtrudedFiniteDifferenceSpace] 3d space to which we are mapping. +""" +function reshape_cgll_sparse_to_field!( + field::Fields.Field, + in_array::SubArray, + R, + ::Spaces.ExtrudedFiniteDifferenceSpace, +) + field_array = parent(field) + + fill!(field_array, zero(eltype(field_array))) + Nf = size(field_array, 4) + Nz = size(field_array, 1) + + # populate the field by iterating over height, then over the sparse vector per face + for z in 1:Nz + for (n, row) in enumerate(R.row_indices) + it, jt, et = (view(R.target_idxs[1], n), view(R.target_idxs[2], n), view(R.target_idxs[3], n)) # cgll_x, cgll_y, elem + for f in 1:Nf + field_array[z, it, jt, f, et] .= in_array[row, z] + end + end + end + # broadcast to the redundant nodes using unweighted dss + space = axes(field) + topology = Spaces.topology(space) + hspace = Spaces.horizontal_space(space) + Spaces.dss!(Fields.field_values(field), topology, hspace.quadrature_style) +end + """ hdwrite_regridfile_rll_to_cgll( FT, @@ -112,10 +156,22 @@ function hdwrite_regridfile_rll_to_cgll( meshfile_overlap = joinpath(REGRID_DIR, outfile_root * "_mesh_overlap.g") weightfile = joinpath(REGRID_DIR, outfile_root * "_remap_weights.nc") - topology = Topologies.Topology2D(space.topology.mesh, Topologies.spacefillingcurve(space.topology.mesh)) - Nq = Spaces.Quadratures.polynomial_degree(space.quadrature_style) + 1 - space_undistributed = Spaces.SpectralElementSpace2D(topology, Spaces.Quadratures.GLL{Nq}()) + if space isa Spaces.ExtrudedFiniteDifferenceSpace + space2d = Spaces.horizontal_space(space) + else + space2d = space + end + topology = Topologies.Topology2D(space2d.topology.mesh, Topologies.spacefillingcurve(space2d.topology.mesh)) + Nq = Spaces.Quadratures.polynomial_degree(space2d.quadrature_style) + 1 + space2d_undistributed = Spaces.SpectralElementSpace2D(topology, Spaces.Quadratures.GLL{Nq}()) + + if space isa Spaces.ExtrudedFiniteDifferenceSpace + vert_center_space = Spaces.CenterFiniteDifferenceSpace(space.vertical_topology.mesh) + space_undistributed = Spaces.ExtrudedFiniteDifferenceSpace(space2d_undistributed, vert_center_space) + else + space_undistributed = space2d_undistributed + end if isfile(datafile_cgll) == false isdir(REGRID_DIR) ? nothing : mkpath(REGRID_DIR) @@ -140,25 +196,16 @@ function hdwrite_regridfile_rll_to_cgll( @warn "Using the existing $datafile_cgll : check topology is consistent" end - function get_time(ds) - if "time" in ds - data_dates = Dates.DateTime.(ds["time"][:]) - elseif "date" in ds - data_dates = TimeManager.strdate_to_datetime.(string.(ds["date"][:])) - else - @warn "No dates available in file $datafile_rll" - data_dates = [Dates.DateTime(0)] - end - end - # read the remapped file with sparse matrices - offline_outvector, times = NCDataset(datafile_cgll, "r") do ds_wt + offline_outvector, coords = NCDataset(datafile_cgll, "r") do ds_wt ( - offline_outvector = ds_wt[varname][:][:, :], # ncol, times - times = get_time(ds_wt), + offline_outvector = Array(ds_wt[varname])[:, :, :], # ncol, z, times + coords = get_coords(ds_wt, space), ) end + times = coords[1] + # weightfile info needed to populate all nodes and save into fields with # sparse matrices _, _, row_indices = NCDataset(weightfile, "r") do ds_wt @@ -166,8 +213,8 @@ function hdwrite_regridfile_rll_to_cgll( end target_unique_idxs = - out_type == "cgll" ? collect(Spaces.unique_nodes(space_undistributed)) : - collect(Spaces.all_nodes(space_undistributed)) + out_type == "cgll" ? collect(Spaces.unique_nodes(space2d_undistributed)) : + collect(Spaces.all_nodes(space2d_undistributed)) target_unique_idxs_i = map(row -> target_unique_idxs[row][1][1], row_indices) target_unique_idxs_j = map(row -> target_unique_idxs[row][1][2], row_indices) target_unique_idxs_e = map(row -> target_unique_idxs[row][2], row_indices) @@ -179,9 +226,16 @@ function hdwrite_regridfile_rll_to_cgll( offline_fields = ntuple(x -> similar(offline_field), length(times)) - ntuple(x -> reshape_cgll_sparse_to_field!(offline_fields[x], offline_outvector[:, x], R), length(times)) + ntuple( + x -> reshape_cgll_sparse_to_field!( + offline_fields[x], + selectdim(offline_outvector, length(coords) + 1, x), + R, + space, + ), + length(times), + ) - # TODO: extend write! to handle time-dependent fields map( x -> write_to_hdf5( REGRID_DIR, @@ -196,6 +250,41 @@ function hdwrite_regridfile_rll_to_cgll( jldsave(joinpath(REGRID_DIR, hd_outfile_root * "_times.jld2"); times = times) end +""" + get_coords(ds, ::Spaces.ExtrudedFiniteDifferenceSpace) + get_coords(ds, ::Spaces.SpectralElementSpace2D) + +Extracts the coordinates from a NetCDF file `ds`. The coordinates are +returned as a tuple of arrays, one for each dimension. The dimensions are +determined by the space type. +""" +function get_coords(ds, ::Spaces.ExtrudedFiniteDifferenceSpace) + data_dates = get_time(ds) + z = ds["z"][:] + return (data_dates, z) +end +function get_coords(ds, ::Spaces.SpectralElementSpace2D) + data_dates = get_time(ds) + return (data_dates,) +end + +""" + get_time(ds) + +Extracts the time information from a NetCDF file `ds`. +""" +function get_time(ds) + if "time" in ds + data_dates = Dates.DateTime.(ds["time"][:]) + elseif "date" in ds + data_dates = TimeManager.strdate_to_datetime.(string.(Int.(ds["date"][:]))) + else + @warn "No dates available in file $datafile_rll" + data_dates = [Dates.DateTime(0)] + end + return data_dates +end + """ write_to_hdf5(REGRID_DIR, hd_outfile_root, time, field, varname, comms_ctx) @@ -486,4 +575,49 @@ function combine_surfaces_from_sol!(combined_field::Fields.Field, fractions::Nam warn_nans && @warn "NaNs were detected and converted to zeros." end +""" + read_remapped_field(name::Symbol, datafile_latlon::String, lev_name = "z") + +Extract data and coordinates from `datafile_latlon`. +""" +function read_remapped_field(name::Symbol, datafile_latlon::String, lev_name = "z") + out = NCDataset(datafile_latlon, "r") do nc + lon = nc["lon"][:] + lat = nc["lat"][:] + lev = lev_name in keys(nc) ? nc[lev_name][:] : Float64(-999) + var = nc[name][:] + coords = (; lon = lon, lat = lat, lev = lev) + + (var, coords) + end + + return out +end + + +""" + function cgll2latlonz(field; DIR = "cgll2latlonz_dir", nlat = 360, nlon = 720, clean_dir = true) + +Regrids a field from CGLL to an RLL array using TempestRemap. It can hanlde multiple other dimensions, such as time and level. + +# Arguments +- `field`: [Fields.Field] to be remapped. +- `DIR`: [String] directory used for remapping. +- `nlat`: [Int] number of latitudes of the regridded array. +- `nlon`: [Int] number of longitudes of the regridded array. +- `clean_dir`: [Bool] flag to delete the temporary directory after remapping. + +# Returns +- Tuple containing the remapped field and its coordinates. +""" +function cgll2latlonz(field; DIR = "cgll2latlonz_dir", nlat = 360, nlon = 720, clean_dir = true) + isdir(DIR) ? nothing : mkpath(DIR) + datafile_latlon = DIR * "/remapped_" * string(name) * ".nc" + remap_field_cgll_to_rll(:var, field, DIR, datafile_latlon, nlat = nlat, nlon = nlon) + new_data, coords = read_remapped_field(:var, datafile_latlon) + clean_dir ? rm(DIR; recursive = true) : nothing + return new_data, coords +end + + end # Module diff --git a/test/Project.toml b/test/Project.toml index a4fdf69d42..77bf40a187 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -35,6 +35,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" ArtifactWrappers = "0.2" Aqua = "0.8" ClimaCoupler = "0.1" +Dates = "1" Downloads = "1" HDF5_jll = "=1.12.2" IntervalSets = "0.5, 0.6, 0.7" @@ -42,5 +43,6 @@ MPI = "0.20" Pkg = "1" PrettyTables = "2" StaticArrays = "1" +Statistics = "1" Unitful = "1" julia = "1.8" diff --git a/test/regridder_tests.jl b/test/regridder_tests.jl index bdb653896f..b8dbf50698 100644 --- a/test/regridder_tests.jl +++ b/test/regridder_tests.jl @@ -8,7 +8,7 @@ using Test using NCDatasets using Dates -using ClimaCoupler: Interfacer, Regridder, TestHelper +using ClimaCoupler: Interfacer, Regridder, TestHelper, TimeManager, PostProcessor import ClimaCoupler.Interfacer: get_field, name, SurfaceModelSimulation, SurfaceStub, update_field! REGRID_DIR = @isdefined(REGRID_DIR) ? REGRID_DIR : joinpath("", "regrid_tmp/") @@ -236,5 +236,79 @@ for FT in (Float32, Float64) # Delete testing directory and files rm(REGRID_DIR; recursive = true, force = true) end + + @testset "test hdwrite_regridfile_rll_to_cgll 3d space for FT=$FT" begin + # Test setup + R = FT(6371e3) + space = TestHelper.create_space(FT, nz = 2, ne = 16, R = R) + + ispath(REGRID_DIR) ? rm(REGRID_DIR; recursive = true, force = true) : nothing + mkpath(REGRID_DIR) + + # lat-lon dataset + data = ones(720, 360, 2, 3) # (lon, lat, z, time) + time = [19000101.0, 19000201.0, 19000301.0] + lats = collect(range(-90, 90, length = 360)) + lons = collect(range(-180, 180, length = 720)) + z = [1000.0, 2000.0] + data = reshape(sin.(lats * π / 90)[:], 1, :, 1, 1) .* data + varname = "sinlat" + + # save the lat-lon data to a netcdf file in the required format for TempestRemap + datafile_rll = joinpath(REGRID_DIR, "lat_lon_data.nc") + NCDataset(datafile_rll, "c") do ds + + defDim(ds, "lat", size(lats)...) + defDim(ds, "lon", size(lons)...) + defDim(ds, "z", size(z)...) + defDim(ds, "date", size(time)...) + + defVar(ds, "lon", lons, ("lon",)) + defVar(ds, "lat", lats, ("lat",)) + defVar(ds, "z", z, ("z",)) + defVar(ds, "date", time, ("date",)) + + defVar(ds, varname, data, ("lon", "lat", "z", "date")) + + end + + hd_outfile_root = "data_cgll_test" + Regridder.hdwrite_regridfile_rll_to_cgll( + FT, + REGRID_DIR, + datafile_rll, + varname, + space, + mono = true, + hd_outfile_root = hd_outfile_root, + ) + + # read in data on CGLL grid from the last saved date + date1 = TimeManager.strdate_to_datetime.(string(Int(time[end]))) + cgll_path = joinpath(REGRID_DIR, "$(hd_outfile_root)_$date1.hdf5") + hdfreader = InputOutput.HDF5Reader(cgll_path) + T_cgll = InputOutput.read_field(hdfreader, varname) + Base.close(hdfreader) + + # regrid back to lat-lon + T_rll, _ = Regridder.cgll2latlonz(T_cgll) + + # check consistency across z-levels + @test T_rll[:, :, 1] == T_rll[:, :, 2] + + # check consistency of CGLL remapped data with original data + @test all(isapprox.(extrema(data), extrema(parent(T_cgll)), atol = 1e-2)) + + # check consistency of lat-lon remapped data with original data + @test all(isapprox.(extrema(data), extrema(T_rll), atol = 1e-3)) + + # visual inspection + # Plots.plot(T_cgll) # using ClimaCorePlots + # Plots.contourf(T_rll[:,:,1]) + + # Delete testing directory and files + rm(REGRID_DIR; recursive = true, force = true) + + end end end