def test_SimulationContext_finalize(base_config, components):
sim = SimulationContext(base_config, components)
listener = [c for c in components if 'listener' in c.args][0]
sim.setup()
assert not listener.simulation_end_called
sim.finalize()
assert listener.simulation_end_called
def test_SimulationContext_setup_default(base_config, components):
sim = SimulationContext(base_config, components)
listener = [c for c in components if 'listener' in c.args][0]
assert not listener.post_setup_called
sim.setup()
unpacked_components = []
for c in components:
unpacked_components.append(c)
if hasattr(c, 'sub_components'):
unpacked_components.extend(c.sub_components)
unpacked_components.append(Metrics())
for a, b in zip(sim._component_manager._components, unpacked_components):
assert type(a) == type(b)
if hasattr(a, 'args'):
assert a.args == b.args
assert is_same_object_method(sim.simulant_creator, sim._population._create_simulants)
assert sim.time_step_events == ['time_step__prepare', 'time_step', 'time_step__cleanup', 'collect_metrics']
for k in sim.time_step_emitters.keys():
assert is_same_object_method(sim.time_step_emitters[k], sim._events._event_types[k].emit)
assert is_same_object_method(sim.end_emitter, sim._events._event_types['simulation_end'].emit)
assert listener.post_setup_called
def test_SimulationContext_setup_default(base_config, components):
sim = SimulationContext(base_config, components)
listener = [c for c in components if 'listener' in c.args][0]
assert not listener.post_setup_called
sim.setup()
def unpack(component):
if isinstance(component, MockComponentB) and len(component.args) > 1:
return [component
] + [MockComponentB(arg) for arg in component.args]
return [component]
unpacked_components = [
c for component in components for c in unpack(component)
]
unpacked_components.insert(len(components), Metrics())
for a, b in zip(sim.component_manager._components, unpacked_components):
assert type(a) == type(b)
if hasattr(a, 'args'):
assert a.args == b.args
assert is_same_object_method(sim.simulant_creator,
sim.population._create_simulants)
assert sim.time_step_events == [
'time_step__prepare', 'time_step', 'time_step__cleanup',
'collect_metrics'
]
for k in sim.time_step_emitters.keys():
assert is_same_object_method(sim.time_step_emitters[k],
sim.events._event_types[k].emit)
assert is_same_object_method(
sim.end_emitter, sim.events._event_types['simulation_end'].emit)
assert listener.post_setup_called
def test_SimulationContext_report(base_config, components):
sim = SimulationContext(base_config, components)
sim.setup()
sim.initialize_simulants()
sim.run()
sim.finalize()
metrics = sim.report()
assert metrics['test'] == len([c for c in sim._component_manager._components if isinstance(c, MockComponentB)])
def test_SimulationContext_initialize_simulants(base_config, components):
sim = SimulationContext(base_config, components)
sim.setup()
pop_size = sim.configuration.population.population_size
current_time = sim.clock.time
assert sim.population.population.empty
sim.initialize_simulants()
assert len(sim.population.population) == pop_size
assert sim.clock.time == current_time
def setup_simulation(components, population_size=100, start=None):
simulation = SimulationContext(components)
simulation.setup()
if start:
simulation.current_time = start
else:
year_start = config.simulation_parameters.year_start
simulation.current_time = datetime(year_start, 1, 1)
if 'initial_age' in config.simulation_parameters:
simulation.population._create_simulants(
population_size,
population_configuration={
'initial_age': config.simulation_parameters.initial_age
})
else:
simulation.population._create_simulants(population_size)
return simulation
def run_timed_simulation(model_specification_file):
simulation = SimulationContext(model_specification_file, None, None, None)
setup_start_time = time()
simulation.setup()
setup_time = time() - setup_start_time
initialize_start_time = time()
simulation.initialize_simulants()
initialize_time = time() - initialize_start_time
run_start_time = time()
simulation.run()
run_time = time() - run_start_time
finalize_start_time = time()
simulation.finalize()
finalize_time = time() - finalize_start_time
return setup_time, initialize_time, run_time, finalize_time
def test_SimulationContext_step(log, base_config, components):
sim = SimulationContext(base_config, components)
sim.setup()
current_time = sim.clock.time
step_size = sim.clock.step_size
listener = [c for c in components if 'listener' in c.args][0]
assert not listener.time_step__prepare_called
assert not listener.time_step_called
assert not listener.time_step__cleanup_called
assert not listener.collect_metrics_called
sim.step()
assert log.debug.called_once_with(current_time)
assert listener.time_step__prepare_called
assert listener.time_step_called
assert listener.time_step__cleanup_called
assert listener.collect_metrics_called
assert sim.clock.time == current_time + step_size
def worker(parameters: Mapping):
node = f"{os.environ['SGE_CLUSTER_NAME']}:{os.environ['HOSTNAME']}"
job = f"{os.environ['JOB_NAME']}: {os.environ['JOB_ID']}:{os.environ['SGE_TASK_ID']}"
input_draw = parameters['input_draw']
random_seed = parameters['random_seed']
model_specification_file = parameters['model_specification_file']
branch_config = parameters['branch_configuration']
logger.info(f'Launching new job {job} on {node}')
logger.info('Starting job: {}'.format(
(input_draw, random_seed, model_specification_file, branch_config)))
try:
np.random.seed([input_draw, random_seed])
worker_ = get_current_job().id
from vivarium.framework.engine import SimulationContext
from vivarium.framework.utilities import collapse_nested_dict
configuration = {}
run_key = {'input_draw': input_draw, 'random_seed': random_seed}
if branch_config is not None:
configuration.update(dict(branch_config))
run_key.update(dict(branch_config))
configuration.update({
'run_configuration': {
'input_draw_number': input_draw,
'run_id': str(worker_) + '_' + str(time()),
'results_directory': parameters['results_path'],
'run_key': run_key,
},
'randomness': {
'random_seed': random_seed,
'additional_seed': input_draw,
},
'input_data': {
'input_draw_number': input_draw,
}
})
sim = SimulationContext(model_specification_file,
configuration=configuration)
logger.info('Simulation configuration:')
logger.info(str(sim.configuration))
start = time()
logger.info('Beginning simulation setup.')
sim.setup()
sim.initialize_simulants()
logger.info(
f'Simulation setup complete in {(time() - start)/60} minutes.')
sim_start = time()
logger.info('Starting main simulation loop.')
sim.run()
sim.finalize()
metrics = sim.report()
end = time()
start_time = pd.Timestamp(**sim.configuration.time.start.to_dict())
end_time = pd.Timestamp(**sim.configuration.time.end.to_dict())
step_size = pd.Timedelta(days=sim.configuration.time.step_size)
num_steps = int(math.ceil((end_time - start_time) / step_size))
logger.info(
f'Simulation main loop completed in {(end - sim_start)/60} minutes.'
)
logger.info(
f'Average step length was {(end - sim_start)/num_steps} seconds.')
logger.info(f'Total simulation run time {(end - start) / 60} minutes.')
idx = pd.MultiIndex.from_tuples(
[(input_draw, random_seed)],
names=['input_draw_number', 'random_seed'])
output_metrics = pd.DataFrame(metrics, index=idx)
for k, v in collapse_nested_dict(run_key):
output_metrics[k] = v
output = [output_metrics.to_msgpack()]
return output
except Exception:
logger.exception('Unhandled exception in worker')
job = get_current_job()
job.meta['root_exception'] = format_exc()
job.save_meta()
raise
finally:
logger.info('Exiting job: {}'.format(
(input_draw, random_seed, model_specification_file,
branch_config)))
def work_horse(job_parameters: dict) -> pd.DataFrame:
node = f"{ENV_VARIABLES.HOSTNAME.value}"
job = f"{ENV_VARIABLES.JOB_ID.value}:{ENV_VARIABLES.TASK_ID.value}"
job_parameters = JobParameters(**job_parameters)
logger.info(f"Launching new job {job} on {node}")
logger.info(f"Starting job: {job_parameters}")
try:
configuration = job_parameters.branch_configuration
# TODO: Need to test serialization of an empty dict, then this
# can go away. If you're successfully running code and this
# assert is still here, delete it.
assert configuration is not None
configuration.update({
"run_configuration": {
"run_id": str(get_current_job().id) + "_" + str(time()),
"results_directory": job_parameters.results_path,
"run_key": job_parameters.job_specific,
},
"randomness": {
"random_seed": job_parameters.random_seed,
"additional_seed": job_parameters.input_draw,
},
"input_data": {
"input_draw_number": job_parameters.input_draw,
},
})
sim = SimulationContext(job_parameters.model_specification,
configuration=configuration)
logger.info("Simulation configuration:")
logger.info(str(sim.configuration))
start_time = pd.Timestamp(**sim.configuration.time.start.to_dict())
end_time = pd.Timestamp(**sim.configuration.time.end.to_dict())
step_size = pd.Timedelta(days=sim.configuration.time.step_size)
num_steps = int(math.ceil((end_time - start_time) / step_size))
start_snapshot = CounterSnapshot()
event = {"start": time()} # timestamps of application events
logger.info("Beginning simulation setup.")
sim.setup()
event["simulant_initialization_start"] = time()
exec_time = {
"setup_minutes":
(event["simulant_initialization_start"] - event["start"]) / 60
} # execution event
logger.info(
f'Simulation setup completed in {exec_time["setup_minutes"]:.3f} minutes.'
)
sim.initialize_simulants()
event["simulation_start"] = time()
exec_time["simulant_initialization_minutes"] = (
event["simulation_start"] -
event["simulant_initialization_start"]) / 60
logger.info(
f'Simulant initialization completed in {exec_time["simulant_initialization_minutes"]:.3f} minutes.'
)
logger.info(
f"Starting main simulation loop with {num_steps} time steps")
sim.run()
event["results_start"] = time()
exec_time["main_loop_minutes"] = (event["results_start"] -
event["simulation_start"]) / 60
exec_time["step_mean_seconds"] = (
event["results_start"] - event["simulation_start"]) / num_steps
logger.info(
f'Simulation main loop completed in {exec_time["main_loop_minutes"]:.3f} minutes.'
)
logger.info(
f'Average step length was {exec_time["step_mean_seconds"]:.3f} seconds.'
)
sim.finalize()
metrics = sim.report(print_results=False)
event["end"] = time()
end_snapshot = CounterSnapshot()
do_sim_epilogue(start_snapshot, end_snapshot, event, exec_time,
job_parameters)
idx = pd.MultiIndex.from_tuples(
[(job_parameters.input_draw, job_parameters.random_seed)],
names=["input_draw_number", "random_seed"],
)
output_metrics = pd.DataFrame(metrics, index=idx)
for k, v in collapse_nested_dict(job_parameters.branch_configuration):
output_metrics[k] = v
return output_metrics
except Exception:
logger.exception("Unhandled exception in worker")
job = get_current_job()
job.meta["root_exception"] = format_exc()
job.save_meta()
raise
finally:
logger.info(f"Exiting job: {job_parameters}")