Improve nan handling in surface plots (#4735)

* cleanup indices and uvs in surface shader * fix nan normals * extend test to cover nans in x/y/z/color individually * cleanup position_calc, normal_calc, light_calc * run heatmaps & surface test in CairoMakie * add changelog entry * add example for NaN handling, 2D surface and a few words on argument types * fix wrong example categorization --------- Co-authored-by: Simon <[email protected]>
MakieOrg · Feb 7, 2025 · 0f095bd · 0f095bd
1 parent 7bb01af
commit 0f095bd
Show file tree

Hide file tree

Showing 7 changed files with 203 additions and 121 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,12 +12,13 @@
 - Added option `lowres_background=true` to Resampler, and renamed `resolution` to `max_resolution` [#4745](https://github.com/MakieOrg/Makie.jl/pull/4745).
 - Added option `throttle=0.0` to `async_latest`, to allow throttling while skipping latest updates [#4745](https://github.com/MakieOrg/Makie.jl/pull/4745).
 - Fixed issue with `WGLMakie.voxels` not rendering on linux with firefox [#4756](https://github.com/MakieOrg/Makie.jl/pull/4756)
+- Cleaned up surface handling in GLMakie: Surface cells are now discarded when there is a nan in x, y or z. Fixed incorrect normal if x or y is nan [#4735](https://github.com/MakieOrg/Makie.jl/pull/4735)
+- Cleaned up `volume` plots: Added `:indexedabsorption` and `:additive` to WGLMakie, generalized `:mip` to include negative values, fixed missing conversions for rgba algorithms (`:additive`, `:absorptionrgba`), fixed missing conversion for `absorption` attribute & extended it to `:indexedabsorption` and `absorptionrgba`, added tests and improved docs. [#4726](https://github.com/MakieOrg/Makie.jl/pull/4726)
 
 ## [0.22.1] - 2025-01-17
 
 - Allow volume textures for mesh color, to e.g. implement a performant volume slice display [#2274](https://github.com/MakieOrg/Makie.jl/pull/2274).
 - Fixed `alpha` use in legends and some CairoMakie cases [#4721](https://github.com/MakieOrg/Makie.jl/pull/4721).
-- Cleaned up `volume` plots: Added `:indexedabsorption` and `:additive` to WGLMakie, generalized `:mip` to include negative values, fixed missing conversions for rgba algorithms (`:additive`, `:absorptionrgba`), fixed missing conversion for `absorption` attribute & extended it to `:indexedabsorption` and `absorptionrgba`, added tests and improved docs. [#4726](https://github.com/MakieOrg/Makie.jl/pull/4726)
 
 ## [0.22.0] - 2024-12-12
 

diff --git a/CairoMakie/test/runtests.jl b/CairoMakie/test/runtests.jl
@@ -187,7 +187,6 @@ excludes = Set([
     "fast pixel marker",
     "scatter with glow", # some are missing
     "scatter with stroke", # stroke acts inward in CairoMakie, outwards in W/GLMakie
-    "heatmaps & surface", # different nan_colors in surface
     "Textured meshscatter", # not yet implemented
     "Voxel - texture mapping", # not yet implemented
     "Miter Joints for line rendering", # CairoMakie does not show overlap here

diff --git a/GLMakie/assets/shader/surface.vert b/GLMakie/assets/shader/surface.vert
@@ -31,6 +31,13 @@ vec4 color_lookup(float intensity, sampler1D color, vec2 norm);
 
 uniform vec3 scale;
 uniform mat4 view, model, projection;
+uniform bool invert_normals;
+
+uniform uint objectid;
+flat out uvec2 o_id;
+flat out int o_InstanceID; // dummy for compat with meshscatter in mesh.frag
+out vec4 o_color;
+out vec3 o_uv;
 
 // See util.vert for implementations
 void render(vec4 position_world, vec3 normal, mat4 view, mat4 projection);
@@ -85,18 +92,18 @@ vec3 normal_from_points(
     // isnan checks should avoid darkening around NaN positions but may not
     // work with all systems
     if (!isnan(s0.z)) {
-        bool check1 = isinbounds(off1, size) && !isnan(s1.z);
-        bool check2 = isinbounds(off2, size) && !isnan(s2.z);
-        bool check3 = isinbounds(off3, size) && !isnan(s3.z);
-        bool check4 = isinbounds(off4, size) && !isnan(s4.z);
+        bool check1 = isinbounds(off1, size) && !isnan(s1.x) && !isnan(s1.y) && !isnan(s1.z);
+        bool check2 = isinbounds(off2, size) && !isnan(s2.x) && !isnan(s2.y) && !isnan(s2.z);
+        bool check3 = isinbounds(off3, size) && !isnan(s3.x) && !isnan(s3.y) && !isnan(s3.z);
+        bool check4 = isinbounds(off4, size) && !isnan(s4.x) && !isnan(s4.y) && !isnan(s4.z);
         if (check1 && check2) result += cross(s2-s0, s1-s0);
         if (check2 && check3) result += cross(s3-s0, s2-s0);
         if (check3 && check4) result += cross(s4-s0, s3-s0);
         if (check4 && check1) result += cross(s1-s0, s4-s0);
     }
     // normal should be zero, but needs to be here, because the dead-code
-    // elimanation of GLSL is overly enthusiastic
-    return normalize(result);
+    // elimination of GLSL is overly enthusiastic
+    return (invert_normals ? -1.0 : 1.0) * normalize(result);
 }
 
 // Overload for surface(Matrix, Matrix, Matrix)
@@ -157,31 +164,77 @@ vec3 getnormal(Nothing pos, sampler1D xs, sampler1D ys, sampler2D zs, ivec2 uv){
     return normal_from_points(s0, s1, s2, s3, s4, off1, off2, off3, off4, textureSize(zs, 0));
 }
 
-uniform uint objectid;
-flat out uvec2 o_id;
-flat out int o_InstanceID; // dummy for compat with meshscatter in mesh.frag
-out vec4 o_color;
-out vec3 o_uv;
+// See main() for more information
+vec3 getposition(Nothing grid, sampler2D x, sampler2D y, sampler2D z, ivec2 idx, vec2 uv) {
+    vec3 center = vec3(texture(x, uv).x, texture(y, uv).x, texture(z, uv).x);
+    if (isnan(center.x) || isnan(center.y) || isnan(center.z)) {
+        return vec3(0);
+    } else {
+        return vec3(texelFetch(x, idx, 0).x, texelFetch(y, idx, 0).x, texelFetch(z, idx, 0).x);
+    }
+}
+vec3 getposition(Nothing grid, sampler1D x, sampler1D y, sampler2D z, ivec2 idx, vec2 uv) {
+    vec3 center = vec3(texture(x, uv.x).x, texture(y, uv.y).x, texture(z, uv).x);
+    if (isnan(center.x) || isnan(center.y) || isnan(center.z)) {
+        return vec3(0);
+    } else {
+        return vec3(texelFetch(x, idx.x, 0).x, texelFetch(y, idx.y, 0).x, texelFetch(z, idx, 0).x);
+    }
+}
+vec3 getposition(Grid2D grid, Nothing x, Nothing y, sampler2D z, ivec2 idx, vec2 uv) {
+    float center = texture(z, uv).x;
+    if (isnan(center)) {
+        return vec3(0);
+    } else {
+        ivec2 size = textureSize(z, 0);
+        return vec3(grid_pos(grid, idx, size), texelFetch(z, idx, 0).x);
+    }
+}
+
+
 
 void main()
 {
+    // We have (N, M) = dims positions with (N-1, M-1) rects between them. Each
+    // instance refers to one rect. Each rect has vertices (0..1, 0..1) so we
+    // can do `base_idx + vertex` to index positions if base_idx refers to a
+    // (N-1, M-1) matrix.
     int index = gl_InstanceID;
-    vec2 offset = vertices;
-    ivec2 offseti = ivec2(offset);
     ivec2 dims = textureSize(position_z, 0);
-    vec3 pos;
-    {{position_calc}}
+    ivec2 base_idx = ind2sub(dims - 1, index);
+    ivec2 vertex_index = base_idx + ivec2(vertices);
+
+    /*
+    When using uv coordinates to access textures here, we need to be careful with
+    how we calculate the uvs. 0, 1 refer to the far edges of the texture:
+
+    0    1/N   2/N   3/N  ... N/N
+    |_____|_____|_____|_      _|
+    |     |     |     |        |
+    |_____|_____|_____|_      _|
+    |     |     |     |        |
+
+    Our textures contain one pixel per vertex though, so that the pixel centers
+    correspond to vertices. I.e. we want to access (0.5/N, 0.5/M) for the first
+    vertex, corresponding to index (0, 0).
+    */
+
+    // Discard rects containing a nan value by making their size zero. For this
+    // we get the value at the center of the rect, which mixes all 4 vertex
+    // values via texture interpolation. (If nan is part of the interpolation
+    // the result will also be nan.)
+    vec2 center_uv = (base_idx + vec2(1)) / dims;
+    vec3 pos = getposition(position, position_x, position_y, position_z, vertex_index, center_uv);
+    vec3 normalvec = getnormal(position, position_x, position_y, position_z, vertex_index);
+
+    // Colors should correspond to vertices, so they need the 0.5 shift discussed
+    // above
+    vec2 vertex_uv = vec2(vertex_index + 0.5) / vec2(dims);
+    o_uv = apply_uv_transform(uv_transform, vec2(vertex_uv.x, 1 - vertex_uv.y));
 
+    o_color = vec4(0.0);
     o_id = uvec2(objectid, 0); // calculated from uv in mesh.frag
-    o_InstanceID = 0;
-    // match up with mesh
-    o_uv = apply_uv_transform(uv_transform, vec2(index01.x, 1 - index01.y));
-    vec3 normalvec = {{normal_calc}};
+    o_InstanceID = 0; // for meshscatter uv_transforms, irrelevant here
 
-    o_color = vec4(0.0);
-    // we still want to render NaN values... TODO: make value customizable?
-    if (isnan(pos.z)) {
-        pos.z = 0.0;
-    }
     render(model * vec4(pos, 1), normalvec, view, projection);
 }
diff --git a/GLMakie/src/glshaders/particles.jl b/GLMakie/src/glshaders/particles.jl
@@ -57,6 +57,20 @@ function intensity_convert_tex(context, intensity::VecOrSignal{T}, verts) where
     end
 end
 
+function position_calc(
+        position_xyz::VectorTypes{T}, target::Type{TextureBuffer}
+    ) where T <: StaticVector
+    return "pos = texelFetch(position, index).xyz;"
+end
+
+function position_calc(
+        position_xyz::VectorTypes{T}, target::Type{GLBuffer}
+    ) where T <: StaticVector
+    len = length(T)
+    filler = join(ntuple(x->0, 3-len), ", ")
+    needs_comma = len != 3 ? ", " : ""
+    return "pos = vec3(position $needs_comma $filler);"
+end
 
 
 @nospecialize
@@ -116,7 +130,7 @@ function draw_mesh_particle(screen, p, data)
             "util.vert", "particles.vert",
             "fragment_output.frag", "lighting.frag", "mesh.frag",
             view = Dict(
-                "position_calc" => position_calc(position, nothing, nothing, nothing, TextureBuffer),
+                "position_calc" => position_calc(position, TextureBuffer),
                 "shading" => light_calc(shading),
                 "MAX_LIGHTS" => "#define MAX_LIGHTS $(screen.config.max_lights)",
                 "MAX_LIGHT_PARAMETERS" => "#define MAX_LIGHT_PARAMETERS $(screen.config.max_light_parameters)",
@@ -270,7 +284,7 @@ function draw_scatter(screen, (marker, position), data)
             "fragment_output.frag", "util.vert", "sprites.geom",
             "sprites.vert", "distance_shape.frag",
             view = Dict(
-                "position_calc" => position_calc(position, nothing, nothing, nothing, GLBuffer),
+                "position_calc" => position_calc(position, GLBuffer),
                 "buffers" => output_buffers(screen, to_value(transparency)),
                 "buffer_writes" => output_buffer_writes(screen, to_value(transparency))
             )

diff --git a/GLMakie/src/glshaders/surface.jl b/GLMakie/src/glshaders/surface.jl
@@ -1,23 +1,3 @@
-
-function position_calc(x...)
-    _position_calc(Iterators.filter(x->!isa(x, Nothing), x)...)
-end
-
-function normal_calc(x::Bool, invert_normals::Bool = false)
-    i = invert_normals ? "-" : ""
-    if x
-        return "$(i)getnormal(position, position_x, position_y, position_z, index2D);"
-    else
-        return "vec3(0, 0, $(i)1);"
-    end
-end
-
-# TODO this shouldn't be necessary
-function light_calc(x::Bool)
-    @error "shading::Bool is deprecated. Use `NoShading` instead of `false` and `FastShading` or `MultiLightShading` instead of true."
-    return light_calc(ifelse(x, FastShading, NoShading))
-end
-
 function light_calc(x::Makie.MakieCore.ShadingAlgorithm)
     if x === NoShading
         return "#define NO_SHADING"
@@ -32,67 +12,6 @@ function light_calc(x::Makie.MakieCore.ShadingAlgorithm)
     end
 end
 
-function _position_calc(
-        position_x::MatTypes{T}, position_y::MatTypes{T}, position_z::MatTypes{T}, target::Type{Texture}
-    ) where T<:AbstractFloat
-    """
-    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
-    ivec2 index2D = ind2sub(dims, index1D);
-    vec2 index01 = (vec2(index2D) + 0.5) / (vec2(dims));
-
-    pos = vec3(
-        texelFetch(position_x, index2D, 0).x,
-        texelFetch(position_y, index2D, 0).x,
-        texelFetch(position_z, index2D, 0).x
-    );
-    """
-end
-
-function _position_calc(
-        position_x::VectorTypes{T}, position_y::VectorTypes{T}, position_z::MatTypes{T},
-        target::Type{Texture}
-    ) where T<:AbstractFloat
-    """
-    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
-    ivec2 index2D = ind2sub(dims, index1D);
-    vec2 index01 = (vec2(index2D) + 0.5) / (vec2(dims));
-
-    pos = vec3(
-        texelFetch(position_x, index2D.x, 0).x,
-        texelFetch(position_y, index2D.y, 0).x,
-        texelFetch(position_z, index2D, 0).x
-    );
-    """
-end
-
-function _position_calc(
-        position_xyz::VectorTypes{T}, target::Type{TextureBuffer}
-    ) where T <: StaticVector
-    "pos = texelFetch(position, index).xyz;"
-end
-
-function _position_calc(
-        position_xyz::VectorTypes{T}, target::Type{GLBuffer}
-    ) where T <: StaticVector
-    len = length(T)
-    filler = join(ntuple(x->0, 3-len), ", ")
-    needs_comma = len != 3 ? ", " : ""
-    "pos = vec3(position $needs_comma $filler);"
-end
-
-function _position_calc(
-        grid::Grid{2}, position_z::MatTypes{T}, target::Type{Texture}
-    ) where T<:AbstractFloat
-    """
-    int index1D = index + offseti.x + offseti.y * dims.x; // + (index/(dims.x-1));
-    ivec2 index2D = ind2sub(dims, index1D);
-    vec2 index01 = (vec2(index2D) + 0.5) / (vec2(dims));
-
-    float height = texelFetch(position_z, index2D, 0).x;
-    pos = vec3(grid_pos(position, index01), height);
-    """
-end
-
 @nospecialize
 # surface(::Matrix, ::Matrix, ::Matrix)
 function draw_surface(screen, main::Tuple{MatTypes{T}, MatTypes{T}, MatTypes{T}}, data::Dict) where T <: AbstractFloat
@@ -150,8 +69,6 @@ function draw_surface(screen, main, data::Dict)
             "util.vert", "surface.vert",
             "fragment_output.frag", "lighting.frag", "mesh.frag",
             view = Dict(
-                "position_calc" => position_calc(position, position_x, position_y, position_z, Texture),
-                "normal_calc" => normal_calc(normal, to_value(invert_normals)),
                 "shading" => light_calc(shading),
                 "picking_mode" => "#define PICKING_INDEX_FROM_UV",
                 "MAX_LIGHTS" => "#define MAX_LIGHTS $(screen.config.max_lights)",

diff --git a/ReferenceTests/src/tests/primitives.jl b/ReferenceTests/src/tests/primitives.jl
@@ -636,16 +636,57 @@ end
 end
 
 @reference_test "Surface with NaN points" begin
-    # prepare surface data
-    zs = [x^2 + y^2 for x in range(-2, 0, length=10), y in range(-2, 0, length=10)]
-    ns = copy(zs)
-    ns[4, 3:6] .= NaN
-    # plot surface
-    f, a, p = surface(1..10, 1..10, ns, colormap = [:lightblue, :lightblue])
-    # plot a wireframe so we can see what's going on, and in which cells.
-    m = Makie.surface2mesh(to_value.(p.converted)...)
-    scatter!(a, m.position, color = isnan.(m.normal), depth_shift = -1f-3)
-    wireframe!(a, m, depth_shift = -1f-3, color = :black)
+    # This is supposed to verify a couple of things:
+    # - cells with nan in positions are not drawn
+    # - colors align to cell centers (via color checkerboard)
+    # - all normals are valid and interpolate correctly (lighting)
+    data = [x^2 + y^2 for x in range(-2, 0, length=11), y in range(-2, 0, length=11)]
+    cs = reshape([(:red, :blue)[mod1(i, 2)] for i in eachindex(data)], size(data))
+
+    f = Figure(size = (500, 1000), backgroundcolor = RGBf(0.3, 0.3, 0.3), figure_padding = (0,0,0,0))
+
+    # Test NaN in positions
+    for i in 1:3, j in 1:2
+        if j == 1
+            xs = collect(1.0:11.0)
+            ys = collect(1.0:11.0)
+        else
+            xs = Float32[x for x in 1:11, y in 1:11]
+            ys = Float32[y for x in 1:11, y in 1:11]
+        end
+        zs = copy(data)
+
+        # shift to second row if matrix
+        if i == 1
+            xs[(3:6) .+ (j-1) * 44] .= NaN
+        elseif i == 2
+            ys[(3:6) .+ (j-1) * 44] .= NaN
+        elseif i == 3
+            zs[4, 3:6] .= NaN
+        end
+
+        a, p = surface(f[i, j], xs, ys, zs, color = cs, nan_color = :red, axis = (show_axis = false,))
+        a.scene.lights = [AmbientLight(RGBf(0, 0, 0)), DirectionalLight(RGBf(2,2,2), Vec3f(0.5, -1, -0.8))]
+        # plot a wireframe so we can see what's going on, and in which cells.
+        m = Makie.surface2mesh(to_value.(p.converted)...)
+        wireframe!(a, m, depth_shift = -1f-3, color = RGBf(0,0.9,0), linewidth = 1)
+    end
+
+    # Test NaN in color
+    cs = copy(data)
+    cs[4, 3:6] .= NaN
+    for i in 1:2
+        nan_color = ifelse(i == 1, :transparent, :red)
+        a, p = surface(f[4, i], 1..11, 1..11, data, color = cs, colormap = [:white, :white];
+            nan_color, axis = (show_axis = false,))
+        a.scene.lights = [AmbientLight(RGBf(0, 0, 0)), DirectionalLight(RGBf(2,2,2), Vec3f(0.5, -1, -0.8))]
+        m = Makie.surface2mesh(to_value.(p.converted)...)
+        wireframe!(a, m, depth_shift = -1f-3, color = RGBf(0,0.9,0), linewidth = 1)
+    end
+
+    colgap!(f.layout, 0.0)
+    rowgap!(f.layout, 0.0)
+
     f
 end