def compute_past_series(df,
new_cases_col,
pastdays_start,
pastdays_end,
draws,
alpha,
beta,
trend_alpha,
lower_ratio,
upper_ratio,
pickleprefix,
rt_col_prefix='smooth_deseas'):
for pastdays in range(pastdays_start, pastdays_end - 1, -1):
print(f'\npastdays: {pastdays}')
if pastdays == 0:
sl = np.s_[:]
else:
sl = np.s_[:-pastdays]
new_cases = df[new_cases_col].to_numpy()[sl]
new_cases_expanded = draw_expanded_series(new_cases,
draws=draws,
season_period=7,
trend_alpha=trend_alpha,
difference_degree=2,
alpha=alpha,
beta=beta,
lower_ratio=lower_ratio,
upper_ratio=upper_ratio,
truncate=False)
print('pre smooth_and_drop')
print(f'new_cases_expanded.shape: {new_cases_expanded.shape}')
new_cases_smoothed, rel_eps, padding_left = smooth_and_drop(
new_cases_expanded,
season_period=7,
trend_alpha=100.,
difference_degree=2)
print('post smooth_and_drop')
print(f'new_cases_smoothed.shape: {new_cases_smoothed.shape}')
print(f'rel_eps.shape: {rel_eps.shape}')
print(f'padding_left: {padding_left}')
simulations = []
for new_cases_s, rel_eps_s in zip(new_cases_smoothed, rel_eps):
print(
f'new_cases_s cut left: {new_cases_s[~np.isnan(new_cases_s)].shape}'
)
print(
f'rel_eps_s cut left: {rel_eps_s[~np.isnan(rel_eps_s)].shape}')
try:
model_, trace_ = MCMC_sample(
onset=new_cases_s[~np.isnan(new_cases_s)],
alpha=alpha,
beta=beta,
rel_eps=rel_eps_s[~np.isnan(rel_eps_s)],
start=0,
window=None,
chains=4,
tune=500,
draws=500,
cores=4,
target_accept=0.95,
dry=False,
progressbar=False)
simulations.append(trace_)
except Exception as ex:
print(ex)
print(f'skipping pastdays {pastdays:03d}')
with open(
os.path.join(
BASE_DATA_PATH,
f'computed/WIP/{pickleprefix}_MCMC_simulations_pastdays_{pastdays_start:03d}_{pastdays_end:03d}.pickle'
), 'wb') as handle:
pickle.dump(simulations, handle)
sampled_Rt = np.vstack(
[t['r_t'][~t.diverging, :] for t in simulations])
combined_trace = {'r_t': sampled_Rt}
save_MCMC_sampling(df,
f'{new_cases_col}_{rt_col_prefix}',
combined_trace,
pastdays,
interval=0.95,
start=padding_left + 1)
df.to_pickle(
os.path.join(
BASE_DATA_PATH,
f'computed/WIP/{pickleprefix}_MCMC_Rt_pastdays_{pastdays_start:03d}_{pastdays_end:03d}.pickle'
))
def compute_past_series(
df,
new_cases_col,
startday,
pastdays_start,
pastdays_end,
draws,
alpha,
beta,
trend_alpha,
lower_ratio,
upper_ratio,
pickleprefix,
mctune=1000,
mcdraws=500,
mccores=8,
mctargetaccept=0.95,
rt_col_prefix="smooth_deseas",
use_rel_res=True,
new_cases_local_col=None,
debug_mode=False,
):
for pastdays in range(pastdays_start, pastdays_end - 1, -1):
logger.info(f"\npastdays: {pastdays}")
if pastdays == 0:
sl = np.s_[startday:]
else:
sl = np.s_[startday:-pastdays]
new_cases = df[new_cases_col].to_numpy()[sl]
new_cases_expanded = draw_expanded_series(
new_cases,
draws=draws,
season_period=7,
trend_alpha=trend_alpha,
difference_degree=2,
alpha=alpha,
beta=beta,
lower_ratio=lower_ratio,
upper_ratio=upper_ratio,
truncate=False,
)
logger.info("pre smooth_and_drop")
logger.info(f"new_cases_expanded.shape: {new_cases_expanded.shape}")
new_cases_smoothed, rel_eps, padding_left = smooth_and_drop(
new_cases_expanded, season_period=7, trend_alpha=100.0, difference_degree=2
)
logger.info("post smooth_and_drop")
logger.info(f"new_cases_smoothed.shape: {new_cases_smoothed.shape}")
logger.info(f"rel_eps.shape: {rel_eps.shape}")
logger.info(f"padding_left: {padding_left}")
# local management
if new_cases_local_col:
new_cases_local = df[new_cases_local_col].to_numpy()[sl]
new_cases_local_expanded = draw_expanded_series(
new_cases_local,
draws=draws,
season_period=7,
trend_alpha=trend_alpha,
difference_degree=2,
alpha=alpha,
beta=beta,
lower_ratio=lower_ratio,
upper_ratio=upper_ratio,
truncate=False,
)
new_cases_local_smoothed, rel_eps_local, padding_left_local = smooth_and_drop(
new_cases_local_expanded, season_period=7, trend_alpha=100.0, difference_degree=2
)
else:
new_cases_local_smoothed = [None] * draws
simulations = []
for new_cases_s, rel_eps_s, new_cases_local_s in zip(new_cases_smoothed, rel_eps, new_cases_local_smoothed):
if new_cases_local_s is not None:
new_cases_local_s = new_cases_local_s[~np.isnan(new_cases_local_s)]
logger.info(
f"new_cases_s cut left: {new_cases_s[~np.isnan(new_cases_s)].shape}"
)
logger.info(f"rel_eps_s cut left: {rel_eps_s[~np.isnan(rel_eps_s)].shape}")
try:
model_, trace_ = MCMC_sample(
onset=new_cases_s[~np.isnan(new_cases_s)],
alpha=alpha,
beta=beta,
rel_eps=rel_eps_s[~np.isnan(rel_eps_s)] if use_rel_res else None,
onset_local=new_cases_local_s,
start=0,
window=None,
chains=mccores,
tune=mctune,
draws=mcdraws,
cores=mccores,
target_accept=mctargetaccept,
dry=False,
progressbar=False,
)
simulations.append(trace_)
except Exception as ex:
logger.info(ex)
logger.info(f"skipping pastdays {pastdays:03d}")
raise ex
# prepare target dir path and verify if exists
TARGET_RESULT_DIR = os.path.join(settings.BASE_DATA_PATH, "computed/WIP/")
if not os.path.exists(TARGET_RESULT_DIR):
os.makedirs(TARGET_RESULT_DIR, exist_ok=True)
with open(
os.path.join(
TARGET_RESULT_DIR,
f"{pickleprefix}_MCMC_simulations_pastdays_{pastdays_start:03d}_{pastdays_end:03d}.pickle",
),
"wb",
) as handle:
pickle.dump(simulations, handle)
if debug_mode:
# use all sampled Rt, including diverging ones
sampled_Rt = np.vstack([t["r_t"] for t in simulations])
else:
sampled_Rt = np.vstack([t["r_t"][~t.diverging, :] for t in simulations])
combined_trace = {"r_t": sampled_Rt}
process_MCMC_sampling(
df,
f"{new_cases_col}_{rt_col_prefix}",
combined_trace,
pastdays,
interval=0.95,
start=padding_left + 1 + startday,
)
df.to_pickle(
os.path.join(
TARGET_RESULT_DIR,
f"{pickleprefix}_MCMC_Rt_pastdays_{pastdays_start:03d}_{pastdays_end:03d}.pickle",
)
)