test apply_transfer_function_filter

mehta-lab · Feb 1, 2025 · 2d8f6b6 · 2d8f6b6
1 parent 60882e9
commit 2d8f6b6
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_filter.py b/tests/test_filter.py
@@ -3,6 +3,20 @@
 import pytest
 
 
+def test_apply_transfer_function_filter():
+    input_array = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    transfer_function = torch.tensor([[1, 0], [0, 0]])
+    result = filter.apply_transfer_function_filter(transfer_function, input_array)
+    expected = torch.tensor([[10, 10], [10, 10]])/4
+    assert torch.allclose(result, expected)
+
+    # Test with incompatible shapes
+    input_array = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    transfer_function = torch.tensor([[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]])
+    with pytest.raises(ValueError):
+        filter.apply_transfer_function_filter(transfer_function, input_array)
+
+
 def test_stretched_multiply():
     small_array = torch.tensor([[1, 2], [3, 4]])
     large_array = torch.tensor(