Melvin Avarez
305 lines
10 KiB
Python
305 lines
10 KiB
Python
"""
|
||
Hidden Markov Model implementation
|
||
"""
|
||
|
||
import warnings
|
||
from typing import Optional
|
||
import numpy as np
|
||
|
||
# Try to import Rust backend
|
||
try:
|
||
from optimizr._core import fit_hmm as _rust_fit_hmm, viterbi_decode as _rust_viterbi
|
||
RUST_AVAILABLE = True
|
||
except ImportError:
|
||
RUST_AVAILABLE = False
|
||
|
||
|
||
class HMM:
|
||
"""
|
||
Hidden Markov Model for regime detection and sequence analysis.
|
||
|
||
Uses the Baum-Welch algorithm (EM) for parameter estimation and
|
||
Viterbi algorithm for finding the most likely state sequence.
|
||
|
||
Parameters
|
||
----------
|
||
n_states : int
|
||
Number of hidden states
|
||
|
||
Attributes
|
||
----------
|
||
transition_matrix_ : np.ndarray or None
|
||
Learned transition probabilities (n_states × n_states)
|
||
emission_means_ : np.ndarray or None
|
||
Mean of Gaussian emission for each state
|
||
emission_stds_ : np.ndarray or None
|
||
Std dev of Gaussian emission for each state
|
||
|
||
Examples
|
||
--------
|
||
>>> import numpy as np
|
||
>>> from optimizr import HMM
|
||
>>>
|
||
>>> # Generate data with regime changes
|
||
>>> returns = np.concatenate([
|
||
... np.random.normal(0.01, 0.02, 500), # Bull market
|
||
... np.random.normal(-0.01, 0.03, 500), # Bear market
|
||
... ])
|
||
>>>
|
||
>>> # Fit HMM
|
||
>>> hmm = HMM(n_states=2)
|
||
>>> hmm.fit(returns, n_iterations=100)
|
||
>>>
|
||
>>> # Decode states
|
||
>>> states = hmm.predict(returns)
|
||
>>> print(f"Detected states: {np.unique(states)}")
|
||
"""
|
||
|
||
def __init__(self, n_states: int = 2):
|
||
if n_states < 2:
|
||
raise ValueError("n_states must be at least 2")
|
||
|
||
self.n_states = n_states
|
||
self.transition_matrix_: np.ndarray = np.zeros((n_states, n_states))
|
||
self.emission_means_: np.ndarray = np.zeros(n_states)
|
||
self.emission_stds_: np.ndarray = np.ones(n_states)
|
||
self._params = None
|
||
|
||
def fit(self, X: np.ndarray, n_iterations: int = 100, tolerance: float = 1e-6) -> 'HMM':
|
||
"""
|
||
Fit HMM parameters using Baum-Welch algorithm.
|
||
|
||
Parameters
|
||
----------
|
||
X : np.ndarray
|
||
Time series observations (1D array)
|
||
n_iterations : int, default=100
|
||
Maximum number of EM iterations
|
||
tolerance : float, default=1e-6
|
||
Convergence threshold for log-likelihood change
|
||
|
||
Returns
|
||
-------
|
||
self : HMM
|
||
Fitted model
|
||
"""
|
||
X = np.asarray(X).flatten()
|
||
|
||
if len(X) == 0:
|
||
raise ValueError("X cannot be empty")
|
||
|
||
if RUST_AVAILABLE:
|
||
# Use Rust implementation
|
||
self._params = _rust_fit_hmm(
|
||
observations=X.tolist(),
|
||
n_states=self.n_states,
|
||
n_iterations=n_iterations,
|
||
tolerance=tolerance
|
||
)
|
||
|
||
self.transition_matrix_ = np.array(self._params.transition_matrix)
|
||
self.emission_means_ = np.array(self._params.emission_means)
|
||
self.emission_stds_ = np.array(self._params.emission_stds)
|
||
else:
|
||
# Pure Python fallback
|
||
warnings.warn(
|
||
"Rust backend not available. Using slower Python implementation.",
|
||
RuntimeWarning
|
||
)
|
||
self._fit_python(X, n_iterations, tolerance)
|
||
|
||
return self
|
||
|
||
def predict(self, X: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Predict most likely state sequence using Viterbi algorithm.
|
||
|
||
Parameters
|
||
----------
|
||
X : np.ndarray
|
||
Time series observations (1D array)
|
||
|
||
Returns
|
||
-------
|
||
states : np.ndarray
|
||
Most likely state at each time step
|
||
"""
|
||
if self.transition_matrix_ is None:
|
||
raise ValueError("Model must be fitted before prediction")
|
||
|
||
X = np.asarray(X).flatten()
|
||
|
||
if RUST_AVAILABLE and self._params is not None:
|
||
states = _rust_viterbi(X.tolist(), self._params)
|
||
return np.array(states)
|
||
else:
|
||
return self._viterbi_python(X)
|
||
|
||
def score(self, X: np.ndarray) -> float:
|
||
"""
|
||
Compute log-likelihood of observations.
|
||
|
||
Parameters
|
||
----------
|
||
X : np.ndarray
|
||
Time series observations
|
||
|
||
Returns
|
||
-------
|
||
log_likelihood : float
|
||
Log P(X | model)
|
||
"""
|
||
if self.transition_matrix_ is None:
|
||
raise ValueError("Model must be fitted before scoring")
|
||
|
||
X = np.asarray(X).flatten()
|
||
alpha = self._forward_python(X)
|
||
return np.log(np.sum(alpha[-1]))
|
||
|
||
def _fit_python(self, X: np.ndarray, n_iterations: int, tolerance: float):
|
||
"""Pure Python implementation of Baum-Welch"""
|
||
n_obs = len(X)
|
||
|
||
# Initialize parameters
|
||
self.transition_matrix_ = np.ones((self.n_states, self.n_states)) / self.n_states
|
||
|
||
# Initialize emissions based on quantiles
|
||
quantiles = np.linspace(0, 1, self.n_states + 1)
|
||
self.emission_means_ = np.zeros(self.n_states)
|
||
self.emission_stds_ = np.ones(self.n_states)
|
||
|
||
for i in range(self.n_states):
|
||
mask = (X >= np.quantile(X, quantiles[i])) & (X < np.quantile(X, quantiles[i+1]))
|
||
if np.any(mask):
|
||
self.emission_means_[i] = np.mean(X[mask])
|
||
self.emission_stds_[i] = max(np.std(X[mask]), 1e-6)
|
||
|
||
prev_ll = float('-inf')
|
||
|
||
# EM iterations
|
||
for _ in range(n_iterations):
|
||
# E-step
|
||
alpha = self._forward_python(X)
|
||
beta = self._backward_python(X)
|
||
gamma = self._compute_gamma_python(alpha, beta)
|
||
xi = self._compute_xi_python(X, alpha, beta)
|
||
|
||
# M-step
|
||
self._update_parameters_python(X, gamma, xi)
|
||
|
||
# Check convergence
|
||
ll = np.log(np.sum(alpha[-1]))
|
||
if abs(ll - prev_ll) < tolerance:
|
||
break
|
||
prev_ll = ll
|
||
|
||
def _forward_python(self, X: np.ndarray) -> np.ndarray:
|
||
"""Forward algorithm"""
|
||
n_obs = len(X)
|
||
alpha = np.zeros((n_obs, self.n_states))
|
||
|
||
# Initialize
|
||
for s in range(self.n_states):
|
||
alpha[0, s] = (1.0 / self.n_states) * self._emission_prob(X[0], s)
|
||
|
||
alpha[0] /= np.sum(alpha[0])
|
||
|
||
# Recursion
|
||
for t in range(1, n_obs):
|
||
for s in range(self.n_states):
|
||
alpha[t, s] = np.sum(alpha[t-1] * self.transition_matrix_[:, s]) * self._emission_prob(X[t], s)
|
||
alpha[t] /= max(np.sum(alpha[t]), 1e-10)
|
||
|
||
return alpha
|
||
|
||
def _backward_python(self, X: np.ndarray) -> np.ndarray:
|
||
"""Backward algorithm"""
|
||
n_obs = len(X)
|
||
beta = np.zeros((n_obs, self.n_states))
|
||
beta[-1] = 1.0
|
||
|
||
for t in range(n_obs - 2, -1, -1):
|
||
for s in range(self.n_states):
|
||
beta[t, s] = np.sum(
|
||
self.transition_matrix_[s] *
|
||
np.array([self._emission_prob(X[t+1], s2) for s2 in range(self.n_states)]) *
|
||
beta[t+1]
|
||
)
|
||
beta[t] /= max(np.sum(beta[t]), 1e-10)
|
||
|
||
return beta
|
||
|
||
def _compute_gamma_python(self, alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
|
||
"""Compute state occupation probabilities"""
|
||
gamma = alpha * beta
|
||
gamma /= np.sum(gamma, axis=1, keepdims=True)
|
||
return gamma
|
||
|
||
def _compute_xi_python(self, X: np.ndarray, alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
|
||
"""Compute transition probabilities"""
|
||
n_obs = len(X)
|
||
xi = np.zeros((n_obs - 1, self.n_states, self.n_states))
|
||
|
||
for t in range(n_obs - 1):
|
||
for i in range(self.n_states):
|
||
for j in range(self.n_states):
|
||
xi[t, i, j] = (alpha[t, i] * self.transition_matrix_[i, j] *
|
||
self._emission_prob(X[t+1], j) * beta[t+1, j])
|
||
xi[t] /= max(np.sum(xi[t]), 1e-10)
|
||
|
||
return xi
|
||
|
||
def _update_parameters_python(self, X: np.ndarray, gamma: np.ndarray, xi: np.ndarray):
|
||
"""M-step: update parameters"""
|
||
# Update transitions
|
||
for i in range(self.n_states):
|
||
denom = np.sum(gamma[:-1, i])
|
||
for j in range(self.n_states):
|
||
numer = np.sum(xi[:, i, j])
|
||
self.transition_matrix_[i, j] = numer / max(denom, 1e-10)
|
||
|
||
# Update emissions
|
||
for s in range(self.n_states):
|
||
weights = gamma[:, s]
|
||
sum_weights = np.sum(weights)
|
||
|
||
if sum_weights > 1e-10:
|
||
self.emission_means_[s] = np.sum(weights * X) / sum_weights
|
||
self.emission_stds_[s] = max(
|
||
np.sqrt(np.sum(weights * (X - self.emission_means_[s])**2) / sum_weights),
|
||
1e-6
|
||
)
|
||
|
||
def _emission_prob(self, obs: float, state: int) -> float:
|
||
"""Gaussian emission probability"""
|
||
mean = self.emission_means_[state]
|
||
std = self.emission_stds_[state]
|
||
z = (obs - mean) / std
|
||
return max(np.exp(-0.5 * z**2) / (std * np.sqrt(2 * np.pi)), 1e-10)
|
||
|
||
def _viterbi_python(self, X: np.ndarray) -> np.ndarray:
|
||
"""Viterbi decoding"""
|
||
n_obs = len(X)
|
||
delta = np.full((n_obs, self.n_states), float('-inf'))
|
||
psi = np.zeros((n_obs, self.n_states), dtype=int)
|
||
|
||
# Initialize
|
||
for s in range(self.n_states):
|
||
delta[0, s] = np.log(1.0 / self.n_states) + np.log(self._emission_prob(X[0], s))
|
||
|
||
# Recursion
|
||
for t in range(1, n_obs):
|
||
for s in range(self.n_states):
|
||
trans_probs = delta[t-1] + np.log(self.transition_matrix_[:, s] + 1e-10)
|
||
psi[t, s] = np.argmax(trans_probs)
|
||
delta[t, s] = trans_probs[psi[t, s]] + np.log(self._emission_prob(X[t], s))
|
||
|
||
# Backtrack
|
||
path = np.zeros(n_obs, dtype=int)
|
||
path[-1] = np.argmax(delta[-1])
|
||
|
||
for t in range(n_obs - 2, -1, -1):
|
||
path[t] = psi[t + 1, path[t + 1]]
|
||
|
||
return path
|