def run_state(state, states_only=False, with_age_structure=False):
"""
Run the fitter for each county in a state.
Parameters
----------
state: str
State to run against.
states_only: bool
If True only run the state level.
with_age_structure: bool
If True run model with age structure.
"""
state_obj = us.states.lookup(state)
logging.info(f"Running MLE fitter for state {state_obj.name}")
model_fitter = ModelFitter.run_for_fips(
fips=state_obj.fips, with_age_structure=with_age_structure)
df_whitelist = load_data.load_whitelist()
df_whitelist = df_whitelist[df_whitelist["inference_ok"] == True]
output_path = get_run_artifact_path(state_obj.fips,
RunArtifact.MLE_FIT_RESULT)
data = pd.DataFrame(model_fitter.fit_results, index=[state_obj.fips])
data.to_json(output_path)
with open(get_run_artifact_path(state_obj.fips, RunArtifact.MLE_FIT_MODEL),
"wb") as f:
pickle.dump(model_fitter.mle_model, f)
# Run the counties.
if not states_only:
# TODO: Replace with build_county_list
df_whitelist = load_data.load_whitelist()
df_whitelist = df_whitelist[df_whitelist["inference_ok"] == True]
all_fips = df_whitelist[df_whitelist["state"].str.lower() ==
state_obj.name.lower()].fips.values
if len(all_fips) > 0:
p = Pool()
fitters = p.map(ModelFitter.run_for_fips, all_fips)
p.close()
county_output_file = get_run_artifact_path(
all_fips[0], RunArtifact.MLE_FIT_RESULT)
data = pd.DataFrame([fit.fit_results for fit in fitters if fit])
data.to_json(county_output_file)
# Serialize the model results.
for fips, fitter in zip(all_fips, fitters):
if fitter:
with open(
get_run_artifact_path(fips,
RunArtifact.MLE_FIT_MODEL),
"wb") as f:
pickle.dump(fitter.mle_model, f)
def __init__(
self,
state,
output_interval_days=4,
run_mode="can-before",
output_dir=None,
jhu_dataset=None,
cds_dataset=None,
include_imputed=False,
):
self.output_interval_days = output_interval_days
self.state = state
self.run_mode = RunMode(run_mode)
self.include_imputed = include_imputed
self.state_abbreviation = us.states.lookup(state).abbr
self.population_data = FIPSPopulation.local().population()
self.output_dir = output_dir
self.jhu_local = jhu_dataset or JHUDataset.local()
self.cds_dataset = cds_dataset or CDSDataset.local()
self.county_timeseries = build_aggregate_county_data_frame(self.jhu_local, self.cds_dataset)
self.county_timeseries["date"] = self.county_timeseries["date"].dt.normalize()
self.state_timeseries = self.jhu_local.timeseries().state_data
self.state_timeseries["date"] = self.state_timeseries["date"].dt.normalize()
self.df_whitelist = load_data.load_whitelist()
self.df_whitelist = self.df_whitelist[self.df_whitelist["inference_ok"] == True]
def build_county_list(state):
"""
Build the and return the fips list
"""
state_obj = us.states.lookup(state)
log.info(f"Get fips list for state {state_obj.name}")
df_whitelist = load_data.load_whitelist()
df_whitelist = df_whitelist[df_whitelist["inference_ok"] == True]
all_fips = df_whitelist[df_whitelist["state"].str.lower() ==
state_obj.name.lower()].fips.tolist()
return all_fips
