test reducers

t-vi · t-vi · commit 51ad92fb1cdf · 2021-02-19T14:03:13.000+01:00
diff --git a/test/test_reducers.py b/test/test_reducers.py
@@ -0,0 +1,28 @@
+import pytest
+import torchdrift
+import torch
+import sklearn.decomposition
+
+def test_pca():
+    pca = torchdrift.reducers.PCAReducer(n_components=2)
+    assert 'n_components' in str(pca)
+    a = torch.randn(100, 50, dtype=torch.double)
+    red = pca.fit(a)
+    pca_ref = sklearn.decomposition.PCA(n_components=2)
+    red_ref = torch.from_numpy(pca_ref.fit_transform(a))
+    # need to find a way to deal with signs
+    torch.testing.assert_allclose(red.abs(), red_ref.abs())
+    b = torch.randn(25, 50, dtype=torch.double)
+    red2 = pca(b)
+    red2_ref = torch.from_numpy(pca_ref.transform(b))
+
+def test_reducer():
+    x = torch.randn(5, 5)
+    r = torchdrift.reducers.Reducer()
+    with pytest.raises(NotImplementedError):
+        r.fit(x)
+    with pytest.raises(NotImplementedError):
+        r(x)
+
+if __name__ == "__main__":
+    pytest.main([__file__])
