Merge pull request #506 from CliMA/ln/3d-regrid

Enable prescribing 3d fields from files

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

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,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,34 +196,25 @@ 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
         (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,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

test/Project.toml

@@ -35,12 +35,14 @@ 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"
 MPI = "0.20"
 Pkg = "1"
 PrettyTables = "2"
 StaticArrays = "1"
+Statistics = "1"
 Unitful = "1"
 julia = "1.8"