Compact MPS/MPO printing (#166)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorMPS"
 uuid = "0d1a4710-d33b-49a5-8f18-73bdf49b47e2"
 authors = ["Matthew Fishman <[email protected]>", "Miles Stoudenmire <[email protected]>"]
-version = "0.3.24"
+version = "0.3.25"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/abstractmps.jl

@@ -2434,20 +2434,21 @@ end
 # Printing functions
 #
 
+# Used to print `#undef` for unassigned MPS tensors.
+# Maybe there's a better way to do this?
+struct Undef end
+Base.show(io::IO, ::Undef) = print(io, "#undef")
+
 function Base.show(io::IO, M::AbstractMPS)
-    print(io, "$(typeof(M))")
-    (length(M) > 0) && print(io, "\n")
-    for i in eachindex(M)
-        if !isassigned(M, i)
-            println(io, "#undef")
-        else
-            A = M[i]
-            if order(A) != 0
-                println(io, "[$i] $(inds(A))")
-            else
-                println(io, "[$i] ITensor()")
-            end
-        end
+    print(io, "$(typeof(M))($(length(M)))")
+    return nothing
+end
+function Base.show(io::IO, ::MIME"text/plain", M::AbstractMPS)
+    println(io, "$(length(M))-element $(typeof(M)):")
+    inds_vec = map(eachindex(M)) do j
+        !isassigned(M, j) && return Undef()
+        return inds(M[j])
     end
+    Base.print_array(io, inds_vec)
     return
 end

test/base/test_mpo.jl

@@ -39,9 +39,13 @@ end
     O = MPO(sites)
     @test length(O) == N
 
-    str = split(sprint(show, O), '\n')
-    @test endswith(str[1], "MPO")
-    @test length(str) == length(O) + 2
+    @test sprint(show, O) in ("MPO(6)", "ITensorMPS.MPO(6)")
+    strs = split(sprint(show, MIME"text/plain"(), O), "\n")
+    @test length(strs) == length(O) + 1
+    @test strs[1] in ("6-element MPO:", "6-element ITensorMPS.MPO:")
+    for str in strs[2:end]
+        @test startswith(str, " ((dim=")
+    end
 
     O[1] = ITensor(sites[1], prime(sites[1]))
     @test hasind(O[1], sites[1])

test/base/test_mps.jl

@@ -24,9 +24,13 @@ include(joinpath(@__DIR__, "utils", "util.jl"))
     @test linkdims(psi) == fill(3, length(psi) - 1)
     @test isnothing(flux(psi))
 
-    str = split(sprint(show, psi), '\n')
-    @test endswith(str[1], "MPS")
-    @test length(str) == length(psi) + 2
+    @test sprint(show, psi) in ("MPS(10)", "ITensorMPS.MPS(10)")
+    strs = split(sprint(show, MIME"text/plain"(), psi), "\n")
+    @test length(strs) == length(psi) + 1
+    @test strs[1] in ("10-element MPS:", "10-element ITensorMPS.MPS:")
+    for str in strs[2:end]
+        @test startswith(str, " ((dim=")
+    end
 
     @test siteind(psi, 2) == sites[2]
     @test findfirstsiteind(psi, sites[2]) == 2