Merge pull request #3 from JuliaDiffEq/master

Rebase
SciML · Jul 3, 2019 · aa3ed39 · aa3ed39
2 parents 9e3111a + 7213994
commit aa3ed39
Show file tree

Hide file tree

Showing 17 changed files with 927 additions and 105 deletions.
diff --git a/CITATION.bib b/CITATION.bib
@@ -1,7 +1,7 @@
 @article{DifferentialEquations.jl-2017,
  author = {Rackauckas, Christopher and Nie, Qing},
  doi = {10.5334/jors.151},
- journal = {},
+ journal = {The Journal of Open Source Software},
  keywords = {Applied Mathematics},
  note = {Exported from https://app.dimensions.ai on 2019/05/05},
  number = {1},

diff --git a/Project.toml b/Project.toml
@@ -11,6 +11,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+
 
 [compat]
 julia = "1"

diff --git a/src/DiffEqOperators.jl b/src/DiffEqOperators.jl
@@ -4,7 +4,7 @@ import Base: +, -, *, /, \, size, getindex, setindex!, Matrix, convert
 using DiffEqBase, StaticArrays, LinearAlgebra
 import LinearAlgebra: mul!, ldiv!, lmul!, rmul!, axpy!, opnorm, factorize, I
 import DiffEqBase: AbstractDiffEqLinearOperator, update_coefficients!, is_constant
-using SparseArrays, ForwardDiff, BandedMatrices
+using SparseArrays, ForwardDiff, BandedMatrices, NNlib
 
 abstract type AbstractDerivativeOperator{T} <: AbstractDiffEqLinearOperator{T} end
 abstract type AbstractDiffEqCompositeOperator{T} <: AbstractDiffEqLinearOperator{T} end
@@ -29,6 +29,9 @@ include("derivative_operators/derivative_operator.jl")
 include("derivative_operators/abstract_operator_functions.jl")
 include("derivative_operators/convolutions.jl")
 include("derivative_operators/concretization.jl")
+include("derivative_operators/ghost_derivative_operator.jl")
+include("derivative_operators/derivative_operator_functions.jl")
+
 
 ### Composite Operators
 include("composite_operators.jl")
@@ -45,4 +48,5 @@ export DiffEqScalar, DiffEqArrayOperator, DiffEqIdentity, JacVecOperator, getops
 export AbstractDerivativeOperator, DerivativeOperator,
        CenteredDifference, UpwindDifference
 export RobinBC, GeneralBC
+export GhostDerivativeOperator
 end # module
diff --git a/src/common_defaults.jl b/src/common_defaults.jl
@@ -8,24 +8,24 @@ convert(::Type{AbstractArray}, L::AbstractDiffEqLinearOperator) = convert(Abstra
 size(L::AbstractDiffEqLinearOperator, args...) = size(convert(AbstractMatrix,L), args...)
 opnorm(L::AbstractDiffEqLinearOperator, p::Real=2) = opnorm(convert(AbstractMatrix,L), p)
 getindex(L::AbstractDiffEqLinearOperator, i::Int) = convert(AbstractMatrix,L)[i]
-getindex(L::AbstractDiffEqLinearOperator, I::Vararg{Int, N}) where {N} = 
+getindex(L::AbstractDiffEqLinearOperator, I::Vararg{Int, N}) where {N} =
   convert(AbstractMatrix,L)[I...]
 for op in (:*, :/, :\)
-  @eval $op(L::AbstractDiffEqLinearOperator, x::Union{AbstractVecOrMat,Number}) = $op(convert(AbstractMatrix,L), x)
+  @eval $op(L::AbstractDiffEqLinearOperator, x::Union{AbstractArray,Number}) = $op(convert(AbstractMatrix,L), x)
   @eval $op(x::Union{AbstractVecOrMat,Number}, L::AbstractDiffEqLinearOperator) = $op(x, convert(AbstractMatrix,L))
 end
-mul!(Y::AbstractVecOrMat, L::AbstractDiffEqLinearOperator, B::AbstractVecOrMat) =
+mul!(Y::AbstractArray, L::AbstractDiffEqLinearOperator, B::AbstractArray) =
   mul!(Y, convert(AbstractMatrix,L), B)
 ldiv!(Y::AbstractVecOrMat, L::AbstractDiffEqLinearOperator, B::AbstractVecOrMat) =
   ldiv!(Y, convert(AbstractMatrix,L), B)
 for pred in (:isreal, :issymmetric, :ishermitian, :isposdef)
   @eval LinearAlgebra.$pred(L::AbstractDiffEqLinearOperator) = $pred(convert(AbstractArray, L))
 end
-factorize(L::AbstractDiffEqLinearOperator) = 
+factorize(L::AbstractDiffEqLinearOperator) =
   FactorizedDiffEqArrayOperator(factorize(convert(AbstractMatrix, L)))
-for fact in (:lu, :lu!, :qr, :qr!, :cholesky, :cholesky!, :ldlt, :ldlt!, 
+for fact in (:lu, :lu!, :qr, :qr!, :cholesky, :cholesky!, :ldlt, :ldlt!,
   :bunchkaufman, :bunchkaufman!, :lq, :lq!, :svd, :svd!)
-  @eval LinearAlgebra.$fact(L::AbstractDiffEqLinearOperator, args...) = 
+  @eval LinearAlgebra.$fact(L::AbstractDiffEqLinearOperator, args...) =
     FactorizedDiffEqArrayOperator($fact(convert(AbstractMatrix, L), args...))
 end
 

diff --git a/src/derivative_operators/abstract_operator_functions.jl b/src/derivative_operators/abstract_operator_functions.jl
@@ -48,15 +48,32 @@ end
 @inline getindex(A::AbstractDerivativeOperator, ::Colon, ::Colon) = Array(A)
 
 @inline function getindex(A::AbstractDerivativeOperator, ::Colon, j)
-    return BandedMatrix(A)[:,j]
+    T = eltype(A.stencil_coefs)
+    v = zeros(T, A.len)
+    v[j] = one(T)
+    copyto!(v, A*v)
+    return v
 end
 
-
-# symmetric right now
 @inline function getindex(A::AbstractDerivativeOperator, i, ::Colon)
-    return BandedMatrix(A)[i,:]
-end
+    @boundscheck checkbounds(A, i, 1)
+    T = eltype(A.stencil_coefs)
+    v = zeros(T, A.len+2)
+
+    bpc = A.boundary_point_count
+    N = A.len
+    bsl = A.boundary_stencil_length
+    slen = A.stencil_length
 
+    if bpc > 0 && 1<=i<=bpc
+        v[1:bsl] .= A.low_boundary_coefs[i]
+    elseif bpc > 0 && (N-bpc)<i<=N
+         v[1:bsl]  .= A.high_boundary_coefs[i-(N-1)]
+    else
+        v[i-bpc:i-bpc+slen-1] .= A.stencil_coefs
+    end
+    return v
+end
 
 # UnitRanges
 @inline function getindex(A::AbstractDerivativeOperator, rng::UnitRange{Int}, ::Colon)
@@ -86,22 +103,6 @@ end
     return BandedMatrix(A)[rng,cng]
 end
 
-#=
-    This definition of the mul! function makes it possible to apply the LinearOperator on
-    a matrix and not just a vector. It basically transforms the rows one at a time.
-=#
-function LinearAlgebra.mul!(x_temp::AbstractArray{T,2}, A::AbstractDerivativeOperator{T}, M::AbstractMatrix{T}) where T<:Real
-    if size(x_temp) == reverse(size(M))
-        for i = 1:size(M,1)
-            mul!(view(x_temp,i,:), A, view(M,i,:))
-        end
-    else
-        for i = 1:size(M,2)
-            mul!(view(x_temp,:,i), A, view(M,:,i))
-        end
-    end
-end
-
 # Base.length(A::AbstractDerivativeOperator) = A.stencil_length
 Base.ndims(A::AbstractDerivativeOperator) = 2
 Base.size(A::AbstractDerivativeOperator) = (A.len, A.len + 2)
@@ -140,20 +141,6 @@ end
 
 get_type(::AbstractDerivativeOperator{T}) where {T} = T
 
-function *(A::AbstractDerivativeOperator,x::AbstractVector)
-    y = zeros(promote_type(eltype(A),eltype(x)), length(x)-2)
-    LinearAlgebra.mul!(y, A::AbstractDerivativeOperator, x::AbstractVector)
-    return y
-end
-
-
-function *(A::AbstractDerivativeOperator,M::AbstractMatrix)
-    y = zeros(promote_type(eltype(A),eltype(M)), size(A,1), size(M,2))
-    LinearAlgebra.mul!(y, A::AbstractDerivativeOperator, M::AbstractMatrix)
-    return y
-end
-
-
 function *(M::AbstractMatrix,A::AbstractDerivativeOperator)
     y = zeros(promote_type(eltype(A),eltype(M)), size(M,1), size(A,2))
     LinearAlgebra.mul!(y, M, BandedMatrix(A))

diff --git a/src/derivative_operators/convolutions.jl b/src/derivative_operators/convolutions.jl
@@ -65,13 +65,13 @@ function convolve_interior!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
     x = _x.u
     mid = div(A.stencil_length,2) + 1
     # Just do the middle parts
-    for i in (1+A.boundary_point_count) : (length(x_temp)-A.boundary_point_count)
+    for i in (2+A.boundary_point_count) : (length(x_temp)-A.boundary_point_count)-1
         xtempi = zero(T)
         cur_stencil = eltype(stencil) <: AbstractVector ? stencil[i-A.boundary_point_count] : stencil
         cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i-A.boundary_point_count] : true
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
         @inbounds for idx in 1:A.stencil_length
-            xtempi += cur_coeff * cur_stencil[idx] * x[i - (mid-idx) + 1]
+            xtempi += cur_coeff * cur_stencil[idx] * x[(i-1) - (mid-idx) + 1]
         end
         x_temp[i] = xtempi
     end
@@ -81,29 +81,55 @@ function convolve_BC_left!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector,
     stencil = A.low_boundary_coefs
     coeff   = A.coefficients
     for i in 1 : A.boundary_point_count
-        xtempi = stencil[i][1]*_x.l
         cur_stencil = stencil[i]
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i-A.boundary_point_count] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i] : true
+        xtempi = cur_coeff*cur_stencil[1]*_x.l
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
-        @inbounds for idx in 2:A.stencil_length
+        @inbounds for idx in 2:A.boundary_stencil_length
             xtempi += cur_coeff * cur_stencil[idx] * _x.u[idx-1]
         end
         x_temp[i] = xtempi
     end
+    # need to account for x.l in first interior
+    mid = div(A.stencil_length,2) + 1
+    x = _x.u
+    i = 1 + A.boundary_point_count
+    xtempi = zero(T)
+    cur_stencil = eltype(A.stencil_coefs) <: AbstractVector ? A.stencil_coefs[i-A.boundary_point_count] : A.stencil_coefs
+    cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i-A.boundary_point_count] : true
+    cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
+    xtempi = cur_coeff*cur_stencil[1]*_x.l
+    @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
 end
 
 function convolve_BC_right!(x_temp::AbstractVector{T}, _x::BoundaryPaddedVector, A::DerivativeOperator) where {T<:Real}
