Move main functionality to top-level of module + add example notebooks to CI

.github/workflows/ci.yml

@@ -70,6 +70,15 @@ jobs:
       - name: Upload coverage report
         uses: codecov/[email protected]
 
+      - name: Run examples
+        run: |
+          cd examples
+          python -m pip install -r requirements.txt
+          python -m pip install jupyter nbclient
+          jupyter execute estimate_hurst.ipynb
+          jupyter execute generative_model.ipynb
+          jupyter execute inversion.ipynb
+
   dist:
     needs: [pre-commit]
     name: Distribution build

.gitignore

@@ -156,3 +156,7 @@ Thumbs.db
 # Common editor files
 *~
 *.swp
+
+# notebooks
+Untitiled.ipynb
+t.ipynb

README.md

@@ -31,8 +31,8 @@ Basic usage
 ```python
 import jax
 import jax.random as jrandom
+import signax
 
-from signax.signature import signature
 
 key = jrandom.PRNGKey(0)
 depth = 3
@@ -41,17 +41,17 @@ depth = 3
 length = 100
 dim = 20
 path = jrandom.normal(shape=(length, dim), key=key)
-output = signature(path, depth)
+output = signax.signature(path, depth)
 # output is a list of array representing tensor algebra
 
 # compute signature for batches (multiple) of paths
 # this is done via `jax.vmap`
 batch_size = 20
 path = jrandom.normal(shape=(batch_size, length, dim), key=key)
-output = jax.vmap(lambda x: signature(x, depth))(path)
+output = jax.vmap(lambda x: signax.signature(x, depth))(path)
 ```
 
-Integrate with [equinox](https://github.com/patrick-kidger/equinox) library
+Integrate with the [equinox](https://github.com/patrick-kidger/equinox) library
 
 ```python
 import equinox as eqx

examples/compare.ipynb

@@ -26,10 +26,7 @@
     "import numpy as np\n",
     "import signatory\n",
     "import torch\n",
-    "from signax.signature import signature\n",
-    "\n",
-    "# set white background\n",
-    "plt.rcParams[\"figure.facecolor\"] = \"white\""
+    "from signax import signature"
    ]
   },
   {

examples/nets.py

@@ -7,8 +7,8 @@
 import jax.numpy as jnp
 import jax.random as jrandom
 
+from signax import signature, signature_combine
 from signax.module import SignatureTransform
-from signax.signature import signature, signature_combine
 from signax.utils import flatten
 
 
@@ -70,6 +70,8 @@ def __init__(
     def __call__(
         self,
         x: jnp.ndarray,
+        *,
+        key=None,
     ):
         """x size (length, dim)"""
         length, _ = x.shape
@@ -164,7 +166,7 @@ def __init__(self, length, adjusted_length, signature_depth=2) -> None:
         self.adjusted_length = adjusted_length
         self.signature_depth = signature_depth
 
-    def __call__(self, x):
+    def __call__(self, x, *, key=None):
         """
         Transform input `x` into a smaller window.
         Each window starts at index 0 with increasing size according

examples/requirements.txt

@@ -0,0 +1,2 @@
+matplotlib
+optax

pyproject.toml

@@ -150,3 +150,5 @@ isort.required-imports = ["from __future__ import annotations"]
 [tool.ruff.per-file-ignores]
 "tests/**" = ["T20"]
 "noxfile.py" = ["T20"]
+"src/signax/module.py" = ["ARG002"]  # unused argument key in __call__ of Module (equinox)
+"examples/nets.py" = ["ARG002"]      # unused argument key in __call__ of Module (equinox)

src/signax/__init__.py

@@ -2,4 +2,25 @@
 
 __version__ = "0.1.2"
 
-__all__ = ("__version__",)
+__all__ = (
+    "__version__",
+    "module",
+    "utils",
+    "tensor_ops",
+    "signature",
+    "logsignature",
+    "signature_combine",
+    "signature_to_logsignature",
+    "multi_signature_combine",
+    "signature_batch",
+)
+
+from signax import module, tensor_ops, utils
+from signax.signatures import (
+    logsignature,
+    multi_signature_combine,
+    signature,
+    signature_batch,
+    signature_combine,
+    signature_to_logsignature,
+)

src/signax/module.py

@@ -1,9 +1,17 @@
 from __future__ import annotations
 
+from typing import Any
+
+__all__ = (
+    "SignatureTransform",
+    "SignatureCombine",
+)
+
 import equinox as eqx
 import jax
 
-from signax.signature_flattened import signature, signature_combine
+from signax.signatures import signature, signature_combine
+from signax.utils import flatten, unravel_signature
 
 
 class SignatureTransform(eqx.Module):
@@ -15,8 +23,10 @@ def __init__(self, depth: int):
     def __call__(
         self,
         path: jax.Array,
+        *,
+        key: Any | None = None,
     ) -> jax.Array:
-        return signature(path, self.depth)
+        return flatten(signature(path, self.depth))
 
 
 class SignatureCombine(eqx.Module):
@@ -28,4 +38,6 @@ def __init__(self, dim: int, depth: int):
         self.depth = depth
 
     def __call__(self, signature1: jax.Array, signature2: jax.Array):
-        return signature_combine(signature1, signature2, self.dim, self.depth)
+        sig1 = unravel_signature(signature1, self.dim, self.depth)
+        sig2 = unravel_signature(signature2, self.dim, self.depth)
+        return flatten(signature_combine(sig1, sig2))

src/signax/signature.py → src/signax/signatures.py

@@ -1,12 +1,21 @@
 from __future__ import annotations
 
+__all__ = (
+    "signature",
+    "logsignature",
+    "signature_batch",
+    "signature_combine",
+    "signature_to_logsignature",
+    "multi_signature_combine",
+)
+
 from functools import partial
 
 import jax
 import jax.numpy as jnp
 
-from .tensor_ops import log, mult, mult_fused_restricted_exp, restricted_exp
-from .utils import compress, lyndon_words
+from signax.tensor_ops import log, mult, mult_fused_restricted_exp, restricted_exp
+from signax.utils import compress, lyndon_words
 
 
 @partial(jax.jit, static_argnames="depth")

src/signax/utils.py

@@ -1,5 +1,14 @@
 from __future__ import annotations
 
+__all__ = (
+    "index_select",
+    "lyndon_words",
+    "compress",
+    "unravel_signature",
+    "flatten",
+    "term_at",
+)
+
 from collections import defaultdict
 from functools import partial
 from typing import cast

tests/test_signature.py

@@ -6,7 +6,7 @@
 import torch
 from numpy.random import default_rng
 
-from signax.signature import (
+from signax import (
     multi_signature_combine,
     signature,
     signature_batch,

tests/test_tensor_ops.py

@@ -9,7 +9,7 @@
 import torch
 from numpy.random import default_rng
 
-from signax.signature import signature, signature_to_logsignature
+from signax import signature, signature_to_logsignature
 from signax.tensor_ops import (
     addcmul,
     mult,