Merge pull request #19 from PALEOtoolkit/plot_coords

get_array can now modify / replace coordinates
PALEOtoolkit · Jul 31, 2022 · 75cb77e · 75cb77e · sjdaines · Jul 31, 2022
2 parents 64865fe + 2bc7397
commit 75cb77e
Show file tree

Hide file tree

Showing 6 changed files with 344 additions and 67 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "PALEOmodel"
 uuid = "bf7b4fbe-ccb1-42c5-83c2-e6e9378b660c"
 authors = ["Stuart Daines <[email protected]>"]
-version = "0.15.3"
+version = "0.15.4"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

diff --git a/docs/src/HOWTOshowmodelandoutput.md b/docs/src/HOWTOshowmodelandoutput.md
@@ -120,26 +120,36 @@ The output can be plotted using the Julia Plots.jl package, see [Plotting output
 
 ## Spatial or wavelength-dependent output
 
-To analyze spatial or eg wavelength-dependent output (eg time series from a 1D column or 3D general circulation model, or quantities that are a function of wavelength or frequency), [`PALEOmodel.get_array`](@ref) takes additional arguments to take 1D or 2D slices from the spatial, spectral and timeseries data. The [`PALEOmodel.FieldArray`](@ref) returned includes coordinates to plot column (1D) and heatmap (2D) data.
+To analyze spatial or eg wavelength-dependent output (eg time series from a 1D column or 3D general circulation model, or quantities that are a function of wavelength or frequency), [`PALEOmodel.get_array`](@ref) takes an additional `selectargs::NamedTuple` argument to take 1D or 2D slices from the spatial, spectral and timeseries data. The [`PALEOmodel.FieldArray`](@ref) returned includes default coordinates to plot column (1D) and heatmap (2D) data, these can be overridden by supplying the optional `coords` keyword argument.
 
 ### Examples for a column-based model
 
 Visualisation of spatial and wavelength-dependent output from the PALEOdev.jl ozone photochemistry example (a single 1D atmospheric column):
 
 #### 1D column data
     julia> plot(title="O3 mixing ratio", output, "atm.O3_mr", (tmodel=[0.0, 0.1, 1.0, 10.0, 100.0, 1000.0], column=1),
-                swap_xy=true, xaxis=:log, labelattribute=:filter_records) # plots O3 vs height
+                swap_xy=true, xscale=:log10, labelattribute=:filter_records) # plots O3 vs default height coordinate
 
-Here the `labelattribute=:filter_records` keyword argument is used to generate plot labels from the `:filter_records` FieldArray attribute, which contains the `tmodel` values used to select the timeseries records.  The plot recipe expands
+Here the optional `labelattribute=:filter_records` keyword argument is used to generate plot labels from the `:filter_records` FieldArray attribute, which contains the `tmodel` values used to select the timeseries records.  The plot recipe expands
 the Vector-valued `tmodel` argument to overlay a sequence of plots.
 
 This is equivalent to first creating and then plotting a sequence of `FieldArray` objects:
-
-    julia> O3_mr = PALEOmodel.get_array(run.output, "atm.O3_mr", tmodel=0.0, column=1)
-    julia> plot(title="O3 mixing ratio", O3_mr, swap_xy=true, xaxis=:log, labelattribute=:filter_records)
-    julia> O3_mr = PALEOmodel.get_array(run.output, "atm.O3_mr", tmodel=0.1, column=1)
+    julia> O3_mr = PALEOmodel.get_array(run.output, "atm.O3_mr", (tmodel=0.0, column=1))
+    julia> plot(title="O3 mixing ratio", O3_mr, swap_xy=true, xscale=:log10, labelattribute=:filter_records)
+    julia> O3_mr = PALEOmodel.get_array(run.output, "atm.O3_mr", (tmodel=0.1, column=1))
     julia> plot!(O3_mr, swap_xy=true, labelattribute=:filter_records)
 
