Skip to content

Wu/reverse accumulate#2531

Draft
cflsepyt wants to merge 5 commits into
mainfrom
wu/reverse_accumulate
Draft

Wu/reverse accumulate#2531
cflsepyt wants to merge 5 commits into
mainfrom
wu/reverse_accumulate

Conversation

@cflsepyt

@cflsepyt cflsepyt commented Jun 30, 2026

Copy link
Copy Markdown

Add a reverse option for column_accumulate!()

  • Code follows the style guidelines OR N/A.
  • Unit tests are included OR N/A.
  • Code is exercised in an integration test OR N/A.
  • Documentation has been added/updated OR N/A.

@nefrathenrici nefrathenrici left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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),

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The tests looks adequate, but I'm not sure that they pass in the current state.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran the tests locally and they all passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants