def main(unused_argv):
logger = logging.getLogger("Covid-19-estimation")
input_data = np.load(FLAGS.input_path)
if FLAGS.loss_prop == 0:
logger.warning(
"Warning: the death data are not used for model training.")
if FLAGS.loss_prop == 1:
logger.warning(
"Warning: the daily observed cases are not used for model training."
)
logger.info("Loaded the training data.")
parsed_knots, parsed_knots_connect, parsed_initial_a = io_utils.parse_knots(
FLAGS.knots, FLAGS.knots_connect, FLAGS.initial_a)
estimator = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
loss_prop=FLAGS.loss_prop,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs
},
initial_guess_a=np.array(parsed_initial_a))
estimator.fit(data=input_data, min_t0=FLAGS.min_t0, max_t0=FLAGS.max_t0)
logger.info("Completed model estimation.")
io_utils.parse_estimated_model(estimator)
io_utils.export_estimation_and_prediction(
estimator=estimator,
test_duration=FLAGS.test_duration,
output_path=FLAGS.output_path,
suffix=FLAGS.file_index,
flatten_future=FLAGS.flatten_future,
to_json=FLAGS.output_to_json)
logger.info("Done.")
def main(unused_argv):
"""Runs model estimation in parallel."""
logger = logging.getLogger("covid-19-estimation")
input_data = np.load(FLAGS.input_path)
if FLAGS.loss_prop == 0:
logger.warning("The death data are not used for model training.")
if FLAGS.loss_prop == 1:
logger.warning(
"The daily observed cases are not used for model training.")
logger.info("Loaded the training data.")
parsed_knots, parsed_knots_connect, parsed_initial_a = io_utils.parse_knots(
FLAGS.knots, FLAGS.knots_connect, FLAGS.initial_a)
pool_workers = multiprocessing.Pool(FLAGS.n_process)
if sum(parsed_knots) != input_data.shape[0]:
raise ValueError(
"The elements in vector 'knots' should sum up to the the length of "
f"data {input_data.shape[0]}, but got value {sum(parsed_knots)}.")
target_fn = functools.partial(train_single_model,
data=input_data,
knots=parsed_knots,
knots_connect=parsed_knots_connect,
learning_rate=FLAGS.learning_rate,
loss_prop=FLAGS.loss_prop,
max_epochs=FLAGS.max_epochs,
initial_a=parsed_initial_a,
enable_tensorboard=FLAGS.enable_tensorboard,
tensorboard_logdir=FLAGS.tensorboard_logdir)
all_t0 = list(range(FLAGS.min_t0, FLAGS.max_t0 + 1))
all_results = pool_workers.map(func=target_fn, iterable=all_t0)
logger.info("Completed model estimation.")
trained_weights, trained_losses = zip(*(all_results))
remove_nan_losses = np.array(trained_losses)
remove_nan_losses[np.isnan(remove_nan_losses)] = np.Inf
best_weights = trained_weights[remove_nan_losses.argmin()]
estimator_template = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
estimator_args={})
selected_t0 = all_t0[remove_nan_losses.argmin()]
model_template = combined_model.Covid19CombinedPredictModel(
n_weights=len(parsed_initial_a),
t0=selected_t0,
len_inputs=len(input_data) + selected_t0)
model_template.load_pretrained_weights(best_weights)
estimator_template._final_model = model_template
estimator_template._final_loss = min(remove_nan_losses)
io_utils.parse_estimated_model(estimator_template)
io_utils.export_estimation_and_prediction(
estimator=estimator_template,
test_duration=FLAGS.test_duration,
output_path=FLAGS.output_path,
flatten_future=FLAGS.flatten_future,
to_json=FLAGS.output_to_json)
logger.info("Done.")
def main(unused_argv):
logger = logging.getLogger("Covid-19-estimation")
input_data = np.load(FLAGS.input_path)
parsed_knots, parsed_knots_connect, parsed_initial_a = io_utils.parse_knots(
FLAGS.knots, FLAGS.knots_connect, FLAGS.initial_a)
# Step 1 estimation (infection cases only).
estimator = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
loss_prop=0,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs
},
initial_guess_a=np.array(parsed_initial_a),
variable_death_rate_trainable=False)
estimator.fit(data=input_data, min_t0=FLAGS.min_t0, max_t0=FLAGS.max_t0)
# Save the estimated weights from step 1 (will be fixed in step 2).
stage1_estimated_a, stage1_estimated_t0 = estimator.final_model.a.numpy(
), estimator.final_model.t0
logger.info("First stage estimation done.")
# Step 2 estimation (death only).
estimator = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
loss_prop=1,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs
},
initial_guess_a=stage1_estimated_a,
variable_a_trainable=False)
estimator.fit(data=input_data,
min_t0=stage1_estimated_t0,
max_t0=stage1_estimated_t0)
io_utils.parse_estimated_model(estimator)
io_utils.export_estimation_and_prediction(
estimator=estimator,
test_duration=FLAGS.test_duration,
output_path=FLAGS.output_path,
suffix=FLAGS.file_index,
flatten_future=FLAGS.flatten_future,
to_json=True)
logger.info("Second stage estimation done.")
def main(unused_argv):
logger = logging.getLogger("Covid-19-estimation")
job_metadata = read_metadata(FLAGS.metadata, NUM_STATES)
job_file_index = str(int(FLAGS.file_index_plus_one) - 1)
state_key, resample_key = map_keys(
job_file_index, NUM_STATES, FLAGS.resample_jobs)
job_state = job_metadata[state_key]["state"]
logger.info(f"Create estimation job for the state {job_state}.")
# Parse the required data & args for running the estimation model.
job_knots = job_metadata[state_key]["knot"]
if FLAGS.resample_jobs:
job_residuals = resample.read_residuals(
FLAGS.residuals, NUM_STATES)
job_input_data = resample.get_resampled_input(
job_residuals[state_key]["fitted"],
job_residuals[state_key]["residual"],
int(resample_key))
else:
job_input_data = job_metadata[state_key]["data"]
job_knots_connect = [1] * len(job_knots)
job_initial_a = [0.2] * (len(job_knots) + 1)
# Step 1 estimation (infection cases only).
estimator = combined_model.Covid19CombinedEstimator(
knots=job_knots,
knots_connect=job_knots_connect,
loss_prop=0,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs},
initial_guess_a=np.array(job_initial_a),
variable_death_rate_trainable=False
)
estimator.fit(data=job_input_data, min_t0=FLAGS.min_t0, max_t0=FLAGS.max_t0)
# Save the estimated weights from step 1 (will be fixed in step 2).
stage1_estimated_a, stage1_estimated_t0 = estimator.final_model.a.numpy(), estimator.final_model.t0
logger.info("First stage estimation done.")
# Step 2 estimation (death only).
estimator = combined_model.Covid19CombinedEstimator(
knots=job_knots,
knots_connect=job_knots_connect,
loss_prop=1,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs},
initial_guess_a=stage1_estimated_a,
variable_a_trainable=False
)
estimator.fit(data=job_input_data, min_t0=stage1_estimated_t0, max_t0=stage1_estimated_t0)
io_utils.parse_estimated_model(estimator)
if FLAGS.resample_jobs:
job_suffix = "state_" + state_key + "resample_" + resample_key
else:
job_suffix = "state_" + state_key
io_utils.export_estimation_and_prediction(
estimator=estimator,
test_duration=FLAGS.test_duration,
output_path=FLAGS.output_path,
suffix=job_suffix,
flatten_future=FLAGS.flatten_future,
to_json=FLAGS.output_to_json
)
logger.info("Second stage estimation done.")
def main(unused_argv):
"""Runs model estimation in parallel."""
logger = logging.getLogger("covid-19-estimation")
input_data = np.load(FLAGS.input_path)
logger.info("Loaded the training data.")
parsed_knots, parsed_knots_connect, parsed_initial_a = io_utils.parse_knots(
FLAGS.knots, FLAGS.knots_connect, FLAGS.initial_a)
pool_workers = multiprocessing.Pool(FLAGS.n_process)
if sum(parsed_knots) != input_data.shape[0]:
raise ValueError(
"The elements in vector 'knots' should sum up to the the length of "
f"data {input_data.shape[0]}, but got value {sum(parsed_knots)}.")
# Step 1 estimation (cases only).
target_fn = functools.partial(train_single_model,
data=input_data,
knots=parsed_knots,
knots_connect=parsed_knots_connect,
learning_rate=FLAGS.learning_rate,
max_epochs=FLAGS.max_epochs,
initial_a=parsed_initial_a,
enable_tensorboard=FLAGS.enable_tensorboard,
tensorboard_logdir=FLAGS.tensorboard_logdir)
# Note that tensorboard only monitors the first stage.
all_t0 = list(range(FLAGS.min_t0, FLAGS.max_t0 + 1))
all_results = pool_workers.map(func=target_fn, iterable=all_t0)
logger.info("Completed model estimation.")
trained_weights, trained_losses = zip(*(all_results))
remove_nan_losses = np.array(trained_losses)
remove_nan_losses[np.isnan(remove_nan_losses)] = np.Inf
best_weights = trained_weights[remove_nan_losses.argmin()]
estimator_template = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
estimator_args={})
selected_t0 = all_t0[remove_nan_losses.argmin()]
model_template = combined_model.Covid19CombinedPredictModel(
n_weights=len(parsed_initial_a),
t0=selected_t0,
len_inputs=len(input_data) + selected_t0)
model_template.load_pretrained_weights(best_weights)
stage1_estimated_a = model_template.a.numpy()
logger.info("First stage estimation done.")
# Step 2 estimation (death only).
estimator = combined_model.Covid19CombinedEstimator(
knots=parsed_knots,
knots_connect=parsed_knots_connect,
loss_prop=1,
estimator_args={
"learning_rate": FLAGS.learning_rate,
"epochs": FLAGS.max_epochs
},
initial_guess_a=stage1_estimated_a,
variable_a_trainable=False)
estimator.fit(data=input_data, min_t0=selected_t0, max_t0=selected_t0)
io_utils.parse_estimated_model(estimator)
io_utils.export_estimation_and_prediction(
estimator=estimator,
test_duration=FLAGS.test_duration,
output_path=FLAGS.output_path,
flatten_future=FLAGS.flatten_future,
to_json=FLAGS.output_to_json)
logger.info("Second stage estimation done.")