+The default height coordinate from the model grid can be replaced using the optional `coords` keyword argument, eg
+    julia> plot(title="O3 mixing ratio", output, "atm.O3_mr", (tmodel=[0.0, 0.1, 1.0, 10.0, 100.0, 1000.0], column=1),
+                coords=["p"=>("atm.pmid", "atm.plower", "atm.pupper")],
+                swap_xy=true, xscale=:log10, yflip=true, yscale=:log10, labelattribute=:filter_records) # plots O3 vs pressure
+
+The current values in the `modeldata` struct can also be plotted using the same syntax, eg
+
+    julia> plot(title="O2, O3 mixing ratios", modeldata, ["atm.O2_mr", "atm.O3_mr"], (column=1,),
+                swap_xy=true, xscale=:log10) # plot current value of O2, O3 vs height
+
+
 #### Wavelength-dependent data
     julia> plot(title="direct transmittance", output, ["atm.direct_trans"], (tmodel=1e12, column=1, cell=[1, 80]),
                 ylabel="fraction", labelattribute=:filter_region) # plots vs wavelength

diff --git a/src/FieldArray.jl b/src/FieldArray.jl
@@ -35,6 +35,23 @@ function Base.:*(fa_in::FieldArray, a::Real)
 end
 Base.:*(a::Real, fa_in::FieldArray) = fa_in*a
 
+# default name from attributes
+default_fieldarray_name(attributes::Nothing) = ""
+
+function default_fieldarray_name(attributes::Dict)
+    name = get(attributes, :domain_name, "")
+    name *= isempty(name) ? "" : "."
+    name *= get(attributes, :var_name, "")
+
+    selectargs_records = get(attributes, :filter_records, NamedTuple())
+    selectargs_region = get(attributes, :filter_region, NamedTuple())
+    if !isempty(selectargs_region) || !isempty(selectargs_records)
+        name *= "(" * join(["$k=$v" for (k, v) in Dict(pairs(merge(selectargs_records, selectargs_region)))], ", ") * ")"
+    end
+
+    return name
+end
+
 #############################################################
 # Create from PALEO objects
 #############################################################
@@ -47,7 +64,7 @@ Get FieldArray from PALEO object `obj`
 function get_array end
 
 """
-    get_array(f::Field; [attributes=nothing], [selectargs...]) -> FieldArray
+    get_array(f::Field [, selectargs::NamedTuple]; [attributes=nothing]) -> FieldArray
 
 Return a [`FieldArray`](@ref) containing `f::Field` data values and
 any attached coordinates, for the spatial region defined by `selectargs`.
@@ -62,26 +79,29 @@ function get_array(
     attributes=nothing,
 ) where {D, V, N, M}
 
-    return FieldArray("", f.values, f.data_dims, attributes)
+    return FieldArray(default_fieldarray_name(attributes), f.values, f.data_dims, attributes)
 end
 
 function get_array(
-    f::PB.Field{D, PB.CellSpace, V, 0, M};
-    attributes=nothing,
-    selectargs...
+    f::PB.Field{D, PB.CellSpace, V, 0, M}, selectargs::NamedTuple=NamedTuple();
+    attributes=nothing,    
 ) where {D, V, M}
 
     values, dims = PB.Grids.get_region(f.mesh, f.values; selectargs...)
 
-    return FieldArray("", values, dims, attributes)
+    if !isnothing(attributes) && !isempty(selectargs)
+        attributes = copy(attributes)
+        attributes[:filter_region] = selectargs
+    end
+
+    return FieldArray(default_fieldarray_name(attributes), values, dims, attributes)
 end
 
 # single data dimension
 # TODO generalize this to arbitrary data dimensions
 function get_array(
-    f::PB.Field{D, PB.CellSpace, V, 1, M};
+    f::PB.Field{D, PB.CellSpace, V, 1, M}, selectargs::NamedTuple=NamedTuple();
     attributes=nothing,
-    selectargs...
 ) where {D, V, M}
 
     f_space_dims_colons = ntuple(i->Colon(), ndims(f.values) - 1)
@@ -105,23 +125,107 @@ function get_array(
         end
     end
 
-    return FieldArray("", values, (dims..., f.data_dims...), attributes)
+    if !isnothing(attributes) && !isempty(selectargs)
+        attributes = copy(attributes)
+        attributes[:filter_region] = selectargs
+    end
+
+    return FieldArray(default_fieldarray_name(attributes), values, (dims..., f.data_dims...), attributes)
 end
 
 
-"""
-    get_array(model::PB.Model, modeldata, varnamefull; selectargs...) -> FieldArray
 
