Resource allocation colormap (#3453)

* Update to DB qiita.slurm_resource_allocations

* added colormap, create equation table

* Fixed styling

* Using qiita.allocation_equations table in util.py

* Debug

* Updates to @antgonza comments

* Changes to @antgonza comments

* Back to np.log

Author: Gossty

qiita_db/meta_util.py

@@ -593,7 +593,7 @@ def update_resource_allocation_redis(active=True):
                 if len(df) == 0:
                     continue
 
-                fig, axs = resource_allocation_plot(df, cname, sname, col_name)
+                fig, axs = resource_allocation_plot(df, col_name)
                 titles = [0, 0]
                 images = [0, 0]
 
@@ -605,21 +605,18 @@ def update_resource_allocation_redis(active=True):
                     # only time
                     new_fig = plt.figure()
                     new_ax = new_fig.add_subplot(111)
-
-                    scatter_data = ax.collections[0]
-                    new_ax.scatter(scatter_data.get_offsets()[:, 0],
-                                   scatter_data.get_offsets()[:, 1],
-                                   s=scatter_data.get_sizes(), label="data")
-
                     line = ax.lines[0]
                     new_ax.plot(line.get_xdata(), line.get_ydata(),
                                 linewidth=1, color='orange')
-
-                    if len(ax.collections) > 1:
-                        failure_data = ax.collections[1]
-                        new_ax.scatter(failure_data.get_offsets()[:, 0],
-                                       failure_data.get_offsets()[:, 1],
-                                       color='red', s=3, label="failures")
+                    handles, labels = ax.get_legend_handles_labels()
+                    for handle, label, scatter_data in zip(handles,
+                                                           labels,
+                                                           ax.collections):
+                        color = handle.get_facecolor()
+                        new_ax.scatter(scatter_data.get_offsets()[:, 0],
+                                       scatter_data.get_offsets()[:, 1],
+                                       s=scatter_data.get_sizes(), label=label,
+                                       color=color)
 
                     new_ax.set_xscale('log')
                     new_ax.set_yscale('log')

qiita_db/support_files/patches/94.sql

@@ -0,0 +1,7 @@
+-- Jan 13, 2025
+-- Adding a table for formulas for resource allocations
+CREATE TABLE qiita.allocation_equations (
+  equation_id     SERIAL PRIMARY KEY,
+  equation_name   TEXT NOT NULL,
+  expression      TEXT NOT NULL
+ );

qiita_db/support_files/patches/test_db_sql/94.sql

@@ -0,0 +1,10 @@
+INSERT INTO qiita.allocation_equations(equation_name, expression)
+        VALUES 
+        ('mem_model1', '(k * (np.log(x))) + (x * a) + b'),
+('mem_model2', '(k * (np.log(x))) + (b * ((np.log(x))**2)) + a'),
+('mem_model3', '(k * (np.log(x))) + (b * ((np.log(x))**2)) + (a * ((np.np.log(x))**3))'),
+('mem_model4', '(k * (np.log(x))) + (b * ((np.log(x))**2)) + (a * ((np.log(x))**2.5))'),
+('time_model1', 'a + b + ((np.log(x)) * k)'),
+('time_model2', 'a + (b * x) + ((np.log(x)) * k)'),
+('time_model3', 'a + (b * ((np.log(x))**2)) + ((np.log(x)) * k)'),
+('time_model4', '(a * ((np.log(x))**3)) + (b * ((np.log(x))**2)) + ((np.log(x)) * k)');

qiita_db/test/test_meta_util.py

@@ -529,16 +529,19 @@ def test_update_resource_allocation_redis(self):
             cname, sname, version, col_name, 'title_mem')
         title_mem = str(r_client.get(title_mem_str))
         self.assertTrue(
-            "model: "
-            "k * log(x) + "
-            "b * log(x)^2 + "
-            "a * log(x)^3" in title_mem
+            "model: (k * (np.log(x))) + "
+            "(b * ((np.log(x))**2)) + "
+            "(a * ((np.log(x))**2.5))" in title_mem
         )
 
         title_time_str = 'resources$#%s$#%s$#%s$#%s:%s' % (
                         cname, sname, version, col_name, 'title_time')
         title_time = str(r_client.get(title_time_str))
-        self.assertTrue("model: a + b + log(x) * k" in title_time)
+        self.assertTrue(
+            "model: (a * ((np.log(x))**3)) + "
+            "(b * ((np.log(x))**2)) + "
+            "((np.log(x)) * k)" in title_time
+        )
 
 
 if __name__ == '__main__':

