Source code for idmtools_platform_comps.utils.sweeping

"""
idmtools utility.

Copyright 2025, Gates Foundation. All rights reserved.
"""
import numpy as np
from typing import Dict, Any, List
from idmtools.entities.simulation import Simulation
from logging import getLogger, DEBUG

logger = getLogger()


##################################################
# Sweeping utility functions
##################################################
[docs]def set_param(simulation: Simulation, param: str, value: Any) -> Dict[str, Any]: """ Set a specific parameter value on the simulation task config. Args: simulation (Simulation): idmtools Simulation object. param (str): Name of the parameter to modify. value (Any): New value to set. Returns: Dict[str, Any]: A dictionary containing the parameter name and value. """ try: return simulation.task.set_parameter(param, value) except ValueError: if "parameters" in simulation.task.config: config = simulation.task.config.parameters else: config = simulation.task.config config[param] = value return {param: value}
[docs]def sweep_functions(simulation: Simulation, func_list: List) -> Dict[str, Any]: """ Apply a list of sweep functions to a simulation. Args: simulation (Simulation): The simulation to update. func_list (List[Callable]): List of functions that apply sweeps to the simulation. Returns: Dict[str, Any]: A dictionary of aggregated metadata from each sweep function. """ tags_updated = {} for func in func_list: tags = func(simulation) if tags: tags_updated.update(tags) return tags_updated
[docs]class ItvFn: """ Sweeping utility for modifying interventions (campaign layer) during sweeps. Requirements: - func must accept an emod-api campaign object as its first argument. - func must return a dictionary of metadata (i.e: tags). Returns: Dict[str, Any]: Metadata returned by the intervention function, with numpy types cast to Python types. """
[docs] def __init__(self, func, *args, **kwargs): # noqa: D107 self.func = func self.args = args self.kwargs = kwargs
def __call__(self, simulation: Simulation): # noqa: D102 import emod_api.campaign as campaign campaign.reset() matadata = self.func(campaign, *self.args, **self.kwargs) # Add new campaign events events = campaign.campaign_dict["Events"] simulation.task.campaign.add_events(events) # Handle adhoc (custom) individual events adhoc_events = campaign.get_adhocs() if len(adhoc_events) > 0: if logger.isEnabledFor(DEBUG): logger.debug("Found adhoc events in campaign. Needs some special processing behind the scenes.") if "Custom_Individual_Events" in simulation.task.config.parameters: ev_exist = set(simulation.task.config.parameters.Custom_Individual_Events) ev_addhoc = set(adhoc_events.keys()) simulation.task.config.parameters.Custom_Individual_Events = list(ev_exist.union(ev_addhoc)) else: simulation.task.config.parameters.Report_Event_Recorder_Events.extend(list(set(adhoc_events.keys()))) # Convert numpy types if matadata: for k, v in matadata.items(): if isinstance(v, (np.int64, np.float64, np.float32, np.uint32, np.int16, np.int32)): matadata[k] = v.item() return matadata
[docs]class CfgFn: """ Sweeping utility for modifying simulation configuration parameters. Requirements: - func must accept simulation.task.config as the first parameter. - func must return a dictionary for tagging. Returns: Dict[str, Any]: Metadata dictionary with Python-native types. """
[docs] def __init__(self, func, *args, **kwargs): # noqa: D107 self.func = func self.args = args self.kwargs = kwargs
def __call__(self, simulation: Simulation): # noqa: D102 matadata = self.func(simulation.task.config, *self.args, **self.kwargs) # Make sure we cast numpy types into normal system types if matadata: for k, v in matadata.items(): if isinstance(v, (np.int64, np.float64, np.float32, np.uint32, np.int16, np.int32)): matadata[k] = v.item() return matadata
[docs]class SwpFn: """ Sweeping utility for modifying task-level elements (e.g., reports, demographics, climate). Requirements: - func must accept simulation.task as the first parameter. - func must return a metadata dictionary (tags). Returns: Dict[str, Any]: Metadata with numpy types cast to Python-native types. """
[docs] def __init__(self, func, *args, **kwargs): # noqa: D107 self.func = func self.args = args self.kwargs = kwargs
def __call__(self, simulation: Simulation): # noqa: D102 matadata = self.func(simulation.task, *self.args, **self.kwargs) # Make sure we cast numpy types into normal system types if matadata: for k, v in matadata.items(): if isinstance(v, (np.int64, np.float64, np.float32, np.uint32, np.int16, np.int32)): matadata[k] = v.item() return matadata