# Model parameters
A = 2.3
f = 1
X0 = 2.5

# Noise parameter
sigma = 0.6

N = 300
t = np.sort(np.random.uniform(0, 10, N))
X = X0 + A * np.sin(2 * np.pi * f * t) + np.random.normal(0, sigma, N)

plt.plot(t, X, "o")
plt.xlabel("$t$")
plt.ylabel(r"$X_{\mathrm{obs}}$")
plt.show()


"""
Explanation: Searching for signals in noise
In this post we will cover some simple ideas that can be used to 
search for a signal in noise. Typically these methods are most 
useful when we are in the low signal to noise (SNR) enviroment.
The aim to show a use of the excellent emcee package
written by Dan Foreman-Mackey.
This post is available as an ipython notebook and can be downloaded/viewed
in nbviewer directly
$\newcommand{\data}{\mathbf{X}_{\textrm{obs}}}$
The basic idea
At discreet times $t_{i}$, we observe some data $\data(t_{i})$. The bold
font here indicates the data might be a vector of $N$ individual observations. We
hopefully have some ideas about the model which might have produced the data. 
Lets imagine we think this is the sum of a deterministic 'signal model'
and some noise process, that is
$$ X(t; \vec{\theta}) = f(t; \vec{\theta}) + n(t) $$ 
Here we use $\vec{\theta}$ a vector of length $d$ for the number of model
parameters. This is chosen on purpose to show the different between the
vectors of data, and those of parameters. 
The trick, which I learned from the gravitational-wave search community
is to make an assumption about the noise $n(t)$. Let's assume it is
stationary and normal such that the probability density function for 
$n(t)$ is the central normal distribution:
$$ N(0, \sigma) = \frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left{ \frac{-x^{2}}{2\sigma^{2}}\right} $$
Is this is the case, then if we knew the deterministic signal model
$f(t\; \vec{\theta}$ for the observed data, subtracting this from the
data should leave only the normally distributed noise.
$$ X(t; \vec{\theta}) - f(t; \vec{\theta}) \sim N(0, \sigma) $$ 
Therefore, given a single observed data point $X_{\textrm{obs}}(t_{i})$ and a
set of model parameters $\vec{\theta}$ we can compute the likelihood of that
data point:
$$ \mathcal{L} = p(X_{\textrm{obs}}(t_{i})| \vec{\theta}, \sigma) = \frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left{ \frac{-\left(X_{\textrm{obs}}(t_{i}) - f(t_{i}; \vec{\theta})\right)^{2}}{2\sigma^{2}}\right} $$
and so, provided all the data points are independant, the likelihood of all the
data points is given by their product
$$ \mathcal{L} = p(\data| \vec{\theta}, \sigma) = \prod_{i}^{N} \frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left{ \frac{-\left(X_{\textrm{obs}}(t_{i}) - f(t_{i}; \vec{\lambda})\right)^{2}}{2\sigma^{2}}\right} $$
Writing this as the log-likelilhood
\begin{equation}
\log\left(\mathcal{L}\right) = -\frac{1}{2}\sum_{i=1}^{N}\left(
\frac{\left(X_{\textrm{obs}}(t_{i})-f(t_{i}; \vec{\theta})\right)^{2}}{\sigma^{2}} + \log(2\pi\sigma^{2})
\right)
\end{equation}
Parameter estimation
Having obtained an expression for the likelihood of some data given some fixed 
parameters we are now in a position to use Bayes theorem
to estimate which parameter best fit the data. This is done by writting 
down the proportional version of Bayes rule:
$$ p(\vec{\theta}| \data) \propto p(\data| \vec{\theta}) p(\vec{\theta}) $$
Notice that we have a method to compute the likelihood term 
$p(\data| \vec{\theta}) = \mathcal{L}$ above. All that is left is
for us to specify the 'prior' $p(\vec{\theta}$: for now we will use
non-informative uniform priors. 
We are now in a possition to use the emcee package to estimate the posterior
parameter probabilties. See the example in the docs
for a far better explanation of all that I've just said!
Example
To illustrate the idea here is a simple example. Lets define the deterministic
model to be
$$ f(t; X_{0}, A, f) = X_{0} + A * \sin(2\pi f t) $$
Fake data
First we define the 'real' parameters, usually these would be
precisely what we don't know and want to specify.
End of explanation
"""


