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