# Expectation-Maximization (EM Algorithm)

# Purpose
The EM-algorithm together with a Kalman Filter seems to be able to estimate parameters in a state space model. This will be investigated here.

# Methodology
Implement some examples.

# Setup


```python
# %load imports.py
## Local packages:

%matplotlib inline
%load_ext autoreload
%autoreload 2
%config Completer.use_jedi = False  ## (To fix autocomplete)

## External packages:
import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)

import numpy as np
import os
import matplotlib.pyplot as plt
#if os.name == 'nt':
#    plt.style.use('presentation.mplstyle')  # Windows

import plotly.express as px 
import plotly.graph_objects as go

import seaborn as sns
import sympy as sp
from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame,
                                      Particle, Point)
from sympy.physics.vector.printing import vpprint, vlatex
from IPython.display import display, Math, Latex
from src.substitute_dynamic_symbols import run, lambdify

import pyro

import sklearn
import pykalman
from statsmodels.sandbox.regression.predstd import wls_prediction_std
import statsmodels.api as sm

from scipy.integrate import solve_ivp

## Local packages:
from src.data import mdl

from src.symbols import *
from src.parameters import *
import src.symbols as symbols
from src import prime_system
from src.models import regression
from src.visualization.regression import show_pred
from src.visualization.plot import track_plot

## Load models:
# (Uncomment these for faster loading):
import src.models.vmm_abkowitz  as vmm_abkowitz 

```

### Example
Let's exemplify this with a simple problem, simulated a ball being dropped in air.
The forces acting on this ball will be the drag from the air, which is modelled as: $C_d \cdot \dot{x}^2$ and the gravity, which is modelled as $g \cdot m$:


```python
x = dynamicsymbols('x')
m,F,C_d,g = sp.symbols('m F C_d g')

eq_kinetics = sp.Eq(F, -m*g + C_d*x.diff()**2)
Math(vlatex(eq_kinetics))
```

The kinematics connecting the forces to motion can be described with Newtons 2nd law:


```python
eq_kinematics = sp.Eq(F, m*x.diff().diff())
Math(vlatex(eq_kinematics))
```

The acceleration can then be calculated using these two equations:


```python
eqs=[eq_kinetics,eq_kinematics]
solution = sp.solve(eqs, x.diff().diff(), m, F, dict=True)[0]
eq_acc = sp.Eq(x.diff().diff(),solution[x.diff().diff()])
Math(vlatex(eq_acc))
```

This equation can be used to simulate the motion of the ball, by doing a time step integration of this initial value problem.


```python
from scipy.integrate import solve_ivp

acceleration_lambda = lambdify(eq_acc.rhs)

inputs={
'C_d' : 0.00,
'g' : 9.81,
'm' : 2,
}

def step(t,states, inputs):

    x1d = states[1]
    x2d = acceleration_lambda(**inputs, x1d=x1d)

    dstates = [x1d, x2d]
    return dstates

t_ = np.linspace(0,10,100)
y0 = [
    0,0
]
solution = solve_ivp(fun=step, y0=y0, t_span=[t_[0],t_[-1]], t_eval=t_, args=(inputs,))

df_result = pd.DataFrame(solution.y.T, index=solution.t, columns=['x','x1d'])
df_result['x2d'] = acceleration_lambda(**inputs, x1d=df_result['x1d'])

fig,axes=plt.subplots(nrows=3)
ax=axes[0]
df_result.plot(y='x', label='$x$ [m]', ax=ax)

ax.set_title('Ball position')

ax=axes[1]
df_result.plot(y='x1d', label='$\dot{x}$ [m/s]', ax=axes[1])
ax.set_title('Ball velocity [m/s]')
ax.set_xlabel('time [s]')

ax=axes[2]
df_result.plot(y='x2d', label='$\ddot{x}$ [m/s^2]', ax=axes[2])
ax.set_title('Ball acceleration [m/s]')
ax.set_xlabel('time [s]')

for ax in axes:
    ax.grid(True)
plt.tight_layout()
```


```python
df_result.iloc[0]
```


```python
df_measure = df_result.copy()
df_measure['x1d'] = np.NaN
df_measure['x2d'] = np.NaN

df_measure['x']+=np.random.normal(loc=0, scale=10, size=len(df_measure))
```