qiita_db/test/test_util.py

@@ -1329,8 +1329,7 @@ def setUp(self):
 
     def test_plot_return(self):
         # check the plot returns correct objects
-        fig1, axs1 = qdb.util.resource_allocation_plot(
-            self.df, self.cname, self.sname, self.col_name)
+        fig1, axs1 = qdb.util.resource_allocation_plot(self.df, self.col_name)
         self.assertIsInstance(
             fig1, Figure,
             "Returned object fig1 is not a Matplotlib Figure")
@@ -1346,46 +1345,46 @@ def test_minimize_const(self):
         self.df[self.col_name] = self.df.samples * self.df['columns']
         fig, axs = plt.subplots(ncols=2, figsize=(10, 4), sharey=False)
 
-        bm, options = qdb.util._resource_allocation_plot_helper(
-            self.df, axs[0], self.cname, self.sname, 'MaxRSSRaw',
-            qdb.util.MODELS_MEM, self.col_name)
+        mem_models, time_models = qdb.util.retrieve_equations()
+        bm_name, bm, options = qdb.util._resource_allocation_plot_helper(
+            self.df, axs[0], 'MaxRSSRaw', mem_models, self.col_name)
         # check that the algorithm chooses correct model for MaxRSSRaw and
         # has 0 failures
         k, a, b = options.x
-        failures_df = qdb.util._resource_allocation_failures(
-            self.df, k, a, b, bm, self.col_name, 'MaxRSSRaw')
+        failures_df = qdb.util._resource_allocation_success_failures(
+            self.df, k, a, b, bm, self.col_name, 'MaxRSSRaw')[-1]
         failures = failures_df.shape[0]
-        self.assertEqual(bm, qdb.util.mem_model3,
+
+        self.assertEqual(bm_name, 'mem_model4',
+                         msg=f"""Best memory model
+                         doesn't match
+                         {bm_name} != 'mem_model4'""")
+        self.assertEqual(bm, mem_models['mem_model4']['equation'],
                          msg=f"""Best memory model
                                  doesn't match
                                  Coefficients:{k} {a} {b}
-                                 {qdb.util.mem_model1}, "qdb.util.mem_model1"
-                                 {qdb.util.mem_model2}, "qdb.util.mem_model2"
-                                 {qdb.util.mem_model3}, "qdb.util.mem_model3"
-                                 {qdb.util.mem_model4}, "qdb.util.mem_model4"
                             """)
         self.assertEqual(failures, 0, "Number of failures must be 0")
 
         # check that the algorithm chooses correct model for ElapsedRaw and
         # has 1 failure
-        bm, options = qdb.util._resource_allocation_plot_helper(
-            self.df, axs[1], self.cname, self.sname, 'ElapsedRaw',
-            qdb.util.MODELS_TIME, self.col_name)
+        bm_name, bm, options = qdb.util._resource_allocation_plot_helper(
+            self.df, axs[1], 'ElapsedRaw', time_models, self.col_name)
         k, a, b = options.x
-        failures_df = qdb.util._resource_allocation_failures(
-            self.df, k, a, b, bm, self.col_name, 'ElapsedRaw')
+        failures_df = qdb.util._resource_allocation_success_failures(
+            self.df, k, a, b, bm, self.col_name, 'ElapsedRaw')[-1]
         failures = failures_df.shape[0]
+        self.assertEqual(bm_name, 'time_model4',
+                         msg=f"""Best time model
+                         doesn't match
+                         {bm_name} != 'time_model4'""")
 
-        self.assertEqual(bm, qdb.util.time_model1,
+        self.assertEqual(bm, time_models[bm_name]['equation'],
                          msg=f"""Best time model
                                 doesn't match
                                 Coefficients:{k} {a} {b}
-                                 {qdb.util.time_model1}, "qdb.util.time_model1"
-                                 {qdb.util.time_model2}, "qdb.util.time_model2"
-                                 {qdb.util.time_model3}, "qdb.util.time_model3"
-                                 {qdb.util.time_model4}, "qdb.util.time_model4"
                                 """)
-        self.assertEqual(failures, 1, "Number of failures must be 1")
+        self.assertEqual(failures, 0, "Number of failures must be 0")
 
     def test_MaxRSS_helper(self):
         tests = [