idmtools.builders.arm_simulation_builder module

idmtools arm builder definition.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.

class idmtools.builders.arm_simulation_builder.ArmType(value)

Bases: Enum

ArmTypes.

cross = 0
pair = 1
class idmtools.builders.arm_simulation_builder.SweepArm(type=ArmType.cross, funcs: Optional[List[Tuple[Callable, Iterable]]] = None)

Bases: object

Class that represents a parameter arm.

__init__(type=ArmType.cross, funcs: Optional[List[Tuple[Callable, Iterable]]] = None)

Constructor.

Parameters:
  • type – Type of Arm(Cross or Pair)

  • funcs – Functions to add as sweeps

add_sweep_definition(func: Callable, values: Iterable[Any])

Add Sweep definition.

Parameters:
  • func – Sweep callback

  • values – Values to Sweep

Returns:

None

get_max_values_count()

Get the max values count from different sweep functions.

Returns:

Max values

adjust_values_length()

Adjust values length.

Returns:

None

class idmtools.builders.arm_simulation_builder.ArmSimulationBuilder

Bases: SimulationBuilder

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
arm.add_sweep_definition(partial(update_parameter, parameter='enable_births'), [True])
arm.add_sweep_definition(partial(update_parameter, parameter='birth_rate'), [0.01, 0.1])
builder.add_arm(arm)

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

print(tabulate([s.task.config for s in list(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('BELEGOST'):
        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__()

Constructor.

add_arm(arm)

Add arm sweep definition.

Parameters:

arm – Arm to add

Returns:

None