issue with ext loading

Author: juliasloan25

Project.toml

@@ -34,10 +34,12 @@ UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
+ConservativeRegridding = "b4b54eb9-2b73-55e3-98c1-00e3baba4ed5"
 
 [extensions]
 ClimaCoreCUDAExt = "CUDA"
 KrylovExt = "Krylov"
+ClimaCoreConservativeRegriddingExt = "ConservativeRegridding"
 
 [compat]
 Adapt = "3.2.0, 4"
@@ -50,6 +52,7 @@ BlockArrays = "1"
 CUDA = "5.5"
 ClimaComms = "0.6.2"
 ClimaCoreMakie = "0.4.6"
+ConservativeRegridding = "0.1"
 CountFlops = "0.1"
 CubedSphere = "0.2, 0.3"
 DataStructures = "0.18.13, 0.19"
@@ -90,6 +93,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 AssociatedLegendrePolynomials = "2119f1ac-fb78-50f5-8cc0-dda848ebdb19"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ConservativeRegridding = "b4b54eb9-2b73-55e3-98c1-00e3baba4ed5"
 CountFlops = "1db9610d-79e1-487a-8d40-77f3295c7593"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
@@ -105,4 +109,4 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "ArgParse", "AssociatedLegendrePolynomials", "BenchmarkTools", "CountFlops", "Dates", "FastBroadcast", "Krylov", "JET", "LazyBroadcast", "Logging", "PrettyTables", "Random", "SafeTestsets", "StatsBase", "TerminalLoggers", "Test"]
+test = ["Aqua", "ArgParse", "AssociatedLegendrePolynomials", "BenchmarkTools", "ConservativeRegridding", "CountFlops", "Dates", "FastBroadcast", "Krylov", "JET", "LazyBroadcast", "Logging", "PrettyTables", "Random", "SafeTestsets", "StatsBase", "TerminalLoggers", "Test"]

ext/ClimaCoreConservativeRegriddingExt.jl

@@ -1,5 +1,12 @@
 module ClimaCoreConservativeRegriddingExt
 
+using ConservativeRegridding, ClimaCore
+import ConservativeRegridding: Regridder, regrid!
+import ClimaCore: Spaces, Meshes, Quadratures, Fields, RecursiveApply, Spaces
+
+export get_element_vertices, integrate_each_element, get_value_per_element!,
+    set_value_per_element!, Regridder, regrid!
+
 """
     get_element_vertices(space::SpectralElementSpace2D)
 
@@ -32,7 +39,7 @@ function get_element_vertices(space)
     # Put each polygon into a vector, with the first coordinate pair repeated at the end
     vertices = collect(Iterators.partition(vertex_coords, 5))
 
-    # Check for zero area polygons
+    # Check for zero area polygons (all latitude or longitude values are the same)
     for polygon in vertices
         if allequal(first.(polygon)) || allequal(last.(polygon))
             @error "Zero area polygon found in vertices" polygon
@@ -82,9 +89,13 @@ to the number of elements in the space.
 Here `ones_field` is a field on the same space as `field` with all
 values set to 1.
 """
-function get_value_per_element!(value_per_element, field, ones_field)
-    integral_each_element = Remapping.integrate_each_element(field)
-    area_each_element = Remapping.integrate_each_element(ones_field)
+function get_value_per_element!(
+    value_per_element,
+    field,
+    ones_field,
+)
+    integral_each_element = integrate_each_element(field)
+    area_each_element = integrate_each_element(ones_field)
     value_per_element .= integral_each_element ./ area_each_element
     return nothing
 end
@@ -117,5 +128,101 @@ function set_value_per_element!(field, value_per_element)
     return field
 end
 
