add named curve support

README.md

@@ -0,0 +1,9 @@
+
+The following example will call `some_method` with the argument `1000`,  then `2000`, etc.
+```
+    vis = Vis()
+    vis.visualize(
+        some_method,
+        [1000,2000,4000,8000]
+    )
+```

runtime_vis/curves.py

@@ -36,8 +36,16 @@ def _curve_n3(x, a, b):
     return a * x ** 2 + b
 
 
-curve_n = Curve('n', _curve_n, 'gray')
-curve_logn = Curve('logn', _curve_logn, 'green')
-curve_nlogn = Curve('nlogn', _curve_nlogn, 'blue')
-curve_n2 = Curve('n2', _curve_n2, 'orange')
-curve_n3 = Curve('n3', _curve_n3, 'red')
+fit_n = Curve('n', _curve_n, 'gray')
+fit_logn = Curve('logn', _curve_logn, 'green')
+fit_nlogn = Curve('nlogn', _curve_nlogn, 'blue')
+fit_n2 = Curve('n2', _curve_n2, 'orange')
+fit_n3 = Curve('n3', _curve_n3, 'red')
+
+curves = {
+    'n': fit_n,
+    'logn': fit_logn,
+    'nlogn': fit_nlogn,
+    'n2': fit_n2,
+    'n3': fit_n3,
+}

runtime_vis/vis.py

@@ -8,37 +8,49 @@
 import warnings
 from typing import Callable, Any, Iterable
 
-from runtime_vis.curves import curve_n, curve_logn, curve_nlogn
+from runtime_vis.curves import curve_n, curve_logn, curve_nlogn, curves
 
 
 class Vis:
-    def __init__(self):
+    def __init__(self, fit: List[str] = None):
         self.initial_guess = {
             'n': [1e-7, 1],
             'logn': [1e-7, 1],
             'nlogn': [1e-7, 1],
             'n2': [1e-10, 1],
         }
-        self.curves = [
-            curve_n,
-            curve_logn,
-            curve_nlogn,
-            curve_n2,
-        ]
+        self.times = []
 
-    def profile_batch(self, function: Callable[[int], Any], name: str, size: int, dataset: pd.DataFrame):
+        if fit:
+            self.curves = []
+            for name in fit:
+                if name in curves:
+                    self.curves.append(curves[name])
+        else:
+            self.curves = [
+                fit_n,
+                fit_logn,
+                fit_nlogn,
+                fit_n2,
+            ]
+
+
+    def profile_batch(self, function: Callable[[int], Any], size: int, dataset: pd.DataFrame):
         p = cProfile.Profile()
         p.enable()
         function(size)
         p.create_stats()
         stats = pstats.Stats(p)
+        time = 0
         for stat in stats.stats.items():
-            if stat[0][2] == name:
+            if stat[1][3] > max:
                 time = stat[1][3]
-                row = {'N': size, 'Time': time}
-                dataset.loc[len(dataset)] = [size, time]
-                render_plot(dataset)
-                break
+
+        row = {'N': size, 'Time': time}
+        dataset.loc[len(dataset)] = [size, time]
+        render_plot(dataset)
+        self.times.append((size,time))
+
         p.disable()
         print(f"{size}\t{time:4.2f}")
 
@@ -68,17 +80,16 @@ def fit_curve(self, dataset: pd.DataFrame):
     def visualize(
         self,
         function: Callable[[int], Any],
-        name: str,
         series: Iterable[int]
-    ):
+    ) -> List[tuple[int,int]]:
         dataset = pd.DataFrame(columns=['N', 'Time'])
         plt.figure()
         plt.ion()
         plt.show()
         plt.xlabel('N')
         plt.ylabel('Time (ms)')
         for size in series:
-            self.profile_batch(function, name, size, dataset)
+            self.profile_batch(function, size, dataset)
             self.fit_curve(dataset)
         plt.ioff()
         plt.show()