From 805d62408582b628bf7e15d250b513ce02248959 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 26 Mar 2026 14:56:51 +0100 Subject: [PATCH] Fix inpainting robustness: correct NaN propagation and fill value MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs in the nearest-neighbor inpainting algorithm: 1. `_propagate_field!` used `value == 0` instead of `donors == 0` to detect cells with no valid neighbors. When valid neighbors average to zero (e.g. temperature near 0°C), cells were incorrectly marked as NaN. 2. `_fill_nans!` converted remaining NaN cells to 0, which is catastrophic for fields like salinity (~34 psu). Now fills with the mean of valid interior data instead. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/DataWrangling/inpainting.jl | 13 +++++++++---- test/runtests_setup.jl | 3 ++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/DataWrangling/inpainting.jl b/src/DataWrangling/inpainting.jl index dde0f1f91..a081fa1d7 100644 --- a/src/DataWrangling/inpainting.jl +++ b/src/DataWrangling/inpainting.jl @@ -51,7 +51,12 @@ function propagate_horizontally!(inpainting::NearestNeighborInpainting, field, m iter += 1 end - launch!(arch, grid, size(field), _fill_nans!, field) + # Fill any remaining NaN values with the mean of valid data. + # Using 0 would be catastrophic for fields like salinity (~34 psu). + valid_sum = sum(x -> ifelse(isnan(x), zero(x), x), field; condition=interior(mask)) + valid_count = sum(x -> !isnan(x), field; condition=interior(mask)) + fill_value = convert(eltype(field), valid_sum / valid_count) + launch!(arch, grid, size(field), _fill_nans!, field, fill_value) fill_halo_regions!(field) return field @@ -80,7 +85,7 @@ end end FT_NaN = convert(FT, NaN) - @inbounds substituting_field[i, j, k] = ifelse(value == 0, FT_NaN, value / donors) + @inbounds substituting_field[i, j, k] = ifelse(donors == 0, FT_NaN, value / donors) end @kernel function _substitute_values!(field, substituting_field) @@ -97,9 +102,9 @@ end @inbounds field[i, j, k] = ifelse(mask[i, j, k], FT_NaN, field[i, j, k]) end -@kernel function _fill_nans!(field) +@kernel function _fill_nans!(field, fill_value) i, j, k = @index(Global, NTuple) - @inbounds field[i, j, k] *= !isnan(field[i, j, k]) + @inbounds field[i, j, k] = ifelse(isnan(field[i, j, k]), fill_value, field[i, j, k]) end """ diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index b3e0fae9d..77ff379a1 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -309,7 +309,8 @@ function test_inpainting_algorithm(arch, dataset, start_date, inpainting; partially_inpainted_interior = on_architecture(CPU(), interior(partially_inpainted_field)) @test all(fully_inpainted_interior .!= 0) - @test any(partially_inpainted_interior .== 0) + # Partially inpainted data should have no NaN (remaining NaN is filled with the field mean) + @test !any(isnan.(partially_inpainted_interior)) end return nothing end