+"""
+    get_array(modeldata, varnamefull [, selectargs::NamedTuple] [; coords::AbstractVector]) -> FieldArray
+   
 Get [`FieldArray`](@ref) by Variable name, for spatial region defined by `selectargs`
 (which are passed to `PB.Grids.get_region`).
+
+Optional argument `coords` can be used to supply plot coordinates from Variable in output.
+Format is a Vector of Pairs of "coord_name"=>("var_name1", "var_name2", ...)
+
+Example: to replace a 1D column default pressure coordinate with a z coordinate:
+ 
+    coords=["z"=>("atm.zmid", "atm.zlower", "atm.zupper")]
+
+NB: the coordinates will be generated by applying `selectargs`,
+so the supplied coordinate Variables must have the same dimensionality as `vars`.
 """
-function get_array(model::PB.Model, modeldata::PB.AbstractModelData, varnamefull::AbstractString; selectargs...)
-    var = PB.get_variable(model, varnamefull)
+function get_array(
+    modeldata::PB.AbstractModelData, varnamefull::AbstractString, selectargs::NamedTuple=NamedTuple();
+    coords=nothing,
+)
+    var = PB.get_variable(modeldata.model, varnamefull)
     !isnothing(var) ||
         throw(ArgumentError("Variable $varnamefull not found"))
     f = PB.get_field(var, modeldata)
-    attrbs = deepcopy(var.attributes)
-    attrbs[:var_name] = var.name
-    attrbs[:domain_name] = var.domain.name
-    return get_array(f; attributes=attrbs, selectargs...)
+    attributes = copy(var.attributes)
+    attributes[:var_name] = var.name
+    attributes[:domain_name] = var.domain.name    
+
+    varray = get_array(f, selectargs; attributes=attributes)
+
+    if isnothing(coords)
+        # keep original coords (if any)
+    else
+        check_coords_argument(coords) ||
+            error("argument coords should be a Vector of Pairs of \"coord_name\"=>(\"var_name1\", \"var_name2\", ...), eg: [\"z\"=>(\"atm.zmid\", \"atm.zlower\", \"atm.zupper\"), ...]")
+
+        vec_coords_arrays = [
+            coord_name => Tuple(get_array(modeldata, cvn, selectargs) for cvn in coord_varnames) 
+            for (coord_name, coord_varnames) in coords
+        ]
+
+        varray = update_coordinates(varray, vec_coords_arrays) 
+    end
+
+    return varray
+end
+
+# check 'coords' of form [] or ["z"=>[ ... ], ] or ["z"=>(...),]
+check_coords_argument(coords) =
+    isa(coords, AbstractVector) && (
+        isempty(coords) || (
+            isa(coords, AbstractVector{<:Pair}) &&
+            isa(first(first(coords)), AbstractString) &&
+            isa(last(first(coords)), Union{AbstractVector, Tuple})
+        )
+    )
+
+"""
+    update_coordinates(varray::FieldArray, vec_coord_arrays::AbstractVector) -> FieldArray
+
+Replace or add coordinates `vec_coord_arrays` to `varray`.
+
+`new_coord_arrays` is a Vector of Pairs of "coord_name"=>(var1::FieldArray, var2::FieldArray, ...)
+
+Example: to replace a 1D column default pressure coordinate with a z coordinate:
+ 
+    coords=["z"=>(zmid::FieldArray, zlower::FieldArray, atm.zupper::FieldArray)]
+"""
+function update_coordinates(varray::FieldArray, vec_coord_arrays::AbstractVector)
+
+    check_coords_argument(vec_coord_arrays) || 
+        error("argument vec_coords_arrays should be a Vector of Pairs of \"coord_name\"=>(var1::FieldArray, var2::FieldArray, ...), eg: [\"z\"=>(zmid::FieldArray, zlower::FieldArray, atm.zupper::FieldArray)]")
+
+    # generate Vector of NamedDimensions to use as new coordinates
+    named_dimensions = PB.NamedDimension[]
+    for (coord_name, coord_arrays) in vec_coord_arrays
+        fixed_coords = []
+        for coord_array in coord_arrays
+            push!(fixed_coords, PB.FixedCoord(get(coord_array.attributes, :var_name, ""), coord_array.values, coord_array.attributes))
+        end
+        push!(
+            named_dimensions, PB.NamedDimension(
+                coord_name,
+                length(first(fixed_coords).values),
+                fixed_coords,
+            )
+        )
+    end
+
+    # replace coordinates
+    varray_newcoords = FieldArray(varray.name, varray.values, Tuple(named_dimensions), varray.attributes)
+
+    return varray_newcoords
 end
