diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 8a95b7bc3..938c2a203 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -13,7 +13,7 @@ body: - bug - batch - config - - flepiconfig + - dependency - documentation - docstring - duplicate @@ -25,8 +25,12 @@ body: - meta/workflow - operations - performance + - plotting - post-processing - quick issue + - r-flepicommon + - r-flepiconfig + - r-inference - seeding - submission - wontfix diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 1ac8bf159..e95f65fea 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -13,7 +13,7 @@ body: - bug - batch - config - - flepiconfig + - dependency - documentation - docstring - duplicate @@ -25,8 +25,12 @@ body: - meta/workflow - operations - performance + - plotting - post-processing - quick issue + - r-flepicommon + - r-flepiconfig + - r-inference - seeding - submission - wontfix diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 000000000..e2449af27 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,41 @@ +name: Lint + +on: + workflow_dispatch: + push: + paths: + - 'flepimop/gempyor_pkg/**/*.py' + pull_request: + paths: + - '**/*.py' + branches: + - main + +jobs: + black-for-python: + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.draft == false + env: + BLACK_LINE_LENGTH: 92 + BLACK_EXTEND_EXCLUDE: 'flepimop/gempyor_pkg/src/gempyor/steps_rk4.py' + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + lfs: true + sparse-checkout: | + * + !documentation/ + sparse-checkout-cone-mode: false + - name: Determine Source + run: | + if [ ${{ github.event_name }} == "push" ]; then + echo "BLACK_SRC=flepimop/gempyor_pkg/" >> $GITHUB_ENV + else + echo "BLACK_SRC=." >> $GITHUB_ENV + fi + - name: Black Formatter Check + uses: psf/black@stable + with: + src: ${{ env.BLACK_SRC }} + options: "--line-length ${{ env.BLACK_LINE_LENGTH }} --extend-exclude '${{ env.BLACK_EXTEND_EXCLUDE }}' --check --verbose" diff --git a/batch/inference_job_launcher.py b/batch/inference_job_launcher.py index b69d223a5..884f210ca 100755 --- a/batch/inference_job_launcher.py +++ b/batch/inference_job_launcher.py @@ -403,23 +403,37 @@ def launch_batch( if "scenarios" in config["outcome_modifiers"]: outcome_modifiers_scenarios = config["outcome_modifiers"]["scenarios"] - handler.launch(job_name, config_filepath, seir_modifiers_scenarios, outcome_modifiers_scenarios) + handler.launch( + job_name, config_filepath, seir_modifiers_scenarios, outcome_modifiers_scenarios + ) # Set job_name as environmental variable so it can be pulled for pushing to git os.environ["job_name"] = job_name # Set run_id as environmental variable so it can be pulled for pushing to git TODO - (rc, txt) = subprocess.getstatusoutput(f"git checkout -b run_{job_name}") # TODO: cd ... + (rc, txt) = subprocess.getstatusoutput( + f"git checkout -b run_{job_name}" + ) # TODO: cd ... print(txt) return rc -def autodetect_params(config, data_path, *, num_jobs=None, sims_per_job=None, num_blocks=None, batch_system=None): +def autodetect_params( + config, + data_path, + *, + num_jobs=None, + sims_per_job=None, + num_blocks=None, + batch_system=None, +): if num_jobs and sims_per_job and num_blocks: return (num_jobs, sims_per_job, num_blocks) if "inference" not in config or "iterations_per_slot" not in config["inference"]: - raise click.UsageError("inference::iterations_per_slot undefined in config, can't autodetect parameters") + raise click.UsageError( + "inference::iterations_per_slot undefined in config, can't autodetect parameters" + ) iterations_per_slot = int(config["inference"]["iterations_per_slot"]) if num_jobs is None: @@ -429,11 +443,17 @@ def autodetect_params(config, data_path, *, num_jobs=None, sims_per_job=None, nu if sims_per_job is None: if num_blocks is not None: sims_per_job = int(math.ceil(iterations_per_slot / num_blocks)) - print(f"Setting number of blocks to {num_blocks} [via num_blocks (-k) argument]") - print(f"Setting sims per job to {sims_per_job} [via {iterations_per_slot} iterations_per_slot in config]") + print( + f"Setting number of blocks to {num_blocks} [via num_blocks (-k) argument]" + ) + print( + f"Setting sims per job to {sims_per_job} [via {iterations_per_slot} iterations_per_slot in config]" + ) else: if "data_path" in config: - raise ValueError("The config has a data_path section. This is no longer supported.") + raise ValueError( + "The config has a data_path section. This is no longer supported." + ) geodata_fname = pathlib.Path(data_path) / config["subpop_setup"]["geodata"] with open(geodata_fname) as geodata_fp: num_subpops = sum(1 for line in geodata_fp) @@ -458,7 +478,9 @@ def autodetect_params(config, data_path, *, num_jobs=None, sims_per_job=None, nu if num_blocks is None: num_blocks = int(math.ceil(iterations_per_slot / sims_per_job)) - print(f"Setting number of blocks to {num_blocks} [via {iterations_per_slot} iterations_per_slot in config]") + print( + f"Setting number of blocks to {num_blocks} [via {iterations_per_slot} iterations_per_slot in config]" + ) return (num_jobs, sims_per_job, num_blocks) @@ -478,7 +500,9 @@ def get_aws_job_queues(job_queue_prefix): return sorted(queues_with_jobs, key=queues_with_jobs.get) -def aws_countfiles_autodetect_runid(s3_bucket, restart_from_location, restart_from_run_id, num_jobs, strict=False): +def aws_countfiles_autodetect_runid( + s3_bucket, restart_from_location, restart_from_run_id, num_jobs, strict=False +): import boto3 s3 = boto3.resource("s3") @@ -487,15 +511,21 @@ def aws_countfiles_autodetect_runid(s3_bucket, restart_from_location, restart_fr all_files = list(bucket.objects.filter(Prefix=prefix)) all_files = [f.key for f in all_files] if restart_from_run_id is None: - print("WARNING: no --restart_from_run_id specified, autodetecting... please wait querying S3 👀🔎...") + print( + "WARNING: no --restart_from_run_id specified, autodetecting... please wait querying S3 👀🔎..." + ) restart_from_run_id = all_files[0].split("/")[3] - if user_confirmation(question=f"Auto-detected run_id {restart_from_run_id}. Correct ?", default=True): + if user_confirmation( + question=f"Auto-detected run_id {restart_from_run_id}. Correct ?", default=True + ): print(f"great, continuing with run_id {restart_from_run_id}...") else: raise ValueError(f"Abording, please specify --restart_from_run_id manually.") final_llik = [f for f in all_files if ("llik" in f) and ("final" in f)] - if len(final_llik) == 0: # hacky: there might be a bucket with no llik files, e.g if init. + if ( + len(final_llik) == 0 + ): # hacky: there might be a bucket with no llik files, e.g if init. final_llik = [f for f in all_files if ("init" in f) and ("final" in f)] if len(final_llik) != num_jobs: @@ -583,8 +613,12 @@ def build_job_metadata(self, job_name): manifest = {} manifest["cmd"] = " ".join(sys.argv[:]) manifest["job_name"] = job_name - manifest["data_sha"] = subprocess.getoutput("cd {self.data_path}; git rev-parse HEAD") - manifest["flepimop_sha"] = subprocess.getoutput(f"cd {self.flepi_path}; git rev-parse HEAD") + manifest["data_sha"] = subprocess.getoutput( + "cd {self.data_path}; git rev-parse HEAD" + ) + manifest["flepimop_sha"] = subprocess.getoutput( + f"cd {self.flepi_path}; git rev-parse HEAD" + ) # Save the manifest file to S3 with open("manifest.json", "w") as f: @@ -594,17 +628,25 @@ def build_job_metadata(self, job_name): # need these to be uploaded so they can be executed. this_file_path = os.path.dirname(os.path.realpath(__file__)) self.save_file( - source=os.path.join(this_file_path, "AWS_inference_runner.sh"), destination=f"{job_name}-runner.sh" + source=os.path.join(this_file_path, "AWS_inference_runner.sh"), + destination=f"{job_name}-runner.sh", ) self.save_file( - source=os.path.join(this_file_path, "AWS_inference_copy.sh"), destination=f"{job_name}-copy.sh" + source=os.path.join(this_file_path, "AWS_inference_copy.sh"), + destination=f"{job_name}-copy.sh", ) tarfile_name = f"{job_name}.tar.gz" self.tar_working_dir(tarfile_name=tarfile_name) - self.save_file(source=tarfile_name, destination=f"{job_name}.tar.gz", remove_source=True) + self.save_file( + source=tarfile_name, destination=f"{job_name}.tar.gz", remove_source=True + ) - self.save_file(source="manifest.json", destination=f"{job_name}/manifest.json", remove_source=True) + self.save_file( + source="manifest.json", + destination=f"{job_name}/manifest.json", + remove_source=True, + ) def tar_working_dir(self, tarfile_name): # this tar file always has the structure: @@ -616,10 +658,13 @@ def tar_working_dir(self, tarfile_name): or q == "covid-dashboard-app" or q == "renv.cache" or q == "sample_data" - or q == "renv" # joseph: I added this to fix a bug, hopefully it doesn't break anything + or q + == "renv" # joseph: I added this to fix a bug, hopefully it doesn't break anything or q.startswith(".") ): - tar.add(os.path.join(self.flepi_path, q), arcname=os.path.join("flepiMoP", q)) + tar.add( + os.path.join(self.flepi_path, q), arcname=os.path.join("flepiMoP", q) + ) elif q == "sample_data": for r in os.listdir(os.path.join(self.flepi_path, "sample_data")): if r != "united-states-commutes": @@ -629,10 +674,17 @@ def tar_working_dir(self, tarfile_name): ) # tar.add(os.path.join("flepiMoP", "sample_data", r)) for p in os.listdir(self.data_path): - if not (p.startswith(".") or p.endswith("tar.gz") or p in self.outputs or p == "flepiMoP"): + if not ( + p.startswith(".") + or p.endswith("tar.gz") + or p in self.outputs + or p == "flepiMoP" + ): tar.add( p, - filter=lambda x: None if os.path.basename(x.name).startswith(".") else x, + filter=lambda x: ( + None if os.path.basename(x.name).startswith(".") else x + ), ) tar.close() @@ -656,7 +708,13 @@ def save_file(self, source, destination, remove_source=False, prefix=""): if remove_source: os.remove(source) - def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_modifiers_scenarios): + def launch( + self, + job_name, + config_filepath, + seir_modifiers_scenarios, + outcome_modifiers_scenarios, + ): s3_results_path = f"s3://{self.s3_bucket}/{job_name}" if self.batch_system == "slurm": @@ -676,7 +734,10 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo ## TODO: check how each of these variables are used downstream base_env_vars = [ {"name": "BATCH_SYSTEM", "value": self.batch_system}, - {"name": "S3_MODEL_PROJECT_PATH", "value": f"s3://{self.s3_bucket}/{job_name}.tar.gz"}, + { + "name": "S3_MODEL_PROJECT_PATH", + "value": f"s3://{self.s3_bucket}/{job_name}.tar.gz", + }, {"name": "DVC_OUTPUTS", "value": " ".join(self.outputs)}, {"name": "S3_RESULTS_PATH", "value": s3_results_path}, {"name": "FS_RESULTS_PATH", "value": fs_results_path}, @@ -700,14 +761,22 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo }, {"name": "FLEPI_STOCHASTIC_RUN", "value": str(self.stochastic)}, {"name": "FLEPI_RESET_CHIMERICS", "value": str(self.reset_chimerics)}, - {"name": "FLEPI_MEM_PROFILE", "value": str(os.getenv("FLEPI_MEM_PROFILE", default="FALSE"))}, - {"name": "FLEPI_MEM_PROF_ITERS", "value": str(os.getenv("FLEPI_MEM_PROF_ITERS", default="50"))}, + { + "name": "FLEPI_MEM_PROFILE", + "value": str(os.getenv("FLEPI_MEM_PROFILE", default="FALSE")), + }, + { + "name": "FLEPI_MEM_PROF_ITERS", + "value": str(os.getenv("FLEPI_MEM_PROF_ITERS", default="50")), + }, {"name": "SLACK_CHANNEL", "value": str(self.slack_channel)}, ] with open(config_filepath) as f: config = yaml.full_load(f) - for ctr, (s, d) in enumerate(itertools.product(seir_modifiers_scenarios, outcome_modifiers_scenarios)): + for ctr, (s, d) in enumerate( + itertools.product(seir_modifiers_scenarios, outcome_modifiers_scenarios) + ): cur_job_name = f"{job_name}_{s}_{d}" # Create first job cur_env_vars = base_env_vars.copy() @@ -719,7 +788,9 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo cur_env_vars.append({"name": "FLEPI_BLOCK_INDEX", "value": "1"}) cur_env_vars.append({"name": "FLEPI_RUN_INDEX", "value": f"{self.run_id}"}) if not (self.restart_from_location is None): - cur_env_vars.append({"name": "LAST_JOB_OUTPUT", "value": f"{self.restart_from_location}"}) + cur_env_vars.append( + {"name": "LAST_JOB_OUTPUT", "value": f"{self.restart_from_location}"} + ) cur_env_vars.append( { "name": "OLD_FLEPI_RUN_INDEX", @@ -732,8 +803,18 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo if self.continuation: cur_env_vars.append({"name": "FLEPI_CONTINUATION", "value": f"TRUE"}) - cur_env_vars.append({"name": "FLEPI_CONTINUATION_RUN_ID", "value": f"{self.continuation_run_id}"}) - cur_env_vars.append({"name": "FLEPI_CONTINUATION_LOCATION", "value": f"{self.continuation_location}"}) + cur_env_vars.append( + { + "name": "FLEPI_CONTINUATION_RUN_ID", + "value": f"{self.continuation_run_id}", + } + ) + cur_env_vars.append( + { + "name": "FLEPI_CONTINUATION_LOCATION", + "value": f"{self.continuation_location}", + } + ) cur_env_vars.append( { "name": "FLEPI_CONTINUATION_FTYPE", @@ -814,7 +895,9 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo postprod_command, command_name="sbatch postprod", fail_on_fail=True ) postprod_job_id = stdout.decode().split(" ")[-1][:-1] - print(f">>> SUCCESS SCHEDULING POST-PROCESSING JOB. Slurm job id is {postprod_job_id}") + print( + f">>> SUCCESS SCHEDULING POST-PROCESSING JOB. Slurm job id is {postprod_job_id}" + ) elif self.batch_system == "local": cur_env_vars.append({"name": "JOB_NAME", "value": f"{cur_job_name}"}) @@ -831,12 +914,24 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo cur_env_vars = base_env_vars.copy() cur_env_vars.append({"name": "FLEPI_SEIR_SCENARIOS", "value": s}) cur_env_vars.append({"name": "FLEPI_OUTCOME_SCENARIOS", "value": d}) - cur_env_vars.append({"name": "FLEPI_PREFIX", "value": f"{config['name']}_{s}_{d}"}) - cur_env_vars.append({"name": "FLEPI_BLOCK_INDEX", "value": f"{block_idx+1}"}) - cur_env_vars.append({"name": "FLEPI_RUN_INDEX", "value": f"{self.run_id}"}) - cur_env_vars.append({"name": "OLD_FLEPI_RUN_INDEX", "value": f"{self.run_id}"}) - cur_env_vars.append({"name": "LAST_JOB_OUTPUT", "value": f"{s3_results_path}/"}) - cur_env_vars.append({"name": "JOB_NAME", "value": f"{cur_job_name}_block{block_idx}"}) + cur_env_vars.append( + {"name": "FLEPI_PREFIX", "value": f"{config['name']}_{s}_{d}"} + ) + cur_env_vars.append( + {"name": "FLEPI_BLOCK_INDEX", "value": f"{block_idx+1}"} + ) + cur_env_vars.append( + {"name": "FLEPI_RUN_INDEX", "value": f"{self.run_id}"} + ) + cur_env_vars.append( + {"name": "OLD_FLEPI_RUN_INDEX", "value": f"{self.run_id}"} + ) + cur_env_vars.append( + {"name": "LAST_JOB_OUTPUT", "value": f"{s3_results_path}/"} + ) + cur_env_vars.append( + {"name": "JOB_NAME", "value": f"{cur_job_name}_block{block_idx}"} + ) cur_job = batch_client.submit_job( jobName=f"{cur_job_name}_block{block_idx}", jobQueue=cur_job_queue, @@ -895,19 +990,29 @@ def launch(self, job_name, config_filepath, seir_modifiers_scenarios, outcome_mo em = "" if self.resume_discard_seeding: em = f", discarding seeding results." - print(f" >> Resuming from run id is {self.restart_from_run_id} located in {self.restart_from_location}{em}") + print( + f" >> Resuming from run id is {self.restart_from_run_id} located in {self.restart_from_location}{em}" + ) if self.batch_system == "aws": print(f" >> Final output will be: {s3_results_path}/model_output/") elif self.batch_system == "slurm": print(f" >> Final output will be: {fs_results_path}/model_output/") if self.s3_upload: - print(f" >> Final output will be uploaded to {s3_results_path}/model_output/") + print( + f" >> Final output will be uploaded to {s3_results_path}/model_output/" + ) if self.continuation: - print(f" >> Continuing from run id is {self.continuation_run_id} located in {self.continuation_location}") + print( + f" >> Continuing from run id is {self.continuation_run_id} located in {self.continuation_location}" + ) print(f" >> Run id is {self.run_id}") print(f" >> config is {config_filepath.split('/')[-1]}") - flepimop_branch = subprocess.getoutput(f"cd {self.flepi_path}; git rev-parse --abbrev-ref HEAD") - data_branch = subprocess.getoutput(f"cd {self.data_path}; git rev-parse --abbrev-ref HEAD") + flepimop_branch = subprocess.getoutput( + f"cd {self.flepi_path}; git rev-parse --abbrev-ref HEAD" + ) + data_branch = subprocess.getoutput( + f"cd {self.data_path}; git rev-parse --abbrev-ref HEAD" + ) data_hash = subprocess.getoutput(f"cd {self.data_path}; git rev-parse HEAD") flepimop_hash = subprocess.getoutput(f"cd {self.flepi_path}; git rev-parse HEAD") print(f""" >> FLEPIMOP branch is {flepimop_branch} with hash {flepimop_hash}""") diff --git a/batch/scenario_job.py b/batch/scenario_job.py index 1961974eb..62907a370 100755 --- a/batch/scenario_job.py +++ b/batch/scenario_job.py @@ -196,13 +196,17 @@ def launch_job_inner( tarfile_name = f"{job_name}.tar.gz" tar = tarfile.open(tarfile_name, "w:gz") for p in os.listdir("."): - if not (p.startswith(".") or p.endswith("tar.gz") or p in dvc_outputs or p == "batch"): + if not ( + p.startswith(".") or p.endswith("tar.gz") or p in dvc_outputs or p == "batch" + ): tar.add(p, filter=lambda x: None if x.name.startswith(".") else x) tar.close() # Upload the tar'd contents of this directory and the runner script to S3 runner_script_name = f"{job_name}-runner.sh" - local_runner_script = os.path.join(os.path.dirname(os.path.realpath(__file__)), "AWS_scenario_runner.sh") + local_runner_script = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "AWS_scenario_runner.sh" + ) s3_client = boto3.client("s3") s3_client.upload_file(local_runner_script, s3_input_bucket, runner_script_name) s3_client.upload_file(tarfile_name, s3_input_bucket, tarfile_name) @@ -219,7 +223,9 @@ def launch_job_inner( {"name": "S3_RESULTS_PATH", "value": results_path}, {"name": "SLOTS_PER_JOB", "value": str(slots_per_job)}, ] - s3_cp_run_script = f"aws s3 cp s3://{s3_input_bucket}/{runner_script_name} $PWD/run-flepimop-inference" + s3_cp_run_script = ( + f"aws s3 cp s3://{s3_input_bucket}/{runner_script_name} $PWD/run-flepimop-inference" + ) command = ["sh", "-c", f"{s3_cp_run_script}; /bin/bash $PWD/run-flepimop-inference"] container_overrides = { "vcpus": vcpu, @@ -246,7 +252,9 @@ def launch_job_inner( containerOverrides=container_overrides, ) - print(f"Batch job with id {resp['jobId']} launched; output will be written to {results_path}") + print( + f"Batch job with id {resp['jobId']} launched; output will be written to {results_path}" + ) def get_dvc_outputs(): diff --git a/bin/lint b/bin/lint new file mode 100755 index 000000000..50b31272b --- /dev/null +++ b/bin/lint @@ -0,0 +1,5 @@ +#!/usr/bin/env bash + +black --line-length 92 \ + --extend-exclude 'flepimop/gempyor_pkg/src/gempyor/steps_rk4.py' \ + --verbose . diff --git a/bin/pre-commit b/bin/pre-commit new file mode 100755 index 000000000..f9fdd15e6 --- /dev/null +++ b/bin/pre-commit @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +if which black > /dev/null 2>&1; then + black --line-length 92 \ + --extend-exclude 'flepimop/gempyor_pkg/src/gempyor/steps_rk4.py' \ + --check --verbose . +else + echo "'black' is not available so python files will not be checked." +fi diff --git a/documentation/gitbook/development/python-guidelines-for-developers.md b/documentation/gitbook/development/python-guidelines-for-developers.md index 59f2d3c67..75a7e60ab 100644 --- a/documentation/gitbook/development/python-guidelines-for-developers.md +++ b/documentation/gitbook/development/python-guidelines-for-developers.md @@ -61,21 +61,23 @@ and to run just some subset of the tests (e.g here just the outcome tests), use: pytest -vvvv -k outcomes ``` -{% hint style="danger" %} -Before committing, make sure you **format your code** using black (see below) and that the **test passes** (see above). -{% endhint %} +For more details on how to use `pytest` please refer to their [usage guide](https://docs.pytest.org/en/latest/how-to/usage.html). ### Formatting -We try to remain close to Python conventions and to follow the updated rules and best practices. For formatting, we use [black](https://github.com/psf/black), the _Uncompromising Code Formatter_ before submitting pull requests. It provides a consistent style, which is useful when diffing. We use a custom length of 120 characters as the baseline is short for scientific code. Here is the line to use to format your code: +We try to remain close to Python conventions and to follow the updated rules and best practices. For formatting, we use [black](https://github.com/psf/black), the _Uncompromising Code Formatter_ before submitting pull requests. It provides a consistent style, which is useful when diffing. We use a custom length of 92 characters as the baseline is short for scientific code. Here is the line to use to format your code: ```bash -black --line-length 120 . --exclude renv* +black --line-length 92 \ + --extend-exclude 'flepimop/gempyor_pkg/src/gempyor/steps_rk4.py' \ + --verbose . ``` -{% hint style="warning" %} -Please use type-hints as much as possible, as we are trying to move towards static checks. -{% endhint %} +For those using a Mac or Linux system for development this command is also available for use by calling `./dev/lint`. Similarly, you can take advantage of the formatting pre-commit hook found at `bin/pre-commit`. To start using it copy this file to your git hooks folder: + +```bash +cp -f bin/pre-commit .git/hooks/ +``` #### Structure of the main classes diff --git a/examples/test_cli.py b/examples/test_cli.py index 349f7219a..857ad8249 100644 --- a/examples/test_cli.py +++ b/examples/test_cli.py @@ -1,4 +1,3 @@ - from click.testing import CliRunner from gempyor.simulate import simulate import os @@ -6,33 +5,35 @@ # See here to test click application https://click.palletsprojects.com/en/8.1.x/testing/ # would be useful to also call the command directly + def test_config_sample_2pop(): - os.chdir(os.path.dirname(__file__) + "/tutorials") - runner = CliRunner() - result = runner.invoke(simulate, ['-c', 'config_sample_2pop.yml']) - print(result.output) # useful for debug - print(result.exit_code) # useful for debug - print(result.exception) # useful for debug - assert result.exit_code == 0 - assert 'completed in' in result.output + os.chdir(os.path.dirname(__file__) + "/tutorials") + runner = CliRunner() + result = runner.invoke(simulate, ["-c", "config_sample_2pop.yml"]) + print(result.output) # useful for debug + print(result.exit_code) # useful for debug + print(result.exception) # useful for debug + assert result.exit_code == 0 + assert "completed in" in result.output def test_sample_2pop_modifiers(): - os.chdir(os.path.dirname(__file__) + "/tutorials") - runner = CliRunner() - result = runner.invoke(simulate, ['-c', 'config_sample_2pop_modifiers.yml']) - print(result.output) # useful for debug - print(result.exit_code) # useful for debug - print(result.exception) # useful for debug - assert result.exit_code == 0 - assert 'completed in' in result.output + os.chdir(os.path.dirname(__file__) + "/tutorials") + runner = CliRunner() + result = runner.invoke(simulate, ["-c", "config_sample_2pop_modifiers.yml"]) + print(result.output) # useful for debug + print(result.exit_code) # useful for debug + print(result.exception) # useful for debug + assert result.exit_code == 0 + assert "completed in" in result.output + def test_simple_usa_statelevel(): - os.chdir(os.path.dirname(__file__) + "/simple_usa_statelevel") - runner = CliRunner() - result = runner.invoke(simulate, ['-c', 'simple_usa_statelevel.yml', '-n', '1']) - print(result.output) # useful for debug - print(result.exit_code) # useful for debug - print(result.exception) # useful for debug - assert result.exit_code == 0 - assert 'completed in' in result.output \ No newline at end of file + os.chdir(os.path.dirname(__file__) + "/simple_usa_statelevel") + runner = CliRunner() + result = runner.invoke(simulate, ["-c", "simple_usa_statelevel.yml", "-n", "1"]) + print(result.output) # useful for debug + print(result.exit_code) # useful for debug + print(result.exception) # useful for debug + assert result.exit_code == 0 + assert "completed in" in result.output diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py index ecbbba962..84414e7ff 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py @@ -21,7 +21,8 @@ def __init__( name=getattr( npi_config, "key", - (npi_config["scenario"].exists() and npi_config["scenario"].get()) or "unknown", + (npi_config["scenario"].exists() and npi_config["scenario"].get()) + or "unknown", ) ) @@ -68,14 +69,20 @@ def __init__( # if parameters are exceeding global start/end dates, index of parameter df will be out of range so check first if self.sanitize: - too_early = min([min(i) for i in self.parameters["start_date"]]) < self.start_date + too_early = ( + min([min(i) for i in self.parameters["start_date"]]) < self.start_date + ) too_late = max([max(i) for i in self.parameters["end_date"]]) > self.end_date if too_early or too_late: - raise ValueError("at least one period start or end date is not between global dates") + raise ValueError( + "at least one period start or end date is not between global dates" + ) for grp_config in npi_config["groups"]: affected_subpops_grp = self.__get_affected_subpops_grp(grp_config) - for sub_index in range(len(self.parameters["start_date"][affected_subpops_grp[0]])): + for sub_index in range( + len(self.parameters["start_date"][affected_subpops_grp[0]]) + ): period_range = pd.date_range( self.parameters["start_date"][affected_subpops_grp[0]][sub_index], self.parameters["end_date"][affected_subpops_grp[0]][sub_index], @@ -111,7 +118,9 @@ def __checkErrors(self): ) if not (self.parameters["start_date"] <= self.parameters["end_date"]).all(): - raise ValueError(f"at least one period_start_date is greater than the corresponding period end date") + raise ValueError( + f"at least one period_start_date is greater than the corresponding period end date" + ) for n in self.affected_subpops: if n not in self.subpops: @@ -153,7 +162,9 @@ def __createFromConfig(self, npi_config): else: start_dates = [self.start_date] end_dates = [self.end_date] - this_spatial_group = helpers.get_spatial_groups(grp_config, affected_subpops_grp) + this_spatial_group = helpers.get_spatial_groups( + grp_config, affected_subpops_grp + ) self.spatial_groups.append(this_spatial_group) # print(self.name, this_spatial_groups) @@ -209,7 +220,9 @@ def __createFromDf(self, loaded_df, npi_config): else: start_dates = [self.start_date] end_dates = [self.end_date] - this_spatial_group = helpers.get_spatial_groups(grp_config, affected_subpops_grp) + this_spatial_group = helpers.get_spatial_groups( + grp_config, affected_subpops_grp + ) self.spatial_groups.append(this_spatial_group) for subpop in this_spatial_group["ungrouped"]: @@ -227,7 +240,9 @@ def __createFromDf(self, loaded_df, npi_config): for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - self.parameters.at[subpop, "value"] = loaded_df.at[",".join(group), "value"] + self.parameters.at[subpop, "value"] = loaded_df.at[ + ",".join(group), "value" + ] else: dist = npi_config["value"].as_random_distribution() drawn_value = dist(size=1) @@ -258,11 +273,16 @@ def __get_affected_subpops(self, npi_config): affected_subpops_grp += [str(n.get()) for n in grp_config["subpop"]] affected_subpops = set(affected_subpops_grp) if len(affected_subpops) != len(affected_subpops_grp): - raise ValueError(f"In NPI {self.name}, some subpops belong to several groups. This is unsupported.") + raise ValueError( + f"In NPI {self.name}, some subpops belong to several groups. This is unsupported." + ) return affected_subpops def get_default(self, param): - if param in self.pnames_overlap_operation_sum or param in self.pnames_overlap_operation_reductionprod: + if ( + param in self.pnames_overlap_operation_sum + or param in self.pnames_overlap_operation_reductionprod + ): return 0.0 else: return 1.0 @@ -278,7 +298,9 @@ def getReductionToWrite(self): # self.parameters.index is a list of subpops for this_spatial_groups in self.spatial_groups: # spatially ungrouped dataframe - df_ungroup = self.parameters[self.parameters.index.isin(this_spatial_groups["ungrouped"])].copy() + df_ungroup = self.parameters[ + self.parameters.index.isin(this_spatial_groups["ungrouped"]) + ].copy() df_ungroup.index.name = "subpop" df_ungroup["start_date"] = df_ungroup["start_date"].apply( lambda l: ",".join([d.strftime("%Y-%m-%d") for d in l]) @@ -301,7 +323,9 @@ def getReductionToWrite(self): "start_date": df_group["start_date"].apply( lambda l: ",".join([d.strftime("%Y-%m-%d") for d in l]) ), - "end_date": df_group["end_date"].apply(lambda l: ",".join([d.strftime("%Y-%m-%d") for d in l])), + "end_date": df_group["end_date"].apply( + lambda l: ",".join([d.strftime("%Y-%m-%d") for d in l]) + ), "value": df_group["value"], } ).set_index("subpop") diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py index e078ddeba..cdda3c4b9 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py @@ -21,7 +21,8 @@ def __init__( name=getattr( npi_config, "key", - (npi_config["scenario"].exists() and npi_config["scenario"].get()) or "unknown", + (npi_config["scenario"].exists() and npi_config["scenario"].get()) + or "unknown", ) ) @@ -60,15 +61,22 @@ def __init__( self.__createFromConfig(npi_config) # if parameters are exceeding global start/end dates, index of parameter df will be out of range so check first - if self.parameters["start_date"].min() < self.start_date or self.parameters["end_date"].max() > self.end_date: - raise ValueError(f"""{self.name} : at least one period start or end date is not between global dates""") + if ( + self.parameters["start_date"].min() < self.start_date + or self.parameters["end_date"].max() > self.end_date + ): + raise ValueError( + f"""{self.name} : at least one period start or end date is not between global dates""" + ) # for index in self.parameters.index: # period_range = pd.date_range(self.parameters["start_date"][index], self.parameters["end_date"][index]) ## This the line that does the work # self.npi_old.loc[index, period_range] = np.tile(self.parameters["value"][index], (len(period_range), 1)).T - period_range = pd.date_range(self.parameters["start_date"].iloc[0], self.parameters["end_date"].iloc[0]) + period_range = pd.date_range( + self.parameters["start_date"].iloc[0], self.parameters["end_date"].iloc[0] + ) self.npi.loc[self.parameters.index, period_range] = np.tile( self.parameters["value"][:], (len(period_range), 1) ).T @@ -90,7 +98,9 @@ def __checkErrors(self): ) if not (self.parameters["start_date"] <= self.parameters["end_date"]).all(): - raise ValueError(f"at least one period_start_date is greater than the corresponding period end date") + raise ValueError( + f"at least one period_start_date is greater than the corresponding period end date" + ) for n in self.affected_subpops: if n not in self.subpops: @@ -122,13 +132,19 @@ def __createFromConfig(self, npi_config): self.parameters["modifier_name"] = self.name self.parameters["start_date"] = ( - npi_config["period_start_date"].as_date() if npi_config["period_start_date"].exists() else self.start_date + npi_config["period_start_date"].as_date() + if npi_config["period_start_date"].exists() + else self.start_date ) self.parameters["end_date"] = ( - npi_config["period_end_date"].as_date() if npi_config["period_end_date"].exists() else self.end_date + npi_config["period_end_date"].as_date() + if npi_config["period_end_date"].exists() + else self.end_date ) self.parameters["parameter"] = self.param_name - self.spatial_groups = helpers.get_spatial_groups(npi_config, list(self.affected_subpops)) + self.spatial_groups = helpers.get_spatial_groups( + npi_config, list(self.affected_subpops) + ) if self.spatial_groups["ungrouped"]: self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = self.dist( size=len(self.spatial_groups["ungrouped"]) @@ -153,10 +169,14 @@ def __createFromDf(self, loaded_df, npi_config): # self.parameters = loaded_df[["modifier_name", "start_date", "end_date", "parameter", "value"]].copy() # dates are picked from config self.parameters["start_date"] = ( - npi_config["period_start_date"].as_date() if npi_config["period_start_date"].exists() else self.start_date + npi_config["period_start_date"].as_date() + if npi_config["period_start_date"].exists() + else self.start_date ) self.parameters["end_date"] = ( - npi_config["period_end_date"].as_date() if npi_config["period_end_date"].exists() else self.end_date + npi_config["period_end_date"].as_date() + if npi_config["period_end_date"].exists() + else self.end_date ) ## This is more legible to me, but if we change it here, we should change it in __createFromConfig as well # if npi_config["period_start_date"].exists(): @@ -175,17 +195,24 @@ def __createFromDf(self, loaded_df, npi_config): # TODO: to be consistent with MTR, we want to also draw the values for the subpops # that are not in the loaded_df. - self.spatial_groups = helpers.get_spatial_groups(npi_config, list(self.affected_subpops)) + self.spatial_groups = helpers.get_spatial_groups( + npi_config, list(self.affected_subpops) + ) if self.spatial_groups["ungrouped"]: self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = loaded_df.loc[ self.spatial_groups["ungrouped"], "value" ] if self.spatial_groups["grouped"]: for group in self.spatial_groups["grouped"]: - self.parameters.loc[group, "value"] = loaded_df.loc[",".join(group), "value"] + self.parameters.loc[group, "value"] = loaded_df.loc[ + ",".join(group), "value" + ] def get_default(self, param): - if param in self.pnames_overlap_operation_sum or param in self.pnames_overlap_operation_reductionprod: + if ( + param in self.pnames_overlap_operation_sum + or param in self.pnames_overlap_operation_reductionprod + ): return 0.0 else: return 1.0 @@ -198,7 +225,9 @@ def getReduction(self, param): def getReductionToWrite(self): # spatially ungrouped dataframe - df = self.parameters[self.parameters.index.isin(self.spatial_groups["ungrouped"])].copy() + df = self.parameters[ + self.parameters.index.isin(self.spatial_groups["ungrouped"]) + ].copy() df.index.name = "subpop" df["start_date"] = df["start_date"].astype("str") df["end_date"] = df["end_date"].astype("str") diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py index 489a48fbb..6cf178735 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py @@ -47,7 +47,9 @@ def __init__( if isinstance(scenario, str): settings = modifiers_library.get(scenario) if settings is None: - raise RuntimeError(f"couldn't find scenario in config file [got: {scenario}]") + raise RuntimeError( + f"couldn't find scenario in config file [got: {scenario}]" + ) # via profiling: faster to recreate the confuse view than to fetch+resolve due to confuse isinstance # checks scenario_npi_config = confuse.RootView([settings]) @@ -68,12 +70,16 @@ def __init__( ) new_params = sub_npi.param_name # either a list (if stacked) or a string - new_params = [new_params] if isinstance(new_params, str) else new_params # convert to list + new_params = ( + [new_params] if isinstance(new_params, str) else new_params + ) # convert to list # Add each parameter at first encounter, with a neutral start for new_p in new_params: if new_p not in self.param_name: self.param_name.append(new_p) - if new_p in pnames_overlap_operation_sum: # re.match("^transition_rate [1234567890]+$",new_p): + if ( + new_p in pnames_overlap_operation_sum + ): # re.match("^transition_rate [1234567890]+$",new_p): self.reductions[new_p] = 0 else: # for the reductionprod and product method, the initial neutral is 1 ) self.reductions[new_p] = 1 @@ -81,7 +87,9 @@ def __init__( for param in self.param_name: # Get reduction return a neutral value for this overlap operation if no parameeter exists reduction = sub_npi.getReduction(param) - if param in pnames_overlap_operation_sum: # re.match("^transition_rate [1234567890]+$",param): + if ( + param in pnames_overlap_operation_sum + ): # re.match("^transition_rate [1234567890]+$",param): self.reductions[param] += reduction elif param in pnames_overlap_operation_reductionprod: self.reductions[param] *= 1 - reduction @@ -104,7 +112,9 @@ def __init__( self.reduction_params.clear() for param in self.param_name: - if param in pnames_overlap_operation_reductionprod: # re.match("^transition_rate \d+$",param): + if ( + param in pnames_overlap_operation_reductionprod + ): # re.match("^transition_rate \d+$",param): self.reductions[param] = 1 - self.reductions[param] # check that no NPI is called several times, and retourn them @@ -124,7 +134,10 @@ def __checkErrors(self): # ) def get_default(self, param): - if param in self.pnames_overlap_operation_sum or param in self.pnames_overlap_operation_reductionprod: + if ( + param in self.pnames_overlap_operation_sum + or param in self.pnames_overlap_operation_reductionprod + ): return 0.0 else: return 1.0 diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/helpers.py b/flepimop/gempyor_pkg/src/gempyor/NPI/helpers.py index f964d4c6e..e18f43f28 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/helpers.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/helpers.py @@ -43,16 +43,22 @@ def get_spatial_groups(grp_config, affected_subpops: list) -> dict: else: spatial_groups["grouped"] = grp_config["subpop_groups"].get() spatial_groups["ungrouped"] = list( - set(affected_subpops) - set(flatten_list_of_lists(spatial_groups["grouped"])) + set(affected_subpops) + - set(flatten_list_of_lists(spatial_groups["grouped"])) ) # flatten the list of lists of grouped subpops, so we can do some checks flat_grouped_list = flatten_list_of_lists(spatial_groups["grouped"]) # check that all subpops are either grouped or ungrouped if set(flat_grouped_list + spatial_groups["ungrouped"]) != set(affected_subpops): - print("set of grouped and ungrouped subpops", set(flat_grouped_list + spatial_groups["ungrouped"])) + print( + "set of grouped and ungrouped subpops", + set(flat_grouped_list + spatial_groups["ungrouped"]), + ) print("set of affected subpops ", set(affected_subpops)) - raise ValueError(f"The two above sets are differs for for intervention with config \n {grp_config}") + raise ValueError( + f"The two above sets are differs for for intervention with config \n {grp_config}" + ) if len(set(flat_grouped_list + spatial_groups["ungrouped"])) != len( flat_grouped_list + spatial_groups["ungrouped"] ): diff --git a/flepimop/gempyor_pkg/src/gempyor/calibrate.py b/flepimop/gempyor_pkg/src/gempyor/calibrate.py index e5cb287aa..875db12ca 100644 --- a/flepimop/gempyor_pkg/src/gempyor/calibrate.py +++ b/flepimop/gempyor_pkg/src/gempyor/calibrate.py @@ -158,7 +158,9 @@ def calibrate( # TODO here for resume if resume or resume_location is not None: - print("Doing a resume, this only work with the same number of slot and parameters right now") + print( + "Doing a resume, this only work with the same number of slot and parameters right now" + ) p0 = None if resume_location is not None: backend = emcee.backends.HDFBackend(resume_location) @@ -195,7 +197,10 @@ def calibrate( # plotting the chain sampler = emcee.backends.HDFBackend(filename, read_only=True) gempyor.postprocess_inference.plot_chains( - inferpar=gempyor_inference.inferpar, sampler_output=sampler, sampled_slots=None, save_to=f"{run_id}_chains.pdf" + inferpar=gempyor_inference.inferpar, + sampler_output=sampler, + sampled_slots=None, + save_to=f"{run_id}_chains.pdf", ) print("EMCEE Run done, doing sampling") @@ -203,11 +208,14 @@ def calibrate( shutil.rmtree(project_path + "model_output/", ignore_errors=True) max_indices = np.argsort(sampler.get_log_prob()[-1, :])[-nsamples:] - samples = sampler.get_chain()[-1, max_indices, :] # the last iteration, for selected slots + samples = sampler.get_chain()[ + -1, max_indices, : + ] # the last iteration, for selected slots gempyor_inference.set_save(True) with multiprocessing.Pool(ncpu) as pool: results = pool.starmap( - gempyor_inference.get_logloss_as_single_number, [(samples[i, :],) for i in range(len(max_indices))] + gempyor_inference.get_logloss_as_single_number, + [(samples[i, :],) for i in range(len(max_indices))], ) # results = [] # for fn in gempyor.utils.list_filenames(folder="model_output/", filters=[run_id, "hosp.parquet"]): diff --git a/flepimop/gempyor_pkg/src/gempyor/compartments.py b/flepimop/gempyor_pkg/src/gempyor/compartments.py index cff807559..445759cbc 100644 --- a/flepimop/gempyor_pkg/src/gempyor/compartments.py +++ b/flepimop/gempyor_pkg/src/gempyor/compartments.py @@ -13,7 +13,13 @@ class Compartments: # Minimal object to be easily picklable for // runs - def __init__(self, seir_config=None, compartments_config=None, compartments_file=None, transitions_file=None): + def __init__( + self, + seir_config=None, + compartments_config=None, + compartments_file=None, + transitions_file=None, + ): self.times_set = 0 ## Something like this is needed for check script: @@ -29,7 +35,7 @@ def __init__(self, seir_config=None, compartments_config=None, compartments_file return def constructFromConfig(self, seir_config, compartment_config): - """ + """ This method is called by the constructor if the compartments are not loaded from a file. It will parse the compartments and transitions from the configuration files. It will populate self.compartments and self.transitions. @@ -43,7 +49,7 @@ def __eq__(self, other): ).all().all() def parse_compartments(self, seir_config, compartment_config): - """ Parse the compartments from the configuration file: + """Parse the compartments from the configuration file: seir_config: the configuration file for the SEIR model compartment_config: the configuration file for the compartments Example: if config says: @@ -75,12 +81,16 @@ def parse_compartments(self, seir_config, compartment_config): else: compartment_df = pd.merge(compartment_df, tmp, on="key") compartment_df = compartment_df.drop(["key"], axis=1) - compartment_df["name"] = compartment_df.apply(lambda x: reduce(lambda a, b: a + "_" + b, x), axis=1) + compartment_df["name"] = compartment_df.apply( + lambda x: reduce(lambda a, b: a + "_" + b, x), axis=1 + ) return compartment_df def parse_transitions(self, seir_config, fake_config=False): rc = reduce( - lambda a, b: pd.concat([a, self.parse_single_transition(seir_config, b, fake_config)]), + lambda a, b: pd.concat( + [a, self.parse_single_transition(seir_config, b, fake_config)] + ), seir_config["transitions"], pd.DataFrame(), ) @@ -93,12 +103,17 @@ def check_transition_element(self, single_transition_config, problem_dimension=N def check_transition_elements(self, single_transition_config, problem_dimension): return True - def access_original_config_by_multi_index(self, config_piece, index, dimension=None, encapsulate_as_list=False): + def access_original_config_by_multi_index( + self, config_piece, index, dimension=None, encapsulate_as_list=False + ): if dimension is None: dimension = [None for i in index] tmp = [y for y in zip(index, range(len(index)), dimension)] tmp = zip(index, range(len(index)), dimension) - tmp = [list_access_element_safe(config_piece[x[1]], x[0], x[2], encapsulate_as_list) for x in tmp] + tmp = [ + list_access_element_safe(config_piece[x[1]], x[0], x[2], encapsulate_as_list) + for x in tmp + ] return tmp def expand_transition_elements(self, single_transition_config, problem_dimension): @@ -108,7 +123,9 @@ def expand_transition_elements(self, single_transition_config, problem_dimension # replace "source" by the actual source from the config for p_idx in range(proportion_size): if new_transition_config["proportional_to"][p_idx] == "source": - new_transition_config["proportional_to"][p_idx] = new_transition_config["source"] + new_transition_config["proportional_to"][p_idx] = new_transition_config[ + "source" + ] temp_array = np.zeros(problem_dimension) @@ -117,13 +134,21 @@ def expand_transition_elements(self, single_transition_config, problem_dimension new_transition_config["rate"] = np.zeros(problem_dimension, dtype=object) new_transition_config["proportional_to"] = np.zeros(problem_dimension, dtype=object) - new_transition_config["proportion_exponent"] = np.zeros(problem_dimension, dtype=object) + new_transition_config["proportion_exponent"] = np.zeros( + problem_dimension, dtype=object + ) - it = np.nditer(temp_array, flags=["multi_index"]) # it is an iterator that will go through all the indexes of the array + it = np.nditer( + temp_array, flags=["multi_index"] + ) # it is an iterator that will go through all the indexes of the array for x in it: try: - new_transition_config["source"][it.multi_index] = list_recursive_convert_to_string( - self.access_original_config_by_multi_index(single_transition_config["source"], it.multi_index) + new_transition_config["source"][it.multi_index] = ( + list_recursive_convert_to_string( + self.access_original_config_by_multi_index( + single_transition_config["source"], it.multi_index + ) + ) ) except Exception as e: print(f"Error {e}:") @@ -134,8 +159,12 @@ def expand_transition_elements(self, single_transition_config, problem_dimension raise e try: - new_transition_config["destination"][it.multi_index] = list_recursive_convert_to_string( - self.access_original_config_by_multi_index(single_transition_config["destination"], it.multi_index) + new_transition_config["destination"][it.multi_index] = ( + list_recursive_convert_to_string( + self.access_original_config_by_multi_index( + single_transition_config["destination"], it.multi_index + ) + ) ) except Exception as e: print(f"Error {e}:") @@ -144,10 +173,14 @@ def expand_transition_elements(self, single_transition_config, problem_dimension print(f">>> this transition destination is: '{single_transition_config['destination']}'") print(f"transition_dimension: '{problem_dimension}'") raise e - + try: - new_transition_config["rate"][it.multi_index] = list_recursive_convert_to_string( - self.access_original_config_by_multi_index(single_transition_config["rate"], it.multi_index) + new_transition_config["rate"][it.multi_index] = ( + list_recursive_convert_to_string( + self.access_original_config_by_multi_index( + single_transition_config["rate"], it.multi_index + ) + ) ) except Exception as e: print(f"Error {e}:") @@ -178,23 +211,27 @@ def expand_transition_elements(self, single_transition_config, problem_dimension print(f">>> this transition destination is: '{single_transition_config['destination']}'") print(f"transition_dimension: '{problem_dimension}'") raise e - - if "proportion_exponent" in single_transition_config: # if proportion_exponent is not defined, it is set to 1 + + if ( + "proportion_exponent" in single_transition_config + ): # if proportion_exponent is not defined, it is set to 1 try: self.access_original_config_by_multi_index( single_transition_config["proportion_exponent"][0], it.multi_index, problem_dimension, ) - new_transition_config["proportion_exponent"][it.multi_index] = list_recursive_convert_to_string( - [ - self.access_original_config_by_multi_index( - single_transition_config["proportion_exponent"][p_idx], - it.multi_index, - problem_dimension, - ) - for p_idx in range(proportion_size) - ] + new_transition_config["proportion_exponent"][it.multi_index] = ( + list_recursive_convert_to_string( + [ + self.access_original_config_by_multi_index( + single_transition_config["proportion_exponent"][p_idx], + it.multi_index, + problem_dimension, + ) + for p_idx in range(proportion_size) + ] + ) ) except Exception as e: print(f"Error {e}:") @@ -204,12 +241,19 @@ def expand_transition_elements(self, single_transition_config, problem_dimension print(f"transition_dimension: '{problem_dimension}'") raise e else: - new_transition_config["proportion_exponent"][it.multi_index] = ["1"] * proportion_size + new_transition_config["proportion_exponent"][it.multi_index] = [ + "1" + ] * proportion_size return new_transition_config def format_source(self, source_column): - rc = [y for y in map(lambda x: reduce(lambda a, b: str(a) + "_" + str(b), x), source_column)] + rc = [ + y + for y in map( + lambda x: reduce(lambda a, b: str(a) + "_" + str(b), x), source_column + ) + ] return rc def unformat_source(self, source_column): @@ -231,7 +275,12 @@ def unformat_destination(self, destination_column): return rc def format_rate(self, rate_column): - rc = [y for y in map(lambda x: reduce(lambda a, b: str(a) + "%*%" + str(b), x), rate_column)] + rc = [ + y + for y in map( + lambda x: reduce(lambda a, b: str(a) + "%*%" + str(b), x), rate_column + ) + ] return rc def unformat_rate(self, rate_column, compartment_dimension): @@ -251,7 +300,9 @@ def format_proportional_to(self, proportional_to_column): lambda x: reduce( lambda a, b: str(a) + "_" + str(b), map( - lambda x: reduce(lambda a, b: str(a) + "+" + str(b), as_list(x)), + lambda x: reduce( + lambda a, b: str(a) + "+" + str(b), as_list(x) + ), x, ), ), @@ -284,7 +335,9 @@ def format_proportion_exponent(self, proportion_exponent_column): ] return rc - def unformat_proportion_exponent(self, proportion_exponent_column, compartment_dimension): + def unformat_proportion_exponent( + self, proportion_exponent_column, compartment_dimension + ): rc = [x.split("%*%") for x in proportion_exponent_column] for row in range(len(rc)): rc[row] = [x.split("*", maxsplit=compartment_dimension - 1) for x in rc[row]] @@ -293,18 +346,28 @@ def unformat_proportion_exponent(self, proportion_exponent_column, compartment_d elem.append(1) return rc - def parse_single_transition(self, seir_config, single_transition_config, fake_config=False): + def parse_single_transition( + self, seir_config, single_transition_config, fake_config=False + ): ## This method relies on having run parse_compartments if not fake_config: single_transition_config = single_transition_config.get() self.check_transition_element(single_transition_config["source"]) self.check_transition_element(single_transition_config["destination"]) - source_dimension = [get_list_dimension(x) for x in single_transition_config["source"]] - destination_dimension = [get_list_dimension(x) for x in single_transition_config["destination"]] - problem_dimension = reduce(lambda x, y: max(x, y), (source_dimension, destination_dimension)) + source_dimension = [ + get_list_dimension(x) for x in single_transition_config["source"] + ] + destination_dimension = [ + get_list_dimension(x) for x in single_transition_config["destination"] + ] + problem_dimension = reduce( + lambda x, y: max(x, y), (source_dimension, destination_dimension) + ) self.check_transition_elements(single_transition_config, problem_dimension) - transitions = self.expand_transition_elements(single_transition_config, problem_dimension) + transitions = self.expand_transition_elements( + single_transition_config, problem_dimension + ) tmp_array = np.zeros(problem_dimension) it = np.nditer(tmp_array, flags=["multi_index"]) @@ -317,7 +380,9 @@ def parse_single_transition(self, seir_config, single_transition_config, fake_co "destination": [transitions["destination"][it.multi_index]], "rate": [transitions["rate"][it.multi_index]], "proportional_to": [transitions["proportional_to"][it.multi_index]], - "proportion_exponent": [transitions["proportion_exponent"][it.multi_index]], + "proportion_exponent": [ + transitions["proportion_exponent"][it.multi_index] + ], }, index=[0], ) @@ -328,7 +393,10 @@ def parse_single_transition(self, seir_config, single_transition_config, fake_co return rc def toFile( - self, compartments_file="compartments.parquet", transitions_file="transitions.parquet", write_parquet=True + self, + compartments_file="compartments.parquet", + transitions_file="transitions.parquet", + write_parquet=True, ): out_df = self.compartments.copy() if write_parquet: @@ -342,7 +410,9 @@ def toFile( out_df["destination"] = self.format_destination(out_df["destination"]) out_df["rate"] = self.format_rate(out_df["rate"]) out_df["proportional_to"] = self.format_proportional_to(out_df["proportional_to"]) - out_df["proportion_exponent"] = self.format_proportion_exponent(out_df["proportion_exponent"]) + out_df["proportion_exponent"] = self.format_proportion_exponent( + out_df["proportion_exponent"] + ) if write_parquet: pa_df = pa.Table.from_pandas(out_df, preserve_index=False) pa.parquet.write_table(pa_df, transitions_file) @@ -355,9 +425,15 @@ def fromFile(self, compartments_file, transitions_file): self.transitions = pq.read_table(transitions_file).to_pandas() compartment_dimension = self.compartments.shape[1] - 1 self.transitions["source"] = self.unformat_source(self.transitions["source"]) - self.transitions["destination"] = self.unformat_destination(self.transitions["destination"]) - self.transitions["rate"] = self.unformat_rate(self.transitions["rate"], compartment_dimension) - self.transitions["proportional_to"] = self.unformat_proportional_to(self.transitions["proportional_to"]) + self.transitions["destination"] = self.unformat_destination( + self.transitions["destination"] + ) + self.transitions["rate"] = self.unformat_rate( + self.transitions["rate"], compartment_dimension + ) + self.transitions["proportional_to"] = self.unformat_proportional_to( + self.transitions["proportional_to"] + ) self.transitions["proportion_exponent"] = self.unformat_proportion_exponent( self.transitions["proportion_exponent"], compartment_dimension ) @@ -371,7 +447,9 @@ def get_comp_idx(self, comp_dict: dict, error_info: str = "no information") -> i :param comp_dict: :return: """ - mask = pd.concat([self.compartments[k] == v for k, v in comp_dict.items()], axis=1).all(axis=1) + mask = pd.concat( + [self.compartments[k] == v for k, v in comp_dict.items()], axis=1 + ).all(axis=1) comp_idx = self.compartments[mask].index.values if len(comp_idx) != 1: raise ValueError( @@ -384,10 +462,11 @@ def get_comp_idx(self, comp_dict: dict, error_info: str = "no information") -> i def get_ncomp(self) -> int: return len(self.compartments) - def get_transition_array(self): with Timer("SEIR.compartments"): - transition_array = np.zeros((self.transitions.shape[1], self.transitions.shape[0]), dtype="int64") + transition_array = np.zeros( + (self.transitions.shape[1], self.transitions.shape[0]), dtype="int64" + ) for cit, colname in enumerate(("source", "destination")): for it, elem in enumerate(self.transitions[colname]): elem = reduce(lambda a, b: a + "_" + b, elem) @@ -418,8 +497,12 @@ def get_transition_array(self): # parenthesis are now supported # assert reduce(lambda a, b: a and b, [(x.find("(") == -1) for x in unique_strings]) # assert reduce(lambda a, b: a and b, [(x.find(")") == -1) for x in unique_strings]) - assert reduce(lambda a, b: a and b, [(x.find("%") == -1) for x in unique_strings]) - assert reduce(lambda a, b: a and b, [(x.find(" ") == -1) for x in unique_strings]) + assert reduce( + lambda a, b: a and b, [(x.find("%") == -1) for x in unique_strings] + ) + assert reduce( + lambda a, b: a and b, [(x.find(" ") == -1) for x in unique_strings] + ) for it, elem in enumerate(self.transitions["rate"]): candidate = reduce(lambda a, b: a + "*" + b, elem) @@ -457,8 +540,12 @@ def get_transition_array(self): # rc = compartment # if rc == -1: # raise ValueError(f"Could not find match for {elem3} in compartments") - proportion_info[0][current_proportion_sum_it] = current_proportion_sum_start - proportion_info[1][current_proportion_sum_it] = current_proportion_sum_start + len(elem_tmp) + proportion_info[0][ + current_proportion_sum_it + ] = current_proportion_sum_start + proportion_info[1][current_proportion_sum_it] = ( + current_proportion_sum_start + len(elem_tmp) + ) current_proportion_sum_it += 1 current_proportion_sum_start += len(elem_tmp) proportion_compartment_index = 0 @@ -536,18 +623,24 @@ def get_transition_array(self): def parse_parameters(self, parameters, parameter_names, unique_strings): # parsed_parameters_old = self.parse_parameter_strings_to_numpy_arrays(parameters, parameter_names, unique_strings) - parsed_parameters = self.parse_parameter_strings_to_numpy_arrays_v2(parameters, parameter_names, unique_strings) + parsed_parameters = self.parse_parameter_strings_to_numpy_arrays_v2( + parameters, parameter_names, unique_strings + ) # for i in range(len(unique_strings)): # print(unique_strings[i], (parsed_parameters[i]==parsed_parameters_old[i]).all()) return parsed_parameters - def parse_parameter_strings_to_numpy_arrays_v2(self, parameters, parameter_names, string_list): + def parse_parameter_strings_to_numpy_arrays_v2( + self, parameters, parameter_names, string_list + ): # is using eval a better way ??? import sympy as sp # Validate input lengths if len(parameters) != len(parameter_names): - raise ValueError("Number of parameter values does not match the number of parameter names.") + raise ValueError( + "Number of parameter values does not match the number of parameter names." + ) # Define the symbols used in the formulas symbolic_parameters_namespace = {name: sp.symbols(name) for name in parameter_names} @@ -568,7 +661,9 @@ def parse_parameter_strings_to_numpy_arrays_v2(self, parameters, parameter_names raise (e) # Print the error message for debugging # the list order needs to be right. - parameter_values = {param: value for param, value in zip(symbolic_parameters, parameters)} + parameter_values = { + param: value for param, value in zip(symbolic_parameters, parameters) + } parameter_values_list = [parameter_values[param] for param in symbolic_parameters] # Create a lambdify function for substitution @@ -583,7 +678,9 @@ def parse_parameter_strings_to_numpy_arrays_v2(self, parameters, parameter_names if not isinstance(substituted_formulas[i], np.ndarray): for k in range(len(substituted_formulas)): if isinstance(substituted_formulas[k], np.ndarray): - substituted_formulas[i] = substituted_formulas[i] * np.ones_like(substituted_formulas[k]) + substituted_formulas[i] = substituted_formulas[i] * np.ones_like( + substituted_formulas[k] + ) return np.array(substituted_formulas) @@ -633,17 +730,23 @@ def parse_parameter_strings_to_numpy_arrays( if not is_totally_resolvable: not_resolvable_indices = [it for it, x in enumerate(is_resolvable) if not x] - tmp_rc[not_resolvable_indices] = self.parse_parameter_strings_to_numpy_arrays( - parameters, - parameter_names, - [string[not is_resolvable]], - operator_reduce_lambdas, - operators[1:], + tmp_rc[not_resolvable_indices] = ( + self.parse_parameter_strings_to_numpy_arrays( + parameters, + parameter_names, + [string[not is_resolvable]], + operator_reduce_lambdas, + operators[1:], + ) ) for numeric_index in [x for x in range(len(is_numeric)) if is_numeric[x]]: tmp_rc[numeric_index] = parameters[0] * 0 + float(string[numeric_index]) for parameter_index in [x for x in range(len(is_parameter)) if is_parameter[x]]: - parameter_name_index = [it for it, x in enumerate(parameter_names) if x == string[parameter_index]] + parameter_name_index = [ + it + for it, x in enumerate(parameter_names) + if x == string[parameter_index] + ] tmp_rc[parameter_index] = parameters[parameter_name_index] rc[sit] = reduce(operator_reduce_lambdas[operators[0]], tmp_rc) @@ -658,7 +761,9 @@ def get_compartments_explicitDF(self): df = df.rename(columns=rename_dict) return df - def plot(self, output_file="transition_graph", source_filters=[], destination_filters=[]): + def plot( + self, output_file="transition_graph", source_filters=[], destination_filters=[] + ): """ if source_filters is [["age0to17"], ["OMICRON", "WILD"]], it means filter all transitions that have as source age0to17 AND (OMICRON OR WILD). @@ -729,7 +834,6 @@ def list_access_element_safe(thing, idx, dimension=None, encapsulate_as_list=Fal ) - def list_access_element(thing, idx, dimension=None, encapsulate_as_list=False): """ This function is used to access elements in a list or a single element. @@ -767,7 +871,9 @@ def compartments(): def plot(): assert config["compartments"].exists() assert config["seir"].exists() - comp = Compartments(seir_config=config["seir"], compartments_config=config["compartments"]) + comp = Compartments( + seir_config=config["seir"], compartments_config=config["compartments"] + ) # TODO: this should be a command like build compartments. ( @@ -786,7 +892,9 @@ def plot(): def export(): assert config["compartments"].exists() assert config["seir"].exists() - comp = Compartments(seir_config=config["seir"], compartments_config=config["compartments"]) + comp = Compartments( + seir_config=config["seir"], compartments_config=config["compartments"] + ) ( unique_strings, transition_array, diff --git a/flepimop/gempyor_pkg/src/gempyor/config_validator.py b/flepimop/gempyor_pkg/src/gempyor/config_validator.py index 93aa4d6ff..74909631c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/config_validator.py +++ b/flepimop/gempyor_pkg/src/gempyor/config_validator.py @@ -1,33 +1,60 @@ import yaml -from pydantic import BaseModel, ValidationError, model_validator, Field, AfterValidator, validator +from pydantic import ( + BaseModel, + ValidationError, + model_validator, + Field, + AfterValidator, + validator, +) from datetime import date from typing import Dict, List, Union, Literal, Optional, Annotated, Any from functools import partial from gempyor import compartments + def read_yaml(file_path: str) -> dict: - with open(file_path, 'r') as stream: + with open(file_path, "r") as stream: config = yaml.safe_load(stream) - + return CheckConfig(**config).model_dump() - + + def allowed_values(v, values): assert v in values return v + # def parse_value(cls, values): # value = values.get('value') # parsed_val = compartments.Compartments.parse_parameter_strings_to_numpy_arrays_v2(value) # return parsed_val - + + class SubpopSetupConfig(BaseModel): geodata: str mobility: Optional[str] selected: List[str] = Field(default_factory=list) # state_level: Optional[bool] = False # pretty sure this doesn't exist anymore + class InitialConditionsConfig(BaseModel): - method: Annotated[str, AfterValidator(partial(allowed_values, values=['Default', 'SetInitialConditions', 'SetInitialConditionsFolderDraw', 'InitialConditionsFolderDraw', 'FromFile', 'plugin']))] = 'Default' + method: Annotated[ + str, + AfterValidator( + partial( + allowed_values, + values=[ + "Default", + "SetInitialConditions", + "SetInitialConditionsFolderDraw", + "InitialConditionsFolderDraw", + "FromFile", + "plugin", + ], + ) + ), + ] = "Default" initial_file_type: Optional[str] = None initial_conditions_file: Optional[str] = None proportional: Optional[bool] = None @@ -36,7 +63,7 @@ class InitialConditionsConfig(BaseModel): ignore_population_checks: Optional[bool] = None plugin_file_path: Optional[str] = None - @model_validator(mode='before') + @model_validator(mode="before") def validate_initial_file_check(cls, values): method = values.get('method') initial_conditions_file = values.get('initial_conditions_file') @@ -46,8 +73,8 @@ def validate_initial_file_check(cls, values): if method in {'InitialConditionsFolderDraw','SetInitialConditionsFolderDraw'} and not initial_file_type: raise ValueError(f"An `initial_file_type` is required when method is '{method}'.") return values - - @model_validator(mode='before') + + @model_validator(mode="before") def plugin_filecheck(cls, values): method = values.get('method') plugin_file_path = values.get('plugin_file_path') @@ -57,13 +84,27 @@ def plugin_filecheck(cls, values): class SeedingConfig(BaseModel): - method: Annotated[str, AfterValidator(partial(allowed_values, values=['NoSeeding', 'PoissonDistributed', 'FolderDraw', 'FromFile', 'plugin']))] = 'NoSeeding' # note: removed NegativeBinomialDistributed because no longer supported + method: Annotated[ + str, + AfterValidator( + partial( + allowed_values, + values=[ + "NoSeeding", + "PoissonDistributed", + "FolderDraw", + "FromFile", + "plugin", + ], + ) + ), + ] = "NoSeeding" # note: removed NegativeBinomialDistributed because no longer supported lambda_file: Optional[str] = None seeding_file_type: Optional[str] = None seeding_file: Optional[str] = None plugin_file_path: Optional[str] = None - @model_validator(mode='before') + @model_validator(mode="before") def validate_seedingfile(cls, values): method = values.get('method') lambda_file = values.get('lambda_file') @@ -76,8 +117,8 @@ def validate_seedingfile(cls, values): if method == 'FromFile' and not seeding_file: raise ValueError("A `seeding_file` is required when method is 'FromFile'.") return values - - @model_validator(mode='before') + + @model_validator(mode="before") def plugin_filecheck(cls, values): method = values.get('method') plugin_file_path = values.get('plugin_file_path') @@ -87,28 +128,35 @@ def plugin_filecheck(cls, values): f"was given '{plugin_file_path!r}'." ) return values - + + class IntegrationConfig(BaseModel): - method: Annotated[str, AfterValidator(partial(allowed_values, values=['rk4', 'rk4.jit', 'best.current', 'legacy']))] = 'rk4' + method: Annotated[ + str, + AfterValidator( + partial(allowed_values, values=["rk4", "rk4.jit", "best.current", "legacy"]) + ), + ] = "rk4" dt: float = 2.0 + class ValueConfig(BaseModel): - distribution: str = 'fixed' - value: Optional[float] = None # NEED TO ADD ABILITY TO PARSE PARAMETERS + distribution: str = "fixed" + value: Optional[float] = None # NEED TO ADD ABILITY TO PARSE PARAMETERS mean: Optional[float] = None sd: Optional[float] = None a: Optional[float] = None b: Optional[float] = None - @model_validator(mode='before') + @model_validator(mode="before") def check_distr(cls, values): - distr = values.get('distribution') - value = values.get('value') - mean = values.get('mean') - sd = values.get('sd') - a = values.get('a') - b = values.get('b') - if distr != 'fixed': + distr = values.get("distribution") + value = values.get("value") + mean = values.get("mean") + sd = values.get("sd") + a = values.get("a") + b = values.get("b") + if distr != "fixed": if not mean and not sd: raise ValueError("Mean and sd must be provided for non-fixed distributions.") if distr == 'truncnorm' and not a and not b: @@ -117,29 +165,38 @@ def check_distr(cls, values): raise ValueError("Value must be provided for fixed distributions.") return values + class BaseParameterConfig(BaseModel): value: Optional[ValueConfig] = None modifier_parameter: Optional[str] = None - name: Optional[str] = None # this is only for outcomes, to build outcome_prevalence_name (how to restrict this?) + name: Optional[str] = ( + None # this is only for outcomes, to build outcome_prevalence_name (how to restrict this?) + ) + class SeirParameterConfig(BaseParameterConfig): value: Optional[ValueConfig] = None - stacked_modifier_method: Annotated[str, AfterValidator(partial(allowed_values, values=['sum', 'product', 'reduction_product']))] = None + stacked_modifier_method: Annotated[ + str, + AfterValidator( + partial(allowed_values, values=["sum", "product", "reduction_product"]) + ), + ] = None rolling_mean_windows: Optional[float] = None timeseries: Optional[str] = None - @model_validator(mode='before') + @model_validator(mode="before") def which_value(cls, values): - value = values.get('value') is not None - timeseries = values.get('timeseries') is not None + value = values.get("value") is not None + timeseries = values.get("timeseries") is not None if value and timeseries: raise ValueError( f"Parsed a parameter with both a timeseries, '{values.get('timeseries')!r}', and a value, '{values.get('value')!r}', please choose one." ) return values - - -class TransitionConfig(BaseModel): + + +class TransitionConfig(BaseModel): # !! sometimes these are lists of lists and sometimes they are lists... how to deal with this? source: List[List[str]] destination: List[List[str]] @@ -147,11 +204,15 @@ class TransitionConfig(BaseModel): proportion_exponent: List[List[str]] proportional_to: List[str] + class SeirConfig(BaseModel): - integration: IntegrationConfig # is this Optional? - parameters: Dict[str, SeirParameterConfig] # there was a previous issue that gempyor doesn't work if there are no parameters (eg if just numbers are used in the transitions) - do we want to get around this? + integration: IntegrationConfig # is this Optional? + parameters: Dict[ + str, SeirParameterConfig + ] # there was a previous issue that gempyor doesn't work if there are no parameters (eg if just numbers are used in the transitions) - do we want to get around this? transitions: List[TransitionConfig] + class SinglePeriodModifierConfig(BaseModel): method: Literal["SinglePeriodModifier"] parameter: str @@ -162,15 +223,18 @@ class SinglePeriodModifierConfig(BaseModel): value: ValueConfig perturbation: Optional[ValueConfig] = None + class MultiPeriodDatesConfig(BaseModel): start_date: date end_date: date - + + class MultiPeriodGroupsConfig(BaseModel): subpop: List[str] subpop_groups: Optional[str] = None periods: List[MultiPeriodDatesConfig] + class MultiPeriodModifierConfig(BaseModel): method: Literal["MultiPeriodModifier"] parameter: str @@ -178,35 +242,43 @@ class MultiPeriodModifierConfig(BaseModel): value: ValueConfig perturbation: Optional[ValueConfig] = None + class StackedModifierConfig(BaseModel): method: Literal["StackedModifier"] modifiers: List[str] + class ModifiersConfig(BaseModel): scenarios: List[str] modifiers: Dict[str, Any] - + @field_validator("modifiers") def validate_data_dict(cls, value: Dict[str, Any]) -> Dict[str, Any]: errors = [] for key, entry in value.items(): method = entry.get("method") - if method not in {"SinglePeriodModifier", "MultiPeriodModifier", "StackedModifier"}: + if method not in { + "SinglePeriodModifier", + "MultiPeriodModifier", + "StackedModifier", + }: errors.append(f"Invalid modifier method: {method}") if errors: raise ValueError("Errors in modifiers:\n" + "\n".join(errors)) return value -class SourceConfig(BaseModel): # set up only for incidence or prevalence. Can this be any name? i don't think so atm +class SourceConfig( + BaseModel +): # set up only for incidence or prevalence. Can this be any name? i don't think so atm incidence: Dict[str, str] = None - prevalence: Dict[str, str] = None + prevalence: Dict[str, str] = None # note: these dictionaries have to have compartment names... more complicated to set this up - @model_validator(mode='before') + @model_validator(mode="before") def which_source(cls, values): - incidence = values.get('incidence') - prevalence = values.get('prevalence') + incidence = values.get("incidence") + prevalence = values.get("prevalence") if incidence and prevalence: raise ValueError( f"Parsed a source with both an incidence, '{values.get('incidence')!r}', and prevalence, '{values.get('prevalence')!r}', please choose one." @@ -221,12 +293,13 @@ def which_source(cls, values): # source_names.append(key) # return source_names # Access keys using a loop + class DelayFrameConfig(BaseModel): source: Optional[SourceConfig] = None probability: Optional[BaseParameterConfig] = None delay: Optional[BaseParameterConfig] = None duration: Optional[BaseParameterConfig] = None - sum: Optional[List[str]] = None # only for sums of other outcomes + sum: Optional[List[str]] = None # only for sums of other outcomes # @validator("sum") # def validate_sum_elements(cls, value: Optional[List[str]]) -> Optional[List[str]]: @@ -240,10 +313,10 @@ class DelayFrameConfig(BaseModel): # return value # note: ^^ this doesn't work yet because it needs to somehow be a level above? to access all OTHER source names - @model_validator(mode='before') + @model_validator(mode="before") def check_outcome_type(cls, values): - sum_present = values.get('sum') is not None - source_present = values.get('source') is not None + sum_present = values.get("sum") is not None + source_present = values.get("source") is not None if sum_present and source_present: raise ValueError("Both 'sum' and 'source' are present, please choose one.") @@ -251,54 +324,77 @@ def check_outcome_type(cls, values): raise ValueError("Neither 'sum' nor 'source' is present, please choose one.") return values + class OutcomesConfig(BaseModel): - method: Literal["delayframe"] # Is this required? I don't see it anywhere in the gempyor code + method: Literal[ + "delayframe" + ] # Is this required? I don't see it anywhere in the gempyor code param_from_file: Optional[bool] = None param_subpop_file: Optional[str] = None outcomes: Dict[str, DelayFrameConfig] - - @model_validator(mode='before') + + @model_validator(mode="before") def check_paramfromfile_type(cls, values): - param_from_file = values.get('param_from_file') is not None - param_subpop_file = values.get('param_subpop_file') is not None + param_from_file = values.get("param_from_file") is not None + param_subpop_file = values.get("param_subpop_file") is not None if param_from_file and not param_subpop_file: raise ValueError("A `param_subpop_file` is required when `param_from_file` is 'True'.") return values + class ResampleConfig(BaseModel): aggregator: Optional[str] = None freq: Optional[str] = None skipna: Optional[bool] = False + class LikelihoodParams(BaseModel): scale: float # are there other options here? + class LikelihoodReg(BaseModel): - name: str + name: str + class LikelihoodConfig(BaseModel): - dist: Annotated[str, AfterValidator(partial(allowed_values, values=['pois', 'norm', 'norm_cov', 'nbinom', 'rmse', 'absolute_error']))] = None + dist: Annotated[ + str, + AfterValidator( + partial( + allowed_values, + values=["pois", "norm", "norm_cov", "nbinom", "rmse", "absolute_error"], + ) + ), + ] = None params: Optional[LikelihoodParams] = None + class StatisticsConfig(BaseModel): name: str sim_var: str data_var: str regularize: Optional[LikelihoodReg] = None resample: Optional[ResampleConfig] = None - scale: Optional[float] = None # is scale here or at likelihood level? - zero_to_one: Optional[bool] = False # is this the same as add_one? remove_na? + scale: Optional[float] = None # is scale here or at likelihood level? + zero_to_one: Optional[bool] = False # is this the same as add_one? remove_na? likelihood: LikelihoodConfig + class InferenceConfig(BaseModel): - method: Annotated[str, AfterValidator(partial(allowed_values, values=['emcee', 'default', 'classical']))] = 'default' # for now - i can only see emcee as an option here, otherwise ignored in classical - need to add these options - iterations_per_slot: Optional[int] # i think this is optional because it is also set in command line?? - do_inference: bool + method: Annotated[ + str, + AfterValidator(partial(allowed_values, values=["emcee", "default", "classical"])), + ] = "default" # for now - i can only see emcee as an option here, otherwise ignored in classical - need to add these options + iterations_per_slot: Optional[ + int + ] # i think this is optional because it is also set in command line?? + do_inference: bool gt_data_path: str statistics: Dict[str, StatisticsConfig] + class CheckConfig(BaseModel): name: str setup_name: Optional[str] = None @@ -319,13 +415,13 @@ class CheckConfig(BaseModel): outcome_modifiers: Optional[ModifiersConfig] = None inference: Optional[InferenceConfig] = None -# add validator for if modifiers exist but seir/outcomes do not - -# there is an error in the one below - @model_validator(mode='before') + # add validator for if modifiers exist but seir/outcomes do not + + # there is an error in the one below + @model_validator(mode="before") def verify_inference(cls, values): - inference_present = values.get('inference') is not None - start_date_groundtruth = values.get('start_date_groundtruth') is not None + inference_present = values.get("inference") is not None + start_date_groundtruth = values.get("start_date_groundtruth") is not None if inference_present and not start_date_groundtruth: raise ValueError("Inference mode is enabled, but no groundtruth dates are provided. Please provide groundtruth dates.") elif start_date_groundtruth and not inference_present: @@ -333,19 +429,19 @@ def verify_inference(cls, values): "Groundtruth dates are provided, but inference mode is not enabled. Please enable inference mode." ) return values - - @model_validator(mode='before') + + @model_validator(mode="before") def check_dates(cls, values): - start_date = values.get('start_date') - end_date = values.get('end_date') + start_date = values.get("start_date") + end_date = values.get("end_date") if start_date and end_date: if end_date <= start_date: raise ValueError( f"`end_date` ('{end_date}') must be later than `start_date` ('{start_date}'))." ) return values - - @model_validator(mode='before') + + @model_validator(mode="before") def init_or_seed(cls, values): init = values.get('initial_conditions') seed = values.get('seeding') diff --git a/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py b/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py index a847a5f90..c484a3980 100644 --- a/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py +++ b/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py @@ -48,7 +48,9 @@ modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -81,7 +83,12 @@ True, ) df = seir.states2Df(modinf, states) -assert df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "20002"] > 1 +assert ( + df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[ + str(modinf.tf), "20002" + ] + > 1 +) print(df) ts = df cp = "R" diff --git a/flepimop/gempyor_pkg/src/gempyor/dev/steps.py b/flepimop/gempyor_pkg/src/gempyor/dev/steps.py index 43066e5ee..1bbf5b207 100644 --- a/flepimop/gempyor_pkg/src/gempyor/dev/steps.py +++ b/flepimop/gempyor_pkg/src/gempyor/dev/steps.py @@ -53,7 +53,11 @@ def ode_integration( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -65,7 +69,9 @@ def ode_integration( def rhs(t, x, today): print("rhs.t", t) states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -80,9 +86,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -91,41 +101,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -142,9 +167,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -168,18 +199,22 @@ def rhs(t, x, today): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -257,7 +292,11 @@ def rk4_integration1( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -267,7 +306,9 @@ def rk4_integration1( def rhs(t, x, today): states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -282,9 +323,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -293,41 +338,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -344,9 +404,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -380,18 +446,22 @@ def rk4_integrate(t, x, today): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -447,7 +517,11 @@ def rk4_integration2( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -458,7 +532,9 @@ def rk4_integration2( @jit(nopython=True) def rhs(t, x, today): states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -473,9 +549,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -484,41 +564,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -535,9 +630,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -572,18 +673,22 @@ def rk4_integrate(t, x, today): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -644,7 +749,11 @@ def rk4_integration3( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -655,7 +764,9 @@ def rk4_integration3( @jit(nopython=True) def rhs(t, x, today): states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -670,9 +781,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -681,41 +796,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -732,9 +862,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -753,16 +889,18 @@ def rk4_integrate(t, x, today): @jit(nopython=True) def day_wrapper_rk4(today, states_next): x_ = np.zeros((2, ncompartments, nspatial_nodes)) - for seeding_instance_idx in range(day_start_idx_dict[today], day_start_idx_dict[today + 1]): + for seeding_instance_idx in range( + day_start_idx_dict[today], day_start_idx_dict[today + 1] + ): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_subpops_dict[seeding_instance_idx] seeding_sources = seeding_sources_dict[seeding_instance_idx] seeding_destinations = seeding_destinations_dict[seeding_instance_idx] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][ + seeding_subpops + ] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts # ADD TO cumulative, this is debatable, @@ -838,7 +976,11 @@ def rk4_integration4( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -849,7 +991,9 @@ def rk4_integration4( @jit(nopython=True) # , fastmath=True, parallel=True) def rhs(t, x, today): states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -864,9 +1008,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -875,41 +1023,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -926,9 +1089,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -963,18 +1132,22 @@ def rk4_integrate(t, x, today): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -1031,7 +1204,11 @@ def rk4_integration5( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -1058,18 +1235,22 @@ def rk4_integration5( this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -1092,7 +1273,9 @@ def rk4_integration5( x = x_ + kx[i - 1] * rk_coefs[i] states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -1107,45 +1290,62 @@ def rk4_integration5( proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] - # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][ - today + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] ] + # exponents should not be a proportion, since we don't sum them over sum compartments + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( - transitions[transition_proportion_start_col][transition_index] + 1 - ) == transitions[transition_proportion_stop_col][transition_index] + transitions[transition_proportion_start_col][ + transition_index + ] + + 1 + ) == transitions[transition_proportion_stop_col][ + transition_index + ] first_proportion = False source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[ + spatial_node + ] : mobility_data_indices[spatial_node + 1] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][ - spatial_node - ] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[ + spatial_node + ] : mobility_data_indices[spatial_node + 1] ] rate_change_compartment = proportion_change_compartment @@ -1157,7 +1357,9 @@ def rk4_integration5( rate_change_compartment *= parameters[ transitions[transition_rate_col][transition_index] ][today][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -1174,8 +1376,12 @@ def rk4_integration5( # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move states_diff[ 1, transitions[transition_destination_col][transition_index], : ] += number_move # Cumumlative @@ -1234,7 +1440,11 @@ def rk4_integration2_smart( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -1248,7 +1458,9 @@ def rhs(t, x): if (today) > ndays: today = ndays - 1 states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -1263,9 +1475,13 @@ def rhs(t, x): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -1274,41 +1490,56 @@ def rhs(t, x): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -1325,9 +1556,15 @@ def rhs(t, x): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return np.reshape(states_diff, states_diff.size) # return a 1D vector @@ -1374,18 +1611,24 @@ def rk4_integrate(today, x): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) - states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) + states_next[seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape @@ -1450,9 +1693,12 @@ def rk4_integrate(today, x): ## Return "UniTuple(float64[:, :, :], 2) (" ## return states and cumlative states, both [ ndays x ncompartments x nspatial_nodes ] ## Dimensions - "int32," "int32," "int32," ## ncompartments ## nspatial_nodes ## Number of days + "int32," + "int32," + "int32," ## ncompartments ## nspatial_nodes ## Number of days ## Parameters - "float64[:, :, :]," "float64," ## Parameters [ nparameters x ndays x nspatial_nodes] ## dt + "float64[:, :, :]," + "float64," ## Parameters [ nparameters x ndays x nspatial_nodes] ## dt ## Transitions "int64[:, :]," ## transitions [ [source, destination, proportion_start, proportion_stop, rate] x ntransitions ] "int64[:, :]," ## proportions_info [ [sum_starts, sum_stops, exponent] x ntransition_proportions ] @@ -1504,7 +1750,11 @@ def rk4_integration_aot( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -1515,7 +1765,9 @@ def rk4_integration_aot( def rhs(t, x, today): # states_current = np.reshape(x, (2, ncompartments, nspatial_nodes))[0] states_current = x[0] - states_diff = np.zeros((2, ncompartments, nspatial_nodes)) # first dim: 0 -> states_diff, 1: states_cum + states_diff = np.zeros( + (2, ncompartments, nspatial_nodes) + ) # first dim: 0 -> states_diff, 1: states_cum for transition_index in range(ntransitions): total_rate = np.ones((nspatial_nodes)) @@ -1530,9 +1782,13 @@ def rhs(t, x, today): proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -1541,41 +1797,56 @@ def rhs(t, x, today): source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) # compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -1592,9 +1863,15 @@ def rhs(t, x, today): # if number_move[spatial_node] > states_current[transitions[transition_source_col][transition_index]][spatial_node]: # number_move[spatial_node] = states_current[transitions[transition_source_col][transition_index]][spatial_node] # Not possible to enforce this anymore, but it shouldn't be a problem or maybe ? # TODO - states_diff[0, transitions[transition_source_col][transition_index]] -= number_move - states_diff[0, transitions[transition_destination_col][transition_index]] += number_move - states_diff[1, transitions[transition_destination_col][transition_index], :] += number_move # Cumumlative + states_diff[ + 0, transitions[transition_source_col][transition_index] + ] -= number_move + states_diff[ + 0, transitions[transition_destination_col][transition_index] + ] += number_move + states_diff[ + 1, transitions[transition_destination_col][transition_index], : + ] += number_move # Cumumlative # states_current = states_next.copy() return states_diff # return a 1D vector @@ -1628,18 +1905,22 @@ def rk4_integrate(t, x, today): this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts ### Shape diff --git a/flepimop/gempyor_pkg/src/gempyor/inference.py b/flepimop/gempyor_pkg/src/gempyor/inference.py index 25289349d..aae76eaff 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference.py @@ -82,11 +82,17 @@ def simulation_atomic( np.random.seed(int.from_bytes(os.urandom(4), byteorder="little")) random_id = np.random.randint(0, 1e8) - npi_seir = seir.build_npi_SEIR(modinf=modinf, load_ID=False, sim_id2load=None, config=config, bypass_DF=snpi_df_in) + npi_seir = seir.build_npi_SEIR( + modinf=modinf, load_ID=False, sim_id2load=None, config=config, bypass_DF=snpi_df_in + ) if modinf.npi_config_outcomes: npi_outcomes = outcomes.build_outcome_modifiers( - modinf=modinf, load_ID=False, sim_id2load=None, config=config, bypass_DF=hnpi_df_in + modinf=modinf, + load_ID=False, + sim_id2load=None, + config=config, + bypass_DF=hnpi_df_in, ) else: npi_outcomes = None @@ -94,10 +100,14 @@ def simulation_atomic( # reduce them parameters = modinf.parameters.parameters_reduce(p_draw, npi_seir) # Parse them - parsed_parameters = modinf.compartments.parse_parameters(parameters, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + parameters, modinf.parameters.pnames, unique_strings + ) # Convert the seeding data dictionnary to a numba dictionnary - seeding_data_nbdict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type=nb.types.int64[:]) + seeding_data_nbdict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.int64[:] + ) for k, v in seeding_data.items(): seeding_data_nbdict[k] = np.array(v, dtype=np.int64) @@ -151,7 +161,9 @@ def get_static_arguments(modinf): ) = modinf.compartments.get_transition_array() outcomes_parameters = outcomes.read_parameters_from_config(modinf) - npi_seir = seir.build_npi_SEIR(modinf=modinf, load_ID=False, sim_id2load=None, config=config) + npi_seir = seir.build_npi_SEIR( + modinf=modinf, load_ID=False, sim_id2load=None, config=config + ) if modinf.npi_config_outcomes: npi_outcomes = outcomes.build_outcome_modifiers( modinf=modinf, @@ -162,7 +174,9 @@ def get_static_arguments(modinf): else: npi_outcomes = None - p_draw = modinf.parameters.parameters_quick_draw(n_days=modinf.n_days, nsubpops=modinf.nsubpops) + p_draw = modinf.parameters.parameters_quick_draw( + n_days=modinf.n_days, nsubpops=modinf.nsubpops + ) initial_conditions = modinf.initial_conditions.get_from_config(sim_id=0, modinf=modinf) seeding_data, seeding_amounts = modinf.seeding.get_from_config(sim_id=0, modinf=modinf) @@ -170,7 +184,9 @@ def get_static_arguments(modinf): # reduce them parameters = modinf.parameters.parameters_reduce(p_draw, npi_seir) # Parse them - parsed_parameters = modinf.compartments.parse_parameters(parameters, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + parameters, modinf.parameters.pnames, unique_strings + ) if real_simulation: states = seir.steps_SEIR( @@ -198,7 +214,9 @@ def get_static_arguments(modinf): subpop=modinf.subpop_struct.subpop_names, ) - zeros = np.zeros((len(coords["date"]), len(coords["mc_name"][1]), len(coords["subpop"]))) + zeros = np.zeros( + (len(coords["date"]), len(coords["mc_name"][1]), len(coords["subpop"])) + ) states = xr.Dataset( data_vars=dict( prevalence=(["date", "compartment", "subpop"], zeros), @@ -261,7 +279,9 @@ def autodetect_scenarios(config): outcome_modifiers_scenarios = None if config["outcomes"].exists() and config["outcome_modifiers"].exists(): if config["outcome_modifiers"]["scenarios"].exists(): - outcome_modifiers_scenarios = config["outcome_modifiers"]["scenarios"].as_str_seq() + outcome_modifiers_scenarios = config["outcome_modifiers"][ + "scenarios" + ].as_str_seq() outcome_modifiers_scenarios = as_list(outcome_modifiers_scenarios) seir_modifiers_scenarios = as_list(seir_modifiers_scenarios) @@ -275,41 +295,41 @@ def autodetect_scenarios(config): return seir_modifiers_scenarios[0], outcome_modifiers_scenarios[0] + # rewrite the get log loss functions as single functions, not in a class. This is not faster # def get_logloss(proposal, inferpar, logloss, static_sim_arguments, modinf, silent=True, save=False): # if not inferpar.check_in_bound(proposal=proposal): # if not silent: # print("OUT OF BOUND!!") # return -np.inf, -np.inf, -np.inf -# +# # snpi_df_mod, hnpi_df_mod = inferpar.inject_proposal( # proposal=proposal, # snpi_df=static_sim_arguments["snpi_df_ref"], # hnpi_df=static_sim_arguments["hnpi_df_ref"], # ) -# +# # ss = copy.deepcopy(static_sim_arguments) # ss["snpi_df_in"] = snpi_df_mod # ss["hnpi_df_in"] = hnpi_df_mod # del ss["snpi_df_ref"] # del ss["hnpi_df_ref"] -# +# # outcomes_df = simulation_atomic(**ss, modinf=modinf, save=save) -# +# # ll_total, logloss, regularizations = logloss.compute_logloss( # model_df=outcomes_df, subpop_names=modinf.subpop_struct.subpop_names # ) # if not silent: # print(f"llik is {ll_total}") -# +# # return ll_total, logloss, regularizations -# +# # def get_logloss_as_single_number(proposal, inferpar, logloss, static_sim_arguments, modinf, silent=True, save=False): # ll_total, logloss, regularizations = get_logloss(proposal, inferpar, logloss, static_sim_arguments, modinf, silent, save) # return ll_total - class GempyorInference: def __init__( self, @@ -333,12 +353,20 @@ def __init__( config.set_file(os.path.join(path_prefix, config_filepath)) - self.seir_modifiers_scenario, self.outcome_modifiers_scenario = autodetect_scenarios(config) + self.seir_modifiers_scenario, self.outcome_modifiers_scenario = ( + autodetect_scenarios(config) + ) if run_id is None: run_id = file_paths.run_id() if prefix is None: - prefix = config["name"].get() + f"_{self.seir_modifiers_scenario}_{self.outcome_modifiers_scenario}" + "/" + run_id + "/" + prefix = ( + config["name"].get() + + f"_{self.seir_modifiers_scenario}_{self.outcome_modifiers_scenario}" + + "/" + + run_id + + "/" + ) in_run_id = run_id if out_run_id is None: out_run_id = in_run_id @@ -387,7 +415,8 @@ def __init__( self.do_inference = True self.inference_method = "emcee" self.inferpar = inference_parameter.InferenceParameters( - global_config=config, subpop_names=self.modinf.subpop_struct.subpop_names + global_config=config, + subpop_names=self.modinf.subpop_struct.subpop_names, ) self.logloss = logloss.LogLoss( inference_config=config["inference"], @@ -412,7 +441,14 @@ def set_save(self, save): def get_all_sim_arguments(self): # inferpar, logloss, static_sim_arguments, modinf, proposal, silent, save - return [self.inferpar, self.logloss, self.static_sim_arguments, self.modinf, self.silent, self.save] + return [ + self.inferpar, + self.logloss, + self.static_sim_arguments, + self.modinf, + self.silent, + self.save, + ] def get_logloss(self, proposal): if not self.inferpar.check_in_bound(proposal=proposal): @@ -479,11 +515,15 @@ def update_run_id(self, new_run_id, new_out_run_id=None): else: self.modinf.out_run_id = new_out_run_id - def one_simulation_legacy(self, sim_id2write: int, load_ID: bool = False, sim_id2load: int = None): + def one_simulation_legacy( + self, sim_id2write: int, load_ID: bool = False, sim_id2load: int = None + ): sim_id2write = int(sim_id2write) if load_ID: sim_id2load = int(sim_id2load) - with Timer(f">>> GEMPYOR onesim {'(loading file)' if load_ID else '(from config)'}"): + with Timer( + f">>> GEMPYOR onesim {'(loading file)' if load_ID else '(from config)'}" + ): with Timer("onerun_SEIR"): seir.onerun_SEIR( sim_id2write=sim_id2write, @@ -533,7 +573,9 @@ def one_simulation( sim_id2load = int(sim_id2load) self.lastsim_sim_id2load = sim_id2load - with Timer(f">>> GEMPYOR onesim {'(loading file)' if load_ID else '(from config)'}"): + with Timer( + f">>> GEMPYOR onesim {'(loading file)' if load_ID else '(from config)'}" + ): if not self.already_built and self.modinf.outcomes_config is not None: self.outcomes_parameters = outcomes.read_parameters_from_config(self.modinf) @@ -541,10 +583,24 @@ def one_simulation( npi_seir = None if parallel: with Timer("//things"): - with ProcessPoolExecutor(max_workers=max(mp.cpu_count(), 3)) as executor: - if self.modinf.seir_config is not None and self.modinf.npi_config_seir is not None: - ret_seir = executor.submit(seir.build_npi_SEIR, self.modinf, load_ID, sim_id2load, config) - if self.modinf.outcomes_config is not None and self.modinf.npi_config_outcomes: + with ProcessPoolExecutor( + max_workers=max(mp.cpu_count(), 3) + ) as executor: + if ( + self.modinf.seir_config is not None + and self.modinf.npi_config_seir is not None + ): + ret_seir = executor.submit( + seir.build_npi_SEIR, + self.modinf, + load_ID, + sim_id2load, + config, + ) + if ( + self.modinf.outcomes_config is not None + and self.modinf.npi_config_outcomes + ): ret_outcomes = executor.submit( outcomes.build_outcome_modifiers, self.modinf, @@ -553,7 +609,9 @@ def one_simulation( config, ) if not self.already_built: - ret_comparments = executor.submit(self.modinf.compartments.get_transition_array) + ret_comparments = executor.submit( + self.modinf.compartments.get_transition_array + ) # print("expections:", ret_seir.exception(), ret_outcomes.exception(), ret_comparments.exception()) @@ -565,18 +623,33 @@ def one_simulation( self.proportion_info, ) = ret_comparments.result() self.already_built = True - if self.modinf.seir_config is not None and self.modinf.npi_config_seir is not None: + if ( + self.modinf.seir_config is not None + and self.modinf.npi_config_seir is not None + ): npi_seir = ret_seir.result() - if self.modinf.outcomes_config is not None and self.modinf.npi_config_outcomes: + if ( + self.modinf.outcomes_config is not None + and self.modinf.npi_config_outcomes + ): npi_outcomes = ret_outcomes.result() else: if not self.already_built: self.build_structure() - if self.modinf.seir_config is not None and self.modinf.npi_config_seir is not None: + if ( + self.modinf.seir_config is not None + and self.modinf.npi_config_seir is not None + ): npi_seir = seir.build_npi_SEIR( - modinf=self.modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config + modinf=self.modinf, + load_ID=load_ID, + sim_id2load=sim_id2load, + config=config, ) - if self.modinf.outcomes_config is not None and self.modinf.npi_config_outcomes: + if ( + self.modinf.outcomes_config is not None + and self.modinf.npi_config_outcomes + ): npi_outcomes = outcomes.build_outcome_modifiers( modinf=self.modinf, load_ID=load_ID, @@ -602,8 +675,12 @@ def one_simulation( with Timer("onerun_SEIR.seeding"): if load_ID: - initial_conditions = self.modinf.initial_conditions.get_from_file(sim_id2load, modinf=self.modinf) - seeding_data, seeding_amounts = self.modinf.seeding.get_from_file(sim_id2load, modinf=self.modinf) + initial_conditions = self.modinf.initial_conditions.get_from_file( + sim_id2load, modinf=self.modinf + ) + seeding_data, seeding_amounts = self.modinf.seeding.get_from_file( + sim_id2load, modinf=self.modinf + ) else: initial_conditions = self.modinf.initial_conditions.get_from_config( sim_id2write, modinf=self.modinf @@ -649,7 +726,7 @@ def one_simulation( parameters=self.outcomes_parameters, loaded_values=loaded_values, npi=npi_outcomes, - bypass_seir_xr=states + bypass_seir_xr=states, ) self.lastsim_outcomes_df = outcomes_df self.lastsim_hpar_df = hpar_df @@ -664,14 +741,18 @@ def one_simulation( ) return 0 - def plot_transition_graph(self, output_file="transition_graph", source_filters=[], destination_filters=[]): + def plot_transition_graph( + self, output_file="transition_graph", source_filters=[], destination_filters=[] + ): self.modinf.compartments.plot( output_file=output_file, source_filters=source_filters, destination_filters=destination_filters, ) - def get_outcome_npi(self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None): + def get_outcome_npi( + self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None + ): npi_outcomes = None if self.modinf.npi_config_outcomes: npi_outcomes = outcomes.build_outcome_modifiers( @@ -695,7 +776,9 @@ def get_seir_npi(self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_F ) return npi_seir - def get_seir_parameters(self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None): + def get_seir_parameters( + self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None + ): param_df = None if bypass_DF is not None: param_df = bypass_DF @@ -716,7 +799,9 @@ def get_seir_parameters(self, load_ID=False, sim_id2load=None, bypass_DF=None, b ) return p_draw - def get_seir_parametersDF(self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None): + def get_seir_parametersDF( + self, load_ID=False, sim_id2load=None, bypass_DF=None, bypass_FN=None + ): p_draw = self.get_seir_parameters( load_ID=load_ID, sim_id2load=sim_id2load, @@ -771,7 +856,9 @@ def get_parsed_parameters_seir( if not self.already_built: self.build_structure() - npi_seir = seir.build_npi_SEIR(modinf=self.modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config) + npi_seir = seir.build_npi_SEIR( + modinf=self.modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config + ) p_draw = self.get_seir_parameters(load_ID=load_ID, sim_id2load=sim_id2load) parameters = self.modinf.parameters.parameters_reduce(p_draw, npi_seir) @@ -788,7 +875,9 @@ def get_reduced_parameters_seir( # bypass_DF=None, # bypass_FN=None, ): - npi_seir = seir.build_npi_SEIR(modinf=self.modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config) + npi_seir = seir.build_npi_SEIR( + modinf=self.modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config + ) p_draw = self.get_seir_parameters(load_ID=load_ID, sim_id2load=sim_id2load) parameters = self.modinf.parameters.parameters_reduce(p_draw, npi_seir) @@ -809,7 +898,9 @@ def paramred_parallel(run_spec, snpi_fn): seir_modifiers_scenario="inference", # NPIs scenario to use outcome_modifiers_scenario="med", # Outcome scenario to use stoch_traj_flag=False, - path_prefix=run_spec["geodata"], # prefix where to find the folder indicated in subpop_setup$ + path_prefix=run_spec[ + "geodata" + ], # prefix where to find the folder indicated in subpop_setup$ ) snpi = pq.read_table(snpi_fn).to_pandas() @@ -820,7 +911,9 @@ def paramred_parallel(run_spec, snpi_fn): params_draw_arr = gempyor_inference.get_seir_parameters( bypass_FN=snpi_fn.replace("snpi", "spar") ) # could also accept (load_ID=True, sim_id2load=XXX) or (bypass_DF=) or (bypass_FN=) - param_reduc_from = gempyor_inference.get_seir_parameter_reduced(npi_seir=npi_seir, p_draw=params_draw_arr) + param_reduc_from = gempyor_inference.get_seir_parameter_reduced( + npi_seir=npi_seir, p_draw=params_draw_arr + ) return param_reduc_from @@ -835,7 +928,9 @@ def paramred_parallel_config(run_spec, dummy): seir_modifiers_scenario="inference", # NPIs scenario to use outcome_modifiers_scenario="med", # Outcome scenario to use stoch_traj_flag=False, - path_prefix=run_spec["geodata"], # prefix where to find the folder indicated in subpop_setup$ + path_prefix=run_spec[ + "geodata" + ], # prefix where to find the folder indicated in subpop_setup$ ) npi_seir = gempyor_inference.get_seir_npi() @@ -843,6 +938,8 @@ def paramred_parallel_config(run_spec, dummy): params_draw_arr = ( gempyor_inference.get_seir_parameters() ) # could also accept (load_ID=True, sim_id2load=XXX) or (bypass_DF=) or (bypass_FN=) - param_reduc_from = gempyor_inference.get_seir_parameter_reduced(npi_seir=npi_seir, p_draw=params_draw_arr) + param_reduc_from = gempyor_inference.get_seir_parameter_reduced( + npi_seir=npi_seir, p_draw=params_draw_arr + ) return param_reduc_from diff --git a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py index 8f1b0fc53..e6a795192 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py @@ -36,9 +36,14 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): """ # identify spatial group affected_subpops = set(subpops) - if parameter_config["subpop"].exists() and parameter_config["subpop"].get() != "all": + if ( + parameter_config["subpop"].exists() + and parameter_config["subpop"].get() != "all" + ): affected_subpops = {str(n.get()) for n in parameter_config["subpop"]} - spatial_groups = NPI.helpers.get_spatial_groups(parameter_config, list(affected_subpops)) + spatial_groups = NPI.helpers.get_spatial_groups( + parameter_config, list(affected_subpops) + ) # ungrouped subpop (all affected subpop by default) have one parameter per subpop if spatial_groups["ungrouped"]: @@ -87,7 +92,9 @@ def build_from_config(self, global_config, subpop_names): for config_part in ["seir_modifiers", "outcome_modifiers"]: if global_config[config_part].exists(): for npi in global_config[config_part]["modifiers"].get(): - if global_config[config_part]["modifiers"][npi]["perturbation"].exists(): + if global_config[config_part]["modifiers"][npi][ + "perturbation" + ].exists(): self.add_modifier( pname=npi, ptype=config_part, diff --git a/flepimop/gempyor_pkg/src/gempyor/initial_conditions.py b/flepimop/gempyor_pkg/src/gempyor/initial_conditions.py index eff95b6b8..b86226cca 100644 --- a/flepimop/gempyor_pkg/src/gempyor/initial_conditions.py +++ b/flepimop/gempyor_pkg/src/gempyor/initial_conditions.py @@ -46,19 +46,30 @@ def __init__( if self.initial_conditions_config is not None: if "ignore_population_checks" in self.initial_conditions_config.keys(): - self.ignore_population_checks = self.initial_conditions_config["ignore_population_checks"].get(bool) + self.ignore_population_checks = self.initial_conditions_config[ + "ignore_population_checks" + ].get(bool) if "allow_missing_subpops" in self.initial_conditions_config.keys(): - self.allow_missing_subpops = self.initial_conditions_config["allow_missing_subpops"].get(bool) + self.allow_missing_subpops = self.initial_conditions_config[ + "allow_missing_subpops" + ].get(bool) if "allow_missing_compartments" in self.initial_conditions_config.keys(): - self.allow_missing_compartments = self.initial_conditions_config["allow_missing_compartments"].get(bool) + self.allow_missing_compartments = self.initial_conditions_config[ + "allow_missing_compartments" + ].get(bool) # TODO: add check, this option onlywork with tidy dataframe if "proportional" in self.initial_conditions_config.keys(): - self.proportional_ic = self.initial_conditions_config["proportional"].get(bool) + self.proportional_ic = self.initial_conditions_config["proportional"].get( + bool + ) def get_from_config(self, sim_id: int, modinf) -> np.ndarray: method = "Default" - if self.initial_conditions_config is not None and "method" in self.initial_conditions_config.keys(): + if ( + self.initial_conditions_config is not None + and "method" in self.initial_conditions_config.keys() + ): method = self.initial_conditions_config["method"].as_str() if method == "Default": @@ -70,10 +81,13 @@ def get_from_config(self, sim_id: int, modinf) -> np.ndarray: if method == "SetInitialConditions" or method == "SetInitialConditionsFolderDraw": # TODO Think about - Does not support the new way of doing compartment indexing if method == "SetInitialConditionsFolderDraw": - ic_df = modinf.read_simID(ftype=self.initial_conditions_config["initial_file_type"], sim_id=sim_id) + ic_df = modinf.read_simID( + ftype=self.initial_conditions_config["initial_file_type"], sim_id=sim_id + ) else: ic_df = read_df( - self.path_prefix / self.initial_conditions_config["initial_conditions_file"].get(), + self.path_prefix + / self.initial_conditions_config["initial_conditions_file"].get(), ) y0 = read_initial_condition_from_tidydataframe( ic_df=ic_df, @@ -86,11 +100,13 @@ def get_from_config(self, sim_id: int, modinf) -> np.ndarray: elif method == "InitialConditionsFolderDraw" or method == "FromFile": if method == "InitialConditionsFolderDraw": ic_df = modinf.read_simID( - ftype=self.initial_conditions_config["initial_file_type"].get(), sim_id=sim_id + ftype=self.initial_conditions_config["initial_file_type"].get(), + sim_id=sim_id, ) elif method == "FromFile": ic_df = read_df( - self.path_prefix / self.initial_conditions_config["initial_conditions_file"].get(), + self.path_prefix + / self.initial_conditions_config["initial_conditions_file"].get(), ) y0 = read_initial_condition_from_seir_output( @@ -103,7 +119,9 @@ def get_from_config(self, sim_id: int, modinf) -> np.ndarray: raise NotImplementedError(f"Unknown initial conditions method [received: '{method}'].") # check that the inputed values sums to the subpop population: - check_population(y0=y0, modinf=modinf, ignore_population_checks=self.ignore_population_checks) + check_population( + y0=y0, modinf=modinf, ignore_population_checks=self.ignore_population_checks + ) return y0 @@ -149,7 +167,9 @@ def read_initial_condition_from_tidydataframe( states_pl = ic_df[ic_df["subpop"] == pl] for comp_idx, comp_name in modinf.compartments.compartments["name"].items(): if "mc_name" in states_pl.columns: - ic_df_compartment_val = states_pl[states_pl["mc_name"] == comp_name]["amount"] + ic_df_compartment_val = states_pl[states_pl["mc_name"] == comp_name][ + "amount" + ] else: filters = modinf.compartments.compartments.iloc[comp_idx].drop("name") ic_df_compartment_val = states_pl.copy() @@ -207,7 +227,9 @@ def read_initial_condition_from_tidydataframe( return y0 -def read_initial_condition_from_seir_output(ic_df, modinf, allow_missing_subpops, allow_missing_compartments): +def read_initial_condition_from_seir_output( + ic_df, modinf, allow_missing_subpops, allow_missing_compartments +): """ Read the initial conditions from the SEIR output. @@ -232,10 +254,13 @@ def read_initial_condition_from_seir_output(ic_df, modinf, allow_missing_subpops ic_df["date"] = ic_df["date"].dt.date ic_df["date"] = ic_df["date"].astype(str) - ic_df = ic_df[(ic_df["date"] == str(modinf.ti)) & (ic_df["mc_value_type"] == "prevalence")] + ic_df = ic_df[ + (ic_df["date"] == str(modinf.ti)) & (ic_df["mc_value_type"] == "prevalence") + ] if ic_df.empty: raise ValueError(f"No entry provided for initial time `ti` in the `initial_conditions::states_file.` " - f"`ti`: '{modinf.ti}'.") + f"`ti`: '{modinf.ti}'." + ) y0 = np.zeros((modinf.compartments.compartments.shape[0], modinf.nsubpops)) for comp_idx, comp_name in modinf.compartments.compartments["name"].items(): @@ -245,7 +270,9 @@ def read_initial_condition_from_seir_output(ic_df, modinf, allow_missing_subpops filters = modinf.compartments.compartments.iloc[comp_idx].drop("name") ic_df_compartment = ic_df.copy() for mc_name, mc_value in filters.items(): - ic_df_compartment = ic_df_compartment[ic_df_compartment["mc_" + mc_name] == mc_value] + ic_df_compartment = ic_df_compartment[ + ic_df_compartment["mc_" + mc_name] == mc_value + ] if len(ic_df_compartment) > 1: # ic_df_compartment = ic_df_compartment.iloc[0] @@ -256,7 +283,9 @@ def read_initial_condition_from_seir_output(ic_df, modinf, allow_missing_subpops ) elif ic_df_compartment.empty: if allow_missing_compartments: - ic_df_compartment = pd.DataFrame(0, columns=ic_df_compartment.columns, index=[0]) + ic_df_compartment = pd.DataFrame( + 0, columns=ic_df_compartment.columns, index=[0] + ) else: raise ValueError( f"Initial Conditions: could not set compartment '{comp_name}' (id: '{comp_idx}') in subpop '{pl}' (id: '{pl_idx}'). " diff --git a/flepimop/gempyor_pkg/src/gempyor/logloss.py b/flepimop/gempyor_pkg/src/gempyor/logloss.py index 95892dac0..09993421f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/logloss.py +++ b/flepimop/gempyor_pkg/src/gempyor/logloss.py @@ -17,7 +17,13 @@ class LogLoss: - def __init__(self, inference_config: confuse.ConfigView, subpop_struct, time_setup, path_prefix: str = "."): + def __init__( + self, + inference_config: confuse.ConfigView, + subpop_struct, + time_setup, + path_prefix: str = ".", + ): # TODO: bad format for gt because each date must have a value for each column, but if it doesn't and you add NA # then this NA has a meaning that depends on skip NA, which is annoying. # A lot of things can go wrong here, in the previous approach where GT was cast to xarray as @@ -35,20 +41,30 @@ def __init__(self, inference_config: confuse.ConfigView, subpop_struct, time_set # made the controversial choice of storing the gt as an xarray dataset instead of a dictionary # of dataframes - self.gt_xr = xr.Dataset.from_dataframe(self.gt.reset_index().set_index(["date", "subpop"])) + self.gt_xr = xr.Dataset.from_dataframe( + self.gt.reset_index().set_index(["date", "subpop"]) + ) # Very important: subsample the subpop in the population, in the right order, and sort by the date index. - self.gt_xr = self.gt_xr.sortby("date").reindex({"subpop": subpop_struct.subpop_names}) + self.gt_xr = self.gt_xr.sortby("date").reindex( + {"subpop": subpop_struct.subpop_names} + ) # This will force at 0, if skipna is False, data of some variable that don't exist if iother exist # and damn python datetime types are ugly... - self.first_date = max(pd.to_datetime(self.gt_xr.date[0].values).date(), time_setup.ti) - self.last_date = min(pd.to_datetime(self.gt_xr.date[-1].values).date(), time_setup.tf) + self.first_date = max( + pd.to_datetime(self.gt_xr.date[0].values).date(), time_setup.ti + ) + self.last_date = min( + pd.to_datetime(self.gt_xr.date[-1].values).date(), time_setup.tf + ) self.statistics = {} for key, value in inference_config["statistics"].items(): self.statistics[key] = statistics.Statistic(key, value) - def plot_gt(self, ax=None, subpop=None, statistic=None, subplot=False, filename=None, **kwargs): + def plot_gt( + self, ax=None, subpop=None, statistic=None, subplot=False, filename=None, **kwargs + ): """Plots ground truth data. Args: @@ -68,7 +84,10 @@ def plot_gt(self, ax=None, subpop=None, statistic=None, subplot=False, filename= fig, axes = plt.subplots( len(self.gt["subpop"].unique()), len(self.gt.columns.drop("subpop")), - figsize=(4 * len(self.gt.columns.drop("subpop")), 3 * len(self.gt["subpop"].unique())), + figsize=( + 4 * len(self.gt.columns.drop("subpop")), + 3 * len(self.gt["subpop"].unique()), + ), dpi=250, sharex=True, ) @@ -81,7 +100,9 @@ def plot_gt(self, ax=None, subpop=None, statistic=None, subplot=False, filename= subpops = [subpop] if statistic is None: - statistics = self.gt.columns.drop("subpop") # Assuming other columns are statistics + statistics = self.gt.columns.drop( + "subpop" + ) # Assuming other columns are statistics else: statistics = [statistic] @@ -107,7 +128,10 @@ def plot_gt(self, ax=None, subpop=None, statistic=None, subplot=False, filename= plt.savefig(filename, **kwargs) # Save the figure if subplot: - return fig, axes # Return figure and subplots for potential further customization + return ( + fig, + axes, + ) # Return figure and subplots for potential further customization else: return ax # Optionally return the axis @@ -121,7 +145,9 @@ def compute_logloss(self, model_df, subpop_names): coords = {"statistic": list(self.statistics.keys()), "subpop": subpop_names} logloss = xr.DataArray( - np.zeros((len(coords["statistic"]), len(coords["subpop"]))), dims=["statistic", "subpop"], coords=coords + np.zeros((len(coords["statistic"]), len(coords["subpop"]))), + dims=["statistic", "subpop"], + coords=coords, ) regularizations = 0 diff --git a/flepimop/gempyor_pkg/src/gempyor/model_info.py b/flepimop/gempyor_pkg/src/gempyor/model_info.py index 54f981d8f..35502aadf 100644 --- a/flepimop/gempyor_pkg/src/gempyor/model_info.py +++ b/flepimop/gempyor_pkg/src/gempyor/model_info.py @@ -1,6 +1,13 @@ import pandas as pd import datetime, os, logging, pathlib, confuse -from . import seeding, subpopulation_structure, parameters, compartments, file_paths, initial_conditions +from . import ( + seeding, + subpopulation_structure, + parameters, + compartments, + file_paths, + initial_conditions, +) from .utils import read_df, write_df logger = logging.getLogger(__name__) @@ -11,7 +18,9 @@ def __init__(self, config: confuse.ConfigView): self.ti = config["start_date"].as_date() self.tf = config["end_date"].as_date() if self.tf <= self.ti: - raise ValueError("tf (time to finish) is less than or equal to ti (time to start)") + raise ValueError( + "tf (time to finish) is less than or equal to ti (time to start)" + ) self.n_days = (self.tf - self.ti).days + 1 self.dates = pd.date_range(start=self.ti, end=self.tf, freq="D") @@ -29,7 +38,7 @@ class ModelInfo: seeding # One of seeding or initial_conditions is required when running seir outcomes # Required if running outcomes seir_modifiers # Not required. If exists, every modifier will be applied to seir parameters - outcomes_modifiers # Not required. If exists, every modifier will be applied to outcomes + outcomes_modifiers # Not required. If exists, every modifier will be applied to outcomes inference # Required if running inference ``` """ @@ -94,7 +103,9 @@ def __init__( # 3. What about subpopulations subpop_config = config["subpop_setup"] if "data_path" in config: - raise ValueError("The config has a data_path section. This is no longer supported.") + raise ValueError( + "The config has a data_path section. This is no longer supported." + ) self.path_prefix = pathlib.Path(path_prefix) self.subpop_struct = subpopulation_structure.SubpopulationStructure( @@ -112,7 +123,9 @@ def __init__( self.seir_config = config["seir"] self.parameters_config = config["seir"]["parameters"] self.initial_conditions_config = ( - config["initial_conditions"] if config["initial_conditions"].exists() else None + config["initial_conditions"] + if config["initial_conditions"].exists() + else None ) self.seeding_config = config["seeding"] if config["seeding"].exists() else None @@ -130,7 +143,9 @@ def __init__( subpop_names=self.subpop_struct.subpop_names, path_prefix=self.path_prefix, ) - self.seeding = seeding.SeedingFactory(config=self.seeding_config, path_prefix=self.path_prefix) + self.seeding = seeding.SeedingFactory( + config=self.seeding_config, path_prefix=self.path_prefix + ) self.initial_conditions = initial_conditions.InitialConditionsFactory( config=self.initial_conditions_config, path_prefix=self.path_prefix ) @@ -144,11 +159,19 @@ def __init__( self.npi_config_seir = None if config["seir_modifiers"].exists(): if config["seir_modifiers"]["scenarios"].exists(): - self.npi_config_seir = config["seir_modifiers"]["modifiers"][seir_modifiers_scenario] - self.seir_modifiers_library = config["seir_modifiers"]["modifiers"].get() + self.npi_config_seir = config["seir_modifiers"]["modifiers"][ + seir_modifiers_scenario + ] + self.seir_modifiers_library = config["seir_modifiers"][ + "modifiers" + ].get() else: - self.seir_modifiers_library = config["seir_modifiers"]["modifiers"].get() - raise ValueError("Not implemented yet") # TODO create a Stacked from all + self.seir_modifiers_library = config["seir_modifiers"][ + "modifiers" + ].get() + raise ValueError( + "Not implemented yet" + ) # TODO create a Stacked from all elif self.seir_modifiers_scenario is not None: raise ValueError( "An seir modifiers scenario was provided to ModelInfo but no 'seir_modifiers' sections in config" @@ -157,7 +180,9 @@ def __init__( logging.info("Running ModelInfo with seir but without SEIR Modifiers") elif self.seir_modifiers_scenario is not None: - raise ValueError("A seir modifiers scenario was provided to ModelInfo but no 'seir:' sections in config") + raise ValueError( + "A seir modifiers scenario was provided to ModelInfo but no 'seir:' sections in config" + ) else: logging.critical("Running ModelInfo without SEIR") @@ -167,11 +192,19 @@ def __init__( self.npi_config_outcomes = None if config["outcome_modifiers"].exists(): if config["outcome_modifiers"]["scenarios"].exists(): - self.npi_config_outcomes = config["outcome_modifiers"]["modifiers"][self.outcome_modifiers_scenario] - self.outcome_modifiers_library = config["outcome_modifiers"]["modifiers"].get() + self.npi_config_outcomes = config["outcome_modifiers"]["modifiers"][ + self.outcome_modifiers_scenario + ] + self.outcome_modifiers_library = config["outcome_modifiers"][ + "modifiers" + ].get() else: - self.outcome_modifiers_library = config["outcome_modifiers"]["modifiers"].get() - raise ValueError("Not implemented yet") # TODO create a Stacked from all + self.outcome_modifiers_library = config["outcome_modifiers"][ + "modifiers" + ].get() + raise ValueError( + "Not implemented yet" + ) # TODO create a Stacked from all ## NEED TO IMPLEMENT THIS -- CURRENTLY CANNOT USE outcome modifiers elif self.outcome_modifiers_scenario is not None: @@ -182,7 +215,9 @@ def __init__( else: self.outcome_modifiers_scenario = None else: - logging.info("Running ModelInfo with outcomes but without Outcomes Modifiers") + logging.info( + "Running ModelInfo with outcomes but without Outcomes Modifiers" + ) elif self.outcome_modifiers_scenario is not None: raise ValueError( "An outcome modifiers scenario was provided to ModelInfo but no 'outcomes:' sections in config" @@ -228,7 +263,9 @@ def __init__( os.makedirs(datadir, exist_ok=True) if self.write_parquet and self.write_csv: - print("Confused between reading .csv or parquet. Assuming input file is .parquet") + print( + "Confused between reading .csv or parquet. Assuming input file is .parquet" + ) if self.write_parquet: self.extension = "parquet" elif self.write_csv: @@ -252,7 +289,9 @@ def get_output_filename(self, ftype: str, sim_id: int, extension_override: str = extension_override=extension_override, ) - def get_filename(self, ftype: str, sim_id: int, input: bool, extension_override: str = ""): + def get_filename( + self, ftype: str, sim_id: int, input: bool, extension_override: str = "" + ): """return a CSP formated filename.""" if extension_override: # empty strings are Falsy @@ -281,7 +320,9 @@ def get_filename(self, ftype: str, sim_id: int, input: bool, extension_override: def get_setup_name(self): return self.setup_name - def read_simID(self, ftype: str, sim_id: int, input: bool = True, extension_override: str = ""): + def read_simID( + self, ftype: str, sim_id: int, input: bool = True, extension_override: str = "" + ): fname = self.get_filename( ftype=ftype, sim_id=sim_id, diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py index 5563f4d85..e5f09c1d2 100644 --- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py +++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py @@ -22,7 +22,9 @@ def run_parallel_outcomes(modinf, *, sim_id2write, nslots=1, n_jobs=1): sim_id2writes = np.arange(sim_id2write, sim_id2write + modinf.nslots) loaded_values = None - if (n_jobs == 1) or (modinf.nslots == 1): # run single process for debugging/profiling purposes + if (n_jobs == 1) or ( + modinf.nslots == 1 + ): # run single process for debugging/profiling purposes for sim_offset in np.arange(nslots): onerun_delayframe_outcomes( sim_id2write=sim_id2writes[sim_offset], @@ -100,7 +102,9 @@ def onerun_delayframe_outcomes( npi_outcomes = None if modinf.npi_config_outcomes: - npi_outcomes = build_outcome_modifiers(modinf=modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config) + npi_outcomes = build_outcome_modifiers( + modinf=modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config + ) loaded_values = None if load_ID: @@ -117,7 +121,13 @@ def onerun_delayframe_outcomes( ) with Timer("onerun_delayframe_outcomes.postprocess"): - postprocess_and_write(sim_id=sim_id2write, modinf=modinf, outcomes_df=outcomes_df, hpar=hpar, npi=npi_outcomes) + postprocess_and_write( + sim_id=sim_id2write, + modinf=modinf, + outcomes_df=outcomes_df, + hpar=hpar, + npi=npi_outcomes, + ) def read_parameters_from_config(modinf: model_info.ModelInfo): @@ -129,7 +139,10 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): if modinf.outcomes_config["param_from_file"].exists(): if modinf.outcomes_config["param_from_file"].get(): # Load the actual csv file - branching_file = modinf.path_prefix / modinf.outcomes_config["param_subpop_file"].as_str() + branching_file = ( + modinf.path_prefix + / modinf.outcomes_config["param_subpop_file"].as_str() + ) branching_data = pa.parquet.read_table(branching_file).to_pandas() if "relative_probability" not in list(branching_data["quantity"]): raise ValueError( @@ -142,14 +155,18 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): "", end="", ) - branching_data = branching_data[branching_data["subpop"].isin(modinf.subpop_struct.subpop_names)] + branching_data = branching_data[ + branching_data["subpop"].isin(modinf.subpop_struct.subpop_names) + ] print( "Intersect with seir simulation: ", len(branching_data.subpop.unique()), "kept", ) - if len(branching_data.subpop.unique()) != len(modinf.subpop_struct.subpop_names): + if len(branching_data.subpop.unique()) != len( + modinf.subpop_struct.subpop_names + ): raise ValueError( f"Places in seir input files does not correspond to subpops in outcome probability file {branching_file}" ) @@ -170,10 +187,14 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): f"unsure how to read outcome {new_comp}: not a str, nor an incidence or prevalence: {src_name}" ) - parameters[new_comp]["probability"] = outcomes_config[new_comp]["probability"]["value"] + parameters[new_comp]["probability"] = outcomes_config[new_comp][ + "probability" + ]["value"] if outcomes_config[new_comp]["probability"]["modifier_parameter"].exists(): parameters[new_comp]["probability::npi_param_name"] = ( - outcomes_config[new_comp]["probability"]["modifier_parameter"].as_str().lower() + outcomes_config[new_comp]["probability"]["modifier_parameter"] + .as_str() + .lower() ) logging.debug( f"probability of outcome {new_comp} is affected by intervention " @@ -181,13 +202,19 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): f"instead of {new_comp}::probability" ) else: - parameters[new_comp]["probability::npi_param_name"] = f"{new_comp}::probability".lower() + parameters[new_comp][ + "probability::npi_param_name" + ] = f"{new_comp}::probability".lower() if outcomes_config[new_comp]["delay"].exists(): - parameters[new_comp]["delay"] = outcomes_config[new_comp]["delay"]["value"] + parameters[new_comp]["delay"] = outcomes_config[new_comp]["delay"][ + "value" + ] if outcomes_config[new_comp]["delay"]["modifier_parameter"].exists(): parameters[new_comp]["delay::npi_param_name"] = ( - outcomes_config[new_comp]["delay"]["modifier_parameter"].as_str().lower() + outcomes_config[new_comp]["delay"]["modifier_parameter"] + .as_str() + .lower() ) logging.debug( f"delay of outcome {new_comp} is affected by intervention " @@ -195,18 +222,28 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): f"instead of {new_comp}::delay" ) else: - parameters[new_comp]["delay::npi_param_name"] = f"{new_comp}::delay".lower() + parameters[new_comp][ + "delay::npi_param_name" + ] = f"{new_comp}::delay".lower() else: logging.critical(f"No delay for outcome {new_comp}, using a 0 delay") outcomes_config[new_comp]["delay"] = {"value": 0} - parameters[new_comp]["delay"] = outcomes_config[new_comp]["delay"]["value"] - parameters[new_comp]["delay::npi_param_name"] = f"{new_comp}::delay".lower() + parameters[new_comp]["delay"] = outcomes_config[new_comp]["delay"][ + "value" + ] + parameters[new_comp][ + "delay::npi_param_name" + ] = f"{new_comp}::delay".lower() if outcomes_config[new_comp]["duration"].exists(): - parameters[new_comp]["duration"] = outcomes_config[new_comp]["duration"]["value"] + parameters[new_comp]["duration"] = outcomes_config[new_comp][ + "duration" + ]["value"] if outcomes_config[new_comp]["duration"]["modifier_parameter"].exists(): parameters[new_comp]["duration::npi_param_name"] = ( - outcomes_config[new_comp]["duration"]["modifier_parameter"].as_str().lower() + outcomes_config[new_comp]["duration"]["modifier_parameter"] + .as_str() + .lower() ) logging.debug( f"duration of outcome {new_comp} is affected by intervention " @@ -214,7 +251,9 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): f"instead of {new_comp}::duration" ) else: - parameters[new_comp]["duration::npi_param_name"] = f"{new_comp}::duration".lower() + parameters[new_comp][ + "duration::npi_param_name" + ] = f"{new_comp}::duration".lower() if outcomes_config[new_comp]["duration"]["name"].exists(): parameters[new_comp]["outcome_prevalence_name"] = ( @@ -231,14 +270,22 @@ def read_parameters_from_config(modinf: model_info.ModelInfo): & (branching_data["quantity"] == "relative_probability") ].copy(deep=True) if len(rel_probability) > 0: - logging.debug(f"Using 'param_from_file' for relative probability in outcome {new_comp}") + logging.debug( + f"Using 'param_from_file' for relative probability in outcome {new_comp}" + ) # Sort it in case the relative probablity file is mispecified - rel_probability.subpop = rel_probability.subpop.astype("category") - rel_probability.subpop = rel_probability.subpop.cat.set_categories( - modinf.subpop_struct.subpop_names + rel_probability.subpop = rel_probability.subpop.astype( + "category" + ) + rel_probability.subpop = ( + rel_probability.subpop.cat.set_categories( + modinf.subpop_struct.subpop_names + ) ) rel_probability = rel_probability.sort_values(["subpop"]) - parameters[new_comp]["rel_probability"] = rel_probability["value"].to_numpy() + parameters[new_comp]["rel_probability"] = rel_probability[ + "value" + ].to_numpy() else: logging.debug( f"*NOT* Using 'param_from_file' for relative probability in outcome {new_comp}" @@ -348,7 +395,9 @@ def compute_all_multioutcomes( outcome_name=new_comp, ) else: - raise ValueError(f"Unknown type for seir simulation provided, got f{type(seir_sim)}") + raise ValueError( + f"Unknown type for seir simulation provided, got f{type(seir_sim)}" + ) # we don't keep source in this cases else: # already defined outcomes if source_name in all_data: @@ -358,16 +407,22 @@ def compute_all_multioutcomes( f"ERROR with outcome {new_comp}: the specified source {source_name} is not a dictionnary (for seir outcome) nor an existing pre-identified outcomes." ) - if (loaded_values is not None) and (new_comp in loaded_values["outcome"].values): + if (loaded_values is not None) and ( + new_comp in loaded_values["outcome"].values + ): ## This may be unnecessary probabilities = loaded_values[ - (loaded_values["quantity"] == "probability") & (loaded_values["outcome"] == new_comp) + (loaded_values["quantity"] == "probability") + & (loaded_values["outcome"] == new_comp) + ]["value"].to_numpy() + delays = loaded_values[ + (loaded_values["quantity"] == "delay") + & (loaded_values["outcome"] == new_comp) ]["value"].to_numpy() - delays = loaded_values[(loaded_values["quantity"] == "delay") & (loaded_values["outcome"] == new_comp)][ - "value" - ].to_numpy() else: - probabilities = parameters[new_comp]["probability"].as_random_distribution()( + probabilities = parameters[new_comp][ + "probability" + ].as_random_distribution()( size=len(modinf.subpop_struct.subpop_names) ) # one draw per subpop if "rel_probability" in parameters[new_comp]: @@ -378,8 +433,12 @@ def compute_all_multioutcomes( ) # one draw per subpop probabilities[probabilities > 1] = 1 probabilities[probabilities < 0] = 0 - probabilities = np.repeat(probabilities[:, np.newaxis], len(dates), axis=1).T # duplicate in time - delays = np.repeat(delays[:, np.newaxis], len(dates), axis=1).T # duplicate in time + probabilities = np.repeat( + probabilities[:, np.newaxis], len(dates), axis=1 + ).T # duplicate in time + delays = np.repeat( + delays[:, np.newaxis], len(dates), axis=1 + ).T # duplicate in time delays = np.round(delays).astype(int) # Write hpar before NPI subpop_names_len = len(modinf.subpop_struct.subpop_names) @@ -387,7 +446,7 @@ def compute_all_multioutcomes( { "subpop": 2 * modinf.subpop_struct.subpop_names, "quantity": (subpop_names_len * ["probability"]) - + (subpop_names_len * ["delay"]), + + (subpop_names_len * ["delay"]), "outcome": 2 * subpop_names_len * [new_comp], "value": np.concatenate( ( @@ -402,42 +461,59 @@ def compute_all_multioutcomes( if npi is not None: delays = NPI.reduce_parameter( parameter=delays, - modification=npi.getReduction(parameters[new_comp]["delay::npi_param_name"].lower()), + modification=npi.getReduction( + parameters[new_comp]["delay::npi_param_name"].lower() + ), ) delays = np.round(delays).astype(int) probabilities = NPI.reduce_parameter( parameter=probabilities, - modification=npi.getReduction(parameters[new_comp]["probability::npi_param_name"].lower()), + modification=npi.getReduction( + parameters[new_comp]["probability::npi_param_name"].lower() + ), ) # Create new compartment incidence: all_data[new_comp] = np.empty_like(source_array) # Draw with from source compartment if modinf.stoch_traj_flag: - all_data[new_comp] = np.random.binomial(source_array.astype(np.int32), probabilities) + all_data[new_comp] = np.random.binomial( + source_array.astype(np.int32), probabilities + ) else: - all_data[new_comp] = source_array * (probabilities * np.ones_like(source_array)) + all_data[new_comp] = source_array * ( + probabilities * np.ones_like(source_array) + ) # Shift to account for the delay ## stoch_delay_flag is whether to use stochastic delays or not stoch_delay_flag = False - all_data[new_comp] = multishift(all_data[new_comp], delays, stoch_delay_flag=stoch_delay_flag) + all_data[new_comp] = multishift( + all_data[new_comp], delays, stoch_delay_flag=stoch_delay_flag + ) # Produce a dataframe an merge it - df_p = dataframe_from_array(all_data[new_comp], modinf.subpop_struct.subpop_names, dates, new_comp) + df_p = dataframe_from_array( + all_data[new_comp], modinf.subpop_struct.subpop_names, dates, new_comp + ) outcomes = pd.merge(outcomes, df_p) # Make duration if "duration" in parameters[new_comp]: - if (loaded_values is not None) and (new_comp in loaded_values["outcome"].values): + if (loaded_values is not None) and ( + new_comp in loaded_values["outcome"].values + ): durations = loaded_values[ - (loaded_values["quantity"] == "duration") & (loaded_values["outcome"] == new_comp) + (loaded_values["quantity"] == "duration") + & (loaded_values["outcome"] == new_comp) ]["value"].to_numpy() else: durations = parameters[new_comp]["duration"].as_random_distribution()( size=len(modinf.subpop_struct.subpop_names) ) # one draw per subpop - durations = np.repeat(durations[:, np.newaxis], len(dates), axis=1).T # duplicate in time + durations = np.repeat( + durations[:, np.newaxis], len(dates), axis=1 + ).T # duplicate in time durations = np.round(durations).astype(int) hpar = pd.DataFrame( data={ @@ -458,7 +534,9 @@ def compute_all_multioutcomes( # print(f"{new_comp}-duration".lower(), npi.getReduction(f"{new_comp}-duration".lower())) durations = NPI.reduce_parameter( parameter=durations, - modification=npi.getReduction(parameters[new_comp]["duration::npi_param_name"].lower()), + modification=npi.getReduction( + parameters[new_comp]["duration::npi_param_name"].lower() + ), ) # npi.getReduction(f"{new_comp}::duration".lower())) durations = np.round(durations).astype(int) # plt.imshow(durations) @@ -492,7 +570,9 @@ def compute_all_multioutcomes( for cmp in parameters[new_comp]["sum"]: sum_outcome += all_data[cmp] all_data[new_comp] = sum_outcome - df_p = dataframe_from_array(sum_outcome, modinf.subpop_struct.subpop_names, dates, new_comp) + df_p = dataframe_from_array( + sum_outcome, modinf.subpop_struct.subpop_names, dates, new_comp + ) outcomes = pd.merge(outcomes, df_p) # Concat our hpar dataframes hpar = ( @@ -554,7 +634,9 @@ def filter_seir_xr(diffI, dates, subpops, filters, outcome_name) -> np.ndarray: if isinstance(mc_value, str): mc_value = [mc_value] # Filter data along the specified mc_type dimension - diffI_filtered = diffI_filtered.where(diffI_filtered[f"mc_{mc_type}"].isin(mc_value), drop=True) + diffI_filtered = diffI_filtered.where( + diffI_filtered[f"mc_{mc_type}"].isin(mc_value), drop=True + ) # Sum along the compartment dimension incidI_arr += diffI_filtered[vtype].sum(dim="compartment") @@ -626,7 +708,9 @@ def multishift(arr, shifts, stoch_delay_flag=True): # for k,case in enumerate(cases): # results[i+k][j] = cases[k] else: - for i in range(arr.shape[0]): # numba nopython does not allow iterating over 2D array + for i in range( + arr.shape[0] + ): # numba nopython does not allow iterating over 2D array for j in range(arr.shape[1]): if i + shifts[i, j] < arr.shape[0]: result[i + shifts[i, j], j] += arr[i, j] diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index c013d7388..0a45883cb 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -95,20 +95,14 @@ def __init__( # Parameter characterized by it's distribution if self.pconfig[pn]["value"].exists(): - self.pdata[pn]["dist"] = self.pconfig[pn][ - "value" - ].as_random_distribution() + self.pdata[pn]["dist"] = self.pconfig[pn]["value"].as_random_distribution() # Parameter given as a file elif self.pconfig[pn]["timeseries"].exists(): - fn_name = os.path.join( - path_prefix, self.pconfig[pn]["timeseries"].get() - ) + fn_name = os.path.join(path_prefix, self.pconfig[pn]["timeseries"].get()) df = utils.read_df(fn_name).set_index("date") df.index = pd.to_datetime(df.index) - if ( - len(df.columns) == 1 - ): # if only one ts, assume it applies to all subpops + if len(df.columns) == 1: # if only one ts, assume it applies to all subpops df = pd.DataFrame( pd.concat([df] * len(subpop_names), axis=1).values, index=df.index, @@ -166,9 +160,9 @@ def __init__( "rolling_mean_windows" ].get() - self.stacked_modifier_method[ - self.pdata[pn]["stacked_modifier_method"] - ].append(pn.lower()) + self.stacked_modifier_method[self.pdata[pn]["stacked_modifier_method"]].append( + pn.lower() + ) logging.debug(f"We have {self.npar} parameter: {self.pnames}") logging.debug(f"Data to sample is: {self.pdata}") @@ -316,9 +310,7 @@ def getParameterDF(self, p_draw: ndarray) -> pd.DataFrame: if "dist" in self.pdata[pn] ], columns=["value"], - index=[ - pn for idx, pn in enumerate(self.pnames) if "dist" in self.pdata[pn] - ], + index=[pn for idx, pn in enumerate(self.pnames) if "dist" in self.pdata[pn]], ) out_df["parameter"] = out_df.index return out_df diff --git a/flepimop/gempyor_pkg/src/gempyor/postprocess_inference.py b/flepimop/gempyor_pkg/src/gempyor/postprocess_inference.py index 2b42e2944..684fa886c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/postprocess_inference.py +++ b/flepimop/gempyor_pkg/src/gempyor/postprocess_inference.py @@ -21,7 +21,15 @@ import pandas as pd import pyarrow.parquet as pq import xarray as xr -from gempyor import config, model_info, outcomes, seir, inference_parameter, logloss, inference +from gempyor import ( + config, + model_info, + outcomes, + seir, + inference_parameter, + logloss, + inference, +) from gempyor.inference import GempyorInference import tqdm import os @@ -37,7 +45,9 @@ def find_walkers_to_sample(inferpar, sampler_output, nsamples, nwalker, nthin): last_llik = sampler_output.get_log_prob()[-1, :] sampled_slots = last_llik > (last_llik.mean() - 1 * last_llik.std()) - print(f"there are {sampled_slots.sum()}/{len(sampled_slots)} good walkers... keeping these") + print( + f"there are {sampled_slots.sum()}/{len(sampled_slots)} good walkers... keeping these" + ) # TODO this function give back good_samples = sampler.get_chain()[:, sampled_slots, :] @@ -50,9 +60,9 @@ def find_walkers_to_sample(inferpar, sampler_output, nsamples, nwalker, nthin): ] # parentesis around i//(sampled_slots.sum() are very important - - -def plot_chains(inferpar, chains, llik, save_to, sampled_slots=None, param_gt=None, llik_gt=None): +def plot_chains( + inferpar, chains, llik, save_to, sampled_slots=None, param_gt=None, llik_gt=None +): """ Plot the chains of the inference :param inferpar: the inference parameter object @@ -113,24 +123,44 @@ def plot_single_chain(frompt, ax, chain, label, gt=None): for sp in tqdm.tqdm(set(inferpar.subpops)): # find unique supopulation these_pars = inferpar.get_parameters_for_subpop(sp) - fig, axes = plt.subplots(max(len(these_pars), 2), 2, figsize=(6, (len(these_pars) + 1) * 2)) + fig, axes = plt.subplots( + max(len(these_pars), 2), 2, figsize=(6, (len(these_pars) + 1) * 2) + ) for idx, par_id in enumerate(these_pars): - plot_single_chain(first_thresh, axes[idx, 0], chains[:, :, par_id], labels[par_id], gt=param_gt[par_id] if param_gt is not None else None) - plot_single_chain(second_thresh, axes[idx, 1], chains[:, :, par_id], labels[par_id], gt=param_gt[par_id] if param_gt is not None else None) + plot_single_chain( + first_thresh, + axes[idx, 0], + chains[:, :, par_id], + labels[par_id], + gt=param_gt[par_id] if param_gt is not None else None, + ) + plot_single_chain( + second_thresh, + axes[idx, 1], + chains[:, :, par_id], + labels[par_id], + gt=param_gt[par_id] if param_gt is not None else None, + ) fig.tight_layout() pdf.savefig(fig) plt.close(fig) + def plot_fit(modinf, loss): subpop_names = modinf.subpop_struct.subpop_names fig, axes = plt.subplots( - len(subpop_names), len(loss.statistics), figsize=(3 * len(loss.statistics), 3 * len(subpop_names)), sharex=True + len(subpop_names), + len(loss.statistics), + figsize=(3 * len(loss.statistics), 3 * len(subpop_names)), + sharex=True, ) for j, subpop in enumerate(modinf.subpop_struct.subpop_names): gt_s = loss.gt[loss.gt["subpop"] == subpop].sort_index() first_date = max(gt_s.index.min(), results[0].index.min()) last_date = min(gt_s.index.max(), results[0].index.max()) - gt_s = gt_s.loc[first_date:last_date].drop(["subpop"], axis=1).resample("W-SAT").sum() + gt_s = ( + gt_s.loc[first_date:last_date].drop(["subpop"], axis=1).resample("W-SAT").sum() + ) for i, (stat_name, stat) in enumerate(loss.statistics.items()): ax = axes[j, i] diff --git a/flepimop/gempyor_pkg/src/gempyor/seeding.py b/flepimop/gempyor_pkg/src/gempyor/seeding.py index fe58657c0..53b81c8dc 100644 --- a/flepimop/gempyor_pkg/src/gempyor/seeding.py +++ b/flepimop/gempyor_pkg/src/gempyor/seeding.py @@ -17,9 +17,13 @@ def _DataFrame2NumbaDict(df, amounts, modinf) -> nb.typed.Dict: if not df["date"].is_monotonic_increasing: - raise ValueError("_DataFrame2NumbaDict got an unsorted dataframe, exposing itself to non-sense") + raise ValueError( + "_DataFrame2NumbaDict got an unsorted dataframe, exposing itself to non-sense" + ) - cmp_grp_names = [col for col in modinf.compartments.compartments.columns if col != "name"] + cmp_grp_names = [ + col for col in modinf.compartments.compartments.columns if col != "name" + ] seeding_dict: nb.typed.Dict = nb.typed.Dict.empty( key_type=nb.types.unicode_type, value_type=nb.types.int64[:], @@ -45,16 +49,25 @@ def _DataFrame2NumbaDict(df, amounts, modinf) -> nb.typed.Dict: nb_seed_perday[(row["date"].date() - modinf.ti).days] = ( nb_seed_perday[(row["date"].date() - modinf.ti).days] + 1 ) - source_dict = {grp_name: row[f"source_{grp_name}"] for grp_name in cmp_grp_names} - destination_dict = {grp_name: row[f"destination_{grp_name}"] for grp_name in cmp_grp_names} + source_dict = { + grp_name: row[f"source_{grp_name}"] for grp_name in cmp_grp_names + } + destination_dict = { + grp_name: row[f"destination_{grp_name}"] for grp_name in cmp_grp_names + } seeding_dict["seeding_sources"][idx] = modinf.compartments.get_comp_idx( - source_dict, error_info=f"(seeding source at idx={idx}, row_index={row_index}, row=>>{row}<<)" + source_dict, + error_info=f"(seeding source at idx={idx}, row_index={row_index}, row=>>{row}<<)", + ) + seeding_dict["seeding_destinations"][idx] = ( + modinf.compartments.get_comp_idx( + destination_dict, + error_info=f"(seeding destination at idx={idx}, row_index={row_index}, row=>>{row}<<)", + ) ) - seeding_dict["seeding_destinations"][idx] = modinf.compartments.get_comp_idx( - destination_dict, - error_info=f"(seeding destination at idx={idx}, row_index={row_index}, row=>>{row}<<)", + seeding_dict["seeding_subpops"][idx] = ( + modinf.subpop_struct.subpop_names.index(row["subpop"]) ) - seeding_dict["seeding_subpops"][idx] = modinf.subpop_struct.subpop_names.index(row["subpop"]) seeding_amounts[idx] = amounts[idx] # id_seed+=1 else: @@ -97,7 +110,9 @@ def get_from_config(self, sim_id: int, modinf) -> nb.typed.Dict: ) dupes = seeding[seeding.duplicated(["subpop", "date"])].index + 1 if not dupes.empty: - raise ValueError(f"Repeated subpop-date in rows {dupes.tolist()} of seeding::lambda_file.") + raise ValueError( + f"Repeated subpop-date in rows {dupes.tolist()} of seeding::lambda_file." + ) elif method == "FolderDraw": seeding = pd.read_csv( self.path_prefix @@ -127,7 +142,9 @@ def get_from_config(self, sim_id: int, modinf) -> nb.typed.Dict: # print(seeding.shape) seeding = seeding.sort_values(by="date", axis="index").reset_index() # print(seeding) - mask = (seeding["date"].dt.date > modinf.ti) & (seeding["date"].dt.date <= modinf.tf) + mask = (seeding["date"].dt.date > modinf.ti) & ( + seeding["date"].dt.date <= modinf.tf + ) seeding = seeding.loc[mask].reset_index() # print(seeding.shape) # print(seeding) @@ -138,7 +155,9 @@ def get_from_config(self, sim_id: int, modinf) -> nb.typed.Dict: if method == "PoissonDistributed": amounts = np.random.poisson(seeding["amount"]) elif method == "NegativeBinomialDistributed": - raise ValueError("Seeding method 'NegativeBinomialDistributed' is not supported by flepiMoP anymore.") + raise ValueError( + "Seeding method 'NegativeBinomialDistributed' is not supported by flepiMoP anymore." + ) elif method == "FolderDraw" or method == "FromFile": amounts = seeding["amount"] else: diff --git a/flepimop/gempyor_pkg/src/gempyor/seir.py b/flepimop/gempyor_pkg/src/gempyor/seir.py index 4e59761f2..5ea236c98 100644 --- a/flepimop/gempyor_pkg/src/gempyor/seir.py +++ b/flepimop/gempyor_pkg/src/gempyor/seir.py @@ -41,7 +41,9 @@ def build_step_source_arg( else: integration_method = "rk4.jit" dt = 2.0 - logging.info(f"Integration method not provided, assuming type {integration_method} with dt=2") + logging.info( + f"Integration method not provided, assuming type {integration_method} with dt=2" + ) ## The type is very important for the call to the compiled function, and e.g mixing an int64 for an int32 can ## result in serious error. Note that "In Microsoft C, even on a 64 bit system, the size of the long int data type @@ -58,7 +60,10 @@ def build_step_source_arg( assert type(transition_array[0][0]) == np.int64 assert type(proportion_array[0]) == np.int64 assert type(proportion_info[0][0]) == np.int64 - assert initial_conditions.shape == (modinf.compartments.compartments.shape[0], modinf.nsubpops) + assert initial_conditions.shape == ( + modinf.compartments.compartments.shape[0], + modinf.nsubpops, + ) assert type(initial_conditions[0][0]) == np.float64 # Test of empty seeding: assert len(seeding_data.keys()) == 4 @@ -162,7 +167,9 @@ def steps_SEIR( f"with method {integration_method}, only deterministic " f"integration is possible (got stoch_straj_flag={modinf.stoch_traj_flag}" ) - seir_sim = steps_experimental.ode_integration(**fnct_args, integration_method=integration_method) + seir_sim = steps_experimental.ode_integration( + **fnct_args, integration_method=integration_method + ) elif integration_method == "rk4.jit1": seir_sim = steps_experimental.rk4_integration1(**fnct_args) elif integration_method == "rk4.jit2": @@ -200,7 +207,9 @@ def steps_SEIR( **compartment_coords, subpop=modinf.subpop_struct.subpop_names, ), - attrs=dict(description="Dynamical simulation results", run_id=modinf.in_run_id), # TODO add more information + attrs=dict( + description="Dynamical simulation results", run_id=modinf.in_run_id + ), # TODO add more information ) return states @@ -223,8 +232,12 @@ def build_npi_SEIR(modinf, load_ID, sim_id2load, config, bypass_DF=None, bypass_ modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, loaded_df=loaded_df, - pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method[ + "sum" + ], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) else: npi = NPI.NPIBase.execute( @@ -232,8 +245,12 @@ def build_npi_SEIR(modinf, load_ID, sim_id2load, config, bypass_DF=None, bypass_ modinf=modinf, modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, - pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method[ + "sum" + ], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) return npi @@ -248,7 +265,9 @@ def onerun_SEIR( np.random.seed() npi = None if modinf.npi_config_seir: - npi = build_npi_SEIR(modinf=modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config) + npi = build_npi_SEIR( + modinf=modinf, load_ID=load_ID, sim_id2load=sim_id2load, config=config + ) with Timer("onerun_SEIR.compartments"): ( @@ -260,11 +279,19 @@ def onerun_SEIR( with Timer("onerun_SEIR.seeding"): if load_ID: - initial_conditions = modinf.initial_conditions.get_from_file(sim_id2load, modinf=modinf) - seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id2load, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_file( + sim_id2load, modinf=modinf + ) + seeding_data, seeding_amounts = modinf.seeding.get_from_file( + sim_id2load, modinf=modinf + ) else: - initial_conditions = modinf.initial_conditions.get_from_config(sim_id2write, modinf=modinf) - seeding_data, seeding_amounts = modinf.seeding.get_from_config(sim_id2write, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id2write, modinf=modinf + ) + seeding_data, seeding_amounts = modinf.seeding.get_from_config( + sim_id2write, modinf=modinf + ) with Timer("onerun_SEIR.parameters"): # Draw or load parameters @@ -275,14 +302,18 @@ def onerun_SEIR( nsubpops=modinf.nsubpops, ) else: - p_draw = modinf.parameters.parameters_quick_draw(n_days=modinf.n_days, nsubpops=modinf.nsubpops) + p_draw = modinf.parameters.parameters_quick_draw( + n_days=modinf.n_days, nsubpops=modinf.nsubpops + ) # reduce them parameters = modinf.parameters.parameters_reduce(p_draw, npi) log_debug_parameters(p_draw, "Parameters without seir_modifiers") log_debug_parameters(parameters, "Parameters with seir_modifiers") # Parse them - parsed_parameters = modinf.compartments.parse_parameters(parameters, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + parameters, modinf.parameters.pnames, unique_strings + ) log_debug_parameters(parsed_parameters, "Unique Parameters used by transitions") with Timer("onerun_SEIR.compute"): @@ -310,7 +341,13 @@ def run_parallel_SEIR(modinf, config, *, n_jobs=1): if n_jobs == 1: # run single process for debugging/profiling purposes for sim_id in tqdm.tqdm(sim_ids): - onerun_SEIR(sim_id2write=sim_id, modinf=modinf, load_ID=False, sim_id2load=None, config=config) + onerun_SEIR( + sim_id2write=sim_id, + modinf=modinf, + load_ID=False, + sim_id2load=None, + config=config, + ) else: tqdm.contrib.concurrent.process_map( onerun_SEIR, @@ -322,7 +359,9 @@ def run_parallel_SEIR(modinf, config, *, n_jobs=1): max_workers=n_jobs, ) - logging.info(f""">> {modinf.nslots} seir simulations completed in {time.monotonic() - start:.1f} seconds""") + logging.info( + f""">> {modinf.nslots} seir simulations completed in {time.monotonic() - start:.1f} seconds""" + ) def states2Df(modinf, states): @@ -337,12 +376,17 @@ def states2Df(modinf, states): # states_diff = np.diff(states_diff, axis=0) ts_index = pd.MultiIndex.from_product( - [pd.date_range(modinf.ti, modinf.tf, freq="D"), modinf.compartments.compartments["name"]], + [ + pd.date_range(modinf.ti, modinf.tf, freq="D"), + modinf.compartments.compartments["name"], + ], names=["date", "mc_name"], ) # prevalence data, we use multi.index dataframe, sparring us the array manipulation we use to do prev_df = pd.DataFrame( - data=states["prevalence"].to_numpy().reshape(modinf.n_days * modinf.compartments.get_ncomp(), modinf.nsubpops), + data=states["prevalence"] + .to_numpy() + .reshape(modinf.n_days * modinf.compartments.get_ncomp(), modinf.nsubpops), index=ts_index, columns=modinf.subpop_struct.subpop_names, ).reset_index() @@ -355,12 +399,17 @@ def states2Df(modinf, states): prev_df.insert(loc=0, column="mc_value_type", value="prevalence") ts_index = pd.MultiIndex.from_product( - [pd.date_range(modinf.ti, modinf.tf, freq="D"), modinf.compartments.compartments["name"]], + [ + pd.date_range(modinf.ti, modinf.tf, freq="D"), + modinf.compartments.compartments["name"], + ], names=["date", "mc_name"], ) incid_df = pd.DataFrame( - data=states["incidence"].to_numpy().reshape(modinf.n_days * modinf.compartments.get_ncomp(), modinf.nsubpops), + data=states["incidence"] + .to_numpy() + .reshape(modinf.n_days * modinf.compartments.get_ncomp(), modinf.nsubpops), index=ts_index, columns=modinf.subpop_struct.subpop_names, ).reset_index() @@ -384,7 +433,9 @@ def write_spar_snpi(sim_id, modinf, p_draw, npi): if npi is not None: modinf.write_simID(ftype="snpi", sim_id=sim_id, df=npi.getReductionDF()) # Parameters - modinf.write_simID(ftype="spar", sim_id=sim_id, df=modinf.parameters.getParameterDF(p_draw=p_draw)) + modinf.write_simID( + ftype="spar", sim_id=sim_id, df=modinf.parameters.getParameterDF(p_draw=p_draw) + ) def write_seir(sim_id, modinf, states): diff --git a/flepimop/gempyor_pkg/src/gempyor/simulate.py b/flepimop/gempyor_pkg/src/gempyor/simulate.py index 34fdd9d4b..d97b94a93 100644 --- a/flepimop/gempyor_pkg/src/gempyor/simulate.py +++ b/flepimop/gempyor_pkg/src/gempyor/simulate.py @@ -299,23 +299,31 @@ def simulate( seir_modifiers_scenarios = None if config["seir_modifiers"].exists(): if config["seir_modifiers"]["scenarios"].exists(): - seir_modifiers_scenarios = config["seir_modifiers"]["scenarios"].as_str_seq() + seir_modifiers_scenarios = config["seir_modifiers"][ + "scenarios" + ].as_str_seq() # Model Info handles the case of the default scneario if not outcome_modifiers_scenarios: outcome_modifiers_scenarios = None if config["outcomes"].exists() and config["outcome_modifiers"].exists(): if config["outcome_modifiers"]["scenarios"].exists(): - outcome_modifiers_scenarios = config["outcome_modifiers"]["scenarios"].as_str_seq() + outcome_modifiers_scenarios = config["outcome_modifiers"][ + "scenarios" + ].as_str_seq() outcome_modifiers_scenarios = as_list(outcome_modifiers_scenarios) seir_modifiers_scenarios = as_list(seir_modifiers_scenarios) print(outcome_modifiers_scenarios, seir_modifiers_scenarios) - scenarios_combinations = [[s, d] for s in seir_modifiers_scenarios for d in outcome_modifiers_scenarios] + scenarios_combinations = [ + [s, d] for s in seir_modifiers_scenarios for d in outcome_modifiers_scenarios + ] print("Combination of modifiers scenarios to be run: ") print(scenarios_combinations) for seir_modifiers_scenario, outcome_modifiers_scenario in scenarios_combinations: - print(f"seir_modifier: {seir_modifiers_scenario}, outcomes_modifier:{outcome_modifiers_scenario}") + print( + f"seir_modifier: {seir_modifiers_scenario}, outcomes_modifier:{outcome_modifiers_scenario}" + ) if not nslots: nslots = config["nslots"].as_number() @@ -354,7 +362,9 @@ def simulate( if config["seir"].exists(): seir.run_parallel_SEIR(modinf, config=config, n_jobs=jobs) if config["outcomes"].exists(): - outcomes.run_parallel_outcomes(sim_id2write=first_sim_index, modinf=modinf, nslots=nslots, n_jobs=jobs) + outcomes.run_parallel_outcomes( + sim_id2write=first_sim_index, modinf=modinf, nslots=nslots, n_jobs=jobs + ) print( f">>> {seir_modifiers_scenario}_{outcome_modifiers_scenario} completed in {time.monotonic() - start:.1f} seconds" ) diff --git a/flepimop/gempyor_pkg/src/gempyor/steps_source.py b/flepimop/gempyor_pkg/src/gempyor/steps_source.py index b8af1d493..ea6f7ad21 100644 --- a/flepimop/gempyor_pkg/src/gempyor/steps_source.py +++ b/flepimop/gempyor_pkg/src/gempyor/steps_source.py @@ -30,9 +30,12 @@ ## Return "UniTuple(float64[:, :, :], 2) (" ## return states and cumlative states, both [ ndays x ncompartments x nspatial_nodes ] ## Dimensions - "int32," "int32," "int32," ## ncompartments ## nspatial_nodes ## Number of days + "int32," + "int32," + "int32," ## ncompartments ## nspatial_nodes ## Number of days ## Parameters - "float64[:, :, :]," "float64," ## Parameters [ nparameters x ndays x nspatial_nodes] ## dt + "float64[:, :, :]," + "float64," ## Parameters [ nparameters x ndays x nspatial_nodes] ## dt ## Transitions "int64[:, :]," ## transitions [ [source, destination, proportion_start, proportion_stop, rate] x ntransitions ] "int64[:, :]," ## proportions_info [ [sum_starts, sum_stops, exponent] x ntransition_proportions ] @@ -84,7 +87,11 @@ def steps_SEIR_nb( percent_day_away = 0.5 for spatial_node in range(nspatial_nodes): percent_who_move[spatial_node] = min( - mobility_data[mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1]].sum() + mobility_data[ + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] + ].sum() / population[spatial_node], 1, ) @@ -111,18 +118,22 @@ def steps_SEIR_nb( this_seeding_amounts = seeding_amounts[seeding_instance_idx] seeding_subpops = seeding_data["seeding_subpops"][seeding_instance_idx] seeding_sources = seeding_data["seeding_sources"][seeding_instance_idx] - seeding_destinations = seeding_data["seeding_destinations"][seeding_instance_idx] + seeding_destinations = seeding_data["seeding_destinations"][ + seeding_instance_idx + ] # this_seeding_amounts = this_seeding_amounts < states_next[seeding_sources] ? this_seeding_amounts : states_next[seeding_instance_idx] states_next[seeding_sources][seeding_subpops] -= this_seeding_amounts - states_next[seeding_sources][seeding_subpops] = states_next[seeding_sources][seeding_subpops] * ( - states_next[seeding_sources][seeding_subpops] > 0 - ) + states_next[seeding_sources][seeding_subpops] = states_next[ + seeding_sources + ][seeding_subpops] * (states_next[seeding_sources][seeding_subpops] > 0) states_next[seeding_destinations][seeding_subpops] += this_seeding_amounts total_seeded += this_seeding_amounts times_seeded += 1 # ADD TO cumulative, this is debatable, - states_daily_incid[today][seeding_destinations][seeding_subpops] += this_seeding_amounts + states_daily_incid[today][seeding_destinations][ + seeding_subpops + ] += this_seeding_amounts total_infected = 0 for transition_index in range(ntransitions): @@ -138,9 +149,13 @@ def steps_SEIR_nb( proportion_info[proportion_sum_starts_col][proportion_index], proportion_info[proportion_sum_stops_col][proportion_index], ): - relevant_number_in_comp += states_current[transition_sum_compartments[proportion_sum_index]] + relevant_number_in_comp += states_current[ + transition_sum_compartments[proportion_sum_index] + ] # exponents should not be a proportion, since we don't sum them over sum compartments - relevant_exponent = parameters[proportion_info[proportion_exponent_col][proportion_index]][today] + relevant_exponent = parameters[ + proportion_info[proportion_exponent_col][proportion_index] + ][today] if first_proportion: only_one_proportion = ( transitions[transition_proportion_start_col][transition_index] + 1 @@ -149,41 +164,56 @@ def steps_SEIR_nb( source_number = relevant_number_in_comp if source_number.max() > 0: total_rate[source_number > 0] *= ( - source_number[source_number > 0] ** relevant_exponent[source_number > 0] + source_number[source_number > 0] + ** relevant_exponent[source_number > 0] / source_number[source_number > 0] ) if only_one_proportion: - total_rate *= parameters[transitions[transition_rate_col][transition_index]][today] + total_rate *= parameters[ + transitions[transition_rate_col][transition_index] + ][today] else: for spatial_node in range(nspatial_nodes): - proportion_keep_compartment = 1 - percent_day_away * percent_who_move[spatial_node] + proportion_keep_compartment = ( + 1 - percent_day_away * percent_who_move[spatial_node] + ) proportion_change_compartment = ( percent_day_away * mobility_data[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] / population[spatial_node] ) rate_keep_compartment = ( proportion_keep_compartment - * relevant_number_in_comp[spatial_node] ** relevant_exponent[spatial_node] + * relevant_number_in_comp[spatial_node] + ** relevant_exponent[spatial_node] / population[spatial_node] - * parameters[transitions[transition_rate_col][transition_index]][today][spatial_node] + * parameters[ + transitions[transition_rate_col][transition_index] + ][today][spatial_node] ) visiting_compartment = mobility_row_indices[ - mobility_data_indices[spatial_node] : mobility_data_indices[spatial_node + 1] + mobility_data_indices[spatial_node] : mobility_data_indices[ + spatial_node + 1 + ] ] rate_change_compartment = proportion_change_compartment rate_change_compartment *= ( - relevant_number_in_comp[visiting_compartment] ** relevant_exponent[visiting_compartment] + relevant_number_in_comp[visiting_compartment] + ** relevant_exponent[visiting_compartment] ) rate_change_compartment /= population[visiting_compartment] - rate_change_compartment *= parameters[transitions[transition_rate_col][transition_index]][ - today - ][visiting_compartment] - total_rate[spatial_node] *= rate_keep_compartment + rate_change_compartment.sum() + rate_change_compartment *= parameters[ + transitions[transition_rate_col][transition_index] + ][today][visiting_compartment] + total_rate[spatial_node] *= ( + rate_keep_compartment + rate_change_compartment.sum() + ) compound_adjusted_rate = 1.0 - np.exp(-dt * total_rate) @@ -221,14 +251,20 @@ def steps_SEIR_nb( for spatial_node in range(nspatial_nodes): if ( number_move[spatial_node] - > states_next[transitions[transition_source_col][transition_index]][spatial_node] - ): - number_move[spatial_node] = states_next[transitions[transition_source_col][transition_index]][ + > states_next[transitions[transition_source_col][transition_index]][ spatial_node ] + ): + number_move[spatial_node] = states_next[ + transitions[transition_source_col][transition_index] + ][spatial_node] states_next[transitions[transition_source_col][transition_index]] -= number_move - states_next[transitions[transition_destination_col][transition_index]] += number_move - states_daily_incid[today, transitions[transition_destination_col][transition_index], :] += number_move + states_next[ + transitions[transition_destination_col][transition_index] + ] += number_move + states_daily_incid[ + today, transitions[transition_destination_col][transition_index], : + ] += number_move states_current = states_next.copy() diff --git a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py index 1e2b9b8de..7c46a9960 100644 --- a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py +++ b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py @@ -27,7 +27,9 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): self.setup_name = setup_name self.data = pd.read_csv( - geodata_file, converters={subpop_names_key: lambda x: str(x).strip()}, skipinitialspace=True + geodata_file, + converters={subpop_names_key: lambda x: str(x).strip()}, + skipinitialspace=True, ) # subpops and populations, strip whitespaces self.nsubpops = len(self.data) # K = # of locations @@ -44,7 +46,9 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): # subpop_names_key is the name of the column in geodata_file with subpops if subpop_names_key not in self.data: - raise ValueError(f"subpop_names_key: {subpop_names_key} does not correspond to a column in geodata.") + raise ValueError( + f"subpop_names_key: {subpop_names_key} does not correspond to a column in geodata." + ) self.subpop_names = self.data[subpop_names_key].tolist() if len(self.subpop_names) != len(set(self.subpop_names)): raise ValueError(f"There are duplicate subpop_names in geodata.") @@ -53,7 +57,9 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): mobility_file = path_prefix / subpop_config["mobility"].get() mobility_file = pathlib.Path(mobility_file) if mobility_file.suffix == ".txt": - print("Mobility files as matrices are not recommended. Please switch soon to long form csv files.") + print( + "Mobility files as matrices are not recommended. Please switch soon to long form csv files." + ) self.mobility = scipy.sparse.csr_matrix( np.loadtxt(mobility_file), dtype=int ) # K x K matrix of people moving @@ -64,7 +70,11 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): ) elif mobility_file.suffix == ".csv": - mobility_data = pd.read_csv(mobility_file, converters={"ori": str, "dest": str}, skipinitialspace=True) + mobility_data = pd.read_csv( + mobility_file, + converters={"ori": str, "dest": str}, + skipinitialspace=True, + ) nn_dict = {v: k for k, v in enumerate(self.subpop_names)} mobility_data["ori_idx"] = mobility_data["ori"].apply(nn_dict.__getitem__) mobility_data["dest_idx"] = mobility_data["dest"].apply(nn_dict.__getitem__) @@ -115,7 +125,9 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): ) else: logging.critical("No mobility matrix specified -- assuming no one moves") - self.mobility = scipy.sparse.csr_matrix(np.zeros((self.nsubpops, self.nsubpops)), dtype=int) + self.mobility = scipy.sparse.csr_matrix( + np.zeros((self.nsubpops, self.nsubpops)), dtype=int + ) if subpop_config["selected"].exists(): selected = subpop_config["selected"].get() @@ -129,4 +141,6 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")): self.subpop_names = selected self.nsubpops = len(self.data) # TODO: this needs to be tested - self.mobility = self.mobility[selected_subpop_indices][:, selected_subpop_indices] + self.mobility = self.mobility[selected_subpop_indices][ + :, selected_subpop_indices + ] diff --git a/flepimop/gempyor_pkg/src/gempyor/utils.py b/flepimop/gempyor_pkg/src/gempyor/utils.py index 25c33d574..377dcdee5 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils.py @@ -99,7 +99,9 @@ def read_df( ) -def command_safe_run(command: str, command_name: str="mycommand", fail_on_fail: bool=True) -> tuple[int, str, str]: +def command_safe_run( + command: str, command_name: str = "mycommand", fail_on_fail: bool = True +) -> tuple[int, str, str]: """ Runs a shell command and prints diagnostics if command fails. @@ -110,14 +112,16 @@ def command_safe_run(command: str, command_name: str="mycommand", fail_on_fail: Returns: As a tuple; the return code, the standard output, and standard error from running the command. - + Raises: Exception: If fail_on_fail=True and the command fails, an exception will be thrown. """ import subprocess import shlex # using shlex to split the command because it's not obvious https://docs.python.org/3/library/subprocess.html#subprocess.Popen - sr = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE) + sr = subprocess.Popen( + shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) (stdout, stderr) = sr.communicate() if sr.returncode != 0: print(f"{command_name} failed failed with returncode {sr.returncode}") @@ -143,7 +147,7 @@ def add_method(cls): Args: cls: The class you want to add a method to. - + Returns: decorator: The decorator. """ @@ -159,7 +163,9 @@ def wrapper(*args, **kwargs): return decorator -def search_and_import_plugins_class(plugin_file_path: str, path_prefix: str, class_name: str, **kwargs: dict[str, Any]) -> Any: +def search_and_import_plugins_class( + plugin_file_path: str, path_prefix: str, class_name: str, **kwargs: dict[str, Any] +) -> Any: """ Function serving to create a class that finds and imports the necessary modules. @@ -174,11 +180,11 @@ def search_and_import_plugins_class(plugin_file_path: str, path_prefix: str, cla Examples: Suppose there is a module called `my_plugin.py with a class `MyClass` located at `/path/to/plugin/`. - + Dynamically import and instantiate the class: >>> instance = search_and_import_plugins_class('/path/to/plugin', path_prefix, 'MyClass', **params) - + View the instance: >>> print(instance) @@ -206,28 +212,33 @@ def search_and_import_plugins_class(plugin_file_path: str, path_prefix: str, cla from functools import wraps -def profile(output_file: str = None, sort_by: str = "cumulative", lines_to_print: int = None, strip_dirs: bool = False): +def profile( + output_file: str = None, + sort_by: str = "cumulative", + lines_to_print: int = None, + strip_dirs: bool = False, +): """ A time profiler decorator. Inspired by and modified the profile decorator of Giampaolo Rodola: http://code.activestate.com/recipes/577817-profile-decorator/ Args: - output_file: + output_file: Path of the output file. If only name of the file is given, it's saved in the current directory. If it's None, the name of the decorated function is used. - sort_by: + sort_by: Sorting criteria for the Stats object. For a list of valid string and SortKey refer to: https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats - lines_to_print: + lines_to_print: Number of lines to print. Default (None) is for all the lines. This is useful in reducing the size of the printout, especially that sorting by 'cumulative', the time consuming operations are printed toward the top of the file. - strip_dirs: - Whether to remove the leading path info from file names. + strip_dirs: + Whether to remove the leading path info from file names. Returns: Profile of the decorated function. @@ -238,7 +249,7 @@ def profile(output_file: str = None, sort_by: str = "cumulative", lines_to_print # Function body content pass >>> my_function() - After running ``my_function``, a file named ``my_function.prof`` will be created in the current WD. + After running ``my_function``, a file named ``my_function.prof`` will be created in the current WD. This file contains the profiling data. """ @@ -282,6 +293,7 @@ class Timer(object): name: Name of event. tstart: Time start. """ + def __init__(self, name): self.name = name @@ -297,6 +309,7 @@ class ISO8601Date(confuse.Template): """ Reads in config dates into datetimes.dates. """ + def convert(self, value: any, view: confuse.ConfigView): """ Converts the given value to a datetime.date object. @@ -331,7 +344,7 @@ def as_date(self) -> datetime.date: @add_method(confuse.ConfigView) def as_evaled_expression(self): """ - Evaluates an expression string, returning a float. + Evaluates an expression string, returning a float. Returns: A float data type of the value associated with the object. @@ -373,7 +386,7 @@ def get_truncated_normal( Returns: rv_frozen: A frozen instance of the truncated normal distribution with the specified parameters. - + Examples: Create a truncated normal distribution with specified parameters (truncated between 1 and 10): >>> truncated_normal_dist = get_truncated_normal(mean=5, sd=2, a=1, b=10) @@ -402,7 +415,7 @@ def get_log_normal( Returns: rv_frozen: A frozen instance of the log normal distribution with the specified parameters. - + Examples: Create a log-normal distribution with specified parameters: >>> log_normal_dist = get_log_normal(meanlog=1, sdlog=0.5) @@ -415,17 +428,17 @@ def get_log_normal( def random_distribution_sampler( distribution: Literal[ "fixed", "uniform", "poisson", "binomial", "truncnorm", "lognorm" - ], - **kwargs: dict[str, Any] + ], + **kwargs: dict[str, Any], ) -> Callable[[], float | int]: """ Create function to sample from a random distribution. - + Args: distribution: The type of distribution to generate a sampling function for. **kwargs: Further parameters that are passed to the underlying function for the given distribution. - + Notes: The further args expected by each distribution type are: - fixed: value, @@ -434,14 +447,14 @@ def random_distribution_sampler( - binomial: n, p, - truncnorm: mean, sd, a, b, - lognorm: meanlog, sdlog. - + Returns: A function that can be called to sample from that distribution. - + Raises: ValueError: If `distribution` is 'binomial' the given `p` must be in (0,1). NotImplementedError: If `distribution` is not one of the type hinted options. - + Examples: >>> import numpy as np >>> np.random.seed(123) @@ -454,15 +467,15 @@ def random_distribution_sampler( if distribution == "fixed": # Fixed value is the same as uniform on [a, a) return functools.partial( - np.random.uniform, - kwargs.get("value"), + np.random.uniform, + kwargs.get("value"), kwargs.get("value"), ) elif distribution == "uniform": # Uniform on [low, high) return functools.partial( - np.random.uniform, - kwargs.get("low"), + np.random.uniform, + kwargs.get("low"), kwargs.get("high"), ) elif distribution == "poisson": @@ -476,9 +489,9 @@ def random_distribution_sampler( elif distribution == "truncnorm": # Truncated normal with mean, sd on interval [a, b] return get_truncated_normal( - mean=kwargs.get("mean"), - sd=kwargs.get("sd"), - a=kwargs.get("a"), + mean=kwargs.get("mean"), + sd=kwargs.get("sd"), + a=kwargs.get("a"), b=kwargs.get("b"), ).rvs elif distribution == "lognorm": @@ -497,8 +510,8 @@ def as_random_distribution(self): Returns: A partial object containing the random distribution. - - Raises: + + Raises: ValueError: When values are out of range. NotImplementedError: If an unknown distribution is found. @@ -515,7 +528,7 @@ def as_random_distribution(self): "distribution": "truncnorm", "mean": 0 "sd": 1, - "a": -1, + "a": -1, "b": 1 }) >>> truncnorm_dist_function = config_truncnorm.as_random_distribution() @@ -528,11 +541,15 @@ def as_random_distribution(self): dist = self["distribution"].get() if dist == "fixed": return functools.partial( - np.random.uniform, self["value"].as_evaled_expression(), self["value"].as_evaled_expression(), + np.random.uniform, + self["value"].as_evaled_expression(), + self["value"].as_evaled_expression(), ) elif dist == "uniform": return functools.partial( - np.random.uniform, self["low"].as_evaled_expression(), self["high"].as_evaled_expression(), + np.random.uniform, + self["low"].as_evaled_expression(), + self["high"].as_evaled_expression(), ) elif dist == "poisson": return functools.partial(np.random.poisson, self["lam"].as_evaled_expression()) @@ -557,13 +574,18 @@ def as_random_distribution(self): ).rvs elif dist == "lognorm": return get_log_normal( - meanlog=self["meanlog"].as_evaled_expression(), sdlog=self["sdlog"].as_evaled_expression(), + meanlog=self["meanlog"].as_evaled_expression(), + sdlog=self["sdlog"].as_evaled_expression(), ).rvs else: raise NotImplementedError(f"unknown distribution [got: {dist}]") else: # we allow a fixed value specified directly: - return functools.partial(np.random.uniform, self.as_evaled_expression(), self.as_evaled_expression(),) + return functools.partial( + np.random.uniform, + self.as_evaled_expression(), + self.as_evaled_expression(), + ) def list_filenames( @@ -578,9 +600,9 @@ def list_filenames( in the filters will be returned. Args: - folder: + folder: The directory to search for files. Defaults to the current directory. - filters: + filters: A string or a list of strings to filter filenames. Only files containing all the provided substrings will be returned. Defaults to an empty list. @@ -652,14 +674,14 @@ def rolling_mean_pad( [22.6, 23.6, 24.6, 25.6]]) ``` """ - weights = (1. / window) * np.ones(window) + weights = (1.0 / window) * np.ones(window) output = scipy.ndimage.convolve1d(data, weights, axis=0, mode="nearest") if window % 2 == 0: rows, cols = data.shape i = rows - 1 - output[i, :] = 0. + output[i, :] = 0.0 window -= 1 - weight = 1. / window + weight = 1.0 / window for l in range(-((window - 1) // 2), 1 + (window // 2)): i_star = min(max(i + l, 0), i) for j in range(cols): @@ -705,7 +727,12 @@ def bash(command: str) -> str: def create_resume_out_filename( - flepi_run_index: str, flepi_prefix: str, flepi_slot_index: str, flepi_block_index: str, filetype: str, liketype: str + flepi_run_index: str, + flepi_prefix: str, + flepi_slot_index: str, + flepi_block_index: str, + filetype: str, + liketype: str, ) -> str: """ Compiles run output information. @@ -717,10 +744,10 @@ def create_resume_out_filename( flepi_block_index: Index of the block. filetype: File type. liketype: Chimeric or global. - + Returns: The path to a corresponding output file. - + Examples: Generate an output file with specified parameters: >>> filename = create_resume_out_filename( @@ -753,7 +780,11 @@ def create_resume_out_filename( def create_resume_input_filename( - resume_run_index: str, flepi_prefix: str, flepi_slot_index: str, filetype: str, liketype: str + resume_run_index: str, + flepi_prefix: str, + flepi_slot_index: str, + filetype: str, + liketype: str, ) -> str: """ Compiles run input information. @@ -764,10 +795,10 @@ def create_resume_input_filename( flepi_slot_index: Index of the slot. filetype: File type. liketype: Chimeric or global. - + Returns: The path to the a corresponding input file. - + Examples: Generate an input file with specified parameters: >>> filename = create_resume_input_filename( @@ -796,10 +827,12 @@ def create_resume_input_filename( ) -def get_filetype_for_resume(resume_discard_seeding: str, flepi_block_index: str) -> list[str]: +def get_filetype_for_resume( + resume_discard_seeding: str, flepi_block_index: str +) -> list[str]: """ Retrieves a list of parquet file types that are relevant for resuming a process based on - specific environment variable settings. + specific environment variable settings. This function dynamically determines the list based on the current operational context given by the environment. @@ -809,7 +842,7 @@ def get_filetype_for_resume(resume_discard_seeding: str, flepi_block_index: str) Returns: List of file types. - + Examples: Determine file types for block index 1 with seeding data NOT discarded: >>> filetypes = get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") @@ -855,7 +888,7 @@ def create_resume_file_names_map( Returns: A dictionary where keys are input file paths and values are corresponding - output file paths. + output file paths. The mappings depend on: - Parquet file types appropriate for resuming a process, as determined by the environment. @@ -870,7 +903,7 @@ def create_resume_file_names_map( No explicit exceptions are raised within the function, but it relies heavily on external functions and environment variables which if improperly configured could lead to unexpected behavior. - + Examples: Generate a mapping of file names for a given resume process: >>> file_names_map = create_resume_file_names_map( @@ -935,7 +968,7 @@ def download_file_from_s3(name_map: dict[str, str]) -> None: local path, and handles errors if the S3 URI format is incorrect or if the download fails. Args: - name_map: + name_map: A dictionary where keys are S3 URIs (strings) and values are the local file paths (strings) where the files should be saved. @@ -969,10 +1002,12 @@ def download_file_from_s3(name_map: dict[str, str]) -> None: import boto3 from botocore.exceptions import ClientError except ModuleNotFoundError: - raise ModuleNotFoundError(( - "No module named 'boto3', which is required for " - "gempyor.utils.download_file_from_s3. Please install the aws target." - )) + raise ModuleNotFoundError( + ( + "No module named 'boto3', which is required for " + "gempyor.utils.download_file_from_s3. Please install the aws target." + ) + ) s3 = boto3.client("s3") first_output_filename = next(iter(name_map.values())) output_dir = os.path.dirname(first_output_filename) @@ -994,9 +1029,9 @@ def download_file_from_s3(name_map: dict[str, str]) -> None: def move_file_at_local(name_map: dict[str, str]) -> None: """ Moves files locally according to a given mapping. - This function takes a dictionary where the keys are source file paths and - the values are destination file paths. It ensures that the destination - directories exist and then copies the files from the source paths to the + This function takes a dictionary where the keys are source file paths and + the values are destination file paths. It ensures that the destination + directories exist and then copies the files from the source paths to the destination paths. Args: diff --git a/flepimop/gempyor_pkg/tests/file_paths/test_run_id.py b/flepimop/gempyor_pkg/tests/file_paths/test_run_id.py index 2e3f1c79b..cb8513ad6 100644 --- a/flepimop/gempyor_pkg/tests/file_paths/test_run_id.py +++ b/flepimop/gempyor_pkg/tests/file_paths/test_run_id.py @@ -37,9 +37,7 @@ def test_get_run_id_default_timestamp(self) -> None: datetime(2023, 8, 9, 16, 0, 0), ], ) - def test_get_run_id_user_provided_timestamp( - self, timestamp: None | datetime - ) -> None: + def test_get_run_id_user_provided_timestamp(self, timestamp: None | datetime) -> None: # Setup rid = run_id(timestamp=timestamp) diff --git a/flepimop/gempyor_pkg/tests/npi/test_SinglePeriodModifier.py b/flepimop/gempyor_pkg/tests/npi/test_SinglePeriodModifier.py index 7e3c2bc59..704949cfb 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/tests/npi/test_SinglePeriodModifier.py @@ -49,7 +49,9 @@ def test_SinglePeriodModifier_start_date_fail(self): config.clear() config.read(user=False) config.set_file(f"{DATA_DIR}/config_test.yml") - with pytest.raises(ValueError, match=r".*at least one period start or end date is not between.*"): + with pytest.raises( + ValueError, match=r".*at least one period start or end date is not between.*" + ): s = model_info.ModelInfo( setup_name="test_seir", config=config, @@ -72,7 +74,9 @@ def test_SinglePeriodModifier_end_date_fail(self): config.clear() config.read(user=False) config.set_file(f"{DATA_DIR}/config_test.yml") - with pytest.raises(ValueError, match=r".*at least one period start or end date is not between.*"): + with pytest.raises( + ValueError, match=r".*at least one period start or end date is not between.*" + ): s = model_info.ModelInfo( setup_name="test_seir", config=config, diff --git a/flepimop/gempyor_pkg/tests/npi/test_npis.py b/flepimop/gempyor_pkg/tests/npi/test_npis.py index bad306b1e..c2a4047ac 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_npis.py +++ b/flepimop/gempyor_pkg/tests/npi/test_npis.py @@ -47,12 +47,18 @@ def test_full_npis_read_write(): # inference_simulator.s, load_ID=False, sim_id2load=None, config=config # ) - inference_simulator.modinf.write_simID(ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF()) + inference_simulator.modinf.write_simID( + ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF() + ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() hnpi_read["value"] = np.random.random(len(hnpi_read)) * 2 - 1 out_hnpi = pa.Table.from_pandas(hnpi_read, preserve_index=False) - pa.parquet.write_table(out_hnpi, file_paths.create_file_name(105, "", 1, "hnpi", "parquet")) + pa.parquet.write_table( + out_hnpi, file_paths.create_file_name(105, "", 1, "hnpi", "parquet") + ) import random random.seed(10) @@ -74,10 +80,16 @@ def test_full_npis_read_write(): npi_outcomes = outcomes.build_outcome_modifiers( inference_simulator.modinf, load_ID=True, sim_id2load=1, config=config ) - inference_simulator.modinf.write_simID(ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF()) + inference_simulator.modinf.write_simID( + ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF() + ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() # runs with the new, random NPI @@ -98,10 +110,16 @@ def test_full_npis_read_write(): npi_outcomes = outcomes.build_outcome_modifiers( inference_simulator.modinf, load_ID=True, sim_id2load=1, config=config ) - inference_simulator.modinf.write_simID(ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF()) + inference_simulator.modinf.write_simID( + ftype="hnpi", sim_id=1, df=npi_outcomes.getReductionDF() + ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() @@ -116,10 +134,15 @@ def test_spatial_groups(): ) # Test build from config, value of the reduction array - npi = seir.build_npi_SEIR(inference_simulator.modinf, load_ID=False, sim_id2load=None, config=config) + npi = seir.build_npi_SEIR( + inference_simulator.modinf, load_ID=False, sim_id2load=None, config=config + ) # all independent: r1 - assert len(npi.getReduction("r1")["2021-01-01"].unique()) == inference_simulator.modinf.nsubpops + assert ( + len(npi.getReduction("r1")["2021-01-01"].unique()) + == inference_simulator.modinf.nsubpops + ) assert npi.getReduction("r1").isna().sum().sum() == 0 # all the same: r2 @@ -127,7 +150,10 @@ def test_spatial_groups(): assert npi.getReduction("r2").isna().sum().sum() == 0 # two groups: r3 - assert len(npi.getReduction("r3")["2020-04-15"].unique()) == inference_simulator.modinf.nsubpops - 2 + assert ( + len(npi.getReduction("r3")["2020-04-15"].unique()) + == inference_simulator.modinf.nsubpops - 2 + ) assert npi.getReduction("r3").isna().sum().sum() == 0 assert len(npi.getReduction("r3").loc[["01000", "02000"], "2020-04-15"].unique()) == 1 assert len(npi.getReduction("r3").loc[["04000", "06000"], "2020-04-15"].unique()) == 1 @@ -160,7 +186,9 @@ def test_spatial_groups(): # all the same: r2 df = npi_df[npi_df["modifier_name"] == "all_together"] assert len(df) == 1 - assert set(df["subpop"].iloc[0].split(",")) == set(inference_simulator.modinf.subpop_struct.subpop_names) + assert set(df["subpop"].iloc[0].split(",")) == set( + inference_simulator.modinf.subpop_struct.subpop_names + ) assert len(df["subpop"].iloc[0].split(",")) == inference_simulator.modinf.nsubpops # two groups: r3 @@ -175,7 +203,9 @@ def test_spatial_groups(): df = npi_df[npi_df["modifier_name"] == "mt_reduce"] assert len(df) == 4 assert df.subpop.to_list() == ["09000,10000", "02000", "06000", "01000,04000"] - assert df[df["subpop"] == "09000,10000"]["start_date"].iloc[0] == "2020-12-01,2021-12-01" + assert ( + df[df["subpop"] == "09000,10000"]["start_date"].iloc[0] == "2020-12-01,2021-12-01" + ) assert ( df[df["subpop"] == "01000,04000"]["start_date"].iloc[0] == df[df["subpop"] == "06000"]["start_date"].iloc[0] @@ -194,15 +224,21 @@ def test_spatial_groups(): ) # Test build from config, value of the reduction array - npi = seir.build_npi_SEIR(inference_simulator.modinf, load_ID=False, sim_id2load=None, config=config) + npi = seir.build_npi_SEIR( + inference_simulator.modinf, load_ID=False, sim_id2load=None, config=config + ) npi_df = npi.getReductionDF() inference_simulator.modinf.write_simID(ftype="snpi", sim_id=1, df=npi_df) - snpi_read = pq.read_table(f"{config_filepath_prefix}model_output/snpi/000000001.105.snpi.parquet").to_pandas() + snpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/snpi/000000001.105.snpi.parquet" + ).to_pandas() snpi_read["value"] = np.random.random(len(snpi_read)) * 2 - 1 out_snpi = pa.Table.from_pandas(snpi_read, preserve_index=False) - pa.parquet.write_table(out_snpi, file_paths.create_file_name(106, "", 1, "snpi", "parquet")) + pa.parquet.write_table( + out_snpi, file_paths.create_file_name(106, "", 1, "snpi", "parquet") + ) inference_simulator = gempyor.GempyorInference( config_filepath=f"{config_filepath_prefix}config_test_spatial_group_npi.yml", @@ -213,11 +249,19 @@ def test_spatial_groups(): out_run_id=107, ) - npi_seir = seir.build_npi_SEIR(inference_simulator.modinf, load_ID=True, sim_id2load=1, config=config) - inference_simulator.modinf.write_simID(ftype="snpi", sim_id=1, df=npi_seir.getReductionDF()) + npi_seir = seir.build_npi_SEIR( + inference_simulator.modinf, load_ID=True, sim_id2load=1, config=config + ) + inference_simulator.modinf.write_simID( + ftype="snpi", sim_id=1, df=npi_seir.getReductionDF() + ) - snpi_read = pq.read_table(f"{config_filepath_prefix}model_output/snpi/000000001.106.snpi.parquet").to_pandas() - snpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/snpi/000000001.107.snpi.parquet").to_pandas() + snpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/snpi/000000001.106.snpi.parquet" + ).to_pandas() + snpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/snpi/000000001.107.snpi.parquet" + ).to_pandas() # now the order can change, so we need to sort by subpop and start_date snpi_wrote = snpi_wrote.sort_values(by=["subpop", "start_date"]).reset_index(drop=True) @@ -225,10 +269,18 @@ def test_spatial_groups(): assert (snpi_read == snpi_wrote).all().all() npi_read = seir.build_npi_SEIR( - inference_simulator.modinf, load_ID=False, sim_id2load=1, config=config, bypass_DF=snpi_read + inference_simulator.modinf, + load_ID=False, + sim_id2load=1, + config=config, + bypass_DF=snpi_read, ) npi_wrote = seir.build_npi_SEIR( - inference_simulator.modinf, load_ID=False, sim_id2load=1, config=config, bypass_DF=snpi_wrote + inference_simulator.modinf, + load_ID=False, + sim_id2load=1, + config=config, + bypass_DF=snpi_wrote, ) assert (npi_read.getReductionDF() == npi_wrote.getReductionDF()).all().all() diff --git a/flepimop/gempyor_pkg/tests/outcomes/make_seir_test_file.py b/flepimop/gempyor_pkg/tests/outcomes/make_seir_test_file.py index 56df652cf..7a192ce13 100644 --- a/flepimop/gempyor_pkg/tests/outcomes/make_seir_test_file.py +++ b/flepimop/gempyor_pkg/tests/outcomes/make_seir_test_file.py @@ -54,7 +54,9 @@ diffI = np.arange(5) * 2 date_data = datetime.date(2020, 4, 15) for i in range(5): - b.loc[(b["mc_value_type"] == "incidence") & (b["date"] == str(date_data)), subpop[i]] = diffI[i] + b.loc[ + (b["mc_value_type"] == "incidence") & (b["date"] == str(date_data)), subpop[i] + ] = diffI[i] pa_df = pa.Table.from_pandas(b, preserve_index=False) pa.parquet.write_table(pa_df, "new_test_no_vacc.parquet") diff --git a/flepimop/gempyor_pkg/tests/outcomes/test_outcomes.py b/flepimop/gempyor_pkg/tests/outcomes/test_outcomes.py index 35f490920..eba5e9ad3 100644 --- a/flepimop/gempyor_pkg/tests/outcomes/test_outcomes.py +++ b/flepimop/gempyor_pkg/tests/outcomes/test_outcomes.py @@ -40,87 +40,140 @@ def test_outcome(): stoch_traj_flag=False, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False + ) - hosp = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.1.hosp.parquet").to_pandas() + hosp = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.1.hosp.parquet" + ).to_pandas() hosp.set_index("date", drop=True, inplace=True) for i, place in enumerate(subpop): for dt in hosp.index: if dt.date() == date_data: assert hosp[hosp["subpop"] == place]["incidI"][dt] == diffI[i] - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == diffI[i] * 0.01 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 * 0.4 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] + == diffI[i] * 0.01 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 * 0.4 + ) for j in range(7): - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + j)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + 8)] == 0 + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + j) + ] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + 8) + ] + == 0 + ) elif dt.date() < date_data: - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidI"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] + == 0 + ) elif dt.date() > (date_data + datetime.timedelta(7)): assert hosp[hosp["subpop"] == place]["incidH"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidICU"][dt] == 0 - hpar = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.1.hpar.parquet").to_pandas() + hpar = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.1.hpar.parquet" + ).to_pandas() for i, place in enumerate(subpop): assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.1 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 7 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "duration")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "duration") + ]["value"].iloc[0] ) == 7 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.01 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.4 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 0 ) @@ -136,10 +189,16 @@ def test_outcome_modifiers_scenario_with_load(): stoch_traj_flag=False, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False + ) - hpar_config = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.1.hpar.parquet").to_pandas() - hpar_rel = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.2.hpar.parquet").to_pandas() + hpar_config = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.1.hpar.parquet" + ).to_pandas() + hpar_rel = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.2.hpar.parquet" + ).to_pandas() for out in ["incidH", "incidD", "incidICU"]: for i, place in enumerate(subpop): @@ -171,16 +230,30 @@ def test_outcomes_read_write_hpar(): stoch_traj_flag=False, out_run_id=3, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.2.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.3.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.2.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.3.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.2.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.3.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.2.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.3.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.2.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.3.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.2.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.3.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() @@ -201,7 +274,9 @@ def test_multishift_notstochdelays(): [36, 29], ] ) - shifts = np.array([[1, 0], [2, 1], [1, 0], [2, 2], [1, 2], [0, 1], [1, 1], [1, 2], [1, 2], [1, 0]]) + shifts = np.array( + [[1, 0], [2, 1], [1, 0], [2, 2], [1, 2], [0, 1], [1, 1], [1, 2], [1, 2], [1, 0]] + ) expected = np.array( [ [0, 39], @@ -232,87 +307,138 @@ def test_outcomes_npi(): ) outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf) - hosp = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet").to_pandas() + hosp = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet" + ).to_pandas() hosp.set_index("date", drop=True, inplace=True) # same as config.yaml (doubled, then NPI halve it) for i, place in enumerate(subpop): for dt in hosp.index: if dt.date() == date_data: assert hosp[hosp["subpop"] == place]["incidI"][dt] == diffI[i] - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == diffI[i] * 0.01 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 * 0.4 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] + == diffI[i] * 0.01 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 * 0.4 + ) for j in range(7): - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + j)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + 8)] == 0 + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + j) + ] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + 8) + ] + == 0 + ) elif dt.date() < date_data: - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidI"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] + == 0 + ) elif dt.date() > (date_data + datetime.timedelta(7)): assert hosp[hosp["subpop"] == place]["incidH"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidICU"][dt] == 0 - hpar = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet").to_pandas() + hpar = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet" + ).to_pandas() # Doubled everything from previous config.yaml for i, place in enumerate(subpop): assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.1 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 7 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "duration")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "duration") + ]["value"].iloc[0] ) == 7 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.01 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 2 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.4 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 0 * 2 ) @@ -330,17 +456,31 @@ def test_outcomes_read_write_hnpi(): out_run_id=106, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() @@ -356,17 +496,27 @@ def test_outcomes_read_write_hnpi2(): out_run_id=106, ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() hnpi_read["value"] = np.random.random(len(hnpi_read)) * 2 - 1 out_hnpi = pa.Table.from_pandas(hnpi_read, preserve_index=False) - pa.parquet.write_table(out_hnpi, file_paths.create_file_name(105, "", 1, "hnpi", "parquet")) + pa.parquet.write_table( + out_hnpi, file_paths.create_file_name(105, "", 1, "hnpi", "parquet") + ) import random random.seed(10) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() # runs with the new, random NPI @@ -378,16 +528,30 @@ def test_outcomes_read_write_hnpi2(): stoch_traj_flag=False, out_run_id=107, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.107.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.107.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.107.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.107.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() @@ -402,89 +566,142 @@ def test_outcomes_npi_custom_pname(): stoch_traj_flag=False, out_run_id=105, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False, sim_id2load=1 + ) - hosp = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet").to_pandas() + hosp = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet" + ).to_pandas() hosp.set_index("date", drop=True, inplace=True) # same as config.yaml (doubled, then NPI halve it) for i, place in enumerate(subpop): for dt in hosp.index: if dt.date() == date_data: assert hosp[hosp["subpop"] == place]["incidI"][dt] == diffI[i] - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == diffI[i] * 0.01 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == diffI[i] * 0.1 * 0.4 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] + == diffI[i] * 0.01 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == diffI[i] * 0.1 * 0.4 + ) for j in range(7): - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + j)] == diffI[i] * 0.1 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7 + 8)] == 0 + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + j) + ] + == diffI[i] * 0.1 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][ + dt + datetime.timedelta(7 + 8) + ] + == 0 + ) elif dt.date() < date_data: - assert hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidH"][dt + datetime.timedelta(7)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidI"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 - assert hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt + datetime.timedelta(2)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidICU"][dt + datetime.timedelta(7)] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["hosp_curr"][dt + datetime.timedelta(7)] + == 0 + ) elif dt.date() > (date_data + datetime.timedelta(7)): assert hosp[hosp["subpop"] == place]["incidH"][dt] == 0 - assert hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + assert ( + hosp[hosp["subpop"] == place]["incidI"][dt - datetime.timedelta(7)] == 0 + ) + assert ( + hosp[hosp["subpop"] == place]["incidD"][dt - datetime.timedelta(4)] == 0 + ) assert hosp[hosp["subpop"] == place]["incidICU"][dt] == 0 - hpar = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet").to_pandas() + hpar = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet" + ).to_pandas() # Doubled everything from previous config.yaml for i, place in enumerate(subpop): assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.1 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 7 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidH") & (hpar["quantity"] == "duration")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidH") + & (hpar["quantity"] == "duration") + ]["value"].iloc[0] ) == 7 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.01 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidD") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidD") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 2 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "probability")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "probability") + ]["value"].iloc[0] ) == 0.4 * 2 ) assert ( float( - hpar[(hpar["subpop"] == place) & (hpar["outcome"] == "incidICU") & (hpar["quantity"] == "delay")][ - "value" - ].iloc[0] + hpar[ + (hpar["subpop"] == place) + & (hpar["outcome"] == "incidICU") + & (hpar["quantity"] == "delay") + ]["value"].iloc[0] ) == 0 * 2 ) @@ -502,16 +719,30 @@ def test_outcomes_read_write_hnpi_custom_pname(): out_run_id=106, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.105.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.105.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() @@ -520,10 +751,14 @@ def test_outcomes_read_write_hnpi2_custom_pname(): prefix = "" - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() hnpi_read["value"] = np.random.random(len(hnpi_read)) * 2 - 1 out_hnpi = pa.Table.from_pandas(hnpi_read, preserve_index=False) - pa.parquet.write_table(out_hnpi, file_paths.create_file_name(105, prefix, 1, "hnpi", "parquet")) + pa.parquet.write_table( + out_hnpi, file_paths.create_file_name(105, prefix, 1, "hnpi", "parquet") + ) import random random.seed(10) @@ -537,10 +772,16 @@ def test_outcomes_read_write_hnpi2_custom_pname(): out_run_id=106, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.105.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() # runs with the new, random NPI @@ -553,16 +794,30 @@ def test_outcomes_read_write_hnpi2_custom_pname(): out_run_id=107, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.107.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.106.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.107.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.106.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.107.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.107.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.106.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.107.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() @@ -580,7 +835,9 @@ def test_outcomes_pcomp(): ) p_compmult = [1, 3] - seir = pq.read_table(f"{config_filepath_prefix}model_output/seir/000000001.105.seir.parquet").to_pandas() + seir = pq.read_table( + f"{config_filepath_prefix}model_output/seir/000000001.105.seir.parquet" + ).to_pandas() seir2 = seir.copy() seir2["mc_vaccination_stage"] = "first_dose" @@ -591,10 +848,16 @@ def test_outcomes_pcomp(): seir2[pl] = seir2[pl] * p_compmult[1] new_seir = pd.concat([seir, seir2]) out_df = pa.Table.from_pandas(new_seir, preserve_index=False) - pa.parquet.write_table(out_df, file_paths.create_file_name(110, prefix, 1, "seir", "parquet")) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False) + pa.parquet.write_table( + out_df, file_paths.create_file_name(110, prefix, 1, "seir", "parquet") + ) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=False + ) - hosp_f = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.111.hosp.parquet").to_pandas() + hosp_f = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.111.hosp.parquet" + ).to_pandas() hosp_f.set_index("date", drop=True, inplace=True) # same as config.yaml (doubled, then NPI halve it) for k, p_comp in enumerate(["0dose", "1dose"]): @@ -602,42 +865,90 @@ def test_outcomes_pcomp(): for i, place in enumerate(subpop): for dt in hosp.index: if dt.date() == date_data: - assert hosp[hosp["subpop"] == place][f"incidI_{p_comp}"][dt] == diffI[i] * p_compmult[k] assert ( - hosp[hosp["subpop"] == place][f"incidH_{p_comp}"][dt + datetime.timedelta(7)] + hosp[hosp["subpop"] == place][f"incidI_{p_comp}"][dt] + == diffI[i] * p_compmult[k] + ) + assert ( + hosp[hosp["subpop"] == place][f"incidH_{p_comp}"][ + dt + datetime.timedelta(7) + ] - diffI[i] * 0.1 * p_compmult[k] < 1e-8 ) assert ( - hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][dt + datetime.timedelta(2)] + hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][ + dt + datetime.timedelta(2) + ] - diffI[i] * 0.01 * p_compmult[k] < 1e-8 ) assert ( - hosp[hosp["subpop"] == place][f"incidICU_{p_comp}"][dt + datetime.timedelta(7)] + hosp[hosp["subpop"] == place][f"incidICU_{p_comp}"][ + dt + datetime.timedelta(7) + ] - diffI[i] * 0.1 * 0.4 * p_compmult[k] < 1e-8 ) for j in range(7): assert ( - hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][dt + datetime.timedelta(7 + j)] + hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][ + dt + datetime.timedelta(7 + j) + ] - diffI[i] * 0.1 * p_compmult[k] < 1e-8 ) - assert hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][dt + datetime.timedelta(7 + 8)] == 0 + assert ( + hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][ + dt + datetime.timedelta(7 + 8) + ] + == 0 + ) elif dt.date() < date_data: - assert hosp[hosp["subpop"] == place][f"incidH_{p_comp}"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place][f"incidH_{p_comp}"][ + dt + datetime.timedelta(7) + ] + == 0 + ) assert hosp[hosp["subpop"] == place][f"incidI_{p_comp}"][dt] == 0 - assert hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][dt + datetime.timedelta(2)] == 0 - assert hosp[hosp["subpop"] == place][f"incidICU_{p_comp}"][dt + datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][dt + datetime.timedelta(7)] == 0 + assert ( + hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][ + dt + datetime.timedelta(2) + ] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place][f"incidICU_{p_comp}"][ + dt + datetime.timedelta(7) + ] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place][f"incidH_{p_comp}_curr"][ + dt + datetime.timedelta(7) + ] + == 0 + ) elif dt.date() > (date_data + datetime.timedelta(7)): assert hosp[hosp["subpop"] == place][f"incidH_{p_comp}"][dt] == 0 - assert hosp[hosp["subpop"] == place][f"incidI_{p_comp}"][dt - datetime.timedelta(7)] == 0 - assert hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][dt - datetime.timedelta(4)] == 0 + assert ( + hosp[hosp["subpop"] == place][f"incidI_{p_comp}"][ + dt - datetime.timedelta(7) + ] + == 0 + ) + assert ( + hosp[hosp["subpop"] == place][f"incidD_{p_comp}"][ + dt - datetime.timedelta(4) + ] + == 0 + ) assert hosp[hosp["subpop"] == place][f"incidICU_{p_comp}"][dt] == 0 - hpar_f = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.111.hpar.parquet").to_pandas() + hpar_f = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.111.hpar.parquet" + ).to_pandas() # Doubled everything from previous config.yaml # for k, p_comp in enumerate(["unvaccinated", "first_dose"]): for k, p_comp in enumerate(["0dose", "1dose"]): @@ -727,16 +1038,30 @@ def test_outcomes_pcomp_read_write(): out_run_id=112, ) - outcomes.onerun_delayframe_outcomes(sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1) + outcomes.onerun_delayframe_outcomes( + sim_id2write=1, modinf=inference_simulator.modinf, load_ID=True, sim_id2load=1 + ) - hpar_read = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.111.hpar.parquet").to_pandas() - hpar_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hpar/000000001.112.hpar.parquet").to_pandas() + hpar_read = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.111.hpar.parquet" + ).to_pandas() + hpar_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hpar/000000001.112.hpar.parquet" + ).to_pandas() assert (hpar_read == hpar_wrote).all().all() - hnpi_read = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.111.hnpi.parquet").to_pandas() - hnpi_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hnpi/000000001.112.hnpi.parquet").to_pandas() + hnpi_read = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.111.hnpi.parquet" + ).to_pandas() + hnpi_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hnpi/000000001.112.hnpi.parquet" + ).to_pandas() assert (hnpi_read == hnpi_wrote).all().all() - hosp_read = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.111.hosp.parquet").to_pandas() - hosp_wrote = pq.read_table(f"{config_filepath_prefix}model_output/hosp/000000001.112.hosp.parquet").to_pandas() + hosp_read = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.111.hosp.parquet" + ).to_pandas() + hosp_wrote = pq.read_table( + f"{config_filepath_prefix}model_output/hosp/000000001.112.hosp.parquet" + ).to_pandas() assert (hosp_read == hosp_wrote).all().all() diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 09b278914..e18cdaca2 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -316,9 +316,9 @@ def test_parameters_instance_attributes( assert set(params.pdata.keys()) == set(mock_inputs.config.keys()) for param_name, param_conf in mock_inputs.config.items(): assert params.pdata[param_name]["idx"] == params.pnames2pindex[param_name] - assert params.pdata[param_name][ - "stacked_modifier_method" - ] == param_conf.get("stacked_modifier_method", "product") + assert params.pdata[param_name]["stacked_modifier_method"] == param_conf.get( + "stacked_modifier_method", "product" + ) if "timeseries" in param_conf: assert params.pdata[param_name]["ts"].equals( mock_inputs.get_timeseries_df(param_name) @@ -348,8 +348,7 @@ def test_parameters_instance_attributes( }, ) assert ( - params.pdata[param_name]["dist"].__self__.kwds - == expected.__self__.kwds + params.pdata[param_name]["dist"].__self__.kwds == expected.__self__.kwds ) assert ( params.pdata[param_name]["dist"].__self__.support() @@ -446,9 +445,7 @@ def test_get_pnames2pindex( # Assertions assert params.get_pnames2pindex() == params.pnames2pindex - assert params.pnames2pindex == { - p: params.pnames.index(p) for p in params.pnames - } + assert params.pnames2pindex == {p: params.pnames.index(p) for p in params.pnames} @pytest.mark.parametrize( "factory,n_days,nsubpops", diff --git a/flepimop/gempyor_pkg/tests/seir/dev_new_test0.py b/flepimop/gempyor_pkg/tests/seir/dev_new_test0.py index 1e0915b82..53c34e039 100644 --- a/flepimop/gempyor_pkg/tests/seir/dev_new_test0.py +++ b/flepimop/gempyor_pkg/tests/seir/dev_new_test0.py @@ -42,7 +42,9 @@ def test_parameters_from_timeserie_file(): ) # p = inference_simulator.s.parameters - p_draw = p.parameters_quick_draw(n_days=inference_simulator.s.n_days, nnodes=inference_simulator.s.nnodes) + p_draw = p.parameters_quick_draw( + n_days=inference_simulator.s.n_days, nnodes=inference_simulator.s.nnodes + ) p_df = p.getParameterDF(p_draw)["parameter"] diff --git a/flepimop/gempyor_pkg/tests/seir/test_compartments.py b/flepimop/gempyor_pkg/tests/seir/test_compartments.py index 1d4319e3b..492d37400 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_compartments.py +++ b/flepimop/gempyor_pkg/tests/seir/test_compartments.py @@ -24,7 +24,9 @@ def test_check_transitions_parquet_creation(): config.set_file(f"{DATA_DIR}/config_compartmental_model_format.yml") original_compartments_file = f"{DATA_DIR}/parsed_compartment_compartments.parquet" original_transitions_file = f"{DATA_DIR}/parsed_compartment_transitions.parquet" - lhs = compartments.Compartments(seir_config=config["seir"], compartments_config=config["compartments"]) + lhs = compartments.Compartments( + seir_config=config["seir"], compartments_config=config["compartments"] + ) rhs = compartments.Compartments( seir_config=config["seir"], compartments_file=original_compartments_file, @@ -43,10 +45,16 @@ def test_check_transitions_parquet_writing_and_loading(): config.clear() config.read(user=False) config.set_file(f"{DATA_DIR}/config_compartmental_model_format.yml") - lhs = compartments.Compartments(seir_config=config["seir"], compartments_config=config["compartments"]) + lhs = compartments.Compartments( + seir_config=config["seir"], compartments_config=config["compartments"] + ) temp_compartments_file = f"{DATA_DIR}/parsed_compartment_compartments.test.parquet" temp_transitions_file = f"{DATA_DIR}/parsed_compartment_transitions.test.parquet" - lhs.toFile(compartments_file=temp_compartments_file, transitions_file=temp_transitions_file, write_parquet=True) + lhs.toFile( + compartments_file=temp_compartments_file, + transitions_file=temp_transitions_file, + write_parquet=True, + ) rhs = compartments.Compartments( seir_config=config["seir"], compartments_file=temp_compartments_file, @@ -86,4 +94,3 @@ def test_ModelInfo_has_compartments_component(): ) assert type(s.compartments) == compartments.Compartments assert type(s.compartments) == compartments.Compartments - diff --git a/flepimop/gempyor_pkg/tests/seir/test_ic.py b/flepimop/gempyor_pkg/tests/seir/test_ic.py index b4cd240ee..e9c275347 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_ic.py +++ b/flepimop/gempyor_pkg/tests/seir/test_ic.py @@ -21,7 +21,9 @@ def test_IC_success(self): outcome_modifiers_scenario=None, write_csv=False, ) - sic = initial_conditions.InitialConditionsFactory(config=s.initial_conditions_config) + sic = initial_conditions.InitialConditionsFactory( + config=s.initial_conditions_config + ) assert sic.initial_conditions_config == s.initial_conditions_config def test_IC_allow_missing_node_compartments_success(self): @@ -40,7 +42,9 @@ def test_IC_allow_missing_node_compartments_success(self): s.initial_conditions_config["allow_missing_nodes"] = True s.initial_conditions_config["allow_missing_compartments"] = True - sic = initial_conditions.InitialConditionsFactory(config=s.initial_conditions_config) + sic = initial_conditions.InitialConditionsFactory( + config=s.initial_conditions_config + ) sic.get_from_config(sim_id=100, modinf=s) def test_IC_IC_notImplemented_fail(self): @@ -58,6 +62,8 @@ def test_IC_IC_notImplemented_fail(self): write_csv=False, ) s.initial_conditions_config["method"] = "unknown" - sic = initial_conditions.InitialConditionsFactory(config=s.initial_conditions_config) + sic = initial_conditions.InitialConditionsFactory( + config=s.initial_conditions_config + ) sic.get_from_config(sim_id=100, modinf=s) diff --git a/flepimop/gempyor_pkg/tests/seir/test_parameters.py b/flepimop/gempyor_pkg/tests/seir/test_parameters.py index 9e03bf87d..384c3eaa2 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_parameters.py +++ b/flepimop/gempyor_pkg/tests/seir/test_parameters.py @@ -65,7 +65,9 @@ def test_parameters_from_config_plus_read_write(): tf=s.tf, subpop_names=s.subpop_struct.subpop_names, ) - p_load = rhs.parameters_load(param_df=read_df("test_pwrite.parquet"), n_days=n_days, nsubpops=nsubpops) + p_load = rhs.parameters_load( + param_df=read_df("test_pwrite.parquet"), n_days=n_days, nsubpops=nsubpops + ) assert (p_draw == p_load).all() @@ -122,7 +124,10 @@ def test_parameters_quick_draw_old(): assert ((2 <= R0s) & (R0s <= 3)).all() assert sigma.shape == (modinf.n_days, modinf.nsubpops) - assert (sigma == config["seir"]["parameters"]["sigma"]["value"]["value"].as_evaled_expression()).all() + assert ( + sigma + == config["seir"]["parameters"]["sigma"]["value"]["value"].as_evaled_expression() + ).all() assert gamma.shape == (modinf.n_days, modinf.nsubpops) assert len(np.unique(gamma)) == 1 @@ -174,6 +179,8 @@ def test_parameters_from_timeseries_file(): tf=s.tf, subpop_names=s.subpop_struct.subpop_names, ) - p_load = rhs.parameters_load(param_df=read_df("test_pwrite.parquet"), n_days=n_days, nsubpops=nsubpops) + p_load = rhs.parameters_load( + param_df=read_df("test_pwrite.parquet"), n_days=n_days, nsubpops=nsubpops + ) assert (p_draw == p_load).all() diff --git a/flepimop/gempyor_pkg/tests/seir/test_seir.py b/flepimop/gempyor_pkg/tests/seir/test_seir.py index 99a4cc236..d9d6d696f 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_seir.py +++ b/flepimop/gempyor_pkg/tests/seir/test_seir.py @@ -74,7 +74,9 @@ def test_constant_population_legacy_integration(): integration_method = "legacy" seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -82,7 +84,9 @@ def test_constant_population_legacy_integration(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -94,7 +98,9 @@ def test_constant_population_legacy_integration(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) states = seir.steps_SEIR( modinf, @@ -142,8 +148,12 @@ def test_constant_population_rk4jit_integration_fail(): ) modinf.seir_config["integration"]["method"] = "rk4.jit" - seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + seeding_data, seeding_amounts = modinf.seeding.get_from_file( + sim_id=100, modinf=modinf + ) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -151,7 +161,9 @@ def test_constant_population_rk4jit_integration_fail(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -163,7 +175,9 @@ def test_constant_population_rk4jit_integration_fail(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) states = seir.steps_SEIR( modinf, @@ -213,7 +227,9 @@ def test_constant_population_rk4jit_integration(): assert modinf.seir_config["integration"]["method"].get() == "rk4" seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -221,7 +237,9 @@ def test_constant_population_rk4jit_integration(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -233,7 +251,9 @@ def test_constant_population_rk4jit_integration(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) states = seir.steps_SEIR( modinf, parsed_parameters, @@ -281,7 +301,9 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): ) seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -289,7 +311,9 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -301,7 +325,9 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) for i in range(5): states = seir.steps_SEIR( @@ -316,11 +342,15 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): ) df = seir.states2Df(modinf, states) assert ( - df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "10001"] + df[ + (df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R") + ].loc[str(modinf.tf), "10001"] > 1 ) assert ( - df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "20002"] + df[ + (df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R") + ].loc[str(modinf.tf), "20002"] > 1 ) @@ -336,11 +366,23 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): ) df = seir.states2Df(modinf, states) assert ( - df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "20002"] + df[ + (df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R") + ].loc[str(modinf.tf), "20002"] > 1 ) - assert df[(df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1")].max()["20002"] > 0 - assert df[(df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1")].max()["10001"] > 0 + assert ( + df[ + (df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1") + ].max()["20002"] + > 0 + ) + assert ( + df[ + (df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1") + ].max()["10001"] + > 0 + ) def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): @@ -367,7 +409,9 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): ) seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -375,7 +419,9 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -387,7 +433,9 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) for i in range(5): states = seir.steps_SEIR( @@ -402,8 +450,18 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): ) df = seir.states2Df(modinf, states) - assert df[(df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1")].max()["20002"] > 0 - assert df[(df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1")].max()["10001"] > 0 + assert ( + df[ + (df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1") + ].max()["20002"] + > 0 + ) + assert ( + df[ + (df["mc_value_type"] == "incidence") & (df["mc_infection_stage"] == "I1") + ].max()["10001"] + > 0 + ) def test_steps_SEIR_no_spread(): @@ -427,7 +485,9 @@ def test_steps_SEIR_no_spread(): ) seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) modinf.mobility.data = modinf.mobility.data * 0 @@ -437,7 +497,9 @@ def test_steps_SEIR_no_spread(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -449,7 +511,9 @@ def test_steps_SEIR_no_spread(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) for i in range(10): states = seir.steps_SEIR( @@ -464,7 +528,9 @@ def test_steps_SEIR_no_spread(): ) df = seir.states2Df(modinf, states) assert ( - df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "20002"] + df[ + (df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R") + ].loc[str(modinf.tf), "20002"] == 0.0 ) @@ -480,7 +546,9 @@ def test_steps_SEIR_no_spread(): ) df = seir.states2Df(modinf, states) assert ( - df[(df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R")].loc[str(modinf.tf), "20002"] + df[ + (df["mc_value_type"] == "prevalence") & (df["mc_infection_stage"] == "R") + ].loc[str(modinf.tf), "20002"] == 0.0 ) @@ -515,7 +583,9 @@ def test_continuation_resume(): seir.onerun_SEIR(sim_id2write=int(sim_id2write), modinf=modinf, config=config) states_old = pq.read_table( - file_paths.create_file_name(modinf.in_run_id, modinf.in_prefix, 100, "seir", "parquet"), + file_paths.create_file_name( + modinf.in_run_id, modinf.in_prefix, 100, "seir", "parquet" + ), ).to_pandas() states_old = states_old[states_old["date"] == "2020-03-15"].reset_index(drop=True) @@ -547,7 +617,9 @@ def test_continuation_resume(): seir.onerun_SEIR(sim_id2write=sim_id2write, modinf=modinf, config=config) states_new = pq.read_table( - file_paths.create_file_name(modinf.in_run_id, modinf.in_prefix, sim_id2write, "seir", "parquet"), + file_paths.create_file_name( + modinf.in_run_id, modinf.in_prefix, sim_id2write, "seir", "parquet" + ), ).to_pandas() states_new = states_new[states_new["date"] == "2020-03-15"].reset_index(drop=True) assert ( @@ -560,10 +632,16 @@ def test_continuation_resume(): ) seir.onerun_SEIR( - sim_id2write=sim_id2write + 1, modinf=modinf, sim_id2load=sim_id2write, load_ID=True, config=config + sim_id2write=sim_id2write + 1, + modinf=modinf, + sim_id2load=sim_id2write, + load_ID=True, + config=config, ) states_new = pq.read_table( - file_paths.create_file_name(modinf.in_run_id, modinf.in_prefix, sim_id2write + 1, "seir", "parquet"), + file_paths.create_file_name( + modinf.in_run_id, modinf.in_prefix, sim_id2write + 1, "seir", "parquet" + ), ).to_pandas() states_new = states_new[states_new["date"] == "2020-03-15"].reset_index(drop=True) for path in ["model_output/seir", "model_output/snpi", "model_output/spar"]: @@ -604,7 +682,9 @@ def test_inference_resume(): seir.onerun_SEIR(sim_id2write=int(sim_id2write), modinf=modinf, config=config) npis_old = pq.read_table( - file_paths.create_file_name(modinf.in_run_id, modinf.in_prefix, sim_id2write, "snpi", "parquet") + file_paths.create_file_name( + modinf.in_run_id, modinf.in_prefix, sim_id2write, "snpi", "parquet" + ) ).to_pandas() config.clear() @@ -632,10 +712,16 @@ def test_inference_resume(): ) seir.onerun_SEIR( - sim_id2write=sim_id2write + 1, modinf=modinf, sim_id2load=sim_id2write, load_ID=True, config=config + sim_id2write=sim_id2write + 1, + modinf=modinf, + sim_id2load=sim_id2write, + load_ID=True, + config=config, ) npis_new = pq.read_table( - file_paths.create_file_name(modinf.in_run_id, modinf.in_prefix, sim_id2write + 1, "snpi", "parquet") + file_paths.create_file_name( + modinf.in_run_id, modinf.in_prefix, sim_id2write + 1, "snpi", "parquet" + ) ).to_pandas() assert npis_old["modifier_name"].isin(["None", "Wuhan", "KansasCity"]).all() @@ -675,7 +761,9 @@ def test_parallel_compartments_with_vacc(): ) seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -683,7 +771,9 @@ def test_parallel_compartments_with_vacc(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -695,7 +785,9 @@ def test_parallel_compartments_with_vacc(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) for i in range(5): states = seir.steps_SEIR( @@ -762,7 +854,9 @@ def test_parallel_compartments_no_vacc(): ) seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id=100, modinf=modinf) - initial_conditions = modinf.initial_conditions.get_from_config(sim_id=100, modinf=modinf) + initial_conditions = modinf.initial_conditions.get_from_config( + sim_id=100, modinf=modinf + ) npi = NPI.NPIBase.execute( npi_config=modinf.npi_config_seir, @@ -770,7 +864,9 @@ def test_parallel_compartments_no_vacc(): modifiers_library=modinf.seir_modifiers_library, subpops=modinf.subpop_struct.subpop_names, pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"], - pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method["reduction_product"], + pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[ + "reduction_product" + ], ) params = modinf.parameters.parameters_quick_draw(modinf.n_days, modinf.nsubpops) @@ -782,7 +878,9 @@ def test_parallel_compartments_no_vacc(): proportion_array, proportion_info, ) = modinf.compartments.get_transition_array() - parsed_parameters = modinf.compartments.parse_parameters(params, modinf.parameters.pnames, unique_strings) + parsed_parameters = modinf.compartments.parse_parameters( + params, modinf.parameters.pnames, unique_strings + ) for i in range(5): states = seir.steps_SEIR( diff --git a/flepimop/gempyor_pkg/tests/seir/test_subpopulationstructure.py b/flepimop/gempyor_pkg/tests/seir/test_subpopulationstructure.py index 34df630c3..b2161cc9b 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_subpopulationstructure.py +++ b/flepimop/gempyor_pkg/tests/seir/test_subpopulationstructure.py @@ -135,7 +135,9 @@ def test_subpopulation_structure_mobility_shape_fail(): temp_file.close() # Ensure the file is closed config.set_file(temp_file.name) # Load from the temporary file path - with pytest.raises(ValueError, match=r"mobility data must have dimensions of length of geodata.*"): + with pytest.raises( + ValueError, match=r"mobility data must have dimensions of length of geodata.*" + ): subpop_struct = subpopulation_structure.SubpopulationStructure( setup_name=TEST_SETUP_NAME, subpop_config=config["subpop_setup"] ) @@ -155,7 +157,9 @@ def test_subpopulation_structure_mobility_fluxes_same_ori_and_dest_fail(): temp_file.close() # Ensure the file is closed config.set_file(temp_file.name) # Load from the temporary file path - with pytest.raises(ValueError, match=r"Mobility fluxes with same origin and destination.*"): + with pytest.raises( + ValueError, match=r"Mobility fluxes with same origin and destination.*" + ): subpop_struct = subpopulation_structure.SubpopulationStructure( setup_name=TEST_SETUP_NAME, subpop_config=config["subpop_setup"] ) @@ -175,7 +179,9 @@ def test_subpopulation_structure_mobility_npz_shape_fail(): temp_file.close() # Ensure the file is closed config.set_file(temp_file.name) # Load from the temporary file path - with pytest.raises(ValueError, match=r"mobility data must have dimensions of length of geodata.*"): + with pytest.raises( + ValueError, match=r"mobility data must have dimensions of length of geodata.*" + ): subpop_struct = subpopulation_structure.SubpopulationStructure( setup_name=TEST_SETUP_NAME, subpop_config=config["subpop_setup"] ) @@ -216,7 +222,8 @@ def test_subpopulation_structure_mobility_exceed_source_node_pop_fail(): config.set_file(temp_file.name) # Load from the temporary file path with pytest.raises( - ValueError, match=r"The following entries in the mobility data exceed the source subpop populations.*" + ValueError, + match=r"The following entries in the mobility data exceed the source subpop populations.*", ): subpop_struct = subpopulation_structure.SubpopulationStructure( setup_name=TEST_SETUP_NAME, subpop_config=config["subpop_setup"] @@ -238,7 +245,8 @@ def test_subpopulation_structure_mobility_rows_exceed_source_node_pop_fail(): config.set_file(temp_file.name) # Load from the temporary file path with pytest.raises( - ValueError, match=r"The following entries in the mobility data exceed the source subpop populations.*" + ValueError, + match=r"The following entries in the mobility data exceed the source subpop populations.*", ): subpop_struct = subpopulation_structure.SubpopulationStructure( setup_name=TEST_SETUP_NAME, subpop_config=config["subpop_setup"] diff --git a/flepimop/gempyor_pkg/tests/statistics/test_statistic_class.py b/flepimop/gempyor_pkg/tests/statistics/test_statistic_class.py index e18861e9e..4843986c1 100644 --- a/flepimop/gempyor_pkg/tests/statistics/test_statistic_class.py +++ b/flepimop/gempyor_pkg/tests/statistics/test_statistic_class.py @@ -445,9 +445,7 @@ def test_apply_transforms(self, factory: Callable[[], MockStatisticInput]) -> No ) if (scale_func := mock_inputs.config.get("scale")) is not None: # Scale config - expected_transformed_data = getattr(np, scale_func)( - expected_transformed_data - ) + expected_transformed_data = getattr(np, scale_func)(expected_transformed_data) assert transformed_data.identical(expected_transformed_data) @pytest.mark.parametrize("factory", all_valid_factories) @@ -464,8 +462,7 @@ def test_llik(self, factory: Callable[[], MockStatisticInput]) -> None: assert isinstance(log_likelihood, xr.DataArray) assert ( - log_likelihood.dims - == mock_inputs.gt_data[mock_inputs.config["data_var"]].dims + log_likelihood.dims == mock_inputs.gt_data[mock_inputs.config["data_var"]].dims ) assert log_likelihood.coords.identical( mock_inputs.gt_data[mock_inputs.config["data_var"]].coords @@ -523,9 +520,7 @@ def test_compute_logloss_data_misshape_value_error( mock_inputs = factory() statistic = mock_inputs.create_statistic_instance() - model_rows, model_cols = mock_inputs.model_data[ - mock_inputs.config["sim_var"] - ].shape + model_rows, model_cols = mock_inputs.model_data[mock_inputs.config["sim_var"]].shape gt_rows, gt_cols = mock_inputs.gt_data[mock_inputs.config["data_var"]].shape expected_match = ( rf"^{mock_inputs.name} Statistic error\: data and groundtruth do not have " diff --git a/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py b/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py index 367a7f550..b12989ec4 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py +++ b/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py @@ -7,7 +7,7 @@ class TestGetLogNormal: """Unit tests for the `gempyor.utils.get_log_normal` function.""" - + @pytest.mark.parametrize( "meanlog,sdlog", [ @@ -22,15 +22,15 @@ class TestGetLogNormal: ], ) def test_construct_distribution( - self, - meanlog: float | int, - sdlog: float | int, + self, + meanlog: float | int, + sdlog: float | int, ) -> None: """Test the construction of a log normal distribution. - This test checks whether the `get_log_normal` function correctly constructs - a log normal distribution with the specified parameters. It verifies that - the returned object is an instance of `rv_frozen`, and that its support and + This test checks whether the `get_log_normal` function correctly constructs + a log normal distribution with the specified parameters. It verifies that + the returned object is an instance of `rv_frozen`, and that its support and parameters (log mean and log standard deviation) are correctly set. Args: diff --git a/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py b/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py index 23e4fad58..c3ccbca79 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py +++ b/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py @@ -7,7 +7,7 @@ class TestGetTruncatedNormal: """Unit tests for the `gempyor.utils.get_truncated_normal` function.""" - + @pytest.mark.parametrize( "mean,sd,a,b", [ @@ -21,10 +21,10 @@ class TestGetTruncatedNormal: ], ) def test_construct_distribution( - self, - mean: float | int, - sd: float | int, - a: float | int, + self, + mean: float | int, + sd: float | int, + a: float | int, b: float | int, ) -> None: """Test the construction of a truncated normal distribution. diff --git a/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py b/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py index 64efc98b1..4ca9c0ac7 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py +++ b/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py @@ -51,9 +51,7 @@ def test_binomial_p_value_error(self, p: float) -> None: def test_output_validation(self, distribution: str, kwargs: dict[str, Any]) -> None: actual = random_distribution_sampler(distribution, **kwargs) if distribution == "fixed": - expected = partial( - np.random.uniform, kwargs.get("value"), kwargs.get("value") - ) + expected = partial(np.random.uniform, kwargs.get("value"), kwargs.get("value")) assert partials_are_similar(actual, expected) elif distribution == "uniform": expected = partial(np.random.uniform, kwargs.get("low"), kwargs.get("high")) @@ -68,12 +66,12 @@ def test_output_validation(self, distribution: str, kwargs: dict[str, Any]) -> N assert inspect.ismethod(actual) assert actual.__self__.kwds.get("loc") == kwargs.get("mean") assert actual.__self__.kwds.get("scale") == kwargs.get("sd") - assert actual.__self__.a == ( - kwargs.get("a") - kwargs.get("mean") - ) / kwargs.get("sd") - assert actual.__self__.b == ( - kwargs.get("b") - kwargs.get("mean") - ) / kwargs.get("sd") + assert actual.__self__.a == (kwargs.get("a") - kwargs.get("mean")) / kwargs.get( + "sd" + ) + assert actual.__self__.b == (kwargs.get("b") - kwargs.get("mean")) / kwargs.get( + "sd" + ) elif distribution == "lognorm": assert inspect.ismethod(actual) assert actual.__self__.kwds.get("s") == kwargs.get("sdlog") diff --git a/flepimop/gempyor_pkg/tests/utils/test_read_df.py b/flepimop/gempyor_pkg/tests/utils/test_read_df.py index 7a0a0c581..48f05dbd5 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_read_df.py +++ b/flepimop/gempyor_pkg/tests/utils/test_read_df.py @@ -137,8 +137,7 @@ def test_subpop_is_cast_as_str(self) -> None: temp_path = Path(temp_file.name) assert temp_path.stat().st_size == 0 assert ( - self.subpop_df.to_parquet(temp_path, engine="pyarrow", index=False) - is None + self.subpop_df.to_parquet(temp_path, engine="pyarrow", index=False) is None ) assert temp_path.stat().st_size > 0 test_df = read_df(fname=temp_path) diff --git a/flepimop/gempyor_pkg/tests/utils/test_rolling_mean_pad.py b/flepimop/gempyor_pkg/tests/utils/test_rolling_mean_pad.py index 94be3394a..665f34159 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_rolling_mean_pad.py +++ b/flepimop/gempyor_pkg/tests/utils/test_rolling_mean_pad.py @@ -112,9 +112,7 @@ def test_rolling_mean_pad( rolling_mean_data = rolling_mean_pad(test_data, window) rolling_mean_reference = self._rolling_mean_pad_reference(test_data, window) assert rolling_mean_data.shape == expected_shape - assert np.isclose( - rolling_mean_data, rolling_mean_reference, equal_nan=True - ).all() + assert np.isclose(rolling_mean_data, rolling_mean_reference, equal_nan=True).all() def _rolling_mean_pad_reference( self, data: npt.NDArray[np.number], window: int diff --git a/flepimop/gempyor_pkg/tests/utils/test_utils.py b/flepimop/gempyor_pkg/tests/utils/test_utils.py index 768451b2a..058e2401b 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_utils.py +++ b/flepimop/gempyor_pkg/tests/utils/test_utils.py @@ -10,7 +10,11 @@ @pytest.mark.parametrize( - ("fname", "extension"), [("mobility", "csv"), ("usa-geoid-params-output", "parquet"),], + ("fname", "extension"), + [ + ("mobility", "csv"), + ("usa-geoid-params-output", "parquet"), + ], ) def test_read_df_and_write_success(fname, extension): os.chdir(tmp_path) @@ -29,7 +33,9 @@ def test_read_df_and_write_success(fname, extension): assert os.path.isfile(tmp_path + "/data/" + fname + "." + extension) -@pytest.mark.parametrize(("fname", "extension"), [("mobility", "csv"), ("usa-geoid-params-output", "parquet")]) +@pytest.mark.parametrize( + ("fname", "extension"), [("mobility", "csv"), ("usa-geoid-params-output", "parquet")] +) def test_read_df_and_write_fail(fname, extension): with pytest.raises(NotImplementedError, match=r".*Invalid.*extension.*Must.*"): os.chdir(tmp_path) @@ -41,7 +47,9 @@ def test_read_df_and_write_fail(fname, extension): assert df2.equals(df1) utils.write_df(tmp_path + "/data/" + fname, df2, extension="") elif extension == "parquet": - df2 = pa.parquet.read_table(f"{DATA_DIR}/" + fname + "." + extension).to_pandas() + df2 = pa.parquet.read_table( + f"{DATA_DIR}/" + fname + "." + extension + ).to_pandas() assert df2.equals(df1) utils.write_df(tmp_path + "/data/" + fname, df2, extension="") @@ -91,9 +99,7 @@ def test_create_resume_out_filename(): filetype="spar", liketype="global", ) - expected_filename = ( - "model_output/output/123/spar/global/intermediate/000000002.000000001.000000001.123.spar.parquet" - ) + expected_filename = "model_output/output/123/spar/global/intermediate/000000002.000000001.000000001.123.spar.parquet" assert result == expected_filename result2 = utils.create_resume_out_filename( @@ -111,14 +117,22 @@ def test_create_resume_out_filename(): def test_create_resume_input_filename(): result = utils.create_resume_input_filename( - flepi_slot_index="2", resume_run_index="321", flepi_prefix="output", filetype="spar", liketype="global" + flepi_slot_index="2", + resume_run_index="321", + flepi_prefix="output", + filetype="spar", + liketype="global", ) expect_filename = "model_output/output/321/spar/global/final/000000002.321.spar.parquet" assert result == expect_filename result2 = utils.create_resume_input_filename( - flepi_slot_index="2", resume_run_index="321", flepi_prefix="output", filetype="seed", liketype="chimeric" + flepi_slot_index="2", + resume_run_index="321", + flepi_prefix="output", + filetype="seed", + liketype="chimeric", ) expect_filename2 = "model_output/output/321/seed/chimeric/final/000000002.321.seed.csv" assert result2 == expect_filename2 @@ -126,17 +140,26 @@ def test_create_resume_input_filename(): def test_get_filetype_resume_discard_seeding_true_flepi_block_index_1(): expected_types = ["spar", "snpi", "hpar", "hnpi", "init"] - assert utils.get_filetype_for_resume(resume_discard_seeding="true", flepi_block_index="1") == expected_types + assert ( + utils.get_filetype_for_resume(resume_discard_seeding="true", flepi_block_index="1") + == expected_types + ) def test_get_filetype_resume_discard_seeding_false_flepi_block_index_1(): expected_types = ["seed", "spar", "snpi", "hpar", "hnpi", "init"] - assert utils.get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") == expected_types + assert ( + utils.get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") + == expected_types + ) def test_get_filetype_flepi_block_index_2(): expected_types = ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] - assert utils.get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="2") == expected_types + assert ( + utils.get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="2") + == expected_types + ) def test_create_resume_file_names_map(): diff --git a/flepimop/gempyor_pkg/tests/utils/test_utils2.py b/flepimop/gempyor_pkg/tests/utils/test_utils2.py index 4b0ae59ba..0822604ed 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_utils2.py +++ b/flepimop/gempyor_pkg/tests/utils/test_utils2.py @@ -26,7 +26,9 @@ class SampleClass: def __init__(self): self.value = 11 - @utils.profile(output_file="get_value.prof", sort_by="time", lines_to_print=10, strip_dirs=True) + @utils.profile( + output_file="get_value.prof", sort_by="time", lines_to_print=10, strip_dirs=True + ) def get_value(self): return self.value @@ -198,7 +200,9 @@ def test_as_random_distribution_binomial_w_fraction_error(config): def test_as_random_distribution_truncnorm(config): - config.add({"value": {"distribution": "truncnorm", "mean": 0, "sd": 1, "a": -1, "b": 1}}) + config.add( + {"value": {"distribution": "truncnorm", "mean": 0, "sd": 1, "a": -1, "b": 1}} + ) dist = config["value"].as_random_distribution() rvs = dist(size=1000) assert len(rvs) == 1000 diff --git a/postprocessing/postprocess_auto.py b/postprocessing/postprocess_auto.py index aaf4a0bff..7cb5edfcf 100644 --- a/postprocessing/postprocess_auto.py +++ b/postprocessing/postprocess_auto.py @@ -26,7 +26,9 @@ def __init__(self, run_id, config_filepath=None, folder_path=None): self.folder_path = folder_path -def get_all_filenames(file_type, all_runs, finals_only=False, intermediates_only=False, ignore_chimeric=True) -> dict: +def get_all_filenames( + file_type, all_runs, finals_only=False, intermediates_only=False, ignore_chimeric=True +) -> dict: """ return dictionanary for each run name """ @@ -159,7 +161,14 @@ def slack_multiple_files_v2(slack_token, message, file_list, channel): help="Maximum number of files to load for in depth plot and individual sim plot", ) def generate_pdf( - config_filepath, run_id, job_name, fs_results_path, slack_token, slack_channel, max_files, max_files_deep + config_filepath, + run_id, + job_name, + fs_results_path, + slack_token, + slack_channel, + max_files, + max_files_deep, ): print("Generating plots") print(f">> config {config_filepath} for run_id {run_id}") @@ -217,7 +226,9 @@ def generate_pdf( for filename in file_list: slot = int(filename.split("/")[-1].split(".")[0]) block = int(filename.split("/")[-1].split(".")[1]) - sim_str = filename.split("/")[-1].split(".")[2] # not necessarily a sim number now + sim_str = filename.split("/")[-1].split(".")[ + 2 + ] # not necessarily a sim number now if sim_str.isdigit(): sim = int(sim_str) if block == 1 and (sim == 1 or sim % 5 == 0): ## first block, only one @@ -238,7 +249,9 @@ def generate_pdf( # In[23]: - fig, axes = plt.subplots(len(node_names) + 1, 4, figsize=(4 * 4, len(node_names) * 3), sharex=True) + fig, axes = plt.subplots( + len(node_names) + 1, 4, figsize=(4 * 4, len(node_names) * 3), sharex=True + ) colors = ["b", "r", "y", "c"] icl = 0 @@ -255,32 +268,60 @@ def generate_pdf( lls = lls.cumsum() feature = "accepts, cumulative" axes[idp, ift].fill_between( - lls.index, lls.quantile(0.025, axis=1), lls.quantile(0.975, axis=1), alpha=0.1, color=colors[icl] + lls.index, + lls.quantile(0.025, axis=1), + lls.quantile(0.975, axis=1), + alpha=0.1, + color=colors[icl], ) axes[idp, ift].fill_between( - lls.index, lls.quantile(0.25, axis=1), lls.quantile(0.75, axis=1), alpha=0.1, color=colors[icl] + lls.index, + lls.quantile(0.25, axis=1), + lls.quantile(0.75, axis=1), + alpha=0.1, + color=colors[icl], + ) + axes[idp, ift].plot( + lls.index, lls.median(axis=1), marker="o", label=run_id, color=colors[icl] ) - axes[idp, ift].plot(lls.index, lls.median(axis=1), marker="o", label=run_id, color=colors[icl]) axes[idp, ift].plot(lls.index, lls.iloc[:, 0:max_files_deep], color="k", lw=0.3) axes[idp, ift].set_title(f"National, {feature}") axes[idp, ift].grid() for idp, nn in enumerate(node_names): idp = idp + 1 - all_nn = full_df[full_df["subpop"] == nn][["sim", "slot", "ll", "accept", "accept_avg", "accept_prob"]] + all_nn = full_df[full_df["subpop"] == nn][ + ["sim", "slot", "ll", "accept", "accept_avg", "accept_prob"] + ] for ift, feature in enumerate(["ll", "accept", "accept_avg", "accept_prob"]): lls = all_nn.pivot(index="sim", columns="slot", values=feature) if feature == "accept": lls = lls.cumsum() feature = "accepts, cumulative" axes[idp, ift].fill_between( - lls.index, lls.quantile(0.025, axis=1), lls.quantile(0.975, axis=1), alpha=0.1, color=colors[icl] + lls.index, + lls.quantile(0.025, axis=1), + lls.quantile(0.975, axis=1), + alpha=0.1, + color=colors[icl], ) axes[idp, ift].fill_between( - lls.index, lls.quantile(0.25, axis=1), lls.quantile(0.75, axis=1), alpha=0.1, color=colors[icl] + lls.index, + lls.quantile(0.25, axis=1), + lls.quantile(0.75, axis=1), + alpha=0.1, + color=colors[icl], + ) + axes[idp, ift].plot( + lls.index, + lls.median(axis=1), + marker="o", + label=run_id, + color=colors[icl], + ) + axes[idp, ift].plot( + lls.index, lls.iloc[:, 0:max_files_deep], color="k", lw=0.3 ) - axes[idp, ift].plot(lls.index, lls.median(axis=1), marker="o", label=run_id, color=colors[icl]) - axes[idp, ift].plot(lls.index, lls.iloc[:, 0:max_files_deep], color="k", lw=0.3) axes[idp, ift].set_title(f"{nn}, {feature}") axes[idp, ift].grid() if idp == len(node_names) - 1: @@ -292,8 +333,11 @@ def generate_pdf( pass import gempyor.utils - llik_filenames = gempyor.utils.list_filenames(folder="model_output/", filters=["final", "llik" , ".parquet"]) - #get_all_filenames("llik", fs_results_path, finals_only=True, intermediates_only=False) + + llik_filenames = gempyor.utils.list_filenames( + folder="model_output/", filters=["final", "llik", ".parquet"] + ) + # get_all_filenames("llik", fs_results_path, finals_only=True, intermediates_only=False) # In[7]: resultST = [] for filename in llik_filenames: diff --git a/utilities/clean_s3.py b/utilities/clean_s3.py index 2998c65c2..08982b2b2 100644 --- a/utilities/clean_s3.py +++ b/utilities/clean_s3.py @@ -9,7 +9,9 @@ s3 = boto3.client("s3") paginator = s3.get_paginator("list_objects_v2") -pages = paginator.paginate(Bucket=bucket, Prefix="", Delimiter="/") # needs paginator cause more than 1000 files +pages = paginator.paginate( + Bucket=bucket, Prefix="", Delimiter="/" +) # needs paginator cause more than 1000 files to_prun = [] # folders: diff --git a/utilities/copy_for_continuation.py b/utilities/copy_for_continuation.py index 33d9da40f..05b7a803a 100644 --- a/utilities/copy_for_continuation.py +++ b/utilities/copy_for_continuation.py @@ -77,10 +77,14 @@ def detect_old_run_id(fp): fn = files[0] old_run_id = detect_old_run_id(fn) new_name = ( - fn.replace("seir", "cont").replace(f"{input_folder}/model_output", "model_output").replace(old_run_id, run_id) + fn.replace("seir", "cont") + .replace(f"{input_folder}/model_output", "model_output") + .replace(old_run_id, run_id) ) - print(f"detected old_run_id: {old_run_id} which will be replaced by user provided run_id: {run_id}") + print( + f"detected old_run_id: {old_run_id} which will be replaced by user provided run_id: {run_id}" + ) empty_str = "°" * len(input_folder) print(f"file: \n OLD NAME: {fn}\n NEW NAME: {empty_str}{new_name}") for fn in tqdm.tqdm(files): diff --git a/utilities/prune_by_llik.py b/utilities/prune_by_llik.py index 5b1f3224b..78ef6c8af 100644 --- a/utilities/prune_by_llik.py +++ b/utilities/prune_by_llik.py @@ -11,7 +11,11 @@ def get_all_filenames( - file_type, fs_results_path="to_prune/", finals_only=False, intermediates_only=True, ignore_chimeric=True + file_type, + fs_results_path="to_prune/", + finals_only=False, + intermediates_only=True, + ignore_chimeric=True, ) -> dict: """ return dictionary for each run name @@ -113,14 +117,18 @@ def get_all_filenames( if fill_missing: # Extract the numbers from the filenames numbers = [int(os.path.basename(filename).split(".")[0]) for filename in all_files] - missing_numbers = [num for num in range(fill_from_min, fill_from_max + 1) if num not in numbers] + missing_numbers = [ + num for num in range(fill_from_min, fill_from_max + 1) if num not in numbers + ] if missing_numbers: missing_filenames = [] for num in missing_numbers: filename = os.path.basename(all_files[0]) filename_prefix = re.search(r"^.*?(\d+)", filename).group() filename_suffix = re.search(r"(\..*?)$", filename).group() - missing_filename = os.path.join(os.path.dirname(all_files[0]), f"{num:09d}{filename_suffix}") + missing_filename = os.path.join( + os.path.dirname(all_files[0]), f"{num:09d}{filename_suffix}" + ) missing_filenames.append(missing_filename) print("The missing filenames with full paths are:") for missing_filename in missing_filenames: @@ -143,7 +151,7 @@ def copy_path(src, dst): file_types = [ "llik", - #"seed", + # "seed", "init", "snpi", "hnpi", diff --git a/utilities/prune_by_llik_and_proj.py b/utilities/prune_by_llik_and_proj.py index 53e623224..585473bf7 100644 --- a/utilities/prune_by_llik_and_proj.py +++ b/utilities/prune_by_llik_and_proj.py @@ -11,7 +11,11 @@ def get_all_filenames( - file_type, fs_results_path="to_prune/", finals_only=False, intermediates_only=True, ignore_chimeric=True + file_type, + fs_results_path="to_prune/", + finals_only=False, + intermediates_only=True, + ignore_chimeric=True, ) -> dict: """ return dictionary for each run name @@ -23,7 +27,7 @@ def get_all_filenames( l = [] for f in Path(str(fs_results_path + "model_output")).rglob(f"*.{ext}"): f = str(f) - + if file_type in f: print(f) if ( @@ -61,7 +65,9 @@ def get_all_filenames( fs_results_path = "to_prune/" best_n = 200 -llik_filenames = get_all_filenames("llik", fs_results_path, finals_only=True, intermediates_only=False) +llik_filenames = get_all_filenames( + "llik", fs_results_path, finals_only=True, intermediates_only=False +) # In[7]: resultST = [] for filename in llik_filenames: @@ -100,7 +106,6 @@ def get_all_filenames( print(f" - {slot:4}, llik: {sorted_llik.loc[slot]['ll']:0.3f}") - #### RERUN FROM HERE TO CHANGE THE REGULARIZATION files_to_keep = list(full_df.loc[best_slots]["filename"].unique()) # important to sort by llik @@ -109,8 +114,9 @@ def get_all_filenames( files_to_keep = [] for fn in all_files: if fn in files_to_keep3: - outcome_fn = fn.replace("llik", "hosp") + outcome_fn = fn.replace("llik", "hosp") import gempyor.utils + outcomes_df = gempyor.utils.read_df(outcome_fn) outcomes_df = outcomes_df.set_index("date") reg = 1.5 @@ -118,19 +124,27 @@ def get_all_filenames( this_bad = 0 bad_subpops = [] for sp in outcomes_df["subpop"].unique(): - max_fit = outcomes_df[outcomes_df["subpop"]==sp]["incidC"][:"2024-04-08"].max() - max_summer = outcomes_df[outcomes_df["subpop"]==sp]["incidC"]["2024-04-08":"2024-09-30"].max() - if max_summer > max_fit*reg: + max_fit = outcomes_df[outcomes_df["subpop"] == sp]["incidC"][ + :"2024-04-08" + ].max() + max_summer = outcomes_df[outcomes_df["subpop"] == sp]["incidC"][ + "2024-04-08":"2024-09-30" + ].max() + if max_summer > max_fit * reg: this_bad += 1 - max_reg = max(max_reg, max_summer/max_fit) + max_reg = max(max_reg, max_summer / max_fit) bad_subpops.append(sp) - #print(f"changing {sp} because max_summer max_summer={max_summer:.1f} > reg*max_fit={max_fit:.1f}, diff {max_fit/max_summer*100:.1f}%") - #print(f">>> MULT BY {max_summer/max_fit*mult:2f}") - #outcomes_df.loc[outcomes_df["subpop"]==sp, ["incidH", "incidD"]] = outcomes_df.loc[outcomes_df["subpop"]==sp, ["incidH", "incidD"]]*max_summer/max_fit*mult - if this_bad>4 or max_reg>4: - print(f"{outcome_fn.split('/')[-1].split('.')[0]} >>> BAAD: {this_bad} subpops AND max_ratio={max_reg:.1f}, sp with max_summer > max_fit*{reg} {bad_subpops}") + # print(f"changing {sp} because max_summer max_summer={max_summer:.1f} > reg*max_fit={max_fit:.1f}, diff {max_fit/max_summer*100:.1f}%") + # print(f">>> MULT BY {max_summer/max_fit*mult:2f}") + # outcomes_df.loc[outcomes_df["subpop"]==sp, ["incidH", "incidD"]] = outcomes_df.loc[outcomes_df["subpop"]==sp, ["incidH", "incidD"]]*max_summer/max_fit*mult + if this_bad > 4 or max_reg > 4: + print( + f"{outcome_fn.split('/')[-1].split('.')[0]} >>> BAAD: {this_bad} subpops AND max_ratio={max_reg:.1f}, sp with max_summer > max_fit*{reg} {bad_subpops}" + ) else: - print(f"{outcome_fn.split('/')[-1].split('.')[0]} >>> GOOD: {this_bad} subpops AND max_ratio={max_reg:.1f}, sp with max_summer > max_fit*{reg} {bad_subpops}") + print( + f"{outcome_fn.split('/')[-1].split('.')[0]} >>> GOOD: {this_bad} subpops AND max_ratio={max_reg:.1f}, sp with max_summer > max_fit*{reg} {bad_subpops}" + ) files_to_keep.append(fn) print(len(files_to_keep)) ### END OF CODE @@ -146,14 +160,18 @@ def get_all_filenames( if fill_missing: # Extract the numbers from the filenames numbers = [int(os.path.basename(filename).split(".")[0]) for filename in all_files] - missing_numbers = [num for num in range(fill_from_min, fill_from_max + 1) if num not in numbers] + missing_numbers = [ + num for num in range(fill_from_min, fill_from_max + 1) if num not in numbers + ] if missing_numbers: missing_filenames = [] for num in missing_numbers: filename = os.path.basename(all_files[0]) filename_prefix = re.search(r"^.*?(\d+)", filename).group() filename_suffix = re.search(r"(\..*?)$", filename).group() - missing_filename = os.path.join(os.path.dirname(all_files[0]), f"{num:09d}{filename_suffix}") + missing_filename = os.path.join( + os.path.dirname(all_files[0]), f"{num:09d}{filename_suffix}" + ) missing_filenames.append(missing_filename) print("The missing filenames with full paths are:") for missing_filename in missing_filenames: