From 8d0daa3236e387acb00afc1181047a7d0b0964d7 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Fri, 20 Dec 2019 23:23:00 +0100 Subject: [PATCH] added some more tests --- test/modules/test_distribution_output.py | 122 ++++++++++++++++++----- 1 file changed, 99 insertions(+), 23 deletions(-) diff --git a/test/modules/test_distribution_output.py b/test/modules/test_distribution_output.py index de1ad8c..7a8f68d 100644 --- a/test/modules/test_distribution_output.py +++ b/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}"