diff --git a/src/FieldRecord.jl b/src/FieldRecord.jl
@@ -128,12 +128,14 @@ function Base.copy(fr::FieldRecord{D, S, V, N, M, R}) where {D, S, V, N, M, R}
     )
 end
 
-"""
-    get_array(fr::FieldRecord; [recordarg] [,selectargs...]) -> FieldArray
+"""    
+    get_array(fr::FieldRecord [, allselectargs::NamedTuple] [; coords::AbstractVector]) -> FieldArray
+    [deprecated] get_array(fr::FieldRecord [; coords::AbstractVector] [; allselectargs...]) -> FieldArray
 
 Return a [`FieldArray`](@ref) containing `fr::FieldRecord` data values and
-any attached coordinates, for records defined by `recordarg` and the 
-spatial region defined by `selectargs`.
+any attached coordinates, for records and spatial region defined by `allselectargs`.
+
+`allselectarg` may include `recordarg` to define records, `selectargs` to define a spatial region.
 
 `recordarg` can be one of:
 - `records=r::Int` to select a single record, or `records = first:last` to select a range.
@@ -144,16 +146,65 @@ spatial region defined by `selectargs`.
 
 Available `selectargs` depend on the grid `fr.mesh`, and 
 are passed to `PB.Grids.get_region`.
