From 59456d3f0f37094e95267c5a27405f1637afedbf Mon Sep 17 00:00:00 2001
From: Simon <sdanisch@protonmail.com>
Date: Wed, 2 Oct 2024 15:02:59 +0200
Subject: [PATCH] fix heatmapshader for ranges (#4434)

* fix heatmapshader for ranges

* fix more edge cases

* add tests for range arguments
---
 ReferenceTests/src/tests/examples2d.jl |  4 +++
 src/basic_recipes/datashader.jl        | 40 +++++++++++++++++++-------
 test/conversions.jl                    |  7 +++++
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/ReferenceTests/src/tests/examples2d.jl b/ReferenceTests/src/tests/examples2d.jl
index 32837693d1f..e69c726558b 100644
--- a/ReferenceTests/src/tests/examples2d.jl
+++ b/ReferenceTests/src/tests/examples2d.jl
@@ -1543,6 +1543,10 @@ end
     ax3, pl3 = heatmap(f[2, 1], Resampler(data2))
     ax4, pl4 = heatmap(f[2, 2], Resampler(data2))
     limits!(ax4, 3000, 3090, 3460, 3500)
+    heatmap(f[3, 1], (1000, 2000), (500, 1000), Resampler(data2))
+    ax = Axis(f[3, 2])
+    limits!(ax, (0, 1), (0, 1))
+    heatmap!(ax, (1, 2), (1, 2), Resampler(data2))
     Colorbar(f[:, 3], pl1)
     sleep(1) # give the async operations some time
     f
diff --git a/src/basic_recipes/datashader.jl b/src/basic_recipes/datashader.jl
index 0c362ca43c7..860cb238de3 100644
--- a/src/basic_recipes/datashader.jl
+++ b/src/basic_recipes/datashader.jl
@@ -564,10 +564,21 @@ function conversion_trait(::Type{<:Heatmap}, ::Resampler)
     return HeatmapShaderConversion()
 end
 
-function Makie.MakieCore.types_for_plot_arguments(::Type{<:Heatmap}, ::HeatmapShaderConversion)
+function MakieCore.types_for_plot_arguments(::Type{<:Heatmap}, ::HeatmapShaderConversion)
     return Tuple{EndPoints{Float32},EndPoints{Float32},<:Resampler}
 end
 
+function data_limits(p::HeatmapShader)
+    x, y = p[1][], p[2][]
+    mini = Vec3f(x[1], y[1], 0)
+    widths = Vec3f(x[2] - x[1], y[2] - y[1], 0)
+    return Rect3f(mini, widths)
+end
+
+function boundingbox(p::HeatmapShader, space::Symbol=:data)
+    return apply_transform_and_model(p, data_limits(p))
+end
+
 function calculated_attributes!(::Type{Heatmap}, plot::HeatmapShader)
     return
 end
@@ -597,8 +608,9 @@ function resample_image(x, y, image, max_resolution, limits)
         nrange = ranges[i]
         si = imgsize[i]
         indices = ((nrange .- data_min[i]) ./ data_width[i]) .* (si .- 1) .+ 1
-        resolution = round(Int, indices[2] - indices[1])
-        return LinRange(max(1, indices[1]), min(indices[2], si), min(resolution, max_resolution[i]))
+        resolution = max(2, round(Int, indices[2] - indices[1]))
+        len = min(resolution, max_resolution[i])
+        return LinRange(max(1, indices[1]), min(indices[2], si), len)
     end
     if isempty(x_index_range) || isempty(y_index_range)
         return nothing
@@ -615,7 +627,7 @@ end
 
 function convert_arguments(::Type{Heatmap}, x, y, image::Resampler)
     x, y, img = convert_arguments(Heatmap, x, y, image.data)
-    return (x, y, Resampler(img))
+    return (EndPoints{Float32}(x...), EndPoints{Float32}(y...), Resampler(img))
 end
 
 function empty_channel!(channel::Channel)
@@ -650,7 +662,9 @@ function Makie.plot!(p::HeatmapShader)
 
     x, y = p.x, p.y
     max_resolution = lift(p, p.values, scene.viewport) do resampler, viewport
-        return resampler.max_resolution isa Automatic ? widths(viewport) : ntuple(x-> resampler.max_resolution, 2)
+        res = resampler.max_resolution isa Automatic ? widths(viewport) :
+              ntuple(x -> resampler.max_resolution, 2)
+        return max.(res, 512) # Not sure why, but viewport can become (1, 1)
     end
     image = lift(x-> x.data, p, p.values)
     image_area = lift(xy_to_rect, x, y; ignore_equal_values=true)
@@ -676,16 +690,22 @@ function Makie.plot!(p::HeatmapShader)
     translate!(lp, 0, 0, -1)
 
     first_downsample = resample_image(x[], y[], image[], max_resolution[], limits[])
-
-    # Make sure we don't trigger an event if the image/xy stays the same
-    args = map(arg -> lift(identity, p, arg; ignore_equal_values=true), (x, y, image, max_resolution))
-
-    CT = Tuple{typeof.(to_value.(args))..., typeof(limits_slow[])}
     # We hide the image when outside of the limits, but we also want to correctly forward, p.visible
     visible = Observable(p.visible[]; ignore_equal_values=true)
     on(p, p.visible) do v
         visible[] = v
+        return
     end
+    # if nothing, the image is currently not visible in the limits chosen
+    if isnothing(first_downsample)
+        visible[] = false
+        first_downsample = EndPoints{Float32}(0, 1), EndPoints{Float32}(0, 1), zeros(Float32, 2, 2)
+    end
+    # Make sure we don't trigger an event if the image/xy stays the same
+    args = map(arg -> lift(identity, p, arg; ignore_equal_values=true), (x, y, image, max_resolution))
+
+    CT = Tuple{typeof.(to_value.(args))..., typeof(limits_slow[])}
+
     # To actually run the resampling on another thread, we need another channel:
     # To make this threadsafe, the observable needs to be triggered from the current thread
     # So we need another channel for the result (unbuffered, so `put!` blocks while the image is being updated)
diff --git a/test/conversions.jl b/test/conversions.jl
index 118dde2ad66..fe725b1c9ae 100644
--- a/test/conversions.jl
+++ b/test/conversions.jl
@@ -16,6 +16,13 @@ using Makie.MakieCore: plotfunc, plotfunc!, func2type
 
 end
 
+@testset "Heatmapshader with ranges" begin
+    hm = Heatmap(((0, 1), (0, 1), Resampler(zeros(4, 4))), Dict{Symbol,Any}())
+    hm.converted[1][] isa Makie.EndPoints{Float32}
+    hm.converted[2][] isa Makie.EndPoints{Float32}
+    hm.converted[3][].data == Resampler(zeros(4, 4)).data
+end
+
 @testset "changing input types" begin
     input = Observable{Any}(decompose(Point2f, Circle(Point2f(0), 2f0)))
     f, ax, pl = mesh(input)