-    stencil = A.low_boundary_coefs
+    stencil = A.high_boundary_coefs
     coeff   = A.coefficients
-    bc_start = length(_x.u) - A.stencil_length
+    bc_start = length(_x.u) - A.boundary_point_count
+    # need to account for _x.r in last interior convolution
+    mid = div(A.stencil_length,2) + 1
+    x = _x.u
+    i = length(x_temp)-A.boundary_point_count
+    xtempi = zero(T)
+    cur_stencil = eltype(A.stencil_coefs) <: AbstractVector ? A.stencil_coefs[i-A.boundary_point_count] : A.stencil_coefs
+    cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i-A.boundary_point_count] : true
+    cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
+    xtempi = cur_coeff*cur_stencil[end]*_x.r
+    @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
     for i in 1 : A.boundary_point_count
-        xtempi = stencil[i][end]*_x.r
         cur_stencil = stencil[i]
-        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[i-A.boundary_point_count] : true
+        cur_coeff   = typeof(coeff)   <: AbstractVector ? coeff[bc_start + i] : true
+        xtempi = cur_coeff*cur_stencil[end]*_x.r
         cur_stencil = use_winding(A) && cur_coeff < 0 ? reverse(cur_stencil) : cur_stencil
-        @inbounds for idx in A.stencil_length:-1:2
+        @inbounds for idx in A.stencil_length:-1:1
             xtempi += cur_coeff * cur_stencil[end-idx] * _x.u[end-idx+1]
         end
