added test for distribution output

pts/modules/distribution_output.py

@@ -82,6 +82,6 @@ class StudentTOutput(DistributionOutput):
 
     @classmethod
     def domain_map(cls, df, loc, scale):
-        scale = nn.Softplus(scale)
-        df = 2.0 + nn.Softplus(df)
+        scale = F.softplus(scale)
+        df = 2.0 + F.softplus(df)
         return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)

test/modules/test_distribution_output.py

@@ -1,7 +1,14 @@
 import pytest
 from typing import Iterable, List, Tuple
 
+import numpy as np
+
 import torch
+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 pts.modules import DistributionOutput, StudentTOutput
 
@@ -10,19 +17,87 @@ BATCH_SIZE = 32
 TOL = 0.3
 START_TOL_MULTIPLE = 1
 
+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[torch.Tensor] = None,
+    init_biases: List[np.ndarray] = None,
     num_epochs: int = 5,
     learning_rate: float = 1e-2
 ):
-    device = torch.device("cpu")
-
+    distr_output.in_features = 1
     arg_proj = distr_output.get_args_proj()
-    arg_proj.initialize()
+    
+    if init_biases is not None:
+        for param, bias in zip(arg_proj.proj, init_biases):
+            nn.init.constant_(param.bias, bias)
+    
+    dummy_data = torch.ones((len(samples),1))
+    dataset = TensorDataset(dummy_data, samples)
+    train_data = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    optimizer = SGD(arg_proj.parameters(), lr=learning_rate)
+
+    for e in range(num_epochs):
+        cumulative_loss = 0
+        num_batches = 0
+
+        for i, (data, sample_label) in enumerate(train_data):
+            optimizer.zero_grad()
+            distr_args = arg_proj(data)
+            distr = distr_output.distribution(distr_args)
+            loss = -distr.log_prob(sample_label).mean()
+            loss.backward()
+            clip_grad_norm_(arg_proj.parameters(), 10.0)
+            optimizer.step()
+
+            num_batches += 1
+            cumulative_loss += loss.item()
+        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[0].detach().numpy() for param in arg_proj(torch.ones((1,1)))
+    ]
 
 
-@pytest.mark.parametrize("loc, scale, df", [(2.3, 0.7, 6.0)])
-def test_studentT_likelihood():
-    pass
+
+@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
+
+    distr = StudentT(df=dfs, loc=locs, scale=scales)
+    samples = distr.sample()
+    
+    init_bias = [
+        inv_softplus(df - 2),
+        loc - START_TOL_MULTIPLE * TOL * loc,
+        inv_softplus(scale - START_TOL_MULTIPLE * TOL * scale),
+    ]
+
+    df_hat, loc_hat, scale_hat = maximum_likelihood_estimate_sgd(
+        StudentTOutput(),
+        samples,
+        init_biases=init_bias,
+        num_epochs=10,
+        learning_rate=1e-2
+    )
+
+    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}"