diff --git a/src/MatrixFields/band_matrix_row.jl b/src/MatrixFields/band_matrix_row.jl index d8d4b727e8..053680d363 100644 --- a/src/MatrixFields/band_matrix_row.jl +++ b/src/MatrixFields/band_matrix_row.jl @@ -162,3 +162,14 @@ is_auto_broadcastable(::Type{BMR}) where {BMR <: BandMatrixRow} = is_auto_broadcastable(eltype(BMR)) add_auto_broadcasters(row::BandMatrixRow) = map(add_auto_broadcasters, row) drop_auto_broadcasters(row::BandMatrixRow) = map(drop_auto_broadcasters, row) + +function (row::BandMatrixRow{ld, bw, T})(args::Vararg{Any, N}) where {ld, bw, T, N} + N == bw || error( + "BandMatrixRow with bandwidth $bw expected $bw arguments, but got $N", + ) + + result = unrolled_sum(Iterators.zip(row.entries, args)) do (entry, arg) + entry * arg + end + return result +end diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index 57ccaed22f..dcb879a575 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -27,5 +27,6 @@ include("finitedifference.jl") include("remapping.jl") include("integrals.jl") include("columnwise.jl") +include("column_convolve.jl") end # module diff --git a/src/Operators/column_convolve.jl b/src/Operators/column_convolve.jl new file mode 100644 index 0000000000..315a7d15ee --- /dev/null +++ b/src/Operators/column_convolve.jl @@ -0,0 +1,60 @@ +import UnrolledUtilities: unrolled_sum + +""" +Represents a 1D stencil to be applied by a convolution operator. + +At the moment the functionality is doubled wrt to BandMatrixRow, but the latter +cannot be used due to the module order (Operators are defined before Matrix Fields) + +""" +struct ConvolutionKernel{ld, N, T} + entries::NTuple{N, T} + ConvolutionKernel{ld, N, T}(entries::NTuple{N, Any}) where {ld, N, T} = + new{ld, N, T}(entries) +end + +ConvolutionKernel{ld}(entries::Vararg{Any, N}) where {ld, N} = + ConvolutionKernel{ld, N}(entries...) + +ConvolutionKernel{ld, N}(entries::Vararg{Any, N}) where {ld, N} = + ConvolutionKernel{ld, N, promote_type(map(typeof, entries)...)}(entries) + +function (row::ConvolutionKernel{ld, bw, T})(args::Vararg{Any, N}) where {ld, bw, T, N} + N == bw || error( + "BandMatrixRow with bandwidth $bw expected $bw arguments, but got $N", + ) + + result = unrolled_sum(Iterators.zip(row.entries, args)) do (entry, arg) + entry * arg + end + return result +end + +# Some indexing wrapper to allow extrapolation +# Will dispatch on the BC (and operator?) in the future +Base.@propagate_inbounds getidx_extrapolated(x, idx, value) = + idx < 1 ? value : idx > length(x) ? value : x[idx] + +""" + apply_kernel!(out, kernel::ConvolutionKernel, in) + +Applies a 1D convolution kernel on a column view + +`in` and `out` are DataLayouts.VF representing single columns. + + - We may need a `getindex` wrapper that will handle BCs + +""" +function apply_kernel!(out, kernel::ConvolutionKernel{ld, bw}, in) where {ld, bw} + + # Should we have `eachindex` here. Should DataLayouts implement it? will they after + # refactor? + @inbounds begin + for i in 1:length(in) + args = unrolled_map(1:bw) do j + getidx_extrapolated(in, i + ld + j - 1, zero(eltype(in))) + end + out[i] = kernel(args...) + end + end +end diff --git a/test/MatrixFields/band_matrix_row.jl b/test/MatrixFields/band_matrix_row.jl index affe69667b..0c751fb0ac 100644 --- a/test/MatrixFields/band_matrix_row.jl +++ b/test/MatrixFields/band_matrix_row.jl @@ -59,4 +59,11 @@ include("matrix_field_test_utils.jl") @test zero(typeof(DiagonalMatrixRow(1))) == DiagonalMatrixRow(0) @test eltype(typeof(DiagonalMatrixRow(1))) == Int @test inv(DiagonalMatrixRow(1)) == DiagonalMatrixRow(float(1)) + + # Test callable behaviour + @test DiagonalMatrixRow(2.0)(10.0) == 20.0 + @test BidiagonalMatrixRow(1.0, 2.0)(10.0, 5.0) == 20.0 + @test TridiagonalMatrixRow(-1.0, 2.0, -1.0)(10.0, 25.0, 12.0) == 28.0 + @test_throws ErrorException TridiagonalMatrixRow(1, 2, 3)(10.0, 5.0) + @test_throws ErrorException TridiagonalMatrixRow(1, 2, 3)(10.0, 5.0, 1.0, 2.0) end diff --git a/test/Operators/finitedifference/unit_column_convolve.jl b/test/Operators/finitedifference/unit_column_convolve.jl new file mode 100644 index 0000000000..171f07e05a --- /dev/null +++ b/test/Operators/finitedifference/unit_column_convolve.jl @@ -0,0 +1,74 @@ +# test/Operators/finitedifference/unit_column_convolve.jl + +using Test +using ClimaComms +import ClimaCore: DataLayouts, Operators, Geometry +using ClimaCore.DataLayouts: VF + +ClimaComms.@import_required_backends + +device = ClimaComms.device() +ArrayType = ClimaComms.array_type(device) + +const getidx_extrapolated = Operators.getidx_extrapolated +const ConvolutionKernel = Operators.ConvolutionKernel + +@testset "apply_kernel! — scalar VF" begin + FT = Float64 + Nv = 10 + + kernel = ConvolutionKernel{-1}(FT(0.25), FT(0.5), FT(0.25)) + + # Build input VF with known values + in_data = VF{FT}(ArrayType{FT}, zeros; Nv) + parent(in_data)[:, 1] .= FT.(1:Nv) + out_data = similar(in_data) + + + # Expected with zero-padding at boundaries: + expected = ArrayType{FT}(undef, Nv) + v = Float64[0; 1:Nv; 0] # zero-padded values + for i in 1:Nv + expected[i] = 0.25 * v[i] + 0.5 * v[i + 1] + 0.25 * v[i + 2] + end + + Operators.apply_kernel!(out_data, kernel, in_data) + @test parent(out_data)[:, 1] ≈ expected +end + +@testset "apply_kernel! — UVWVector VF" begin + FT = Float64 + Nv = 10 + S = Geometry.UVWVector{FT} + + kernel = ConvolutionKernel{-1}(FT(0.25), FT(0.5), FT(0.25)) + + # Build input VF with known UVWVector values + # Each level i has u=i, v=2i, w=3i + in_data = VF{S}(ArrayType{FT}, zeros; Nv) + for i in 1:Nv + in_data[i] = S(FT(i), FT(2i), FT(3i)) + end + out_data = similar(in_data) + + # Expected with zero-padding at boundaries: + # i=1: 0.25*zero(S) + 0.5*S(1,2,3) + 0.25*S(2,4,6) = S(1.25, 2.5, 3.75) + # i=2: 0.25*S(1,2,3) + 0.5*S(2,4,6) + 0.25*S(3,6,9) = S(2.25, 4.5, 6.75) + # i=Nv: 0.25*S(Nv-1,2(Nv-1),3(Nv-1)) + 0.5*S(Nv,2Nv,3Nv) + 0.25*zero(S) + expected = Vector{S}(undef, Nv) + v = Vector{S}(undef, Nv + 2) + v[1] = zero(S) # left padding + for i in 1:Nv + v[i + 1] = S(FT(i), FT(2i), FT(3i)) + end + v[Nv + 2] = zero(S) # right padding + + for i in 1:Nv + expected[i] = FT(0.25) * v[i] + FT(0.5) * v[i + 1] + FT(0.25) * v[i + 2] + end + + Operators.apply_kernel!(out_data, kernel, in_data) + for i in 1:Nv + @test out_data[i] ≈ expected[i] + end +end