From 8611f8fc30ab43f9bcd8646682b9ecce9eb3ff69 Mon Sep 17 00:00:00 2001 From: Junpeng Lao Date: Wed, 30 Oct 2024 07:51:20 +0100 Subject: [PATCH 1/2] Add logjac to logdensity_fn --- docs/examples/howto_sample_multiple_chains.md | 3 ++- docs/examples/quickstart.md | 3 ++- tests/smc/__init__.py | 2 +- tests/smc/test_inner_kernel_tuning.py | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/examples/howto_sample_multiple_chains.md b/docs/examples/howto_sample_multiple_chains.md index a5b6566f8..c2947e29f 100644 --- a/docs/examples/howto_sample_multiple_chains.md +++ b/docs/examples/howto_sample_multiple_chains.md @@ -57,8 +57,9 @@ observed = np.random.normal(loc, scale, size=1_000) def logdensity_fn(loc, log_scale, observed=observed): """Univariate Normal""" scale = jnp.exp(log_scale) + logjac = log_scale logpdf = stats.norm.logpdf(observed, loc, scale) - return jnp.sum(logpdf) + return logjac + jnp.sum(logpdf) def logdensity(x): diff --git a/docs/examples/quickstart.md b/docs/examples/quickstart.md index 870e5df9a..a290bfdad 100644 --- a/docs/examples/quickstart.md +++ b/docs/examples/quickstart.md @@ -48,8 +48,9 @@ observed = np.random.normal(loc, scale, size=1_000) def logdensity_fn(loc, log_scale, observed=observed): """Univariate Normal""" scale = jnp.exp(log_scale) + logjac = log_scale logpdf = stats.norm.logpdf(observed, loc, scale) - return jnp.sum(logpdf) + return logjac + jnp.sum(logpdf) logdensity = lambda x: logdensity_fn(**x) diff --git a/tests/smc/__init__.py b/tests/smc/__init__.py index 006d7ba38..71d59e529 100644 --- a/tests/smc/__init__.py +++ b/tests/smc/__init__.py @@ -9,7 +9,7 @@ def logdensity_by_observation(self, log_scale, coefs, preds, x): scale = jnp.exp(log_scale) y = jnp.dot(x, coefs) logpdf = stats.norm.logpdf(preds, y, scale) - return logpdf + return log_scale + logpdf def logdensity_fn(self, log_scale, coefs, preds, x): """Linear regression""" diff --git a/tests/smc/test_inner_kernel_tuning.py b/tests/smc/test_inner_kernel_tuning.py index 7d6190af5..080a02749 100644 --- a/tests/smc/test_inner_kernel_tuning.py +++ b/tests/smc/test_inner_kernel_tuning.py @@ -69,7 +69,7 @@ def logdensity_fn(self, log_scale, coefs, preds, x): scale = jnp.exp(log_scale) y = jnp.dot(x, coefs) logpdf = stats.norm.logpdf(preds, y, scale) - return jnp.sum(logpdf) + return log_scale + jnp.sum(logpdf) def test_smc_inner_kernel_adaptive_tempered(self): self.smc_inner_kernel_tuning_test_case( From e63b057962d92dd45340650b4ae08881b0653e87 Mon Sep 17 00:00:00 2001 From: Junpeng Lao Date: Wed, 30 Oct 2024 11:00:42 +0100 Subject: [PATCH 2/2] Refactor logprior_fn in SMCLinearRegressionTestCase --- tests/smc/__init__.py | 13 ++++++------- tests/smc/test_inner_kernel_tuning.py | 2 +- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/smc/__init__.py b/tests/smc/__init__.py index 71d59e529..8f5a39d6f 100644 --- a/tests/smc/__init__.py +++ b/tests/smc/__init__.py @@ -9,13 +9,16 @@ def logdensity_by_observation(self, log_scale, coefs, preds, x): scale = jnp.exp(log_scale) y = jnp.dot(x, coefs) logpdf = stats.norm.logpdf(preds, y, scale) - return log_scale + logpdf + return logpdf def logdensity_fn(self, log_scale, coefs, preds, x): """Linear regression""" logpdf = self.logdensity_by_observation(log_scale, coefs, preds, x) return jnp.sum(logpdf) + def logprior_fn(self, log_scale, coefs): + return log_scale + stats.norm.logpdf(log_scale) + stats.norm.logpdf(coefs) + def observations(self): num_particles = 100 @@ -27,9 +30,7 @@ def observations(self): def particles_prior_loglikelihood(self): observations, num_particles = self.observations() - logprior_fn = lambda x: stats.norm.logpdf(x["log_scale"]) + stats.norm.logpdf( - x["coefs"] - ) + logprior_fn = lambda x: self.logprior_fn(**x) loglikelihood_fn = lambda x: self.logdensity_fn(**x, **observations) log_scale_init = np.random.randn(num_particles) @@ -45,9 +46,7 @@ def partial_posterior_test_case(self): y_data = 3 * x_data + np.random.normal(size=x_data.shape) observations = {"x": x_data, "preds": y_data} - logprior_fn = lambda x: stats.norm.logpdf(x["log_scale"]) + stats.norm.logpdf( - x["coefs"] - ) + logprior_fn = lambda x: self.logprior_fn(**x) log_scale_init = np.random.randn(num_particles) coeffs_init = np.random.randn(num_particles) diff --git a/tests/smc/test_inner_kernel_tuning.py b/tests/smc/test_inner_kernel_tuning.py index 080a02749..7d6190af5 100644 --- a/tests/smc/test_inner_kernel_tuning.py +++ b/tests/smc/test_inner_kernel_tuning.py @@ -69,7 +69,7 @@ def logdensity_fn(self, log_scale, coefs, preds, x): scale = jnp.exp(log_scale) y = jnp.dot(x, coefs) logpdf = stats.norm.logpdf(preds, y, scale) - return log_scale + jnp.sum(logpdf) + return jnp.sum(logpdf) def test_smc_inner_kernel_adaptive_tempered(self): self.smc_inner_kernel_tuning_test_case(