Skip to content

Commit

Permalink
Improve operators
Browse files Browse the repository at this point in the history
  • Loading branch information
@dlfivefifty
dlfivefifty committed Jan 28, 2025
1 parent d4641e0 commit 74dc23f
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 45 deletions.
21 changes: 18 additions & 3 deletions src/bases/bases.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -662,8 +662,8 @@ cumsum_layout(::ExpansionLayout, A, dims) = cumsum_layout(ApplyLayout{typeof(*)}
###
# diff
###
diff_layout(::AbstractBasisLayout, Vm; dims...) = error("Overload diff(::$(typeof(Vm)))")
function diff_layout(::AbstractBasisLayout, a, order; dims...)
diff_layout(::AbstractBasisLayout, Vm, order...; dims...) = error("Overload diff(::$(typeof(Vm)))")
function diff_layout(::AbstractBasisLayout, a, order::Int; dims...)
order < 0 && throw(ArgumentError("order must be non-negative"))
order == 0 && return a
isone(order) ? diff(a) : diff(diff(a), order-1)
Expand Down Expand Up @@ -730,9 +730,11 @@ abslaplacian_axis(::Inclusion{<:Number}, A, order=1; dims...) = -diff(A, 2order;

laplacian(A, order...; dims...) = laplacian_layout(MemoryLayout(A), A, order...; dims...)
laplacian_layout(layout, A, order...; dims...) = laplacian_axis(axes(A,1), A, order...; dims...)
laplacian_axis(::Inclusion{<:Number}, A, order...; dims...) = -abslaplacian(A, order...)
laplacian_axis(::Inclusion{<:Number}, A, order=1; dims...) = diff(A, 2order; dims...)


laplacian_layout(::ExpansionLayout, A, order...; dims...) = laplacian_layout(ApplyLayout{typeof(*)}(), A, order...; dims...)
abslaplacian_layout(::ExpansionLayout, A, order...; dims...) = abslaplacian_layout(ApplyLayout{typeof(*)}(), A, order...; dims...)

function abslaplacian_layout(::SubBasisLayout, Vm, order...; dims::Integer=1)
dims == 1 || error("not implemented")
Expand All @@ -744,6 +746,19 @@ function laplacian_layout(::SubBasisLayout, Vm, order...; dims::Integer=1)
laplacian(parent(Vm), order...)[:,parentindices(Vm)[2]]
end

function laplacian_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVecOrMat, order...; dims=1)
a = arguments(LAY, V)
dims == 1 || throw(ArgumentError("cannot take laplacian a vector along dimension $dims"))
*(laplacian(a[1], order...), tail(a)...)
end

function abslaplacian_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVecOrMat, order...; dims=1)
a = arguments(LAY, V)
dims == 1 || throw(ArgumentError("cannot take abslaplacian a vector along dimension $dims"))
*(abslaplacian(a[1], order...), tail(a)...)
end



