added some more tests

test/modules/test_distribution_output.py

@@ -8,9 +8,14 @@ import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
 from torch.utils.data import TensorDataset, DataLoader
 from torch.optim import SGD
-from torch.distributions import StudentT
+from torch.distributions import StudentT, Beta, NegativeBinomial
 
-from pts.modules import DistributionOutput, StudentTOutput
+from pts.modules import (
+    DistributionOutput,
+    StudentTOutput,
+    BetaOutput,
+    NegativeBinomialOutput,
+)
 
 NUM_SAMPLES = 2000
 BATCH_SIZE = 32
@@ -22,11 +27,13 @@ def inv_softplus(y: np.ndarray) -> np.ndarray:
     return np.log(np.exp(y) - 1)
 
 
-def maximum_likelihood_estimate_sgd(distr_output: DistributionOutput,
-                                    samples: torch.Tensor,
-                                    init_biases: List[np.ndarray] = None,
-                                    num_epochs: int = 5,
-                                    learning_rate: float = 1e-2):
+def maximum_likelihood_estimate_sgd(
+    distr_output: DistributionOutput,
+    samples: torch.Tensor,
+    init_biases: List[np.ndarray] = None,
+    num_epochs: int = 5,
+    learning_rate: float = 1e-2,
+):
     arg_proj = distr_output.get_args_proj(in_features=1)
 
     if init_biases is not None:
@@ -57,21 +64,86 @@ def maximum_likelihood_estimate_sgd(distr_output: DistributionOutput,
         print("Epoch %s, loss: %s" % (e, cumulative_loss / num_batches))
 
     if len(distr_args[0].shape) == 1:
-        return [
-            param.detach().numpy() for param in arg_proj(torch.ones((1, 1)))
-        ]
+        return [param.detach().numpy() for param in arg_proj(torch.ones((1, 1)))]
 
-    return [
-        param[0].detach().numpy() for param in arg_proj(torch.ones((1, 1)))
+    return [param[0].detach().numpy() for param in arg_proj(torch.ones((1, 1)))]
+
+
+@pytest.mark.parametrize("concentration1, concentration0", [(3.75, 1.25)])
+def test_beta_likelihood(concentration1: float, concentration0: float) -> None:
+    """
+    Test to check that maximizing the likelihood recovers the parameters
+    """
+
+    # generate samples
+    concentration1s = torch.zeros((NUM_SAMPLES,)) + concentration1
+    concentration0s = torch.zeros((NUM_SAMPLES,)) + concentration0
+
+    distr = Beta(concentration1s, concentration0s)
+    samples = distr.sample()
+
+    init_biases = [
+        inv_softplus(concentration1 - START_TOL_MULTIPLE * TOL * concentration1),
+        inv_softplus(concentration0 - START_TOL_MULTIPLE * TOL * concentration0),
     ]
 
+    concentration1_hat, concentration0_hat = maximum_likelihood_estimate_sgd(
+        BetaOutput(),
+        samples,
+        init_biases=init_biases,
+        learning_rate=0.05,
+        num_epochs=10,
+    )
+
+    print("concentration1:", concentration1_hat, "concentration0:", concentration0_hat)
+    assert (
+        np.abs(concentration1_hat - concentration1) < TOL * concentration1
+    ), f"concentration1 did not match: concentration1 = {concentration1}, concentration1_hat = {concentration1_hat}"
+    assert (
+        np.abs(concentration0_hat - concentration0) < TOL * concentration0
+    ), f"concentration0 did not match: concentration0 = {concentration0}, concentration0_hat = {concentration0_hat}"
+
+
+@pytest.mark.parametrize("mu_alpha", [(2.5, 0.7)])
+def test_neg_binomial(mu_alpha: Tuple[float, float]) -> None:
+    """
+    Test to check that maximizing the likelihood recovers the parameters
+    """
+    # test instance
+    mu, alpha = mu_alpha
+
+    # generate samples
+    mus = torch.zeros((NUM_SAMPLES,)) + mu
+    alphas = torch.zeros((NUM_SAMPLES,)) + alpha
+
+    neg_bin_distr = NegativeBinomial(
+        total_count=1.0 / alphas, probs=mus * alphas / (1.0 + mus * alphas)
+    )
+    samples = neg_bin_distr.sample()
+
+    init_biases = [
+        inv_softplus(mu - START_TOL_MULTIPLE * TOL * mu),
+        inv_softplus(alpha + START_TOL_MULTIPLE * TOL * alpha),
+    ]
+
+    mu_hat, alpha_hat = maximum_likelihood_estimate_sgd(
+        NegativeBinomialOutput(), samples, init_biases=init_biases, num_epochs=15,
+    )
+
+    assert (
+        np.abs(mu_hat - mu) < TOL * mu
+    ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
+    assert (
+        np.abs(alpha_hat - alpha) < TOL * alpha
+    ), f"alpha did not match: alpha = {alpha}, alpha_hat = {alpha_hat}"
+
 
 @pytest.mark.parametrize("df, loc, scale,", [(6.0, 2.3, 0.7)])
 def test_studentT_likelihood(df: float, loc: float, scale: float):
 
-    dfs = torch.zeros((NUM_SAMPLES, )) + df
-    locs = torch.zeros((NUM_SAMPLES, )) + loc
-    scales = torch.zeros((NUM_SAMPLES, )) + scale
+    dfs = torch.zeros((NUM_SAMPLES,)) + df
+    locs = torch.zeros((NUM_SAMPLES,)) + loc
+    scales = torch.zeros((NUM_SAMPLES,)) + scale
 
     distr = StudentT(df=dfs, loc=locs, scale=scales)
     samples = distr.sample()
@@ -86,12 +158,16 @@ def test_studentT_likelihood(df: float, loc: float, scale: float):
         StudentTOutput(),
         samples,
         init_biases=init_bias,
-        num_epochs=10,
-        learning_rate=1e-2)
+        num_epochs=15,
+        learning_rate=1e-3,
+    )
 
-    assert (np.abs(df_hat - df) < TOL * df
-            ), f"df did not match: df = {df}, df_hat = {df_hat}"
-    assert (np.abs(loc_hat - loc) < TOL * loc
-            ), f"loc did not match: loc = {loc}, loc_hat = {loc_hat}"
-    assert (np.abs(scale_hat - scale) < TOL * scale
-            ), f"scale did not match: scale = {scale}, scale_hat = {scale_hat}"
+    assert (
+        np.abs(df_hat - df) < TOL * df
+    ), f"df did not match: df = {df}, df_hat = {df_hat}"
+    assert (
+        np.abs(loc_hat - loc) < TOL * loc
+    ), f"loc did not match: loc = {loc}, loc_hat = {loc_hat}"
+    assert (
+        np.abs(scale_hat - scale) < TOL * scale
+    ), f"scale did not match: scale = {scale}, scale_hat = {scale_hat}"