diff --git a/src/DimensionalData.jl b/src/DimensionalData.jl index 1f188c6cf..a0cebadc9 100644 --- a/src/DimensionalData.jl +++ b/src/DimensionalData.jl @@ -62,6 +62,8 @@ export Begin, End export AbstractDimArray, DimArray +export CoordArray + export AbstractDimVector, AbstractDimMatrix, AbstractDimVecOrMat, DimVector, DimMatrix, DimVecOrMat export AbstractDimStack, DimStack @@ -101,6 +103,7 @@ include("array/methods.jl") include("array/matmul.jl") include("array/broadcast.jl") include("array/show.jl") +include("coordarray/coordarray.jl") # Stacks include("stack/stack.jl") include("stack/indexing.jl") diff --git a/src/coordarray/coordarray.jl b/src/coordarray/coordarray.jl new file mode 100644 index 000000000..960c6cfb4 --- /dev/null +++ b/src/coordarray/coordarray.jl @@ -0,0 +1,117 @@ +using DimensionalData: to_indices, slicedims, dimnum, show_main, show_after, print_block_separator + +""" + CoordArray <: AbstractDimArray + +A dimensional array type that supporting both dimension coordinates and non-dimension coordinates. + +## Example + +```julia +using DimensionalData + +# Create dimension coordinates +time_dim = Ti(DateTime(2020, 1, 1):Day(1):DateTime(2020, 1, 10)) +x_dim = X(1:5) +y_dim = Y(1:3) +z_dim = Z(1:4) + +# Create non-dimension coordinates +lat = rand(5, 3) # 2D coordinate +lon = rand(5, 3) # 2D coordinate +elevation = rand(4) # 1D coordinate along X dimension + +# Create the CoordArray +data = rand(10, 5, 3, 4) +da = CoordArray( + data, + (time_dim, x_dim, y_dim, z_dim), + coords=(; + latitude = ((X, Y), lat), + longitude = ((X, Y), lon), + elevation = ((Z,), elevation) + ), + name="temperature" +) + +da[Ti(At(DateTime(2020, 1, 5))), X(2:4), Y(1)] +``` +""" + +abstract type AbstractCoordArray{T,N,D,C,A} <: AbstractDimArray{T,N,D,A} end + +struct CoordArray{T,N,D<:Tuple,R<:Tuple,A<:AbstractArray{T,N},C,Na,Me} <: AbstractCoordArray{T,N,D,C,A} + data::A + dims::D + refdims::R + coords::C # Non-dimension coordinates: Symbol => AbstractArray + name::Na + metadata::Me + + function CoordArray( + data::A, dims::D, refdims::R, coords::C, name::Na, metadata::Me + ) where {D<:Tuple,R<:Tuple,A<:AbstractArray{T,N},C,Na,Me} where {T,N} + checkdims(data, dims) + checkcoords(coords, dims) + new{T,N,D,R,A,C,Na,Me}(data, dims, refdims, coords, name, metadata) + end +end + +# Constructors +function CoordArray( + data::AbstractArray, dims; refdims=(), coords=(;), name=NoName(), metadata=NoMetadata() +) + # Process coordinates: convert (dims, data) tuples to DimArrays + dims = format(dims, data) + processed_coords = _process_coords(coords, dims) + CoordArray(data, dims, refdims, processed_coords, name, metadata) +end + +function _process_coords(specs, parent_dims) + map(specs) do spec + _process_coord_spec(spec, parent_dims) + end +end + +_process_coord_spec(spec, parent_dims) = spec + +function _process_coord_spec(spec::Tuple{Tuple,AbstractArray}, parent_dims) + coord_dims, coord_data = spec + DimArray(coord_data, dims(parent_dims, coord_dims)) +end + + +# Interface methods +coords(_) = nothing +coords(x::CoordArray) = getfield(x, :coords) + +# Override rebuildsliced to handle coordinate slicing +function rebuildsliced(f::Function, A::CoordArray, data::AbstractArray, I::Tuple) + # Get the sliced dimensions and refdims using DimensionalData's logic + I1 = to_indices(A, I) + newdims, newrefdims = slicedims(f, A, I1) + coords = slicecoords(f, A, I1) + rebuild(A; data, dims=newdims, refdims=newrefdims, coords) +end + +function slicecoords(f::Function, A::CoordArray, indices::Tuple) + Adims = dims(A) + return map(coords(A)) do coord + # For DimArray coordinates, slice them according to the indexing pattern + # For other coordinates, preserve as-is + isa(coord, AbstractDimArray) ? _slice_coordinate(f, coord, indices, Adims) : coord + end +end + +function _slice_coordinate(f::Function, coord::AbstractDimArray, indices, parent_dims) + # Map the indices to the coordinate's dimensions + coord_indices = map(dims(coord)) do dim + indices[dimnum(parent_dims, dim)] + end + return f(coord, coord_indices...) +end + +function checkcoords(coords, dims) +end + +include("show.jl") \ No newline at end of file diff --git a/src/coordarray/show.jl b/src/coordarray/show.jl new file mode 100644 index 000000000..073a550c0 --- /dev/null +++ b/src/coordarray/show.jl @@ -0,0 +1,58 @@ +# Extend Base.show to display coordinate information + +function Base.show(io::IO, mime::MIME"text/plain", A::CoordArray) + lines, blockwidth, istop = show_main(io, mime, A::AbstractBasicDimArray) + # Add coordinate block + coord_lines, blockwidth, istop = print_coord_block(io, mime, coords(A); + blockwidth, displaywidth=displaysize(io)[2], separatorwidth=blockwidth, istop + ) + lines += coord_lines + + # Printing the array data is optional, subtypes can + # show other things here instead. + ds = displaysize(io) + ctx = IOContext(io, + :blockwidth => blockwidth, + :displaysize => (ds[1] - lines, ds[2]), + :isblocktop => istop + ) + show_after(ctx, mime, A) + return nothing +end + +function print_coord_block(io, mime, coordinates; + blockwidth=0, displaywidth, separatorwidth=blockwidth, istop=false, label="coords" +) + lines = 0 + if isempty(coordinates) + new_blockwidth = blockwidth + stilltop = istop + else + # Calculate the maximum width needed for coordinate display + coord_lines = [] + for (name, coord) in pairs(coordinates) + coord_str = if isa(coord, AbstractDimArray) + dim_names = join([string(nameof(typeof(d))) for d in dims(coord)], ", ") + " $name: ($dim_names) $(summary(coord))" + else + " $name: $(summary(coord))" + end + push!(coord_lines, coord_str) + end + + coord_width = maximum(textwidth, coord_lines) + new_blockwidth = max(blockwidth, min(displaywidth - 2, coord_width)) + new_blockwidth = print_block_separator(io, label, separatorwidth, new_blockwidth; istop) + println(io) + + # Print each coordinate + for coord_line in coord_lines + println(io, coord_line) + lines += 1 + end + + lines += 2 # for separator and final newline + stilltop = false + end + return lines, new_blockwidth, stilltop +end \ No newline at end of file diff --git a/test/coordarray.jl b/test/coordarray.jl new file mode 100644 index 000000000..86469c391 --- /dev/null +++ b/test/coordarray.jl @@ -0,0 +1,75 @@ +using DimensionalData, Test, Dates +using DimensionalData: coords +using DimensionalData: At, Between, Near, NoName, NoMetadata, format, checkdims + +@testset "CoordArray" begin + + @testset "CoordArray construction" begin + # Test the new (dimensions, data) syntax + data = rand(3, 4, 5) + lat = rand(3, 4) + lon = rand(3, 4) + elevation = rand(5) + + xdim = X(1:3) + ydim = Y(1:4) + zdim = Z(1:5) + coords1 = (; + latitude=((X, Y), lat), + longitude=((X, Y), lon), + elevation=((Z,), elevation) + ) + coords2 = (; + latitude=DimArray(lat, (xdim, ydim)), + longitude=DimArray(lon, (xdim, ydim)), + elevation=DimArray(elevation, zdim) + ) + + for coordspec in (coords1, coords2) + DimArray(data, (xdim, ydim, zdim)) + da = CoordArray(data, (xdim, ydim, zdim); coords=coordspec, name="temperature") + + @test coords(da).latitude isa DimArray + @test coords(da).longitude isa DimArray + @test coords(da).elevation isa DimArray + @test coords(da).latitude.dims == da.dims[1:2] + @test coords(da).longitude.dims == da.dims[1:2] + @test coords(da).elevation.dims == (zdim,) + end + + end + + @testset "CoordArray indexing" begin + temperature = 15 .+ 8 .* randn(2, 3, 4) + lon = [[42.25 42.21 42.63]; [42.63 42.59 42.59]]# 2x3 matrix + lat = [[-99.83 -99.32]; [-99.79 -99.23]; [-99.79 -99.23]] # 3x2 matrix + time_vals = DateTime(2014, 9, 6):Day(1):DateTime(2014, 9, 9) + reference_time = DateTime(2014, 9, 5) + + x_dim = X(1:2) + y_dim = Y(1:3) + time_dim = Ti(time_vals) + + lon = DimArray(lon, (x_dim, y_dim)) + lat = DimArray(lat, (y_dim, x_dim)) + + # Create CoordArray with non-dimension coordinates + da = CoordArray( + temperature, + (x_dim, y_dim, time_dim); + coords=(; lon, lat, reference_time), + ) + + # Test coordinate access + @test haskey(coords(da), :lon) + @test haskey(coords(da), :lat) + @test haskey(coords(da), :reference_time) + + # Test indexing with coordinates + @test da[Ti(At(time_vals[3])), X(2), Y(:)] == temperature[2, :, 3] + new_da = da[Ti(At(time_vals[2])), X(2), Y(:)] + @test new_da.coords.lon == lon[2, :] + @test new_da.coords.lat == lat[:, 2] + @test new_da.coords.reference_time == reference_time + end +end \ No newline at end of file