From 807af04429c33f42b128671a5d1613316e93285b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 3 Aug 2022 20:11:45 +0100
Subject: [PATCH 1/2] Allow 2d coordinates for heatmap plots

Allows plotting on a non-square grid for backends that support it.
---
 src/axes.jl  |  6 +++++-
 src/utils.jl | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/axes.jl b/src/axes.jl
index 38cd72579..490cf342b 100644
--- a/src/axes.jl
+++ b/src/axes.jl
@@ -423,7 +423,7 @@ function expand_extrema!(axis::Axis, v::Tuple{MIN,MAX}) where {MIN<:Number,MAX<:
     ex.emax = isfinite(v[2]) ? max(v[2], ex.emax) : ex.emax
     ex
 end
-function expand_extrema!(axis::Axis, v::AVec{N}) where {N<:Number}
+function expand_extrema!(axis::Axis, v::Union{AVec{N},AMat{N}}) where {N<:Number}
     ex = axis[:extrema]
     for vi in v
         expand_extrema!(ex, vi)
@@ -515,6 +515,10 @@ function expand_extrema!(sp::Subplot, plotattributes::AKW)
     if plotattributes[:seriestype] === :heatmap
         for letter in (:x, :y)
             data = plotattributes[letter]
+            if typeof(data) <: Surface
+                # Was given 2d coordinates, need to extract from Surface struct
+                data = data.surf
+            end
             axis = sp[get_attr_symbol(letter, :axis)]
             scale = get(plotattributes, get_attr_symbol(letter, :scale), :identity)
             expand_extrema!(axis, heatmap_edges(data, scale))
diff --git a/src/utils.jl b/src/utils.jl
index 0fa166614..9feb056eb 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -273,10 +273,14 @@ function _heatmap_edges(v::AVec, isedges::Bool = false, ispolar::Bool = false)
     extra_max = (v[end] - v[end - 1]) / 2
     vcat(vmin - extra_min, 0.5 * (v[1:(end - 1)] + v[2:end]), vmax + extra_max)
 end
+function _heatmap_edges(v::AMat, isedges::Bool = false, ispolar::Bool = false)
+    # For now, don't handle edges here and just defer to the backend
+    v
+end
 
 "create an (n+1) list of the outsides of heatmap rectangles"
 function heatmap_edges(
-    v::AVec,
+    v::Union{AVec,AMat},
     scale::Symbol = :identity,
     isedges::Bool = false,
     ispolar::Bool = false,
@@ -307,6 +311,32 @@ function heatmap_edges(
     return x, y
 end
 
+function heatmap_edges(
+    x::AMat,
+    xscale::Symbol,
+    y::AMat,
+    yscale::Symbol,
+    z_size::Tuple{Int,Int},
+    ispolar::Bool = false,
+)
+    x_size = size(x)
+    y_size = size(y)
+    if x_size != y_size
+        error("""When x and y are 2d, they must have the same size""")
+    end
+    ismidpoints = z_size == x_size
+    isedges = (z_size[1] + 1, z_size[2] + 1) == x_size
+    if !ismidpoints && !isedges
+        println("z_size=$z_size, x_size=$x_size, y_size=$y_size")
+        println("x $(typeof(x)), y $(typeof(y))")
+        error("""Size of x & y does not match the size of z.
+                Must be either `size(z) == size(x) == size(y))` (x & y define midpoints)
+                or `(size(z)[1]+1,size(z)[2]+1) == size(x) == size(y)` (x & y define edges).""")
+    end
+    x, y = heatmap_edges(x, xscale, isedges), heatmap_edges(y, yscale, isedges, ispolar) # special handle for `r` in polar plots
+    return x, y
+end
+
 function is_uniformly_spaced(v; tol = 1e-6)
     dv = diff(v)
     maximum(dv) - minimum(dv) < tol * mean(abs.(dv))

From 3133ee88b797c2eac6d8fddbdcec48b6362ec1ec Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 4 Aug 2022 09:58:38 +0100
Subject: [PATCH 2/2] Fix formatting

---
 src/utils.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index 9feb056eb..bf0675fb1 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -329,9 +329,11 @@ function heatmap_edges(
     if !ismidpoints && !isedges
         println("z_size=$z_size, x_size=$x_size, y_size=$y_size")
         println("x $(typeof(x)), y $(typeof(y))")
-        error("""Size of x & y does not match the size of z.
-                Must be either `size(z) == size(x) == size(y))` (x & y define midpoints)
-                or `(size(z)[1]+1,size(z)[2]+1) == size(x) == size(y)` (x & y define edges).""")
+        error(
+            """Size of x & y does not match the size of z.
+              Must be either `size(z) == size(x) == size(y))` (x & y define midpoints)
+              or `(size(z)[1]+1,size(z)[2]+1) == size(x) == size(y)` (x & y define edges).""",
+        )
     end
     x, y = heatmap_edges(x, xscale, isedges), heatmap_edges(y, yscale, isedges, ispolar) # special handle for `r` in polar plots
     return x, y