More compact show for block sparse arrays (#42)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.2.15"
+version = "0.2.16"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockSparseArrays.jl

@@ -25,9 +25,9 @@ include("blocksparsearrayinterface/cat.jl")
 
 # functions defined for any abstractblocksparsearray
 include("abstractblocksparsearray/abstractblocksparsearray.jl")
-include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
 include("abstractblocksparsearray/abstractblocksparsematrix.jl")
 include("abstractblocksparsearray/abstractblocksparsevector.jl")
+include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
 include("abstractblocksparsearray/views.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/sparsearrayinterface.jl")

src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -77,3 +77,42 @@ function Base.setindex!(
   blocks(a)[Int.(I)...] = value
   return a
 end
+
+using TypeParameterAccessors: unspecify_type_parameters
+function show_typeof_blocksparse(io::IO, a::AbstractBlockSparseArray)
+  Base.show(io, unspecify_type_parameters(typeof(a)))
+  print(io, '{')
+  show(io, eltype(a))
+  print(io, ", ")
+  show(io, ndims(a))
+  print(io, ", ")
+  show(io, blocktype(a))
+  print(io, ", …")
+  print(io, '}')
+  return nothing
+end
+
+# Copied from `BlockArrays.jl`.
+block2string(b, s) = string(join(map(string, b), '×'), "-blocked ", Base.dims2string(s))
+
+function summary_blocksparse(io::IO, a::AbstractArray)
+  print(io, block2string(blocksize(a), size(a)))
+  print(io, ' ')
+  show_typeof_blocksparse(io, a)
+  return nothing
+end
+
+function Base.summary(io::IO, a::AbstractBlockSparseArray)
+  summary_blocksparse(io, a)
+  return nothing
+end
+
+function Base.showarg(io::IO, a::AbstractBlockSparseArray, toplevel::Bool)
+  if toplevel
+    show_typeof_blocksparse(io, a)
+  else
+    print(io, "::")
+    show_typeof_blocksparse(io, a)
+  end
+  return nothing
+end

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -10,6 +10,7 @@ using BlockArrays:
   mortar,
   unblock
 using DerivableInterfaces: DerivableInterfaces, @interface
+using GPUArraysCore: @allowscalar
 using SplitApplyCombine: groupcount
 using TypeParameterAccessors: similartype
 
@@ -390,6 +391,6 @@ function Base.show(io::IO, mime::MIME"text/plain", a::AnyAbstractBlockSparseArra
   print(io, ":")
   println(io)
   a′ = ReplacedUnstoredBlockSparseArray(a, GetUnstoredBlockShow(axes(a)))
-  Base.print_array(io, a′)
+  @allowscalar Base.print_array(io, a′)
   return nothing
 end

src/blocksparsearray/blocksparsearray.jl

@@ -199,3 +199,40 @@ blockstype(arraytype::Type{<:BlockSparseArray}) = SparseArrayDOK{AbstractArray}
 ##   # TODO: Preserve GPU data!
 ##   return BlockSparseArray{elt}(undef, axes)
 ## end
+
+# TypeParameterAccessors.jl interface
+using TypeParameterAccessors: TypeParameterAccessors, Position, set_type_parameters
+TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(eltype)) = Position(1)
+TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(ndims)) = Position(2)
+TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blocktype)) = Position(3)
+function TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blockstype))
+  return Position(4)
+end
+
+# TODO: Make this generic to `AbstractBlockSparseVector` using
+# TypeParameterAccessors.jl, for example using:
+# `set_ndims(unspecify_type_parameters(typeof(a)), 1)`.
+function show_typeof_blocksparse(io::IO, a::BlockSparseVector)
+  print(io, "BlockSparseVector")
+  print(io, '{')
+  show(io, eltype(a))
+  print(io, ", ")
+  show(io, blocktype(a))
+  print(io, ", …")
+  print(io, '}')
+  return nothing
+end
+
+# TODO: Make this generic to `AbstractBlockSparseMatrix` using
+# TypeParameterAccessors.jl, for example using:
+# `set_ndims(unspecify_type_parameters(typeof(a)), 2)`.
+function show_typeof_blocksparse(io::IO, a::BlockSparseMatrix)
+  print(io, "BlockSparseMatrix")
+  print(io, '{')
+  show(io, eltype(a))
+  print(io, ", ")
+  show(io, blocktype(a))
+  print(io, ", …")
+  print(io, '}')
+  return nothing
+end

