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 train_single_model(
t0: int, data: np.ndarray, knots: List[int], knots_connect: List[int],
learning_rate: np.float64, loss_prop: np.float64, max_epochs: int,
initial_a: List[np.float64], enable_tensorboard: bool,
tensorboard_logdir: Text
) -> Tuple[combined_model.Covid19CombinedPredictModel,
infection_model.TensorType]:
"""Trains single Covid19CombinedEstimator with one worker.
Args:
t0: specifies the number of days between the occurrence of the first
infected case (patient zero) and the first observed case.
data: training data (number of daily new confirmed cases) in a 1d array.
knots: a list of integers in which each represents the length of one piece
in the piecewise linear infection rate model. These integers should sum
up to the length of training data.
knots_connect:.
learning_rate : the learning rate in an optimizer.
loss_prop:.
max_epochs : the number of steps to run trainer.
initial_a : a list of float numbers indicating the infection rate at each
knot. The length equals the length of 'knots' + 1.
Returns:
weights: the weights in the best model after training with t0.
loss: the loss of the best model after training with t0.
"""
estimator = combined_model.Covid19CombinedEstimator(
knots=knots,
knots_connect=knots_connect,
loss_prop=loss_prop,
estimator_args={
"learning_rate": learning_rate,
"epochs": max_epochs
},
initial_guess_a=np.array(initial_a))
model, loss = estimator._fit_with_t0(data, t0, "", enable_tensorboard,
tensorboard_logdir)
return model.weights, loss
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.")