def get_calibrated_params_limited_iters(country, area, multi_beta_calibration,
maxiters):
"""
Returns calibrated parameters using only the first `maxiters` iterations of BO.
"""
state = load_state(calibration_states[country][area])
train_G = state['train_G']
train_G = train_G[:min(maxiters, len(train_G))]
train_theta = state['train_theta']
mob_settings = calibration_mob_paths[country][area][0]
with open(mob_settings, 'rb') as fp:
mob_kwargs = pickle.load(fp)
mob = MobilitySimulator(**mob_kwargs)
data_start_date = calibration_start_dates[country][area]
data_end_date = calibration_lockdown_dates[country]['end']
unscaled_area_cases = collect_data_from_df(
country=country,
area=area,
datatype='new',
start_date_string=data_start_date,
end_date_string=data_end_date)
assert (len(unscaled_area_cases.shape) == 2)
# Scale down cases based on number of people in town and region
sim_cases = downsample_cases(unscaled_area_cases, mob_kwargs)
n_days, n_age = sim_cases.shape
G_obs = torch.tensor(sim_cases).reshape(1, n_days * n_age)
G_obs_aggregate = torch.tensor(sim_cases).sum(dim=-1)
def objective(G):
return -(G - G_obs_aggregate).pow(2).sum(dim=-1) / n_days
train_G_objectives = objective(train_G)
best_observed_idx = train_G_objectives.argmax()
best_observed_obj = train_G_objectives[best_observed_idx].item()
param_bounds = (calibration_model_param_bounds_multi
if multi_beta_calibration else
calibration_model_param_bounds_single)
sim_bounds = pdict_to_parr(pdict=param_bounds,
multi_beta_calibration=multi_beta_calibration).T
normalized_calibrated_params = train_theta[best_observed_idx]
calibrated_params = transforms.unnormalize(normalized_calibrated_params,
sim_bounds)
calibrated_params = parr_to_pdict(
parr=calibrated_params, multi_beta_calibration=multi_beta_calibration)
return calibrated_params
def get_test_capacity(country, area, mob_settings, end_date_string='2021-01-01'):
'''
Computes heuristic test capacity in `country` and `area` based
on true case data by determining the maximum daily increase
in positive cases.
'''
unscaled_area_cases = collect_data_from_df(
country=country, area=area, datatype='new',
start_date_string='2020-01-01', end_date_string=end_date_string)
sim_cases = downsample_cases(unscaled_area_cases, mob_settings)
daily_increase = sim_cases.sum(axis=1)[1:] - sim_cases.sum(axis=1)[:-1]
test_capacity = int(np.round(daily_increase.max()))
return test_capacity
def make_bayes_opt_functions(args):
'''
Generates and returns functions used to run Bayesian optimization
Argument:
args: Keyword arguments specifying exact settings for optimization
Returns:
objective : objective maximized for BO
generate_initial_observations : function to generate initial observations
initialize_model : function to initialize GP
optimize_acqf_and_get_observation : function to optimize acquisition function based on model
case_diff : computes case difference between prediction array and ground truth at t=T
unnormalize_theta : converts BO params to simulation params (unit cube to real parameters)
header : header lines to be printed to log file
'''
header = []
# depending on mode, set parameter bounds
if args.measures_optimized:
param_bounds = settings_measures_param_bounds
else:
param_bounds = settings_model_param_bounds
# remember line executed
header.append('=' * 100)
header.append(datetime.now().strftime("%d/%m/%Y %H:%M:%S"))
header.append('python ' + ' '.join(sys.argv))
header.append('=' * 100)
mob_settings = args.mob
data_area = args.area
data_country = args.country
# initialize mobility object to obtain information (no trace generation yet)
with open(mob_settings, 'rb') as fp:
kwargs = pickle.load(fp)
mob = MobilitySimulator(**kwargs)
# data settings
verbose = not args.not_verbose
use_households = not args.no_households
data_start_date = args.start
data_end_date = args.end
debug_simulation_days = args.endsimat
# simulation settings
n_init_samples = args.ninit
n_iterations = args.niters
simulation_roll_outs = args.rollouts
cpu_count = args.cpu_count
dynamic_tracing = not args.no_dynamic_tracing
load_observations = args.load
# set testing parameters
testing_params = settings_testing_params
# BO acquisition function optimization (Knowledge gradient)
acqf_opt_num_fantasies = args.acqf_opt_num_fantasies
acqf_opt_num_restarts = args.acqf_opt_num_restarts
acqf_opt_raw_samples = args.acqf_opt_raw_samples
acqf_opt_batch_limit = args.acqf_opt_batch_limit
acqf_opt_maxiter = args.acqf_opt_maxiter
"""
Bayesian optimization pipeline
"""
# Import Covid19 data
# Shape (max_days, num_age_groups)
new_cases_ = collect_data_from_df(country=data_country, area=data_area, datatype='new',
start_date_string=data_start_date, end_date_string=data_end_date)
assert(len(new_cases_.shape) == 2)
if new_cases_[0].sum() == 0:
print('No positive cases at provided start time; cannot seed simulation.\n'
'Consider setting a later start date for calibration using the "--start" flag.')
exit(0)
# Scale down cases based on number of people in town, region, and downsampling
new_cases = np.ceil(
(new_cases_ * mob.num_people_unscaled) /
(mob.downsample * mob.region_population))
num_age_groups = new_cases.shape[1]
header.append('Downsampling : ' + str(mob.downsample))
header.append('Town population: ' + str(mob.num_people))
header.append('Town population (unscaled): ' + str(mob.num_people_unscaled))
header.append('Region population : ' + str(mob.region_population))
# Set test capacity per day as (a) command line; or (b) maximum daily positive case increase over observed period
if args.testingcap:
testing_params['tests_per_batch'] = (args.testingcap / mob.num_people_unscaled)
else:
daily_increase = new_cases.sum(axis=1)[1:] - new_cases.sum(axis=1)[:-1]
testing_params['tests_per_batch'] = int(daily_increase.max())
test_lag_days = int(testing_params['test_reporting_lag'] / TO_HOURS)
assert(int(testing_params['test_reporting_lag']) % 24 == 0)
# generate initial seeds based on case numbers
initial_seeds = gen_initial_seeds(new_cases)
header.append('Initial seed counts : ' + str(initial_seeds))
# in debug mode, shorten time of simulation, shorten time
if debug_simulation_days:
new_cases = new_cases[:debug_simulation_days]
# Maximum time fixed by real data, init mobility simulator simulation
# maximum time to simulate, in hours
max_time = int(new_cases.shape[0] * TO_HOURS)
max_time += TO_HOURS * test_lag_days # longer due to test lag in simulations
testing_params['testing_t_window'] = [0.0, max_time]
mob.simulate(max_time=max_time, dynamic_tracing=True)
header.append(
'Daily test capacity in sim.: ' + str(testing_params['tests_per_batch']))
header.append(
'Max time T (days): ' + str(new_cases.shape[0]))
header.append(
'Target cases per age group at t=0: ' + str(list(map(int, new_cases[0].tolist()))))
header.append(
'Target cases per age group at t=T: ' + str(list(map(int, new_cases[-1].tolist()))))
# instantiate correct distributions
distributions = CovidDistributions(country=args.country)
# set Bayesian optimization target as positive cases
n_days, n_age = new_cases.shape
G_obs = torch.tensor(new_cases).reshape(n_days * n_age) # flattened
sim_bounds = pdict_to_parr(param_bounds, measures_optimized=args.measures_optimized).T
n_params = sim_bounds.shape[1]
header.append(f'Parameters : {n_params}')
header.append('Parameter bounds: ' + str(parr_to_pdict(sim_bounds.T, measures_optimized=args.measures_optimized)))
# extract lockdown period
sim_start_date = pd.to_datetime(args.start)
sim_end_date = sim_start_date + timedelta(days=int(max_time / TO_HOURS))
lockdown_start_date = pd.to_datetime(
settings_lockdown_dates[args.country]['start'])
lockdown_end_date = pd.to_datetime(
settings_lockdown_dates[args.country]['end'])
days_until_lockdown_start = (lockdown_start_date - sim_start_date).days
days_until_lockdown_end = (lockdown_end_date - sim_start_date).days
header.append(f'Simulation starts at : {sim_start_date}')
header.append(f' ends at : {sim_end_date}')
header.append(f'Lockdown starts at : {lockdown_start_date}')
header.append(f' ends at : {lockdown_end_date}')
# create settings dictionary for simulations
launch_kwargs = dict(
mob_settings=mob_settings,
distributions=distributions,
random_repeats=simulation_roll_outs,
cpu_count=cpu_count,
initial_seeds=initial_seeds,
testing_params=testing_params,
max_time=max_time,
num_people=mob.num_people,
num_sites=mob.num_sites,
home_loc=mob.home_loc,
site_loc=mob.site_loc,
dynamic_tracing=dynamic_tracing,
verbose=False)
'''
Define central functions for optimization
'''
G_obs = torch.tensor(new_cases).reshape(1, n_days * n_age)
def composite_squared_loss(G):
'''
Objective function
Note: in BO, objectives are maximized
'''
return - (G - G_obs).pow(2).sum(dim=-1)
# select objective
objective = GenericMCObjective(composite_squared_loss)
def case_diff(preds):
'''
Computes case difference of predictions and ground truth at t=T
'''
return preds.reshape(n_days, n_age)[-1].sum() - torch.tensor(new_cases)[-1].sum()
def unnormalize_theta(theta):
'''
Computes unnormalized parameters
'''
return transforms.unnormalize(theta, sim_bounds)
def composite_simulation(norm_params):
"""
Takes a set of normalized (unit cube) BO parameters
and returns simulator output means and standard errors based on multiple
random restarts. This corresponds to the black-box function.
"""
# un-normalize normalized params to obtain simulation parameters
params = transforms.unnormalize(norm_params, sim_bounds)
# finalize settings based which parameters are calibrated
kwargs = copy.deepcopy(launch_kwargs)
if args.measures_optimized:
'''
Measures are calibrated
'''
measure_params = parr_to_pdict(params, measures_optimized=args.measures_optimized)
# social distancing measures: calibration is only done for `SocialDistancingForAllMeasure` for now
measure_list_ = [
SocialDistancingForPositiveMeasure(
t_window=Interval(0.0, max_time), p_stay_home=1.0),
SocialDistancingForPositiveMeasureHousehold(
t_window=Interval(0.0, max_time), p_isolate=1.0),
SocialDistancingForAllMeasure(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
p_stay_home=measure_params['p_stay_home']),
]
# close sites if specified
if args.measures_close:
beta_multipliers = {'education': 1.0, 'social': 1.0,
'bus_stop': 1.0, 'office': 1.0, 'supermarket': 1.0}
for category in args.measures_close:
if category in beta_multipliers.keys():
beta_multipliers[category] = 0.0
else:
raise ValueError(f'Site type `{category}` passed in `--measures_close` is invalid.\n'
f'Available are {str(list(beta_multipliers.keys()))}')
measure_list_.append(BetaMultiplierMeasureByType(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
beta_multiplier=beta_multipliers
))
kwargs['measure_list'] = MeasureList(measure_list_)
# get optimized model paramters for this country and area
calibrated_model_params = settings_optimized_town_params[args.country][args.area]
if calibrated_model_params is None:
raise ValueError(f'Cannot optimize measures for {args.country}-{args.area} because model parameters '
'have not been fitted yet. Set values in `calibration_settings.py`')
kwargs['params'] = calibrated_model_params
else:
'''
Model parameters calibrated
'''
kwargs['measure_list'] = MeasureList([
SocialDistancingForPositiveMeasure(
t_window=Interval(0.0, max_time), p_stay_home=1.0),
SocialDistancingForPositiveMeasureHousehold(
t_window=Interval(0.0, max_time), p_isolate=1.0),
])
kwargs['params'] = parr_to_pdict(params, measures_optimized=args.measures_optimized)
# run simulation in parallel,
summary = launch_parallel_simulations(**kwargs)
# (random_repeats, n_people)
posi_started = torch.tensor(summary.state_started_at['posi'])
posi_started -= test_lag_days * TO_HOURS # account for test lag
# (random_repeats, n_days)
age_groups = torch.tensor(summary.people_age)
posi_cumulative = convert_timings_to_cumulative_daily(
timings=posi_started, age_groups=age_groups, time_horizon=n_days * TO_HOURS)
if posi_cumulative.shape[0] <= 1:
raise ValueError('Must run at least 2 random restarts per setting to get estimate of noise in observation.')
# compute mean and standard error of means
G = torch.mean(posi_cumulative, dim=0)
G_sem = torch.std(posi_cumulative, dim=0) / math.sqrt(posi_cumulative.shape[0])
# make sure noise is not zero for non-degerateness
G_sem = torch.max(G_sem, MIN_NOISE)
# flatten
G = G.reshape(1, n_days * n_age)
G_sem = G_sem.reshape(1, n_days * n_age)
return G, G_sem
def generate_initial_observations(n, logger):
"""
Takes an integer `n` and generates `n` initial observations
from the black box function using Sobol random parameter settings
in the unit cube. Returns parameter setting and black box function outputs
"""
if n <= 0:
raise ValueError(
'qKnowledgeGradient and GP needs at least one observation to be defined properly.')
# sobol sequence
# new_thetas: [n, n_params]
new_thetas = torch.tensor(
sobol_seq.i4_sobol_generate(n_params, n), dtype=torch.float)
# simulator observations
# new_G, new_G_sem: [n, n_days * n_age] (flattened outputs)
new_G = torch.zeros((n, n_days * n_age), dtype=torch.float)
new_G_sem = torch.zeros((n, n_days * n_age), dtype=torch.float)
for i in range(n):
t0 = time.time()
# get mean and standard error of mean (sem) of every simulation output
G, G_sem = composite_simulation(new_thetas[i, :])
new_G[i, :] = G
new_G_sem[i, :] = G_sem
# log
G_objectives = objective(new_G[:i+1])
best_idx = G_objectives.argmax()
best = G_objectives[best_idx].item()
current = objective(G).item()
case_diff = (
G.reshape(n_days, n_age)[-1].sum()
- G_obs.reshape(n_days, n_age)[-1].sum())
t1 = time.time()
logger.log(
i=i - n,
time=t1 - t0,
best=best,
objective=current,
case_diff=case_diff,
theta=transforms.unnormalize(new_thetas[i, :].detach().squeeze(), sim_bounds)
)
# save state
state = {
'train_theta': new_thetas[:i+1],
'train_G': new_G[:i+1],
'train_G_sem': new_G_sem[:i+1],
'best_observed_obj': best,
'best_observed_idx': best_idx,
}
save_state(state, logger.filename + '_init')
# compute best objective from simulations
f = objective(new_G)
best_f_idx = f.argmax()
best_f = f[best_f_idx].item()
return new_thetas, new_G, new_G_sem, best_f, best_f_idx
def initialize_model(train_x, train_y, train_y_sem):
"""
Defines a GP given X, Y, and noise observations (standard error of mean)
"""
train_ynoise = train_y_sem.pow(2.0) # noise is in variance units
# standardize outputs to zero mean, unit variance to have good hyperparameter tuning
model = FixedNoiseGP(train_x, train_y, train_ynoise, outcome_transform=Standardize(m=n_days * n_age))
# "Loss" for GPs - the marginal log likelihood
mll = ExactMarginalLogLikelihood(model.likelihood, model)
return mll, model
# Model initialization
# parameters used in BO are always in unit cube for optimal hyperparameter tuning of GPs
bo_bounds = torch.stack([torch.zeros(n_params), torch.ones(n_params)])
def optimize_acqf_and_get_observation(acq_func, args):
"""
Optimizes the acquisition function, and returns a new candidate and a noisy observation.
botorch defaults: num_restarts=10, raw_samples=256, batch_limit=5, maxiter=200
"""
batch_initial_conditions = gen_one_shot_kg_initial_conditions(
acq_function=acq_func,
bounds=bo_bounds,
q=1,
num_restarts=args.acqf_opt_num_restarts,
raw_samples=args.acqf_opt_raw_samples,
options={"batch_limit": args.acqf_opt_batch_limit,
"maxiter": args.acqf_opt_maxiter},
)
# optimize acquisition function
candidates, _ = optimize_acqf(
acq_function=acq_func,
bounds=bo_bounds,
q=1,
num_restarts=args.acqf_opt_num_restarts,
raw_samples=args.acqf_opt_raw_samples, # used for intialization heuristic
options={"batch_limit": args.acqf_opt_batch_limit,
"maxiter": args.acqf_opt_maxiter},
batch_initial_conditions=batch_initial_conditions
)
# proposed evaluation
new_theta = candidates.detach()
# observe new noisy function evaluation
new_G, new_G_sem = composite_simulation(new_theta.squeeze())
return new_theta, new_G, new_G_sem
# return functions
return (
objective,
generate_initial_observations,
initialize_model,
optimize_acqf_and_get_observation,
case_diff,
unnormalize_theta,
header,
)
def add(self,
*,
simulation_info,
country,
area,
measure_list,
full_scale=True,
test_update=None,
seed_summary_path=None,
set_calibrated_params_to=None,
set_initial_seeds_to=None,
expected_daily_base_expo_per100k=0,
beacon_config=None,
thresholds_roc=None,
estimate_mobility_reduction=False,
store_mob=False):
# Set time window based on experiment start and end date
sim_days = (pd.to_datetime(self.end_date) -
pd.to_datetime(self.start_date)).days
max_time = TO_HOURS * sim_days # in hours
# extract lockdown period
lockdown_start_date = pd.to_datetime(
calibration_lockdown_dates[country]['start'])
lockdown_end_date = pd.to_datetime(
calibration_lockdown_dates[country]['end'])
days_until_lockdown_start = (lockdown_start_date -
pd.to_datetime(self.start_date)).days
days_until_lockdown_end = (lockdown_end_date -
pd.to_datetime(self.start_date)).days
# Load mob settings
mob_settings_file = calibration_mob_paths[country][area][
1 if full_scale else 0]
with open(mob_settings_file, 'rb') as fp:
mob_settings = pickle.load(fp)
num_age_groups = len(mob_settings['mob_rate_per_age_per_type'])
# Obtain COVID19 case date for country and area to estimate testing capacity and heuristic seeds if necessary
unscaled_area_cases = collect_data_from_df(
country=country,
area=area,
datatype='new',
start_date_string=self.start_date,
end_date_string=self.end_date)
assert (len(unscaled_area_cases.shape) == 2)
# Scale down cases based on number of people in town and region
sim_cases = downsample_cases(unscaled_area_cases, mob_settings)
# Instantiate correct state transition distributions (estimated from literature)
distributions = CovidDistributions(country=country)
# Expected base rate infections
if expected_daily_base_expo_per100k > 0.0:
# Scale expectation to simulation size
num_people = len(mob_settings['home_loc'])
lambda_base_expo_population = expected_daily_base_expo_per100k * (
num_people / 100000)
# Convert to individual base rate by dividing by population size; priority queue handles superposition
lambda_base_expo_indiv = lambda_base_expo_population / num_people
# Poisson process with rate lambda: interarrival times are Exponential r.v. with mean = 1 / lambda
# Hence set rate of Expo r.v.s to 1 / (1 / lambda) = lambda
distributions.lambda_0 = lambda_base_expo_indiv
# Get initial seeds for simulation
# (a) Define heuristically based on true cases and literature distribution estimates
if seed_summary_path is None:
# Generate initial seeds based on unscaled case numbers in town
initial_seeds = gen_initial_seeds(sim_cases, day=0)
if sum(initial_seeds.values()) == 0:
print(
'No states seeded at start time; cannot start simulation.\n'
'Consider setting a later start date for calibration using the "--start" flag.'
)
sys.exit(0)
# (b) Define based state of previous batch of simulations,
# using the random rollout that best matched the true cases in terms of squared error
else:
seed_summary_ = load_summary(seed_summary_path)
seed_day_ = seed_summary_.max_time # take seeds at the end of simulation
initial_seeds = extract_seeds_from_summary(seed_summary_,
seed_day_, sim_cases)
if set_initial_seeds_to is not None:
initial_seeds = set_initial_seeds_to
if set_calibrated_params_to is not None:
calibrated_params = set_calibrated_params_to
else:
# Load calibrated model parameters for this area
calibrated_params = get_calibrated_params(
country=country,
area=area,
multi_beta_calibration=self.multi_beta_calibration,
estimate_mobility_reduction=estimate_mobility_reduction)
if self.multi_beta_calibration:
betas = calibrated_params['betas']
else:
betas = {
'education': calibrated_params['beta_site'],
'social': calibrated_params['beta_site'],
'bus_stop': calibrated_params['beta_site'],
'office': calibrated_params['beta_site'],
'supermarket': calibrated_params['beta_site'],
}
model_params = {
'betas': betas,
'beta_household': calibrated_params['beta_household'],
}
# Add standard measure of positives staying isolated
measure_list += [
# standard behavior of positively tested: full isolation
SocialDistancingForPositiveMeasure(t_window=Interval(
0.0, max_time),
p_stay_home=1.0),
SocialDistancingForPositiveMeasureHousehold(t_window=Interval(
0.0, max_time),
p_isolate=1.0),
]
measure_list = MeasureList(measure_list)
testing_params = copy.deepcopy(calibration_testing_params)
testing_params['testing_t_window'] = [0.0, max_time]
if test_update:
testing_params = test_update(testing_params)
# store simulation
sim_kwargs = dict(
# Generic information
experiment_info=self.experiment_info,
simulation_info=simulation_info,
start_date=self.start_date,
end_date=self.end_date,
sim_days=sim_days,
country=country,
area=area,
random_repeats=self.random_repeats,
# Mobility and measures
mob_settings_file=mob_settings_file,
full_scale=full_scale,
measure_list=measure_list,
testing_params=testing_params,
store_mob=store_mob,
# Model
model_params=model_params,
distributions=distributions,
initial_seeds=initial_seeds,
)
# Beacon
# fields are added here (even though defaulting to `None`) to double check backwards compatibility
# with stored `Result` objects prior to implementing beacon functionality
if beacon_config is not None:
sim_kwargs['beacon_config'] = beacon_config
if thresholds_roc is not None:
sim_kwargs['thresholds_roc'] = thresholds_roc
sim = Simulation(**sim_kwargs)
if self.continued_run and self.check_summary_existence(sim):
if self.verbose:
print(f'[Skipped Sim] {self.get_sim_path(sim)}')
else:
self.sims.append(sim)
if self.verbose:
print(f'[Added Sim] {self.get_sim_path(self.sims[-1])}')
def make_bayes_opt_functions(args):
'''
Generates and returns functions used to run Bayesian optimization
Argument:
args: Keyword arguments specifying exact settings for optimization
Returns:
objective : objective maximized for BO
generate_initial_observations : function to generate initial observations
initialize_model : function to initialize GP
optimize_acqf_and_get_observation : function to optimize acquisition function based on model
case_diff : computes case difference between prediction array and ground truth at t=T
unnormalize_theta : converts BO params to simulation params (unit cube to real parameters)
header : header lines to be printed to log file
'''
header = []
# set parameter bounds based on calibration mode (single beta vs multiple beta)
multi_beta_calibration = args.multi_beta_calibration
if multi_beta_calibration:
param_bounds = calibration_model_param_bounds_multi
else:
param_bounds = calibration_model_param_bounds_single
# remember line executed
header.append('=' * 100)
header.append(datetime.now().strftime("%d/%m/%Y %H:%M:%S"))
header.append('python ' + ' '.join(sys.argv))
header.append('=' * 100)
data_country = args.country
data_area = args.area
mob_settings = args.mob or calibration_mob_paths[data_country][data_area][0] # 0: downscaled, 1: full scale
# initialize mobility object to obtain information (no trace generation yet)
with open(mob_settings, 'rb') as fp:
mob_kwargs = pickle.load(fp)
mob = MobilitySimulator(**mob_kwargs)
# data settings
verbose = not args.not_verbose
use_households = not args.no_households
data_start_date = args.start or calibration_start_dates[data_country][data_area]
data_end_date = args.end or calibration_lockdown_dates[args.country]['end']
per_age_group_objective = args.per_age_group_objective
# simulation settings
n_init_samples = args.ninit
n_iterations = args.niters
simulation_roll_outs = args.rollouts
cpu_count = args.cpu_count
lazy_contacts = not args.no_lazy_contacts
load_observations = args.load
# set testing parameters
testing_params = calibration_testing_params
# BO acquisition function optimization (Knowledge gradient)
acqf_opt_num_fantasies = args.acqf_opt_num_fantasies
acqf_opt_num_restarts = args.acqf_opt_num_restarts
acqf_opt_raw_samples = args.acqf_opt_raw_samples
acqf_opt_batch_limit = args.acqf_opt_batch_limit
acqf_opt_maxiter = args.acqf_opt_maxiter
"""
Bayesian optimization pipeline
"""
# Import Covid19 data
# Shape (max_days, num_age_groups)
unscaled_area_cases = collect_data_from_df(country=data_country, area=data_area, datatype='new',
start_date_string=data_start_date, end_date_string=data_end_date)
assert(len(unscaled_area_cases.shape) == 2)
# Scale down cases based on number of people in town and region
sim_cases = downsample_cases(unscaled_area_cases, mob_kwargs)
# Generate initial seeds based on unscaled case numbers in town
initial_seeds = gen_initial_seeds(
sim_cases, day=0)
if sum(initial_seeds.values()) == 0:
print('No states seeded at start time; cannot start simulation.\n'
'Consider setting a later start date for calibration using the "--start" flag.')
exit(0)
num_age_groups = sim_cases.shape[1]
header.append('Downsampling : {}'.format(mob.downsample))
header.append('Simulation population: {}'.format(mob.num_people))
header.append('Simulation population (unscaled): {}'.format(mob.num_people_unscaled))
header.append('Area population : {}'.format(mob.region_population))
header.append('Initial seed counts : {}'.format(initial_seeds))
scaled_test_capacity = get_test_capacity(
country=data_country, area=data_area,
mob_settings=mob_kwargs, end_date_string=data_end_date)
testing_params['tests_per_batch'] = scaled_test_capacity
test_lag_days = int(testing_params['test_reporting_lag'] / TO_HOURS)
assert(int(testing_params['test_reporting_lag']) % 24 == 0)
# Maximum time fixed by real data, init mobility simulator simulation
# maximum time to simulate, in hours
max_time = int(sim_cases.shape[0] * TO_HOURS)
max_time += TO_HOURS * test_lag_days # simulate longer due to test lag in simulations
testing_params['testing_t_window'] = [0.0, max_time]
mob.simulate(max_time=max_time, lazy_contacts=True)
header.append(
'Target cases per age group at t=0: {} {}'.format(sim_cases[0].sum().item(), list(sim_cases[0].tolist())))
header.append(
'Target cases per age group at t=T: {} {}'.format(sim_cases[-1].sum().item(), list(sim_cases[-1].tolist())))
header.append(
'Daily test capacity in sim.: {}'.format(testing_params['tests_per_batch']))
# instantiate correct distributions
distributions = CovidDistributions(country=args.country)
# set Bayesian optimization target as positive cases
n_days, n_age = sim_cases.shape
sim_bounds = pdict_to_parr(
pdict=param_bounds,
multi_beta_calibration=multi_beta_calibration
).T
n_params = sim_bounds.shape[1]
header.append(f'Parameters : {n_params}')
header.append('Parameter bounds: {}'.format(parr_to_pdict(parr=sim_bounds.T, multi_beta_calibration=multi_beta_calibration)))
# extract lockdown period
sim_start_date = pd.to_datetime(data_start_date)
sim_end_date = sim_start_date + timedelta(days=int(max_time / TO_HOURS))
lockdown_start_date = pd.to_datetime(
calibration_lockdown_dates[args.country]['start'])
lockdown_end_date = pd.to_datetime(
calibration_lockdown_dates[args.country]['end'])
days_until_lockdown_start = (lockdown_start_date - sim_start_date).days
days_until_lockdown_end = (lockdown_end_date - sim_start_date).days
header.append(f'Simulation starts at : {sim_start_date}')
header.append(f' ends at : {sim_end_date}')
header.append(f'Lockdown starts at : {lockdown_start_date}')
header.append(f' ends at : {lockdown_end_date}')
header.append(f'Cases compared until : {pd.to_datetime(data_end_date)}')
header.append(f' for days : {sim_cases.shape[0]}')
# create settings dictionary for simulations
launch_kwargs = dict(
mob_settings=mob_settings,
distributions=distributions,
random_repeats=simulation_roll_outs,
cpu_count=cpu_count,
initial_seeds=initial_seeds,
testing_params=testing_params,
max_time=max_time,
num_people=mob.num_people,
num_sites=mob.num_sites,
home_loc=mob.home_loc,
site_loc=mob.site_loc,
lazy_contacts=lazy_contacts,
verbose=False)
'''
Define central functions for optimization
'''
G_obs = torch.tensor(sim_cases).reshape(1, n_days * n_age)
G_obs_aggregate = torch.tensor(sim_cases).sum(dim=-1)
'''
Objective function
Note: in BO and botorch, objectives are maximized
'''
if per_age_group_objective:
def composite_squared_loss(G):
return - (G - G_obs).pow(2).sum(dim=-1) / n_days
else:
def composite_squared_loss(G):
return - (G - G_obs_aggregate).pow(2).sum(dim=-1) / n_days
# select objective function
objective = GenericMCObjective(composite_squared_loss)
def case_diff(preds):
'''
Computes aggregate case difference of predictions and ground truth at t=T
'''
if per_age_group_objective:
return preds[-1].sum(dim=-1) - G_obs_aggregate[-1]
else:
return preds[-1] - G_obs_aggregate[-1]
def unnormalize_theta(theta):
'''
Computes unnormalized parameters
'''
return transforms.unnormalize(theta, sim_bounds)
def composite_simulation(norm_params):
"""
Takes a set of normalized (unit cube) BO parameters
and returns simulator output means and standard errors based on multiple
random restarts. This corresponds to the black-box function.
"""
# un-normalize normalized params to obtain simulation parameters
params = transforms.unnormalize(norm_params, sim_bounds)
# finalize model parameters based on given parameters and calibration mode
kwargs = copy.deepcopy(launch_kwargs)
all_params = parr_to_pdict(parr=params, multi_beta_calibration=multi_beta_calibration)
if multi_beta_calibration:
betas = all_params['betas']
else:
betas = {
'education': all_params['beta_site'],
'social': all_params['beta_site'],
'bus_stop': all_params['beta_site'],
'office': all_params['beta_site'],
'supermarket': all_params['beta_site'],
}
model_params = {
'betas' : betas,
'beta_household' : all_params['beta_household'],
}
# set exposure parameters
kwargs['params'] = model_params
# set measure parameters
kwargs['measure_list'] = MeasureList([
# standard behavior of positively tested: full isolation
SocialDistancingForPositiveMeasure(
t_window=Interval(0.0, max_time), p_stay_home=1.0),
SocialDistancingForPositiveMeasureHousehold(
t_window=Interval(0.0, max_time), p_isolate=1.0),
# social distancing factor during lockdown: calibrated
SocialDistancingForAllMeasure(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
p_stay_home=all_params['p_stay_home']),
# site specific measures: fixed in advance, outside of calibration
BetaMultiplierMeasureByType(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
beta_multiplier=calibration_lockdown_beta_multipliers)
])
# run simulation in parallel,
summary = launch_parallel_simulations(**kwargs)
# (random_repeats, n_people)
posi_started = torch.tensor(summary.state_started_at['posi'])
posi_started -= test_lag_days * TO_HOURS # account for test lag in objective computation
# (random_repeats, n_days)
age_groups = torch.tensor(summary.people_age)
# (random_repeats, n_days, n_age_groups)
posi_cumulative = convert_timings_to_cumulative_daily(
timings=posi_started, age_groups=age_groups, time_horizon=n_days * TO_HOURS)
if posi_cumulative.shape[0] <= 1:
raise ValueError('Must run at least 2 random restarts per setting to get estimate of noise in observation.')
# compute aggregate if not using objective per age-group
if not per_age_group_objective:
posi_cumulative = posi_cumulative.sum(dim=-1)
# compute mean and standard error of means
G = torch.mean(posi_cumulative, dim=0)
G_sem = torch.std(posi_cumulative, dim=0) / math.sqrt(posi_cumulative.shape[0])
# make sure noise is not zero for non-degenerateness
G_sem = torch.max(G_sem, MIN_NOISE)
# flatten
if per_age_group_objective:
G = G.reshape(n_days * n_age)
G_sem = G_sem.reshape(n_days * n_age)
return G, G_sem
def generate_initial_observations(n, logger, loaded_init_theta=None, loaded_init_G=None, loaded_init_G_sem=None):
"""
Takes an integer `n` and generates `n` initial observations
from the black box function using Sobol random parameter settings
in the unit cube. Returns parameter setting and black box function outputs.
If `loaded_init_theta/G/G_sem` are specified, initialization is loaded (possibly partially, in which
case the initialization using the Sobol random sequence is continued where left off).
"""
if n <= 0:
raise ValueError(
'qKnowledgeGradient and GP needs at least one observation to be defined properly.')
# sobol sequence proposal points
# new_thetas: [n, n_params]
new_thetas = torch.tensor(
sobol_seq.i4_sobol_generate(n_params, n), dtype=torch.float)
# check whether initial observations are loaded
loaded = (loaded_init_theta is not None
and loaded_init_G is not None
and loaded_init_G_sem is not None)
if loaded:
n_loaded = loaded_init_theta.shape[0] # loaded no. of observations total
n_loaded_init = min(n_loaded, n) # loaded no. of quasi-random initialization observations
n_init = max(n_loaded, n) # final no. of observations returned, at least quasi-random initializations
# check whether loaded proposal points are same as without loading observations
try:
assert(np.allclose(loaded_init_theta[:n_loaded_init], new_thetas[:n_loaded_init]))
except AssertionError:
print(
'\n\n\n===> Warning: parameters of loaded inital observations '
'do not coincide with initialization that would have been done. '
'Double check simulation, ninit, and parameter bounds, which could change '
'the initial random Sobol sequence. \nThe loaded parameter settings are used. \n\n\n'
)
if n_init > n:
new_thetas = loaded_init_theta # size of tensor increased to `n_init`, as more than Sobol init points loaded
else:
n_loaded = 0 # loaded no. of observations total
n_loaded_init = 0 # loaded no. of quasi-random initialization observations
n_init = n # final no. of observations returned, at least quasi-random initializations
# instantiate simulator observation tensors
if per_age_group_objective:
# new_G, new_G_sem: [n_init, n_days * n_age] (flattened outputs)
new_G = torch.zeros((n_init, n_days * n_age), dtype=torch.float)
new_G_sem = torch.zeros((n_init, n_days * n_age), dtype=torch.float)
else:
# new_G, new_G_sem: [n_init, n_days]
new_G = torch.zeros((n_init, n_days), dtype=torch.float)
new_G_sem = torch.zeros((n_init, n_days), dtype=torch.float)
# generate `n` initial evaluations at quasi random settings; if applicable, skip and load expensive evaluation result
for i in range(n_init):
# if loaded, use initial observation for this parameter settings
if loaded and i <= n_loaded - 1:
new_thetas[i] = loaded_init_theta[i]
G, G_sem = loaded_init_G[i], loaded_init_G_sem[i]
walltime = 0.0
# if not loaded, evaluate as usual
else:
t0 = time.time()
G, G_sem = composite_simulation(new_thetas[i])
walltime = time.time() - t0
new_G[i] = G
new_G_sem[i] = G_sem
# log
G_objectives = objective(new_G[:i+1])
best_idx = G_objectives.argmax()
best = G_objectives[best_idx].item()
current = objective(G).item()
if per_age_group_objective:
case_diff = G.reshape(n_days, n_age)[-1].sum() - G_obs_aggregate[-1]
else:
case_diff = G[-1] - G_obs_aggregate[-1]
logger.log(
i=i - n,
time=walltime,
best=best,
objective=current,
case_diff=case_diff,
theta=transforms.unnormalize(new_thetas[i, :].detach().squeeze(), sim_bounds)
)
# save state
state = {
'train_theta': new_thetas[:i+1],
'train_G': new_G[:i+1],
'train_G_sem': new_G_sem[:i+1],
'best_observed_obj': best,
'best_observed_idx': best_idx,
}
save_state(state, logger.filename)
# compute best objective from simulations
f = objective(new_G)
best_f_idx = f.argmax()
best_f = f[best_f_idx].item()
return new_thetas, new_G, new_G_sem, best_f, best_f_idx
def initialize_model(train_x, train_y, train_y_sem):
"""
Defines a GP given X, Y, and noise observations (standard error of mean)
"""
train_ynoise = train_y_sem.pow(2.0) # noise is in variance units
# standardize outputs to zero mean, unit variance to have good hyperparameter tuning
outcome_transform = Standardize(m=n_days * n_age if per_age_group_objective else n_days)
model = FixedNoiseGP(train_x, train_y, train_ynoise, outcome_transform=outcome_transform)
# "Loss" for GPs - the marginal log likelihood
mll = ExactMarginalLogLikelihood(model.likelihood, model)
return mll, model
# Model initialization
# parameters used in BO are always in unit cube for optimal hyperparameter tuning of GPs
bo_bounds = torch.stack([torch.zeros(n_params), torch.ones(n_params)])
def optimize_acqf_and_get_observation(acq_func, args):
"""
Optimizes the acquisition function, and returns a new candidate and a noisy observation.
botorch defaults: num_restarts=10, raw_samples=256, batch_limit=5, maxiter=200
"""
batch_initial_conditions = gen_one_shot_kg_initial_conditions(
acq_function=acq_func,
bounds=bo_bounds,
q=1,
num_restarts=args.acqf_opt_num_restarts,
raw_samples=args.acqf_opt_raw_samples,
options={"batch_limit": args.acqf_opt_batch_limit,
"maxiter": args.acqf_opt_maxiter},
)
# optimize acquisition function
candidates, _ = optimize_acqf(
acq_function=acq_func,
bounds=bo_bounds,
q=1,
num_restarts=args.acqf_opt_num_restarts,
raw_samples=args.acqf_opt_raw_samples, # used for intialization heuristic
options={"batch_limit": args.acqf_opt_batch_limit,
"maxiter": args.acqf_opt_maxiter},
batch_initial_conditions=batch_initial_conditions
)
# proposed evaluation
new_theta = candidates.detach().squeeze()
# observe new noisy function evaluation
new_G, new_G_sem = composite_simulation(new_theta)
return new_theta, new_G, new_G_sem
# return functions
return (
objective,
generate_initial_observations,
initialize_model,
optimize_acqf_and_get_observation,
case_diff,
unnormalize_theta,
header,
)
def get_unique_calibration_params(*, country, area, multi_beta_calibration, maxiters=None):
"""
Returns all unique parameter settings that ** improved ** the objective
during calibration for a `country` and an `area`
"""
param_bounds = (
calibration_model_param_bounds_multi
if multi_beta_calibration else
calibration_model_param_bounds_single)
sim_bounds = pdict_to_parr(
pdict=param_bounds,
multi_beta_calibration=multi_beta_calibration
).T
state = load_state(calibration_states[country][area])
train_theta = state['train_theta']
train_G = state['train_G']
mob_settings = calibration_mob_paths[country][area][0]
with open(mob_settings, 'rb') as fp:
mob_kwargs = pickle.load(fp)
mob = MobilitySimulator(**mob_kwargs)
data_start_date = calibration_start_dates[country][area]
data_end_date = calibration_lockdown_dates[country]['end']
unscaled_area_cases = collect_data_from_df(country=country, area=area, datatype='new',
start_date_string=data_start_date, end_date_string=data_end_date)
assert (len(unscaled_area_cases.shape) == 2)
# Scale down cases based on number of people in town and region
sim_cases = downsample_cases(unscaled_area_cases, mob_kwargs)
n_days, n_age = sim_cases.shape
G_obs = torch.tensor(sim_cases).reshape(1, n_days * n_age)
G_obs_aggregate = torch.tensor(sim_cases).sum(dim=-1)
def objective(G):
return - (G - G_obs_aggregate).pow(2).sum(dim=-1) / n_days
# if maxiters provided, select submatrix of state
if maxiters:
train_theta = train_theta[:min(maxiters, train_theta.shape[0])]
train_G = train_G[:min(maxiters, train_G.shape[0])]
# extract all parameter settings that improved
best = - 99999999999999
t = 0
all_params = []
while t < train_theta.shape[0]:
theta = train_theta[t]
G = train_G[t]
obj = objective(G).item()
if obj > best:
best = obj
calibrated_params = transforms.unnormalize(theta, sim_bounds)
all_params.append(
(t, parr_to_pdict(parr=calibrated_params, multi_beta_calibration=multi_beta_calibration)))
t += 1
return all_params
def add(self, *,
simulation_info,
country,
area,
measure_list,
lockdown_measures_active=True,
full_scale=True,
test_update=None,
seed_summary_path=None,
set_calibrated_params_to=None,
set_initial_seeds_to=None,
expected_daily_base_expo_per100k=0,
store_mob=False):
# Set time window based on experiment start and end date
sim_days = (pd.to_datetime(self.end_date) - pd.to_datetime(self.start_date)).days
max_time = TO_HOURS * sim_days # in hours
# extract lockdown period
lockdown_start_date = pd.to_datetime(
calibration_lockdown_dates[country]['start'])
lockdown_end_date = pd.to_datetime(
calibration_lockdown_dates[country]['end'])
days_until_lockdown_start = (lockdown_start_date - pd.to_datetime(self.start_date)).days
days_until_lockdown_end = (lockdown_end_date - pd.to_datetime(self.start_date)).days
# Load mob settings
mob_settings_file = calibration_mob_paths[country][area][1 if full_scale else 0]
with open(mob_settings_file, 'rb') as fp:
mob_settings = pickle.load(fp)
# Obtain COVID19 case date for country and area to estimate testing capacity and heuristic seeds if necessary
unscaled_area_cases = collect_data_from_df(country=country, area=area, datatype='new',
start_date_string=self.start_date, end_date_string=self.end_date)
assert(len(unscaled_area_cases.shape) == 2)
# Scale down cases based on number of people in town and region
sim_cases = downsample_cases(unscaled_area_cases, mob_settings)
# Instantiate correct state transition distributions (estimated from literature)
distributions = CovidDistributions(country=country)
# Expected base rate infections
if expected_daily_base_expo_per100k > 0.0:
# Scale expectation to simulation size
num_people = len(mob_settings['home_loc'])
lambda_base_expo_population = expected_daily_base_expo_per100k * (num_people / 100000)
# Convert to individual base rate by dividing by population size; priority queue handles superposition
lambda_base_expo_indiv = lambda_base_expo_population / num_people
distributions.lambda_0 = lambda_base_expo_indiv
# Get initial seeds for simulation
# (a) Define heuristically based on true cases and literature distribution estimates
if seed_summary_path is None:
# Generate initial seeds based on unscaled case numbers in town
initial_seeds = gen_initial_seeds(
sim_cases, day=0)
if sum(initial_seeds.values()) == 0:
print('No states seeded at start time; cannot start simulation.\n'
'Consider setting a later start date for calibration using the "--start" flag.')
sys.exit(0)
# (b) Define based state of previous batch of simulations,
# using the random rollout that best matched the true cases in terms of squared error
else:
seed_summary_ = load_summary(seed_summary_path)
seed_day_ = seed_summary_.max_time # take seeds at the end of simulation
initial_seeds = extract_seeds_from_summary(
seed_summary_, seed_day_, sim_cases)
if set_initial_seeds_to is not None:
initial_seeds = set_initial_seeds_to
# Load calibrated model parameters for this area
calibrated_params = get_calibrated_params(
country=country, area=area, multi_beta_calibration=self.multi_beta_calibration)
if set_calibrated_params_to is not None:
calibrated_params = set_calibrated_params_to
p_stay_home_calibrated = calibrated_params['p_stay_home']
if self.multi_beta_calibration:
betas = calibrated_params['betas']
else:
betas = {
'education': calibrated_params['beta_site'],
'social': calibrated_params['beta_site'],
'bus_stop': calibrated_params['beta_site'],
'office': calibrated_params['beta_site'],
'supermarket': calibrated_params['beta_site'],
}
model_params = {
'betas': betas,
'beta_household': calibrated_params['beta_household'],
}
# Add standard measure of positives staying isolated
measure_list += [
# standard behavior of positively tested: full isolation
SocialDistancingForPositiveMeasure(
t_window=Interval(0.0, max_time), p_stay_home=1.0),
SocialDistancingForPositiveMeasureHousehold(
t_window=Interval(0.0, max_time), p_isolate=1.0),
]
# Add standard measures if simulation is happening during lockdown
# Set lockdown_measures_active to False to explore counterfactual scenarios
if lockdown_measures_active:
measure_list += [
# social distancing factor during lockdown: calibrated
SocialDistancingForAllMeasure(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
p_stay_home=p_stay_home_calibrated),
# site specific measures: fixed in advance, outside of calibration
BetaMultiplierMeasureByType(
t_window=Interval(TO_HOURS * days_until_lockdown_start,
TO_HOURS * days_until_lockdown_end),
beta_multiplier=calibration_lockdown_beta_multipliers)
]
measure_list = MeasureList(measure_list)
# Set testing conditions
scaled_test_capacity = get_test_capacity(
country, area, mob_settings, end_date_string=self.end_date)
testing_params = copy.deepcopy(calibration_testing_params)
testing_params['tests_per_batch'] = scaled_test_capacity
testing_params['testing_t_window'] = [0.0, max_time]
if test_update:
testing_params = test_update(testing_params)
# store simulation
self.sims.append(Simulation(
# Generic information
experiment_info=self.experiment_info,
simulation_info=simulation_info,
start_date=self.start_date,
end_date=self.end_date,
sim_days=sim_days,
country=country,
area=area,
random_repeats=self.random_repeats,
# Mobility and measures
mob_settings_file=mob_settings_file,
full_scale=full_scale,
measure_list=measure_list,
testing_params=testing_params,
store_mob=store_mob,
# Model
model_params=model_params,
distributions=distributions,
initial_seeds=initial_seeds,
))
if self.verbose:
print(f'[Added Sim] {self.get_sim_path(self.sims[-1])}')
for summary_path, exp in zip(summary_paths, exps.values()):
try:
_, filename = os.path.split(summary_path)
print('Plotting: ' + filename)
resulttuple = load_summary(summary_path + '.pk')
summary = resulttuple[1]
mob_settings_paths = calibration_mob_paths[country][area][
1 if exp['full_scale'] else 0]
with open(mob_settings_paths, 'rb') as fp:
mob_settings = pickle.load(fp)
area_cases = collect_data_from_df(country=country,
area=area,
datatype='new',
start_date_string=start_date,
end_date_string=end_date)
sim_cases = downsample_cases(
area_cases,
mob_settings) # only downscaling due LK data for cities
plotter = Plotter()
plotter.plot_positives_vs_target(summary,
sim_cases.sum(axis=1),
title='Calibration period',
filename=filename,
figsize=(6, 4),
start_date=start_date,
errorevery=1,