revs

docs/src/regridder.md

@@ -30,4 +30,6 @@ 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/Regridder.jl

@@ -34,7 +34,7 @@ export write_to_hdf5,
 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`.
@@ -44,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::SubArray, R, space::Spaces.SpectralElementSpace2D)
+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
@@ -65,21 +67,35 @@ function reshape_cgll_sparse_to_field!(field::Fields.Field, in_array::SubArray,
     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,
-    space::Spaces.ExtrudedFiniteDifferenceSpace,
+    ::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))
+            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
@@ -180,23 +196,6 @@ function hdwrite_regridfile_rll_to_cgll(
         @warn "Using the existing $datafile_cgll : check topology is consistent"
     end
 
-    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.(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, coords = NCDataset(datafile_cgll, "r") do ds_wt
         (
@@ -237,7 +236,6 @@ function hdwrite_regridfile_rll_to_cgll(
         length(times),
     )
 
-    # TODO: extend write! to handle time-dependent fields
     map(
         x -> write_to_hdf5(
             REGRID_DIR,
@@ -252,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)
 

test/Project.toml

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

test/regridder_tests.jl

@@ -241,20 +241,20 @@ for FT in (Float32, Float64)
             # 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)
+            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