+"""
+    Regridder(dst_space, src_space; kwargs...)
+
+Create a regridder between two ClimaCore Spaces.
+This works by finding the vertices of the elements of the source and
+destination spaces, and then computing the areas of their intersections.
+
+This is currently only defined for 2D spaces, but could be extended to
+3D spaces.
+
+This function is intended to be used with the finite volume approximation
+of the ClimaCore spectral element space.
+"""
+function Regridder(
+    dst_space::Spaces.SpectralElementSpace2D,
+    src_space::Spaces.SpectralElementSpace2D;
+    kwargs...,
+)
+    dst_vertices = get_element_vertices(dst_space)
+    src_vertices = get_element_vertices(src_space)
+    return ConservativeRegridding.Regridder(dst_vertices, src_vertices; kwargs...)
+end
+Regridder(
+    dst_field::Fields.Field,
+    src_field::Fields.Field;
+    kwargs...,
+) = Regridder(axes(dst_field), axes(src_field); kwargs...)
+
+"""
+    regrid!(dst_field, regridder_tuple, src_field)
+
+Perform conservative regridding from `src_field` to `dst_field` using a
+Regridder object and pre-allocated buffers.
+
+The `regridder_tuple` should be a NamedTuple containing:
+- `regridder`: The ConservativeRegridding.Regridder object
+- `value_per_element_src`: Pre-allocated buffer for source values
+- `value_per_element_dst`: Pre-allocated buffer for destination values
+- `ones_src`: Pre-allocated field of ones on the source space
+```
+"""
+function regrid!(
+    dst_field::Fields.Field,
+    regridder_tuple::NamedTuple,
+    src_field::Fields.Field,
+)
+    @assert eltype(dst_field) isa Number && eltype(src_field) isa Number "Regridding is only supported for scalar fields"
+    @assert eltype(dst_field) == eltype(src_field) "Source and destination fields must have the same element type"
+
+    # Use pre-allocated buffers from the regridder tuple
+    (; value_per_element_src, value_per_element_dst, ones_src, regridder) = regridder_tuple
+
+    # Get one value per element in the source field, equal to the quadrature-weighted average of the
+    # values at nodes of the element
+    get_value_per_element!(value_per_element_src, src_field, ones_src)
+
+    # Perform the regridding
+    ConservativeRegridding.regrid!(value_per_element_dst, regridder, value_per_element_src)
+
+    # Now that we have our regridded vector, put it onto a field on the second space
+    set_value_per_element!(dst_field, value_per_element_dst)
+    return nothing
+end
+
+"""
+    regrid!(dst_field, regridder, src_field)
+
+Perform conservative regridding from `src_field` to `dst_field` using a
+Regridder object.
+
+This is a convenience function that allocates the buffers for you.
+Note that this is not efficient for repeated regridding with the same regridder,
+but it may be helpful for one-off regriddings or testing/debugging.
+"""
+function regrid!(
+    dst_field::Fields.Field,
+    regridder::ConservativeRegridding.Regridder,
+    src_field::Fields.Field,
+)
+    # Allocate space for the buffers
+    value_per_element_src =
+        zeros(Float64, Meshes.nelements(axes(src_field).grid.topology.mesh))
+    value_per_element_dst =
+        zeros(Float64, Meshes.nelements(axes(dst_field).grid.topology.mesh))
+    ones_src = ones(axes(src_field))
+    regridder_tuple = (;
+        regridder,
+        value_per_element_src,
+        value_per_element_dst,
+        ones_src,
+    )
+
+    # Perform the regridding
+    regrid!(dst_field, regridder_tuple, src_field)
+    return nothing
+end
 
 end

test/ClimaCoreConservativeRegriddingExt/test_conservative_regridding.jl

@@ -1,61 +1,128 @@
+using Test
+using ClimaCore
 using ClimaCore:
     CommonSpaces, Remapping, Fields, Spaces, RecursiveApply, Meshes, Quadratures
 using ConservativeRegridding
 
