Add VI to SIR_model_with_age_groups_fixed_C.ipynb

Priesemann-Group · May 26, 2021 · 59447c8 · 59447c8
1 parent b970c7b
commit 59447c8
Show file tree

Hide file tree

Showing 4 changed files with 236 additions and 1,762 deletions.
diff --git a/covid19_npis/model/approximate_posterior.py b/covid19_npis/model/approximate_posterior.py
@@ -56,7 +56,9 @@ def build_iaf(values_iaf_dict, order_list, values_exclude_dict=None):
         k: i for k, i in zip(values_iaf_dict.keys(), range(len(values_iaf_dict)))
     }
 
-    size_iaf = sum([int(np.prod(tensor.shape)) for tensor in values_iaf_dict.values()])
+    size_iaf = 2 * sum(
+        [int(np.prod(tensor.shape)) for tensor in values_iaf_dict.values()]
+    )
 
     size_splits = [int(np.prod(v.shape)) for v in values_iaf_dict.values()]
 

diff --git a/covid19_npis/model/likelihood.py b/covid19_npis/model/likelihood.py
@@ -76,7 +76,7 @@ def _studentT_positive_tests(modelParams, pos_tests):
     # Scale of the likelihood sigma
     sigma = yield HalfCauchy(
         name="sigma_likelihood_pos_tests",
-        scale=50.0,
+        scale=1,
         event_stack=modelParams.num_countries,
         conditionally_independent=True,
         transform=transformations.SoftPlus(),
@@ -169,7 +169,7 @@ def _studentT_total_tests(modelParams, total_tests):
     # Scale of the likelihood sigma for each country
     sigma = yield HalfCauchy(
         name="sigma_likelihood_total_tests",
-        scale=50.0,
+        scale=1,
         event_stack=modelParams.num_countries,
         conditionally_independent=True,
         transform=transformations.SoftPlus(),
@@ -227,7 +227,7 @@ def _studentT_deaths(modelParams, deaths):
     # Scale of the likelihood sigma for each country
     sigma = yield HalfCauchy(
         name="sigma_likelihood_deaths",
-        scale=5,
+        scale=1,
         event_stack=modelParams.num_countries,
         conditionally_independent=True,
         transform=transformations.SoftPlus(),

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 #Own pymc4 version for progressbar and enhancments PR pending
-git+https://github.com/jdehning/pymc4.git@b6b26f9c32e1fda02970a3de4229b62d218f67e4
+git+https://github.com/jdehning/pymc4.git@7b92752e661579e87c7a390bccf9a49c1e4e8166
 #Fix black version
 git+https://github.com/psf/black@19.10b0
 # Own arviz version till PR is integrated to version 0.12.0

diff --git a/scripts/notebooks/SIR_model_with_age_groups_fixed_C.ipynb b/scripts/notebooks/SIR_model_with_age_groups_fixed_C.ipynb