[WIP] More matrix functions, broadcasting, constructors

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "DiagonalArrays"
 uuid = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.3.17"
+version = "0.3.18"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/abstractdiagonalarray/abstractdiagonalarray.jl

@@ -4,6 +4,30 @@ abstract type AbstractDiagonalArray{T,N} <: AbstractSparseArray{T,N} end
 const AbstractDiagonalMatrix{T} = AbstractDiagonalArray{T,2}
 const AbstractDiagonalVector{T} = AbstractDiagonalArray{T,1}
 
+# Define for type stability, for some reason the generic versions
+# in SparseArraysBase.jl is not type stable.
+# TODO: Investigate type stability of `iszero` in SparseArraysBase.jl.
+function Base.iszero(a::AbstractDiagonalArray)
+  return iszero(diagview(a))
+end
+
+using FillArrays: AbstractFill, getindex_value
+using LinearAlgebra: norm
+# TODO: `_norm` works around:
+# https://github.com/JuliaArrays/FillArrays.jl/issues/417
+# Change back to `norm` when that is fixed.
+_norm(a, p::Int=2) = norm(a, p)
+function _norm(a::AbstractFill, p::Int=2)
+  nrm1 = norm(getindex_value(a))
+  return (length(a))^(1/oftype(nrm1, p)) * nrm1
+end
+function LinearAlgebra.norm(a::AbstractDiagonalArray, p::Int=2)
+  # TODO: `_norm` works around:
+  # https://github.com/JuliaArrays/FillArrays.jl/issues/417
+  # Change back to `norm` when that is fixed.
+  return _norm(diagview(a), p)
+end
+
 using LinearAlgebra: LinearAlgebra, ishermitian, isposdef, issymmetric
 LinearAlgebra.ishermitian(a::AbstractDiagonalMatrix{<:Real}) = issquare(a)
 function LinearAlgebra.ishermitian(a::AbstractDiagonalMatrix{<:Number})

src/abstractdiagonalarray/diagonalarraydiaginterface.jl

@@ -99,25 +99,3 @@ function Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:DiagonalArrayStyle}
   copyto!(diagview(dest), broadcasted_diagview(bc))
   return dest
 end
-
-## SparseArraysBase.StorageIndex(i::DiagIndex) = StorageIndex(index(i))
-
-## function Base.getindex(a::AbstractDiagonalArray, i::DiagIndex)
-##   return a[StorageIndex(i)]
-## end
-
-## function Base.setindex!(a::AbstractDiagonalArray, value, i::DiagIndex)
-##   a[StorageIndex(i)] = value
-##   return a
-## end
-
-## SparseArraysBase.StorageIndices(i::DiagIndices) = StorageIndices(indices(i))
-
-## function Base.getindex(a::AbstractDiagonalArray, i::DiagIndices)
-##   return a[StorageIndices(i)]
-## end
-
-## function Base.setindex!(a::AbstractDiagonalArray, value, i::DiagIndices)
-##   a[StorageIndices(i)] = value
-##   return a
-## end

src/diaginterface/diaginterface.jl

@@ -97,6 +97,7 @@ function setdiagindex!(a::AbstractArray, v, i::Integer)
 end
 
 function getdiagindices(a::AbstractArray, I)
+  # TODO: Should this be a view?
   return @view diagview(a)[I]
 end
 

src/diagonalarray/diagonalarray.jl

@@ -20,10 +20,58 @@ SparseArraysBase.unstored(a::DiagonalArray) = a.unstored
 Base.size(a::DiagonalArray) = size(unstored(a))
 Base.axes(a::DiagonalArray) = axes(unstored(a))
 
+function DiagonalArray(diag::AbstractVector, unstored::Unstored)
+  return _DiagonalArray(diag, parent(unstored))
+end
 function DiagonalArray(::UndefInitializer, unstored::Unstored)
