Merge pull request #129 from JuliaDiffEq/nnlib_convolutions

Fixed dispatch 2D/3D multiplication for DerivativeOperator

Author: ChrisRackauckas

src/derivative_operators/derivative_operator_functions.jl

@@ -27,7 +27,7 @@ end
 
 for MT in [2,3]
     @eval begin
-        function LinearAlgebra.mul!(x_temp::AbstractArray{T,$MT}, A::DerivativeOperator{T,N,Wind,T2,S1}, M::AbstractArray{T,$MT}) where {T<:Real,N,Wind,T2,SL,S1<:SArray{Tuple{SL},T,1,SL}}
+        function LinearAlgebra.mul!(x_temp::AbstractArray{T,$MT}, A::DerivativeOperator{T,N,false,T2,S1}, M::AbstractArray{T,$MT}) where {T<:Real,N,T2,SL,S1<:SArray{Tuple{SL},T,1,SL}}
 
             # Check that x_temp has correct dimensions
             v = zeros(ndims(x_temp))
@@ -38,52 +38,49 @@ for MT in [2,3]
             ndimsM = ndims(M)
             @assert N <= ndimsM
 
-            # Respahe x_temp for NNlib.conv!
-            new_size = Any[size(x_temp)...]
+            # Determine padding for NNlib.conv!
             bpc = A.boundary_point_count
-            setindex!(new_size, new_size[N]- 2*bpc, N)
-            new_shape = []
-            for i in 1:ndimsM
-                if i != N
-                    push!(new_shape,:)
-                else
-                    push!(new_shape,bpc+1:new_size[N]+bpc)
-                end
-             end
-             _x_temp = reshape(view(x_temp, new_shape...), (new_size...,1,1))
+            pad = zeros(Int64,ndimsM)
+            pad[N] = bpc
+
+            # Reshape x_temp for NNlib.conv!
+            _x_temp = reshape(x_temp, (size(x_temp)...,1,1))
 
             # Reshape M for NNlib.conv!
             _M = reshape(M, (size(M)...,1,1))
-            s = A.stencil_coefs
-            sl = A.stencil_length
 
             # Setup W, the kernel for NNlib.conv!
+            s = A.stencil_coefs
+            sl = A.stencil_length
             Wdims = ones(Int64, ndims(_x_temp))
             Wdims[N] = sl
             W = zeros(Wdims...)
             Widx = Any[Wdims...]
             setindex!(Widx,:,N)
-            W[Widx...] = s ./ A.dx^A.derivative_order # this will change later 
-            cv = DenseConvDims(_M, W)
+            W[Widx...] = s
 
+            cv = DenseConvDims(_M, W, padding=pad)
             conv!(_x_temp, _M, W, cv)
 
             # Now deal with boundaries
-            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_BC_left!(view(x_temp, idx...), view(M, idx...), A)
-                convolve_BC_right!(view(x_temp, idx...), view(M, idx...), A)
+            if bpc > 0
+                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_BC_left!(view(x_temp, idx...), view(M, idx...), A)
+                    convolve_BC_right!(view(x_temp, idx...), view(M, idx...), A)
+                end
             end
+            mul!(x_temp,x_temp,1/A.dx^A.derivative_order)
         end
     end
 end

test/differentiation_dimension.jl

@@ -21,7 +21,46 @@ function second_derivative_stencil(N)
   A
 end
 