"""
weaklaplacian(A)
Expand Down
88 changes: 46 additions & 42 deletions src/operators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,13 +108,15 @@ show(io::IO, ::MIME"text/plain", δ::DiracDelta) = show(io, δ)
# Differentiation
#########

abstract type AbstractDifferentialQuasiMatrix{T} <: LazyQuasiMatrix{T} end

"""
Derivative(axis)
represents the differentiation (or finite-differences) operator on the
specified axis.
"""
struct Derivative{T,D<:Inclusion,Order} <: LazyQuasiMatrix{T}
struct Derivative{T,D<:Inclusion,Order} <: AbstractDifferentialQuasiMatrix{T}
axis::D
order::Order
end
Expand All @@ -125,7 +127,7 @@ Laplacian(axis)
represents the laplacian operator `Δ` on the
specified axis.
"""
struct Laplacian{T,D<:Inclusion,Order} <: LazyQuasiMatrix{T}
struct Laplacian{T,D<:Inclusion,Order} <: AbstractDifferentialQuasiMatrix{T}
axis::D
order::Order
end
Expand All @@ -136,14 +138,42 @@ AbsLaplacian(axis)
represents the positive-definite/negative/absolute-value laplacian operator `|Δ| ≡ -Δ` on the
specified axis.
"""
struct AbsLaplacian{T,D<:Inclusion,Order} <: LazyQuasiMatrix{T}
struct AbsLaplacian{T,D<:Inclusion,Order} <: AbstractDifferentialQuasiMatrix{T}
axis::D
order::Order
end

_operatororder(D) = something(D.order, 1)
operatororder(D) = something(D.order, 1)

show(io::IO, a::AbstractDifferentialQuasiMatrix) = summary(io, a)
axes(D::AbstractDifferentialQuasiMatrix) = (D.axis, D.axis)
==(a::AbstractDifferentialQuasiMatrix, b::AbstractDifferentialQuasiMatrix) = a.axis == b.axis && operatororder(a) == operatororder(b)
copy(D::AbstractDifferentialQuasiMatrix) = D



@simplify function *(D::AbstractDifferentialQuasiMatrix, B::AbstractQuasiVecOrMat)
T = typeof(zero(eltype(D)) * zero(eltype(B)))
operatorcall(D, convert(AbstractQuasiArray{T}, B), D.order)
end

^(D::AbstractDifferentialQuasiMatrix{T}, k::Integer) where T = similaroperator(D, D.axis, operatororder(D) .* k)

function view(D::AbstractDifferentialQuasiMatrix, kr::Inclusion, jr::Inclusion)
@boundscheck axes(D,1) == kr == jr || throw(BoundsError(D,(kr,jr)))
D
end

operatorcall(D::AbstractDifferentialQuasiMatrix, B, order) = operatorcall(D)(B, order)
operatorcall(D::AbstractDifferentialQuasiMatrix, B, ::Nothing) = operatorcall(D)(B)


operatorcall(::Derivative) = diff
operatorcall(::Laplacian) = laplacian
operatorcall(::AbsLaplacian) = abslaplacian

for (Op, op) in ((:Derivative, :diff), (:Laplacian, :laplacian), (:AbsLaplacian, :abslaplacian))

for Op in (:Derivative, :Laplacian, :AbsLaplacian)
@eval begin
$Op{T, D}(axis::D, order) where {T,D<:Inclusion} = $Op{T,D,typeof(order)}(axis, order)
$Op{T, D}(axis::D) where {T,D<:Inclusion} = $Op{T,D,Nothing}(axis, nothing)
Expand All @@ -153,42 +183,20 @@ for (Op, op) in ((:Derivative, :diff), (:Laplacian, :laplacian), (:AbsLaplacian,
$Op(ax::AbstractQuasiVector{T}, order...) where T = $Op{eltype(eltype(ax))}(ax, order...)
$Op(L::AbstractQuasiMatrix, order...) = $Op(axes(L,1), order...)

show(io::IO, a::$Op) = summary(io, a)
function summary(io::IO, D::$Op{<:Any,<:Inclusion,Nothing})
print(io, "$($Op)(")
summary(io, D.axis)
print(io,")")
end
similaroperator(D::$Op, ax, ord) = $Op{eltype(D)}(ax, ord)

simplifiable(::typeof(*), A::$Op, B::$Op) = Val(true)
*(D1::$Op, D2::$Op) = similaroperator(convert(AbstractQuasiMatrix{promote_type(eltype(D1),eltype(D2))}, D1), D1.axis, operatororder(D1)+operatororder(D2))


function summary(io::IO, D::$Op)
print(io, "$($Op)(")
summary(io, D.axis)
print(io, ", ")
print(io, D.order)
print(io,")")
end

axes(D::$Op) = (D.axis, D.axis)
==(a::$Op, b::$Op) = a.axis == b.axis && _operatororder(a) == _operatororder(b)
copy(D::$Op) = D


@simplify function *(D::$Op, B::AbstractQuasiVecOrMat)
T = typeof(zero(eltype(D)) * zero(eltype(B)))
if D.order isa Nothing
$op(convert(AbstractQuasiArray{T}, B))
else
$op(convert(AbstractQuasiArray{T}, B), D.order)
if !isnothing(D.order)
print(io, ", ")
print(io, D.order)
end
end

^(D::$Op{T}, k::Integer) where T = $Op{T}(D.axis, _operatororder(D) .* k)

@simplify *(D1::$Op, D2::$Op) = $Op{promote_type(eltype(D1),eltype(D2))}(D1.axis, _operatororder(D1)+_operatororder(D2))

function view(D::$Op, kr::Inclusion, jr::Inclusion)
@boundscheck axes(D,1) == kr == jr || throw(BoundsError(D,(kr,jr)))
D
print(io,")")
end
end
end
Expand All @@ -199,12 +207,8 @@ end
# end


const Identity{T,D} = QuasiDiagonal{T,Inclusion{T,D}}

Identity(d::Inclusion) = QuasiDiagonal(d)

struct OperatorLayout <: AbstractLazyLayout end
MemoryLayout(::Type{<:Derivative}) = OperatorLayout()
MemoryLayout(::Type{<:AbstractDifferentialQuasiMatrix}) = OperatorLayout()
# copy(M::Mul{OperatorLayout, <:ExpansionLayout}) = simplify(M)
# copy(M::Mul{OperatorLayout, <:AbstractLazyLayout}) = M.A * expand(M.B)

Expand All @@ -215,4 +219,4 @@ abs(Δ::Laplacian{T}) where T = AbsLaplacian{T}(axes(Δ,1), Δ.order)
-(Δ::Laplacian{<:Any,<:Any,Nothing}) = abs(Δ)
-(Δ::AbsLaplacian{T,<:Any,Nothing}) where T = Laplacian{T}(Δ.axis)

^(Δ::AbsLaplacian{T}, k::Real) where T = AbsLaplacian{T}(Δ.axis, _operatororder(Δ)*k)
^(Δ::AbsLaplacian{T}, k::Real) where T = AbsLaplacian{T}(Δ.axis, operatororder(Δ)*k)

0 comments on commit 74dc23f

Please sign in to comment.