add a tvb_kernels package

.github/workflows/build.yml

@@ -1,5 +1,7 @@
 name: Test Py
-on: [push]
+on:
+  push:
+    branches-ignore: [tvb_kernels]
 
 jobs:
   build:

.github/workflows/contrib-tests.yml

@@ -1,5 +1,7 @@
 name: Check tvb-contrib
-on: [push]
+on:
+  push:
+    branches-ignore: [tvb_kernels]
 
 jobs:
   build:

.github/workflows/kernels.yml

@@ -0,0 +1,80 @@
+name: build & test kernel library
+on:
+  push:
+
+# only run latest push
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  pip-tvbk:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: true
+      matrix:
+        python-version: [ "3.10", ]
+
+    steps:
+
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        id: setPy
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: install tools and dependencies
+        run: python3 -m pip install -U pip pytest scipy ctypesgen
+
+      - name: build & teste
+        run: |
+          cd tvb_kernels
+          pip install .
+          pytest
+
+
+  spack-tvbk:
+    runs-on: ubuntu-latest
+    steps:
+
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: setup spack
+        uses: spack/setup-spack@v2
+        with:
+          ref: develop
+          buildcache: true
+          color: true
+          path: spack
+
+      - name: install deps w/ spack
+        shell: spack-bash {0}
+        run: |
+          cd tvb_kernels
+          spack -e . concretize
+          spack -e . install
+          spack env activate .
+          spack env status
+          pip install pytest ctypesgen
+
+      - name: build & test
+        shell: spack-bash {0}
+        run: |
+          cd tvb_kernels
+          spack env activate .
+          pip install .
+          pytest
+
+      # https://github.com/spack/setup-spack?tab=readme-ov-file#example-caching-your-own-binaries-for-public-repositories
+      - name: Push packages and update index
+        run: |
+          cd tvb_kernels
+          spack -e . mirror set --push --oci-username ${{ github.actor }} --oci-password "${{ secrets.GITHUB_TOKEN }}" local-buildcache
+          spack -e . buildcache push --base-image ubuntu:22.04 --update-index local-buildcache
+        if: ${{ !cancelled() }}

.github/workflows/notebooks.yml

@@ -1,5 +1,7 @@
 name: Test Notebooks
-on: [push]
+on:
+  push:
+    branches-ignore: [tvb_kernels]
 
 jobs:
   build:
@@ -59,4 +61,4 @@ jobs:
 
       - name: run notebooks
         run: |
-          python ./tvb_build/notebook_runner.py ./tvb_documentation/demos
+          python ./tvb_build/notebook_runner.py ./tvb_documentation/demos

.github/workflows/pg-tests.yml

@@ -1,5 +1,7 @@
 name: Test PG
-on: [push]
+on:
+  push:
+    branches-ignore: [tvb_kernels]
 
 jobs:
   build:

.github/workflows/win-tests.yml

@@ -1,5 +1,7 @@
 name: Test Win
-on: [push]
+on:
+  push:
+    branches-ignore: [tvb_kernels]
 
 jobs:
   build:

tvb_kernels/.gitignore

@@ -0,0 +1,14 @@
+.cache
+build
+*.jpg
+build
+.spack-env
+spack.lock
+*.whl
+*.so
+*.o
+tvb_kernels/_ctg_tvbk.py
+tvbk.js
+tvbk.wasm
+compile_commands.json
+_ctg_*.py

tvb_kernels/README.md

@@ -0,0 +1,39 @@
+# tvb_kernels
+
+This is a library of computational kernels for TVB.
+
+## scope
+
+in order of priority
+
+- [x] sparse delay coupling functions
+- [ ] fused heun neural mass model step functions
+- [ ] neural ODE
+- [ ] bold / tavg monitors
+
+which will then be used internally by TVB.
+
+## building
+
+No stable versions are available, so users can install from
+source with `pip install .` or to just get a wheel `pip wheel .`
+
+For development, prefer
+```
+pip install nanobind scikit-build-core[pyproject]
+pip install --no-build-isolation -Ceditable.rebuild=true -ve .
+```
+
+## roadmap
+
+- improve api (ownership, checking etc)
+- variants
+  - single
+  - batches
+  - spatialized parameters
+- cuda ports for Jax, CuPy and Torch users 
+- mex functions?
+- scalable components
+  - domains
+  - projections
+  - small vm to automate inner loop

tvb_kernels/ghci-build.yml

@@ -0,0 +1 @@
+../.github/workflows/kernels.yml

