optimiz-rs/examples/notebooks/01_hmm_tutorial.ipynb

In [ ]:

import numpy as np
import matplotlib.pyplot as plt
from optimizr import HMM

# Set random seed for reproducibility
np.random.seed(42)

print("OptimizR HMM Module Loaded Successfully!")

Example 1: Market Regime Detection

We'll model financial returns with 3 hidden states:

• State 0: Bull Market (high mean, low volatility)
• State 1: Bear Market (negative mean, high volatility)
• State 2: Sideways/Neutral (zero mean, medium volatility)

In [ ]:

def generate_regime_data(n_samples=500, seed=42):
    """
    Generate synthetic market returns with 3 regimes.
    """
    np.random.seed(seed)
    
    # Define true regime parameters
    true_means = np.array([0.08, -0.06, 0.01])     # Bull, Bear, Sideways
    true_stds = np.array([0.02, 0.05, 0.03])       # Volatilities
    
    # Transition matrix (tend to stay in same regime)
    transition_matrix = np.array([
        [0.85, 0.10, 0.05],  # Bull -> Bull, Bear, Sideways
        [0.10, 0.80, 0.10],  # Bear -> ...
        [0.15, 0.15, 0.70]   # Sideways -> ...
    ])
    
    # Generate state sequence
    true_states = [0]  # Start in bull market
    for _ in range(n_samples - 1):
        current_state = true_states[-1]
        next_state = np.random.choice(3, p=transition_matrix[current_state])
        true_states.append(next_state)
    
    true_states = np.array(true_states)
    
    # Generate observations
    returns = np.zeros(n_samples)
    for t in range(n_samples):
        state = true_states[t]
        returns[t] = np.random.normal(true_means[state], true_stds[state])
    
    return returns, true_states, true_means, true_stds

# Generate data
returns, true_states, true_means, true_stds = generate_regime_data()

print(f"Generated {len(returns)} return observations")
print(f"True means: {true_means}")
print(f"True stds: {true_stds}")
print(f"State distribution: {np.bincount(true_states)}")

Visualize the Generated Data

In [ ]:

fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)

# Plot returns with color-coded regimes
colors = ['green', 'red', 'gray']
regime_names = ['Bull', 'Bear', 'Sideways']

for state in range(3):
    mask = true_states == state
    axes[0].scatter(np.where(mask)[0], returns[mask], 
                   c=colors[state], label=regime_names[state], alpha=0.6, s=20)

axes[0].axhline(y=0, color='black', linestyle='--', alpha=0.3)
axes[0].set_ylabel('Returns', fontsize=12)
axes[0].set_title('Synthetic Market Returns (Color = True Regime)', fontsize=14, fontweight='bold')
axes[0].legend()
axes[0].grid(alpha=0.3)

# Plot cumulative returns
cumulative = np.cumsum(returns)
axes[1].plot(cumulative, linewidth=2, color='blue')
axes[1].set_xlabel('Time', fontsize=12)
axes[1].set_ylabel('Cumulative Return', fontsize=12)
axes[1].set_title('Cumulative Returns', fontsize=14, fontweight='bold')
axes[1].grid(alpha=0.3)

plt.tight_layout()
plt.show()

Fit the HMM Model

Now we'll use the Baum-Welch algorithm to learn the parameters from data.

In [ ]:

# Create and fit HMM
hmm = HMM(n_states=3, random_state=42)

print("Fitting HMM with Baum-Welch algorithm...")
hmm.fit(returns, n_iterations=100, tolerance=1e-6)

print("\nLearned Parameters:")
print(f"Transition Matrix:\n{hmm.transition_matrix_}")
print(f"\nEmission Means: {hmm.emission_means_}")
print(f"Emission Stds: {hmm.emission_stds_}")

Decode States with Viterbi Algorithm

In [ ]:

# Predict states using Viterbi
predicted_states = hmm.predict(returns)

print(f"Predicted state distribution: {np.bincount(predicted_states)}")

Evaluate Model Performance

Since HMM states are unlabeled, we need to find the best permutation mapping.

In [ ]:

from itertools import permutations

