add tests; fix transforms new operatordataset (#57)

src/nos/transforms/center_quantile_scaler.py

@@ -28,12 +28,13 @@ def __init__(
 
         assert 1.0 >= center_interval > 0
 
-        self.n_dim = src.size(-1)
+        self.n_dim = src.size(1)
 
         q = torch.tensor([center_interval / 2, 0.5, 1 - center_interval / 2])
 
         # source "distribution"
-        quantile_points = torch.quantile(src.view(-1, self.n_dim), q, dim=0, interpolation="linear")
+        src = src.transpose(1, -1)
+        quantile_points = torch.quantile(src.reshape(-1, self.n_dim), q, dim=0, interpolation="linear")
         self.median = nn.Parameter(quantile_points[1])
         self.delta = nn.Parameter(quantile_points[2] - quantile_points[0])
 
@@ -52,10 +53,11 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         Returns:
             The transformed tensor, scaled to the target distribution.
         """
-        out = (tensor - self.median) / self.delta
+
+        out = (tensor.transpose(1, -1) - self.median) / self.delta
         out = out * self.target_delta + self.target_median
 
-        return out
+        return out.transpose(1, -1)
 
     def undo(self, tensor: torch.Tensor) -> torch.Tensor:
         """Reverses the transformation applied by the forward method, mapping the tensor back to its original
@@ -67,6 +69,6 @@ def undo(self, tensor: torch.Tensor) -> torch.Tensor:
         Returns:
             The tensor with the quantile scaling transformation reversed according to the src distribution.
         """
-        out = (tensor - self.target_median) / self.target_delta
+        out = (tensor.transpose(1, -1) - self.target_median) / self.target_delta
         out = out * self.delta + self.median
-        return out
+        return out.transpose(1, -1)

src/nos/transforms/median_peak_scaler_without_shift.py

@@ -11,9 +11,11 @@ def __init__(
         src: torch.Tensor,
     ):
         super().__init__()
+
+        src = src.transpose(1, -1)
         peaks, _ = torch.max(torch.abs(src), dim=1, keepdim=True)
         medians, _ = torch.median(peaks, dim=0, keepdim=True)
-        self.medians = nn.Parameter(medians)
+        self.medians = nn.Parameter(medians.transpose(1, -1))
 
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         return tensor / self.medians

src/nos/transforms/quantile_scaler.py

@@ -55,12 +55,12 @@ def __init__(
         self.n_quantile_intervals = n_quantile_intervals
         self.n_q_points = n_quantile_intervals + 2  # n intervals have n + 2 edges
 
-        self.n_dim = src.size(-1)
+        self.n_dim = src.size(1)
 
         # source "distribution"
         self.quantile_fractions = torch.linspace(0, 1, self.n_q_points)
         quantile_points = torch.quantile(
-            src.view(-1, self.n_dim),
+            src.transpose(1, -1).reshape(-1, self.n_dim),
             self.quantile_fractions,
             dim=0,
             interpolation="linear",
@@ -129,6 +129,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         Returns:
             The transformed tensor, scaled to the target distribution.
         """
+        tensor = tensor.transpose(1, -1)
         indices = self._get_scaling_indices(tensor, self.quantile_points)
         # Scale input tensor to the unit interval based on source quantiles
         p_min = self.quantile_points[indices].view(tensor.shape)
@@ -142,7 +143,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         out = out * delta_t
         out = out + p_t_min
 
-        return out
+        return out.transpose(1, -1)
 
     def undo(self, tensor: torch.Tensor) -> torch.Tensor:
         """Reverses the transformation applied by the forward method, mapping the tensor back to its original
@@ -154,6 +155,7 @@ def undo(self, tensor: torch.Tensor) -> torch.Tensor:
         Returns:
             The tensor with the quantile scaling transformation reversed according to the src distribution.
         """
+        tensor = tensor.transpose(1, -1)
         indices = self._get_scaling_indices(tensor, self.target_quantile_points)
 
         # Scale input tensor to the unit interval based on the target distribution
@@ -168,4 +170,4 @@ def undo(self, tensor: torch.Tensor) -> torch.Tensor:
         out = out * delta
         out = out + p_min
 
-        return out
+        return out.transpose(1, -1)

tests/transforms/test_center_quantile_scaler.py

@@ -6,8 +6,8 @@
 
 
 class TestCenterQuantileScaler:
-    x = torch.rand(7, 89, 13)
-    y = torch.rand(17, 97, 13)
+    x = torch.rand(7, 13, 89)
+    y = torch.rand(17, 13, 97)
 
     def test_can_init(self):
         transform = CenterQuantileScaler(self.x)

tests/transforms/test_median_peak_scaler_without_shift.py

@@ -6,8 +6,8 @@
 
 
 class TestMedianPeakScaler:
-    x = torch.rand(7, 11, 13)
-    y = torch.rand(17, 19, 13)
+    x = torch.rand(7, 13, 11)
+    y = torch.rand(17, 13, 19)
 
     def test_can_initialize(self):
         transform = MedianPeak(self.x)

tests/transforms/test_min_max_scaler.py

@@ -0,0 +1,42 @@
+import torch
+
+from nos.transforms import (
+    MinMaxScale,
+)
+
+
+class TestMinMaxScale:
+    def test_can_initialize(self):
+        scaler = MinMaxScale(torch.zeros(1, 2, 3), torch.ones(1, 2, 3))
+        assert isinstance(scaler, MinMaxScale)
+
+    def test_can_forward(self):
+        scaler = MinMaxScale(torch.zeros(2, 2), torch.ones(2, 2))
+
+        x = torch.rand(2, 2)
+
+        out = scaler(x)
+
+        assert isinstance(out, torch.Tensor)
+
+    def test_forward_correct(self):
+        scaler = MinMaxScale(-torch.ones(1, 2, 1), 2 * torch.ones(1, 2, 1))
+
+        x = torch.linspace(-1, 2, 30)
+        x = x.reshape(1, 2, -1)
+
+        out = scaler(x)
+
+        assert torch.all(out >= -1.0)
+        assert torch.all(out <= 1.0)
+        assert torch.sum(torch.isclose(out, -torch.ones(out.shape))) == 1
+        assert torch.sum(torch.isclose(out, torch.ones(out.shape))) == 1
+
+    def test_can_undo(self):
+        scaler = MinMaxScale(-torch.ones(1, 2, 1), 2 * torch.ones(1, 2, 1))
+
+        x = torch.linspace(-1, 2, 30)
+
+        out = scaler.undo(x)
+
+        assert isinstance(out, torch.Tensor)

tests/transforms/test_quantile_scaler.py

@@ -6,8 +6,8 @@
 
 
 class TestQuantileScaler:
-    x = torch.rand(7, 11, 13)
-    y = torch.rand(17, 19, 13)
+    x = torch.rand(7, 13, 11)
+    y = torch.rand(17, 13, 19)
 
     def test_can_init(self):
         transform = QuantileScaler(self.x)