-  return _DiagonalArray(
-    Vector{eltype(unstored)}(undef, minimum(size(unstored))), parent(unstored)
-  )
+  return DiagonalArray(Vector{eltype(unstored)}(undef, minimum(size(unstored))), unstored)
+end
+
+# This helps to support diagonals where the elements are known
+# from the types, for example diagonals that are `Zeros` and `Ones`.
+function DiagonalArray{T,N,D,U}(
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}
+) where {T,N,D<:AbstractVector{T},U<:AbstractArray{T,N}}
+  return DiagonalArray(D((Base.OneTo(minimum(length, ax)),)), Unstored(U(ax)))
+end
+function DiagonalArray{T,N,D,U}(
+  ax1::AbstractUnitRange{<:Integer}, ax_rest::Vararg{AbstractUnitRange{<:Integer}}
+) where {T,N,D<:AbstractVector{T},U<:AbstractArray{T,N}}
+  return DiagonalArray{T,N,D,U}((ax1, ax_rest...))
+end
+function DiagonalArray{T,N,D,U}(
+  sz::Tuple{Integer,Vararg{AbstractUnitRange{<:Integer}}}
+) where {T,N,D<:AbstractVector{T},U<:AbstractArray{T,N}}
+  return DiagonalArray{T,N,D,U}(Base.OneTo.(sz))
+end
+function DiagonalArray{T,N,D,U}(
+  sz1::Integer, sz_rest::Vararg{Integer}
+) where {T,N,D<:AbstractVector{T},U<:AbstractArray{T,N}}
+  return DiagonalArray{T,N,D,U}((sz1, sz_rest...))
+end
+
+# This helps to support diagonals where the elements are known
+# from the types, for example diagonals that are `Zeros` and `Ones`.
+# These versions use the default unstored type `Zeros{T,N}`.
+function DiagonalArray{T,N,D}(
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}
+) where {T,N,D<:AbstractVector{T}}
+  return DiagonalArray{T,N,D,Zeros{T,N}}(ax)
+end
+function DiagonalArray{T,N,D}(
+  ax1::AbstractUnitRange{<:Integer}, ax_rest::Vararg{AbstractUnitRange{<:Integer}}
+) where {T,N,D<:AbstractVector{T}}
+  return DiagonalArray{T,N,D,Zeros{T,N}}(ax1, ax_rest...)
+end
+function DiagonalArray{T,N,D}(
+  sz::Tuple{Integer,Vararg{AbstractUnitRange{<:Integer}}}
+) where {T,N,D<:AbstractVector{T}}
+  return DiagonalArray{T,N,D,Zeros{T,N}}(sz)
+end
+function DiagonalArray{T,N,D}(
+  sz1::Integer, sz_rest::Vararg{Integer}
+) where {T,N,D<:AbstractVector{T}}
+  return DiagonalArray{T,N,D,Zeros{T,N}}(sz1, sz_rest...)
 end
 
 # Constructors accepting axes.
@@ -32,7 +80,7 @@ function DiagonalArray{T,N}(
   ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 ) where {T,N}
   N == length(ax) || throw(ArgumentError("Wrong number of axes"))
