From a5a14a0325880ae1ccd541b516d0c6cab8e93aeb Mon Sep 17 00:00:00 2001 From: LenkaNovak Date: Thu, 16 Nov 2023 19:23:34 -0800 Subject: [PATCH] regrid 3d init 2d works 4d works add to src anf test clean up add deps typo revs regrid 3d init 2d works 4d works add to src anf test clean up add deps typo revs fix rebase --- docs/src/postprocessor.md | 11 +-- docs/src/regridder.md | 4 + src/PostProcessor.jl | 25 +----- src/Regridder.jl | 184 ++++++++++++++++++++++++++++++++------ test/Project.toml | 2 + test/regridder_tests.jl | 76 +++++++++++++++- 6 files changed, 246 insertions(+), 56 deletions(-) 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