sigtorch 0.1.0__py3-none-any.whl

sigtorch/__init__.py

@@ -0,0 +1,23 @@
+"""sigtorch: differentiable signatures and logsignatures in pure PyTorch, with the API of signatory.
+
+``import sigtorch as signatory`` is intended to work as a drop-in replacement for code written against signatory 1.2.
+"""
+
+from .augment import Augment
+from .deprecated import max_parallelism
+from .logsignature_module import (LogSignature, Logsignature, SignatureToLogSignature, SignatureToLogsignature,
+                                  logsignature, logsignature_channels, signature_to_logsignature)
+from .path import Path
+from .signature_inversion_module import invert_signature
+from .signature_module import (Signature, extract_signature_term, multi_signature_combine, signature,
+                               signature_channels, signature_combine)
+from .utility import all_words, lyndon_brackets, lyndon_words
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Augment", "LogSignature", "Logsignature", "Path", "Signature", "SignatureToLogSignature",
+    "SignatureToLogsignature", "all_words", "extract_signature_term", "invert_signature", "logsignature",
+    "logsignature_channels", "lyndon_brackets", "lyndon_words", "max_parallelism", "multi_signature_combine",
+    "signature", "signature_channels", "signature_combine", "signature_to_logsignature",
+]

sigtorch/_backends.py

@@ -0,0 +1,89 @@
+"""Signature levels of a piecewise-linear path, given its increments.
+
+Depths 1-3 use closed forms built from cumulative sums and batched matrix products, which never materialise the
+per-step level-2/3 tensors of the non-stream signature. Deeper signatures combine per-segment exponentials with Chen's
+identity in a balanced tree (non-stream) or a parallel prefix scan (stream), so the number of sequential steps grows
+with log(length) rather than length.
+"""
+
+import torch
+
+from . import _tensor_algebra as ta
+
+
+def signature_levels(dX: torch.Tensor, depth: int, stream: bool) -> ta.Levels:
+    """Signature levels of the path with increments ``dX`` of shape (B, N, C).
+
+    Returns levels of shape (B, C**k), or (B, N, C**k) holding the signature of every prefix if ``stream``.
+    """
+    if depth <= 3:
+        return _closed_form_stream(dX, depth) if stream else _closed_form(dX, depth)
+    return _scan(dX, depth) if stream else _tree(dX, depth)
+
+
+def _exclusive_cumsum(x: torch.Tensor) -> torch.Tensor:
+    """Cumulative sum over dim 1 of everything strictly before each step."""
+    return x.cumsum(dim=1) - x
+
+
+def _closed_form(dX: torch.Tensor, depth: int) -> ta.Levels:
+    batch, _, channels = dX.shape
+    levels = [dX.sum(dim=1)]
+    if depth >= 2:
+        before = _exclusive_cumsum(dX)  # S1 up to the start of each segment
+        # sum_t S1_before[t] (x) dX[t]  +  sum_t dX[t] (x) dX[t] / 2
+        S2 = torch.bmm(before.transpose(1, 2), dX) + 0.5 * torch.bmm(dX.transpose(1, 2), dX)
+        levels.append(S2.reshape(batch, channels ** 2))
+    if depth >= 3:
+        outer = torch.einsum('bti,btj->btij', dX, dX)
+        S2_before = _exclusive_cumsum(torch.einsum('bti,btj->btij', before, dX) + 0.5 * outer)
+        # Chen: S3 += S2_before (x) dX + S1_before (x) dX^2 / 2 + dX^3 / 6, summed over segments
+        S3 = (torch.einsum('btij,btk->bijk', S2_before, dX)
+              + 0.5 * torch.einsum('bti,btjk->bijk', before, outer)
+              + torch.einsum('bti,btjk->bijk', dX, outer) / 6)
+        levels.append(S3.reshape(batch, channels ** 3))
+    return levels
+
+
+def _closed_form_stream(dX: torch.Tensor, depth: int) -> ta.Levels:
+    batch, length, channels = dX.shape
+    S1 = dX.cumsum(dim=1)
+    levels = [S1]
+    if depth >= 2:
+        before = S1 - dX
+        outer = torch.einsum('bti,btj->btij', dX, dX)
+        S2 = (torch.einsum('bti,btj->btij', before, dX) + 0.5 * outer).cumsum(dim=1)
+        levels.append(S2.reshape(batch, length, channels ** 2))
+    if depth >= 3:
+        S2_before = torch.cat([torch.zeros_like(S2[:, :1]), S2[:, :-1]], dim=1)
+        S3 = (torch.einsum('btij,btk->btijk', S2_before, dX)
+              + 0.5 * torch.einsum('bti,btjk->btijk', before, outer)
+              + torch.einsum('bti,btjk->btijk', dX, outer) / 6).cumsum(dim=1)
+        levels.append(S3.reshape(batch, length, channels ** 3))
+    return levels
+
+
+def _tree(dX: torch.Tensor, depth: int) -> ta.Levels:
+    """Multiplies the segment exponentials pairwise, halving their number each round."""
+    levels = ta.restricted_exp(dX, depth)
+    while levels[0].size(1) > 1:
+        tail = None
+        if levels[0].size(1) % 2:  # an odd one out joins the next round at the end
+            tail = [x[:, -1:] for x in levels]
+            levels = [x[:, :-1] for x in levels]
+        levels = ta.mult([x[:, 0::2] for x in levels], [x[:, 1::2] for x in levels])
+        if tail is not None:
+            levels = [torch.cat([x, t], dim=1) for x, t in zip(levels, tail)]
+    return [x[:, 0] for x in levels]
+
+
+def _scan(dX: torch.Tensor, depth: int) -> ta.Levels:
+    """Inclusive prefix products of the segment exponentials (Hillis-Steele scan)."""
+    levels = ta.restricted_exp(dX, depth)
+    length = levels[0].size(1)
+    offset = 1
+    while offset < length:
+        combined = ta.mult([x[:, :-offset] for x in levels], [x[:, offset:] for x in levels])
+        levels = [torch.cat([x[:, :offset], c], dim=1) for x, c in zip(levels, combined)]
+        offset *= 2
+    return levels

