improve show of confusion matrix by finding optimal column width (cw)

it is claulcated based on the maximal values and lebel width while ensuring the table header ("Gorubnd Truth") has enough width and the first column (firstcw) has enough space for "Predicted" header
JuliaAI · Aug 14, 2023 · 9ad78a3 · 9ad78a3
1 parent 7cea929
commit 9ad78a3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/src/measures/confusion_matrix.jl b/src/measures/confusion_matrix.jl
@@ -164,28 +164,31 @@ splitw(w::Int) = (sp1 = div(w, 2); sp2 = w - sp1; (sp1, sp2))
 
 function Base.show(stream::IO, m::MIME"text/plain", cm::ConfusionMatrixObject{C}
                    ) where C
+    mincw    = ceil(Int, 12/C)
     width    = displaysize(stream)[2]
-    cw       = 13
+    cw       = max(length(string(maximum(cm.mat))),maximum(length.(cm.labels)),mincw)
+    firstcw  = max(length(string(maximum(cm.mat))),maximum(length.(cm.labels)),9)
     textlim  = 9
     totalwidth = cw * (C+1) + C + 2
     width < totalwidth && (show(stream, m, cm.mat); return)
 
     iob     = IOBuffer()
     wline   = s -> write(iob, s * "\n")
     splitcw = s -> (w = cw - length(s); splitw(w))
+    splitfirstcw = s -> (w = firstcw - length(s); splitw(w))
     cropw   = s -> length(s) > textlim ? s[1:prevind(s, textlim)] * "…" : s
 
     # 1.a top box
-    " "^(cw+1) * "┌" * "─"^((cw + 1) * C - 1) * "┐" |> wline
+    " "^(firstcw+1) * "┌" * "─"^((cw + 1) * C - 1) * "┐" |> wline
     gt = "Ground Truth"
     w  = (cw + 1) * C - 1 - length(gt)
     sp1, sp2 = splitw(w)
-    " "^(cw+1) * "│" * " "^sp1 * gt * " "^sp2 * "│" |> wline
+    " "^(firstcw+1) * "│" * " "^sp1 * gt * " "^sp2 * "│" |> wline
     # 1.b separator
-    "┌" * "─"^cw * "┼" * ("─"^cw * "┬")^(C-1) * "─"^cw * "┤" |> wline
+    "┌" * "─"^firstcw * "┼" * ("─"^cw * "┬")^(C-1) * "─"^cw * "┤" |> wline
     # 2.a description line
     pr = "Predicted"
-    sp1, sp2 = splitcw(pr)
+    sp1, sp2 = splitfirstcw(pr)
     partial = "│" * " "^sp1 * pr * " "^sp2 * "│"
     for c in 1:C
         # max = 10
@@ -195,12 +198,12 @@ function Base.show(stream::IO, m::MIME"text/plain", cm::ConfusionMatrixObject{C}
     end
     partial |> wline
     # 2.b separating line
-    "├" * "─"^cw * "┼" * ("─"^cw * "┼")^(C-1) * ("─"^cw * "┤") |> wline
+    "├" * "─"^firstcw * "┼" * ("─"^cw * "┼")^(C-1) * ("─"^cw * "┤") |> wline
     # 2.c line by line
     for c in 1:C
         # line
         s  = cm.labels[c] |> cropw
-        sp1, sp2 = splitcw(s)
+        sp1, sp2 = splitfirstcw(s)
         partial = "│" * " "^sp1 * s * " "^sp2 * "│"
         for r in 1:C
             e = string(cm[c, r])
@@ -210,11 +213,11 @@ function Base.show(stream::IO, m::MIME"text/plain", cm::ConfusionMatrixObject{C}
         partial |> wline
         # separator
         if c < C
-            "├" * "─"^cw * "┼" * ("─"^cw * "┼")^(C-1) * ("─"^cw * "┤") |> wline
+            "├" * "─"^firstcw * "┼" * ("─"^cw * "┼")^(C-1) * ("─"^cw * "┤") |> wline
         end
     end
     # 2.d final line
-    "└" * "─"^cw * "┴" * ("─"^cw * "┴")^(C-1) * ("─"^cw * "┘") |> wline
+    "└" * "─"^firstcw * "┴" * ("─"^cw * "┴")^(C-1) * ("─"^cw * "┘") |> wline
     write(stream, take!(iob))
 end
 

diff --git a/test/measures/confusion_matrix.jl b/test/measures/confusion_matrix.jl
@@ -78,17 +78,17 @@ end
     Base.show(iob, MIME("text/plain"), MLJBase._confmat(ŷ, y))
     siob = String(take!(iob))
     @test strip(siob) == strip("""
-                         ┌─────────────────────────────────────────┐
-                         │              Ground Truth               │
-           ┌─────────────┼─────────────┬─────────────┬─────────────┤
-           │  Predicted  │      1      │      2      │      3      │
-           ├─────────────┼─────────────┼─────────────┼─────────────┤
-           │      1      │      3      │      0      │      0      │
-           ├─────────────┼─────────────┼─────────────┼─────────────┤
-           │      2      │      0      │      3      │      0      │
-           ├─────────────┼─────────────┼─────────────┼─────────────┤
-           │      3      │      0      │      0      │      3      │
-           └─────────────┴─────────────┴─────────────┴─────────────┘""")
+              ┌──────────────┐
+              │ Ground Truth │
+    ┌─────────┼────┬────┬────┤
+    │Predicted│ 1  │ 2  │ 3  │
+    ├─────────┼────┼────┼────┤
+    │    1    │ 3  │ 0  │ 0  │
+    ├─────────┼────┼────┼────┤
+    │    2    │ 0  │ 3  │ 0  │
+    ├─────────┼────┼────┼────┤
+    │    3    │ 0  │ 0  │ 3  │
+    └─────────┴────┴────┴────┘""")
 end
 
 @testset "ConfusionMatrix measure" begin