optim.py

from autograd.scipy import stats
from scipy.optimize import minimize

import autograd.numpy as anp
from autograd import hessian, value_and_grad


y = np.random.normal(scale=[1,4], size=(100,2)).T.flatten()


def likelihood_AD(gamma):
    # Parameters
    p00, p11 = 1/(1+anp.exp(-gamma[:2]))
    mu       = gamma[2:4]
    sigma2   = anp.exp(gamma[4:])
    T        = len(y)
    
    # Transition Matrix
    P = anp.array([[p00, 1-p11],[1-p00, p11]])
    
    # Bookkeeping
    global xi_10, xi_11
    xi_10 = [anp.zeros(shape=(1,2)) for _ in range(T+1)]    # Predictive
    xi_11 = [anp.zeros(shape=(1,2)) for _ in range(T)]     # Filtered
    lik   = 0
    
    # Initialize to OLS estimates:
    A  = anp.row_stack([anp.identity(2) - P, anp.ones(2)])
    xi_10[0] = anp.linalg.inv(A.T@A)@A.T@np.concatenate([np.zeros(2),[1]])
    
    # Forward filter recursion
    for t in range(T):
        # State densities
        eta        = stats.norm.pdf(y[t], mu, sigma2)

        # Likelihood
        lik += anp.log(eta @ xi_10[t])

        # Filtering
        xi_11[t]   = (eta * xi_10[t])  / (eta@xi_10[t])

        # Prediction
        xi_10[t+1] = P@xi_11[t]

    # Return likelihood-value
    return -lik



# Optimization
results = minimize(
                value_and_grad(likelihood_AD),
                x0=anp.random.normal(size=6),
                jac=True,
                method='BFGS',
                options={'gtol': 1e-7, 'maxiter': 20000, 'disp': True}, 
            )