sigtorch/_tensor_algebra.py

@@ -0,0 +1,75 @@
+"""Operations on elements of the truncated tensor algebra.
+
+An element is a list ``[x_1, ..., x_depth]`` where ``x_k`` has shape ``(..., C**k)``: the level-k tensor flattened in
+row-major order, after any number of leading batch dimensions. The scalar term is implicit: 1 for signatures (group
+elements), 0 for log-signatures (Lie elements). All operations broadcast over the leading dimensions.
+"""
+
+from typing import List, Optional
+
+import torch
+
+Levels = List[torch.Tensor]
+
+
+def split_levels(flat: torch.Tensor, channels: int, depth: int) -> Levels:
+    """Splits a ``(..., C + C**2 + ... + C**depth)`` tensor into its levels."""
+    return list(torch.split(flat, [channels ** k for k in range(1, depth + 1)], dim=-1))
+
+
+def _outer(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """(..., m) x (..., n) -> (..., m * n), the flattened tensor product."""
+    return (a.unsqueeze(-1) * b.unsqueeze(-2)).flatten(-2)
+
+
+def restricted_exp(x: torch.Tensor, depth: int) -> Levels:
+    """exp(x) for a level-1 element ``x`` of shape (..., C): the signature of a single straight segment."""
+    levels = [x]
+    for k in range(2, depth + 1):
+        levels.append(_outer(levels[-1], x) / k)
+    return levels
+
+
+def _product(a: List[Optional[torch.Tensor]], b: List[Optional[torch.Tensor]]) -> List[Optional[torch.Tensor]]:
+    """Truncated product of two elements with zero scalar term. ``None`` marks a level known to be zero."""
+    depth = len(a)
+    out: List[Optional[torch.Tensor]] = [None] * depth
+    for k in range(1, depth):  # output level k + 1
+        for i in range(k):  # a at level i + 1, b at level k - i
+            if a[i] is not None and b[k - 1 - i] is not None:
+                term = _outer(a[i], b[k - 1 - i])
+                out[k] = term if out[k] is None else out[k] + term
+    return out
+
+
+def mult(a: Levels, b: Levels) -> Levels:
+    """(1 + a) * (1 + b) - 1: the product of two group elements, e.g. Chen's identity for signatures."""
+    cross = _product(a, b)
+    return [x + y if c is None else x + y + c for x, y, c in zip(a, b, cross)]
+
+
+def log(a: Levels) -> Levels:
+    """log(1 + a) = a - a^2 / 2 + a^3 / 3 - ..., truncated."""
+    result = list(a)
+    power: List[Optional[torch.Tensor]] = list(a)
+    for n in range(2, len(a) + 1):
+        power = _product(power, a)
+        coeff = (-1) ** (n + 1) / n
+        result = [r if p is None else r + coeff * p for r, p in zip(result, power)]
+    return result
+
+
+def inverse(a: Levels) -> Levels:
+    """Inverse of a group-like element such as a signature: reverse every word and multiply level k by (-1)^k.
+
+    This equals the signature of the time-reversed path.
+    """
+    channels = a[0].size(-1)
+    out = []
+    for k, x in enumerate(a, start=1):
+        batch_shape = x.shape[:-1]
+        nb = len(batch_shape)
+        x = x.reshape(*batch_shape, *([channels] * k))
+        x = x.permute(*range(nb), *range(nb + k - 1, nb - 1, -1)).reshape(*batch_shape, channels ** k)
+        out.append(x if k % 2 == 0 else -x)
+    return out

sigtorch/augment.py

@@ -0,0 +1,69 @@
+"""A torch.nn.Module for augmenting streams of data before taking signatures.
+
+Adapted from signatory's augment.py (Copyright 2019 Patrick Kidger, Apache License 2.0, see LICENSE-signatory);
+rewritten for sigtorch's pure-PyTorch implementation.
+"""
+
+from typing import Callable, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class Augment(nn.Module):
+    """Sweeps a small feedforward network along a stream, as in Bonnier et al. 2019, "Deep Signature Transforms".
+
+    The first layer is a ``Conv1d`` with the given kernel, stride, padding and dilation; later layers are pointwise
+    ``Conv1d`` layers with ``activation`` in between (so the network's output is linear). Input has shape
+    (batch, stream, in_channels); the output has the network's output channels, preceded by the original stream (if
+    ``include_original``) and a time channel running from 0 to 1 (if ``include_time``). Those two are aligned with
+    the last point under each kernel window, so they assume stride 1 and no padding.
+    """
+
+    def __init__(self, in_channels: int, layer_sizes: Union[int, Tuple[int, ...]], kernel_size: int, stride: int = 1,
+                 padding: int = 0, dilation: int = 1, bias: bool = True,
+                 activation: Callable[[torch.Tensor], torch.Tensor] = F.relu, include_original: bool = True,
+                 include_time: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        if isinstance(layer_sizes, int):
+            layer_sizes = (layer_sizes,)
+        self.in_channels = in_channels
+        self.layer_sizes = layer_sizes
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.bias = bias
+        self.activation = activation
+        self.include_original = include_original
+        self.include_time = include_time
+
+        self.convs = nn.ModuleList()
+        if layer_sizes:
+            self.convs.append(nn.Conv1d(in_channels, layer_sizes[0], kernel_size=kernel_size, stride=stride,
+                                        padding=padding, dilation=dilation, bias=bias))
+            for in_size, out_size in zip(layer_sizes[:-1], layer_sizes[1:]):
+                self.convs.append(nn.Conv1d(in_size, out_size, kernel_size=1, bias=bias))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 3:
+            raise RuntimeError(f"Argument x should have three dimensions, (batch, stream, channel). Given shape "
+                               f"{tuple(x.shape)}.")
+        length = x.size(1) - self.kernel_size + 1
+        pieces = []
+        if self.include_original:
+            pieces.append(x.narrow(1, self.kernel_size - 1, length))
+        if self.include_time:
+            time = torch.linspace(0, 1, length, dtype=x.dtype, device=x.device)
+            pieces.append(time.view(1, length, 1).expand(x.size(0), length, 1))
+        if self.layer_sizes:
+            out = self.convs[0](x.transpose(1, 2))
+            for conv in self.convs[1:]:
+                out = conv(self.activation(out))
+            pieces.append(out.transpose(1, 2))
+        return torch.cat(pieces, dim=2)
+
+    def extra_repr(self) -> str:
+        return (f'activation={self.activation}, include_original={self.include_original}, '
+                f'include_time={self.include_time}')

sigtorch/deprecated.py

@@ -0,0 +1,21 @@
+"""Deprecated interfaces, kept so that code written for signatory keeps running.
+
+Adapted from signatory's deprecated.py (Copyright 2019 Patrick Kidger, Apache License 2.0, see LICENSE-signatory);
+rewritten for sigtorch's pure-PyTorch implementation.
+"""
+
+import warnings
+from typing import Optional
+
+_max_parallelism = [8]
+
+
+def max_parallelism(value: Optional[int] = None) -> int:
+    """Deprecated and does nothing (as in signatory). Use torch.set_num_threads / torch.get_num_threads instead."""
+    warnings.warn("max_parallelism is deprecated and now does nothing. Use torch.set_num_threads and "
+                  "torch.get_num_threads to control the amount of parallelism.", DeprecationWarning)
+    if value == -1:
+        value = float('inf')
+    if value is not None:
+        _max_parallelism[0] = value
+    return _max_parallelism[0]

sigtorch/logsignature_module.py

@@ -0,0 +1,145 @@
+"""The logsignature transform."""
+
+from typing import Union
+
+import torch
+from torch import nn
+
+from . import _tensor_algebra as ta
+from . import signature_module as smodule
+from . import utility
+
+_MODES = ("words", "brackets", "expand")
+
+
+def _check_mode(mode: str) -> None:
+    if mode not in _MODES:
+        raise ValueError("Invalid values for argument 'mode'. Valid values are 'expand', 'brackets', or 'words'.")
+
+
+def _mobius(n: int) -> int:
+    result, p = 1, 2
+    while p * p <= n:
+        if n % p == 0:
+            n //= p
+            if n % p == 0:
+                return 0
+            result = -result
+        p += 1
+    return -result if n > 1 else result
+
+
+def logsignature_channels(in_channels: int, depth: int) -> int:
+    """Number of channels of a logsignature with ``mode`` "words" or "brackets" (the dimension of the truncated free
+    Lie algebra, from Witt's formula)."""
+    if in_channels < 1:
+        raise ValueError("in_channels must be at least 1")
+    if depth < 1:
+        raise ValueError("depth must be at least 1")
+    return sum(sum(_mobius(k // d) * in_channels ** d for d in range(1, k + 1) if k % d == 0) // k
+               for k in range(1, depth + 1))
+
+
+def _check_signature_tensor(signature: torch.Tensor, channels: int, depth: int, stream: bool,
+                            scalar_term: bool) -> None:
+    ndim, dims = (3, "(batch, stream, channel)") if stream else (2, "(batch, channel)")
+    if signature.ndim != ndim:
+        raise ValueError(f"Argument 'signature' must be a {ndim}-dimensional tensor, with dimensions corresponding "
+                         f"to {dims} respectively.")
+    if 0 in signature.shape:
+        raise ValueError("Argument 'signature' cannot have dimensions of size zero.")
+    if signature.size(-1) != smodule.signature_channels(channels, depth, scalar_term):
+        raise ValueError("Argument 'signature' has the wrong number of channels for the specified channels and depth.")
+    if not signature.is_floating_point():
+        raise ValueError("Argument 'signature' must be of floating point type.")
+
+
+class SignatureToLogSignature(nn.Module):
+    """:class:`torch.nn.Module` wrapper around :func:`sigtorch.signature_to_logsignature`."""
+
+    def __init__(self, channels: int, depth: int, stream: bool = False, mode: str = "words",
+                 scalar_term: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        _check_mode(mode)
+        utility._check_channels_depth(channels, depth)
+        self._channels = channels
+        self._depth = depth
+        self._stream = stream
+        self._mode = mode
+        self._scalar_term = scalar_term
+
+    def forward(self, signature: torch.Tensor) -> torch.Tensor:
+        _check_signature_tensor(signature, self._channels, self._depth, self._stream, self._scalar_term)
+        levels = ta.split_levels(signature[..., int(self._scalar_term):], self._channels, self._depth)
+        levels = ta.log(levels)
+        if self._mode == "expand":
+            return torch.cat(levels, dim=-1)
+        out = []
+        for length, level in enumerate(levels, start=1):
+            words = level[..., utility.lyndon_flat_indices(self._channels, length, level.device)]
+            if self._mode == "brackets":
+                words = utility.words_to_brackets(words, self._channels, length)
+            out.append(words)
+        return torch.cat(out, dim=-1)
+
+    def extra_repr(self) -> str:
+        return (f'channels={self._channels}, depth={self._depth}, stream={self._stream}, mode={self._mode!r}')
+
+
+SignatureToLogsignature = SignatureToLogSignature
+
+
+def signature_to_logsignature(signature: torch.Tensor, channels: int, depth: int, stream: bool = False,
+                              mode: str = "words", scalar_term: bool = False) -> torch.Tensor:
+    """The logsignature corresponding to a signature computed by :func:`sigtorch.signature` with the same
+    ``channels``, ``depth``, ``stream`` and ``scalar_term``. ``mode`` as :func:`sigtorch.logsignature`."""
+    return SignatureToLogSignature(channels, depth, stream, mode, scalar_term)(signature)
+
+
+class LogSignature(nn.Module):
+    """:class:`torch.nn.Module` wrapper around :func:`sigtorch.logsignature`."""
+
+    def __init__(self, depth: int, stream: bool = False, inverse: bool = False, mode: str = "words", **kwargs):
+        super().__init__(**kwargs)
+        _check_mode(mode)
+        self._depth = depth
+        self._stream = stream
+        self._inverse = inverse
+        self._mode = mode
+
+    def prepare(self, in_channels: int) -> None:
+        """Builds the Lyndon tables for paths with ``in_channels`` channels ahead of the first call."""
+        for length in range(1, self._depth + 1):
+            utility._lyndon_flat_indices_cpu(in_channels, length)
+            if self._mode == "brackets":
+                utility._words_to_brackets_cpu(in_channels, length)
+
+    def forward(self, path: torch.Tensor, basepoint: Union[bool, torch.Tensor] = False) -> torch.Tensor:
+        signature = smodule.signature(path, self._depth, stream=self._stream, basepoint=basepoint,
+                                      inverse=self._inverse)
+        return signature_to_logsignature(signature, path.size(-1), self._depth, self._stream, self._mode)
+
+    def extra_repr(self) -> str:
+        return (f'depth={self._depth}, stream={self._stream}, inverse={self._inverse}, mode={self._mode!r}')
+
+
+Logsignature = LogSignature
+
+
+def logsignature(path: torch.Tensor, depth: int, stream: bool = False, basepoint: Union[bool, torch.Tensor] = False,
+                 inverse: bool = False, mode: str = "words") -> torch.Tensor:
+    """Logsignature transform of a batch of paths, as ``signatory.logsignature``.
+
+    Arguments ``path``, ``depth``, ``stream``, ``basepoint`` and ``inverse`` are as in :func:`sigtorch.signature`.
+    ``mode`` selects the representation:
+
+    - ``"words"`` (default): coefficients of the Lyndon words in the tensor-algebra logarithm, one channel per Lyndon
+      word (see :func:`sigtorch.lyndon_words`).
+    - ``"brackets"``: coefficients in the Lyndon basis of the free Lie algebra (see :func:`sigtorch.lyndon_brackets`).
+    - ``"expand"``: every coefficient of the tensor-algebra logarithm; same shape as the signature.
+
+    Returns:
+        Tensor of shape (batch, logsignature_channels(channels, depth)) for "words"/"brackets", with a stream
+        dimension inserted as for :func:`sigtorch.signature` if ``stream``.
+    """
+    return LogSignature(depth, stream=stream, inverse=inverse, mode=mode)(path, basepoint=basepoint)