idmtools.builders.arm_simulation_builder module#

idmtools arm builder definition.

class idmtools.builders.arm_simulation_builder.ArmSimulationBuilder[source]#

Bases: object

Class that represents an experiment builder.

This particular sweep builder build sweeps in “ARMS”. This is particular useful in situations where you want to sweep parameters that have branches of parameters. For Example, let’s say we have a model with the following parameters: * population * susceptible * recovered * enable_births * birth_rate

Enable births controls an optional feature that is controlled by the birth_rate parameter. If we want to sweep a set of parameters on population, susceptible with enabled_births set to off but also want to sweep the birth_rate we could do that like so

###################################
# This example provides how you can check if your sweeps are working as expected
###################################
from functools import partial
from idmtools.builders import ArmSimulationBuilder, SweepArm, ArmType
from idmtools.entities.command_task import CommandTask
from idmtools.entities.templated_simulation import TemplatedSimulations
from tabulate import tabulate


def update_parameter(simulation, parameter, value):
    simulation.task.config[parameter] = value


base_task = CommandTask('example')
base_task.config = dict(enable_births=False)
builder = ArmSimulationBuilder()
# Define our first set of sweeps
arm = SweepArm(type=ArmType.cross)
arm.add_sweep_definition(partial(update_parameter, parameter='population'), [500, 1000])
arm.add_sweep_definition(partial(update_parameter, parameter='susceptible'), [0.5, 0.9])
builder.add_arm(arm)
# Now add the sweeps with the birth_rate as well
arm2 = SweepArm(type=ArmType.cross)
arm2.add_sweep_definition(partial(update_parameter, parameter='enable_births'), [True])
arm2.add_sweep_definition(partial(update_parameter, parameter='birth_rate'), [0.01, 0.1])
builder.add_arm(arm2)

sims = TemplatedSimulations(base_task=base_task)
sims.add_builder(builder)

print(tabulate([s.task.config for s in sims], headers="keys"))

This would result in the output

Arm Example Values#
enable_births	population	susceptible	birth_rate
False	500	0.5
False	500	0.9
False	1000	0.5
False	1000	0.9
True	500	0.5	0.01
True	500	0.5	0.1
True	500	0.9	0.01
True	500	0.9	0.1
True	1000	0.5	0.01
True	1000	0.5	0.1
True	1000	0.9	0.01
True	1000	0.9	0.1

Examples

"""
        This file demonstrates how to use ArmExperimentBuilder in PythonExperiment's builder.
        We are then adding the builder to PythonExperiment.

        |__sweep arm1
            |_ a = 1
            |_ b = [2,3]
            |_ c = [4,5]
        |__ sweep arm2
            |_ a = [6,7]
            |_ b = 2
        Expect sims with parameters:
            sim1: {a:1, b:2, c:4}
            sim2: {a:1, b:2, c:5}
            sim3: {a:1, b:3, c:4}
            sim4: {a:1, b:3, c:5}
            sim5: {a:6, b:2}
            sim6: {a:7, b:2}
        Note:
            arm1 and arm2 are adding to total simulations
"""
import os
import sys
from functools import partial

from idmtools.builders import SweepArm, ArmType, ArmSimulationBuilder
from idmtools.core.platform_factory import platform
from idmtools.entities.experiment import Experiment
from idmtools.entities.templated_simulation import TemplatedSimulations
from idmtools_models.python.json_python_task import JSONConfiguredPythonTask
from idmtools_test import COMMON_INPUT_PATH

# define specific callbacks for a, b, and c
setA = partial(JSONConfiguredPythonTask.set_parameter_sweep_callback, param="a")
setB = partial(JSONConfiguredPythonTask.set_parameter_sweep_callback, param="b")
setC = partial(JSONConfiguredPythonTask.set_parameter_sweep_callback, param="c")


if __name__ == "__main__":
    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)

        # define our first sweep Sweep Arm
        arm1 = SweepArm(type=ArmType.cross)
        builder = ArmSimulationBuilder()
        arm1.add_sweep_definition(setA, 1)
        arm1.add_sweep_definition(setB, [2, 3])
        arm1.add_sweep_definition(setC, [4, 5])
        builder.add_arm(arm1)

        # adding more simulations with sweeping
        arm2 = SweepArm(type=ArmType.cross)
        arm2.add_sweep_definition(setA, [6, 7])
        arm2.add_sweep_definition(setB, [2])
        builder.add_arm(arm2)

        # add our builders to our template
        ts.add_builder(builder)

        # create experiment from the template
        experiment = Experiment.from_template(ts, name=os.path.split(sys.argv[0])[1],
                                              tags={"string_tag": "test", "number_tag": 123, "KeyOnly": None})
        # run the experiment
        experiment.run()
        # in most real scenarios, you probably do not want to wait as this will wait until all simulations
        # associated with an experiment are done. We do it in our examples to show feature and to enable
        # testing of the scripts
        experiment.wait()
        # use system status as the exit code
        sys.exit(0 if experiment.succeeded else -1)

__init__()[source]#: Constructor.

property count#: Total simulations to be built by builder. :returns: arm count

add_arm(arm)[source]#

Add arm sweep definition. :param arm: Arm to add

Returns:: None