from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
import pymc as pm

%pylab inline

# The state of the stystem
true_pos = 0.
true_vel = 1.
# Assume linear decrease of the speed of the system
true_damping = -1.5

# timestep
dt = 0.01

p = true_pos
v = true_vel
positions, velocities = [], []
while True:
    # euler integration
    p = p + v * dt
    v = v + true_damping * dt
    
    positions.append(p)
    velocities.append(v)
    
    if v <= 0:
        break
    
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(positions, label="pos")
ax.plot(velocities, label="vel")
ax.legend()
ax.grid();

# draf the graphical model with daft
# TODO LaTeX/TikZ is nicer!
from matplotlib import rc
import daft

model = daft.PGM([5.5, 5.5], origin=[-1., -5])

model.add_node(daft.Node("damping", r"$d$", .5, 0))

for i in range(3):
    model.add_node(daft.Node("v" + str(i), r"$v_%d$" % i, i, -2))
    model.add_node(daft.Node("p" + str(i), r"$p_%d$" % i, i, -3))
    model.add_node(daft.Node("o" + str(i), r"$p_%d$" % i, i, -4, observed=True))
    
    model.add_edge("damping", "v" + str(i))
    model.add_edge("v" + str(i), "p" + str(i))
    model.add_edge("p" + str(i), "o" + str(i))

for i in range(2):
    model.add_edge("v" + str(i), "v" + str(i+1))
    model.add_edge("p" + str(i), "p" + str(i+1))


model.render();

import pymc as pm

# let's treat the real positions as observations
observations = positions[:]

# PRIORS
damping = pm.Normal("damping", mu=0, tau=1/2.**2)
#damping = pm.Uniform("damping", lower=-4, upper=4)

# we assume little system noise
tau_system_noise = (1 / 0.5**2)

# the state consist of pos and vel
# vel: we can't judge the initial velocity --> assume it's 0 with big std
vel_states = [pm.Normal("v0", mu=-4, tau=(1 / 2**2))]
# pos: the first pos is just the observation
pos_states = [pm.Normal("p0", mu=observations[0], tau=tau_system_noise)]

for i in range(1, len(observations)):
    new_vel = pm.Normal("v" + str(i),
                        mu=vel_states[-1] + damping * dt,
                        tau=tau_system_noise)
    new_pos = pm.Normal("s" + str(i),
                        mu=pos_states[-1] + new_vel * dt,
                        tau=tau_system_noise)
    vel_states.append(new_vel)
    pos_states.append(new_pos)

# some observation noise
tau_observation_noise = (1 / 0.5**2)
obs = pm.Normal("obs", mu=pos_states, tau=tau_observation_noise, value=observations, observed=True)

mcmc = pm.MCMC([damping, vel_states, pos_states])
mcmc.sample(10000, 5000)
pm.Matplot.plot(mcmc.get_node("damping"))
damping_samples = mcmc.trace("damping")[:]
print "damping -- mean:%f; std:%f" % (mean(damping_samples), std(damping_samples))
print "real damping -- %f" % true_damping

# let's treat the real positions as observations
observations = positions[:]

# PRIORS
damping = pm.Normal("damping", mu=0, tau=1/2.**2)
#damping = pm.Uniform("damping", lower=-4, upper=4)

# we assume little system noise
tau_system_noise = (1 / 0.5**2)

# the state consist of pos and vel
# vel: we can't judge the initial velocity --> assume it's 0 with big std
vel_states = [pm.Normal("v0", mu=-4, tau=(1 / 2**2))]
# pos: the first pos is just the observation
pos_states = [pm.Normal("p0", mu=observations[0], tau=tau_system_noise)]

for i in range(1, len(observations)):
    new_vel = pm.Normal("v" + str(i),
                        mu=vel_states[-1] + damping * dt,
                        tau=tau_system_noise)
    new_pos = pm.Normal("s" + str(i),
                        mu=pos_states[-1] + new_vel * dt,
                        tau=tau_system_noise)
    vel_states.append(new_vel)
    pos_states.append(new_pos)
