fixes

Author: juliasloan25

experiments/ClimaEarth/components/ocean/climaocean_helpers.jl

@@ -14,10 +14,44 @@ function surface_flux(f::OC.AbstractField)
     end
 end
 
-function Interfacer.remap(operation::OC.AbstractOperations.AbstractOperation, target_space)
+function Interfacer.remap(
+    operation::OC.AbstractOperations.AbstractOperation,
+    remapping,
+    target_space,
+)
+    evaluated_field = OC.Field(operation)
+    OC.compute!(evaluated_field)
+    return Interfacer.remap(evaluated_field, remapping, target_space)
+end
+
+function Interfacer.remap!(
+    target_field,
+    operation::OC.AbstractOperations.AbstractOperation,
+    remapping,
+)
     evaluated_field = OC.Field(operation)
     OC.compute!(evaluated_field)
-    return Interfacer.remap(evaluated_field, target_space)
+    return Interfacer.remap!(target_field, evaluated_field, remapping)
+end
+
+# Handle the case when remap is called on an operation without remapping argument
+# This happens when the generic remap! calls remap(source, target_space)
+# This should not happen for Oceananigans - operations always need remapping
+function Interfacer.remap(operation::OC.AbstractOperations.AbstractOperation, target_space)
+    error(
+        "Cannot remap Oceananigans AbstractOperation to target space without remapping object. " *
+        "Use the specialized get_field! method for OceananigansSimulation instead.",
+    )
+end
+
+# Handle the case when remap is called on an Oceananigans Field without remapping argument
+# This happens when the generic remap! calls remap(source, target_space)
+function Interfacer.remap(src_field::OC.Field, target_space)
+    error(
+        "Cannot remap Oceananigans Field to target space without remapping object. " *
+        "Oceananigans Fields require a remapping object to convert to ClimaCore spaces. " *
+        "Use the specialized get_field! method for OceananigansSimulation instead.",
+    )
 end
 
 """
@@ -108,11 +142,11 @@ function compute_cell_matrix(
     FT = eltype(grid)
 
     vertices_per_cell = 5 # convention: [sw, nw, ne, se, sw]
-    ArrayType = Oceananigans.Architectures.array_type(arch)
+    ArrayType = OC.Architectures.array_type(arch)
     cell_matrix = ArrayType{Tuple{FT, FT}}(undef, vertices_per_cell, Fx * Fy)
 
     arch = grid.architecture
-    Oceananigans.Utils.launch!(
+    OC.Utils.launch!(
         arch,
         grid,
         (Fx, Fy),
@@ -127,13 +161,13 @@ function compute_cell_matrix(
     return cell_matrix
 end
 
-flip(::Face) = Center()
-flip(::Center) = Face()
+flip(::OC.Face) = OC.Center()
+flip(::OC.Center) = OC.Face()
 
-left_index(i, ::Center) = i
-left_index(i, ::Face) = i - 1
-right_index(i, ::Center) = i + 1
-right_index(i, ::Face) = i
+left_index(i, ::OC.Center) = i
+left_index(i, ::OC.Face) = i - 1
+right_index(i, ::OC.Center) = i + 1
+right_index(i, ::OC.Face) = i
 
 @kernel function _compute_cell_matrix!(cell_matrix, Fx, ℓx, ℓy, grid)
     i, j = @index(Global, NTuple)
@@ -153,17 +187,17 @@ right_index(i, ::Face) = i
     ise = right_index(i, ℓx)
     jse = left_index(j, ℓy)
 
-    xsw = ξnode(isw, jsw, 1, grid, vx, vy, nothing)
-    ysw = ηnode(isw, jsw, 1, grid, vx, vy, nothing)
+    xsw = OC.ξnode(isw, jsw, 1, grid, vx, vy, nothing)
+    ysw = OC.ηnode(isw, jsw, 1, grid, vx, vy, nothing)
 
-    xnw = ξnode(inw, jnw, 1, grid, vx, vy, nothing)
-    ynw = ηnode(inw, jnw, 1, grid, vx, vy, nothing)
+    xnw = OC.ξnode(inw, jnw, 1, grid, vx, vy, nothing)
+    ynw = OC.ηnode(inw, jnw, 1, grid, vx, vy, nothing)
 
-    xne = ξnode(ine, jne, 1, grid, vx, vy, nothing)
-    yne = ηnode(ine, jne, 1, grid, vx, vy, nothing)
+    xne = OC.ξnode(ine, jne, 1, grid, vx, vy, nothing)
+    yne = OC.ηnode(ine, jne, 1, grid, vx, vy, nothing)
 
-    xse = ξnode(ise, jse, 1, grid, vx, vy, nothing)
-    yse = ηnode(ise, jse, 1, grid, vx, vy, nothing)
+    xse = OC.ξnode(ise, jse, 1, grid, vx, vy, nothing)
+    yse = OC.ηnode(ise, jse, 1, grid, vx, vy, nothing)
 
     linear_idx = i + (j - 1) * Fx
     @inbounds begin

experiments/ClimaEarth/components/ocean/oceananigans.jl

@@ -1,14 +1,14 @@
 import Oceananigans as OC
 import ClimaOcean as CO
 import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities
-import Interfacer: remap, remap!, get_field
+import ClimaCoupler.Interfacer: remap, remap!, get_field
 import ClimaComms
 import ClimaCore as CC
 import Thermodynamics as TD
 import ClimaParams as CP
 import ClimaOcean.EN4: download_dataset
 using KernelAbstractions: @kernel, @index, @inbounds
-import ConservativeRemapping as CR
+import ConservativeRegridding as CR
 
 include("climaocean_helpers.jl")
 
@@ -76,7 +76,7 @@ function OceananigansSimulation(
     bottom_height = CO.regrid_bathymetry(
         underlying_grid;
         minimum_depth = 30,
-        interpolation_passes = 20,
+        interpolation_passes = 1,
         major_basins = 1,
     )
     grid = OC.ImmersedBoundaryGrid(
@@ -183,7 +183,7 @@ To regrid from ClimaCore to Oceananigans, use `CR.regrid!(dest_vector, transpose
 """
 function construct_remappers(grid_oc, boundary_space)
     # Get the vector of polygons for Oceananigans and ClimaCore spaces
