royerlab · JoOkuma · May 27, 2025 · May 27, 2025 · May 27, 2025 · May 27, 2025
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -0,0 +1,52 @@
+name: Benchmarks
+
+on:
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0  # Important for asv to access the git history
+        ref: ${{ github.ref }}
+
+    - name: Fetch main branch
+      run: |
+        git fetch origin main
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        python-version: "3.12"
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv pip install 'asv<0.6.2' virtualenv
+        uv pip install "."
+
+    - name: Run benchmarks
+      id: benchmark_output
+      run: |
+        cd benchmarks
+        asv machine --yes
+        asv run --python=same > benchmark_results.txt
+        cat benchmark_results.txt
+        echo "benchmark-results<<EOF" >> $GITHUB_OUTPUT
+        cat benchmark_results.txt >> $GITHUB_OUTPUT
+        echo "EOF" >> $GITHUB_OUTPUT
+
+    - name: Comment PR with benchmark results
+      uses: peter-evans/create-or-update-comment@v4
+      with:
+        issue-number: ${{ github.event.pull_request.number }}
+        body: |
+          ### Benchmark Results
+
+          ```
+          ${{ steps.benchmark_output.outputs.benchmark-results }}
+          ```
diff --git a/.gitignore b/.gitignore
@@ -181,3 +181,6 @@ man_track.txt
 
 # macOS files
 .DS_Store
+
+# ASV
+.asv/
diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json
@@ -0,0 +1,17 @@
+{
+    "version": 1,
+    "project": "tracksdata",
+    "project_url": "https://github.com/royerlab/tracksdata",
+    "repo": "..",
+    "branches": ["main"],
+    "environment_type": "virtualenv",
+    "show_commit_url": "https://github.com/royerlab/tracksdata/commit/",
+    "pythons": ["3.12"],
+    "benchmark_dir": "benchmarks",
+    "env_dir": ".asv/env",
+    "results_dir": ".asv/results",
+        "html_dir": ".asv/html",
+    "build_command": [
+        "pip install ."
+    ]
+}
diff --git a/benchmarks/benchmarks/__init__.py b/benchmarks/benchmarks/__init__.py
diff --git a/benchmarks/benchmarks/graph_backends.py b/benchmarks/benchmarks/graph_backends.py
@@ -0,0 +1,41 @@
+from tracksdata.edges._distance_edges import DistanceEdgesOperator
+from tracksdata.graph._base_graph import BaseGraphBackend
+from tracksdata.graph._rustworkx_graph import RustWorkXGraphBackend
+from tracksdata.nodes._random import RandomNodesOperator
+
+
+class GraphBackendSuite:
+    """
+    Benchmark suite for graph backend operations.
+    """
+
+    params = (
+        (RustWorkXGraphBackend,),
+        (1_000, 10_000, 100_000),
+    )
+    timeout = 1800  # 30 minutes
+    param_names = ("backend", "n_nodes")
+
+    def setup(
+        self,
+        backend: BaseGraphBackend,
+        n_nodes: int,
+    ) -> None:
+        self.graph = backend()
+        self.nodes_operator = RandomNodesOperator(
+            n_time_points=50,
+            n_nodes=(int(n_nodes * 0.95), int(n_nodes * 1.05)),
+            n_dim=3,
+            show_progress=False,
+        )
+        self.edges_operator = DistanceEdgesOperator(
+            distance_threshold=10,
+            n_neighbors=3,
+            show_progress=False,
+        )
+
+    def time_simple_workflow(self, *args, **kwargs) -> None:
+        # add nodes
+        self.nodes_operator.add_nodes(self.graph)
+        # add edges
+        self.edges_operator.add_edges(self.graph)
diff --git a/src/tracksdata/nodes/_random.py b/src/tracksdata/nodes/_random.py
@@ -0,0 +1,71 @@
+from typing import Any, Literal
+
+import numpy as np
+
+from tracksdata.graph._base_graph import BaseGraphBackend
+from tracksdata.nodes._base_nodes import BaseNodesOperator
+from tracksdata.utils._processing import maybe_show_progress
+
+
+class RandomNodesOperator(BaseNodesOperator):
+    def __init__(
+        self,
+        n_time_points: int,
+        n_nodes: tuple[int, int],
+        n_dim: Literal[2, 3] = 3,
+        random_state: int = 0,
+        show_progress: bool = False,
+    ):
8000        self.n_time_points = n_time_points
+        self.n_nodes = n_nodes
+
+        if n_dim == 2:
+            self.spatial_cols = ["x", "y"]
+        elif n_dim == 3:
+            self.spatial_cols = ["x", "y", "z"]
+        else:
+            raise ValueError(f"Invalid number of dimensions: {n_dim}")
+
+        self.rng = np.random.default_rng(random_state)
+        self._show_progress = show_progress
+
+    def add_nodes(
+        self,
+        graph: BaseGraphBackend,
+        *,
+        t: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if t is None:
+            for t in maybe_show_progress(
+                range(self.n_time_points),
+                desc="Processing time points",
+                show_progress=self._show_progress,
+            ):
+                self.add_nodes(graph, t=t, **kwargs)
+            return
+
+        if self.spatial_cols not in graph.node_features_keys:
+            graph.add_node_feature_key(self.spatial_cols, None)
+
+        n_nodes_at_t = self.rng.integers(
+            self.n_nodes[0],
+            self.n_nodes[1],
+        )
+
+        node_ids = []
+        coords = self.rng.uniform(
+            low=0,
+            high=1,
+            size=(n_nodes_at_t, len(self.spatial_cols)),
+        )
+
+        for c in coords:
+            node_id = graph.add_node(
+                {"t": t, **dict(zip(self.spatial_cols, c, strict=True))},
+                validate_keys=False,
+                **kwargs,
+            )
+            node_ids.append(node_id)
+
+        return graph
-Original file line number
+Diff line change
@@ Expand Up / @@ -181,3 +181,6 @@ man_track.txt @@
     # macOS files
     .DS_Store
+    # ASV
+    .asv/