event2vector 0.3.1__tar.gz → 0.4.0__tar.gz

{event2vector-0.3.1 → event2vector-0.4.0}/CHANGELOG.md

@@ -4,6 +4,30 @@ All notable changes to `event2vector` are documented here. The format follows
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the project uses
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.0] - 2026-10-15
+
+### Added
+
+- `geometry="affine"` (`AffineModel`): each event is an invertible affine map
+  `x -> exp(M_e) x + v_e` applied on the left of an affine state, so an event can
+  act on the history (shrink, stretch or rotate it) instead of only adding to it.
+  With `M_e = 0`, the initialisation, the point is exactly the Euclidean sum. The
+  step is recovered exactly by the reconstruction loss, and subsequences compose
+  into one affine map:
+  - `AffineModel.point(state)` returns the plotted point `b` from a state;
+  - `AffineModel.compose(event_ids)` returns one `(A, v)` for a subsequence;
+  - `AffineModel.override_strength(event_ids)` returns the singular values of its
+    linear part.
+- Estimator options `affine_readout` (`"point"` or `"full"`) and
+  `affine_max_log_scale` (bounds how much one event can stretch or shrink the state).
+- Tests for the affine geometry (`tests/test_affine.py`).
+
+### Changed
+
+- The MPS-safe `matrix_exp` is now a shared helper used by both the `rotation` and
+  `affine` geometries, and automatic device selection falls back to the CPU for
+  `affine` as it already did for `rotation`.
+
 ## [0.3.1] - 2026-10-11
 
 ### Fixed
