Add a basic convolution function

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

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

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

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

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