```python
dt = t_[1]-t_[0]

A = np.array([[1, dt, 0.5 * (dt**2)],
              [0, 1,            dt],
              [0, 0,            1]])
```


```python
from pykalman import KalmanFilter

kf = KalmanFilter(transition_matrices=A,
                 initial_state_mean = [df_result.iloc[0]['x'], df_result.iloc[0]['x1d'], df_result.iloc[0]['x2d']],
    
    #data.observation_matrix,
    #data.initial_transition_covariance,
    #data.initial_observation_covariance,
    #data.transition_offsets,
    #data.observation_offset,
    #data.initial_state_mean,
    #data.initial_state_covariance,
    em_vars=[
      'transition_matrices', 
      'observation_matrices',
      'transition_covariance', 
      'observation_covariance',
      'observation_offsets', 
      'initial_state_mean',
      'initial_state_covariance'
    ]  # variables to perform EM over. Any variable not appearing here is left untouched
)
observations = df_measure['x']
```


```python
kf.transition_matrices
```


```python
# Learn good values for parameters named in `em_vars` using the EM algorithm
loglikelihoods = np.zeros(10)
for i in range(len(loglikelihoods)):
    kf = kf.em(X=observations, n_iter=1)
    loglikelihoods[i] = kf.loglikelihood(observations)
```


```python
kf.transition_matrices
```


```python
# Estimate the state without using any observations.  This will let us see how
# good we could do if we ran blind.
n_dim_state = A.shape[0]
n_timesteps = observations.shape[0]
blind_state_estimates = np.zeros((n_timesteps, n_dim_state))
for t in range(n_timesteps - 1):
    if t == 0:
        blind_state_estimates[t] = kf.initial_state_mean
    
    blind_state_estimates[t + 1] = (
      np.dot(kf.transition_matrices, blind_state_estimates[t])
      #+ kf.transition_offsets[t]
    )

df_blind = pd.DataFrame(data=blind_state_estimates, index=df_result.index, columns=['x','x1d','x2d'])
```


```python
states_pred = kf.em(observations).smooth(observations)[0]
```


```python
df_pred = pd.DataFrame(data=states_pred, index=df_result.index, columns=['x','x1d','x2d'])
```


```python
fig,axes=plt.subplots(nrows=3)
fig.set_size_inches(10,10)

ax=axes[0]
df_result.plot(y='x', label='raw', ax=ax)
df_measure.plot(y='x', label='measure', ax=ax)
#df_pred.plot(y='x', label='filter', style='--', ax=ax)
df_blind.plot(y='x', label='blind', style=':', ax=ax)


ax.set_title('Ball position')

ax=axes[1]
df_result.plot(y='x1d', label='raw', ax=axes[1])
#df_pred.plot(y='x1d', label='filter', style='--', ax=axes[1])
df_blind.plot(y='x1d', label='blind', style=':', ax=axes[1])

ax.set_title('Ball velocity [m/s]')
ax.set_xlabel('time [s]')
plt.tight_layout()

ax=axes[2]
df_result.plot(y='x2d', label='raw', ax=axes[2])
#df_pred.plot(y='x2d', label='filter', style='--', ax=axes[2])
df_blind.plot(y='x2d', label='blind', style=':', ax=axes[2])

ax.set_title('Ball acceleration [m/s2]')
ax.set_xlabel('time [s]')
plt.tight_layout()
```


```python

```

5866f264ff931ce40e9dc8f52ff87efb12d59cf1
ipynb
Jupyter Notebook
notebooks/14.01_EM_algorithm.ipynb
martinlarsalbert/wPCC
16e0d4cc850d503247916c9f5bd9f0ddb07f8930
MIT
notebooks/14.01_EM_algorithm.ipynb
martinlarsalbert/wPCC
16e0d4cc850d503247916c9f5bd9f0ddb07f8930
MIT
notebooks/14.01_EM_algorithm.ipynb
martinlarsalbert/wPCC
16e0d4cc850d503247916c9f5bd9f0ddb07f8930
MIT
Qwen/Qwen-72B
1. YES
2. YES