Merge pull request #572 from SpeedyWeather/mg/ringgrids-from-array

RingGrids: Fix RingGrids Array conversion

.gitignore

@@ -37,9 +37,13 @@ run_*/
 # no video outputs
 *.mp4
 
-# JOSS paper
+# JOSS paper
 docs/joss/media
 docs/joss/paper.jats
 
 # not sure where they are from
 default.profraw
+
+# vs code
+.vscode
+settings.json

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
-* Fixed a bug in `RingGrids`, so that now broadcasts are defined even when the dimensions mismatch in some cases 
+- RingGrids: To wrap an Array with the horizontal dimension in matrix shape into a full grid, one has to use e.g. `FullGaussianGrid(map, input_as=Matrix)` now. [#572](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/572)
+* RingGrids: Fixed a bug in `RingGrids`, so that now broadcasts are defined even when the dimensions mismatch in some cases 
 
 ## v0.11.0

docs/src/examples_2D.md

@@ -142,11 +142,13 @@ and then plotted via `heatmap(lon, lat, vor)`. While you can do that to give you
 on the plotting, SpeedyWeather.jl also defines an extension for Makie.jl, see [Extensions](@ref).
 Because if our matrix `vor` here was an `AbstractGrid` (see [RingGrids](@ref)) then all
 its geographic information (which grid point is where) would directly be encoded in the type.
-From the netCDF file you need to use the longitude and latitude dimensions.
+From the netCDF file, however, you would need to use the longitude and latitude dimensions.
+
+So we can also just do (`input_as=Matrix` here as all our grids use and expect a horizontal dimension
+flattened into a vector by default)
 
-So we can also just do
 ```@example galewsky_setup
-vor_grid = FullGaussianGrid(vor)
+vor_grid = FullGaussianGrid(vor, input_as=Matrix)
 
 using CairoMakie    # this will load the extension so that Makie can plot grids directly
 heatmap(vor_grid, title="Relative vorticity [1/s]")
@@ -319,4 +321,4 @@ Can you spot the Himalayas or the Andes?
 
 ## References
 
-[^G04]: Galewsky, Scott, Polvani, 2004. *An initial-value problem for testing numerical models of the global shallow-water equations*, Tellus A. DOI: [10.3402/tellusa.v56i5.14436](https://doi.org/10.3402/tellusa.v56i5.14436)
+[^G04]: Galewsky, Scott, Polvani, 2004. *An initial-value problem for testing numerical models of the global shallow-water equations*, Tellus A. DOI: [10.3402/tellusa.v56i5.14436](https://doi.org/10.3402/tellusa.v56i5.14436)

docs/src/ringgrids.md

@@ -42,8 +42,12 @@ map = randn(Float32, 8, 4)
 ```
 
 ```@example ringgrids
-grid = FullGaussianGrid(map)
+grid = FullGaussianGrid(map, input_as=Matrix)
 ```
+Note that `input_as=Matrix` is necessary as, RingGrids have a flattened horizontal dimension into a vector.
+To distinguish the 2nd horizontal dimension from a possible vertical dimension the keyword argument here
+is required.
+
 A full Gaussian grid has always ``2N`` x ``N`` grid points, but a `FullClenshawGrid` has ``2N`` x ``N-1``,
 if those dimensions don't match, the creation will throw an error. To reobtain the data from a grid,
 you can access its `data` field which returns a normal `Vector`
@@ -273,4 +277,4 @@ but the longitudes can be different for all four, a, b, c, d.
 
 ```@autodocs
 Modules = [SpeedyWeather.RingGrids]
-```
+```

docs/src/speedytransforms.md

@@ -178,7 +178,7 @@ to
 We now wrap this matrix
 therefore to associate it with the necessary grid information
 ```@example speedytransforms
-map = FullClenshawGrid(m)
+map = FullClenshawGrid(m, input_as=Matrix)
 
 using CairoMakie
 heatmap(map)

src/RingGrids/full_grids.jl

@@ -19,12 +19,25 @@ matrix_size(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) =
 
 ## CONVERSION
 # convert an AbstractMatrix to the full grids, and vice versa
-(Grid::Type{<:AbstractFullGrid})(M::AbstractMatrix) = Grid(vec(M))
+"""
+($TYPEDSIGNATURES)
+Initialize an instance of the grid from an Array. For keyword argument `input_as=Vector` (default)
+the leading dimension is interpreted as a flat vector of all horizontal entries in one layer.
+For `input_as==Matrx` the first two leading dimensions are interpreted as longitute and latitude.
+This is only possible for full grids that are a subtype of `AbstractFullGridArray`.
+"""
+(Grid::Type{<:AbstractFullGridArray})(M::AbstractArray; input_as=Vector) = Grid(M, input_as)
+
+function (Grid::Type{<:AbstractFullGridArray})(M::AbstractArray, input_as::Type{Matrix})
+    # flatten the two horizontal dimensions into one, identical to vec(M) for M <: AbstractMatrix
+    M_flat = reshape(M, :, size(M)[3:end]...)
+    Grid(M_flat)
+end 
+
 Base.Array(grid::AbstractFullGridArray) = Array(reshape(grid.data, :, get_nlat(grid), size(grid.data)[2:end]...))
 Base.Matrix(grid::AbstractFullGridArray) = Array(grid)
 
 ## INDEXING
-
 """$(TYPEDSIGNATURES) `UnitRange` for every grid point of grid `Grid` of resolution `nlat_half`
 on ring `j` (`j=1` is closest ring around north pole, `j=nlat` around south pole)."""
 function each_index_in_ring(

src/RingGrids/general.jl

@@ -117,12 +117,18 @@ function (::Type{Grid})(data::AbstractArray, nlat_half::Integer) where {Grid<:Ab
 end
 
 # if no nlat_half provided calculate it
-function (::Type{Grid})(data::AbstractArray) where {Grid<:AbstractGridArray}
+(::Type{Grid})(M::AbstractArray; input_as=Vector) where Grid<:AbstractGridArray = Grid(M, input_as)
+
+function (::Type{Grid})(data::AbstractArray, input_as::Type{Vector}) where Grid<:AbstractGridArray
     npoints2D = size(data, 1)                   # from 1st dim of data
     nlat_half = get_nlat_half(Grid, npoints2D)  # get nlat_half of Grid
     return Grid(data, nlat_half)
 end
 
+function (::Type{Grid})(data::AbstractArray, input_as::Type{Matrix})  where Grid<:AbstractGridArray
+    error("Only full grids can be created from matrix input")
+end 
+
 for f in (:zeros, :ones, :rand, :randn)
     @eval begin
         # general version with ArrayType(zeros(...)) conversion

src/dynamics/orography.jl

@@ -156,7 +156,7 @@ function initialize!(   orog::EarthOrography,
 
     # height [m], wrap matrix into a grid
     # TODO also read lat, lon from file and flip array in case it's not as expected
-    orography_highres = orog.file_Grid(ncfile["orog"][:, :])
+    orography_highres = orog.file_Grid(ncfile["orog"].var[:, :], input_as=Matrix)
 
     # Interpolate/coarsen to desired resolution
     interpolate!(orography, orography_highres)

src/output/output.jl

@@ -275,7 +275,7 @@ function initialize!(
         orog_grid = model.orography.orography
         orog_matrix = output.u
         output.as_matrix && (orog_matrix = Matrix(orog_grid))
-        output.as_matrix || RingGrids.interpolate!(output_Grid(output.u), orog_grid, output.interpolator)
+        output.as_matrix || RingGrids.interpolate!(output_Grid(output.u, input_as=Matrix), orog_grid, output.interpolator)
         defVar(dataset, "orography", orog_matrix, (lon_name, lat_name), attrib=orog_attribs)
     end
 
@@ -425,12 +425,12 @@ function write_netcdf_variables!(   output::OutputWriter,
             RingGrids.Matrix!(MGs...; quadrant_rotation, matrix_quadrant)
 
         else                    # or interpolate onto a full grid
-            :u in output_vars       && RingGrids.interpolate!(output_Grid(u),    u_grid, interpolator)
-            :v in output_vars       && RingGrids.interpolate!(output_Grid(v),    v_grid, interpolator)
-            :vor in output_vars     && RingGrids.interpolate!(output_Grid(vor),  vor_grid, interpolator)
-            :div in output_vars     && RingGrids.interpolate!(output_Grid(div),  div_grid, interpolator)
-            :temp in output_vars    && RingGrids.interpolate!(output_Grid(temp), temp_grid, interpolator)
-            :humid in output_vars   && RingGrids.interpolate!(output_Grid(humid), humid_grid, interpolator)
+            :u in output_vars       && RingGrids.interpolate!(output_Grid(u, input_as=Matrix),    u_grid, interpolator)
+            :v in output_vars       && RingGrids.interpolate!(output_Grid(v, input_as=Matrix),    v_grid, interpolator)
+            :vor in output_vars     && RingGrids.interpolate!(output_Grid(vor, input_as=Matrix),  vor_grid, interpolator)
+            :div in output_vars     && RingGrids.interpolate!(output_Grid(div, input_as=Matrix),  div_grid, interpolator)
+            :temp in output_vars    && RingGrids.interpolate!(output_Grid(temp, input_as=Matrix), temp_grid, interpolator)
+            :humid in output_vars   && RingGrids.interpolate!(output_Grid(humid, input_as=Matrix), humid_grid, interpolator)
         end
 
         # UNSCALE THE SCALED VARIABLES
@@ -466,11 +466,11 @@ function write_netcdf_variables!(   output::OutputWriter,
             RingGrids.Matrix!(MGs...; quadrant_rotation, matrix_quadrant)
         end
     else
-        :pres in output_vars && RingGrids.interpolate!(output_Grid(pres), pres_grid, interpolator)
-        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_cond), precip_large_scale, interpolator)
-        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv), precip_convection, interpolator)
-        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv), precip_convection, interpolator)
-        :cloud in output_vars && RingGrids.interpolate!(output_Grid(cloud), cloud_top, interpolator)
+        :pres in output_vars && RingGrids.interpolate!(output_Grid(pres, input_as=Matrix), pres_grid, interpolator)
+        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_cond, input_as=Matrix), precip_large_scale, interpolator)
+        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv, input_as=Matrix), precip_convection, interpolator)
+        :precip in output_vars && RingGrids.interpolate!(output_Grid(precip_conv, input_as=Matrix), precip_convection, interpolator)
+        :cloud in output_vars && RingGrids.interpolate!(output_Grid(cloud, input_as=Matrix), cloud_top, interpolator)
     end
 
     # after output set precip accumulators back to zero

