Allow CartesianIndex and CartesianIndices with load!

[haakon-e]

Describe the bug

Reading a dataset normally, you can index it with CartesianIndex and CartesianIndices, which is often very useful!

For performance reasons it is sometimes useful to preallocate the output array, and therefore use load!. However, it appears this kind of indexing does not work with load!.

To Reproduce

# create dataset
ds = NCDataset("/tmp/test.nc","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("/tmp/test.nc")

# 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, 1:10,10:30)  # works
NCDatasets.load!(variable(ds, "temperature"), v, CartesianIndices((1:10,10:30)))  # does not work

vv = [1.0f0]
NCDatasets.load!(variable(ds, "temperature"), vv, 1, 1)  # works
NCDatasets.load!(variable(ds, "temperature"), vv, CartesianIndex(5,5))  # does not work

stack traces

julia> NCDatasets.load!(variable(ds, "temperature"), v, CartesianIndices((1:10,10:30))) 
ERROR: MethodError: no method matching load!(::NCDatasets.Variable{Float32, 2, NCDataset{Nothing, Missing}}, ::Matrix{Float32}, ::CartesianIndices{2, Tuple{UnitRange{Int64}, UnitRange{Int64}}})

Closest candidates are:
  load!(::NCDatasets.Variable{T, 2}, ::AbstractArray{T}, ::Colon, ::UnitRange) where T
   @ NCDatasets ~/.julia/packages/NCDatasets/ivzVX/src/variable.jl:148
  load!(::NCDatasets.Variable{T, N}, ::AbstractArray{T}, ::Union{Colon, Integer, StepRange, UnitRange}...) where {T, N}
   @ NCDatasets ~/.julia/packages/NCDatasets/ivzVX/src/variable.jl:140
  load!(::CommonDataModel.MemoryVariable, ::Any, ::Any...)
   @ CommonDataModel ~/.julia/packages/CommonDataModel/GGvMn/src/memory_dataset.jl:57
  ...



julia> NCDatasets.load!(variable(ds, "temperature"), vv, CartesianIndex(5,5))  # does not work
ERROR: MethodError: no method matching load!(::NCDatasets.Variable{Float32, 2, NCDataset{Nothing, Missing}}, ::Vector{Float32}, ::CartesianIndex{2})

Closest candidates are:
  load!(::NCDatasets.Variable{T, 2}, ::AbstractArray{T}, ::Colon, ::UnitRange) where T
   @ NCDatasets ~/.julia/packages/NCDatasets/ivzVX/src/variable.jl:148
  load!(::NCDatasets.Variable{T, N}, ::AbstractArray{T}, ::Union{Colon, Integer, StepRange, UnitRange}...) where {T, N}
   @ NCDatasets ~/.julia/packages/NCDatasets/ivzVX/src/variable.jl:140
  load!(::CommonDataModel.MemoryVariable, ::Any, ::Any...)
   @ CommonDataModel ~/.julia/packages/CommonDataModel/GGvMn/src/memory_dataset.jl:57
  ...

Expected behavior

Cartesian indexing to work with load!

Environment

OS info

julia> versioninfo()
Julia Version 1.10.2
Commit bd47eca2c8a (2024-03-01 10:14 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 8 × Apple M1
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 4 virtual cores)
Environment:
  JULIA_NUM_THREADS = auto
  JULIA_EDITOR = code

(examples) pkg> st NCDatasets
Status `~/Documents/examples/Project.toml`
  [85f8d34a] NCDatasets v0.14.2

[haakon-e]

PS: To get around this, in my internals, I currently do variants of:

inds = CartesianIndices((1:10,10:30))
NCDatasets.load!(variable(ds, "temperature"), v, inds.indices...)

ind = CartesianIndex(5,5)
NCDatasets.load!(variable(ds, "temperature"), vv, ind.I...)

and I imagine a simple implementation of this could simply expand the indices in a similar fashion.

[haakon-e]

I feel like a fix should be within reach by adding methods to

NCDatasets.jl/src/variable.jl

Lines 455 to 465 in ce245f0

	_normalizeindex(n,ind::Base.OneTo) = 1:1:ind.stop
	_normalizeindex(n,ind::Colon) = 1:1:n
	_normalizeindex(n,ind::Integer) = ind:1:ind
	_normalizeindex(n,ind::UnitRange) = StepRange(ind)
	_normalizeindex(n,ind::StepRange) = ind
	_normalizeindex(n,ind) = error("unsupported index")
	# indexes can be longer than sz
	function normalizeindexes(sz,indexes)
	return ntuple(i -> _normalizeindex(sz[i],indexes[i]), length(sz))
	end

of the form

_normalizeindex(n,ind::CartesianIndex) = ind
_normalizeindex(n,ind::CartesianIndices) = ind

This doesn't immediately works, e.g. for CartesianIndex(1,1) because then sz and indexes doesn't have the same length. However, sz is currently only used for parsing :, so perhaps this could be fixed some other way?

Certainly, the to_indices method here expands CartesianIndex correctly:

NCDatasets.jl/src/variable.jl

Lines 95 to 108 in ce245f0

	@inline function unsafe_load!(ncvar::Variable, data, indices::Union{Integer, UnitRange, StepRange, Colon}...)
	sizes = size(ncvar)
	normalizedindices = normalizeindexes(sizes, indices)
	ind = to_indices(ncvar,normalizedindices)
	start,count,stride,jlshape = ncsub(ncvar,ind)
	@boundscheck begin
	checkbounds(ncvar,indices...)
	checkbuffer(prod(count),data)
	end
	nc_get_vars!(ncvar.ds.ncid,ncvar.varid,start,count,stride,data)
	end

[Alexander-Barth]

Thanks a lot for your detailed issue report.
Yes, it makes totally sense to support CartesianIndex and CartesianIndices with load!!.
The latest commit 7460265 seems to solve your test case.

(it seems that we do not need normalizeindexes at all and can just rely on to_indices instead.)

[haakon-e]

Thanks a lot! This seems to work great.

A small aside: In the future, I'd greatly appreciate if my contributions are explicitly acknowledged in the commit(s) (see how) with: Co-authored-by: Haakon Ludvig Langeland Ervik <[email protected]>. I'll try to remember noting this next time too.

Thanks again for the quick turnaround! Would you mind pushing a new release?

[Alexander-Barth]

No problem. I changed the commit message to explicit acknowledge you.

0335d26
0284393

Do not hesitate to make a PR; this will sort out the attribution automatically :-)

I will make a release shorty!