From 6b99aba15bee609b57e138965c611e09fb49a2a9 Mon Sep 17 00:00:00 2001
From: alexmul1114 <96952477+alexmul1114@users.noreply.github.com>
Date: Tue, 31 Oct 2023 16:29:44 -0400
Subject: [PATCH] Change plotting of mttkrp to use previous results

---
 benchmark/gcp.jl         | 28 +++++++++++++++++++++++-----
 benchmark/mttkrp.jl      | 19 -------------------
 benchmark/plot_mttkrp.jl | 24 +++++++++++++-----------
 3 files changed, 36 insertions(+), 35 deletions(-)
 delete mode 100644 benchmark/mttkrp.jl

diff --git a/benchmark/gcp.jl b/benchmark/gcp.jl
index cbb3c1d..68e01d0 100644
--- a/benchmark/gcp.jl
+++ b/benchmark/gcp.jl
@@ -6,14 +6,16 @@ using Distributions
 
 const SUITE = BenchmarkGroup()
 
-bench = SUITE["gcp"] = BenchmarkGroup()
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1
+
+bench_gcp = SUITE["gcp"] = BenchmarkGroup()
 
 # Benchmark least squares loss 
 for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
     Random.seed!(0)
     M = CPD(ones(r), rand.(sz, r))
     X = [M[I] for I in CartesianIndices(size(M))]
-    bench["least-squares-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, LeastSquaresLoss()) 
+    bench_gcp["least-squares-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, LeastSquaresLoss()) 
 end
 
 # Benchmark Poisson loss
@@ -21,7 +23,7 @@ for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
     Random.seed!(0)
     M = CPD(fill(10.0, r), rand.(sz, r))
     X = [rand(Poisson(M[I])) for I in CartesianIndices(size(M))]
-    bench["poisson-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, PoissonLoss())
+    bench_gcp["poisson-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, PoissonLoss())
 end
 
 # Benchmark Gamma loss
@@ -30,7 +32,7 @@ for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
     M = CPD(ones(r), rand.(sz, r))
     k = 1.5
     X = [rand(Gamma(k, M[I]/k)) for I in CartesianIndices(size(M))]
-    bench["gamma-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, GammaLoss())
+    bench_gcp["gamma-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, GammaLoss())
 end
 
 # Benchmark BernoulliOdds Loss
@@ -38,5 +40,21 @@ for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
     Random.seed!(0)
     M = CPD(ones(r), rand.(sz, r))
     X = [rand(Bernoulli(M[I]/(M[I] + 1))) for I in CartesianIndices(size(M))]
-    bench["bernoulliOdds-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, BernoulliOddsLoss())
+    bench_gcp["bernoulliOdds-size(X)=$sz, rank(X)=$r"] = @benchmarkable gcp($X, $r, BernoulliOddsLoss())
+end
+
+
+# MTTKRP benchmarks
+bench_mttkrp = SUITE["mttkrp"] = BenchmarkGroup()
+
+szs = [10,30,50,80,120,200]
+shapes = [(sz, sz, sz) for sz in szs]
+n = 1
+rs = 20:20:200
+
+for sz in shapes, r in rs
+    Random.seed!(0)
+    X = randn(sz)
+    U = [randn(Ik,r) for Ik in sz]
+    bench_mttkrp["mttkrp-size(X)=$sz, rank(X)=$r"] = @benchmarkable GCPDecompositions.mttkrp($X, $U, $n)
 end
\ No newline at end of file
diff --git a/benchmark/mttkrp.jl b/benchmark/mttkrp.jl
deleted file mode 100644
index 2b465a7..0000000
--- a/benchmark/mttkrp.jl
+++ /dev/null
@@ -1,19 +0,0 @@
-using BenchmarkTools
-
-using Random
-using GCPDecompositions
-
-const SUITE = BenchmarkGroup()
-
-bench = SUITE["mttkrp"] = BenchmarkGroup()
-
-szs = [(15, 20, 25), (30, 40, 50), (60, 70, 80), (80, 90, 100), (100, 110, 120)]
-for sz in szs, r in 1:7
-
-    Random.seed!(0)
-    X = randn(sz)
-    U = [randn(Ik,r) for Ik in sz]
-    n = 1
-
-    bench["mttkrp-size(X)=$sz, rank(X)=$r"] = @benchmarkable GCPDecompositions.mttkrp($X, $U, $n)
-end
\ No newline at end of file
diff --git a/benchmark/plot_mttkrp.jl b/benchmark/plot_mttkrp.jl
index ebc57ef..cdbfa7f 100644
--- a/benchmark/plot_mttkrp.jl
+++ b/benchmark/plot_mttkrp.jl
@@ -3,25 +3,27 @@ using GCPDecompositions
 using BenchmarkTools
 using UnicodePlots
 
-BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1
+# Load benchmark results from benchmark_results.txt
+benchmark_results = readresults("benchmark_results.txt")
+mttkrp_results = benchmark_results.benchmarkgroup["mttkrp"]
 
-szs = [15, 30, 45, 60, 75, 90]
+szs = [10,30,50,80,120,200]
 shapes = [(sz, sz, sz) for sz in szs]
-n = 1
-rs = 1:5
+rs = 20:20:200
 results = zeros((size(szs)[1], size(rs)[1]))
 
-for (idx, r) in enumerate(rs)
-    Random.seed!(0)
-    Xs = [randn(sz) for sz in shapes];
-    Us = [[randn(Ik,r) for Ik in sz] for sz in shapes];
-    times = [@belapsed GCPDecompositions.mttkrp($X, $U, $n) for (X,U) in zip(Xs, Us)] 
-    results[:, idx] = times
+# Collect results in array where columns are ranks and rows are sizes
+for (col_idx, r) in enumerate(rs)
+    for (row_idx, sz) in enumerate(szs)
+        # Get median runtime (in milliseconds)
+        median_time = median(mttkrp_results["mttkrp-size(X)=($sz, $sz, $sz), rank(X)=$r"]).time / 10^6
+        results[row_idx, col_idx] = median_time
+    end
 end
 
 first_r = rs[1]
 plt = lineplot(szs, results[:, 1], title="MTTKRP runtime vs. size", 
-            xlabel="Size", ylabel="Runtime(s)", name="r = $first_r");
+            xlabel="Size", ylabel="Runtime(ms)", name="r = $first_r")
 for (idx, r) in enumerate(rs[2:end])
     lineplot!(plt, szs, results[:, idx + 1], name="r = $r")
 end