allow CartesianIndex and CartesianIndices with load! (issue #250)

JuliaGeo · Mar 4, 2024 · 7460265 · 7460265
1 parent ce245f0
commit 7460265
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/src/variable.jl b/src/variable.jl
@@ -92,10 +92,9 @@ function checkbuffer(len,data)
     end
 end
 
-@inline function unsafe_load!(ncvar::Variable, data, indices::Union{Integer, UnitRange, StepRange, Colon}...)
+@inline function unsafe_load!(ncvar::Variable, data, indices::Union{Integer, UnitRange, StepRange, CartesianIndex, CartesianIndices, Colon}...)
     sizes = size(ncvar)
-    normalizedindices = normalizeindexes(sizes, indices)
-    ind = to_indices(ncvar,normalizedindices)
+    ind = to_indices(ncvar,indices)
 
     start,count,stride,jlshape = ncsub(ncvar,ind)
 
@@ -137,7 +136,7 @@ load!(ds["temp"].var,data,:,1) # loads the 1st column
     array must be `UInt8` and cannot be the julia `Char` type, because the
     julia `Char` type uses 4 bytes and the NetCDF `NC_CHAR` only 1 byte.
 """
-@inline function load!(ncvar::Variable{T,N}, data::AbstractArray{T}, indices::Union{Integer, UnitRange, StepRange, Colon}...) where {T,N}
+@inline function load!(ncvar::Variable{T,N}, data::AbstractArray{T}, indices::Union{Integer, UnitRange, StepRange, CartesianIndex, CartesianIndices, Colon}...) where {T,N}
     unsafe_load!(ncvar, data, indices...)
 end
 

diff --git a/test/test_variable.jl b/test/test_variable.jl
@@ -288,3 +288,31 @@ data2 = zeros(Int,10)
 # asking too many elements
 @test_throws BoundsError NCDatasets.load!(ds["data"].var,data2,1:10)
 close(ds)
+
+# issue 250
+
+fname = tempname()
+ds = NCDataset(fname,"c")
+defDim(ds,"lon",100)
+defDim(ds,"lat",110)
+v = defVar(ds,"temperature",Float32,("lon","lat"))
+data = [Float32(i+j) for i = 1:100, j = 1:110];
+v[:,:] = data;
+close(ds)
+
+# reading
+ds = NCDataset(fname)
+
+# works:
+ds["temperature"][CartesianIndices((1:10,10:30))]
+ds["temperature"][CartesianIndex(1,1)]
+
+# read in-place
+v = zeros(Float32, 10, 21);
+NCDatasets.load!(variable(ds, "temperature"), v, CartesianIndices((1:10,10:30)))
+@test v[:,:] == data[CartesianIndices((1:10,10:30))]
+
+
+vv = [1.0f0]
+NCDatasets.load!(variable(ds, "temperature"), vv, CartesianIndex(5,5))
+@test vv[1] == data[CartesianIndex(5,5)]