Source code for idmtools_calibra.calib_base_app

import os
from pathlib import Path

from idmtools.assets import AssetCollection
from idmtools_calibra.calib_manager import CalibManager
from idmtools_calibra.algorithms.optim_tool import OptimTool
from idmtools_calibra.plotters.likelihood_plotter import LikelihoodPlotter
from idmtools_calibra.plotters.optim_tool_plotter import OptimToolPlotter
from idmtools_calibra.plotters.site_data_plotter import SiteDataPlotter
from idmtools.core.platform_factory import Platform
from idmtools.entities import CommandLine
from idmtools_calibra.singularity_json_python_task import SingularityJSONConfiguredPythonTask as Task
from idmtools_models.python.json_python_task import JSONConfiguredPythonTask

params = None  # hack block
campaign_builder_fn = None
demog_builder_fn = None

def constrain_sample(sample):
    """
    This function is called on every samples and allow the user to edit them before they are passed
    to the map_sample_to_model_input function.
    It is useful to round some parameters as demonstrated below.
    #Can do much more here, e.g. for
    Clinical Fever Threshold High <  MSP1 Merozoite Kill Fraction
    if 'Clinical Fever Threshold High' and "MSP1 Merozoite Kill Fraction" in sample:
    sample['Clinical Fever Threshold High'] = \
        min( sample['Clinical Fever Threshold High'], sample['MSP1 Merozoite Kill Fraction'] ).
    You can omit this function by not specifying it in the OptimTool constructor call below.
    Args:
        sample: The sample coming from the next point algorithm

    Returns: The sample with constrained values

    """
    return sample

campaign_mapper = None
demog_mapper = None

def map_sample_to_model_input(simulation, sample):
    """
    This method maps the samples generated by the next point algorithm to the model inputs (via the simulation object).
    All parameters specified for dynamic calibration (above) that do not have a MapTo value must have associated mapping
    logic in this method.
    Args:
        simulation: idmtools simulation
        sample: The sample containing a values for all the params. e.g. {'Clinical Fever Threshold High':1, ... }

    Returns: A dictionary containing the tags that will be attached to the simulation
    """

    tags = {}
    global params
    from functools import partial

    build_camp_actual = None
    if campaign_builder_fn:
        build_camp_actual = partial( campaign_builder_fn )
    build_demog_actual = None
    if demog_builder_fn:
        build_demog_actual = partial( demog_builder_fn )

    for p in params.CALIBRATION_PARAMETERS:
        if "MapTo" not in p:
            raise Exception(
                "Calibration parameter dicts must include a MapTo (model param name) entry. Missing from parameter: %s"
                % p["Name"]
            )
        if p["Name"] not in sample:
            print(
                "Warning: %s not in sample, perhaps resuming previous iteration"
                % p["Name"]
            )
            continue
        value = sample.pop(p["Name"])
        mapto_key = p["MapTo"]
        if mapto_key.startswith( "campaign:" ):
            if not campaign_mapper:
                raise ValueError( "No campaign mapper function defined." )
            build_camp_actual = campaign_mapper( build_camp_actual, mapto_key, value )
            tags[mapto_key] = f"{value}"
        elif mapto_key.startswith( "demog:" ):
            if not demog_mapper:
                raise ValueError( "No demog mapper function defined." )
            build_demog_actual = demog_mapper( build_demog_actual, mapto_key, value )
            tags[mapto_key] = f"{value}"
        else:
            tags.update(simulation.task.set_parameter(mapto_key.replace( "config:", "" ), value))

    for name, value in sample.items():
        print("UNUSED PARAMETER:" + name)
    assert len(sample) == 0  # All params used

    if build_camp_actual:
        simulation.task.create_campaign_from_callback( builder=build_camp_actual )

    if build_demog_actual:
        simulation.task.create_demog_from_callback( builder=build_demog_actual, from_sweep=True )

    return tags


def init(settings, site, task=None, platform=None):
    # the object representing the resource we will use for running simulations
    global params
    params = settings
    if not platform:
        print( "No platform specified. Creating Platform with Normal proirity." )
        platform = Platform(settings.LOCALE, priority="Normal")

    # If any directories of files were specified to be added as assets of the simulations, add them now.
    assets = AssetCollection()
    for directory in settings.INPUT_DIRS:
        assets.add_directory(directory, relative_path=os.path.basename(directory))

    if task is None:
        command = None
        if "SIF" in settings.__dict__:
            assets.add_assets(AssetCollection.from_id(item_id=settings.SIF))

            command = CommandLine(
                f"singularity exec ./Assets/{settings.SIF_FILENAME} python3 Assets/{Path(settings.MODEL_DRIVER).name}"
            )
            task = Task(
                provided_command=command,
                script_path=str(settings.MODEL_DRIVER),
                common_assets=assets,
                config_file_name=settings.CONFIG_FILENAME
            )
        else:
            #print( '"SIF" not found in settings."' )
            command = CommandLine(
                f"python3 Assets/{Path(settings.MODEL_DRIVER).name}"
            )
            task = JSONConfiguredPythonTask(script_path=str(settings.MODEL_DRIVER), common_assets=assets, 
                config_file_name=settings.CONFIG_FILENAME)
            task.provided_command = command

        #print( str( command ) )

        # The task object defines what to run, how, and what assets will be associated with an experiment of simulations
        # Each calibration iteration will run one experiment (group of simulations)

    # The default plotters used in an Optimization with OptimTool
    plotters = [
        LikelihoodPlotter(combine_sites=True),
        SiteDataPlotter(num_to_plot=settings.num_to_plot, combine_sites=True),
        OptimToolPlotter(),  # OTP must be last because it calls gc.collect()
    ]

    # desired fraction of N-sphere area to unit cube area for numerical derivative (automatic radius scaling with N)
    n_dynamic_parameters = len(
        [p for p in settings.CALIBRATION_PARAMETERS if p["Dynamic"]]
    )

    if n_dynamic_parameters == 0:
        warning_note = """
            /!\\ WARNING /!\\ OptimTool requires at least one dynamic parameter ('Dynamic' set to True). Exiting ...                  
            """
        print(warning_note)
        exit()

    r = OptimTool.get_r(n_dynamic_parameters, settings.volume_fraction)

    # Here we combine all the inputs above to configure OptimTool (our next point algorithm) and the CalibManager
    # (the object that will drive the calibration).
    optimtool = OptimTool(
        settings.CALIBRATION_PARAMETERS,
        constrain_sample,
        rsquared_thresh=0.81,
        mu_r=r,  # <-- radius for numerical derivatve. CAREFUL not to go too small with integer parameters
        sigma_r=r / 10.0,  # <-- stdev of radius
        # this many samples will be repetitions of the existing best-sample (with
        # different run numbers)
        center_repeats=2,
        # Samples per iteration, includes center repeats.
        # Actual number of sims run is this number times number of replicates.
        samples_per_iteration=settings.N_SAMPLES,
    )

    calib_manager = CalibManager(
        name=settings.CALIBRATION_NAME,
        task=task,
        map_sample_to_model_input_fn=map_sample_to_model_input,
        sites=[site],
        next_point=optimtool,
        sim_runs_per_param_set=settings.N_REPLICATES,
        # <-- Number of times to replicate the center (current guess). For comparing intrinsic and extrinsic variability.
        max_iterations=settings.N_ITERATIONS,  # <-- number of iterations to run
        plotters=plotters,
    )
    calib_manager.platform = platform
    return calib_manager


def go(calib_manager, **kwargs):
    calib_manager.run_calibration(**kwargs)