augment_regime_detection

augment_regime_detection(
    data,
    date_column,
    close_column,
    window=252,
    n_regimes=2,
    method='hmm',
    step_size=1,
    n_iter=100,
    n_jobs=-1,
    reduce_memory=False,
    engine='auto',
)

Detect regimes in a financial time series using a specified method (e.g., HMM).

Parameters

Name Type Description Default
data DataFrame or GroupBy(pandas or polars) Input time-series data. Grouped inputs are processed per group before the regime labels are appended. required
date_column str Column name containing dates or timestamps. required
close_column str Column name with closing prices for regime detection. required
window Union[int, Tuple[int, int], List[int]] Size of the rolling window to fit the regime detection model. Default is 252. 252
n_regimes int Number of regimes to detect (e.g., 2 for bull/bear). Default is 2. 2
method str Method for regime detection. Currently supports β€˜hmm’. Default is β€˜hmm’. 'hmm'
step_size int Step size between HMM fits (e.g., 10 fits every 10 rows). Default is 1. 1
n_iter int Number of iterations for HMM fitting. Default is 100. 100
n_jobs int Number of parallel jobs for group processing (-1 uses all cores). Default is -1. -1
reduce_memory bool If True, reduces memory usage. Default is False. False
engine (auto, pandas, polars) Execution engine. "auto" (default) infers the backend from the input data while allowing explicit overrides. "auto"

Returns

Name Type Description
DataFrame DataFrame with added columns: - {close_column}regime{window}: Integer labels for detected regimes (e.g., 0, 1).

Notes

  • Uses Hidden Markov Model (HMM) to identify latent regimes based on log returns.
  • Regimes reflect distinct statistical states (e.g., high/low volatility, trending).
  • Requires β€˜hmmlearn’ package. Install with pip install hmmlearn.

Examples

import pandas as pd
import polars as pl
import pytimetk as tk

df = tk.load_dataset('stocks_daily', parse_dates=['date'])

# Example 1 - Single stock regime detection with pandas engine
# Requires hmmlearn: pip install hmmlearn
regime_df = (
    df.query("symbol == 'AAPL'")
    .augment_regime_detection(
        date_column='date',
        close_column='close',
        window=252,
        n_regimes=2
    )
)
regime_df.head().glimpse()
# Example 2 - Multiple stocks with groupby using pandas engine
# Requires hmmlearn: pip install hmmlearn
regime_df = (
    df.groupby('symbol')
    .augment_regime_detection(
        date_column='date',
        close_column='close',
        window=[252, 504],  # One year and two years
        n_regimes=3
    )
)
regime_df.groupby('symbol').tail(1).glimpse()
# Example 3 - Single stock regime detection with polars engine
# Requires hmmlearn: pip install hmmlearn
pl_single = pl.from_pandas(df.query("symbol == 'AAPL'"))
regime_df = pl_single.tk.augment_regime_detection(
    date_column='date',
    close_column='close',
    window=252,
    n_regimes=2
)
regime_df.glimpse()
# Example 4 - Multiple stocks with groupby using polars engine
# Requires hmmlearn: pip install hmmlearn
pl_df = pl.from_pandas(df)
regime_df = (
    pl_df.group_by('symbol')
    .tk.augment_regime_detection(
        date_column='date',
        close_column='close',
        window=504,
        n_regimes=3,
    )
)
regime_df.groupby('symbol').tail(1).glimpse()