Using persis_info to speed up alloc_func

libensemble/alloc_funcs/persistent_aposmm_alloc.py

@@ -6,45 +6,47 @@
 def persistent_aposmm_alloc(W, H, sim_specs, gen_specs, alloc_specs, persis_info):
     """
     This allocation function will give simulation work if possible, but
-    otherwise start up to 1 persistent generator.  If all points requested by
+    otherwise start a persistent APOSMM generator.  If all points requested by
     the persistent generator have been returned from the simulation evaluation,
     then this information is given back to the persistent generator.
 
+    This function assumes that one persistent APOSMM will be started and never
+    stopped (until some exit_criterion is satisfied).
+
     :See:
-        ``/libensemble/tests/regression_tests/test_6-hump_camel_persistent_uniform_sampling.py``
+        ``/libensemble/tests/regression_tests/test_6-hump_camel_persistent_aposmm.py``
     """
 
     Work = {}
-    gen_count = count_persis_gens(W)
+    if 'next_to_give' not in persis_info:
+        persis_info['next_to_give'] = 0
+
 
-    # If i is in persistent mode, and any of its calculated values have
-    # returned, give them back to i. Otherwise, give nothing to i
+    # If any persistent worker's calculated values have returned, give them back.
     for i in avail_worker_ids(W, persistent=True):
-        if sum(H['returned']) < gen_specs['initial_sample_size']:
+        if persis_info.get('sample_done') or sum(H['returned']) >= gen_specs['initial_sample_size']:
             # Don't return if the initial sample is not complete
-            continue
+            persis_info['sample_done'] = True
 
-        gen_inds = (H['gen_worker'] == i)
-        returned_but_not_given = np.logical_and(H['returned'][gen_inds], ~H['given_back'][gen_inds])
-        if np.any(returned_but_not_given):
-            inds_to_give = np.where(returned_but_not_given)[0]
+            returned_but_not_given = np.logical_and(H['returned'], ~H['given_back'])
+            if np.any(returned_but_not_given):
+                inds_to_give = np.where(returned_but_not_given)[0]
 
-            gen_work(Work, i, [n[0] for n in sim_specs['out']] + [n[0] for n in gen_specs['out']],
-                     np.atleast_1d(inds_to_give), persis_info[i], persistent=True)
+                gen_work(Work, i, [n[0] for n in sim_specs['out']] + [n[0] for n in gen_specs['out']],
+                         np.atleast_1d(inds_to_give), persis_info[i], persistent=True)
 
-            H['given_back'][inds_to_give] = True
+                H['given_back'][inds_to_give] = True
 
-    task_avail = ~H['given']
     for i in avail_worker_ids(W, persistent=False):
-        if np.any(task_avail):
+        if persis_info['next_to_give'] < len(H):
             # perform sim evaluations (if they exist in History).
-            sim_ids_to_send = np.nonzero(task_avail)[0][0]  # oldest point
-            sim_work(Work, i, sim_specs['in'], np.atleast_1d(sim_ids_to_send), persis_info[i])
-            task_avail[sim_ids_to_send] = False
+            sim_work(Work, i, sim_specs['in'], np.atleast_1d(persis_info['next_to_give']), persis_info[i])
+            persis_info['next_to_give'] += 1
 
-        elif gen_count == 0:
+        elif persis_info.get('gen_started') is None:
             # Finally, call a persistent generator as there is nothing else to do.
-            gen_count += 1
+            persis_info['gen_started'] = True
+
             gen_work(Work, i, gen_specs['in'], [], persis_info[i],
                      persistent=True)