Merge pull request #132 from JuliaDiffEq/nnlib_convolutions

2d multiplication for combinations

Author: ChrisRackauckas

src/derivative_operators/convolutions.jl

@@ -1,14 +1,14 @@
 # mul! done by convolutions
-function LinearAlgebra.mul!(x_temp::AbstractVector{T}, A::DerivativeOperator, x::AbstractVector{T}) where T<:Real
-    convolve_BC_left!(x_temp, x, A)
-    convolve_interior!(x_temp, x, A)
-    convolve_BC_right!(x_temp, x, A)
+function LinearAlgebra.mul!(x_temp::AbstractVector{T}, A::DerivativeOperator, x::AbstractVector{T}; overwrite = true) where T<:Real
+    convolve_BC_left!(x_temp, x, A, overwrite = overwrite)
+    convolve_interior!(x_temp, x, A, overwrite = overwrite)
+    convolve_BC_right!(x_temp, x, A, overwrite = overwrite)
 end
 
 ################################################
 
 # Against a standard vector, assume already padded and just apply the stencil
-function convolve_interior!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator) where {T<:Real}
+function convolve_interior!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator; overwrite = true) where {T<:Real}
     @assert length(x_temp)+2 == length(x)
     stencil = A.stencil_coefs
     coeff   = A.coefficients
@@ -21,11 +21,11 @@ function convolve_interior!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::
         for idx in 1:A.stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[i - mid + idx]
         end
-        x_temp[i] = xtempi
+        x_temp[i] = xtempi + !overwrite*x_temp[i]
     end
 end
 
-function convolve_BC_left!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator) where {T<:Real}
+function convolve_BC_left!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator; overwrite = true) where {T<:Real}
     stencil = A.low_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
@@ -36,11 +36,11 @@ function convolve_BC_left!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::D
         for idx in 2:A.boundary_stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[idx]
         end
-        x_temp[i] = xtempi
+        x_temp[i] = xtempi + !overwrite*x_temp[i]
     end
 end
 
-function convolve_BC_right!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator) where {T<:Real}
+function convolve_BC_right!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator; overwrite = true) where {T<:Real}
     stencil = A.high_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
@@ -51,14 +51,14 @@ function convolve_BC_right!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::
         for idx in (A.boundary_stencil_length-1):-1:1
             xtempi += cur_coeff * cur_stencil[end-idx] * x[end-idx]
         end
-        x_temp[end-A.boundary_point_count+i] = xtempi
+        x_temp[end-A.boundary_point_count+i] = xtempi + !overwrite*x_temp[end-A.boundary_point_count+i]
     end
 end
 
 ###########################################
 
 # Against A BC-padded vector, specialize the computation to explicitly use the left, right, and middle parts
-function convolve_interior!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator) where {T<:Real}
+function convolve_interior!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator; overwrite = true) where {T<:Real}
     stencil = A.stencil_coefs
     coeff   = A.coefficients
     x = _x.u
@@ -72,11 +72,11 @@ function convolve_interior!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
         @inbounds for idx in 1:A.stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * x[(i-1) - (mid-idx) + 1]
         end
-        x_temp[i] = xtempi
+        x_temp[i] = xtempi + !overwrite*x_temp[i]
     end
 end
 
-function convolve_BC_left!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator) where {T<:Real}
+function convolve_BC_left!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator; overwrite = true) where {T<:Real}
     stencil = A.low_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
@@ -87,7 +87,7 @@ function convolve_BC_left!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
         @inbounds for idx in 2:A.boundary_stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * _x.u[idx-1]
         end
-        x_temp[i] = xtempi
+        x_temp[i] = xtempi + !overwrite*x_temp[i]
     end
     # need to account for x.l in first interior
     mid = div(A.stencil_length,2) + 1
@@ -101,10 +101,10 @@ function convolve_BC_left!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
     @inbounds for idx in 2:A.stencil_length
         xtempi += cur_coeff * cur_stencil[idx] * x[(i-1) - (mid-idx) + 1]
     end
-    x_temp[i] = xtempi
+    x_temp[i] = xtempi + !overwrite*x_temp[i]
 end
 
-function convolve_BC_right!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator) where {T<:Real}
+function convolve_BC_right!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator; overwrite = true) where {T<:Real}
     stencil = A.high_boundary_coefs
     coeff   = A.coefficients
     bc_start = length(_x.u) - A.boundary_point_count
@@ -120,7 +120,7 @@ function convolve_BC_right!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
     @inbounds for idx in 1:A.stencil_length-1
         xtempi += cur_coeff * cur_stencil[idx] * x[(i-1) - (mid-idx) + 1]
     end
-    x_temp[i] = xtempi
+    x_temp[i] = xtempi + !overwrite*x_temp[i]
     for i in 1 : A.boundary_point_count
         cur_stencil = stencil[i]
         cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[bc_start + i] : coeff isa Number ? coeff : true
@@ -129,6 +129,25 @@ function convolve_BC_right!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
         @inbounds for idx in A.stencil_length:-1:1
             xtempi += cur_coeff * cur_stencil[end-idx] * _x.u[end-idx+1]
         end
-        x_temp[bc_start + i] = xtempi
+        x_temp[bc_start + i] = xtempi + !overwrite*x_temp[bc_start + i]
+    end
+end
+
+###########################################
+
+function convolve_interior_add_range!(x_temp::AbstractVector{T}, x::AbstractVector{T}, A::DerivativeOperator, offset::Int) where {T<:Real}
+    @assert length(x_temp)+2 == length(x)
+    stencil = A.stencil_coefs
+    coeff   = A.coefficients
+    mid = div(A.stencil_length,2)
+    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] : 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]
+        end
+        x_temp[i] += xtempi
     end
 end

src/derivative_operators/derivative_operator.jl

@@ -166,3 +166,4 @@ end
 CenteredDifference(args...) = CenteredDifference{1}(args...)
 UpwindDifference(args...) = UpwindDifference{1}(args...)
 use_winding(A::DerivativeOperator{T,N,Wind}) where {T,N,Wind} = Wind
+diff_axis(A::DerivativeOperator{T,N}) where {T,N} = N