tvb_kernels/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["hatchling", "hatch-build-scripts", "ctypesgen"]
+build-backend = "hatchling.build"
+
+[project]
+name = "tvb_kernels"
+version = "0.0.1"
+description = "TVB computational kernels"
+readme = "README.md"
+requires-python = ">=3.8"
+authors = [
+    { name = "Marmaduke Woodman", email = "[email protected]" },
+]
+classifiers = [ ]
+dependencies = [ "numpy" ]
+
+[project.urls]
+Homepage = "https://github.com/the-virtual-brain/tvb-root"
+
+[tool.cibuildwheel]
+build-verbosity = 1
+test-command = "pytest"
+test-requires = "pytest"
+
+skip = ["pp*"] # don't build pypy wheels
+archs = ["auto"]
+
+# [tool.hatch.version]
+# path = "tvb_kernels/version.py"
+
+[tool.hatch.build]
+include = [
+    "tvb_kernels/*.py",
+]
+
+[[tool.hatch.build.hooks.build-scripts.scripts]]
+work_dir = "src"
+out_dir = "tvb_kernels"
+commands = [ 
+    "make -B libtvbk.so",
+    "ctypesgen -l libtvbk.so tvbk.h > _ctg_tvbk.py",
+]
+artifacts = [ 
+    "libtvbk.so",
+    "_ctg_tvbk.py",
+]

tvb_kernels/spack.yaml

@@ -0,0 +1,14 @@
+spack:
+
+  specs:
+  - py-numpy
+  - py-scipy
+  - py-pip
+
+  concretizer:
+    unify: true
+
+  mirrors:
+    local-buildcache:
+      url: oci://ghcr.io/maedoc/spack-buildcache
+      signed: false

tvb_kernels/src/Makefile

@@ -0,0 +1,19 @@
+# run `bear -- make -B -j` to get compile_commands.json
+
+CC = gcc
+CFLAGS = -msse4.2 -O3 -funroll-loops -fopenmp-simd
+SRC = tvbk_conn.c
+OBJ = $(patsubst %.c,%.o,$(SRC))
+SO = libtvbk.so
+ctg = ../tvb_kernels/_ctg_tvbk.py
+
+all: $(SO) $(ctg)
+
+$(SO) : $(OBJ)
+	$(CC) -shared $< -o $@
+
+$(ctg): tvbk.h
+	ctypesgen -l $(SO) $< > $@
+
+tvbk.wasm: em.cpp tvbk.h $(SRC)
+	emcc -lembind -o tvbk.js em.cpp $(SRC)

tvb_kernels/src/em.cpp

@@ -0,0 +1,19 @@
+#include <emscripten/bind.h>
+
+using namespace emscripten;
+
+extern "C" {
+#include "tvbk.h"
+}
+
+class Conn {
+  const tvbk_conn conn;
+
+public:
+  Conn() : conn({0}) {}
+  void cx_nop(uint32_t t) { tvbk_cx_nop(&conn, t); }
+};
+
+EMSCRIPTEN_BINDINGS(tvb_kernels) {
+  class_<Conn>("Conn").constructor().function("cx_nop", &Conn::cx_nop);
+}

tvb_kernels/src/tvbk.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <stdint.h>
+
+typedef struct tvbk_params tvbk_params;
+
+struct tvbk_params {
+  const uint32_t count;
+  const float *const values;
+};
+
+/* a afferent coupling buffer into which the cx functions
+   accumulate their results */
+typedef struct tvbk_cx tvbk_cx;
+
+struct tvbk_cx {
+  /* values for 1st and 2nd Heun stage respectively.
+     each shaped (num_node, ) */
+  float *const cx1;
+  float *const cx2;
+  /* delay buffer (num_node, num_time)*/
+  float *const buf;
+  const uint32_t num_node;
+  const uint32_t num_time; // horizon, power of 2
+};
+
+typedef struct tvbk_conn tvbk_conn;
+
+struct tvbk_conn {
+  const int num_node;
+  const int num_nonzero;
+  const int num_cvar;
+  const float *const weights;    // (num_nonzero,)
+  const uint32_t *const indices; // (num_nonzero,)
+  const uint32_t *const indptr;  // (num_nodes+1,)
+  const uint32_t *const idelays; // (num_nonzero,)
+};
+
+/* not currently used */
+typedef struct tvbk_sim tvbk_sim;
+struct tvbk_sim {
+  // keep invariant stuff at the top, per sim stuff below
+  const int rng_seed;
+  const int num_node;
+  const int num_svar;
+  const int num_time;
+  const int num_params;
+  const int num_spatial_params;
+  const float dt;
+  const int oversample; // TODO "oversample" for stability,
+  const int num_skip;   // skip per output sample
+  float *z_scale;       // (num_svar), sigma*sqrt(dt)
+
+  // parameters
+  const tvbk_params global_params;
+  const tvbk_params spatial_params;
+
+  float *state_trace; // (num_time//num_skip, num_svar, num_nodes)
+  float *states;      // full states (num_svar, num_nodes)
+
+  const tvbk_conn conn;
+};
+
+void tvbk_cx_j(const tvbk_cx *cx, const tvbk_conn *conn, uint32_t t);
+void tvbk_cx_i(const tvbk_cx *cx, const tvbk_conn *conn, uint32_t t);
+void tvbk_cx_nop(const tvbk_cx *cx, const tvbk_conn *conn, uint32_t t);