test/Project.toml

@@ -19,3 +19,4 @@ SymmetrySectors = "f8a8ad64-adbc-4fce-92f7-ffe2bb36a86e"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"

test/test_basics.jl

@@ -38,6 +38,7 @@ using SparseArraysBase: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, stored
 using TensorAlgebra: contract
 using Test: @test, @test_broken, @test_throws, @testset, @inferred
 using TestExtras: @constinferred
+using TypeParameterAccessors: TypeParameterAccessors, Position
 include("TestBlockSparseArraysUtils.jl")
 
 arrayts = (Array, JLArray)
@@ -1084,13 +1085,53 @@ arrayts = (Array, JLArray)
     @test storedlength(b) == 17
   end
   @testset "show" begin
+    vectort_elt = arrayt{elt,1}
+    matrixt_elt = arrayt{elt,2}
+    arrayt_elt = arrayt{elt,3}
+
+    a = BlockSparseVector{elt,arrayt{elt,1}}([2, 2])
+    # Either option is possible depending on namespacing.
+    @test (
+      sprint(summary, a) ==
+      "2-blocked 4-element BlockSparseVector{$(elt), $(vectort_elt), …}"
+    ) || (
+      sprint(summary, a) ==
+      "2-blocked 4-element BlockSparseArrays.BlockSparseVector{$(elt), $(vectort_elt), …}"
+    )
+
+    a = BlockSparseMatrix{elt,arrayt{elt,2}}([2, 2], [2, 2])
+    # Either option is possible depending on namespacing.
+    @test (
+      sprint(summary, a) == "2×2-blocked 4×4 BlockSparseMatrix{$(elt), $(matrixt_elt), …}"
+    ) || (
+      sprint(summary, a) ==
+      "2×2-blocked 4×4 BlockSparseArrays.BlockSparseMatrix{$(elt), $(matrixt_elt), …}"
+    )
+
+    a = BlockSparseArray{elt,3,arrayt{elt,3}}([2, 2], [2, 2], [2, 2])
+
+    # Either option is possible depending on namespacing.
+    @test (
+      sprint(summary, a) ==
+      "2×2×2-blocked 4×4×4 BlockSparseArray{$(elt), 3, $(arrayt_elt), …}"
+    ) || (
+      sprint(summary, a) ==
+      "2×2×2-blocked 4×4×4 BlockSparseArrays.BlockSparseArray{$(elt), 3, $(arrayt_elt), …}"
+    )
+
     if elt === Float64
       # Not testing other element types since they change the
       # spacing so it isn't easy to make the test general.
-      a = BlockSparseArray{elt}([2, 2], [2, 2])
-      a[1, 2] = 12
+      a = BlockSparseMatrix{elt,arrayt{elt,2}}([2, 2], [2, 2])
+      @allowscalar a[1, 2] = 12
       @test sprint(show, "text/plain", a) ==
         "$(summary(a)):\n $(zero(eltype(a)))  $(eltype(a)(12))  │  .  .\n $(zero(eltype(a)))   $(zero(eltype(a)))  │  .  .\n ───────────┼──────\n  .     .   │  .  .\n  .     .   │  .  ."
     end
   end
+  @testset "TypeParameterAccessors.position" begin
+    @test TypeParameterAccessors.position(BlockSparseArray, eltype) == Position(1)
+    @test TypeParameterAccessors.position(BlockSparseArray, ndims) == Position(2)
+    @test TypeParameterAccessors.position(BlockSparseArray, blocktype) == Position(3)
+    @test TypeParameterAccessors.position(BlockSparseArray, blockstype) == Position(4)
+  end
 end