From d5e69c7416d83fef133d091700cad29cf12ad27d Mon Sep 17 00:00:00 2001
From: Rafael Schouten <rafaelschouten@gmail.com>
Date: Wed, 27 Sep 2023 20:50:07 +0200
Subject: [PATCH] fix missing plots and colorbar

---
 ext/DimensionalDataMakie.jl | 22 ++++++++++++++--------
 test/plotrecipes.jl         | 31 +++++++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/ext/DimensionalDataMakie.jl b/ext/DimensionalDataMakie.jl
index dcb87f786..ebc27a608 100644
--- a/ext/DimensionalDataMakie.jl
+++ b/ext/DimensionalDataMakie.jl
@@ -143,9 +143,9 @@ for (f1, f2) in _paired(:plot => :heatmap, :heatmap, :image, :contour, :contourf
             end
             # Add a Colorbar for heatmaps and contourf
             if $(f1 in (:plot, :heatmap, :contourf)) 
-                Colorbar(p.figure[1, 2];
+                Colorbar(p.figure[1, 2], p.plot;
                     label=DD.label(A), colorbarkw...
-                ) 
+                )
             end
             return p
         end
@@ -344,8 +344,8 @@ Makie.plottype(A::AbstractDimArray{<:Any,3}) = Makie.Volume
 # Conversions
 function Makie.convert_arguments(t::Makie.PointBased, A::AbstractDimArray{<:Any,1})
     A = _prepare_for_makie(A)
-    xs = lookup(A, 1)
-    return Makie.convert_arguments(t, parent(xs), parent(A))
+    xs = parent(lookup(A, 1))
+    return Makie.convert_arguments(t, xs, _float32ornan(parent(A)))
 end
 function Makie.convert_arguments(t::Makie.PointBased, A::AbstractDimArray{<:Number,2})
     return Makie.convert_arguments(t, parent(A))
@@ -353,19 +353,19 @@ end
 function Makie.convert_arguments(t::Makie.SurfaceLike, A::AbstractDimArray{<:Any,2})
     A1 = _prepare_for_makie(A)
     xs, ys = map(parent, lookup(A1))
-    return Makie.convert_arguments(t, xs, ys, parent(A1))
+    return xs, ys, last(Makie.convert_arguments(t, parent(A1)))
 end
 function Makie.convert_arguments(
     t::Makie.DiscreteSurface, A::AbstractDimArray{<:Any,2}
 )
     A1 = _prepare_for_makie(A)
     xs, ys = map(parent, lookup(A1))
-    return xs, ys, parent(A1)
+    return xs, ys, last(Makie.convert_arguments(t, parent(A1)))
 end
 function Makie.convert_arguments(t::Makie.VolumeLike, A::AbstractDimArray{<:Any,3}) 
     A1 = _prepare_for_makie(A)
     xs, ys, zs = map(parent, lookup(A1))
-    return xs, ys, zs, parent(A1)
+    return xs, ys, zs, last(Makie.convert_arguments(t, parent(A1)))
 end
 # fallbacks with descriptive error messages
 function Makie.convert_arguments(t::Makie.ConversionTrait, A::AbstractDimArray{<:Any,N}) where {N}
@@ -428,7 +428,9 @@ function _split_attributes(dim::Dimension)
     return attributes, dims[1]
 end
 
-_prepare_for_makie(A, replacements=()) = _permute_xyz(A, replacements) |> _reorder
+function _prepare_for_makie(A, replacements=()) 
+    A1 = _permute_xyz(A, replacements) |> _reorder
+end
 
 # Permute the data after replacing the dimensions with X/Y/Z
 _permute_xyz(A::AbstractDimArray, replacements::Pair) = _permute_xyz(A, (replacements,))
@@ -486,4 +488,8 @@ function _keywords2dimpairs(x, y)
     end
 end
 
+_float32ornan(A::AbstractArray{Float32}) = A
+_float32ornan(A::AbstractArray) = _float32ornan.(A)
+_float32ornan(x) = ismissing(x) ? NaN32 : Float32(x)
+
 end
diff --git a/test/plotrecipes.jl b/test/plotrecipes.jl
index 6ec14edae..02868264f 100644
--- a/test/plotrecipes.jl
+++ b/test/plotrecipes.jl
@@ -158,31 +158,52 @@ nothing
 # da_im2 |> plot
 
 if !haskey(ENV, "CI")
-    using GLMakie
+
+using GLMakie
+@testset "Makie" begin
     # 1d
     A1 = rand(X('a':'e'); name=:test)
+    A1m = rand([missing, (1:3)...], X('a':'e'); name=:test)
+    A1m[3] = missing
     plot(A1)
+    plot(A1m)
     scatter(A1)
+    scatter(A1m)
     lines(A1)
+    lines(A1m)
     scatterlines(A1)
+    scatterlines(A1m)
     stairs(A1)
+    stairs(A1m)
     stem(A1)
+    stem(A1m)
     barplot(A1)
+    barplot(A1m)
     waterfall(A1)
+    waterfall(A1m)
     # 2d
     A2 = rand(X(10:10:100), Y(['a', 'b', 'c']))
     A2r = rand(Y(10:10:100), X(['a', 'b', 'c']))
+    A2m = rand([missing, (1:5)...], Y(10:10:100), X(['a', 'b', 'c']))
+    A2m[3] = missing
     plot(A2)
+    plot(A2m)
     # Categorical wins: it's on x, even though its Y
     boxplot(A2)
     boxplot(A2r)
+    @test_throws ArgumentError boxplot(A2m)
     violin(A2)
+    violin(A2r)
+    @test_throws ArgumentError violin(A2m)
     rainclouds(A2)
+    @test_throws ErrorException rainclouds(A2m)
     surface(A2)
+    surface(A2m)
     # Series also puts Categories in the legend no matter where they are
     series(A2)
     series(A2r)
     series(A2r; labeldim=Y)
+    series(A2m)
     @test_throws ArgumentError plot(A2; y=:c)
     # x/y can be specified
     A2ab = DimArray(rand(6, 10), (:a, :b); name=:stuff)
@@ -199,10 +220,16 @@ if !haskey(ENV, "CI")
     series(A2ab; labeldim=:b)
     # 3d
     A3 = rand(X(7), Z(10), Y(5))
+    A3m = rand([missing, (1:7)...], X(7), Z(10), Y(5))
+    A3m[3] = missing
     volume(A3)
+    volume(A3m)
     volumeslices(A3)
+    volumeslices(A3m)
     # x/y/z can be specified
     A3abc = DimArray(rand(10, 10, 7), (:a, :b, :c); name=:stuff)
     volume(A3abc; x=:c)
-    volumeslices(A3; z=:a)
+    volumeslices(A3abc; z=:a)
+end
+
 end