-@testset "Differenting first three dimensions" begin
+@testset "Differentiation on 2D array" begin
+    M = zeros(22,22)
+    M_temp = zeros(20,22)
+    indices = Iterators.drop(CartesianIndices((22,22)), 0)
+    for idx in indices
+        M[idx] = sin(idx[1]*0.1)
+    end
+    L = CenteredDifference(4,4,0.1,20)
+
+    #test mul!
+    mul!(M_temp, L, M)
+
+    correct_row = 10000.0*fourth_deriv_approx_stencil(20)*M[:,1]
+
+    for i in 1:22
+        @test M_temp[:,i] ≈ correct_row
+    end
+
+    # Test that * agrees will mul!
+    @test M_temp == L*M
+
+    # Differentiation along second dimension
+    L = CenteredDifference{2}(4,4,0.1,20)
+    M_temp_2 = zeros(22,20)
+    indices = Iterators.drop(CartesianIndices((22,22)), 0)
+    for idx in indices
+        M[idx] = sin(idx[2]*0.1)
+    end
+
+    #test mul!
+    mul!(M_temp_2, L, M)
+    for i in 1:22
+        @test M_temp_2[i,:] ≈ correct_row
+    end
+
+    # Test that * agrees will mul!
+    @test M_temp_2 == L*M
+end
+
+@testset "Differenting on 3D array with L2" begin
     M = zeros(22,22,22)
     M_temp = zeros(20,22,22)
     indices = Iterators.drop(CartesianIndices((22,22,22)), 0)
@@ -83,21 +122,81 @@ end
     @test M_temp_3 == L*M
 end
 
+@testset "Differentiation on 3D array with L4" begin
+    M = zeros(22,22,22)
+    M_temp = zeros(20,22,22)
+    indices = Iterators.drop(CartesianIndices((22,22,22)), 0)
+    for idx in indices
+        M[idx] = sin(idx[1]*0.1)
+    end
+    L = CenteredDifference(4,4,0.1,20)
+
+    #test mul!
+    mul!(M_temp, L, M)
+
+    correct_row = 10000.0*fourth_deriv_approx_stencil(20)*M[:,1,1]
+
+    for i in 1:22
+        for j in 1:22
+            @test M_temp[:,i,j] ≈ correct_row
+        end
+    end
+
+    # Test that * agrees will mul!
+    @test M_temp == L*M
+
+    # Differentiation along second dimension
+    L = CenteredDifference{2}(4,4,0.1,20)
+    M_temp_2 = zeros(22,20,22)
+    for idx in indices
+        M[idx] = sin(idx[2]*0.1)
+    end
+
+    #test mul!
+    mul!(M_temp_2, L, M)
+    for i in 1:22
+        for j in 1:22
+            @test M_temp_2[i,:,j] ≈ correct_row
+        end
+    end
+
+    # Test that * agrees will mul!
+    @test M_temp_2 == L*M
+
+    # Differentiation along third dimension
+    L = CenteredDifference{3}(4,4,0.1,20)
+    M_temp_3 = zeros(22,22,20)
+    for idx in indices
+        M[idx] = sin(idx[3]*0.1)
+    end
+
+    #test mul!
+    mul!(M_temp_3, L, M)
+    for i in 1:22
+        for j in 1:22
+            @test M_temp_3[i,j,:] ≈ correct_row
+        end
+    end
+
+    # Test that * agrees will mul!
+    @test M_temp_3 == L*M
+end
+
 @testset "Differentiating an arbitrary higher dimension" begin
     N = 6
     L = CenteredDifference{N}(4,4,0.1,30)
-    M = zeros(5,5,5,5,5,32)
-    M_temp = zeros(5,5,5,5,5,30)
-    indices = Iterators.drop(CartesianIndices((5,5,5,5,5,32)), 0)
+    M = zeros(5,5,5,5,5,32,5);
+    M_temp = zeros(5,5,5,5,5,30,5);
+    indices = Iterators.drop(CartesianIndices((5,5,5,5,5,32,5)), 0);
     for idx in indices
         M[idx] = cos(idx[N]*0.1)
     end
 
-    correct_row = (10.0^4)*fourth_deriv_approx_stencil(30)*M[1,1,1,1,1,:]
+    correct_row = (10.0^4)*fourth_deriv_approx_stencil(30)*M[1,1,1,1,1,:,1]
 
     #test mul!
     mul!(M_temp, L, M)
-    indices = Iterators.drop(CartesianIndices((5,5,5,5,5,30)), 0)
+    indices = Iterators.drop(CartesianIndices((5,5,5,5,5,30,5)), 0);
     for idx in indices
         @test M_temp[idx] ≈ correct_row[idx[N]]
     end