Trajectory Metrics: Aggregation Trajectory Metric

This example demonstrates trajectory-level metrics using AggregationTrajectoryMetric, which computes distances between trajectories that contain multiple sequence types (e.g., states, events, or intervals).

Setup

import polars as pl
import matplotlib.pyplot as plt

from tanat import build_states, build_events, build_trajectories
from tanat.dataset import simulate_trajectories, simulate_static
from tanat.metric.entity import HammingEntityMetric
from tanat.metric.sequence import (
    EditSequenceMetric,
    LCSSequenceMetric,
)
from tanat.metric import AggregationTrajectoryMetric
from tanat.metric.static import StaticMetric

Generate synthetic trajectory data

N_TRAJ = 100
SEED = 42

raw = simulate_trajectories(
    sequences={
        "states": {
            "type": "state",
            "n_ids": N_TRAJ,
            "seq_length_range": (3, 8),
            "features": ["score", "status"],
        },
        "events": {
            "type": "event",
            "n_ids": N_TRAJ,
            "seq_length_range": (2, 6),
            "features": ["score", "status"],
        },
    },
    shared_ids=True,
    seed=SEED,
)

# Build pools for each sequence type
states_pool = build_states(
    temporal_data=raw["states"],
    id_column="id",
    start_column="start",
    end_column="end",
)
events_pool = build_events(
    temporal_data=raw["events"],
    id_column="id",
    time_column="time",
)

# Build trajectory pool
traj_pool = build_trajectories(pools={"states": states_pool, "events": events_pool})

┌─ State SequenceStore
│
│ Step 1/4: Sorting & preparing data
│
│ Step 2/4: Building sequence index
│
│ Step 3/4: Writing entity & time index features
│
│ Step 4/4: Computing & writing metadata
│
└─ Done (100 sequences · 531 entities · 0.01s)
┌─ Event SequenceStore
│
│ Step 1/4: Sorting & preparing data
│
│ Step 2/4: Building sequence index
│
│ Step 3/4: Writing entity & time index features
│
│ Step 4/4: Computing & writing metadata
│
└─ Done (100 sequences · 383 entities · 0.00s)
┌─ TrajectoryStore
│
│ Step 1/2: Linking pools: states, events
│
│ Step 2/2: Building trajectory index & metadata
│
└─ Done (100 trajectories · 2 pool(s) · 0.00s)

# Cast features to categorical
for sp in traj_pool.sequence_pools.values():
    sp.cast_features({"status": pl.Categorical})

print(traj_pool)

┌────────────────────────────────────────────────┐
│             TrajectoryPool Summary             │
└────────────────────────────────────────────────┘

Overview
─────────────────────────
  Trajectories       100
  Store              /home/runner/.tanat/_quick_trajectory_9f00a010
  id_column          id

Time Index
─────────────────────────
  Type               Datetime(time_unit='us', time_zone=None) [2000-01-04 17:05:08.109495 → 2025-01-01 00:00:00]
  t0                 position=0, anchor=start

Sequences (2)
─────────────────────────
  • states              StateSequencePool(n=100, entity_features=2, static_features=0, store='/home/runner/.tanat/_quick_state_d2bae59e')
  • events              EventSequencePool(n=100, entity_features=2, static_features=0, store='/home/runner/.tanat/_quick_event_eaab17c0')

Define trajectory metric

hamming = HammingEntityMetric(entity_feature="status")

# Use different metrics per alias (sequence type)
agg = AggregationTrajectoryMetric(
    sequence_metrics={
        "events": LCSSequenceMetric(entity_metric=hamming, mode="normalized"),
        "states": EditSequenceMetric(entity_metric=hamming, normalize=True),
    },
    agg_fun="mean",
)
print(agg)

AggregationTrajectoryMetric(settings=AggregationSettings(sequence_metrics={'events': LCSSequenceMetric(settings=LCSSettings(entity_metric=HammingEntityMetric(settings=HammingSettings(entity_feature='status', cost=None, mismatch_cost=1.0)), equality_threshold=0.0, mode='normalized')), 'states': EditSequenceMetric(settings=EditSettings(entity_metric=HammingEntityMetric(settings=HammingSettings(entity_feature='status', cost=None, mismatch_cost=1.0)), indel_cost=1.0, normalize=True))}, agg_fun='mean', weights=None))

Compute distance between a single pair

traj_ids = traj_pool.unique_ids
traj_a = traj_pool[traj_ids[0]]
traj_b = traj_pool[traj_ids[1]]

dist = agg(traj_a, traj_b)
print(f"\nDistance between {traj_ids[0]} and {traj_ids[1]}: {dist:.4f}")

Distance between 1 and 2: 0.8750

Compute full pairwise distance matrix

matrix = agg.compute_matrix(traj_pool)
print(f"\nDistance matrix shape: {matrix.shape}")
print(f"Mean distance: {matrix.to_numpy()[matrix.to_numpy() > 0].mean():.4f}")

┌─ AggregationTrajectoryMetric
│
│   ┌─ EditSequenceMetric
│   │

│   │ Chunks:   0%|          | 0/1 [00:00<?, ?it/s]
│   │ Chunks: 100%|██████████| 1/1 [00:00<00:00, 1062.93it/s]
│   │
│   └─ Done (100 sequences · 0.00s)
│
│   ┌─ LCSSequenceMetric
│   │

│   │ Chunks:   0%|          | 0/1 [00:00<?, ?it/s]
│   │ Chunks: 100%|██████████| 1/1 [00:00<00:00, 1246.82it/s]
│   │
│   └─ Done (100 sequences · 0.00s)
│
└─ Done (100 trajectories · 0.02s)

Distance matrix shape: (100, 100)
Mean distance: 0.6860

Visualize trajectory distances

fig, ax = plt.subplots(figsize=(8, 6))
arr = matrix.to_numpy()
im = ax.imshow(arr, cmap="viridis", aspect="auto")
ax.set_title(
    "Trajectory distances\n(Edit distance for states, LCS for events)",
    fontsize=12,
    fontweight="bold",
)
ax.set_xlabel("Trajectory index")
ax.set_ylabel("Trajectory index")
cbar = plt.colorbar(im, ax=ax)
cbar.set_label("Distance")
plt.tight_layout()
plt.show()

Use static metric in AggregationMetric

A StaticMetric can also be used in an AggregationMetric. This trajectory metric defines a distance as an aggregate of sequence metrics. In addition, it is possible to define a static metric as an additional metric to aggregate (with it own weight).

# first add simulated static data to the trajectory pool
static_data = simulate_static(n_ids=N_TRAJ, features=["age", "group"])
traj_pool.add_static_features(static_data, id_column="id")


# define a static metric based on `age` feature
def age_cmp(s1, s2):
    """Comparison of age static features"""
    return abs(s1["age"] - s2["age"])


custom_static_metric = StaticMetric(cmp_fnct=age_cmp)

# The trajectory metric is defined by the sequence metric required
# for the `events` type, and we add a `static_metric` with a very small weight
# (the comparison of ages leads to values between 0 and 120).

traj_metric = AggregationTrajectoryMetric(
    sequence_metrics={
        "events": LCSSequenceMetric(entity_metric=hamming, mode="normalized")
    },
    static_metric=custom_static_metric,
    static_metric_weight=0.01,
)

traj_ids = traj_pool.unique_ids
traj_a = traj_pool[traj_ids[0]]
traj_b = traj_pool[traj_ids[1]]

dist = traj_metric(traj_a, traj_b)
print(f"\nDistance between {traj_ids[0]} and {traj_ids[1]}: {dist:.1f}")

Distance between 1 and 2: 1.5