Merge branch 'main' of github.com:aidos-lab/dect

Author: ErnstRoell

README.md

@@ -1,5 +1,5 @@
 # DECT - Differentiable Euler Characteristic Transform
-[![arXiv](https://img.shields.io/badge/arXiv-2310.07630-b31b1b.svg)](https://arxiv.org/abs/2310.07630) ![GitHub contributors](https://img.shields.io/github/contributors/aidos-lab/dect) ![GitHub](https://img.shields.io/github/license/aidos-lab/dect) [![Maintainability](https://api.codeclimate.com/v1/badges/82f86d7e2f0aae342055/maintainability)](https://codeclimate.com/github/aidos-lab/dect/maintainability)
+[![arXiv](https://img.shields.io/badge/arXiv-2310.07630-b31b1b.svg)](https://arxiv.org/abs/2310.07630) ![GitHub contributors](https://img.shields.io/github/contributors/aidos-lab/dect) ![GitHub](https://img.shields.io/github/license/aidos-lab/dect) [![Maintainability](https://qlty.sh/badges/b7958e48-8382-4fb0-ac0b-63b95b7a5426/maintainability.svg)](https://qlty.sh/gh/aidos-lab/projects/dect)
 
 This is the official implementation for the **Differentiable Euler Characteristic
 Transform**, a geometrical-topological method for shape classification. Our

dect/ect.py

@@ -78,6 +78,10 @@ def compute_ect(
     """
 
     # ecc.shape[0], index.max().item() + 1, ecc.shape[2],
+
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -165,6 +169,10 @@ def compute_ect_point_cloud(
         point clouds (thus ECT's), N is the number of direction and R is the
         resolution.
     """
+
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     lin = torch.linspace(
         start=-radius, end=radius, steps=resolution, device=x.device
     ).view(-1, 1, 1)
@@ -208,6 +216,9 @@ def compute_ect_points(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -273,6 +284,9 @@ def compute_ect_edges(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:
@@ -357,6 +371,9 @@ def compute_ect_mesh(
         The index tensor is assumed to start at 0.
     """
 
+    # ensure that the scale is in the right device
+    scale = torch.tensor([scale], device=x.device)
+
     if index is not None:
         batch_len = int(index.max() + 1)
     else:

tests/test_ect_edges.py

@@ -1,8 +1,40 @@
 """
 Tests the core functions for computing the
-ECT over a point cloud.
+ECT over a edges.
 """
 
+import pytest
+import torch
+import warnings
 
-def test_true():
-    pass
+from dect.directions import generate_uniform_directions
+from dect.ect import compute_ect_edges
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"])
+def test_compute_ect_edges_noindex(device):
+    """
+    Test the `compute_ect` function for point clouds.
+    """
+
+    # Check if device is available, else skip the test.
+    if not getattr(torch, device).is_available():
+        warnings.warn(f"Device {device} not available, skipping the tests.")
+        return
+
+    seed = 2024
+    ambient_dimension = 5
+    num_points = 10
+    v = generate_uniform_directions(
+        num_thetas=17, seed=seed, device=device, d=ambient_dimension
+    ).to(device)
+    x = torch.rand(size=(num_points, ambient_dimension), device=device)
+    edge_index = torch.tensor([[0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6]], device=device)
+    ect = compute_ect_edges(
+        x, edge_index=edge_index, v=v, radius=1, resolution=13, scale=10
+    )
+
+    assert ect.device.type == device
+
+    # TODO: Implement proper tests that the ect has been computed correctly.
+    # Most likely a parametrized set of fixtures are the way to go here.

tests/test_ect_faces.py

@@ -1,8 +1,47 @@
 """
 Tests the core functions for computing the
-ECT over a point cloud.
+ECT over a edges.
 """
 
+import pytest
+import torch
 
-def test_true():
-    pass
+from dect.directions import generate_uniform_directions
+from dect.ect import compute_ect_mesh
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"])
+def test_compute_ect_mesh_noindex(device):
+    """
+    Test the `compute_ect` function for point clouds.
+    """
+
+    # Check if device is available, else skip the test.
+    if not getattr(torch, device).is_available():
+        return
+
+    seed = 2024
+    ambient_dimension = 5
+    num_points = 10
+    v = generate_uniform_directions(
+        num_thetas=17, seed=seed, device=device, d=ambient_dimension
+    ).to(device)
+    x = torch.rand(size=(num_points, ambient_dimension), device=device)
+    edge_index = torch.tensor([[0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6]], device=device)
+    face_index = torch.tensor(
+        [[0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6], [2, 3, 4, 5, 6, 7]], device=device
+    )
+    ect = compute_ect_mesh(
+        x,
+        edge_index=edge_index,
+        face_index=face_index,
+        v=v,
+        radius=1,
+        resolution=13,
+        scale=10,
+    )
+
+    assert ect.device.type == device
+
+    # TODO: Implement proper tests that the ect has been computed correctly.
+    # Most likely a parametrized set of fixtures are the way to go here.

tests/test_ect_point_clouds.py

@@ -6,57 +6,48 @@
 third is the ambient dimension.
 Note that in this format each point in the point
 cloud has to have the same cardinality.
+
+
+NOTE: This will run on the devices available on
+the machine when executed. Currently it is only
+tested on GPU and CPU and MPS is skipped.
+If you run the tests on MPS and tests fail, please
+contact me.
 """
 
+import pytest
 import torch
 
-from dect.directions import generate_2d_directions, generate_uniform_directions
+from dect.directions import generate_uniform_directions
 from dect.ect import compute_ect_point_cloud
 
 
-def test_compute_ect_point_cloud_case_cpu():
+@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"])
+def test_compute_ect_point_cloud(device):
     """
     Tests the ECT computation of a point cloud.
     Differs in that it expects an input shape of
     size [B,N,D], where B is the batch size,
     N is the number of points and D is the ambient
     dimension.
     """
-    ambient_dimension = 4
-    num_pts = 100
-    batch_size = 8
-    num_thetas = 100
-    seed = 0
-    x = torch.rand(size=(batch_size, num_pts, ambient_dimension))
-    v = generate_uniform_directions(
-        num_thetas=num_thetas, d=ambient_dimension, device="cpu", seed=seed
-    )
-
-    ect = compute_ect_point_cloud(x, v, radius=10, resolution=30, scale=500)
-
-    assert ect[0].max() == num_pts
-    assert ect[0].min() == 0
 
+    if not getattr(torch, device).is_available():
+        return
 
-def test_compute_ect_point_cloud_case_cuda():
-    """
-    Tests the ECT computation of a point cloud.
-    Differs in that it expects an input shape of
-    size [B,N,D], where B is the batch size,
-    N is the number of points and D is the ambient
-    dimension.
-    """
     ambient_dimension = 4
     num_pts = 100
     batch_size = 8
     num_thetas = 100
     seed = 0
-    x = torch.rand(size=(batch_size, num_pts, ambient_dimension), device="cuda")
+    x = torch.rand(size=(batch_size, num_pts, ambient_dimension), device=device)
     v = generate_uniform_directions(
-        num_thetas=num_thetas, d=ambient_dimension, device="cuda", seed=seed
+        num_thetas=num_thetas, d=ambient_dimension, device=device, seed=seed
     )
 
     ect = compute_ect_point_cloud(x, v, radius=10, resolution=30, scale=500)
 
+    assert ect.device.type == device
+
     assert ect[0].max() == num_pts
     assert ect[0].min() == 0

tests/test_ect_points.py

@@ -3,53 +3,38 @@
 ECT over a point cloud.
 """
 
+import pytest
 import torch
 
-from dect.directions import generate_2d_directions, generate_uniform_directions
+from dect.directions import generate_uniform_directions
 from dect.ect import compute_ect_points
 
 """
 Test the ECT for 
 """
 
 
-def test_compute_ect_case_points():
+@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"])
+def test_compute_ect_case_points_noindex_cpu(device):
     """
     Test the `compute_ect` function for point clouds.
     """
-    device = "cpu"
+
+    # Check if device is available, else skip the test.
+    if not getattr(torch, device).is_available():
+        return
+
     # 2D Case
     seed = 2024
     ambient_dimension = 5
     num_points = 10
     v = generate_uniform_directions(
         num_thetas=17, seed=seed, device=device, d=ambient_dimension
     ).to(device)
-    x = torch.rand(size=(num_points, ambient_dimension))
-    ect = compute_ect_points(x, v=v, radius=1, resolution=13, scale=10)
-
-    assert ect.get_device() == -1
-
-    # Check that min and max are 0 and num_pts
-    torch.testing.assert_close(ect.max(), torch.tensor(num_points, dtype=torch.float32))
-    torch.testing.assert_close(ect.min(), torch.tensor(0.0, dtype=torch.float32))
-
-
-def test_compute_ect_case_points_cuda():
-    """
-    Test the `compute_ect` function for point clouds on the gpu.
-    """
-    if not torch.cuda.is_available():
-        return
-
-    device = "cuda:0"
-    # 2D Case
-    num_points = 10
-    v = generate_2d_directions(num_thetas=17).to(device)
-    x = torch.rand(size=(num_points, 2), device=device)
+    x = torch.rand(size=(num_points, ambient_dimension), device=device)
     ect = compute_ect_points(x, v=v, radius=1, resolution=13, scale=10)
 
-    assert ect.get_device() == 0  # 0 indicates cuda.
+    assert ect.device.type == device
 
     # Check that min and max are 0 and num_pts
     torch.testing.assert_close(