def best_permutation_accuracy(true_states, predicted_states, n_states=3):
    """
    Find best permutation mapping and compute accuracy.
    """
    best_acc = 0
    best_perm = None
    
    for perm in permutations(range(n_states)):
        mapped = np.array([perm[s] for s in predicted_states])
        acc = np.mean(mapped == true_states)
        if acc > best_acc:
            best_acc = acc
            best_perm = perm
    
    return best_acc, best_perm

accuracy, best_mapping = best_permutation_accuracy(true_states, predicted_states)

print(f"Best accuracy: {accuracy:.2%}")
print(f"Best mapping: {best_mapping}")
print(f"Interpretation: Predicted state {best_mapping[0]} = Bull")
print(f"                Predicted state {best_mapping[1]} = Bear")
print(f"                Predicted state {best_mapping[2]} = Sideways")

Visualize Results

In [ ]:

# Apply best mapping
mapped_predictions = np.array([best_mapping[s] for s in predicted_states])

fig, axes = plt.subplots(3, 1, figsize=(14, 10), sharex=True)

# Plot 1: True states
for state in range(3):
    mask = true_states == state
    axes[0].scatter(np.where(mask)[0], returns[mask],
                   c=colors[state], label=regime_names[state], alpha=0.6, s=20)
axes[0].set_ylabel('Returns', fontsize=12)
axes[0].set_title('True Hidden States', fontsize=14, fontweight='bold')
axes[0].legend()
axes[0].grid(alpha=0.3)

# Plot 2: Predicted states
for state in range(3):
    mask = mapped_predictions == state
    axes[1].scatter(np.where(mask)[0], returns[mask],
                   c=colors[state], label=f'Predicted {regime_names[state]}', alpha=0.6, s=20)
axes[1].set_ylabel('Returns', fontsize=12)
axes[1].set_title(f'Predicted States (Accuracy: {accuracy:.2%})', fontsize=14, fontweight='bold')
axes[1].legend()
axes[1].grid(alpha=0.3)

# Plot 3: Errors
errors = true_states != mapped_predictions
axes[2].scatter(np.where(errors)[0], returns[errors], 
               c='red', marker='x', s=100, label='Misclassified', alpha=0.7)
axes[2].scatter(np.where(~errors)[0], returns[~errors],
               c='green', marker='.', s=20, label='Correct', alpha=0.3)
axes[2].set_xlabel('Time', fontsize=12)
axes[2].set_ylabel('Returns', fontsize=12)
axes[2].set_title('Classification Errors', fontsize=14, fontweight='bold')
axes[2].legend()
axes[2].grid(alpha=0.3)

plt.tight_layout()
plt.show()

Confusion Matrix

In [ ]:

from sklearn.metrics import confusion_matrix
import seaborn as sns

cm = confusion_matrix(true_states, mapped_predictions)

plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=regime_names, yticklabels=regime_names)
plt.xlabel('Predicted State', fontsize=12)
plt.ylabel('True State', fontsize=12)
plt.title('Confusion Matrix', fontsize=14, fontweight='bold')
plt.show()

print("\nPer-State Accuracy:")
for i, name in enumerate(regime_names):
    acc = cm[i, i] / cm[i].sum()
    print(f"{name}: {acc:.2%}")

Example 2: Comparing Rust vs Python Performance

In [ ]:

import time

# Generate larger dataset
large_returns, _, _, _ = generate_regime_data(n_samples=5000)

# Time the fitting process
hmm_bench = HMM(n_states=3, random_state=42)

start = time.time()
hmm_bench.fit(large_returns, n_iterations=50)
rust_time = time.time() - start

print(f"Rust-accelerated fitting time: {rust_time:.3f} seconds")
print(f"For {len(large_returns)} observations with 50 iterations")
print(f"\nEstimated pure Python time: ~{rust_time * 50:.1f}s (50-100x slower)")

Key Takeaways

1. HMMs model sequential data with hidden states and observable outputs
2. Baum-Welch (EM) learns parameters from unlabeled data
3. Viterbi finds the most likely state sequence
4. OptimizR provides 50-100x speedup over pure Python for large datasets
5. Applications: Finance, speech, biology, weather, NLP

Further Reading

• Rabiner, L. R. (1989). "A tutorial on hidden Markov models and selected applications in speech recognition."
• Murphy, K. P. (2012). "Machine Learning: A Probabilistic Perspective" - Chapter 17