Merge pull request #2184 from ZDu-IDM/fix_arm_builder

Improved Simulation builders (fixed several issues)

examples/builders/arm_experiment_builder_python.py

@@ -37,7 +37,7 @@
 
 
 if __name__ == "__main__":
-    with platform('BELEGOST'):
+    with platform('Calculon'):
         base_task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"))
         # define that we are going to create multiple simulations from this task
         ts = TemplatedSimulations(base_task=base_task)

examples/builders/csv_builder_python.py

@@ -48,7 +48,7 @@
 if __name__ == "__main__":
     # define what platform we want to use. Here we use a context manager but if you prefer you can
     # use objects such as Platform('COMPS') instead
-    with platform('BELEGOST'):
+    with platform('Calculon'):
         # define our base task
         base_task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"),
                                              parameters=dict(c='c-value'))

examples/builders/experiment_builder_python.py

@@ -37,7 +37,7 @@ def param_update_ab(simulation, param, value):
 if __name__ == "__main__":
     # define what platform we want to use. Here we use a context manager but if you prefer you can
     # use objects such as Platform('BELEGOST') instead
-    with platform('BELEGOST'):
+    with platform('Calculon'):
         # define our base task
         base_task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"),
                                              parameters=dict(c='c-value'))

examples/builders/manual_building.py

@@ -19,7 +19,7 @@
 if __name__ == "__main__":
 
     # define our platform
-    platform = Platform('BELEGOST')
+    platform = Platform('Calculon')
 
     # create experiment  object and define some extra assets
     assets_path = os.path.join(COMMON_INPUT_PATH, "python", "Assets")

examples/builders/multi_argument_sweep.py

@@ -32,14 +32,16 @@ def param_update(simulation, a_value, b_value, c_value):
 
 if __name__ == "__main__":
     # define what platform we want to use. Here we use a context manager but if you prefer you can
-    # use objects such as Platform('BELEGOST') instead
-    with Platform('BELEGOST'):
+    # use objects such as Platform('Calculon') instead
+    with Platform('Calculon'):
         # define our base task
         base_task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"),
                                              parameters=dict())
         # define our input csv sweep
         builder = SimulationBuilder()
