Add runtime tests

.github/workflows/regression_tests.yml

@@ -0,0 +1,39 @@
+# .github/workflows/regression_tests.yml
+name: Regression Tests
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  regression_tests:
+    name: regression_tests
+    runs-on: ubuntu-20.04
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        lfs: true
+        fetch-depth: 0 # This ensures we can checkout main branch too
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+        architecture: 'x64'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e ".[dev]"
+
+    - name: Run benchmarks and compare to baseline
+      if: github.event.pull_request.base.ref == 'main'
+      run: |
+        # Check if regression test results exist in main branch
+        if [ -f 'git cat-file -e main:tests/regression_test_baselines.json' ]; then
+          git checkout main tests/regression_test_baselines.json
+        else
+          echo "No regression test results found in main branch"
+        fi
+        pytest -m regression

.github/workflows/tests.yml

@@ -39,4 +39,4 @@ jobs:
     - name: Test with pytest
       run: |
         pip install pytest pytest-cov
-        pytest tests/ --cov=jaxley --cov-report=xml
+        pytest tests/ -m "not regression" --cov=jaxley --cov-report=xml

.github/workflows/update_regression_baseline.yml

@@ -0,0 +1,88 @@
+# .github/workflows/update_regression_tests.yml
+
+# for details on triggering a workflow from a comment, see:
+# https://dev.to/zirkelc/trigger-github-workflow-for-comment-on-pull-request-45l2
+name: Update Regression Baseline
+
+on:
+  issue_comment: # trigger from comment; event runs on the default branch
+    types: [created]
+  # pull_request:
+  #   branches:
+  #     - main
+
+jobs:
+  update_regression_tests:
+    name: update_regression_tests
+    runs-on: ubuntu-20.04
+    # Trigger from a comment
+    # if: github.event.issue.pull_request && contains(github.event.comment.body, '/update_regression_baselines')
+    permissions:
+      contents: write
+      pull-requests: write
+
+    steps:
+    - name: Get PR branch
+      uses: xt0rted/pull-request-comment-branch@v1
+      id: comment-branch
+
+    - name: Checkout PR branch
+      uses: actions/checkout@v3
+      with:
+        # ref: ${{ steps.comment-branch.outputs.head_sha }} # using head_sha vs. head_ref makes this work for forks
+        lfs: true
+        fetch-depth: 0 # This ensures we can checkout main branch too
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+        architecture: 'x64'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e ".[dev]"
+
+    - name: Update baseline
+      if: github.event.pull_request.base.ref == 'main'
+      run: |
+        git config --global user.name '${{ github.event.pull_request.user.login }}'
+        git config --global user.email '${{ github.event.pull_request.user.login }}@users.noreply.github.com'
+        # Check if regression test results exist in main branch
+        if [ -f 'git cat-file -e main:tests/regression_test_baselines.json' ]; then
+          git checkout main tests/regression_test_baselines.json
+        else
+          echo "No regression test results found in main branch"
+        fi
+        NEW_BASELINE=1 pytest -m regression
+
+    - name: Add Baseline update report to PR comment
+      uses: actions/github-script@v7
+      if: github.event.pull_request.base.ref == 'main' # might need `always()` to work for comments
+      with:
+        github-token: ${{ secrets.GITHUB_TOKEN }}
+        script: |
+          const fs = require('fs');
+          const TestReport = fs.readFileSync('tests/regression_test_report.txt', 'utf8');
+
+          await github.rest.issues.createComment({
+            issue_number: context.issue.number,
+            owner: context.repo.owner,
+            repo: context.repo.repo,
+            body: `## New Baselines \n\`\`\`\n${TestReport}\n\`\`\``
+          });
+
+    - name: Commit and push
+      if: github.event.pull_request.base.ref == 'main'
+      run: |
+        git add -f tests/regression_test_baselines.json # since it's in .gitignore
+        git commit -m "Update regression test baselines"
+        git push origin HEAD:${{ github.head_ref }}
+
+    # Needed when workflow is triggered from a comment
+    # - name: Commit and push
+    #   if: github.event.pull_request.base.ref == 'main'
+    #   run: |
+    #     git add -f tests/regression_test_baselines.json # since it's in .gitignore
+    #     git commit -m "Update regression test baselines"
+    #     git push origin HEAD:${{ steps.comment-branch.outputs.head_sha }}

.gitignore

@@ -55,6 +55,8 @@ coverage.xml
 *.py,cover
 .hypothesis/
 .pytest_cache/
+tests/regression_test_results.json
+tests/regression_test_baselines.json
 
 # Translations
 *.mo

pyproject.toml

