# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from aberrations import make_atm_data, make_1D_vibe_data, make_vibe_params
from observer import make_state_transition_vibe, make_kfilter_vibe, KFilter, make_impulse_from_tt, make_kfilter_turb
from matplotlib import pyplot as plt
import numpy as np
from scipy import signal, linalg
from copy import deepcopy
# %matplotlib inline

f_sampling = 1000
N_vibe = 0
use_turb = True
steps = 5000
noise = 0.06
delay = 1

rms = lambda data: np.sqrt(np.mean(data ** 2))
# -

if use_turb:
    turb = make_atm_data(steps)[:,0]

true_vibes = np.zeros((N_vibe, steps))
vibe_errs = np.zeros((N_vibe, steps))
vibe_params = np.zeros((N_vibe, 4))
# let's add in process noise so that there's something for Q to fit to!
process_vars = (10**-np.random.uniform(2, 3, (N_vibe,)))**2

process_vars

# +
for i in range(N_vibe):
    pars = make_vibe_params(N=1)
    true_vibes[i] = make_1D_vibe_data(steps, vib_params=pars, N=1)
    vibe_errs[i] = np.random.normal(0, process_vars[i], (steps,))
    vibe_params[i] = pars
    
vibes = true_vibes + vibe_errs

# +
if use_turb:
    truths = np.sum(vibes, axis=0) + turb
else:
    truths = np.sum(vibes, axis=0)
    
measurements = truths + np.random.normal(0, noise, (truths.size,))
# -

plt.semilogy(*signal.periodogram(truths, fs=f_sampling))
plt.ylim(1e-7)

kfilter_v = make_kfilter_vibe(vibe_params[:,1:3], process_vars)
kfilter_t = make_kfilter_turb(make_impulse_from_tt(truths[:500] + np.random.normal(0, 0.06, (500,))))
kfilter_t.iters = max(kfilter_t.iters, 500)

kfilter = kfilter_v + kfilter_t

# +
# something seems very wrong with my control code, so I'm just ditching it.

actions = np.zeros(steps - kfilter.iters,)
for i in range(kfilter.iters, steps - 1):
    kfilter.update(measurements[i])
    actions[i - kfilter.iters + delay] = kfilter.H.dot((kfilter.A**(delay)).dot(kfilter.state))
    kfilter.predict()
# -

print(rms(measurements))
print(rms(actions - truths[kfilter.iters:]))

ft, Pt = signal.periodogram(truths, fs=f_sampling)
fres, Pres = signal.periodogram(truths[kfilter.iters:] - actions, fs=f_sampling)
plt.loglog(ft, Pt)
plt.loglog(fres, Pres)
plt.ylim(1e-7)
plt.show()



StarcoderdataJupyterScriptsDedupFiltered