Skip to content

Commit

Permalink
feat: type-stable variable
Browse files Browse the repository at this point in the history
  • Loading branch information
haakon-e committed Apr 25, 2024
1 parent 5271611 commit 0768d5c
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 1 deletion.
3 changes: 2 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"

[targets]
test = ["Dates", "Test", "Random", "Printf", "IntervalSets"]
test = ["Dates", "Test", "Random", "Printf", "IntervalSets", "JET"]
21 changes: 21 additions & 0 deletions src/variable.jl
Original file line number Diff line number Diff line change
Expand Up @@ -75,14 +75,35 @@ end

"""
v = variable(ds::NCDataset,varname::String)
v = variable(ds, varname, DType, dimnames)
Return the NetCDF variable `varname` in the dataset `ds` as a
`NCDataset.Variable`. No scaling or other transformations are applied when the
variable `v` is indexed.
With the second form, supplying the correct data `DType` (e.g. `Float32`)
and `dimnames` (a tuple, e.g. `("lat", "lon")`) can lets you read in a type-stable way.
!!! note
Incorrectly specified `DType` or `dimnames` can lead to incorrect
results (e.g. "garbage" output) with no errors, so proceed with caution.
"""
variable(ds::NCDataset,varname::AbstractString) = _variable(ds,varname)
variable(ds::NCDataset,varname::Symbol) = _variable(ds,varname)

function variable(
ds::NCDataset,
varname::Union{AbstractString, Symbol},
DType::DataType,
dimnames::NTuple{N, AbstractString},
) where {N}
dimids = nc_inq_dimid.(ds.ncid, dimnames)
varid = nc_inq_varid(ds.ncid, varname)
Variable{DType, N, typeof(ds)}(ds, varid, dimids)
end


export variable


Expand Down
17 changes: 17 additions & 0 deletions test/test_variable.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ using Dates
using Printf
using NCDatasets
using DataStructures
import JET

sz = (4,5)
filename = tempname()
Expand Down Expand Up @@ -316,3 +317,19 @@ NCDatasets.load!(variable(ds, "temperature"), v, CartesianIndices((1:10,10:30)))
vv = [1.0f0]
NCDatasets.load!(variable(ds, "temperature"), vv, CartesianIndex(5,5))
@test vv[1] == data[CartesianIndex(5,5)]

### Test type stability of `variable(ds, varname, DType, dimnames)`
JET.@test_opt variable(ds, "temperature", Float32, ("lon", "lat"))
# Note: Incorrect order of `dimnames` doesn't error if slices are within bounds...
var1 = variable(ds, "temperature", Float32, ("lon", "lat"))[1:5, 1:10]
var2 = variable(ds, "temperature", Float32, ("lat", "lon"))[1:5, 1:10] # TODO: Ensure this errors
@test var1 == var2
# ... but incorrect order of `dimnames` and out-of-bounds slices errors
@test_throws variable(ds, "temperature", Float32, ("lat", "lon"))[:, 1] # `:` is actually `1:110`
# incorrect `dimnames` errors
@test_throws variable(ds, "temperature", Float32, ("lon", "lat", "time"))
@test_throws variable(ds, "temperature", Float32, ("lon",))
# Incorrect `DType` doesn't error, but gives incorrect result
d1 = variable(ds, "temperature", Float32, ("lon", "lat"))[1:5, 1:3]
d2 = variable(ds, "temperature", Float64, ("lon", "lat"))[1:5, 1:3]
@test d1 != d2 # TODO: Ensure this errors

0 comments on commit 0768d5c

Please sign in to comment.