Construct DiagonalArray and delta from axes (#32)

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.9"
+version = "0.3.10"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -15,5 +15,5 @@ ArrayLayouts = "1.10.4"
 DerivableInterfaces = "0.5"
 FillArrays = "1.13.0"
 LinearAlgebra = "1.10.0"
-SparseArraysBase = "0.7.1"
+SparseArraysBase = "0.7.2"
 julia = "1.10"

src/diagonalarray/delta.jl

@@ -1,5 +1,39 @@
 using FillArrays: Ones
 
+function delta(
+  elt::Type, ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}
+)
+  return DiagonalArray(Ones{elt}(minimum(length, ax)), ax)
+end
+function δ(
+  elt::Type, ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}
+)
+  return delta(elt, ax)
+end
+function delta(ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}})
+  return delta(Float64, ax)
+end
+function δ(ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}})
+  return delta(Float64, ax)
+end
+
+function delta(
+  elt::Type, ax1::AbstractUnitRange{<:Integer}, axs::AbstractUnitRange{<:Integer}...
+)
+  return delta(elt, (ax1, axs...))
+end
+function δ(
+  elt::Type, ax1::AbstractUnitRange{<:Integer}, axs::AbstractUnitRange{<:Integer}...
+)
+  return delta(elt, (ax1, axs...))
+end
+function delta(ax1::AbstractUnitRange{<:Integer}, axs::AbstractUnitRange{<:Integer}...)
+  return delta(Float64, (ax1, axs...))
+end
+function δ(ax1::AbstractUnitRange{<:Integer}, axs::AbstractUnitRange{<:Integer}...)
+  return delta(Float64, (ax1, axs...))
+end
+
 function delta(elt::Type, size::Tuple{Vararg{Int}})
   return DiagonalArray(Ones{elt}(minimum(size)), size)
 end

src/diagonalarray/diagonalarray.jl

@@ -21,11 +21,79 @@ Base.size(a::DiagonalArray) = size(unstored(a))
 Base.axes(a::DiagonalArray) = axes(unstored(a))
 
 function DiagonalArray(::UndefInitializer, unstored::Unstored)
-  return _DiagonalArray(Vector{eltype(unstored)}(undef, ndims(unstored)), parent(unstored))
-end
-
-function DiagonalArray{T,N}(diag::AbstractVector, unstored::AbstractArray) where {T,N}
-  return _DiagonalArray(convert(AbstractVector{T}, diag), dims, getunstored)
+  return _DiagonalArray(
+    Vector{eltype(unstored)}(undef, minimum(size(unstored))), parent(unstored)
+  )
+end
+
+# Constructors accepting axes.
+function DiagonalArray{T,N}(
+  diag::AbstractVector,
+  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))
+end
+function DiagonalArray{T,N}(
+  diag::AbstractVector,
+  ax1::AbstractUnitRange{<:Integer},
+  axs::AbstractUnitRange{<:Integer}...,
+) where {T,N}
+  return DiagonalArray{T,N}(diag, (ax1, axs...))
+end
+function DiagonalArray{T}(
+  diag::AbstractVector,
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
+) where {T}
+  return DiagonalArray{T,length(ax)}(diag, ax)
+end
+function DiagonalArray{T}(
+  diag::AbstractVector,
+  ax1::AbstractUnitRange{<:Integer},
+  axs::AbstractUnitRange{<:Integer}...,
+) where {T}
+  return DiagonalArray{T}(diag, (ax1, axs...))
+end
+function DiagonalArray(
+  diag::AbstractVector{T},
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
+) where {T}
+  return DiagonalArray{T,length(ax)}(diag, ax)
+end
+function DiagonalArray(
+  diag::AbstractVector,
+  ax1::AbstractUnitRange{<:Integer},
+  axs::AbstractUnitRange{<:Integer}...,
+)
+  return DiagonalArray(diag, (ax1, axs...))
+end
+
+# undef constructors accepting axes.
+function DiagonalArray{T,N}(
+  ::UndefInitializer,
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
+) where {T,N}
+  return DiagonalArray{T,N}(Vector{T}(undef, minimum(length, ax)), ax)
+end
+function DiagonalArray{T,N}(
+  ::UndefInitializer,
+  ax1::AbstractUnitRange{<:Integer},
+  axs::AbstractUnitRange{<:Integer}...,
+) where {T,N}
+  return DiagonalArray{T,N}(undef, (ax1, axs...))
+end
+function DiagonalArray{T}(
+  ::UndefInitializer,
+  ax::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
+) where {T}
+  return DiagonalArray{T,length(ax)}(undef, ax)
+end
+function DiagonalArray{T}(
+  ::UndefInitializer,
+  ax1::AbstractUnitRange{<:Integer},
+  axs::AbstractUnitRange{<:Integer}...,
+) where {T}
+  return DiagonalArray{T}(undef, (ax1, axs...))
 end
 
 function DiagonalArray{T,N}(diag::AbstractVector, dims::Dims{N}) where {T,N}

