diff --git a/pysindy/optimizers/sbr.py b/pysindy/optimizers/sbr.py index 72f6a4e3f..dea3b222f 100644 --- a/pysindy/optimizers/sbr.py +++ b/pysindy/optimizers/sbr.py @@ -111,9 +111,13 @@ def _numpyro_model(self, x, y): n_features = x.shape[1] n_targets = y.shape[1] - # sample the hyperparameters. - tau, c_sq = sample_reg_horseshoe_hyper( - self.sparsity_coef_tau0, self.slab_shape_nu, self.slab_shape_s + # sample the horseshoe hyperparameters. + tau = numpyro.sample("tau", HalfCauchy(self.sparsity_coef_tau0)) + c_sq = numpyro.sample( + "c_sq", + InverseGamma( + self.slab_shape_nu / 2, self.slab_shape_nu / 2 * self.slab_shape_s**2 + ), ) # sample the parameters compute the predicted outputs. @@ -141,18 +145,6 @@ def _run_mcmc(self, x, y, **kwargs): return mcmc -def sample_reg_horseshoe_hyper(tau0=0.1, nu=4, s=2): - """ - For details on this prior, please refer to: - Hirsh, S. M., Barajas-Solano, D. A., & Kutz, J. N. (2021). - parsifying Priors for Bayesian Uncertainty Quantification in - Model Discovery (arXiv:2107.02107). arXiv. http://arxiv.org/abs/2107.02107 - """ - tau = numpyro.sample("tau", HalfCauchy(tau0)) - c_sq = numpyro.sample("c_sq", InverseGamma(nu / 2, nu / 2 * s**2)) - return tau, c_sq - - def _sample_reg_horseshoe(tau, c_sq, shape): lamb = numpyro.sample("lambda", HalfCauchy(1.0), sample_shape=shape) lamb_squiggle = jnp.sqrt(c_sq) * lamb / jnp.sqrt(c_sq + tau**2 * lamb**2)