# use pymc.Container for *_states
vel_states = pm.Container(vel_states)
pos_states = pm.Container(pos_states)

# some observation noise
tau_observation_noise = (1 / 0.5**2)
obs = pm.Normal("obs", mu=pos_states, tau=tau_observation_noise, value=observations, observed=True)

# SAMPLE
mcmc = pm.MCMC([damping, obs, vel_states, pos_states])
mcmc.sample(10000, 5000)
pm.Matplot.plot(mcmc.get_node("damping"))
damping_samples = mcmc.trace("damping")[:]
print "damping -- mean:%f; std:%f" % (mean(damping_samples), std(damping_samples))
print "real damping -- %f" % true_damping

# PRIORS
damping = pm.Normal("damping", mu=0, tau=1/2.**2)
damping = pm.Uniform("damping", lower=-4, upper=4)

# we assume little system noise
tau_system_noise = (1 / 0.5**2)

# the state consist of pos and vel
# vel: we can't judge the initial velocity --> assume it's 0 with big std
vel_states = [pm.Normal("v0", mu=-4, tau=(1 / 2**2))]
# pos: the first pos is just the observation
pos_states = [pm.Normal("p0", mu=observations[0], tau=tau_system_noise)]

for i in range(1, len(observations)):
    new_vel = pm.Normal("v" + str(i),
                        mu=vel_states[-1] + damping * dt,
                        tau=tau_system_noise)
    new_pos = pm.Normal("s" + str(i),
                        mu=pos_states[-1] + new_vel * dt,
                        tau=tau_system_noise)
    vel_states.append(new_vel)
    pos_states.append(new_pos)
# use pymc.Container for *_states
vel_states = pm.Container(vel_states)
pos_states = pm.Container(pos_states)

# some observation noise
tau_observation_noise = (1 / 0.5**2)

# SAMPLE
mcmc = pm.MCMC([damping, obs, vel_states, pos_states])
mcmc.use_step_method(pm.AdaptiveMetropolis, vel_states)
mcmc.use_step_method(pm.AdaptiveMetropolis, pos_states)

mcmc.sample(10000, 5000)
pm.Matplot.plot(mcmc.get_node("damping"))
damping_samples = mcmc.trace("damping")[:]
print "damping -- mean:%f; std:%f" % (mean(damping_samples), std(damping_samples))
print "real damping -- %f" % true_damping

# PRIORS
damping = pm.Normal("damping", mu=0, tau=1/2.**2)
damping = pm.Uniform("damping", lower=-4, upper=4)

# we assume little system noise
tau_system_noise = (1 / 0.5**2)
mus_vel=[-4]
mus_pos=[observations[0]]
taus=[1 / 2**2]

for i in range(1, len(observations)):
    mus_vel.append(mus_vel[-1] + damping * dt)
    mus_pos.append(mus_pos[-1] + mus_vel[-1] * dt)
    taus.append(tau_system_noise)
vel_states = pm.Normal("vel_states", mu=mus_vel, tau=taus)
pos_states = pm.Normal("pos_states", mu=mus_pos, tau=taus)

# SAMPLE
mcmc = pm.MCMC([damping, obs, vel_states, pos_states])
mcmc.use_step_method(pm.AdaptiveMetropolis, vel_states)
mcmc.use_step_method(pm.AdaptiveMetropolis, pos_states)

mcmc.sample(10000, 5000)
pm.Matplot.plot(mcmc.get_node("damping"))
damping_samples = mcmc.trace("damping")[:]
print "damping -- mean:%f; std:%f" % (mean(damping_samples), std(damping_samples))
print "real damping -- %f" % true_damping

#######################################################################
# Generative Model

# The initial state of the stystem
true_pos = 0.0
true_vel = 1.0
 
# Assume linear decrease of the speed of the system
true_damping = -.5
# Number of steps
N = 167
# Timestep
dt = 0.01
 
