feat: Add CoordArray dimensional arrays with coordinate support

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")

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")

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

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