src/physics/albedo.jl

@@ -56,6 +56,6 @@ function initialize!(albedo::AlbedoClimatology, model::PrimitiveEquation)
     end
     ncfile = NCDataset(path)
 
-    a = albedo.file_Grid(ncfile[albedo.varname][:, :])
+    a = albedo.file_Grid(ncfile[albedo.varname].var[:, :], input_as=Matrix)
     interpolate!(albedo.albedo, a)
 end

src/physics/land.jl

@@ -13,7 +13,7 @@ Base.@kwdef struct SeasonalLandTemperature{NF, Grid<:AbstractGrid{NF}} <: Abstra
     "number of latitudes on one hemisphere, Equator included"
     nlat_half::Int
 
-    # OPTIONS
+    # OPTIONS
     "Time step used to update land surface temperatures"
     Δt::Dates.Day = Dates.Day(3)
 
@@ -32,7 +32,7 @@ Base.@kwdef struct SeasonalLandTemperature{NF, Grid<:AbstractGrid{NF}} <: Abstra
     "The missing value in the data respresenting ocean"
     missing_value::NF = NF(NaN)
 
-    # to be filled from file
+    # to be filled from file
     "Monthly land surface temperatures [K], interpolated onto Grid"
     monthly_temperature::Vector{Grid} = [zeros(Grid, nlat_half) for _ in 1:12]
 end