-  return _DiagonalArray(convert(AbstractVector{T}, diag), Zeros{T}(ax))
+  return DiagonalArray(convert(AbstractVector{T}, diag), Unstored(Zeros{T}(ax)))
 end
 function DiagonalArray{T,N}(
   diag::AbstractVector,
@@ -97,7 +145,7 @@ function DiagonalArray{T}(
 end
 
 function DiagonalArray{T,N}(diag::AbstractVector, dims::Dims{N}) where {T,N}
-  return _DiagonalArray(convert(AbstractVector{T}, diag), Zeros{T}(dims))
+  return DiagonalArray(convert(AbstractVector{T}, diag), Unstored(Zeros{T}(dims)))
 end
 
 function DiagonalArray{T,N}(diag::AbstractVector, dims::Vararg{Int,N}) where {T,N}
@@ -161,6 +209,28 @@ function DiagonalArray{T}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N}
   return DiagonalArray{T,N}(undef, dims)
 end
 
+# 0-dim limit.
+function DiagonalArray{T,0,D}(
+  ::UndefInitializer, ax::Tuple{}
+) where {T,D<:AbstractVector{T}}
+  return DiagonalArray{T,0,D}(D(undef, 0), ax)
+end
+function DiagonalArray{T,0,D}(::UndefInitializer) where {T,D<:AbstractVector{T}}
+  return DiagonalArray{T,0,D}(undef, ())
+end
+function DiagonalArray{T,0}(::UndefInitializer, ax::Tuple{}) where {T}
+  return DiagonalArray{T,0,Vector{T}}(undef, ax)
+end
+function DiagonalArray{T,0}(::UndefInitializer) where {T}
+  return DiagonalArray{T,0}(undef, ())
+end
+function DiagonalArray{T}(::UndefInitializer, axes::Tuple{}) where {T}
+  return DiagonalArray{T,0}(undef, ())
+end
+function DiagonalArray{T}(::UndefInitializer) where {T}
+  return DiagonalArray{T}(undef, ())
+end
+
 # Axes version
 function DiagonalArray{T}(::UndefInitializer, axes::NTuple{N,Base.OneTo{Int}}) where {T,N}
   return DiagonalArray{T,N}(undef, length.(axes))
@@ -197,3 +267,109 @@ function DerivableInterfaces.permuteddims(a::DiagonalArray, perm)
   # Unlike `permutedims(::Diagonal, perm)`, we copy here.
   return DiagonalArray(diagview(a), ax_perm)
 end
+
+# Scalar indexing.
+using DerivableInterfaces: @interface, interface
+one_based_range(r) = false
+one_based_range(r::Base.OneTo) = true
+one_based_range(r::Base.Slice) = true
+function _diag_axes(a::DiagonalArray, I...)
+  return map(ntuple(identity, ndims(a))) do d
+    return Base.axes1(axes(a, d)[I[d]])
+  end
+end
+# A view that preserves the diagonal structure.
+function _view_diag(a::DiagonalArray, I...)
+  ax = _diag_axes(a, I...)
+  return DiagonalArray(view(diagview(a), Base.OneTo(minimum(length, I))), ax)
+end
+# A slice that preserves the diagonal structure.
+function _getindex_diag(a::DiagonalArray, I...)
+  ax = _diag_axes(a, I...)
+  return DiagonalArray(diagview(a)[Base.OneTo(minimum(length, I))], ax)
+end
+function Base.view(a::DiagonalArray, I...)
+  I′ = to_indices(a, I)
+  return if all(one_based_range, I′)
+    _view_diag(a, I′...)
+  else
+    invoke(view, Tuple{AbstractArray,Vararg}, a, I′...)
+  end
+end
+function Base.getindex(a::DiagonalArray, I::Int...)
+  return @interface interface(a) a[I...]
+end
+function Base.getindex(a::DiagonalArray, I::DiagIndex)
+  return getdiagindex(a, index(I))
+end
+function Base.getindex(a::DiagonalArray, I::DiagIndices)
+  # TODO: Should this be a view?
+  return @view diagview(a)[indices(I)]
+end
+function Base.getindex(a::DiagonalArray, I...)
+  I′ = to_indices(a, I)
+  return if all(i -> i isa Real, I′)
+    # Catch scalar indexing case.
+    @interface interface(a) a[I...]
+  elseif all(one_based_range, I′)
+    _getindex_diag(a, I′...)
+  else
+    copy(view(a, I′...))
+  end
+end
+
+# Define in order to preserve immutable diagonals such as FillArrays.
+function DiagonalArray{T,N}(a::DiagonalArray{T,N}) where {T,N}
+  # TODO: Should this copy? This matches the design of `LinearAlgebra.Diagonal`:
+  # https://github.com/JuliaLang/LinearAlgebra.jl/blob/release-1.12/src/diagonal.jl#L110-L112
+  return a
+end
+function DiagonalArray{T,N}(a::DiagonalArray{<:Any,N}) where {T,N}
+  return DiagonalArray{T,N}(diagview(a))
+end
+function DiagonalArray{T}(a::DiagonalArray) where {T}
+  return DiagonalArray{T,ndims(a)}(a)
+end
+function DiagonalArray(a::DiagonalArray)
+  return DiagonalArray{eltype(a),ndims(a)}(a)
+end
+function Base.AbstractArray{T,N}(a::DiagonalArray{<:Any,N}) where {T,N}
+  return DiagonalArray{T,N}(a)
+end
+
+# TODO: These definitions work around this issue:
+# https://github.com/JuliaArrays/FillArrays.jl/issues/416
+# when the diagonal is a FillArrays.Ones or Zeros.
+using Base.Broadcast: Broadcast, broadcast, broadcasted
+using FillArrays: AbstractFill, Ones, Zeros
+_broadcasted(f::F, a::AbstractArray) where {F} = broadcasted(f, a)
+_broadcasted(::typeof(identity), a::Ones) = a
+_broadcasted(::typeof(identity), a::Zeros) = a
+_broadcasted(::typeof(complex), a::Ones) = Ones{complex(eltype(a))}(axes(a))
+_broadcasted(::typeof(complex), a::Zeros) = Zeros{complex(eltype(a))}(axes(a))
+_broadcasted(elt::Type, a::Ones) = Ones{elt}(axes(a))
+_broadcasted(elt::Type, a::Zeros) = Zeros{elt}(axes(a))
+_broadcasted(::typeof(inv), a::Ones) = _broadcasted(typeof(inv(oneunit(eltype(a)))), a)
+using LinearAlgebra: pinv
+_broadcasted(::typeof(pinv), a::Ones) = _broadcasted(typeof(inv(oneunit(eltype(a)))), a)
+_broadcasted(::typeof(sqrt), a::Ones) = _broadcasted(typeof(sqrt(one(eltype(a)))), a)
+_broadcasted(::typeof(sqrt), a::Zeros) = _broadcasted(typeof(sqrt(zero(eltype(a)))), a)
+_broadcasted(::typeof(cbrt), a::Ones) = _broadcasted(typeof(cbrt(one(eltype(a)))), a)
+_broadcasted(::typeof(cbrt), a::Zeros) = _broadcasted(typeof(cbrt(zero(eltype(a)))), a)
+_broadcasted(::typeof(exp), a::Zeros) = Ones{typeof(exp(zero(eltype(a))))}(axes(a))
+_broadcasted(::typeof(cis), a::Zeros) = Ones{typeof(cis(zero(eltype(a))))}(axes(a))
+_broadcasted(::typeof(log), a::Ones) = Zeros{typeof(log(one(eltype(a))))}(axes(a))
+_broadcasted(::typeof(cos), a::Zeros) = Ones{typeof(cos(zero(eltype(a))))}(axes(a))
+_broadcasted(::typeof(sin), a::Zeros) = _broadcasted(typeof(sin(zero(eltype(a)))), a)
+_broadcasted(::typeof(tan), a::Zeros) = _broadcasted(typeof(tan(zero(eltype(a)))), a)
+_broadcasted(::typeof(sec), a::Zeros) = Ones{typeof(sec(zero(eltype(a))))}(axes(a))
+_broadcasted(::typeof(cosh), a::Zeros) = Ones{typeof(cosh(zero(eltype(a))))}(axes(a))
+# Eager version of `_broadcasted`.
+_broadcast(f::F, a::AbstractArray) where {F} = copy(_broadcasted(f, a))
+
+function Broadcast.broadcasted(
+  ::DiagonalArrayStyle{N}, f::F, a::DiagonalArray{T,N,Diag}
+) where {F,T,N,Diag<:AbstractFill{T}}
+  # TODO: Check that `f` preserves zeros?
+  return DiagonalArray(_broadcasted(f, diagview(a)), axes(a))
+end

src/diagonalarray/diagonalmatrix.jl

@@ -58,3 +58,91 @@ function LinearAlgebra.mul!(
   d_dest .= d1 .* d2 .* α .+ d_dest .* β
   return a_dest
 end
+
+# Adapted from https://github.com/JuliaLang/LinearAlgebra.jl/blob/release-1.12/src/diagonal.jl#L866-L928.
+function LinearAlgebra.tr(a::DiagonalMatrix)
+  checksquare(a)
+  # TODO: Define as `sum(tr, diagview(a))` like LinearAlgebra.jl?
+  return sum(diagview(a))
+end
+# TODO: Special case for FillArrays diagonals.
+function LinearAlgebra.det(a::DiagonalMatrix)
+  checksquare(a)
+  # TODO: Define as `prod(det, diagview(a))` like LinearAlgebra.jl?
+  return prod(diagview(a))
+end
+# TODO: Special case for FillArrays diagonals.
+function LinearAlgebra.logabsdet(a::DiagonalMatrix)
+  checksquare(a)
+  return mapreduce(((d1, s1), (d2, s2)) -> (d1 + d2, s1 * s2), diagview(a)) do x
+    return (log(abs(x)), sign(x))
+  end
+end
+# TODO: Special case for FillArrays diagonals.
+function LinearAlgebra.logdet(a::DiagonalMatrix{<:Complex})
+  checksquare(a)
+  z = sum(log, diagview(a))
+  return complex(real(z), rem2pi(imag(z), RoundNearest))
+end
+
+# Matrix functions
+for f in (
+  :exp,
+  :cis,
+  :log,
+  :sqrt,
+  :cos,
+  :sin,
+  :tan,
+  :csc,
+  :sec,
+  :cot,
+  :cosh,
+  :sinh,
+  :tanh,
+  :csch,
+  :sech,
+  :coth,
+  :acos,
+  :asin,
+  :atan,
+  :acsc,
+  :asec,
+  :acot,
+  :acosh,
+  :asinh,
+  :atanh,
+  :acsch,
+  :asech,
+  :acoth,
+)
+  @eval begin
+    function Base.$f(a::DiagonalMatrix)
+      checksquare(a)
+      return DiagonalMatrix(_broadcast($f, diagview(a)), axes(a))
+    end
+  end
+end
+
+# Cube root of a real-valued diagonal matrix
+function Base.cbrt(a::DiagonalMatrix{<:Real})
+  checksquare(a)
+  return DiagonalMatrix(_broadcast(cbrt, diagview(a)), axes(a))
+end
+
+function LinearAlgebra.inv(a::DiagonalMatrix)
+  checksquare(a)
+  # `DiagonalArrays._broadcast` works around issues like https://github.com/JuliaArrays/FillArrays.jl/issues/416
+  # when the diagonal is a FillArray or similar lazy array.
+  d⁻¹ = _broadcast(inv, diagview(a))
+  any(isinf, d⁻¹) && error("Singular Exception")
+  return DiagonalMatrix(d⁻¹, axes(a))
+end
+
+# TODO: Support `atol` and `rtol` keyword arguments:
+# https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.pinv
+using LinearAlgebra: pinv
+function LinearAlgebra.pinv(a::DiagonalMatrix)
+  checksquare(a)
+  return DiagonalMatrix(_broadcast(pinv, diagview(a)), axes(a))
+end

src/dual.jl

@@ -1,3 +1,11 @@
 # TODO: Define `TensorProducts.dual`.
 dual(x) = x
 issquare(a::AbstractMatrix) = (axes(a, 1) == dual(axes(a, 2)))
+# Like `LinearAlgebra.checksquare` but based on `DiagonalArrays.issquare`,
+# which checks the axes and allows customizing to check that the
+# codomain is the dual of the domain.
+# Returns the codomain if the check passes.
+function checksquare(a::AbstractMatrix)
+  issquare(a) || throw(DimensionMismatch(lazy"matrix is not square: axes are $(axes(a))"))
+  return axes(a, 1)
+end