Spaces:
Sleeping
Sleeping
File size: 4,031 Bytes
558db1e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import numpy as np
import pandas as pd
import json
import os
from config import logger, OUTPUT_DIR, Color
class NormalInverseWishartUpdater:
"""
Maintains a Normal-Inverse-Wishart (NIW) prior over the expected returns and covariance.
Allows for online, incremental updates of beliefs without full ML model retraining.
"""
def __init__(self, tickers, initial_mu=None, initial_cov=None, kappa_0=10.0, nu_0=None):
self.tickers = tickers
self.d = len(tickers)
self.kappa = float(kappa_0)
self.nu = float(nu_0 if nu_0 is not None else self.d + 2)
if initial_mu is not None:
self.mu = initial_mu.copy()
else:
self.mu = np.zeros(self.d)
if initial_cov is not None:
# Scale matrix Lambda = nu * Sigma
self.Lambda = initial_cov * self.nu
else:
self.Lambda = np.eye(self.d) * self.nu
self.state_path = os.path.join(OUTPUT_DIR, "niw_prior_state.json")
def save_state(self):
state = {
'tickers': self.tickers,
'kappa': self.kappa,
'nu': self.nu,
'mu': self.mu.tolist(),
'Lambda': self.Lambda.tolist()
}
with open(self.state_path, 'w') as f:
json.dump(state, f)
def load_state(self):
if os.path.exists(self.state_path):
try:
with open(self.state_path, 'r') as f:
state = json.load(f)
if state['tickers'] == self.tickers:
self.kappa = state['kappa']
self.nu = state['nu']
self.mu = np.array(state['mu'])
self.Lambda = np.array(state['Lambda'])
return True
except Exception as e:
logger.warning(f"Failed to load NIW state: {e}")
return False
def update(self, x_new):
"""
Updates the NIW parameters given a new observation x_new (1D array of returns).
"""
x = np.asarray(x_new)
if x.shape[0] != self.d:
raise ValueError("Observation dimension mismatch.")
# Update formulas for single observation (n=1)
kappa_n = self.kappa + 1.0
nu_n = self.nu + 1.0
diff = x - self.mu
mu_n = (self.kappa * self.mu + x) / kappa_n
# Rank-1 update to the scale matrix
Lambda_n = self.Lambda + (self.kappa / kappa_n) * np.outer(diff, diff)
# To prevent kappa and nu from growing infinitely (which would freeze the prior),
# we can apply a slight exponential decay to the weights (rolling window effect).
# We cap kappa and nu at a rolling window equivalent of 252 days.
max_memory = 252.0
if kappa_n > max_memory:
decay = max_memory / kappa_n
kappa_n *= decay
nu_n *= decay
Lambda_n *= decay
self.kappa = kappa_n
self.nu = nu_n
self.mu = mu_n
self.Lambda = Lambda_n
def get_posterior(self):
"""Returns posterior expected returns and covariance matrix."""
# The expected value of the covariance matrix under Inverse-Wishart is Lambda / (nu - d - 1)
# Using a safer denominator max(1, nu - d - 1)
denom = max(1.0, self.nu - self.d - 1)
cov_posterior = self.Lambda / denom
return self.mu, cov_posterior
def compute_divergence(self, new_mu, new_cov):
"""
Computes the Mahalanobis distance between the new ML predictions and the NIW prior.
"""
try:
cov_inv = np.linalg.inv(new_cov)
except np.linalg.LinAlgError:
cov_inv = np.linalg.pinv(new_cov)
diff = new_mu - self.mu
# Simple Mahalanobis distance
dist = np.sqrt(np.dot(diff.T, np.dot(cov_inv, diff)))
return float(dist)
|