Implement oneelement in SparseArraysBaseExt

Author: mtfishman

Project.toml

@@ -16,9 +16,11 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
+SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 
 [extensions]
 ITensorBaseDiagonalArraysExt = "DiagonalArrays"
+ITensorBaseSparseArraysBaseExt = ["NamedDimsArrays", "SparseArraysBase"]
 
 [compat]
 Accessors = "0.1.39"
@@ -28,6 +30,7 @@ FillArrays = "1.13.0"
 LinearAlgebra = "1.10"
 MapBroadcast = "0.1.5"
 NamedDimsArrays = "0.4"
+SparseArraysBase = "0.2.10"
 UnallocatedArrays = "0.1.1"
 UnspecifiedTypes = "0.1.1"
 VectorInterface = "0.5.0"

ext/ITensorBaseSparseArraysBaseExt/ITensorBaseSparseArraysBaseExt.jl

@@ -0,0 +1,13 @@
+module ITensorBaseSparseArraysBaseExt
+
+using ITensorBase: ITensor, Index
+using NamedDimsArrays: dename
+using SparseArraysBase: SparseArraysBase, oneelement
+
+function SparseArraysBase.oneelement(
+  value, index::NTuple{N,Int}, ax::NTuple{N,Index}
+) where {N}
+  return ITensor(oneelement(value, index, only.(axes.(dename.(ax)))), ax)
+end
+
+end

src/abstractitensor.jl

@@ -91,24 +91,6 @@ function ITensor(parent::AbstractArray)
   return ITensor(parent, ())
 end
 
-# TODO:
-# 1. Generalize this to arbitrary dimensions.
-# 2. Define a basic non-ITensor version in `SparseArraysBase.jl`,
-#    and this constructor can wrap that one. It could construct
-#    a `OneElement` sparse object like `FillArrays.OneElement`
-#    (https://juliaarrays.github.io/FillArrays.jl/stable/#FillArrays.OneElement).
-# 3. Define `oneelement(value, index::Tuple{Vargarg{Int}}, axes::Tuple{Vararg{Index}})`,
-#    where the pair version calls out to that one.
-function oneelement(elt::Type{<:Number}, iv::Pair{<:Index,<:Int})
-  a = ITensor(first(iv))
-  a[last(iv)] = one(elt)
-  return a
-end
-# TODO: The non-ITensor version should default to `Float64`,
-# like `FillArrays.OneElement`.
-# TODO: Make the element type `UnspecifiedOne`.
-oneelement(iv::Pair{<:Index,<:Int}) = oneelement(Bool, iv)
-
 using Accessors: @set
 setdenamed(a::ITensor, denamed) = (@set a.parent = denamed)
 setdenamed!(a::ITensor, denamed) = (a.parent = denamed)

test/Project.toml

@@ -5,6 +5,7 @@ ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

test/test_basics.jl

@@ -5,14 +5,14 @@ using ITensorBase:
   gettag,
   hastag,
   inds,
-  oneelement,
   plev,
   prime,
   settag,
   tags,
   unsettag
 using DiagonalArrays: δ, delta, diagview
 using NamedDimsArrays: dename, name, named
+using SparseArraysBase: oneelement
 using Test: @test, @test_broken, @testset
 
 @testset "ITensorBase" begin
@@ -46,6 +46,24 @@ using Test: @test, @test_broken, @testset
     @test plev(i) == 0
     @test length(tags(i)) == 0
   end
+  @testset "delta" begin
+    i, j = Index.((2, 2))
+    for a in (
+      delta(i, j),
+      delta(Bool, i, j),
+      delta((i, j)),
+      delta(Bool, (i, j)),
+      δ(i, j),
+      δ(Bool, i, j),
+      δ((i, j)),
+      δ(Bool, (i, j)),
+    )
+      @test eltype(a) === Bool
+      # TODO: Fix this.
+      @test_broken diagview(a)
+      @test diagview(dename(a)) == ones(2)
+    end
+  end
   @testset "oneelement" begin
     i = Index(3)
     a = oneelement(i => 2)
@@ -67,22 +85,4 @@ using Test: @test, @test_broken, @testset
     @test a[2] == 1
     @test a[3] == 0
   end
-  @testset "delta" begin
-    i, j = Index.((2, 2))
-    for a in (
-      delta(i, j),
-      delta(Bool, i, j),
-      delta((i, j)),
-      delta(Bool, (i, j)),
-      δ(i, j),
-      δ(Bool, i, j),
-      δ((i, j)),
-      δ(Bool, (i, j)),
-    )
-      @test eltype(a) === Bool
-      # TODO: Fix this.
-      @test_broken diagview(a)
-      @test diagview(dename(a)) == ones(2)
-    end
-  end
 end