diff --git a/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl b/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl index 5927ab5327..956c44c619 100644 --- a/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl +++ b/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl @@ -9,26 +9,26 @@ care about the surface. @inline to_node(pt::CC.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) """ - map_interpolate(points, oc_field::OC.Field) + map_interpolate!(target_field::CC.Fields.Field, source_field::OC.Field) -Interpolate the given 3D field onto the target points. +Interpolate the given 3D field onto the target field, modifying the target field in place. -If the underlying grid does not contain a given point, return 0 instead. +If the underlying grid does not contain a given point, writes 0 instead. -Note: `map_interpolate` does not support interpolation from `Field`s defined on +Note: `map_interpolate!` does not support interpolation from `Field`s defined on `OrthogononalSphericalShellGrids` such as the `TripolarGrid`. - -TODO: Use a non-allocating version of this function (simply replace `map` with `map!`) """ -function map_interpolate(points, oc_field::OC.Field) - loc = map(L -> L(), OC.Fields.location(oc_field)) - grid = oc_field.grid - data = oc_field.data +function map_interpolate!(target_field::CC.Fields.Field, source_field::OC.Field) + points = CC.Fields.coordinate_field(axes(target_field)) + loc = map(L -> L(), OC.Fields.location(source_field)) + grid = source_field.grid + data = source_field.data # TODO: There has to be a better way min_lat, max_lat = extrema(OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center())) - map(points) do pt + # Use map! on the field directly, which handles complex data layouts correctly + map!(target_field, points) do pt FT = eltype(pt) # The oceananigans grid does not cover the entire globe, so we should not @@ -36,8 +36,9 @@ function map_interpolate(points, oc_field::OC.Field) min_lat < pt.lat < max_lat || return FT(0) fᵢ = OC.Fields.interpolate(to_node(pt), data, loc, grid) - convert(FT, fᵢ)::FT + return convert(FT, fᵢ)::FT end + return nothing end """ @@ -54,14 +55,33 @@ function surface_flux(f::OC.AbstractField) end end -function Interfacer.remap(field::OC.Field, target_space) - return map_interpolate(CC.Fields.coordinate_field(target_space), field) +function Interfacer.remap!(target_field::CC.Fields.Field, source_field::OC.Field) + map_interpolate!(target_field, source_field) + return nothing end -function Interfacer.remap(operation::OC.AbstractOperations.AbstractOperation, target_space) +function Interfacer.remap!( + target_field::CC.Fields.Field, + operation::OC.AbstractOperations.AbstractOperation, +) evaluated_field = OC.Field(operation) OC.compute!(evaluated_field) - return Interfacer.remap(evaluated_field, target_space) + Interfacer.remap!(target_field, evaluated_field) + return nothing +end + +function Interfacer.remap(target_space, field::OC.Field) + # Allocate target field and call remap! + target_field = CC.Fields.zeros(target_space) + Interfacer.remap!(target_field, field) + return target_field +end + +function Interfacer.remap(target_space, operation::OC.AbstractOperations.AbstractOperation) + # Allocate target field and call remap! + target_field = CC.Fields.zeros(target_space) + Interfacer.remap!(target_field, operation) + return target_field end """ diff --git a/src/Interfacer.jl b/src/Interfacer.jl index 8bd1e65cb1..61bc4f7c2a 100644 --- a/src/Interfacer.jl +++ b/src/Interfacer.jl @@ -458,15 +458,14 @@ abstract type SlabplanetTerraMode <: AbstractSlabplanetSimulationMode end Remap the given `field` onto the `target_space`. If the field is already on the target space or a compatible one, it is returned unchanged. -Note that this method has a lot of allocations and is not efficient. +This is a convenience wrapper around `remap!` that allocates the output field. Non-ClimaCore fields should provide a method to this function. """ function remap end -function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) - source_space = axes(field) - comms_ctx = ClimaComms.context(source_space) +function remap(source_field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) + source_space = axes(source_field) # Check if the source and target spaces are compatible spaces_are_compatible = @@ -475,13 +474,59 @@ function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) CC.Spaces.issubspace(target_space, source_space) # TODO: Handle remapping of Vectors correctly - if hasproperty(field, :components) - @assert length(field.components) == 1 "Can only work with simple vectors" - field = field.components.data.:1 + if hasproperty(source_field, :components) + @assert length(source_field.components) == 1 "Can only work with simple vectors" + source_field = source_field.components.data.:1 end # If the spaces are the same or one is a subspace of the other, we can just return the input field - spaces_are_compatible && return field + spaces_are_compatible && return source_field + + # Allocate target field and call remap! + target_field = CC.Fields.zeros(target_space) + remap!(target_field, source_field) + return target_field +end + +function remap(source_field::Number, target_space::CC.Spaces.AbstractSpace) + # Allocate target field and call remap! + target_field = CC.Fields.zeros(target_space) + remap!(target_field, source_field) + return target_field +end + +""" + remap!(target_field, source_field) + +Remap the given `source_field` onto the `target_field`. This is the core non-allocating +implementation. + +Non-ClimaCore fields should provide a method to this function. +""" +function remap! end + +function remap!(target_field::CC.Fields.Field, source_field::CC.Fields.Field) + source_space = axes(source_field) + target_space = axes(target_field) + comms_ctx = ClimaComms.context(source_space) + + # Check if the source and target spaces are compatible + spaces_are_compatible = + source_space == target_space || + CC.Spaces.issubspace(source_space, target_space) || + CC.Spaces.issubspace(target_space, source_space) + + # TODO: Handle remapping of Vectors correctly + if hasproperty(source_field, :components) + @assert length(source_field.components) == 1 "Can only work with simple vectors" + source_field = source_field.components.data.:1 + end + + # If the spaces are the same or one is a subspace of the other, we can just copy + if spaces_are_compatible + target_field .= source_field + return nothing + end # Get vector of LatLongPoints for the target space to get the hcoords # Copy target coordinates to CPU if they are on GPU @@ -493,10 +538,10 @@ function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) # Remap the field, using MPI if applicable if comms_ctx isa ClimaComms.SingletonCommsContext # Remap source field to target space as an array - remapped_array = CC.Remapping.interpolate(field, hcoords, []) + remapped_array = CC.Remapping.interpolate(source_field, hcoords, []) - # Convert remapped array to a field in the target space - return CC.Fields.array2field(remapped_array, target_space) + # Copy remapped array into target field + target_field .= CC.Fields.array2field(remapped_array, target_space) else # Gather then broadcast the global hcoords and offsets offset = [length(hcoords)] @@ -506,33 +551,22 @@ function remap(field::CC.Fields.Field, target_space::CC.Spaces.AbstractSpace) # Interpolate on root and broadcast to all processes remapper = CC.Remapping.Remapper(source_space; target_hcoords = all_hcoords) remapped_array = - ClimaComms.bcast(comms_ctx, CC.Remapping.interpolate(remapper, field)) + ClimaComms.bcast(comms_ctx, CC.Remapping.interpolate(remapper, source_field)) my_ending_offset = sum(all_offsets[1:ClimaComms.mypid(comms_ctx)]) my_starting_offset = my_ending_offset - offset[] - # Convert remapped array to a field in the target space on each process - return CC.Fields.array2field( + # Copy remapped array into target field on each process + target_field .= CC.Fields.array2field( remapped_array[(1 + my_starting_offset):my_ending_offset], target_space, ) end + return nothing end -function remap(num::Number, target_space::CC.Spaces.AbstractSpace) - return num -end - -""" - remap!(target_field, source) - -Remap the given `source` onto the `target_field`. - -Non-ClimaCore fields should provide a method to [`Interfacer.remap`](@ref), or directly to this -function. -""" -function remap!(target_field, source) - target_field .= remap(source, axes(target_field)) +function remap!(target_field::CC.Fields.Field, source::Number) + fill!(target_field, source) return nothing end