Skip to content

Make remap! not allocate #1596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
juliasloan25 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Make remap! not allocate #1596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ac5690
make allocating remap call non-allocating remap
juliasloan25 Dec 5, 2025
65a44ca
fix CMIP remapping
juliasloan25 Dec 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 36 additions & 16 deletions experiments/ClimaEarth/components/ocean/climaocean_helpers.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,35 +9,36 @@ 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
# interpolate outside of its latitude bounds. Instead we return 0
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

"""
Expand All @@ -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

"""
Expand Down
90 changes: 62 additions & 28 deletions src/Interfacer.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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
Expand All @@ -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)]
Expand All @@ -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

Expand Down
Loading