test/test_basics.jl

@@ -69,6 +69,38 @@ using LinearAlgebra: Diagonal
       @test diagindices(IndexCartesian(), a) ==
         CartesianIndex.(Iterators.zip(1:3, 1:3, 1:3))
     end
+    @testset "DiagonalArray constructors" begin
+      v = randn(elt, 2)
+      @test DiagonalArray(v, 2, 2) ≡
+        DiagonalArray(v, (2, 2)) ≡
+        DiagonalArray(v, Base.OneTo(2), Base.OneTo(2)) ≡
+        DiagonalArray(v, (Base.OneTo(2), Base.OneTo(2))) ≡
+        DiagonalArray{elt}(v, 2, 2) ≡
+        DiagonalArray{elt}(v, (2, 2)) ≡
+        DiagonalArray{elt}(v, Base.OneTo(2), Base.OneTo(2)) ≡
+        DiagonalArray{elt}(v, (Base.OneTo(2), Base.OneTo(2))) ≡
+        DiagonalArray{elt,2}(v, 2, 2) ≡
+        DiagonalArray{elt,2}(v, (2, 2)) ≡
+        DiagonalArray{elt,2}(v, Base.OneTo(2), Base.OneTo(2)) ≡
+        DiagonalArray{elt,2}(v, (Base.OneTo(2), Base.OneTo(2)))
+      @test size(DiagonalArray{elt}(undef, 2, 2)) ≡
+        size(DiagonalArray{elt}(undef, (2, 2))) ≡
+        size(DiagonalArray{elt}(undef, Base.OneTo(2), Base.OneTo(2))) ≡
+        size(DiagonalArray{elt}(undef, (Base.OneTo(2), Base.OneTo(2)))) ≡
+        size(DiagonalArray{elt,2}(undef, 2, 2)) ≡
+        size(DiagonalArray{elt,2}(undef, (2, 2))) ≡
+        size(DiagonalArray{elt,2}(undef, Base.OneTo(2), Base.OneTo(2))) ≡
+        size(DiagonalArray{elt,2}(undef, (Base.OneTo(2), Base.OneTo(2))))
+      @test elt ≡
+        eltype(DiagonalArray{elt}(undef, 2, 2)) ≡
+        eltype(DiagonalArray{elt}(undef, (2, 2))) ≡
+        eltype(DiagonalArray{elt}(undef, Base.OneTo(2), Base.OneTo(2))) ≡
+        eltype(DiagonalArray{elt}(undef, (Base.OneTo(2), Base.OneTo(2)))) ≡
+        eltype(DiagonalArray{elt,2}(undef, 2, 2)) ≡
+        eltype(DiagonalArray{elt,2}(undef, (2, 2))) ≡
+        eltype(DiagonalArray{elt,2}(undef, Base.OneTo(2), Base.OneTo(2))) ≡
+        eltype(DiagonalArray{elt,2}(undef, (Base.OneTo(2), Base.OneTo(2))))
+    end
     @testset "Matrix multiplication" begin
       a1 = DiagonalArray{elt}(undef, (2, 3))
       a1[1, 1] = 11
@@ -120,13 +152,21 @@ using LinearAlgebra: Diagonal
     @testset "delta" begin
       for (a, elt′) in (
         (delta(2, 3), Float64),
+        (delta(Base.OneTo(2), Base.OneTo(3)), Float64),
         (δ(2, 3), Float64),
+        (δ(Base.OneTo(2), Base.OneTo(3)), Float64),
         (delta((2, 3)), Float64),
+        (delta(Base.OneTo.((2, 3))), Float64),
         (δ((2, 3)), Float64),
+        (δ(Base.OneTo.((2, 3))), Float64),
         (delta(Bool, 2, 3), Bool),
+        (delta(Bool, Base.OneTo(2), Base.OneTo(3)), Bool),
         (δ(Bool, 2, 3), Bool),
+        (δ(Bool, Base.OneTo(2), Base.OneTo(3)), Bool),
         (delta(Bool, (2, 3)), Bool),
+        (delta(Bool, Base.OneTo.((2, 3))), Bool),
         (δ(Bool, (2, 3)), Bool),
+        (δ(Bool, Base.OneTo.((2, 3))), Bool),
       )
         @test eltype(a) === elt′
         @test diaglength(a) == 2