Feature/flatten asktell dict values

libensemble/gen_classes/aposmm.py

@@ -19,8 +19,6 @@ def __init__(
         from libensemble.gen_funcs.persistent_aposmm import aposmm
 
         gen_specs["gen_f"] = aposmm
-        if len(kwargs) > 0:  # so user can specify aposmm-specific parameters as kwargs to constructor
-            gen_specs["user"] = kwargs
         if not gen_specs.get("out"):  # gen_specs never especially changes for aposmm even as the problem varies
             n = len(kwargs["lb"]) or len(kwargs["ub"])
             gen_specs["out"] = [
@@ -34,7 +32,7 @@ def __init__(
         if not persis_info:
             persis_info = add_unique_random_streams({}, 4, seed=4321)[1]
             persis_info["nworkers"] = 4
-        super().__init__(gen_specs, History, persis_info, libE_info)
+        super().__init__(gen_specs, History, persis_info, libE_info, **kwargs)
         self.all_local_minima = []
         self.results_idx = 0
         self.last_ask = None

libensemble/gen_classes/sampling.py

@@ -31,11 +31,10 @@ class UniformSample(SampleBase):
     mode by adjusting the allocation function.
     """
 
-    def __init__(self, _, persis_info, gen_specs, libE_info=None):
-        self.persis_info = persis_info
-        self.gen_specs = gen_specs
-        self.libE_info = libE_info
+    def __init__(self, _=[], persis_info={}, gen_specs={}, libE_info=None, **kwargs):
+        super().__init__(gen_specs, _, persis_info, libE_info, **kwargs)
         self._get_user_params(self.gen_specs["user"])
+        self.gen_specs["out"] = [("x", float, (self.n,))]
 
     def ask_numpy(self, n_trials):
         H_o = np.zeros(n_trials, dtype=self.gen_specs["out"])
@@ -57,11 +56,11 @@ class UniformSampleDicts(Generator):
     mode by adjusting the allocation function.
     """
 
-    def __init__(self, _, persis_info, gen_specs, libE_info=None):
-        self.persis_info = persis_info
+    def __init__(self, _, persis_info, gen_specs, libE_info=None, **kwargs):
         self.gen_specs = gen_specs
-        self.libE_info = libE_info
+        self.persis_info = persis_info
         self._get_user_params(self.gen_specs["user"])
+        self.gen_specs["out"] = [("x", float, (self.n,))]
 
     def ask(self, n_trials):
         H_o = []

libensemble/gen_funcs/persistent_gen_wrapper.py

@@ -1,10 +1,8 @@
 import inspect
 
-import numpy as np
-
 from libensemble.message_numbers import EVAL_GEN_TAG, FINISHED_PERSISTENT_GEN_TAG, PERSIS_STOP, STOP_TAG
 from libensemble.tools.persistent_support import PersistentSupport
-from libensemble.utils.misc import np_to_list_dicts
+from libensemble.utils.misc import list_dicts_to_np, np_to_list_dicts
 
 
 def persistent_gen_f(H, persis_info, gen_specs, libE_info):
@@ -24,11 +22,7 @@ def persistent_gen_f(H, persis_info, gen_specs, libE_info):
     while tag not in [STOP_TAG, PERSIS_STOP]:
         H_o = gen.ask(b)
         if isinstance(H_o, list):
-            H_o_arr = np.zeros(len(H_o), dtype=gen_specs["out"])
-            for i in range(len(H_o)):
-                for key in H_o[0].keys():
-                    H_o_arr[i][key] = H_o[i][key]
-            H_o = H_o_arr
+            H_o = list_dicts_to_np(H_o)
         tag, Work, calc_in = ps.send_recv(H_o)
         gen.tell(np_to_list_dicts(calc_in))
 

libensemble/generators.py

@@ -8,6 +8,7 @@
 from libensemble.comms.comms import QComm, QCommThread
 from libensemble.executors import Executor
 from libensemble.message_numbers import EVAL_GEN_TAG, PERSIS_STOP
+from libensemble.tools.tools import add_unique_random_streams
 from libensemble.utils.misc import list_dicts_to_np, np_to_list_dicts
 
 """
@@ -90,6 +91,18 @@ class LibensembleGenerator(Generator):
     ``ask_numpy/tell_numpy`` methods communicate numpy arrays containing the same data.
     """
 
+    def __init__(
+        self, gen_specs: dict = {}, History: npt.NDArray = [], persis_info: dict = {}, libE_info: dict = {}, **kwargs
+    ):
+        self.gen_specs = gen_specs
+        if len(kwargs) > 0:  # so user can specify gen-specific parameters as kwargs to constructor
+            self.gen_specs["user"] = kwargs
+        if not persis_info:
+            self.persis_info = add_unique_random_streams({}, 4, seed=4321)[1]
+            self.persis_info["nworkers"] = 4
+        else:
+            self.persis_info = persis_info
+
     @abstractmethod
     def ask_numpy(self, num_points: Optional[int] = 0) -> npt.NDArray:
         """Request the next set of points to evaluate, as a NumPy array."""
@@ -113,10 +126,10 @@ class LibensembleGenThreadInterfacer(LibensembleGenerator):
     """
 
     def __init__(
-        self, gen_specs: dict, History: npt.NDArray = [], persis_info: dict = {}, libE_info: dict = {}
+        self, gen_specs: dict, History: npt.NDArray = [], persis_info: dict = {}, libE_info: dict = {}, **kwargs
     ) -> None:
+        super().__init__(gen_specs, History, persis_info, libE_info, **kwargs)
         self.gen_f = gen_specs["gen_f"]
-        self.gen_specs = gen_specs
         self.History = History
         self.persis_info = persis_info
         self.libE_info = libE_info

libensemble/tests/unit_tests/test_asktell.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+from libensemble.tools.tools import add_unique_random_streams
+from libensemble.utils.misc import list_dicts_to_np
+
+
+def test_asktell_sampling_and_utils():
+    from libensemble.gen_classes.sampling import UniformSample
+
+    persis_info = add_unique_random_streams({}, 5, seed=1234)
+    gen_specs = {
+        "out": [("x", float, (2,))],
+        "user": {
+            "lb": np.array([-3, -2]),
+            "ub": np.array([3, 2]),
+        },
+    }
+
+    # Test initialization with libensembley parameters
+    gen = UniformSample(None, persis_info[1], gen_specs, None)
+    assert len(gen.ask(10)) == 10
+
+    # Test initialization gen-specific keyword args
+    gen = UniformSample(lb=np.array([-3, -2]), ub=np.array([3, 2]))
+    assert len(gen.ask(10)) == 10
+
+    out_np = gen.ask_numpy(3)  # should get numpy arrays, non-flattened
+    out = gen.ask(3)  # needs to get dicts, 2d+ arrays need to be flattened
+    assert all([len(x) == 2 for x in out])  # np_to_list_dicts is now tested
+
+    # now we test list_dicts_to_np directly
+    out_np = list_dicts_to_np(out)
+
+    # check combined values resemble flattened list-of-dicts values
+    assert out_np.dtype.names == ("x",)
+    for i, entry in enumerate(out):
+        for j, value in enumerate(entry.values()):
+            assert value == out_np["x"][i][j]
+
+
+if __name__ == "__main__":
+    test_asktell_sampling_and_utils()

libensemble/utils/misc.py

@@ -81,33 +81,51 @@ def specs_checker_setattr(obj, key, value):
         obj.__dict__[key] = value
 
 
-def _copy_data(array, list_dicts):
-    for i, entry in enumerate(list_dicts):
-        for field in entry.keys():
-            array[field][i] = entry[field]
-    return array
+def _decide_dtype(name: str, entry, size: int) -> tuple:
+    if size == 1 or not size:
+        return (name, type(entry))
+    else:
+        return (name, type(entry), (size,))
 
 
-def _decide_dtype(name, entry):
-    if hasattr(entry, "shape") and len(entry.shape):  # numpy type
-        return (name, entry.dtype, entry.shape)
-    else:
-        return (name, type(entry))
+def _combine_names(names: list) -> list:
+    """combine fields with same name *except* for final digit"""
+    # how many final digits could possibly be in each name?
+    #  do we have to iterate through negative-indexes until we reach a non-digit?
+    return list(set(i[:-1] if i[-1].isdigit() else i for i in names))
 
 
 def list_dicts_to_np(list_dicts: list) -> npt.NDArray:
     if list_dicts is None:
         return None
 
-    first = list_dicts[0]
-    new_dtype_names = [i for i in first.keys()]
+    first = list_dicts[0]  # for determining dtype of output np array
+    new_dtype_names = _combine_names([i for i in first.keys()])  # -> ['x', 'y']
+    combinable_names = []  # [['x0', 'x1'], ['y0', 'y1', 'y2'], []]
+    for name in new_dtype_names:
+        combinable_names.append([i for i in first.keys() if i[:-1] == name])
+
     new_dtype = []
-    for i, entry in enumerate(first.values()):  # must inspect values to get presumptive types
+
+    for i, entry in enumerate(combinable_names):
         name = new_dtype_names[i]
-        new_dtype.append(_decide_dtype(name, entry))
+        size = len(combinable_names[i])
+        if len(entry):  # combinable names detected, e.g. x0, x1
+            new_dtype.append(_decide_dtype(name, first[entry[0]], size))
+        else:  # only a single name, e.g. local_pt
+            new_dtype.append(_decide_dtype(name, first[name], size))
 
     out = np.zeros(len(list_dicts), dtype=new_dtype)
-    return _copy_data(out, list_dicts)
+
+    for i, group in enumerate(combinable_names):
+        new_dtype_name = new_dtype_names[i]
+        for j, input_dict in enumerate(list_dicts):
+            if not len(group):  # only a single name, e.g. local_pt
+                out[new_dtype_name][j] = input_dict[new_dtype_name]
+            else:  # combinable names detected, e.g. x0, x1
+                out[new_dtype_name][j] = tuple([input_dict[name] for name in group])
+
+    return out
 
 
 def np_to_list_dicts(array: npt.NDArray) -> List[dict]:
@@ -117,6 +135,10 @@ def np_to_list_dicts(array: npt.NDArray) -> List[dict]:
     for row in array:
         new_dict = {}
         for field in row.dtype.names:
-            new_dict[field] = row[field]
+            if hasattr(row[field], "__len__") and len(row[field]) > 1:
+                for i, x in enumerate(row[field]):
+                    new_dict[field + str(i)] = x
+            else:
+                new_dict[field] = row[field]
         out.append(new_dict)
     return out