From 56116cc3ec07356a7e5ebc614ecc27a542dd3051 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 26 Sep 2025 13:49:14 +0200 Subject: [PATCH 1/6] Refactor regrid function for better handling of arrays Refactor regrid function to improve clarity and performance. --- src/XESMF.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/XESMF.jl b/src/XESMF.jl index 84f90d6..806980e 100644 --- a/src/XESMF.jl +++ b/src/XESMF.jl @@ -42,8 +42,9 @@ function sparse_regridder_weights(FT, regridder) end # Generic regridding function that does not check the dimensions of the `src` and -# `dst` arrays -function regrid!(dst::AbstractVector, regridder::Regridder, src::AbstractVector) +# `dst` arrays. For general vectors that might be discontinuous in memory, we need +# to broadcast the value to a `DenseArray` before performing the sparse matrix multiply +function (regridder::Regridder)(dst::AbstractVector, src::AbstractVector) regridder.src_temp .= src LinearAlgebra.mul!(regridder.dst_temp, regridder.weights, regridder.src_temp) dst .= regridder.dst_temp @@ -51,6 +52,9 @@ function regrid!(dst::AbstractVector, regridder::Regridder, src::AbstractVector) return dst end +# The easy case, for dense vectors, the regridding operation defaults to a matrix multiplication +(regridder::Regridder)(dst::DenseVector, src::DenseVector) = LinearAlgebra.mul!(dst, regridder.weights, src) + sparse_regridder_weights(regridder) = sparse_regridder_weights(Float64, regridder) """ From 3a916b7d49689021c4b03e5465508a663e90917c Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 26 Sep 2025 14:02:53 +0200 Subject: [PATCH 2/6] Update XESMF.jl --- src/XESMF.jl | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/XESMF.jl b/src/XESMF.jl index 806980e..f02b5ce 100644 --- a/src/XESMF.jl +++ b/src/XESMF.jl @@ -41,6 +41,9 @@ function sparse_regridder_weights(FT, regridder) return weights end +# The easy case, for dense vectors, the regridding operation defaults to a matrix multiplication +(regridder::Regridder)(dst::DenseVector, src::DenseVector) = LinearAlgebra.mul!(dst, regridder.weights, src) + # Generic regridding function that does not check the dimensions of the `src` and # `dst` arrays. For general vectors that might be discontinuous in memory, we need # to broadcast the value to a `DenseArray` before performing the sparse matrix multiply @@ -48,12 +51,20 @@ function (regridder::Regridder)(dst::AbstractVector, src::AbstractVector) regridder.src_temp .= src LinearAlgebra.mul!(regridder.dst_temp, regridder.weights, regridder.src_temp) dst .= regridder.dst_temp - return dst end -# The easy case, for dense vectors, the regridding operation defaults to a matrix multiplication -(regridder::Regridder)(dst::DenseVector, src::DenseVector) = LinearAlgebra.mul!(dst, regridder.weights, src) +# Mixed cases +function (regridder::Regridder)(dst::DenseVector, src::AbstractVector) + regridder.src_temp .= src + return LinearAlgebra.mul!(dst, regridder.weights, regridder.src_temp) +end + +function (regridder::Regridder)(dst::AbstractVector, src::DenseVector) + LinearAlgebra.mul!(regridder.dst_temp, regridder.weights, src) + dst .= regridder.dst_temp + return dst +end sparse_regridder_weights(regridder) = sparse_regridder_weights(Float64, regridder) From e8deb9dc74c7717203f6efbf98c9847c05c49582 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 26 Sep 2025 10:45:16 -0400 Subject: [PATCH 3/6] Bump version from 0.1.3 to 0.1.4 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f77b0bc..6ed35ab 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "XESMF" uuid = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" authors = ["NumericalEarth and contributors"] -version = "0.1.3" +version = "0.1.4" [deps] CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" From a9ca0c3e16b2a4af29e3a7975a4d6dbeb3926e88 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 26 Sep 2025 16:50:27 +0200 Subject: [PATCH 4/6] add a test --- test/test_oceananigans.jl | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/test_oceananigans.jl b/test/test_oceananigans.jl index 095def5..4bbb39b 100644 --- a/test/test_oceananigans.jl +++ b/test/test_oceananigans.jl @@ -96,5 +96,22 @@ end weights = XESMF.sparse_regridder_weights(regridder) @test weights isa SparseMatrixCSC + + # test that the regridder works with dense and strided arrays + dense_tg = zeros(prod(size(tg))) + dense_ll = zeros(prod(size(ll))) + + strided_tg = vec(view(zeros(size(tg, 1)+5, size(tg, 2)+5), 1:size(tg, 1), 1:size(tg, 2))) + strided_ll = vec(view(zeros(size(ll, 1)+5, size(ll, 2)+5), 1:size(ll, 1), 1:size(ll, 2))) + + rand_tg = rand(prod(size(tg))) + rand_ll = rand(prod(size(ll))) + + dense_tg .= rand_tg + strided_tg .= rand_tg + regridder(dense_ll, dense_tg) + regridder(strided_ll, strided_tg) + + @test all(dense_ll .== strided_ll) end end From 8905a144df959c52857c2b32f40cbeb028c548af Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sat, 27 Sep 2025 06:34:41 -0400 Subject: [PATCH 5/6] update --- src/XESMF.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/XESMF.jl b/src/XESMF.jl index f02b5ce..b25694c 100644 --- a/src/XESMF.jl +++ b/src/XESMF.jl @@ -12,7 +12,7 @@ struct Regridder{S, M, V1, V2} dst_temp :: V2 end -Base.summary(r::Regridder{S, M, V1, V2}) where {S, M, V1, V2} = "$(r.method) Regridder" +Base.summary(r::Regridder{S, M, V1, V2}) where {S, M, V1, V2} = "$(uppercasefirst(r.method)) Regridder" function Base.show(io::IO, r::Regridder) print(io, summary(r), '\n') @@ -54,7 +54,7 @@ function (regridder::Regridder)(dst::AbstractVector, src::AbstractVector) return dst end -# Mixed cases +# Mixed cases function (regridder::Regridder)(dst::DenseVector, src::AbstractVector) regridder.src_temp .= src return LinearAlgebra.mul!(dst, regridder.weights, regridder.src_temp) From 3e7ea2b4334339099b03c759b58f556a7870b1a7 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sat, 27 Sep 2025 06:55:05 -0400 Subject: [PATCH 6/6] fix tests --- test/runtests.jl | 2 +- test/setup_runtests.jl | 1 - test/test_oceananigans.jl | 44 +++++++++++++++++++++------------------ 3 files changed, 25 insertions(+), 22 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index d6fb54c..60c158f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,7 +4,7 @@ include("setup_runtests.jl") @testset "Unit Tests" begin include("test_unit.jl") end - + @testset "Oceananigans Integration Tests" begin include("test_oceananigans.jl") end diff --git a/test/setup_runtests.jl b/test/setup_runtests.jl index 9a42277..562e8de 100644 --- a/test/setup_runtests.jl +++ b/test/setup_runtests.jl @@ -1,4 +1,3 @@ using Test -using Oceananigans using PythonCall using XESMF \ No newline at end of file diff --git a/test/test_oceananigans.jl b/test/test_oceananigans.jl index 4bbb39b..9f6daaf 100644 --- a/test/test_oceananigans.jl +++ b/test/test_oceananigans.jl @@ -1,25 +1,31 @@ include("setup_runtests.jl") + +using Oceananigans +using Oceananigans.Fields: AbstractField using SparseArrays -x_node_array(x::AbstractVector, Nx, Ny) = view(x, 1:Nx) |> Array -y_node_array(x::AbstractVector, Nx, Ny) = view(x, 1:Ny) |> Array -x_node_array(x::AbstractMatrix, Nx, Ny) = view(x, 1:Nx, 1:Ny) |> Array +function x_node_array(x::AbstractVector, Nx, Ny) + return Array(repeat(view(x, 1:Nx), 1, Ny))' +end +function y_node_array(x::AbstractVector, Nx, Ny) + return Array(repeat(view(x, 1:Ny)', Nx, 1))' +end +x_node_array(x::AbstractMatrix, Nx, Ny) = Array(view(x, 1:Nx, 1:Ny))' -x_vertex_array(x::AbstractVector, Nx, Ny) = view(x, 1:Nx+1) |> Array -y_vertex_array(x::AbstractVector, Nx, Ny) = view(x, 1:Ny+1) |> Array -x_vertex_array(x::AbstractMatrix, Nx, Ny) = view(x, 1:Nx+1, 1:Ny+1) |> Array +function x_vertex_array(x::AbstractVector, Nx, Ny) + return Array(repeat(view(x, 1:Nx+1), 1, Ny+1))' +end +function y_vertex_array(x::AbstractVector, Nx, Ny) + return Array(repeat(view(x, 1:Ny+1)', Nx+1, 1))' +end +x_vertex_array(x::AbstractMatrix, Nx, Ny) = Array(view(x, 1:Nx+1, 1:Ny+1))' y_node_array(x::AbstractMatrix, Nx, Ny) = x_node_array(x, Nx, Ny) y_vertex_array(x::AbstractMatrix, Nx, Ny) = x_vertex_array(x, Nx, Ny) -function regridding_weights(dst_field, src_field) +function extract_xesmf_coordinates_structure(dst_field::AbstractField, src_field::AbstractField) ℓx, ℓy, ℓz = Oceananigans.Fields.instantiated_location(src_field) - # We only support regridding between centered fields. - @assert ℓx isa Center - @assert ℓy isa Center - @assert (ℓx, ℓy, ℓz) == Oceananigans.Fields.instantiated_location(dst_field) - dst_grid = dst_field.grid src_grid = src_field.grid @@ -35,7 +41,7 @@ function regridding_weights(dst_field, src_field) λvˢ = λnodes(src_grid, Face(), Face(), ℓz, with_halos=true) φvˢ = φnodes(src_grid, Face(), Face(), ℓz, with_halos=true) - # Build data structures expected by XESMF. + # Build data structures expected by xESMF Nˢx, Nˢy, Nˢz = size(src_field) Nᵈx, Nᵈy, Nᵈz = size(dst_field) @@ -59,7 +65,7 @@ function regridding_weights(dst_field, src_field) "lat_b" => φvˢ, "lon_b" => λvˢ) - return src_coordinates, dst_coordinates + return dst_coordinates, src_coordinates end @testset "Oceananigans Integration Tests" begin @@ -72,7 +78,7 @@ end cll = CenterField(ll) # Test that we can create the coordinate structures - src_coordinates, dst_coordinates = regridding_weights(ctg, cll) + dst_coordinates, src_coordinates = extract_xesmf_coordinates_structure(cll, ctg) # Verify coordinate structures are valid @test haskey(src_coordinates, "lat") @@ -89,13 +95,11 @@ end @test size(tg) == (360, 170, 1) @test size(ll) == (360, 180, 1) - xesmf = XESMF.xesmf - periodic = Oceananigans.Grids.topology(ctg.grid, 1) === Periodic ? PythonCall.pybuiltins.True : pybuiltins.False + periodic = Oceananigans.Grids.topology(ctg.grid, 1) === Periodic ? true : false method = "conservative" - regridder = xesmf.Regridder(src_coordinates, dst_coordinates, method; periodic) - weights = XESMF.sparse_regridder_weights(regridder) + regridder = XESMF.Regridder(src_coordinates, dst_coordinates; method, periodic) - @test weights isa SparseMatrixCSC + @test regridder.weights isa SparseMatrixCSC # test that the regridder works with dense and strided arrays dense_tg = zeros(prod(size(tg)))