[NDTensors] More compact tensor printing (#1680)

Author: mtfishman

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.4.15"
+version = "0.4.16"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/dense/densetensor.jl

@@ -319,5 +319,7 @@ end
 
 function Base.show(io::IO, mime::MIME"text/plain", T::DenseTensor)
     summary(io, T)
-    return print_tensor(io, T)
+    print_tensor(io, T)
+    println(io)
+    return nothing
 end

NDTensors/src/diag/diagtensor.jl

@@ -217,5 +217,6 @@ end
 function Base.show(io::IO, mime::MIME"text/plain", T::DiagTensor)
     summary(io, T)
     print_tensor(io, T)
+    println(io)
     return nothing
 end

NDTensors/src/tensor/tensor.jl

@@ -495,5 +495,9 @@ end
 # Printing
 #
 
-print_tensor(io::IO, T::Tensor) = Base.print_array(io, expose(T))
-print_tensor(io::IO, T::Tensor{<:Number, 1}) = Base.print_array(io, reshape(T, (dim(T), 1)))
+function print_tensor(io::IO, T::Tensor)
+    return Base.print_array(IOContext(io, :limit => true), expose(T))
+end
+function print_tensor(io::IO, T::Tensor{<:Number, 1})
+    return Base.print_array(IOContext(io, :limit => true), reshape(T, (dim(T), 1)))
+end

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensors"
 uuid = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 authors = ["Matthew Fishman <[email protected]>", "Miles Stoudenmire <[email protected]>"]
-version = "0.9.14"
+version = "0.9.15"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/itensor.jl

@@ -2025,8 +2025,16 @@ end
 # TODO: make a specialized printing from Diag
 # that emphasizes the missing elements
 function show(io::IO, T::ITensor)
-    println(io, "ITensor ord=$(order(T))")
-    return show(io, MIME"text/plain"(), tensor(T))
+    if get(io, :compact, false)
+        # Just show the indices in compact view, used in some
+        # cases when printing arrays of ITensors (similar to
+        # printing of MPS in ITensorMPS.jl).
+        show(io, inds(T))
+    else
+        println(io, "ITensor ord=$(order(T))")
+        show(io, MIME"text/plain"(), tensor(T))
+    end
+    return nothing
 end
 
 function show(io::IO, mime::MIME"text/plain", T::ITensor)

test/base/test_itensor.jl

@@ -7,6 +7,7 @@ using ITensors:
     ITensor,
     Order,
     QN,
+    TagSet,
     ⊕,
     δ,
     addtags,
@@ -102,6 +103,24 @@ end
             @test storage(A) isa NDTensors.EmptyStorage{NDTensors.EmptyNumber}
         end
 
+        @testset "show" begin
+            i = Index{Int}(1, 2, ITensors.Neither, TagSet(), 0)
+            a = ITensor([1, 2], i)
+            function res_show(prefix = "")
+                return "ITensor ord=1\nDim 1: (dim=2|id=1)\n" *
+                    "$(prefix)Dense{Float64, Vector{Float64}}\n" *
+                    " 2-element\n 1.0\n 2.0\n"
+            end
+            @test sprint(show, a) in (res_show(), res_show("NDTensors."))
+            @test sprint(show, a; context = :compact => true) == "((dim=2|id=1),)"
+            function res_show_text(prefix = "")
+                return "ITensor ord=1 (dim=2|id=1)\n" *
+                    "$(prefix)Dense{Float64, Vector{Float64}}"
+            end
+            @test sprint(show, MIME"text/plain"(), a) in
+                (res_show_text(), res_show_text("NDTensors."))
+        end
+
         @testset "diag" for ElType in (Float32, Float64, ComplexF32, ComplexF64)
             i, j = Index.(2, ("i", "j"))
             A = random_itensor(ElType, i, j)