-        x_temp[i] = xtempi
+        x_temp[bc_start + i] = xtempi
     end
 end
diff --git a/src/derivative_operators/derivative_operator.jl b/src/derivative_operators/derivative_operator.jl
@@ -1,3 +1,5 @@
+index(i::Int, N::Int) = i + div(N, 2) + 1
+
 struct DerivativeOperator{T<:Real,N,Wind,T2,S1,S2<:SVector,T3,F} <: AbstractDerivativeOperator{T}
     derivative_order        :: Int
     approximation_order     :: Int
@@ -52,14 +54,30 @@ function CenteredDifference{N}(derivative_order::Int,
 
     stencil_length          = derivative_order + approximation_order - 1 + (derivative_order+approximation_order)%2
     boundary_stencil_length = derivative_order + approximation_order
-    dummy_x                 = -div(stencil_length,2) : div(stencil_length,2)
-    boundary_x              = -boundary_stencil_length+1:0
+    stencil_x               = zeros(T, stencil_length)
     boundary_point_count    = div(stencil_length,2) - 1 # -1 due to the ghost point
+
+    interior_x              = boundary_point_count+2:len+1-boundary_point_count
+    dummy_x                 = -div(stencil_length,2) : div(stencil_length,2)-1
+    boundary_x              = -boundary_stencil_length+1:0
+
     # Because it's a N x (N+2) operator, the last stencil on the sides are the [b,0,x,x,x,x] stencils, not the [0,x,x,x,x,x] stencils, since we're never solving for the derivative at the boundary point.
     deriv_spots             = (-div(stencil_length,2)+1) : -1
     boundary_deriv_spots    = boundary_x[2:div(stencil_length,2)]
 
-    stencil_coefs           = [convert(SVector{stencil_length, T}, calculate_weights(derivative_order, zero(T), dummy_x)) for i in 1:len(dx)]
+    function generate_coordinates(i, stencil_x, dummy_x, dx)
+        len = length(stencil_x)
+        stencil_x .= stencil_x.*zero(T)
+        for idx in 1:div(len,2)
+            shifted_idx1 = index(idx, len)
+            shifted_idx2 = index(-idx, len)
+            stencil_x[shifted_idx1] = stencil_x[shifted_idx1-1] + dx[i+idx-1]
+            stencil_x[shifted_idx2] = stencil_x[shifted_idx2+1] - dx[i-idx]
+        end
+        return stencil_x
+    end
+
+    stencil_coefs           = convert(SVector{length(interior_x)}, [convert(SVector{stencil_length, T}, calculate_weights(derivative_order, zero(T), generate_coordinates(i, stencil_x, dummy_x, dx))) for i in interior_x])
     _low_boundary_coefs     = SVector{boundary_stencil_length, T}[convert(SVector{boundary_stencil_length, T}, calculate_weights(derivative_order, oneunit(T)*x0, boundary_x)) for x0 in boundary_deriv_spots]
     low_boundary_coefs      = convert(SVector{boundary_point_count},_low_boundary_coefs)
     high_boundary_coefs     = convert(SVector{boundary_point_count},reverse(SVector{boundary_stencil_length, T}[reverse(low_boundary_coefs[i]) for i in 1:boundary_point_count]))
@@ -75,7 +93,7 @@ function CenteredDifference{N}(derivative_order::Int,
         boundary_stencil_length,
         boundary_point_count,
         low_boundary_coefs,
-        high_boundary_coefs,coefficients,coeff_func,false
+        high_boundary_coefs,coefficients,coeff_func
         )
 end
 

diff --git a/src/derivative_operators/derivative_operator_functions.jl b/src/derivative_operators/derivative_operator_functions.jl
@@ -0,0 +1,97 @@
+function LinearAlgebra.mul!(x_temp::AbstractArray{T}, A::DerivativeOperator{T,N}, M::AbstractArray{T}) where {T<:Real,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)
+        mul!(view(x_temp, idx...), A, view(M, idx...))
+    end
+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}}
+
+            # 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
+
+            # Respahe x_temp for NNlib.conv!
+            new_size = Any[size(x_temp)...]
+            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))
+
+            # 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!
+            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)
+
+            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)
+            end
+        end
+    end
+end
+
+function *(A::DerivativeOperator{T,N},M::AbstractArray{T}) where {T<:Real,N}
+    size_x_temp = [size(M)...]
+    size_x_temp[N] -= 2
+    x_temp = zeros(promote_type(eltype(A),eltype(M)), size_x_temp...)
+    LinearAlgebra.mul!(x_temp, A, M)
+    return x_temp
+end