Fixed norm const for SBVM

examples/ssbvm_mixture.py

@@ -133,18 +133,15 @@ def ss_model(data, num_data, num_mix_comp=2):
         sine = SineBivariateVonMises(
             phi_loc=phi_loc[assign],
             psi_loc=psi_loc[assign],
-            # These concentrations are an order of magnitude lower than expected (550-1000)!
-            phi_concentration=70 * phi_conc[assign],
-            psi_concentration=70 * psi_conc[assign],
+            phi_concentration=1000 * phi_conc[assign],
+            psi_concentration=1000 * psi_conc[assign],
             weighted_correlation=corr_scale[assign],
         )
         return numpyro.sample("phi_psi", SineSkewed(sine, skewness[assign]), obs=data)
 
 
 def run_hmc(rng_key, model, data, num_mix_comp, args, bvm_init_locs):
-    kernel = NUTS(
-        model, init_strategy=init_to_value(values=bvm_init_locs), max_tree_depth=7
-    )
+    kernel = NUTS(model, init_strategy=init_to_value(values=bvm_init_locs))
     mcmc = MCMC(kernel, num_samples=args.num_samples, num_warmup=args.num_warmup)
     mcmc.run(rng_key, data, len(data), num_mix_comp)
     mcmc.print_summary()
@@ -162,7 +159,7 @@ def num_mix_comps(amino_acid):
     return num_mix.get(amino_acid, 9)
 
 
-def ramachandran_plot(data, pred_data, aas, file_name="ssbvm_mixture.pdf"):
+def ramachandran_plot(data, pred_data, aas, file_name="ssbvm_mixture.png"):
     amino_acids = {"S": "Serine", "P": "Proline", "G": "Glycine"}
     fig, axss = plt.subplots(2, len(aas))
     cdata = data

numpyro/distributions/directional.py

@@ -308,10 +308,12 @@ class SineBivariateVonMises(Distribution):
     .. note:: Sample efficiency drops as
 
         .. math::
-            \frac{\rho}{\kappa_1\kappa_2} \rightarrow 1
+            \frac{\rho^2}{\kappa_1\kappa_2} \rightarrow 1
 
-        because the distribution becomes increasingly bimodal. To avoid bimodality use the `weighted_correlation`
-        parameter with a skew away from one (e.g., Beta(1,3)). The `weighted_correlation` should be in [0,1].
+        because the distribution becomes increasingly bimodal. To avoid inefficient sampling use the
+        `weighted_correlation` parameter with a skew away from one (e.g.,
+        `TransformedDistribution(Beta(5,5), AffineTransform(loc=-1, scale=2))`).  The `weighted_correlation`
+        should be in [-1,1].
 
     .. note:: The correlation and weighted_correlation params are mutually exclusive.
 
@@ -404,7 +406,8 @@ def norm_const(self):
             jnp.log(jnp.clip(corr**2, jnp.finfo(jnp.result_type(float)).tiny))
             - jnp.log(4 * jnp.prod(conc, axis=-1))
         )
-        fs += log_I1(49, conc, terms=51).sum(-1)
+        num_I1terms = 10_001
+        fs += log_I1(49, conc, terms=num_I1terms).sum(-1)
         norm_const = 2 * jnp.log(jnp.array(2 * pi)) + logsumexp(fs, 0)
         return norm_const.reshape(jnp.shape(self.phi_loc))
 

test/test_distributions.py

@@ -3464,3 +3464,16 @@ def test_gaussian_random_walk_linear_recursive_equivalence():
     x2 = dist2.sample(random.PRNGKey(7))
     assert jnp.allclose(x1, x2.squeeze())
     assert jnp.allclose(dist1.log_prob(x1), dist2.log_prob(x2))
+
+
+@pytest.mark.parametrize("conc", [1.0, 10.0, 1000.0, 10000.0])
+def test_sine_bivariate_von_mises_norm(conc):
+    dist = SineBivariateVonMises(0, 0, conc, conc, 0.0)
+    num_samples = 500
+    x = jnp.linspace(-jnp.pi, jnp.pi, num_samples)
+    y = jnp.linspace(-jnp.pi, jnp.pi, num_samples)
+    mesh = jnp.stack(jnp.meshgrid(x, y), axis=-1)
+    integral_torus = (
+        jnp.exp(dist.log_prob(mesh)) * (2 * jnp.pi) ** 2 / num_samples**2
+    ).sum()
+    assert jnp.allclose(integral_torus, 1.0, rtol=1e-2)