-    vertices_oc = compute_cell_matrix(grid_oc)
+    vertices_oc = compute_cell_matrix(grid_oc.underlying_grid)
     vertices_cc = CC.Remapping.get_element_vertices(boundary_space)
 
     remapper_oc_to_cc = CR.Regridder(vertices_cc, vertices_oc)
@@ -210,12 +210,15 @@ end
 
 # Non-allocating ClimaCore -> Oceananigans remap
 function Interfacer.remap!(dst_field::OC.Field, src_field::CC.Fields.Field, remapping)
-    value_per_element_cc =
-        CC.Remappping.get_value_per_element(src_field, remapping.field_ones_cc)
+    CC.Remapping.get_value_per_element!(
+        remapping.value_per_element_cc,
+        src_field,
+        remapping.field_ones_cc,
+    )
     CR.regrid!(
-        vec(interior(dst_field)),
+        vec(OC.interior(dst_field, :, :, 1)),
         transpose(remapping.remapper_oc_to_cc),
-        value_per_element_cc,
+        remapping.value_per_element_cc,
     )
     return nothing
 end
@@ -235,13 +238,18 @@ function Interfacer.remap!(dst_field::CC.Fields.Field, src_field::OC.Field, rema
     CR.regrid!(
         remapping.value_per_element_cc,
         remapping.remapper_oc_to_cc,
-        vec(interior(src_field)),
+        vec(OC.interior(src_field, :, :, 1)),
     )
 
     # Convert the vector of remapped values to a ClimaCore Field with one value per element
     CC.Remapping.set_value_per_element!(dst_field, remapping.value_per_element_cc)
     return nothing
 end
+# Handle the case of remapping the area fraction field, which is a ClimaCore Field
+Interfacer.remap!(dst_field::CC.Fields.Field, src_field::CC.Fields.Field, remapping) =
+    Interfacer.remap!(dst_field, src_field)
+Interfacer.remap!(dst_field::CC.Fields.Field, src_field::Number, remapping) =
+    Interfacer.remap!(dst_field, src_field)
 # Allocating Oceananigans -> ClimaCore remap
 function Interfacer.remap(
     src_field::OC.Field,
@@ -253,10 +261,22 @@ function Interfacer.remap(
     return dst_field
 end
 
+# Handle the case of remapping a scalar number to a ClimaCore space
+Interfacer.remap(num::Number, remapping, target_space::CC.Spaces.AbstractSpace) =
+    Interfacer.remap(num, target_space)
+
 # Extend Interfacer.get_field to allow automatic remapping to the target space
 # TODO see if we can remove this
-function Interfacer.get_field(sim, quantity, target_space)
-    return remap(Interfacer.get_field(sim, quantity), sim.remapping, target_space)
+function Interfacer.get_field(sim::OceananigansSimulation, quantity, target_space)
+    return Interfacer.remap(
+        Interfacer.get_field(sim, quantity),
+        sim.remapping,
+        target_space,
+    )
+end
+function Interfacer.get_field!(target_field, sim::OceananigansSimulation, quantity)
+    Interfacer.remap!(target_field, Interfacer.get_field(sim, quantity), sim.remapping)
+    return nothing
 end
 
 """
@@ -363,21 +383,23 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi
     Interfacer.remap!(sim.remapping.scratch_field_oc2, F_turb_ρτyz, sim.remapping) # meridional momentum flux
 
     # Rename for clarity; these are now Center, Center Oceananigans fields
-    F_turb_ρτxz_oc = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1)
-    F_turb_ρτyz_oc = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1)
+    oc_F_turb_ρτxz = sim.remapping.scratch_field_oc1
+    oc_F_turb_ρτyz = sim.remapping.scratch_field_oc2
 
     # Weight by (1 - sea ice concentration)
-    OC.interior(F_turb_ρτxz_oc, :, :, 1) .= F_turb_ρτxz_oc .* (1.0 .- ice_concentration)
-    OC.interior(F_turb_ρτyz_oc, :, :, 1) .= F_turb_ρτyz_oc .* (1.0 .- ice_concentration)
+    OC.interior(oc_F_turb_ρτxz, :, :, 1) .=
+        OC.interior(oc_F_turb_ρτxz, :, :, 1) .* (1.0 .- ice_concentration)
+    OC.interior(oc_F_turb_ρτyz, :, :, 1) .=
+        OC.interior(oc_F_turb_ρτyz, :, :, 1) .* (1.0 .- ice_concentration)
 
     # Set the momentum flux BCs at the correct locations using the remapped scratch fields
     oc_flux_u = surface_flux(sim.ocean.model.velocities.u)
     oc_flux_v = surface_flux(sim.ocean.model.velocities.v)
     set_from_extrinsic_vector!(
         (; u = oc_flux_u, v = oc_flux_v),
         grid,
-        F_turb_ρτxz_cc,
-        F_turb_ρτyz_cc,
+        oc_F_turb_ρτxz,
+        oc_F_turb_ρτyz,
     )
 
     (; reference_density, heat_capacity, fresh_water_density) = sim.ocean_properties