<jupyter_start><jupyter_text>KF EM RTS Nomoto<jupyter_code>%load_ext autoreload
%autoreload 2

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from src.extended_kalman_filter import extended_kalman_filter, simulate, rts_smoother
import src.extended_kalman_filter as ekf
from src.visualization.plot import plot


import src.visualization.book_format as book_format
book_format.set_style()
def lambda_f_constructor(K, T_1):
    def lambda_f(x, input):
        delta = input["delta"]
        f = np.array([[x[1], (K * delta - x[1]) / T_1]]).T
        return f

    return lambda_f
np.random.seed(42)

T_1 = 1.8962353076056344
K = 0.17950970687951323
h = 0.02

lambda_f = lambda_f_constructor(K=K, T_1=T_1)
def do_simulation(h_, lambda_f, N_=4000):

    t_ = np.arange(0, N_ * h_, h_)

    us = np.deg2rad(
        np.concatenate(
            (
                -10 * np.ones(int(N_ / 4)),
                10 * np.ones(int(N_ / 4)),
                -10 * np.ones(int(N_ / 4)),
                10 * np.ones(int(N_ / 4)),
            )
        )
    )
    data = pd.DataFrame(index=t_)
    data["delta"] = us

    np.random.seed(42)
    E = np.array([[0, 1]]).T
    process_noise = np.deg2rad(0.1)
    ws = np.random.normal(scale=process_noise, size=N_)

    data["psi"] = 0
    data["r"] = 0
    df = simulate(data=data, lambda_f=lambda_f, E=E, ws=ws, state_columns=["psi", "r"])

    measurement_noise = np.deg2rad(1)
    df["epsilon"] = np.random.normal(scale=measurement_noise, size=N_)
    df["psi_measure"] = df["psi"] + df["epsilon"]
    df["psi_deg"] = np.rad2deg(df["psi"])
    df["psi_measure_deg"] = np.rad2deg(df["psi_measure"])
    df["delta_deg"] = np.rad2deg(df["delta"])

    return df
df = do_simulation(lambda_f=lambda_f, h_=h)
data = pd.DataFrame(index=df.index)
data['psi'] = df['psi_measure']
data['r'] = df['r']
data['delta'] = df['delta']
plot(dataframes={'sim':data}, keys=["psi",'r'], fig_size=(10,3));
P_prd = np.diag([np.deg2rad(1), np.deg2rad(0.1)])

Qd = np.deg2rad(np.diag([0.1]))

Rd = np.diag([np.deg2rad(1)])

E = np.array(
    [[0], [1]],
)

C_ = np.array([[1, 0]])

Cd = C_
def lambda_jacobian_constructor(h, K, T_1):
    def lambda_jacobian(x, input):

        delta = input["delta"]
        r = x[1]

        jac = np.array(
            [
                [1, h],
                [0, 1 - h / T_1],
            ]
        )
        return jac

    return lambda_jacobian
lambda_jacobian = lambda_jacobian_constructor(h=h, K=K, T_1=T_1)
state_columns = ['psi','r']
time_steps = extended_kalman_filter(no_states=2, 
                                    no_measurement_states=1, 
                                    P_prd=P_prd, 
                                    lambda_f=lambda_f, 
                                    lambda_jacobian=lambda_jacobian, 
                                    data=data, 
                                    Qd=Qd, 
                                    Rd=Rd, 
                                    E=E, 
                                    Cd=Cd, 
                                    state_columns = state_columns,
                                    measurement_columns=['psi'])
df_kalman = ekf.time_steps_to_df(time_steps, state_columns=state_columns)
time_steps_smooth = ekf.rts_smoother(
            time_steps=time_steps,
            lambda_jacobian=lambda_jacobian,
            Qd=Qd,
            lambda_f=lambda_f,
            E=E,
        )
df_smooth = ekf.time_steps_to_df(time_steps_smooth, state_columns=state_columns)
fig_size=(10,3)
plot(dataframes={'real':df, 
                 'measure':data, 
                 #'kalman':df_kalman, 
                 'smooth':df_smooth}, 
                 keys=["psi",'r'], fig_size=fig_size);<jupyter_output><empty_output><jupyter_text>Parameter variation<jupyter_code>def vary(K,T_1):
    
    # Vary:
    lambda_f = lambda_f_constructor(K=K, T_1=T_1)
    lambda_jacobian = lambda_jacobian_constructor(h=h, K=K, T_1=T_1)
    
    # Kalman filter:
    time_steps = extended_kalman_filter(no_states=2, 
                                    no_measurement_states=1, 
                                    P_prd=P_prd, 
                                    lambda_f=lambda_f, 
                                    lambda_jacobian=lambda_jacobian, 
                                    data=data, 
                                    Qd=Qd, 
                                    Rd=Rd, 
                                    E=E, 
                                    Cd=Cd, 
                                    state_columns = state_columns,
                                    measurement_columns=['psi'])
    
    #return time_steps
    
    # RTS smoother:
    time_steps_smooth = ekf.rts_smoother(
            time_steps=time_steps,
            lambda_jacobian=lambda_jacobian,
            Qd=Qd,
            lambda_f=lambda_f,
            E=E,
        )
    
    return time_steps_smooth
    
