regrid 3d init

2d works

4d works

add to src anf test

clean up

add deps

typo

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
-```

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,5 @@ 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
 ```

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

src/Regridder.jl

@@ -26,7 +26,8 @@ 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. =#
@@ -44,7 +45,7 @@ Redundant nodes are populated using `dss` operations.
 - `in_array`: [Array] input used to fill `field`.
 - `R`: [NamedTuple] containing `target_idxs` and `row_indices` used for indexing.
 """
-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, space::Spaces.SpectralElementSpace2D)
     field_array = parent(field)
 
     fill!(field_array, zero(eltype(field_array)))
@@ -64,6 +65,33 @@ function reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::Array, R)
     Spaces.dss!(Fields.field_values(field), topology, hspace.quadrature_style)
 end
 
+function reshape_cgll_sparse_to_field!(
+    field::Fields.Field,
+    in_array::SubArray,
+    R,
+    space::Spaces.ExtrudedFiniteDifferenceSpace,
+)
+    field_array = parent(field)
+
+    fill!(field_array, zero(eltype(field_array)))
+    Nf = size(field_array, 4)
+    Nz = size(field_array, 1)
+
+    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))
+            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 +140,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,34 +180,42 @@ function hdwrite_regridfile_rll_to_cgll(
         @warn "Using the existing $datafile_cgll : check topology is consistent"
     end
 
-    function get_time(ds)
+    function get_coords(ds, space)
         if "time" in ds
             data_dates = Dates.DateTime.(ds["time"][:])
         elseif "date" in ds
-            data_dates = TimeManager.strdate_to_datetime.(string.(ds["date"][:]))
+            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
+        if space isa Spaces.ExtrudedFiniteDifferenceSpace
+            z = ds["z"][:]
+            return (data_dates, z)
+        else
+            return (data_dates,)
+        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
         (Array(ds_wt["S"]), Array(ds_wt["col"]), Array(ds_wt["row"]))
     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,7 +227,15 @@ 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(
@@ -486,4 +542,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

test/Project.toml

@@ -36,6 +36,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 ArtifactWrappers = "0.2"
 Aqua = "0.8"
 ClimaCoupler = "0.1"
+Dates = "1.9"
 Downloads = "1"
 HDF5_jll = "=1.12.2"
 IntervalSets = "0.5, 0.6, 0.7"
@@ -45,5 +46,6 @@ Pkg = "1"
 PrettyTables = "2"
 SimpleNonlinearSolve = "=0.1.23"
 StaticArrays = "1"
+Statistics = "1.9"
 Unitful = "1"
 julia = "1.8"

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)
+            # space2d = TestHelper.create_space(FT, nz = 1) # TODO
+
+            ispath(REGRID_DIR) ? rm(REGRID_DIR; recursive = true, force = true) : nothing
+            mkpath(REGRID_DIR)
+
+            # lat-lon dataset
+            data = ones(720, 360, 2, 3)
+            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"
+
+            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