class PointMass(object):
    """
    Point-mass
    
     - The state consist of position and velocity.
     - The speed is decreased in a linear fashion by a damping factor.
    """ 
    def __init__(self, damping):
        self.damping = damping
 
    def integrate(self, p0, v0, N, dt):
        # The state consist of position and velocity
        positions = np.zeros((N,))
        velocities = np.zeros((N,))
        positions[0] = p0
        velocities[0] = v0
        for n in range(N - 1):
            # Euler integration
            positions[n + 1] = positions[n] + velocities[n] * dt
            velocities[n + 1] = velocities[n] + self.damping * dt
        return (positions, velocities)
 
point_mass = PointMass(true_damping)
 
positions, velocities = point_mass.integrate(true_pos, true_vel, N, dt)
 
# Assume little system noise
std_system_noise = 0.05
tau_system_noise = 1.0/std_system_noise**2
 
# Treat the real positions as observations
observations = positions + np.random.randn(N,) * std_system_noise

#######################################################################
# PyMC MODEL
# Damping is modelled with a Uniform prior
damping = pm.Uniform("damping", lower=-4.0, upper=4.0, value=-0.5)
 
# Initial position & velocity unknown -> assume Uniform priors
init_pos = pm.Uniform("init_pos", lower=-1.0, upper=1.0, value=0.5)
init_vel = pm.Uniform("init_vel", lower=0.0, upper=2.0, value=1.5)
 
@pm.deterministic
def det_pos(damping=damping, init_pos=init_pos, init_vel=init_vel):
    # Apply damping, init_pos and init_vel estimates and integrate
    point_mass.damping = damping.item()
    pos, vel = point_mass.integrate(init_pos, init_vel, N, dt)
    return pos
 
# Standard deviation is modelled with a Uniform prior
std_pos = pm.Uniform("std", lower=0.0, upper=1.0, value=0.5)
 
@pm.deterministic
def det_prec_pos(std_pos=std_pos):
    # Precision, based on standard deviation
    return 1.0/std_pos**2
 
# The observations are based on the estimated positions and precision
obs_pos = pm.Normal("obs", mu=det_pos, tau=det_prec_pos, value=observations, observed=True)

#######################################################################
# SAMPLE
mcmc = pm.MCMC([damping, init_pos, init_vel, det_prec_pos, obs_pos])
mcmc.sample(50000, 25000)

#######################################################################
# ANALYSIS
#mcmc.summary()
pm.Matplot.plot(mcmc.get_node("damping"))
pm.Matplot.plot(mcmc.get_node("init_pos"))
pm.Matplot.plot(mcmc.get_node("init_vel"))
 
damping_samples = mcmc.trace("damping")[:]
est_damping = np.mean(damping_samples)
est_std_damping = np.std(damping_samples)
 
init_pos_samples = mcmc.trace("init_pos")[:]
est_init_pos = np.mean(init_pos_samples)
est_std_init_pos = np.std(init_pos_samples)
 
init_vel_samples = mcmc.trace("init_vel")[:]
est_init_vel = np.mean(init_vel_samples)
est_std_init_vel = np.std(init_vel_samples)
 
print "           REAL    EST   STD"
print "damping:  %5.2f  %5.2f (%4.2f)" % (true_damping, est_damping, est_std_damping)
print "position: %5.2f  %5.2f (%4.2f)" % (true_pos, est_init_pos, est_std_init_pos)
print "velocity: %5.2f  %5.2f (%4.2f)" % (true_vel, est_init_vel, est_std_init_vel)
 
# Apply the mean estimates and integrate
pm.damping = est_damping
est_positions, est_velocities = point_mass.integrate(est_init_pos, est_init_vel, N, dt)
 
# Plot the system response
fig, (ax1, ax2) = plt.subplots(2, figsize=(12, 9))

ax1.plot(observations, label="true")
ax1.plot(est_positions, label="estimate")
ax1.legend()
ax1.set_title("positions")

ax2.plot(velocities, alpha=0.5, label="true (unobserved)")
ax2.plot(est_velocities, label="estimate")
ax2.legend()
ax2.set_title("velocities")
ax2.grid();