From 18207aeb725fa9eedd997b3c1c3c89743be1e76b Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Wed, 15 Jul 2020 15:22:25 -0700 Subject: [PATCH 1/2] Tweak defaults for CompartmentalModel.fit_mcmc() --- pyro/contrib/epidemiology/compartmental.py | 11 +++--- tests/contrib/epidemiology/test_models.py | 45 +++++++++++----------- 2 files changed, 29 insertions(+), 27 deletions(-) diff --git a/pyro/contrib/epidemiology/compartmental.py b/pyro/contrib/epidemiology/compartmental.py index 413812a889..3d37a1a06a 100644 --- a/pyro/contrib/epidemiology/compartmental.py +++ b/pyro/contrib/epidemiology/compartmental.py @@ -469,9 +469,10 @@ def fit_mcmc(self, **options): Note that computational cost is exponential in `num_quant_bins`. Defaults to 1 for relaxed inference. :param bool haar: Whether to use a Haar wavelet reparameterizer. + Defaults to True. :param int haar_full_mass: Number of low frequency Haar components to - include in the full mass matrix. If nonzero this implies - ``haar=True``. + include in the full mass matrix. If ``haar=False`` then this is + ignored. Defaults to 10. :param int heuristic_num_particles: Passed to :meth:`heuristic` as ``num_particles``. Defaults to 1024. :returns: An MCMC object for diagnostics, e.g. ``MCMC.summary()``. @@ -489,14 +490,14 @@ def fit_mcmc(self, **options): # Setup Haar wavelet transform. haar = options.pop("haar", False) - haar_full_mass = options.pop("haar_full_mass", 0) + haar_full_mass = options.pop("haar_full_mass", 10) full_mass = options.pop("full_mass", self.full_mass) assert isinstance(haar, bool) assert isinstance(haar_full_mass, int) and haar_full_mass >= 0 assert isinstance(full_mass, (bool, list)) haar_full_mass = min(haar_full_mass, self.duration) - if haar_full_mass: - haar = True + if not haar: + haar_full_mass = 0 if full_mass is True: haar_full_mass = 0 # No need to split. elif haar_full_mass >= self.duration: diff --git a/tests/contrib/epidemiology/test_models.py b/tests/contrib/epidemiology/test_models.py index 2aea21c6d6..c6b5a0fd5d 100644 --- a/tests/contrib/epidemiology/test_models.py +++ b/tests/contrib/epidemiology/test_models.py @@ -27,17 +27,18 @@ ("svi", {"guide_rank": 2}), ("svi", {"guide_rank": "full"}), ("mcmc", {}), - ("mcmc", {"haar": True}), + ("mcmc", {"haar": False}), + ("mcmc", {"haar_full_mass": 0}), ("mcmc", {"haar_full_mass": 2}), ("mcmc", {"num_quant_bins": 2}), ("mcmc", {"num_quant_bins": 4}), ("mcmc", {"num_quant_bins": 8}), ("mcmc", {"num_quant_bins": 12}), ("mcmc", {"num_quant_bins": 16}), - ("mcmc", {"num_quant_bins": 2, "haar": True}), + ("mcmc", {"num_quant_bins": 2, "haar": False}), ("mcmc", {"arrowhead_mass": True}), ("mcmc", {"jit_compile": True}), - ("mcmc", {"jit_compile": True, "haar_full_mass": 2}), + ("mcmc", {"jit_compile": True, "haar_full_mass": 0}), ("mcmc", {"jit_compile": True, "num_quant_bins": 2}), ("mcmc", {"num_chains": 2, "mp_context": "spawn"}), ("mcmc", {"num_chains": 2, "mp_context": "spawn", "num_quant_bins": 2}), @@ -77,8 +78,8 @@ def test_simple_sir_smoke(duration, forecast, options, algo): ("svi", {}), ("svi", {"haar": False}), ("mcmc", {}), - ("mcmc", {"haar": True}), - ("mcmc", {"haar_full_mass": 2}), + ("mcmc", {"haar": False}), + ("mcmc", {"haar_full_mass": 0}), ("mcmc", {"num_quant_bins": 2}), ], ids=str) def test_simple_seir_smoke(duration, forecast, options, algo): @@ -116,7 +117,7 @@ def test_simple_seir_smoke(duration, forecast, options, algo): @pytest.mark.parametrize("algo,options", [ ("svi", {}), ("mcmc", {}), - ("mcmc", {"haar_full_mass": 2}), + ("mcmc", {"haar_full_mass": 0}), ], ids=str) def test_simple_seird_smoke(duration, forecast, options, algo): population = 100 @@ -155,8 +156,8 @@ def test_simple_seird_smoke(duration, forecast, options, algo): @pytest.mark.parametrize("forecast", [7]) @pytest.mark.parametrize("options", [ {}, + {"haar": False}, {"num_quant_bins": 2}, - {"haar_full_mass": 2}, ], ids=str) def test_overdispersed_sir_smoke(duration, forecast, options): population = 100 @@ -186,7 +187,7 @@ def test_overdispersed_sir_smoke(duration, forecast, options): @pytest.mark.parametrize("forecast", [7]) @pytest.mark.parametrize("options", [ {}, - {"haar_full_mass": 2}, + {"haar": False}, {"num_quant_bins": 2}, ], ids=str) def test_overdispersed_seir_smoke(duration, forecast, options): @@ -221,8 +222,8 @@ def test_overdispersed_seir_smoke(duration, forecast, options): @pytest.mark.parametrize("forecast", [0, 7]) @pytest.mark.parametrize("options", [ {}, - {"haar": True}, - {"haar_full_mass": 2}, + {"haar": False}, + {"haar_full_mass": 0}, {"num_quant_bins": 2}, ], ids=str) def test_superspreading_sir_smoke(duration, forecast, options): @@ -253,8 +254,8 @@ def test_superspreading_sir_smoke(duration, forecast, options): @pytest.mark.parametrize("forecast", [0, 7]) @pytest.mark.parametrize("options", [ {}, - {"haar": True}, - {"haar_full_mass": 2}, + {"haar": False}, + {"haar_full_mass": 0}, {"num_quant_bins": 2}, ], ids=str) def test_superspreading_seir_smoke(duration, forecast, options): @@ -335,7 +336,7 @@ def test_coalescent_likelihood_smoke(duration, forecast, options, algo): ("svi", {}), ("svi", {"haar": False}), ("mcmc", {}), - ("mcmc", {"haar_full_mass": 2}), + ("mcmc", {"haar": False}), ("mcmc", {"num_quant_bins": 2}), ], ids=str) def test_heterogeneous_sir_smoke(duration, forecast, options, algo): @@ -368,8 +369,8 @@ def test_heterogeneous_sir_smoke(duration, forecast, options, algo): @pytest.mark.parametrize("options", [ xfail_param({}, reason="Delta is incompatible with relaxed inference"), {"num_quant_bins": 2}, - {"num_quant_bins": 2, "haar": True}, - {"num_quant_bins": 2, "haar_full_mass": 3}, + {"num_quant_bins": 2, "haar": False}, + {"num_quant_bins": 2, "haar_full_mass": 0}, {"num_quant_bins": 4}, ], ids=str) def test_sparse_smoke(duration, forecast, options): @@ -411,8 +412,8 @@ def test_sparse_smoke(duration, forecast, options): @pytest.mark.parametrize("forecast", [0, 7]) @pytest.mark.parametrize("options", [ {}, - {"haar": True}, - {"haar_full_mass": 4}, + {"haar": False}, + {"haar_full_mass": 0}, {"num_quant_bins": 2}, ], ids=str) def test_unknown_start_smoke(duration, pre_obs_window, forecast, options): @@ -459,8 +460,8 @@ def test_unknown_start_smoke(duration, pre_obs_window, forecast, options): ("svi", {}), ("svi", {"haar": False}), ("mcmc", {}), - ("mcmc", {"haar": True}), - ("mcmc", {"haar_full_mass": 2}), + ("mcmc", {"haar": False}), + ("mcmc", {"haar_full_mass": 0}), ("mcmc", {"num_quant_bins": 2}), ], ids=str) def test_regional_smoke(duration, forecast, options, algo): @@ -500,11 +501,11 @@ def test_regional_smoke(duration, forecast, options, algo): ("svi", {}), ("svi", {"haar": False}), ("mcmc", {}), - ("mcmc", {"haar": True}), - ("mcmc", {"haar_full_mass": 2}), + ("mcmc", {"haar": False}), + ("mcmc", {"haar_full_mass": 0}), ("mcmc", {"num_quant_bins": 2}), ("mcmc", {"jit_compile": True}), - ("mcmc", {"jit_compile": True, "haar_full_mass": 2}), + ("mcmc", {"jit_compile": True, "haar": False}), ("mcmc", {"jit_compile": True, "num_quant_bins": 2}), ], ids=str) def test_hetero_regional_smoke(duration, forecast, options, algo): From 6ea5bd2d090c758396cc4edef4622006e3b5e6ec Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Wed, 15 Jul 2020 15:28:54 -0700 Subject: [PATCH 2/2] Simplify tutorial --- tutorial/source/epi_intro.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tutorial/source/epi_intro.ipynb b/tutorial/source/epi_intro.ipynb index fe01e87cf4..b1a716d4f6 100644 --- a/tutorial/source/epi_intro.ipynb +++ b/tutorial/source/epi_intro.ipynb @@ -592,7 +592,7 @@ "%%time\n", "model = SimpleSIRModel(population, recovery_time, obs)\n", "mcmc = model.fit_mcmc(num_samples=4 if smoke_test else 400,\n", - " haar_full_mass=10, jit_compile=True)" + " jit_compile=True)" ] }, {