def Generic_lnuniformprior(theta, theta_lims):
    """ Generic uniform priors on theta 
    
    Parameters
    ----------
    theta : array_like 
        Value of the parameters
    theta_lims: array_like of shape (len(theta), 2)
        Array of pairs [min, max] for each theta paramater
        
    """
    
    theta_lims = np.array(theta_lims)
    
    if all(theta - theta_lims[:, 0] > 0 ) and all(theta_lims[:, 1] - theta > 0):
        return np.prod(1.0/np.diff(theta_lims, axis=1))
    return -np.inf

def Generic_lnlike(theta, t, x, model):
    """ Generic likelihood function for signal in Gaussian noise
    
    Parameters
    ----------
    theta : array_like
        Value of the parameters, the noise strength `sigma` should ALWAYS be
        the last element in the list
    t : array_like
        The independant variable
    x : array_like
        The observed dependent variable
    model : func
        Signal model, calling `model(theta[:-1], t)` should
        produced the corresponding value of the signal 
        alone `x_val`, without noise. 
    
    """
    
    sigma2 = theta[-1]**2
    return -0.5*(np.sum((x-model(theta[:-1], t))**2 / sigma2 + np.log(2*np.pi*sigma2)))   
    

"""
Explanation: Marginalisation
We will now use the MCMC algorithms of emcee to explore the 
joint probability distributions of the parameters. 
First we must define log-likelihood and log-prior in code. We
will introduce two generic functions which can be used for any
sigal model. If this is unfamiliar, you can define the functions
lnprior and lnlike directly as in the emcee documentation.
End of explanation
"""


# Prior
theta_lims = [[2.0, 2.5], 
              [0.99, 1.01], 
              [1.0, 5.0],
              [0.3, 0.9]
              ]
lnprior = lambda theta: Generic_lnuniformprior(theta, theta_lims) 

# Likelihood
def model(theta, t):
    A, f, X0 = theta 
    return X0 + A * np.sin(2 * np.pi * f * t)
    
lnlike = lambda theta, x, y: Generic_lnlike(theta, x, y, model)       

# Tie the two together
def lnprob(theta, x, y):
    """ Generic function ties together the lnprior and lnlike """ 
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, x, y)

"""
Explanation: We now use pythons lambda functions to generate the specific prior and likelihoods
for this signal model. It is worth noting that all the parameters are wrapped up in 
a list theta. The order must be the same for all functions, and the noise parameter
sigma should be the last element of the list.
End of explanation
"""


ndim = 4
nwalkers = 100

# Initialise the walkers
walker_pos = [[np.random.uniform(tl[0], tl[1]) for tl in theta_lims] 
                                                        for i in range(nwalkers)] 

# Run the sampler
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(t, X))
out = sampler.run_mcmc(walker_pos, 2000)

"""
Explanation: We now have a choice about how to initialise the walkers. In the emcee docs the 
result from the maximum likelihood result is used. I prefer to start them by randomly
covering the uniform prior parameter space. The choice of this parameter space is critical
to the success of the algorithm. Typically the data will be used to help specify
feasible values.
End of explanation
"""


def PlotWalkers(sampler, symbols=None, alpha=0.4, color="k"):
    """ Plot all the chains from a sampler """
    nwalkers, nsteps, ndim = sampler.chain.shape

    fig, axes = plt.subplots(ndim, 1, sharex=True, figsize=(8, 9))

    for i in range(ndim):
        axes[i].plot(sampler.chain[:, :, i].T, color="k", alpha=alpha)
        if symbols:
            axes[i].set_xlabel(symbols[i])
            
symbols = ["$A$", "$f$", "$X_{0}$", "$\sigma$"]            
PlotWalkers(sampler, symbols=symbols)
plt.show()

"""
Explanation: We can plot the progress of each walker during the MCMC exploration of 
parameter space. Here is a function which will automate the process
End of explanation
"""


samples = sampler.chain[:, 1000:, :].reshape((-1, ndim))

fig = triangle.corner(samples, labels=symbols,
                      truths=[A, f, X0, sigma]
                      )


axes = fig.get_axes()
ax = axes[0]
def prior(x):
    return np.zeros(len(x)) + 7419.1
xlims = ax.get_xlim()
x = np.linspace(xlims[0], xlims[1], 100)
ax.plot(x, prior(x), "r", lw=10)
print ax.get_ylim(), ax.get_xlim()

plt.show()

"""
Explanation: We now can plot the marginal distributions in a so-called triangle plot. First "burning"
the first several hundred points before it seems to have 'converged'
End of explanation
"""


