Sequence Metrics: DTW

This example demonstrates DTWSequenceMetric, which computes Dynamic Time Warping distance allowing flexible temporal alignment between sequences of different lengths.

Setup

import matplotlib.pyplot as plt
import polars as pl

from tanat import build_states
from tanat.dataset import simulate_states
from tanat.metric.entity import HammingEntityMetric
from tanat.metric.sequence import DTWSequenceMetric

Generate synthetic data

SEED = 42
N_IDS = 80

raw_df = simulate_states(
    n_ids=N_IDS,
    seq_length_range=(3, 8),
    features=["score", "status"],
    seed=SEED,
)
pool = build_states(raw_df, id_column="id", start_column="start", end_column="end")

┌─ 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 (80 sequences · 450 entities · 0.00s)

# Cast features to categorical
pool.cast_features({"status": pl.Categorical})
print(pool)

┌────────────────────────────────────────────────┐
│           StateSequencePool Summary            │
└────────────────────────────────────────────────┘

Overview
─────────────────────────
  Sequences          80
  Store              /home/runner/.tanat/_quick_state_a118dc98
  id_column          id

Time Index
─────────────────────────
  Type               Datetime(time_unit='us', time_zone=None) [2000-01-12 06:14:52.240595 → 2025-05-09 18:37:51.409412]
  Columns            ['start', 'end']
  t0                 position=0, anchor=start

Entity Features (2)
─────────────────────────
  • score               Numerical [1 → 100]
  • status              Categorical (5 categories)

Define metric

hamming = HammingEntityMetric(entity_feature="status")
metric = DTWSequenceMetric(entity_metric=hamming, normalize=True)
print(metric)

DTWSequenceMetric(settings=DTWSettings(entity_metric=HammingEntityMetric(settings=HammingSettings(entity_feature='status', cost=None, mismatch_cost=1.0)), window=None, normalize=True))

Compute distance between a single pair

ids = pool.unique_ids
dist = metric(pool[ids[0]], pool[ids[1]])
print(f"Distance between {ids[0]} and {ids[1]}: {dist:.4f}")

Distance between 1 and 2: 0.3000

Compute full pairwise distance matrix

dm = metric.compute_matrix(pool)
print(f"Distance matrix shape: {dm.shape}")

┌─ DTWSequenceMetric
│

│ Chunks:   0%|          | 0/1 [00:00<?, ?it/s]
│ Chunks: 100%|██████████| 1/1 [00:01<00:00,  1.25s/it]
│ Chunks: 100%|██████████| 1/1 [00:01<00:00,  1.25s/it]
│
└─ Done (80 sequences · 1.26s)
Distance matrix shape: (80, 80)

Visualize distances

arr = dm.to_numpy()
fig, ax = plt.subplots(figsize=(6.5, 5.5))
im = ax.imshow(arr, cmap="viridis_r", vmin=0, vmax=1)
ax.set_title("DTW distance matrix", fontsize=12, fontweight="bold")
ax.set_xlabel("Sequence index")
ax.set_ylabel("Sequence index")
cbar = plt.colorbar(im, ax=ax)
cbar.set_label("Distance")
plt.tight_layout()
plt.show()