+
+Optional argument `coords` can be used to supply coordinates from additional `FieldRecords`, replacing any coordinates
+attached to `fr`. Format is a Vector of Pairs of "coord_name"=>(cr1::FieldRecord, cr2::FieldRecord, ...)
+
+Example: to replace a 1D column default pressure coordinate with a z coordinate:
+
+    coords=["z"=>(zmid::FieldRecord, zlower::FieldRecord, zupper::FieldRecord)]
+
+NB: the coordinates will be generated by applying `allselectargs`,
+so the supplied coordinate FieldRecords must have the same dimensionality as `fr`.
 """
 function get_array(
-    fr::FieldRecord{D, S, V, N, M, R}; selectargs...
+    fr::FieldRecord; 
+    coords=nothing,
+    allselectargs...
+)
+    isempty(allselectargs) ||
+        Base.depwarn(
+            "allselectargs... will be deprecated in a future release.  Please use allselectargs::NamedTuple instead",
+            :get_array,
+        )
+
+    return get_array(fr, NamedTuple(allselectargs); coords=coords)
+end
+
+function get_array(
+    fr::FieldRecord, allselectargs::NamedTuple; # allselectargs::NamedTuple=NamedTuple() creates a method ambiguity with deprecated form above
+    coords::Union{Nothing, AbstractVector}=nothing
+)
+
+    # get data NB: with original coords or no coords
+    varray = _get_array(fr, allselectargs)
+
+    if !isnothing(coords)
+        check_coords_argument(coords) ||
+            error("coords argument should be of form 'coords=[\"z\"=>zmid::FieldRecord, zlower::FieldRecord, zupper::FieldRecord), ...]")
+
+        # get arrays for replacement coordinates
+        vec_coord_arrays = [
+            coord_name => Tuple(_get_array(coord_fr, allselectargs) for coord_fr in coord_fieldrecords) 
+            for (coord_name, coord_fieldrecords) in coords
+        ]
+
+        # replace coordinates
+        varray = update_coordinates(varray, vec_coord_arrays)
+    end
+
+    return varray
+end
+
+function _get_array(
+    fr::FieldRecord{D, S, V, N, M, R}, allselectargs::NamedTuple, 
 ) where {D, S, V, N, M, R}
 
     # select records to use and create PB.NamedDimension
     ridx = 1:length(fr)
-    selectargs_region = Dict()
+    selectargs_region = NamedTuple()
     selectargs_records = NamedTuple()
-    for (k, v) in selectargs
+    for (k, v) in zip(keys(allselectargs), allselectargs)
         if k==:records
             if v isa Integer
                 ridx = [v]                
@@ -170,23 +221,18 @@ function get_array(
                 end
             end            
         else
-            selectargs_region[k] = v
+            selectargs_region = merge(selectargs_region, NamedTuple((k=>v,)))
         end
     end
     records_dim = PB.NamedDimension("records", length(ridx), PB.get_region(fr.coords_record, ridx))
 
     # add attributes for selection used
     attributes = copy(fr.attributes)
-    attributes[:filter_records] = selectargs_records
-    attributes[:filter_region] = NamedTuple(selectargs_region)
+    isempty(selectargs_records) || (attributes[:filter_records] = selectargs_records;)
+    isempty(selectargs_region) || (attributes[:filter_region] = selectargs_region;)
 
     # Generate name from attributes
-    name = get(attributes, :domain_name, "")
-    name *= isempty(name) ? "" : "."
-    name *= get(attributes, :var_name, "")
-    if !isempty(selectargs_region) || !isempty(selectargs_records)
-        name *= "(" * join(["$k=$v" for (k, v) in merge(Dict(pairs(selectargs_records)), selectargs_region)], ", ") * ")"
-    end
+    name = default_fieldarray_name(attributes)
 
     # Select selectargs_region 
     if field_single_element(fr)
@@ -213,7 +259,10 @@ function get_array(
         # pass through to Field
 
         # get FieldArray from first Field record and use this to figure out array shapes etc
-        far = get_array(fr[first(ridx)]; selectargs_region...)        
+        far = get_array(fr[first(ridx)], selectargs_region)
+        # TODO - Julia bug ? length(far.dims) returning wrong value, apparently triggered by this line
+        #        attributes = isnothing(attributes) ? Dict{Symbol, Any}() : copy(attributes)
+        # in get_array
         if length(ridx) > 1
             # add additional (last) dimension for records
             favalues = Array{eltype(far.values), length(far.dims)+1}(undef, size(far.values)..., length(ridx))
@@ -242,7 +291,7 @@ function get_array(
         # fill with values from FieldArrays for Fields for remaining records
         if length(ridx) > 1
             for (i, r) in enumerate(ridx[2:end])
-                far = get_array(fr[r]; selectargs_region...)
+                far = get_array(fr[r], selectargs_region)
                 if isempty(far.dims)
                     fa.values[i+1] = far.values
                 else