def test_parse_configuration(self):
"""
Check covid19sim/utils.py/parse_configuration
"""
toy_conf = OmegaConf.create({
"start_time": "2020-02-28 00:00:00",
})
with self.assertRaises(AssertionError):
parse_configuration(toy_conf)
toy_conf["RISK_MODEL"] = ""
parsed_conf = parse_configuration(toy_conf)
self.assertEqual(parsed_conf["start_time"],
datetime.datetime(2020, 2, 28, 0, 0, 0))
def test_dump_conf(self):
"""
asserts that a conf that is dumped and parsed again yields identical results
"""
conf = get_test_conf("naive_local.yaml")
with TemporaryDirectory() as d:
dump_conf(conf, Path(d) / "dumped.yaml")
with (Path(d) / "dumped.yaml").open("r") as f:
loaded_conf = yaml.safe_load(f)
parsed_conf = parse_configuration(loaded_conf)
def __init__(self, debug_data_path: str):
self.path = debug_data_path
self.zipfd = None
if self.path.endswith(".zip"):
self.zipfd = zipfile.ZipFile(self.path)
fileinfos = self.zipfd.infolist()
assert len(fileinfos) == 1, "unexpected hdf5 archive structure"
raise NotImplementedError # TODO LATER @@@@
else:
self.hdf5fd = h5py.File(self.path, "r")
self.conf = parse_configuration(json.loads(self.hdf5fd.attrs["config"]))
self.dataset = self.hdf5fd["dataset"]
self.is_delta = self.hdf5fd["is_delta"]
assert self.dataset.shape == self.is_delta.shape
self.simulation_days = self.dataset.shape[0]
assert self.dataset.shape[1] == 24 # 24 potential timeslots per day
self.n_people = self.dataset.shape[2]
def get_test_conf(conf_name):
"""
Loads the default configurations in configs and overwrites it
with values in test_configs/`conf_name`
conf_name **must** be in `tests/test_configs/`
Args:
conf_name (str): name of the configuration to load in `tests/test_configs/`
Returns:
dict: full overwritten config, using covid19sim.utils.parse_configuration
"""
config_path = Path(__file__).parent.parent / "test_configs" / conf_name
config_path = config_path.resolve()
assert config_path.suffix == ".yaml"
assert config_path.exists()
config = HYDRA_SIM_PATH / "config.yaml"
with config.open("r") as f:
defaults = yaml.safe_load(f)["defaults"]
default_confs = []
for d in defaults:
if type(d) == str:
conf = OmegaConf.load(str(HYDRA_SIM_PATH / (d + ".yaml")))
else:
conf = OmegaConf.load(
str(HYDRA_SIM_PATH / list(d.keys())[0] /
(list(d.values())[0] + ".yaml")))
default_confs.append(conf)
conf = OmegaConf.merge(*default_confs, OmegaConf.load(str(config_path)))
return parse_configuration(conf)
def main(conf: DictConfig):
"""
Enables command line execution of the simulator.
Args:
conf (DictConfig): yaml configuration file
"""
# -------------------------------------------------
# ----- Load the experimental configuration -----
# -------------------------------------------------
conf = parse_configuration(conf)
# -------------------------------------
# ----- Create Output Directory -----
# -------------------------------------
if conf["outdir"] is None:
conf["outdir"] = str(Path(__file__) / "output")
timenow = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
conf[
"outdir"
] = "{}/sim_v2_people-{}_days-{}_init-{}_uptake-{}_seed-{}_{}_{}".format(
conf["outdir"],
conf["n_people"],
conf["simulation_days"],
conf["init_fraction_sick"],
conf["APP_UPTAKE"],
conf["seed"],
timenow,
str(time.time_ns())[-6:],
)
if Path(conf["outdir"]).exists():
out_path = Path(conf["outdir"])
out_idx = 1
while (out_path.parent / (out_path.name + f"_{out_idx}")).exists():
out_idx += 1
conf["outdir"] = str(out_path.parent / (out_path.name + f"_{out_idx}"))
os.makedirs(conf["outdir"])
logfile = f"{conf['outdir']}/log_{timenow}.txt"
outfile = os.path.join(conf["outdir"], "data")
# ---------------------------------
# ----- Filter-Out Warnings -----
# ---------------------------------
import warnings
# warnings.filterwarnings("ignore")
# ----------------------------
# ----- Run Simulation -----
# ----------------------------
# correctness of configuration file
assert not conf['RISK_MODEL'] != "" or conf['INTERVENTION_DAY'] >= 0, "risk model is given, but no intervnetion day specified"
assert conf['N_BEHAVIOR_LEVELS'] >= 2, "At least 2 behavior levels are required to model behavior changes"
if conf['TRACE_SYMPTOMS']:
warnings.warn("TRACE_SYMPTOMS: True hasn't been implemented. It will have no affect.")
log(f"RISK_MODEL = {conf['RISK_MODEL']}", logfile)
log(f"INTERVENTION_DAY = {conf['INTERVENTION_DAY']}", logfile)
log(f"seed: {conf['seed']}", logfile)
# complete decsription of intervention
type_of_run = _get_intervention_string(conf)
conf['INTERVENTION'] = type_of_run
log(f"Type of run: {type_of_run}", logfile)
if conf['COLLECT_TRAINING_DATA']:
data_output_path = os.path.join(conf["outdir"], "train.zarr")
collection_server = DataCollectionServer(
data_output_path=data_output_path,
config_backup=conf,
human_count=conf['n_people'],
simulation_days=conf['simulation_days'],
)
collection_server.start()
else:
collection_server = None
conf["outfile"] = outfile
city = simulate(
n_people=conf["n_people"],
init_fraction_sick=conf["init_fraction_sick"],
start_time=conf["start_time"],
simulation_days=conf["simulation_days"],
outfile=conf["outfile"],
out_chunk_size=conf["out_chunk_size"],
seed=conf["seed"],
conf=conf,
logfile=logfile
)
# write the full configuration file along with git commit hash
dump_conf(city.conf, "{}/full_configuration.yaml".format(city.conf["outdir"]))
# log the simulation statistics
city.tracker.write_metrics()
# (baseball-cards) write full simulation data
if hasattr(city, "tracker") and \
hasattr(city.tracker, "collection_server") and \
isinstance(city.tracker.collection_server, DataCollectionServer) and \
city.tracker.collection_server is not None:
city.tracker.collection_server.stop_gracefully()
city.tracker.collection_server.join()
# if COLLECT_TRAINING_DATA is true
if not conf["tune"]:
# ----------------------------------------------
# ----- Not Tune: Collect Training Data -----
# ----------------------------------------------
# write values to train with
train_priors = os.path.join(f"{conf['outdir']}/train_priors.pkl")
city.tracker.write_for_training(city.humans, train_priors, conf)
timenow = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
log("Dumping Tracker Data in {}".format(conf["outdir"]), logfile)
Path(conf["outdir"]).mkdir(parents=True, exist_ok=True)
filename = f"tracker_data_n_{conf['n_people']}_seed_{conf['seed']}_{timenow}.pkl"
data = extract_tracker_data(city.tracker, conf)
dump_tracker_data(data, conf["outdir"], filename)
else:
# ------------------------------------------------------
# ----- Tune: Write logs And Tacker Data -----
# ------------------------------------------------------
timenow = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
log("Dumping Tracker Data in {}".format(conf["outdir"]), logfile)
Path(conf["outdir"]).mkdir(parents=True, exist_ok=True)
filename = f"tracker_data_n_{conf['n_people']}_seed_{conf['seed']}_{timenow}.pkl"
data = extract_tracker_data(city.tracker, conf)
dump_tracker_data(data, conf["outdir"], filename)
# Shutdown the data collection server if one's running
if collection_server is not None:
collection_server.stop_gracefully()
collection_server.join()
# Remove the IPCs if they were stored somewhere custom
if os.environ.get("COVID19SIM_IPC_PATH", None) is not None:
print("<<<<<<<< Cleaning Up >>>>>>>>")
for file in Path(os.environ.get("COVID19SIM_IPC_PATH")).iterdir():
if file.name.endswith(".ipc"):
print(f"Removing {str(file)}...")
os.remove(str(file))
return conf