-space1 = CommonSpaces.CubedSphereSpace(;
+# Extensions load lazily in Julia - they need to be explicitly triggered
+# This ensures the extension is loaded and its methods are available
+const ClimaConservativeRegrid =
+    Base.get_extension(ClimaCore, :ClimaCoreConservativeRegriddingExt)
+@assert !isnothing(ClimaConservativeRegrid) "ClimaCoreConservativeRegriddingExt extension not loaded"
+
+const src_space = CommonSpaces.CubedSphereSpace(;
     radius = 10,
     n_quad_points = 3,
     h_elem = 8,
 )
-space2 = CommonSpaces.CubedSphereSpace(;
+const dst_space = CommonSpaces.CubedSphereSpace(;
     radius = 10,
     n_quad_points = 4,
     h_elem = 6,
 )
 
-vertices1 = Remapping.get_element_vertices(space1)
-vertices2 = Remapping.get_element_vertices(space2)
-
-# Pass in destination vertices first, source vertices second
-# TODO open issue in CR.jl about ordering of inputs source/dest
-regridder_1_to_2 = ConservativeRegridding.Regridder(vertices2, vertices1)
-regridder_2_to_1 = ConservativeRegridding.Regridder(vertices1, vertices2)
-
-# Define a field on the first space, to use as our source field
-field1 = Fields.coordinate_field(space1).lat
-ones_field1 = Fields.ones(space1)
-
-# Check that integrating over each element and summing gives the same result as integrating over the whole domain
-@assert isapprox(sum(Remapping.integrate_each_element(field1)), sum(field1), atol = 1e-12)
-# Check that integrating 1 over each element and summing gives the same result as integrating 1 over the whole domain
-@assert sum(Remapping.integrate_each_element(ones_field1)) ≈ sum(ones_field1)
-
-# Get one value per element in the field, equal to the average of the values at nodes of the element
-value_per_element1 = zeros(Float64, Meshes.nelements(space1.grid.topology.mesh))
-Remapping.get_value_per_element!(value_per_element1, field1, ones_field1)
-
-# Allocate a vector with length equal to the number of elements in the target space
-value_per_element2 = zeros(Float64, Meshes.nelements(space2.grid.topology.mesh))
-ConservativeRegridding.regrid!(value_per_element2, regridder_1_to_2, value_per_element1)
-
-# Now that we have our regridded vector, put it onto a field on the second space
-field2 = Fields.zeros(space2)
-Remapping.set_value_per_element!(field2, value_per_element2)
-field1_one_value_per_element = Fields.zeros(space1)
-Remapping.set_value_per_element!(field1_one_value_per_element, value_per_element1)
-
-# # Plot the fields
-# using ClimaCoreMakie
-# using GLMakie
-# fig = ClimaCoreMakie.fieldheatmap(field1)
-# save("field1.png", fig)
-# fig = ClimaCoreMakie.fieldheatmap(field1_one_value_per_element)
-# save("field1_one_value_per_element.png", fig)
-# fig = ClimaCoreMakie.fieldheatmap(field2)
-# save("field2.png", fig)
-
-# Check the conservation error
-abs_error = abs(sum(field1) - sum(field2))
-@assert abs_error < 1e-12
-abs_error_one_value_per_element = abs(sum(field1_one_value_per_element) - sum(field2))
-@assert abs_error_one_value_per_element < 2e-12
+@testset "test get_element_vertices" begin
+    vertices = ClimaConservativeRegrid.get_element_vertices(src_space)
+    @test length(vertices) == Meshes.nelements(src_space.grid.topology.mesh)
+
+    # Check that there are 5 vertices per element (quadrilaterals with repeated first vertex)
+    @test all(length(vertex) == 5 for vertex in vertices)
+    @test all(vertex[1] == vertex[5] for vertex in vertices)
+end
+
+@testset "test integrate_each_element" begin
+    # Test integrating a field of ones
+    ones_field = Fields.ones(src_space)
+    integral_each_element =
+        ClimaConservativeRegrid.integrate_each_element(ones_field)
+    @test isapprox(sum(integral_each_element), sum(ones_field), atol = 1e-11)
+
+    # Test integrating a field of latitude
+    field = Fields.coordinate_field(src_space).lat
+    integral_each_element = ClimaConservativeRegrid.integrate_each_element(field)
+    @test isapprox(sum(integral_each_element), sum(field), atol = 1e-12)
+end
+
+@testset "test get_value_per_element!" begin
+    field = Fields.coordinate_field(src_space).lat
+    ones_field = Fields.ones(src_space)
+    value_per_element = zeros(Float64, Meshes.nelements(src_space.grid.topology.mesh))
+    ClimaConservativeRegrid.get_value_per_element!(
+        value_per_element,
+        field,
+        ones_field,
+    )
+
+    @test isapprox(sum(value_per_element), sum(field), atol = 1e-12)
+end
+
+@testset "test set_value_per_element!" begin
+    field = Fields.coordinate_field(src_space).lat
+    value_per_element = ones(Float64, Meshes.nelements(src_space.grid.topology.mesh))
+    ClimaConservativeRegrid.set_value_per_element!(field, value_per_element)
+
+    @test isapprox(sum(field), sum(value_per_element), atol = 1e-12)
+    @test all(field .== 1.0)
+end
+
+@testset "test Regridder constructor" begin
+    regridder = ClimaConservativeRegrid.Regridder(dst_space, src_space)
+    @test regridder isa ConservativeRegridding.Regridder
+end
+
+@testset "test regrid!" begin
+    src_field = Fields.coordinate_field(src_space).lat
+    dst_field = Fields.zeros(dst_space)
+
+    # Test regrid! without pre-allocated buffers
+    regridder = ClimaConservativeRegrid.Regridder(dst_space, src_space)
+    ClimaConservativeRegrid.regrid!(dst_field, regridder, src_field)
+    @test isapprox(sum(dst_field), sum(src_field), atol = 1e-12)
+
+    # Test regrid! with pre-allocated buffers
+    value_per_element_src = zeros(Float64, Meshes.nelements(src_space.grid.topology.mesh))
+    value_per_element_dst = zeros(Float64, Meshes.nelements(dst_space.grid.topology.mesh))
+    ones_src = ones(src_space)
+    regridder_tuple = (;
+        regridder,
+        value_per_element_src,
+        value_per_element_dst,
+        ones_src,
+    )
+    ClimaConservativeRegrid.regrid!(dst_field, regridder_tuple, src_field)
+    @test isapprox(sum(dst_field), sum(src_field), atol = 1e-12)
+end
+
+@testset "test regrid! onto the same space" begin
+    src_field = Fields.coordinate_field(src_space).lat
+    dst_field = Fields.zeros(src_space)
+
+    # Test regrid! without pre-allocated buffers
+    regridder = ClimaConservativeRegrid.Regridder(src_space, src_space)
+    ClimaConservativeRegrid.regrid!(dst_field, regridder, src_field)
+    @test isapprox(sum(dst_field), sum(src_field), atol = 1e-12)
+end
+
+@testset "test regrid! of a constant field" begin
+    src_field = ones(src_space)
+    dst_field = Fields.zeros(src_space)
+
+    # Test regrid! without pre-allocated buffers
+    regridder = ClimaConservativeRegrid.Regridder(src_space, src_space)
+    ClimaConservativeRegrid.regrid!(dst_field, regridder, src_field)
+    @test isapprox(sum(dst_field), sum(src_field), atol = 1e-12)
+end
+
+@testset "test regrid! from source to destination and back" begin
+    src_field = Fields.coordinate_field(src_space).lat
+    dst_field = Fields.zeros(dst_space)
+
+    # Regrid from source to destination
+    regridder = ClimaConservativeRegrid.Regridder(dst_space, src_space)
+    ClimaConservativeRegrid.regrid!(dst_field, regridder, src_field)
+    @test isapprox(sum(dst_field), sum(src_field), atol = 1e-12)
+
+    # Regrid from destination to source using the transpose of the regridder
+    ClimaConservativeRegrid.regrid!(src_field, transpose(regridder), dst_field)
+    @test isapprox(sum(src_field), sum(dst_field), atol = 1e-12)
+end