now test in test_sampling_asktell_gen

Libensemble · Aug 9, 2024 · e27487d · e27487d
1 parent d14f4d2
commit e27487d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/libensemble/gen_classes/sampling.py b/libensemble/gen_classes/sampling.py
@@ -15,7 +15,6 @@ class SampleBase(LibensembleGenerator):
 
     def _get_user_params(self, user_specs):
         """Extract user params"""
-        # b = user_specs["initial_batch_size"]
         self.ub = user_specs["ub"]
         self.lb = user_specs["lb"]
         self.n = len(self.lb)  # dimension
@@ -32,7 +31,7 @@ class UniformSample(SampleBase):
     mode by adjusting the allocation function.
     """
 
-    def __init__(self, _, persis_info, gen_specs, libE_info=None) -> list:
+    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
@@ -63,7 +62,7 @@ class UniformSampleDicts(Generator):
     mode by adjusting the allocation function.
     """
 
-    def __init__(self, _, persis_info, gen_specs, libE_info=None) -> list:
+    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

diff --git a/libensemble/tests/functionality_tests/test_sampling_asktell_gen.py b/libensemble/tests/functionality_tests/test_sampling_asktell_gen.py
@@ -42,8 +42,10 @@ def sim_f(In):
     gen_specs = {
         "persis_in": ["x", "f", "grad", "sim_id"],
         "out": [("x", float, (2,))],
+        "initial_batch_size": 20,
+        "batch_size": 10,
         "user": {
-            "initial_batch_size": 20,
+            "initial_batch_size": 20,  # for wrapper
             "lb": np.array([-3, -2]),
             "ub": np.array([3, 2]),
         },