coefficient handling for compositions

Author: ajozefiak

src/derivative_operators/derivative_operator_functions.jl

@@ -26,6 +26,36 @@ function LinearAlgebra.mul!(x_temp::AbstractArray{T}, A::DerivativeOperator{T,N}
     end
 end
 
+# Additive mul! that is used for handling compositions
+function mul_add!(x_temp::AbstractArray{T}, A::DerivativeOperator{T,N}, M::AbstractArray{T}) where {T,N}
+
+    # Check that x_temp has correct dimensions
+    v = zeros(ndims(x_temp))
+    v[N] = 2
+    @assert [size(x_temp)...]+v == [size(M)...]
+
+    # Check that axis of differentiation is in the dimensions of M and x_temp
+    ndimsM = ndims(M)
+    @assert N <= ndimsM
+
+    dimsM = [axes(M)...]
+    alldims = [1:ndims(M);]
+    otherdims = setdiff(alldims, N)
+
+    idx = Any[first(ind) for ind in axes(M)]
+    itershape = tuple(dimsM[otherdims]...)
+    nidx = length(otherdims)
+    indices = Iterators.drop(CartesianIndices(itershape), 0)
+
+    setindex!(idx, :, N)
+    for I in indices
+        Base.replace_tuples!(nidx, idx, idx, otherdims, I)
+        convolve_interior_add!(view(x_temp, idx...), view(M, idx...), A)
+        convolve_BC_right_add!(view(x_temp, idx...), view(M, idx...), A)
+        convolve_BC_left_add!(view(x_temp, idx...), view(M, idx...), A)
+    end
+end
+
 # A more efficient mul! implementation for a single, regular-grid, centered difference,
 # scalar coefficient DerivativeOperator operating on a 2D or 3D AbstractArray
 for MT in [2,3]
@@ -145,11 +175,10 @@ function LinearAlgebra.mul!(x_temp::AbstractArray{T,2}, A::AbstractDiffEqComposi
         sl = L.stencil_length
         axis = typeof(L).parameters[2]
         offset = convert(Int64,(Wdims[axis] - sl)/2)
+        coeff = L.coefficients isa Number ? L.coefficients : true
         for i in offset+1:Wdims[axis]-offset
             idx[axis]=i
-
-            W[idx...] += s[i-offset]
-
+            W[idx...] += coeff*s[i-offset]
             idx[axis] = mid_Wdims[axis]
         end
     end
@@ -271,7 +300,7 @@ function convolve_interior_add!(x_temp::AbstractVector{T}, x::AbstractVector{T},
     for i in (1+A.boundary_point_count) : (length(x_temp)-A.boundary_point_count)
         xtempi = zero(T)
         cur_stencil = eltype(stencil) <: AbstractVector ? stencil[i] : stencil
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : coeff isa Number ? coeff : true
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
         for idx in 1:A.stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[i - mid + idx]
@@ -288,7 +317,7 @@ function convolve_interior_add_range!(x_temp::AbstractVector{T}, x::AbstractVect
     for i in [(1+A.boundary_point_count):(A.boundary_point_count+offset); (length(x_temp)-A.boundary_point_count-offset+1):(length(x_temp)-A.boundary_point_count)]
         xtempi = zero(T)
         cur_stencil = eltype(stencil) <: AbstractVector ? stencil[i] : stencil
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : coeff isa Number ? coeff : true
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
         for idx in 1:A.stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[i - mid + idx]
@@ -301,10 +330,10 @@ function convolve_BC_left_add!(x_temp::AbstractVector{T}, x::AbstractVector{T},
     stencil = A.low_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
-        xtempi = stencil[i][1]*x[1]
         cur_stencil = stencil[i]
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : coeff isa Number ? coeff : true
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
+        xtempi = cur_coeff*stencil[i][1]*x[1]
         for idx in 2:A.boundary_stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[idx]
         end
@@ -316,10 +345,10 @@ function convolve_BC_right_add!(x_temp::AbstractVector{T}, x::AbstractVector{T},
     stencil = A.high_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
-        xtempi = stencil[i][end]*x[end]
         cur_stencil = stencil[i]
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : coeff isa Number ? coeff : true
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
+        xtempi = cur_coeff*stencil[i][end]*x[end]
         for idx in (A.boundary_stencil_length-1):-1:1
             xtempi += cur_coeff * cur_stencil[end-idx] * x[end-idx]
         end

test/2D_3D_fast_multiplication.jl

@@ -1,26 +1,5 @@
 using LinearAlgebra, DiffEqOperators, Random, Test, BandedMatrices, SparseArrays
 
-function fourth_deriv_approx_stencil(N)
-    A = zeros(N,N+2)
-    A[1,1:8] = [3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3]
-    A[2,1:8] = [2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6]
-    A[N-1,N-5:end] = reverse([2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6], dims=2)
-    A[N,N-5:end] = reverse([3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3], dims=2)
-    for i in 3:N-2
-        A[i,i-2:i+4] = [-1/6 2.0 -13/2 28/3 -13/2 2.0 -1/6]
-    end
-    return A
-end
-
-function second_derivative_stencil(N)
-  A = zeros(N,N+2)
-  for i in 1:N, j in 1:N+2
-      (j-i==0 || j-i==2) && (A[i,j]=1)
-      j-i==1 && (A[i,j]=-2)
-  end
-  A
-end
-
 @testset "2D Multiplication with no boundary points and dx = 1.0" begin
 
     # Test (Lxx + Lyy)*M, dx = 1.0, no coefficient
@@ -409,3 +388,64 @@ end
     @test M_temp ≈ ((Lx2*M)[1:N,2:N+1]+(Ly2*M)[2:N+1,1:N]+(Lx3*M)[1:N,2:N+1] +(Ly3*M)[2:N+1,1:N] + (Lx4*M)[1:N,2:N+1] +(Ly4*M)[2:N+1,1:N])
 
 end
+
+# THis testset uses the last testset which has a several non-trivial cases,
+# and additionally tests coefficient handling. All operators are handled by the
+# fast 2D/3D dispatch
+@testset "2D coefficient handling" begin
+
+    dx = 0.1
+    dy = 0.25
+    N = 100
+    M = zeros(N+2,N+2)
+    M_temp = zeros(N,N+2)
+    for i in 1:N+2
+        for j in 1:N+2
+            M[i,j] = cos(dx*i)+sin(dy*j)
+        end
+    end
+
+    # Lx2 has 0 boundary points
+    Lx2 = 5.5*CenteredDifference{1}(2,2,dx,N)
+    # Lx3 has 1 boundary point
+    Lx3 = 1.45*CenteredDifference{1}(3,3,dx,N)
+    # Lx4 has 2 boundary points
+    Lx4 = 0.5*CenteredDifference{1}(4,4,dx,N)
+
+    # Test a single axis, multiple operators: (Lxx+Lxxxx)*M, dx = 1.0
+    A = Lx2+Lx4
+    mul!(M_temp, A, M)
+    @test M_temp ≈ ((Lx2*M) + (Lx4*M))
+
+    # Test a single axis, multiple operators: (Lxx++Lxxx+Lxxxx)*M, dx = 1.0
+    A += Lx3
+    mul!(M_temp, A, M)
+    @test M_temp ≈ ((Lx2*M) + (Lx3*M) + (Lx4*M))
+
+
+    # Ly2 has 0 boundary points
+    Ly2 = 8.14*CenteredDifference{2}(2,2,dy,N)
+    # Ly3 has 1 boundary point
+    Ly3 = 2.0*CenteredDifference{2}(3,3,dy,N)
+    # Ly4 has 2 boundary points
+    Ly4 = 4.567*CenteredDifference{2}(4,4,dy,N)
+    M_temp = zeros(N+2,N)
+
+    # Test a single axis, multiple operators: (Lyy+Lyyyy)*M, dx = 1.0
+    A = Ly2+Ly4
+    mul!(M_temp, A, M)
+    @test M_temp ≈ ((Ly2*M) + (Ly4*M))
+
+    # Test a single axis, multiple operators: (Lyy++Lyyy+Lyyyy)*M, dx = 1.0
+    A += Ly3
+    mul!(M_temp, A, M)
+    @test M_temp ≈ ((Ly2*M) + (Ly3*M) + (Ly4*M))
+
+
+    # Test multiple operators on both axis: (Lxx + Lyy + Lxxx + Lyyy + Lxxxx + Lyyyy)*M, no coefficient
+    A = Lx2 + Ly2 + Lx3 + Ly3 + Lx4 + Ly4
+    M_temp = zeros(100,100)
+    mul!(M_temp, A, M)
+
+    @test M_temp ≈ ((Lx2*M)[1:N,2:N+1]+(Ly2*M)[2:N+1,1:N]+(Lx3*M)[1:N,2:N+1] +(Ly3*M)[2:N+1,1:N] + (Lx4*M)[1:N,2:N+1] +(Ly4*M)[2:N+1,1:N])
+end