Data Assimilation

import jax
import flax

import jax.numpy as jnp
from jax import jit
from jax import lax

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

key = jax.random.key(0)
key, subkey = jax.random.split(key)

@jit
def ikeda_update(point):
    u = 0.9
    x, y = point[...,0], point[...,1]
    t = 0.4 - (6 / (1 + x ** 2 + y ** 2))
    x_new = 1 + u * (x * jnp.cos(t) - y * jnp.sin(t))
    y_new = u * (x * jnp.sin(t) + y * jnp.cos(t))
    return jnp.stack((x_new, y_new), axis=-1)

ikeda_jacobian = jit(jax.jacfwd(ikeda_update))

@jit
def iterate(points):
    def body_fn(i, val):
        return ikeda_update(val)
    return lax.fori_loop(0, 50, body_fn, points)

@jit
def generate_ikeda(subkey, size=10**5):
    return iterate(jax.random.uniform(subkey, shape=(size, 2), minval=0, maxval=0.5))

%timeit generate_ikeda(subkey)
key, subkey = jax.random.split(key)

def plot_ikeda_points(points = []):
    """
    Plots extra given points against an ikeda background
    Points is a list(tuple(x,y,color))
    """
    attractor = generate_ikeda(jax.random.key(0))
    plt.scatter(attractor[:3000, 0], attractor[:3000, 1], label='Attractor Points')
    
    for point in points:
        plt.scatter(**point)
        
    plt.legend(bbox_to_anchor=(1,1))
    return None

plot_ikeda_points()

from matplotlib.animation import FuncAnimation

time_units = 10
true_state = jnp.array([1.25, 0])

fig, ax = plt.subplots(1,1)

attractor = generate_ikeda(subkey)

class ikeda_attractor:

    def __init__(self, true_state : jax.numpy.array):
        self.true_state = true_state

    def update(self, i):
        
        return None

def animate(i):
    global true_state
    ax.clear()
    ax.scatter(attractor[:3000, 0], attractor[:3000, 1], label='Attractor Points')
    ax.scatter(*true_state, label='True State')
    true_state = ikeda_update(true_state)

anim = FuncAnimation(fig, animate, frames=time_units, interval=50)

plt.show()

# Prior state estimate: \hat{x}^{-}
true_state = jnp.array([1.25, 0])


# Prior covariance: P^{-}
prior_covariance = 1/16 * jnp.eye(2)
prior_estimate = true_state + 1/4 * jax.random.normal(subkey, shape=(2,))
key, subkey = jax.random.split(key)

# Nonlinear Measurement 'h'
@jit
def measure(true_state):
    return jnp.linalg.norm(true_state)

# Measurement Covariance: r
measurement_covariance = jnp.array([1/4])

# Measurement: y = h(x) + eta
measurement = measure(true_state) + jnp.sqrt(measurement_covariance) * jax.random.normal(subkey)
key, subkey = jax.random.split(key)

def track_point(time_units, true_state, prior_estimate, prior_covariance, update_method, update_method_kwargs):
    """
    Given a starting (true) point, a prior estimate of the state, a prior covariance of the state, a 
    """
    for time_unit in range(time_units):
        plot_ikeda_points([
             {'x':0,'y':0,'c':'red','label':'Sensor Position','s':100},
             {'x':true_state[0],'y':true_state[1],'c':'lime','label':'Point to Track','s':100, 'alpha':0.5},
             {'x':prior_estimate[..., 0],'y':prior_estimate[..., 1],'c':'purple','label':'Point to Track','s':100, 'alpha':0.5},  
        ])
        plt.show()
        true_state = ikeda_update(true_state)
        print(prior_estimate)
        prior_estimate, prior_covariance = update_method(prior_estimate, prior_covariance, **update_method_kwargs)

track_point(time_units=10,
            true_state=jnp.array([1.25, 0]),
            prior_estimate=jnp.array([1.25, 0]), #+ jax.random.uniform(minval=-0.1, maxval=0.1, key=key, shape=(100,2)),
            prior_covariance=jnp.eye(3),
            update_method=lambda x_1, x_2, **x: (ikeda_update(x_1)+ 0.01, x_2),
            update_method_kwargs={})

def bayesian_recursive_update(prior_estimate, prior_covariance, measurement, measurement_covariance, num_steps):
    state_estimate = prior_estimate
    covariance = ikeda_jacobian(prior_estimate) @ prior_covariance @ ikeda_jacobian(prior_estimate).T

    measurement = prior_estimate
    for _ in range(num_steps):
        measurement_jacobian = jax.grad(measure)(state_estimate)
        kalman_gain = covariance @ measurement_jacobian.T * (measurement_jacobian @ covariance @ measurement_jacobian.T + num_steps * measurement_covariance) ** (-1)
        state_estimate = state_estimate + kalman_gain * (measurement - measure(state_estimate))
        covariance = (jnp.eye(2) - kalman_gain @ measurement_jacobian) @ covariance
    print(state_estimate)
    return state_estimate, covariance

track_point(time_units=10,
            true_state=jnp.array([1.25, 0]),
            prior_estimate=jnp.array([1.24, 0]),# + jax.random.uniform(minval=-0.1, maxval=0.1, key=key, shape=(100,2)),
            prior_covariance=jnp.eye(2),
            update_method=bayesian_recursive_update,
            update_method_kwargs={'measurement':measure(jnp.array([1.25, 0])), 'measurement_covariance':measurement_covariance, 'num_steps':20})

#key, subkey = jax.random.split(key)
#ensemble = jnp.array([1.25, 0]) + jax.random.uniform(minval=-0.5, maxval=0.5, key=key, shape=(100,2))
#ensemble.mean(axis=0)

jnp.linalg.inv(jnp.array([[2]]))


def func(x1, x2):
    pass

func(**{'x1':1, 'x2':2})