Wu/reverse accumulate#2531
Draft
cflsepyt wants to merge 5 commits into
Draft
Conversation
nefrathenrici
left a comment
Member
There was a problem hiding this comment.
This looks good, but the reverse implementation is lacking a few details.
| is_c2c_or_f2f = Spaces.staggering(space) == Spaces.staggering(axes(output)) | ||
| is_c2f = !is_c2c_or_f2f && Spaces.staggering(space) == Spaces.CellCenter() | ||
| is_f2c = !is_c2c_or_f2f && !is_c2f | ||
| @inbounds if init == UnspecifiedInit() |
Member
There was a problem hiding this comment.
I realized that it is not enough to just reverse the indices. We need to put the starting boundary to the top level, flip the step direction, and flip the sign of the center/face half-level shift used for staggered accumulation.
I sketched out the needed changes
function single_column_accumulate!(
f::F,
transform::T,
_output,
_input,
init,
space,
reverse,
) where {F, T}
device = ClimaComms.device(space)
first_level = left_idx(space)
last_level = right_idx(space)
output = unstrip_space(_output, space)
is_c2c_or_f2f = Spaces.staggering(space) == Spaces.staggering(axes(output))
is_c2f = !is_c2c_or_f2f && Spaces.staggering(space) == Spaces.CellCenter()
is_f2c = !is_c2c_or_f2f && !is_c2f
# When reverse = true, the starting boundary becomes the top level, the step
# direction flips, and the center/face half-level shift used for staggered
# accumulation flips sign.
start_level, stop_level, direction =
reverse ? (last_level, first_level, -1) : (first_level, last_level, 1)
stagger = reverse ? -half : half
@inbounds if init == UnspecifiedInit()
@assert !is_c2f
accumulated_value = get_level_value(space, _input, start_level)
next_level = start_level + direction
init_output_level = is_c2c_or_f2f ? start_level : nothing
else
accumulated_value =
is_f2c ? f(init, get_level_value(space, _input, start_level)) : init
next_level = is_f2c ? start_level + direction : start_level
init_output_level = is_c2f ? start_level - stagger : nothing
end
@inbounds if !isnothing(init_output_level)
Fields.level(output, init_output_level)[] = transform(accumulated_value)
end
# Iterate levels from `next_level` to `stop_level` stepping by `direction`
n_steps = direction * (stop_level - next_level) + 1
@inbounds for i in 1:n_steps
level = next_level + direction * (i - 1)
accumulated_value =
f(accumulated_value, get_level_value(space, _input, level))
output_level =
is_c2c_or_f2f ? level : (is_c2f ? level + stagger : level - stagger)
Fields.level(output, output_level)[] = transform(accumulated_value)
end
end| (ᶜwhole_number_reverse, ᶜoutput, motzkin_number.(ᶜwhole_number_reverse), true), | ||
| (ᶠwhole_number_reverse, ᶠoutput, motzkin_number.(ᶠwhole_number_reverse), true), | ||
| (ᶠwhole_number_reverse, ᶜoutput, motzkin_number.(ᶜwhole_number_reverse .+ 1), true), | ||
| (ᶜwhole_number_reverse, ᶠoutput, motzkin_number.(ᶠwhole_number_reverse .- 1), true), |
Member
There was a problem hiding this comment.
The tests looks adequate, but I'm not sure that they pass in the current state.
Author
There was a problem hiding this comment.
I ran the tests locally and they all passed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add a reverse option for column_accumulate!()