test files rework + kl_div import

src/timediffusion/__init__.py

@@ -1,10 +1,11 @@
-from .utils import count_params, DimUniversalStandardScaler
+from .utils import count_params, DimUniversalStandardScaler, kl_div
 from .models import TimeDiffusionProjector, TimeDiffusion
 from .frameworks import TD
 
 __all__ = [
     # useful functions
     "count_params",
+    "kl_div",
     # data processing
     "DimUniversalStandardScaler",
     # models

src/timediffusion/utils.py

@@ -37,7 +37,7 @@ def kl_div(x: Union[np.array, torch.Tensor], y: Union[np.array, torch.Tensor],
     if type(x) is not type(y):
         raise ValueError(f"input arrays for kl_div should be same type, got {type(x)} and {type(y)}")
 
-    if type(x) is np.array:
+    if type(x) is np.ndarray:
         clip = lambda arr: np.clip(arr, a_min=eps, a_max=1)
         log = np.log
     elif type(x) is torch.Tensor:

tests/test_frameworks.py

@@ -0,0 +1,70 @@
+import pytest
+
+import numpy as np
+
+import torch
+from torch import nn
+
+from timediffusion import TD
+
+
+@pytest.mark.parametrize("dims", [[1, 35], [1, 7, 7], [1, 5, 5, 5], [2, 35], [2, 7, 7], [2, 5, 5, 5]])
+@pytest.mark.parametrize("mask_dropout", [None, 0.2])
+class TestTD:
+    def test_fit(self, dims, mask_dropout):
+        model = TD(input_dims=dims)
+
+        data = np.ones(dims)
+        if mask_dropout is None:
+            mask = None
+        else:
+            np.random.seed(42)
+            mask = np.random.uniform(low=0., high=1.0, size=data.shape) < mask_dropout
+
+        try:
+            model.fit(data, mask=mask, epochs=1, batch_size=1, steps_per_epoch=2)
+        except Exception as e:
+            pytest.fail(f"TD fit with {dims = } failed with exception: {e}")
+
+    def test_restore(self, dims, mask_dropout):
+        model = TD(input_dims=dims)
+
+        data = np.ones(dims)
+        if mask_dropout is None:
+            mask = None
+        else:
+            np.random.seed(42)
+            mask = np.random.uniform(low=0., high=1.0, size=data.shape) < mask_dropout
+
+        try:
+            model.fit(example=data, mask=mask, epochs=1, batch_size=1, steps_per_epoch=2)
+            res = model.restore(data, mask=mask, steps=2)
+        except Exception as e:
+            pytest.fail(f"TD restore with {dims = } failed with exception: {e}")
+
+        if mask is not None:
+            assert np.allclose(res.numpy()[mask], data[mask])
+
+    def test_forecast(self, dims, mask_dropout):
+        if len(dims) > 2:
+            return
+        horizon = 3
+
+        model = TD(input_dims=dims)
+
+        data = np.ones(dims)
+        if mask_dropout is None:
+            mask = None
+        else:
+            np.random.seed(42)
+            mask = np.random.uniform(low=0., high=1.0, size=data.shape) < mask_dropout
+
+        try:
+            model.fit(example=data, mask=mask, epochs=1, batch_size=1, steps_per_epoch=2)
+            res = model.forecast(horizon, steps=2)
+        except Exception as e:
+            pytest.fail(f"TD restore with {dims = } failed with exception: {e}")
+
+        assert len(res.shape) == 2
+        assert res.shape[0] == dims[0]
+        assert res.shape[1] == horizon

tests/test_models.py

@@ -0,0 +1,50 @@
+import pytest
+
+import numpy as np
+
+import torch
+from torch import nn
+
+from timediffusion import TimeDiffusionProjector, TimeDiffusion
+
+
+@pytest.mark.parametrize("model_init", [TimeDiffusionProjector, TimeDiffusion])
+@pytest.mark.parametrize("dims", [[1, 35], [1, 7, 7], [1, 5, 5, 5], [2, 35], [2, 7, 7], [2, 5, 5, 5]])
+class TestTimeDiffusion:
+    def test_forward_pass(self, model_init, dims):
+        model = model_init(input_dims=dims)
+
+        # unbatched forward pass
+        data = torch.ones(*dims)
+        try:
+            res = model(data)
+        except Exception as e:
+            pytest.fail(f"Unbatched forward pass of {type(model).__name__} with {dims = } failed with exception: {e}")
+        assert data.shape == res.shape
+
+        # batched forward pass
+        data = torch.ones(1, *dims)
+        try:
+            res = model(data)
+        except Exception as e:
+            pytest.fail(f"Batched forward pass of {type(model).__name__} with {dims = } failed with exception: {e}")
+        assert data.shape == res.shape
+
+    def test_backward_pass(self, model_init, dims):
+        model = model_init(input_dims=dims)
+
+        # unbatched backward pass
+        data = torch.ones(*dims)
+        try:
+            res = model(data)
+            loss = (res - 1).mean().backward()
+        except Exception as e:
+            pytest.fail(f"Unbatched backward pass of {type(model).__name__} with {dims = } failed with exception: {e}")
+
+        # batched backward pass
+        data = torch.ones(1, *dims)
+        try:
+            res = model(data)
+            loss = (res - 1).mean().backward()
+        except Exception as e:
+            pytest.fail(f"Batched backward pass of {type(model).__name__} with {dims = } failed with exception: {e}")

tests/test_utils.py

@@ -0,0 +1,41 @@
+import pytest
+
+import numpy as np
+
+import torch
+from torch import nn
+
+from timediffusion import count_params, kl_div
+
+
+@pytest.mark.parametrize(
+    "x,y",
+    [
+        (np.sin(np.arange(10)), np.arange(10)),
+        (torch.sin(torch.arange(10)), torch.arange(10)),
+        pytest.param(np.arange(10), torch.arange(10), marks=pytest.mark.xfail)
+    ],
+)
+def test_kl_div(x, y):
+    kl_div(x, y)
+
+@pytest.mark.parametrize(
+    "in_features,out_features",
+    [(100, 100), (200, 200)],
+)
+def test_count_params_linear(in_features, out_features):
+    linear = nn.Linear(in_features, out_features)
+    assert count_params(linear) == in_features * out_features + out_features
+
+@pytest.mark.parametrize(
+    "in_channels,out_channels,kernel_size,groups",
+    [
+        (3, 24, 3, 1),
+        (1, 10, 2, 1),
+        (3, 12, 4, 3)
+     ],
+)
+def test_count_params_conv(in_channels, out_channels, kernel_size, groups):
+    conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels,
+         kernel_size=kernel_size, groups=groups)
+    assert count_params(conv) == out_channels * (in_channels // groups) * kernel_size + out_channels