T_1s = np.linspace(0.5*T_1, 1.5*T_1, 5)

time_steps_vary = {}
for T_1_ in T_1s:
    time_steps_vary[T_1_] = vary(K=K, T_1=T_1_)
dataframes_vary = {f'T:{np.round(T_1_,2)}':ekf.time_steps_to_df(time_steps_, state_columns=state_columns) for T_1_,time_steps_ in time_steps_vary.items()}

styles = {f'T:{np.round(T_1_,2)}':{'style':'-', 'alpha':0.5, 'lw':0.5} for T_1_,time_steps_ in time_steps_vary.items()}

dataframes_vary['real'] = df
styles['real'] = {
    'lw':1,
}
plot(dataframes_vary, 
     keys=["psi",'r'], 
     styles=styles,
     fig_size=fig_size);
    
    
dataframes_vary_variances = {f'T:{np.round(T_1_,2)}':pd.DataFrame(ekf.variance(time_steps_).T, columns=state_columns).iloc[1000:-1000] for T_1_,time_steps_ in time_steps_vary.items()}

fig = plot(dataframes_vary_variances, 
                 keys=["psi",'r'], fig_size=fig_size);

fig.suptitle('Variances');
plt.tight_layout()

from scipy.stats import multivariate_normal

def _loglikelihood(time_step):
    
    mean = time_step['x_hat'].flatten()
    cov = time_step['P_hat']
    x_prd = time_step['x_prd']
    #cov = time_step['P_prd']
    rv = multivariate_normal(mean=mean, cov=cov)
    return rv.logpdf(x_prd.flatten())
    
    
def loglikelihood(time_steps):
    
    #x_hats = ekf.x_hat(time_steps)
    #x_prds = ekf.x_prd(time_steps)
    #cov = time_steps[0]['P_hat']
    #
    #epsilon = x_prds - x_hats 
    #rv = multivariate_normal(mean=np.zeros(x_hats.shape[0]), cov=cov)
    #loglikelihood = np.sum(rv.logpdf(epsilon.T))    
    
    loglikelihood=0
    for time_step in time_steps:
        loglikelihood+=_loglikelihood(time_step)
        
    
        
    return loglikelihood

x_hats = ekf.x_hat(time_steps)
np.zeros(x_hats.shape[0])
loglikelihood(time_steps)
loglikelihoods = np.zeros(len(time_steps_vary))
for i,(T_1_, time_steps_) in enumerate(time_steps_vary.items()):
    loglikelihoods[i] = loglikelihood(time_steps_)
    
fig,ax=plt.subplots()
ax.plot(T_1s, loglikelihoods,'.-', label='parameter variation');
ax.plot([T_1,T_1],[np.min(loglikelihoods), np.max(loglikelihoods)],'r--', label=r'real $T_1$')
ax.set_ylabel('likelihood')
ax.set_xlabel(r'$T_1$');
ax.legend();


from scipy.optimize import minimize

def fun(x):
    T_1_ = x[0]
    K_ = x[1]
    
    time_steps = vary(K=K_, T_1=T_1_)
    return -loglikelihood(time_steps)
x0=[T_1*1.5, K*1.5]
res = minimize(fun, x0=x0, bounds=[(0.01,np.inf), (0.01,np.inf)], tol=0.001)
res
T_1
K
time_steps_optimized = vary(K=res.x[1], T_1=res.x[0])
df_sim_identified = simulate(data=data, lambda_f=lambda_f_constructor(K=res.x[1], T_1=res.x[0]), state_columns=["psi", "r"], hidden_state_columns=["r"])

fig_size=(10,3)
plot(dataframes={
                 'measure':data,
                 'real':df,
                 'filter identified': ekf.time_steps_to_df(time_steps_optimized, state_columns=state_columns),
                 'sim identified': df_sim_identified,
                },
                keys=["psi",'r'], fig_size=fig_size);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13735152893606267, 'fraction_numerical': 0.024639878695981804, 'mean_word_length': 2.28150912106136, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '446009', 'license': 'MIT', 'n_tokens_mistral': 3092, 'n_tokens_neox': 2903, 'n_words': 427, 'path': 'notebooks/15.50_KF_EM_RTS_nomoto.ipynb', 'repo_name': 'martinlarsalbert/wPCC', 'size': '7914'}