@@ -70,7 +70,7 @@ function land_timestep!(land::PrognosticVariablesLand{NF},
                         initialize::Bool = false) where NF
 
     # escape immediately if Δt of land model hasn't passed yet
-    # unless the land hasn't been initialized yet
+    # unless the land hasn't been initialized yet
     initialize || (time - land.time) < land_model.Δt && return false    # = executed
 
     # otherwise update land prognostic variables:
@@ -79,7 +79,7 @@ function land_timestep!(land::PrognosticVariablesLand{NF},
     next_month = (this_month % 12) + 1      # mod for dec 12 -> jan 1
 
     # linear interpolation weight between the two months
-    # TODO check whether this shifts the climatology by 1/2 a month
+    # TODO check whether this shifts the climatology by 1/2 a month
     weight = convert(NF, Dates.days(time-Dates.firstdayofmonth(time))/Dates.daysinmonth(time))
 
     (; monthly_temperature) = land_model
@@ -104,7 +104,7 @@ Base.@kwdef struct SeasonalSoilMoisture{NF, Grid<:AbstractGrid{NF}} <: AbstractS
     "number of latitudes on one hemisphere, Equator included"
     nlat_half::Int
 
-    # OPTIONS
+    # OPTIONS
     "Depth of top soil layer [m]"
     D_top::Float64 = 0.07
 
@@ -117,7 +117,7 @@ Base.@kwdef struct SeasonalSoilMoisture{NF, Grid<:AbstractGrid{NF}} <: AbstractS
     "Soil wetness at wilting point [volume fraction]"
     W_wilt::Float64 = 0.17
 
-    # READ CLIMATOLOGY FROM FILE
+    # READ CLIMATOLOGY FROM FILE
     "path to the folder containing the soil moisture file, pkg path default"
     path::String = "SpeedyWeather.jl/input_data"
 
@@ -134,7 +134,7 @@ Base.@kwdef struct SeasonalSoilMoisture{NF, Grid<:AbstractGrid{NF}} <: AbstractS
     "The missing value in the data respresenting ocean"
     missing_value::NF = NF(NaN)
 
-    # to be filled from file
+    # to be filled from file
     "Monthly soil moisture volume fraction [1], top layer, interpolated onto Grid"
     monthly_soil_moisture_layer1::Vector{Grid} = [zeros(Grid, nlat_half) for _ in 1:12]
 
@@ -161,7 +161,7 @@ function soil_timestep!(land::PrognosticVariablesLand{NF},
     next_month = (this_month % 12) + 1      # mod for dec 12 -> jan 1
 
     # linear interpolation weight between the two months
-    # TODO check whether this shifts the climatology by 1/2 a month
+    # TODO check whether this shifts the climatology by 1/2 a month
     weight = convert(NF, Dates.days(time-Dates.firstdayofmonth(time))/Dates.daysinmonth(time))
 
     (; monthly_soil_moisture_layer1, monthly_soil_moisture_layer2) = soil_model
@@ -182,7 +182,7 @@ Base.@kwdef struct VegetationClimatology{NF, Grid<:AbstractGrid{NF}} <: Abstract
     "number of latitudes on one hemisphere, Equator included"
     nlat_half::Int
 
-    # OPTIONS
+    # OPTIONS
     "Combine high and low vegetation factor, a in high + a*low [1]"
     low_veg_factor::Float64 = 0.8
 
@@ -202,7 +202,7 @@ Base.@kwdef struct VegetationClimatology{NF, Grid<:AbstractGrid{NF}} <: Abstract
     "The missing value in the data respresenting ocean"
     missing_value::NF = NF(NaN)
 
-    # to be filled from file
+    # to be filled from file
     "High vegetation cover [1], interpolated onto Grid"
     high_cover::Grid = zeros(Grid, nlat_half)
 
@@ -227,8 +227,8 @@ function initialize!(vegetation::VegetationClimatology, model::PrimitiveEquation
     ncfile = NCDataset(path)
 
     # high and low vegetation cover
-    vegh = vegetation.file_Grid(ncfile[vegetation.varname_vegh][:, :])
-    vegl = vegetation.file_Grid(ncfile[vegetation.varname_vegl][:, :])
+    vegh = vegetation.file_Grid(ncfile[vegetation.varname_vegh].var[:, :], input_as=Matrix)
+    vegl = vegetation.file_Grid(ncfile[vegetation.varname_vegl].var[:, :], input_as=Matrix)
 
     # interpolate onto grid
     high_vegetation_cover = vegetation.high_cover
@@ -277,7 +277,7 @@ function soil_moisture_availability!(
                             soil_moisture_layer2,
                             high_cover, low_cover)
 
-        # Fortran SPEEDY source/land_model.f90 line 111 origin unclear
+        # Fortran SPEEDY source/land_model.f90 line 111 origin unclear
         veg = max(0, high_cover[ij] + low_veg_factor*low_cover[ij])
 
         # Fortran SPEEDY documentation eq. 51

src/physics/land_sea_mask.jl

@@ -80,7 +80,7 @@ function initialize!(land_sea_mask::LandSeaMask, model::PrimitiveEquation)
     ncfile = NCDataset(path)
 
     # high resolution land-sea mask
-    lsm_highres = file_Grid(ncfile["lsm"][:, :])
+    lsm_highres = file_Grid(ncfile["lsm"].var[:, :], input_as=Matrix)
 
     # average onto grid cells of the model
     RingGrids.grid_cell_average!(land_sea_mask.mask, lsm_highres)

test/grids.jl

@@ -203,6 +203,28 @@ end
     end
 end
 
+@testset "FullGrids conversions to/from Arrays" begin 
+    for idims in ((), (5,), (5,5))
+        NF = Float64
+        N = length(idims)+1
+        data = rand(8,4, idims...)
+        grid = FullGaussianArray(data, input_as=Matrix)
+        @test Array(grid) == data
+
+        data = rand(8,3, idims...)
+        grid = FullClenshawArray(data, input_as=Matrix)
+        @test Array(grid) == data
+
+        data = rand(8,3, idims...)
+        grid = FullHEALPixArray(data, input_as=Matrix)
+        @test Array(grid) == data
+
+        data = rand(8,3, idims...)
+        grid = FullOctaHEALPixArray(data, input_as=Matrix)
+        @test Array(grid) == data
+    end   
+end 
+
 @testset "Grid indices" begin
     for G in (  FullClenshawGrid,
                 FullGaussianGrid,