@@ -69,6 +69,7 @@ dev = [
 [tool.pytest.ini_options]
 markers = [
     "slow: marks tests as slow (T > 10s)",
+    "regression: marks regression tests",
 ]
 
 [tool.isort]

tests/conftest.py

@@ -1,6 +1,7 @@
 # This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
 # licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import json
 import os
 from copy import deepcopy
 from typing import Optional
@@ -9,6 +10,7 @@
 
 import jaxley as jx
 from jaxley.synapses import IonotropicSynapse
+from tests.test_regression import generate_regression_report, load_json
 
 
 @pytest.fixture(scope="session")
@@ -202,3 +204,44 @@ def get_or_compute_swc2jaxley_params(
 
     yield get_or_compute_swc2jaxley_params
     params = {}
+
+
+@pytest.fixture(scope="session", autouse=True)
+def print_session_report(request, pytestconfig):
+    """Cleanup a testing directory once we are finished."""
+    NEW_BASELINE = os.environ["NEW_BASELINE"] if "NEW_BASELINE" in os.environ else 0
+
+    dirname = os.path.dirname(__file__)
+    baseline_fname = os.path.join(dirname, "regression_test_baselines.json")
+    results_fname = os.path.join(dirname, "regression_test_results.json")
+
+    collected_regression_tests = [
+        item for item in request.session.items if item.get_closest_marker("regression")
+    ]
+
+    def update_baseline():
+        if NEW_BASELINE:
+            results = load_json(results_fname)
+            with open(baseline_fname, "w") as f:
+                json.dump(results, f, indent=2)
+            os.remove(results_fname)
+
+    def print_regression_report():
+        baselines = load_json(baseline_fname)
+        results = load_json(results_fname)
+
+        report = generate_regression_report(baselines, results)
+        # "No baselines found. Run `git checkout main;UPDATE_BASELINE=1 pytest -m regression; git checkout -`"
+        with open(dirname + "/regression_test_report.txt", "w") as f:
+            f.write(report)
+
+        # the following allows to print the report to the console despite pytest
+        # capturing the output and without specifying the "-s" flag
+        capmanager = request.config.pluginmanager.getplugin("capturemanager")
+        with capmanager.global_and_fixture_disabled():
+            print("\n\n\nRegression Test Report\n----------------------\n")
+            print(report)
+
+    if len(collected_regression_tests) > 0:
+        request.addfinalizer(update_baseline)
+        request.addfinalizer(print_regression_report)

tests/regression_test_baselines.json

@@ -0,0 +1,92 @@
+{
+  "ec3a4fad11d2bfb1bc5f8f10529cb06f2ff9919b377e9c0a3419c7f7f237f06e": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 1,
+      "artificial": false,
+      "connect": false,
+      "connection_prob": 0.0,
+      "voltage_solver": "jaxley.stone"
+    },
+    "runtimes": {
+      "build_time": 0.5663588841756185,
+      "compile_time": 18.334174553553265,
+      "run_time": 2.8035006523132324
+    }
+  },
+  "128cfe30d4ffb9c1abd9dc0fa25b0e86940437b3eb1d46584e21f2c780ed78e8": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 1,
+      "artificial": false,
+      "connect": false,
+      "connection_prob": 0.0,
+      "voltage_solver": "jax.sparse"
+    },
+    "runtimes": {
+      "build_time": 0.39449914296468097,
+      "compile_time": 3.041297356287638,
+      "run_time": 2.4506944020589194
+    }
+  },
+  "45cb5fa937517154a8d7bd2ac6d4542ff66c7cd3f5199976706ae44134eec301": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 10,
+      "artificial": false,
+      "connect": true,
+      "connection_prob": 0.1,
+      "voltage_solver": "jaxley.stone"
+    },
+    "runtimes": {
+      "build_time": 1.836039702097575,
+      "compile_time": 29.307872613271076,
+      "run_time": 18.758193572362263
+    }
+  },
+  "872ba2d409d18daf5e0e947953385c3d0967087ed122f72ba01990397429318e": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 10,
+      "artificial": false,
+      "connect": true,
+      "connection_prob": 0.1,
+      "voltage_solver": "jax.sparse"
+    },
+    "runtimes": {
+      "build_time": 1.2564191023508708,
+      "compile_time": 26.581148147583008,
+      "run_time": 25.064125458399456
+    }
+  },
+  "da2f14fe319cf40d2ec65fdde6f3e0c997ef803e637d1ae7d2f2846c2369dbb2": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 1000,
+      "artificial": true,
+      "connect": true,
+      "connection_prob": 0.001,
+      "voltage_solver": "jaxley.stone"
+    },
+    "runtimes": {
+      "build_time": 109.68021249771118,
+      "compile_time": 45.64337452252706,
+      "run_time": 41.48241376876831
+    }
+  },
+  "4a250131b8b31132e19d9f82ececa4d2dc26b0678326089f8a2c9de3696418fc": {
+    "test_name": "test_runtime",
+    "input_kwargs": {
+      "num_cells": 1000,
+      "artificial": true,
+      "connect": true,
+      "connection_prob": 0.001,
+      "voltage_solver": "jax.sparse"
+    },
+    "runtimes": {
+      "build_time": 108.68635201454163,
+      "compile_time": 61.08988928794861,
+      "run_time": 58.32885948816935
+    }
+  }
+}