VAR(1) with Frobenius penalty on weight matrix

var1-frobenius.py

import numpy as np
import pymc3 as pm
import theano.tensor as tt

def frobenius_norm(X):
    return tt.sum(tt.nlinalg.trace(A@A.T))**0.5

# Create simulated data via forward evolution of system
K        = 3
T        = 10
error_sd_true = 0.5

A_true = np.random.randn(K,K)

x  = np.zeros([T, K])
x[0] = np.random.randn(K)

for t in range(1, T):
    jump = np.random.randn(K)
    x[t] = A_true@x[t-1] + jump*error_sd_true
    
penalty_weight = 0.1

with pm.Model():
        
    A_scale = pm.InverseGamma('A_scale', alpha=1, beta=1)
    A       = pm.Normal('A', sd=A_scale, shape=[K,K])
    sd      = pm.HalfCauchy('error_sd', beta=1)
    
    pm.Normal('x_end', mu=x[0:-1]@A, observed=x[1:], sd=sd)
    pm.Potential('A_penalty', frobenius_norm(A) * penalty_weight)
        
    trace = pm.sample()