Merge pull request #23 from cics-nd/issue-22

Fix issue 22, tests for example scripts

Author: sdatkinson

examples/regression_1d.py

@@ -8,6 +8,7 @@
 
 import os
 import sys
+from argparse import ArgumentParser
 from time import time
 
 import torch
@@ -30,7 +31,7 @@ def f(x):
     return np.sin(2.0 * np.pi * x) + np.cos(3.5 * np.pi * x) - 3.0 * x + 5.0
 
 
-if __name__ == "__main__":
+def main(args):
     # Create data:
     n = 100
     x = np.linspace(0, 1, n).reshape((-1, 1))
@@ -41,9 +42,12 @@ def f(x):
     kern = kernels.Linear(1) + kernels.Rbf(1) + kernels.Constant(1)
 
     # Try different models:
-    model = GPR(x, y, kern)
-    # model = VFE(x, y, kern)
-    # model.cuda()  # If you want to use GPU
+    if args.model_type == "GPR":
+        model = GPR(x, y, kern)
+    elif args.model_type == "VFE":
+        model = VFE(x, y, kern)
+    if args.cuda:
+        model.cuda()  # If you want to use GPU
 
     # Train
     model.optimize(method="L-BFGS-B", max_iter=100)
@@ -62,17 +66,30 @@ def f(x):
     # Show prediction
     x_test = x_test.flatten()
     plt.figure()
-    plt.fill_between(x_test, (mu - unc).flatten(), (mu + unc).flatten(), 
-        color=(0.9,) * 3)
+    plt.fill_between(
+        x_test, (mu - unc).flatten(), (mu + unc).flatten(), color=(0.9,) * 3
+    )
     plt.plot(x_test, mu)
     plt.plot(x_test, f(x_test))
     for y_samp_i in y_samp:
         plt.plot(x_test, y_samp_i, color=(0.4, 0.7, 1.0), alpha=0.5)
     plt.plot(x, y, "o")
     if hasattr(model, "Z"):
         plt.plot(
-            model.Z.detach().cpu().numpy(), 
-            1.0 + plt.ylim()[0] * np.ones(model.Z.shape[0]), 
-            "+"
+            model.Z.detach().cpu().numpy(),
+            1.0 + plt.ylim()[0] * np.ones(model.Z.shape[0]),
+            "+",
         )
-    plt.show()
+    if args.no_plot:
+        plt.close()
+    else:
+        plt.show()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--model-type", type=str, choices=("GPR", "VFE"), default="GPR")
+    parser.add_argument("--cuda", action="store_true")
+    parser.add_argument("--no-plot", action="store_true")
+
+    main(parser.parse_args())

gptorch/mean_functions.py

@@ -11,8 +11,6 @@
 
 from .util import torch_dtype
 
-torch.set_default_dtype(torch_dtype)
-
 
 class Constant(torch.nn.Module):
     """
@@ -26,13 +24,13 @@ def __init__(self, dy: int, val: torch.Tensor = None):
                 raise ValueError("Provided val doesn't match output dimension")
             val = val.clone()
         else:
-            val = torch.zeros(dy)
+            val = torch.zeros(dy, dtype=torch_dtype)
 
         self._dy = dy
         self.val = torch.nn.Parameter(val)
 
     def forward(self, x):
-        output = torch.zeros(x.shape[0], self._dy)
+        output = torch.zeros(x.shape[0], self._dy, dtype=torch_dtype)
         if self._is_cuda():
             output = output.cuda()
         return output + self.val

test/test_base.py

@@ -1,3 +1,6 @@
+# File: test_base.py
+# File Created: Saturday, 30th October 2021 7:39:23 am
+# Author: Steven Atkinson (steven@atkinson.mn)
 
 """
 Basic tests for the repo
@@ -10,8 +13,10 @@ def test_torch_dtype():
     """
 
     import torch
+
     dtype = torch.get_default_dtype()
-    import gptorch
+    import gptorch  # noqa F401
+
     new_dtype = torch.get_default_dtype()
 
     assert new_dtype == dtype

test/test_examples.py

@@ -0,0 +1,25 @@
+# File: test_examples.py
+# File Created: Saturday, 30th October 2021 9:04:58 am
+# Author: Steven Atkinson (steven@atkinson.mn)
+
+import os
+from subprocess import check_call
+from sys import executable
+
+import pytest
+
+args = (
+    ("regression_1d.py", "--no-plot"),
+    ("regression_1d.py", "--no-plot", "--model-type", "VFE"),
+)
+
+
+@pytest.mark.parametrize("args", args)
+def test_example(args):
+    basename, args = args[0], args[1:]
+    script_path = os.path.join(os.path.dirname(__file__), "..", "examples", basename)
+    check_call((executable, script_path) + args)
+
+
+if __name__ == "__main__":
+    pytest.main()

test/test_models/common.py

@@ -9,11 +9,13 @@
 from gptorch.models.base import GPModel
 from gptorch.util import torch_dtype
 
-torch.set_default_dtype(torch_dtype)
 
-
-def gaussian_predictions(model: GPModel, x_test: torch.Tensor, 
-        expected_mu: np.ndarray, expected_s: np.ndarray):
+def gaussian_predictions(
+    model: GPModel,
+    x_test: torch.Tensor,
+    expected_mu: np.ndarray,
+    expected_s: np.ndarray,
+):
     """
     Every GP model with a gaussian likelihood needs the same set of tests run on
     its ._predict() method.
@@ -29,10 +31,20 @@ def gaussian_predictions(model: GPModel, x_test: torch.Tensor,
     assert mu_diag.shape[1] == model.Y.shape[1]
     assert all([ss == ms for ss, ms in zip(mu_diag.shape, s_diag.shape)])
 
-    assert all([a == pytest.approx(e) for a, e in zip(
-        mu_diag.detach().numpy().flatten(), expected_mu.flatten())])
-    assert all([a == pytest.approx(e) for a, e in zip(
-        s_diag.detach().numpy().flatten(), expected_s.diagonal().flatten())])
+    assert all(
+        [
+            a == pytest.approx(e)
+            for a, e in zip(mu_diag.detach().numpy().flatten(), expected_mu.flatten())
+        ]
+    )
+    assert all(
+        [
+            a == pytest.approx(e)
+            for a, e in zip(
+                s_diag.detach().numpy().flatten(), expected_s.diagonal().flatten()
+            )
+        ]
+    )
 
     # Predictions with full covariance
     mu_full, s_full = model._predict(x_test, diag=False)
@@ -43,8 +55,16 @@ def gaussian_predictions(model: GPModel, x_test: torch.Tensor,
     assert mu_full.shape[0] == x_test.shape[0]
     assert mu_full.shape[1] == model.Y.shape[1]
     assert all([ss == x_test.shape[0] for ss in s_full.shape])
-    
-    assert all([a == pytest.approx(e) for a, e in zip(
-        mu_full.detach().numpy().flatten(), expected_mu.flatten())])
-    assert all([a == pytest.approx(e) for a, e in zip(
-        s_full.detach().numpy().flatten(), expected_s.flatten())])
+    assert all(
+        [
+            a == pytest.approx(e)
+            for a, e in zip(mu_full.detach().numpy().flatten(), expected_mu.flatten())
+        ]
+    )
+    assert all(
+        [
+            a == pytest.approx(e)
+            for a, e in zip(s_full.detach().numpy().flatten(), expected_s.flatten())
+        ]
+    )
+