-        builder.add_multiple_parameter_sweep_definition(param_update, range(2), range(2), range(2))
+        # we can use add_sweep_definition call to do multiple parameter sweeping now
+        builder.add_sweep_definition(param_update, range(2), range(2), range(2))
+        #builder.add_multiple_parameter_sweep_definition(param_update, range(2), range(2), range(2))
 
         # define our experiment with its metadata
         experiment = Experiment.from_builder(

examples/builders/simulation_builder.py

@@ -9,7 +9,7 @@
 from idmtools_models.python.json_python_task import JSONConfiguredPythonTask
 from idmtools_test import COMMON_INPUT_PATH
 
-with platform('BELEGOST'):
+with platform('Calculon'):
     base_task = JSONConfiguredPythonTask(
         script_path=os.path.join(COMMON_INPUT_PATH, "compsplatform", "working_model.py"),
         # add common assets from existing collection

examples/builders/yaml_builder_python.py

@@ -52,7 +52,7 @@
 if __name__ == "__main__":
     # define what platform we want to use. Here we use a context manager but if you prefer you can
     # use objects such as Platform('BELEGOST') instead
-    with platform('BELEGOST'):
+    with platform('Calculon'):
         # define our base task
         base_task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"),
                                              parameters=dict(c='c-value'))

examples/python_model/multiple_parameter_sweeping.py

@@ -25,7 +25,7 @@ def update_parameter_callback(simulation, pop_size, pop_infected, n_days, rand_s
 ts = TemplatedSimulations(base_task=task)
 
 sb = SimulationBuilder()
-sb.add_multiple_parameter_sweep_definition(update_parameter_callback, pop_size=[10000, 20000], pop_infected=[10, 100],
+sb.add_sweep_definition(update_parameter_callback, pop_size=[10000, 20000], pop_infected=[10, 100],
                                            n_days=[100, 110], rand_seed=[1234, 4567])
 # has to add workorder.json with simulation to update cmd arguments
 # for command line task sweep with workorder, DO NOT use since it will not update arguments:

idmtools_core/idmtools/builders/arm_simulation_builder.py

@@ -3,13 +3,11 @@
 
 Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
 """
-import copy
-import collections
 from enum import Enum
 from itertools import product
-from typing import Tuple, List, Callable, Iterable, Any
-
+from typing import Tuple, List, Callable, Iterable
 from idmtools.builders import SimulationBuilder
+from idmtools.builders.simulation_builder import TSweepFunction
 
 
 class ArmType(Enum):
@@ -28,65 +26,148 @@ class SweepArm:
     def __init__(self, type=ArmType.cross, funcs: List[Tuple[Callable, Iterable]] = None):
         """
         Constructor.
-
         Args:
             type: Type of Arm(Cross or Pair)
             funcs: Functions to add as sweeps
         """
-        if funcs is None:
-            funcs = []
-        self.sweep_functions = []
         self.type = type
+        self.sweeps = []
+        self.__count = 0
 
+        if funcs is None:
+            funcs = []
         for func, values in funcs:
             self.add_sweep_definition(func, values)
 
-    def add_sweep_definition(self, func: Callable, values: Iterable[Any]):  # noqa F821
+    @property
+    def count(self) -> int:
         """
-        Add Sweep definition.
+        Simulation count.
+
+        Returns:
+            count
+        """
+        return self.__count
+
+    @count.setter
+    def count(self, cnt: int):
+        """
+        Set the count property.
 
         Args:
-            func: Sweep callback
-            values: Values to Sweep
+            cnt: count set
 
         Returns:
             None
         """
-        self.sweep_functions.append((func, values if isinstance(values, collections.abc.Iterable) and not (
-            isinstance(values, str)) else [values]))
+        # print('cnt: ', cnt)
+        if self.__count == 0:
+            self.__count = cnt
+        elif self.type == ArmType.cross:
+            self.__count = self.__count * cnt
+        elif self.type == ArmType.pair:
+            if self.__count != cnt:
+                raise ValueError(
+                    f"For pair case, all function inputs must have the save size/length: {cnt} != {self.__count}")
+            else:
+                self.__count = cnt
+
+    def add_sweep_definition(self, function: TSweepFunction, *args, **kwargs):
+        """
+        Add a sweep definition callback that takes possible multiple parameters (None or many).
 
-        if self.type == ArmType.pair:
-            self.adjust_values_length()
+        The sweep will be defined as a cross-product between the parameters passed.
 
-    def get_max_values_count(self):
-        """
-        Get the max values count from different sweep functions.
+        Args:
+            function: The sweep function, which must include a **simulation** parameter (or
+                whatever is specified in :attr:`~idmtools.builders.ExperimentBuilder.SIMULATION_ATTR`).
+            args: List of arguments to be passed
+            kwargs: List of keyword arguments to be passed
 
         Returns:
-            Max values
-        """
-        cnts = [len(values) for _, values in self.sweep_functions]
-        return max(cnts)
+            None. Updates the Sweeps
+
+        Examples:
+            Examples of valid functions::
+
+                # This function takes one parameter
+                def myFunction(simulation, parameter_a):
+                    pass
+
+               # This function takes one parameter with default value
+                def myFunction(simulation, parameter_a=6):
+                    pass
+
+                # This function takes two parameters (parameters may have default values)
+                def myFunction(simulation, parameter_a, parameter_b=9):
+                    pass
+
+                # Function that takes three parameters (parameters may have default values)
+                def three_param_callback(simulation, parameter_a, parameter_b, parameter_c=10):
+                    pass
+
+            Calling Sweeps that take multiple parameters::
+
+                # This example references the above valid function example
+                sb = SimulationBuilder()
+
+                # Add a sweep on the myFunction that takes parameter(s).
+                # Here we sweep the values 1-4 on parameter_a and a,b on parameter_b
+                sb.add_sweep_definition(myFunction, range(1,5), ["a", "b"])
+
+                sb2 = SimulationBuilder()
+                # Example calling using a dictionary instead
+                sb.add_sweep_definition(three_param_callback, dict(parameter_a=range(1,5), parameter_b=["a", "b"], parameter_c=range(4,5))
+                # The following is equivalent
+                sb.add_sweep_definition(three_param_callback, **dict(parameter_a=range(1,5), parameter_b=["a", "b"], parameter_c=range(4,5)))
+
+                # If all parameters have default values, we can even simply do
+                sb.add_sweep_definition(three_param_callback)
+
+                # Remark: in general
+                def my_callback(simulation, parameter_1, parameter_2, ..., parameter_n):
+                    pass
 
-    def adjust_values_length(self):
+                Calling Sweeps that take multiple parameters::
+
+                sb = SimulationBuilder()
+                sb.add_sweep_definition(my_callback, Iterable_1, Iterable_2, ..., Iterable_m)
+
+                # Note:  the # of Iterable object must match the parameters # of my_callback, which don't have default values or use the key (parameter names)
+
+                sb.add_sweep_definition(my_callback, parameter_1=Iterable_1, parameter_2=Iterable_2, ..., parameter_m=Iterable_m)
+
+                # The following is equivalent
+                sb.add_sweep_definition(my_callback, dict(parameter_1=Iterable_1, parameter_2=Iterable_2, ..., parameter_m=Iterable_m))
+                sb.add_sweep_definition(my_callback, **dict(parameter_1=Iterable_1, parameter_2=Iterable_2, ..., parameter_m=Iterable_m))
         """
-        Adjust values length.
+        builder = SimulationBuilder()
+        builder.add_sweep_definition(function, *args, **kwargs)
+        self.sweeps.extend(builder.sweeps)
+        self.count = builder.count
 
-        Returns:
-            None
+    def add_multiple_parameter_sweep_definition(self, function: TSweepFunction, *args, **kwargs):
         """
-        if self.type != ArmType.pair:
-            return
+        Add a sweep definition callback that takes possible multiple parameters (None or many).
+
+        The sweep will be defined as a cross-product between the parameters passed.
 
-        count_max = self.get_max_values_count()
-        temp_sweep_functions = []
-        for func, values in self.sweep_functions:
-            values_new = copy.deepcopy(values)
-            values_new = list(values_new)
-            values_new.extend([values[-1]] * (count_max - len(values)))
-            temp_sweep_functions.append((func, values_new))
+        Args:
+            function: The sweep function, which must include a **simulation** parameter (or
+                whatever is specified in :attr:`~idmtools.builders.ExperimentBuilder.SIMULATION_ATTR`).
+            args: List of arguments to be passed
+            kwargs: List of keyword arguments to be passed
+
+        Returns:
+            None. Updates the Sweeps
 
-        self.sweep_functions = temp_sweep_functions
+        Examples:
+            Refer to the comments in the add_sweep_definition function for examples
+        """
+        builder = SimulationBuilder()
+        builder.add_multiple_parameter_sweep_definition(function, *args, **kwargs)
+        self.sweeps.extend(builder.sweeps)
+        self.count = builder.count
 
 
 class ArmSimulationBuilder(SimulationBuilder):
@@ -178,44 +259,58 @@ def __init__(self):
         self.arms = []
         self.sweep_definitions = []
 
-    def add_arm(self, arm):
+    def add_arm(self, arm: SweepArm):
         """
         Add arm sweep definition.
-
         Args:
             arm: Arm to add
-
         Returns:
             None
         """
-        arm_list = arm if isinstance(arm, collections.abc.Iterable) else [arm]
-        for a in arm_list:
-            self.arms.append(a)
-            self._apply(a)
+        self.arms.append(arm)
+        if arm.type == ArmType.cross:
+            self.sweep_definitions.extend(product(*arm.sweeps))
+        elif arm.type == ArmType.pair:
+            self.sweep_definitions.extend(zip(*arm.sweeps))
+        self.count = sum([arm.count for arm in self.arms])
 
-    def _apply(self, arm):
+    def add_sweep_definition(self, function: TSweepFunction, *args, **kwargs):
         """
-        Apply our arm.
+        Add parameters sweep definition.
 
         Args:
-            arm: Arm to apply
+            function: The sweep function, which must include a **simulation** parameter (or
+                whatever is specified in :attr:`~idmtools.builders.ExperimentBuilder.SIMULATION_ATTR`).
+                The function also must include EXACTLY ONE free parameter, which the values will be passed to.
+                The function can also be a partial--any Callable type will work.
+            args: List of arguments to be passed
+            kwargs: List of keyword arguments to be passed
 
         Returns:
             None
         """
-        self.sweeps = []
-        for func, values in arm.sweep_functions:
-            self.add_sweep_definition(func, values)
+        raise ValueError("Please use SweepArm instead, or use SimulationBuilder directly!")
 
-        if arm.type == ArmType.cross:
-            self.sweep_definitions.extend(product(*self.sweeps))
-        elif arm.type == ArmType.pair:
-            self.sweep_definitions.extend(zip(*self.sweeps))
+    def add_multiple_parameter_sweep_definition(self, function: TSweepFunction, *args, **kwargs):
+        """
+        Add parameters sweep definition.
+
+        Args:
+            function: The sweep function, which must include a **simulation** parameter (or
+                whatever is specified in :attr:`~idmtools.builders.ExperimentBuilder.SIMULATION_ATTR`).
+                The function also must include EXACTLY ONE free parameter, which the values will be passed to.
+                The function can also be a partial--any Callable type will work.
+            args: List of arguments to be passed
+            kwargs: List of keyword arguments to be passed
+
+        Returns:
+            None
+        """
+        raise ValueError("Please use SweepArm instead, or use SimulationBuilder directly!")
 
     def __iter__(self):
         """
         Iterator for the simulations defined.
-
         Returns:
             Iterator
         """