@@ -35,5 +59,6 @@ All notable changes to `event2vector` are documented here. The format follows
 
 Earlier releases are listed on [PyPI](https://pypi.org/project/event2vector/#history).
 
+[0.4.0]: https://pypi.org/project/event2vector/0.4.0/
 [0.3.1]: https://pypi.org/project/event2vector/0.3.1/
 [0.3.0]: https://pypi.org/project/event2vector/0.3.0/

event2vector-0.4.0/CITATION.cff

@@ -0,0 +1,26 @@
+cff-version: 1.2.0
+message: "If you use event2vector, please cite the paper below."
+title: "event2vector: geometric embeddings for event sequences"
+type: software
+authors:
+  - family-names: Sulc
+    given-names: Antonin
+version: 0.4.0
+date-released: 2026-10-15
+license: MIT
+repository-code: "https://github.com/sulcantonin/event2vec_public"
+url: "https://event2vec.com"
+keywords:
+  - sequence-embedding
+  - event-sequences
+  - representation-learning
+  - geometric deep learning
+preferred-citation:
+  type: article
+  title: "Event2Vec: A Geometric Approach to Learning Composable Representations of Event Sequences"
+  authors:
+    - family-names: Sulc
+      given-names: Antonin
+  journal: "arXiv preprint arXiv:2509.12188"
+  year: 2025
+  url: "https://arxiv.org/abs/2509.12188"

{event2vector-0.3.1 → event2vector-0.4.0}/MANIFEST.in

@@ -3,3 +3,4 @@ include CHANGELOG.md
 exclude *.pkl
 exclude *.model
 exclude data/*
+include CITATION.cff

{event2vector-0.3.1/event2vector.egg-info → event2vector-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: event2vector
-Version: 0.3.1
+Version: 0.4.0
 Summary: Scikit-learn-style geometric embeddings for event sequences.
 Author: Antonin Sulc
 License-Expression: MIT
@@ -133,7 +133,7 @@ Key methods:
 - `most_similar`: gensim-style nearest-neighbor lookup over learned event embeddings using tokens or full sequences as queries.
 - `pad_sequences=True`: enables fully vectorized batches with masking for substantial throughput gains on large corpora.
 
-Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` geometry defaults to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for it; pass `device="mps"` to override.
+Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` and `affine` geometries default to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for them; pass `device="mps"` to override.
 
 ## Geometries (v0.2)
 
@@ -147,6 +147,7 @@ Beyond `euclidean` and `hyperbolic`, the estimator exposes the paper's pilot geo
 | `spherical` | κ-gyroaddition (`kappa>0`) | cyclic / modular (periodic) structure |
 | `rotation` | left `SO(n)` product | non-commutative orientation (order matters) |
 | `heisenberg` | Heisenberg group product | order-dependent signed (Lévy) area |
+| `affine` | `x -> exp(M_e) x + v_e` on the left | events that act on the history (shrink, stretch, rotate it), not only add to it |
 | `product` | per-factor update | two structure types at once (e.g. hyperbolic × spherical) |
 
 ```python
@@ -162,6 +163,39 @@ model = Event2Vec(num_event_types=len(vocab), geometry="rotation", rotation_n=3)
 model = Event2Vec(num_event_types=len(vocab), geometry="product", product_kappas=(-1.0, 1.0))
 ```
 
+### Affine geometry (v0.4)
+
+In the additive geometries an event always pushes the same way. With
+`geometry="affine"` each event is an invertible affine map
+`x -> exp(M_e) x + v_e`, so the same event contributes differently depending on
+what came before, while the model stays a group:
+
+- with every `M_e = 0` (the initialisation) the point is exactly the Euclidean sum,
+  so training starts from the additive model;
+- every step is invertible and the reconstruction loss recovers it exactly;
+- any subsequence composes into a single affine map, so its effect on *any*
+  history is one `A x + v`.
+
+The state is the whole map `[vec(A), b]`; the point to plot is `b`, the image of
+the origin (`model.model.point(...)`). The singular values of a subsequence's `A`
+say how much it overrides what came before: `1` passes the history through,
+`< 1` erases it along that direction.
+
+```python
+model = Event2Vec(num_event_types=len(vocab), geometry="affine", embedding_dim=2)
+model.fit(sequences)
+
+states = model.transform(sequences, as_numpy=False)  # [vec(A), b] per sequence
+points = model.model.point(states)                   # (n, 2) points to plot
+
+A, v = model.model.compose([3, 7, 1])                # one map for a subsequence
+print(model.model.override_strength([3, 7, 1]))      # its singular values
+```
+
+Options: `affine_readout="point"` (default; the decoder reads `b`) or `"full"`
+(it also reads `A`), and `affine_max_log_scale` to bound how strongly a single
+event can stretch or shrink the state.
+
 Reproduce the curvature sweep on a real knowledge graph (auto-downloads WN18RR):
 
 ```bash

{event2vector-0.3.1 → event2vector-0.4.0}/README.md

@@ -98,7 +98,7 @@ Key methods:
 - `most_similar`: gensim-style nearest-neighbor lookup over learned event embeddings using tokens or full sequences as queries.
 - `pad_sequences=True`: enables fully vectorized batches with masking for substantial throughput gains on large corpora.
 
-Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` geometry defaults to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for it; pass `device="mps"` to override.
+Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` and `affine` geometries default to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for them; pass `device="mps"` to override.
 
 ## Geometries (v0.2)
 
@@ -112,6 +112,7 @@ Beyond `euclidean` and `hyperbolic`, the estimator exposes the paper's pilot geo
 | `spherical` | κ-gyroaddition (`kappa>0`) | cyclic / modular (periodic) structure |
 | `rotation` | left `SO(n)` product | non-commutative orientation (order matters) |
 | `heisenberg` | Heisenberg group product | order-dependent signed (Lévy) area |
+| `affine` | `x -> exp(M_e) x + v_e` on the left | events that act on the history (shrink, stretch, rotate it), not only add to it |
 | `product` | per-factor update | two structure types at once (e.g. hyperbolic × spherical) |
 
 ```python
@@ -127,6 +128,39 @@ model = Event2Vec(num_event_types=len(vocab), geometry="rotation", rotation_n=3)
 model = Event2Vec(num_event_types=len(vocab), geometry="product", product_kappas=(-1.0, 1.0))
 ```
 
+### Affine geometry (v0.4)
+
+In the additive geometries an event always pushes the same way. With
+`geometry="affine"` each event is an invertible affine map
+`x -> exp(M_e) x + v_e`, so the same event contributes differently depending on
+what came before, while the model stays a group:
+
+- with every `M_e = 0` (the initialisation) the point is exactly the Euclidean sum,
+  so training starts from the additive model;
+- every step is invertible and the reconstruction loss recovers it exactly;
+- any subsequence composes into a single affine map, so its effect on *any*
+  history is one `A x + v`.
+
+The state is the whole map `[vec(A), b]`; the point to plot is `b`, the image of
+the origin (`model.model.point(...)`). The singular values of a subsequence's `A`
+say how much it overrides what came before: `1` passes the history through,
+`< 1` erases it along that direction.
+
+```python
+model = Event2Vec(num_event_types=len(vocab), geometry="affine", embedding_dim=2)
+model.fit(sequences)
+
+states = model.transform(sequences, as_numpy=False)  # [vec(A), b] per sequence
+points = model.model.point(states)                   # (n, 2) points to plot
+
+A, v = model.model.compose([3, 7, 1])                # one map for a subsequence
+print(model.model.override_strength([3, 7, 1]))      # its singular values
+```
+
+Options: `affine_readout="point"` (default; the decoder reads `b`) or `"full"`
+(it also reads `A`), and `affine_max_log_scale` to bound how strongly a single
+event can stretch or shrink the state.
+
 Reproduce the curvature sweep on a real knowledge graph (auto-downloads WN18RR):
 
 ```bash

{event2vector-0.3.1 → event2vector-0.4.0}/event2vector/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.1"
+__version__ = "0.4.0"
 
 from .models import EuclideanModel, HyperbolicModel, HyperbolicUtils, Event2Vec
 from .geometries import (
@@ -6,6 +6,7 @@ from .geometries import (
     StereographicModel,
     RotationModel,
     HeisenbergModel,
+    AffineModel,
     ProductModel,
 )
 from .data import generate_sequences, get_sequences
@@ -20,6 +21,7 @@ __all__ = [
     "StereographicModel",
     "RotationModel",
     "HeisenbergModel",
+    "AffineModel",
     "ProductModel",
     "generate_sequences",
     "get_sequences",

{event2vector-0.3.1 → event2vector-0.4.0}/event2vector/geometries.py

@@ -17,6 +17,10 @@ rotation (SO(n))   left matrix product          non-commutative orientation;
                                                 abelianization
 heisenberg         nilpotent group product      the signed (Levy) area / commutator,
                                                 a purely order-dependent quantity
+affine (Aff+(d))   x -> exp(M_e) x + v_e        events that act on the history
+                   on the left                  (shrink / stretch / rotate it) as
+                                                well as add to it; M_e = 0 is
+                                                exactly the Euclidean sum
 product            per-factor update            two structure types at once
                                                 (e.g. hierarchy H x cycle S)
 ================  ===========================  ===================================
@@ -35,6 +39,22 @@ import torch.nn as nn
 EPS = 1e-5
 
 
+def _matrix_exp(a):
+    """torch.matrix_exp that also works for MPS tensors: MPS has no matrix_exp
+    kernel (pytorch/pytorch#141287), so it is computed on the CPU and moved back.
+    Both transfers are differentiable."""
+    if a.device.type == "mps":
+        return torch.matrix_exp(a.cpu()).to(a.device)
+    return torch.matrix_exp(a)
+
+
+def _solve(a, b):
+    """torch.linalg.solve, computed on the CPU for MPS tensors (as above)."""
+    if a.device.type == "mps":
+        return torch.linalg.solve(a.cpu(), b.cpu()).to(a.device)
+    return torch.linalg.solve(a, b)
+
+
 # --------------------------------------------------------------------------- #
 # kappa-stereographic model: one signed curvature for E / S / H
 # --------------------------------------------------------------------------- #
@@ -216,11 +236,7 @@ class RotationModel(nn.Module):
         a = torch.zeros(b, self.n, self.n, device=gen.device, dtype=gen.dtype)
         a[:, self.triu_i, self.triu_j] = gen
         a = a - a.transpose(1, 2)
-        if a.device.type == "mps":
-            # MPS has no matrix_exp kernel (pytorch/pytorch#141287); compute it on
-            # the CPU and move it back. Both transfers are differentiable.
-            return torch.matrix_exp(a.cpu()).to(a.device)
-        return torch.matrix_exp(a)
+        return _matrix_exp(a)
 
     def zero_state(self, batch, device):
         return torch.eye(self.n, device=device).reshape(1, -1).repeat(batch, 1)
@@ -333,6 +349,152 @@ class HeisenbergModel(nn.Module):
         return ((a - b) ** 2).sum(-1, keepdim=True)
 
 
+# --------------------------------------------------------------------------- #
+# Affine group model: events that act on the history, not only add to it
+# --------------------------------------------------------------------------- #
+class AffineModel(nn.Module):
+    """State is an orientation-preserving affine map ``x -> A x + b`` of R^d,
+    stored flat as ``[vec(A), b]`` and starting at the identity. Each event is the
+    affine map ``x -> exp(M_e) x + v_e`` applied on the left::
+
+        A_t = exp(M_e) A_{t-1}            b_t = exp(M_e) b_{t-1} + v_e
+
+    ``b`` is the sequence's *point* (the image of the origin). With every
+    ``M_e = 0`` it is exactly the Euclidean sum of the ``v_e``, so this geometry
+    contains event2vec's additive model and starts training from it.
+
+    The linear part lets an event act on what came before -- shrink, stretch or
+    rotate it -- so the same event contributes differently after different
+    histories, while the model stays a group: every step is invertible (``exp``
+    of any matrix is), the reconstruction loss recovers it exactly, and any
+    subsequence composes into one affine map (:meth:`compose`). Its singular
+    values (:meth:`override_strength`) say how strongly that subsequence
+    overrides the history: values below 1 erase the history along a direction.
+
+    Args:
+        embedding_dim: dimension ``d`` of the point ``b`` (2 for a planar map).
+            Each event has ``d + d*d`` parameters (``v_e`` and ``M_e``).
+        readout: ``"point"`` decodes the next event from ``b`` only (the plotted
+            point); ``"full"`` also reads ``A``.
+        max_log_scale: if set, ``M_e`` entries are squashed to
+            ``(-max_log_scale, max_log_scale)``, bounding how much one event can
+            stretch or shrink the state (each step's condition number is at most
+            ``exp(2 * d * max_log_scale)``). ``None`` leaves ``M_e`` free.
+    """
+
+    geometry = "affine"
+
+    def __init__(self, num_event_types, embedding_dim=2, dropout_p=0.0,
+                 readout="point", max_log_scale=None):
+        super().__init__()
+        if readout not in ("point", "full"):
+            raise ValueError("readout must be 'point' or 'full'.")
+        d = embedding_dim
+        self.d = d
+        self.point_dim = d
+        self.embedding_dim = d + d * d        # generator: [v_e, vec(M_e)]
+        self.state_dim = d * d + d            # state: [vec(A), b]
+        self.readout = readout
+        self.max_log_scale = max_log_scale
+        self.embedding = nn.Embedding(num_event_types, self.embedding_dim)
+        self.dropout = nn.Dropout(p=dropout_p)
+        self.linear_dec = nn.Linear(d if readout == "point" else self.state_dim,
+                                    num_event_types)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.linear_dec is not None:  # None when used as a ProductModel factor
+            nn.init.xavier_uniform_(self.linear_dec.weight)
+            if self.linear_dec.bias is not None:
+                nn.init.zeros_(self.linear_dec.bias)
+        d = self.d
+        with torch.no_grad():
+            nn.init.normal_(self.embedding.weight[:, :d], 0.0, 0.02)  # v_e as Euclidean
+            self.embedding.weight[:, d:].zero_()                      # M_e = 0: A_e = I
+
+    # -- group elements ----------------------------------------------------- #
+    def _maps(self, gen):
+        """(B, d + d*d) generators -> (A_e (B, d, d), v_e (B, d))."""
+        d = self.d
+        v, m = gen[:, :d], gen[:, d:].reshape(-1, d, d)
+        if self.max_log_scale is not None:
+            c = self.max_log_scale
+            m = c * torch.tanh(m / c)
+        return _matrix_exp(m), v
+
+    def _split(self, h):
+        d = self.d
+        return h[:, : d * d].reshape(-1, d, d), h[:, d * d:]
+
+    @staticmethod
+    def _join(a, b):
+        return torch.cat([a.reshape(a.shape[0], -1), b], dim=-1)
+
+    def zero_state(self, batch, device):
+        a = torch.eye(self.d, device=device).expand(batch, -1, -1)
+        return self._join(a, torch.zeros(batch, self.d, device=device))
+
+    def point(self, h):
+        """The plotted point ``b`` of a state (or of ``transform``'s output)."""
+        return h[..., self.d * self.d:]
+
+    def to_tangent(self, h):
+        return self.point(h) if self.readout == "point" else h
+
+    # -- recurrence ---------------------------------------------------------- #
+    def encoder(self, x, h):
+        a_e, v_e = self._maps(self.dropout(self.embedding(x)))
+        a_old, b_old = self._split(h)
+        a_new = torch.bmm(a_e, a_old)
+        b_new = torch.bmm(a_e, b_old.unsqueeze(-1)).squeeze(-1) + v_e
+        return self._join(a_new, b_new), self._join(a_e, v_e)
+
+    def decoder(self, h):
+        if self.linear_dec is None:
+            raise RuntimeError(
+                "this model is a ProductModel factor; decode through the "
+                "product's shared decoder, not the factor's")
+        return self.linear_dec(self.to_tangent(h))
+
+    def forward(self, x, h=None):
+        if h is None:
+            h = self.zero_state(x.shape[0], x.device)
+        h_next, e = self.encoder(x, h)
+        return self.decoder(h_next), h_next, e
+
+    def recover_increment(self, h_old, h_new):
+        """g_e = g_new . g_old^-1:  A_e = A_new A_old^-1,  v_e = b_new - A_e b_old."""
+        a_old, b_old = self._split(h_old)
+        a_new, b_new = self._split(h_new)
+        # A_e A_old = A_new  <=>  A_old^T A_e^T = A_new^T (solve, don't invert)
+        a_e = _solve(a_old.transpose(1, 2), a_new.transpose(1, 2)).transpose(1, 2)
+        v_e = b_new - torch.bmm(a_e, b_old.unsqueeze(-1)).squeeze(-1)
+        return self._join(a_e, v_e)
+
+    def dist2(self, a, b):
+        return ((a - b) ** 2).sum(-1, keepdim=True)            # Frobenius^2 on [A, b]
+
+    # -- interpretation -------------------------------------------------------- #
+    @torch.no_grad()
+    def compose(self, event_ids):
+        """One affine map (A, v) for a whole subsequence: applying it to any state
+        equals feeding the events one by one (exact; it is a group product)."""
+        dev = self.embedding.weight.device
+        a, v = torch.eye(self.d, device=dev), torch.zeros(self.d, device=dev)
+        for i in event_ids:
+            a_e, v_e = self._maps(self.embedding.weight[int(i)].unsqueeze(0))
+            a, v = a_e[0] @ a, a_e[0] @ v + v_e[0]
+        return a, v
+
+    @torch.no_grad()
+    def override_strength(self, event_ids):
+        """Singular values of the subsequence's linear part, largest first:
+        1 = the history passes through unchanged, < 1 = erased along that
+        direction, > 1 = amplified."""
+        a, _ = self.compose(event_ids)
+        return torch.linalg.svdvals(a)
+
+
 # --------------------------------------------------------------------------- #
 # Product manifold: two structure types at once (e.g. hierarchy H x cycle S)
 # --------------------------------------------------------------------------- #

{event2vector-0.3.1 → event2vector-0.4.0}/event2vector/models.py

@@ -8,6 +8,7 @@ import torch.nn as nn
 from torch.nn.utils.rnn import pad_sequence
 
 from .geometries import (
+    AffineModel,
     HeisenbergModel,
     ProductModel,
     RotationModel,
@@ -196,6 +197,8 @@ class Event2Vec:
         kappa: Optional[float] = None,
         rotation_n: int = 3,
         heisenberg_m: int = 1,
+        affine_readout: str = "point",
+        affine_max_log_scale: Optional[float] = None,
         product_kappas: Sequence[float] = (-1.0, 1.0),
         lambda_reconstruction: float = 0.2,
         lambda_consistency: float = 0.2,
@@ -232,6 +235,8 @@ class Event2Vec:
         self.kappa = kappa
         self.rotation_n = rotation_n
         self.heisenberg_m = heisenberg_m
+        self.affine_readout = affine_readout
+        self.affine_max_log_scale = affine_max_log_scale
         self.product_kappas = tuple(product_kappas)
         self.lambda_reconstruction = lambda_reconstruction
         self.lambda_consistency = lambda_consistency
@@ -249,11 +254,11 @@ class Event2Vec:
             if use_gpu and torch.cuda.is_available():
                 device = torch.device("cuda")
             elif use_gpu and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-                if self.geometry == "rotation":
-                    # MPS has no matrix_exp, so every rotation step would round-trip
-                    # through the CPU; for a model this small the CPU is much faster.
+                if self.geometry in ("rotation", "affine"):
+                    # MPS has no matrix_exp, so every step would round-trip through
+                    # the CPU; for a model this small the CPU is much faster.
                     warnings.warn(
-                        "geometry='rotation' needs torch.matrix_exp, which MPS does not "
+                        f"geometry='{self.geometry}' needs torch.matrix_exp, which MPS does not "
                         "implement; using device='cpu'. Pass device='mps' to run on MPS "
                         "anyway (matrix_exp is then computed on the CPU).",
                         RuntimeWarning,
@@ -318,6 +323,14 @@ class Event2Vec:
                 m=self.heisenberg_m,
                 dropout_p=self.dropout_p,
             )
+        elif self.geometry == "affine":
+            model = AffineModel(
+                num_event_types=self.num_event_types,
+                embedding_dim=self.embedding_dim,
+                dropout_p=self.dropout_p,
+                readout=self.affine_readout,
+                max_log_scale=self.affine_max_log_scale,
+            )
         elif self.geometry == "product":
             factors = [
                 StereographicModel(
@@ -332,7 +345,7 @@ class Event2Vec:
         else:
             raise ValueError(
                 "geometry must be one of: 'euclidean', 'hyperbolic', 'spherical', "
-                "'stereographic', 'rotation', 'heisenberg', 'product'."
+                "'stereographic', 'rotation', 'heisenberg', 'affine', 'product'."
             )
         return model.to(self.device)
 

{event2vector-0.3.1 → event2vector-0.4.0/event2vector.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: event2vector
-Version: 0.3.1
+Version: 0.4.0
 Summary: Scikit-learn-style geometric embeddings for event sequences.
 Author: Antonin Sulc
 License-Expression: MIT
@@ -133,7 +133,7 @@ Key methods:
 - `most_similar`: gensim-style nearest-neighbor lookup over learned event embeddings using tokens or full sequences as queries.
 - `pad_sequences=True`: enables fully vectorized batches with masking for substantial throughput gains on large corpora.
 
-Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` geometry defaults to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for it; pass `device="mps"` to override.
+Device control: set `use_gpu=False` to force CPU even if CUDA/MPS is present, or pass an explicit `device` (e.g., `"cuda:0"` or `"cpu"`). On Apple Silicon the `rotation` and `affine` geometries default to the CPU, because MPS has no `matrix_exp` kernel and the CPU is faster for them; pass `device="mps"` to override.
 
 ## Geometries (v0.2)
 
@@ -147,6 +147,7 @@ Beyond `euclidean` and `hyperbolic`, the estimator exposes the paper's pilot geo
 | `spherical` | κ-gyroaddition (`kappa>0`) | cyclic / modular (periodic) structure |
 | `rotation` | left `SO(n)` product | non-commutative orientation (order matters) |
 | `heisenberg` | Heisenberg group product | order-dependent signed (Lévy) area |
+| `affine` | `x -> exp(M_e) x + v_e` on the left | events that act on the history (shrink, stretch, rotate it), not only add to it |
 | `product` | per-factor update | two structure types at once (e.g. hyperbolic × spherical) |
 
 ```python
@@ -162,6 +163,39 @@ model = Event2Vec(num_event_types=len(vocab), geometry="rotation", rotation_n=3)
 model = Event2Vec(num_event_types=len(vocab), geometry="product", product_kappas=(-1.0, 1.0))
 ```
 
+### Affine geometry (v0.4)
+
+In the additive geometries an event always pushes the same way. With
+`geometry="affine"` each event is an invertible affine map
+`x -> exp(M_e) x + v_e`, so the same event contributes differently depending on
+what came before, while the model stays a group:
+
+- with every `M_e = 0` (the initialisation) the point is exactly the Euclidean sum,
+  so training starts from the additive model;
+- every step is invertible and the reconstruction loss recovers it exactly;
+- any subsequence composes into a single affine map, so its effect on *any*
+  history is one `A x + v`.
+
+The state is the whole map `[vec(A), b]`; the point to plot is `b`, the image of
+the origin (`model.model.point(...)`). The singular values of a subsequence's `A`
+say how much it overrides what came before: `1` passes the history through,
+`< 1` erases it along that direction.
+
+```python
+model = Event2Vec(num_event_types=len(vocab), geometry="affine", embedding_dim=2)
+model.fit(sequences)
+
+states = model.transform(sequences, as_numpy=False)  # [vec(A), b] per sequence
+points = model.model.point(states)                   # (n, 2) points to plot
+
+A, v = model.model.compose([3, 7, 1])                # one map for a subsequence
+print(model.model.override_strength([3, 7, 1]))      # its singular values
+```
+
+Options: `affine_readout="point"` (default; the decoder reads `b`) or `"full"`
+(it also reads `A`), and `affine_max_log_scale` to bound how strongly a single
+event can stretch or shrink the state.
+
 Reproduce the curvature sweep on a real knowledge graph (auto-downloads WN18RR):
 
 ```bash

{event2vector-0.3.1 → event2vector-0.4.0}/event2vector.egg-info/SOURCES.txt

@@ -1,4 +1,5 @@
 CHANGELOG.md
+CITATION.cff
 MANIFEST.in
 README.md
 pyproject.toml
@@ -23,4 +24,5 @@ event2vector/datasets/realdata.py
 event2vector/eval/__init__.py
 event2vector/eval/figures.py
 event2vector/eval/harness.py
+tests/test_affine.py
 tests/test_device.py

{event2vector-0.3.1 → event2vector-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "event2vector"
-version = "0.3.1"
+version = "0.4.0"
 description = "Scikit-learn-style geometric embeddings for event sequences."
 readme = "README.md"
 requires-python = ">=3.9"

event2vector-0.4.0/tests/test_affine.py

@@ -0,0 +1,97 @@
+"""Tests for the affine geometry (``geometry="affine"``, :class:`AffineModel`).
+
+The affine state is a map ``x -> A x + b``; each event applies
+``x -> exp(M_e) x + v_e`` on the left. Run with `pytest tests/ -q`.
+"""
+import math
+
+import numpy as np
+import pytest
+import torch
+
+from event2vector import AffineModel, Event2Vec
+
+SEQS = [[0, 1, 2, 3], [0, 1, 2, 4], [0, 2, 3, 4], [0, 1, 3, 4], [0, 2, 1, 4]] * 8
+
+
+def _affine(dim=2, **kw):
+    am = AffineModel(num_event_types=5, embedding_dim=dim, **kw)
+    am.eval()
+    with torch.no_grad():  # a non-trivial linear part for every event
+        am.embedding.weight.normal_(0.0, 0.3)
+    return am
+
+
+def _feed(model, ids):
+    h = model.zero_state(1, torch.device("cpu"))
+    for i in ids:
+        h = model.forward(torch.tensor([i]), h)[1]
+    return h
+
+
+@pytest.mark.parametrize("kw", [dict(), dict(affine_readout="full",
+                                             affine_max_log_scale=1.0)])
+def test_affine_fits_transforms_and_learns(kw):
+    est = Event2Vec(num_event_types=5, geometry="affine", embedding_dim=2,
+                    dropout_p=0.0, num_epochs=40, learning_rate=5e-3, batch_size=8,
+                    pad_sequences=True, random_state=0, use_gpu=False, **kw)
+    est.fit(SEQS)
+    emb = est.transform(SEQS)
+    assert emb.shape == (len(SEQS), est.model.state_dim)  # [vec(A), b] = 4 + 2
+    assert np.isfinite(emb).all()
+    assert min(est.training_loss_) < est.training_loss_[0] * 0.98
+    points = est.model.point(torch.as_tensor(emb))
+    assert points.shape == (len(SEQS), 2)
+
+
+def test_affine_with_identity_maps_is_the_euclidean_sum():
+    """M_e = 0 (the initialisation) makes the point b exactly the sum of v_e."""
+    am = AffineModel(num_event_types=5, embedding_dim=3)
+    am.eval()
+    ids = [0, 3, 1, 3]
+    v = am.embedding.weight[:, :3].detach()
+    assert torch.allclose(am.point(_feed(am, ids))[0], v[ids].sum(0), atol=1e-6)
+
+
+def test_affine_compose_equals_stepwise_and_is_order_sensitive():
+    am = _affine()
+    ids = [1, 2, 3, 4]
+    a, v = am.compose(ids)
+    assert torch.allclose(am.point(_feed(am, ids))[0], v, atol=1e-5)
+    # applied to any state, the composed map equals feeding the events one by one
+    _, b_stem = am._split(_feed(am, [0, 2]))
+    after = _feed(am, [0, 2] + ids)
+    assert torch.allclose(am.point(after)[0], a @ b_stem[0] + v, atol=1e-5)
+    assert not torch.allclose(am.compose([1, 2])[1], am.compose([2, 1])[1], atol=1e-4)
+
+
+def test_affine_recover_increment_is_exact():
+    """The reconstruction loss's left-cancellation recovers each step exactly."""
+    am = _affine()
+    h0 = _feed(am, [0, 1, 2])
+    h1, step = am.encoder(torch.tensor([3]), h0)
+    assert torch.allclose(am.recover_increment(h0, h1), step, atol=1e-5)
+    assert am.dist2(am.recover_increment(h0, h1), step).item() < 1e-8
+
+
+def test_affine_max_log_scale_bounds_each_step():
+    am = _affine(max_log_scale=0.5)
+    with torch.no_grad():
+        am.embedding.weight.mul_(100.0)
+    s = am.override_strength([2])
+    bound = math.exp(2 * 2 * 0.5)
+    assert s.max().item() <= bound + 1e-4 and s.min().item() >= 1 / bound - 1e-4
+
+
+def test_affine_rejects_unknown_readout():
+    with pytest.raises(ValueError):
+        AffineModel(num_event_types=5, readout="nope")
+
+
+def test_auto_device_skips_mps_for_affine(monkeypatch):
+    """Like rotation, affine needs matrix_exp: auto-selection prefers the CPU."""
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+    with pytest.warns(RuntimeWarning, match="geometry='affine' needs torch.matrix_exp"):
+        est = Event2Vec(num_event_types=3, geometry="affine")
+    assert est.device.type == "cpu"